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

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

[V0 Deprecation] Remove V0 Sequence class & Sampler (#25332)

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

Woosuk Kwon committed Sep 21, 2025 at 08:52 UTC 26e673fe9303ad759a47e19a087764393c69109f

27 files changed +70 -3697

tests/conftest.py

+1 -1

@@ -48,10 +48,10 @@ from vllm.distributed import (cleanup_dist_env_and_memory,

48 initialize_model_parallel)

49 from vllm.inputs import TextPrompt

50 from vllm.logger import init_logger

51 +from vllm.logprobs import Logprob

52 from vllm.multimodal.utils import fetch_image

53 from vllm.outputs import RequestOutput

54 from vllm.sampling_params import BeamSearchParams

54 -from vllm.sequence import Logprob

55 from vllm.transformers_utils.utils import maybe_model_redirect

56 from vllm.utils import set_default_torch_num_threads

tests/models/multimodal/generation/test_granite_speech.py

+1 -1

@@ -7,8 +7,8 @@ from typing import Optional

7 import pytest

8 from transformers import AutoModelForSpeechSeq2Seq

10 +from vllm.logprobs import SampleLogprobs

11 from vllm.lora.request import LoRARequest

11 -from vllm.sequence import SampleLogprobs

13 from ....conftest import (AudioTestAssets, HfRunner, PromptAudioInput,

14 VllmRunner)

tests/models/multimodal/generation/test_phi4mm.py

+1 -1

@@ -12,10 +12,10 @@ from huggingface_hub import snapshot_download

12 from transformers import AutoTokenizer

14 from vllm.assets.image import ImageAsset

15 +from vllm.logprobs import SampleLogprobs

16 from vllm.lora.request import LoRARequest

17 from vllm.multimodal.image import convert_image_mode, rescale_image_size

18 from vllm.platforms import current_platform

18 -from vllm.sequence import SampleLogprobs

20 from ....conftest import (IMAGE_ASSETS, HfRunner, PromptAudioInput,

21 PromptImageInput, VllmRunner)

tests/models/multimodal/generation/test_pixtral.py

+1 -1

@@ -13,8 +13,8 @@ from mistral_common.tokens.tokenizers.multimodal import image_from_chunk

13 from transformers import AutoProcessor

15 from vllm import SamplingParams, TextPrompt, TokensPrompt

16 +from vllm.logprobs import Logprob, SampleLogprobs

17 from vllm.multimodal import MultiModalDataBuiltins

17 -from vllm.sequence import Logprob, SampleLogprobs

19 from ....utils import VLLM_PATH, large_gpu_test

20 from ...utils import check_logprobs_close

tests/models/multimodal/generation/vlm_utils/model_utils.py

+1 -1

@@ -19,7 +19,7 @@ from transformers import (AutoConfig, AutoTokenizer, BatchFeature,

19 GenerationConfig, GenerationMixin)

20 from transformers.video_utils import VideoMetadata

22 -from vllm.sequence import SampleLogprobs

22 +from vllm.logprobs import SampleLogprobs

23 from vllm.utils import is_list_of

25 from .....conftest import HfRunner, ImageAsset, ImageTestAssets

tests/models/multimodal/generation/vlm_utils/types.py

+1 -1

@@ -12,7 +12,7 @@ from transformers import AutoModelForCausalLM

12 from transformers.models.auto.auto_factory import _BaseAutoModelClass

14 from vllm.config import RunnerOption

15 -from vllm.sequence import SampleLogprobs

15 +from vllm.logprobs import SampleLogprobs

16 from vllm.transformers_utils.tokenizer import AnyTokenizer

18 from .....conftest import (AUDIO_ASSETS, IMAGE_ASSETS, HfRunner, ImageAsset,

tests/models/utils.py

+1 -1

@@ -12,7 +12,7 @@ from transformers import PretrainedConfig

13 from vllm.config import ModelConfig, ModelDType, RunnerOption

14 from vllm.inputs import InputContext

15 -from vllm.sequence import Logprob, PromptLogprobs, SampleLogprobs

15 +from vllm.logprobs import Logprob, PromptLogprobs, SampleLogprobs

17 from .registry import HF_EXAMPLE_MODELS

tests/tokenization/test_detokenize.py

+1 -139

@@ -8,10 +8,7 @@ import pytest

8 from transformers import (AutoTokenizer, PreTrainedTokenizer,

9 PreTrainedTokenizerFast)

11 -from vllm.inputs import token_inputs

12 -from vllm.sequence import Logprob, SamplingParams, Sequence, SequenceGroup

13 -from vllm.transformers_utils.detokenizer import Detokenizer

14 -from vllm.transformers_utils.tokenizer import get_tokenizer

11 +from vllm.sampling_params import SamplingParams

12 from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer

13 from vllm.v1.engine import EngineCoreRequest

14 from vllm.v1.engine.detokenizer import (FastIncrementalDetokenizer,

@@ -217,138 +214,3 @@ def test_oov_decode(tokenizer, fast):

214

215 assert decoded_text == ''

216 assert out_ids == [len(tokenizer)]

220 -

221 -

222 -@pytest.fixture

223 -def detokenizer(tokenizer_name: str) -> Detokenizer:

224 - tokenizer = get_tokenizer(

225 - tokenizer_name,

226 - tokenizer_mode="mistral" if "mistral" in tokenizer_name else "auto",

227 - trust_remote_code=False,

228 - revision=None,

229 - )

230 -

231 - return Detokenizer(tokenizer)

232 -

233 -

234 -@pytest.fixture(name="complete_sequence_token_ids")

235 -def create_complete_sequence_token_ids(complete_sequence: str,

236 - tokenizer) -> list[int]:

237 - return tokenizer(complete_sequence, add_special_tokens=False).input_ids

238 -

239 -

240 -def create_sequence(prompt_token_ids=None):

241 - prompt_token_ids = prompt_token_ids or []

242 - return Sequence(

243 - seq_id=0,

244 - inputs=token_inputs(prompt_token_ids),

245 - block_size=16,

246 - )

247 -

248 -

249 -def create_dummy_logprobs(

250 - complete_sequence_token_ids: list[int]) -> list[dict[int, Logprob]]:

251 - return [{

252 - token_id: Logprob(logprob=0.0),

253 - token_id + 1: Logprob(logprob=0.1)

254 - } for token_id in complete_sequence_token_ids]

255 -

256 -

257 -def create_dummy_prompt_logprobs(

258 - complete_sequence_token_ids: list[int]

259 -) -> list[Optional[dict[int, Any]]]:

260 - # logprob for the first prompt token is None.

261 - logprobs: list[Optional[dict[int, Any]]] = [None]

262 - logprobs.extend(create_dummy_logprobs(complete_sequence_token_ids)[1:])

263 - return logprobs

264 -

265 -

266 -@pytest.mark.parametrize("complete_sequence", TRUTH)

267 -@pytest.mark.parametrize("tokenizer_name", TOKENIZERS)

268 -@pytest.mark.parametrize("skip_special_tokens", [True, False], indirect=True)

269 -def test_decode_sequence_logprobs(complete_sequence: str,

270 - complete_sequence_token_ids: list[int],

271 - detokenizer: Detokenizer,

272 - skip_special_tokens: bool):

273 - """Verify Detokenizer decodes logprobs correctly."""

274 - sampling_params = SamplingParams(skip_special_tokens=skip_special_tokens,

275 - logprobs=2)

276 -

277 - # Run sequentially.

278 - seq = create_sequence()

279 - dummy_logprobs = create_dummy_logprobs(complete_sequence_token_ids)

280 - sequential_logprobs_text_chosen_token: list[str] = []

281 - sequential_logprobs_text_other_token: list[str] = []

282 - for new_token, logprobs in zip(complete_sequence_token_ids,

283 - dummy_logprobs):

284 - seq.append_token_id(new_token, logprobs)

285 - detokenizer.decode_sequence_inplace(seq, sampling_params)

286 - sequential_logprobs_text_chosen_token.append(

287 - seq.output_logprobs[-1][new_token].decoded_token)

288 - sequential_logprobs_text_other_token.append(

289 - seq.output_logprobs[-1][new_token + 1].decoded_token)

290 - sequential_result = seq.output_text

291 -

292 - assert sequential_result == "".join(sequential_logprobs_text_chosen_token)

293 - assert sequential_result != "".join(sequential_logprobs_text_other_token)

294 -

295 - if not skip_special_tokens:

296 - # Text for logprobs for the chosen token should be the same as the

297 - # generated text. Note that this will only be true if we skip

298 - # special tokens.

299 - assert sequential_result == complete_sequence

300 -

301 -

302 -@pytest.mark.parametrize("complete_sequence", TRUTH)

303 -@pytest.mark.parametrize("tokenizer_name", TOKENIZERS)

304 -def test_decode_prompt_logprobs(complete_sequence: str,

305 - complete_sequence_token_ids: list[int],

306 - detokenizer: Detokenizer):

307 -

308 - # We want to use skip_special_tokens=False here but Mistral tokenizers

309 - # don't support that.

310 - if complete_sequence not in SPECIAL_TOKS_TRUTH:

311 - skip_special_tokens = True

312 - elif not isinstance(detokenizer.tokenizer, MistralTokenizer):

313 - skip_special_tokens = False

314 - else:

315 - pytest.skip("MistralTokenizers don't support "

316 - "skip_special_tokens=False")

317 - return

318 - """Verify Detokenizer decodes prompt logprobs correctly."""

319 - sampling_params = SamplingParams(skip_special_tokens=skip_special_tokens,

320 - prompt_logprobs=1)

321 -

322 - # Run sequentially.

323 - seq = create_sequence(complete_sequence_token_ids)

324 - seq_group = SequenceGroup(request_id="1",

325 - seqs=[seq],

326 - sampling_params=sampling_params,

327 - arrival_time=0.0)

328 - dummy_logprobs = create_dummy_prompt_logprobs(complete_sequence_token_ids)

329 - detokenizer.decode_prompt_logprobs_inplace(seq_group,

330 - dummy_logprobs,

331 - position_offset=0)

332 - # First logprob is None.

333 - decoded_prompt_logprobs: list[dict[int, Any]] = dummy_logprobs[

334 - 1:] # type: ignore

335 -

336 - # decoded_prompt_logprobs doesn't contain the first token.

337 - token_ids = complete_sequence_token_ids

338 - tokenizer = detokenizer.tokenizer

339 - text_full = tokenizer.decode(token_ids,

340 - skip_special_tokens=skip_special_tokens)

341 - text_first = tokenizer.decode(token_ids[0],

342 - skip_special_tokens=skip_special_tokens)

343 - text = text_full[len(text_first):]

344 -

345 - # Text for logprobs for the chosen token should be the same as the

346 - # prompt text. Note that the first logprob is None.

347 - assert text == "".join([

348 - logprobs[token_id].decoded_token

349 - for token_id, logprobs in zip(token_ids[1:], decoded_prompt_logprobs)

350 - ])

351 - assert text != "".join([

352 - logprobs[token_id + 1].decoded_token

353 - for token_id, logprobs in zip(token_ids[1:], decoded_prompt_logprobs)

354 - ])

tests/tool_use/test_jamba_tool_parser.py

+1 -1

@@ -12,7 +12,7 @@ from partial_json_parser.core.options import Allow

12 from vllm.entrypoints.openai.protocol import (DeltaMessage, FunctionCall,

13 ToolCall)

14 from vllm.entrypoints.openai.tool_parsers import JambaToolParser

15 -from vllm.transformers_utils.detokenizer import detokenize_incrementally

15 +from vllm.transformers_utils.detokenizer_utils import detokenize_incrementally

16 from vllm.transformers_utils.tokenizer import AnyTokenizer, get_tokenizer

18 MODEL = "ai21labs/Jamba-tiny-dev"

tests/tool_use/test_qwen3coder_tool_parser.py

+1 -1

@@ -13,7 +13,7 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

13 ToolCall)

14 from vllm.entrypoints.openai.tool_parsers.qwen3coder_tool_parser import (

15 Qwen3CoderToolParser)

16 -from vllm.transformers_utils.detokenizer import detokenize_incrementally

16 +from vllm.transformers_utils.detokenizer_utils import detokenize_incrementally

17 from vllm.transformers_utils.tokenizer import AnyTokenizer, get_tokenizer

19 MODEL = "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"

tests/tool_use/test_seed_oss_tool_parser.py

+1 -1

@@ -13,7 +13,7 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

13 DeltaMessage, FunctionCall,

14 ToolCall)

15 from vllm.entrypoints.openai.tool_parsers import SeedOssToolParser

16 -from vllm.transformers_utils.detokenizer import detokenize_incrementally

16 +from vllm.transformers_utils.detokenizer_utils import detokenize_incrementally

17 from vllm.transformers_utils.tokenizer import AnyTokenizer, get_tokenizer

19 # Use a common model that is likely to be available

tests/tool_use/test_xlam_tool_parser.py

+1 -1

@@ -11,7 +11,7 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

11 DeltaMessage, FunctionCall,

12 ToolCall)

13 from vllm.entrypoints.openai.tool_parsers import xLAMToolParser

14 -from vllm.transformers_utils.detokenizer import detokenize_incrementally

14 +from vllm.transformers_utils.detokenizer_utils import detokenize_incrementally

15 from vllm.transformers_utils.tokenizer import AnyTokenizer, get_tokenizer

17 # Use a common model that is likely to be available

tests/v1/engine/test_output_processor.py

+1 -1

@@ -12,9 +12,9 @@ from tests.v1.engine.utils import (NUM_PROMPT_LOGPROBS_UNDER_TEST,

12 STOP_STRINGS,

13 DummyOutputProcessorTestVectors,

14 MockEngineCore)

15 +from vllm.logprobs import PromptLogprobs, SampleLogprobs

16 from vllm.outputs import CompletionOutput, RequestOutput

17 from vllm.sampling_params import RequestOutputKind, SamplingParams

17 -from vllm.sequence import PromptLogprobs, SampleLogprobs

18 from vllm.transformers_utils.tokenizer import AnyTokenizer

19 from vllm.v1.engine import EngineCoreRequest

