@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cd4a72a28d

Fork 0

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, 23:30 UTC

[V1][Spec decode] Move drafter to model runner (#13363)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Feb 17, 2025 at 15:40 UTC cd4a72a28d32851c7d188ee4f0a0145f4250e2aa

9 files changed +84 -57

tests/v1/core/test_scheduler.py

@@ -203,6 +203,7 @@ def test_schedule_partial_requests():

203 req_ids=[request.request_id for request in requests],

204 req_id_to_index=req_to_index,

205 sampled_token_ids=[[0] for _ in range(len(requests))],

206 + spec_token_ids=None,

207 logprobs=None,

208 prompt_logprobs_dict={},

209 )

@@ -259,6 +260,7 @@ def test_stop_via_update_from_output():

260 sampled_token_ids=[[EOS_TOKEN_ID],

261 [10,

262 11]], # First request hits EOS, second continues

263 + spec_token_ids=None,

264 logprobs=None,

265 prompt_logprobs_dict={})

266

@@ -307,6 +309,7 @@ def test_stop_via_update_from_output():

309 },

310 sampled_token_ids=[[10, 42, 12],

311 [13, 14]], # First request hits stop token

312 + spec_token_ids=None,

313 logprobs=None,

314 prompt_logprobs_dict={})

315

@@ -354,6 +357,7 @@ def test_stop_via_update_from_output():

357 },

358 sampled_token_ids=[[10, 11, 12],

359 [13]], # First request exceeds max_tokens

360 + spec_token_ids=None,

361 logprobs=None,

362 prompt_logprobs_dict={})

363

@@ -394,6 +398,7 @@ def test_stop_via_update_from_output():

398 req_ids=[requests[0].request_id],

399 req_id_to_index={requests[0].request_id: 0},

400 sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],

401 + spec_token_ids=None,

402 logprobs=None,

403 prompt_logprobs_dict={})

404

@@ -434,6 +439,7 @@ def test_schedule_concurrent_batches():

439 req_ids=[requests[0].request_id],

440 req_id_to_index={requests[0].request_id: 0},

441 sampled_token_ids=[[0]],

442 + spec_token_ids=None,

443 logprobs=None,

444 prompt_logprobs_dict={},

445 )

@@ -450,6 +456,7 @@ def test_schedule_concurrent_batches():

456 req_ids=[requests[1].request_id],

457 req_id_to_index={requests[1].request_id: 0},

458 sampled_token_ids=[[0]],

459 + spec_token_ids=None,

460 logprobs=None,

461 prompt_logprobs_dict={},

462 )

vllm/v1/core/scheduler.py

+4 -7

@@ -474,6 +474,7 @@ class Scheduler:

474 model_runner_output: "ModelRunnerOutput",

475 ) -> EngineCoreOutputs:

476 sampled_token_ids = model_runner_output.sampled_token_ids

477 + spec_token_ids = model_runner_output.spec_token_ids

478 logprobs = model_runner_output.logprobs

479 prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict

480 num_scheduled_tokens = scheduler_output.num_scheduled_tokens

@@ -530,13 +531,9 @@ class Scheduler:

531 self.encoder_cache_manager.free_encoder_input(

532 request, input_id)

533

533 - if request.num_computed_tokens >= request.num_tokens:

534 - # Clear the spec tokens as the request has generated

535 - # a new token. Here, We assume all spec tokens are verified

536 - # if we perform speculative decoding for this request.

537 - # Therefore, we can clear all spec tokens after

538 - # the generation step.

539 - request.clear_spec_tokens()

534 + # Add newly generated spec token ids to the request.

535 + if spec_token_ids is not None:

536 + request.spec_token_ids = spec_token_ids[req_index]

537

538 # Get prompt logprobs for this request.

539 prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)

vllm/v1/engine/core.py

-30

@@ -27,7 +27,6 @@ from vllm.v1.executor.abstract import Executor

27 from vllm.v1.outputs import ModelRunnerOutput

28 from vllm.v1.request import Request, RequestStatus

29 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder

30 -from vllm.v1.spec_decode.ngram_proposer import NgramProposer

30 from vllm.version import __version__ as VLLM_VERSION

32 logger = init_logger(__name__)

@@ -86,15 +85,6 @@ class EngineCore:

85 self.batch_queue_size)

86 self.batch_queue = queue.Queue(self.batch_queue_size)

89 - # Setup speculative decode.

90 - # TODO: find a better way to check if we are using ngram.

91 - self.use_spec_decode = False

92 - if self.scheduler.speculative_config:

93 - assert self.scheduler.speculative_config.ngram_prompt_lookup_min \

94 - , "Only ngram spec decode is supported in V1."

95 - self.proposer = NgramProposer()

96 - self.use_spec_decode = True

97 -

88 def _initialize_kv_caches(self,

89 vllm_config: VllmConfig) -> Tuple[int, int]:

