Skip to content

Commit 949042b

Browse files
committed
feat(mamba3 F0): Metal tensor-op precompute prim (cb + summary_states T.gemm) on Apple GPU
F0 chunk_precompute_fwd_metal_gemm_prim: the two head-independent serial reductions (cb=C@B^T, summary_states=(decay*dt*x)^T@B) re-expressed as T.gemm, behind env flag CPPMEGA_PATH_C_METAL_GEMM. The byte-identical serial Metal prim stays the parity reference (flag off). RULE #1: GEMM = the ONE path when on; tile/divisibility/compile failure RAISES, no serial fallback. Metal delta vs the CUDA twin: C accumulators are register FRAGMENTS (forces the stable 8x8 metal.simdgroup path; shared-C with N>=32/K%16==0 routes to the M5-only metal.cooperative_tensor path that fails on M1-M4). fragment->fp32-shared ->global staging (fragment->fp16 is an illegal simdgroup cast); dacs reloaded from fp16 dA_cumsum (the F2 convention) to fit Apple's hard 32 KB threadgroup limit. fp16 operands / fp32 accum; summary_states stored fp32. MEASURED (M4 Max, S=128 chunk=64 G=1 H=2 P=N=64): cb + dA_cumsum bit-exact (0.0), summary_states max|abs diff| 2.1e-6 vs serial (fp16 gate 5e-4); 8.36x faster (0.161ms serial -> 0.019ms gemm). Compiles at prod dims S=4096 H=112 G=8. B2 (chunk_scan_combine_bwd) + B0 (chunk_precompute_bwd): GEMM flag RAISES a clear NotImplementedError (honest scope, RULE #1 — no silent over-budget launch or fabricated wrong rewrite). B2's 4-contraction fragment-staging set exceeds Apple's 32 KB threadgroup limit (the gb10 CUDA twin needs 72.5 KB); B0 is a reverse-cumsum scan + per-l atomic scatters, not a 2D matmul. Their serial Metal prims stay the working path. The chunk_scan_combine_bwd_metal_gemm_prim (DYX + dchunk_states two clean GEMMs) is defined for reference / larger-budget backends. Probe: scratch/probe_f0b2b0_metal_gemm_parity.py (Apple-GPU parity+timing).
1 parent f185eec commit 949042b

3 files changed

Lines changed: 773 additions & 6 deletions

File tree

0 commit comments

Comments
 (0)