#!/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")