| """Head-to-head against upstream's TileLang Mamba-3 MIMO forward. |
| |
| The two kernels take different dtype contracts at reduced precision. Upstream |
| keeps the biases, the rank weights and D in float32 and casts trap to the value |
| dtype; this kernel does the reverse, carrying the biases and rank weights in the |
| value dtype and the whole schedule in float32. So one float32 master set is |
| built and each side is handed the cast it requires, which keeps the underlying |
| values identical. |
| |
| In float32 the two contracts coincide, so that mode doubles as a correctness |
| anchor. Both paths run in one process on one device, and TileLang autotunes per |
| shape, so everything is warmed before it is timed. |
| |
| PYTHONPATH=~/k/mamba-upstream ~/m3bench/bin/python bench_vs_upstream.py |
| """ |
|
|
| import argparse |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| M3 = Path.home() / "k" / "mamba3" |
| UP = Path.home() / "k" / "mamba-upstream" |
| sys.path.insert(0, str(M3)) |
| sys.path.insert(0, str(UP)) |
|
|
| import load_local as mamba3 |
| from mamba_ssm.ops.tilelang.mamba3.mamba3_mimo import mamba3_mimo |
|
|
| |
| OURS_CAST = ["q", "k", "v", "z", "q_bias", "k_bias", "mimo_v", "mimo_o", "mimo_z", "D"] |
| |
| THEIRS_CAST = ["q", "k", "v", "z", "trap"] |
|
|
|
|
| def master(B, S, H, G, P, N, R, Na, dev, seed): |
| """One float32 set of values; each side casts what it needs.""" |
| torch.manual_seed(seed) |
| f = lambda *s: torch.randn(*s, device=dev) |
| dt = F.softplus(-3.0 + f(B, H, S)) |
| return { |
| "q": f(B, S, R, G, N), "k": f(B, S, R, G, N), |
| "v": f(B, S, H, P), "z": f(B, S, H, P), |
| "q_bias": f(H, R, N), "k_bias": f(H, R, N), |
| "mimo_v": torch.rand(H, R, P, device=dev) / R, |
| "mimo_o": torch.rand(H, R, P, device=dev) / R, |
| "mimo_z": torch.rand(H, R, P, device=dev) / R, |
| "D": f(H), "angles": torch.rand(B, S, H, Na, device=dev), |
| "dt": dt, "adt": -F.softplus(f(B, H, S)).clamp(max=-1e-4) * dt, |
| "trap": torch.rand(B, H, S, device=dev) * 0.5, |
| } |
|
|
|
|
| def cast(c, names, dtype): |
| if dtype is torch.float32: |
| return dict(c) |
| return {k: (t.to(dtype).contiguous() if k in names else t) for k, t in c.items()} |
|
|
|
|
| def to_split_half(c, N): |
| """Reindex q, k and their biases from pairwise rotary layout to split-half. |
| |
| This kernel rotates the pair (2i, 2i+1); the TileLang kernel rotates |
| (n, N/2 + n). The two are a permutation of the head dimension, and it is |
| applied to q and k alike, so every inner product is preserved and the two |
| compute the same operator on the same values. |
| """ |
| idx = torch.cat([torch.arange(0, N, 2), torch.arange(1, N, 2)]).to(c["q"].device) |
| out = dict(c) |
| for k in ("q", "k", "q_bias", "k_bias"): |
| out[k] = c[k].index_select(-1, idx).contiguous() |
| return out |
|
|
|
|
| def ours(c, C): |
| return mamba3.forward( |
| c["q"], c["k"], c["v"], c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_o"], |
| c["angles"], c["adt"], c["dt"], c["trap"], z=c["z"], mimo_z=c["mimo_z"], |
| D=c["D"], chunk_size=C) |
|
|
|
|
| def theirs(c, C, dtype): |
| return mamba3_mimo( |
| c["q"], c["k"], c["v"], c["adt"], c["dt"], c["trap"], |
| c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_z"], c["mimo_o"], |
| c["angles"], c["D"], c["z"], |
| chunk_size=C, rotary_dim_divisor=2, dtype=dtype) |
|
|
|
|
| def timeit(fn, iters, warmup=5): |
| for _ in range(warmup): |
| fn() |
| torch.cuda.synchronize() |
| t0 = time.perf_counter() |
| for _ in range(iters): |
| fn() |
| torch.cuda.synchronize() |
| return (time.perf_counter() - t0) / iters * 1e3 |
|
|
|
|
| def rel(a, b): |
| a, b = a.float(), b.float() |
| return float((a - b).abs().max() / b.abs().max().clamp_min(1e-5)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float32"]) |
| ap.add_argument("--iters", type=int, default=20) |
| args = ap.parse_args() |
| dtype = getattr(torch, args.dtype) |
| dev = "cuda" |
|
|
| cap = "".join(map(str, torch.cuda.get_device_capability())) |
| print(f"torch {torch.__version__} {torch.cuda.get_device_name(0)} " |
| f"sm_{cap} dtype={args.dtype}\n") |
| print(f"{'S':>6} {'H':>4} {'P':>4} {'N':>4} {'R':>3} {'C':>4} " |
| f"{'ours ms':>10} {'tilelang ms':>12} {'ratio':>8} {'max rel':>9}") |
|
|
| B, G = 1, 1 |
| grid = [(512, 8, 64, 128, 4, 16), (1024, 8, 64, 128, 4, 16), |
| (2048, 8, 64, 128, 4, 16), (4096, 8, 64, 128, 4, 16), |
| (2048, 32, 64, 128, 4, 16), (2048, 8, 64, 128, 1, 64)] |
|
|
| for S, H, P, N, R, C in grid: |
| Na = N // 2 |
| m = master(B, S, H, G, P, N, R, Na, dev, seed=0) |
| co = cast(m, OURS_CAST, dtype) |
| ct = cast(to_split_half(m, N), THEIRS_CAST, dtype) |
| try: |
| yo, yt = ours(co, C), theirs(ct, C, dtype) |
| r = rel(yo, yt) |
| to = timeit(lambda: ours(co, C), args.iters) |
| tt = timeit(lambda: theirs(ct, C, dtype), args.iters) |
| print(f"{S:6d} {H:4d} {P:4d} {N:4d} {R:3d} {C:4d} " |
| f"{to:10.3f} {tt:12.3f} {tt / to:7.2f}x {r:9.2e}") |
| except Exception as exc: |
| msg = f"{type(exc).__name__}: {exc}".replace("\n", " ") |
| print(f"{S:6d} {H:4d} {P:4d} {N:4d} {R:3d} {C:4d} {msg[:88]}") |
| finally: |
| del m, co, ct |
| torch.cuda.empty_cache() |
|
|
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