90 start = time.time()

@@ -158,9 +148,6 @@ class EngineCore:

148 return EngineCoreOutputs(

149 outputs=[], scheduler_stats=self.scheduler.make_stats())

150

161 - if self.use_spec_decode:

162 - self.propose_tokens()

163 -

151 scheduler_output = self.scheduler.schedule()

152 output = self.model_executor.execute_model(scheduler_output)

153 engine_core_outputs = self.scheduler.update_from_output(

@@ -221,23 +208,6 @@ class EngineCore:

208 def profile(self, is_start: bool = True):

209 self.model_executor.profile(is_start)

210

224 - def propose_tokens(self):

225 - assert self.scheduler.speculative_config is not None

226 - for req in self.scheduler.running:

227 - # Ignore requests that are doing chunked prefill.

228 - if req.num_computed_tokens < req.num_tokens - 1:

229 - continue

230 - # Ignore requests that already have spec tokens.

231 - if req.spec_token_ids:

232 - continue

233 - spec_tokens = self.proposer.propose(

234 - req.all_token_ids,

235 - self.scheduler.speculative_config.ngram_prompt_lookup_min,

236 - self.scheduler.speculative_config.num_speculative_tokens,

237 - )

238 - if spec_tokens:

239 - req.append_spec_token_ids(spec_tokens)

240 -

211 def reset_prefix_cache(self):

212 self.scheduler.reset_prefix_cache()

213

vllm/v1/outputs.py

@@ -67,6 +67,9 @@ class ModelRunnerOutput:

67 # each request due to speculative/jump decoding.

68 sampled_token_ids: List[List[int]]

70 + # num_reqs x num_spec_tokens

71 + spec_token_ids: Optional[List[List[int]]]

72 +

73 # [num_reqs, max_num_logprobs + 1]

74 # [num_reqs, max_num_logprobs + 1]

75 # [num_reqs]

vllm/v1/request.py

-12

@@ -104,18 +104,6 @@ class Request:

104 self._output_token_ids.extend(token_ids)

105 self._all_token_ids.extend(token_ids)

106

107 - def append_spec_token_ids(

108 - self,

109 - token_ids: Union[int, List[int]],

110 - ) -> None:

111 - if isinstance(token_ids, int):

112 - self.spec_token_ids.append(token_ids)

113 - else:

114 - self.spec_token_ids.extend(token_ids)

115 -

116 - def clear_spec_tokens(self) -> None:

117 - self.spec_token_ids.clear()

118 -

107 @property

108 def num_tokens(self) -> int:

109 return len(self._all_token_ids)

vllm/v1/spec_decode/ngram_proposer.py

+15 -8

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 from typing import List, Optional

4 -from vllm.v1.utils import ConstantList

4 +import numpy as np

7 class NgramProposer:

@@ -9,8 +9,12 @@ class NgramProposer:

9 def __init__(self):

10 pass

12 - def propose(self, context_token_ids: ConstantList[int], n: int,

13 - k: int) -> Optional[List[int]]:

12 + def propose(

13 + self,

14 + context_token_ids: np.ndarray,

15 + n: int,

16 + k: int,

17 + ) -> Optional[np.ndarray]:

18 """Proposes the next sequence of tokens based on n-gram pattern

19 matching in the context. The function finds matches of the last n

20 tokens in the previous context, and returns k tokens that followed

@@ -25,8 +29,8 @@ class NgramProposer:

29 the maximum amount of tokens until the end.

31 Returns:

28 - List[int]: The sequence of tokens that followed

29 - the matched n-gram in the context.

32 + np.ndarray: The sequence of tokens that followed

33 + the matched n-gram in the context.

34 None: If no matching n-gram pattern is found.

36 Example:

@@ -66,9 +70,12 @@ class NgramProposer:

70 return lps

72 @staticmethod

69 - def _find_subarray_kmp(context_token_ids: ConstantList[int], n: int,

70 - k: int) -> Optional[List[int]]:

71 - context_len = len(context_token_ids)

73 + def _find_subarray_kmp(

74 + context_token_ids: np.ndarray,

75 + n: int,

76 + k: int,

77 + ) -> Optional[np.ndarray]:

78 + context_len = context_token_ids.shape[0]

79 assert n > 0

81 pattern = context_token_ids[-n:]

vllm/v1/worker/gpu_input_batch.py

@@ -78,6 +78,7 @@ class InputBatch:

78 )

79 self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()

80 self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)

81 + self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)

82 self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)

83 self.num_computed_tokens_cpu = np.empty(max_num_reqs, dtype=np.int32)

@@ -217,7 +218,11 @@ class InputBatch:

218 end_idx = start_idx + len(request.output_token_ids)

219 self.token_ids_cpu[req_index,

220 start_idx:end_idx] = request.output_token_ids

221 + # Number of token ids in token_ids_cpu.

222 + # NOTE(woosuk): This may include spec decode tokens.

223 self.num_tokens[req_index] = request.num_tokens

224 + # Number of tokens without spec decode tokens.

225 + self.num_tokens_no_spec[req_index] = request.num_tokens

226

227 self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens

228 self.block_table.add_row(req_index, request.block_ids)

@@ -356,6 +361,8 @@ class InputBatch:

361 self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[

362 last_req_index, :num_tokens]

363 self.num_tokens[empty_index] = num_tokens

364 + self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[

365 + last_req_index]

366 self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[

367 last_req_index]

368 self.num_computed_tokens_cpu[

vllm/v1/worker/gpu_model_runner.py

+47

@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

33 from vllm.v1.outputs import LogprobsTensors, ModelRunnerOutput

34 from vllm.v1.sample.metadata import SamplingMetadata

35 from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID

36 +from vllm.v1.spec_decode.ngram_proposer import NgramProposer

37 from vllm.v1.utils import bind_kv_cache

38 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

39 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

@@ -117,6 +118,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):

118 # req_id -> (input_id -> encoder_output)

119 self.encoder_cache: Dict[str, Dict[int, torch.Tensor]] = {}

120

121 + # Set up speculative decoding.

122 + self.use_spec_decode = False

123 + if self.speculative_config:

124 + # TODO: find a better way to check if we are using ngram.

125 + assert self.speculative_config.ngram_prompt_lookup_min, \

126 + "Currently, only ngram spec decode is supported in V1."

127 + self.drafter = NgramProposer()

128 + self.use_spec_decode = True

129 +

130 # Request states.

131 self.requests: Dict[str, CachedRequestState] = {}

132 # Persistent batch.

@@ -367,6 +377,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

377 self.input_batch.token_ids_cpu[

378 req_index,

379 start_token_index:end_token_index] = req_data.new_token_ids

380 + self.input_batch.num_tokens_no_spec[req_index] = end_token_index

381 # Add spec_token_ids to token_ids_cpu.

382 spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(

383 req_id, [])

@@ -1009,15 +1020,51 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1020 for seq in sampled_token_ids[valid_mask].split(gen_lens)

1021 ]

1022

1023 + if not self.use_spec_decode:

1024 + spec_token_ids = None

1025 + else:

1026 + spec_token_ids = self.generate_draft_token_ids(

1027 + valid_sampled_token_ids)

1028 +

1029 model_runner_output = ModelRunnerOutput(

1030 req_ids=req_ids,

1031 req_id_to_index=self.input_batch.req_id_to_index,

1032 sampled_token_ids=valid_sampled_token_ids,

1033 + spec_token_ids=spec_token_ids,

1034 logprobs=logprobs_lists,

1035 prompt_logprobs_dict=prompt_logprobs_dict,

1036 )

1037 return model_runner_output

1038

1039 + def generate_draft_token_ids(

1040 + self,

1041 + sampled_token_ids: List[List[int]],

1042 + ) -> List[List[int]]:

1043 + # TODO(woosuk): Optimize.

1044 + num_reqs = len(sampled_token_ids)

1045 + draft_token_ids: List[List[int]] = []

1046 + for i in range(num_reqs):

1047 + if len(sampled_token_ids[i]) == 0:

1048 + # Skip speculative decoding.

1049 + draft_token_ids.append([])

1050 + continue

1051 +

1052 + # Add sampled_token_ids to token_ids_cpu.

1053 + start_idx = self.input_batch.num_tokens_no_spec[i]

1054 + end_idx = start_idx + len(sampled_token_ids[i])

1055 + self.input_batch.token_ids_cpu[

1056 + i, start_idx:end_idx] = sampled_token_ids[i]

1057 + drafter_output = self.drafter.propose(

1058 + self.input_batch.token_ids_cpu[i, :end_idx],

1059 + self.speculative_config.ngram_prompt_lookup_min,

1060 + self.speculative_config.num_speculative_tokens,

1061 + )

1062 + if drafter_output is None or len(drafter_output) == 0:

1063 + draft_token_ids.append([])

1064 + else:

1065 + draft_token_ids.append(drafter_output.tolist())

1066 + return draft_token_ids

1067 +

1068 def load_model(self) -> None:

1069 logger.info("Starting to load model %s...", self.model_config.model)

1070 with DeviceMemoryProfiler() as m: # noqa: SIM117

vllm/v1/worker/tpu_model_runner.py

@@ -696,6 +696,7 @@ class TPUModelRunner:

696 req_ids=all_req_ids,

697 req_id_to_index=self.input_batch.req_id_to_index,

698 sampled_token_ids=[[token_id] for token_id in sampled_token_ids],

699 + spec_token_ids=None,

700 logprobs=None,

701 prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]

702 )