stage2 logits block load

[yanxuer-999]

Motivation

Under vLLM CUDA Graph capture, the gluon MLA decode kernel's KV-split count
gets frozen to 1. The host derives NUM_KV_SPLITS from min_kv_seq_len, which
vLLM hard-codes to 1 during capture, so every long-context decode replays on a
single workgroup — collapsing throughput (e.g. 100k-token decode runs in
thousands of µs instead of tens). This affects the bh16 decode regimes
(bh16bn64 for nhead ≤ 96 and bh16bn128 fp8) used by DeepSeek/Kimi-style MLA.

This PR makes the split count immune to graph capture and, along the way,
removes a serial dependency chain in the stage-2 reduction.

Technical Details

2. Stage-2 reduction optimization.
   • Replaced the serial, dependency-chained per-split merge with a block-load
     vectorized tree reduction (load BLOCK_S splits as a tile and reduce),
     with BLOCK_S made adaptive (min(64, next_pow2(NUM_KV_SPLITS))) to avoid
     VGPR over-allocation for small split counts.

Test Plan

• Correctness (op_tests/test_mla.py, checkAllclose atol/rtol=0.01 vs torch
  golden) across regimes and sequence lengths (32 / 64 / 128 / 256 / 16k / 100k):
  bh64, bh16bn64 (nhead ≤ 16 and 17–96), bh16bn128 (fp8), plus MTP.
• Graph-freeze reproduction: forced min_kv_seq_len=1 (the vLLM capture case)
  and compared against the device-side policy.
• Regression check on the normal path (no forced freeze).
• Added a host-side unit test for the split-budget policy
  (test_mla_gluon_dynamic_splits.py) verifying the budget matches the ~256-WG
  formula and is structurally independent of sequence length; it lives under
  op_tests/triton_tests/attention/ so the Triton CI selector runs it whenever
  mla_gluon.py changes.

Test Result

• All correctness checks pass across the regimes/lengths above (no regression).
• Immunity to the frozen min_kv_seq_len=1 (100k tokens, batch=1):

  regime	before (frozen)	after	speedup
  bh16bn128 (fp8)	2178 µs	39.5 µs	~55×
  bh16bn64 (nhead=96)	2649 µs	88.6 µs	~30×

• Normal (non-frozen) path unchanged: e.g. bh16bn64 nhead=16 @100k ≈ 45 µs,
  short seq (64) ≈ 11 µs, bh64 @16k/b64 ≈ 521 µs.

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 4555 --add-label <label>

[LiuYinfeng01]

Kimi-K3 TPOT comparison

Speedups are relative to the Fixed-1 baseline. The last column compares AITER PR #4555 directly against vLLM PR #51119 Dynamic.

Input tokens	Fixed-1 TPOT	Fixed-128 TPOT / speedup	PR #51119 Dynamic TPOT / speedup	AITER PR #4555 TPOT / speedup	PR #4555 vs Dynamic
1	20.647 ms	21.001 ms / 0.983×	20.397 ms / 1.012×	20.412 ms / 1.012×	0.07% slower
512	20.937 ms	21.015 ms / 0.996×	20.656 ms / 1.014×	20.463 ms / 1.023×	0.93% faster
1,000	21.248 ms	20.965 ms / 1.014×	20.621 ms / 1.030×	20.468 ms / 1.038×	0.74% faster
100,000	86.751 ms	22.058 ms / 3.933×	22.037 ms / 3.937×	21.149 ms / 4.102×	4.03% faster

[Dewei-Wang-sh]

lgtm