Skip to content

Commit ac862b8

Browse files
apstenku123apstenku123
committed
docs(triton-route): Tri-Dao bwd route+MMA proof + honest parity NO-GO
7/7 Tri-Dao mamba_ssm bwd route to real TF32 MMA on gb10 (288 HMMA.1688.F32.TF32 SASS, sm_121a); FIX#1 output grid-scaling correct at IR (dprev=grid*4096 full); native dstates production (S=4096 H=112)=1.206ms/kernel. End-to-end parity BLOCKED by a remaining input read-buffer grid-scaling defect (root-caused, reproduced); no parity PASS fabricated. RULE#1.
1 parent e48b8b4 commit ac862b8

1 file changed

Lines changed: 146 additions & 0 deletions

File tree

‎docs/TRITON-ROUTE-CORRECT.md‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
# Tri-Dao mamba_ssm bwd through OUR tilelang->tvm stack: route + MMA proof + remaining parity blocker
2+
3+
MEASURED on NVIDIA GB10 (sm_121a, Blackwell). All GPU exec under the gb10 mutex.
4+
tilelang @ `merge/upstream-codegen-reorg` HEAD `838b83b3` (FIX#1 `46b8115d` + FIX#2 `838b83b3`).
5+
venv `/home/dave/cppmega-venv`, source `/home/dave/source/tilelang`.
6+
7+
## GO / NO-GO: NO-GO for end-to-end numeric parity (HONEST)
8+
9+
- ROUTING + MMA codegen: **GO** — all 7/7 Tri-Dao bwd kernels compile end-to-end
10+
tilelang->tvm to real TF32 tensor-core MMA, on the production shape.
11+
- OUTPUT grid-scaling (FIX#1): **GO at IR level** — dprev output buffer is sized to the
12+
FULL grid extent (proven below), not a single tile.
13+
- End-to-end numeric PARITY (routed-vs-native, all elements 1e-3): **NO-GO** — blocked by a
14+
REMAINING frontend defect in INPUT read-buffer grid-scaling (root-caused below with a
15+
reproducing experiment). NO parity PASS is claimed. RULE #1: no strides=0 / single-tile
16+
number is reported as a result.
17+
18+
## 1. Routing: 7/7 to real HMMA TF32 MMA (MEASURED)
19+
20+
`route_all7.py` (compile each captured TTIR via `from_ttir(..., _allow_text_ttir=True)`):
21+
22+
```
23+
ROUTE _chunk_scan_bwd_dstates len= 33776 global=True mma=2 atomicAdd=0
24+
ROUTE _chunk_scan_bwd_dc len= 34395 global=True mma=2 atomicAdd=0
25+
ROUTE _chunk_scan_bwd_dcb len= 48488 global=True mma=2 atomicAdd=0
26+
ROUTE _chunk_scan_bwd_dx len= 55251 global=True mma=2 atomicAdd=1
27+
ROUTE _chunk_state_bwd_db len= 38231 global=True mma=2 atomicAdd=0
28+
ROUTE _chunk_state_bwd_dx len= 49304 global=True mma=2 atomicAdd=3
29+
ROUTE _chunk_state_bwd_ddAcs_stable len= 38991 global=True mma=2 atomicAdd=1
30+
TOTAL_OK_WITH_MMA=7/7
31+
```
32+
33+
CUDA carries `tl::mma_sync<kTensorFloat32,...,16,8,8,...>` (m16n8k8 TF32) +
34+
`#include <tl_templates/cuda/instruction/mma.h>`, `extern "C" __global__ __launch_bounds__(128,1)`.
35+
36+
## 2. SASS: real tensor-core MMA + atomics on the production-shape cubin (MEASURED)
37+
38+
The cubin is shape-agnostic (all dims/strides are runtime kernel args; the same cubin serves any
39+
grid), built for `.target sm_121a` / `.elftype @"ET_EXEC"`. `export_sass` then grep:
40+
41+
```
42+
288 HMMA.1688.F32.TF32 # total across all 7 kernels (real TF32 m16n8k8 tensor-core)
43+
per-kernel HMMA: dstates 32, dc 32, dcb 32, dx 32, db 32, ddAcs 64, state_dx 64
44+
reduce kernels also emit: REDG.E.ADD.F32.FTZ.RN.STRONG.GPU (real global atomic-add, 64 in dx)
45+
```
46+
47+
This is OUR codegen emitting genuine Blackwell tensor-core MMA on the production shape.
48+
49+
## 3. MEASURE: native dstates on the REAL production config (MEASURED)
50+
51+
Config §P1: S=4096, c=64, g=8, H=112, P=64, N=64, bs=1.
52+
`native_prod.py`, grid=(1,64,112), dprev_numel=29,360,128:
53+
54+
```
55+
NATIVE_PROD_MS_PER_KERNEL = 1.20584 ms (200-iter timed avg)
56+
dprev sum=-203350.219 nz=29360127/29360128 (full output)
57+
```
58+
59+
Small representative config (b1 nh8 hd64 ds64 nc8 cs64): native = 0.01470 ms/kernel.
60+
61+
Note on the "~10ms-class": that figure is the FULL mamba_ssm backward (all 7 kernels + the
62+
non-Triton ops), not one kernel. The single dstates kernel native cost is ~1.2 ms at production
63+
shape. The routed kernel ms cannot be reported until parity is fixed (see §5); reporting a routed
64+
ms for a kernel that produces wrong/truncated output would violate RULE #1.
65+
66+
## 4. OUTPUT grid-scaling (FIX#1) — proven at IR level (MEASURED)
67+
68+
Routed PrimFunc declared param shapes (`from_ttir` -> `tvm.ir.save_json`):
69+
70+
```
71+
arg0 (dout) = (2048 * gridDim_1 * gridDim_2 * gridDim_0 * gridDim_0_1,)
72+
arg1 (C) = (2048 * gridDim_1 * gridDim_2 * gridDim_0 * gridDim_0_1,)
73+
arg2 (dprev) = (4096 * gridDim_1 * gridDim_2 * gridDim_0 * gridDim_0_1,) <-- OUTPUT, full grid
74+
arg3 (dA) = (32 * gridDim_1 * gridDim_2 * gridDim_0 * gridDim_0_1,)
75+
```
76+
77+
For the small config grid=(1,8,8): arg2 = 4096*1*1*8*8 = 262144 = exactly
78+
b*nc*nh*hd*ds = 1*8*8*64*64. The single-tile `(4096,)` truncation is gone; the in-kernel store
79+
guard is `idx < grid*4096` (matches the dense output). FIX#1 is correct for the OUTPUT.
80+
81+
## 5. REMAINING PARITY BLOCKER (root-caused, reproducing experiment) — RULE #1 RAISE
82+
83+
The end-to-end routed-vs-native parity FAILS. Honest root cause, proven by a controlled
84+
experiment (`parity_tiny.py`, NOT a strides=0 shortcut):
85+
86+
- **Defect**: INPUT read buffers (arg0 dout, arg3 dA) are grid-scaled with the **per-K-tile-load
87+
footprint** (`flat_tile_extent`: dout 64x32=2048, dA 32) times the grid, in
88+
`poc/triton_frontend/__init__.py::_flat_extent_for_indices` (the
89+
`isinstance(entry, tir.Buffer/LazyTileExpr)` branch, ~line 218:
90+
`extent_expr = flat_tile_extent; for ... : extent_expr *= program_id_extent`).
91+
- For a STRIDED block-pointer input read whose per-block base uses the **seqlen stride** and that
92+
is read across a **K-loop** (cs/BK trips, advancing by `stride*BK` each trip), the true reachable
93+
flat index is `(grid-1)*block_stride + chunk_footprint`, which **exceeds** `flat_tile_extent*grid`.
94+
i.e. the declared input extent UNDER-counts by ~chunk_size/BLOCK_K.
95+
- The lowering then synthesises a matching in-kernel read guard
96+
`idx < gridDim_1*gridDim_2*gridDim_0_1*gridDim_0 * 2048` (and `*32` for dA). Valid strided reads
97+
beyond that bound are **masked to 0** -> the GEMM consumes zeros -> the routed output is
98+
near-empty (wrong), NOT truncated-by-one-tile.
99+
- This is the symmetric twin of FIX#1: FIX#1 fixed the OUTPUT (whose per-program footprint IS one
100+
dense tile, so `grid*tile_numel` is exact); the INPUTS need the per-program **strided** footprint,
101+
which `grid*tile_numel` does not bound.
102+
103+
Reproducing experiment (degenerate config chosen so cs==BLOCK_K==32, a single K-trip, declared
104+
extent == real extent: dout=8192=8192, dA=128=128):
105+
106+
```
107+
=== PARITY DEGENERATE (cs=BK=32, single K-trip) ===
108+
NATIVE nz=16384/16384 sum=-1240.0947
109+
ROUTED nz=128/16384 sum=219.3126
110+
MAXDIFF=1.023045e+02 ALLCLOSE_1e-3=False
111+
FAIL
112+
```
113+
114+
Even here the routed output is ~empty (128/16384 nonzero) because for chunk c>0 the dout per-block
115+
base `c*cs*stride_seqlen` already reaches the `grid*2048` guard boundary and the upper reads mask
116+
to 0. This isolates the bug to the input read-buffer extent / guard — NOT the MMA, NOT the output
117+
store (both proven correct above).
118+
119+
### Why not papered over
120+
A sound fix requires threading the per-block **stride span** of each input block-pointer (available
121+
as PtrState/scalar-param stride Vars when the offset tile is materialized) into
122+
`_flat_extent_for_indices`, so the input read buffer is declared to span
123+
`(grid-1)*block_stride + per_program_footprint` (a symbolic extent over the stride Vars). The offset
124+
tile reaching `_flat_extent_for_indices` is currently an OPAQUE buffer (`carry_index_223` etc.) with
125+
no stride metadata attached, so the bound cannot be derived at that point today. Over-declaring with
126+
a guessed constant could mask a real OOB or silently change results, which RULE #1 forbids. This is
127+
reported as the exact remaining stage rather than shipped as a fabricated PASS.
128+
129+
## Environment gotcha (recorded for the next run)
130+
`from_ttir` (libtriton MLIR parse) and `import torch` (its own static LLVM) **cannot coexist in one
131+
process** — whichever loads second segfaults the interpreter. The parity harness therefore splits:
132+
process A runs `from_ttir` and dumps the PrimFunc via `tvm.ir.save_json`; process B (`import torch`
133+
first, then `import tilelang`, then `tvm.ir.load_json` + `tilelang.compile`) drives the kernel.
134+
TVM-script round-trip is lossy for tilelang tile ops (`T.copy` not in the tirx parser) — use
135+
`save_json`/`load_json`, which preserve the full IR (35 params verified).
136+
137+
## Files (absolute)
138+
- `/Volumes/external/sources/tilelang/poc/triton_frontend/_test_harness/tridao_parity/route_all7.py`
139+
- `/Volumes/external/sources/tilelang/poc/triton_frontend/_test_harness/tridao_parity/phaseA_native.py`
140+
- `/Volumes/external/sources/tilelang/poc/triton_frontend/_test_harness/tridao_parity/native_prod.py`
141+
- `/Volumes/external/sources/tilelang/poc/triton_frontend/_test_harness/tridao_parity/parity_stage2b.py`
142+
- `/Volumes/external/sources/tilelang/poc/triton_frontend/_test_harness/tridao_parity/parity_tiny.py`
143+
- `/Volumes/external/sources/tilelang/poc/triton_frontend/_test_harness/tridao_parity/emit_pf_json.py`
144+
- `/Volumes/external/sources/tilelang/poc/triton_frontend/_test_harness/tridao_parity/sass_all7.py`
145+
- Frontend defect site: `/Volumes/external/sources/tilelang/poc/triton_frontend/__init__.py`
146+
`_flat_extent_for_indices` (~line 218) + `_redecl_input_buffer` (~line 314).

0 commit comments

Comments
 (0)