[PERF] Add opt-in atomic stage2 override

[vorapolsiloai]

Summary

Add an opt-in AITER_FLYDSL_FORCE_ATOMIC=1 override that replaces an eligible tuned FlyDSL stage-2 _reduce kernel with its registered _atomic sibling while constructing MoE metadata.

Safety and precedence are explicit:

• unchanged when the override is unset;
• unchanged when no registered atomic sibling exists;
• ignored for non-FlyDSL, already-atomic, or otherwise ineligible kernels;
• AITER_FLYDSL_FORCE_REDUCE=1 takes precedence when both variables are set;
• output allocation and MoE sorting use the final metadata, so accumulation semantics match the selected epilogue.

Motivation and measured result

Using #4554's Qwen3-VL tuning rows on 8× MI355X, 1024 input/output tokens, concurrency 4:

• tuned _reduce: 375.19 output tok/s, 889.47 total tok/s;
• forced _atomic: 391.46 output tok/s, 928.27 total tok/s;
• reference range: 389.48–391.05 output tok/s, 924.04–928.77 total tok/s.

The opt-in override recovers approximately 4.3% output throughput for this point while retaining safe fallback behavior.

Qwen3-VL integration stack

• vLLM #50212: fused attention prologue; requires AITER main containing the strided-cache support merged via [HIP] [kernel] mrope cache-quant: accept strided flash KV-cache view #4531.
• AITER #4531: merged into AITER main as b0cc5392; provides full strided MRoPE KV-cache support.
• AITER #4554: tuning rows used by the measurement above.
• This PR — AITER #4556: optional atomic stage-2 override.
• AITER #4761: independent Triton unified-attention prefill/decode optimization used by the later Image C stack.

Rebase provenance

• Frozen AITER main: fc2e5d57fb5b8ad8e7e23f7103071dde798ea618
• Current PR head: 5390e8912bcfd58f6132a30428dcb5a2bf7744de

Validation

• pytest -q op_tests/test_fused_moe_kernel_selection.py — 7 passed.
• Coverage includes available/unavailable sibling behavior, ineligible kernels, FORCE_REDUCE precedence, and the full tuned-config lookup → metadata rewrite → MoE-sorting accumulation path.
• Black and Ruff checks — passed.
• Python compile and git diff --check — passed.
• TP8 one-point serving benchmark completed 32/32 requests successfully.
• Upstream CI completed except for an unrelated MI35X test_moe_2stage.py shard that exhausted GPU memory with the opt-in override unset; a maintainer rerun has been requested.

This PR was developed with AI assistance. The human author reviewed the changes, ran the listed validation, and owns the final implementation. Commit metadata contains only human attribution and DCO trailers.

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

[vorapolsiloai]

The rebased branch is mergeable and the PR-specific tests pass, but one standard-test shard needs a maintainer rerun:

• linux-aiter-mi35x-1, shard 4 failed in unrelated op_tests/test_moe_2stage.py after ~9 minutes with torch.OutOfMemoryError (requested 48.19 GiB; 43.60 GiB free; 188.66 GiB already allocated).
• The override in this PR remains opt-in and was unset in that standard-test job.
• I attempted the job rerun API, but fork authors do not have the required repository permission.

Could a maintainer please rerun the failed shard?

[vorapolsiloai]

close this now as alternative fixing at config level #5155 has now merged