File size: 12,427 Bytes
f2c0505 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 | 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()
|