DeepInfer/scripts/tile_kernels_patch/compare_with_cuda.py

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