You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[FlyDSL] Merge gfx1250 PA decode partitions in the tile kernel and speed up dense decode - #7079
gfx1250 PA decode writes per-partition partials and launches a second kernel to merge them. For MiniMax-M3 sparse decode (H16/KV1/D128, BF16 or FP8 page16 caches, 2048 top-k tokens per row) at 64 rows, that reducer takes about 2.9 µs plus a launch gap out of 11.5 µs. This PR merges partitions inside the tile kernel for the two common plan shapes, so those calls launch a single kernel. It also speeds up the dense tile kernels of MiniMax-M3's full-attention layers (page128 caches, decode and MTP4 verify) for FP8 and BF16 KV, and cuts launch and reducer latency for the plans that keep the reducer.
Partition merges
Multi-partition plans merge within a workgroup cluster. With P = min(partition bound, max_partitions) in [2, 16], the tile kernel launches P CTAs per (sequence, KV head) as one cluster (cluster=(P, 1, 1), grid batch * P). CTA i runs the sequence's task i from reduce_info, if it has one, and stores packed partials with plain stores. After a cluster-scope release, the cluster barrier and a cluster-scope acquire, each CTA merges 1/P of the sequence's rows. A cluster shares one XCC and its L2, so the release and acquire lower to s_wait_storecnt 0 and global_inv scope:SCOPE_SE; agent scope would add a device-wide L2 writeback and invalidate per CTA. Plans qualify when batch * P fits the plan capacity or two CTAs per CU, so CTAs idling at the barrier never displace working ones, and when batch * kv_heads stays within 15 clusters per 16 CUs: gfx1250 keeps at most 240 clusters resident on 256 CUs whatever their size, and from the 241st cluster on a second wave makes the merge about 1.5× slower than the reducer.
Plans with one task slot per sequence write final output directly, which covers sparse prefill. A sequence with one task stores its rows, the CTA indexed by an empty sequence clears its rows, and a sequence that owns several slots (freed by empty sequences) merges in its last-arriving CTA. Arrivals are counted in a new optional PADecodePlan.arrival_counters ([batch, kv_heads, 32] int32, one 128-byte line per counter), zeroed when the plan is created and reset by the last arrival, so CUDA-graph replays and plan refreshes need no host work. This merge uses device-scope (SCOPE_DEV) accesses.
Both merges reproduce the reducer's math and partition order. Sinks, head_dim > 512, partition bounds above 16, and gfx942/gfx950 keep the reducer. AITER_FLYDSL_PA_CLUSTER_MERGE=0 restores the reducer for cluster-eligible plans.
Multi-partition decode does not use a last-arriver merge: on gfx1250 each device-scope round trip costs about 4 µs, and merging BF16 sparse decode at 64 rows that way took 17.7 µs in the tile kernel, against 7.45 µs + 2.9 µs for tile and reducer.
Dense tile kernels
FP8: four D128 query tiles over page128 caches give each wave its own 16-row query tile. This is MTP4 at GQA 16, MiniMax-M3's verify step. The previous kernel split each page's tokens across waves for QK and the head dimension for PV, so every query tile exchanged maxima, sums and P through LDS: 12 barriers per 128-token page and 400 VGPRs. Now each wave computes QK over the whole page, and QK reads K rows in the token order that PV's P operand contracts, so P goes from the accumulator to the WMMA operand in registers. One barrier per page remains, for the double-buffered K/V, and the kernel uses 285 VGPRs. Interior pages skip the causal mask, fold the row scale and V normalization into one FMA in the exponent, and sum probabilities with a packed tree. The two row groups of each QK tile read opposite halves of their 16-token chunks, so K loads avoid LDS bank conflicts.
BF16: four D128 query tiles over page128 caches use wave pairs. Each pair shares two 16-row query tiles and each wave covers half of every page, so one K or V operand feeds two WMMAs. QK's accumulator rows already match PV's P-operand split, so P stays in registers, and a pair merges its two online-softmax states once at the end. A page needs one barrier instead of nine and 687 instructions instead of 1548, at 414 VGPRs. The kernel uses LLVM's iterative-ILP scheduler (a per-kernel compile hint), which halves its VGPR bank switches. It issues the second page's DMA right after the query loads. Sliding windows take the path too.
Both per-wave paths support all four output modes: reducer, one-partition direct output, last-arriver merge and cluster merge.
FP8 single-query-tile D128 page128 tiles (decode) double-buffer pages with asynchronous per-token scale copies, the existing path for D256 and multi-tile FP8, instead of one buffer with synchronous scale loads.
Both FP8 pipelines issue the first page's DMA before staging Q, so its latency overlaps Q quantization (the LDS regions are disjoint). This also speeds up sparse decode.
BF16 page128 and multi-tile FP8 pipelines drop the barrier between query tiles: the next tile's max barrier already orders its P and sum writes after every wave's reads.
Launch and reducer latency
Kernel arguments preload into SGPRs. Every kernel started with a dependent scalar load from the kernarg segment, and the tile kernel loaded its remaining arguments only after its work record arrived. The tile kernel's arguments are reordered so that everything the reducer-fed and direct-output specializations read fits in the 30 preloaded user SGPRs (120 bytes). The KV head count is passed explicitly instead of being read from the hidden grid size. The reducer's 108 argument bytes preload as well. Cluster and last-arriver merge kernels are not preloaded: their merge-only arguments don't fit, and partially preloaded arguments made sparse cluster merges 2–6% slower. Preloading follows AITER_FLYDSL_KERNARG_PRELOAD, as the MoE kernels do.
The reducer overlaps its loads. Above 32 live partitions, each group issues its partial loads before the weight barrier, so the statistics and partials arrive in one round trip (128 live partitions: 4.28 → 3.20 µs). Rungs above 16 partitions spread their loads over up to four partition groups, and group 0 still adds the weighted partials in partition order, so outputs match bit for bit (32 live partitions: 3.44 → 2.92 µs). D128 plans averaging at most 32 partitions per sequence now take four groups on gfx1250. gfx942/gfx950 compile to the same reducer ISA as before.
Performance
Measured on gfx1250 (256 CUs, 8 XCCs), FlyDSL 0.4.1, PyTorch 2.11.0+rocm10.1, HIP 7.16, with GPU runs serialized. MiniMax-M3 attention shapes: H16/KV1/D128. Dense rows use page128 caches with one sequence per request; verify rows have 4 causal query tokens. Sparse rows use page16 caches gathered from a 4096-page pool with 2048 top-k tokens per row, and prefill rows see min(row + 1, 2048) tokens. Times include decode and reduction, excluding planning and compilation. Each run takes the minimum of five 10-call CUDA graphs with 5 replays each; the table reports the median over three rounds with rotating tree order. The base is main at 0e53c46, which this branch merges.
Case
main
This PR
Speedup
FP8 dense decode, bs1 8K
9.38 µs
7.32 µs
1.28×
FP8 dense decode, bs4 8K
10.06 µs
7.95 µs
1.27×
FP8 dense decode, bs8 8K
10.79 µs
8.47 µs
1.27×
FP8 dense decode, bs32 8K
14.16 µs
11.30 µs
1.25×
FP8 dense decode, bs4 32K
12.30 µs
9.64 µs
1.28×
FP8 dense decode, bs3 220K
25.12 µs
19.23 µs
1.31×
FP8 dense verify, bs3×4 8K
13.90 µs
10.28 µs
1.35×
FP8 dense verify, bs3×4 32K
16.21 µs
12.47 µs
1.30×
FP8 dense verify, bs8×4 32K
21.12 µs
15.70 µs
1.35×
FP8 dense verify, bs3×4 220K
37.64 µs
26.99 µs
1.39×
FP8 sparse decode, 4 rows
9.45 µs
7.29 µs
1.30×
FP8 sparse decode, 64 rows
11.11 µs
8.32 µs
1.34×
FP8 sparse verify, 3×4 = 12 rows
9.65 µs
7.38 µs
1.31×
FP8 sparse prefill, 1024 rows
24.83 µs
16.54 µs
1.50×
FP8 sparse prefill, 8192 rows
199.41 µs
162.46 µs
1.23×
BF16 dense decode, bs1 8K
9.51 µs
7.60 µs
1.25×
BF16 dense decode, bs4 8K
10.13 µs
8.34 µs
1.21×
BF16 dense decode, bs8 8K
10.83 µs
9.39 µs
1.15×
BF16 dense decode, bs32 8K
15.30 µs
14.15 µs
1.08×
BF16 dense decode, bs4 32K
12.92 µs
11.11 µs
1.16×
BF16 dense decode, bs3 220K
27.92 µs
26.20 µs
1.07×
BF16 dense verify, bs3×4 8K
14.60 µs
9.29 µs
1.57×
BF16 dense verify, bs3×4 32K
18.44 µs
12.71 µs
1.45×
BF16 dense verify, bs8×4 32K
26.03 µs
16.96 µs
1.53×
BF16 dense verify, bs3×4 220K
48–63 µs*
29.33 µs
1.6–2.1×
BF16 sparse decode, 4 rows
9.55 µs
7.16 µs
1.33×
BF16 sparse decode, 64 rows
11.59 µs
8.91 µs
1.30×
BF16 sparse verify, 3×4 = 12 rows
9.92 µs
7.38 µs
1.34×
BF16 sparse prefill, 1024 rows
26.11 µs
17.84 µs
1.46×
BF16 sparse prefill, 4096 rows
131.56 µs
109.50 µs
1.20×
* main is bimodal on this case across runs; this PR holds 29.3–31 µs.
Outputs are bitwise identical to main except for FP8 and BF16 dense verify, where the per-wave paths sum the softmax denominator in a different order and fuse scales into the exponent. Their relative error against an FP32 reference is unchanged to four decimals (FP8 bs3×4 220K: 0.0369 on both). The FP8 bs3×4 220K verify row matches the MiniMax-M3 decode trace (3 requests × 4 tokens at about 221K context). At 64 sparse rows (rocprofv3 kernel trace, medians), the BF16 cluster kernel takes 8.21 µs where main runs a 7.17 µs tile kernel and a 2.92 µs reducer; FP8 takes 7.69 µs against 6.73 µs + 3.00 µs.
The table uses plan_pa_decode's default partition cap. ATOM's MiniMax-M3 sparse path instead caps each call at get_recommended_splits(rows, 1) and passes max_context_length=2048. With those plans, sparse decode is 0.74–0.78× of main's time at 4–64 rows (8 partitions) and 0.80–0.81× at 128 rows (4 partitions), BF16 and FP8 alike. 256 rows (2 partitions, 256 clusters) keeps the reducer, about 5% faster than on main. Sparse prefill already writes output directly on main with a one-partition cap and is unchanged.
ATOM's dense calls pass no max_context_length, so their capture-sized block tables compile the reducer for 256 partitions. With such tables (block tables padded to 256K tokens, median of three alternating rounds), dense decode is 1.25–1.32× faster than main for FP8 KV (bs1 8K: 9.80 → 7.59 µs) and 1.06–1.22× for BF16.
Ten new cases pin sparse decode at 8 partitions (BF16 and every FP8 variant, with CUDA-graph replays) and BF16 at 16 partitions to the cluster path. 13 existing cases already take it at 4 and 5 partitions.
57 FP8 and 7 BF16 cases run the per-wave dense paths in every output mode: reducer with sinks, one-partition direct output, last-arriver merge with CUDA-graph replays, and cluster merge. They also cover 60 rows that pad the last tile (MTP5 at GQA 12), two queries per tile (MTP8 at GQA 8), FP16 queries, a sliding window, probability tails, and varied per-token magnitudes that fail when a scale is misplaced. A spy confirms each case takes the intended path.
A CUDA-graph plan-refresh test replays new lengths through ATOM-style tables with NaN-poisoned scratch and output.
A reducer test checks split-group results bit for bit against single-group ones (head dims 64–512, with and without sinks) and against a float64 reference.
174 more cases, checked against a CPU float64 reference and against the reducer in the same process, all pass. They cover query lengths 1–4, pages 16/64/128, head dims 64–512, GQA 1–16, 1–16 KV heads, sliding windows, empty and short rows, max_partitions and workgroup_budget settings that give cluster sizes 2–16, BF16 and FP16 queries, FP8 per-token KV, CUDA-graph replays and plan refreshes, and one-slot plans that merge. Outputs are bitwise identical to main except the MTP4 page128 cases that now take the per-wave paths: one FP8 case (relative error 0.04201, main 0.04203) and seven BF16 cases (errors within 1e-4 of before).
Differential checks against main, with per-token K/V magnitudes scaled by e^N(0, 1), at query lengths 1–8 with GQA 8–64 and lengths 1 to 33000:
FP8: maximum relative difference 6.8e-4.
BF16 (including windows, FP16 queries and one-partition plans): 3.4e-3, i.e. 1–2 BF16 ulps, with error against the FP32 reference moving by at most 5.5e-5.
Deliberately misplacing V scales (FP8), shifting the causal bound by one, or breaking the wave-pair merge (BF16) gives differences of 0.13 or more.
Mixed-length rounding: in a mixed plan (one 220K sequence among fifteen 8K ones), the long sequence's 256 partitions now add in four blocks. That row differs from before by 2.4e-5 relative, with unchanged error against the FP32 reference.
Race stress: CUDA graphs that fill the partials and the output with NaN before every call. Every call matches the eager output bit for bit, and the sparse outputs equal main's.
Sparse: 20,000 calls each for BF16 64 rows (uniform, and mixed with empty rows), BF16 8 rows and FP8 64 rows.
Dense MTP4, FP8 and BF16: 10,000 calls each under the reducer, a 12-way cluster merge, a last-arriver merge and one-partition direct output.
Dense decode: 10,000 calls each for the reducer paths this PR changes (split rungs, four- and sixteen-group dual CTAs, the LDS branch, mixed plans, direct output).
Resident-cluster limit: with 2 and 4 partitions, 239 and 240 rows take the cluster path and 241, 244 and 256 rows take the reducer.
Black, Ruff and git diff --check pass.
Notes
The cluster merge relies on the gfx1250 memory model's cluster scope (a cluster shares an L2). The one-slot merge relies on device-scope (SCOPE_DEV) hardware semantics, as fused_compress_attn_hca_gfx1250.py does.
Decodes running concurrently must not share a plan that has arrival counters.
VGPRs: the one-slot merge variant of the sparse tile kernel uses 158, against 102. The FP8 per-wave dense kernel uses 285, against 400 for the multi-tile kernel it replaces, and the BF16 one uses 414, against 426.
…or one-slot-per-sequence plans
Plans whose capacity equals the batch (prefill-like and large batches) no
longer launch the partition reducer on gfx1250. A sequence with a single task
stores its final rows directly, the CTA indexed by an empty sequence clears
its rows, and sequences that still own several tasks (slots freed by empty
ones) merge them in their last-arriving CTA.
The merge counts arrivals in a new zero-initialized PADecodePlan field,
arrival_counters [batch, kv_heads, 32] int32 (one 128-byte line per counter),
allocated once with the plan and reset by the last arrival, so CUDA-graph
replays and plan refreshes keep working. Partials, statistics, the counter
and the merge readback use device-scope (SCOPE_DEV) accesses; agent-scope
fences would lower to a device-wide L2 writeback per CTA. The merge
reproduces the register reducer's math and order, so outputs are bitwise
identical.
Multi-partition decode plans keep the separate reducer: merging there costs
two device-scope round trips (BF16 sparse decode 64 rows: 17.7 us merged
kernel vs 7.45 us tile + 2.9 us reducer).
Co-authored-by: Cursor <cursoragent@cursor.com>
Multi-partition gfx1250 decode plans no longer launch the partition reducer
when every sequence's tasks fit one workgroup cluster. The tile kernel runs
P = min(partition bound, max_partitions) CTAs per (sequence, KV head) as one
cluster (P in [2, 16]): CTA i runs the sequence's task i from reduce_info
when it has one, stores packed partials with plain stores, and after a
cluster-scope release, the cluster barrier and a cluster-scope acquire every
CTA of the cluster merges 1/P of the sequence's rows with the reducer's
math and partition order, so outputs stay bitwise identical. A sequence's
single task stores final rows directly and empty sequences clear theirs.
A cluster shares one XCC and L2, so the release/acquire lower to
s_wait_storecnt 0 and global_inv SCOPE_SE without the L2 writeback that
made device-scope merging slow (17.7 us tile kernel at 64 rows).
The plan is unchanged; the grid becomes batch * P CTAs along x, and plans
qualify when that fits the plan capacity or two CTAs per CU, so idle CTAs
waiting at the barrier never displace working ones. Sinks, head_dim > 512,
larger partition bounds and plans with one slot per sequence keep their
paths. AITER_FLYDSL_PA_CLUSTER_MERGE=0 restores the reducer.
Against origin/main, BF16 sparse decode 64 rows (CUDA graph): 11.4 -> 9.0 us;
4-32 rows save 1.9-2.4 us; FP8 sparse decode 64 rows: 11.0 -> 8.6 us.
Co-authored-by: Cursor <cursoragent@cursor.com>
Pin sparse decode rows of up to 2048 page16 tokens (8 partitions, with graph
replays) and 4096-token rows (16 partitions, the widest cluster) to the
cluster-merge path on gfx1250; other architectures keep the reducer.
Co-authored-by: Cursor <cursoragent@cursor.com>
ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm
vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all
All standard extended tests (excludes ci:atom_full)
Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 7079 --add-label <label>
One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.
PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Supported dimensions and valid tensor views can violate cluster-merge load/store alignment contracts and the last-arriver path’s 8-byte atomic alignment requirement.
This PR reduces gfx1250 paged-attention decode launch overhead by merging eligible partition results inside the tile kernel instead of launching a separate reducer.
Changes:
Adds cluster-based merging and a last-arriver merge for one-slot-per-sequence plans.
Extends plan metadata and dispatch while retaining the reducer for unsupported cases.
Adds 8- and 16-partition coverage, including graph replay checks.
…uster limit
gfx1250 keeps at most 240 workgroup clusters resident on its 256 CUs,
whatever the cluster size: sparse decode with 2 or 4 partitions per row runs
at the same speed up to 240 rows and about 1.5x slower from 241 rows, when
the extra clusters wait for a second wave. ATOM caps 241-256 sparse rows at 2
partitions, so those batches ran 1.34x (BF16) to 1.40x (FP8) slower than the
reducer. Cluster merging now also requires sequences x KV heads <= 15 per 16
CUs; larger grids keep the reducer.
Co-authored-by: Cursor <cursoragent@cursor.com>
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Accepted scratch and output buffers may be only element-aligned, yet the merge path assumes 8-byte alignment for atomic and vector accesses, risking invalid memory operations.
The committed decode replay cases do not exercise the last-arriver counter reset, so stale counters or transitions between single-task and multi-task sequences could regress undetected. Add a gfx1250 replay/plan-refresh test with capacity equal to batch and mixed empty/nonempty sequences, asserting that at least one sequence receives multiple tasks and selects the last-arriver path. Check output correctness and zeroed arrival counters after each decode, including after refreshing the plan with changed lengths.
…-tile pages for dense decode
FP8 D128 page128 tiles with four 16-row query tiles (MTP4 at GQA 16) now give
each wave its own query tile. QK reads K rows in PV's contraction order, so P
stays in registers and softmax needs no cross-wave exchange: one barrier per
128-token page instead of twelve. Interior pages skip causal masks and fuse
the row scale and V normalization into the exponent.
Single-query-tile FP8 D128 page128 tiles (decode) now double-buffer pages with
asynchronous per-token scale copies, and both FP8 pipelines issue the first
page's DMA before staging Q. BF16 page128 and multi-tile FP8 pipelines drop the
barrier between query tiles that the next tile's max barrier already orders.
Co-authored-by: Cursor <cursoragent@cursor.com>
Rows 0-7 and 8-15 of each QK tile read opposite halves of their 16-token
chunks, so a K operand load no longer puts both row groups on the same LDS
banks. Half-wave 1 swaps its P pieces to match; scale and mask offsets follow
the same mapping.
Co-authored-by: Cursor <cursoragent@cursor.com>
…A merges
One task slot per sequence with two empty rows gives the long row three tasks,
so the per-wave FP8 query-tile path publishes partials and merges them in its
last-arriving CTA across graph replays.
Co-authored-by: Cursor <cursoragent@cursor.com>
zhiding512
changed the title
[FlyDSL] Merge gfx1250 PA decode partitions in the tile kernel
[FlyDSL] Merge gfx1250 PA decode partitions in the tile kernel and speed up dense FP8 decode
Oct 28, 2026
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Unresolved critical issues include invalid alignment assumptions for supported output and workspace layouts and a fused scaling order that can overflow on finite inputs.
BF16 D128 page128 tiles with four 16-row query tiles (MTP4 at GQA 16, the
MiniMax-M3 verify step) no longer split each page's tokens across waves and
exchange maxima and P through LDS (nine barriers per page). Wave pairs share
two query tiles and each wave covers half of every page, so one K or V operand
feeds two WMMAs. QK's accumulator rows are PV's P operand split, so P stays in
registers and each page needs only its buffer barrier. A pair merges its two
online-softmax states once, through the free page buffer. Interior pages skip
causal masks and fuse the softmax scale into the exponent; windowed tiles mask
their leading pages as well.
All four output modes take the path. New BF16 cases cover direct output,
last-arriver and cluster merges, 60 padded rows, two queries per tile, FP16
queries and a window spanning masked and unmasked pages.
Co-authored-by: Cursor <cursoragent@cursor.com>
…ve ILP
LLVM's iterative ILP scheduler cuts the per-wave BF16 page loop from 776 to
609 instructions, mostly VGPR bank switches (255 to 122 per page), at 414
VGPRs. Dense verify on MiniMax-M3 shapes is 1-4% faster in ABBA rounds;
outputs are unchanged. Other tile kernels keep the default scheduler.
Co-authored-by: Cursor <cursoragent@cursor.com>
Per-wave BF16 tiles now issue the second page's DMA right after the query
loads, so the first page's wait leaves it in flight instead of issuing it
after that wait. Masked pages also fuse a positive softmax scale into the
exponent, as unmasked pages do (masked scores stay -inf), cutting the masked
page loop from 901 to 802 instructions; only masked pages' rounding changes.
Dense verify on MiniMax-M3 shapes: 1.0-2.9% faster in ABBA rounds.
Co-authored-by: Cursor <cursoragent@cursor.com>
…iles and the reducer
Every gfx1250 PA decode kernel started with a scalar load of its first
pointer from the kernarg segment, and the tile kernel loaded its remaining
arguments only after its work record arrived, so two dependent kernarg round
trips preceded the page-table lookup. gfx1250 preloads kernel arguments into
up to 30 user SGPRs (120 bytes).
* The tile kernel's arguments are reordered so that everything the
reducer-fed and direct-output specializations read fits in the preloaded
bytes, and the KV head count is passed explicitly instead of being read
from the hidden grid size. Cluster and last-arriver merge kernels are not
preloaded: their merge-only arguments do not fit, and argument loads that a
preload leaves behind made sparse cluster merges 2-6% slower.
* The reducer's 108 argument bytes all preload, so its first load is
reduce_info.
* The BF16 page128 pipeline keeps its Q loads ahead of the first page's
table lookup.
Dense M3 decode (CUDA graph, 2 ABBA rounds vs 9773c35): BF16 bs1/4/8/32
8K 0.88/0.86/0.90/0.94, bs4 32K 0.92, bs3 220K 0.97; FP8 0.91/0.91/0.92/0.94,
0.91, 0.96. Verify cases 0.93-0.97, sparse cases unchanged. Outputs are
bitwise identical on all 36 M3 cases and the cluster/shape suites.
A new test replays dense page128 decode in a CUDA graph that refreshes its
plan with new lengths, for the reducer and direct-output paths.
Co-authored-by: Cursor <cursoragent@cursor.com>
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
The BF16 four-tile and FP8 fused paths can overflow exp2 for large finite inputs, while the documented cluster-merge opt-out is ineffective and plan-refresh regression coverage is missing.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The documented cluster-merge opt-out is ignored, and removing the planner regression leaves plan refreshes and arrival-counter state without test coverage.
bf16_wave_query_tiles also required head_dim == 128. Packed page128
already bounds D256 to 32 query rows, so four query tiles leave only D64
and D128: keep bf16_page128_packed and m_tiles == NWARP, as suggested in
the #7079 review.
The wave-pair pipeline assumed D128's 2048-byte chunk spacing for both
operands. Packed K keeps it (eight dims over a 128-token page), but
transposed V spaces its eight-token chunks 16 * head_dim bytes apart.
Pass the spacing to the operand loader so D64 reads the right V chunks;
D128 keeps the same offsets and its ISA is unchanged.
Add D64 variants of the four-tile BF16 page128 cases: direct, cluster,
last-arriver merge, sliding window, and padded MTP5 GQA12.
bf16_wave_query_tiles also required head_dim == 128. Packed page128
already bounds D256 to 32 query rows, so four query tiles leave only D64
and D128: keep bf16_page128_packed and m_tiles == NWARP, as suggested in
the #7079 review.
The wave-pair pipeline assumed D128's 2048-byte chunk spacing for both
operands. Packed K keeps it (eight dims over a 128-token page), but
transposed V spaces its eight-token chunks 16 * head_dim bytes apart.
Pass the spacing to the operand loader so D64 reads the right V chunks;
D128 keeps the same offsets and its ISA is unchanged.
Add D64 variants of the four-tile BF16 page128 cases: direct, cluster,
last-arriver merge, sliding window, and padded MTP5 GQA12.
rebklee
added a commit
to rebklee/aiter
that referenced
this pull request
Nov 3, 2026
Conflicts in traits.py, pa_decode_kernel.py, pa_decode.py and
test_flydsl_pa_decode.py: main's ROCm#7079 (gfx1250 partition merges)
added merge_partitions and cluster_merge where this PR adds causal
(Gfx1250Traits, the compile_pa_decode_tile and
compile_gfx1250_pa_decode signatures and the calls between them).
Kept both, ROCm#7079's first. Gfx1250Traits is built positionally, so the
call follows its field order.
The test file keeps both case lists. _check_pipe_policy keeps main's
cluster check and builds the PaDecodeSchedule.select kwargs once,
without cluster_merge, for the MTP4 fused check and the pipe check.
ROCm#7079's per-wave gfx1250 MTP tiles don't read causal; the next commit
handles them.
Signed-off-by: Rebecca Lee <rebecca.lee@amd.com>
rebklee
added a commit
to rebklee/aiter
that referenced
this pull request
Nov 3, 2026
ROCm#7079's per-wave gfx1250 tiles (D128, page 128, four 16-row query
tiles, such as MTP4 at GQA 16) bound every row causally and don't read
causal, so after the merge a causal=False call there compiled the
causal kernel. Non-causal blocks of more than one token now take the
shared gfx1250 tiles, which apply causal=False and its two-sided
window. One-token blocks, which resolve to non-causal, keep the
per-wave tiles: one row's causal bound is the whole context.
Causal kernels are unchanged. Checked compile-only for gfx1250 (no
hardware): the six non-causal per-wave shapes, BF16 and FP8, with and
without a window, now compile the shared tiles; before this change
their kernels were byte-identical to the causal ones. Every other
case's ISA is unchanged.
Signed-off-by: Rebecca Lee <rebecca.lee@amd.com>
Resolve conflicts with #6956 (AOT caches built by replaying the runtime)
and #7079 (gfx1250 PA decode partition merges).
- aiter/aot/flydsl/pa_decode_plan.py now follows the replay protocol
instead of the removed OpKind/compile_one_config dispatch: it replays
plan_pa_decode on FakeTensors at the largest batch of each bucket and
at each capacity where the block size can change, dense and windowed.
Each arch replays with the most CUs of its devices, so the set equals
the previous plan_variants minus the variants no device can reach
(e.g. bucket 1 at 1024 threads). common.py registers it in AOT_MODULES;
arch selection moves from --arch to GPU_ARCHS, as for every module.
- README: main's module table and GPU_ARCHS wording, plus the planner
entry and section.
- op_tests: #7079 dropped test_pa_decode_plan; keep this PR's version
alongside main's new tests.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Resolve conflicts with the gfx1250 in-kernel partition merges (#7079),
MXFP4 K/V (#7056) and the D64 BF16 wave-pair pipeline (#7384):
- Gfx1250Traits keeps context_partition_size alongside merge_partitions
and cluster_merge; the per-wave FP8 and BF16 pipelines from main size
their token ranges by the plan's partition, not KV_COMPUTE_BLOCK.
- Batched FP8 softmax now covers two or three query tiles; four tiles
take main's per-wave pipeline, which is 7-20% faster on gfx1250
(D128/page128/MTP4/GQA16, CUDA graphs). BF16 four-tile plain-V caches
likewise take main's wave-pair pipeline (20-45% faster than the packed
pipeline).
- PADecodePlan keeps arrival_counters and context_partition_size, passed
by keyword; gfx1250-strict supports the same partition sizes as gfx1250.
- pa_decode bounds in-kernel merges and the reducer by the plan's
partition size; MXFP4 accepts only 256-token partitions.
- Tests take main's new merge, split-group and MXFP4 cases; the plan test
stays removed as on main.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
gfx1250 PA decode writes per-partition partials and launches a second kernel to merge them. For MiniMax-M3 sparse decode (H16/KV1/D128, BF16 or FP8 page16 caches, 2048 top-k tokens per row) at 64 rows, that reducer takes about 2.9 µs plus a launch gap out of 11.5 µs. This PR merges partitions inside the tile kernel for the two common plan shapes, so those calls launch a single kernel. It also speeds up the dense tile kernels of MiniMax-M3's full-attention layers (page128 caches, decode and MTP4 verify) for FP8 and BF16 KV, and cuts launch and reducer latency for the plans that keep the reducer.
Partition merges
max_partitions) in [2, 16], the tile kernel launches P CTAs per (sequence, KV head) as one cluster (cluster=(P, 1, 1), gridbatch * P). CTA i runs the sequence's task i fromreduce_info, if it has one, and stores packed partials with plain stores. After a cluster-scope release, the cluster barrier and a cluster-scope acquire, each CTA merges 1/P of the sequence's rows. A cluster shares one XCC and its L2, so the release and acquire lower tos_wait_storecnt 0andglobal_inv scope:SCOPE_SE; agent scope would add a device-wide L2 writeback and invalidate per CTA. Plans qualify whenbatch * Pfits the plan capacity or two CTAs per CU, so CTAs idling at the barrier never displace working ones, and whenbatch * kv_headsstays within 15 clusters per 16 CUs: gfx1250 keeps at most 240 clusters resident on 256 CUs whatever their size, and from the 241st cluster on a second wave makes the merge about 1.5× slower than the reducer.PADecodePlan.arrival_counters([batch, kv_heads, 32]int32, one 128-byte line per counter), zeroed when the plan is created and reset by the last arrival, so CUDA-graph replays and plan refreshes need no host work. This merge uses device-scope (SCOPE_DEV) accesses.AITER_FLYDSL_PA_CLUSTER_MERGE=0restores the reducer for cluster-eligible plans.Multi-partition decode does not use a last-arriver merge: on gfx1250 each device-scope round trip costs about 4 µs, and merging BF16 sparse decode at 64 rows that way took 17.7 µs in the tile kernel, against 7.45 µs + 2.9 µs for tile and reducer.
Dense tile kernels
Launch and reducer latency
AITER_FLYDSL_KERNARG_PRELOAD, as the MoE kernels do.Performance
Measured on gfx1250 (256 CUs, 8 XCCs), FlyDSL 0.4.1, PyTorch 2.11.0+rocm10.1, HIP 7.16, with GPU runs serialized. MiniMax-M3 attention shapes: H16/KV1/D128. Dense rows use page128 caches with one sequence per request; verify rows have 4 causal query tokens. Sparse rows use page16 caches gathered from a 4096-page pool with 2048 top-k tokens per row, and prefill rows see
min(row + 1, 2048)tokens. Times include decode and reduction, excluding planning and compilation. Each run takes the minimum of five 10-call CUDA graphs with 5 replays each; the table reports the median over three rounds with rotating tree order. The base ismainat 0e53c46, which this branch merges.main*
mainis bimodal on this case across runs; this PR holds 29.3–31 µs.Outputs are bitwise identical to
mainexcept for FP8 and BF16 dense verify, where the per-wave paths sum the softmax denominator in a different order and fuse scales into the exponent. Their relative error against an FP32 reference is unchanged to four decimals (FP8 bs3×4 220K: 0.0369 on both). The FP8 bs3×4 220K verify row matches the MiniMax-M3 decode trace (3 requests × 4 tokens at about 221K context). At 64 sparse rows (rocprofv3 kernel trace, medians), the BF16 cluster kernel takes 8.21 µs wheremainruns a 7.17 µs tile kernel and a 2.92 µs reducer; FP8 takes 7.69 µs against 6.73 µs + 3.00 µs.The table uses
plan_pa_decode's default partition cap. ATOM's MiniMax-M3 sparse path instead caps each call atget_recommended_splits(rows, 1)and passesmax_context_length=2048. With those plans, sparse decode is 0.74–0.78× ofmain's time at 4–64 rows (8 partitions) and 0.80–0.81× at 128 rows (4 partitions), BF16 and FP8 alike. 256 rows (2 partitions, 256 clusters) keeps the reducer, about 5% faster than onmain. Sparse prefill already writes output directly onmainwith a one-partition cap and is unchanged.ATOM's dense calls pass no
max_context_length, so their capture-sized block tables compile the reducer for 256 partitions. With such tables (block tables padded to 256K tokens, median of three alternating rounds), dense decode is 1.25–1.32× faster thanmainfor FP8 KV (bs1 8K: 9.80 → 7.59 µs) and 1.06–1.22× for BF16.Validation
python -m pytest -q op_tests/test_flydsl_pa_decode.pyon gfx1250: 203 passed, 5 skipped (main: 119 passed, 5 skipped).max_partitionsandworkgroup_budgetsettings that give cluster sizes 2–16, BF16 and FP16 queries, FP8 per-token KV, CUDA-graph replays and plan refreshes, and one-slot plans that merge. Outputs are bitwise identical tomainexcept the MTP4 page128 cases that now take the per-wave paths: one FP8 case (relative error 0.04201,main0.04203) and seven BF16 cases (errors within 1e-4 of before).main, with per-token K/V magnitudes scaled by e^N(0, 1), at query lengths 1–8 with GQA 8–64 and lengths 1 to 33000:main's.git diff --checkpass.Notes
SCOPE_DEV) hardware semantics, asfused_compress_attn_hca_gfx1250.pydoes.pa_decodecomputes one tight per-call bound,max_partitions = min(plan cap, capacity, context tiles, window tiles), and passes it to the reducer; the cluster and one-slot gates use the same bound.