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

#!/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.