Repository navigation
fix(attention): keep MTP-5 decode on FlyDSL pa_decode past the gluon query-length limit - #2843
ZhiweiYan-96 wants to merge 8 commits into
Conversation
…query-length limit With --num-speculative-tokens 5 the verify step has query length 6. _dispatch_decode sent any query length above PA_GLUON_MAX_QUERY_LEN (4) to paged_attention_unified, which cannot carry per-token FP8 KV scales, so Qwen3.5 MXFP4 graph capture failed and MTP-5 could not start. With ATOM_PA_FLYDSL=1, paged_attention_triton runs aiter FlyDSL pa_decode, which accepts any positive query length. Add flydsl_decode_static_ok(), the layer-static half of _flydsl_pa_decode_num_seqs, and consult it in _dispatch_decode only when gluon_decode_over_limit is already true on a GPU with gluon decode: a layer FlyDSL will serve stays on paged_attention_triton (skipping the ASM arm). Every within-limit route, including MTP-3 (qlen 4) and the draft steps (qlen 1), is decided exactly as before, and a GPU without gluon (gfx1250) keeps reaching FlyDSL through paged_attention_unified. run_pa_decode gains require_flydsl (default False). paged_attention_triton sets it when the call is past the gluon envelope, so on a GPU with gluon a call FlyDSL then rejects on a per-call clause raises instead of running gluon outside its envelope. Other callers are unchanged. Tests: TestFlyDSLPastGluonRouting (19 cases); _route now pins ATOM_PA_FLYDSL and the two FlyDSL capability probes.
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
not ready yet.. the routing idea is fine but the gate should not live here.
atom/model_ops/base_attention.py:363 — this re-implements aiter's pa_decode acceptance rules in ATOM, no kernel specific stuff should be here. expose one predicate from aiter and call only that.
atom/model_ops/attention_mha.py:761 — require_flydsl is re-derived at call time, can disagree with the _dispatch_decode decision. pass the decision down instead.
atom/model_ops/attention_mha.py:1124 — block size read from config but the call uses cache/metadata, can diverge. use the same source as the kernel call.
atom/model_ops/attention_mha.py:1165 — this also bypasses the ASM arm when use_triton_attn is False, not just past-gluon. gate on the actual condition.
atom/model_ops/base_attention.py:854 — plan/scratch built before the raise, leaks on retry. check feasibility first.
atom/model_ops/base_attention.py:400 — is this the real aiter contract or just the shapes you tested? never hard code..
tests/test_paged_attention_dispatch.py:340 — monkeypatched probes, tests never hit the real aiter acceptance. add one that does.
atom/model_ops/attention_mha.py:1117 — gfx1250 + per-token FP8 KV + MTP-5 still hits the same failure, out of scope? leave a comment.
no re-run after the rebase, and no perf/gsm8k on current main. any perf number? gsm8k?
— valar-cop 👮⭐
… statically Address review on #2843: - Drop flydsl_decode_static_ok and the static branch in _dispatch_decode. Past the gluon envelope the call still routes to paged_attention_unified, which already hands FlyDSL the calls it takes on gfx1250. That per-call check (flydsl_decode_seqs, the same gate run_pa_decode applies) now also covers a gluon GPU past the envelope, so no second copy of the rules, no config-derived block size, and no ASM-arm bypass. - Pass the decision down: paged_attention_unified calls paged_attention_triton(require_flydsl=True) instead of re-deriving it. - run_pa_decode checks require_flydsl right after the feasibility check, before any plan or scratch is built; gfx1250 keeps its own error. - Tests: per-call gate on real aiter capability code, decision pass-down, refusal without plan/scratch, and a GPU test that runs aiter's real pa_decode on the Qwen3.5 MTP-5 shape.
|
Thanks @valar-cop. Pushed 7ab3e00, which drops the static gate. Past the gluon envelope, Point by point:
Validation
Perf / GSM8K: these numbers come from the pre-rebase tree (k=5 against k=3, Qwen3.5 TP4, c12, aiperf 1 h). I have not re-run them on current main; the re-run on this head is pending.
🤖 Generated with Claude Code |
_flydsl_pa_decode_num_seqs no longer mirrors aiter's pa_decode checks. It calls aiter.ops.flydsl.pa_decode.can_use_flydsl_pa_decode (ROCm/aiter#7206) with the arguments run_pa_decode then passes to pa_decode: the sliced lengths and block table, the K/V scales in the layout pa_decode receives, and the same context bound. ATOM keeps only its routing choices (sinks, sliding window, non-persistent and BF16 KV stay on gluon) and the sequence-count recovery. ATOM_PA_FLYDSL=1 now requires an aiter with the predicate; without it the first FlyDSL decision raises an error naming the missing function. The mirrored block-size, tile and partition constants are removed. _flydsl_takes_decode and paged_attention_triton share one helper for the K/V scales, so the decision and the call see the same tensors.
|
@valar-cop I followed your suggestion on base_attention.py:363. The decision now comes from aiter, and ATOM only calls it.
Perf and GSM8K on current main are still pending. 🤖 Generated with Claude Code |
|
GPU test results (MI355X): |
|
@valar-cop review this pr Merged current main (c4c7768, no conflicts). Depends on ROCm/aiter#7206, which adds |
| ) from exc | ||
| can_use = getattr(mod, "can_use_flydsl_pa_decode", None) | ||
| if can_use is None: | ||
| raise RuntimeError( |
There was a problem hiding this comment.
docs and envs still say #5809, but now we need #7206.. please update both
There was a problem hiding this comment.
Updated in accf537: atom/utils/envs.py and docs/environment_variables.md now say FlyDSL decode needs aiter #5809 (explicit work plan) and #7206 (can_use_flydsl_pa_decode), and that with an older aiter setting the flag raises an error naming #7206.
| fwd_ctx, | ||
| require_flydsl=True, | ||
| ) | ||
| self._reject_per_token_scales_on_unified(k_scale, v_scale, "decode") |
There was a problem hiding this comment.
default user has no hint here, add the ATOM_PA_FLYDSL=1 reason in this branch?
There was a problem hiding this comment.
Added in accf537: with ATOM_PA_FLYDSL unset, a decode past the gluon envelope (e.g. Qwen3.5 MTP-5, qlen 6 x group 8) logs one info line per process saying it runs on unified_attention and that ATOM_PA_FLYDSL=1 lets FlyDSL take it (needs aiter #7206). aiter is not asked in that case. Test: test_flydsl_off_hints_once_and_skips_the_gate (3 calls -> 1 line, predicate not called).
| if self._flydsl_takes_decode( | ||
| q, k_cache, v_cache, k_scale, v_scale, attn_metadata | ||
| ): | ||
| return self.paged_attention_triton( |
There was a problem hiding this comment.
can we ask the gate once per call instead of twice per layer?
There was a problem hiding this comment.
Done in accf537: _flydsl_takes_decode returns the sequence count, paged_attention_unified passes it down as flydsl_seqs, and run_pa_decode skips its own check when given one. Same arguments as before (the check already used the tensors run_pa_decode receives), so routing is unchanged; other run_pa_decode callers (vLLM/SGLang bridges, sparse sites) still check inside. Test: test_the_gate_is_asked_once_per_call drives unified -> triton -> run_pa_decode to the plan step and asserts the gate ran once. CPU: tests/test_paged_attention_dispatch.py 166 passed with #7206 pa_decode, 165 + 1 skip with the installed aiter; the only failure is main's test_batch_one_asks_for_the_ceiling, which needs a visible GPU.
…in; docs name aiter #7206 - _flydsl_takes_decode returns the sequence count and paged_attention_unified passes it to run_pa_decode as flydsl_seqs, which then skips its own check: one predicate call per layer call instead of two. - With ATOM_PA_FLYDSL unset, a decode past the gluon envelope logs once per process that the flag would route it to FlyDSL (needs aiter #7206). - envs.py and docs/environment_variables.md: FlyDSL decode needs aiter #5809 and #7206.
| plan_step_owner=None, | ||
| plan_cache_owner=None, | ||
| plan_partition_cap: int | None = None, | ||
| require_flydsl: bool = False, |
There was a problem hiding this comment.
require_flydsl is dead now, no caller passes it.. remove the param and the raise branch?
There was a problem hiding this comment.
Removed in 1ef5a2b: the require_flydsl parameter, its raise branch and its docstring, plus test_a_required_flydsl_call_never_reaches_gluon_or_a_plan; the GPU test now passes flydsl_seqs. The no-gluon refusal stays covered by test_no_gluon_fallback_raises_instead.
| fp8 KV with per-token scales, qlen 6 -> 48 rows. _dispatch_decode still | ||
| sends it to paged_attention_unified; the FlyDSL decision is made there on | ||
| the call's own tensors (the same check run_pa_decode applies) and passed | ||
| down as require_flydsl. |
There was a problem hiding this comment.
docstring still says require_flydsl, it's flydsl_seqs now..
| if not _flydsl_off_hint_logged: | ||
| _flydsl_off_hint_logged = True | ||
| logger.info( | ||
| "decode with max_seqlen_q=%d x %d q heads per kv head runs on " |
There was a problem hiding this comment.
this route refuses per-token fp8 kv scales, so the hint reads wrong.. mention that?
There was a problem hiding this comment.
Right, with per-token FP8 KV that call is refused, not run on unified. Fixed in 1ef5a2b: the hint moved to paged_attention_unified. Per-token scales: the refusal itself now ends with "ATOM_PA_FLYDSL=1 lets aiter's FlyDSL pa_decode take it (needs aiter with ROCm/aiter#7206)" and no info line is logged; otherwise the one info line per process says it runs on unified. Neither appears where the flag would not help (ATOM_USE_UNIFIED_ATTN, flash layout, window, inside the gluon envelope). Tests: test_flydsl_off_names_the_opt_in_once[per_token=False/True], test_no_hint_inside_the_gluon_envelope. CPU 166 passed (with #7206 pa_decode); only failure is main's GPU-only test_batch_one_asks_for_the_ceiling.
…he opt-in where it is true - run_pa_decode: remove require_flydsl and its raise branch; no caller passes it since the decision travels as flydsl_seqs. - The ATOM_PA_FLYDSL hint moves to paged_attention_unified. With per-token FP8 KV scales unified refuses the call, so the refusal names the opt-in and no info line claims it runs on unified; otherwise one info line per process. Neither fires where the flag would not help (forced unified, flash layout, window, inside the gluon envelope).
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
a few small things..
perf and gsm8k re-run on current base for mtp-5?
ci green.
| ) | ||
| # A default user gets the opt-in named once: in the per-token FP8 KV | ||
| # refusal (unified cannot carry those scales), else in one info line. | ||
| flydsl_would_take = not envs.ATOM_PA_FLYDSL and self._flydsl_route_open( |
There was a problem hiding this comment.
hint can promise a route the gate then refuses for bf16 kv and sinks, fold those into _flydsl_route_open?
There was a problem hiding this comment.
Folded in ac263bc: _flydsl_route_open now also returns False for sink layers and for BF16 KV on a gluon GPU (the two routing choices _flydsl_pa_decode_num_seqs makes before asking aiter), so the hint only appears where ATOM_PA_FLYDSL=1 would offer the call to FlyDSL. Test: test_no_hint_where_atom_keeps_the_call_off_flydsl[sinks|bf16_kv], with a positive control that the same layer without the condition is open.
| gfx1250 with per-token FP8 KV past qlen 4 is decided by the same check; | ||
| it is not exercised by this path's tests or runs (out of scope here). | ||
| """ | ||
| if not envs.ATOM_PA_FLYDSL or not self._flydsl_route_open(attn_metadata): |
There was a problem hiding this comment.
test becomes vacuous now, _no_gluon sets ATOM_PA_FLYDSL=False so it returns early. also set it True there?
There was a problem hiding this comment.
Fixed in ac263bc: the test sets ATOM_PA_FLYDSL=1 and is parametrized on ATOM_USE_UNIFIED_ATTN: off -> taken with 4 seqs, on -> not taken. Checked by removing the ATOM_USE_UNIFIED_ATTN condition: [True] fails (as does test_declined_calls_stay_on_unified[unified_env]).
| Omitted, ``pa_decode_gluon`` sizes and allocates its own, so a FlyDSL | ||
| call, which uses its plan's cached scratch, allocates none. | ||
|
|
||
| ``flydsl_seqs``: the sequence count the caller already got from |
There was a problem hiding this comment.
pr body still mentions require_flydsl and the removed test, please update it..
There was a problem hiding this comment.
Updated: the body now describes flydsl_seqs (one predicate call per layer call), the opt-in hint and the envs/docs change; require_flydsl and the removed test are gone, and the test counts are for ac263bc (CPU, 181 passed with aiter#7206 pa_decode). Perf and GSM8K on current main for MTP-5 have not been re-run yet; the body still says so.
…e unified-env test - _flydsl_route_open also returns False for sink layers and for BF16 KV on a gluon GPU, which ATOM's FlyDSL routing never takes, so the ATOM_PA_FLYDSL hint only names a route the flag would open. - test_use_unified_attn_keeps_decode_off_flydsl sets ATOM_PA_FLYDSL=1 and checks both sides: taken (4 seqs) without ATOM_USE_UNIFIED_ATTN, not taken with it.
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
a few small things..
mtp-5 perf and gsm8k re-run on current head?
ci green.
Motivation
With
--num-speculative-tokens 5the verify step has query length 6._dispatch_decodesends any query length abovePA_GLUON_MAX_QUERY_LEN(4) topaged_attention_unified, which cannot carry per-token FP8 KV scales. On Qwen3.5-397B-A17B MXFP4 (per-token FP8 KV) on MI355X, CUDA-graph capture therefore fails withlayer 3 takes unified_attention for decode, which cannot carry its per-token KV scales(server log of our MTP-5 launch on 10-28), and MTP-5 cannot start.The gluon limit does not apply when
ATOM_PA_FLYDSL=1:paged_attention_tritonthen runs aiter FlyDSLpa_decode, which accepts any positive query length (rows = query length x GQA group; Qwen3.5 TP4 is 8 q heads / 1 kv head per rank, so qlen 6 is 48 rows). Its plan, scratch and reducer have no query-length bound that MTP-5 reaches (read from the aiter source inrocm/atom-dev:nightly_202610251529, aiter7be01d1c7): dense plans are built withquery_length=1, scratch is sized bymax_seqlen_q, and the reducer holds while batch x (rows x kv_heads)^2 <= 2^30 (1024 x 48^2 = 2.4M).On main,
paged_attention_unifiedalready hands FlyDSL the calls it accepts, but only on a GPU without gluon decode (gfx1250,_flydsl_takes_decode). On gfx942/gfx950 the over-limit call still reaches unified attention and fails as above.Depends on ROCm/aiter#7206
Whether FlyDSL
pa_decodetakes a call is now decided by aiter. ROCm/aiter#7206 addsaiter.ops.flydsl.pa_decode.can_use_flydsl_pa_decode, which runspa_decode's own host-side checks and its compile-time LDS limits. This PR calls it and no longer keeps a copy of those rules in ATOM. WithATOM_PA_FLYDSL=1, ATOM therefore needs an aiter that has the function. On an older aiter the first FlyDSL decision raises aRuntimeErrorthat names the missing function. There is no fallback mirror.ATOM_PA_FLYDSLis off by default, so builds that do not opt in are unaffected.Change
_dispatch_decodeis the same as on main. A call past the gluon envelope still goes topaged_attention_unified, and the ASM and triton arms are unchanged.paged_attention_unified->_flydsl_takes_decode: past the gluon envelope on a gluon GPU (the Qwen3.5 MTP-5 case), and on a GPU without gluon (as on main), it asksflydsl_decode_seqson the same tensors and metadata the kernel call receives. If FlyDSL takes the call, it returns the sequence count andpaged_attention_unifiedcallspaged_attention_triton(..., flydsl_seqs=n), which hands it torun_pa_decode, so aiter's predicate is asked once per layer call._flydsl_takes_decodeandpaged_attention_tritonbuild the K/V scales with the same helper, so the decision and the call see the same scale tensors._flydsl_pa_decode_num_seqsreplaces its copy of aiter's rules with a call tocan_use_flydsl_pa_decode. The call gets the argumentsrun_pa_decodethen passes topa_decode: the lengths and block table sliced to the recovered sequence count, the query length, the K/V scales,sliding_window=0and the same context bound. ATOM keeps two things. It keeps its routing choices: sinks, sliding window, non-persistent calls and BF16 KV on a gluon GPU stay on gluon. It also keeps the recovery of the sequence count from ATOM's padded batch. The mirrored constants (block sizes, tile, partition cap) are removed.run_pa_decode(flydsl_seqs=None): given a count, it skips its ownflydsl_decode_seqscall and runs FlyDSL with it. Other callers (vLLM, SGLang, MiniMax sparse) do not pass it and are checked inside as before.ATOM_PA_FLYDSLunset) are told about the opt-in only where it would open a FlyDSL route (_flydsl_route_open: notATOM_USE_UNIFIED_ATTN, paged layout, no window, no sinks, not BF16 KV on a gluon GPU, past the gluon envelope on a gluon GPU). With per-token FP8 KV the unified refusal namesATOM_PA_FLYDSL=1and aiter#7206; otherwise one info line per process says the call runs on unified attention.atom/utils/envs.pyanddocs/environment_variables.md: FlyDSL decode needs aiter #5809 and #7206.vLLM bridge (not changed)
plugin/vllm/attention/layer_mha.py::_dispatch_decode_backenduses the samegluon_decode_over_limitpredicate but has no unified fallback: past the limit it raisesNotImplementedError("query length ... is past the gluon decode kernel"). A qlen 6 decode there fails loudly rather than computing a wrong answer, so this PR leaves it alone. Opening it to FlyDSL needs its own static gate (block size from the vLLM cache config, a partition count checked against FlyDSL's cap, no ATOM-built work plan) and is a separate change.Why MTP-3 is unchanged
On a gluon GPU,
_flydsl_takes_decodereturns False unlessgluon_decode_over_limitis true. MTP-3 is qlen 4 x group 8, for whichgluon_decode_over_limit(4, 8, 1)is False, so its route is decided by the same code as on main and no FlyDSL count is passed down. Draft steps run qlen 1 and are unaffected for the same reason.test_within_gluon_keeps_its_routechecks qlen 1/2/4.Measured
Qwen3.5-397B-A17B MXFP4 (FP8 attention, per-token FP8 KV), TP4 on MI355X,
rocm/atom-dev:nightly_202610251529(aiter7be01d1c7), AgentX trace at concurrency 12, one 3600 s aiperf run per arm. Acceptance is forced, as in InferenceX; mean acceptance length from the server'sMTP Statsline is 3.39 for k=3 and 4.41 for k=5. Both arms run the same code; k=3 does not take the new path.These runs used this change on an older base (ATOM
daa4f1bb8plus local Qwen3.5 decode commits that are not upstream, including MTP draft mid-step graph capture), not on currentmain, so the absolute numbers include those other commits. The table is aiperf's summary of each run:Paired per request on the 2678 requests both runs served (same conversation and turn, equal output length): the median per-request ITL ratio k=5/k=3 is 0.934 (95% CI 0.928-0.938). k=5 is faster on 1791 requests and slower on 881; sign test p = 3e-69. The public SGLang c12 reference for this model on MI355X is 255.75 tok/s/user, so k=5 moves this point from 254.59 / 255.75 = 0.996x to 273.89 / 255.75 = 1.071x of it.
GSM8K, all 1319 questions, 3-shot, max 2048 generated tokens, real (not forced) acceptance, same older base: strict-match 0.9121 with k=5 against 0.9257 with k=3, flexible-extract 0.8620 against 0.8741. Paired on the same questions, 80 are correct only with k=3 and 62 only with k=5 (exact McNemar p = 0.15 strict, p = 0.34 flexible, computed from the two runs' per-sample files paired by doc_id), so the gap is within run-to-run noise. Real acceptance on GSM8K with k=5 is 4.51 tokens per forward (server log).
Follow-up in aiter (separate, not needed here)
At 48 rows aiter's tuned head-256 FP8 pipe (
D256_PIPE, limited to 32 rows) is not used, so MTP-5 verify attention takes the generic multi-M-tile path. Kernel microbenchmark on MI355X, ctx 150k, per layer, bs 1/2/4/8: 72.3 / 74.2 / 85.7 / 142.6 us generic, against 59.4 / 61.5 / 68.5 / 116.7 us with the pipe limit raised to 48 rows (numerics identical to the generic path). That is a one-line aiter limit plus tests and will go to aiter separately; this PR needs no change for it, the same dispatch picks it up.Tests
tests/test_paged_attention_dispatch.py:TestFlyDSLPastGluon: qlen 5/6/8 take FlyDSL, and aiter is asked with per-token scales that have their trailing singleton dimension. qlen 1/2/4 keep their route and never call aiter. FlyDSL off,ATOM_USE_UNIFIED_ATTN, flash layout, sliding window, BF16 KV and an aiter decline all stay on unified. The decision is passed down asflydsl_seqs(test_unified_passes_the_decision_down), andtest_the_gate_is_asked_once_per_calldrives unified -> triton ->run_pa_decodeto the plan step and checks the gate ran once. The opt-in hint: once per process, in the refusal for per-token FP8 KV, and never inside the gluon envelope or for sinks / BF16 KV (test_flydsl_off_names_the_opt_in_once,test_no_hint_inside_the_gluon_envelope,test_no_hint_where_atom_keeps_the_call_off_flydsl).TestFlyDSLCapabilityGate: aiter's predicate is mocked here, because these tests use meta tensors. The tests check the exact arguments aiter gets (sliced rectangle, scales, window, context bound), that an aiter decline falls back, that ATOM's routing choices never reach aiter, the sequence-count recovery, and the error on an aiter without the predicate.test_the_predicate_is_aiterschecks that the function called is aiter's own when aiter has it.test_qwen35_mtp5_shape_is_accepted_by_real_aiterruns aiter's realpa_decodewith no mocks on gfx942/gfx950. It skips without a GPU and has not run yet.At
ac263bc47(main merged atc4c7768fe), on CPU (HIP_VISIBLE_DEVICES=, arch probe pinned to gfx950) inrocm/atom-dev:nightly_202611011456(aiter0dd271600):test_paged_attention_dispatch.py+test_minimax_m3_sparse_pa_floor.py181 passed, 4 skipped, 1 failed with aiter#7206'spa_decodeloaded over the installed one; 180 passed, 5 skipped, 1 failed without it (the extra skip is the test that needs aiter's predicate). The one failure,TestDenseDecodeSplits::test_batch_one_asks_for_the_ceiling(from main #2179), needs a GPU. black and ruff are clean on the changed files.Perf and GSM8K above have not been re-run on current main.
🤖 Generated with Claude Code