#!/usr/bin/env python3
"""B12X into the sm12x auto-priority list (platforms/cuda.py).

The sm12x branch of _get_backend_priorities only offers
[TRITON_MLA, FLASHINFER_MLA_SPARSE_SM120]. Neither matches GLM-5.3-Flash:
TRITON_MLA rejects sparse, and FLASHINFER_MLA_SPARSE_SM120 only supports
head_size 576 (DeepSeek-V4 MLA), not GLM's 512. The B12X paged/sparse-MLA
backend explicitly supports compute (12, 0) and (12, 1), so it is offered
first. This matters for the MTP draft model: its vllm_config does NOT
inherit the global --attention-backend, so its MLA layers go through
auto-selection and die without this patch.
"""

import importlib
import sys

P = "/usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py"

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,
            ]"""

s = open(P).read()
if NEW in s:
    print("[patch] already applied")
elif OLD in s:
    open(P, "w").write(s.replace(OLD, NEW, 1))
    print("[patch] applied: B12X first in sm12x MLA priorities")
else:
    sys.exit("[patch] sm12x priority block not found (source drift)")