Commit 58684f3
committed
feat(§DYN): Triton-mold static->dynamic B2 batched prim — MEASURED NO-GO (0.750x, static=0 confirmed)
NEW flag-gated CUDA B2 prim chunk_scan_combine_bwd_cuda_prim_gemm_batched_dyn:
the §TB1/§SO1 batched math/grid/in-kernel head-loop, BYTE-FOR-BYTE, with the five
GEMM operand-staging tiles (dY16/opA/opB/store_fp32/dCdiag_sh) moved from explicit
STATIC scope="shared" to the DYNAMIC region (T.alloc_shared default shared.dyn).
This is the Triton mold (Tri-Dao bwd HW proof on the same gb10: STATIC=0,
MAXDYN=101376). Selected by CPPMEGA_PATH_C_B2_GEMM_BATCHED_DYN (mutually exclusive
with the other B2 flags; RAISES on ambiguity). RULE #1: ONE path when flagged,
RAISES on smem/proof/compile failure — no fallback.
MEASURED on gb10 (prod b=1,S=4096,chunk=64,G=8,H=112,P=64,N=64, HEADS_PER_CTA=2):
- BUILD: num_static_shared_decls=0, extern __shared__ buf_dyn_shmem present
=> STATIC smem dropped to ~0, the Triton mold achieved. HPC=2 LAUNCHES with the
full 4-GEMM layout (the §SO1 squeeze need not be reverted; real total 91136 B
FITS the 101376 dynamic opt-in cap).
- TIMING: batched_dyn=1407.187ms vs v1=1054.847ms => 0.750x NO-GO.
vs §TB1 batched_static=1420.256ms => 1.009x (dynamic move ~1% faster).
- PARITY: all 8 grads vs v1 within ~9e-5; chained 8-grad gate PASSES (worst
8.1e-4 < 1e-3).
HONEST NO-GO (matches the DESIGN risks): static->dynamic achieved launch-feasibility
(static=0) but does NOT beat v1. WHY: (1) the dinp 3-index serial term (un-GEMM-able,
untouched by the smem fix) dominates; (2) path_c's dC is per-HEAD, so unlike Tri-Dao
there is NO M-fattening head-accumulation — only operand-staging/sync amortization,
which is smaller; (3) low CTA supply + ~91KB dynamic per CTA caps occupancy to ~1
CTA/SM at bs1. The real fix would be re-architecting dC/dinp to per-GROUP split
buffers (a larger change than the smem flip).
GATE FIX (RULE #1): the §DYN build-site smem gate counts the ACTUAL §SO1 prim allocs
(dY fp16 not fp32; DYX a SINGLE (L,L) tile not an HPC band) => real total 91136 B at
HPC=2 (FITS), vs the §TB1 STATIC gate's stale 123904 B (which wrongly RAISES). An
over-counting gate that blocks a hardware-valid launch is itself a fail-fast defect.
z3: scratch/proof_b2_batched_driver.py extended with the §DYN scope-flip proof — the
static->dynamic move is byte-layout-only (operand maps / single-writer / mask / scale
unchanged); 3 positives assert the §DYN proof verdict is IDENTICAL to the §TB1 static
proof, + a non-vacuous negative (interleaved dynamic head-bands MUST break
single_writer). All pass non-vacuously.
INVARIANTS: v1 / §27 gemm / §TB1 batched / v2 prim bodies BYTE-IDENTICAL (verified by
md5); B1 untouched; the only edits to existing code are the mutual-exclusion guard
(adds the new flag) and the new elif branch. probe gets a --b2-gemm-ab §DYN arm.1 parent 4f5d548 commit 58684f3
3 files changed
Lines changed: 1360 additions & 117 deletions
File tree
- cppmega_mlx/nn/_tilelang
- scratch
0 commit comments