[gbyu-amd]

Runs MLA prefill attention in FP8 on gfx1250 (MI455X) with either of two AITER kernels, each behind its own flag. Both reuse what #2191 built for gfx950: fused QKV/KV per-tensor quantization, Q reuse across cached chunks, the DCP handoff and the cached BF16 fallback.

Flag	Default	Kernel	Serves
ATOM_USE_ASM_FP8_PREFILL_ATTN (new)	0	ASM FP8 varlen FMHA, fmha_fwd_fp8_varlen_asm (ROCm/aiter#7375)	gfx1250 B0+, QK/V head dims 192/128 (Kimi-K3)
ATOM_USE_FLYDSL_FP8_PREFILL_ATTN (#2191)	0	FlyDSL FP8 FMHA: gfx950 as before; on gfx1250 the a8w8 flash_attn_varlen_a8w8_m32x8 (ROCm/aiter#7214)	gfx950, gfx1250

Each layer checks once, before any FP8 preparation, whether the selected kernel serves it; otherwise it runs BF16 attention. Both kernels take per-tensor fp32 Q/K/V descales, write BF16, use bottom-right causal masking and return natural-log FP32 LSE. The ASM kernel writes the LSE as [total_q, heads, 1]; ATOM transposes it to the [heads, total_q] layout the merges read.

Defaults. Both flags default to 0, so BF16 prefill attention stays the default and nothing changes unless a flag is set. Set only one: with both set, ASM is used, ATOM_USE_FLYDSL_FP8_PREFILL_ATTN is ignored, and a warning saying so is logged once per process. Where the selected kernel does not apply (ASM on gfx950, A0 gfx1250-strict, other head dims, an AITER without the kernel), the layer falls back to BF16 — not to the other FP8 kernel — before quantization, and the fallback is logged once per process.

• ASM counts as available when aiter.ops.mha.fmha_fwd_fp8_varlen_asm and aiter.jit.utils.asm_guard.is_gfx1250_asm_supported import. Its check mirrors the FP8 branch of AITER's _select_gfx1250_fmha_fwd_varlen_asm: gfx1250 B0+, 192/192/128 head dims, query heads divisible by KV heads. The per-call conditions (3-D contiguous e4m3 Q/K/V, one-element fp32 descales on the device, BF16 output) hold for what ATOM's quantizers produce.
• The FlyDSL a8w8 entry counts as available only when it accepts q_descale/k_descale/v_descale (#7214). Its check reads shapes and dtypes only, so meta tensors stand in for Q/K/V.

Kernel performance

One MI455X (gfx1250 B0), K3 shapes (12 heads, D192x128), one sequence, kernel time only (quantization excluded); AITER 7061cbd88 with #7214 and #7375 applied.

Causal prefill, Q = K = context:

Context	BF16 ASM	FlyDSL FP8	ASM FP8	ASM / BF16	ASM / FlyDSL
1k	0.058 ms	0.033 ms	0.026 ms	2.24×	1.26×
4k	0.084 ms	0.064 ms	0.047 ms	1.76×	1.35×
8k	0.190 ms	0.140 ms	0.083 ms	2.30×	1.70×
16k	0.549 ms	0.364 ms	0.238 ms	2.31×	1.53×
32k	2.106 ms	1.225 ms	0.854 ms	2.47×	1.43×
64k	8.400 ms	4.414 ms	3.344 ms	2.51×	1.32×
128k	33.815 ms	17.468 ms	13.402 ms	2.52×	1.30×

At 128k that is 1,951 / 3,777 / 4,923 TFLOP/s. With 16k new tokens over a cached context of 1k–128k (non-causal), ASM FP8 is 2.14–2.59× BF16 and 1.27–1.39× FlyDSL FP8. Both FP8 kernels show the same output error against BF16 (relative 0.046–0.054); ASM and FlyDSL differ from each other by 0.017–0.027.

Dependencies

• [ASM] [HIP] [FlyDSL] [gfx1250] asm mha fp8 fwd: add D192x128 per-tensor varlen kernel aiter#7375: the ASM FP8 varlen FMHA — open. Without it ATOM_USE_ASM_FP8_PREFILL_ATTN=1 falls back to BF16.
• [FlyDSL] FMHA gfx1250: per-tensor fp8 descales for the a8w8 kernel aiter#7214: per-tensor descales for the gfx1250 FlyDSL a8w8 FMHA — open. Only needed for ATOM_USE_FLYDSL_FP8_PREFILL_ATTN=1 on gfx1250.

Validation

• Black, Ruff and py_compile pass on the changed files.
• tests/test_mla_flydsl_gather_dispatch.py and tests/test_envs.py: 135 passed, in a container with this AITER and no GPU attached (a test-only pytest plugin fixes AITER's Triton arch and fp8 dtype, which otherwise need a device at import). New cases: both flags default off; with both set, ASM wins and the warning is logged once across layers; ASM serves K3 on gfx1250 and is asked once; where ASM declines (gfx950, A0, other head dims) the layer runs BF16 without consulting FlyDSL and logs the fallback; the ASM launch gets the quantized Q/K/V and descales, and its LSE comes back as [heads, total_q].
• Pending: GSM8K on Kimi-K3 TP8/DCP8 (2 × 4 MI455X) with ASM FP8 and FlyDSL FP8 prefill against BF16, and end-to-end prefill performance. Results will be added here.

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