# BI-V100 patch script for Qwen3.6-27B (Qwen3_5 architecture)
#
# Triton situation on BI-V100:
#   - Standard Triton 2.3.1 is already present in the image.
#   - HAS_TRITON = False (hardcoded in vendor vllm), but Triton is still used
#     for TP-mode cache management (custom_cache_manager / libentry).
#   - The vendor's triton_utils/__init__.py, custom_cache_manager.py, libentry.py
#     are already correct for standard Triton 2.3.1 — do NOT overwrite them.
#   - DO NOT install BI-V150 corex Triton 2.1.0 (pkgs/triton): that causes
#     GPU hang on BI-V100 because the Triton CUDA PTX kernels are incompatible.

# Recommended server start command for TP=4 support 100K, need chunked prefill
# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
#     --model /workspace/models/Qwen3.6-27B --port 1111 --served-model-name llm \
#     --max-model-len 100000 --enforce-eager --trust-remote-code -tp 4 --gpu-memory-utilization 0.95 \
#     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
#     --max-num-batched-tokens 4096 --enable-chunked-prefill
#
# With prefix caching (GDN align-mode, requires chunked prefill):
# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
#     --model /workspace/models/Qwen3.6-35B-A3B --port 1111 --served-model-name llm \
#     --max-model-len 150000 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
#     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
#     --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
#     --max-seq-len-to-capture 32768

# --- environment: corex runtime paths + five optimization env flags ----------
# patch_ops.sh ONLY configures the environment + deploys the patched files.
# The actual server start command lives in ../computility-run.yaml.
# These exports are picked up by the verification steps below; the server process
# inherits them because computility-run.yaml is launched from a shell that has
# run this script (or these same vars are set in the yaml env: block).
export PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages
export LD_LIBRARY_PATH=/usr/local/corex/lib64:/usr/local/iluvatar/lib64:/usr/local/openmpi/lib
export PATH=/usr/local/corex/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/openmpi/bin
export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
export VLLM_ENGINE_ITERATION_TIMEOUT_S=${VLLM_ENGINE_ITERATION_TIMEOUT_S:-3600}

# Five GatedDeltaNet / MoE / attention optimizations (all default ON, each falls
# back to pure PyTorch on any anomaly — see qwen3_5.py / paged_attn.py). Setting
# any to 0 disables it. Mirrored in computility-run.yaml env: block.
export MOE_NATIVE=${MOE_NATIVE:-1}           # MoE sort-by-expert + native ixformer kernels
export CHUNK_PARALLEL=${CHUNK_PARALLEL:-1}   # GDN loop2 linear-matrix recurrence
export PREFIX_FLASH=${PREFIX_FLASH:-1}       # single-pad square native flash for prefix attn
export RMSNORM_NATIVE=${RMSNORM_NATIVE:-1}   # native rms_norm for RMSNormGated
export LOOP1_NATIVE=${LOOP1_NATIVE:-1}       # GDN loop1 triangular solve via solve_triangular

# --- libcusolver.so: make torch.linalg.solve_triangular usable (LOOP1_NATIVE) -
# corex torch has a hardcoded rpath to /opt/sw_home/local/cuda/lib64/libcusolver.so,
# which does not exist on this image. Without it, solve_triangular dlopen fails on
# every call → loop1_native falls back to the 63-step serial Python loop (losing the
# ~4-5x speedup). LD_LIBRARY_PATH does NOT help because corex ignores it for this lib.
# Fix: create the expected path as a symlink to the real corex libcusolver.so.
CUSOLVER_SRC=/usr/local/corex/lib64/libcusolver.so
CUSOLVER_DST=/opt/sw_home/local/cuda/lib64/libcusolver.so
if [ -f "$CUSOLVER_SRC" ]; then
    mkdir -p /opt/sw_home/local/cuda/lib64
    if [ ! -e "$CUSOLVER_DST" ] || [ "$(readlink -f "$CUSOLVER_DST")" != "$(readlink -f "$CUSOLVER_SRC")" ]; then
        ln -sf "$CUSOLVER_SRC" "$CUSOLVER_DST"
        echo "[patch_ops] linked libcusolver.so -> $CUSOLVER_SRC (enables LOOP1_NATIVE)"
    fi
else
    echo "[patch_ops] WARNING: $CUSOLVER_SRC not found; LOOP1_NATIVE will fall back to serial loop" >&2
fi

# --- paged_attn.py: replace forward_prefix with pure-PyTorch fallback -------
# The Triton context_attention_fwd kernel hangs BI-V100 GPUs permanently
# (standard Triton 2.3.1 PTX is not supported by the corex runtime either).
# Our paged_attn.py bypasses it entirely via _forward_prefix_pytorch, which
# utilizes K-tiling techniques, and also have _forward_decode_pytorch to bypass kernel
# when context length is high
cp ./paged_attn.py /usr/local/corex/lib/python3/dist-packages/vllm/attention/ops/paged_attn.py

