You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[FlyDSL] Read (B, next_n) context lengths in the paged MQA logits kernel - #6450
Currently, vLLM passes context_lens to the paged MQA logits kernels as a 2D (B, next_n) table, the same layout DeepGEMM uses. The FlyDSL kernel only reads a 1D [B] vector, so with MTP on, flydsl_fp8_paged_mqa_logits fails with a reshape error:
RuntimeError: shape '[4]' is invalid for input of size 12
This PR lets the FlyDSL paged_mqa_logits kernel accept the 2D context_lens that vLLM sends with MTP, while staying backward compatible with the 1D format.
Split out of #6153 as requested in review (one backend per PR); the Triton and Gluon change is in #6451. The two PRs touch different files and can merge in either order.
Example: 4 requests, next_n=3, context lengths 8192 / 1500 / 3000 / 700, vLLM sends the following 2D array:
b
j=0
j=1
j=2
0
8190
8191
8192
1
1498
1499
1500
2
2998
2999
3000
3
698
699
700
The fix uses the row length (next_n) as the stride and reads the last entry of each row, which gives [8192, 1500, 3000, 700].
What changes
The stride is a build-time constant. The kernel reads request b's length at b * ctx_stride + ctx_stride - 1; stride 1 builds the pre-PR kernel (same code and name), stride n builds a variant named ..._cs{n}, and _compile() caches one build per stride.
1D [B]: stride is 1, so this is context_lens[b], the same load as today.
2D (B, n): stride is n, so this is the last entry of row b, which is the full sequence length.
The wrapper picks the stride from the tensor's rank, and raises ValueError unless context_lens is a contiguous int32 [B] or (B, n) on the same device as q_fp8, with exactly B rows and n equal to 1 or next_n:
a non-contiguous input is rejected rather than copied, so there is no hidden allocation per call. Before, reshape silently copied some strided inputs and kept others as strided views that the kernel read as dense.
a (B, 0) table would have made the kernel read before the start of the tensor, and a (B, n) with n other than next_n read the wrong column as the length.
a wrong row count raised a RuntimeError from reshape before; it is now a ValueError with the expected shape.
the kernel reads context_lens as int32 words, so an int64 table gave wrong lengths before.
Backward compatibility
A contiguous [B] is unchanged, and (B, 1) behaves exactly like [B].
A contiguous (B, next_n) table is the only new input. vLLM already keeps its decode seq_lens contiguous for this.
A strided, non-int32 or off-device context_lens now raises ValueError.
Tests
In op_tests/test_flydsl_fp8_paged_mqa_logits.py:
The context_lens layout is a new axis of the existing sweep, -c/--ctx_lens (default 1d 2d). For 2d, FlyDSL gets the (B, next_n) table vLLM builds, is graded against the same fp32 reference as 1d, and is timed next to it in the summary table.
test_fp8_paged_mqa_logits_mixed_context_lens: B=5, next_n=3, distinct lengths and one empty (CUDA-graph padding) request, graded against the fp32 reference for 1d and 2d, with and without SplitKV. It fails if the kernel reads row 0's length for every request.
test_fp8_paged_mqa_logits_bad_context_lens: too few or too many rows, no columns, a column count other than next_n, an int64 table, a CPU table, a strided [B] and a non-contiguous (B, n) must raise ValueError.
Results
MI355X, Triton 3.8, FlyDSL 0.3.4.1, upstream main at 9fd025efd plus this change:
python op_tests/test_flydsl_fp8_paged_mqa_logits.py: all tables pass. In the default sweep (batch 1-8, next_n 1-8, 1k context) 2d has the same error against the reference as 1d in every case, and its kernel time is within +0.1% to +1.9% of 1d.
The new test fails on main without this change, with the reshape error above.
black and ruff (0.16.0) are clean.
Kernel time at the GLM-5 indexer shape with 8k context: 1D matches the kernel on main (e.g. B=16, next_n=3: 30.6 / 30.2 us vs 30.7 / 30.2 us over two runs), and 2D stays within about 1% of 1D.
ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm
vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
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 6450 --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.
… sweep
-c/--ctx_lens (default 1d 2d) runs the 2D table through the existing
sweep, graded against the fp32 reference and timed next to 1D. The
rejection checks stay as a small separate test.
This causal bound uses the padded next_n, but the kernel uses nn = min(next_n_lens[b], next_n) when next_n_lens is supplied. For example, with next_n=8 and three live rows, row 0 includes position L_b-3, rather than stopping at L_b-8. Document the bound using the live row count so callers get the correct ragged-query semantics.
One case with B=5, next_n=3, distinct lengths and an empty (CUDA-graph
padding) request, graded against the fp32 reference for 1D and 2D, with
and without SplitKV. The 2D table clamps at 0 like vLLM.
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
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.
samremes
left a comment
Motivation
Currently, vLLM passes
context_lensto the paged MQA logits kernels as a 2D(B, next_n)table, the same layout DeepGEMM uses. The FlyDSL kernel only reads a 1D[B]vector, so with MTP on,flydsl_fp8_paged_mqa_logitsfails with areshapeerror:This PR lets the FlyDSL paged_mqa_logits kernel accept the 2D
context_lensthat vLLM sends with MTP, while staying backward compatible with the 1D format.Split out of #6153 as requested in review (one backend per PR); the Triton and Gluon change is in #6451. The two PRs touch different files and can merge in either order.
Example: 4 requests,
next_n=3, context lengths 8192 / 1500 / 3000 / 700, vLLM sends the following 2D array:bj=0j=1j=2The fix uses the row length (
next_n) as the stride and reads the last entry of each row, which gives [8192, 1500, 3000, 700].What changes
b's length atb * ctx_stride + ctx_stride - 1; stride 1 builds the pre-PR kernel (same code and name), stridenbuilds a variant named..._cs{n}, and_compile()caches one build per stride.[B]: stride is 1, so this iscontext_lens[b], the same load as today.(B, n): stride isn, so this is the last entry of rowb, which is the full sequence length.ValueErrorunlesscontext_lensis a contiguous int32[B]or(B, n)on the same device asq_fp8, with exactlyBrows andnequal to 1 ornext_n:reshapesilently copied some strided inputs and kept others as strided views that the kernel read as dense.(B, 0)table would have made the kernel read before the start of the tensor, and a(B, n)withnother thannext_nread the wrong column as the length.RuntimeErrorfromreshapebefore; it is now aValueErrorwith the expected shape.context_lensas int32 words, so an int64 table gave wrong lengths before.Backward compatibility
[B]is unchanged, and(B, 1)behaves exactly like[B].(B, next_n)table is the only new input. vLLM already keeps its decodeseq_lenscontiguous for this.context_lensnow raisesValueError.Tests
In
op_tests/test_flydsl_fp8_paged_mqa_logits.py:context_lenslayout is a new axis of the existing sweep,-c/--ctx_lens(default1d 2d). For2d, FlyDSL gets the(B, next_n)table vLLM builds, is graded against the same fp32 reference as1d, and is timed next to it in the summary table.test_fp8_paged_mqa_logits_mixed_context_lens: B=5,next_n=3, distinct lengths and one empty (CUDA-graph padding) request, graded against the fp32 reference for1dand2d, with and withoutSplitKV. It fails if the kernel reads row 0's length for every request.test_fp8_paged_mqa_logits_bad_context_lens: too few or too many rows, no columns, a column count other thannext_n, an int64 table, a CPU table, a strided[B]and a non-contiguous(B, n)must raiseValueError.Results
MI355X, Triton 3.8, FlyDSL 0.3.4.1, upstream
mainat9fd025efdplus this change:python op_tests/test_flydsl_fp8_paged_mqa_logits.py: all tables pass. In the default sweep (batch 1-8,next_n1-8, 1k context)2dhas the same error against the reference as1din every case, and its kernel time is within +0.1% to +1.9% of1d.mainwithout this change, with thereshapeerror above.blackandruff(0.16.0) are clean.main(e.g. B=16,next_n=3: 30.6 / 30.2 us vs 30.7 / 30.2 us over two runs), and 2D stays within about 1% of 1D.