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