[MLA] Add FlyDSL FP8 prefill attention

[gbyu-amd]

Adds FlyDSL FP8 MLA prefill attention on gfx950. Enable it with ATOM_USE_FLYDSL_FP8_PREFILL_ATTN=1 (default: off).

The FP8 path fuses K NoPE/RoPE concatenation with per-tensor QKV quantization. Q is quantized once and reused across cached chunks. Where supported, FlyDSL gather writes FP8 K/V using descales max(new_token_descale, 1e-6) * 2; these are estimates for cached K/V and larger outliers saturate. BF16 gather outputs use fused K/V dynamic quantization. Both quantizers live in atom/model_ops/triton_fused_qkv_quant.py, and the layer's softmax scale is passed directly to attention.

AITER capability checks are cached per layer. If FP8 MHA is unavailable or the device/model configuration is unsupported, ATOM selects BF16 attention before quantization or FP8 gather and logs a warning once per process. This preserves BF16 Q/K/V for the fallback. FP16 activations and synthetic RoPE-padding configurations also use the normal attention path. If only FP8 gather is unsupported, BF16 gather plus dynamic quantization retains FP8 attention. Kernel execution errors propagate without retry.

ATOM tests its own fused QKV/KV kernels and the capability/fallback policy; AITER kernel correctness remains covered in AITER.

Dependencies

• ROCm/aiter#5400: FP8 output in FlyDSL gather — merged.
• ROCm/aiter#5508: short-query cached-prefix FP8 FMHA split-K — merged.
• ROCm/aiter#5718: public FP8 MHA/gather capability checks — pending. Until these APIs are installed, FP8 prefill selects BF16 attention with a warning.

Performance — 2026-09-15

Kimi-K3, c64, TP8/DCP8 on 8 gfx950 GPUs; LMCache 192 GiB, chunk size 1024, max sequences 128 and max batched tokens 8192. Launch arguments and recorded non-FP8 environment settings match the baseline selected by agentx_best_of.csv: agentx-20260911-115504/c64_lmcache_dcp8. The baseline uses older source revisions.

Both runs use the full 3,600-second profiling phase and the original 40-second completion grace period; 707 warmup requests are excluded. Metrics come from profile_export_aiperf.json; TTFT/ITL percentiles were verified against profile_export.jsonl.

Metric	Baseline	FP8	Change
TTFT P50 (ms)	1,356.06	1,335.62	-1.51%
TTFT P90 (ms)	3,802.34	3,328.99	-12.45%
TTFT P99 (ms)	28,132.76	26,816.95	-4.68%
ITL P50 (ms)	77.11	75.15	-2.53%
ITL P90 (ms)	100.66	96.99	-3.65%
ITL P99 (ms)	168.05	159.08	-5.34%
P90 interactivity (tok/s)	9.934	10.310	+3.79%
Total token throughput (tok/s)	97,767.58	101,323.63	+3.64%

P90 interactivity = 1000 / ITL P90. FP8 completed 3,134 profiling requests with no errors (baseline: 3,049).

The table uses a full repeat with the same-source compiler cache populated; no new FlyDSL cache artifacts were written during profiling.

Tested commits: ATOM ec5165bc11c895f41566a8bcbc8b885470b11b51, AITER 5188a03c3b2955dac21757e14b772148a37da959.

Accuracy — 2026-09-15

Full GSM8K after the performance run, using the same FP8 TP8/DCP8 server configuration: 1,319 questions, 5-shot, concurrency 64, temperature 0, top_p 1, max generation tokens 12,288.

Scoring filter	Accuracy	Correct / Total
strict-match	96.3609%	1271 / 1319
flexible-extract	96.2851%	1270 / 1319

Capability and fallback validation — 2026-09-21

ATOM 55108c6c, AITER 84028a316 (based on main 54ac58bab).

• Combined AITER capability/validation tests and ATOM dispatch regressions: 94 passed, 1 skipped. Metadata tests cover gfx942/gfx1201, unsupported head dimensions and LDS requirements, head grouping, and gather output/scale/device mismatches.
• ATOM regressions verify cached checks, a single warning per process, BF16 data preservation before FP8 preparation in plain/cached-single/cached-chunked paths, and unquantized Q/KV handoff to DCP on fallback.
• Real gfx950 checks passed for FP8 QKV + MHA, graph replay with modified inputs, direct FP8 gather + cached MHA, and BF16 gather + fused KV quantization + FP8 MHA. Simulated gfx942 rejection selects real BF16 attention and exactly matches the BF16 reference; MHA execution exceptions are not retried.
• Black, Ruff and diff checks pass. No full-model evaluation was rerun for this capability change; performance and accuracy figures above are historical runs at the stated revisions.

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