[ROCm] FlyDSL paged MQA-logits kernel for the sparse indexer - #14
Open
amd-sriram wants to merge 2 commits into
Open
amd-sriram wants to merge 2 commits into
amd-sriram wants to merge 2 commits into
Conversation
The decode half of the sparse-attention indexer runs on aiter's Triton/Gluon deepgemm_fp8_paged_mqa_logits. ROCm/aiter#4221 adds a gfx950 FlyDSL implementation of the same kernel, reported at 1.3-4.2x over Gluon on MI355X, and the decode logits are the indexer's dominant per-step cost at long context. Route the gfx950 branch of rocm_fp8_paged_mqa_logits at it when VLLM_ROCM_USE_AITER_FLYDSL_PAGED_MQA_LOGITS is set, so one build can run both arms and any aiter without the PR falls back to the existing path. Off by default while the PR is unmerged. This is a second call site rather than the module swap used for the prefill kernel in the preceding commit. The two launchers take the same tensors in the same order, but the FlyDSL one's remaining knobs are keyword-only with its own tuned defaults, so sharing a caller would mean forcing the Gluon ChunkK=256 on it. KVBlockSize is not optional: the kernel asserts it against the KV cache's own block dimension, so it is passed through along with Preshuffle, exactly as the Gluon call does. The FlyDSL path skips the -inf prefill the kernel's docstring asks for. That request is for a caller that reads the whole row; this one does not. FlyDSL writes every column up to its causal bound `context_lens[b] - next_n + n` and nothing past it (`is_writer = ... & (col <= q_limit)`), and that bound is the decode top-k's read bound, so the tail the fill would cover is never written and never read. The sanitize is kept, and is the same bounded sanitize_decode_logits() the Gluon path calls. Its window is that same causal bound, so it covers exactly the columns FlyDSL wrote: no more work than the Gluon path already does, and both paths hand the top-k a workspace in the same state. Keeping it means the flag toggles the kernel and nothing else, which is what makes the two arms comparable. The FlyDSL kernel is written against gfx950, so _flydsl_paged_mqa_logits_kernel() returns None on every other architecture and the Triton/Gluon kernel is used regardless of the flag. Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
…ts kernel vLLM hands the indexer (B, next_n) per-row causal bounds; the kernel takes one length per sequence, so next_n > 1 raised on the reshape. Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
amd-sriram
force-pushed
the
rocm/flydsl-paged-mqa-logits-nextn
branch
from
September 14, 2026 16:09
48e733d to
5103db9
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Depends on ROCm/aiter#4221 (FlyDSL paged MLA indexer, gfx950).
What this does
rocm_fp8_paged_mqa_logitsat aiter's FlyDSLpaged (decode) MQA-logits kernel when
VLLM_ROCM_USE_AITER_FLYDSL_PAGED_MQA_LOGITSis set. Off by default while[Feature]: Multi-node serving vllm-project/vllm#4221 is unmerged, and any aiter without the PR falls back to Gluon.
Why (2) is needed
The kernel takes
context_lensof shape(B,)and derives each Q row's causalbound itself as
col <= context_len - next_n + row. vLLM instead precomputesthose bounds and passes
(B, next_n)(indexer.py, native spec decode) or(B, 1)otherwise. The kernel'scontext_lens.reshape(batch_size)thereforeraised for any
next_n > 1. The last column of the vLLM tensor is theper-sequence length in both layouts, so that is what gets passed.
Interaction with the next_n cap
At
num_speculative_tokens >= 2the indexer flattens decode on ROCm and thiskernel is called at
next_n = 1regardless. #13 lifts that cap.Both are needed to exercise vllm-project#4221 above
next_n = 2.Not yet verified
main.keeps, so gsm8k/niah/gpqa are still outstanding.