Преглед изворни кода

fix(recipes): GLM-5.3-Flash -> vLLM NVFP4 on 2x Spark with NoPE-MLA mod

Rewrite glm53-flash from SGLang/FP8/4-node to the proven vLLM/NVFP4/2-node
day-0 path (forum 381433 + kingjones30/GLM-5.3-Flash-2x-DGX-Spark):

- model: LibertAIDAI/GLM-5.3-Flash-NVFP4 (weight-only NVFP4-A16; halves the
  footprint so it fits 2 Sparks; sparse+linear KV nearly free)
- engine: vLLM day-0 image vllm/vllm-openai:glm53-flash-arm64-cu130 (arm64/cu130)
- 4-part build-time mod (docker-images/glm53-flash-nvfp4/patches/patch_mla.py,
  verbatim from the reference repo): NoPE-MLA rope-pad into DeepSeek 512+64
  fp8_ds_mla + sm120 sparse-MLA topk 2176->2048 (width, allowlist, compaction).
  Applied at image build; runtime-activated via VLLM_MLA_NOPE_PAD_ROPE=1.
- --moe-backend marlin (fixes the silent FP4 MoE corruption on sm_121)
- --language-model-only (multimodal processor OOM-kills the GB10 front-end)
- VLLM_ENGINE_READY_TIMEOUT_S=3600 (survive the 60+min FlashInfer-autotune boot)
- measured 24.74/30.30/19.58 tok/s at MTP-5, 256K needle-exact (2x Spark)
Alexander Korolev пре 1 месец
родитељ
комит
995e45c6c6

+ 30 - 0
docker-images/glm53-flash-nvfp4/Dockerfile

@@ -0,0 +1,30 @@
+# syntax=docker/dockerfile:1.7
+# GLM-5.3-Flash NVFP4 (LibertAIDAI weight-only NVFP4-A16) on DGX Spark
+# (SM121 / GB10), TP=2.
+#
+# Base: vllm/vllm-openai:glm53-flash-arm64-cu130 — vLLM's day-0 per-model
+# image (arm64/cu130, glm5_next arch support baked in). It CANNOT serve this
+# model on GB10 stock: glm5_next's NoPE-MLA (qk_rope_head_dim=0) + KDA linear
+# attention + the sparse kpool indexer fail every attention backend the sm_121
+# selector offers, and the auto-selected NVFP4 MoE backend silently corrupts
+# output (degenerate token loops, zero errors).
+#   * --moe-backend marlin (runtime, in the recipe) fixes the MoE corruption.
+#   * The 4-part patch below fixes the MLA.
+#
+# The patch (patches/patch_mla.py) dresses GLM's NoPE-MLA (kv_lora 512, rope 0)
+# in DeepSeek's fp8_ds_mla 512+64 layout and compacts the sm120 sparse-MLA
+# topk table from 2176 to the compiled 2048. It is applied at build time;
+# runtime activates the NoPE-pad path via VLLM_MLA_NOPE_PAD_ROPE=1 (recipe env).
+#
+# Provenance + measured numbers (24.74 code / 30.30 structured / 19.58 prose
+# tok/s at MTP-5, 2x Spark, 256K needle-exact):
+#   https://github.com/kingjones30/GLM-5.3-Flash-2x-DGX-Spark
+#   https://forums.developer.nvidia.com/t/381433
+
+ARG BASE_IMAGE
+FROM ${BASE_IMAGE}
+
+COPY patches/patch_mla.py /tmp/patch_mla.py
+RUN set -eux; \
+    python3 /tmp/patch_mla.py; \
+    rm -f /tmp/patch_mla.py

+ 22 - 0
docker-images/glm53-flash-nvfp4/build.sh

@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+# Build the SM121-patched GLM-5.3-Flash NVFP4 image.
+#
+# 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:local bash build.sh).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-glm53-flash-nvfp4:local}"
+
+ARGS=()
+while IFS= read -r line; do
+    case "$line" in
+        ''|\#*) continue ;;
+    esac
+    key="${line%%=*}"; val="${line#*=}"
+    [ -n "$key" ] && ARGS+=(--build-arg "${key}=${val}")
+done < versions.env
+
+echo "=== Building ${TAG} ==="
+DOCKER_BUILDKIT=1 docker build "${ARGS[@]}" --progress=plain -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

+ 126 - 0
docker-images/glm53-flash-nvfp4/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")
+

+ 5 - 0
docker-images/glm53-flash-nvfp4/versions.env

