"""Validate the SM90 dense requant on a real checkpoint tensor (wq_a + shared expert)."""

import json

import torch
from safetensors import safe_open

from deepinfer.h20_ops.quant import requant_dense_fp8, ue8m0_to_fp32

root = "/data/models/DeepSeek-V4.1-Flash-1017"
with open(f"{root}/model.safetensors.index.json") as f:
    idx = json.load(f)["weight_map"]

targets = [
    "layers.0.attn.wq_a.weight",
    "layers.0.ffn.shared_experts.w1.weight",
]
for name in targets:
    with safe_open(f"{root}/{idx[name]}", framework="pt", device="cpu") as f:
        w = f.get_tensor(name)
        s = f.get_tensor(name.replace(".weight", ".scale"))
    dev = "cuda"
    w, s = w.to(dev), s.to(dev)
    ref = w.to(torch.float32) * ue8m0_to_fp32(s).repeat_interleave(32, 0).repeat_interleave(32, 1)
    w8, sf = requant_dense_fp8(w, ue8m0_to_fp32(s), dst_block=128)
    rec = w8.to(torch.float32) * sf.repeat_interleave(128, 0).repeat_interleave(128, 1)
    err = (rec - ref).abs()
    print(
        f"{name}: {tuple(w.shape)} sf {tuple(s.shape)} -> {tuple(sf.shape)} | "
        f"max_abs {err.max().item():.5g} rel {(err.max() / ref.abs().max()).item():.3e}"
    )
print("REAL_WEIGHT_REQUANT_OK")