Fix Step3.5 MTP local argmax draft step - #8
Closed
voipmonitor wants to merge 3 commits into
Closed
Conversation
Automatically apply B12X virtual TP padding when supported B12X attention/MoE paths need compatible TP shapes. Remove the public virtual TP sharding and alignment CLI/config knobs, use fixed B12X capability alignments, and warn when automatic padding is enabled. Co-authored-by: OpenAI Codex <codex@openai.com>
📝 WalkthroughWalkthroughThis PR refactors vLLM's virtual tensor parallelism (TP) configuration and sparse attention indexer handling. It removes explicit virtual-TP sharding CLI parameters and configuration fields, transitioning to a plan-based B12X padding approach. It migrates sparse attention and MLA indexers from compressed/integration APIs to paged-indexer APIs, and adds distributed context parallelism (DCP) support for global sequence length handling in decode tensor preparation. ChangesVirtual TP + Sparse/MLA Indexer Refactoring
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
serve-glm51.sh (1)
20-24: ⚡ Quick winRedundant environment variable assignments.
Several variables are being set twice with conflicting patterns:
OMP_NUM_THREADSis set with a default on line 14, then hard-set on line 23NCCL_IB_DISABLEis set with a default on line 19, then hard-set on line 20The hard-set assignments on lines 20-23 make the conditional defaults on lines 14 and 19 ineffective. Consider removing either the conditional defaults or the hard-set assignments to avoid confusion.
♻️ Proposed fix to remove redundant conditional defaults
export PYTHONPATH="${SCRIPT_DIR}${PYTHONPATH:+:${PYTHONPATH}}" -export OMP_NUM_THREADS="${OMP_NUM_THREADS:-16}" export SAFETENSORS_FAST_GPU="${SAFETENSORS_FAST_GPU:-1}" export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-2,3,4,5,6,7,8,9}" export CUTE_DSL_ARCH="${CUTE_DSL_ARCH:-sm_120a}" export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-32}" -export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-1}" export NCCL_IB_DISABLE=1 export NCCL_P2P_LEVEL=SYS export NCCL_PROTO=LL,LL128,Simple export OMP_NUM_THREADS=16 export LLM_WORKER_MULTIPROC_METHOD=spawn🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@serve-glm51.sh` around lines 20 - 24, The script currently sets NCCL_IB_DISABLE and OMP_NUM_THREADS twice (conditional defaults earlier and hard-set later), causing the defaults to be ineffective; remove the duplicate hard-set exports for NCCL_IB_DISABLE and OMP_NUM_THREADS (or alternatively convert the hard-set lines for NCCL_IB_DISABLE, NCCL_P2P_LEVEL, NCCL_PROTO, OMP_NUM_THREADS, LLM_WORKER_MULTIPROC_METHOD into conditional/parameterized exports) so each environment variable is defined only once and the intended conditional defaults take effect; target the export statements for NCCL_IB_DISABLE and OMP_NUM_THREADS (and optionally NCCL_P2P_LEVEL, NCCL_PROTO, LLM_WORKER_MULTIPROC_METHOD) in the snippet and either delete the redundant hard-set lines or make them conditional (preserving the original default behavior).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@festr.sh`:
- Line 1: This script lacks a shebang which can cause it to run under an
unexpected shell; edit festr.sh to add a POSIX shebang as the very first line
(use env to locate bash or a plain /bin/sh shebang) so the script executes with
the intended shell, and ensure the file is marked executable (chmod +x) after
updating; leave the existing docker run -d \ command unchanged.
---
Nitpick comments:
In `@serve-glm51.sh`:
- Around line 20-24: The script currently sets NCCL_IB_DISABLE and
OMP_NUM_THREADS twice (conditional defaults earlier and hard-set later), causing
the defaults to be ineffective; remove the duplicate hard-set exports for
NCCL_IB_DISABLE and OMP_NUM_THREADS (or alternatively convert the hard-set lines
for NCCL_IB_DISABLE, NCCL_P2P_LEVEL, NCCL_PROTO, OMP_NUM_THREADS,
LLM_WORKER_MULTIPROC_METHOD into conditional/parameterized exports) so each
environment variable is defined only once and the intended conditional defaults
take effect; target the export statements for NCCL_IB_DISABLE and
OMP_NUM_THREADS (and optionally NCCL_P2P_LEVEL, NCCL_PROTO,
LLM_WORKER_MULTIPROC_METHOD) in the snippet and either delete the redundant
hard-set lines or make them conditional (preserving the original default
behavior).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 40c48c03-fd07-4e35-9385-99a9fe1d0d83
📥 Commits
Reviewing files that changed from the base of the PR and between 589045f and 00e862f2b2d844d9aafceac2469436cd6d5abd18.
📒 Files selected for processing (18)
festr.shserve-ds4-flash.shserve-ds4-pro-tp10.shserve-glm51-tp6.shserve-glm51.shtests/config/test_virtual_tp.pytests/engine/test_arg_utils.pytests/model_executor/layers/test_sparse_attn_indexer_b12x.pyvllm/config/__init__.pyvllm/config/parallel.pyvllm/config/speculative.pyvllm/config/virtual_tp.pyvllm/engine/arg_utils.pyvllm/model_executor/layers/sparse_attn_indexer.pyvllm/model_executor/virtual_tp.pyvllm/models/deepseek_v4/attention.pyvllm/v1/attention/backends/mla/indexer.pyvllm/v1/spec_decode/step3p5.py
💤 Files with no reviewable changes (6)
- serve-glm51-tp6.sh
- vllm/config/init.py
- serve-ds4-flash.sh
- vllm/config/speculative.py
- vllm/config/parallel.py
- vllm/engine/arg_utils.py
| @@ -0,0 +1,67 @@ | |||
| docker run -d \ | |||
There was a problem hiding this comment.
Add shebang to make script executable.
The script is missing a shebang line, which means it may not execute correctly or may execute with an unexpected shell. This can lead to portability and compatibility issues.
🔧 Proposed fix to add shebang
+#!/usr/bin/env bash
docker run -d \📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| docker run -d \ | |
| #!/usr/bin/env bash | |
| docker run -d \ |
🧰 Tools
🪛 Shellcheck (0.11.0)
[error] 1-1: Tips depend on target shell and yours is unknown. Add a shebang or a 'shell' directive.
(SC2148)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@festr.sh` at line 1, This script lacks a shebang which can cause it to run
under an unexpected shell; edit festr.sh to add a POSIX shebang as the very
first line (use env to locate bash or a plain /bin/sh shebang) so the script
executes with the intended shell, and ensure the file is marked executable
(chmod +x) after updating; leave the existing docker run -d \ command unchanged.
Source: Linters/SAST tools
Pass spec_step_idx through the local-argmax draft token path instead of calling get_top_tokens() with its default step 0. DeepSeek V4/DS4 MTP has one head per speculative step; using step 0 for later draft positions produced near-zero acceptance and could crash the worker under CUDA graphs. Reuse the existing _model_get_top_tokens helper so models that do not accept spec_step_idx keep the previous call shape.
voipmonitor
force-pushed
the
codex/ds4-step35-mtp-local-argmax-spec-step-20260608
branch
from
June 8, 2026 06:39
00e862f to
84d96db
Compare
lukealonso
force-pushed
the
dev/black-benediction
branch
from
June 8, 2026 16:36
1dd7928 to
bb6c5b7
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
2 participants
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.
Summary
Fix Step3.5 MTP greedy local-argmax draft selection so each speculative step calls the model through
_model_get_top_tokens(hidden_states, spec_step_idx)instead of directly callingmodel.get_top_tokens(hidden_states).The direct call dropped
spec_step_idx, which made multi-step MTP models such as DeepSeek V4 / GLM use the default MTP head for every draft position. Withuse_local_argmax_reduction=true, later draft positions had near-zero acceptance and could crash under CUDA graphs. The helper preserves compatibility with models whoseget_top_tokensdoes not acceptspec_step_idx.Validation
4 / 158, per-position0.051, 0.000, followed by a CUDA illegal memory access on a subsequent smoke request.113 / 146, per-position0.890, 0.658.200; 4/4 smoke prompts coherent withcjk=0; acceptance163 / 206 = 79.1%, per-position0.961,0.621.No production B12X binding/workspace behavior is changed.