| 123456789101112131415161718192021222324252627282930313233343536373839 |
- #!/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)")
|
