Skip to content

Commit ab8332b

Browse files
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

0 commit comments

Comments
 (0)