~/ting-lin
← back to blog

An Optimal Mapper Meets a Real GPU: Hand-Lowering AccelForge's Attention Mappings to Triton

2026-08-21
accelforgemappingfusionGPU kernelscost models

AccelForge's mapper, FFM (arXiv 2602.15166), does something no earlier mapper could: it searches the fused mapspace — where fusion is just a storage-placement decision in a LoopTree, not a graph rewrite — and provably returns optimal mappings, in roughly linear time, where genetic-algorithm and simulated-annealing mappers need ~15,000 CPU-hours to get within even 1–10% of optimal. Impressive, with one asterisk: optimal inside its cost model. I wanted to know what that optimality is worth on a real chip. So over the past weeks I took FFM's attention mappings for an RTX 4090 model, hand-wrote a faithful Triton implementation for each under a strict audit protocol, and measured. This post is the story of that audit — told through the two research documents where we recorded it: a manual lowering audit of forty candidate kernels, and a pairwise validation study that grew out of the wreckage.

Ground rules: what does "faithful" even mean?

Comparing a simulator's ranking against measured kernels is meaningless unless the kernels actually implement what the simulator ranked. So before any timing, every candidate has to pass two gates. Numerical correctness: the Triton kernels run on the GPU against PyTorch SDPA (FP16 in, FP32 accumulate, atol/rtol 2e-2). Static fidelity: a checker compares each kernel's declared pipeline against its exported LoopTree — the sequential split structure, which ranks are spatialized (a spatialized reduction rank means split-k), which tiles are honored. Any drift fails the candidate. Triton autotuning is off; manifests are frozen before measurement; every unavoidable refinement — a register-bounded 256-row tile refined to 128, an L2-persistence window driven through ctypes because Triton won't expose it — is written down, never silent.

The audit's first correction was terminological, and it matters. Every mapping FFM returned is a Sequential split with three branches — QK, softmax, AV — with intermediates held at modeled L2 across the split. The exporter used to label this fusion: fused. It is not single-kernel fusion; it is a three-stage pipeline whose handoff is charged to L2. The label is now kernel_fusion: not_implied_by_looptree_split. And here is the rub: "materialized at L2" is unrealizable in CUDA. L2 is a cache, not an addressable memory; no source code can guarantee an L2-only allocation. You either build a multi-kernel pipeline with globally addressed workspaces and hope cache hits approximate the model, or you write one fused kernel with register-resident scores — a different storage decision than the mapping made. Every lowering is a disclosed approximation; exact fidelity is not available.

Experiment one: ten mappings, forty handwritten kernels

The first attempt was a top-10 ranking study: four SDPA shapes, ten structurally distinct mappings each, one handwritten Triton file per mapping. The four cases (all FP16 in / FP32 accumulate, query length = key/value length):

Case Batch Heads Seq len Head dim Causal
A 1 8 128 64 no
B 1 16 512 64 yes
C 1 32 1024 128 yes
D 8 32 512 128 yes

All 40 passed both gates. Then we timed them — first on shared GPUs, then a full retest on exclusively assigned 4090s, 300 interleaved CUDA-event reps. The predicted ranking did not survive. The shape of that failure is visible directly in the measurements (Fig. 1 and 2 below).

Read B and C first: the model predicts a full tie, the hardware measures a 25–40× spread. The outliers are not noise — b-02 (512 one-key AV split-k slices, 2622 µs) and c-07 (tile-1 outer-product EV chunks, 15.8 ms) faithfully trace the spatial fragmentation of their LoopTrees. Case D, the only case with a real predicted ordering, inverts strongly: the predicted-fastest mapping (split-k over the QK reduction) measures slowest, and the predicted-slowest group — no split-k — measures 6–7× faster. The model's latency axis has no resolution inside structurally distinct strategy families.