@@ -0,0 +1,5 @@
+# Authoritative build inputs for the GLM-5.3-Flash NVFP4 SM121 image.
+# Base: vLLM day-0 per-model image (arm64/cu130, glm5_next arch support
+# baked in). The patches/ dir carries the NoPE-MLA + sm120 sparse-MLA GB10
+# fixes (see Dockerfile).
+BASE_IMAGE=vllm/vllm-openai:glm53-flash-arm64-cu130

+ 72 - 43
recipes/glm53-flash/recipe.toml

@@ -1,67 +1,96 @@
-# GLM-5.3-Flash (SGLang day-0) — 4× DGX Spark / GB10 (sm_121), TP=4.
+# GLM-5.3-Flash NVFP4 (vLLM) — 2× DGX Spark / GB10 (sm_121), TP=2.
 #
-# 320B-A18B native FP8, hybrid sparse + linear attention + mHC. First
-# natively multimodal GLM-5 (text + vision). No MTP/DFlash draft shipped —
-# plain decode until a drafter lands. Model:
-# https://huggingface.co/zai-org/GLM-5.3-Flash
+# 320B/18B-active, hybrid 34 KDA linear-attention + 11 sparse-MLA layers,
+# glm5_next arch (NoPE-MLA, qk_rope_head_dim=0), in-checkpoint MTP head.
+# Text-only here (--language-model-only — the multimodal processor balloons
+# the API front-end to ~15.7 GB and the GB10 kernel OOM-kills it).
 #
-# Engine: SGLang day-0 — no vLLM has glm5_next: local-inference-lab fork
-# idle since Aug 11; aidendle94 / eugr images all pre-release.
+# WHY NVFP4 on 2 nodes (not the official FP8 on 4): LibertAIDAI's weight-only
+# NVFP4-A16 quant halves the footprint (~155 GB) so it fits 2 Sparks; the
+# sparse+linear attention keeps KV nearly free (KV pool ~738K tok at GMU 0.85).
+# Model: https://huggingface.co/LibertAIDAI/GLM-5.3-Flash-NVFP4
 #
-# Image: lmsysorg/sglang:glm-5.3-flash — multi-arch incl. arm64 (digest
-# sha256:989e2fff0…; re-verify at serve time). Day-0 = boots, not tuned.
+# Engine: vLLM day-0 per-model image (glm5_next support baked in). The arch
+# fails every sm_121 attention backend out of the box; a 4-part build-time
+# mod (docker-images/glm53-flash-nvfp4/patches/patch_mla.py) dresses GLM's
+# NoPE-MLA in DeepSeek's fp8_ds_mla 512+64 layout and compacts the sm120
+# sparse-MLA topk table 2176->2048. Applied at image build; this recipe
+# activates it at runtime via VLLM_MLA_NOPE_PAD_ROPE=1.
 #
-# Memory: 305.8 GiB → 76.5 GiB/node on 4× Spark; ~43 GiB/node left for KV.
-# Sparse+linear attention keeps KV small, so long context should be feasible.
+# TWO SILENT GB10 GOTCHAS (from the reference — both already handled here):
+#   1. SILENT FP4 MoE CORRUPTION: the auto-selected FLASHINFER_CUTLASS NVFP4
+#      MoE backend emits degenerate token loops with ZERO errors (model loads,
+#      /health green, every completion garbage). --moe-backend marlin fixes it.
+#      If a GB10 NVFP4 serve "works" but loops, check which MoE backend engaged
+#      before blaming the quant.
+#   2. A 60+ MIN "No available shared memory broadcast block found" loop at
+#      boot is NOT a hang — the worker is FlashInfer-autotuning a kernel flavor
+#      (tune persists to ~/.cache/vllm, pays once). VLLM_ENGINE_READY_TIMEOUT_S
+#      is raised to 3600 so the launcher doesn't give up first.
+#
+# Measured (median of 3, temp 0, single stream, 2x Spark, MTP-5 acceptance ~2.55):
+#   24.74 code / 30.30 structured / 19.58 prose tok/s; no-spec baseline 14.6 flat.
+#   256K needle-exact (30K/100K/200K deep); ~1150 tok/s prefill.
+#
+# Image: glm53-flash-nvfp4:local (built from vllm/vllm-openai:glm53-flash-arm64-cu130
+#   + the patch). `sm121 serve` builds it on the build node if the checksum label
+#   is stale (docker_image_dir below).
 #
 # PREREQUISITES:
