Skip to content

Commit 8f5a2db

Browse files
apstenku123apstenku123
committed
docs(triton-fold): MEASURED-from-HEAD addressing-fold report (ms+SASS vs Triton)
Reproduced FULL transform-1 Coalesce-style addressing fold from tilelang HEAD 67ab579a (== gb10 f5640a3a), fresh OFF/OPT PrimFunc regen + interleaved A/B on GB10 sm_121a: - §P1 dstates: OFF/ON MAXDIFF 4.88e-04 ALLCLOSE, OFF-vs-ON byte-identical (0.0) - small real-strided multi-K-trip (strided dout/C): MAXDIFF 3.05e-05 PASS - median ms/kernel (CUDA events, 6 interleaved rounds): OFF 1467.97 -> ON 1103.26 (1.331x, -364.71ms); native 1.11844ms (remaining ~986x) - SASS sm_121a: local spill 658 -> 306 (-53%); LDGSTS 0 (cp.async not yet); HMMA 32 (== OFF, GEMM tensor-core byte-intact) vs Triton 256 - routed family codegen A/B: 5/7 OK both gates, mma=2 preserved; dcb/dx FAIL pre-existing+identical OFF==ON (not introduced by fold) GO (partial): fold correct+reproducible; remaining gap is LDGSTS coalescing + HMMA under-util, not the eliminated addressing arrays.
1 parent c64b6ca commit 8f5a2db

1 file changed

Lines changed: 122 additions & 0 deletions

File tree

