File size: 6,687 Bytes
6e06900 fd33e97 6e06900 fd33e97 6e06900 | 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 | """Every path the forward launcher can select, against the sequential reference.
launch_impl picks among a generic, a tensor-core and a row-block streamed kernel
for chunk_state and for chunk_scan, on gates over dtype, C*R, the mimo rank, the
head dimension and the state dimension. The tensor-core kernels take the rank
and the chunk length from template arguments and discard the runtime values, so
a geometry that reaches the wrong instantiation does not fall back to a slower
kernel, it contracts with the wrong rank and returns a plausible tensor.
Each row therefore does two things: it asks the extension which path the
geometry takes, through the same gates the launcher selects with, and it checks
the result against mimo_step_ref. The grid is then required to cover every
combination the launcher can produce, so a path cannot stop being exercised
without the run failing.
run.cmd tests/test_dispatch.py
"""
import sys
from pathlib import Path
import torch
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "ref"))
sys.path.insert(0, str(ROOT / "tests"))
import load_local as mamba3 # noqa: E402
from anchor import max_rel_err # noqa: E402
from mamba3_ref import mimo_step_ref # noqa: E402
from test_mamba3 import build, run_fwd # noqa: E402
DEV = "cuda"
GEN, TC, STR = "generic", "tensor-core", "streamed"
# label, (B, S, H, G, P, N, R, Na, C), dtype, expected (state, scan)
GRID = [
# float32 never reaches the bf16 state kernel, but its scan contracts on
# three-term tf32 whenever C*R == 64 and the rank and head dim are in the set
("f32 baseline", (1, 64, 4, 1, 64, 128, 4, 64, 16), torch.float32, (GEN, TC)),
("f32 grouped query", (1, 64, 8, 2, 32, 64, 4, 32, 16), torch.float32, (GEN, TC)),
("f32 ragged S", (1, 70, 4, 1, 64, 128, 4, 64, 16), torch.float32, (GEN, TC)),
("f32 headdim 128", (1, 64, 4, 1, 128, 128, 4, 64, 16), torch.float32, (GEN, TC)),
# both generic: C*R != 64, or a head dim outside the instantiated set
("f32 C*R=12, no tf32", (1, 48, 4, 1, 64, 128, 1, 64, 12), torch.float32, (GEN, GEN)),
("f32 headdim 48", (1, 64, 4, 1, 48, 128, 4, 64, 16), torch.float32, (GEN, GEN)),
# bf16 with a head dimension outside the instantiated set
("bf16 headdim 48", (1, 64, 4, 1, 48, 128, 4, 64, 16), torch.bfloat16, (GEN, GEN)),
# N a multiple of 8 but not of 16: the scan gate holds, the state gate does not
("bf16 N=72", (1, 64, 4, 1, 64, 72, 4, 36, 16), torch.bfloat16, (GEN, TC)),
# C*R a multiple of 16 but not 64: state on tensor cores, scan generic
("bf16 C*R=32", (1, 64, 4, 1, 64, 128, 2, 64, 16), torch.bfloat16, (TC, GEN)),
("bf16 C*R=16", (1, 64, 4, 1, 64, 128, 1, 64, 16), torch.bfloat16, (TC, GEN)),
# C*R = 64, the whole instantiated rank set and every head dimension
("bf16 R=1 C=64", (1, 64, 4, 1, 64, 128, 1, 64, 64), torch.bfloat16, (TC, TC)),
("bf16 R=2 C=32", (1, 64, 4, 1, 64, 128, 2, 64, 32), torch.bfloat16, (TC, TC)),
("bf16 R=4 C=16", (1, 64, 4, 1, 64, 128, 4, 64, 16), torch.bfloat16, (TC, TC)),
("bf16 R=8 C=8", (1, 64, 4, 1, 64, 128, 8, 64, 8), torch.bfloat16, (TC, TC)),
("bf16 headdim 16", (1, 64, 4, 1, 16, 128, 4, 64, 16), torch.bfloat16, (TC, TC)),
("bf16 headdim 32", (1, 64, 4, 1, 32, 128, 4, 64, 16), torch.bfloat16, (TC, TC)),
("bf16 headdim 128", (1, 64, 4, 1, 128, 128, 4, 64, 16), torch.bfloat16, (TC, TC)),
# above C*R = 64 both stages stream, in either dtype
("f32 streamed", (1, 128, 4, 1, 64, 128, 4, 64, 32), torch.float32, (STR, STR)),
("bf16 streamed", (1, 128, 4, 1, 64, 128, 4, 64, 32), torch.bfloat16, (STR, STR)),
("bf16 streamed C*R=256", (1, 128, 4, 1, 64, 128, 8, 64, 32), torch.bfloat16, (STR, STR)),
]
# Ranks the tensor-core scan has no instantiation for. C*R = 64 satisfies the
# rest of its gate, so before the rank predicate was added these dispatched the
# rank-8 kernel and returned garbage; they must be refused, not run.
REFUSED = [(16, 4), (32, 2), (64, 1)]
fails = []
covered = set()
tf32_seen = set()
def check(label, ok, detail=""):
print(f" {label:<26} {'ok' if ok else 'FAIL':<5} {detail}")
if not ok:
fails.append(label)
print(f"torch {torch.__version__} {torch.cuda.get_device_name(0)}\n")
print("1. each dispatched path against mimo_step_ref")
print(f" {'case':<26} {'state':<12} {'scan':<12} {'tf32':<5} {'rel':>10}")
for label, geom, dtype, want in GRID:
B, S, H, G, P, N, R, Na, C = geom
is_bf = dtype is torch.bfloat16
d = mamba3.dispatch_paths(R, N, P, C, is_bf)
state, scan, tf32 = d["state"], d["scan"], d["state_tf32"]
covered.add((state, scan))
if state == GEN:
tf32_seen.add(tf32)
c32 = build(B, S, H, G, P, N, R, Na, DEV, seed=0)
cin = build(B, S, H, G, P, N, R, Na, DEV, seed=0, dtype=dtype) if is_bf else c32
ref, _ = mimo_step_ref(c32["q"], c32["k"], c32["v"], c32["adt"], c32["dt"],
c32["trap"], c32["q_bias"], c32["k_bias"], c32["angles"],
c32["mimo_v"], c32["mimo_o"], D=c32["D"], z=c32["z"],
mimo_z=c32["mimo_z"])
rel = max_rel_err(run_fwd(cin, C), ref)
tol = 3e-2 if is_bf else 5e-5
print(f" {label:<26} {state:<12} {scan:<12} {str(tf32):<5} {rel:>10.3e}")
if (state, scan) != want:
check(label, False, f"took {(state, scan)}, grid says {want}")
elif rel >= tol:
check(label, False, f"rel {rel:.3e} against tol {tol:.0e}")
print("\n2. ranks with no tensor-core instantiation are refused, not dispatched")
for R, C in REFUSED:
B, S, H, G, P, N, Na = 1, 64, 4, 1, 64, 128, 64
c = build(B, S, H, G, P, N, R, Na, DEV, seed=0, dtype=torch.bfloat16)
try:
run_fwd(c, C)
torch.cuda.synchronize()
check(f"R={R} C={C} C*R={C*R}", False, "ran instead of raising")
except RuntimeError as exc:
msg = str(exc)
check(f"R={R} C={C} C*R={C*R}", "mimo_rank" in msg, msg.splitlines()[0][:56])
print("\n3. the grid covers every path the launcher can select")
REACHABLE = {(GEN, GEN), (GEN, TC), (TC, GEN), (TC, TC), (STR, STR)}
missing = REACHABLE - covered
check("every (state, scan) pair exercised", not missing,
"all five" if not missing else f"missing {sorted(missing)}")
check("generic state seen with and without tf32", tf32_seen == {False, True},
f"tf32 flags {sorted(tf32_seen)}")
print("\n" + ("PASS" if not fails else f"FAIL ({len(fails)}): " + "; ".join(fails)))
sys.exit(0 if not fails else 1)
|