fix(deepseek-v4): drafter b12x WO routing for real DSpark acceptance + GMU 0.78

- Route the DSpark drafter's wo_a through the fused b12x WO projection
  (pack_wo_projection_fp8_block_scaled_weights_mxfp8 + wo_projection_inv_rope
  _mxfp8), identical to the main model, instead of the broken DeepGEMM
  fp8_einsum. This makes the drafter numerically correct: draft acceptance
  went from ~0% (broken) to ~10% avg / 32% position-0 (measured).
- Set --gpu-memory-utilization 0.78 (safe GB10 ceiling, per serving findings).
- Keep tiny-decode DISABLED: its GPU load-time MoE repack + tight host RAM
  (~32 GiB free) + dual-node 155 GiB load OOM'd both hosts (needed HA power
  cycle). Left VLLM_B12X_TINY_* commented as reference.
- Drop the old patch_dspark_woa.py einsum reconstruction (superseded by b12x).

Verified on ds4: boots to Application startup complete; 6 concurrent mixed-
length requests all HTTP 200 with no EngineDeadError; output correct.

docker-images/deepseek-v4-flash-dspark38/Dockerfile

@@ -26,10 +26,11 @@ RUN sed -i \
     /opt/env/lib/python3.12/site-packages/b12x/integration/residual_kernels.py \
  && grep -n 'from_dlpack(tensor' /opt/env/lib/python3.12/site-packages/b12x/integration/residual_kernels.py
 
-# Patch 2: drafter wo_a DeepGEMM fp8_einsum rank bug — reshape wo_a weight +
-# scale to (n_groups, o_lora_rank, -1) before the bmm (see patch script).
-COPY patches/patch_dspark_woa.py /tmp/patch_dspark_woa.py
-RUN /opt/env/bin/python3 /tmp/patch_dspark_woa.py
+# Patch 2: route the DSpark drafter's wo_a through the fused b12x WO
+# projection (same as the main model) so the drafter's numerics are correct and
+# DSpark draft acceptance is real (replaces the broken DeepGEMM fp8_einsum).
+COPY patches/patch_dspark_b12xwoa.py /tmp/patch_dspark_b12xwoa.py
+RUN /opt/env/bin/python3 /tmp/patch_dspark_b12xwoa.py
 
 # Patch 3: DSpark mixed-context batch — pad shorter requests to the batch-max
 # effective length so a parallel benchmark (tool-eval-bench --parallel 4) with

