[HIP] [CI] [JIT] Add aiter.mqa_logits with a tuned Triton/FlyDSL table

[amd-sriram]

Motivation

vLLM's sparse indexer prefill calls the Triton/Gluon fp8_mqa_logits directly on gfx950. Moving it to the FlyDSL kernel from ROCm#4538 took a vLLM PR (vllm-project/vllm#59280), and so would every later switch. Like aiter.paged_mqa_logits for decode (ROCm#6140), this adds aiter.mqa_logits: the framework asks for MQA logits, and AITER picks the backend and FlyDSL variant from a tuned table, the way fused_moe does.

What it adds

• mqa_logits(Q, KV, kv_scales, weights, cu_starts, cu_ends, clean_logits=True): takes the same tensors as fp8_mqa_logits and returns the same [seq_len, seq_len_kv] logits.
• A tuned table keyed on (gfx, cu_num, heads, head_dim, seq_len, seq_len_kv). Both lengths are tensor sizes, so the lookup never reads device memory. A call uses the row with the largest tuned seq_len_kv not above its own, then the largest tuned seq_len not above its own; a length below every tuned one uses the smallest. Rows are looked up with the live GPU arch (get_gfx_runtime()).
• Each row picks triton (the existing fp8_mqa_logits, Gluon on gfx950) or flydsl with a variant from KERNEL_VARIANTS.
• A tuner. It checks every candidate against a torch reference before timing it, then re-times the three fastest interleaved and keeps the best median.
• gfx950 rows for the GLM-5 indexer (H=32, D=128): 24 shapes, seq_len 256 to 8192 and seq_len_kv 4096 to 60000, limited to what vLLM produces (seq_len <= seq_len_kv, logits within the default 512 MB VLLM_SPARSE_INDEXER_MAX_LOGITS_MB).

Without a matching row it runs fp8_mqa_logits, so an empty table changes nothing. A FlyDSL row is only used where the kernel can run (gfx942/gfx950, the variant exists, native e4m3fn on gfx950); otherwise the call runs fp8_mqa_logits.

The CSVs are named like fused MoE's (tuned_mqa_logits.csv, model_configs/glm5_tuned_mqa_logits.csv). A mqa_logits_tuned name would also match paged_mqa_logits_tuned_*.csv in the model_configs glob.

Test Result

MI355X (gfx950), AITER 1a37aa008 with this change applied, Triton 3.8.0, FlyDSL 0.3.4.1.

• FlyDSL wins all 24 shapes, mostly with mfma32x32x64_bkv128_r2_w4_lds3 (13 of 24).
• Against Triton/Gluon, what vLLM runs on gfx950 today: -36.6% geomean, from -16.6% to -74.5%.
• Against FlyDSL's built-in variant choice (what [ROCm][Perf] gfx950: use FlyDSL fp8 MQA logits kernel for the sparse indexer prefill vllm-project/vllm#59280 uses): -5.8% geomean, from -15.4% to +3.3%. Below about 40 us the run-to-run spread is about 10%, so those rows are close to noise.
• Every candidate matches torch (exact -inf mask, calc_diff < 1e-3); none was rejected.
• op_tests/test_mqa_logits.py passes, including the live-GPU lookup check (child process started with the other arch in GPU_ARCHS), the dispatched-config check and the FlyDSL fallback.
• --run_config: 24/24 within 10% on one pass, 22/24 on a second (the two smallest kernels, about 20 and 40 us, at +12%). Collision guard: OK.
• With GPU_ARCHS naming another arch, the FlyDSL module builds its variant list for that arch at import, so a gfx950 FlyDSL row falls back to fp8_mqa_logits instead of failing.
• Not measured end to end.

python3 csrc/mqa_logits/mqa_logits_tune.py \
    -i aiter/configs/model_configs/glm5_untuned_mqa_logits.csv \
    -o aiter/configs/model_configs/glm5_tuned_mqa_logits.csv