Fig. 1 — predicted vs. measured, all forty candidates. Each point is one hand-lowered kernel; colors are the model's predicted-latency groups, and the dotted vertical line is the single latency the model assigns to each group (points are nudged horizontally only so they don't stack — the model really does give tied mappings one identical prediction). Cases B and C are full ties on the model axis yet spread 25–40× on hardware; case D slopes the wrong way (ρ = −0.81). Only each case's measured extremes are labeled; hover any point for its candidate id and both numbers.
Fig. 2 — case D, rank by rank. Left: the model's ordering (rank 1 = predicted fastest). Right: the hardware's. The red lines are the split-k candidates the model liked best; all three land at the bottom of the measured ranking, and the predicted-slowest group (no split-k) sweeps the top seven spots.

Four missing cost terms explain essentially all of it. (1) Split-k is free in the model: spatializing a reduction across SMs forces partial-sum materialization plus a cross-CTA reduction kernel — the dominant cost in d-01, unpriced. (2) Tile-1 spatial loops are free parallelism in the model; on hardware they are CTA fragmentation. (3) Softmax is an infinite-throughput elementwise copy in the workload spec — no max, no exp, no sum, no divide — so the mapper never sees the most schedule-sensitive part of attention, and every real kernel contains an unsearched softmax schedule. (4) L2-only intermediates don't exist, as above.

A reproducibility footnote that isn't a footnote

One more finding from this phase, because it changes how anyone should consume mapper output: the search is not deterministic in parallel mode. Two runs with the same hash seed return different mapping counts (958 vs 982 vs 975) and different top-10 strategy sets — the latency groups are stable, but which LoopTrees occupy them is a draw from a distribution. The root cause is in AccelForge, not our scripts: joblib returns pmapping groups in completion order, the concatenated row order feeds Pareto filtering, and the distinct-row filter keeps the first occurrence of model-equal rows — so which structurally distinct survivors you get depends on thread scheduling. Serial execution (set_n_parallel_jobs(1)) reproduces our checkpoint bit-for-bit. Every ranking claim in this post therefore comes with the discipline: manifests frozen from a serial run, hashes recorded.

The protocol grows up: one decision apart

The top-10 failure had a clear moral: don't ask the model to rank across strategy families — ask it the smallest directional question. The redesigned protocol picks two mappings that differ by exactly one decision (an added temporal l tile above the split; QK taking spatial s ownership instead of l) with a strict predicted latency gap. Everything else is held fixed: same pipeline shape, same fusion boundary, frozen manifests, handwritten gated lowerings, 300 interleaved rounds. The only question: does A beat B, as predicted?

On dense SDPA, all three pairs pass (A ≈ 114 µs; B = 228 / 303 / 273 µs against predicted gaps of 1.5× / 2.5× / 4.5×). Sparse top-k attention is where the boundary shows: pair 2 passes, pair 1 fails by 14× — and that failure is diagnostic gold. The model's #1-ranked mapping is the most pathologically fragmented of all 178: one program per score element (~21M tiny CTAs) plus a 16-way AV split-k writing ~1.4 GB of FP32 partials. The model ranks it 30% better than a conventionally blocked mapping that measures 14× faster. Among the non-pathological mappings, the ordering is correct (idx 154 vs 166: 2.66 ms vs 73.8 ms, confirmed). The rule across every pair we ran: the prediction holds exactly when the two mappings share every feature the model doesn't price. Pairs differing only in priced features pass; pairs differing in an unpriced feature invert.

Fig. 3 — the pairwise verdicts, mapped. Axes are the predicted gap and measured ratio (B/A, log–log); the dotted diagonal is magnitude agreement, the solid line the ordering boundary. The three dense pairs land on the correct side, near the diagonal. Sparse pair 1 — the pair that differed in unpriced features (split-k partial traffic, tile-1 fragmentation) — crosses the boundary: predicted 30% better, measured 14× slower.

The uncomfortable audit: is fusion even on the menu?