20 from vllm.v1.engine.output_processor import (OutputProcessor,

vllm/executor/executor_base.py

+1 -1

@@ -15,10 +15,10 @@ from vllm.config import VllmConfig

15 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator

16 from vllm.logger import init_logger

17 from vllm.lora.request import LoRARequest

18 -from vllm.model_executor.layers.sampler import SamplerOutput

18 from vllm.sequence import ExecuteModelRequest, PoolerOutput

19 from vllm.tasks import SupportedTask

20 from vllm.utils import make_async

21 +from vllm.v1.outputs import SamplerOutput

22 from vllm.worker.worker_base import WorkerBase

24 logger = init_logger(__name__)

vllm/executor/ray_distributed_executor.py

+1 -1

@@ -17,12 +17,12 @@ from vllm.executor.msgspec_utils import encode_hook

17 from vllm.executor.ray_utils import (RayWorkerWrapper, initialize_ray_cluster,

18 ray)

19 from vllm.logger import init_logger

20 -from vllm.model_executor.layers.sampler import SamplerOutput

20 from vllm.platforms import current_platform

21 from vllm.ray.ray_env import get_env_vars_to_copy

22 from vllm.sequence import ExecuteModelRequest

23 from vllm.utils import (_run_task_with_lock, get_distributed_init_method,

24 get_ip, get_open_port, make_async)

25 +from vllm.v1.outputs import SamplerOutput

27 if ray is not None:

28 from ray.actor import ActorHandle

vllm/inputs/__init__.py

+1 -12

@@ -7,15 +7,7 @@ from .data import (DataPrompt, DecoderOnlyInputs, EmbedsInputs, EmbedsPrompt,

7 SingletonPrompt, TextPrompt, TokenInputs, TokensPrompt,

8 build_explicit_enc_dec_prompt, embeds_inputs,

9 to_enc_dec_tuple_list, token_inputs, zip_enc_dec_prompts)

10 -from .registry import (DummyData, InputContext, InputProcessingContext,

11 - InputRegistry)

12 -

13 -INPUT_REGISTRY = InputRegistry()

14 -"""

15 -The global [`InputRegistry`][vllm.inputs.registry.InputRegistry] which is used

16 -by [`LLMEngine`][vllm.LLMEngine] to dispatch data processing according to the

17 -target model.

18 -"""

10 +from .registry import InputContext, InputProcessingContext

12 __all__ = [

13 "DataPrompt",

@@ -36,9 +28,6 @@ __all__ = [

28 "build_explicit_enc_dec_prompt",

29 "to_enc_dec_tuple_list",

30 "zip_enc_dec_prompts",

39 - "INPUT_REGISTRY",

40 - "DummyData",

31 "InputContext",

32 "InputProcessingContext",

43 - "InputRegistry",

33 ]

vllm/inputs/registry.py

+1 -66

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Mapping

4 from dataclasses import dataclass

5 -from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union

5 +from typing import TYPE_CHECKING, Any, Union

7 import torch

8 from transformers import BatchFeature, PretrainedConfig, ProcessorMixin

@@ -15,16 +15,9 @@ from vllm.utils.jsontree import JSONTree, json_map_leaves

16 if TYPE_CHECKING:

17 from vllm.config import ModelConfig

18 - from vllm.multimodal import (MultiModalDataDict, MultiModalPlaceholderDict,

19 - MultiModalRegistry)

20 - from vllm.sequence import SequenceData

18 from vllm.transformers_utils.tokenizer import AnyTokenizer

19 else:

20 ModelConfig = Any

24 - MultiModalDataDict = Any

25 - MultiModalPlaceholderDict = Any

26 - MultiModalRegistry = Any

27 - SequenceData = Any

21 AnyTokenizer = Any

23 _T = TypeVar("_T")

@@ -191,61 +184,3 @@ class InputProcessingContext(InputContext):

184 f"on data={data} with kwargs={allowed_kwargs}")

185

186 raise ValueError(msg) from exc

194 -

195 -

196 -class DummyData(NamedTuple):

197 - """

198 - Dummy data used for profiling.

199 -

200 - Note: This is only used in V0.

201 - """

202 -

203 - seq_data: SequenceData

204 - multi_modal_data: Optional[MultiModalDataDict] = None

205 - multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None

206 -

207 -

208 -class InputRegistry:

209 - """

210 - Note: This is only used in V0.

211 - """

212 -

213 - def dummy_data_for_profiling(

214 - self,

215 - model_config: ModelConfig,

216 - seq_len: int,

217 - mm_registry: MultiModalRegistry,

218 - is_encoder_data: bool = False,

219 - ) -> DummyData:

220 - """

221 - Create dummy data for profiling the memory usage of a model.

222 -

223 - The model is identified by ``model_config``.

224 - """

225 - # Avoid circular import

226 - from vllm.multimodal.cache import processor_only_cache_from_config

227 - from vllm.sequence import SequenceData

228 -

229 - if not model_config.is_multimodal_model:

230 - seq_data = SequenceData.from_prompt_token_counts((0, seq_len))

231 - return DummyData(seq_data=seq_data)

232 -

233 - cache = processor_only_cache_from_config(model_config, mm_registry)

234 -

235 - # Encoder dummy data does not contain multi-modal data

236 - if is_encoder_data:

237 - enc_data = mm_registry.get_encoder_dummy_data(model_config,

238 - seq_len,

239 - cache=cache)

240 - seq_data = SequenceData.from_seqs(enc_data.prompt_token_ids)

241 - return DummyData(seq_data=seq_data)

242 -

243 - dec_data = mm_registry.get_decoder_dummy_data(model_config,

244 - seq_len,

245 - cache=cache)

246 -

247 - return DummyData(

248 - seq_data=SequenceData.from_seqs(dec_data.prompt_token_ids),

249 - multi_modal_data=dec_data.multi_modal_data.get_data(),

250 - multi_modal_placeholders=dec_data.multi_modal_placeholders,

251 - )

vllm/model_executor/__init__.py

+1 -3

@@ -3,13 +3,11 @@

4 from vllm.model_executor.parameter import (BasevLLMParameter,

5 PackedvLLMParameter)

6 -from vllm.model_executor.sampling_metadata import (SamplingMetadata,

7 - SamplingMetadataCache)

6 +from vllm.model_executor.sampling_metadata import SamplingMetadata

7 from vllm.model_executor.utils import set_random_seed

9 __all__ = [

10 "SamplingMetadata",

12 - "SamplingMetadataCache",

11 "set_random_seed",

12 "BasevLLMParameter",

13 "PackedvLLMParameter",

vllm/model_executor/layers/logits_processor.py

-91

@@ -1,13 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """A layer that compute logits from hidden_stats."""

4 -import inspect

5 -from concurrent.futures import ThreadPoolExecutor

4 from typing import Optional

6 import torch

10 -import vllm.envs as envs

8 from vllm.distributed import (tensor_model_parallel_all_gather,

9 tensor_model_parallel_gather)

10 from vllm.model_executor.custom_op import CustomOp

@@ -16,11 +13,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

13 from vllm.model_executor.sampling_metadata import SamplingMetadata

14 from vllm.platforms import current_platform

19 -_logits_processor_threadpool: Optional[ThreadPoolExecutor] = None

20 -if envs.VLLM_LOGITS_PROCESSOR_THREADS is not None:

21 - _logits_processor_threadpool = ThreadPoolExecutor(

22 - envs.VLLM_LOGITS_PROCESSOR_THREADS)

23 -

17 @CustomOp.register("logits_processor")

18 class LogitsProcessor(CustomOp):

@@ -60,15 +52,10 @@ class LogitsProcessor(CustomOp):

52 hidden_states: torch.Tensor,

53 sampling_metadata: Optional[SamplingMetadata] = None,

54 embedding_bias: Optional[torch.Tensor] = None,

63 - prune_hidden_states: bool = True,

55 ) -> Optional[torch.Tensor]:

56 if self.logits_as_input:

57 logits = hidden_states

58 else:

68 - if sampling_metadata is not None and prune_hidden_states:

69 - hidden_states = _prune_hidden_states(hidden_states,

70 - sampling_metadata)

71 -

59 # Get the logits for the next tokens.

60 logits = self._get_logits(hidden_states, lm_head, embedding_bias)

61 if logits is not None:

@@ -79,12 +66,6 @@ class LogitsProcessor(CustomOp):

67 if self.scale != 1.0:

68 logits *= self.scale

82 -

83 - # Apply logits processors (if any).

84 - if sampling_metadata is not None and \

85 - sampling_metadata.seq_groups is not None:

86 - logits = _apply_logits_processors(logits, sampling_metadata)

87 -

69 return logits

71 def _gather_logits(self, logits: torch.Tensor) -> torch.Tensor:

@@ -125,75 +106,3 @@ class LogitsProcessor(CustomOp):

106 s += f", org_vocab_size={self.org_vocab_size}"

107 s += f", scale={self.scale}, logits_as_input={self.logits_as_input}"

108 return s

128 -

129 -

130 -def _prune_hidden_states(

131 - hidden_states: torch.Tensor,

132 - sampling_metadata: SamplingMetadata,

133 -) -> torch.Tensor:

134 - # NOTE(kzawora): The if guard is needed for Gaudi - in some scenarios

135 - # (warmup, profile_run) we might not have selected_token_indices,

136 - # so we skip pruning.

137 - if sampling_metadata.selected_token_indices is not None:

138 - return hidden_states.index_select(

139 - 0, sampling_metadata.selected_token_indices)

140 - else:

141 - return hidden_states

142 -

143 -

144 -def _apply_logits_processors(

145 - logits: torch.Tensor,

146 - sampling_metadata: SamplingMetadata,

147 -) -> torch.Tensor:

148 - found_logits_processors = False

149 - logits_processed = 0

150 - logits_row_ids_and_logits_row_futures = []

151 - for seq_group in sampling_metadata.seq_groups:

152 - seq_ids = seq_group.seq_ids

153 - sampling_params = seq_group.sampling_params

154 - logits_processors = sampling_params.logits_processors

155 - if logits_processors:

156 - found_logits_processors = True

157 -

158 - for seq_id, logits_row_idx in zip(seq_ids,

159 - seq_group.sample_indices):

160 - logits_row = logits[logits_row_idx]

161 - past_tokens_ids = seq_group.seq_data[seq_id].output_token_ids

162 - prompt_tokens_ids = seq_group.seq_data[seq_id].prompt_token_ids

163 -

164 - if _logits_processor_threadpool is not None:

165 - logits_row_ids_and_logits_row_futures.append(

166 - (logits_row_idx,

167 - _logits_processor_threadpool.submit(

168 - _apply_logits_processors_single_seq, logits_row,

169 - logits_processors, past_tokens_ids,

170 - prompt_tokens_ids)))

171 - else:

172 - logits[logits_row_idx] = \

173 - _apply_logits_processors_single_seq(

174 - logits_row, logits_processors, past_tokens_ids,

175 - prompt_tokens_ids)

176 -

177 - logits_processed += len(seq_group.sample_indices) + len(

178 - seq_group.prompt_logprob_indices)

179 -

180 - for logits_row_idx, future in logits_row_ids_and_logits_row_futures:

181 - logits[logits_row_idx] = future.result()

182 -

183 - if found_logits_processors:

184 - # verifies that no rows in logits were missed unexpectedly

185 - assert logits_processed == logits.shape[0]

186 - return logits

187 -

188 -

189 -def _apply_logits_processors_single_seq(logits_row, logits_processors,

190 - past_tokens_ids,

191 - prompt_tokens_ids) -> torch.Tensor:

192 - for logits_processor in logits_processors:

193 - parameters = inspect.signature(logits_processor).parameters

194 - if len(parameters) == 3:

195 - logits_row = logits_processor(prompt_tokens_ids, past_tokens_ids,

196 - logits_row)

197 - else:

198 - logits_row = logits_processor(past_tokens_ids, logits_row)

199 - return logits_row

vllm/model_executor/layers/sampler.py deleted

-1198

@@ -1,1198 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""A layer that samples the next tokens from the model's outputs."""

4 -import itertools

5 -from collections.abc import Iterator

6 -from dataclasses import dataclass

7 -from importlib.util import find_spec

8 -from math import inf

9 -from typing import Optional, Union

10 -

11 -import msgspec

12 -import torch

13 -import torch.nn as nn

14 -

15 -import vllm.envs as envs

16 -from vllm.logprobs import Logprob, PromptLogprobs, SampleLogprobs

17 -from vllm.model_executor.layers.utils import apply_penalties

18 -from vllm.model_executor.sampling_metadata import (SamplingMetadata,

19 - SamplingTensors,

20 - SequenceGroupToSample)

21 -from vllm.sampling_params import SamplingType

22 -from vllm.sequence import (VLLM_INVALID_TOKEN_ID,

23 - CompletionSequenceGroupOutput, SequenceOutput)

24 -

25 -if envs.VLLM_USE_FLASHINFER_SAMPLER and find_spec("flashinfer"):

26 - # yapf: disable

27 - from flashinfer.sampling import (

28 - top_k_top_p_sampling_from_probs as flashinfer_top_k_top_p_sampling)

29 -

30 - # yapf: enable

31 -else:

32 - flashinfer_top_k_top_p_sampling = None

33 -

34 -from vllm.logger import init_logger

35 -

36 -logger = init_logger(__name__)

37 -

38 -

39 -def get_sampler() -> torch.nn.Module:

40 - if envs.VLLM_USE_V1:

41 - # Lazy import: the v1 package isn't distributed

42 - from vllm.v1.sample.sampler import Sampler as V1Sampler

43 - return V1Sampler()

44 - return Sampler()

45 -

46 -

47 -# (num_token_ids, num_parent_ids) per sequence group.

48 -SampleResultType = list[tuple[list[int], list[int]]]

49 -

50 -# Types of temporary data structures used for

51 -# computing sample_result

52 -SampleMetadataType = dict[SamplingType, tuple[list[int],

53 - list[SequenceGroupToSample]]]

54 -MultinomialSamplesType = dict[SamplingType, torch.Tensor]

55 -SampleResultsDictType = dict[int, tuple[list[int], list[int]]]

56 -

57 -

58 -# Encapsulates temporary data structures for computing

59 -# sample_result.

60 -#

61 -# * For multi-step scheduling: must be returned

62 -# by `Sampler.forward()` and used later to compute the pythonized

63 -# sample_result

64 -#

65 -# * For single-step scheduling: consumed immediately

66 -# inside `Sampler.forward()` to compute pythonized sample_result.

67 -@dataclass

68 -class SampleResultArgsType:

69 - sample_metadata: SampleMetadataType

70 - multinomial_samples: MultinomialSamplesType

71 - sample_results_dict: SampleResultsDictType

72 - sampling_metadata: SamplingMetadata

73 - greedy_samples: Optional[torch.Tensor]

74 -

75 -

76 -# Union of non-deferred (single-step scheduling)

77 -# vs deferred (multi-step scheduling)

78 -# sample result types

79 -MaybeDeferredSampleResultType = Union[SampleResultType, SampleResultArgsType]

80 -

81 -# Abbreviation of the _sample() return type

82 -SampleReturnType = tuple[MaybeDeferredSampleResultType, Optional[torch.Tensor]]

83 -

84 -

85 -class SamplerOutput(

86 - msgspec.Struct,

87 - omit_defaults=True, # type: ignore[call-arg]

88 - array_like=True): # type: ignore[call-arg]

89 - """For each sequence group, we generate a list of SequenceOutput object,

90 - each of which contains one possible candidate for the next token.

91 -

92 - This data structure implements methods, so it can be used like a list, but

93 - also has optional fields for device tensors.

94 - """

95 -

96 - outputs: list[CompletionSequenceGroupOutput]

97 -

98 - # On-device tensor containing probabilities of each token.

99 - sampled_token_probs: Optional[torch.Tensor] = None

100 -

101 - # On-device tensor containing the logprobs of each token.

102 - logprobs: Optional["torch.Tensor"] = None

103 -

104 - # Holds either (1) the pythonized sampler result (single-step scheduling)

105 - # or (2) what will be arguments for later deferred pythonization of the

106 - # sampler result (muliti-step scheduling)

107 - deferred_sample_results_args: Optional[SampleResultArgsType] = None

108 -

109 - # On-device tensor containing the sampled token ids.

110 - sampled_token_ids: Optional[torch.Tensor] = None

111 - # CPU tensor containing the sampled token ids. Used during multi-step to

112 - # return the sampled token ids from last rank to AsyncLLMEngine to be

113 - # 'broadcasted' to all other PP ranks for next step.

114 - sampled_token_ids_cpu: Optional[torch.Tensor] = None

115 -

116 - # On-device tensor containing the sampled token embeddings (embeddings

117 - # corresponding to the sampled token ids). Used when prompt embeddings are

118 - # specified in lieu of prompt token ids or text.

119 - sampled_token_embeds: Optional[torch.Tensor] = None

120 -

121 - # Optional last hidden states from the model.

122 - hidden_states: Optional[torch.Tensor] = None

123 -

124 - # Optional prefill hidden states from the model

125 - # (used for models like EAGLE).

126 - prefill_hidden_states: Optional[torch.Tensor] = None

127 -

128 - # Time taken in the forward pass for this across all workers

129 - model_forward_time: Optional[float] = None

130 -

131 - # Time taken in the model execute function. This will include model forward,

132 - # block/sync across workers, cpu-gpu sync time and sampling time.

133 - model_execute_time: Optional[float] = None

134 -

135 - def __getitem__(self, idx: int) -> CompletionSequenceGroupOutput:

136 - return self.outputs[idx]

137 -

138 - def __setitem__(self, idx: int, value):

139 - self.outputs[idx] = value

140 -

141 - def __iter__(self) -> Iterator[CompletionSequenceGroupOutput]:

142 - return iter(self.outputs)

143 -

144 - def __len__(self):

145 - return len(self.outputs)

146 -

147 - def __eq__(self, other: object):

148 - return isinstance(other,

149 - self.__class__) and self.outputs == other.outputs

150 -

151 - def __repr__(self) -> str:

152 - """Show the shape of a tensor instead of its values to reduce noise.

153 - """

154 - sampled_token_probs_repr = ("None" if self.sampled_token_probs is None

155 - else self.sampled_token_probs.shape)

156 - sampled_token_ids_repr = ("None" if self.sampled_token_ids is None else

157 - self.sampled_token_ids.shape)

158 - return (f"SamplerOutput(outputs={self.outputs}, "

159 - f"sampled_token_probs={sampled_token_probs_repr}, "

160 - f"sampled_token_ids={sampled_token_ids_repr})")

161 -

162 -

163 -class Sampler(nn.Module):

164 - """Samples the next tokens from the model's outputs.

165 -

166 - This layer does the following:

167 - 1. Discard the hidden states that are not used for sampling (i.e., all

168 - tokens except the final one in each prompt).

169 - 2. Compute the logits for the next tokens.

170 - 3. Apply presence, frequency and repetition penalties.

171 - 4. Apply temperature scaling.

172 - 5. Apply top-p and top-k truncation.

173 - 6. Sample the next tokens.

174 - Here, each sequence group within the batch can have different sampling

175 - parameters (e.g., sampling method, temperature, top-p, top-k, etc.).

176 -

177 - The structure of the logits tensor is coupled with the seq_groups in

178 - sampling_metadata. Typically, each sequence in each seq_group has one row in

179 - logits for the next token to be sampled; however, for a seq_group with a

180 - prompt request with the prompt_logprobs sampling parameter, there are rows

181 - in logits for each token in the input prompt.

182 - """

183 -

184 - def __init__(self):

185 - super().__init__()

186 -

187 - # Whether or not the SamplerOutput should have on-device tensors

188 - # containing the sampled token ids and probabilities. This is used by

189 - # speculative decoding and when prompt embeddings are specified.

190 - self.include_gpu_probs_tensor = False

191 - self.should_modify_greedy_probs_inplace = False

192 -

193 - def _init_sampling_tensors(

194 - self,

195 - logits: torch.Tensor,

196 - sampling_metadata: SamplingMetadata,

197 - ):

198 - """The goal here is to reuse sampling tensors between similar decode

199 - runs. This is possible because sampling logic does not change between

200 - decodes of the same sequences.

201 - """

202 - _, vocab_size = logits.shape

203 -

204 - # First free any existing stored sampling tensors.

205 - # This is necessary because some sampling tensors may

206 - # have pinned memory.

207 - self._sampling_tensors = None

208 -

209 - # Initialize new sampling tensors

210 - (sampling_tensors, do_penalties, do_top_p_top_k,

211 - do_min_p) = SamplingTensors.from_sampling_metadata(

212 - sampling_metadata, vocab_size, logits.device, logits.dtype)

213 -

214 - self._sampling_tensors = sampling_tensors

215 - self._do_penalties = do_penalties

216 - self._do_top_p_top_k = do_top_p_top_k

217 - self._do_min_p = do_min_p

218 -

219 - def forward(

220 - self,

221 - logits: torch.Tensor,

222 - sampling_metadata: SamplingMetadata,

223 - ) -> Optional[SamplerOutput]:

224 - """

225 - Single-step scheduling:

226 - * Perform GPU-side sampling computation & compute

227 - GPU-side logprobs tensor

228 - * Pythonize sampling result & logprobs tensor

229 -

230 - Multi-step scheduling:

231 - * Perform GPU-side sampling computation & compute

232 - GPU-side logprobs tensor

233 - * Defer Pythonization of sampling result & logprobs

234 - tensor

235 - * Encapsulate arguments required for deferred Pythonization

236 - in the

237 - [`SamplerOutput`][vllm.model_executor.layers.sampler.SamplerOutput]

238 - structure

239 -

240 - Args:

241 - logits: (num_tokens, vocab_size).

242 - sampling_metadata: Metadata for sampling.

243 - """

244 - assert logits is not None

245 - _, vocab_size = logits.shape

246 -

247 - # Prepare sampling tensors with pinned memory to avoid blocking.

248 - if not sampling_metadata.reuse_sampling_tensors:

249 - self._init_sampling_tensors(logits, sampling_metadata)

250 - elif self._do_penalties:

251 - # In this case, the sampling tensors logic depends on

252 - # "output_tokens" of a sequence. As a result, we cannot

253 - # reuse sampling tensors, since "output_tokens" changes

254 - # between decode runs.

255 - self._init_sampling_tensors(logits, sampling_metadata)

256 -

257 - assert self._sampling_tensors is not None

258 - sampling_tensors = self._sampling_tensors

259 - do_penalties = self._do_penalties

260 - do_top_p_top_k = self._do_top_p_top_k

261 - do_min_p = self._do_min_p

262 -

263 - logits = _apply_min_tokens_penalty(logits, sampling_metadata)

264 -

265 - # Apply presence and frequency penalties.

266 - if do_penalties:

267 - logits = apply_penalties(logits, sampling_tensors.prompt_tokens,

268 - sampling_tensors.output_tokens,

269 - sampling_tensors.presence_penalties,

270 - sampling_tensors.frequency_penalties,

271 - sampling_tensors.repetition_penalties)

272 -

273 - # Use float32 to apply temperature scaling.

274 - # Use in-place division to avoid creating a new tensor.

275 - logits = logits.to(torch.float)

276 - logits.div_(sampling_tensors.temperatures.unsqueeze(dim=1))

277 -

278 - if do_top_p_top_k and flashinfer_top_k_top_p_sampling is None:

279 - logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps,

280 - sampling_tensors.top_ks)

281 -

282 - if do_min_p:

283 - logits = _apply_min_p(logits, sampling_tensors.min_ps)

284 -

285 - # We use float32 for probabilities and log probabilities.

286 - # Compute the probabilities.

287 - probs = torch.softmax(logits, dim=-1, dtype=torch.float)

288 - # Compute the log probabilities.

289 - logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)

290 -

291 - # Sample the next tokens.

292 - maybe_deferred_sample_results, maybe_sampled_tokens_tensor = _sample(

293 - probs,

294 - logprobs,

295 - sampling_metadata,

296 - sampling_tensors,

297 - include_gpu_probs_tensor=self.include_gpu_probs_tensor,

298 - modify_greedy_probs=self._should_modify_greedy_probs_inplace,

299 - )

300 -

301 - if self.include_gpu_probs_tensor:

302 - # Since we will defer sampler result Pythonization,

303 - # preserve GPU-side tensors in support of later

304 - # deferred pythonization of logprobs

305 - assert maybe_sampled_tokens_tensor is not None

306 - on_device_tensors = (probs, logprobs, maybe_sampled_tokens_tensor)

307 - else:

308 - # Since Pythonization has already happened, don't preserve

309 - # GPU-side tensors.

310 - on_device_tensors = None

311 -

312 - # Get the logprobs query results.

313 - prompt_logprobs = None

314 - sample_logprobs = None

315 - if not sampling_metadata.skip_sampler_cpu_output:

316 - # Pythonize logprobs now (GPU -> CPU); do not defer.

317 - assert not isinstance(maybe_deferred_sample_results,

318 - SampleResultArgsType)

319 - prompt_logprobs, sample_logprobs = get_logprobs(

320 - logprobs, sampling_metadata, maybe_deferred_sample_results)

321 -

322 - return _build_sampler_output(

323 - maybe_deferred_sample_results,

324 - sampling_metadata,

325 - prompt_logprobs,

326 - sample_logprobs,

327 - on_device_tensors=on_device_tensors,

328 - skip_sampler_cpu_output=sampling_metadata.skip_sampler_cpu_output)

329 -

330 - @property

331 - def _should_modify_greedy_probs_inplace(self) -> bool:

332 - """Whether or not the sampler should modify the probability distribution

333 - of greedily-sampled tokens such that multinomial sampling would sample

334 - the greedily-sampled token.

335 -

336 - In other words, if True then we set the probability of the greedily-

337 - sampled token to 1.

338 -

339 - This is used by speculative decoding, which requires that the sampling

340 - method be encoded into the probability distribution.

341 - """

342 - return self.should_modify_greedy_probs_inplace

343 -

344 -

345 -def _apply_min_tokens_penalty(

346 - logits: torch.Tensor,

347 - sampling_metadata: SamplingMetadata,

348 -) -> torch.Tensor:

349 - """Apply min_tokens penalty which sets stop tokens to -inf if min_tokens

350 - have not been generated yet

351 - """

352 - # list of indices in logits that will be set to -inf

353 - logits_to_penalize: list[tuple[int, int]] = []

354 - logits_applied = 0

355 - for seq_group in sampling_metadata.seq_groups:

356 - seq_ids = seq_group.seq_ids

357 - sampling_params = seq_group.sampling_params

358 -

359 - sample_indices = seq_group.sample_indices

360 - logits_applied += len(sample_indices) + len(

361 - seq_group.prompt_logprob_indices)

362 - if not seq_group.do_sample:

363 - continue

364 -

365 - start_idx = sample_indices[0]

366 - min_tokens = sampling_params.min_tokens

367 - token_ids_to_penalize = sampling_params.all_stop_token_ids

368 - if min_tokens > 0 and token_ids_to_penalize:

369 - seqs_to_penalize: list[int] = []

370 - for j, seq_id in enumerate(seq_ids):

371 - seq_data = seq_group.seq_data[seq_id]

372 - if len(seq_data.output_token_ids_array) < min_tokens:

373 - seqs_to_penalize.append(j)

374 -

375 - if seqs_to_penalize:

376 - # convert to the index into logits

377 - seqs_to_penalize = [start_idx + j for j in seqs_to_penalize]

378 - # itertools.product pairs each seq index with every token id

379 - logits_to_penalize.extend(

380 - itertools.product(seqs_to_penalize, token_ids_to_penalize))

381 -

382 - if logits_to_penalize:

383 - # use zip and * to group indices along each dimension

384 - # eg. [ (1,2), (1,3), (5,6) ] -> ( (1,1,5), (2,3,6) )

385 - logits[tuple(zip(*logits_to_penalize))] = -float("inf")

386 -

387 - # verifies that no rows in logits were missed unexpectedly

388 - assert logits_applied == logits.shape[0]

389 - return logits

390 -

391 -

392 -def _apply_top_k_top_p(

393 - logits: torch.Tensor,

394 - p: torch.Tensor,

395 - k: torch.Tensor,

396 -) -> torch.Tensor:

397 - logits_sort, logits_idx = logits.sort(dim=-1, descending=False)

398 -

399 - # Apply top-k.

400 - top_k_mask = logits_sort.size(1) - k.to(torch.long)

401 - # Get all the top_k values.

402 - top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1))

403 - top_k_mask = logits_sort < top_k_mask

404 - logits_sort.masked_fill_(top_k_mask, -float("inf"))

405 -

406 - # Apply top-p.

407 - probs_sort = logits_sort.softmax(dim=-1)

408 - probs_sum = probs_sort.cumsum(dim=-1)

409 - top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)

410 - # at least one

411 - top_p_mask[:, -1] = False

412 - logits_sort.masked_fill_(top_p_mask, -float("inf"))

413 -

414 - # Re-sort the probabilities.

415 - logits = torch.empty_like(logits_sort).scatter_(dim=-1,

416 - index=logits_idx,

417 - src=logits_sort)

