[lirui927]

Changes

Run MiniMax-M3 paged decode on gfx1250 (MI455) through FlyDSL. gfx942/gfx950 are unchanged.

• ATOM_PA_FLYDSL=1 routes dense and sparse paged decode to FlyDSL, as on gfx942/gfx950, for FP8 and BF16 KV. gfx1250 needs it, since it has no gluon decode.
• The V cache goes to FlyDSL as its 5D SHUFFLE view.
• On gfx1250, plans past 4096 sequences set max_batch_size, so large sparse prefill steps stay on FlyDSL. gfx942/gfx950 keep the 4096 cap.
• Sparse rows skip pa_fwd_asm on gfx1250.
• Prefill steps above ATOM_QK_NORM_ROPE_CACHE_FUSION_MAX_TOKENS write a BF16 KV cache without per-token quantization; before, long BF16 prefills crashed (from @XiaobingSuper, fix(attention): run MiniMax-M3 on gfx1250 (unified decode, BF16 KV long prefill) #2594).

This PR does not depend on ROCm/aiter#6699 to merge. The end-to-end results below use it: it fixes fused all-reduce + RMSNorm on gfx1250, which is outside this PR. Without it, MiniMax-M3 TP4 on gfx1250 loses GSM8K accuracy (0.9219 vs 0.9325 on an older aiter), and on current aiter main the server can crash during CUDA graph capture.

Test Result

MiniMax-M3-MXFP4, TP4, MI455, aiter with ROCm/aiter#6699, ATOM_PA_FLYDSL=1. GSM8K 5-shot, full set.

KV	flexible-extract	strict-match
FP8	0.9348	0.9348
BF16	0.9356	0.9363

8K in / 1K out, concurrency 4, mean of two runs per KV dtype:

KV	Output throughput	TPOT	TTFT
FP8	631.1 tok/s	5.81 ms	266 ms
BF16	629.7 tok/s	5.79 ms	296 ms

Unit tests: 168 passed.

[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 2710 --add-label <label>

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[valarLip]

put these into aiter, instead of ATOM

[XiaobingSuper]

Moved the backend capability checks into aiter in ROCm/aiter#7026. Once that API is available, this PR can import the two helpers and drop the local architecture table.

[XiaobingSuper]

Updated in 8a6f9c1: ATOM now imports both capability queries from their owning aiter backends and removes the local architecture table/helpers. ROCm/aiter#7026 has merged.

[XiaobingSuper]

Refined in e753771: each consumer now imports the exact capability helper directly from its owning aiter module; there is no ATOM alias/re-export layer.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[valarLip]

Thanks for moving the capability checks into aiter. I re-read at c0354d4 by reading the source, checking against aiter 7c362a7a04; nothing was run on gfx1250. Item 1 needs fixing before merge, and 2–5 should be addressed.

Must fix

1. On gfx1250, every new sparse-prefill length compiles a planner and leaves a buffer set that is never freed (base_attention.py:511).

• Per-call plans (plan_step_owner=None) are cached in _FLYDSL_RUNTIME_PLANS under a key that includes num_seqs, and are never evicted.
• Each plan gets its own scratch in _FLYDSL_PLAN_SCRATCH, keyed by id(plan): [Hkv, capacity, rows] fp32 twice, plus [Hkv, capacity, rows, D]. These are never evicted either.
• aiter's compile_pa_decode_plan(batch, capacity, max_parts, sliding_window, query_length) is @cached on batch, so each new batch also compiles a planner kernel.

This PR makes sparse rows skip pa_fwd_asm on gfx1250, so every M3 prefill step now enters this path with rows = T * Hkv. In serving, prompt lengths vary, so each new T means a planner JIT (a TTFT spike) and another resident buffer set, about 64 MiB of output for T = 16384, g = 16, plus the fp32 pair. All of it is allocated after the KV pool has been sized from gpu-memory-utilization, so a long-running server accumulates it until it OOMs. On gfx950, rows past 64 go to ASM, so the key set stays bounded, which is why this did not show before. The PR's validation uses a fixed 8K input and so exercises a single length.

Prefill plans need a bounded set of shapes: size the plan by a capacity bucket and reuse its scratch, or scope the plan and scratch to the step, instead of caching one per row count.

Should fix

2. On gfx1250 an unsupported call raises mid-serve instead of being refused at startup (base_attention.py:750). There is no gluon fallback on gfx1250. So a sparse prefill whose rows exceed _flydsl_plan_max_batch() (81792, e.g. 65536 batched tokens × 2 KV heads per rank) makes flydsl_decode_seqs return None and raises "FlyDSL paged decode does not take this call ...", and the worker dies on its first large prefill. Without ATOM_PA_FLYDSL, M3's sparse decode likewise raises at runtime. Both are known from the config, so they should be refused at startup.

3. The error for FP8 dense decode without ATOM_PA_FLYDSL is misleading (attention_mha.py:1116). On gfx1250 every dense decode dispatches to unified. With the env off, _flydsl_takes_decode is False and _reject_per_token_scales_on_unified raises NotImplementedError about per-token KV scales first. The "set ATOM_PA_FLYDSL=1" message in run_pa_decode is never reached, and the same happens for sinks layers and unsupported block sizes.

4. The kernel choice is now made in two places (attention_mha.py:649). _dispatch_decode's docstring says it is "the single place that answers which decode kernel runs. The runners do not re-decide". Now paged_attention_unified re-decides and forwards to paged_attention_triton, so traces show unified wrapping triton while FlyDSL runs. dispatch_backend already has the forward context and this layer's cache, so it can decide once whether FlyDSL takes the call and return a FlyDSL arm explicitly.

5. The docs and docstrings describe a fallback that gfx1250 does not have. paged_attention_triton says "always FlyDSL on gfx1250", but that needs ATOM_PA_FLYDSL=1 and raises otherwise. docs/environment_variables.md says calls FlyDSL does not take "fall back to gluon", and _flydsl_explicit_plan_api warns "using gluon", but there is no gluon on gfx1250; those cases raise.

Smaller

• _flydsl_takes_decode runs the whole FlyDSL capability check and drops the result, then run_pa_decode runs flydsl_decode_seqs again. The V view, dense_decode_splits and _flydsl_context_bound are also computed twice. All of this runs per layer per decode step in the eager region. Returning (seqs, context_bound) once and passing it down removes the repeats.
• _FLYDSL_PLAN_SCAN_LDS_BYTES copies aiter's private _PLAN_SCAN_LDS_BYTES (kernels/pa_decode_plan.py), but unlike the other mirrored constants in that file it is not pinned in TestFlyDSLConstantsMatchAiter. If aiter changes it, ATOM's LDS bound and the planner's own check disagree.
• The FP8 / non-FP8 cache-write branch added to _gemma_norm_rope_then_cache (attention_mha.py:469) repeats the one in rope_cache almost line for line. The BF16 crash this fixes was exactly such a drift. A shared _write_kv_cache(...) would keep the two together.
• Some names no longer match what the code does:
  • _view_v_cache_for_pa_decode_bf16_asm now also serves FlyDSL's FP8 V;
  • _shuffle_v_for_flydsl only takes a view, it does not shuffle;
  • on gfx1250, paged_attention_unified and paged_attention_triton both run FlyDSL.
• On gfx1250, paged_attention_triton and sparse_pa_per_row still allocate gluon-shaped exp_sums / max_logits / temporary_output on every call even though FlyDSL uses its own scratch. For a large sparse prefill that is a [T, 1, splits, g, 128] buffer per layer that nothing reads.