After the ranking experiments we audited the mapping language itself, and found something the benchmarks had been hiding: single-kernel, FlashAttention-style fusion is not representable in a LoopTree. Multi-Einsum nests always split into sequential branches — one compute node per leaf, and no code path interleaves two Einsums' loops in one nest. AccelForge's only fusion notion is weaker: shared outer loops with the intermediate parked in an on-chip storage node above the split — a pipeline with an on-chip handoff, not one kernel. And the mapper never returns even that. The join requires both branches to agree on the shared node and every loop above it; a Scores tile in SharedMemory usually blows the capacity budget; per-Einsum Pareto pruning rarely keeps SMEM-resident outputs. Every returned mapping hands off at L2 or MainMemory.

The objective makes this worse, and the mechanism is worth stating precisely. In this model, fusion never acts on latency directly — it only moves each intermediate's handoff level: DRAM, L2, SMEM, registers. Modeled latency is a max over component latencies dominated by the DRAM term, so once an intermediate fits anywhere on-chip, an L2 handoff is charged barely more than an SMEM one. Among on-chip levels, the latency axis is essentially flat; only energy distinguishes them, by orders of magnitude per access. Fusion is therefore searched along the energy axis — yet all our searches had used LATENCY | RESOURCE_USAGE, the one objective that structurally cannot see these decisions. Rerunning with ENERGY_DELAY_PRODUCT returned exactly one mapping — still an L2 handoff. And when we hand-authored a proper SMEM-handoff fused mapping (scores and weights in SharedMemory above the split, all branches sharing a spatial prefix) and asked the model to score it: 65× worse latency, 18× worse energy than the EDP winner — in the model's own currency.

Why does the model score a faithful fused mapping so badly? An SMEM handoff forces all three Einsums to share one spatial prefix, and the only ranks QK, softmax, and AV have in common are l and h — s can't be spatial above the split, because softmax needs the whole row in one CTA. So a fused mapping cannot split AV's reduction spatially — while the model prices AV's Spatial s split as a huge free win, with no split-k partial cost. Every unfused mapping that spatializes AV's reduction looks artificially better than every fused one. The unpriced split-k benefit that inverts our hardware rankings is the same gap that suppresses fusion inside the model's own objective. The model gap and the fusion gap are one gap.

So, can AccelForge help optimize kernels?

Honestly answered: it can, just not at the job we first gave it. As a ranking oracle for kernel engineers — "which of these schedules is faster on an SM" — it fails wherever the answer is decided by costs below its abstraction floor: split-k partial traffic, CTA fragmentation, softmax arithmetic, launch overhead, occupancy. And as a fusion proposer it is disarmed twice over: the language can't express single-kernel fusion, and the objective can't incentivize even the fusion it can express, as long as split-k is free. But the mapspace itself is not the problem. The winning designs are in it; pairwise orderings hold when pairs differ only in priced features; the latency groups it computes are stable across runs; and its energy-first view of fusion is the architecturally right question — how much buffer buys how much fusion, on hardware that doesn't exist yet.

The fix list is additive, and the audit tells us the order: price split-k partial write + re-read + reduction, price tile-1 spatial overhead, give softmax real reduction semantics, charge launch overhead — then worry about representing interleaved fusion (a deep redesign; the cheap 80% version is modeling fused overlap at evaluation time on top of sequential mappings). Every step on that list can be regression-tested with the pairwise protocol: one decision apart, frozen manifests, directional verdict. That protocol — not the simulator's absolute numbers — is what I'd trust to tell us when the last mile is actually closing.

All experiments: AccelForge rev 2495e289, RTX 4090, handwritten Triton candidates under the two-gate audit. Methodology and raw artifacts: docs/preliminary_research/ (dense_sdpa_manual_lowering_audit.md, topk_fusion_tiling_validation.md, pairwise_mapping_ordering_validation_goal.md) and preliminary_results/.

~