# --- model_runner.py: fix prefix_cache_hit stays True in chunked-prefill chunk 2+ ---
# Bug: _compute_for_prefix_cache_hit Case 1 (prefix_cache_len <= context_len)
# leaves prefix_cache_hit=True. Then _add_seq_group uses block_table=computed_block_nums
# (only the original prefix blocks), ignoring chunk-1 KV cache blocks.
# _forward_prefix_pytorch then gets an undersized block_tables and crashes with
# "amax(): Expected reduction dim -1 to have non-zero size" on the 2nd tile.
# Fix: set prefix_cache_hit=False for Case 1 so the full block_tables is used.
python3 ./patch_model_runner.py

# --- transformers: Qwen3_5 tokenizer / model files --------------------------
pip install transformers==4.55.3 -i https://pypi.tuna.tsinghua.edu.cn/simple
cp -r ./qwen3_5 /usr/local/lib/python3.10/site-packages/transformers/models/
cp -r ./qwen3_5_moe /usr/local/lib/python3.10/site-packages/transformers/models/
python3 ./patch_transformers_qwen3_5.py

# --- vllm model: Qwen3.6-27B (Qwen3_5 arch) --------------------------------
cp ./mamba_cache.py /usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models/
cp ./qwen3_5.py /usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models/qwen3_5.py
python3 ./patch_vllm_qwen3_5.py

# --- sequence.py: fix completion_tokens inflation under chunked prefill ------
# Bug: get_output_token_ids_to_return(delta=True) with num_new_tokens=0
# returns _cached_all_token_ids[-0:] == [0:] (the ENTIRE prompt+output list).
# Each prefill chunk step adds prompt_len to previous_num_tokens, so a 10K
# prompt processed in 3 chunks inflates completion_tokens by ~30K.
# Also adds num_cached_tokens field to RequestMetrics for prefix-cache stats.
cp ./sequence.py /usr/local/corex/lib/python3/dist-packages/vllm/sequence.py

# --- scheduler.py: record num_cached_tokens in RequestMetrics ----------------
# Sets seq_group.metrics.num_cached_tokens = prefix_cache_len on first prefill
# when --enable-prefix-caching is active, so serving_chat.py can report it in
# usage.prompt_tokens_details.cached_tokens (OpenAI-compatible API response).
cp ./scheduler.py /usr/local/corex/lib/python3/dist-packages/vllm/core/scheduler.py

# --- xformers: bypass cudnnFlashAttnForward (head_dim=256 > 128 limit) ------
# Injects _run_sdpa_fallback (pure matmul+softmax) into xformers.py.
# Required because head_dim=256 > 128 and ixformer flash attention either
# crashes (is_causal=True) or produces wrong output (attn_mask path).
# The fallback uses query_start_loc to derive actual query lengths, so it
# works correctly during profiling runs with chunked-prefill-style batches.
# also bypasses auto chunked prefill on
python3 ./patch_xformers_sdpa_seq.py

# --- tool parser: Qwen3 XML tool call format ---------------------------------
# Registers "qwen3_coder" parser for Qwen3.6 XML-style tool calls:
#   <tool_call><function=name><parameter=key>\nvalue\n</parameter></function></tool_call>
# Use at server start: --tool-call-parser qwen3_coder --enable-auto-tool-choice
cp ./qwen3coder_tool_parser.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/tool_parsers/
python3 ./patch_vllm_tool_parser.py

# --- reasoning parser: Qwen3 <think>...</think> split ------------------------
# Adds --reasoning-parser qwen3 support.
# Routes thinking tokens to reasoning_content, rest to content in the delta.
# Works together with --tool-call-parser qwen3_coder (think → tool call flow).
cp -r ./reasoning /usr/local/corex/lib/python3/dist-packages/vllm/
cp ./protocol.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/protocol.py
cp ./cli_args.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/cli_args.py
cp ./serving_chat.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/serving_chat.py
cp ./api_server.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/api_server.py
cp ./chat_utils.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/chat_utils.py

