"""Microbench NVFP4 W4A4 GEMM backends at SenseNova-Loop decoder shapes (run inside the SGLang image)."""
import torch, sys
from flashinfer import fp4_quantize, mm_fp4
torch.manual_seed(0); dev = "cuda"
SHAPES = {"gate_up": (4096, 24576), "down": (12288, 4096), "gdn_qkvz": (4096, 12288), "attn_qkv": (4096, 10240), "o_proj": (4096, 4096)}
MS = [int(m) for m in (sys.argv[1] if len(sys.argv) > 1 else "1,4,8,16,32,64").split(",")]
BACKENDS = ["cudnn", "cutlass", "b12x", "cake", "cutedsl_low_latency", "auto"]
def gs(t): return (6.0 * 448.0 / t.abs().max().float()).float()
def timeit(fn, n=50):
    for _ in range(3): fn()
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        for _ in range(n): fn()
    torch.cuda.synchronize(); s, e = torch.cuda.Event(True), torch.cuda.Event(True)
    s.record(); g.replay(); e.record(); torch.cuda.synchronize(); return s.elapsed_time(e) / n * 1e3
kda = None
try:
    import sglang.kernels.kda_kernels.qwen3x_nvfp4_gemm as q
    q._has_runtime = lambda d: True
    from sglang.kernels.ops.gemm import try_qwen3x_nvfp4_gemm as kda
except Exception as ex: print("kda import fail", ex)
tot = {}
for name, (K, N) in SHAPES.items():
    w = torch.randn(N, K, device=dev, dtype=torch.bfloat16) * 0.02; wg = gs(w)
    wq, ws = fp4_quantize(w, wg, 16, False, True)
    gbytes = N * K * (0.5 + 1 / 16) / 1e9
    for M in MS:
        x = torch.randn(M, K, device=dev, dtype=torch.bfloat16); xg = gs(x)
        xq, xs = fp4_quantize(x, xg, 16, False, True); alpha = (1.0 / (xg * wg)).reshape(1)
        ref = (x.float() @ w.float().t())
        row = []
        for be in BACKENDS + (["kda"] if kda else []):
            try:
                out = torch.empty(M, N, device=dev, dtype=torch.bfloat16)
                if be == "kda":
                    f = lambda: kda(xq, wq, xs, ws, alpha, torch.bfloat16, N)
                    if f() is None: raise RuntimeError("unsupported shape")
                else:
                    f = lambda: mm_fp4(xq, wq.t(), xs, ws.t(), alpha, torch.bfloat16, out, backend=be)
                o = f(); o = out if o is None else o
                err = ((o.float() - ref).norm() / ref.norm()).item()
                us = timeit(f); row.append(f"{be}={us:6.1f}us({gbytes/us*1e6:4.0f}GB/s,e{err:.3f})")
                tot.setdefault((be, M), 0); tot[(be, M)] += us
            except Exception as ex:
                row.append(f"{be}=FAIL[{str(ex)[:40]}]")
        print(f"{name:9s} M={M:3d} " + " ".join(row), flush=True)
print("sum over 5 shapes (us):", {f"{b}@{m}": round(v, 1) for (b, m), v in sorted(tot.items(), key=lambda t: (t[0][1], t[1]))})