Skip to content

[Cherry-Pick][BugFix] fix MTP bugs in TP and overlap(#7172) #7192

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
freeliuzc merged 2 commits into from
Apr 8, 2026
+7 −6
Merged

[Cherry-Pick][BugFix] fix MTP bugs in TP and overlap(#7172) #7192

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
be9287d
fix MTP bugs in TP and overlap
huicongyao Apr 3, 2026
278ba14
fix
huicongyao Apr 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions custom_ops/gpu_ops/speculate_decoding/speculate_save_output.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,9 @@ void SpeculateSaveWithOutputMsg(const paddle::Tensor& accept_tokens,
int msg_queue_id,
int save_each_rank,
bool skip_prefill) {
// printf("enter save output");
if (!save_each_rank && rank_id > 0) {
// NOTE(yaohuicong): Skip non-zero TP ranks — they share identical sampling
// outputs, so only rank 0 needs to send results to the message queue.
if (rank_id > 0) {
return;
}

Expand Down
4 changes: 3 additions & 1 deletion custom_ops/gpu_ops/speculate_decoding/speculate_save_output_with_topk.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,9 @@ void SpeculateSaveOutMmsgTopK(const paddle::Tensor& sampled_token_ids,
int message_flag, // Target: 3, Draft: 4
int64_t rank_id,
bool save_each_rank) {
if (!save_each_rank && rank_id > 0) {
// NOTE(yaohuicong): Skip non-zero TP ranks — they share identical sampling
// outputs, so only rank 0 needs to send results to the message queue.
if (rank_id > 0) {
return;
}

Expand Down
4 changes: 1 addition & 3 deletions fastdeploy/worker/gpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,9 +345,7 @@ def _predict_next_launch_token_num(self) -> int:
is_block_step_cpu = self.share_inputs["is_block_step_cpu"].numpy()
next_real_bsz = (seq_lens_this_time_cpu > 0).sum().item() + (is_block_step_cpu > 0).sum().item()
token_num_one_step = (self.speculative_config.num_speculative_tokens + 1) if self.speculative_decoding else 1
next_launch_token_num = (
seq_lens_this_time_cpu.sum().item() + is_block_step_cpu.sum().item() * token_num_one_step
)
next_launch_token_num = next_real_bsz * token_num_one_step
return next_launch_token_num, next_real_bsz

def only_prefill(self):
Expand Down
Loading