[Bugfix] Fix dp mtp hang (#40589)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: sherryC41 <sherry.c.c41@gmail.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
SherryC41 committed
Jul 7, 2026 at 05:08 UTC
567a78432d8e00ac17e4288a4c15ca795c6a1bb4
1 file changed
+45
-18
vllm/v1/worker/gpu_model_runner.py
+45
-18
@@ -4527,17 +4527,23 @@ class GPUModelRunner(
4527
self._copy_draft_token_ids_to_cpu(scheduler_output)
4528
4529
spec_config = self.speculative_config
4530
- propose_drafts_after_bookkeeping = False
4530
+ draft_after_bookkeeping = False
4531
if spec_config is not None:
4532
# Decide whether to run the drafter or zero out draft tokens.
4533
input_fits_in_drafter = self._input_fits_in_drafter(
4534
spec_decode_common_attn_metadata
4535
)
4536
- use_gpu_toks = (
4536
+ # Whether the drafter runs a GPU model forward (and thus carries
4537
+ # TP/EP/DP collectives), independent of padded-batch timing.
4538
+ drafter_runs_model_forward = (
4539
spec_config.use_eagle()
4540
or spec_config.uses_draft_model()
4541
or spec_config.uses_extract_hidden_states()
4540
- ) and not spec_config.disable_padded_drafter_batch
4542
+ )
4543
+ use_gpu_toks = (
4544
+ drafter_runs_model_forward
4545
+ and not spec_config.disable_padded_drafter_batch
4546
+ )
4547
if use_gpu_toks:
4548
# EAGLE/DraftModel speculative decoding can use the GPU sampled tokens
4549
# as inputs, and does not need to wait for bookkeeping to finish.
@@ -4552,19 +4558,23 @@ class GPUModelRunner(
4558
sampled_token_ids = sampler_output.sampled_token_ids
4559
if input_fits_in_drafter:
4560
propose_draft_token_ids(sampled_token_ids)
4555
- elif self.valid_sampled_token_count_event is not None:
4556
- assert spec_decode_common_attn_metadata is not None
4557
- next_token_ids, valid_sampled_tokens_count = (
4558
- self.drafter.prepare_next_token_ids_padded(
4559
- sampled_token_ids,
4560
- self.requests,
4561
- self.input_batch,
4562
- self.discard_request_mask.gpu,
4561
+ else:
4562
+ if self.valid_sampled_token_count_event is not None:
4563
+ assert spec_decode_common_attn_metadata is not None
4564
+ next_token_ids, valid_sampled_tokens_count = (
4565
+ self.drafter.prepare_next_token_ids_padded(
4566
+ sampled_token_ids,
4567
+ self.requests,
4568
+ self.input_batch,
4569
+ self.discard_request_mask.gpu,
4570
+ )
4571
)
4564
- )
4565
- self._copy_valid_sampled_token_count(
4566
- next_token_ids, valid_sampled_tokens_count
4567
- )
4572
+ self._copy_valid_sampled_token_count(
4573
+ next_token_ids, valid_sampled_tokens_count
4574
+ )
4575
+ if self.parallel_config.data_parallel_size > 1:
4576
+ # Prevent hang when DP ranks disagree on input_fits_in_drafter
4577
+ self.drafter.dummy_run(num_tokens=1)
4578
elif (
4579
spec_config.use_ngram_gpu()
4580
and not spec_config.disable_padded_drafter_batch
@@ -4588,7 +4598,9 @@ class GPUModelRunner(
4598
next_token_ids, valid_sampled_tokens_count
4599
)
4600
else:
4591
- propose_drafts_after_bookkeeping = input_fits_in_drafter
4601
+ # These drafters consume CPU sampled tokens, so they run
4602
+ # after bookkeeping.
4603
+ draft_after_bookkeeping = True
4604
4605
if not input_fits_in_drafter:
4606
# Zero out draft tokens so the scheduler doesn't schedule
@@ -4620,10 +4632,25 @@ class GPUModelRunner(
4632
scheduler_output.total_num_scheduled_tokens,
4633
)
4634
4623
- if propose_drafts_after_bookkeeping:
4635
+ if draft_after_bookkeeping:
4636
# ngram and other speculative decoding methods use the sampled
4637
# tokens on the CPU, so they are run after bookkeeping.
4626
- propose_draft_token_ids(valid_sampled_token_ids)
4638
+ if input_fits_in_drafter:
4639
+ propose_draft_token_ids(valid_sampled_token_ids)
4640
+ elif (
4641
+ drafter_runs_model_forward
4642
+ and self.parallel_config.data_parallel_size > 1
4643
+ ):
4644
+ # Prevent hang when DP ranks disagree on input_fits_in_drafter
4645
+ assert isinstance(
4646
+ self.drafter,
4647
+ EagleProposer
4648
+ | DFlashProposer
4649
+ | DraftModelProposer
4650
+ | ExtractHiddenStatesProposer
4651
+ | Gemma4Proposer,
4652
+ )
4653
+ self.drafter.dummy_run(num_tokens=1)
4654
4655
# Finalize KV connector (wait_for_save + clear metadata) after
4656
# draft model runs. Deferred from target model forward to allow