#!/usr/bin/env python
"""Issue 5 round two: raw-op probe of the cutlass W8A8 int8 path on this stack.

Bypasses every layer/loader and tests exactly what CutlassInt8ScaledMMLinearKernel
calls: ops.scaled_int8_quant (dynamic per-token symmetric) then
ops.cutlass_scaled_mm, against a float32 reference. A failure here means the
cutlass int8 kernels themselves are broken on this stack (CUDA 13.4 / sm_86 /
torch 2.13.0+cu130 / vllm 0.31.0), independent of any checkpoint or loader.

Usage: cutlass_ops_probe.py
"""

from __future__ import annotations

import json

import torch
from vllm import _custom_ops as ops


def rel(a: torch.Tensor, b: torch.Tensor) -> float:
    return float((a - b).norm() / b.norm())


def main() -> int:
    torch.manual_seed(2301)
    device = torch.device("cuda")
    results = []

    for m, k, n in ((64, 256, 128), (512, 4096, 4096), (512, 4096, 12288)):
        w = torch.randn(n, k, device=device)
        w_scale = w.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
        w_q = torch.round(w / w_scale).clamp(-127, 127).to(torch.int8)

        x = torch.randn(m, k, dtype=torch.bfloat16, device=device)
        x_q, x_s, _ = ops.scaled_int8_quant(x.contiguous(), None, None, symmetric=True)

        # Reference: dequantized per-token acts @ dequantized per-channel weights.
        ref = (x_q.float() * x_s.float()) @ (w_q.float() * w_scale.float()).T

        w_q_t = w_q.t()
        out = ops.cutlass_scaled_mm(
            x_q, w_q_t, scale_a=x_s, scale_b=w_scale, out_dtype=torch.bfloat16
        ).float()

        # Also check the quantizer op itself against a float32 reference.
        act_scale_ref = x.float().abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
        x_q_ref = torch.round(x.float() / act_scale_ref).clamp(-127, 127).to(torch.int8)
        quant_match = bool(torch.equal(x_q, x_q_ref)) and bool(
            torch.allclose(x_s.float(), act_scale_ref, rtol=1e-3, atol=1e-7)
        )

        results.append({
            "mk": [m, k, n],
            "quant_op_matches_reference": quant_match,
            "gemm_rel_l2": round(rel(out, ref), 5),
            "out_finite": bool(torch.isfinite(out).all()),
            "scale_a_shape": list(x_s.shape),
            "scale_b_shape": list(w_scale.shape),
        })

    print(json.dumps(results, indent=1))
    return 0


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