418 - return logits

419 -

420 -

421 -def _apply_min_p(

422 - logits: torch.Tensor,

423 - min_p: torch.Tensor,

424 -) -> torch.Tensor:

425 - """

426 - Adapted from

427 - https://github.com/oobabooga/text-generation-webui/blob/3146124ec01f02c8fb1650a6517cf1b60b537aaf/modules/sampler_hijack.py#L16C17-L16C17

428 - """

429 - probs = torch.softmax(logits, dim=-1)

430 - top_probs, _ = probs.max(dim=-1, keepdim=True)

431 - scaled_min_p = min_p.unsqueeze_(dim=1) * top_probs

432 - tokens_to_remove = probs < scaled_min_p

433 - logits = logits.masked_fill_(tokens_to_remove, -float("inf"))

434 -

435 - return logits

436 -

437 -

438 -def _greedy_sample(

439 - selected_seq_groups: list[SequenceGroupToSample],

440 - samples: torch.Tensor,

441 -) -> SampleResultType:

442 - """Run greedy sampling on a given samples.

443 -

444 - Args:

445 - selected_seq_groups: A list of sequence groups batched.

446 - samples: (num_selected_samples,) A tensor of samples. The length of

447 - samples could be smaller than selected_seq_groups if

448 - seq_group.do_sample is False.

449 - Returns:

450 - Tuple of (next_token_ids, parent_ids). The length of returned list is

451 - same as the length of selected_seq_groups. If the corresponding

452 - seq_group has do_sample=False, tuple contains ([], [])

453 - """

454 - samples_lst = samples.tolist()

455 - sample_idx = 0

456 - results: SampleResultType = []

457 - for seq_group in selected_seq_groups:

458 - if not seq_group.do_sample:

459 - results.append(([], []))

460 - continue

461 -

462 - seq_ids = seq_group.seq_ids

463 - num_parent_seqs = len(seq_ids)

464 - assert num_parent_seqs == 1, (

465 - "Greedy sampling should have only one seq.")

466 - parent_ids = list(range(num_parent_seqs))

467 - next_token_ids = [samples_lst[sample_idx]]

468 - results.append((next_token_ids, parent_ids))

469 - sample_idx += num_parent_seqs

470 - return results

471 -

472 -

473 -def _random_sample(

474 - selected_seq_groups: list[SequenceGroupToSample],

475 - random_samples: torch.Tensor,

476 -) -> SampleResultType:

477 - """Run random sampling on a given samples.

478 -

479 - Args:

480 - selected_seq_groups: A list of sequence groups batched.

481 - random_samples: (num_selected_samples,) A tensor of samples. The

482 - length of samples could be smaller than selected_seq_groups if

483 - seq_group.do_sample is False.

484 - Returns:

485 - Tuple of (next_token_ids, parent_ids). The length of returned list is

486 - same as the length of selected_seq_groups. If the corresponding

487 - seq_group has do_sample=False, tuple contains ([], [])

488 - """

489 - # Find the maximum n value of the prompt phase requests.

490 - random_samples = random_samples.cpu()

491 - sample_idx = 0

492 - results: SampleResultType = []

493 - for seq_group in selected_seq_groups:

494 - if not seq_group.do_sample:

495 - results.append(([], []))

496 - continue

497 -

498 - seq_ids = seq_group.seq_ids

499 - sampling_params = seq_group.sampling_params

500 - is_prompt = seq_group.is_prompt

501 - num_parent_seqs = len(seq_ids)

502 - if is_prompt:

503 - # Prompt phase.

504 - parent_ids = [0] * sampling_params.n

505 - next_token_ids = random_samples[

506 - sample_idx, :sampling_params.n].tolist()

507 - else:

508 - # Generation phase.

509 - parent_ids = list(range(num_parent_seqs))

510 - next_token_ids = random_samples[sample_idx:sample_idx +

511 - num_parent_seqs, 0].tolist()

512 - results.append((next_token_ids, parent_ids))

513 - sample_idx += num_parent_seqs

514 - return results

515 -

516 -

517 -# torch.multinomial forces a GPU<->CPU sync.

518 -# Therefore, we use an optimized implementation instead.

519 -# Note that we always sample with replacement.

520 -# probs will be modified in place, but this is fine, as we pass

521 -# in a copy already.

522 -def _multinomial(

523 - probs: torch.Tensor,

524 - num_samples: int,

525 - seq_groups: Optional[list[SequenceGroupToSample]] = None,

526 -) -> torch.Tensor:

527 - if num_samples > 1:

528 - probs = probs.repeat_interleave(num_samples, dim=0)

529 - q = torch.empty_like(probs)

530 - if seq_groups is None:

531 - q.exponential_()

532 - else:

533 - sample_idx = 0

534 - for seq_group in seq_groups:

535 - seq_ids = seq_group.seq_ids

536 - stride = len(seq_ids) * num_samples

537 - assert seq_group.generator is not None

538 - q[sample_idx:sample_idx +

539 - stride].exponential_(generator=seq_group.generator)

540 - sample_idx += stride

541 - return probs.div_(q).argmax(dim=1).view(-1, num_samples)

542 -

543 -

544 -def _top_k_top_p_multinomial_with_flashinfer(

545 - probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor,

546 - num_samples: int, seq_groups: Optional[list[SequenceGroupToSample]]):

547 - if num_samples > 1:

548 - probs = probs.repeat_interleave(num_samples, dim=0)

549 - top_ks = top_ks.repeat_interleave(num_samples)

550 - top_ps = top_ps.repeat_interleave(num_samples)

551 - batch_next_token_ids = flashinfer_top_k_top_p_sampling(

552 - probs,

553 - top_ks,

554 - top_ps,

555 - )

556 - return batch_next_token_ids.view(-1, num_samples)

557 -

558 -

559 -def get_pythonized_sample_results(

560 - sample_result_args: SampleResultArgsType) -> SampleResultType:

561 - '''This function consumes GPU-side sampler results and computes

562 - Pythonized CPU-side sampler results (GPU -> CPU sync.)

563 -

564 - Single-step scheduling: this function is invoked at sampling-time

565 - for immediate Pythonization.

566 -

567 - Multi-step scheduling: Pythonization is deferred until after multiple

568 - GPU-side steps have been completed.

569 -

570 - Args:

571 - sample_result_args: GPU-side inputs to the Pythonization process

572 -

573 - Returns:

574 - Pythonized sampler results

575 - '''

576 -

577 - (

578 - sample_metadata,

579 - sampling_metadata,

580 - greedy_samples,

581 - multinomial_samples,

582 - sample_results_dict,

583 - ) = (

584 - sample_result_args.sample_metadata,

585 - sample_result_args.sampling_metadata,

586 - sample_result_args.greedy_samples,

587 - sample_result_args.multinomial_samples,

588 - sample_result_args.sample_results_dict,

589 - )

590 -

591 - for sampling_type in SamplingType:

592 - if sampling_type not in sample_metadata:

593 - continue

594 - (seq_group_id, seq_groups) = sample_metadata[sampling_type]

595 - if sampling_type == SamplingType.GREEDY:

596 - sample_results = _greedy_sample(seq_groups, greedy_samples)

597 - elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED):

598 - sample_results = _random_sample(seq_groups,

599 - multinomial_samples[sampling_type])

600 - sample_results_dict.update(zip(seq_group_id, sample_results))

601 -

602 - return [

603 - sample_results_dict.get(i, ([], []))

604 - for i in range(len(sampling_metadata.seq_groups))

605 - ]

606 -

607 -

608 -def _sample_with_torch(

609 - probs: torch.Tensor,

610 - logprobs: torch.Tensor,

611 - sampling_metadata: SamplingMetadata,

612 - sampling_tensors: SamplingTensors,

613 - include_gpu_probs_tensor: bool,

614 - modify_greedy_probs: bool,

615 -) -> SampleReturnType:

616 - '''Torch-oriented _sample() implementation.

617 -

618 - Single-step scheduling:

619 - * Perform GPU-side sampling computation

620 - * Immediately Pythonize sampling result

621 -

622 - Multi-step scheduling:

623 - * Perform GPU-side sampling computation

624 - * Defer Pythonization & preserve GPU-side

625 - tensors required for Pythonization

626 - '''

627 -

628 - categorized_seq_group_ids: dict[SamplingType, list[int]] = {

629 - t: []

630 - for t in SamplingType

631 - }

632 - categorized_sample_indices = sampling_metadata.categorized_sample_indices

633 - for i, seq_group in enumerate(sampling_metadata.seq_groups):

634 - sampling_params = seq_group.sampling_params

635 - sampling_type = sampling_params.sampling_type

636 - categorized_seq_group_ids[sampling_type].append(i)

637 -

638 - sample_results_dict: SampleResultsDictType = {}

639 - sample_metadata: SampleMetadataType = {}

640 - multinomial_samples: MultinomialSamplesType = {}

641 - greedy_samples: Optional[torch.Tensor] = None

642 -

643 - # Create output tensor for sampled token ids.

644 - if include_gpu_probs_tensor:

645 - sampled_token_ids_tensor = torch.full((logprobs.shape[0], 1),

646 - VLLM_INVALID_TOKEN_ID,

647 - dtype=torch.long,

648 - device=logprobs.device)

649 - else:

650 - sampled_token_ids_tensor = None

651 -

652 - # Counterintuitively, having two loops here is actually faster.

653 - # The first loop can run without waiting on GPU<->CPU sync.

654 - for sampling_type in SamplingType:

655 - sample_indices = categorized_sample_indices[sampling_type]

656 - num_tokens = len(sample_indices)

657 - if num_tokens == 0:

658 - continue

659 -

660 - seq_group_id = categorized_seq_group_ids[sampling_type]

661 - seq_groups = [sampling_metadata.seq_groups[i] for i in seq_group_id]

662 - sample_metadata[sampling_type] = (seq_group_id, seq_groups)

663 - long_sample_indices = sample_indices.long()

664 - if sampling_type == SamplingType.GREEDY:

665 - greedy_samples = torch.argmax(logprobs[long_sample_indices],

666 - dim=-1)

667 -

668 - if sampled_token_ids_tensor is not None:

669 - # Store sampled tokens in output tensor.

670 - sampled_token_ids_tensor[

671 - long_sample_indices] = greedy_samples.unsqueeze(-1)

672 -

673 - if modify_greedy_probs:

674 - # If required, modify the probabilities such that sampling from

675 - # the modified distribution would always sample the argmax

676 - # token id.

677 - _modify_greedy_probs_inplace(logprobs, probs,

678 - long_sample_indices,

679 - greedy_samples)

680 -

681 - elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED):

682 - max_n_in_batch = 1

683 - for seq_group in seq_groups:

684 - if seq_group.is_prompt:

685 - sampling_params = seq_group.sampling_params

686 - max_n_in_batch = max(max_n_in_batch, sampling_params.n)

687 - seq_groups_arg = (None if sampling_type == SamplingType.RANDOM else

688 - seq_groups)

689 -

690 - if flashinfer_top_k_top_p_sampling is not None:

691 - logger.warning("FlashInfer 0.2.3+ does not support "

692 - "per-request generators. Falling back to "

693 - "PyTorch-native implementation.")

694 -

695 - multinomial_samples[sampling_type] = _multinomial(

696 - probs[long_sample_indices],

697 - max_n_in_batch,

698 - seq_groups=seq_groups_arg)

699 -

700 - if sampled_token_ids_tensor is not None:

701 - # Store sampled tokens in output tensor.

702 - sampled_token_ids_tensor[long_sample_indices] = \

703 - multinomial_samples[sampling_type].to(torch.long)

704 -

705 - else:

706 - raise ValueError(f"Unsupported sampling type: {sampling_type}")

707 -

708 - # Encapsulate arguments for computing Pythonized sampler

709 - # results, whether deferred or otherwise.

710 - maybe_deferred_args = SampleResultArgsType(

711 - sampling_metadata=sampling_metadata,

712 - sample_metadata=sample_metadata,

713 - multinomial_samples=multinomial_samples,

714 - greedy_samples=greedy_samples,

715 - sample_results_dict=sample_results_dict)

716 -

717 - if not sampling_metadata.skip_sampler_cpu_output:

718 - # GPU<->CPU sync happens here.

719 - # This also converts the sampler output to a Python object.

720 - # Return Pythonized sampler result & sampled token ids

721 - return get_pythonized_sample_results(

722 - maybe_deferred_args), sampled_token_ids_tensor

723 - else:

724 - # Defer sampler result Pythonization; return deferred

725 - # Pythonization args & sampled token ids

726 - return (

727 - maybe_deferred_args,

728 - sampled_token_ids_tensor,

729 - )

730 -

731 -

732 -def _sample(

733 - probs: torch.Tensor,

734 - logprobs: torch.Tensor,

735 - sampling_metadata: SamplingMetadata,

736 - sampling_tensors: SamplingTensors,

737 - include_gpu_probs_tensor: bool,

738 - modify_greedy_probs: bool,

739 -) -> SampleReturnType:

740 - """

741 - Args:

742 - probs: (num_query_tokens_in_batch, num_vocab)

743 - logprobs: (num_query_tokens_in_batch, num_vocab)

744 - sampling_metadata: The metadata for a batch for sampling.

745 - sampling_tensors: Tensors that include sampling related metadata.

746 -

747 - Returns:

748 - (next_token_ids, parent_seq_ids) for each seq group in a batch.

749 - If sampling is skipped, it returns ([], [])

750 - sampled_token_ids_tensor: A tensor of sampled token ids.

751 - """

752 - return _sample_with_torch(

753 - probs,

754 - logprobs,

755 - sampling_metadata,

756 - sampling_tensors,

757 - include_gpu_probs_tensor=include_gpu_probs_tensor,

758 - modify_greedy_probs=modify_greedy_probs,

759 - )

760 -

761 -

762 -def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:

763 - """

764 - This function calculates the ranks of the chosen tokens in a logprob tensor.

765 -

766 - Args:

767 - x (torch.Tensor): 2D logprob tensor of shape (N, M)

768 - where N is the no. of tokens and M is the vocab dim.

769 - indices (torch.Tensor): List of chosen token indices.

770 -

771 - Returns:

772 - torch.Tensor: 1D tensor of shape (N,) where N is the no. of tokens.

773 - Each element in the returned tensor represents the rank

774 - of the chosen token in the input logprob tensor.

775 - """

776 - vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype),

777 - indices]

778 - result = (x > vals[:, None])

779 - del vals

780 - return result.sum(1).add_(1)

781 -

782 -

783 -def get_logprobs(

784 - logprobs: torch.Tensor,

785 - sampling_metadata: SamplingMetadata,

786 - sample_results: SampleResultType,

787 -) -> tuple[list[Optional[PromptLogprobs]], list[SampleLogprobs]]:

788 - """Return sample logprobs and prompt logprobs.

789 -

790 - The logic consists of 3 parts.

791 - - Select indices to compute logprob from, ranks of token ids, and

792 - the top k token ids from logprobs.

793 - - Compute prompt logprobs if required.

794 - - Compute sample logprobs if required.

795 -

796 - Args:

797 - logprobs: (num_query_tokens_across_batch, num_vocab). Each query token's

798 - logprob per vocab. Sequence groups' query tokens are batched in a

799 - single flattened tensor. For example, assuming there are N

800 - seq groups, it is sorted by prefill tokens for seq_group_1 (if

801 - prompt logprob is enabled), decode tokens for seq_group_1 (if

802 - sampling is required), prefill tokens for seq_group_2, ...

803 - sampling_metadata: The sampling metadata.

804 - sample_results: (num_seq_groups) The tuple of (next_token_ids,

805 - parent_ids) for each sequence group. When beam search is enabled,

806 - sample_results can contain different number of seq_ids from

807 - sampling_metadata.seq_groups. It is because beam search creates

808 - 2 * BEAM_WIDTH number of samples (whereas there are only up to

809 - BEAM_WIDTH number of seq_ids).

810 -

811 - Returns:

812 - A tuple of prompt and sample logprobs per sequence group in a batch.

813 - """

814 - # The index of query token to calculate logprobs. It includes both

815 - # prompt and sample logprob indices.

816 - query_indices: list[int] = []

817 - # The next token ids to get the logprob value from.

818 - next_token_ids: list[int] = []

819 - # The largest requested number of logprobs. We find logprobs as many as the

820 - # largest num logprobs in this API. If every logprobs is None, it will be

821 - # set to -1.

822 - largest_num_logprobs = -1

823 -

824 - # Select indices to compute logprob from, ranks of token ids, and the top

825 - # k token ids from logprobs.

826 - for (seq_group, sample_result) in zip(sampling_metadata.seq_groups,

827 - sample_results):

828 - sampling_params = seq_group.sampling_params

829 -

830 - # Update indices and tokens for prompt logprobs.

831 - if (seq_group.is_prompt

832 - and sampling_params.prompt_logprobs is not None):

833 - largest_num_logprobs = max(largest_num_logprobs,

834 - sampling_params.prompt_logprobs)

835 - next_prompt_tokens = _get_next_prompt_tokens(seq_group)

836 - query_indices.extend(seq_group.prompt_logprob_indices)

837 - next_token_ids.extend(next_prompt_tokens)

838 -

839 - # Update indices and next tokenes for sample logprob.

840 - if seq_group.do_sample:

841 - token_ids, parent_seq_ids = sample_result

842 - # NOTE: We cannot directly use sample_indices because

843 - # sample_indices only contain parent seq_ids of a previous step.

844 - # The current step may have different number of seq_ids, and

845 - # we can obtain it from `sample_result[1]`.

846 - query_idx = seq_group.sample_indices[0]

847 - query_indices.extend(

848 - [query_idx + parent_id for parent_id in parent_seq_ids])

849 - next_token_ids.extend(token_ids)

850 -

851 - if sampling_params.logprobs is not None:

852 - largest_num_logprobs = max(largest_num_logprobs,

853 - sampling_params.logprobs)

854 -

855 - assert len(next_token_ids) == len(query_indices)

856 -

857 - if len(query_indices) == 0:

858 - empty_sampled_logprob: SampleLogprobs = []

859 - empty_prompt_logprob: Optional[PromptLogprobs] = None

860 - num_seq_groups = len(sampling_metadata.seq_groups)

861 - return [empty_prompt_logprob

862 - ] * num_seq_groups, [empty_sampled_logprob] * num_seq_groups

863 -

864 - selected_logprobs, ranks = None, None

865 - top_logprobs, top_token_ids = None, None

866 -

867 - # If largest_num_logprobs == -1, i.e. no logprobs are requested, we can

868 - # skip the whole logprob calculation.

869 - if largest_num_logprobs >= 0:

870 - query_indices_gpu = torch.tensor(query_indices, device=logprobs.device)

871 - next_token_ids_gpu = torch.tensor(next_token_ids,

872 - device=logprobs.device)

873 -

874 - # (num_selected_query_tokens, num_logprobs). Note that query_indices can

875 - # contain duplicates if beam search is enabled.

876 - selected_logprobs = logprobs[[

877 - query_indices_gpu,

878 - next_token_ids_gpu,

879 - ]]

880 - ranks = _get_ranks(

881 - logprobs[query_indices_gpu],

882 - next_token_ids_gpu,

883 - )

884 - assert selected_logprobs.shape[0] == ranks.shape[0]

885 -

886 - # We need to compute top k only if there exists logprobs > 0.

887 - if largest_num_logprobs > 0:

888 - # Logprobs of topk tokens for a batch of sequence groups.

889 - # (num_query_tokens_across_batch).

890 - top_logprobs, top_token_ids = torch.topk(logprobs,

891 - largest_num_logprobs,

892 - dim=-1)

893 - top_logprobs = top_logprobs.to('cpu')

894 - top_token_ids = top_token_ids.to('cpu')

895 -

896 - selected_logprobs = selected_logprobs.to('cpu')

897 - ranks = ranks.to('cpu')

898 -

899 - # Find prompt/sample logprobs.

900 - prompt_logprobs_per_seq_group: list[Optional[PromptLogprobs]] = []

901 - sample_logprobs_per_seq_group: list[SampleLogprobs] = []

902 - top_logprob_idx = 0

903 - selected_logprobs_idx = 0

904 -

905 - for seq_group, sample_result in zip(sampling_metadata.seq_groups,

906 - sample_results):

907 - (prompt_logprobs, top_logprob_idx,

908 - selected_logprobs_idx) = _get_prompt_logprob_if_needed(

909 - seq_group, selected_logprobs, ranks, top_token_ids, top_logprobs,

910 - selected_logprobs_idx, top_logprob_idx)

911 - prompt_logprobs_per_seq_group.append(prompt_logprobs)

912 -

913 - (sampled_logprobs, top_logprob_idx,

914 - selected_logprobs_idx) = _get_sampled_logprob_if_needed(

915 - seq_group, sample_result, selected_logprobs, ranks, top_token_ids,

916 - top_logprobs, selected_logprobs_idx, top_logprob_idx)

917 - sample_logprobs_per_seq_group.append(sampled_logprobs)

918 -

919 - return prompt_logprobs_per_seq_group, sample_logprobs_per_seq_group

920 -

921 -

922 -def _get_prompt_logprob_if_needed(

923 - seq_group: SequenceGroupToSample,

924 - selected_logprobs: torch.Tensor,

925 - ranks: torch.Tensor,

926 - top_token_ids: torch.Tensor,

927 - top_logprobs: torch.Tensor,

928 - selected_logprobs_idx: int,

929 - top_logprob_idx: int,

930 -):

931 - """Compute the prompt logprob from a sequence group if needed."""

932 - sampling_params = seq_group.sampling_params

