From c645e9a2144f7fede7222c0fc8937a93def04402 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 22 Feb 2026 18:27:12 -0800 Subject: [PATCH] [Model Runner V2] Remove propose_draft method (#35070) Signed-off-by: Woosuk Kwon --- vllm/v1/worker/gpu/model_runner.py | 29 +++++------------------------ 1 file changed, 5 insertions(+), 24 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 8204fd3c3..ccab6cec8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -858,29 +858,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): computed_prefill, self.req_states.prefill_len.np, out=computed_prefill ) - @torch.inference_mode() - def propose_draft( - self, - input_batch: InputBatch, - last_hidden_states: torch.Tensor, - aux_hidden_states: list[torch.Tensor] | None, - num_sampled: torch.Tensor, - num_rejected: torch.Tensor, - ) -> torch.Tensor: - assert self.speculator is not None - draft_tokens = self.speculator.propose( - input_batch, - last_hidden_states, - aux_hidden_states, - num_sampled, - num_rejected, - self.req_states.last_sampled_tokens, - self.req_states.next_prefill_tokens, - self.sampler.sampling_states.temperature.gpu, - self.sampler.sampling_states.seeds.gpu, - ) - return draft_tokens - @torch.inference_mode() def execute_model( self, @@ -1113,12 +1090,16 @@ class GPUModelRunner(LoRAModelRunnerMixin): input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected ) if self.speculator is not None: - draft_tokens = self.propose_draft( + draft_tokens = self.speculator.propose( input_batch, hidden_states, aux_hidden_states, num_sampled, num_rejected, + self.req_states.last_sampled_tokens, + self.req_states.next_prefill_tokens, + self.sampler.sampling_states.temperature.gpu, + self.sampler.sampling_states.seeds.gpu, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)