Skip to content

Commit 39e8875

Browse files
apstenku123apstenku123
committed
docs(triton-route): Tri-Dao mamba_ssm bwd routed through OUR tilelang->tvm stack
3/7 bwd kernels route end-to-end to runnable sm_121 tensor-core CUDA (dstates: PTX .target sm_121a + 32 mma.sync.m16n8k8 tf32; SASS 32 HMMA.1688.F32.TF32, via JITKernel.export_ptx/sass -- proven NOT a native-triton fallback). The other 4 break at named tilelang-lower / emission boundaries after this round's parser fixes pushed them past the TTIR reduce-region + atomic_rmw barrier. Native Tri-Dao fwd+bwd measured 3.78 ms on GB10 (vs path_c v1 905 ms); routed-kernel numeric/timed run blocked on the unbuilt PtrAnalysis C++ shim (MVP scalar-load path) + the strided launch harness -- reported honestly.
1 parent 27a520b commit 39e8875

1 file changed

Lines changed: 125 additions & 0 deletions

File tree

‎docs/TRITON-FRONTEND-ROUTE.md‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
# Tri-Dao mamba_ssm backward Triton kernels routed through OUR tilelang stack
2+
3+
Status: 3 of 7 backward kernels route end-to-end through
4+
`tilelang triton frontend -> tilelang lower -> tvm CUDA codegen -> sm_121`
5+
to runnable tensor-core (HMMA / TF32) machine code on GB10. The other 4
6+
break at precisely-named, deeper boundaries (no silent fallback; RULE #1).
7+
8+
Measured on GB10 (sm_121a), venv `/home/dave/cppmega-venv`, tilelang from
9+
source `/home/dave/source/tilelang`. Date 2026-06-06.
10+
11+
## The route (verbatim path)
12+
13+
```
14+
@triton.jit kernel (mamba_ssm.ops.triton.ssd_chunk_scan / ssd_chunk_state)
15+
-> Triton compile -> TTIR (custom tt.* dialect text)
16+
-> poc.triton_frontend.triton_native_parse.parse_ttir_via_triton
17+
(Triton's OWN libtriton.ir validates + canonically re-prints TTIR;
18+
we parse that text into mlir.ir-shaped adapter ops)
19+
-> poc.triton_frontend.from_ttir (TTIR -> tvm.tir.PrimFunc)
20+
-> tilelang.engine.lower / tilelang.compile -> tvm CUDA codegen
21+
-> nvcc/ptxas -> sm_121a cubin (HMMA.1688.F32.TF32 SASS)
22+
```
23+
24+
Entry points used:
25+
- `tilelang.frontends.triton.compile_ttir(ttir, name=..., target="cuda")`
26+
-> `JITKernel` (the production wrapper; dispatches to `from_ttir`).
27+
- `poc.triton_frontend.from_ttir(...)` -> `PrimFunc` (lower-only path).
28+
29+
## Which kernels route (3/7)
30+
31+
| kernel | stage reached | CUDA src | tensor-core |
32+
|---|---|---|---|
33+
| `_chunk_scan_bwd_dstates` | runnable JITKernel | 33586 B | 32x HMMA.1688.F32.TF32 |
34+
| `_chunk_scan_bwd_dc` | runnable JITKernel | 34128 B | mma_sync |
35+
| `_chunk_scan_bwd_dcb` | runnable JITKernel | 47676 B | mma_sync |
36+
37+
### Proof it is OUR stack (tvm), not a native-triton fallback
38+
39+
`compile_ttir(_chunk_scan_bwd_dstates)` -> JITKernel -> `export_ptx` / `export_sass`:
40+
41+
- PTX header: `.version 9.3`, **`.target sm_121a`** (GB10 native).
42+
- PTX: **32 `mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32`** + 8 `ldmatrix`.
43+
- SASS (ptxas): **32 `HMMA.1688.F32.TF32`** + 8 `LDSM.16.M88.4`.
44+
- Generated `.cu` includes `tl_templates/cuda/instruction/mma.h`,
45+
`tl_templates/cuda/gemm.h` (TileLang/TVM template headers); 0 triton tokens.
46+
47+
That chain (tvm CUDA codegen -> ptxas HMMA for sm_121) cannot be produced by
48+
the native Triton runtime; it is our tilelang->tvm path.
49+
50+
## Which kernels do NOT route yet (4/7) and exactly where
51+
52+
| kernel | break stage | precise cause |
53+
|---|---|---|
54+
| `_chunk_scan_bwd_dx` | tilelang_lower | GEMM `m_warp * n_warp == num_warps` ICHECK (CUDA op/gemm.cc) fails for the reduce+atomic-surrounded dot |
55+
| `_chunk_state_bwd_ddAcs_stable` | tilelang_lower | same `m_warp * n_warp == num_warps` ICHECK |
56+
| `_chunk_state_bwd_db` | tilelang_lower | `variables (arg4,) are used, but are not passed in as API arguments` -- a scalar-load of `dA_cumsum_ptr` re-declares a duplicate `arg4` buffer (shape `(1,)`) that aliases the param buffer instead of reusing it |
57+
| `_chunk_state_bwd_dx` | from_ttir | `Can't cast a handle to other types` during emission (pointer/handle used in a scalar arithmetic context) |
58+
59+
All four FAIL LOUD at the named boundary -- no degraded/empty kernel is emitted.
60+
61+
## OUR-stack fixes made (this round)
62+
63+
All in `/Volumes/external/sources/tilelang/poc/triton_frontend/triton_native_parse.py`:
64+
65+
1. **Generic-form inline region parsing for `tt.reduce` / `tt.scan` combiners.**
66+
Triton prints the combiner region generically:
67+
`%0 = "tt.reduce"(%in) <{axis = 1 : i32}> ({ ^bb0(%a,%b): arith.addf ... tt.reduce.return ... }) : (...) -> R`.
68+
The text parser previously RAISED at the closing `}) : (...) -> ...` line
69+
("hit the close of a GENERIC-FORM inline region"). Added:
70+
- `opens_generic_region` detection (header ends with `({`) distinct from the
71+
scf bare-`{` region opener.
72+
- `_build_generic_region`: consumes `^bbN(args):` block args + inner ops into
73+
`op.regions[].blocks[].operations[]` (the exact structure
74+
`op_emitters.reduction._detect_via_mlir` walks to find the `arith.*`
75+
combiner), and parses the closing `}) : (...) -> R` line for the result type
76+
(which is NOT on the header), rewriting the result Value's type via the new
77+
`_Value.set_type`.
78+
- `_parse_props_attrs` for the `<{axis = ...}>` properties block (surfaces
79+
`axis` to the emitter).
80+
- `_OPNAME_RE` broadened to match multi-dotted op names (`tt.reduce.return`).
81+
82+
2. **`tt.atomic_rmw` custom-form RMW op keyword.** TTIR prints
83+
`tt.atomic_rmw fadd, acq_rel, gpu, %ptr, %val, %mask : ...`; the leading
84+
`fadd` keyword is now mapped to the `rmw_op` attribute the
85+
`map_tt_atomic_rmw` emitter requires (added `tt.atomic_rmw -> rmw_op` to
86+
`_POSITIONAL_KEYWORD_ATTRS`).
87+
88+
Effect: the 4 non-routing kernels moved PAST the TTIR-parse barrier (reduce
89+
region + atomic_rmw now parse) and now break deeper in tilelang lowering /
90+
emission -- a strictly more advanced boundary than before.
91+
92+
Prior-round fixes (commit 5cfb5bc2 / gb10 b0004c30) still in force:
93+
NEW `triton_native_parse.py` provider; `from_ttir` RAISES instead of emitting a
94+
`T.evaluate(0)` stub; `map_tt_dot` allocates the CUDA MMA accumulator C in
95+
`local.fragment`; `_emit_region` propagates `ctx.target`.
96+
97+
## Measured ms (GB10)
98+
99+
- Native Tri-Dao `mamba_chunk_scan_combined` full fwd+bwd (b1 s2048 h8 d64 n64
100+
chunk256): **3.78 ms** on GB10 (the "10ms-class" reference; faster here for
101+
this size on GB10). vs path_c v1 905 ms.
102+
- Routed kernels: COMPILE to runnable sm_121 tensor-core JITKernels
103+
(verified PTX+SASS). A *numerically-correct timed run* of the routed kernels
104+
is BLOCKED on (a) the PtrAnalysis C++ shim being unbuilt on GB10 -- the
105+
frontend uses the MVP scalar-load path, which does not reproduce Tri-Dao's
106+
strided tile addressing -- and (b) the full Tri-Dao strided launch harness.
107+
A degenerate (zero-stride) launch core-dumps (illegal address), confirming
108+
the addressing gap. We report this honestly rather than fabricate a number.
109+
110+
## To finish the remaining 4 + get parity/ms
111+
112+
- Build the PtrAnalysis C++ shim on GB10 (`_triton_frontend_cxx`) in a fresh
113+
process (libtriton-LLVM double-registration aborts otherwise) so tile loads
114+
use the real multi-element pointer analysis -> enables correct numeric parity.
115+
- `_chunk_state_bwd_db`: fix the scalar-pointer-load buffer aliasing in
116+
`op_emitters/memory.py` (`_redeclare_ctx_buffer_1d` must rewrite ALL prior
117+
references, including ones already emitted into statements, or the scalar load
118+
must reuse `ctx.buffers[key]`).
119+
- `_chunk_scan_bwd_dx` / `_chunk_state_bwd_ddAcs_stable`: the GEMM warp-partition
120+
ICHECK in `src/backend/cuda/op/gemm.cc` (`isSquare` yields `best_m=best_n=1`
121+
when no `m*n==num_warps` partition keeps `m_per_warp,n_per_warp >= 1`).
122+
Needs a Python-side GEMM policy/num_warps choice that factors the tile, or a
123+
C++ fallback in the partitioner (a tilelang rebuild).
124+
- `_chunk_state_bwd_dx`: trace the handle-cast in emission (a `!tt.ptr` used in a
125+
scalar arith op needs an address-of/load instead of a value cast).

0 commit comments

Comments
 (0)