Skip to content

Commit 6558907

Browse files
apstenku123apstenku123
committed
docs: TILELANG-AUTO-MONO-FUSION — measured 2-kernel auto-fusion on Apple GPU
Honest writeup of the auto_fuse_chunk_region AUTO-MONO-FUSION pass + the MEASURED Apple-GPU bench: detect+z3-prove+dispatch a producer-consumer region, auto-fuse 2 kernels into ONE smem-resident Metal kernel. REAL: 1.45-1.96x faster (1 dispatch vs 2), bit-exact parity (0.0) every element. DEFERRED (honest): full stateful mamba3 F0/F1/F2 SSD chunk-carry mono-synthesis -> hand-written mono builder. Records the two engineering walls (mx.fast bridge mislaunch; torch.mps fp16-fragment->half4 cast bug) and how each was handled.
1 parent ab8332b commit 6558907

1 file changed

Lines changed: 123 additions & 0 deletions

File tree

‎docs/TILELANG-AUTO-MONO-FUSION.md‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
# TileLang AUTO-MONO-FUSION — multi-kernel chunk region -> ONE kernel
2+
3+
Extension of the `AutoGemmifyReductions` auto-GEMM pass (docs/APPLE-METAL-GEMM-
4+
AUTOPASS-Z3.md, commit 5b30bb10) ONE level up: at the **dataflow-region** level.
5+
Where the auto-GEMM pass detects a single serial reduction and proves it is a
6+
`T.gemm`, this pass detects a multi-kernel producer-consumer region that
7+
round-trips an intermediate through GLOBAL memory and AUTO-FUSES it into ONE
8+
persistent kernel that keeps the intermediate RESIDENT in shared memory.
9+
10+
Files:
11+
- `/Volumes/external/sources/tilelang/tilelang/transform/auto_fuse_chunk_region.py`
12+
— detector + z3 fusion-safety prover + dispatcher (commit a576c9d6).
13+
- `/Volumes/external/sources/tilelang/testing/python/transform/test_auto_fuse_chunk_region.py`
14+
— 17 analysis tests (graph/proof/decline). **17/17 PASS**.
15+
- `/Volumes/external/sources/tilelang/testing/python/transform/bench_auto_fuse_metal.py`
16+
— the COMPILE+RUN harness that actually fuses + measures on the Apple GPU.
17+
18+
Environment of record: this M-series Mac, macOS arm64, tilelang built+loadable
19+
from `/Volumes/external/sources/tilelang/build` (the live `build/` is UNTOUCHED —
20+
the pass is a Python TIR/graph module, no C++ rebuild). z3 4.15.4.
21+
22+
---
23+
24+
## What the pass does (three stages, mirroring the auto-GEMM pass)
25+
26+
1. **DETECT** (`match_fusion_region`). Reuses the production `path_c_fusion`
27+
dataflow graph — `_infer_edges` producer-consumer matching + Kahn topo-sort
28+
`_nodes_in_dependency_order` (imports the real `cppmega_mlx.runtime` module
29+
when present, else a byte-identical local twin). Recognizes a connected
30+
linear/tree producer-consumer chain, classifies INTERNAL (privatizable)
31+
buffers vs region inputs/outputs, and flags ESCAPING buffers.
32+
33+
2. **PROVE** (`prove_fusion`, z3). Extends the auto-GEMM z3 obligation set with
34+
the NEW fusion-safety obligations, each checked NON-VACUOUSLY over the actual
35+
`nchunks` x `state_cells` extents (negate the property, require UNSAT):
36+
- **privatization** — single-writer + single-reader per (chunk, cell), so
37+
promoting the buffer global->smem introduces no cross-threadgroup race;
38+
- **carry-domination** — producer write[c] precedes consumer read[c] in the
39+
fused serial schedule;
40+
- **escape** — no internal buffer is also a region output.
41+
z3 disabled / unavailable / UNKNOWN / SAT -> DECLINE (never fuse without a
42+
proof). Non-vacuity is locked by the disabled-z3 and degenerate-extent tests.
43+
44+
3. **DISPATCH** (`dispatch_region`). On ACCEPT (proved + in-budget + GEMMs kept +
45+
nothing escapes) selects ONE mono kernel to replace the N per-node builds. On
46+
DECLINE leaves the region MULTI-KERNEL (RULE #1, no wrong fusion).
47+
48+
---
49+
50+
## MEASURED on the Apple GPU — the auto-fusion ACTUALLY fires + wins
51+
52+
`bench_auto_fuse_metal.py` closes the loop end-to-end. It builds the smallest
53+
demonstrable producer-consumer GEMM chain that round-trips an intermediate
54+
through GLOBAL memory (the structural skeleton of the mamba3 F0->F1 chunk
55+
hand-off):
56+
57+
```
58+
K_A: Y = X @ W1 (producer — writes Y to GLOBAL)
59+
K_B: Z = Y @ W2 (consumer — re-reads Y from GLOBAL)
60+
```
61+
62+
The harness first consults the pass (`dispatch_region`). It returns
63+
`fused=True replaced=('K_A','K_B')`, privatizing `Y`. Only then does it emit the
64+
fused kernel. The fused kernel keeps `Y` RESIDENT in shared memory across both
65+
`T.gemm` calls (no second metal dispatch, Y never hits global), and **both GEMMs
66+
stay `T.gemm`** (`keeps_gemms=True`) — the cppmega-class recipe, NOT a scalar
67+
fusion.
68+
69+
Compiled via `tilelang.compile(target="metal")`, run on the Apple GPU through the
70+
native torch.mps boundary, 100-iter timing, M=K=N=P=64, fp32:
71+
72+
| path | metal dispatches | ms (typical) | max\|abs vs ref\| |
73+
|------------------------------|-----------------:|-------------:|------------------:|
74+
| MULTI-KERNEL (K_A + K_B) | 2 | 0.079–0.108 | 0.000e+00 |
75+
| **AUTO-FUSED (ONE kernel)** | **1** | **~0.055** | **0.000e+00** |
76+
77+
- **AUTO-fused 2 kernels -> 1** (pass-decided, z3-proved, Y privatized to smem).
78+
- **Speedup multi/fused = 1.45x – 1.96x** across runs (fused is rock-steady at
79+
~0.055 ms; the multi-kernel time varies with the second-dispatch + global-Y
80+
round-trip overhead it removes).
81+
- **Parity bit-exact over EVERY element**: fused-vs-multi = 0.000e+00 and
82+
fused-vs-ref = 0.000e+00. The two paths are algorithmically identical (Y in
83+
fp32), differing ONLY in whether Y lives in global (multi) or smem (fused).
84+
85+
This is a genuine AUTO-fusion win: the pass removed one of two metal dispatches
86+
and the inter-kernel global round-trip of `Y`, with zero numeric change.
87+
88+
---
89+
90+
## Honest scope (RULE #1) — prototype vs general
91+
92+
| claim | status |
93+
|-------|--------|
94+
| detector recognizes a producer-consumer region + topo-sorts + classifies internal/escaping buffers | **REAL** prototype, 17/17 tests |
95+
| z3 fusion-safety proof (privatization + carry-domination + escape), non-vacuous | **REAL** — declines on disabled-z3 / degenerate extents / dropped GEMMs / escape / over-budget |
96+
| AUTO-fuse a **2-kernel** producer-consumer region into ONE compiled Metal kernel, keep the intermediate smem-resident, measure on the Apple GPU | **REAL** — 1.45–1.96x, bit-exact parity, 1 dispatch vs 2 |
97+
| general N-kernel stateful **mamba3 F0/F1/F2 SSD chunk** mono-fusion with the full `state[headdim,dstate]` chunk-axis carry | **DEFERRED** — dispatcher selects the hand-written mono builder name; it does NOT auto-synthesize the full stateful chunk-carry TIR. The 2-kernel GEMM-chain bench is the demonstrated bar. |
98+
99+
The PROTOTYPE BAR (a working 2-kernel auto-fusion + parity) is MET and MEASURED.
100+
The full stateful SSD chunk mono-kernel auto-synthesis is honestly DEFERRED to the
101+
proven hand-written builder — the same honest deferral the auto-GEMM pass makes
102+
(`_emit_gemm_for_match` returns None rather than fabricating a brittle raw-TIR
103+
splice). On a region it cannot safely fuse (backward B0/B2 — B2 per-contraction
104+
GEMM is a measured 0.749x NO-GO, B0 is a reverse-cumsum scatter), the pass
105+
DECLINES and leaves it multi-kernel (`decline_reason=fused_body_drops_gemms_perf_nogo`).
106+
107+
### Two engineering walls hit + how they were handled (honest)
108+
109+
1. The `mx.fast.metal_kernel` MSL-rewrite bridge (`_mlx_runtime.wrap_tilelang_metal_kernel`)
110+
does NOT correctly launch an arbitrary compiled `T.gemm` Metal kernel — it
111+
produced near-zero output (verified: producer GEMM in isolation returned a
112+
~99.8%-zero buffer). The bench instead uses the native torch.mps boundary,
113+
where the fused kernel is **bit-exact** (diff 0.0). Reported honestly rather
114+
than shipping wrong-output timings.
115+
2. The native torch.mps adapter has an fp16 fragment->`half4` C-style-cast
116+
codegen bug (`SyntaxError: cast from simdgroup_float8x8 to half4`) when an
117+
fp32 accumulator fragment is copied to an fp16 shared/output buffer. The bench
118+
uses all-fp32 buffers, which sidesteps it and keeps the result bit-exact. (The
119+
cppmega F0 Track-A path avoids this by storing summary_states fp32.)
120+
121+
No C++ rebuild; the live `/Volumes/external/sources/tilelang/build` is untouched.
122+
The pass is default-OFF (env `TILELANG_ENABLE_AUTO_FUSE_CHUNK` /
123+
PassConfig `tl.auto_fuse_chunk_region`).

0 commit comments

Comments
 (0)