[Triton/Gluon] [HIP] [CK] Read (B, next_n) context lengths in the paged MQA logits kernels

[amd-sriram]

Motivation

vLLM's sparse indexer now builds context lengths as a 2D (B, next_n) per-row table for every caller ((B, 1) without MTP), the layout DeepGEMM's fp8_paged_mqa_logits takes. The AITER paged MQA-logits kernels read one length per sequence (load(context_len_ptr + pid_batch)) and rebuild the per-row bound themselves. With MTP on, the 2D table therefore gives the kernels the wrong lengths: silently for Gluon/Triton, and as a reshape error for FlyDSL.

vllm-project/vllm#59420 works around it on the vLLM side by copying the last column into a 1D tensor before each call. This makes the kernels read the 2D table directly, so vLLM can pass it as-is and nothing is copied.

Technical details

• Each kernel reads sequence b's length at b * ContextLenStride + ContextLenStride - 1: element b of [B] (stride 1, today's code), or the last entry of row b of a contiguous (B, n) table (stride n). The per-row causal bound is still rebuilt from that length, unchanged.
• Gluon (_gluon_deepgemm_fp8_paged_mqa_logits, _preshuffle, _preshuffle_varctx) and Triton (_deepgemm_fp8_paged_mqa_logits, _stage1, _varctx_schedule): new ContextLenStride constexpr, default 1, so the [B] path compiles to the same code as before.
• FlyDSL flydsl_fp8_paged_mqa_logits: new runtime ctx_stride kernel argument. A [B, n] table must be contiguous. A strided [B] is now made contiguous: before, reshape(batch_size) kept the strided view and the kernel read it as dense.
• Wrappers derive the stride from the tensor (_context_len_stride); a 2D table must be contiguous.
• AOT Gluon (Triton < 3.4): a stride other than 1 adds _S<n> to the kernel name, so a missing prebuilt binary fails loudly instead of running the 1D one.
• The ragged-K variants take prefix_sum_context_lens and are unchanged.
• Contract: the per-row table is the standard one, row j of a sequence of length L seeing L - next_n + 1 + j tokens, which is what vLLM builds.

Test Plan

• New test_paged_mqa_logits_2d_context_lens in op_tests/test_pa_mqa_logits_offset.py: for Gluon plain, Gluon preshuffle and FlyDSL (gfx950), a (B, next_n) table must give exactly the logits of the per-sequence lengths. Mixed lengths (3000, 257, 64, 3), next_n=3, KVBlockSize 64.
• Full op_tests/test_pa_mqa_logits_offset.py, 10 times.
• Kernel time, 1D vs 2D, cold cache (run_perftest rotates inputs), GLM-5 indexer shape, 8k context.
• Stage1 and the varctx schedule: 2D vs 1D over 100 repeated calls.

Test Result

MI355X, AITER 48b1365 with this change applied (the touched files are identical to main), Triton 3.8, FlyDSL 0.3.4.1:

• FlyDSL [B] input: a strided tensor now gives the same logits as a contiguous one; (B, 1) matches [B].
• Without the change the new test fails for every entry point. With it: test_pa_mqa_logits_offset.py 108/108, 10/10 runs.
• Raw 2D table vs per-sequence lengths, varied lengths, B=16, next_n=3: bit-identical for Gluon and FlyDSL. Stage1 and varctx: 100/100 calls identical.
• No copy and no measurable cost:

path	B	next_n	1D us	2D us	delta
Gluon preshuffle	16	3	10.02	10.01	-0.1%
Gluon preshuffle	128	3	47.61	47.42	-0.4%
Gluon preshuffle	16	5	14.49	14.47	-0.1%
Gluon preshuffle	128	5	70.50	70.00	-0.7%
FlyDSL	16	3	14.96	15.10	+0.9%
FlyDSL	128	3	45.33	45.24	-0.2%
FlyDSL	16	5	19.22	19.50	+1.5%
FlyDSL	128	5	60.06	60.52	+0.8%

• Not tested: the AOT Gluon path (needs Triton < 3.4).

Found while testing (not changed here)

With the new test's inputs, two paths are not reproducible run to run even with 1D lengths, so the test leaves them out:

• deepgemm_fp8_paged_mqa_logits_stage1: 11/11 repeated calls with identical inputs gave a different -inf mask from the first, on rows 1 and 3, columns 0..71.
• Gluon preshuffle with VarCtxSchedule: up to 5/200 repeated call pairs differ (the schedule value itself is stable at 1).

Neither is on the vLLM gfx942/gfx950 decode path. Worth a separate issue.

[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: f82d757f-364f-4366-b17e-ebc8a181c2d0

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

[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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
multigpu	Aiter multi-GPU tests on the 8-GPU runner
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 9 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.