python3 csrc/mqa_logits/mqa_logits_tune.py --run_config aiter/configs/model_configs/glm5_tuned_mqa_logits.csv
python3 op_tests/test_mqa_logits.py
python3 -m unittest op_tests.tuning_tests.test_config_shape_collision.TestConfigShapeCollision.test_mqa_logits

All 24 shapes (us, single sweep timings; FlyDSL auto = the variant flydsl_fp8_mqa_logits picks by itself)

seq_len	seq_len_kv	Triton/Gluon	FlyDSL auto	picked variant	picked	vs Triton/Gluon	vs FlyDSL auto
256	4096	45.9	21.2	mfma32x32x64_bkv64_r1_w2_lds2	18.4	-59.9%	-13.2%
1024	4096	60.2	38.8	mfma32x32x64_bkv64_r2_w2_lds3	38.8	-35.5%	+0.0%
2048	4096	80.2	54.2	mfma32x32x64_bkv128_r2_w4_lds3	56.0	-30.2%	+3.3%
4096	4096	95.1	70.2	mfma32x32x64_bkv128_r2_w4_lds3	67.2	-29.3%	-4.2%
256	8192	76.5	30.0	mfma32x32x64_bkv64_r2_w2_lds2	29.0	-62.1%	-3.4%
1024	8192	108.5	67.5	mfma32x32x64_bkv128_r2_w4_lds3	62.7	-42.2%	-7.2%
2048	8192	145.8	103.9	mfma32x32x64_bkv64_r2_w4_lds3	98.9	-32.1%	-4.8%
4096	8192	212.7	161.2	mfma32x32x64_bkv128_r2_w4_lds3	161.1	-24.3%	-0.1%
8192	8192	271.9	204.3	mfma32x32x64_bkv64_r2_w4_lds3	204.3	-24.9%	+0.0%
256	16384	147.8	44.5	mfma32x32x64_bkv128_r2_w4_lds3	37.7	-74.5%	-15.4%
1024	16384	200.4	111.8	mfma32x32x64_bkv64_r2_w4_lds3	111.7	-44.2%	-0.1%
2048	16384	271.2	209.9	mfma32x32x64_bkv128_r2_w4_lds3	199.7	-26.4%	-4.8%
4096	16384	440.0	376.1	mfma32x32x64_bkv128_r2_w4_lds3	345.0	-21.6%	-8.3%
8192	16384	668.9	592.8	mfma32x32x64_bkv128_r2_w4_lds3	557.1	-16.7%	-6.0%
256	32768	140.9	69.0	mfma32x32x64_bkv128_r2_w4_lds3	62.2	-55.9%	-9.9%
1024	32768	288.1	218.3	mfma32x32x64_bkv128_r2_w4_lds3	207.9	-27.8%	-4.8%
2048	32768	514.1	417.4	mfma32x32x64_bkv128_r2_w4_lds3	380.6	-26.0%	-8.8%
4096	32768	826.0	778.4	mfma32x32x64_bkv128_r2_w4_lds2	689.1	-16.6%	-11.5%
256	49152	150.1	88.7	mfma32x32x64_bkv64_r2_w4_lds3	89.1	-40.6%	+0.4%
1024	49152	400.4	318.1	mfma32x32x64_bkv128_r2_w4_lds3	297.6	-25.7%	-6.4%
2048	49152	698.0	611.3	mfma32x32x64_bkv128_r2_w4_lds2	557.8	-20.1%	-8.8%
256	60000	193.2	112.3	mfma32x32x64_bkv64_r2_w4_lds3	109.0	-43.6%	-2.9%
1024	60000	476.2	398.8	mfma32x32x64_bkv128_r2_w4_lds3	359.7	-24.5%	-9.8%
2048	60000	820.6	740.6	mfma32x32x64_bkv128_r2_w4_lds2	661.0	-19.4%	-10.8%

[coderabbitai]

Caution

CodeRabbit couldn't post its review summary.

Error details

HttpError

[coderabbitai]

Important

• 🔍 Trigger review

This repository does not receive automatic reviews because it has fewer than 10 stars.

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: f723eb13-0f98-429f-8fb8-fdd97721d602

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}