933 - is_prompt = seq_group.is_prompt

934 -

935 - # Find prompt logprobs

936 - prompt_logprobs: Optional[PromptLogprobs] = None

937 - if is_prompt and sampling_params.prompt_logprobs is not None:

938 - prompt_logprobs = []

939 - num_logprobs = sampling_params.prompt_logprobs

940 - next_prompt_tokens = _get_next_prompt_tokens(seq_group)

941 - # Pre-select indexes and create a list. It is faster than calling .item

942 - # repetitively.

943 - selected_logprob_items = selected_logprobs[

944 - selected_logprobs_idx:selected_logprobs_idx +

945 - len(next_prompt_tokens)].tolist()

946 - rank_items = ranks[selected_logprobs_idx:selected_logprobs_idx +

947 - len(next_prompt_tokens)].tolist()

948 -

949 - for idx, token_id in enumerate(next_prompt_tokens):

950 - # Calculate the prompt logprob of the real prompt tokens.

951 - # {token_id: (logprob, rank_from_vocab)}

952 - prompt_logprobs_dict: dict[int, tuple[float, int]] = {

953 - token_id: (selected_logprob_items[idx], rank_items[idx])

954 - }

955 -

956 - # Add top K prompt logprobs along with its rank.

957 - if num_logprobs > 0:

958 - top_ids = top_token_ids[

959 - top_logprob_idx, :num_logprobs].tolist()

960 - top_probs = top_logprobs[

961 - top_logprob_idx, :num_logprobs].tolist()

962 - # Top K is already sorted by rank, so we can use 1 ~

963 - # num_logprobs + 1 for rank.

964 - top_ranks = range(1, num_logprobs + 1)

965 - prompt_logprobs_dict.update({

966 - top_id: (top_prob, rank)

967 - for top_id, top_prob, rank in zip(top_ids, top_probs,

968 - top_ranks)

969 - })

970 - prompt_logprobs.append({

971 - token_id: Logprob(*logprob_and_rank)

972 - for token_id, logprob_and_rank in prompt_logprobs_dict.items()

973 - })

974 - # + 1 to go to the next prompt token.

975 - top_logprob_idx += 1

976 -

977 - # + len(next_prompt_tokens) to go to the next prompt.

978 - selected_logprobs_idx += len(next_prompt_tokens)

979 - return prompt_logprobs, top_logprob_idx, selected_logprobs_idx

980 -

981 -

982 -def _get_sampled_logprob_if_needed(

983 - seq_group: SequenceGroupToSample,

984 - sample_result: tuple[list[int], list[int]],

985 - selected_logprobs: torch.Tensor,

986 - ranks: torch.Tensor,

987 - top_token_ids: torch.Tensor,

988 - top_logprobs: torch.Tensor,

989 - selected_logprobs_idx: int,

990 - top_logprob_idx: int,

991 -):

992 - """Compute the sample logprob if needed."""

993 - seq_ids = seq_group.seq_ids

994 - num_logprobs = seq_group.sampling_params.logprobs

995 - sampled_logprobs: SampleLogprobs = []

996 - next_token_ids, parent_seq_ids = sample_result

997 -

998 - if seq_group.do_sample:

999 - assert len(next_token_ids) > 0

1000 - if num_logprobs is None:

1001 - for next_token_id in next_token_ids:

1002 - # Use a dummy logprob

1003 - sampled_logprobs.append({next_token_id: Logprob(inf)})

1004 - else:

1005 - # Pre-select items from tensor. tolist() is faster than repetitive

1006 - # `.item()` calls.

1007 - selected_logprob_items = selected_logprobs[

1008 - selected_logprobs_idx:selected_logprobs_idx +

1009 - len(next_token_ids)].tolist()

1010 - rank_items = ranks[selected_logprobs_idx:selected_logprobs_idx +

1011 - len(next_token_ids)].tolist()

1012 - for idx, (next_token_id, parent_id) in enumerate(

1013 - zip(next_token_ids, parent_seq_ids)):

1014 - # Get the logprob of a sampled token.

1015 - sampled_logprobs_dict = {

1016 - next_token_id:

1017 - (selected_logprob_items[idx], rank_items[idx])

1018 - }

1019 - if num_logprobs is not None and num_logprobs > 0:

1020 - # Get top K logprobs.

1021 - top_ids = top_token_ids[top_logprob_idx +

1022 - parent_id, :num_logprobs].tolist()

1023 - top_probs = top_logprobs[

1024 - top_logprob_idx + parent_id, :num_logprobs].tolist()

1025 - # Top K is already sorted by rank, so we can use 1 ~

1026 - # num_logprobs + 1 for rank.

1027 - top_ranks = range(1, num_logprobs + 1)

1028 - sampled_logprobs_dict.update({

1029 - top_id: (top_prob, rank)

1030 - for top_id, top_prob, rank in zip(

1031 - top_ids, top_probs, top_ranks)

1032 - })

1033 -

1034 - sampled_logprobs.append({

1035 - token_id: Logprob(*logprob_and_rank)

1036 - for token_id, logprob_and_rank in

1037 - sampled_logprobs_dict.items()

1038 - })

1039 -

1040 - # NOTE: This part of code is not intuitive. `selected_logprobs` include

1041 - # logprobs for the current step, which has len(next_token_ids) tokens

1042 - # per sequence group. `logprobs` includes logprobs from the previous

1043 - # steps, which has len(seq_ids) tokens per sequence group.

1044 -

1045 - # Iterate to the next sequence group in a batch.

1046 - selected_logprobs_idx += len(next_token_ids)

1047 - # Iterate to the next sequence group in a batch.

1048 - top_logprob_idx += len(seq_ids)

1049 - return sampled_logprobs, top_logprob_idx, selected_logprobs_idx

1050 -

1051 -

1052 -def _modify_greedy_probs_inplace(logprobs: torch.Tensor, probs: torch.Tensor,

1053 - sample_indices: torch.Tensor,

1054 - greedy_samples: torch.Tensor) -> None:

1055 - """Modify the probability distributions of the greedily-sampled tokens such

1056 - that each sampled token has a "probability" of 1.0. This is required by

1057 - speculative decoding, which depends on the sampling method being encoded

1058 - within the probability distribution for correctness.

1059 -

1060 - # Why do we only need to do this for greedy sampling?

1061 -

1062 - vLLM's sampler performs the following steps for greedy or multinomial

1063 - (random) sampling:

1064 - 1. Get logits from model.

1065 - 2. Modify logits according to per-sequence sampling parameters.

1066 - - Multiply by temperature, top-k and top-p masking, penalize tokens

1067 - according to their frequency, etc.

1068 - 3. Sample a token.

1069 - - Random sampling simply samples from the modified probability

1070 - distribution.

1071 - - Greedy sampling performs `argmax` to obtain the token with the

1072 - highest likelihood.

1073 -

1074 - Ignoring greedy sampling for a moment, we find that the computed probability

1075 - distribution has the following property: we can sample from it independently

1076 - and find that the token sampled by the Sampler has a frequency corresponding

1077 - to how often we see it in our sampling. In other words, for tokens sampled

1078 - with vLLM's random SamplingType, the computed probability distribution

1079 - encodes the sampling methodology completely.

1080 -

1081 - Greedy sampling does not normally have this property. vLLM modifies logits

1082 - according to sampling params, then performs `argmax`, then returns the

1083 - sampled token and the computed probability distribution. If we sample from

1084 - the distribution, we'll find the likelihood of the greedily-sampled token

1085 - is not always 1.0.

1086 -

1087 - Since lossless speculative decoding requires that the sampling methodology

1088 - be encoded within the probability distribution, we are motivated to modify

1089 - the probability distribution such that the sampled token has probability 1

1090 - when speculative decoding is used.

1091 -

1092 - NOTE: Alternatively, we could use an extremely low temperature to achieve

1093 - greedy sampling using multinomial computation and unite the codepaths. This

1094 - has implications on the overall design of the sampler, e.g. how to record

1095 - accurate logprobs for the user, so this improvement is deferred to later.

1096 - """

1097 - # NOTE: logprobs are not modified so they can be returned to the user.

1098 - probs[sample_indices, :] = 0

1099 - probs[sample_indices, greedy_samples] = 1.0

1100 -

1101 -

1102 -def _build_sampler_output(

1103 - maybe_deferred_sample_results: MaybeDeferredSampleResultType,

1104 - sampling_metadata: SamplingMetadata,

1105 - prompt_logprobs: Optional[list[Optional[PromptLogprobs]]],

1106 - sample_logprobs: Optional[list[SampleLogprobs]],

1107 - on_device_tensors: Optional[tuple[torch.Tensor, torch.Tensor,

1108 - torch.Tensor]],

1109 - skip_sampler_cpu_output: bool = False,

1110 -) -> SamplerOutput:

1111 - """Construct Python objects with the output of sampling.

1112 -

1113 - Args:

1114 - on_device_tensors: Tuple containing on-device tensors with the

1115 - probabilities used in sampling and the sampled token ids. This

1116 - allows post-processing without copies to CPU/serialization, e.g. in

1117 - speculative decoding rejection sampling.

1118 - """

1119 - sampler_output: list[CompletionSequenceGroupOutput] = []

1120 -

1121 - if skip_sampler_cpu_output:

1122 - assert isinstance(maybe_deferred_sample_results, SampleResultArgsType)

1123 - deferred_sample_results_args = maybe_deferred_sample_results

1124 - else:

1125 - assert prompt_logprobs is not None

1126 - assert sample_logprobs is not None

1127 - assert not isinstance(maybe_deferred_sample_results,

1128 - SampleResultArgsType)

1129 - assert len(sampling_metadata.seq_groups) \

1130 - == len(maybe_deferred_sample_results) \

1131 - == len(prompt_logprobs) \

1132 - == len(sample_logprobs)

1133 - deferred_sample_results_args = None

1134 -

1135 - for (seq_group, sample_result, group_prompt_logprobs,

1136 - group_sample_logprobs) in zip(sampling_metadata.seq_groups,

1137 - maybe_deferred_sample_results,

1138 - prompt_logprobs, sample_logprobs):

1139 - seq_ids = seq_group.seq_ids

1140 - next_token_ids, parent_ids = sample_result

1141 - seq_outputs: list[SequenceOutput] = []

1142 - for parent_id, next_token_id, logprobs in zip(

1143 - parent_ids, next_token_ids, group_sample_logprobs):

1144 - seq_outputs.append(

1145 - SequenceOutput(seq_ids[parent_id], next_token_id,

1146 - logprobs))

1147 - sampler_output.append(

1148 - CompletionSequenceGroupOutput(seq_outputs,

1149 - group_prompt_logprobs))

1150 -

1151 - # If not specified, store None values in SamplerOutput.

1152 - if on_device_tensors is not None:

1153 - (sampled_token_probs, logprobs_tensor,

1154 - sampled_token_ids) = on_device_tensors

1155 - else:

1156 - sampled_token_probs, logprobs_tensor, sampled_token_ids = (None, None,

1157 - None)

1158 -

1159 - return SamplerOutput(

1160 - outputs=sampler_output,

1161 - sampled_token_probs=sampled_token_probs,

1162 - sampled_token_ids=sampled_token_ids,

1163 - logprobs=logprobs_tensor,

1164 - deferred_sample_results_args=deferred_sample_results_args)

1165 -

1166 -

1167 -def _get_next_prompt_tokens(

1168 - seq_group: SequenceGroupToSample) -> tuple[int, ...]:

1169 - """Get a list of next prompt tokens to compute logprob from a

1170 - given sequence group.

1171 -

1172 - It is used to compute prompt logprob. Imagine you have logprob for each

1173 - query token. Query token needs to know the next prompt token id to compute

1174 - prompt logprob. This is a helper to obtain next prompt token ids.

1175 -

1176 - This API has to be used only when the caller knows seq_group is in prefill

1177 - stage.

1178 -

1179 - Returns:

1180 - A list of next prompt tokens to compute logprob.

1181 - """

1182 - assert seq_group.is_prompt, (

1183 - "Caller should ensure the sequence group is in a prefill stage.")

1184 - seq_ids = seq_group.seq_ids

1185 - query_len = seq_group.query_len

1186 - assert query_len is not None

1187 - # prompt has only 1 seq id.

1188 - assert len(seq_ids) == 1

1189 - seq_data = seq_group.seq_data[seq_ids[0]]

1190 - computed_len = seq_data.get_num_computed_tokens()

1191 - prompt_tokens = seq_data.prompt_token_ids

1192 - # +1 because we are looking for a next prompt token.

1193 - next_token_index_start = computed_len + 1

1194 - next_token_index_end = min(computed_len + query_len + 1,

1195 - len(prompt_tokens))

1196 - next_prompt_tokens = prompt_tokens[

1197 - next_token_index_start:next_token_index_end]

1198 - return next_prompt_tokens

vllm/model_executor/models/medusa.py

+4 -56

@@ -2,18 +2,15 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from collections.abc import Iterable

5 -from typing import Optional

6 import torch

7 import torch.nn as nn

9 from vllm.config import VllmConfig

10 from vllm.model_executor.layers.logits_processor import LogitsProcessor

12 -from vllm.model_executor.layers.sampler import SamplerOutput

11 from vllm.model_executor.layers.vocab_parallel_embedding import (

12 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead)

13 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

16 -from vllm.model_executor.sampling_metadata import SamplingMetadata

15 from .utils import maybe_prefix

@@ -105,8 +102,10 @@ class Medusa(nn.Module):

102 return [block(hidden_states) for block in self.blocks]

103

104 def compute_logits(

108 - self, hidden_states: list[torch.Tensor],

109 - sampling_metadata: SamplingMetadata) -> list[torch.Tensor]:

105 + self,

106 + hidden_states: list[torch.Tensor],

107 + sampling_metadata,

108 + ) -> list[torch.Tensor]:

109 logits_lst: list[torch.Tensor] = []

110

111 for hs, lm_head in zip(hidden_states, self.lm_heads):

@@ -130,57 +129,6 @@ class Medusa(nn.Module):

129

130 return logits_lst

131

133 - def sample(

134 - self,

135 - logits: list[torch.Tensor],

136 - sampling_metadata: SamplingMetadata,

137 - ) -> list[SamplerOutput]:

138 - logits = torch.stack(logits, dim=0).float()

139 - logprobs = torch.log_softmax(logits, dim=-1)

140 - token_ids = logits.argmax(-1) # support only top-1 for now

141 - probs = torch.softmax(logits, dim=-1)

142 -

143 - token_id_list = []

144 - token_prob_list = []

145 - token_logprob_list = []

146 -

147 - for idx, seq_group in enumerate(sampling_metadata.seq_groups):

148 - token_id_list.append(token_ids[:, seq_group.sample_indices])

149 - token_prob_list.append(probs[:, seq_group.sample_indices])

150 - token_logprob_list.append(logprobs[:, seq_group.sample_indices])

151 -

152 - outputs: list[Optional[SamplerOutput]] = []

153 - for idx in range(len(sampling_metadata.seq_groups)):

154 - outputs.append(

155 - SamplerOutput(

156 - outputs=None,

157 - sampled_token_probs=token_prob_list[idx].squeeze(1),

158 - logprobs=token_logprob_list[idx].squeeze(1),

159 - sampled_token_ids=token_id_list[idx].squeeze(1),

160 - ))

161 -

162 - return outputs

163 -

164 - def generate_proposals(

165 - self,

166 - previous_hidden_states: torch.Tensor,

167 - sampling_metadata: SamplingMetadata,

168 - ) -> Optional[list[SamplerOutput]]:

169 - # During preemption, we may receive an empty tensor (batch_size=0)

170 - if previous_hidden_states.size(0) == 0:

171 - # Return None to signal the Top1Proposer that no proposals

172 - # were generated for this batch, allowing it to handle this

173 - # special case appropriately

174 - return None

175 -

176 - return self.sample(

177 - logits=self.compute_logits(

178 - hidden_states=self.forward(previous_hidden_states),

179 - sampling_metadata=sampling_metadata,

180 - ),

181 - sampling_metadata=sampling_metadata,

182 - )

183 -

132 def load_weights(self, weights: Iterable[tuple[str,

133 torch.Tensor]]) -> set[str]:

134 params_dict = dict(self.named_parameters())

vllm/model_executor/models/mlp_speculator.py

+40 -40

@@ -8,9 +8,7 @@ import torch

8 import torch.nn as nn

10 from vllm.config import VllmConfig

11 -from vllm.model_executor import SamplingMetadata

11 from vllm.model_executor.layers.logits_processor import LogitsProcessor

13 -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

12 from vllm.model_executor.layers.vocab_parallel_embedding import (

13 ParallelLMHead, VocabParallelEmbedding)

14 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

@@ -141,55 +139,57 @@ class MLPSpeculator(nn.Module):

139 self.config = config

140 self.logits_processor = LogitsProcessor(config.vocab_size,

141 config.vocab_size, 1.0)

144 - self.sampler = get_sampler()

142

146 - def generate_proposals(

147 - self,

148 - input_ids: torch.Tensor,

149 - previous_hidden_states: torch.Tensor,

150 - num_predict_tokens: int,

151 - sampling_metadata: SamplingMetadata,

152 - ) -> list[SamplerOutput]:

153 - if num_predict_tokens > self.max_speculative_tokens:

154 - raise ValueError(f"Max speculative tokens for model is "

155 - f"{self.max_speculative_tokens}, but "

156 - f"{num_predict_tokens} were requested")

157 -

158 - # b x 1 x d

159 - previous_hidden_states = previous_hidden_states.unsqueeze(1)

143 + # NOTE(woosuk): This method is commented out because it is old code

144 + # using V0. We should either port it to V1 or remove it.

145

161 - if self.scale_input:

162 - previous_hidden_states = self.ln0(previous_hidden_states) / SQRT2

146 + # def generate_proposals(

147 + # self,

148 + # input_ids: torch.Tensor,

149 + # previous_hidden_states: torch.Tensor,

150 + # num_predict_tokens: int,

151 + # sampling_metadata: SamplingMetadata,

152 + # ) -> list[SamplerOutput]:

153 + # if num_predict_tokens > self.max_speculative_tokens:

154 + # raise ValueError(f"Max speculative tokens for model is "

155 + # f"{self.max_speculative_tokens}, but "

156 + # f"{num_predict_tokens} were requested")

157 +

158 + # # b x 1 x d

159 + # previous_hidden_states = previous_hidden_states.unsqueeze(1)

160 +

161 + # if self.scale_input:

162 + # previous_hidden_states = self.ln0(previous_hidden_states) / SQRT2

163

164 - # b x 1

165 - last_tokens = input_ids.unsqueeze(1)

164 + # # b x 1

165 + # last_tokens = input_ids.unsqueeze(1)

166

167 - next_tokens = []

167 + # next_tokens = []

168

169 - for head_index in range(num_predict_tokens):

169 + # for head_index in range(num_predict_tokens):

170

171 - # Project and predict

172 - z = self.emb[head_index](last_tokens) # b k d

173 - states = self.proj[head_index](previous_hidden_states)

171 + # # Project and predict

172 + # z = self.emb[head_index](last_tokens) # b k d

173 + # states = self.proj[head_index](previous_hidden_states)

174

175 - # Weighted add of state_weight*state and emb_weight*z

176 - # Let subsequent LN take care of denominator

177 - # state_weight is close to 1, so shouldn't be any precision issues

178 - states.add_(z, alpha=self.emb_weight / self.state_weight)

175 + # # Weighted add of state_weight*state and emb_weight*z

176 + # # Let subsequent LN take care of denominator

177 + # # state_weight is close to 1, so shouldn't be any precision issues

178 + # states.add_(z, alpha=self.emb_weight / self.state_weight)

179

180 - states = self.activation(self.ln[head_index](states)) # b k d

181 - previous_hidden_states = states

182 - # TODO: not yet supporting top_k_tokens_per_head

183 - states = states.flatten(0, 1)

180 + # states = self.activation(self.ln[head_index](states)) # b k d

181 + # previous_hidden_states = states

182 + # # TODO: not yet supporting top_k_tokens_per_head

183 + # states = states.flatten(0, 1)

184

185 - logits = self.logits_processor(self.head[head_index], states,

186 - sampling_metadata)

185 + # logits = self.logits_processor(self.head[head_index], states,

186 + # sampling_metadata)

187

188 - output = self.sampler(logits, sampling_metadata)

189 - last_tokens = output.sampled_token_ids

190 - next_tokens.append(output)

188 + # output = self.sampler(logits, sampling_metadata)

189 + # last_tokens = output.sampled_token_ids

190 + # next_tokens.append(output)

191

192 - return next_tokens

192 + # return next_tokens

193

194 def load_weights(self, weights: Iterable[tuple[str,

195 torch.Tensor]]) -> set[str]:

vllm/model_executor/models/phi4flash.py

+1 -5

@@ -697,16 +697,12 @@ class Phi4FlashForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

697 hidden_states: torch.Tensor,

698 sampling_metadata: SamplingMetadata,

699 ) -> Optional[torch.Tensor]:

700 - # If the shape is the same, it means that we have already

701 - # prune hidden states manually.

702 - prune_hidden_states = hidden_states.size(

703 - 0) != sampling_metadata.selected_token_indices.size(0)

700 processed_logits = self.logits_processor(

701 self.lm_head,

702 hidden_states,

703 sampling_metadata,

704 self.embedding_bias,

709 - prune_hidden_states=prune_hidden_states)

705 + )

706 return processed_logits

707