# --- chat_template: multi_system 兼容 (合并多/乱序 system 到首位) -------------
# 数据集 34 条请求含多个或乱序 system 消息, 原模型 chat_template.jinja 对非首位
# system raise "System message must be at the beginning." → 15 条 4xx, 威胁成功率≥99%。
# 模型目录 (/root/public-storage 或容器内 /model) 是只读挂载, 不能改原模板。
# 改后模板随 qwen3_6_scripts 一起 COPY 进来 (本目录 ./chat_template_multi_system.jinja),
# 这里 cp 到 /workspace/ (Dockerfile 已 mkdir /workspace), 启动时用
# --chat-template /workspace/chat_template_multi_system.jinja 指向它 (见 computility-run.yaml)。
mkdir -p /workspace
CHAT_TMPL=/workspace/chat_template_multi_system.jinja
if [ -f "./chat_template_multi_system.jinja" ]; then
    cp "./chat_template_multi_system.jinja" "$CHAT_TMPL"
    echo "[patch_ops] deployed multi_system chat template → $CHAT_TMPL"
else
    echo "[patch_ops] WARNING: ./chat_template_multi_system.jinja not found in $(pwd)" >&2
fi

# --- verification: five optimization flags present in deployed package -------
VLLM_PKG=/usr/local/corex/lib64/python3/dist-packages/vllm
echo ""
echo "=== patch_ops verification ==="
_fail=0
check() {  # check  <grep_pattern>  <file>  <label>
    if grep -q "$1" "$2" 2>/dev/null; then
        echo "  [ok]   $3"
    else
        echo "  [MISS] $3  ($1 in $2)"; _fail=1
    fi
}
check "MOE_NATIVE"        "$VLLM_PKG/model_executor/models/qwen3_5.py" "MOE_NATIVE (qwen3_5.py)"
check "_native_experts_sorted" "$VLLM_PKG/model_executor/models/qwen3_5.py" "_native_experts_sorted"
check "CHUNK_PARALLEL"   "$VLLM_PKG/model_executor/models/qwen3_5.py" "CHUNK_PARALLEL"
check "LOOP1_NATIVE"     "$VLLM_PKG/model_executor/models/qwen3_5.py" "LOOP1_NATIVE"
check "RMSNORM_NATIVE"   "$VLLM_PKG/model_executor/models/qwen3_5.py" "RMSNORM_NATIVE"
check "PREFIX_FLASH"     "$VLLM_PKG/attention/ops/paged_attn.py"      "PREFIX_FLASH (paged_attn.py)"
check "_forward_prefix_flash" "$VLLM_PKG/attention/ops/paged_attn.py" "_forward_prefix_flash"

# libcusolver link present (else LOOP1_NATIVE silently falls back at runtime)
if [ -e /opt/sw_home/local/cuda/lib64/libcusolver.so ]; then
    echo "  [ok]   libcusolver.so link (LOOP1_NATIVE runtime)"
else
    echo "  [MISS] libcusolver.so link — LOOP1_NATIVE will fall back to serial loop"; _fail=1
fi

# multi_system chat template deployed (else multi-system requests 4xx at runtime)
if [ -f /workspace/chat_template_multi_system.jinja ]; then
    echo "  [ok]   /workspace/chat_template_multi_system.jinja (multi_system fix)"
else
    echo "  [MISS] /workspace/chat_template_multi_system.jinja — multi-system requests will 4xx"; _fail=1
fi

# solve_triangular actually callable (catches a broken cusolver link early)
if python3 -c "import torch,os;os.environ.setdefault('CUDA_VISIBLE_DEVICES','0');a=torch.eye(4,device='cuda');torch.linalg.solve_triangular(-a,a,upper=False,unitriangular=True)" 2>/dev/null; then
    echo "  [ok]   torch.linalg.solve_triangular callable"
else
    echo "  [MISS] torch.linalg.solve_triangular not callable — check libcusolver"; _fail=1
fi

if [ $_fail -eq 0 ]; then
    echo "  => all five optimizations deployed and enabled"
else
    echo "  => WARNING: some checks failed — review above" >&2
fi

# --- server start command lives in computility-run.yaml ----------------------
# patch_ops.sh only configures the environment + deploys patched files. It does
# NOT start the server. The server launch command (python3 -m vllm... with all
# CLI args) and the env block (PYTHONPATH / LD_LIBRARY_PATH / PATH / five
# *_NATIVE flags / VLLM_ENGINE_ITERATION_TIMEOUT_S) are in ../computility-run.yaml.
#
# After running this script, start the service via the computility runner, which
# must run from a NON-qwen3_6_scripts directory (this dir's xformers.py shadows
# the real xformers package).
#
# Sanity check after startup: the log must NOT contain any "* FALLBACK" line
# during cold prefill (prefix_flash FALLBACK on cache-hit requests is expected —
# the pad-ratio router intentionally falls back to pure PyTorch there).
#   grep "FALLBACK" <server_log>   # empty during cold prefill = all five active
# 80K cold prefill TTFT ≈ 57s with all five on (cached=0).
echo ""
echo "=== patch complete. Start the server with the computility runner using: ==="
echo "    ../computility-run.yaml"
echo "  (env + command are defined there; run from outside qwen3_6_scripts/)"