MLA PS mode fp8 -n 16,3 16,4 and nhead=32,64,128 support cp round robin

[minmengdie]

Motivation

This PR extends gfx950 MLA persistent (PS) FP8 support to additional decode qlen/head-count combinations and adds CP round-robin (interleaved KV) coverage, including wiring new asm kernel entries and updating routing/metadata to recognize the new supported shapes.

Technical Details

Broaden gfx950 FP8/F8 PS decode support gating (incl. nhead 32/64/128) and add qlen=3 handling in metadata/dispatch.
Add CP round-robin FP8 test support by threading explicit unit q_scale/kv_scale into the test harness and enabling fp8/fp8 runs.
Update gfx950 MLA asm kernel registry CSV with new CP round-robin kernel variants and remove superseded entries.

Test Plan

python3 op_tests/test_mla_persistent_round_robin.py -d fp8 -kvd fp8 -lse
python3 op_tests/test_mla_persistent.py -d fp8 -kvd fp8 -lse -n 16,3 16,4 16,6 32,2 32,3 32,4 64,1 64,2 64,3 128,1 128,2 -b 32 64 128 -c 1024 5000 8192 16384

Test Result

Check	Result
Total per-rank out / lse checks	15120
Failures	0
has_NaN=True	0
AssertionError (cal_diff hard gate)	0
Worst per-rank absolute delta	0.1875 (fp8 quantization noise level)
Merge lse (reference / kernel)	0 / 0 failures
Merge out (reference / kernel)	276 / 276 failures

nhead,qlen	Round 1	Round 2	Average	Reproducible
16,6	+1.35%	—	+1.35%	Noise level
32,2	+1.45%	—	+1.45%	Noise level
32,3	−11.64%	−11.10%	−11.37%	Yes (faster)
32,4	+1.83%	—	+1.83%	Noise level
64,1	+0.23%	—	+0.23%	Noise level
64,2	+0.05%	—	+0.05%	Noise level
64,3	−1.22%	−0.47%	−0.85%	Noise level
128,1	+0.56%	—	+0.56%	Noise level
128,2	+1.75%	—	+1.75%	Noise level

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 PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4521 --add-label <label>

[Copilot]

Pull request overview

This PR extends gfx950 MLA persistent (PS) FP8 support to additional decode qlen/head-count combinations and adds CP round-robin (interleaved KV) coverage, including wiring new asm kernel entries and updating routing/metadata to recognize the new supported shapes.

Changes:

• Broaden gfx950 FP8/F8 PS decode support gating (incl. nhead 32/64/128) and add qlen=3 handling in metadata/dispatch.
• Add CP round-robin FP8 test support by threading explicit unit q_scale/kv_scale into the test harness and enabling fp8/fp8 runs.
• Update gfx950 MLA asm kernel registry CSV with new CP round-robin kernel variants and remove superseded entries.

Reviewed changes

Copilot reviewed 7 out of 19 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
op_tests/test_mla_persistent.py	Updates PS-mode support gating for gfx950 fp8/fp8 head-count cases.
op_tests/test_mla_persistent_round_robin.py	Enables fp8/fp8 CP round-robin testing and threads q/kv scales.
hsa/gfx950/mla/mla_asm.csv	Registers new gfx950 FP8 CP round-robin asm kernels and removes old entries.
csrc/py_itfs_cu/asm_mla.cu	Adjusts gfx950 FP8 dispatch heuristics (incl. qlen=3) and kernel selection.
csrc/kernels/mla/metadata/v1_2_device.cuh	Broadens native-support gating and adds packed-QoLen special-case for qlen=3.
aiter/ops/attention.py	Aligns metadata selection logic for gfx950 fp8/fp8 qlen=3 and minor comment edits.
aiter/mla.py	Broadens gfx950 fp8/fp8 PS-mode support gating (nhead 32/64/128).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.