Skip to content

Commit 648085a

Browse files
committed
feat(mamba3-pathc): Stage 3 — B0/B1/B2 backward Metal kernels (transpose of F0/F1/F2)
3 single-entry grid kernels in mamba3_chunked_backward_core.py, the exact analytic transpose of the proven forward F0/F1/F2: - B2 mamba3_chunk_scan_combine_bwd: output/gate + Y_diag + Y_off transpose -> dC(off+diag), dx(skip), dz, dchunk_states, dinp(diag), dA_cumsum_y, dD. - B1 mamba3_inter_chunk_recur_bwd: the NEW O(S/C) REVERSE upper-tri inter-chunk combiner (adjoint of the F1 lower-tri recurrence) -> dstates, dh0, dA_cumsum_tail. REUSES the forward-materialized prev_states (no 8x replay). - B0 mamba3_chunk_precompute_bwd: decay_states transpose + dinp assembly -> dx(inp), dB, dlog_decay (cumsum-VJP), ddt (decay+inp paths) + dA scatter. Validated vs the MLX backward proto (itself 1.30e-4 vs serial VJP): worst grad max|abs| = 3.68e-4 < 1e-3 across dz/dx/dC/dB/dlog_decay/ddt/dh0/dD, incl non-zero dh_last and S in {256,512}, chunk=64 (the production tile config). chunk{128,256} parity holds at the MLX fp32-algebra level (proto); the Metal tile pins chunk=64=block_M like the forward F2. scratch/test_b0b1b2_metal_vs_proto.py is the standalone RULE #1 parity gate. Not yet wired into the live region (descriptors/interpose/region-flip next). Flag CPPMEGA_PATH_C_MAMBA3_CHUNKED_SCAN still default OFF (merge-safe).
1 parent d9559d6 commit 648085a

3 files changed

Lines changed: 1285 additions & 0 deletions

File tree

0 commit comments

Comments
 (0)