runtime/scripts/kernel_selftest.py

from __future__ import annotations

import argparse
import math
import time
import json
import hashlib
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import torch
from safetensors import safe_open

from orbitquant_wan_a2.nibbles import pack_uint4, unpack_uint4
from orbitquant_wan_a2.reference import a4_pack_reference, dequant_packed, w4a4_linear_reference
from orbitquant_wan_a2.triton_w4a4 import (
    PackedActivation,
    a4_pack_triton,
    triton_available,
    w4a4_linear_triton,
)


def make_rotation(k: int, h: int, device):
    g = torch.Generator(device="cpu").manual_seed(12000 + k)
    perm = torch.randperm(k, generator=g, dtype=torch.int64).to(device)
    signs = (torch.randint(0, 2, (k,), generator=g, dtype=torch.int8) * 2 - 1).to(device)
    # Strictly increasing symmetric 16-level surrogate.  Kernel parity does not
    # depend on the particular Lloyd-Max values; real runs load the frozen bank.
    cb = torch.linspace(-0.12, 0.12, 16, dtype=torch.float32, device=device)
    return perm, signs, cb


def sync_ms(fn, iters=10):
    for _ in range(2):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) * 1000 / iters


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--allow-no-cuda", action="store_true")
    ap.add_argument("--rows", type=int, default=17)
    ap.add_argument(
        "--gate-file", default="/content/orbitquant_w4a4_cuda_gate.json"
    )
    ap.add_argument(
        "--packed-dir",
        default=None,
        help="Optional packed runtime artifact. When supplied, gate representative real artifact weights too.",
    )
    ap.add_argument(
        "--rotation-bank",
        default="/content/OrbitQuant_Animate2_W4A4_PACKED/orbitquant_rotations.safetensors",
        help="Use the actual frozen Project-A permutation/sign/codebook bank when present",
    )
    args = ap.parse_args()

    if not torch.cuda.is_available() or not triton_available():
        msg = "CUDA + Triton are required for the kernel gate"
        if args.allow_no_cuda:
            print("SKIP:", msg)
            return
        raise RuntimeError(msg)

    device = torch.device("cuda")
    print("GPU:", torch.cuda.get_device_name(device))
    print("CC :", torch.cuda.get_device_capability(device))
    print("Torch:", torch.__version__)

    bank = None
    bank_path = Path(args.rotation_bank)
    if bank_path.is_file():
        from orbitquant_wan_a2.rotation_bank import RotationBank
        bank = RotationBank.load(bank_path)
        print("Rotation bank:", bank_path)

    for k, h in [(5120, 1024), (13824, 512)]:
        torch.manual_seed(100 + k)
        x = torch.randn(args.rows, k, device=device, dtype=torch.bfloat16)
        if bank is not None:
            item = bank.tensors[k]
            perm = item["perm"].to(device=device, dtype=torch.int64)
            signs = item["signs"].to(device=device, dtype=torch.int8)
            cb = item["codebook"].to(device=device, dtype=torch.float32)
            assert int(item["block_size"].item()) == h
        else:
            perm, signs, cb = make_rotation(k, h, device)

        ref_pack, ref_scale = a4_pack_reference(x, perm, signs, h, cb)
        got = a4_pack_triton(x, perm, signs, h, cb)
        torch.cuda.synchronize()

        codes_equal = torch.equal(got.codes.cpu(), ref_pack.cpu())
        scale_err = (got.scale.float() - ref_scale.float()).abs()
        print(f"A4 D={k}: codes_equal={codes_equal} scale_max={scale_err.max().item():.4e} scale_mean={scale_err.mean().item():.4e}")
        if not codes_equal:
            mism = (unpack_uint4(got.codes, k) != unpack_uint4(ref_pack, k)).float().mean().item()
            # Floating reduction order can only be accepted if code decisions
            # still agree to effectively all entries.  A packed runtime must
            # not silently drift across centroid boundaries.
            raise RuntimeError(f"A4 Triton code mismatch for D={k}: fraction={mism:.4e}")
        torch.testing.assert_close(got.scale, ref_scale, rtol=2e-5, atol=2e-5)
        # Stronger gate: after the scale/code representation is dequantized to
        # the BF16 activation consumed by Project-A F.linear, the Triton path
        # must reproduce the PyTorch fake-A4 activation bit-for-bit.
        a_ref_bf16 = dequant_packed(ref_pack, ref_scale, cb, k, torch.bfloat16)
        a_got_bf16 = dequant_packed(got.codes, got.scale, cb, k, torch.bfloat16)
        a_exact = float((a_ref_bf16 == a_got_bf16).float().mean().item())
        print(f"A4 D={k}: BF16_dequant_exact={a_exact:.10f}")
        if a_exact != 1.0:
            raise RuntimeError(f"A4 BF16 dequant mismatch for D={k}: exact_fraction={a_exact:.10f}")

        # Exercise the packed GEMM at the *real* Wan-Animate-2 projection
        # geometries, not a toy N.  This covers every Project-A target shape:
        #   5120 -> 5120    attention + FFN output projection
        #   5120 -> 13824   FFN expansion
        #   13824 -> 5120   FFN contraction
        real_ns = [5120, 13824] if k == 5120 else [5120]
        last = None
        for n in real_ns:
            w_codes = torch.randint(0, 16, (n, k), device=device, dtype=torch.uint8)
            wpack_row = pack_uint4(w_codes)
            wpack = wpack_row.transpose(0, 1).contiguous()
            wscale = torch.rand(n, device=device, dtype=torch.float32) * 3.0 + 0.1
            bias = torch.randn(n, device=device, dtype=torch.bfloat16)
            ref = w4a4_linear_reference(ref_pack, ref_scale, wpack_row, wscale, cb, k, bias)
            out = w4a4_linear_triton(got, wpack, wscale, cb, bias)
            torch.cuda.synchronize()
            diff = (out.float() - ref.float()).abs()
            print(
                f"GEMM M={args.rows} N={n} K={k}: "
                f"max={diff.max().item():.5f} mean={diff.mean().item():.5f}"
            )
            torch.testing.assert_close(out, ref, rtol=2e-2, atol=1.5e-1)
            last = (wpack, wscale, bias, n)
            del w_codes, wpack_row, ref, out, diff

        a_ms = sync_ms(lambda: a4_pack_triton(x, perm, signs, h, cb), iters=5)
        wpack, wscale, bias, n = last
        g_ms = sync_ms(lambda: w4a4_linear_triton(got, wpack, wscale, cb, bias), iters=5)
        print(f"TIMING D={k}: A4_pack={a_ms:.3f} ms  packed_GEMM_N{n}={g_ms:.3f} ms")

    packed_gate = None
    if args.packed_dir is not None:
        pdir = Path(args.packed_dir)
        manifest_path = pdir / "packed_manifest.json"
        if not manifest_path.is_file():
            raise RuntimeError(f"packed manifest not found: {manifest_path}")
        manifest_bytes = manifest_path.read_bytes()
        manifest = json.loads(manifest_bytes)
        if manifest.get("format") != "OrbitQuant_WanAnimate2_direct_source_packed_nonuniform_W4A4_v3":
            raise RuntimeError(f"wrong packed artifact format: {manifest.get('format')}")
        if manifest.get("target_count") != 480 or len(manifest.get("targets", {})) != 480:
            raise RuntimeError("packed artifact must contain exactly 480 W4 targets")
        if manifest.get("weight_storage_layout") != "K_half_by_N":
            raise RuntimeError("packed artifact must use GEMM-native [K/2,N] layout")

        # One real packed target for each target matrix geometry.  This validates
        # artifact layout + stored BF16 row norms + LUT decode against the custom
        # Triton GEMM, not merely random synthetic uint4 matrices.
        selected = {}
        for stored_key, info in manifest["targets"].items():
            shape = (int(info["output_dim"]), int(info["input_dim"]))
            selected.setdefault(shape, (stored_key, info))
        required_shapes = {(5120, 5120), (13824, 5120), (5120, 13824)}
        missing_shapes = required_shapes - set(selected)
        if missing_shapes:
            raise RuntimeError(f"packed artifact is missing target geometries: {sorted(missing_shapes)}")

        actual_results = []
        for n, k in sorted(required_shapes):
            stored_key, info = selected[(n, k)]
            with safe_open(str(pdir / info["shard"]), framework="pt", device="cpu") as sf:
                wp = sf.get_tensor(info["packed_tensor"]).to(device)
                ws = sf.get_tensor(info["scale_tensor"]).to(device)
            if tuple(wp.shape) != (k // 2, n):
                raise RuntimeError(
                    f"actual packed layout mismatch for {stored_key}: {tuple(wp.shape)} != {(k // 2, n)}"
                )
            item = bank.tensors[k] if bank is not None else None
            if item is None:
                h = 1024 if k == 5120 else 512
                perm, signs, cb = make_rotation(k, h, device)
            else:
                h = int(item["block_size"].item())
                perm = item["perm"].to(device=device, dtype=torch.int64)
                signs = item["signs"].to(device=device, dtype=torch.int8)
                cb = item["codebook"].to(device=device, dtype=torch.float32)

            torch.manual_seed(9100 + n + k)
            x = torch.randn(args.rows, k, device=device, dtype=torch.bfloat16)
            ap_ref, as_ref = a4_pack_reference(x, perm, signs, h, cb)
            ap = a4_pack_triton(x, perm, signs, h, cb)
            if not torch.equal(ap.codes, ap_ref):
                raise RuntimeError(f"actual-artifact A4 code mismatch for {stored_key}")
            # Bias is a passthrough tensor when that official linear has one.
            bias = None
            bias_key = stored_key[:-len("weight")] + "bias"
            binfo = manifest.get("passthrough_tensors", {}).get(bias_key)
            if binfo is not None:
                with safe_open(str(pdir / binfo["shard"]), framework="pt", device="cpu") as sf:
                    bias = sf.get_tensor(binfo["tensor"]).to(device=device, dtype=torch.bfloat16)

            ref = w4a4_linear_reference(
                ap_ref, as_ref, wp.transpose(0, 1).contiguous(), ws, cb, k, bias
            )
            out = w4a4_linear_triton(ap, wp, ws, cb, bias)
            torch.cuda.synchronize()
            diff = (out.float() - ref.float()).abs()
            torch.testing.assert_close(out, ref, rtol=2e-2, atol=1.5e-1)
            actual_results.append({
                "target": stored_key,
                "shape_NK": [n, k],
                "max_abs": float(diff.max().item()),
                "mean_abs": float(diff.mean().item()),
            })
            print(
                f"ACTUAL ARTIFACT {stored_key} N={n} K={k}: "
                f"max={diff.max().item():.5f} mean={diff.mean().item():.5f}"
            )
            del wp, ws, x, ap_ref, as_ref, ap, ref, out, diff, bias
            torch.cuda.empty_cache()

        shard_sizes = {name: int((pdir / name).stat().st_size) for name in manifest.get("shards", [])}
        packed_gate = {
            "packed_dir": str(pdir.resolve()),
            "manifest_sha256": hashlib.sha256(manifest_bytes).hexdigest(),
            "shard_sizes": shard_sizes,
            "actual_targets": actual_results,
            "format": manifest["format"],
            "weight_storage_layout": manifest["weight_storage_layout"],
        }
        print("ACTUAL PACKED ARTIFACT CUDA GATE: PASS")

    gate = {
        "status": "PASS",
        "gpu": torch.cuda.get_device_name(device),
        "compute_capability": list(torch.cuda.get_device_capability(device)),
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "dimensions": [5120, 13824],
        "real_projection_shapes": [[5120, 5120], [13824, 5120], [5120, 13824]],
        "rows_tested": int(args.rows),
        "a4": "RPBH -> post-rotation L2 -> Lloyd-Max -> uint4",
        "a4_bf16_dequant_bit_exact": True,
        "w4a4": "packed nonuniform uint4 W [K/2,N] x packed uint4 A -> BF16 tile dequant -> FP32 accumulate",
        "weight_storage_layout": "K_half_by_N",
        "packed_artifact": packed_gate,
    }
    Path(args.gate_file).write_text(json.dumps(gate, indent=2))
    print("\nCUDA KERNEL GATE: PASS")
    print("  RPBH -> post-rotation norm -> Lloyd-Max A4 -> uint4: PASS")
    print("  packed nonuniform W4 x packed A4 Triton GEMM: PASS")
    print("  gate:", args.gate_file)


if __name__ == "__main__":
    main()