708 def load_weights(

vllm/model_executor/sampling_metadata.py

+2 -592

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

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from array import array

5 -from dataclasses import dataclass

6 -from typing import Optional

7 -

8 -import torch

9 -

10 -from vllm.sampling_params import SamplingParams, SamplingType

11 -from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData,

12 - SequenceGroupMetadata)

13 -from vllm.utils import (PyObjectCache, async_tensor_h2d,

14 - is_pin_memory_available, make_tensor_with_pad)

15 -

16 -_SAMPLING_EPS = 1e-5

17 -

18 -

19 -@dataclass

20 -class SequenceGroupToSample:

21 - # |---------- N-1 iteration --------|

22 - # |---------------- N iteration ---------------------|

23 - # |- tokenA -|......................|-- newTokens ---|

24 - # |---------- context_len ----------|

25 - # |-------------------- seq_len ----------------------|

26 - # |-- query_len ---|

27 -

28 - # Sequence ids for the sequence group in a previous step.

29 - seq_ids: list[int]

30 - sampling_params: SamplingParams

31 - # seq_id -> sequence data.

32 - seq_data: dict[int, SequenceData]

33 - # The length of the sequence (all tokens seen in the past + new token to

34 - # compute attention) of the sequence group. None if it is in a decode

35 - # stage.

36 - seq_len: Optional[int]

37 - # The length of new query tokens to compute in the current step. None if it

38 - # is in a decode stage. The length of query_len <= seq_len if chunked

39 - # prefill is enabled.

40 - query_len: Optional[int]

41 - # A random number generator for sampling.

42 - generator: Optional[torch.Generator]

43 - # True if the sequence group is in prefill stage. False if it is in a

44 - # decode stage.

45 - is_prompt: bool

46 - # Query token indices from logits. to compute prompt logprob. Empty if

47 - # prompt logprob is not required.

48 - prompt_logprob_indices: list[int]

49 - # Sample token indices from logits. Empty if sampling is not required.

50 - sample_indices: list[int]

51 -

52 - @property

53 - def do_sample(self):

54 - return len(self.sample_indices) > 0

55 -

56 - def __post_init__(self):

57 - if len(self.prompt_logprob_indices) > 0:

58 - assert self.sampling_params.prompt_logprobs is not None

59 - if self.is_prompt:

60 - assert self.seq_len is not None

61 - assert self.query_len is not None

62 -

63 -

64 -def gen_seq_group_to_sample_builder(num_seqs: int):

65 - return lambda: SequenceGroupToSample(

66 - seq_ids=[0] * num_seqs,

67 - sampling_params=None,

68 - seq_data=None, # type: ignore

69 - seq_len=0,

70 - query_len=0,

71 - generator=None,

72 - is_prompt=True,

73 - prompt_logprob_indices=[],

74 - sample_indices=[],

75 - )

76 -

77 -

78 -class SamplingMetadataCache:

79 - """Used to cache SamplingMetadata objects between scheduler iterations"""

80 -

81 - def __init__(self):

82 - self._seq_group_to_sample_cache: dict[int, PyObjectCache] = {}

83 -

84 - def get_cached_seq_group_to_sample(self, num_seqs):

85 - if num_seqs not in self._seq_group_to_sample_cache:

86 - self._seq_group_to_sample_cache[num_seqs] = PyObjectCache(

87 - gen_seq_group_to_sample_builder(num_seqs))

88 -

89 - obj = self._seq_group_to_sample_cache[num_seqs].get_object()

90 - return obj

91 -

92 - def reset(self):

93 - for cache in self._seq_group_to_sample_cache.values():

94 - cache.reset()

95 -

5 class SamplingMetadata:

98 - """Metadata for input sequences. Used in sampler.

99 -

100 - The usage is as follows;

101 - ```

102 - hidden_states = execute_model(...)

103 - logits = hidden_states[sampling_metadata.selected_token_indices]

104 - sample(logits)

105 -

106 - def sample(logits):

107 - # Use categorized_sample_indices for sampling....

108 - ```

109 -

110 - Args:

111 - seq_groups: List of batched sequence groups.

112 - selected_token_indices: (num_query_tokens_to_logprob). Indices to find

113 - logits from the initial model output hidden states.

114 - categorized_sample_indices: SamplingType -> token indices to sample.

115 - Each token indices is 2D tensor of (num_indices, num_indices) where

116 - the first item means the sample index within the returned logit

117 - (before pruning padding), and the second item means the sample

118 - index after pruning using selected_token_indices.

119 - For example, if the returned logit is [1, 2, 3], and we select

120 - [1, 2] for sampling, the pruned logit will be [2, 3]. In this case,

121 - The first tuple is [1, 2] (sampled index within original logit),

122 - and the second tuple is [0, 1] (sampled index within pruned logit).

123 - num_prompts: Number of prompt sequence groups in seq_groups.

124 - skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU

125 - serialization of token outputs.

126 - reuse_sampling_tensors: Indicates if we want to reuse sampling

127 - tensors that are part of the sampler forward pass. Currently,

128 - it is mainly used for multi-step decode.

129 -

130 - """

131 -

132 - def __init__(

133 - self,

134 - seq_groups: list[SequenceGroupToSample],

135 - selected_token_indices: torch.Tensor,

136 - categorized_sample_indices: dict[SamplingType, torch.Tensor],

137 - num_prompts: int,

138 - skip_sampler_cpu_output: bool = False,

139 - reuse_sampling_tensors: bool = False,

140 - ) -> None:

141 - self.seq_groups = seq_groups

142 - self.selected_token_indices = selected_token_indices

143 - self.categorized_sample_indices = categorized_sample_indices

144 - self.num_prompts = num_prompts

145 - self.skip_sampler_cpu_output = skip_sampler_cpu_output

146 - self.reuse_sampling_tensors = reuse_sampling_tensors

147 -

148 - @staticmethod

149 - def prepare(

150 - seq_group_metadata_list: list[SequenceGroupMetadata],

151 - seq_lens: list[int],

152 - query_lens: list[int],

153 - device: str,

154 - pin_memory: bool,

155 - generators: Optional[dict[str, torch.Generator]] = None,

156 - cache: Optional[SamplingMetadataCache] = None,

157 - ) -> "SamplingMetadata":

158 - (

159 - seq_groups,

160 - selected_token_indices,

161 - categorized_sample_indices,

162 - num_prompts,

163 - ) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens,

164 - device, generators, cache)

165 - selected_token_indices = async_tensor_h2d(

166 - selected_token_indices,

167 - dtype=torch.long,

168 - target_device=device,

169 - pin_memory=pin_memory,

170 - )

171 - categorized_sample_indices = {

172 - t:

173 - async_tensor_h2d(

174 - seq_ids,

175 - dtype=torch.int,

176 - target_device=device,

177 - pin_memory=pin_memory,

178 - )

179 - for t, seq_ids in categorized_sample_indices.items()

180 - }

181 -

182 - sampling_metadata = SamplingMetadata(

183 - seq_groups=seq_groups,

184 - selected_token_indices=selected_token_indices,

185 - categorized_sample_indices=categorized_sample_indices,

186 - num_prompts=num_prompts,

187 - )

188 - return sampling_metadata

189 -

190 - def __repr__(self) -> str:

191 - return (

192 - "SamplingMetadata("

193 - f"seq_groups={self.seq_groups}, "

194 - f"selected_token_indices={self.selected_token_indices}, "

195 - f"categorized_sample_indices={self.categorized_sample_indices})")

196 -

197 -

198 -def _prepare_seq_groups(

199 - seq_group_metadata_list: list[SequenceGroupMetadata],

200 - seq_lens: list[int],

201 - query_lens: list[int],

202 - device: str,

203 - generators: Optional[dict[str, torch.Generator]] = None,

204 - cache: Optional[SamplingMetadataCache] = None,

205 -) -> tuple[

206 - list[SequenceGroupToSample],

207 - list[int],

208 - dict[SamplingType, list[int]],

209 - int,

210 -]:

211 - """Prepare sequence groups and indices for sampling.

212 -

213 - Args:

214 - seq_group_metadata_list: A list of sequence group to batch.

215 - seq_lens: A list of sequence lens per sequence group.

216 - Index of prompt len should match with seq_group_metadata_list.

217 - query_lens: A list of query lengths. Prompt lens include the length

218 - of entire prompt tokens, and it could be shorter.

219 - device: A device to use for random number generators,

220 - `SequenceGroupToSample.generator`.

221 - generators: A store of per-request random number generators used

222 - for seeded requests.

223 -

224 - Returns:

225 - seq_groups: A list of sequence group to sample.

226 - selected_token_indices: See the definition from `SamplingMetadata`.

227 - categorized_sample_indices: See the definition from `SamplingMetadata`.

228 - num_prompts: Total number of prompts from `seq_group_metadata_list`.

229 - """

230 - # Batched sequence groups for the current model forward stsep.

231 - seq_groups: list[SequenceGroupToSample] = []

232 - # A list of token indices to sample/compute logprob. It is used to

233 - # prune the outcome logits from the model for the performance.

234 - selected_token_indices: list[int] = []

235 - # Used for selected_token_indices.

236 - model_output_idx = 0

237 -

238 - # Sampling type -> (

239 - # indices to sample/prompt logprob within pruned output logits,

240 - # indices to sample within pruned logits)

241 - categorized_sample_indices: dict[SamplingType, list[int]] = {

242 - t: []

243 - for t in SamplingType

244 - }

245 - # Index of logits to compute logprob. Logits include both prompt logprob

246 - # and sample logprob indices.

247 - logit_idx = 0

248 - # Total number of prompts from given sequence groups.

249 - num_prompts = 0

250 -

251 - for i, seq_group_metadata in enumerate(seq_group_metadata_list):

252 - seq_ids = seq_group_metadata.seq_data.keys()

253 -

254 - if cache is not None:

255 - sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids))

256 -

257 - for j, seq_id in enumerate(seq_ids):

258 - sample_obj.seq_ids[j] = seq_id

259 -

260 - sample_obj.prompt_logprob_indices.clear()

261 - sample_obj.sample_indices.clear()

262 -

263 - sampling_params = seq_group_metadata.sampling_params

264 - is_prompt = seq_group_metadata.is_prompt

265 - generator: Optional[torch.Generator] = None

266 - # If the current seq group is in decode stage, it is None.

267 - seq_len: Optional[int] = None

268 - query_len: Optional[int] = None

269 - prompt_logprob_indices: list[int] = (sample_obj.prompt_logprob_indices

270 - if cache is not None else [])

271 - sample_indices: list[int] = (sample_obj.sample_indices

272 - if cache is not None else [])

273 - do_sample = seq_group_metadata.do_sample

274 -

275 - if seq_group_metadata.is_prompt:

276 - if sampling_params.seed is not None:

277 - generator = torch.Generator(device=device).manual_seed(

278 - sampling_params.seed)

279 - if generators is not None:

280 - generators[seq_group_metadata.request_id] = generator

281 -

282 - num_prompts += 1

283 - num_prefill_sample = len(seq_ids)

284 - assert num_prefill_sample == 1

285 - assert query_lens is not None and seq_lens is not None

286 - query_len, seq_len = query_lens[i], seq_lens[i]

287 - # If we need sampling, exclude num_prefill_sample tokens from

288 - # prompt logprob.

289 - prompt_logprob_len = (query_len - num_prefill_sample

290 - if do_sample else query_len)

291 - sample_len = num_prefill_sample if do_sample else 0

292 - else:

293 - # Decode

294 - prompt_logprob_len = 0

295 - query_len = query_lens[i] if query_lens is not None and len(

296 - query_lens) > 0 else 1

297 - sample_len = len(seq_ids) * query_len if do_sample else 0

298 -

299 - if sampling_params.seed is not None and generators is not None:

300 - generator = generators.get(seq_group_metadata.request_id)

301 -

302 - # Update indices to select from the model output.

303 - """

304 - This blocks computes selected_token_indices which is used in the

305 - following way.

306 -

307 - hidden_states = model(...)

308 - logits = hidden_states[selected_token_indices]

309 - """

310 -

311 - if sampling_params.prompt_logprobs is not None:

312 - selected_token_indices.extend(

313 - range(model_output_idx, model_output_idx + prompt_logprob_len))

314 - model_output_idx += prompt_logprob_len

315 - if do_sample:

316 - selected_token_indices.extend(

317 - range(model_output_idx, model_output_idx + sample_len))

318 - model_output_idx += sample_len

319 -

320 - # We now find indices for logprob computation and sampling.

321 - """

322 - This block computes categorized_sample_indices which is used in the

323 - following way.

324 -

325 - hidden_states = model(...)

326 - logits = hidden_states[selected_token_indices]

327 - def sample(logits):

328 - # Use categorized_sample_indices for sampling.

329 - # prompt_logprob_indices to find prompt logprob indices.

330 - # sample_indices to find sample indices.

331 - """

332 -

333 - if sampling_params.prompt_logprobs is not None:

334 - prompt_logprob_indices.extend(

335 - range(logit_idx, logit_idx + prompt_logprob_len))

336 - logit_idx += prompt_logprob_len

337 - if do_sample:

338 - sample_indices.extend(range(logit_idx, logit_idx + sample_len))

339 - categorized_sample_indices[sampling_params.sampling_type].extend(

340 - list(range(logit_idx, logit_idx + sample_len)))

341 - logit_idx += sample_len

342 -

343 - if cache is not None:

344 - sample_obj.sampling_params = sampling_params

345 - sample_obj.seq_data = seq_group_metadata.seq_data

346 - sample_obj.seq_len = seq_len

347 - sample_obj.query_len = query_len

348 - sample_obj.generator = generator

349 - sample_obj.is_prompt = is_prompt

350 - else:

351 - sample_obj = SequenceGroupToSample(

352 - seq_ids=list(seq_ids),

353 - sampling_params=sampling_params,

354 - seq_data=seq_group_metadata.seq_data,

355 - seq_len=seq_len,

356 - query_len=query_len,

357 - generator=generator,

358 - is_prompt=is_prompt,

359 - prompt_logprob_indices=list(prompt_logprob_indices),

360 - sample_indices=list(sample_indices),

361 - )

362 -

363 - seq_groups.append(sample_obj)

364 -

365 - if cache is not None:

366 - cache.reset()

367 -

368 - return (seq_groups, selected_token_indices, categorized_sample_indices,

369 - num_prompts)

370 -

371 -

372 -@dataclass

373 -class SamplingTensors:

374 - """Tensors for sampling."""

375 -

376 - temperatures: torch.Tensor

377 - top_ps: torch.Tensor

378 - top_ks: torch.Tensor

379 - min_ps: torch.Tensor

380 - presence_penalties: torch.Tensor

381 - frequency_penalties: torch.Tensor

382 - repetition_penalties: torch.Tensor

383 - prompt_tokens: torch.Tensor

384 - output_tokens: torch.Tensor

385 -

386 - @classmethod

387 - def from_sampling_metadata(

388 - cls,

389 - sampling_metadata: "SamplingMetadata",

390 - vocab_size: int,

391 - device: torch.device,

392 - dtype: torch.dtype,

393 - ) -> tuple["SamplingTensors", bool, bool, bool]:

394 - prompt_tokens: list[array] = []

395 - output_tokens: list[array] = []

396 - top_ks: list[int] = []

397 - temperatures: list[float] = []

398 - top_ps: list[float] = []

399 - min_ps: list[float] = []

400 - presence_penalties: list[float] = []

401 - frequency_penalties: list[float] = []

402 - repetition_penalties: list[float] = []

403 - do_penalties = False

404 - do_top_p_top_k = False

405 - do_min_p = False

406 -

407 - assert sampling_metadata.seq_groups is not None

408 - for seq_group in sampling_metadata.seq_groups:

409 - seq_ids = seq_group.seq_ids

410 - sampling_params = seq_group.sampling_params

411 - temperature = sampling_params.temperature

412 - p = sampling_params.presence_penalty

413 - f = sampling_params.frequency_penalty

414 - r = sampling_params.repetition_penalty

415 - top_p = sampling_params.top_p

416 - min_p = sampling_params.min_p

417 -

418 - # k should not be greater than the vocab size.

419 - top_k = min(sampling_params.top_k, vocab_size)

420 - top_k = vocab_size if top_k < 1 else top_k

421 - if temperature < _SAMPLING_EPS:

422 - # NOTE: Zero temperature means deterministic sampling

423 - # (i.e., greedy sampling or beam search).

424 - # Set the temperature to 1 to avoid division by zero.

425 - temperature = 1.0

426 - if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS

427 - or top_k != vocab_size):

428 - do_top_p_top_k = True

429 - if not do_min_p and min_p > _SAMPLING_EPS:

430 - do_min_p = True

431 - if not do_penalties and (abs(p) >= _SAMPLING_EPS

432 - or abs(f) >= _SAMPLING_EPS

433 - or abs(r - 1.0) >= _SAMPLING_EPS):

434 - do_penalties = True

435 -

436 - is_prompt = seq_group.is_prompt

437 - if is_prompt and sampling_params.prompt_logprobs is not None:

438 - # For tokens in the prompt that we only need to get

439 - # their logprobs

440 - query_len = seq_group.query_len

441 - assert query_len is not None

442 - prefill_len = len(seq_group.prompt_logprob_indices)

443 - temperatures += [temperature] * prefill_len

444 - top_ps += [top_p] * prefill_len

445 - top_ks += [top_k] * prefill_len

446 - min_ps += [min_p] * prefill_len

447 - presence_penalties += [0] * prefill_len

448 - frequency_penalties += [0] * prefill_len

449 - repetition_penalties += [1] * prefill_len

450 -

451 - if seq_group.do_sample:

452 - sample_lens = len(seq_group.sample_indices)

453 - assert sample_lens >= len(seq_ids)

454 - temperatures += [temperature] * sample_lens

455 - top_ps += [top_p] * sample_lens

456 - top_ks += [top_k] * sample_lens

457 - min_ps += [min_p] * sample_lens

458 - presence_penalties += [p] * sample_lens

459 - frequency_penalties += [f] * sample_lens

460 - repetition_penalties += [r] * sample_lens

461 -

462 - if do_penalties:

463 - for seq_group in sampling_metadata.seq_groups:

464 - seq_ids = seq_group.seq_ids

465 - sampling_params = seq_group.sampling_params

466 - if (seq_group.is_prompt

467 - and sampling_params.prompt_logprobs is not None):

468 - prefill_len = len(seq_group.prompt_logprob_indices)

469 - prompt_tokens.extend(

470 - array(VLLM_TOKEN_ID_ARRAY_TYPE)

471 - for _ in range(prefill_len))

472 - output_tokens.extend(

473 - array(VLLM_TOKEN_ID_ARRAY_TYPE)

474 - for _ in range(prefill_len))

475 - if seq_group.do_sample:

476 - for seq_id in seq_ids:

477 - seq_data = seq_group.seq_data[seq_id]

478 - prompt_tokens.append(seq_data.prompt_token_ids_array)

479 - output_tokens.append(seq_data.output_token_ids_array)

480 -

481 - sampling_tensors = SamplingTensors.from_lists(

482 - temperatures,

483 - top_ps,

484 - top_ks,

485 - min_ps,

486 - presence_penalties,

487 - frequency_penalties,

488 - repetition_penalties,

489 - prompt_tokens,

490 - output_tokens,

491 - vocab_size,

492 - device,

493 - dtype,

494 - )

495 - return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p)

496 -

497 - @classmethod

498 - def from_lists(

499 - cls,

500 - temperatures: list[float],

501 - top_ps: list[float],

502 - top_ks: list[int],

503 - min_ps: list[float],

504 - presence_penalties: list[float],

505 - frequency_penalties: list[float],

506 - repetition_penalties: list[float],

507 - prompt_tokens: list[array],

508 - output_tokens: list[array],

509 - vocab_size: int,

510 - device: torch.device,

511 - dtype: torch.dtype,

512 - ) -> "SamplingTensors":

513 - # Note that the performance will be very bad without

514 - # pinned memory.

515 - pin_memory = is_pin_memory_available()

516 -

517 - do_penalties = prompt_tokens or output_tokens

518 -

519 - if do_penalties:

520 - prompt_t = make_tensor_with_pad(

521 - prompt_tokens,

522 - vocab_size,

523 - device="cpu",

524 - dtype=torch.int64,

525 - pin_memory=pin_memory,

526 - )

527 - output_t = make_tensor_with_pad(

528 - output_tokens,

529 - vocab_size,

530 - device="cpu",

531 - dtype=torch.int64,

532 - pin_memory=pin_memory,

533 - )

534 - else:

535 - empty_tensor = torch.empty(0, device=device, dtype=torch.long)

536 - prompt_t = empty_tensor

537 - output_t = empty_tensor

538 -

539 - temperatures_t = torch.tensor(

540 - temperatures,

541 - device="cpu",

542 - dtype=dtype,

543 - pin_memory=pin_memory,

544 - )

545 - top_ps_t = torch.tensor(

546 - top_ps,

547 - device="cpu",

548 - dtype=dtype,

549 - pin_memory=pin_memory,

550 - )

551 - min_ps_t = torch.tensor(

552 - min_ps,

553 - device="cpu",

554 - dtype=dtype,

555 - pin_memory=pin_memory,

556 - )

557 - presence_penalties_t = torch.tensor(

558 - presence_penalties,

559 - device="cpu",

560 - dtype=dtype,

561 - pin_memory=pin_memory,

562 - )

563 - frequency_penalties_t = torch.tensor(

564 - frequency_penalties,

565 - device="cpu",

566 - dtype=dtype,

567 - pin_memory=pin_memory,

568 - )

569 - repetition_penalties_t = torch.tensor(

570 - repetition_penalties,

571 - device="cpu",

572 - dtype=dtype,

573 - pin_memory=pin_memory,

574 - )

