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