Commit 949042b
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
- cppmega_mlx/nn/_tilelang
- scratch
0 commit comments