Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
Three correctness fixes for the SM120/SM121 (consumer Blackwell / GB10) serving path, ported from open upstream PRs: 1. MXFP4 persistent-kernel constraint must cover sm120/sm121 (upstream vllm-project#52850): has_native_mxfp() is cuda_capability_geq(10,0), so gating the is_persistent constraint on the sm100 family alone leaves consumer Blackwell requiring a constraint nothing sets and every native-MXFP matmul raises. Gate on >=100 instead. Adds test_swizzle_mxfp4_sets_persistent_for_native_mxfp. 2. DSpark spec-decode reshape guard in the SM120 FlashInfer sparse decode (upstream vllm-project#52499): flashinfer's dsv4 decode API expects [batch, q_len_per_request, ...] queries; flattened spec-decode batches misroute to the varlen prefill kernel which asserts num_tokens > 64. Reshape to [num_decodes, next_n, ...] when tokens exceed requests, and handle empty/small (<=64) prefill chunks via the uniform decode-form call. 3. Ragged decode batches take the padded indexer path (upstream vllm-project#52500): requires_padding can be computed False for ragged warmup/mixed batches (SM120 TP=2: 8 tokens over 6 seqs), where the uniform reshape would crash. Also pad when num_decode_tokens does not divide the decode batch. Verified: all four files compile; pre-commit clean. test_flashmla_sparse.py failures are pre-existing (no compiled extensions in this venv; identical 4 failures with changes stashed). Co-authored-by: OMP Agent <noreply@omp.local> Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
|
Adoption data point: independently ported into a third-party GB10 (DGX Spark) serving fork alongside #52499 — randomvariable/vllm@07011d8 (2026-08-19). |
Three correctness fixes for the SM120/SM121 (consumer Blackwell / GB10) serving path, ported from open upstream PRs: 1. MXFP4 persistent-kernel constraint must cover sm120/sm121 (upstream vllm-project#52850): has_native_mxfp() is cuda_capability_geq(10,0), so gating the is_persistent constraint on the sm100 family alone leaves consumer Blackwell requiring a constraint nothing sets and every native-MXFP matmul raises. Gate on >=100 instead. Adds test_swizzle_mxfp4_sets_persistent_for_native_mxfp. 2. DSpark spec-decode reshape guard in the SM120 FlashInfer sparse decode (upstream vllm-project#52499): flashinfer's dsv4 decode API expects [batch, q_len_per_request, ...] queries; flattened spec-decode batches misroute to the varlen prefill kernel which asserts num_tokens > 64. Reshape to [num_decodes, next_n, ...] when tokens exceed requests, and handle empty/small (<=64) prefill chunks via the uniform decode-form call. 3. Ragged decode batches take the padded indexer path (upstream vllm-project#52500): requires_padding can be computed False for ragged warmup/mixed batches (SM120 TP=2: 8 tokens over 6 seqs), where the uniform reshape would crash. Also pad when num_decode_tokens does not divide the decode batch. Verified: all four files compile; pre-commit clean. test_flashmla_sparse.py failures are pre-existing (no compiled extensions in this venv; identical 4 failures with changes stashed). Co-authored-by: OMP Agent <noreply@omp.local> Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
Three correctness fixes for the SM120/SM121 (consumer Blackwell / GB10) serving path, ported from open upstream PRs: 1. MXFP4 persistent-kernel constraint must cover sm120/sm121 (upstream vllm-project#52850): has_native_mxfp() is cuda_capability_geq(10,0), so gating the is_persistent constraint on the sm100 family alone leaves consumer Blackwell requiring a constraint nothing sets and every native-MXFP matmul raises. Gate on >=100 instead. Adds test_swizzle_mxfp4_sets_persistent_for_native_mxfp. 2. DSpark spec-decode reshape guard in the SM120 FlashInfer sparse decode (upstream vllm-project#52499): flashinfer's dsv4 decode API expects [batch, q_len_per_request, ...] queries; flattened spec-decode batches misroute to the varlen prefill kernel which asserts num_tokens > 64. Reshape to [num_decodes, next_n, ...] when tokens exceed requests, and handle empty/small (<=64) prefill chunks via the uniform decode-form call. 3. Ragged decode batches take the padded indexer path (upstream vllm-project#52500): requires_padding can be computed False for ragged warmup/mixed batches (SM120 TP=2: 8 tokens over 6 seqs), where the uniform reshape would crash. Also pad when num_decode_tokens does not divide the decode batch. Verified: all four files compile; pre-commit clean. test_flashmla_sparse.py failures are pre-existing (no compiled extensions in this venv; identical 4 failures with changes stashed). Co-authored-by: OMP Agent <noreply@omp.local> Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
Three correctness fixes for the SM120/SM121 (consumer Blackwell / GB10) serving path, ported from open upstream PRs: 1. MXFP4 persistent-kernel constraint must cover sm120/sm121 (upstream vllm-project#52850): has_native_mxfp() is cuda_capability_geq(10,0), so gating the is_persistent constraint on the sm100 family alone leaves consumer Blackwell requiring a constraint nothing sets and every native-MXFP matmul raises. Gate on >=100 instead. Adds test_swizzle_mxfp4_sets_persistent_for_native_mxfp. 2. DSpark spec-decode reshape guard in the SM120 FlashInfer sparse decode (upstream vllm-project#52499): flashinfer's dsv4 decode API expects [batch, q_len_per_request, ...] queries; flattened spec-decode batches misroute to the varlen prefill kernel which asserts num_tokens > 64. Reshape to [num_decodes, next_n, ...] when tokens exceed requests, and handle empty/small (<=64) prefill chunks via the uniform decode-form call. 3. Ragged decode batches take the padded indexer path (upstream vllm-project#52500): requires_padding can be computed False for ragged warmup/mixed batches (SM120 TP=2: 8 tokens over 6 seqs), where the uniform reshape would crash. Also pad when num_decode_tokens does not divide the decode batch. Verified: all four files compile; pre-commit clean. test_flashmla_sparse.py failures are pre-existing (no compiled extensions in this venv; identical 4 failures with changes stashed). Co-authored-by: OMP Agent <noreply@omp.local> Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
Three correctness fixes for the SM120/SM121 (consumer Blackwell / GB10) serving path, ported from open upstream PRs: 1. MXFP4 persistent-kernel constraint must cover sm120/sm121 (upstream vllm-project#52850): has_native_mxfp() is cuda_capability_geq(10,0), so gating the is_persistent constraint on the sm100 family alone leaves consumer Blackwell requiring a constraint nothing sets and every native-MXFP matmul raises. Gate on >=100 instead. Adds test_swizzle_mxfp4_sets_persistent_for_native_mxfp. 2. DSpark spec-decode reshape guard in the SM120 FlashInfer sparse decode (upstream vllm-project#52499): flashinfer's dsv4 decode API expects [batch, q_len_per_request, ...] queries; flattened spec-decode batches misroute to the varlen prefill kernel which asserts num_tokens > 64. Reshape to [num_decodes, next_n, ...] when tokens exceed requests, and handle empty/small (<=64) prefill chunks via the uniform decode-form call. 3. Ragged decode batches take the padded indexer path (upstream vllm-project#52500): requires_padding can be computed False for ragged warmup/mixed batches (SM120 TP=2: 8 tokens over 6 seqs), where the uniform reshape would crash. Also pad when num_decode_tokens does not divide the decode batch. Verified: all four files compile; pre-commit clean. test_flashmla_sparse.py failures are pre-existing (no compiled extensions in this venv; identical 4 failures with changes stashed). Co-authored-by: OMP Agent <noreply@omp.local> Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
There was a problem hiding this comment.
looks like we also need to add num_decode_tokens % decode_lens.shape[0] != 0 for unpacking
There was a problem hiding this comment.
Addressed in 2c4e78e, now rebased onto current main. The padding predicate is computed once:
needs_padded_path = decode_metadata.requires_padding or (
num_decode_tokens % decode_lens.shape[0] != 0
)Both the packing branch and the matching unpack_seq_triton branch use needs_padded_path. For 8 tokens across 6 requests, the divisibility fallback selects the padded path for both operations, preventing the padded-pack/unpadded-unpack reshape mismatch you identified.
The rebase preserved both patches unchanged; syntax and applicable pre-commit checks passed. The separate evenly divisible [1, 2, 3] case still relies on correct requires_padding metadata, as acknowledged in my earlier comment.
| elif decode_metadata.requires_padding or ( | ||
| num_decode_tokens % decode_lens.shape[0] != 0 |
There was a problem hiding this comment.
what if the decode_lens=[1, 2, 3], batch_size = 3, 6 % 3 == 0? In this case, it still needs padding.
There was a problem hiding this comment.
You're right that 6 % 3 == 0 does not establish uniformity. My earlier reply overstated the gap for this specific example: with consistent metadata, the existing preparation handles it.
- Native preparation sets
requires_padding = min_decode_len != max_decode_len. For[1, 2, 3], this isTrue, so the sharedneeds_padded_pathflag selects both padded packing and unpacking even though the modulo term is false. - Flattened preparation converts the requests into six one-token rows, with
decode_lens = [1, 1, 1, 1, 1, 1]. In that representation,requires_padding=Falseand an unpadded reshape are appropriate.
See the metadata preparation on the rebased head.
I checked seven isolated CPU cases using the production _prepare_decode_tensors method with real PyTorch tensors, including native/flattened [1, 2, 3], padded token budgets, and eight tokens across six requests. All passed. These validate metadata preparation only; GPU inference was not rerun.
The modulo condition remains an additional fallback, not a replacement for the existing uniformity check. I haven't reproduced a current failure for [1, 2, 3] with consistent metadata, so I haven't added another condition for that example.
…and unpack The pack site checked 'requires_padding or tokens % batch_size != 0' but the matching unpack site only checked 'requires_padding' - so a batch that took the padded pack path via the divisibility fallback would then take the WRONG unpack path, since the two conditions could disagree. Factor the check into one needs_padded_path shared by both sites so they can never disagree. Reported by @Jie-Fang on vllm-project#52500. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01C7YveGPhnT8v6qrAEucmVW Signed-off-by: pavelzak <pavel.zakharov@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 10 included reviews per hour; 8 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe decode path now selects padded packing when decode token counts are not divisible by batch size. The same condition controls unpacking, keeping both branches aligned for ragged warmup and mixed batches. ChangesDecode padding handling
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to Ragged decode batches that previously could fail the uniform reshape now use matching padded packing and unpacking, with no remaining merge-blocking risk identified. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
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. Comment |
|
@Jie-Fang thanks for catching this. Confirmed and reproduced directly: with the old asymmetric check, pack takes the padded path for an 8-tokens-over-6-seqs ragged batch ( On your |
|
LGTM, thanks! |
…dexer decode_metadata.requires_padding can be computed False for ragged mixed/warmup decode batches (observed on SM120 TP=2: 8 tokens over 6 requests), in which case the uniform `num_decode_tokens // batch` reshape below crashes. Also take the padded path whenever the decode token count does not divide evenly over the decode batch. Validated on 2x DGX Spark (GB10) TP=2 serving DeepSeek-V4-Flash-0731. Signed-off-by: pavelzak <pavel.zakharov@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…and unpack The pack site checked 'requires_padding or tokens % batch_size != 0' but the matching unpack site only checked 'requires_padding' - so a batch that took the padded pack path via the divisibility fallback would then take the WRONG unpack path, since the two conditions could disagree. Factor the check into one needs_padded_path shared by both sites so they can never disagree. Reported by @Jie-Fang on vllm-project#52500. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01C7YveGPhnT8v6qrAEucmVW Signed-off-by: pavelzak <pavel.zakharov@gmail.com>
ad44d02 to
2c4e78e
Compare
github-actions
Bot
added
the
ready
jeejeelee
added
the
verified
|
/ci run |
|
✅ Triggered Buildkite CI #88582 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #88601 for commit |
Include the upstream MRV2 warmup rollback (vllm-project#56654) before rerunning CI. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: pavelzak <pavel.zakharov@gmail.com>
|
/ci run |
|
✅ Triggered Buildkite CI #88661 for commit |
3 participants
Purpose
Fix a crash in
sparse_attn_indexerfor ragged decode batches.decode_metadata.requires_paddingcan be computedFalsefor ragged mixed/warmup decode batches (observed on SM120/GB10 TP=2: 8 decode tokens over 6 requests). The non-padded path then performs a uniformnum_decode_tokens // batchreshape, which crashes when the token count does not divide evenly over the decode batch.Fix: also take the padded path whenever
num_decode_tokens % decode_lens.shape[0] != 0. This is a strictly-safer superset of the current condition; evenly-divisible batches keep the fast path.Not a duplicate: no open PR touches this condition in
vllm/model_executor/layers/sparse_attn_indexer.py; #47808 (DSpark confidence-scheduled verification) restructured other parts of this file but kept therequires_padding-only condition. The metadata-side computation (requires_padding = min_decode_len != max_decode_leninbackends/mla/indexer.py) is unchanged onmain, so the ragged case is still reachable.Test Plan
python -m py_compile vllm/model_executor/layers/sparse_attn_indexer.pyTest Result
AI assistance was used for this PR (rebasing the fix from a v0.26.0-based production branch onto
main); every changed line was reviewed and validated end-to-end by the submitter on the hardware above.