Fix(attention): adapt the FlyDSL paged decode to aiter #5809 (explicit work plans)

[yitingw1]

Motivation

ROCm/aiter#5809 removes FlyDSL's static (no-plan) decode path: every pa_decode call now needs an explicit work plan. ATOM still takes the static path wherever there is no capture-owned plan: dense decode at an uncaptured batch size, MiniMax-M3 sparse decode, and the vLLM bridge. As a result, ATOM_PA_FLYDSL=1 on an aiter with #5809 fails at the first FlyDSL decode.

This PR gives every FlyDSL call an explicit plan with aiter's default workgroup budget (2×CU). Calls the planner cannot take go to gluon instead of raising. It also removes ATOM_PA_FLYDSL_PLAN: without a planner-less path, setting it to 0 could only mean gluon, which is what ATOM_PA_FLYDSL=0 already does. A leftover setting is ignored with a warning.

Merge after aiter #5809; upgrade ATOM and aiter together. With ATOM_PA_FLYDSL=0 (the default) nothing changes. Results: GSM8K 0.963–0.968; agentic E2E is neutral against the pre-#5809 FlyDSL path.

Design

Call site	main (aiter without #5809)	This PR
Dense decode, captured batch	Capture-owned plan, refreshed once per forward	Unchanged; default budget and partition ceiling; target and draft share it
Dense decode, uncaptured batch / eager	Static FlyDSL	Runtime plan, default ceiling, refreshed per call
MiniMax-M3 sparse decode	Static FlyDSL, split cap 8	Runtime plan, cap 8, refreshed once per step, including inside graph capture
vLLM bridge	Static FlyDSL	Runtime plan, static split count as cap, refreshed per call
> 4096 rows, or context bound < query window	Static FlyDSL / n/a	gluon
ATOM_PA_FLYDSL_PLAN	0 = static FlyDSL	Removed (warning if set)

 run_pa_decode
   │ ATOM_PA_FLYDSL? in FlyDSL's domain? rows <= 4096? bound >= ql?  ── no ──► gluon
   ▼ yes
   │ capture-owned plan fits?  ── yes ──► use it
   ▼ no
   runtime plan, keyed (stream, rows, Hkv, ql, cap, window, device, step-scoped)
     refresh: once per (plan_step_owner, capturing) if an owner is given, else every call
   ▼
   aiter pa_decode(..., work_plan=plan, max_context_length=bound)

Default budget. #5809 ships no tuned budgets, and its autotune key does not include context length. On the agentic context distribution, the default 2×CU was the fastest budget at every dense and draft batch size we measured.

Once-per-step sparse refresh. All 57 sparse layers in a step select the same per-row lengths. The local tail block is always selected, so each length depends only on seq_len; we observed 0 differences over 9,000 steps. One planner launch per step is therefore enough. Two details make this correct:

• The "already refreshed" marker lives on the caller's per-step decode metadata (plan_step_owner), not on ATOM's ForwardContext. The vLLM plugin never sets that context, so a marker stored there would outlive the step.
• The marker is keyed by capture state. Warmup and capture share the same metadata and stream, so without that key the captured graph would record no refresh, and every replay would run on the capture-time lengths. On an earlier revision this produced garbled output, shown below.

Partition ceiling. Native dense runtime plans use aiter's default ceiling, as capture-owned plans do. Sparse keeps its cap of 8, because each row reads a bounded top-k window. The vLLM bridge keeps its static split count until that path is measured.

Test Result

Config: MI355X, MiniMax-M3-MXFP4 TP4 + EAGLE3-GQA, ATOM_PA_FLYDSL=1. aiter is #5809 820e5e3df merged with main b1cdc19eb.

Unit tests: tests/test_paged_attention_dispatch.py, 142 passed. Each behavioral fix has a test that fails when that fix is reverted.

Accuracy (GSM8K 5-shot, 1319 questions, real acceptance):

Revision	conc 8	conc 32
This PR	0.9629 / 0.9621	0.9682 / 0.9682
Earlier revision without the capture-state key	garbled; stopped at 116/1319	0.9591 / 0.9583

Concurrency 8 keeps decode batches between 5 and 16, so sparse attention runs through FlyDSL under graph replay. At concurrency 32 most steps go to ASM, which is why that column did not expose the bug.

E2E (agentic, 3600 s per point, same node). This PR vs the pre-#5809 FlyDSL path:

	Throughput	Interactivity p50	Interactivity p90
C24 (ABBA)	−0.08%	+0.62%	+0.49%
C32 (single pair)	+0.64%	+0.15%	−2.0%

All differences are within the run-to-run spread, so the change is performance-neutral. The kernel microbench shows #5809's dense decode 5.6–8.9% faster at batch 16–40, but this does not show up end to end: dense attention is only 3 layers plus the draft.

Known Limitations

1. Runtime plans are never released. A new uncaptured row count compiles the planner once (~70 ms cold, then served from Triton's disk cache) and keeps the plan and its scratch for the life of the process, outside the KV-cache budget. Measured on native ATOM over 1-hour agentic runs: 22–26 runtime plans per rank, 33–75 MiB of scratch (C24 / C32). The scratch per plan is at most 2 MiB for sparse (FlyDSL only takes ≤ 64 rows, so ≤ 64 plans) and about 8 MiB for dense (only uncaptured batches, bounded by max_num_seqs). The vLLM bridge is the exposed case: its mixed prefill+decode steps do not pad the decode count, so with ATOM_PA_FLYDSL=1 (off by default) it can accumulate one ~8 MiB plan per distinct decode count up to max_num_seqs; this path has not been run. Calls without a per-step owner (the vLLM bridge, native dense eager) also refresh their plan on every layer.
2. The once-per-step sparse refresh assumes every sparse layer uses the same top-k, init-block and local-block settings, as MiniMax-M3 does; it is only enabled when sparse_local_block >= 1 (otherwise every call refreshes). With these settings, per-layer lengths matched in every step: 0 differences over 10,200 steps (571K layer comparisons) with ATOM_M3_INDEXER_CP=1 at C15. In that run, 88% of requests exceeded the 36,864-token CP long-path threshold.
3. vLLM FULL-graph and TBO capture of runtime plans have not been traced end to end.
4. CI skips tests/test_paged_attention_dispatch.py because the non-GPU job has no aiter; this is the same on main. Run the tests locally.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2514 --add-label <label>

[valarLip]

I reviewed this at 457c34a9e. Giving every FlyDSL call an explicit plan makes sense, and I confirmed several parts are right:

• Target and draft get separate scratch when they share a plan.
• _flydsl_context_bound returns a Python int on every path.
• Every max_seqlen_k source bounds the GPU lengths.
• The SGLang bridge never builds a plan.

The issues below are mostly in how runtime plans are keyed and cached.

Correctness

TBO ubatches can share a runtime plan and use each other's lengths. The runtime plan key starts with stream_id, and its comment says "TBO ubatches must not share a plan". But make_tbo_contexts gives every ubatch the same compute_stream (ubatching.py:617-629), so two ubatches with the same row count get the same plan_key. The "already refreshed this step" mark lives on each ubatch's own decode_md. As a result, ub0's first layer refreshes plan P with ub0's lengths, then ub1's first layer (no mark on its owner yet) refreshes P with ub1's lengths. ub0's later layers find their own mark and decode with ub1's work_info, and nothing raises. I found no gate that keeps decode TBO off for MiniMax-M3: the config only refuses it with DCP plus speculative decoding. Adding the ubatch id to the key, or attaching the plan to its owner, would fix it.

The sparse "refresh once per step" relies on sparse_local_block >= 1 without enforcing it. Whether bt_ctx includes the tail depends on whether this layer selected the tail block (index_topk.py:154-158), so all sparse layers agree only when the tail is forced. The model's default is sparse_local_block=0 (model.py:466). The shipped MiniMax-M3-MXFP4 config sets it to 1, so this does not trigger today. An assertion when local_blocks < 1 would keep a future config from silently reusing another layer's lengths.

Runtime plans are never released

_FLYDSL_RUNTIME_PLANS and _FLYDSL_PLAN_SCRATCH are module-level dicts that are only ever inserted into, and their keys include num_seqs. The builder's own comment (aiter_attention.py:836-842) explains why it mints plans only during capture: a new batch value is a 65-72 ms cold planner specialization, and the plan can never be freed. run_pa_decode now does exactly that at runtime anyway.

The vLLM bridge is hit hardest. layer_mha.py:556 passes neither a plan nor an owner, and in mixed prefill+decode steps the decode count is not padded. Every new count therefore stalls serving for one compile and pins a plan plus its scratch for the life of the process, outside the KV budget. Known Limitations #1 covers the compile, but not the memory and not the bridge.

Two related points:

• Plans are refreshed on every layer. With no plan_step_owner (the whole vLLM bridge, and native dense eager or uncaptured batches), plan_pa_decode is re-issued on every layer, although all layers of a forward share the same context_lens.
• The plan logic exists twice. refresh_flydsl_plan in the builder and the runtime path in run_pa_decode contradict each other. If the builder owned all native plans (dense and sparse) on the per-step metadata, and the runtime fallback served only the bridges, the global dicts, the stream key and the capture-state mark would all go away. That would also remove the need to write _flydsl_refreshed_plans onto a caller-owned metadata object.

Robustness

• No check for the aiter API version. The gate checks shapes only. plan_pa_decode is imported inside the function, and the call now uses #5809's keyword-only max_context_length and a required work_plan. With an older aiter and ATOM_PA_FLYDSL=1, the first FlyDSL decode raises instead of falling back to gluon. Probing for the API once in the gate would restore the fallback.
• _flydsl_context_bound raises. It raises on a bool, a non-integer or a negative value, while the gate's contract is to return None and fall back. Nothing in the tree produces such a value today.
• The window logic is dead code. The gate rejects sliding_window > 0, so window is always 0 when the plan is built, and the pa_decode call hard-codes sliding_window=0 (line 630). The window part of the key and the related existing_plan checks are therefore dead. Once the window gate is lifted, the plan and the kernel would disagree.

Docs and recipe

• The recipe numbers predate this change. The p90 numbers in the recipe (+4.0% / +11.8% / +20.5% vs gluon) come from [MiniMax-M3] Route the paged decode to aiter's FlyDSL kernel with work planner #2366. At that time the planner was deliberately kept off the two sparse call sites "where it measures a net loss". This PR routes sparse through the planner and moves to #5809's API, but its end-to-end comparison is against the pre-#5809 FlyDSL path, not gluon. Please either re-measure against gluon or note which configuration the numbers belong to.
• The 4096 limit counts sequences, not rows. docs/environment_variables.md and the comment at line 519 say calls over 4096 rows fall back, but the code compares the sequence count. With ql=4, 2048 sequences are 8192 rows and still take FlyDSL.
• Stale comments. Line 566 still says "the sparse sites and the vLLM/SGLang bridges do not" pass a plan. attention_mha.py:628 still says dense is "the only one a work plan can rebalance". aiter_attention.py:798 still refers to aiter #5546.

Smaller items

• Unused scratch on the FlyDSL path. attention_mha.py still allocates gluon-shaped exp_sums / max_logits / temporary_output on every call, while the FlyDSL branch uses _flydsl_plan_scratch, so these three are never read when FlyDSL is on. Allocating them in the gluon branch of run_pa_decode would avoid it. The gate also still checks max_context_partition_num.
• Tests. Several tests assert on inspect.getsource of the code under test (for example "work_plan=work_plan" in dispatch). The whole group is behind skipif(not _HAS_FLYDSL_PA), so CI never runs it. test_scratch_is_per_plan_not_per_shape already notes that two TBO ubatches must not share scratch, but no test covers two owners mapping to one plan_key, which is the first issue above.

Not verified

The new max_context_length=attn_metadata.max_seqlen_k is max_model_len on the graph path and the step's actual maximum in eager mode, and aiter documents it as the scheduling bound. Graph and eager runs of the same batch may therefore pick different kernel specializations, and eager could cold-compile mid-serving when contexts cross a threshold. My aiter does not include #5809, so I couldn't check PaDecodeSchedule.select itself.

[yitingw1]

Fixed in a79c0c5: TBO ubatch id in the plan key, once-per-step refresh only with sparse_local_block >= 1, an aiter API probe that falls back to gluon, one block-table bound for graph and eager, dead window code, and stale docs and comments. Runtime-plan lifetime (measured 33–75 MiB per rank per hour) and builder-owned plans are out of scope and listed under Known Limitations.