575 - top_ks_t = torch.tensor(

576 - top_ks,

577 - device="cpu",

578 - dtype=torch.int,

579 - pin_memory=pin_memory,

580 - )

581 - # Because the memory is pinned, we can do non-blocking

582 - # transfer to device.

583 -

584 - return cls(

585 - temperatures=temperatures_t.to(device=device, non_blocking=True),

586 - top_ps=top_ps_t.to(device=device, non_blocking=True),

587 - top_ks=top_ks_t.to(device=device, non_blocking=True),

588 - min_ps=min_ps_t.to(device=device, non_blocking=True),

589 - presence_penalties=presence_penalties_t.to(device=device,

590 - non_blocking=True),

591 - frequency_penalties=frequency_penalties_t.to(device=device,

592 - non_blocking=True),

593 - repetition_penalties=repetition_penalties_t.to(device=device,

594 - non_blocking=True),

595 - prompt_tokens=prompt_t.to(device=device, non_blocking=True),

596 - output_tokens=output_t.to(device=device, non_blocking=True),

597 - )

6 + # Placeholder until it can be safely removed.

7 + pass

vllm/sequence.py

+4 -1318

@@ -1,28 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Sequence and its related classes."""

4 -import copy

5 -import enum

6 -from abc import ABC, abstractmethod

7 -from array import array

8 -from collections import defaultdict

9 -from collections.abc import Mapping

10 -from collections.abc import Sequence as GenericSequence

11 -from dataclasses import dataclass, field

12 -from functools import reduce

13 -from typing import TYPE_CHECKING, Any, Callable, Optional, Union

4 +from dataclasses import dataclass

5 +from typing import TYPE_CHECKING, Any, Optional, Union

7 import msgspec

8 import torch

18 -from vllm.inputs import SingletonInputs

19 -from vllm.logprobs import Logprob, PromptLogprobs, SampleLogprobs

20 -from vllm.multimodal import MultiModalKwargs, MultiModalPlaceholderDict

21 -from vllm.pooling_params import PoolingParams

22 -from vllm.sampling_params import RequestOutputKind, SamplingParams

23 -

10 if TYPE_CHECKING:

25 - from vllm.lora.request import LoRARequest

11 from vllm.v1.worker.kv_connector_model_runner_mixin import (

12 KVConnectorOutput)

13 else:

@@ -34,50 +19,6 @@ VLLM_TOKEN_ID_ARRAY_TYPE = "l"

19 VLLM_INVALID_TOKEN_ID = -1

37 -def array_full(token_id: int, count: int):

38 - """[`array`][] equivalent of [numpy.full][]."""

39 - return array(VLLM_TOKEN_ID_ARRAY_TYPE, [token_id]) * count

40 -

41 -

42 -class SequenceStatus(enum.IntEnum):

43 - """Status of a sequence."""

44 - WAITING = 0

45 - RUNNING = 1

46 - SWAPPED = 2

47 - # Note: anything after SWAPPED (2) will be considered

48 - # as a finished status.

49 - FINISHED_STOPPED = 3

50 - FINISHED_LENGTH_CAPPED = 4

51 - FINISHED_ABORTED = 5

52 - FINISHED_IGNORED = 6

53 -

54 - @staticmethod

55 - def is_finished(status: "SequenceStatus") -> bool:

56 - return status > SequenceStatus.SWAPPED

57 -

58 - @staticmethod

59 - def get_finished_reason(status: "SequenceStatus") -> Union[str, None]:

60 - if status == SequenceStatus.FINISHED_STOPPED:

61 - finish_reason = "stop"

62 - elif status == SequenceStatus.FINISHED_LENGTH_CAPPED:

63 - finish_reason = "length"

64 - elif status == SequenceStatus.FINISHED_ABORTED:

65 - finish_reason = "abort"

66 - elif status == SequenceStatus.FINISHED_IGNORED:

67 - # The ignored sequences are the sequences whose prompt lengths

68 - # are longer than the model's length cap. Therefore, the stop

69 - # reason should also be "length" as in OpenAI API.

70 - finish_reason = "length"

71 - else:

72 - finish_reason = None

73 - return finish_reason

74 -

75 -

76 -class SequenceStage(enum.Enum):

77 - PREFILL = enum.auto()

78 - DECODE = enum.auto()

79 -

80 -

22 @dataclass

23 class RequestMetrics:

24 """Metrics associated with a request.

@@ -107,971 +48,12 @@ class RequestMetrics:

48 model_execute_time: Optional[float] = None

110 -class SequenceDataDelta(

111 - msgspec.Struct,

112 - array_like=True, # type: ignore[call-arg]

113 - omit_defaults=True): # type: ignore[call-arg]

114 - """Delta SequenceData to send to workers per step."""

115 - # A new token to be appended to existing SequenceData.

116 - new_output_token_ids: list[int]

117 - # Overwriting existing `cumulative_logprob`

118 - new_cumulative_logprob: float

119 - # Overwriting existing `num_computed_tokens`.

120 - new_num_computed_tokens: int

121 - # Overwriting existing `stage`.

122 - new_stage: SequenceStage

123 -

124 -

125 -class SequenceData(msgspec.Struct,

126 - omit_defaults=True): # type: ignore[call-arg]

127 - """Data associated with a sequence."""

128 - # NOTE: we cannot use Union[list, array] because msgspec cannot support

129 - # union of 2 list types.

130 - _prompt_token_ids: array

131 - _output_token_ids: array = msgspec.field(

132 - default_factory=lambda: array(VLLM_TOKEN_ID_ARRAY_TYPE, []))

133 -

134 - _prompt_embeds: Optional[torch.Tensor] = None

135 - _output_embeds: Optional[torch.Tensor] = None

136 -

137 - ### The below fields should not be passed as an argument ###

138 - _cumulative_logprob: float = 0.0

139 - _prompt_token_ids_tuple: tuple[int,

140 - ...] = msgspec.field(default_factory=tuple)

141 - # The number of tokens that are computed (that run against the model).

142 - _num_computed_tokens: int = 0

143 - # The number of tokens with prefix cache hit.

144 - _num_cached_tokens: int = 0

145 - _stage: SequenceStage = SequenceStage.PREFILL

146 - _cached_all_token_ids: list[int] = msgspec.field(default_factory=list)

147 - _cached_all_token_embeds: Optional[torch.Tensor] = None

148 -

149 - # It is used to get delta input. It is reset when `get_delta_and_reset`

150 - # is called.

151 - _new_appended_tokens: list[int] = msgspec.field(default_factory=list)

152 -

153 - # It is used to compute mrope_position_ids.

154 - _mrope_position_delta: Optional[int] = None

155 -

156 - @staticmethod

157 - def from_prompt_token_counts(

158 - *token_counts: tuple[int, int]) -> "SequenceData":

159 - """

160 - Construct a [`SequenceData`][vllm.sequence.SequenceData] instance

161 - by concatenating prompt token sequences.

162 -

163 - Each tuple represents one token sequence, expressed in the form

164 - `(token_id, count)`.

165 - """

166 - if len(token_counts) == 0:

167 - return SequenceData.from_seqs([])

168 -

169 - prompt_token_ids_arr = reduce(

170 - array.__iadd__,

171 - (array_full(token_id, count) for token_id, count in token_counts),

172 - )

173 -

174 - return SequenceData(prompt_token_ids_arr)

175 -

176 - @staticmethod

177 - def from_seqs(

178 - prompt_token_ids: GenericSequence[int],

179 - output_token_ids: Optional[GenericSequence[int]] = None,

180 - *,

181 - prompt_embeds: Optional[torch.Tensor] = None,

182 - ) -> "SequenceData":

183 - """

184 - Construct a [`SequenceData`][vllm.sequence.SequenceData] instance

185 - from prompt and output token sequences.

186 - """

187 - prompt_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE,

188 - prompt_token_ids)

189 -

190 - if output_token_ids is None:

191 - return SequenceData(prompt_token_ids_arr,

192 - _prompt_embeds=prompt_embeds)

193 -

194 - output_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE,

195 - output_token_ids)

196 -

197 - return SequenceData(prompt_token_ids_arr,

198 - _output_token_ids=output_token_ids_arr,

199 - _prompt_embeds=prompt_embeds)

200 -

201 - def __post_init__(self) -> None:

202 - assert self._prompt_token_ids.typecode == "l"

203 - assert self._output_token_ids.typecode == "l"

204 - self._prompt_token_ids_tuple: tuple[int, ...] = tuple(

205 - self._prompt_token_ids)

206 - self._update_cached_all_tokens()

207 - if self._prompt_embeds is not None:

208 - self._update_cached_all_token_embeds()

209 -

210 - def _update_cached_all_tokens(self):

211 - assert isinstance(self._prompt_token_ids, array)

212 - assert isinstance(self._output_token_ids, array)

213 - self._cached_all_token_ids: list[int] = list(self._prompt_token_ids +

214 - self._output_token_ids)

215 -

216 - def _update_cached_all_token_embeds(self):

217 - assert isinstance(self._prompt_embeds, torch.Tensor)

218 - self._cached_all_token_embeds: torch.Tensor = self._prompt_embeds

219 - if self._output_embeds is not None:

220 - self._cached_all_token_embeds = torch.cat(

221 - (self._cached_all_token_embeds, self._output_embeds), dim=0)

222 -

223 - @property

224 - def cumulative_logprob(self) -> float:

225 - """The cumulative log probability of the output."""

226 - return self._cumulative_logprob

227 -

228 - @property

229 - def prompt_token_ids(self) -> tuple[int, ...]:

230 - """The token IDs of the prompt."""

231 - return self._prompt_token_ids_tuple

232 -

233 - @prompt_token_ids.setter

234 - def prompt_token_ids(self, new_prompt_token_ids) -> None:

235 - raise NotImplementedError

236 -

237 - @property

238 - def prompt_token_ids_array(self) -> array:

239 - """Return the prompt token ids in array type.

240 -

241 - Note that the array is in "I" type, and it is not compatible

242 - with torch.long (2 bytes vs 4 bytes). So beware of the usage.

243 - """

244 - return self._prompt_token_ids

245 -

246 - @property

247 - def output_token_ids(self) -> tuple[int, ...]:

248 - """The token IDs of the output."""

249 - return tuple(self._output_token_ids)

250 -

251 - @output_token_ids.setter

252 - def output_token_ids(self,

253 - new_output_token_ids: GenericSequence[int]) -> None:

254 - self._output_token_ids = array(VLLM_TOKEN_ID_ARRAY_TYPE,

255 - new_output_token_ids)

256 - self._update_cached_all_tokens()

257 -

258 - @property

259 - def output_embeds(self) -> Optional[torch.Tensor]:

260 - return self._output_embeds

261 -

262 - @output_embeds.setter

263 - def output_embeds(self, new_output_token_embeds: torch.Tensor) -> None:

264 - self._output_token_embeds = new_output_token_embeds

265 - self._update_cached_all_token_embeds()

266 -

267 - @property

268 - def output_token_ids_array(self) -> array:

269 - """Return the prompt token ids in array type.

270 -

271 - Note that the array is in "I" type, and it is not compatible

272 - with torch.long (2 bytes vs 4 bytes). So beware of the usage.

273 - """

274 - assert isinstance(self._output_token_ids, array)

275 - return self._output_token_ids

276 -

277 - @property

278 - def prompt_embeds(self) -> Optional[torch.Tensor]:

279 - return self._prompt_embeds

280 -

281 - @prompt_embeds.setter

282 - def prompt_embeds(self, prompt_embeds: torch.Tensor) -> None:

283 - self._prompt_embeds = prompt_embeds

284 - self._update_cached_all_token_embeds()

285 -

286 - @property

287 - def mrope_position_delta(self) -> Optional[int]:

288 - return self._mrope_position_delta

289 -

290 - @mrope_position_delta.setter

291 - def mrope_position_delta(self, new_mrope_position_delta):

292 - self._mrope_position_delta = new_mrope_position_delta

293 -

294 - def append_token_id(self,

295 - token_id: int,

296 - logprob: float,

297 - token_embed: Optional[torch.Tensor] = None) -> None:

298 - self._output_token_ids.append(token_id)

299 - self._new_appended_tokens.append(token_id)

300 - self._cached_all_token_ids.append(token_id)

301 - self._cumulative_logprob += logprob

302 - if token_embed is not None:

303 - # Do not pass in with batch or sequence dimensions

304 - assert token_embed.ndim == 1

305 - token_embed = token_embed.detach().cpu().unsqueeze(0)

306 - if self._output_embeds is None:

307 - self._output_embeds = token_embed

308 - else:

309 - self._output_embeds = torch.cat(

310 - (self._output_embeds, token_embed), dim=0)

311 - assert self._cached_all_token_embeds is not None

312 - self._cached_all_token_embeds = torch.cat(

313 - (self._cached_all_token_embeds,

314 - token_embed.to(device=self._cached_all_token_embeds.device)),

315 - dim=0)

316 -

317 - def get_len(self) -> int:

318 - return len(self._output_token_ids) + len(self._prompt_token_ids)

319 -

320 - def get_prompt_len(self) -> int:

321 - return len(self._prompt_token_ids)

322 -

323 - def get_output_len(self) -> int:

324 - return len(self._output_token_ids)

325 -

326 - def get_token_ids(self) -> list[int]:

327 - return self._cached_all_token_ids

328 -

329 - def get_token_embeddings(self) -> Optional[torch.Tensor]:

330 - return self._cached_all_token_embeds

331 -

332 - def get_prefix_token_ids(

333 - self, num_tokens: int

334 - ) -> tuple[tuple[int, ...], Optional[tuple[int, ...]]]:

335 - """Get prefix tokens, and make the return value hashable"""

336 - prompt_length = self.get_prompt_len()

337 - if num_tokens > prompt_length:

338 - return (self._prompt_token_ids_tuple,

339 - tuple(self._output_token_ids[:num_tokens - prompt_length]))

340 - else:

341 - return (self._prompt_token_ids_tuple[:num_tokens], None)

342 -

343 - def get_num_computed_tokens(self) -> int:

344 - """Return the number of prefill tokens that are already computed."""

345 - return self._num_computed_tokens

346 -

347 - def update_num_computed_tokens(self, num_new_computed_tokens: int):

348 - """Update number of tokens computed so far."""

349 - self._num_computed_tokens += num_new_computed_tokens

350 - assert self._num_computed_tokens <= self.get_len(), (

351 - self._num_computed_tokens, self.get_len())

352 - # If all tokens are computed, it means it is in decoding phase.

353 - if self.get_num_uncomputed_tokens() == 0:

354 - self._stage = SequenceStage.DECODE

355 -

356 - def get_num_cached_tokens(self) -> int:

357 - """Return the number of tokens with prefix cache hit."""

358 - return self._num_cached_tokens

359 -

360 - def update_num_cached_tokens(self, num_cached_tokens: int):

361 - """Update the number of tokens with prefix cache hit."""

362 - self._num_cached_tokens = num_cached_tokens

363 -

364 - def reset_state_for_recompute(self) -> None:

365 - """Reset the number of computed tokens from this sequence. It is

366 - supposed to be called when a sequence needs to be started from

367 - the beginning again (e.g., sequence is preempted).

368 - """

369 - self._num_computed_tokens = 0

370 - self._stage = SequenceStage.PREFILL

371 - self._new_appended_tokens = []

372 -

373 - def get_num_uncomputed_tokens(self) -> int:

374 - """Return the number of prefill tokens that are not computed."""

375 - # we use `get_len()` which includes prompt_len + output_len instead

376 - # of prompt_len here. This is because during recompute we need to

377 - # prefill for both prompt and output.

378 - return self.get_len() - self.get_num_computed_tokens()

379 -

380 - def get_last_token_id(self) -> int:

381 - if not self._output_token_ids:

382 - return self._prompt_token_ids[-1]

383 - return self._output_token_ids[-1]

384 -

385 - def get_prompt_token_ids(self) -> tuple[int, ...]:

386 - return self.prompt_token_ids

387 -

388 - def get_output_token_ids(self) -> tuple[int, ...]:

389 - return self.output_token_ids

390 -

391 - def get_delta_and_reset(self) -> SequenceDataDelta:

392 - delta = SequenceDataDelta(self._new_appended_tokens,

393 - self._cumulative_logprob,

394 - self.get_num_computed_tokens(), self.stage)

395 - # Reset delta state.

396 - self._new_appended_tokens = []

397 - return delta

398 -

399 - def apply_delta(self, delta: SequenceDataDelta):

400 - self._num_computed_tokens = delta.new_num_computed_tokens

401 - self._cumulative_logprob = delta.new_cumulative_logprob

402 - self._stage = delta.new_stage

403 - self._output_token_ids.extend(delta.new_output_token_ids)

404 - self._cached_all_token_ids.extend(delta.new_output_token_ids)

405 -

406 - @property

407 - def stage(self) -> SequenceStage:

408 - return self._stage

409 -

410 - def __repr__(self) -> str:

411 - return (f"SequenceData("

412 - f"prompt_token_ids={self._prompt_token_ids}, "

413 - f"prompt_embeds.shape="

414 - f"{getattr(self._prompt_embeds, 'shape', None)}, "

415 - f"output_token_ids={self.output_token_ids}, "

416 - f"cumulative_logprob={self.cumulative_logprob}, "

417 - f"get_num_computed_tokens={self.get_num_computed_tokens()})")

418 -

419 -

420 -class Sequence:

421 - """Stores the data, status, and block information of a sequence.

422 -

423 - The sequence is constructed from the

424 - [`DecoderOnlyInputs`][vllm.inputs.data.DecoderOnlyInputs] (for decoder-only)

425 - or [`EncoderDecoderInputs`][vllm.inputs.data.EncoderDecoderInputs]

426 - (for encoder-decoder) instance passed in through the `inputs`

427 - constructor argument.

428 -

429 - Args:

430 - seq_id: The ID of the sequence.

431 - inputs: The inputs of the sequence.

432 - block_size: The block size of the sequence. Should be the same as the

433 - block size used by the block manager and cache engine.

434 - eos_token_id: The end-of-sequence (EOS) token id recognized by this LLM.

435 - lora_request: LoRA request.

436 - """

437 -

438 - def __init__(

439 - self,

440 - seq_id: int,

441 - inputs: SingletonInputs,

442 - block_size: int,

443 - eos_token_id: Optional[int] = None,

444 - lora_request: Optional[LoRARequest] = None,

445 - ) -> None:

446 - self.seq_id = seq_id

447 - self.inputs = inputs

448 - self.block_size = block_size

449 - self.eos_token_id = eos_token_id

450 - self.lora_request = lora_request

451 -

452 - self.data = SequenceData.from_seqs(

453 - self.prompt_token_ids,

454 - prompt_embeds=self.inputs["prompt_embeds"]

455 - if self.inputs["type"] == "embeds" else None)

456 - self.output_logprobs: SampleLogprobs = []

457 - self.output_text = ""

458 -

459 - self.status = SequenceStatus.WAITING

460 - self.stop_reason: Union[int, str, None] = None

461 -

462 - # These are used to keep track of delta outputs

463 - self._last_output_token_ids_offset: int = 0

464 - self._last_output_text_offset: int = 0

465 -

466 - # Used for incremental detokenization

467 - self.prefix_offset = 0

468 - self.read_offset = 0

469 - # Input + output tokens

470 - self.tokens: Optional[list[str]] = None

471 -

472 - @property

473 - def n_blocks(self) -> int:

474 - return (self.get_len() + self.block_size - 1) // self.block_size

475 -

476 - @property

477 - def prompt(self) -> Optional[str]:

478 - if self.inputs["type"] == "embeds":

479 - return None

480 - return self.inputs.get("prompt")

481 -

482 - @property

483 - def prompt_token_ids(self) -> list[int]:

484 - if self.inputs["type"] == "embeds":

485 - return [0] * len(self.inputs["prompt_embeds"])

486 - return self.inputs["prompt_token_ids"]

487 -

488 - @property

489 - def multi_modal_data(self) -> MultiModalKwargs:

490 - if self.inputs["type"] == "multimodal":

491 - return self.inputs["mm_kwargs"].get_data()

492 -

493 - return MultiModalKwargs()

494 -

495 - @property

496 - def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:

497 - if self.inputs["type"] == "multimodal":

498 - return self.inputs["mm_placeholders"]

499 -

500 - return {}

501 -

502 - @property

503 - def lora_int_id(self) -> int:

504 - return self.lora_request.lora_int_id if self.lora_request else 0

505 -

506 - def get_output_text_to_return(self, buffer_length: int,

507 - delta: bool) -> str:

508 - """If delta is True, only new text since the last call to

509 - this method is returned"""

510 -

511 - # We return the full output text if the sequence is finished.

512 - truncate = buffer_length and not self.is_finished()

513 - if not delta:

514 - return self.output_text[:-buffer_length] if truncate else (

515 - self.output_text)

516 - length = len(self.output_text)

517 - if truncate:

518 - length -= buffer_length

519 - last_offset = self._last_output_text_offset

520 - if last_offset < length:

521 - self._last_output_text_offset = length

522 - return self.output_text[last_offset:length]

523 - return ""

524 -

525 - def get_output_token_ids_to_return(

526 - self, delta: bool) -> Union[GenericSequence[int], int]:

527 - """If delta is True, only new tokens since the last call to

528 - this method are returned"""

529 - if not delta:

530 - return self.get_output_token_ids()

531 -

532 - output_len = self.get_output_len()

533 -

534 - # Get the number of new tokens

535 - num_new_tokens = output_len - self._last_output_token_ids_offset

