A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 03:30 UTC
[Spec Decode] Refactor spec decoding into a separate function (#20238)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jun 30, 2025 at 08:13 UTC
2062c0723d38a8f4a8a7565b61a99e8c81b5cacd
1 file changed
+60
-33
vllm/v1/worker/gpu_model_runner.py
+60
-33
@@ -1388,6 +1388,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1388
hidden_states, aux_hidden_states = model_output
1389
else:
1390
hidden_states = model_output
1391
+ aux_hidden_states = None
1392
+
1393
# Broadcast PP output for external_launcher (torchrun)
1394
# to make sure we are synced across pp ranks
1395
# TODO: Support overlapping mirco-batches
@@ -1510,25 +1512,67 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1512
if not self.speculative_config:
1513
# Speculative decoding is not enabled.
1514
spec_token_ids = None
1513
- elif self.speculative_config.method == "ngram":
1515
+ else:
1516
+ spec_token_ids = self.propose_draft_token_ids(
1517
+ scheduler_output,
1518
+ valid_sampled_token_ids,
1519
+ sampling_metadata,
1520
+ hidden_states,
1521
+ sample_hidden_states,
1522
+ aux_hidden_states,
1523
+ spec_decode_metadata,
1524
+ attn_metadata,
1525
+ )
1526
+
1527
+ # Clear KVConnector state after all KVs are generated.
1528
+ if has_kv_transfer_group():
1529
+ get_kv_transfer_group().clear_connector_metadata()
1530
+
1531
+ self.eplb_step()
1532
+
1533
+ return ModelRunnerOutput(
1534
+ req_ids=self.input_batch.req_ids,
1535
+ req_id_to_index=self.input_batch.req_id_to_index,
1536
+ sampled_token_ids=valid_sampled_token_ids,
1537
+ spec_token_ids=spec_token_ids,
1538
+ logprobs=logprobs_lists,
1539
+ prompt_logprobs_dict=prompt_logprobs_dict,
1540
+ pooler_output=[],
1541
+ finished_sending=finished_sending,
1542
+ finished_recving=finished_recving,
1543
+ num_nans_in_logits=num_nans_in_logits,
1544
+ )
1545
+
1546
+ def propose_draft_token_ids(
1547
+ self,
1548
+ scheduler_output: "SchedulerOutput",
1549
+ sampled_token_ids: list[list[int]],
1550
+ sampling_metadata: SamplingMetadata,
1551
+ hidden_states: torch.Tensor,
1552
+ sample_hidden_states: torch.Tensor,
1553
+ aux_hidden_states: Optional[torch.Tensor],
1554
+ spec_decode_metadata: Optional[SpecDecodeMetadata],
1555
+ attn_metadata: dict[str, Any],
1556
+ ) -> list[list[int]]:
1557
+ num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1558
+ if self.speculative_config.method == "ngram":
1559
assert isinstance(self.drafter, NgramProposer)
1515
- spec_token_ids = self.generate_draft_token_ids(
1516
- valid_sampled_token_ids, sampling_metadata)
1560
+ spec_token_ids = self.propose_ngram_draft_token_ids(
1561
+ sampled_token_ids)
1562
elif self.speculative_config.method == "medusa":
1563
assert isinstance(self.drafter, MedusaProposer)
1519
- if max_gen_len == 1:
1564
+ if sample_hidden_states.shape[0] == len(sampled_token_ids):
1565
+ # The input to the target model does not include draft tokens.
1566
hidden_states = sample_hidden_states
1567
else:
1568
indices = []
1569
offset = 0
1570
for num_draft, tokens in zip(
1571
spec_decode_metadata.num_draft_tokens,
1526
- valid_sampled_token_ids):
1572
+ sampled_token_ids):
1573
indices.append(offset + len(tokens) - 1)
1574
offset += num_draft + 1
1529
-
1530
- indices = torch.tensor(indices,
1531
- device=sample_hidden_states.device)
1575
+ indices = torch.tensor(indices, device=self.device)
1576
hidden_states = sample_hidden_states[indices]
1577
1578
spec_token_ids = self.drafter.propose(
@@ -1539,7 +1583,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1583
assert isinstance(self.drafter, EagleProposer)
1584
# TODO(woosuk): Refactor the loop.
1585
next_token_ids: list[int] = []
1542
- for i, token_ids in enumerate(valid_sampled_token_ids):
1586
+ for i, token_ids in enumerate(sampled_token_ids):
1587
if token_ids:
1588
# Common case.
1589
next_token_id = token_ids[-1]
@@ -1569,7 +1613,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1613
if spec_decode_metadata is None:
1614
# input_ids can be None for multimodal models.
1615
target_token_ids = self.input_ids[:num_scheduled_tokens]
1572
- target_positions = positions[:num_scheduled_tokens]
1616
+ # TODO(woosuk): Support M-RoPE.
1617
+ target_positions = self.positions[:num_scheduled_tokens]
1618
if self.use_aux_hidden_state_outputs:
1619
target_hidden_states = torch.cat(
1620
[h[:num_scheduled_tokens] for h in aux_hidden_states],
@@ -1582,7 +1627,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1627
# TODO(woosuk): Refactor this.
1628
num_draft_tokens = spec_decode_metadata.num_draft_tokens
1629
num_rejected_tokens = [
1585
- n + 1 - len(valid_sampled_token_ids[i]) if n > 0 else 0
1630
+ n + 1 - len(sampled_token_ids[i]) if n > 0 else 0
1631
for i, n in enumerate(num_draft_tokens)
1632
]
1633
num_rejected_tokens_tensor = async_tensor_h2d(
@@ -1597,7 +1642,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1642
num_tokens,
1643
)
1644
target_token_ids = self.input_ids[token_indices]
1600
- target_positions = positions[token_indices]
1645
+ # TODO(woosuk): Support M-RoPE.
1646
+ target_positions = self.positions[token_indices]
1647
if self.use_aux_hidden_state_outputs:
1648
target_hidden_states = torch.cat(
1649
[h[token_indices] for h in aux_hidden_states], dim=-1)
@@ -1616,25 +1662,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1662
sampling_metadata=sampling_metadata,
1663
)
1664
spec_token_ids = draft_token_ids.tolist()
1619
-
1620
- # Clear KVConnector state after all KVs are generated.
1621
- if has_kv_transfer_group():
1622
- get_kv_transfer_group().clear_connector_metadata()
1623
-
1624
- self.eplb_step()
1625
-
1626
- return ModelRunnerOutput(
1627
- req_ids=self.input_batch.req_ids,
1628
- req_id_to_index=self.input_batch.req_id_to_index,
1629
- sampled_token_ids=valid_sampled_token_ids,
1630
- spec_token_ids=spec_token_ids,
1631
- logprobs=logprobs_lists,
1632
- prompt_logprobs_dict=prompt_logprobs_dict,
1633
- pooler_output=[],
1634
- finished_sending=finished_sending,
1635
- finished_recving=finished_recving,
1636
- num_nans_in_logits=num_nans_in_logits,
1637
- )
1665
+ return spec_token_ids
1666
1667
def kv_connector_no_forward(
1668
self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
@@ -1682,10 +1710,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1710
scheduler_output.finished_req_ids)
1711
return None, None
1712
1685
- def generate_draft_token_ids(
1713
+ def propose_ngram_draft_token_ids(
1714
self,
1715
sampled_token_ids: list[list[int]],
1688
- sampling_metadata: SamplingMetadata,
1716
) -> list[list[int]]:
1717
# TODO(woosuk): Optimize.
1718
draft_token_ids: list[list[int]] = []