#!/bin/bash
# The sm12x auto-priority list in platforms/cuda.py only offers
# [TRITON_MLA, FLASHINFER_MLA_SPARSE_SM120] — neither matches GLM-5.3's
# head-512 MLA (SM120 flashinfer only supports head 576). The MTP draft
# model's vllm_config does NOT inherit the global --attention-backend, so
# its MLA layers auto-select and die. The B12X paged/sparse-MLA backend
# (which the main model uses via the global flag) supports (12,0)/(12,1),
# so offer it first.
set -euo pipefail
P=/usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py
python3 - <<'PY'
p = "/usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py"
s = open(p).read()
old = """        elif device_capability.major == 12:
            return [
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
            ]"""
new = """        elif device_capability.major == 12:
            return [
                AttentionBackendEnum.B12X,
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
            ]"""
if "AttentionBackendEnum.B12X,\n                AttentionBackendEnum.TRITON_MLA" in s:
    print("[b12x-sm12x] already patched")
else:
    assert old in s, "sm12x priority block not found (source drift)"
    open(p, "w").write(s.replace(old, new, 1))
    print("[b12x-sm12x] patched: B12X first in sm12x MLA priorities")
import importlib, vllm.platforms.cuda
importlib.reload(vllm.platforms.cuda)
print("[b12x-sm12x] OK")
PY
find /usr/local/lib/python3.12/dist-packages/vllm -name "__pycache__" -type d -prune -exec rm -rf {} + 2>/dev/null || true