Commit 2b884ce
committed
feat(path_c/B2): add dstate-split B2 v2 cuda prim (track 3, env-gated A/B)
Add chunk_scan_combine_bwd_cuda_prim_v2 — the §17-GO v1 backward B2 math with
the dstate (N) axis SPLIT across a NEW third grid dim bz (extent dstate_split/KN):
- Grid T.Kernel(batch*nchunks, nheads, dstate_split); kn=bz owns disjoint
N-range [n0,n0+n_per) so dC/dchunk_states/dinp writes are race-free (no atomics).
- Per-tg shared (dacs/dAcs_acc/dY/XT/DYX) is dstate-independent => ZERO added
shared (occupancy not lost to shared pressure); N-serial reductions shrink by KN.
- dz/dx writes + dD atomic gated to bz==0 (uniform kn-branch / kn_is_zero factor);
the N-independent dseg segsum-VJP counted ONCE on bz==0 (branchless tri*kn_is_zero
mask — avoids the §17 barrier-in-if codegen bug).
- dA_cumsum_y becomes a per-bz PARTIAL of shape (b,H,nchunks,dstate_split,chunk)
at the SAME out_idx slot 16; summing the KN axis recovers v1 dA_cumsum_y exactly.
Selected CUDA-only behind env CPPMEGA_PATH_C_B2_V2 (+ CPPMEGA_PATH_C_B2_DSTATE_SPLIT,
default 2) in build_chunk_scan_combine_bwd_metal — NO try/except so a compile/parity
failure PROPAGATES (RULE #1: the ONE path, never falls back to v1/Metal/numpy).
dstate%dstate_split RAISES at build. v1 cuda prim + Metal prim + B0/B1 UNCHANGED
(byte-identical); KN==1 is the v1 body bit-for-bit.
Probe: --b2-v2-ab times BOTH v1 and v2 (KN sweep via CPPMEGA_PATH_C_B2_AB_KNS,
default 2,4) in one gb10 run, printing medians+speedup LABELLED MEASURED and the
v2-vs-v1 math-equivalence max|abs|; a slower v2 prints NO-GO (never silently used).
When CPPMEGA_PATH_C_B2_V2 is set the dA_y buffer gains the KN axis and is summed
before B0; off = v1 byte-identical chain. Local checks: py_compile + AST cross-symbol
(v1/Metal byte-shape preserved, out_idx [11..17], no try/except in builder, v2 RAISES
on bad split). No GPU on this Mac — Profile agent does the gb10 compile/run.1 parent 5258b80 commit 2b884ce
2 files changed
Lines changed: 491 additions & 10 deletions
File tree
- cppmega_mlx/nn/_tilelang
- scratch
0 commit comments