|
| 1 | +# Triton-Mold path_c B2 Backward (§DYN static→dynamic) — MEASURED VERDICT |
| 2 | + |
| 3 | +Status: **DECIDED — HONEST NO-GO (0.745x vs in-process v1).** The static→dynamic |
| 4 | +"Triton mold" achieves launch-feasibility (driver `STATIC=0`, HPC=2 launches with |
| 5 | +the full 4-GEMM layout) but does **not** beat the v1 threaded B2 prim. |
| 6 | + |
| 7 | +Hardware: NVIDIA GB10 (sm_121, unified 121 GB). Prod cfg: `b=1, S=4096, chunk=64, |
| 8 | +G=8, H=112, P=64, N=64`, `HEADS_PER_CTA=2` (the only valid HPC for `nheads/ngroups |
| 9 | += 14`; see §2). Toolchain: CUDA 13.3, tilelang from source, torch 2.13 dev cu132. |
| 10 | +Date: 2026-06-06. |
| 11 | + |
| 12 | +--- |
| 13 | + |
| 14 | +## 1. What was built (the Triton mold) |
| 15 | + |
| 16 | +NEW flag-gated CUDA prim `chunk_scan_combine_bwd_cuda_prim_gemm_batched_dyn` |
| 17 | +(`cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py:2699`). It is the |
| 18 | +§TB1/§SO1 batched math/grid/in-kernel-head-loop **byte-for-byte**, with the five |
| 19 | +GEMM operand-staging tiles (`dY16`, `opA`, `opB`, `store_fp32`, `dCdiag_sh`) moved |
| 20 | +from explicit **STATIC** `scope="shared"` to the **DYNAMIC** region (the |
| 21 | +`T.alloc_shared` default `shared.dyn`). tilelang lowers `shared.dyn` into |
| 22 | +`extern __shared__ __align__(1024) uchar buf_dyn_shmem[]`, so the compiler reserves |
| 23 | +~0 static smem and the driver grants the full per-block dynamic opt-in. This is the |
| 24 | +Tri-Dao Triton bwd HW mold (same gb10: STATIC=0, MAXDYN opt-in). |
| 25 | + |
| 26 | +Wired behind `CPPMEGA_PATH_C_B2_GEMM_BATCHED_DYN`, mutually exclusive with the other |
| 27 | +B2 flags (RAISES on ambiguity). Own dynamic-smem gate against the 101376 B opt-in |
| 28 | +cap (RAISES over-cap). z3/TLA proof gate (RAISES `GemmRewriteNotProven` on |
| 29 | +sat/unknown). RULE #1: ONE path when flagged, no fallback. **Default path (no flag) |
| 30 | +is unchanged v1** (`else: prim = chunk_scan_combine_bwd_cuda_prim`). |
| 31 | + |
| 32 | +--- |
| 33 | + |
| 34 | +## 2. cuFuncGetAttribute — STATIC≈0 PROVEN at the driver level |
| 35 | + |
| 36 | +The cubin (nvcc `-arch=sm_121a` on the exact emitted `device_kernel.cu`) loaded via |
| 37 | +the CUDA driver API (`cuda.bindings.driver`): |
| 38 | + |
| 39 | +``` |
| 40 | +cuFuncGetAttribute(main_kernel): |
| 41 | + CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES = 0 <-- STATIC = 0 (mold achieved) |
| 42 | + CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES (default) = 49152 |
| 43 | + CU_FUNC_ATTRIBUTE_NUM_REGS = 192 |
| 44 | +cuFuncSetAttribute(MAX_DYNAMIC = 91136) = GRANTED -> now 91136 (HPC=2 LAUNCHES) |
| 45 | +device MAX_SHARED_MEMORY_PER_BLOCK_OPTIN = 101376 |
| 46 | +``` |
| 47 | + |
| 48 | +cuobjdump cross-check on the same cubin: |
| 49 | +``` |
| 50 | +Function main_kernel: REG:192 SHARED:1024 LOCAL:0 STACK:0 (smem=1024) |
| 51 | +.nv.shared.main_kernel size 0x400 (=1024 B) <-- just the __align__(1024) slot for |
| 52 | + the extern dyn array, NOT staging |
| 53 | +``` |
| 54 | + |
| 55 | +**STATIC dropped from 57344 B (the §40a26d44 per-head GEMM staging) to 0.** The only |
| 56 | +"static" smem is the 1024 B alignment placeholder for the dynamic `buf_dyn_shmem[]`. |
| 57 | +The 91136 B dynamic request (HPC=2) fits the 101376 B sm_121 opt-in cap → the driver |
| 58 | +grants it → **HPC=2 LAUNCHES** with the full 4-GEMM layout. |
| 59 | + |
| 60 | +`HEADS_PER_CTA=4` is **INVALID** at prod: `nheads/ngroups = 112/8 = 14`, and 4 neither |
| 61 | +divides nor is a multiple of 14, so a CTA head-band would straddle a group whose |
| 62 | +`C`/`B` operands are per-GROUP. `_b2_batched_heads_per_cta` RAISES (RULE #1: no silent |
| 63 | +straddle). The valid HPC set is {1, 2, 7, 14}; HPC=2 is the measured point. |
| 64 | + |
| 65 | +--- |
| 66 | + |
| 67 | +## 3. TIMING — MEASURED NO-GO |
| 68 | + |
| 69 | +In ONE gb10 process (`probe_chunked_backward_cuda_gb10.py --prod --b2-gemm-ab`, |
| 70 | +`CPPMEGA_PATH_C_B2_HEADS_PER_CTA=2`): |
| 71 | + |
| 72 | +| prim | ms (median) | vs in-process v1 | |
| 73 | +| ---------------------------- | ----------: | ---------------: | |
| 74 | +| v1 threaded (`prim` default) | 1052.074 | 1.000x | |
| 75 | +| single-tile GEMM (§27) | 1370.047 | 0.768x | |
| 76 | +| batched-large-tile (§TB1) | 1418.405 | 0.742x | |
| 77 | +| **§DYN batched_dyn (mold)** | **1412.696** | **0.745x NO-GO** | |
| 78 | + |
| 79 | +§DYN vs batched_static = 1.004x (the static→dynamic scope flip is ~0.4% faster — the |
| 80 | +smem move is free but does not change the dominant cost). **0.745x < 1.0x → NO-GO.** |
| 81 | + |
| 82 | +NOTE on baselines: this probe's in-process v1 measures **1052 ms**, not the historical |
| 83 | +905.4 ms reference (an earlier wf machine/thermal state). All §DYN ratios above are vs |
| 84 | +the **same-process 1052 ms v1**, which is the apples-to-apples comparison. Against the |
| 85 | +historical references: §DYN 1412.7 ms is 0.64x of 905.4 ms (v1_threaded), 0.87x of |
| 86 | +1231 ms (§27), and 141x of 10 ms (Tri-Dao Triton). Tri-Dao remains ~140x faster. |
| 87 | + |
| 88 | +Backward chain (B2 + B1 + B0), this run: **B2 1051.1 + B1 5.55 + B0 266.2 ms**. |
| 89 | +(The 447.8 ms reference is the v1 B2 from an earlier wf; B2 alone is 1051 ms here.) |
| 90 | + |
| 91 | +--- |
| 92 | + |
| 93 | +## 4. SASS — real HMMA, but PER-HEAD relabel (NOT fattened-M head-accum) |
| 94 | + |
| 95 | +`cuobjdump --dump-sass` on the §DYN cubin: |
| 96 | +``` |
| 97 | +128 HMMA.16816.F32 (m16n8k16 tensor-core MMA, F32 accumulate) |
| 98 | + 0 OMMA / IMMA / BMMA |
| 99 | +``` |
| 100 | + |
| 101 | +These are real tensor-core HMMA tiles (`HMMA.16816.F32 R84, R96, R124, RZ ; ...`). |
| 102 | +But the M-dimension of every GEMM is the **per-head** tile (L=64 or P=64), issued |
| 103 | +inside `for hh in T.serial(HPC)` loops — the four contractions (DYX = dY@xᵀ, dC_off = |
| 104 | +dY@prev_states, dC_diag masked, dchunk = (decay·dY)ᵀ@C) each run once **per head** with |
| 105 | +operands that differ on BOTH sides. The 128 HMMA = 4 GEMMs × HPC=2 heads × the 16-row |
| 106 | +fragmentation of the 64-tiles. **This is a per-64 (per-head) relabel, NOT a fattened-M |
| 107 | +head-accumulation.** Tri-Dao fattens M by looping `nheads_per_program` heads that SHARE |
| 108 | +a per-GROUP operand (B/C) into ONE accumulator tile; path_c cannot, because dC and dinp |
| 109 | +are **per-HEAD** (distinct operands both sides) — a dense tall-M over heads would compute |
| 110 | +HPC² cross-head garbage (block-diagonal violation). The mold therefore amortizes only |
| 111 | +operand-staging/sync, never the tensor-core M-setup. |
| 112 | + |
| 113 | +--- |
| 114 | + |
| 115 | +## 5. PARITY — all 8 grads PASS |
| 116 | + |
| 117 | +§DYN vs v1 (in-process, max|abs| over ALL elements): |
| 118 | +``` |
| 119 | +dC 6.21e-05 dx 0.00e+00 dz 0.00e+00 dchunk 9.30e-05 dinp 6.41e-05 |
| 120 | +dA_y 3.15e-05 dD 4.41e-06 |
| 121 | +``` |
| 122 | +Chained 8-grad gate (B2→B1→B0, GATE 1e-3, ALL elements, default v1 chain): |
| 123 | +``` |
| 124 | +dz 1.73e-04 dx 8.10e-04 dC 5.03e-05 dB 1.09e-05 dlog_decay 6.67e-04 |
| 125 | +ddt 1.50e-04 dh0 1.84e-04 dD 2.57e-05 -> WORST 8.099e-04 < 1e-03 PASS |
| 126 | +``` |
| 127 | +dD uses the fp16-cache; the §DYN deltas are fp16-rounding (≤9.3e-5), well within gate. |
| 128 | + |
| 129 | +--- |
| 130 | + |
| 131 | +## 6. z3 PROOF — non-vacuous |
| 132 | + |
| 133 | +`proof_b2_batched_driver.py`, `dyn_scope_flip`: |
| 134 | +- POSITIVES dchunk / dcoff / dcdiag: `z3_proved=True, single_writer=True, |
| 135 | + operand_maps_match=True, scale_equiv=True, verdict_identical_to_static=True`. |
| 136 | +- NEGATIVE `interleaved_dyn_bands` (m_stride = L//HPC, overlapping head rows): |
| 137 | + `z3_proved=False, single_writer=False` — correctly fails (non-vacuous). |
| 138 | +- VERDICT: **ALL_POSITIVES_PROVED_AND_NON_VACUOUS.** |
| 139 | + |
| 140 | +The static→dynamic move is byte-layout-only: operand maps, single-writer band |
| 141 | +disjointness (m_blocks=HPC), mask, and scale-fold are identical to the §TB1 static |
| 142 | +proof. The dynamic region does NOT relax single-writer. |
| 143 | + |
| 144 | +--- |
| 145 | + |
| 146 | +## 7. Step tok/s — UNCHANGED (NO-GO consequence) |
| 147 | + |
| 148 | +§DYN is NO-GO, so it is **never on the production step path**: the default B2 branch |
| 149 | +(no flag) is v1, and RULE #1 makes §DYN opt-in only (RAISES on smem/proof/compile |
| 150 | +failure, no silent use). Production step tok/s @8L/@28L, bs1 AND bs4, is therefore the |
| 151 | +**v1 baseline** — wiring the 0.745x §DYN prim into the step would only regress it. No |
| 152 | +step-tok/s gain is claimed or measured for the mold; the gap vs Megatron 3399 tok/s is |
| 153 | +the pre-existing v1 gap, unchanged by this work. |
| 154 | + |
| 155 | +--- |
| 156 | + |
| 157 | +## 8. HONEST NO-GO — WHY static→dynamic still can't beat v1 |
| 158 | + |
| 159 | +The mold removed the smem launch-blocker (STATIC 57344→0, HPC=2 launches), but the |
| 160 | +dominant costs are untouched: |
| 161 | + |
| 162 | +1. **The dinp 3-index serial term is un-GEMM-able.** `dinp[b,s,h,p,n]` is a per-element |
| 163 | + outer-product zeroed and written by a serial thread-strided loop; the smem flip does |
| 164 | + nothing for it. It dominates the kernel. |
| 165 | +2. **path_c dC is per-HEAD → NO M-fattening.** Unlike Tri-Dao there is no per-group |
| 166 | + head-accumulation to amortize tensor-core setup (§4); the mold only amortizes |
| 167 | + operand-staging/sync, a small fraction. |
| 168 | +3. **Occupancy cap.** ~91 KB dynamic per CTA + 192 regs pins ~1 CTA/SM at bs1; the |
| 169 | + batched layout has fewer, fatter CTAs than v1's threaded grid, so it loses the |
| 170 | + latency-hiding v1 gets from many light CTAs. |
| 171 | + |
| 172 | +The real fix is re-architecting dC/dinp into per-GROUP split buffers (so a genuine |
| 173 | +fattened-M head-accumulation becomes valid) — a far larger change than the smem flip, |
| 174 | +and out of scope here. **Verdict: ship nothing to the default path; keep v1. The §DYN |
| 175 | +prim stays as a flag-gated, proven, launch-feasible NO-GO reference.** |
| 176 | + |
| 177 | +--- |
| 178 | + |
| 179 | +## Reproduce |
| 180 | + |
| 181 | +``` |
| 182 | +# timing + parity (HPC=2 mandatory: 14 = nheads/ngroups) |
| 183 | +CPPMEGA_PATH_C_B2_HEADS_PER_CTA=2 python scratch/probe_chunked_backward_cuda_gb10.py --prod --b2-gemm-ab |
| 184 | +# driver attributes (STATIC=0, MAXDYN grant) + SASS |
| 185 | +HPC=2 python scratch/wf1_cufunc_attr_b2dyn.py # source-regex + SASS HMMA |
| 186 | +# then: nvcc -arch=sm_121a -cubin <cache>/device_kernel.cu ... -o /tmp/wf1_dyn.cubin |
| 187 | +python scratch/wf1_cubin_attr.py /tmp/wf1_dyn.cubin # cuFuncGetAttribute STATIC/MAXDYN |
| 188 | +cuobjdump --dump-resource-usage /tmp/wf1_dyn.cubin # SHARED:1024, REG:192 |
| 189 | +cuobjdump --dump-sass /tmp/wf1_dyn.cubin | grep -c HMMA # 128 HMMA.16816.F32 |
| 190 | +# z3 |
| 191 | +python scratch/proof_b2_batched_driver.py # ALL_POSITIVES_PROVED_AND_NON_VACUOUS |
| 192 | +``` |
0 commit comments