Skip to content

Commit 283889b

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

0 commit comments

Comments
 (0)