Commit 283889b
committed
feat(mamba3 B2): tensor-core CUDA backward prim (DYX/dC_off/dC_diag/dchunk_states T.gemm) on sm_121
Adds chunk_scan_combine_bwd_cuda_prim_gemm — the §26 backward lever. Converts B2's
four GEMM-able contractions from threaded-serial reductions to tensor-core T.gemm,
mirroring the proven F0 14.5x precedent (chunk_precompute_fwd_cuda_prim) and the F2
masked-intra-chunk-GEMM (chunk_scan_fwd_cuda_prim):
(A) DYX[l,s] = T.gemm(dY16, XT16, transpose_B) M=L,N=L,K=P
(B) dC_off[l,n] = T.gemm(dY16, ps16) M=L,N=N,K=P (sd folds post-GEMM)
(C) dC_diag[l,n] = T.gemm(M16(masked), B16) M=L,N=N,K=L (F2 causal lower-tri mask)
(D) dchunk[p,n] = T.gemm(dY_sd16, C16, transpose_A) M=P,N=N,K=L (decay sd folded into dY operand, F0 pattern)
dinp (3-index, s-dependent Lmat) + dz/dx/dD + dseg STAY threaded/byte-identical
(documented scope boundary; dinp does not collapse to a single GEMM). fp16 operands
/ fp32 register-fragment accum, plain scope=shared (no swizzle), disable_tma=True;
smem-minimal reuse tiling (opA/opB/store_fp32 reused across phases, XT dropped) =
72.5KB at prod L=P=N=64 (<99KB gb10 budget).
Selected by env CPPMEGA_PATH_C_B2_GEMM (CUDA-only, mutually exclusive with _V2 — both
set RAISES). MMA-divisibility + smem-budget RAISE gates (RULE #1: the ONE cuda path,
no threaded-serial fallback). Metal B2 prim, v1/v2 cuda prims, B1, B0, all grids
BYTE-IDENTICAL (AST-verified). out_idx [11..17] + dA_cumsum_y shape unchanged.
probe: --b2-gemm-ab times v1 vs GEMM in one process + GEMM-vs-v1 max|abs|; the env
flag routes the chained 8-grad gate (dD via honest fp16-cache gold) through the GEMM
prim. NO GPU this round (py_compile+ast+cross-symbol+smem-byte-budget on Mac;
MEASURE on gb10 under the GPU mutex).1 parent e335358 commit 283889b
2 files changed
Lines changed: 541 additions & 15 deletions
File tree
- cppmega_mlx/nn/_tilelang
- scratch
0 commit comments