Skip to content

Commit 2b884ce

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

0 commit comments

Comments
 (0)