docker-images: deepseek-v4-flash-dspark38

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

@@ -0,0 +1,39 @@
+# syntax=docker/dockerfile:1.7
+# DeepSeek-V4-Flash-DSpark on aidendle94 sparkrun-vllm-ds4-gb10:production-3.8,
+# with a root-cause patch for BUG-2026-08-03T204037.
+#
+# The REAL blocker (found by boot-testing on ds4) is NOT the DeepGEMM
+# `t.dim() == N` assertion from the original bug write-up. On this image the
+# DSV4 target + drafter load cleanly via the b12x MoE backend — DeepGEMM scale
+# packing never fires. Boot instead crashes during KV-cache profiling
+# (profile_run -> dummy forward, under torch.compile/AOT), inside
+# b12x/integration/residual_kernels.py `_to_kernel_tensor`:
+#
+#   from_dlpack(tensor, ...) -> tensor.__dlpack__()
+#   BufferError: Can't export tensors that require gradient, use tensor.detach()
+#
+# Under torch.compile the norm parameter carries requires_grad=True (inductor
+# keeps an autograd graph), and cutlass's from_dlpack refuses grad-requiring
+# tensors. Fix: detach the tensor before dlpack export — a no-op for the
+# requires_grad=False case, correct for the compiled case.
+
+FROM aidendle94/sparkrun-vllm-ds4-gb10:production-3.8
+
+# Patch 1: b12x cutlass from_dlpack rejects grad-requiring tensors under
+# torch.compile (BufferError during KV-cache profiling). detach() before export.
+RUN sed -i \
+    's/from_dlpack(tensor, assumed_align=assumed_align)/from_dlpack(tensor.detach(), assumed_align=assumed_align)/' \
+    /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: 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
+# divergent request contexts doesn't kill the engine (see patch script).
+COPY patches/patch_dspark_uniform.py /tmp/patch_dspark_uniform.py
+RUN /opt/env/bin/python3 /tmp/patch_dspark_uniform.py