536 - self._last_output_token_ids_offset = output_len

537 -

538 - # Return new tokens

539 - if num_new_tokens == 1:

540 - # Optimization for single decode token case

541 - # (which is what we have most of the time)

542 - return self.data._cached_all_token_ids[-1]

543 -

544 - if num_new_tokens == 0:

545 - return []

546 -

547 - return self.data._cached_all_token_ids[-num_new_tokens:]

548 -

549 - def hash_of_block(self, logical_idx: int) -> int:

550 - # TODO This can produce incorrect hash when block size > prompt size

551 -

552 - # Compute the number of tokens in the sequence

553 - # TODO: The current hashing function is O(L^2). We should optimize

554 - # this in the future.

555 - num_tokens = self.num_hashed_tokens_of_block(logical_idx)

556 - hashed_tokens = self.data.get_prefix_token_ids(num_tokens)

557 - return hash((hashed_tokens, self.lora_int_id))

558 -

559 - def extra_hash(self) -> Optional[int]:

560 - """

561 - This function computes an extra hash for a sequence, specifically

562 - designed for prefix caching mode. The final sequence hash is determined

563 - by applying token_ids from the sequence's blocks.

564 - """

565 - if self.lora_int_id == 0:

566 - return None

567 -

568 - # NOTE: If there are additional factors influencing the block aside from

569 - # token_ids, include them as input parameters to the hash.

570 - return hash(self.lora_int_id)

571 -

572 - def num_hashed_tokens_of_block(self, logical_idx: int):

573 - return logical_idx * self.block_size + self.block_size

574 -

575 - def reset_state_for_recompute(self):

576 - """Reset the sequence states for recomputation."""

577 - self.data.reset_state_for_recompute()

578 -

579 - def append_token_id(self,

580 - token_id: int,

581 - logprobs: dict[int, Logprob],

582 - token_embed: Optional[torch.Tensor] = None) -> None:

583 - assert token_id in logprobs

584 - self.output_logprobs.append(logprobs)

585 - self.data.append_token_id(token_id, logprobs[token_id].logprob,

586 - token_embed)

587 -

588 - def get_len(self) -> int:

589 - return self.data.get_len()

590 -

591 - def get_prompt_len(self) -> int:

592 - return self.data.get_prompt_len()

593 -

594 - def get_output_len(self) -> int:

595 - return self.data.get_output_len()

596 -

597 - def get_token_ids(self) -> list[int]:

598 - return self.data.get_token_ids()

599 -

600 - def get_prompt_token_ids(self) -> tuple[int, ...]:

601 - return self.data.get_prompt_token_ids()

602 -

603 - def get_last_token_id(self) -> int:

604 - return self.data.get_last_token_id()

605 -

606 - def get_output_token_ids(self) -> tuple[int, ...]:

607 - return self.data.get_output_token_ids()

608 -

609 - def get_cumulative_logprob(self) -> float:

610 - return self.data.cumulative_logprob

611 -

612 - def is_finished(self) -> bool:

613 - return SequenceStatus.is_finished(self.status)

614 -

615 - def fork(self, new_seq_id: int) -> "Sequence":

616 - new_seq = copy.deepcopy(self)

617 - new_seq.seq_id = new_seq_id

618 - return new_seq

619 -

620 - def get_num_new_tokens(self) -> int:

621 - """Get the number of new tokens to be computed.

622 -

623 - Returns:

624 - The new number of tokens to be computed. I.e., 1 for decode, or

625 - the remaining prompt size for prefill.

626 - """

627 - if self.data.stage == SequenceStage.DECODE:

628 - return 1

629 - return self.data.get_num_uncomputed_tokens()

630 -

631 - def get_num_computed_tokens(self) -> int:

632 - return self.data.get_num_computed_tokens()

633 -

634 - def is_prefill(self) -> bool:

635 - return self.data.stage == SequenceStage.PREFILL

636 -

637 - def __repr__(self) -> str:

638 - return (f"Sequence(seq_id={self.seq_id}, "

639 - f"status={self.status.name}, "

640 - f"num_blocks={self.n_blocks})")

641 -

642 -

643 -class SequenceGroupState(msgspec.Struct,

644 - omit_defaults=True): # type: ignore[call-arg]

645 - """Mutable state tied to a specific sequence group"""

646 -

647 - # for multi-step decoding

648 - num_steps: int = 1

649 - current_step: int = 0

650 -

651 - @property

652 - def remaining_steps(self) -> int:

653 - return self.num_steps - self.current_step

654 -

655 -

656 -class SequenceGroup:

657 - """A group of sequences that are generated from the same prompt.

658 -

659 - Args:

660 - request_id: The ID of the request.

661 - seqs: The list of sequences.

662 - sampling_params: The sampling parameters used to generate the outputs.

663 - arrival_time: The arrival time of the request.

664 - lora_request: LoRA request.

665 - pooling_params: The parameters used to generate the pooler

666 - for a pooling model.

667 - pooled_data: The extracted hidden states from a pooling model.

668 - encoder_seq: Optional, the single encoder sequence. Should be None

669 - unless you are working with an encoder/decoder model.

670 - trace_headers: OpenTelemetry trace headers.

671 - priority: User-defined priority of the request.

672 - draft_size: The number of speculative tokens plus one from the target

673 - model; equal to max number of tokens a step can generate

674 - for single-draft speculative decoding but larger than

675 - that for multi-draft SD (currently not supported).

676 - """

677 -

678 - def __init__(self,

679 - request_id: str,

680 - seqs: list[Sequence],

681 - arrival_time: float,

682 - sampling_params: Optional[SamplingParams] = None,

683 - lora_request: Optional[LoRARequest] = None,

684 - pooling_params: Optional[PoolingParams] = None,

685 - pooled_data: Optional[torch.Tensor] = None,

686 - encoder_seq: Optional[Sequence] = None,

687 - trace_headers: Optional[Mapping[str, str]] = None,

688 - priority: int = 0,

689 - draft_size: int = 1) -> None:

690 - self.request_id = request_id

691 - self.seqs = seqs

692 - self.first_seq = seqs[0]

693 - self.arrival_time = arrival_time

694 - self.is_single_seq = len(seqs) == 1

695 - self.seqs_dict = {seq.seq_id: seq for seq in seqs}

696 -

697 - self.sampling_params = sampling_params

698 - self.metrics = RequestMetrics(arrival_time=arrival_time,

699 - last_token_time=arrival_time,

700 - first_scheduled_time=None,

701 - first_token_time=None,

702 - time_in_queue=None)

703 - self.last_token_latency = 0.0

704 - self.lora_request = lora_request

705 - self.prompt_logprobs: Optional[PromptLogprobs] = None

706 - self.state = SequenceGroupState()

707 - self.pooling_params = pooling_params

708 - self.pooled_data = pooled_data

709 - self.encoder_seq = encoder_seq

710 - self.trace_headers = trace_headers

711 - self.priority = priority

712 -

713 - self.cached_request_output = None

714 -

715 - @property

716 - def prompt(self) -> Optional[str]:

717 - return self.first_seq.prompt

718 -

719 - @property

720 - def prompt_token_ids(self) -> list[int]:

721 - return self.first_seq.prompt_token_ids

722 -

723 - @property

724 - def encoder_prompt(self) -> Optional[str]:

725 - # There are either 0 or 1 encoder sequences

726 - # If one is present, its prompt is distinct

727 - # from the decoder's.

728 - return (self.encoder_seq.prompt

729 - if self.encoder_seq is not None else None)

730 -

731 - @property

732 - def encoder_prompt_token_ids(self) -> Optional[list[int]]:

733 - # There are either 0 or 1 encoder sequences

734 - # If one is present, its prompt token ids are

735 - # distinct from the decoder's.

736 - return (self.encoder_seq.prompt_token_ids

737 - if self.encoder_seq is not None else None)

738 -

739 - @property

740 - def multi_modal_data(self) -> MultiModalKwargs:

741 - if self.first_seq.multi_modal_data:

742 - return self.first_seq.multi_modal_data

743 - elif self.encoder_seq is not None:

744 - return self.encoder_seq.multi_modal_data

745 - return MultiModalKwargs()

746 -

747 - @property

748 - def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:

749 - if self.first_seq.multi_modal_data:

750 - return self.first_seq.multi_modal_placeholders

751 - elif self.encoder_seq is not None:

752 - return self.encoder_seq.multi_modal_placeholders

753 - return {}

754 -

755 - @property

756 - def lora_int_id(self) -> int:

757 - return self.lora_request.lora_int_id if self.lora_request else 0

758 -

759 - def set_last_token_time(self, now: float) -> None:

760 - """Sets the last token time for Request level timings."""

761 - # If still in prefill phase, assertion fails.

762 - assert not self.is_prefill(), (

763 - "seq_group.set_last_token_time() should not be called "

764 - "if the seq_group is in prefill phase.")

765 - self.last_token_latency = now - self.metrics.last_token_time

766 - self.metrics.last_token_time = now

767 -

768 - def get_last_token_latency(self) -> float:

769 - """Returns the latency of the last token."""

770 - assert not self.is_prefill(), (

771 - "seq_group.get_last_token_latency() should not be called "

772 - "if the seq_group is in prefill phase.")

773 - return self.last_token_latency

774 -

775 - def maybe_set_first_token_time(self, time: float) -> None:

776 - """Sets the first token time for Request level timings."""

777 - # Note: in a case where a sequence_group is swapped and

778 - # recomputed, the time between iterations is counted

779 - # in TPOT, rather than recalculating TTFT (since from the )

780 - # POV of the user, there is simply a long generation delay.

781 - if (self.metrics.first_token_time is None

782 - and self.first_seq.get_output_len() == 1):

783 - self.metrics.first_token_time = time

784 -

785 - def maybe_set_first_scheduled_time(self, time: float) -> None:

786 - """Sets the first scheduled time and time in queue for Request

787 - level timings."""

788 - if self.metrics.first_scheduled_time is None:

789 - self.metrics.first_scheduled_time = time

790 - self.metrics.time_in_queue = time - self.metrics.arrival_time

791 -

792 - def set_finished_time(self, time: Optional[float]) -> None:

793 - """Sets the finished time for Request level timings."""

794 - self.metrics.finished_time = time

795 -

796 - def get_max_num_running_seqs(self) -> int:

797 - """The maximum number of sequences running in parallel in the remaining

798 - lifetime of the request."""

799 - if self.is_single_seq:

800 - return 0 if self.first_seq.is_finished() else 1

801 - return self.num_seqs() - self.num_finished_seqs()

802 -

803 - def get_seqs(

804 - self,

805 - status: Optional[SequenceStatus] = None,

806 - ) -> list[Sequence]:

807 - if status is None:

808 - return self.seqs

809 -

810 - if self.is_single_seq:

811 - return self.seqs if self.first_seq.status == status else []

812 -

813 - return [seq for seq in self.seqs if seq.status == status]

814 -

815 - def is_encoder_decoder(self) -> bool:

816 - return self.encoder_seq is not None

817 -

818 - def get_encoder_seq(self) -> Optional[Sequence]:

819 - return self.encoder_seq

820 -

821 - def get_finished_seqs(self) -> list[Sequence]:

822 - if self.is_single_seq:

823 - return self.seqs if self.first_seq.is_finished() else []

824 -

825 - return [seq for seq in self.seqs if seq.is_finished()]

826 -

827 - def update_num_computed_tokens(self, num_new_computed_tokens: int):

828 - """Update number of tokens computed so far."""

829 - for seq in self.seqs:

830 - if not seq.is_finished():

831 - seq.data.update_num_computed_tokens(num_new_computed_tokens)

832 -

833 - def get_num_uncomputed_tokens(self) -> int:

834 - num_uncomputed_tokens = 0

835 - for seq in self.seqs:

836 - if not seq.is_finished():

837 - num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()

838 - return num_uncomputed_tokens

839 -

840 - def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:

841 - # Optimization. We don't need to call get_seqs if we don't need to

842 - # filter by states.

843 - if status is None:

844 - return len(self.seqs)

845 -

846 - if self.is_single_seq:

847 - return 1 if self.seqs[0].status == status else 0

848 -

849 - return len(self.get_seqs(status))

850 -

851 - def num_finished_seqs(self) -> int:

852 - if self.is_single_seq:

853 - return 1 if self.seqs[0].is_finished() else 0

854 - return len(self.get_finished_seqs())

855 -

856 - def is_finished(self) -> bool:

857 - if self.is_single_seq:

858 - return self.first_seq.is_finished()

859 - return all(seq.is_finished() for seq in self.seqs)

860 -

861 - def is_prefill(self) -> bool:

862 - return self.first_seq.is_prefill()

863 -

864 - def __repr__(self) -> str:

865 - return (f"SequenceGroup(request_id={self.request_id}, "

866 - f"sampling_params={self.sampling_params}, "

867 - f"num_seqs={len(self.seqs)})")

868 -

869 - def uses_prompt_embeds(self) -> bool:

870 - """Returns True if the sequence group uses input embeds."""

871 - return any(seq.data.prompt_embeds is not None for seq in self.seqs)

872 -

873 -

874 -class SequenceGroupMetadataDelta(

875 - msgspec.Struct,

876 - tag=True, # type: ignore[call-arg]

877 - array_like=True, # type: ignore[call-arg]

878 - omit_defaults=True): # type: ignore[call-arg]

879 - """Delta of SequenceGroupMetadata.

880 -

881 - After sending the first SequenceGroupMetadata, vLLM scheduler

882 - only sends delta to reduce the data payload size.

883 - """

884 - seq_data_delta: dict[int, SequenceDataDelta]

885 - request_id: str

886 - block_tables: dict[int, list[int]]

887 - is_prompt: bool

888 - do_sample: bool = True

889 - token_chunk_size: Optional[int] = None

890 - computed_block_nums: Optional[list[int]] = None

891 - state: Optional[SequenceGroupState] = msgspec.field(

892 - default_factory=lambda: SequenceGroupState())

893 -

894 -

895 -class SequenceGroupMetadata(

896 - msgspec.Struct,

897 - tag=True, # type: ignore[call-arg]

898 - array_like=True, # type: ignore[call-arg]

899 - omit_defaults=True): # type: ignore[call-arg]

900 - """Metadata for a sequence group. Used to create `AttentionMetadata`.

901 -

902 - Attributes:

903 - request_id: The ID of the request.

904 - is_prompt: Whether the request is at prompt stage.

905 - seq_data: The sequence data. (Seq id -> sequence data)

906 - sampling_params: The sampling parameters used to generate the outputs.

907 - block_tables: The block tables. (Seq id -> list of physical block

908 - numbers)

909 - do_sample: True if sampling is required. Sampling is not required when

910 - e.g., prefill is chunked, and the current iteration only computes

911 - query tokens for prefill, we don't need sampling.

912 - pooling_params: Pooling parameters.

913 - lora_request: LoRA request.

914 - computed_block_nums: The block numbers that are already computed,

915 - used in prefix caching.

916 - state: Internal state tied to this sequence group.

917 - token_type_ids: Token type IDs.

918 - multi_modal_data: Multi modal data.

919 - multi_modal_placeholders: Multi modal placeholders.

920 - encoder_seq_data: Optional sequence data for encoder prompt

921 - (SequenceGroup.encoder_seq). Should be None

922 - unless you are working with an encoder/decoder

923 - model.

924 - cross_block_table: Optional cross-attention block table associated

925 - with the encoder prompt

926 - (SequenceGroup.encoder_seq). Should be None

927 - unless you are working with an encoder/decoder

928 - model.

929 - """

930 -

931 - request_id: str

932 - is_prompt: bool

933 - seq_data: dict[int, SequenceData]

934 - sampling_params: Optional[SamplingParams]

935 - block_tables: dict[int, list[int]]

936 - do_sample: bool = True

937 - pooling_params: Optional[PoolingParams] = None

938 - lora_request: Optional[LoRARequest] = None

939 - computed_block_nums: Optional[list[int]] = None

940 - state: Optional[SequenceGroupState] = msgspec.field(

941 - default_factory=lambda: SequenceGroupState())

942 - multi_modal_data: Optional[MultiModalKwargs] = None

943 - multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None

944 - encoder_seq_data: Optional[SequenceData] = None

945 - cross_block_table: Optional[list[int]] = None

946 - token_chunk_size: Optional[int] = None

947 -

948 - ### Stateful fields that are lazily defined. ###

949 - # The number of speculative tokens adopted in this request.

950 - # None means specuative decoding is not used.

951 - # Zero means speculative decoding is disabled for some reasons.

952 - # TODO: We should maintain this states out of the sequence group.

953 - num_speculative_tokens: Optional[int] = None

954 -

955 - def __post_init__(self):

956 - if self.seq_data is not None and self.token_chunk_size is None:

957 - if self.is_prompt:

958 - self.token_chunk_size = next(iter(

959 - self.seq_data.values())).get_len()

960 - else:

961 - self.token_chunk_size = 1

962 -

963 - @property

964 - def lora_int_id(self) -> int:

965 - return self.lora_request.lora_int_id if self.lora_request else 0

966 -

967 - # Multi-Step Chunked-Prefill property

968 - @property

969 - def is_single_step_prompt(self) -> bool:

970 - # do_sample is true, only when the token_chunk_size matches the

971 - # num_uncomputed_tokens of the sequence. This indicates that

972 - # the prompt will finish processing in a single `execute_model`

973 - # step.

974 - return self.is_prompt and self.do_sample

975 -

976 - def get_first_seq_id(self) -> int:

977 - # This is an efficient way of fetching the seq_id when

978 - # we know this SequenceGroup has only one sequence.

979 - return next(iter(self.seq_data))

980 -

981 - def apply_delta(self,

982 - sequence_group_metadata_delta: SequenceGroupMetadataDelta):

983 - for id, delta in sequence_group_metadata_delta.seq_data_delta.items():

984 - self.seq_data[id].apply_delta(delta)

985 - assert self.request_id == sequence_group_metadata_delta.request_id

986 - self.block_tables = sequence_group_metadata_delta.block_tables

987 - self.token_chunk_size = sequence_group_metadata_delta.token_chunk_size

988 - self.do_sample = sequence_group_metadata_delta.do_sample

989 - self.is_prompt = sequence_group_metadata_delta.is_prompt

990 -

991 - def finish_step(self) -> None:

992 - assert self.state is not None

993 - assert self.state.current_step < self.state.num_steps, \

994 - f"current step {self.state.current_step}, num_steps {self.state.num_steps}" # noqa

995 - self.state.current_step += 1

996 -

997 -

998 -class SequenceOutput(

999 - msgspec.Struct,

1000 - omit_defaults=True, # type: ignore[call-arg]

1001 - array_like=True): # type: ignore[call-arg]

1002 - """The model output associated with a sequence.

1003 -

1004 - Attributes:

1005 - parent_seq_id: The ID of the parent sequence (for forking in beam

1006 - search).

1007 - output_token: The output token ID.

1008 - logprobs: The logprobs of the output token.

1009 - (Token id -> logP(x_i+1 | x_0, ..., x_i))

1010 - output_embed: Optional output embedding tensor.

1011 - """

1012 - parent_seq_id: int

1013 - output_token: int

1014 - logprobs: dict[int, Logprob]

1015 - output_embed: Optional[torch.Tensor] = None

1016 -

1017 - def __repr__(self) -> str:

1018 - output_embed_shape = \

1019 - self.output_embed.shape if self.output_embed is not None else None

1020 - return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "

1021 - f"output_token={self.output_token}, "

1022 - f"output_embed.shape={output_embed_shape}, "

1023 - f"logprobs={self.logprobs})")

1024 -

1025 - def __eq__(self, other: object) -> bool:

1026 - if not isinstance(other, SequenceOutput):

1027 - raise NotImplementedError()

1028 - equal = (self.parent_seq_id == other.parent_seq_id

1029 - and self.output_token == other.output_token)

1030 - log_probs_equal = other.logprobs == self.logprobs

1031 - return equal and log_probs_equal

1032 -

1033 -

1034 -class SequenceGroupOutput(ABC):

1035 - """The base class for model outputs associated with a sequence group."""

1036 -

1037 - @abstractmethod

1038 - def __repr__(self) -> str:

1039 - pass

1040 -

1041 - @abstractmethod

1042 - def __eq__(self, other: object) -> bool:

1043 - pass

1044 -

1045 -

1046 -class CompletionSequenceGroupOutput(

1047 - msgspec.Struct,

1048 - omit_defaults=True, # type: ignore[call-arg]

1049 - array_like=True): # type: ignore[call-arg]

1050 - """The model output associated with a completion sequence group."""

1051 - __metaclass__ = SequenceGroupOutput

1052 - samples: list[SequenceOutput]

1053 - # Prompt logprob for each prompt query token.

1054 - prompt_logprobs: Optional[PromptLogprobs]

1055 - step_index: Optional[int] = 0

1056 -

1057 - def __repr__(self) -> str:

1058 - return (f"CompletionSequenceGroupOutput(samples={self.samples}, "

1059 - f"prompt_logprobs={self.prompt_logprobs})")

1060 -

1061 - def __eq__(self, other: object) -> bool:

1062 - if not isinstance(other, CompletionSequenceGroupOutput):

1063 - raise NotImplementedError()

1064 - return (self.samples == other.samples

1065 - and self.prompt_logprobs == other.prompt_logprobs)

1066 -

1067 -

51 class PoolingSequenceGroupOutput(

52 msgspec.Struct,

53 omit_defaults=True, # type: ignore[call-arg]

54 array_like=True, # type: ignore[call-arg]

55 ):

56 """The model output associated with a pooling sequence group."""