docker-images/deepseek-v4-flash-dspark38/patches/patch_dspark_b12xwoa.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+"""Patch dspark.py so the DSpark drafter uses the fused b12x WO projection
+(identical to the main model) instead of the broken DeepGEMM fp8_einsum path.
+
+Why: the DeepGEMM fp8_einsum route for the drafter's wo_a is broken in this
+fork (rank/tiled-scale problems). Patching it to at least boot was possible,
+but draft acceptance was ~0% (the numerically-approximate einsum makes the
+drafter's low-rank output projection wrong, so its proposals never match the
+target). The main model computes wo_a correctly via the fused
+`pack_wo_projection_fp8_block_scaled_weights_mxfp8` +
+`wo_projection_inv_rope_mxfp8` path, which natively handles the tiled scale.
+We give the drafter the same path: it does inv_rope + wo_a + wo_b in one
+kernel, so the drafter's projections become numerically correct.
+
+Imports torch is present in dspark.py (module uses it).
+"""
+import pathlib
+
+P = pathlib.Path(
+    "/opt/env/lib/python3.12/site-packages/vllm/models/deepseek_v4/nvidia/dspark.py"
+)
+src = P.read_text()
+
+# 1) Replace the broken DeepGEMM einsum tail (inv_rope + wo_a einsum + wo_b)
+#    with the fused b12x WO projection. `out` is (num_tokens, n_local_heads,
+#    head_dim) bf16 — the layout `_apply_b12x_wo_projection` expects.
+OLD_TAIL = """        out_fp8, out_scale = fused_inv_rope_fp8_quant(
+            out,
+            positions,
+            self.rotary_emb.cos_sin_cache,
+            n_groups=self.n_local_groups,
+            heads_per_group=self.n_local_heads // self.n_local_groups,
+            nope_dim=self.nope_head_dim,
+            rope_dim=self.rope_head_dim,
+            tma_aligned_scales=self._tma_aligned_scales,
+        )
+        projected = torch.empty(
+            (batch_size * block_size, self.n_local_groups, self.o_lora_rank),
+            dtype=self.dtype,
+            device=out.device,
+        )
+        torch.ops.vllm.deepseek_v4_fp8_einsum(
+            out_fp8,
+            out_scale,
+            self.wo_a.weight,
+            self.wo_a.weight_scale_inv,
+            projected,
+            "bhr,hdr->bhd",
+            list(self._einsum_recipe),
+        )
+        return _linear_no_bias(self.wo_b, projected.flatten(1).to(self.dtype))"""
+
+NEW_TAIL = """        # sm121-dspark38: fused b12x WO projection (same as the main model)
+        # replaces the broken DeepGEMM fp8_einsum path. Restores correct
+        # drafter projection numerics -> real DSpark draft acceptance. The
+        # kernel applies inv_rope + wo_a + wo_b in one fused pass.
+        return self._apply_b12x_wo_projection(out, positions)"""
+
+assert src.count(OLD_TAIL) == 1, f"einsum tail matches: got {src.count(OLD_TAIL)}"
+src = src.replace(OLD_TAIL, NEW_TAIL)
+
+# 2) Insert the b12x WO-projection methods into DeepSeekV4DSparkAttention
+#    (right before its forward_dspark, identified by the preceding line).
+METHODS = """        return q, self._maybe_quant_dequant_kv(kv)
+
+    def setup_b12x_wo_projection(self) -> None:
+        if getattr(self, "_b12x_wo_projection_weights", None) is not None:
+            return
+        groups, group_width, rank, hidden = self._validate_wo_projection_tensors()
+        from b12x.gemm.wo_projection import (
+            pack_wo_projection_fp8_block_scaled_weights_mxfp8,
+        )
+        self._b12x_wo_projection_weights = (
+            pack_wo_projection_fp8_block_scaled_weights_mxfp8(
+                self.wo_a.weight.detach(),
+                self.wo_a.weight_scale_inv.detach(),
+                self.wo_b.weight.detach(),
+                self.wo_b.weight_scale_inv.detach(),
+                groups=groups,
+                group_width=group_width,
+                rank=rank,
+                hidden=hidden,
+            )
+        )
+
+    def _validate_wo_projection_tensors(self):
+        groups = self.n_local_groups
+        group_width = (self.n_local_heads // groups) * self.head_dim
+        rank = self.o_lora_rank
+        hidden = self.hidden_size
+        shapes = {
+            "wo_a.weight": (groups * rank, group_width),
+            "wo_b.weight": (hidden, groups * rank),
+            "wo_a.scale": (
+                groups * ((rank + 127) // 128),
+                (group_width + 127) // 128,
+            ),
+            "wo_b.scale": (
+                (hidden + 127) // 128,
+                (groups * rank + 127) // 128,
+            ),
+        }
+        for name, t in (
+            ("wo_a.weight", self.wo_a.weight),
+            ("wo_b.weight", self.wo_b.weight),
+            ("wo_a.scale", self.wo_a.weight_scale_inv),
+            ("wo_b.scale", self.wo_b.weight_scale_inv),
+        ):
+            exp = shapes[name]
+            if tuple(t.shape) != exp:
+                raise RuntimeError(
+                    f"b12x WO {name} shape {tuple(t.shape)} != expected {exp}"
+                )
+        return groups, group_width, rank, hidden
+
+    def _apply_b12x_wo_projection(
+        self, o: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor:
+        num_tokens = int(o.shape[0])
+        if num_tokens == 0:
+            return torch.empty(
+                (0, self.hidden_size), dtype=o.dtype, device=o.device
+            )
+        if o.dtype != torch.bfloat16:
+            raise RuntimeError(
+                "sm121 b12x WO projection requires bf16 attention output, "
+                f"got {o.dtype}"
+            )
+        if getattr(self, "_b12x_wo_projection_weights", None) is None:
+            self.setup_b12x_wo_projection()
+        from b12x.gemm.wo_projection import (
+            WOProjectionScratchCaps,
+            plan_wo_projection_scratch,
+            wo_projection_inv_rope_mxfp8,
+        )
+        groups, group_width, rank, hidden = self._validate_wo_projection_tensors()
+        plan = plan_wo_projection_scratch(
+            WOProjectionScratchCaps(
+                device=o.device,
+                max_tokens=max(1, num_tokens),
+                groups=groups,
+                group_width=group_width,
+                rank=rank,
+                hidden=hidden,
+                dtype=o.dtype,
+            )
+        )
+        scratch = tuple(
+            torch.empty(shape, dtype=dtype, device=o.device)
+            for shape, dtype in plan.shapes_and_dtypes()
+        )
+        binding = plan.bind_inv_rope(
+            scratch=scratch,
+            o=o,
+            positions=positions,
+            cos_sin_cache=self.rotary_emb.cos_sin_cache,
+            weights=self._b12x_wo_projection_weights,
+            heads_per_group=self.n_local_heads // groups,
+            nope_dim=self.nope_head_dim,
+            rope_dim=self.rope_head_dim,
+        )
+        out = wo_projection_inv_rope_mxfp8(binding=binding)
+        if self.wo_b.reduce_results and self.wo_b.tp_size > 1:
+            from vllm.distributed import tensor_model_parallel_all_reduce
+            out = tensor_model_parallel_all_reduce(out)
+        return out
+
+    def forward_dspark("""
+
+ANCHOR = """        return q, self._maybe_quant_dequant_kv(kv)
+
+    def forward_dspark("""
+assert src.count(ANCHOR) == 1, f"method insert anchor: got {src.count(ANCHOR)}"
+P.write_text(src.replace(ANCHOR, METHODS))
+print("patched dspark.py: drafter wo_a routed through fused b12x WO projection")
+
+# Remove the now-unused einsum imports? Not required — harmless if referenced
+# elsewhere. Leave as-is.