‎docs/TRITON-ADDRESSING-FOLD.md‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# Triton Addressing-Fold (FULL transform-1, Coalesce-style)
2+
3+
Kill the spilled `[64]`/`[2048]`/`[4096]` index/mask arrays in the routed-Triton
4+
prologue by folding addressing/mask tiles into the strided `T.copy` load body
5+
instead of materializing them in any scope.
6+
7+
- tilelang HEAD: `67ab579a` (local `merge/upstream-codegen-reorg`) == gb10 `f5640a3a`
8+
(six fold files byte-identical between the two: `mlir_walker.py`, `op_mapping.py`,
9+
`__init__.py`, `op_emitters/{arith,memory,control}.py`).
10+
- Gate: `from_ttir(..., prologue_opt=True)` (default ON for the routed path;
11+
`prologue_opt=False` reproduces the pre-fold serial prologue for A/B).
12+
- Hardware: NVIDIA GB10, `sm_121a`, aarch64-linux.
13+
- All numbers below are REGENERATED FRESH FROM HEAD: both fold-OFF and fold-ON
14+
PrimFunc JSONs are re-emitted from the `_chunk_scan_bwd_dstates` TTIR via
15+
`from_ttir(prologue_opt=False|True)`, then compiled and run interleaved.
16+
Fold-OFF JSON = 497670 bytes, fold-ON JSON = 270624 bytes (the materialized
17+
index/mask arrays disappear from the serialized PrimFunc).
18+
19+
## What folds
20+
21+
`build_addressing_fold_set` (`poc/triton_frontend/mlir_walker.py:582`) reconstructs
22+
the MLIR use-def graph by **Value identity** (the native parser leaves
23+
`value.uses=None` and prints the same value differently as operand vs result, so a
24+
name-keyed map silently fails to connect a load's mask operand back to its
25+
`arith.cmpi`). A producer result is fold-eligible by fixpoint iff EVERY use is
26+
either a load/store addressing/mask SINK or an addressing PRODUCER whose own
27+
results are eligible. `scf.for`/`scf.yield` are transparent for non-accumulator
28+
(pointer/int) loop-carried slots; the f32 accumulator slot and any `tt.dot` /
29+
store-VALUE / reduce / return use disqualify (RULE #1: never fold a data tile).
30+
31+
Eligible `tile_binop` / `bcast` / `make_range` / `expand_dims` / bool-mask tiles
32+
are kept as `LazyTileExpr` and consumed per-lane inside the copy loop body
33+
(`_resolve_lane_operand` / `_scalarize_tile_index_base` / `_read_lane` /
34+
strided-store rhs already read a `LazyTileExpr` per-lane). The arrays are never
35+
materialized -> no local spill, no shared overflow, nothing to thread-distribute.
36+
The cooperative GEMM half (`dot_a_*`/`ptx_ldmatrix`/`mma_sync<TF32,...>`) is untouched.
37+
38+
## Correctness (§P1 + small real-strided multi-K-trip)
39+
40+
Controlled A/B, native Triton `_chunk_scan_bwd_dstates_kernel` as reference.
41+
42+
### §P1 production config (b1 nh112 hd64 ds64 nc64 cs64, grid (1,64,112), numel 29,360,128, multi-K-trip)
43+
44+
| build | MAXDIFF vs native | ALLCLOSE 1e-3 |
45+
|-------|-------------------|---------------|
46+
| fold OFF | 4.882812e-04 | PASS |
47+
| fold ON | 4.882812e-04 | PASS |
48+
| OFF vs ON | **0.000000e+00 (byte-identical)** | — |
49+
50+
### Small real-strided multi-K-trip (b1 nh4 hd64 ds64 nc4 cs64, 2 K-trips)
51+
52+
Real non-default strides via padded-then-sliced tensors:
53+
`dout.stride=[73728,288,72,1]`, `C.stride=[36864,144,72,1]` (note 72 != hd=64).
54+
55+
| build | MAXDIFF vs native | ALLCLOSE 1e-3 |
56+
|-------|-------------------|---------------|
57+
| fold ON | 3.051758e-05 | PASS |
58+
59+
## Performance (CUDA events, N=50/round, 6 interleaved OFF/OPT rounds)
60+
61+
| build | median ms/kernel |
62+
|-------|------------------|
63+
| fold OFF (pre-fold serial prologue) | **1467.97** |
64+
| fold ON (addressing-fold) | **1103.26** |
65+
| native Triton (reference) | **1.11844** |
66+
67+
- Fold speedup: **1.331x**, **-364.71 ms** (per-round spread <2 ms; very stable).
68+
- Remaining gap to native: OPT/native = **~986x** (was ~1312x at fold-OFF).
69+
- The drop is real but the kernel is still ~1000x off native: the residual cost
70+
is the un-coalesced loads (LDGSTS=0) + the tensor-core under-utilization
71+
(HMMA 32 vs Triton 256), NOT the prologue arrays anymore.
72+
73+
## SASS (cuobjdump `sm_121a`, regenerated from HEAD)
74+
75+
| metric | fold OFF | fold ON | Triton (ref, prior MEASURED) |
76+
|--------|----------|---------|------------------------------|
77+
| local spill STL+LDL | 658 (412+246) | **306 (168+138)** | 0 |
78+
| coalesced LDGSTS | 0 | 0 | 75 |
79+
| HMMA | 32 | 32 | 256 |
80+
| total instr (this counter) | 17528 | 8932 | 1981 |
81+
82+
- **Local spill halved: 658 -> 306 (-53%).** The fold removes the index/mask
83+
arrays that were spilling; the residual 306 is the GEMM accumulator/fragment
84+
traffic, not addressing arrays.
85+
- HMMA unchanged (32) OFF==ON: the tensor-core GEMM path is byte-intact through
86+
the fold.
87+
88+
## Other-paths no-regression (routed family, codegen A/B OFF vs ON)
89+
90+
| kernel | OFF | ON | note |
91+
|--------|-----|----|------|
92+
| `_chunk_scan_bwd_dstates` | OK len37174 mma2 | OK len22883 mma2 | folded (shorter, GEMM intact) |
93+
| `_chunk_scan_bwd_dc` | OK len38271 mma2 | OK len23943 mma2 | folded |
94+
| `_chunk_state_bwd_db` | OK len42823 mma2 | OK len27468 mma2 | folded |
95+
| `_chunk_state_bwd_dx` | OK len54284 mma2 | OK len39413 mma2 | folded |
96+
| `_chunk_state_bwd_ddAcs_stable` | OK len42922 mma2 | OK len31133 mma2 | folded |
97+
| `_chunk_scan_bwd_dcb` | FAIL (PtrState unresolved SSA) | FAIL (same) | **pre-existing at HEAD, identical OFF==ON** |
98+
| `_chunk_scan_bwd_dx` | FAIL (undefined.size()==0) | FAIL (same) | **pre-existing at HEAD, identical OFF==ON** |
99+
100+
5/7 codegen OK in both gates; `mma=2` preserved in every OK kernel (GEMM
101+
cooperative half intact); the 2 failures are pre-existing and identical in both
102+
gates (NOT introduced by the fold). Native dstates baseline reproduced at
103+
1.11844 ms.
104+
105+
## Honest GO / NO-GO
106+
107+
**GO (partial, as designed).** The fold landed and is correct + reproducible from
108+
HEAD: byte-identical output OFF==ON, real 1.331x / -365 ms speedup, local spill
109+
658 -> 306. What did NOT land in this transform (explicitly out of scope here and
110+
flagged for follow-up):
111+
112+
- **LDGSTS still 0.** Coalesced `cp.async` global->shared is a separate rewrite
113+
(the loads still go through per-lane strided `T.copy`, not a coalesced LDGSTS).
114+
This is the bulk of the remaining ~986x gap.
115+
- **HMMA still 32 vs Triton 256.** Tensor-core under-utilization is independent of
116+
the prologue fold.
117+
- **Spill 306 not 0.** The residual is GEMM accumulator/fragment traffic, not the
118+
eliminated addressing arrays.
119+
120+
RULE #1 honored: the fold either applies or the kernel keeps the intact GEMM
121+
cooperative path; no silent serial fallback was introduced; the 2 routed failures
122+
are pre-existing and surfaced (not papered over).

0 commit comments

Comments
 (0)