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()