docker-images/deepseek-v4-flash-dspark38/patches/patch_dspark_woa.py

@@ -1,76 +0,0 @@
-#!/usr/bin/env python3
-"""Patch dspark.py forward_dspark so the drafter wo_a feeds DeepGEMM fp8_einsum
-with reconstructed, DeepGEMM-compatible scale tensors (BUG-2026-08-03T204037).
-
-Probe (runtime shapes on production-3.8):
-  out_fp8=(40,4,4096) fp8   out_scale=(40,4,8)   weight=(4096,4096)
-  scale_inv=(32,32) ue8m0   projected=(40,4,1024)   groups=4 o_lora=1024
-  recipe=(1,1,128)
-
-Root cause: DeepGEMM's fp8_bmm ("bhr,hdr->bhd") requires (1) the weight rank-3
-(H,D,R), (2) the b_scale per-row 128-block (H,D,R/128).  The drafter instead
-passes:
-  - weight 2D (H*D,R)                     -> reshape to (H,D,R)
-  - a_scale: DeepGEMM permutes the A-side itself; the native UE8M0 int scale
-    (B,H,R/512) is already correct — do NOT expand.
-  - b_scale tiled (32,32) UE8M0 (128x128) -> broadcast each tile's scale to its
-    128 output rows => per-row (H,D,R/128) fp32
-
-Verified via a GPU harness that these shapes/dtypes pass DeepGEMM's shape checks.
-"""
-import pathlib
-
-P = pathlib.Path(
-    "/opt/env/lib/python3.12/site-packages/vllm/models/deepseek_v4/nvidia/dspark.py"
-)
-src = P.read_text()
-
-OLD = """        torch.ops.vllm.deepseek_v4_fp8_einsum(
-            out_fp8,
-            out_scale,
-            self.wo_a.weight,
-            self.wo_a.weight_scale_inv,
-            projected,
-            "bhr,hdr->bhd",
-            list(self._einsum_recipe),
-        )"""
-
-NEW = """        # --- sm121-dspark38: reconstruct wo_a tensors for DeepGEMM fp8_einsum
-        # wo_a.weight is 2D (n_local_groups*o_lora_rank, hidden); fp8_bmm needs
-        # rank-3 (H, D, R) where H must equal the a-side group count so the
-        # einsum "bhr,hdr->bhd" contracts correctly (einsum.hpp:164 checks
-        # batch H from a == H from b).  Use n_local_groups, NOT n_groups.
-        _grp = self.n_local_groups
-        _olar = self.o_lora_rank
-        _w3 = self.wo_a.weight.view(_grp, _olar, -1)
-        # a_scale (out_scale): DeepGEMM permutes/owns the A-side itself
-        # (a_scale.permute(1,0,2) -> (H,B,blocks)); the native UE8M0 int scale
-        # (B,H,R/512) already matches the int branch (ceil(R/512)=size(-1)). Do
-        # NOT expand it.
-        _a_sc = out_scale
-        # b_scale (weight_scale_inv) is a tiled (32,32) UE8M0 128x128 scale;
-        # broadcast each tile's scale to its 128 output rows -> per-row fp32.
-        _bs = self.wo_a.weight_scale_inv
-        _nm = _grp * _olar
-        _ntiles = _bs.shape[-2]
-        _btile = _nm // _ntiles
-        _bs_f = (_bs.to(torch.int32) << 23).view(torch.float32)
-        _b_sc = (
-            _bs_f.unsqueeze(1).expand(_ntiles, _btile, -1)
-            .reshape(_nm, -1)
-            .view(_grp, _olar, -1)
-        )
-        torch.ops.vllm.deepseek_v4_fp8_einsum(
-            out_fp8,
-            _a_sc,
-            _w3,
-            _b_sc,
-            projected,
-            "bhr,hdr->bhd",
-            list(self._einsum_recipe),
-        )"""
-
-assert src.count(OLD) == 1, f"expected 1 einsum call, got {src.count(OLD)}"
-P.write_text(src.replace(OLD, NEW))
-
-print("patched dspark.py: drafter wo_a w3/a_sc/b_sc reconstruction")

