Skip to content

Commit c64b6ca

Browse files
apstenku123apstenku123
committed
docs(triton-prologue-opt): MEASURED §P1 prologue-opt results (transforms 1+3 land, no perf delta; T2 blocked on shared budget)
1 parent 294b515 commit c64b6ca

1 file changed

Lines changed: 144 additions & 0 deletions

File tree

‎docs/TRITON-PROLOGUE-OPT.md‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
# Triton Prologue Optimizer — routed Tri-Dao chunk kernels
2+
3+
Gated prologue optimizer for the routed-triton path in the TileLang triton
4+
frontend (`poc/triton_frontend`). Three transforms were scoped; this records
5+
which landed, with MEASURED numbers reproducible from committed HEAD on gb10
6+
(sm_121, aarch64-linux).
7+
8+
Source HEAD (tilelang): `29888fee` — "feat(triton-frontend): prologue-opt
9+
transforms 1+3 (drop int32-overflow guards, fold masks) + gated
10+
thread-distribute machinery". The gb10 source tree
11+
(`/home/dave/source/tilelang`) working copy of the 4 touched files is
12+
byte-identical (shasum) to this commit; the measurements below regenerate the
13+
PrimFunc + CUDA from HEAD and re-run.
14+
15+
## The gap (recap)
16+
17+
Routed `_chunk_scan_bwd_dstates` is CORRECT through our stack at §P1 but
18+
~1290x slower than native (≈1473 ms vs 1.15 ms/kernel). strip-and-time:
19+
~100% of the ms is the PROLOGUE; the cooperative GEMM half (T.copy/T.gemm,
20+
tensor-core, threadIdx-swizzled) is ~0 ms. The prologue is serial scalar
21+
index/mask arrays that every thread of every block rebuilds element-by-element.
22+
23+
## What landed
24+
25+
| Transform | Status | Evidence |
26+
|---|---|---|
27+
| (1) Fold address arithmetic into T.copy region | PARTIAL — mask/guard-chain folding only | `_emit_andi` folds `constant-true & x -> x` (gated). Deeper region-index folding NOT done. |
28+
| (2) Thread-distribute the prologue | IMPLEMENTED, gated OFF | machinery works (184 threadIdx when forced ON) but ptxas: too much shared data. |
29+
| (3) Drop int32-overflow guards | LANDED | `arith.cmpi addr<=INT32_MAX / INT32_MIN<=addr -> constant-true tile`. |
30+
31+
## MEASURED — IR / CUDA shape (regen from HEAD, gb10)
32+
33+
Generated CUDA for `_chunk_scan_bwd_dstates` (de-monomorphized §P1 PrimFunc,
34+
symbolic flat-arg extents):
35+
36+
| build | json bytes | .cu bytes | for-loops | int32 guards | threadIdx | mma |
37+
|---|---|---|---|---|---|---|
38+
| `prologue_opt=False` (baseline) | 497670 | 37174 | 158 | 28 | 25 | 2 |
39+
| `prologue_opt=True` (HEAD) | 332061 | 26685 | 107 | 0 | 25 | 2 |
40+
41+
Transforms 1+3 remove **51 for-loops and all 28 int32-overflow-guard
42+
references** (the `2147483647` / `-2147483648` literals + their bool[] AND
43+
loops). threadIdx count is unchanged (25, all inside the GEMM) — confirming
44+
transform 2 did NOT alter thread distribution in the shipped build.
45+
46+
## MEASURED — §P1 parity (gb10, CUDA_LAUNCH_BLOCKING)
47+
48+
`parity_prod_dstates.py`, real production config
49+
(b1 nh112 hd64 ds64 nc64 cs64, numel 29,360,128, grid (1,64,112)),
50+
2 K-trips/chunk (cs64 / BK32):
51+
52+
```
53+
MAXDIFF = 4.882812e-04 ALLCLOSE_1e-3 = True PASS
54+
```
55+
56+
Bit-identical to the pre-change baseline (the gate is a numeric no-op).
57+
58+
Small real-strided multi-K-trip (b1 nh8 s128 nc2, cs64/BK32 → 2 trips,
59+
vs native mamba_ssm kernel):
60+
61+
```
62+
MAXDIFF = 4.577637e-05 ALLCLOSE_1e-3 = True PASS
63+
```
64+
65+
## MEASURED — §P1 EXEC ms (routed, interleaved OFF/OPT/OFF/OPT)
66+
67+
```
68+
OFF (prologue_opt=False): 1468.21, 1477.25 ms mean ~1472.7
69+
OPT (prologue_opt=True): 1469.99, 1476.92 ms mean ~1473.5
70+
```
71+
72+
**Per-transform delta (1+3 together): +0.8 ms (~0.05%) — within run-to-run
73+
noise. No measurable §P1 speedup.** Ratio vs native 1.15 ms: ~1281x (was
74+
~1290x). The 51 dropped serial loops are cheap relative to the dominant cost;
75+
the remaining serial prologue (materialized [64]/[2048]/[4096] index+mask
76+
arrays and per-lane redundant address arithmetic) is untouched by 1+3. Only
77+
transform 2 would move the needle, and it is blocked (below).
78+
79+
## Transform 2 blocker (MEASURED, RULE #1 fail-closed)
80+
81+
Forcing `thread_distribute=True` (prologue_opt=True): the walker DOES
82+
thread-distribute — threadIdx jumps 25 → 184, the [64]/[4096] tiles become
83+
cooperative shared-fill loops with `__syncthreads`. But it promotes 63 prologue
84+
tiles to `__shared__` (multiple [4096] int64/int/bool arrays), and NVRTC/ptxas
85+
fails:
86+
87+
```
88+
ptxas error : Entry function '_chunk_scan_bwd_dstates_kernel' uses too much
89+
shared data (0x1b460 bytes, 0x18c00 max)
90+
```
91+
92+
111712 bytes/block vs the 101376-byte (99 KiB) sm_121 limit. Distribution
93+
requires shared scope (cooperative fill + barrier; a thread-local distributed
94+
write would leave 127/128 slots per lane uninitialized — forbidden). The
95+
correct fix is full transform (1): fold the addressing into the T.copy region
96+
indices/predicate so these tiles are never materialized in ANY scope; the small
97+
residue can then be thread-distributed in shared within budget. Until
98+
address-folding lands, distribution is gated OFF — it RAISES (ptxas error)
99+
rather than emit a racy local write or silently overflow. This is the honest
100+
blocker, not a silent fallback.
101+
102+
## MEASURED — other-paths no-regression
103+
104+
- **fla** (`fla_chunk_delta_h_real_ttir.mlir`): generated CUDA is
105+
BYTE-IDENTICAL with prologue_opt on/off (sha `581db03dbb91`, len 29515). The
106+
gate is a no-op — fla has no int32-overflow guards / constant-true ANDs.
107+
- **matmul** (M=N=K=64): the .cu DIFFERS (opt_ON 14727 B vs opt_OFF 22126 B —
108+
opt also folds matmul's addressing guards), but the OUTPUT is **bit-identical
109+
on vs off (maxdelta 0.000e+00)** and **EXEC ms is identical: 1.343 ms (opt)
110+
vs 1.343/1.346 ms (off)**. The .cu size delta is dead-code prologue arrays
111+
ptxas eliminates; the GEMM compute is unchanged. No regression.
112+
113+
Note: the IMPLEMENT-note claim "non-routed matmul/fla paths default OFF … and
114+
are untouched" is imprecise — any kernel routed through `from_ttir` gets
115+
`prologue_opt=True` by default, so the gate DOES fire on matmul's guards. But
116+
it is a proven numeric + perf no-op there (bit-identical output, identical ms).
117+
118+
## Gate
119+
120+
`from_ttir(prologue_opt=bool)` (default True). `prologue_opt=False` reproduces
121+
the 497670-byte / 37174-byte / 158-loop / 28-guard baseline exactly.
122+
`routed_triton_thread_distribute` (transform 2) is a distinct sub-gate, OFF by
123+
default and only set via `from_ttir(..., thread_distribute=True)` for dev.
124+
125+
## GO / NO-GO
126+
127+
**Partial GO.** Transforms (3) + (1-partial) LANDED, correct, committed, and
128+
§P1 bit-correct (MAXDIFF 4.88e-04) + small multi-K-trip (4.58e-05). They strip
129+
51 serial loops + 28 guards from the IR with ZERO §P1 perf regression and ZERO
130+
other-path regression. They do NOT yet move the §P1 ms (still ~1473 ms) because
131+
the dominant serial prologue survives. Transform (2) is implemented but blocked
132+
on shared-mem budget; landing the perf win requires full transform (1)
133+
region-index folding first. Honest: no §P1 speedup yet; the cleanup is real and
134+
safe; the speedup path is identified and gated correctly.
135+
136+
## Reproduce from HEAD (gb10)
137+
138+
```
139+
cd /home/dave/source/tilelang # working files == tilelang 29888fee (shasum-verified)
140+
source /home/dave/cppmega-venv/bin/activate
141+
# regen §P1 PrimFunc (prologue_opt=True default) + parity + timing:
142+
python poc/triton_frontend/_test_harness/tridao_parity/emit_pf_json.py _chunk_scan_bwd_dstates
143+
python poc/triton_frontend/_test_harness/tridao_parity/parity_prod_dstates.py
144+
```

0 commit comments

Comments
 (0)