"""Compare tile_kernels' compiled kernels against the pure-torch references. The point is to find out whether `tile_kernels.torch` is a usable stand-in on a card where the kernels cannot be compiled, so equality (bitwise for quantised bytes, a tight tolerance for floats) is the only acceptable answer. A "close enough" cast changes tokens downstream, and a stand-in that changes tokens is not a stand-in. Shapes and flags are the ones DeepInfer actually passes on SM90 (`get_fp8_cast_params`: fmt="e4m3", num_per_channels=128, round_sf=True, use_packed_ue8m0=False, use_tma_aligned_col_major_sf=True), not invented ones. python tk_ab.py [op ...] """ import sys import traceback import tile_kernels.engram as tk_engram import tile_kernels.quant as tk_quant import tile_kernels.torch as ref import tile_kernels.transform as tk_tr import torch torch.manual_seed(0) DEV = "cuda" # The quantisation recipe the engine uses on Hopper. CAST = dict(fmt="e4m3", num_per_channels=128) CAST_FLAGS = dict(round_sf=True, use_packed_ue8m0=False, use_tma_aligned_col_major_sf=True) FAILED = [] def report(name: str, a: object, b: object, *, exact: bool, rtol: float = 1e-6) -> None: if isinstance(a, tuple) and isinstance(b, tuple): if len(a) != len(b): print(f" FAIL {name}: {len(a)} vs {len(b)} outputs") FAILED.append(name) return for i, (x, y) in enumerate(zip(a, b, strict=True)): report(f"{name}[{i}]", x, y, exact=exact, rtol=rtol) return if a is None or b is None: ok = a is None and b is None print(f" {'ok ' if ok else 'FAIL'} {name}: one side None ({a is None} vs {b is None})") if not ok: FAILED.append(name) return if not isinstance(a, torch.Tensor): a = torch.as_tensor(a, device=DEV) if not isinstance(b, torch.Tensor): b = torch.as_tensor(b, device=DEV) if a.shape != b.shape: print(f" FAIL {name}: shape {tuple(a.shape)} vs {tuple(b.shape)}") FAILED.append(name) return if a.dtype != b.dtype: print(f" FAIL {name}: dtype {a.dtype} vs {b.dtype}") FAILED.append(name) return if exact or not a.is_floating_point(): if torch.equal(a, b): print(f" ok {name}: exact") return bad = int((a != b).sum()) idx = (a != b).nonzero()[0].tolist() print(f" FAIL {name}: {bad}/{a.numel()} differ, first {idx}: {a[tuple(idx)].item()} vs {b[tuple(idx)].item()}") FAILED.append(name) return diff = (a.float() - b.float()).abs() denom = b.float().abs().clamp_min(1e-6) rel = (diff / denom).max().item() ok = rel <= rtol print(f" {'ok ' if ok else 'FAIL'} {name}: max rel {rel:.3e}") if not ok: FAILED.append(name) def op_per_token_cast() -> None: for m, n in [(128, 4096), (1, 2048), (257, 128), (64, 7168)]: x = (torch.randn((m, n), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous() got = tk_quant.per_token_cast(x, num_per_channels=CAST["num_per_channels"], **CAST_FLAGS, fmt="e4m3") want = ref.cast( x, "e4m3", block_size=(1, CAST["num_per_channels"]), **CAST_FLAGS, ) print(f"per_token_cast m={m} n={n}") report("cast", tuple(got) if isinstance(got, tuple) else got, tuple(want) if isinstance(want, tuple) else want, exact=True) def op_cast_back() -> None: # Round-trip: quantise with the kernel, then dequantise with each side. `fmt` here is the # *output* type, not the quantised format. for m, n in [(128, 4096), (64, 2048)]: x = (torch.randn((m, n), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous() q = tk_quant.per_token_cast(x, num_per_channels=CAST["num_per_channels"], **CAST_FLAGS, fmt="e4m3") got = tk_quant.cast_back(tuple(q), "bf16", (1, CAST["num_per_channels"])) want = ref.cast_back(tuple(q), "bf16", block_size=(1, CAST["num_per_channels"])) print(f"cast_back m={m} n={n}") report("dequant", got, want, exact=False, rtol=1e-3) def op_swiglu() -> None: # The kernel fuses activation and cast; the reference separates them, so compose it. for m, n in [(128, 4096), (33, 512)]: x = (torch.randn((m, n), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous() got = tk_quant.swiglu_forward_and_per_token_cast(x, num_per_channels=CAST["num_per_channels"], **CAST_FLAGS, fmt="e4m3") activated = ref.swiglu_forward(x) if isinstance(activated, tuple): activated = activated[0] want = ref.cast(activated, "e4m3", block_size=(1, CAST["num_per_channels"]), **CAST_FLAGS) print(f"swiglu_forward_and_per_token_cast m={m} n={n}") report("swiglu", tuple(got) if isinstance(got, tuple) else got, tuple(want) if isinstance(want, tuple) else want, exact=True) def op_rotary() -> None: # The kernel takes (seqlen, num_heads, rot_dim); rot_dim is 64 or 128 only. seqlen, heads, rot_dim = 64, 8, 128 q = torch.randn((seqlen, heads, rot_dim), device=DEV, dtype=torch.bfloat16) k = torch.randn((seqlen, heads, rot_dim), device=DEV, dtype=torch.bfloat16) positions = torch.arange(seqlen, device=DEV, dtype=torch.int64) angles = torch.rand((seqlen, rot_dim // 2), device=DEV, dtype=torch.float32) * 6.28 cos = angles.cos() sin = angles.sin() # The kernel reads a (max_pos, rot_dim) cache of cos|sin halves. cache = torch.cat([cos, sin], dim=-1).contiguous() # The kernel rotates in place and returns None, so it is applied to a copy. got_q = q.clone() got_k = k.clone() tk_tr.apply_rotary(got_q, cache, got_k, positions, interleaved=False) want_q = ref.rotary_embedding_ref(q, cos, sin, positions, interleaved=False) want_k = ref.rotary_embedding_ref(k, cos, sin, positions, interleaved=False) print(f"apply_rotary seqlen={seqlen} heads={heads} rot_dim={rot_dim}") # Both are bf16, but the two implementations may round a tie in the last bit differently. # Measure the disagreement in bf16 representable steps: at most 1 ULP is a rounding tie, and # 8 bits of mantissa means most elements have far fewer than 3 significant decimals. for label, got, want in (("q", got_q, want_q), ("k", got_k, want_k)): a = got.float() b = want.float().to(torch.bfloat16).float() # Relative to the value's own magnitude: |a-b| / ulp(b), where ulp is 2^-8 of the magnitude. ulp = b.abs() * 2.0**-8 steps = ((a - b).abs() / ulp.clamp_min(torch.finfo(torch.float32).tiny)).nan_to_num(0) worst = steps.max().item() # One bf16 last-bit disagreement on a tie is rounding, not a logic difference. Allow up to # two steps and demand it stays rare, the same rule the kernel A/B harness uses for # outputs whose answer sits on a rounding boundary. over = int((steps > 2.0).sum()) frac = float((steps > 1.0).float().mean()) ok = over == 0 and frac <= 1e-3 print(f" {'ok ' if ok else 'FAIL'} {label}: max {worst:.2f} ULP, {frac * 100:.4f}% beyond 1 ULP") if not ok: FAILED.append(label) def op_engram_hash() -> None: # vocab_sizes is (num_ngram_layers, max_ngram_size - 1, num_embed_table_per_ngram). layers, ngram, tables, vocab = 2, 4, 3, 1000 # Each of these has its own dtype: ids are int32, and the pool arithmetic is int64. ids = torch.randint(0, vocab, (256, ngram), device=DEV, dtype=torch.int32) # dtypes come from the kernel's docstring: multipliers is (layers, max_ngram_size) int64, # the table metadata is int32. multipliers = torch.randint(1, 2**31, (layers, ngram), device=DEV, dtype=torch.int64) vocab_sizes = torch.full((layers, ngram - 1, tables), vocab, device=DEV, dtype=torch.int32) offsets = torch.zeros((layers, (ngram - 1) * tables), device=DEV, dtype=torch.int32) mask = torch.zeros((256,), device=DEV, dtype=torch.bool) got = tk_engram.engram_hash(ids, multipliers, vocab_sizes, offsets, mask) want = ref.engram_hash_ref(ids, multipliers, vocab_sizes, offsets, mask) print("engram_hash") report("hash", got, want, exact=True) OPS = { "per_token_cast": op_per_token_cast, "cast_back": op_cast_back, "swiglu": op_swiglu, "rotary": op_rotary, "engram_hash": op_engram_hash, } def main() -> int: selected = sys.argv[1:] or list(OPS) unknown = [n for n in selected if n not in OPS] if unknown: print(f"unknown op(s): {unknown}; have {list(OPS)}", file=sys.stderr) return 2 for name in selected: print(f"[{name}]") try: OPS[name]() except Exception: print(" ERROR") traceback.print_exc() FAILED.append(name) print() if FAILED: print(f"MISMATCH {len(FAILED)}: {sorted(set(FAILED))}") return 1 print("ALL_MATCH") return 0 if __name__ == "__main__": sys.exit(main())