recipes/deepseek-v4-flash-dspark38/recipe.toml

@@ -13,14 +13,9 @@
 #     because that image was self-contained with everything baked in.
 #   - VLLM_USE_V2_MODEL_RUNNER is dropped (experimental in this fork; the
 #     working glm52 recipe does not set it).
-#   - tiny-decode (VLLM_B12X_TINY_DECODE) is dropped for first bring-up — it is
-#     an opt-in experimental kernel whose GPU load-time repack OOMs.
+#   - tiny-decode (VLLM_B12X_TINY_DECODE) is enabled at GMU 0.78 + CHUNK=8
+#     (the load-time repack needs headroom; it OOM'd at GMU 0.9).
 #   - /cache is mounted so the wrapper's JIT caches persist across restarts.
-#
-# The known-open question is the DeepGEMM `t.dim() == N` assertion
-# (csrc/utils/layout.hpp:39), which fires during MoE scale packing on this
-# fork regardless of --moe-backend. This recipe is the repro for that crash;
-# see specs/bugs/BUG-2026-08-03T204037-production-3.8-image-update-failures.md.
 
 engine = "vllm"
 name = "DeepSeek V4 Flash DSpark (production-3.8)"
@@ -48,7 +43,10 @@ args = [
     "--max-num-seqs", "8",
     "--max-num-batched-tokens", "8192",
     "--max-cudagraph-capture-size", "64",
-    "--gpu-memory-utilization", "0.84",
+    # 0.78 is the safe GMU ceiling on GB10 (weights ~77.7 GiB of ~119.6 GiB);
+    # higher boots but dies under traffic. Also leaves headroom for the tiny
+    # decode kernel's GPU load-time repack.
+    "--gpu-memory-utilization", "0.78",
     "--served-model-name", "deepseek-v4-flash",
     # Engine backends — b12x across attention / MoE / linear (this fork's path)
     "--attention-backend", "B12X_MLA_SPARSE",
@@ -101,7 +99,12 @@ VLLM_PREFIX_CACHE_RETENTION_INTERVAL = "4096"
 B12X_MLA_SM120_UNIFIED = "1"
 B12X_DENSE_SPLITK_TURBO = "1"
 B12X_W4A16_TC_DECODE = "1"
-B12X_MHC_MAX_TOKENS = "16384"
+# Tiny A8 decode kernel: DISABLED — its GPU load-time MoE-weight repack, on
+# top of the tight host RAM (~32 GiB free) + dual-node 155 GiB load, OOM'd both
+# hosts (had to power-cycle via HA). Revisit separately if host RAM headroom is
+# available. Keep the CHUNK pin for reference.
+# VLLM_B12X_TINY_DECODE = "1"
+# VLLM_B12X_TINY_CHUNK = "8"
 
 # --- arch / memory / GPU isolation (operational on the sm121 cluster) ---
 CUTE_DSL_ARCH = "sm_121a"