[Bugfix][Spec Decode] Fix DeepSeek/GLM MTP with sequence-parallel MoE and warm spec-sized logits all-gather#48572
Open
haosdent wants to merge 1 commit into
Open
[Bugfix][Spec Decode] Fix DeepSeek/GLM MTP with sequence-parallel MoE and warm spec-sized logits all-gather#48572haosdent wants to merge 1 commit into
haosdent wants to merge 1 commit into
Conversation
mergify
Bot
added
deepseek
… and warm spec-sized logits all-gather Fixes DeepSeek/GLM-5.2 MTP speculative decoding (vllm-project#48568): 1. Restore full tokens in DeepSeekMultiTokenPredictorLayer.forward under sequence-parallel MoE (mtp_block reduce-scatters its output), mirroring DeepseekV2Model.forward and the Qwen MTP fix (vllm-project#48429). 2. Warm compute_logits at the spec-decode row count during init so the larger vocab-parallel all-gather connects its RCCL channels before a live decode step ("no transport for peer on channel N"). Signed-off-by: haosdent <haosdent@gmail.com>
Contributor
Author
|
@AndreasKaratzas can you help to take a look, this is to fix MTP issues when I run it in AMD Mi300x |
haosdent
changed the title
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
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.
Purpose
Two platform-independent issues prevent DeepSeek/GLM-5.2 MTP speculative decoding from working (surfaced on MI300X, #48568):
Sequence-parallel MoE token restore. Under sequence-parallel MoE (expert-parallel + TP>1),
mtp_blockreduce-scatters its output, butDeepSeekMultiTokenPredictorLayer.forwardnever restored full tokens beforeshared_head/compute_logits, so ranks fed mismatched token counts into the vocab-parallel collective and it deadlocked. Restore full tokens, mirroringDeepseekV2Model.forwardand the Qwen MTP fix ([BugFix] Restore full tokens for Qwen MTP When MoE SP #48429).Spec-sized logits all-gather warmup. Warmup only ran
compute_logits' vocab-parallel all-gather at 1 row per request, but spec decode verifiesmax_num_reqs*(1+num_speculative_tokens)rows. On ROCm/RCCL the larger message connects extra channels lazily on first use; when that first use is a live decode step the connect fails (no transport for peer on channel N). Warmcompute_logitsat the spec row count during init.Test Plan
Serve GLM-5.2-FP8 with MTP speculative decoding on 8x MI300X (gfx942):
Test Result
Validated on 8x MI300X: the server starts and serves normally with MTP speculative decoding active, producing coherent completions. Without this change the engine hangs in the logits all-gather on the first speculative-decode step.