code/tools/verify_sm120.py

#!/usr/bin/env python3
"""Verify an SM120 build against the SM100-captured golden bundle.

Usage (on the SM120 box, with the patched vLLM venv active):
  python verify_sm120.py /path/to/glm52-sm120-golden [--ckpt /path/to/model]

Stages (each independent; run what's available):
  1. kernels   - replay kernel_vectors.pt through the JIT extension; the
                 custom kernels are plain CUDA and must match to ~1e-2 rel.
  2. moe       - replay moe_layer*.pt through HybridExpertsMoEMethod using
                 real layer weights from --ckpt (loads one layer per file).
  3. attention - replay attn_layer*.pt through the ACTIVE sparse-MLA
                 backend if a standalone replay is wired up for it; else
                 prints the tensor contract so a new kernel can be tested
                 directly against q/kv_pages/topk_indices -> attn_out.
  4. e2e       - after `vllm serve` is up on this box, re-run the prompts
                 in e2e_goldens.pt (greedy) and compare generated ids and
                 top-50 logprob overlap per step.

fp8_ds_mla page layout (656 B/token, block=64 tokens):
  [0:512]   c_kv latent, fp8 e4m3 (512 dims)
  [512:528] 4x fp32 scales (one per 128-dim group of c_kv)
  [528:656] k_pe rope part, bf16 (64 dims)
Attention contract: out[t] = softmax(q[t] . K[topk_indices[t]] / sqrt(576))
  . V[topk_indices[t]] over the 2048 selected tokens (MQA over 576-dim
  latent+rope, 512-dim value = c_kv), invalid indices (-1) masked.
"""
import argparse
import glob
import os
import sys

import torch

TOLS = {"gemv": 2e-2, "dequant": 2e-3}


def rel_err(a, b):
    a, b = a.float(), b.float()
    return ((a - b).norm() / b.norm().clamp_min(1e-9)).item()


def check(name, got, ref, tol):
    e = rel_err(got, ref)
    ok = e <= tol
    print(f"  {'PASS' if ok else 'FAIL'} {name}: rel_err={e:.5f} (tol {tol})")
    return ok


def stage_kernels(gold):
    from vllm.model_executor.layers.quantization.nvfp4_aqlm_hybrid import (
        _get_ext,
    )

    ext = _get_ext()
    vecs = torch.load(os.path.join(gold, "kernel_vectors.pt"),
                      map_location="cpu", weights_only=True)
    print("== stage 1: kernels (needs --ckpt for weights) ==")
    ok = True
    from verify_sm120 import _load_layer  # self-import for reuse

    for li, rec in vecs.items():
        t = _load_layer(ARGS.ckpt, li, "cuda:0")
        x_h = rec["x_h"].cuda()
        x_i = rec["x_i"].cuda()
        aq = rec["aq_ids"].cuda()
        nv = rec["nv_ids"].cuda()
        pairs = [
            ("aqlm_gemv_w13", ext.aqlm_moe_gemv(
                x_h, t["w13_codes"], t["w13_codebooks"], t["w13_scales"], aq),
             "gemv"),
            ("aqlm_gemv_w2c", ext.aqlm_moe_gemv(
                x_i, t["w2c_codes"], t["w2c_codebooks"], t["w2c_scales"], aq),
             "gemv"),
            ("aqlm_dequant_w13", ext.aqlm_moe_dequant(
                t["w13_codes"], t["w13_codebooks"], t["w13_scales"],
                rec["deq_aq"].cuda()), "dequant"),
            ("nvfp4_gemv_w13", ext.nvfp4_moe_gemv(
                x_h, t["nvfp4_w13_packed"], t["nvfp4_w13_bscale"],
                t["nvfp4_w13_scale2"], nv), "gemv"),
            ("nvfp4_gemv_w2", ext.nvfp4_moe_gemv(
                x_i, t["nvfp4_w2_packed"], t["nvfp4_w2_bscale"],
                t["nvfp4_w2_scale2"], nv), "gemv"),
            ("nvfp4_dequant_w2", ext.nvfp4_moe_dequant(
                t["nvfp4_w2_packed"], t["nvfp4_w2_bscale"],
                t["nvfp4_w2_scale2"], rec["deq_nv"].cuda()), "dequant"),
        ]
        print(f" layer {li}:")
        for name, got, kind in pairs:
            ok &= check(name, got.cpu(), rec[name], TOLS[kind])
    return ok


