#!/usr/bin/env python3 """SM121 kernel patches for GLM-5.3-Flash-NVFP4 on GB10. This Spark recipe runs this once while building glm53-flash-sm121:v8 (GB10 / SM121). The serving tag on top is MM + Ray TP2 with the OpenAI API on :8888. Stock glm53-flash-arm64-cu130 only lists SM120 packed MLA on capability 12; our checkpoint is NoPE (pe_dim=0), so this image selects SM90 sparse-MLA + FA2 instead. We did not author FA2 MLA or FlashInfer 0.6.18 — we only gate what is already in the tree onto GB10. Modes: sm90 default. Seven steps the Dockerfile runs with no args. sm120 unused packed-cache path (fp8_ds_mla pad / topk 2048 / skip warmup). Kept so it is not a second file. Do not bake. python3 glm53-flash_SM121.py python3 glm53-flash_SM121.py sm90 python3 glm53-flash_SM121.py sm120 python3 glm53-flash_SM121.py --legacy-sm120 Unknown args are refused. Each replacement must match stock vLLM / FI 0.6.18 exactly once or we abort. Already-applied trees are skipped. """ from __future__ import annotations import sys from pathlib import Path VLLM = Path("/usr/local/lib/python3.12/dist-packages/vllm") FI = Path("/usr/local/lib/python3.12/dist-packages/flashinfer") LOG = "[mia-sm121]" USAGE = ( "usage: python3 glm53-flash_SM121.py [sm90|sm120|--legacy-sm120]\n" " sm90 (default): seven GB10 steps baked by Dockerfile\n" " sm120 / --legacy-sm120: unused packed-cache path (do not bake)" ) # --------------------------------------------------------------------------- # CLI (parse first so the bake helpers can sit in FI-then-vLLM source order) # --------------------------------------------------------------------------- def parse_mode(argv: list[str]) -> str: if not argv: return "sm90" if len(argv) != 1: raise SystemExit("refuse extra args\n" + USAGE) arg = argv[0] if arg == "sm90": return "sm90" if arg in ("sm120", "--legacy-sm120"): return "sm120" raise SystemExit("unknown mode %r\n%s" % (arg, USAGE)) def apply_once(path: Path, old: str, new: str, label: str) -> str: text = path.read_text() n_old = text.count(old) n_new = text.count(new) if n_old == 0 and n_new == 1: return "skipped" if n_old != 1: raise SystemExit( "%s refuse %s (old=%d new=%d); stock tree changed" % (LOG, label, n_old, n_new) ) path.write_text(text.replace(old, new, 1)) return "applied" def announce(title: str, results: list[str]) -> None: if all(r == "skipped" for r in results): print("%s skip %s (already in this image)" % (LOG, title)) return print("%s %s (%s)" % (LOG, title, " ".join(results))) # --------------------------------------------------------------------------- # FlashInfer 0.6.18 — files exist in this image after the Dockerfile pip. # Defined first (unlike a v1-then-v8 changelog). apply_sm90 still runs them # after the vLLM tree so a rebuild matches the same end state. # --------------------------------------------------------------------------- def bake_fa2_fp8_cta_tile() -> None: old = " constexpr uint32_t EFF_CTA_TILE_KV = std::is_same_v ? 32 : CTA_TILE_KV;\n" new = " constexpr uint32_t EFF_CTA_TILE_KV = std::is_same_v ? (CTA_TILE_KV < 32u ? CTA_TILE_KV : 32u) : CTA_TILE_KV;\n" r = apply_once( FI / "data/include/flashinfer/attention/mla.cuh", old, new, "mla.cuh fp8 tile", ) announce("FA2 fp8 CTA tile cap for GB10 ~101KB smem", [r]) def bake_fa2_fp8_sm12_gate() -> None: old = " major, minor = get_compute_capability(self.device)\n if major != 9:\n" new = " major, minor = get_compute_capability(self.device)\n if major not in (9, 12):\n" r = apply_once(FI / "mla/_core.py", old, new, "_core.py fp8 SM12 gate") announce("FA2 fp8 gate major in (9, 12)", [r]) # --------------------------------------------------------------------------- # vLLM tree — SM90 NoPE + FA2 on capability 12, PDL off, indexer/kpool # --------------------------------------------------------------------------- def bake_sm90_on_capability_12() -> None: old = """ elif device_capability.major == 12: return [ AttentionBackendEnum.TRITON_MLA, AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120, ]""" new = """ elif device_capability.major == 12: return [ AttentionBackendEnum.TRITON_MLA, AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90, AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120, ]""" r = apply_once(VLLM / "platforms/cuda.py", old, new, "cuda.py SM90 list") announce("cuda.py lists SM90 on capability 12", [r]) def bake_sm90_wrapper_for_gb10() -> None: path = VLLM / "v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py" cap_old = ( " def supports_compute_capability(cls, capability: DeviceCapability) -> bool:\n" " return capability.major == 9\n" ) cap_new = ( " def supports_compute_capability(cls, capability: DeviceCapability) -> bool:\n" " return capability.major in (9, 12)\n" ) fa_old = ' backend="fa3",\n' fa_new = ( ' backend=("fa3" if torch.cuda.get_device_capability()[0] == 9 else "fa2"),\n' ) gate_old = """ if not has_flashinfer_sm90_nope_mla(): return ( "FLASHINFER_MLA_SPARSE_SM90 requires FlashInfer with SM90 " "MLA support (ckv_scale_arr in " "BatchMLAPagedAttentionWrapper.run, FlashInfer >= 0.6.18)" )""" gate_new = """ if kv_cache_dtype in ("fp8", "fp8_e4m3") and not has_flashinfer_sm90_nope_mla(): return ( "FLASHINFER_MLA_SPARSE_SM90 fp8 KV requires FlashInfer with " "SM90 MLA support (ckv_scale_arr in " "BatchMLAPagedAttentionWrapper.run, FlashInfer >= 0.6.18)" )""" results = [ apply_once(path, cap_old, cap_new, "sm90 capability gate"), apply_once(path, fa_old, fa_new, "sm90 FA2 wrapper"), apply_once(path, gate_old, gate_new, "sm90 fp8 version gate"), ] announce("sm90 cap (9,12) + FA2 off-Hopper + fp8 gate", results) def bake_pdl_off_sm12() -> None: old = """ @classmethod def is_arch_support_pdl(cls) -> bool: try: device = torch.cuda.current_device() major, _ = torch.cuda.get_device_capability(device) except Exception: return False return major >= 9 """ new = """ @classmethod def is_arch_support_pdl(cls) -> bool: try: device = torch.cuda.current_device() major, _ = torch.cuda.get_device_capability(device) except Exception: return False # PDL lowering is unvalidated on SM12x (GB10) and races on KDA # state kernels there; keep it to Hopper/Blackwell-datacenter. return major in (9, 10) """ r = apply_once(VLLM / "platforms/cuda.py", old, new, "cuda.py PDL") announce("PDL gated off on SM12x (KDA race on this kit)", [r]) def bake_indexer_topk_neg1() -> None: path = VLLM / "model_executor/layers/sparse_attn_indexer_kpool.py" prefill_old = ( " pool_topk = torch.empty(\n" " (num_rows, select_k), dtype=torch.int32, device=logits.device\n" " )\n" ) prefill_new = ( " pool_topk = torch.full(\n" " (num_rows, select_k), -1, dtype=torch.int32, device=logits.device\n" " )\n" ) decode_old = ( " pool_topk = torch.empty(\n" " (num_rows, select_k), dtype=torch.int32, device=logits.device\n" " )\n" ) decode_new = ( " pool_topk = torch.full(\n" " (num_rows, select_k), -1, dtype=torch.int32, device=logits.device\n" " )\n" ) results = [ apply_once(path, prefill_old, prefill_new, "indexer prefill -1"), apply_once(path, decode_old, decode_new, "indexer decode -1"), ] announce("indexer top-k -1 init (prefill + decode)", results) def bake_kpool_pid_clamp() -> None: old = " hist_out = tl.where(pid >= 0, hist_val, -1)\n" new = " hist_out = tl.where((pid >= 0) & (pid < pool_len), hist_val, -1)\n" r = apply_once( VLLM / "models/glm5next/nvidia/ops/kpool_compress.py", old, new, "kpool pid clamp", ) announce("kpool pid clamp to pool_len", [r]) def apply_sm90() -> None: bake_sm90_on_capability_12() bake_sm90_wrapper_for_gb10() bake_pdl_off_sm12() bake_indexer_topk_neg1() bake_kpool_pid_clamp() bake_fa2_fp8_cta_tile() bake_fa2_fp8_sm12_gate() # --------------------------------------------------------------------------- # sm120 — unused packed-cache path (python3 glm53-flash_SM121.py sm120) # Marker-based _once, not apply_once. Not invoked by the Dockerfile. # --------------------------------------------------------------------------- MARKER = "GLM53_NOPE_FP8_DS_MLA_PAD" TOPK_MARKER = "GLM53_SM120_TOPK2048" ROOT: Path | None = None def _once(path: Path, needle: str, replacement: str, marker: str) -> str: text = path.read_text() if marker in text: return "already" if needle not in text: raise SystemExit(f"{path}: expected snippet not found") path.write_text(text.replace(needle, replacement, 1)) return "patched" def packed_cache_write() -> str: assert ROOT is not None path = ROOT / "v1/attention/backend.py" needle = """ from vllm import _custom_ops as ops ops.concat_and_cache_mla( kv_c_normed, k_pe.squeeze(1), kv_cache, slot_mapping.flatten(), kv_cache_dtype=kv_cache_dtype, scale=k_scale, ) """ replacement = f""" from vllm import _custom_ops as ops # {MARKER}: fp8_ds_mla pages always store 64-dim RoPE. NoPE models # (qk_rope_head_dim=0) pad zeros so concat_and_cache_mla accepts them. pe = k_pe.squeeze(1) if kv_cache_dtype == "fp8_ds_mla" and pe.shape[-1] == 0: pe = pe.new_zeros(*pe.shape[:-1], 64) ops.concat_and_cache_mla( kv_c_normed, pe, kv_cache, slot_mapping.flatten(), kv_cache_dtype=kv_cache_dtype, scale=k_scale, ) """ return _once(path, needle, replacement, MARKER) def packed_sm120_decode() -> str: assert ROOT is not None path = ROOT / "v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py" needle = """ if isinstance(q, tuple): q = torch.cat(q, dim=-1) num_actual_toks = q.shape[0] """ replacement = f""" if isinstance(q, tuple): q = torch.cat(q, dim=-1) # {MARKER}: match the 656 B fp8_ds_mla page (ckv + 64-dim kpe). rope_dim = self.qk_rope_head_dim if rope_dim == 0: q = torch.nn.functional.pad(q, (0, 64)) rope_dim = 64 num_actual_toks = q.shape[0] """ status = _once(path, needle, replacement, MARKER) if status == "already": return status needle2 = """ qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, """ replacement2 = """ qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=rope_dim, """ text = path.read_text() if needle2 not in text: raise SystemExit(f"{path}: flashinfer rope-dim call site not found") path.write_text(text.replace(needle2, replacement2, 1)) return "patched" def packed_sm120_topk() -> str: """Unused packed path: SM120 helper still wants topk=2048.""" assert ROOT is not None path = ROOT / "v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py" needle = """ output = q.new_empty( (num_actual_toks, self.num_heads, self.kv_lora_rank), dtype=q.dtype, ) """ replacement = f""" # {TOPK_MARKER}: FlashInfer SM120 v32/GLM wants block_tables (N, 1, 2048). if topk_indices_physical.shape[-1] != 2048: if topk_indices_physical.shape[-1] > 2048: topk_indices_physical = topk_indices_physical[..., :2048].contiguous() else: pad = 2048 - topk_indices_physical.shape[-1] topk_indices_physical = torch.nn.functional.pad( topk_indices_physical, (0, pad), value=-1 ) output = q.new_empty( (num_actual_toks, self.num_heads, self.kv_lora_rank), dtype=q.dtype, ) """ status = _once(path, needle, replacement, TOPK_MARKER) if status == "already": return status text = path.read_text() text = text.replace( "max_seq_len=attn_metadata.topk_tokens,", "max_seq_len=2048,", 1, ) text = text.replace( "sparse_mla_top_k=attn_metadata.topk_tokens,", "sparse_mla_top_k=2048,", 1, ) path.write_text(text) return "patched" def packed_skip_fi_warmup() -> str: """Unused packed path: fused_moe gemm1 autotune OOMs rank 0 on GB10.""" assert ROOT is not None path = ROOT / "model_executor/warmup/kernel_warmup.py" needle = """ flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) """ marker = "GLM53_SKIP_FI_SPARSE_WARMUP" replacement = f""" # {marker}: skip FlashInfer SM120 sparse-MLA autotune (kills rank 0 on GB10). deepseek_v4_sparse_mla_attention_warmup(worker) """ return _once(path, needle, replacement, marker) def packed_skip_fi_autotune() -> str: assert ROOT is not None path = ROOT / "model_executor/warmup/kernel_warmup.py" marker = "GLM53_SKIP_FI_AUTOTUNE" needle = ''' from flashinfer.autotuner import AutoTuner, set_autotune_process_group ''' replacement = f''' # {marker}: fused_moe gemm1/gemm2 autotune kills rank 0 on GB10. logger.info_once("Skipping FlashInfer autotune on SM121") return from flashinfer.autotuner import AutoTuner, set_autotune_process_group ''' return _once(path, needle, replacement, marker) def apply_sm120() -> None: global ROOT import vllm ROOT = Path(vllm.__file__).resolve().parent cache = packed_cache_write() decode = packed_sm120_decode() topk = packed_sm120_topk() warmup = packed_skip_fi_warmup() autotune = packed_skip_fi_autotune() print( "%s packed-path (not baked) cache_write=%s sm120_decode=%s " "sm120_topk=%s skip_fi_warmup=%s skip_fi_autotune=%s" % (LOG, cache, decode, topk, warmup, autotune) ) def main(argv: list[str] | None = None) -> None: mode = parse_mode(sys.argv[1:] if argv is None else argv) if mode == "sm90": apply_sm90() return apply_sm120() if __name__ == "__main__": main()