docker-images: glm53-flash-nvfp4-lab

docker-images/glm53-flash-nvfp4-lab/Dockerfile

@@ -0,0 +1,53 @@
+# syntax=docker/dockerfile:1.7
+# GLM-5.3-Flash NVFP4 (local-inference-lab quant) on DGX Spark (SM121/GB10), TP=2.
+#
+# Same day-0 base + NoPE-MLA mod as docker-images/glm53-flash-nvfp4, PLUS the
+# four lab-checkpoint patches from FujitsuPolycom/glm53-flash-tp2-spark
+# (Apache-2.0, attribution in patches/UPSTREAM.md):
+#   * model.patch               — checkpoint naming shim (attn_hc.* submodules,
+#                                 forget_gate flatten, fused [3C,1,4] conv1d split)
+#   * modelopt.patch            — MTP MIXED_PRECISION fix: the draft model
+#                                 resolves quantization across the checkpoint /
+#                                 target-mapper namespaces (without it:
+#                                 KeyError model.layers.45.mtp_block...w2_weight_scale)
+#   * sparse_attn_indexer*.patch — disable persistent_topk / cooperative
+#                                 workspace on CC 12.x (requests 62 blocks vs 48,
+#                                 128 KB smem vs GB10's 99 KB)
+#
+# NOT vendored (base-image or fabric specific, see recipe header): the upstream
+# image-layer patches (flashkda prefill, mm-renderer), their bundled NCCL 2.30.7
+# (direct-cabled ring), and the instanttensor loader (optional speed-up).
+#
+# Verification markers: `[quantprobe]` in modelopt.py (runtime log must show
+# algo=MXFP8 for the MTP layer), `hc_attn_base` in model.py.
+
+ARG BASE_IMAGE
+FROM ${BASE_IMAGE}
+
+# 1) NoPE-MLA rope-pad + sm120 sparse-MLA topk mod (anchors self-assert).
+COPY patches/patch_mla.py /tmp/patch_mla.py
+
+# 2) FujitsuPolycom lab-checkpoint patches (unmodified upstream .patch files).
+COPY patches/model.patch patches/modelopt.patch \
+     patches/sparse_attn_indexer.patch patches/sparse_attn_indexer_kpool.patch /tmp/
+
+RUN set -eux; \
+    B=/usr/local/lib/python3.12/dist-packages/vllm; \
+    python3 /tmp/patch_mla.py; \
+    patch "$B/models/glm5next/nvidia/model.py" /tmp/model.patch; \
+    patch "$B/model_executor/layers/quantization/modelopt.py" /tmp/modelopt.patch; \
+    patch "$B/model_executor/layers/sparse_attn_indexer.py" /tmp/sparse_attn_indexer.patch; \
+    patch "$B/model_executor/layers/sparse_attn_indexer_kpool.py" /tmp/sparse_attn_indexer_kpool.patch; \
+    grep -q quantprobe "$B/model_executor/layers/quantization/modelopt.py"; \
+    grep -q hc_attn_base "$B/models/glm5next/nvidia/model.py"; \
+    python3 -m py_compile \
+        "$B/models/glm5next/nvidia/model.py" \
+        "$B/model_executor/layers/quantization/modelopt.py" \
+        "$B/model_executor/layers/sparse_attn_indexer.py" \
+        "$B/model_executor/layers/sparse_attn_indexer_kpool.py"; \
+    rm -f /tmp/patch_mla.py /tmp/*.patch
+
+# 3) instanttensor loader (recipe header previously deferred this as optional;
+#    ~30 s weight load vs ~700 s sequential). Wheel is py3-none-any; --no-deps
+#    keeps the day-0 NCCL pins (its resolver would otherwise downgrade NCCL).
+RUN python3 -m pip install --no-deps --no-build-isolation instanttensor==0.1.9

docker-images/glm53-flash-nvfp4-lab/build.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Build the SM121 GLM-5.3-Flash NVFP4 (lab quant) image: day-0 base +
+# NoPE-MLA mod + FujitsuPolycom lab-checkpoint patches.
+#
+# Generic driver: forwards every KEY=VAL pin in versions.env as a --build-arg
+# to the Dockerfile. TAG is set by the sm121 launcher to the recipe's
+# docker_image (manual build: TAG=glm53-flash-nvfp4-lab:local bash build.sh).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-glm53-flash-nvfp4-lab:local}"
+
+ARGS=()
+while IFS= read -r line; do
+    case "$line" in
+        ''|\#*) continue ;;
+    esac
+    key="${line%%=*}"; val="${line#*=}"
+    val="${val%\"}"; val="${val#\"}"   # versions.env values may be shell-quoted
+    [ -n "$key" ] && ARGS+=(--build-arg "${key}=${val}")
+done < versions.env
+
+echo "=== Building ${TAG} ==="
+# Persistent buildkit cache (per node, shared across image dirs): the
+# version-independent stages (apt, torch, CUDA base) are reused across
+# rebuilds — only the pinned source stages (vllm/b12x/…) re-download and
+# re-compile. Override the dir with SM121_BUILDKIT_CACHE if desired.
+CACHE_DIR="${SM121_BUILDKIT_CACHE:-$HOME/.cache/sm121-buildkit}"
+mkdir -p "${CACHE_DIR}"
+docker buildx build "${ARGS[@]}" --cache-from "type=local,src=${CACHE_DIR}" \
+    --cache-to "type=local,dest=${CACHE_DIR},mode=max" \
+    --progress=plain --load -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

docker-images/glm53-flash-nvfp4-lab/patches/UPSTREAM.md

@@ -0,0 +1,19 @@
+# Vendored patches — provenance
+
+`model.patch`, `modelopt.patch`, `sparse_attn_indexer.patch`,
+`sparse_attn_indexer_kpool.patch`: unmodified copies from
+[FujitsuPolycom/glm53-flash-tp2-spark](https://github.com/FujitsuPolycom/glm53-flash-tp2-spark)
+(master @ vendoring date 2026-08-27), Apache License 2.0. See their
+`THIRD_PARTY_NOTICES.md` as well. Verified to apply with zero fuzz against
+`vllm/vllm-openai:glm53-flash-arm64-cu130`.
+
+`patch_mla.py`: the NoPE-MLA / sm120 sparse-MLA mod from
+[kingjones30/GLM-5.3-Flash-2x-DGX-Spark](https://github.com/kingjones30/GLM-5.3-Flash-2x-DGX-Spark)
+(see also `docker-images/glm53-flash-nvfp4/`).
+
+Not vendored from FujitsuPolycom: `image/patches/glm5next-kda.patch`
+(requires their base's chriswritescode SM120 overlay; enables
+`--kda-prefill-backend flashkda`), `image/patches/renderers-base.patch`
+(multimodal renderer; we serve `--language-model-only`), and
+`image/nccl/libnccl.so.2.30.7` (direct-cabled-ring fabric; our clusters use
+the stock NCCL with per-cluster `[nccl.*]` config).

docker-images/glm53-flash-nvfp4-lab/patches/model.patch

@@ -0,0 +1,36 @@
+--- a/vllm/model.py
++++ b/vllm/model.py
+@@ -770,6 +770,33 @@
+         for args in weights:
+             name, loaded_weight = args[:2]
+             kwargs: dict = args[2] if len(args) > 2 else {}
++            # Checkpoint layout shim (NVFP4 export uses newer serialization):
++            # renames verified shape/byte-equivalent against the fp8 export.
++            for _hc_from, _hc_to in (
++                ("attn_hc.base", "hc_attn_base"),
++                ("attn_hc.fn", "hc_attn_fn"),
++                ("attn_hc.scale", "hc_attn_scale"),
++                ("ffn_hc.base", "hc_ffn_base"),
++                ("ffn_hc.fn", "hc_ffn_fn"),
++                ("ffn_hc.scale", "hc_ffn_scale"),
++                (".self_attn.forget_gate.A_log", ".self_attn.A_log"),
++                (".self_attn.forget_gate.dt_bias", ".self_attn.dt_bias"),
++                (".self_attn.forget_gate.f_a_proj.", ".self_attn.f_a_proj."),
++                (".self_attn.forget_gate.f_b_proj.", ".self_attn.f_b_proj."),
++            ):
++                if _hc_from in name:
++                    name = name.replace(_hc_from, _hc_to)
++                    break
++            # Fused linear-attn conv: split [3*C,1,4] into q/k/v convs.
++            if ".self_attn.conv1d.weight" in name:
++                _c = loaded_weight.shape[0] // 3
++                for _j, _nm in enumerate(("q_conv1d", "k_conv1d", "v_conv1d")):
++                    _sub = name.replace(".conv1d.weight", ".%s.weight" % _nm)
++                    _p = params_dict[_sub]
++                    _wl = getattr(_p, "weight_loader", default_weight_loader)
++                    _wl(_p, loaded_weight[_j * _c : (_j + 1) * _c])
++                    loaded_params.add(_sub)
++                continue
+             if "rotary_emb.inv_freq" in name:
+                 continue
+ 

docker-images/glm53-flash-nvfp4-lab/patches/modelopt.patch

@@ -0,0 +1,39 @@
+--- a/vllm/modelopt.py
++++ b/vllm/modelopt.py
+@@ -2356,6 +2356,22 @@
+     def _quantized_layer_prefix_candidates(prefix: str) -> tuple[str, ...]:
+         candidates = [prefix]
+ 
++        # MTP draft naming (see nvfp4 stand-up notes): strip mtp_block,
++        # then re-root over every naming the shared quantized_layers dict
++        # can be in - checkpoint form (model.language_model.*) or the
++        # target-tree form its mapper rewrites to (language_model.model.*).
++        _bases = {prefix, prefix.replace(".mtp_block.", ".")}
++        for _b in list(_bases):
++            for _root in ("model.", "model.language_model.", "language_model.model."):
++                if _b.startswith(_root):
++                    _bases.add(_b[len(_root) :])
++        for _b in list(_bases):
++            if _b.startswith("layers."):
++                _bases.add("model." + _b)
++                _bases.add("model.language_model." + _b)
++                _bases.add("language_model.model." + _b)
++        candidates.extend(_bases - {prefix})
++
+         if prefix.endswith(".lm_head"):
+             candidates.append("lm_head")
+ 
+@@ -2387,6 +2403,13 @@
+             return None
+ 
+         quant_algo = self._resolve_quant_algo(prefix)
++        import sys as _sys
++        print(
++            "[quantprobe] layer=%s prefix=%s algo=%s"
++            % (type(layer).__name__, prefix, quant_algo),
++            file=_sys.stderr,
++            flush=True,
++        )
+ 
+         if isinstance(layer, (LinearBase, ParallelLMHead)):
+             if quant_algo == "FP8":

docker-images/glm53-flash-nvfp4-lab/patches/patch_mla.py

@@ -0,0 +1,126 @@
+"""NoPE-MLA rope pad (env-gated): make GLM-5.3-style qk_rope_head_dim=0 models
+ride the DeepSeek 512+64 fp8_ds_mla path on SM120/121 by (a) constructing the
+inner MLAAttention with rope=pe_pad so the KV-cache spec/kernels see head 576,
+and (b) zero-padding q and k_pe just before the attention call. Zeros are
+rotation-free and contribute exactly 0 to every logit — mathematically a no-op.
+Active only when VLLM_MLA_NOPE_PAD_ROPE=1 and the layer is genuinely NoPE."""
+import re, sys
+P = "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/mla.py"
+src = open(P).read()
+if "VLLM_MLA_NOPE_PAD_ROPE" in src:
+    print("already patched"); sys.exit(0)
+
+a1 = "        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim\n"
+r1 = a1 + (
+"        import os as _os\n"
+"        self.pe_pad = 0\n"
+"        if qk_rope_head_dim == 0 and _os.environ.get(\"VLLM_MLA_NOPE_PAD_ROPE\", \"0\") == \"1\":\n"
+"            self.pe_pad = 64\n"
+)
+assert src.count(a1) == 1, f"anchor1 count {src.count(a1)}"
+src = src.replace(a1, r1)
+
+a2 = """            qk_nope_head_dim=self.qk_nope_head_dim,
+            qk_rope_head_dim=self.qk_rope_head_dim,"""
+r2 = """            qk_nope_head_dim=self.qk_nope_head_dim,
+            qk_rope_head_dim=self.qk_rope_head_dim + self.pe_pad,"""
+assert src.count(a2) == 1, f"anchor2 count {src.count(a2)}"
+src = src.replace(a2, r2)
+
+a3 = "        attn_out = self.mla_attn(\n"
+r3 = (
+"        if self.pe_pad:\n"
+"            q = torch.nn.functional.pad(q, (0, self.pe_pad))\n"
+"            k_pe = q.new_zeros((k_pe.shape[0], 1, self.pe_pad))\n"
+) + a3
+assert src.count(a3) == 1, f"anchor3 count {src.count(a3)}"
+src = src.replace(a3, r3)
+
+open(P, "w").write(src)
+import py_compile; py_compile.compile(P, doraise=True)
+print("mla.py NoPE-pad patch applied + compiles")
+
+# ---- part 2: SM120 sparse decode must be told the REAL topk table width ----
+# glm5_next kpool indexer builds tables of width index_topk + kpool alignment
+# (2048+128=2176); the impl passed config topk (2048), tripping flashinfer
+# wrapper shape check. Pass the actual table width for both top_k and
+# max_seq_len (padding entries are -1 and skipped, same as short-seq padding).
+P2 = "/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py"
+src2 = open(P2).read()
+if "topk_indices_physical.shape[-1]" in src2:
+    print("sm120 width patch already applied")
+else:
+    b1 = "            max_seq_len=attn_metadata.topk_tokens,"
+    n1 = "            max_seq_len=topk_indices_physical.shape[-1],"
+    b2 = "            sparse_mla_top_k=attn_metadata.topk_tokens,"
+    n2 = "            sparse_mla_top_k=topk_indices_physical.shape[-1],"
+    assert src2.count(b1) == 1, f"sm120 anchor1 count {src2.count(b1)}"
+    assert src2.count(b2) == 1, f"sm120 anchor2 count {src2.count(b2)}"
+    src2 = src2.replace(b1, n1).replace(b2, n2)
+    open(P2, "w").write(src2)
+    py_compile.compile(P2, doraise=True)
+    print("sm120 topk width patch applied + compiles")
+
+# ---- part 3: allowlist the glm5_next topk width in the SM120 decode dispatch ----
+# glm5_next tables are index_topk + kpool tail rounded to BLOCK_N=128 => 2176.
+# The dsv3_2 decode binding derives topk and num_splits dynamically from
+# indices.shape[-1] (2176 = 34*64, tile-aligned); the frozenset is a
+# tested-shape allowlist, not a compiled-variant limit. Extend it.
+P3 = "/usr/local/lib/python3.12/dist-packages/flashinfer/mla/_sparse_mla_sm120.py"
+src3 = open(P3).read()
+if "(8, 2176)" in src3 or "GLM5_NEXT_WIDTH" in src3:
+    print("sm120 dispatch allowlist already extended")
+else:
+    a = "_DECODE_DSV3_2_PAGE_BLOCK_SIZE = 64"
+    add = ("# GLM5_NEXT_WIDTH: 2048 topk + kpool tail rounded to 128 -> 2176 (34x64 tiles)\n"
+           "_DECODE_DSV3_2_DISPATCH = _DECODE_DSV3_2_DISPATCH | frozenset(\n"
+           "    {(h, 2176) for h in (8, 16, 32, 64, 128)}\n"
+           ")\n")
+    assert src3.count(a) == 1, f"fi anchor count {src3.count(a)}"
+    src3 = src3.replace(a, add + a)
+    open(P3, "w").write(src3)
+    py_compile.compile(P3, doraise=True)
+    print("sm120 dispatch allowlist extended + compiles")
+
+# ---- part 4: compact glm5_next topk table to the compiled 2048 width ----
+# Backend validation pins index_topk=2048; buffer is 2048 + (kpool-1) tail
+# rounded to 2176. Compiled decode kernel only takes 2048. Keep the top
+# (2048 - tail) ranked entries + the always-select tail (never duplicated:
+# tail tokens come from the incomplete pool, excluded from pool selection).
+# Attention is permutation-invariant over the gathered set, so order is free.
+src4 = open(P2).read()
+if "_glm_kpool_tail" in src4:
+    print("sm120 tail-compaction already applied")
+else:
+    c1 = "        self.kv_scale_format = _kv_scale_format_for_model(model_type)\n"
+    r1 = c1 + (
+"        _kp = 1\n"
+"        if vllm_config.model_config is not None:\n"
+"            _kp = int(getattr(vllm_config.model_config.hf_text_config, \"index_kpool\", 1) or 1)\n"
+"        self._glm_kpool_tail = max(0, _kp - 1)\n"
+)
+    assert src4.count(c1) == 1, f"c1 count {src4.count(c1)}"
+    src4 = src4.replace(c1, r1)
+    c2 = """                NUM_TOPK_TOKENS=topk_indices.shape[1],
+            ),
+        )
+"""
+    r2 = c2 + (
+"        _w = topk_indices_physical.shape[-1]\n"
+"        _k = attn_metadata.topk_tokens\n"
+"        if _w > _k:\n"
+"            _t = min(self._glm_kpool_tail, _w - _k)\n"
+"            topk_indices_physical = torch.cat(\n"
+"                (\n"
+"                    topk_indices_physical[..., : _k - _t],\n"
+"                    topk_indices_physical[..., _k : _k + _t],\n"
+"                ),\n"
+"                dim=-1,\n"
+"            ).contiguous()\n"
+)
+    assert src4.count(c2) == 1, f"c2 count {src4.count(c2)}"
+    src4 = src4.replace(c2, r2)
+    open(P2, "w").write(src4)
+    py_compile.compile(P2, doraise=True)
+    print("sm120 tail-compaction applied + compiles")
+

docker-images/glm53-flash-nvfp4-lab/patches/sparse_attn_indexer.patch

@@ -0,0 +1,17 @@
+--- a/vllm/sparse_attn_indexer.py
++++ b/vllm/sparse_attn_indexer.py
+@@ -622,10 +622,10 @@
+             and current_platform.has_device_capability(90)
+             and not current_platform.is_device_capability_family(120)
+         )
+-        use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
+-            512,
+-            1024,
+-            2048,
++        use_persistent_topk = (
++            current_platform.is_cuda()
++            and topk_tokens in (512, 1024, 2048)
++            and not current_platform.is_device_capability_family(120)
+         )
+         if use_cooperative_topk:
+             workspace_manager = current_workspace_manager()

docker-images/glm53-flash-nvfp4-lab/patches/sparse_attn_indexer_kpool.patch

@@ -0,0 +1,15 @@
+--- a/vllm/sparse_attn_indexer_kpool.py
++++ b/vllm/sparse_attn_indexer_kpool.py
+@@ -807,7 +807,11 @@
+         else:
+             topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
+ 
+-        if current_platform.is_cuda() and select_k in (512, 1024, 2048):
++        if (
++            current_platform.is_cuda()
++            and select_k in (512, 1024, 2048)
++            and not current_platform.is_device_capability_family(120)
++        ):
+             workspace_manager = current_workspace_manager()
+             (topk_workspace,) = workspace_manager.get_simultaneous(
+                 ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),

docker-images/glm53-flash-nvfp4-lab/versions.env

@@ -0,0 +1,6 @@
+# Authoritative build inputs for the GLM-5.3-Flash NVFP4 (local-inference-lab)
+# SM121 image: day-0 vLLM image + the NoPE-MLA GB10 mod + the lab-checkpoint
+# patches (see Dockerfile and patches/UPSTREAM.md).
+# digest-pinned: tested base; tag is mutable (uncomment to float)
+#BASE_IMAGE="vllm/vllm-openai:glm53-flash-arm64-cu130"
+BASE_IMAGE="vllm/vllm-openai:glm53-flash-arm64-cu130@sha256:905c02933be6021301db2dc284e24e3727467aa3a0f63b41d609885778a07bce"