Skip to content

Commit e48b8b4

Browse files
apstenku123apstenku123
committed
docs(triton-frontend): STEP3 route+measure on gb10 — 3/7 compile to real TVM tensor-core MMA, 1 launches, parity NO-GO (output-buffer single-tile truncation bug, RULE#1 reported)
1 parent 39e8875 commit e48b8b4

1 file changed

Lines changed: 119 additions & 0 deletions

File tree

‎docs/TRITON-FRONTEND-ROUTE-FIXED.md‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
# Tri-Dao mamba_ssm bwd Triton kernels through tilelang -> tvm on gb10 (STEP 3 route+measure)
2+
3+
Date: 2026-06-06. Host: NVIDIA GB10 (sm_121), CUDA 13.2, driver 595.71.05.
4+
venv: `/home/dave/cppmega-venv` loading tilelang+tvm from SOURCE `/home/dave/source/tilelang`
5+
(commit `b5e4fed3`). triton 3.7.0, mamba_ssm 2.2.x.
6+
7+
## Honest GO / NO-GO
8+
9+
PARTIAL GO on COMPILE; **NO-GO on PARITY** (output is single-tile truncated — see bug below).
10+
11+
- 3/7 group-A bwd kernels COMPILE END-TO-END through OUR frontend
12+
(`from_ttir`/`compile_ttir`) -> tilelang lower -> **TVM CUDA codegen** to a real
13+
`__global__` tensor-core kernel. NOT a stub, NOT native-triton fallback.
14+
- The routed `_chunk_scan_bwd_dstates` kernel **LAUNCHES on gb10** (executes, returns).
15+
- BUT the routed output is WRONG: only one program tile (4096 = hdim*dstate = 64*64) of
16+
the full `dprev_states` tensor is written. Parity vs native triton FAILS. Root cause is
17+
a real frontend output-buffer size-inference defect (below). Per RULE#1 this is reported
18+
loudly, not papered over.
19+
20+
## Evidence the lowered PrimFunc is NON-EMPTY (real tt.dot body)
21+
22+
`from_ttir(_chunk_scan_bwd_dstates, _allow_text_ttir=True, target=cuda)`:
23+
- PrimFunc length: **29448 chars**; `T.evaluate(0)` count: **0**; has `T.gemm(...)`: YES; has loops: YES.
24+
- Real body contains: `T.gemm(T.region(dot_a_shared...), T.region(tile_load_305...), T.region(dot_c_frag...))`,
25+
masked `tile_load_*[i]=T.if_then_else(mask, arg[idx], 0.0)`, `T.exp(...)` for dA_cumsum,
26+
full int64 index arithmetic with overflow guards.
27+
28+
This comes from the NATIVE-PARSE provider (`triton_native_parse.py`, Provider 0), which uses
29+
the already-installed `triton._C.libtriton.ir.parse_mlir_module` to validate + canonically
30+
re-print the TTIR and adapt it for the existing OP_TABLE walker. NO C++ shim, NO mlir.ir/IREE
31+
build needed (the C++ PtrAnalysis shim is unavailable in-process because libtriton's static LLVM
32+
is already loaded; the frontend warns and uses the MVP scalar load/store path).
33+
34+
## Evidence of REAL tilelang->tvm CUDA codegen (real HMMA from OUR codegen)
35+
36+
`compile_ttir(...).get_kernel_source()` for `_chunk_scan_bwd_dstates_kernel` (33586 chars):
37+
```
38+
#include <tl_templates/cuda/instruction/mma.h>
39+
...
40+
tl::mma_sync<tl::DataType::kTensorFloat32, tl::DataType::kTensorFloat32,
41+
tl::DataType::kFloat32, 16, 8, 8, false, true>(
42+
reinterpret_cast<float*>(dot_c_frag_308 + ...),
43+
reinterpret_cast<const uint32_t*>(A_local + ...),
44+
reinterpret_cast<const uint32_t*>(B_local + ...));
45+
```
46+
`tl::mma_sync<...kTensorFloat32...16,8,8...>` = m16n8k8 TF32 HMMA emitted by OUR tilelang
47+
CUDA template (`tl_templates/cuda/instruction/mma.h`), driven by TVM codegen. This is OUR
48+
codegen, NOT the triton runtime, NOT a stub.
49+
50+
3/7 compile to this form: each `mma_sync`x2, `__global__`x2:
51+
- `_chunk_scan_bwd_dstates` (33586 chars), `_chunk_scan_bwd_dc` (34128), `_chunk_scan_bwd_dcb` (47676).
52+
53+
## N / 7 status (group-A first)
54+
55+
| kernel | TTIR capture | tilelang->tvm CUDA (real MMA) | launches on gb10 | correct output |
56+
|---|---|---|---|---|
57+
| _chunk_scan_bwd_dstates (A) | YES | YES (mma=2) | YES | NO (1-tile truncation) |
58+
| _chunk_scan_bwd_dc (A) | YES | YES (mma=2) | not run | NO (same bug) |
59+
| _chunk_scan_bwd_dcb (A) | YES | YES (mma=2) | not run | NO (same bug) |
60+
| _chunk_state_bwd_db (A) | YES | FAIL: `variables (arg4,) are used but not passed in` (tilelang var-binding) | - | - |
61+
| _chunk_scan_bwd_dx (B) | YES | FAIL: `m_warp*n_warp==num_warps` num_warps=0 | - | - |
62+
| _chunk_state_bwd_dx (B) | YES | FAIL: `Can't cast a handle to other types` | - | - |
63+
| _chunk_state_bwd_ddAcs_stable (B) | YES | FAIL: num_warps=0 | - | - |
64+
65+
**3/7 compile to real tensor-core MMA CUDA through OUR stack. 1/3 confirmed to launch.
66+
0/7 confirmed correct (output truncation bug blocks parity).**
67+
68+
## MEASURED (gb10, representative SMALL config)
69+
70+
Config: batch=1, nheads=8, headdim=64, dstate=64, nchunks=8, chunk_size=64, ngroups=1,
71+
seqlen=512, BLOCK_M=BLOCK_N=64, BLOCK_K=32, HAS_SEQ_IDX=False.
72+
73+
- NATIVE triton `_chunk_scan_bwd_dstates` (inner JITFunction, autotune stripped, BLOCKS pinned),
74+
per-kernel wall time: **0.0167 ms** (this small config is tiny; the manifest's ~10ms-class
75+
figure is for the full production-size workload, not this representative slice).
76+
- ROUTED tilelang->tvm kernel, per-launch wall time: **8.19 ms** — dominated by Python/FFI
77+
per-call overhead + the MVP SCALAR load/store path (C++ PtrAnalysis shim is disabled
78+
in-process), NOT a tensor-core-bound number. This is the launch of a kernel that only
79+
writes ONE tile, so it is NOT a fair perf comparison and MUST NOT be read as "the routed
80+
kernel is 8ms".
81+
- path_c v1 905ms: not a meaningful comparison at this stage (different scope; routed kernel
82+
is per-kernel and output-incorrect).
83+
84+
## Parity vs native (1e-3, all elements): FAIL
85+
86+
Routed output `out_buf`: 32 / 4096 elements nonzero; native fills the full
87+
`(batch, nchunks, nheads, headdim, dstate) = 262144`-element tensor. allclose(1e-3) = False.
88+
The routed kernel cannot write past one 64x64 tile.
89+
90+
## ROOT CAUSE (real frontend defect — reported loudly, NOT patched here)
91+
92+
The output buffer `dprev_states_ptr` is size-inferred by the walker as a SINGLE program tile
93+
`T.Buffer((4096,))` (4096 = hdim*dstate), whereas the INPUT buffers are correctly grid-scaled
94+
`(2048 * gridDim_0*gridDim_0_1*gridDim_0_2*gridDim_0_3,)`. Two consequences in the generated CUDA:
95+
1. A spurious store guard `if (idx < 4096) arg2[idx] = ...` (mma.cu line ~607) clamps every
96+
write whose per-block base offset (`blockIdx.x*arg22 + ...`) lands beyond the first tile —
97+
so all blocks except the first tile are silently dropped.
98+
2. The undersized `[4096]` output param vs the real 262144-element tensor causes an FFI/launch
99+
segfault when a correctly-sized output buffer is passed; passing exactly 4096 avoids the
100+
crash but yields the truncated result above.
101+
102+
Native triton stores with mask `(offs_m < hdim) & (offs_n < dstate)` ONLY and a per-block base
103+
pointer offset — it has NO `< 4096` whole-buffer clamp. The clamp is purely an artifact of OUR
104+
output-buffer size inference treating a masked store's decl_buffer tile-extent as the whole-buffer
105+
extent. Fix belongs in the frontend output-buffer shape inference (grid-scale the store target
106+
like the load sources), then re-verify parity. This is a substantive frontend change, out of
107+
scope for this measure step, and is left UNPATCHED rather than worked around.
108+
109+
## Reproduce
110+
111+
- Capture TTIR (pre-tvm, isolates libtriton): `/tmp/ttir7/*.ttir` (7 kernels).
112+
- PrimFunc dump: `/tmp/verify_route.py <name>`.
113+
- Compile + CUDA dump: `/tmp/run_route.py <name>` -> `/tmp/route_cuda/<name>_compiled.cu`.
114+
- Native ref (isolated proc): `/tmp/phaseA_native.py` -> `/tmp/dstates_io.pt`.
115+
- Routed launch + measure (isolated proc): `/tmp/phaseB2.py`.
116+
117+
NOTE: native-triton and the routed tvm path CANNOT coexist in one process (both link LLVM
118+
statically; loading both -> duplicate cl::opt -> segfault). All measurement is two-phase:
119+
native in one subprocess (saves tensors), routed in a fresh subprocess.

0 commit comments

Comments
 (0)