-#   Weights cached on every node (~/.cache/huggingface → NAS, 328 GB):
-#       huggingface-cli download zai-org/GLM-5.3-Flash
+#   Weights cached on both nodes (NAS ~/.cache/huggingface):
+#       huggingface-cli download LibertAIDAI/GLM-5.3-Flash-NVFP4
 #
-# REFERENCE:
-#   https://forums.developer.nvidia.com/t/glm-5-3-flash-320b-total-parameters-18b-active/381350
+# REFERENCE (provenance):
+#   https://github.com/kingjones30/GLM-5.3-Flash-2x-DGX-Spark
+#   https://forums.developer.nvidia.com/t/381433
 
-engine = "sglang"
-name = "GLM-5.3-Flash (SGLang day-0)"
-docker_image = "lmsysorg/sglang:glm-5.3-flash"
-model = "zai-org/GLM-5.3-Flash"
-required_nodes = 4
-tensor_parallel_size = 4
+engine = "vllm"
+name = "GLM-5.3-Flash NVFP4 (vLLM, 2x Spark)"
+docker_image = "glm53-flash-nvfp4:local"
+docker_image_dir = "glm53-flash-nvfp4"
+model = "LibertAIDAI/GLM-5.3-Flash-NVFP4"
+required_nodes = 2
+tensor_parallel_size = 2
 pipeline_parallel_size = 1
 distributed_backend = "mp"
-shm_size = "64gb"
+shm_size = "32gb"
 
-entrypoint = ["python3", "-m", "sglang.launch_server"]
+entrypoint = ["vllm", "serve"]
 
 args = [
-    # Memory — weights ~76.5 GiB of 119.6 GiB/node; KV gets the rest.
-    "--mem-fraction-static", "0.87",
-    "--max-running-requests", "10",
-    # Context — 262K start; model card claims up to 1M, scale after measuring.
+    # GLM tokenizer
+    "--trust-remote-code",
+    # Executor (mp, no ray)
+    "--distributed-executor-backend", "mp",
+    # Memory — weights ~77.5 GiB of 119.6 GiB/node; sparse+linear KV nearly free
+    "--gpu-memory-utilization", "0.85",
     "--max-model-len", "262144",
-    # Attention / KV
-    "--attention-backend", "flashinfer",
-    "--kv-cache-dtype", "fp8_e4m3",
-    # Chunked prefill
-    "--chunked-prefill-size", "8192",
-    # Compile + decode graphs
-    "--enable-torch-compile", "--torch-compile-max-bs", "4",
-    "--cuda-graph-max-bs-decode", "4",
-    "--num-continuous-decode-steps", "2",
-    # Tokenizer / reasoning / tool calling (GLM family; parser names carry
-    # over from the vLLM glm52 recipes — verify they exist in this SGLang)
-    "--reasoning-parser", "glm45",
+    "--max-num-seqs", "1",
+    # DeepSeek-MLA KV layout the NoPE-pad mod dresses the attention in
+    "--kv-cache-dtype", "fp8_ds_mla",
+    # Text-only (multimodal processor OOM-kills the GB10 front-end)
+    "--language-model-only",
+    # GOTCHA 1: marlin avoids the silent FP4 MoE corruption
+    "--moe-backend", "marlin",
+    # MTP=5 (in-checkpoint drafter; acceptance length ~2.55)
+    "--speculative-config", """
+'{"method":"mtp","num_speculative_tokens":5}'""",
+    # Tokenizer / reasoning / tool calling (GLM family)
     "--tool-call-parser", "glm47",
-    # Observability
-    "--enable-metrics",
-    "--enable-cache-report",
+    "--enable-auto-tool-choice",
+    "--reasoning-parser", "glm45",
 ]
 
 [env]
 # Offline — weights must be cached on each node
 HF_HUB_OFFLINE = "1"
+# Activates the build-time NoPE-MLA rope-pad mod (mathematically a no-op:
+# zeros are rotation-free and add exactly 0 to every logit)
+VLLM_MLA_NOPE_PAD_ROPE = "1"
+# GOTCHA 2: survive the 60+ min FlashInfer-autotune boot window
+VLLM_ENGINE_READY_TIMEOUT_S = "3600"
 
-# Architecture
+# Architecture (aarch64 / sm_121a DGX Spark)
 TORCH_CUDA_ARCH_LIST = "12.1a"
 CUTE_DSL_ARCH = "sm_121a"
 OMP_NUM_THREADS = "4"