|
| 1 | +"""GB10/sm_121 measurement + parity harness for the FUSED SSD forward (§23). |
| 2 | +
|
| 3 | +Builds the NEW fused ``ssd_fused_fwd_cuda_prim`` (F0+F1+F2 in ONE smem-resident |
| 4 | +CUDA kernel) at the prod cfg (S=4096 c=64 g=8 H=112 P=64 N=64 bs1) and: |
| 5 | +
|
| 6 | + 1. reports whether it LAUNCHES at N=64 (the smem-fit question). |
| 7 | + 2. times it (warm median + min over N timed dispatches, device-resident). |
| 8 | + 3. checks PARITY of the fused output + final_state vs: |
| 9 | + (a) the un-fused F0->F1->F2 CUDA chain (the real production reference), |
| 10 | + (b) the SERIAL per-timestep reference (the contract anchor), |
| 11 | + over ALL elements (max|abs|), gate fp16 < 5e-4. |
| 12 | + 4. also builds + times the un-fused F0/F1/F2 chain at the SAME cfg so we have |
| 13 | + the apples-to-apples comparison (fused ms vs 16.37+1.24+F2 ms). |
| 14 | +
|
| 15 | +RULE #1: every kernel goes through its ONE builder; a compile/run/parity failure |
| 16 | +is SURFACED (traceback printed, returns False) — never swallowed, never falls |
| 17 | +back to a serial path silently. |
| 18 | +
|
| 19 | +Run on gb10 (same env as scratch/probe_chunked_scan_cuda_gb10.py): |
| 20 | + PYTHONPATH=/home/dave/source/cppmega_mlx:/home/dave/source/tilelang/3rdparty/tvm/python:/home/dave/source/tilelang/3rdparty/tvm/3rdparty/tvm-ffi/python \ |
| 21 | + TVM_LIBRARY_PATH=/home/dave/source/tilelang/build/lib \ |
| 22 | + /home/dave/cppmega-venv/bin/python scratch/probe_ssd_fused_fwd_gb10.py [--prod] |
| 23 | +""" |
| 24 | + |
| 25 | +import sys |
| 26 | +import time |
| 27 | +import traceback |
| 28 | + |
| 29 | +# RULE #1: fix the NVRTC builtins loader BEFORE torch/tilelang import (gb10 |
| 30 | +# sm_121 CUDA codegen needs the nvrtc builtins path for exp2/etc.). No-op + |
| 31 | +# byte-identical to the fp8 bench's ordering. |
| 32 | +try: |
| 33 | + from cppmega_mlx._gb10_nvrtc_env import ensure_nvrtc_builtins_path |
| 34 | + |
| 35 | + ensure_nvrtc_builtins_path() |
| 36 | +except Exception as _e: # surfaced, not swallowed — print so it's visible |
| 37 | + print(f"[nvrtc-env] ensure_nvrtc_builtins_path unavailable: {_e!r}") |
| 38 | + |
| 39 | +import torch |
| 40 | +from einops import rearrange, repeat |
| 41 | + |
| 42 | +DEV = "cuda" |
| 43 | + |
| 44 | + |
| 45 | +# --------------------------------------------------------------------------- # |
| 46 | +# References (identical algebra to tests/test_mamba3_chunk_scan_combine_f2 and # |
| 47 | +# scratch/probe_chunked_scan_cuda_gb10.py — the SSD numerical ground truth). # |
| 48 | +# --------------------------------------------------------------------------- # |
| 49 | +def serial_full_forward(C, Bmat, x, A, dt, h0, D): |
| 50 | + """SERIAL per-timestep diagonal forward over the FULL sequence (fp32). |
| 51 | +
|
| 52 | + h[t] = exp(A[h]*dt[t]) * h[t-1] + dt[t] * (x[t] outer B[t]) |
| 53 | + y[t] = sum_n h[t]*C[t] + D*x[t] (seeded by h0). Returns (out, h_last). |
| 54 | + """ |
| 55 | + batch, seqlen, ngroups, dstate = C.shape |
| 56 | + _, _, nheads, headdim = x.shape |
| 57 | + h = nheads // ngroups |
| 58 | + Cf = repeat(C.float(), "b l g n -> b l (g h) n", h=h) |
| 59 | + Bf = repeat(Bmat.float(), "b l g n -> b l (g h) n", h=h) |
| 60 | + xf = x.float() |
| 61 | + Af = A.float() |
| 62 | + dtf = dt.float() |
| 63 | + state = h0.float().clone() |
| 64 | + out = torch.zeros(batch, seqlen, nheads, headdim, device=x.device) |
| 65 | + for t in range(seqlen): |
| 66 | + decay = torch.exp(Af.view(1, nheads) * dtf[:, t]) |
| 67 | + inp = dtf[:, t][:, :, None, None] * ( |
| 68 | + xf[:, t][:, :, :, None] * Bf[:, t][:, :, None, :] |
| 69 | + ) |
| 70 | + state = decay[:, :, None, None] * state + inp |
| 71 | + y = torch.einsum("bhpn,bhn->bhp", state, Cf[:, t]) |
| 72 | + out[:, t] = y + D.float().view(1, nheads, 1) * xf[:, t] |
| 73 | + return out, state |
| 74 | + |
| 75 | + |
| 76 | +def eager_precompute(C, Bmat, x, A, dt, h0, chunk_size): |
| 77 | + """F0+F1 reference -> (cb fp16, dA_cumsum fp16, prev_states fp16).""" |
| 78 | + batch, seqlen, ngroups, dstate = C.shape |
| 79 | + _, _, nheads, headdim = x.shape |
| 80 | + nchunks = seqlen // chunk_size |
| 81 | + Cf, Bf, xf, Af, dtf = C.float(), Bmat.float(), x.float(), A.float(), dt.float() |
| 82 | + a = Af.view(1, 1, nheads) * dtf |
| 83 | + a_c = rearrange(a, "b (c l) h -> b h c l", c=nchunks) |
| 84 | + dA_cumsum = torch.cumsum(a_c, dim=-1) |
| 85 | + Cc = rearrange(Cf, "b (c l) g n -> b c l g n", c=nchunks) |
| 86 | + Bc = rearrange(Bf, "b (c s) g n -> b c s g n", c=nchunks) |
| 87 | + cb = torch.einsum("bclgn,bcsgn->bcgls", Cc, Bc) |
| 88 | + h = nheads // ngroups |
| 89 | + Bexp = rearrange(Bf, "b (c s) g n -> b c s g n", c=nchunks).repeat_interleave(h, dim=3) |
| 90 | + xexp = rearrange(xf, "b (c s) hh p -> b c s hh p", c=nchunks) |
| 91 | + dtc = rearrange(dtf, "b (c s) hh -> b hh c s", c=nchunks) |
| 92 | + decay_states = torch.exp(dA_cumsum[:, :, :, -1:] - dA_cumsum) |
| 93 | + states = torch.einsum("bhcs,bhcs,bcshp,bcshn->bchpn", decay_states, dtc, xexp, Bexp) |
| 94 | + chunk_tail = dA_cumsum[:, :, :, -1] |
| 95 | + init = h0.float().unsqueeze(1) |
| 96 | + states_cat = torch.cat([init, states], dim=1) |
| 97 | + pad = torch.nn.functional.pad(chunk_tail, (1, 0)) |
| 98 | + cc = pad.shape[-1] |
| 99 | + csum = torch.cumsum(pad, dim=-1) |
| 100 | + seg = csum[:, :, :, None] - csum[:, :, None, :] |
| 101 | + mask = torch.tril(torch.ones(cc, cc, device=pad.device, dtype=torch.bool)) |
| 102 | + seg = seg.masked_fill(~mask, float("-inf")) |
| 103 | + decay_chunk = torch.exp(seg) |
| 104 | + new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states_cat) |
| 105 | + prev_states = new_states[:, :-1] |
| 106 | + return ( |
| 107 | + cb.half().contiguous(), |
| 108 | + dA_cumsum.half().contiguous(), |
| 109 | + prev_states.half().contiguous(), |
| 110 | + ) |
| 111 | + |
| 112 | + |
| 113 | +def _time(fn, n=20): |
| 114 | + for _ in range(3): |
| 115 | + fn() |
| 116 | + torch.cuda.synchronize() |
| 117 | + ts = [] |
| 118 | + for _ in range(n): |
| 119 | + torch.cuda.synchronize() |
| 120 | + t0 = time.perf_counter() |
| 121 | + fn() |
| 122 | + torch.cuda.synchronize() |
| 123 | + ts.append((time.perf_counter() - t0) * 1e3) |
| 124 | + ts.sort() |
| 125 | + return ts[len(ts) // 2], ts[0] |
| 126 | + |
| 127 | + |
| 128 | +def run_cfg(batch, seqlen, chunk, ngroups, nheads, headdim, dstate): |
| 129 | + from cppmega_mlx.nn._tilelang.mamba3_ssd_fused_fwd import ( |
| 130 | + build_ssd_fused_fwd, |
| 131 | + ssd_fused_fwd_grid, |
| 132 | + ssd_fused_fwd_smem_budget_bytes, |
| 133 | + ) |
| 134 | + from cppmega_mlx.nn._tilelang.mamba3_chunked_precompute_core import ( |
| 135 | + build_chunk_precompute_metal, |
| 136 | + build_inter_chunk_recur_metal, |
| 137 | + ) |
| 138 | + from cppmega_mlx.nn._tilelang.mamba3_chunked_scan_core import ( |
| 139 | + build_chunk_scan_combine_metal, |
| 140 | + ) |
| 141 | + |
| 142 | + nchunks = seqlen // chunk |
| 143 | + tg, grid = ssd_fused_fwd_grid(batch, seqlen, chunk, ngroups, nheads, headdim, dstate) |
| 144 | + budget = ssd_fused_fwd_smem_budget_bytes(chunk, headdim, dstate) |
| 145 | + print(f"\n=== cfg b={batch} S={seqlen} c={chunk} g={ngroups} H={nheads} " |
| 146 | + f"P={headdim} N={dstate} ===") |
| 147 | + print(f"[FUSED] grid={grid} tg={tg} smem_budget={budget['total']} B " |
| 148 | + f"({budget['total']/1024:.2f} KiB) cap=101376 B") |
| 149 | + |
| 150 | + torch.manual_seed(0) |
| 151 | + C = (torch.randn(batch, seqlen, ngroups, dstate, device=DEV) * 0.1).half() |
| 152 | + Bmat = (torch.randn(batch, seqlen, ngroups, dstate, device=DEV) * 0.1).half() |
| 153 | + x = (torch.randn(batch, seqlen, nheads, headdim, device=DEV) * 0.1).half() |
| 154 | + A = -torch.rand(nheads, device=DEV).half() |
| 155 | + dt = (torch.rand(batch, seqlen, nheads, device=DEV) * 0.05).half() |
| 156 | + D = torch.randn(nheads, device=DEV).half() |
| 157 | + h0 = (torch.randn(batch, nheads, headdim, dstate, device=DEV) * 0.1).float() |
| 158 | + |
| 159 | + # ---- SERIAL ground truth (CPU, fp32) ---- |
| 160 | + out_serial, hlast_serial = serial_full_forward( |
| 161 | + C.cpu(), Bmat.cpu(), x.cpu(), A.cpu(), dt.cpu(), h0.cpu(), D.cpu()) |
| 162 | + |
| 163 | + # =================== FUSED kernel =================== |
| 164 | + try: |
| 165 | + kf = build_ssd_fused_fwd( |
| 166 | + batch, seqlen, chunk, ngroups, nheads, headdim, dstate, target="cuda") |
| 167 | + except Exception: |
| 168 | + print("[FUSED] COMPILE FAILED:") |
| 169 | + traceback.print_exc() |
| 170 | + return False |
| 171 | + print("[FUSED] COMPILE ok") |
| 172 | + |
| 173 | + out_f = torch.zeros(batch, seqlen, nheads, headdim, device=DEV, dtype=torch.float16) |
| 174 | + fst_f = torch.zeros(batch, nheads, headdim, dstate, device=DEV, dtype=torch.float32) |
| 175 | + fused_args = ( |
| 176 | + x.contiguous(), Bmat.contiguous(), C.contiguous(), A.contiguous(), |
| 177 | + dt.contiguous(), D.contiguous(), h0.contiguous(), |
| 178 | + ) |
| 179 | + try: |
| 180 | + kf(*fused_args, out_f, fst_f) |
| 181 | + torch.cuda.synchronize() |
| 182 | + except Exception: |
| 183 | + print("[FUSED] RUN FAILED (LAUNCH at N=%d):" % dstate) |
| 184 | + traceback.print_exc() |
| 185 | + return False |
| 186 | + print("[FUSED] RUN ok (LAUNCHES at N=%d)" % dstate) |
| 187 | + |
| 188 | + f_med, f_min = _time(lambda: kf(*fused_args, out_f, fst_f)) |
| 189 | + print(f"[FUSED] TIMING median={f_med:.3f} ms/call min={f_min:.3f} ms") |
| 190 | + |
| 191 | + nan_f = bool(torch.isnan(out_f).any()) or bool(torch.isnan(fst_f).any()) |
| 192 | + fs_out_serial = float((out_f.float().cpu() - out_serial).abs().max()) |
| 193 | + fs_hlast_serial = float((fst_f.float().cpu() - hlast_serial).abs().max()) |
| 194 | + print(f"[FUSED vs SERIAL] out max|abs|={fs_out_serial:.3e} " |
| 195 | + f"final_state max|abs|={fs_hlast_serial:.3e} NaN={nan_f}") |
| 196 | + |
| 197 | + # =================== UN-FUSED F0 -> F1 -> F2 chain =================== |
| 198 | + try: |
| 199 | + k0 = build_chunk_precompute_metal( |
| 200 | + batch, seqlen, chunk, ngroups, nheads, headdim, dstate, target="cuda") |
| 201 | + k1 = build_inter_chunk_recur_metal( |
| 202 | + batch, seqlen, chunk, ngroups, nheads, headdim, dstate, target="cuda") |
| 203 | + k2 = build_chunk_scan_combine_metal( |
| 204 | + batch, seqlen, chunk, ngroups, nheads, headdim, dstate, target="cuda") |
| 205 | + except Exception: |
| 206 | + print("[UNFUSED] COMPILE FAILED:") |
| 207 | + traceback.print_exc() |
| 208 | + return False |
| 209 | + print("[UNFUSED] F0/F1/F2 COMPILE ok") |
| 210 | + |
| 211 | + cb = torch.zeros(batch, nchunks, ngroups, chunk, chunk, device=DEV, dtype=torch.float16) |
| 212 | + dA_cumsum = torch.zeros(batch, nheads, nchunks, chunk, device=DEV, dtype=torch.float16) |
| 213 | + summary_states = torch.zeros(batch, nchunks, nheads, headdim, dstate, device=DEV, dtype=torch.float32) |
| 214 | + prev_states = torch.zeros(batch, nchunks, nheads, headdim, dstate, device=DEV, dtype=torch.float32) |
| 215 | + final_state_u = torch.zeros(batch, nheads, headdim, dstate, device=DEV, dtype=torch.float32) |
| 216 | + out_u = torch.zeros(batch, seqlen, nheads, headdim, device=DEV, dtype=torch.float16) |
| 217 | + dt_k = rearrange(dt, "b (c s) hh -> b hh c s", c=nchunks).contiguous() |
| 218 | + |
| 219 | + def _run_f0(): |
| 220 | + k0(x.contiguous(), Bmat.contiguous(), C.contiguous(), A.contiguous(), |
| 221 | + dt.contiguous(), cb, dA_cumsum, summary_states) |
| 222 | + |
| 223 | + def _run_f1(): |
| 224 | + k1(summary_states.contiguous(), dA_cumsum.contiguous(), h0.contiguous(), |
| 225 | + prev_states, final_state_u) |
| 226 | + |
| 227 | + def _run_f2(): |
| 228 | + k2(cb.contiguous(), x.contiguous(), dt_k.contiguous(), dA_cumsum.contiguous(), |
| 229 | + C.contiguous(), prev_states.contiguous(), D.contiguous(), out_u) |
| 230 | + |
| 231 | + try: |
| 232 | + _run_f0(); torch.cuda.synchronize() |
| 233 | + _run_f1(); torch.cuda.synchronize() |
| 234 | + _run_f2(); torch.cuda.synchronize() |
| 235 | + except Exception: |
| 236 | + print("[UNFUSED] RUN FAILED:") |
| 237 | + traceback.print_exc() |
| 238 | + return False |
| 239 | + print("[UNFUSED] F0/F1/F2 RUN ok (F2 LAUNCHES at N=%d)" % dstate) |
| 240 | + |
| 241 | + t0, t0min = _time(_run_f0) |
| 242 | + t1, t1min = _time(_run_f1) |
| 243 | + # F2 must be re-fed prev_states each time, but prev_states is stable post-F1. |
| 244 | + t2, t2min = _time(_run_f2) |
| 245 | + unfused_sum = t0 + t1 + t2 |
| 246 | + print(f"[UNFUSED] F0={t0:.3f} F1={t1:.3f} F2={t2:.3f} " |
| 247 | + f"SUM={unfused_sum:.3f} ms/call (medians)") |
| 248 | + |
| 249 | + # parity of UN-FUSED chain vs serial (so we know the reference itself is sound) |
| 250 | + uo_serial = float((out_u.float().cpu() - out_serial).abs().max()) |
| 251 | + uh_serial = float((final_state_u.float().cpu() - hlast_serial).abs().max()) |
| 252 | + print(f"[UNFUSED vs SERIAL] out max|abs|={uo_serial:.3e} " |
| 253 | + f"final_state max|abs|={uh_serial:.3e}") |
| 254 | + |
| 255 | + # parity of FUSED vs UN-FUSED (the production reference) over ALL elements |
| 256 | + fu_out = float((out_f.float() - out_u.float()).abs().max()) |
| 257 | + fu_fst = float((fst_f - final_state_u).abs().max()) |
| 258 | + print(f"[FUSED vs UNFUSED] out max|abs|={fu_out:.3e} " |
| 259 | + f"final_state max|abs|={fu_fst:.3e}") |
| 260 | + |
| 261 | + # ===== verdicts ===== |
| 262 | + gate = 5e-4 |
| 263 | + parity_ok = ( |
| 264 | + fs_out_serial < gate and fs_hlast_serial < gate |
| 265 | + and fu_out < gate and fu_fst < gate and not nan_f |
| 266 | + ) |
| 267 | + print(f"\n[SUMMARY] FUSED median={f_med:.3f} ms UNFUSED sum={unfused_sum:.3f} ms " |
| 268 | + f"speedup={unfused_sum/f_med:.2f}x cppmega_fused_ref=3.110 ms " |
| 269 | + f"(fused/cppmega={f_med/3.110:.2f}x)") |
| 270 | + print(f"[SUMMARY] parity_all_elements={'PASS' if parity_ok else 'FAIL'} " |
| 271 | + f"(gate fp16<{gate})") |
| 272 | + print(f"[CSV] cfg=S{seqlen}c{chunk}g{ngroups}H{nheads}P{headdim}N{dstate} " |
| 273 | + f"fused_ms={f_med:.4f} fused_min={f_min:.4f} " |
| 274 | + f"f0={t0:.4f} f1={t1:.4f} f2={t2:.4f} unfused_sum={unfused_sum:.4f} " |
| 275 | + f"fused_vs_unfused={unfused_sum/f_med:.4f}x " |
| 276 | + f"fused_vs_cppmega={f_med/3.110:.4f}x " |
| 277 | + f"fs_out={fs_out_serial:.3e} fs_fst={fs_hlast_serial:.3e} " |
| 278 | + f"fu_out={fu_out:.3e} fu_fst={fu_fst:.3e} parity={'PASS' if parity_ok else 'FAIL'}") |
| 279 | + return parity_ok |
| 280 | + |
| 281 | + |
| 282 | +def main(): |
| 283 | + prod = "--prod" in sys.argv |
| 284 | + print("=== FUSED SSD forward probe (gb10 sm_121) ===") |
| 285 | + print("torch", torch.__version__, "cuda_avail", torch.cuda.is_available(), |
| 286 | + "dev", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "NONE") |
| 287 | + if not torch.cuda.is_available(): |
| 288 | + print("NO CUDA DEVICE") |
| 289 | + sys.exit(2) |
| 290 | + if prod: |
| 291 | + cfgs = [dict(batch=1, seqlen=4096, chunk=64, ngroups=8, nheads=112, |
| 292 | + headdim=64, dstate=64)] |
| 293 | + else: |
| 294 | + cfgs = [ |
| 295 | + dict(batch=1, seqlen=256, chunk=64, ngroups=1, nheads=2, headdim=64, dstate=16), |
| 296 | + dict(batch=1, seqlen=512, chunk=64, ngroups=8, nheads=8, headdim=64, dstate=64), |
| 297 | + ] |
| 298 | + ok = True |
| 299 | + for cfg in cfgs: |
| 300 | + try: |
| 301 | + ok = run_cfg(**cfg) and ok |
| 302 | + except Exception: |
| 303 | + print("CFG CRASHED:") |
| 304 | + traceback.print_exc() |
| 305 | + ok = False |
| 306 | + print("\nALL PARITY PASS" if ok else "\nPARITY/RUN FAILURE (see above)") |
| 307 | + sys.exit(0 if ok else 1) |
| 308 | + |
| 309 | + |
| 310 | +if __name__ == "__main__": |
| 311 | + main() |
0 commit comments