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 7, 22:45 UTC
[Model Runner V2] Remove propose_draft method (#35070)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 22, 2026 at 18:27 UTC
c645e9a2144f7fede7222c0fc8937a93def04402
1 file changed
+5
-24
vllm/v1/worker/gpu/model_runner.py
+5
-24
@@ -858,29 +858,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
858
computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
859
)
860
861
- @torch.inference_mode()
862
- def propose_draft(
863
- self,
864
- input_batch: InputBatch,
865
- last_hidden_states: torch.Tensor,
866
- aux_hidden_states: list[torch.Tensor] | None,
867
- num_sampled: torch.Tensor,
868
- num_rejected: torch.Tensor,
869
- ) -> torch.Tensor:
870
- assert self.speculator is not None
871
- draft_tokens = self.speculator.propose(
872
- input_batch,
873
- last_hidden_states,
874
- aux_hidden_states,
875
- num_sampled,
876
- num_rejected,
877
- self.req_states.last_sampled_tokens,
878
- self.req_states.next_prefill_tokens,
879
- self.sampler.sampling_states.temperature.gpu,
880
- self.sampler.sampling_states.seeds.gpu,
881
- )
882
- return draft_tokens
883
-
861
@torch.inference_mode()
862
def execute_model(
863
self,
@@ -1113,12 +1090,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1090
input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
1091
)
1092
if self.speculator is not None:
1116
- draft_tokens = self.propose_draft(
1093
+ draft_tokens = self.speculator.propose(
1094
input_batch,
1095
hidden_states,
1096
aux_hidden_states,
1097
num_sampled,
1098
num_rejected,
1099
+ self.req_states.last_sampled_tokens,
1100
+ self.req_states.next_prefill_tokens,
1101
+ self.sampler.sampling_states.temperature.gpu,
1102
+ self.sampler.sampling_states.seeds.gpu,
1103
)
1104
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1105
self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)