def _load_layer(ckpt, li, device):
    import json

    from safetensors import safe_open

    idx = json.load(open(f"{ckpt}/model.safetensors.index.json"))
    wm = idx["weight_map"]
    p = f"model.layers.{li}.mlp.experts"
    names = ["hyb_kind", "w13_codes", "w13_codebooks", "w13_scales",
             "w2m_codes", "w2m_codebooks", "w2m_scales",
             "w2c_codes", "w2c_codebooks", "w2c_scales",
             "nvfp4_w13_packed", "nvfp4_w13_bscale", "nvfp4_w13_scale2",
             "nvfp4_w2_packed", "nvfp4_w2_bscale", "nvfp4_w2_scale2"]
    t, opened = {}, {}
    for n in names:
        shard = wm[f"{p}.{n}"]
        if shard not in opened:
            opened[shard] = safe_open(f"{ckpt}/{shard}", framework="pt")
        t[n] = opened[shard].get_tensor(f"{p}.{n}").to(device)
    return t


def stage_moe(gold):
    print("== stage 2: hybrid MoE method replay ==")
    from vllm.model_executor.layers.quantization.nvfp4_aqlm_hybrid import (
        HybridExpertsMoEMethod,
    )

    ok = True
    for f in sorted(glob.glob(os.path.join(gold, "moe_layer*_call*.pt"))):
        rec = torch.load(f, map_location="cpu", weights_only=True)
        li = rec["layer_idx"]
        t = _load_layer(ARGS.ckpt, li, "cuda:0")

        class L:
            activation = "silu"
        for k, v in t.items():
            setattr(L, k, v)

        class Moe:
            num_experts = 256

            class moe_parallel_config:
                tp_size = 1
                ep_size = 1

        m = HybridExpertsMoEMethod.__new__(HybridExpertsMoEMethod)
        m.n_nvfp4, m.n_base, m.n_cold = (rec["n_nvfp4"], rec["n_base"],
                                         rec["n_cold"])
        m.moe = Moe()
        m.layer_idx = li
        m._stats_dir = None
        HybridExpertsMoEMethod.process_weights_after_loading(m, L)
        x = rec["x"].cuda()
        out = m._apply_gemv(L, x, rec["topk_weights"].cuda(),
                            rec["topk_ids"].cuda()).to(rec["out"].dtype)
        ok &= check(os.path.basename(f), out.cpu(), rec["out"], 3e-2)
    return ok


def stage_attention(gold):
    print("== stage 3: sparse-MLA attention vectors ==")
    files = sorted(glob.glob(os.path.join(gold, "attn_layer*_call*.pt")))
    for f in files:
        rec = torch.load(f, map_location="cpu", weights_only=True)
        print(f"  {os.path.basename(f)}: q{tuple(rec['q'].shape)} "
              f"pages{tuple(rec['kv_pages'].shape)} "
              f"topk{tuple(rec['topk_indices'].shape)} "
              f"-> out{tuple(rec['attn_out'].shape)}")
    print("  (contract in module docstring; wire your SM120 kernel's "
          "replay here and compare with rel_err <= 3e-2)")
    return True


def stage_e2e(gold, port):
    print("== stage 4: e2e goldens vs running server ==")
    import json
    import urllib.request

    g = torch.load(os.path.join(gold, "e2e_goldens.pt"),
                   map_location="cpu", weights_only=True)
    ok = True
    for i, case in enumerate(g["goldens"]):
        body = {"model": ARGS.ckpt,
                "prompt": case["prompt_token_ids"],
                "max_tokens": len(case["generated_token_ids"]),
                "temperature": 0.0, "logprobs": 50}
        req = urllib.request.Request(
            f"http://localhost:{port}/v1/completions",
            data=json.dumps(body).encode(),
            headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=600) as r:
            resp = json.load(r)
        got_text = resp["choices"][0]["text"]
        match = got_text == case["generated_text"]
        # token-identical is the ideal; tiny kernel diffs may flip late
        # tokens, so also report first-divergence step
        print(f"  case {i}: {'PASS (exact)' if match else 'text differs'}")
        if not match:
            print(f"    golden: {case['generated_text'][:60]!r}")
            print(f"    got:    {got_text[:60]!r}")
            ok = False
    return ok


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("gold")
    ap.add_argument("--ckpt", default="/data/glm52")
    ap.add_argument("--port", type=int, default=0,
                    help="if set, run e2e stage against a live server")
    ap.add_argument("--stages", default="kernels,moe,attention")
    ARGS = ap.parse_args()
    sys.modules["verify_sm120"] = sys.modules["__main__"]

    results = {}
    for s in ARGS.stages.split(","):
        fn = {"kernels": stage_kernels, "moe": stage_moe,
              "attention": stage_attention}.get(s)
        if fn:
            results[s] = fn(ARGS.gold)
    if ARGS.port:
        results["e2e"] = stage_e2e(ARGS.gold, ARGS.port)
    print("\nsummary:", {k: "PASS" if v else "FAIL" for k, v in results.items()})
    sys.exit(0 if all(results.values()) else 1)