[amd-sriram]

Purpose

Review feedback on vllm-project#59422: vLLM should ask AITER for paged MQA logits and let AITER pick the kernel, instead of carrying a Gluon/FlyDSL switch that changes every time one of the kernels improves. This calls AITER's paged_mqa_logits (from aiter.paged_mqa_logits import paged_mqa_logits, imported lazily like aiter.fused_moe) for the sparse indexer decode logits on gfx942/gfx950. AITER picks Gluon or FlyDSL, and their launch knobs, from its tuned table (ROCm/aiter#6140), the same model as its fused-MoE dispatch.

vLLM passes the same tensors it passes to the Gluon kernel today and no backend knobs (ChunkK, WavePerEU, SplitKV). When the installed AITER has no aiter.paged_mqa_logits module, the existing Gluon call runs unchanged, so this is a no-op until vLLM's AITER pin includes the op.

Dependency

ROCm/aiter#6140 (aiter/paged_mqa_logits.py, its tuned table and tuner).

Details

• The 2D (B, next_n) seq_lens table is passed through as-is: aiter.paged_mqa_logits takes [B] or [B, next_n] and reduces to one length per sequence itself, so this path does not need [ROCm][Bugfix] Pass per-sequence context lengths to the AITER paged MQA-logits kernels vllm-project/vllm#59420.
• No change to the C4A / Triton / torch paths or to non-gfx942/gfx950 archs.

Test Plan

• pytest tests/v1/attention/test_rocm_aiter_paged_mqa_logits.py (new): the AITER op gets the caller's tensors, the 2D seq_lens and only Preshuffle / KVBlockSize; the Gluon call runs when AITER lacks the op.
• MI355X, real tensors through rocm_fp8_paged_mqa_logits with the 2D seq_lens table, against the torch reference, with [HIP] [FlyDSL] [CI] Add aiter.paged_mqa_logits with a tuned Gluon/FlyDSL table ROCm/aiter#6140 installed. Shapes picked to hit both backends in the tuned table.

Test Result

• Unit test: 2/2 pass on MI355X.
• Real tensors (AITER_LOG_TUNED_CONFIG=1 shows the pick):

batch	next_n	context	AITER pick	calc_diff vs torch
16	3	8192	gluon ChunkK=128 WavePerEU=4	9.2e-11
256	3	8192	flydsl	8.5e-11
8	1	32768	gluon ChunkK=256 WavePerEU=4	7.5e-11
64	8	8192	flydsl	8.5e-11

• Kernel level (from [HIP] [FlyDSL] [CI] Add aiter.paged_mqa_logits with a tuned Gluon/FlyDSL table ROCm/aiter#6140, MI355X, GLM-5 indexer shape, batch 1..256, next_n 1..8): the tuned pick vs the Gluon ChunkK=256, WavePerEU=2 this file passes today is -10.8% geomean at 8k context and -17.1% at 32k; FlyDSL is picked on 16 of 72 shapes, all at batch >= 32.
• End-to-end serving benchmark and accuracy: not run yet.

[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: 34637253-d135-4ea2-960a-2e40a1b22e1b

• 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.}