#!/usr/bin/env python3
"""
patch_dev145_bf16_oproj.py

Fixes: AttributeError: 'ColumnParallelLinear' object has no attribute
'weight_scale' in vllm/models/deepseek_v4/nvidia/ops/o_proj.py:68 when
running --speculative-config '{"method":"mtp",...}' on
canada-quant/DeepSeek-V4-Flash-W4A16-FP8-MTP with the dev145 SM89 wheel.

Root cause: deep_gemm_fp8_o_proj() unconditionally assumes an FP8-quantized
wo_a (reads weight_scale_inv / weight_scale). The target model's attention is
FP8 (compressed-tensors) and sails through; the checkpoint's MTP drafter head
is BF16/unquantized, so its wo_a is a plain ColumnParallelLinear with no
scale attribute. The fork author's spec-decode validation used the separate
DSpark checkpoint, never method=mtp with a BF16 head.

Fix: insert an early branch for scale-less wo_a that replicates the fused
op's math in plain torch: inverse RoPE (interleaved even/odd pairs on the
LAST rope_dim dims of each head, rotation by -theta, mirroring
_fused_inv_rope_fp8_quant_per_head in
common/ops/fused_inv_rope_fp8_quant.py), grouped bf16 einsum against
wo_a.weight viewed [n_groups, o_lora_rank, heads_per_group*head_dim], then
wo_b. Unfused and bf16, but the drafter is a single tiny layer, so the cost
is noise.

Usage:
    # inside the serving venv
    python patch_dev145_bf16_oproj.py            # patches installed vllm
    python patch_dev145_bf16_oproj.py <path>     # patches an explicit file

Idempotent (marker-guarded); writes <file>.bak_bf16 once before modifying.
"""
import pathlib
import py_compile
import sys

MARKER = "SM89 patch: BF16/unquantized wo_a"

OLD = '''    Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` /
    ``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``.
    """
    o_fp8, o_scale = fused_inv_rope_fp8_quant(
'''

NEW = '''    Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` /
    ``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``.
    """
    if (
        getattr(wo_a, "weight_scale_inv", None) is None
        and getattr(wo_a, "weight_scale", None) is None
    ):
        # ---- SM89 patch: BF16/unquantized wo_a (e.g. the BF16 MTP drafter
        # head in canada-quant/DeepSeek-V4-Flash-W4A16-FP8-MTP). The fused
        # FP8 path below requires block scales; replicate its math in plain
        # torch: inverse RoPE (interleaved even/odd pairs on the LAST
        # ``rope_dim`` dims, rotation by -theta, mirroring
        # _fused_inv_rope_fp8_quant_per_head), grouped bf16 einsum, wo_b.
        num_tokens, num_heads, head_dim = o.shape
        of = o.to(torch.float32)
        cs = cos_sin_cache[positions.to(torch.long)]
        half = rope_dim // 2
        cos = cs[:, :half].unsqueeze(1)
        sin = cs[:, half:].unsqueeze(1)
        rope = of[..., nope_dim:]
        x1 = rope[..., 0::2]
        x2 = rope[..., 1::2]
        rope_inv = torch.stack(
            (x1 * cos + x2 * sin, x2 * cos - x1 * sin), dim=-1
        ).flatten(-2)
        o_inv = torch.cat((of[..., :nope_dim], rope_inv), dim=-1)
        o_grouped = o_inv.to(torch.bfloat16).reshape(
            num_tokens, n_groups, heads_per_group * head_dim
        )
        w = wo_a.weight.to(torch.bfloat16).view(
            n_groups, o_lora_rank, heads_per_group * head_dim
        )
        z_bf16 = torch.einsum("bhr,hdr->bhd", o_grouped, w)
        return wo_b(z_bf16.flatten(1))
    o_fp8, o_scale = fused_inv_rope_fp8_quant(
'''


def resolve_installed_target() -> pathlib.Path:
    import vllm  # noqa: PLC0415

    return (
        pathlib.Path(vllm.__file__).parent
        / "models" / "deepseek_v4" / "nvidia" / "ops" / "o_proj.py"
    )


def main() -> int:
    if len(sys.argv) > 1:
        target = pathlib.Path(sys.argv[1])
    else:
        target = resolve_installed_target()

    if not target.is_file():
        print(f"XX target not found: {target}", file=sys.stderr)
        return 2

    src = target.read_text()

    if MARKER in src:
        print(f"OK already patched, nothing to do: {target}")
        return 0

    if OLD not in src:
        print(
            "XX expected code block not found -- file differs from the "
            "dev145 (g8c631d45e) layout this patch targets. Refusing to "
            f"guess. File: {target}",
            file=sys.stderr,
        )
        return 2

    backup = target.with_suffix(target.suffix + ".bak_bf16")
    if not backup.exists():
        backup.write_text(src)
        print(f"OK backup written: {backup}")

    target.write_text(src.replace(OLD, NEW, 1))
    py_compile.compile(str(target), doraise=True)
    print(f"OK patched + compiled: {target}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())