Commit ab8332b
committed
feat(mamba3 mono-bwd): cppmega mono-chunk PORT — B2 GEMM + dinp_diag-B0 fused, flag-gated, gb10-only
ONE kernel per (chunk,head) keeping the chunk state resident, the four §27
contractions as internal T.gemm (DYX/dC_off/dC_diag/dchunk_states), and the
dinp_diag-dependent HALF of B0 (dx_diag/dB_diag/ddt_diag) folded BEFORE any
dinp_diag global round-trip (the cppmega 'reuse from shared before any global
write' move). New: bwd_mono_cuda_prim / bwd_mono_metal_prim / build_bwd_mono /
bwd_mono_grid + MAMBA3_BWD_MONO_OP_NAME; probe --bwd-mono-ab / --nano.
Honest scope (RULE #1): B1 (cross-chunk reverse adjoint) + the dstates-coupled
B0 remainder CANNOT fuse into a per-CTA kernel -> stay separate; the math is
linear in dinp so the split is exact (mono-mode B0 takes dinp_diag=0).
SMEM wall (MEASURED structural NO-GO at prod): resident DINP[L,P*N] fp32 tile is
L*P*N*4 = 1,048,576 B at L=P=N=64 -> total 1.139 MB >> gb10 ~99 KB; build RAISES
with where+what (no over-budget launch, no fallback). Fits only at small L*P*N
(L16P32N32=79KB). Metal build RAISES unconditionally (gb10-only; 32KB cap).
Predicted NO-GO for throughput (§27 four-GEMM B2 0.749x; cppmega mono-chunk
3.0x slower; §18 occupancy-saturated). 6-kernel §17 447.8ms chain byte-identical
when mono not invoked (additive entry point, 6 builders/grids/op-names untouched).
Mac-verified: prim constructs at prod dims; MMA-divisibility + gb10 + Apple
budget gates all RAISE; body lowers to a race-free Metal JITKernel at the 32KB-fit
config; dinp_diag-B0 fold algebra matches B0's ddt_inp=sum dinp*x*B exactly. CUDA
build+measure is the GB10 phase (do not push; GB10 pushes).1 parent 73e2184 commit ab8332b
3 files changed
Lines changed: 836 additions & 3 deletions
File tree
- cppmega_mlx/nn/_tilelang
- docs
- scratch
0 commit comments