1074 - __metaclass__ = SequenceGroupOutput

57 # Annotated as Any to be compatible with msgspec

58 # The actual type is in SequenceGroup.pooled_data

59 data: Any

@@ -1161,305 +143,9 @@ class PoolerOutput(

143 self.__class__) and self.outputs == other.outputs

144

145

1164 -def get_all_seq_ids(

1165 - seq_group_metadata_list: list[SequenceGroupMetadata]) -> list[int]:

1166 - """Given a list of SequenceGroupMetadata, create a list of all

1167 - sequence ids.

1168 - """

1169 - return [seq_id for sg in seq_group_metadata_list for seq_id in sg.seq_data]

1170 -

1171 -

1172 -def get_all_seq_ids_and_request_ids(

1173 - seq_group_metadata_list: list[SequenceGroupMetadata]

1174 -) -> tuple[list[int], dict[str, set[int]]]:

1175 - """Given a list of SequenceGroupMetadata, create a list of all

1176 - sequence ids.

1177 - """

1178 - seq_ids: list[int] = []

1179 - request_id_seq_ids_mapping: defaultdict[str, set[int]] = defaultdict(set)

1180 - for sg in seq_group_metadata_list:

1181 - for seq_id in sg.seq_data:

1182 - seq_ids.append(seq_id)

1183 - request_id_seq_ids_mapping[sg.request_id].add(seq_id)

1184 - return seq_ids, request_id_seq_ids_mapping

1185 -

1186 -

1187 -class HiddenStates(msgspec.Struct, array_like=True,

1188 - omit_defaults=True): # type: ignore[call-arg]

1189 - """Hidden states corresponding to in-progress sequences.

1190 - Used in speculative decoding to pass hidden states from

1191 - the target model to the proposer model.

1192 -

1193 - seq_ids are the sequence ids of each entry of the batch

1194 - dimension of the hidden_states tensor"""

1195 - # Scorer hidden states. For prefill step, it is used for hidden states of

1196 - # all tokens, whereas for decode step, it is used for last accepted tokens.

1197 - hidden_states: torch.Tensor

1198 - # The sequence group metadata list. Only needed for decode step.

1199 - seq_group_metadata_list: Optional[list[SequenceGroupMetadata]] = None

1200 - # Scorer hidden states of the 2nd last token proposed by the proposer (

1201 - # irrespective of whether it was accepted or not). Only used for cases when

1202 - # last proposed token is accepted (i.e., in case of bonus tokens). For the

1203 - # case of no bonus tokens, these are ignored.

1204 - second_last_token_hidden_states: Optional[torch.Tensor] = None

1205 -

1206 - _seq_ids: list[int] = msgspec.field(default_factory=list)

1207 -

1208 - def __post_init__(self):

1209 - if self.seq_group_metadata_list is not None:

1210 - assert len(self.seq_group_metadata_list) == len(self.hidden_states)

1211 - self._seq_ids = get_all_seq_ids(self.seq_group_metadata_list)

1212 -

1213 - @property

1214 - def seq_ids(self) -> list[int]:

1215 - return self._seq_ids

1216 -

1217 - def update(self,

1218 - hidden_states: torch.Tensor,

1219 - seq_group_metadata_list: list[SequenceGroupMetadata],

1220 - second_last_token_hidden_states: Optional[torch.Tensor] = None):

1221 - """Update hidden states from target model invocation. Only used for

1222 - decode steps"""

1223 - assert len(seq_group_metadata_list) == len(hidden_states)

1224 - self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list))

1225 - self.hidden_states = torch.cat([self.hidden_states, hidden_states])

1226 -

1227 - if self.second_last_token_hidden_states is not None:

1228 - # Adding dummy hidden_states to this to maintain same shape

1229 - self.second_last_token_hidden_states = torch.cat([

1230 - self.second_last_token_hidden_states,

1231 - torch.zeros_like(hidden_states)

1232 - if second_last_token_hidden_states is None else

1233 - second_last_token_hidden_states

1234 - ])

1235 -

1236 - def prune(self,

1237 - seq_group_metadata_list: list[SequenceGroupMetadata]) -> None:

1238 - """Prune to provided list of sequence ids. Only used for decode steps.

1239 - """

1240 - # Currently this prunes all seq_ids not present in

1241 - # seq_group_metadata_list which might cause problems where a sequence

1242 - # may be "paused" then "resumed" later. This should only prune sequences

1243 - # which are confirmed to be aborted.

1244 - seq_ids = get_all_seq_ids(seq_group_metadata_list)

1245 - # Only keep sequence IDs that exist in self._seq_ids

1246 - seq_ids = [seq_id for seq_id in seq_ids if seq_id in self._seq_ids]

1247 - if seq_ids != self._seq_ids:

1248 - # Batch contents changed - prune removed sequences.

1249 - index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]

1250 - self.hidden_states = self.hidden_states[index]

1251 - if self.second_last_token_hidden_states is not None:

1252 - self.second_last_token_hidden_states = self\

1253 - .second_last_token_hidden_states[index]

1254 - self._seq_ids = seq_ids

1255 -

1256 - def expand_with_bonus_tokens(

1257 - self, seq_with_bonus_token_in_last_step: set) -> None:

1258 - """Expand hidden states for sequences with bonus tokens. This is in

1259 - alignment with `MultiStepWorker._expand_execute_model_request`."""

1260 - if self.second_last_token_hidden_states is None \

1261 - or not seq_with_bonus_token_in_last_step:

1262 - return

1263 -

1264 - index = []

1265 - for seq_id in self._seq_ids:

1266 - i = self._seq_ids.index(seq_id)

1267 - if seq_id in seq_with_bonus_token_in_last_step:

1268 - index.append(i + len(self._seq_ids))

1269 - index.append(i)

1270 -

1271 - self.hidden_states = torch.cat(

1272 - [self.hidden_states, self.second_last_token_hidden_states])[index]

1273 -

1274 -

146 class ExecuteModelRequest(

147 msgspec.Struct,

148 array_like=True, # type: ignore[call-arg]

149 omit_defaults=True): # type: ignore[call-arg]

1279 - """The model execution request, containing CPU metadata only. The LLM

1280 - engine should create an instance of this class for each request batch."""

1281 - # The sequence group metadata list.

1282 - seq_group_metadata_list: list[Union[SequenceGroupMetadata,

1283 - SequenceGroupMetadataDelta]]

1284 - # Blocks to swap in. List of CPU -> GPU block number.

1285 - blocks_to_swap_in: list[tuple[int,

1286 - int]] = msgspec.field(default_factory=list)

1287 - # Blocks to swap out. List of GPU -> CPU block number.

1288 - blocks_to_swap_out: list[tuple[int,

1289 - int]] = msgspec.field(default_factory=list)

1290 - # Blocks to copy. Source to dest block.

1291 - blocks_to_copy: list[tuple[int, int]] = msgspec.field(default_factory=list)

1292 - # Virtual engine ID for pipeline parallel.

1293 - virtual_engine: int = 0

1294 - # The number of slots for lookahead decoding.

1295 - num_lookahead_slots: int = 0

1296 - # The number of requests in the running queue.

1297 - running_queue_size: int = 0

1298 - # Optional hidden states from prior step.

1299 - previous_hidden_states: Optional[HiddenStates] = None

1300 - # The number of forward steps to run.

1301 - num_steps: int = 1

1302 - # Finished request ids since last step.

1303 - finished_requests_ids: list[str] = msgspec.field(default_factory=list)

1304 - # The last sampled token ids for multi step decoding.

1305 - last_sampled_token_ids: Optional[torch.Tensor] = None

1306 - # Async callback

1307 - async_callback: Optional[Callable] = None

1308 -

1309 - @property

1310 - def is_last_step(self) -> bool:

1311 - # TODO(will) make this be able to handle batches with variable number of

1312 - # steps

1313 - assert len(self.seq_group_metadata_list) > 0

1314 - first_seq_group = self.seq_group_metadata_list[0]

1315 - assert first_seq_group.state is not None

1316 - return first_seq_group.state.remaining_steps == 1

1317 -

1318 - @property

1319 - def current_step(self) -> int:

1320 - # TODO(will) make this be able to handle batches with variable number of

1321 - # steps

1322 - assert len(self.seq_group_metadata_list) > 0

1323 - state = self.seq_group_metadata_list[0].state

1324 - assert state is not None

1325 - return state.current_step

1326 -

1327 - def clone(

1328 - self, seq_group_metadata_list: list[Union[SequenceGroupMetadata,

1329 - SequenceGroupMetadataDelta]]

1330 - ) -> "ExecuteModelRequest":

1331 - """Clone the request with a new sequence group metadata list."""

1332 - return ExecuteModelRequest(

1333 - seq_group_metadata_list=seq_group_metadata_list,

1334 - blocks_to_swap_in=self.blocks_to_swap_in.copy(),

1335 - blocks_to_swap_out=self.blocks_to_swap_out.copy(),

1336 - blocks_to_copy=self.blocks_to_copy.copy(),

1337 - virtual_engine=self.virtual_engine,

1338 - num_lookahead_slots=self.num_lookahead_slots,

1339 - running_queue_size=self.running_queue_size,

1340 - previous_hidden_states=self.previous_hidden_states,

1341 - num_steps=self.num_steps,

1342 - finished_requests_ids=self.finished_requests_ids,

1343 - last_sampled_token_ids=self.last_sampled_token_ids.clone()

1344 - if self.last_sampled_token_ids is not None else None,

1345 - async_callback=self.async_callback)

1346 -

1347 -

1348 -@dataclass

1349 -class SequenceGroupBase:

1350 - group_id: str # the original request id before splitting

1351 -

1352 - assembled_seq_group: Optional[SequenceGroup] = None

1353 -

1354 - # seq id to a unique index inside this group

1355 - seq_id_to_index: dict[str, int] = field(default_factory=dict)

1356 -

1357 - # seq ids to be finished

1358 - to_be_finished: dict[str, SequenceGroup] = field(default_factory=dict)

1359 -

1360 - # seq id to finished sequences

1361 - finished_reqs: dict[str, SequenceGroup] = field(default_factory=dict)

1362 -

1363 - streaming: bool = False

1364 -

1365 - output_produced: bool = False

1366 -

1367 - @staticmethod

1368 - def add_request(request_id: str, engine, params, *args, **kwargs):

1369 - """When we are ready to add a request with request_id and params

1370 - into the engine, we can split the request into multiple requests.

1371 - """

1372 - raise NotImplementedError

1373 -

1374 - def finish_seq(self, seq: SequenceGroup):

1375 - """The sequence `seq` finishes, we should record the information.

1376 - """

1377 - del self.to_be_finished[seq.request_id]

1378 - self.finished_reqs[seq.request_id] = seq

1379 -

1380 - def maybe_assemble_group(

1381 - self, seq_group: SequenceGroup) -> Optional[SequenceGroup]:

1382 - """Assemble the sequence group, for producing the final

1383 - output, or adding request in the engine again.

1384 - """

1385 - raise NotImplementedError

1386 -

1387 -

1388 -class ParallelSampleSequenceGroup(SequenceGroupBase):

1389 -

1390 - @staticmethod

1391 - def add_request(request_id: str, engine, params, **kwargs):

1392 - original_params = params

1393 - group = ParallelSampleSequenceGroup(request_id)

1394 - seqs = []

1395 - for i in range(original_params.n):

1396 - request_id_i = f"{request_id}_parallel_sample_{i}"

1397 - group.seq_id_to_index[request_id_i] = i

1398 - params = original_params.clone()

1399 - params.n = 1

1400 - if params.seed is not None:

1401 - params.seed += i

1402 - seq_group = engine._add_processed_request(

1403 - request_id_i,

1404 - params=params,

1405 - **kwargs,

1406 - ) # type: ignore

1407 - assert seq_group is not None

1408 - engine.seq_id_to_seq_group[request_id_i] = group

1409 - group.to_be_finished[request_id_i] = seq_group

1410 - seqs.append(seq_group.seqs[0])

1411 -

1412 - # for parallel sampling, the `assembled_seq_group` is always

1413 - # available, since we have all the sequences ready, and they

1414 - # will not change.

1415 - group.assembled_seq_group = SequenceGroup(

1416 - request_id=request_id,

1417 - seqs=seqs,

1418 - arrival_time=seq_group.arrival_time,

1419 - sampling_params=original_params,

1420 - lora_request=seq_group.lora_request,

1421 - pooling_params=seq_group.pooling_params,

1422 - pooled_data=seq_group.pooled_data,

1423 - encoder_seq=seq_group.encoder_seq,

1424 - trace_headers=seq_group.trace_headers,

1425 - priority=seq_group.priority,

1426 - )

1427 -

1428 - group.streaming = params.output_kind == RequestOutputKind.DELTA

1429 - group.output_produced = False

1430 -

1431 - def maybe_assemble_group(

1432 - self, seq_group: SequenceGroup) -> Optional[SequenceGroup]:

1433 -

1434 - # in the streaming mode, we will return the assembled sequence

1435 - # for the first remaining sequence, and then return None for the

1436 - # rest of sequences

1437 - if self.streaming:

1438 - first_remaining_id = next(iter(self.to_be_finished))

1439 - if seq_group.request_id == first_remaining_id:

1440 - return self.assembled_seq_group

1441 - return None

1442 -

1443 - # in the non-streaming mode, we will return the assembled sequence

1444 - # when the last sequences finishes, and then return None for the

1445 - # rest of the time

1446 - if (len(self.to_be_finished) == 1

1447 - and seq_group.request_id in self.to_be_finished

1448 - and seq_group.is_finished()):

1449 - assert self.assembled_seq_group is not None

1450 - params = self.assembled_seq_group.sampling_params

1451 - assert isinstance(params, SamplingParams)

1452 - if not self.output_produced:

1453 - self.output_produced = True

1454 - if params._real_n is not None:

1455 - # Get the top-n sequences.

1456 - n = params._real_n or params.n

1457 - seqs = self.assembled_seq_group.seqs

1458 - sorting_key = lambda seq: seq.get_cumulative_logprob()

1459 - sorted_seqs = sorted(seqs, key=sorting_key, reverse=True)

1460 - top_n_seqs = sorted_seqs[:n]

1461 - self.assembled_seq_group.seqs = top_n_seqs

1462 - return self.assembled_seq_group

1463 - if self.output_produced:

1464 - return None

1465 - return None

150 + # Placeholder. Remove.

151 + pass

vllm/transformers_utils/detokenizer.py deleted

-162

@@ -1,162 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -from typing import Optional

5 -

6 -from vllm.logprobs import Logprob

7 -from vllm.sequence import (VLLM_INVALID_TOKEN_ID, SamplingParams, Sequence,

8 - SequenceGroup)

9 -

10 -from .detokenizer_utils import (convert_prompt_ids_to_tokens,

11 - detokenize_incrementally)

12 -from .tokenizer import AnyTokenizer

13 -

14 -

15 -class Detokenizer:

16 - """Provides methods to decode the output of a model into text."""

17 -

18 - def __init__(self, tokenizer: AnyTokenizer):

19 - self.tokenizer = tokenizer

20 -

21 - def decode_prompt_logprobs_inplace(self, seq_group: SequenceGroup,

22 - prompt_logprobs: list[Optional[dict[

23 - int, Logprob]]],

24 - position_offset: int) -> None:

25 - """Decodes the logprobs for the prompt of a sequence group.

26 -

27 - Args:

28 - seq_group: The sequence group to decode.

29 - prompt_logprobs: The logprobs to decode.

30 - position_offset: Offset of the first index of the logprobs

31 - relative to the start of the sequence (for chunked prefill).

32 -

33 - Returns:

34 - The prompt logprobs with the decoded tokens.

35 - """

36 - prms = seq_group.sampling_params

37 - assert prms is not None

38 -

39 - # We can pick any sequence for the prompt.

40 - seq = seq_group.get_seqs()[0]

41 - # Only prompt, without the generated token.

42 - all_token_ids = seq.get_token_ids()

43 - prompt_token_ids = all_token_ids[:-1]

44 - prefix_offset = 0

45 - read_offset = 0

46 - next_iter_prefix_offset = 0

47 - next_iter_read_offset = 0

48 - next_iter_tokens: list[str] = []

49 - prev_tokens = None

50 -

51 - for token_position_in_logprob, prompt_logprobs_for_token in enumerate(

52 - prompt_logprobs):

53 -

54 - # Absolute token position equals the index in the logprobs

55 - # list plus the offset of the entire logprobs list relative

56 - # to the start of the sequence.

57 - token_position = token_position_in_logprob + position_offset

58 - if not prompt_logprobs_for_token:

59 - continue

60 - for token_id, sample_logprob in prompt_logprobs_for_token.items():

61 - if (sample_logprob.decoded_token is None

62 - and token_id != VLLM_INVALID_TOKEN_ID):

63 - prompt_token_ids_with_token = (

64 - prompt_token_ids[:token_position] + [token_id])

65 - (new_tokens, new_text, new_prefix_offset,

66 - new_read_offset) = detokenize_incrementally(

67 - tokenizer=self.tokenizer,

68 - all_input_ids=prompt_token_ids_with_token,

69 - prev_tokens=prev_tokens,

70 - prefix_offset=prefix_offset,

71 - read_offset=read_offset,

72 - skip_special_tokens=prms.skip_special_tokens,

73 - spaces_between_special_tokens=prms.

74 - spaces_between_special_tokens,

75 - )

76 -

77 - sample_logprob.decoded_token = new_text

78 -

79 - # Use the offsets & prev tokens corresponding to

80 - # real tokens to ensure detokenization is consistent

81 - # actual with prompt.

82 - if token_id == all_token_ids[token_position]:

83 - next_iter_prefix_offset = new_prefix_offset

84 - next_iter_read_offset = new_read_offset

85 - next_iter_tokens = new_tokens

86 -

87 - # Advance to the next token position.

88 - prefix_offset = next_iter_prefix_offset

89 - read_offset = next_iter_read_offset

90 - if prev_tokens is None:

91 - prev_tokens = next_iter_tokens.copy()

92 - else:

93 - prev_tokens.extend(next_iter_tokens)

94 -

95 - def decode_sequence_inplace(self, seq: Sequence,

96 - prms: SamplingParams) -> int:

97 - """Decodes the new token for a sequence. In-place operation.

98 -

99 - Args:

100 - seq: The sequence to decode.

101 - prms: The sampling parameters used to generate the sequence.

102 -

103 - Returns:

104 - The number of characters added to the output text.

105 - """

106 - all_input_ids = seq.get_token_ids()

107 - token_id_generated_this_iteration = all_input_ids[-1]

108 -

109 - # Convert prompt token IDs to tokens if necessary.

110 - # Do it here so that we don't have to repeat this

111 - # computation for each logprob.

112 - if seq.tokens is None:

113 - (seq.tokens, seq.prefix_offset,

114 - seq.read_offset) = convert_prompt_ids_to_tokens(

115 - tokenizer=self.tokenizer,

116 - prompt_ids=all_input_ids[:-1],

117 - skip_special_tokens=prms.skip_special_tokens,

118 - )

119 -

120 - (new_tokens, new_decoded_token_text, prefix_offset,

121 - read_offset) = detokenize_incrementally(

122 - tokenizer=self.tokenizer,

123 - all_input_ids=all_input_ids,

124 - prev_tokens=seq.tokens,

125 - prefix_offset=seq.prefix_offset,

126 - read_offset=seq.read_offset,

127 - skip_special_tokens=prms.skip_special_tokens,

128 - spaces_between_special_tokens=prms.spaces_between_special_tokens,

129 - )

130 -

131 - # Decode logprobs

132 - logprobs = seq.output_logprobs[-1]

133 - if logprobs:

134 - previous_tokens = all_input_ids[:-1]

135 - for token_id, sample_logprob in logprobs.items():

136 - # If the token was generated this iteration,

137 - # use the provided text.

138 - if token_id == token_id_generated_this_iteration:

139 - sample_logprob.decoded_token = new_decoded_token_text

140 - continue

141 -

142 - if (sample_logprob.decoded_token is None

143 - and token_id != VLLM_INVALID_TOKEN_ID):

144 - all_input_ids_with_logprob = previous_tokens + [token_id]

145 - (_, new_text, _, _) = detokenize_incrementally(

146 - tokenizer=self.tokenizer,

147 - all_input_ids=all_input_ids_with_logprob,

148 - prev_tokens=seq.tokens,

149 - prefix_offset=seq.prefix_offset,

150 - read_offset=seq.read_offset,

151 - skip_special_tokens=prms.skip_special_tokens,

152 - spaces_between_special_tokens=prms.

153 - spaces_between_special_tokens,

154 - )

155 - sample_logprob.decoded_token = new_text

156 -

157 - seq.tokens.extend(new_tokens)

158 - seq.prefix_offset = prefix_offset

159 - seq.read_offset = read_offset

160 - seq.output_text += new_decoded_token_text

161 -

162 - return len(new_decoded_token_text)

vllm/worker/worker_base.py

+1 -1

@@ -11,12 +11,12 @@ import torch.nn as nn

11 from vllm.config import VllmConfig, set_current_vllm_config

12 from vllm.logger import init_logger

13 from vllm.lora.request import LoRARequest

14 -from vllm.model_executor.layers.sampler import SamplerOutput

14 from vllm.sequence import ExecuteModelRequest

15 from vllm.utils import (enable_trace_function_call_for_thread,

16 resolve_obj_by_qualname, run_method,

17 update_environment_variables,

18 warn_for_unimplemented_methods)

19 +from vllm.v1.outputs import SamplerOutput

21 logger = init_logger(__name__)