docker-images/deepseek-v4-flash-dspark38/build.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+# Build the patched production-3.8 overlay for deepseek-v4-flash-dspark38.
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-sparkrun-vllm-ds4-gb10:dspark38}"
+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 --cache-from "type=local,src=${CACHE_DIR}" \
+    --cache-to "type=local,dest=${CACHE_DIR},mode=max" \
+    --progress=plain --load -t "${TAG}" . 2>&1
+echo "=== Done: ${TAG} ==="

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_uniform.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""Patch dspark_proposer.py so a DSpark batch with mixed per-request context
+lengths does not kill the engine (observed under `tool-eval-bench --parallel 4`).
+
+Stock behaviour: `_trim_rejected_target_context` + `_view_by_request` require
+every request in a DSpark batch to have the SAME flattened length. When
+requests have divergent context lengths (e.g. [1, 2052]) this raises ValueError
+in the worker -> EngineCore dies -> EngineDeadError on every later request.
+
+There are two separate uniform checks that both crash:
+  1. `_trim_rejected_target_context` requires uniform *effective* lengths after
+     rejection-trimming (raises only when any token was rejected).
+  2. `_view_by_request` requires the flattened input to be a multiple of
+     batch_size (crashes when all drafts are ACCEPTED and contexts differ,
+     because the raw per-request lengths are returned unpadded).
+
+Fix: restructure `_trim_rejected_target_context` to ALWAYS split the flattened
+per-request chunks via `query_start_loc`, trim rejected tokens (when provided),
+and pad every shorter request to the batch-max effective length by repeating its
+last valid token (hidden state + position). This keeps the batch rectangular /
+uniform for `_view_by_request` and `prefill_main` while the draft continuation
+point `[:, -1]` stays the real last token. Only falls back to returning the raw
+tensor if per-request boundaries are unavailable.
+"""
+import pathlib
+
+P = pathlib.Path(
+    "/opt/env/lib/python3.12/site-packages/vllm/v1/spec_decode/dspark_proposer.py"
+)
+src = P.read_text()
+
+OLD = """        if num_rejected_tokens_gpu is None or common_attn_metadata is None:
+            return target_hidden_states, target_positions
+
+        rejected = num_rejected_tokens_gpu.detach().cpu().tolist()
+        if not any(int(value) for value in rejected):
+            return target_hidden_states, target_positions
+
+        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
+        if query_start_loc_cpu is None:
+            query_start_loc_cpu = common_attn_metadata.query_start_loc.detach().cpu()
+        query_starts = query_start_loc_cpu.tolist()
+        if len(query_starts) != len(rejected) + 1:
+            raise ValueError(
+                "DSpark rejected-context trimming requires query_start_loc "
+                "to have batch_size + 1 entries; got "
+                f"{len(query_starts)} starts for {len(rejected)} requests."
+            )
+
+        hidden_chunks: list[torch.Tensor] = []
+        position_chunks: list[torch.Tensor] = []
+        effective_lengths: list[int] = []
+        flat_positions = target_positions.reshape(-1)
+        for req_index, num_rejected in enumerate(rejected):
+            start = int(query_starts[req_index])
+            end = int(query_starts[req_index + 1]) - int(num_rejected)
+            if end <= start:
+                raise ValueError(
+                    "DSpark rejected-context trimming removed every token for "
+                    f"request {req_index}: start={start}, end={end}."
+                )
+            hidden_chunks.append(target_hidden_states[start:end])
+            position_chunks.append(flat_positions[start:end])
+            effective_lengths.append(end - start)
+
+        if len(set(effective_lengths)) != 1:
+            raise ValueError(
+                "DSpark currently requires uniform effective per-request "
+                "target context lengths after rejection trimming; got "
+                f"{effective_lengths}."
+            )
+
+        return torch.cat(hidden_chunks, dim=0), torch.cat(position_chunks, dim=0)"""
+
+NEW = """        # sm121-dspark38: if we cannot determine per-request boundaries, fall
+        # back to returning the tensors untouched.
+        qsl_cpu = getattr(common_attn_metadata, "query_start_loc_cpu", None)
+        if common_attn_metadata is None or (
+            qsl_cpu is None
+            and getattr(common_attn_metadata, "query_start_loc", None) is None
+        ):
+            return target_hidden_states, target_positions
+        if qsl_cpu is None:
+            qsl_cpu = common_attn_metadata.query_start_loc.detach().cpu()
+        query_starts = qsl_cpu.tolist()
+        nreq = len(query_starts) - 1
+        if nreq < 1:
+            return target_hidden_states, target_positions
+
+        rejected = (
+            num_rejected_tokens_gpu.detach().cpu().tolist()
+            if num_rejected_tokens_gpu is not None
+            else [0] * nreq
+        )
+        if len(query_starts) != len(rejected) + 1:
+            raise ValueError(
+                "DSpark rejected-context trimming requires query_start_loc "
+                "to have batch_size + 1 entries; got "
+                f"{len(query_starts)} starts for {len(rejected)} requests."
+            )
+
+        hidden_chunks: list[torch.Tensor] = []
+        position_chunks: list[torch.Tensor] = []
+        effective_lengths: list[int] = []
+        flat_positions = target_positions.reshape(-1)
+        for req_index, num_rejected in enumerate(rejected):
+            start = int(query_starts[req_index])
+            end = int(query_starts[req_index + 1]) - int(num_rejected)
+            if end <= start:
+                raise ValueError(
+                    "DSpark rejected-context trimming removed every token for "
+                    f"request {req_index}: start={start}, end={end}."
+                )
+            hidden_chunks.append(target_hidden_states[start:end])
+            position_chunks.append(flat_positions[start:end])
+            effective_lengths.append(end - start)
+
+        # sm121-dspark38: pad every shorter request to the batch-max effective
+        # length by repeating its last valid token (hidden state + position),
+        # so the batch stays rectangular/uniform for _view_by_request and
+        # prefill_main while the draft continuation point ([:, -1]) remains the
+        # real last token. The stock fork hard-rejects non-uniform lengths.
+        max_eff = max(effective_lengths)
+        for i, (eff, chunk, pos_chunk) in enumerate(
+            zip(effective_lengths, hidden_chunks, position_chunks)
+        ):
+            if eff < max_eff:
+                pad = max_eff - eff
+                hidden_chunks[i] = torch.cat(
+                    [chunk, chunk[-1:].expand(pad, -1).contiguous()], dim=0
+                )
+                position_chunks[i] = torch.cat(
+                    [pos_chunk, pos_chunk[-1:].expand(pad).contiguous()], dim=0
+                )
+        effective_lengths[:] = [max_eff] * len(effective_lengths)
+
+        return torch.cat(hidden_chunks, dim=0), torch.cat(position_chunks, dim=0)"""
+
+assert src.count(OLD) == 1, f"expected 1 block, got {src.count(OLD)}"
+P.write_text(src.replace(OLD, NEW))
+print("patched dspark_proposer.py: always-pad to uniform context")