| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Datastructures defining a GPU input batch |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | from typing import cast |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | |
| 11 | from vllm.config.reasoning import ReasoningConfig |
| 12 | from vllm.lora.request import LoRARequest |
| 13 | from vllm.multimodal.inputs import MultiModalFeatureSpec |
| 14 | from vllm.pooling_params import PoolingParams |
| 15 | from vllm.sampling_params import SamplingParams, SamplingType |
| 16 | from vllm.utils import length_from_prompt_token_ids_or_embeds |
| 17 | from vllm.utils.collection_utils import swap_dict_values |
| 18 | from vllm.utils.torch_utils import PIN_MEMORY |
| 19 | from vllm.v1.outputs import LogprobsTensors |
| 20 | from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates |
| 21 | from vllm.v1.sample.logits_processor import ( |
| 22 | BatchUpdateBuilder, |
| 23 | LogitsProcessors, |
| 24 | MoveDirectionality, |
| 25 | ) |
| 26 | from vllm.v1.sample.metadata import SamplingMetadata |
| 27 | from vllm.v1.sample.thinking_budget_state import ( |
| 28 | maybe_create_thinking_budget_state_holder, |
| 29 | ) |
| 30 | from vllm.v1.utils import copy_slice |
| 31 | from vllm.v1.worker.block_table import MultiGroupBlockTable, SlotMappingMode |
| 32 | |
| 33 | |
| 34 | @dataclass |
| 35 | class CachedRequestState: |
| 36 | req_id: str |
| 37 | prompt_token_ids: list[int] | None |
| 38 | mm_features: list[MultiModalFeatureSpec] |
| 39 | sampling_params: SamplingParams | None |
| 40 | generator: torch.Generator | None |
| 41 | |
| 42 | block_ids: tuple[list[int], ...] |
| 43 | num_computed_tokens: int |
| 44 | output_token_ids: list[int] |
| 45 | |
| 46 | mrope_positions: torch.Tensor | None = None |
| 47 | mrope_position_delta: int | None = None |
| 48 | |
| 49 | lora_request: LoRARequest | None = None |
| 50 | prompt_embeds: torch.Tensor | None = None |
| 51 | # To accumulate prompt logprobs tensor chunks across prefill steps. |
| 52 | in_progress_prompt_logprobs_cpu: LogprobsTensors | None = None |
| 53 | |
| 54 | # Per-position mask for mixed-mode inputs (e.g chat completion with |
| 55 | # prompt_embeds content parts). See `Request.prompt_is_token_ids`. |
| 56 | prompt_is_token_ids: list[bool] | None = None |
| 57 | |
| 58 | # Used when both async_scheduling and spec_decode are enabled. |
| 59 | prev_num_draft_len: int = 0 |
| 60 | |
| 61 | # for pooling models |
| 62 | pooling_params: PoolingParams | None = None |
| 63 | pooling_states: PoolingStates | None = None |
| 64 | |
| 65 | def __post_init__(self): |
| 66 | self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds( |
| 67 | self.prompt_token_ids, self.prompt_embeds |
| 68 | ) |
| 69 | |
| 70 | if self.pooling_params is not None: |
| 71 | self.pooling_states = PoolingStates() |
| 72 | |
| 73 | @property |
| 74 | def num_tokens(self) -> int: |
| 75 | return self.num_prompt_tokens + len(self.output_token_ids) |
| 76 | |
| 77 | def get_token_id(self, idx: int) -> int: |
| 78 | if idx < self.num_prompt_tokens: |
| 79 | if self.prompt_token_ids is None: |
| 80 | raise ValueError( |
| 81 | f"Tried to access token index {idx}, but that token was " |
| 82 | "provided via prompt_embeds, and its ID is unknown." |
| 83 | ) |
| 84 | return self.prompt_token_ids[idx] |
| 85 | if idx - self.num_prompt_tokens < len(self.output_token_ids): |
| 86 | return self.output_token_ids[idx - self.num_prompt_tokens] |
| 87 | return -1 |
| 88 | |
| 89 | |
| 90 | class InputBatch: |
| 91 | def __init__( |
| 92 | self, |
| 93 | max_num_reqs: int, |
| 94 | max_model_len: int, |
| 95 | max_num_batched_tokens: int, |
| 96 | device: torch.device, |
| 97 | vocab_size: int, |
| 98 | block_sizes: list[int], # The block_size of each kv cache group |
| 99 | kernel_block_sizes: list[int], |
| 100 | max_num_blocks_per_req: list[int], |
| 101 | logitsprocs: LogitsProcessors | None = None, |
| 102 | logitsprocs_need_output_token_ids: bool = False, |
| 103 | num_spec_tokens: int = 0, |
| 104 | is_pooling_model: bool = False, |
| 105 | cp_kv_cache_interleave_size: int = 1, |
| 106 | reasoning_config: ReasoningConfig | None = None, |
| 107 | use_replayssm: bool = False, |
| 108 | slot_mapping_modes: list[SlotMappingMode] | None = None, |
| 109 | ): |
| 110 | self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder( |
| 111 | reasoning_config, |
| 112 | max_num_reqs, |
| 113 | num_spec_tokens, |
| 114 | device, |
| 115 | PIN_MEMORY, |
| 116 | ) |
| 117 | self.thinking_token_budget_reqs: set[str] = set() |
| 118 | self.is_pooling_model = is_pooling_model |
| 119 | self.max_num_reqs = max_num_reqs |
| 120 | self.max_model_len = max_model_len |
| 121 | self.max_num_batched_tokens = max_num_batched_tokens |
| 122 | self.device = device |
| 123 | self.vocab_size = vocab_size |
| 124 | |
| 125 | self._req_ids: list[str | None] = [] |
| 126 | self.req_id_to_index: dict[str, int] = {} |
| 127 | |
| 128 | # TODO(woosuk): This buffer could be too large if max_model_len is big. |
| 129 | # Find a way to reduce the CPU memory usage. |
| 130 | # This buffer is not directly transferred to the GPU, so it does not |
| 131 | # need to be pinned. |
| 132 | self.token_ids_cpu_tensor = torch.zeros( |
| 133 | (max_num_reqs, max_model_len), |
| 134 | device="cpu", |
| 135 | dtype=torch.int32, |
| 136 | pin_memory=False, |
| 137 | ) |
| 138 | self.token_ids_cpu = self.token_ids_cpu_tensor.numpy() |
| 139 | self.is_token_ids_tensor = torch.zeros( |
| 140 | (max_num_reqs, max_model_len), |
| 141 | device="cpu", |
| 142 | dtype=bool, |
| 143 | pin_memory=False, |
| 144 | ) |
| 145 | self.is_token_ids = self.is_token_ids_tensor.numpy() |
| 146 | # Store prompt embeddings per request to avoid OOM from large upfront |
| 147 | # allocation if max_model_len is big. |
| 148 | # Maps req_index -> tensor of shape (num_prompt_tokens, hidden_size) |
| 149 | self.req_prompt_embeds: dict[int, torch.Tensor] = {} |
| 150 | self.num_tokens_no_spec_cpu_tensor = torch.zeros( |
| 151 | (max_num_reqs,), |
| 152 | device="cpu", |
| 153 | dtype=torch.int32, |
| 154 | pin_memory=PIN_MEMORY, |
| 155 | ) |
| 156 | self.num_tokens_no_spec = self.num_tokens_no_spec_cpu_tensor.numpy() |
| 157 | self.num_prompt_tokens_cpu_tensor = torch.zeros( |
| 158 | (max_num_reqs,), |
| 159 | device="cpu", |
| 160 | dtype=torch.int32, |
| 161 | pin_memory=PIN_MEMORY, |
| 162 | ) |
| 163 | self.num_prompt_tokens = self.num_prompt_tokens_cpu_tensor.numpy() |
| 164 | self.num_computed_tokens_cpu_tensor = torch.zeros( |
| 165 | (max_num_reqs,), |
| 166 | device="cpu", |
| 167 | dtype=torch.int32, |
| 168 | pin_memory=PIN_MEMORY, |
| 169 | ) |
| 170 | self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy() |
| 171 | |
| 172 | # Mamba2 ReplaySSM decode ring origin (num_computed at each request's |
| 173 | # last full-state write); populated only when the feature is on. |
| 174 | self.use_replayssm = use_replayssm |
| 175 | self.replayssm_decode_base_cpu_tensor = torch.zeros( |
| 176 | (max_num_reqs,), |
| 177 | device="cpu", |
| 178 | dtype=torch.int32, |
| 179 | pin_memory=PIN_MEMORY, |
| 180 | ) |
| 181 | self.replayssm_decode_base = self.replayssm_decode_base_cpu_tensor.numpy() |
| 182 | |
| 183 | # Block table. |
| 184 | self.block_table = MultiGroupBlockTable( |
| 185 | max_num_reqs=max_num_reqs, |
| 186 | max_num_batched_tokens=max_num_batched_tokens, |
| 187 | pin_memory=PIN_MEMORY, |
| 188 | device=device, |
| 189 | block_sizes=block_sizes, |
| 190 | kernel_block_sizes=kernel_block_sizes, |
| 191 | max_num_blocks=max_num_blocks_per_req, |
| 192 | cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, |
| 193 | slot_mapping_modes=slot_mapping_modes, |
| 194 | ) |
| 195 | |
| 196 | # Sampling-related. |
| 197 | self.temperature = torch.empty( |
| 198 | (max_num_reqs,), dtype=torch.float32, device=device |
| 199 | ) |
| 200 | self.temperature_cpu_tensor = torch.empty( |
| 201 | (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=PIN_MEMORY |
| 202 | ) |
| 203 | self.temperature_cpu = self.temperature_cpu_tensor.numpy() |
| 204 | self.greedy_reqs: set[str] = set() |
| 205 | self.random_reqs: set[str] = set() |
| 206 | |
| 207 | self.top_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) |
| 208 | self.top_p_cpu_tensor = torch.empty( |
| 209 | (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=PIN_MEMORY |
| 210 | ) |
| 211 | self.top_p_cpu = self.top_p_cpu_tensor.numpy() |
| 212 | self.top_p_reqs: set[str] = set() |
| 213 | |
| 214 | self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device) |
| 215 | self.top_k_cpu_tensor = torch.empty( |
| 216 | (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY |
| 217 | ) |
| 218 | self.top_k_cpu = self.top_k_cpu_tensor.numpy() |
| 219 | self.top_k_reqs: set[str] = set() |
| 220 | |
| 221 | # Frequency penalty related data structures |
| 222 | self.frequency_penalties = torch.empty( |
| 223 | (max_num_reqs,), dtype=torch.float, device=device |
| 224 | ) |
| 225 | self.frequency_penalties_cpu_tensor = torch.empty( |
| 226 | (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY |
| 227 | ) |
| 228 | self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy() |
| 229 | self.frequency_penalties_reqs: set[str] = set() |
| 230 | |
| 231 | # Presence penalty related data structures |
| 232 | self.presence_penalties = torch.empty( |
| 233 | (max_num_reqs,), dtype=torch.float, device=device |
| 234 | ) |
| 235 | self.presence_penalties_cpu_tensor = torch.empty( |
| 236 | (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY |
| 237 | ) |
| 238 | self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy() |
| 239 | self.presence_penalties_reqs: set[str] = set() |
| 240 | |
| 241 | # Repetition penalty related data structures |
| 242 | self.repetition_penalties = torch.empty( |
| 243 | (max_num_reqs,), dtype=torch.float, device=device |
| 244 | ) |
| 245 | self.repetition_penalties_cpu_tensor = torch.empty( |
| 246 | (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY |
| 247 | ) |
| 248 | self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy() |
| 249 | self.repetition_penalties_reqs: set[str] = set() |
| 250 | |
| 251 | # Speculative decoding |
| 252 | self.num_accepted_tokens_cpu_tensor = torch.ones( |
| 253 | (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY |
| 254 | ) |
| 255 | self.num_accepted_tokens_cpu = self.num_accepted_tokens_cpu_tensor.numpy() |
| 256 | |
| 257 | # lora related |
| 258 | self.request_lora_mapping = np.zeros((self.max_num_reqs,), dtype=np.int64) |
| 259 | self.lora_id_to_request_ids: dict[int, set[str]] = {} |
| 260 | self.lora_id_to_lora_request: dict[int, LoRARequest] = {} |
| 261 | |
| 262 | # req_index -> generator |
| 263 | # NOTE(woosuk): The indices of the requests that do not have their own |
| 264 | # generator should not be included in the dictionary. |
| 265 | self.generators: dict[int, torch.Generator] = {} |
| 266 | |
| 267 | self.num_logprobs: dict[str, int] = {} |
| 268 | |
| 269 | # req_id -> list of specific token IDs to compute logprobs for |
| 270 | # More efficient than num_logprobs=-1 when only a few tokens are needed |
| 271 | self.logprob_token_ids: dict[str, list[int]] = {} |
| 272 | |
| 273 | # Internal representation of per-step batch state changes, used for |
| 274 | # reordering persistent batch and generating logitsprocs batch state |
| 275 | # updates. Should reset each step. |
| 276 | self.batch_update_builder = BatchUpdateBuilder() |
| 277 | |
| 278 | # TODO convert this to LogitsProcessor |
| 279 | self.has_allowed_token_ids: set[str] = set() |
| 280 | # NOTE(lufang): In the mask tensor, if the corresponding token allowed, |
| 281 | # the value is False. Since we use masked_fill_ to set -inf. |
| 282 | self.allowed_token_ids_mask: torch.Tensor | None = None |
| 283 | self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None |
| 284 | |
| 285 | # req_index -> bad_words_token_ids |
| 286 | self.bad_words_token_ids: dict[int, list[list[int]]] = {} |
| 287 | |
| 288 | self.logits_processing_needs_token_ids = np.zeros(max_num_reqs, dtype=bool) |
| 289 | |
| 290 | self.req_output_token_ids: list[list[int] | None] = [] |
| 291 | |
| 292 | # Store provided logitsprocs. If none are provided, initialize empty |
| 293 | # data structure |
| 294 | self.logitsprocs = logitsprocs or LogitsProcessors() |
| 295 | self.logitsprocs_need_output_token_ids = logitsprocs_need_output_token_ids |
| 296 | |
| 297 | # Store last speculative tokens for sampler. |
| 298 | self.spec_token_ids: list[list[int]] = [[] for _ in range(max_num_reqs)] |
| 299 | |
| 300 | # This is updated each time the batch constituents change. |
| 301 | self.sampling_metadata = self._make_sampling_metadata() |
| 302 | |
| 303 | # for pooling models |
| 304 | self.pooling_params: dict[str, PoolingParams] = {} |
| 305 | self.pooling_states: dict[str, PoolingStates] = {} |
| 306 | |
| 307 | # Cached reference to the GPU tensor of previously sampled tokens |
| 308 | self.prev_sampled_token_ids: torch.Tensor | None = None |
| 309 | self.prev_req_id_to_index: dict[str, int] | None = None |
| 310 | # These are used to update output_token_ids with real sampled |
| 311 | # ids from prior step, if required by current sampling params |
| 312 | # (e.g. penalties). |
| 313 | self.sampled_token_ids_cpu: torch.Tensor | None = None |
| 314 | self.async_copy_ready_event: torch.Event | None = None |
| 315 | |
| 316 | @property |
| 317 | def req_ids(self) -> list[str]: |
| 318 | # None elements should only be present transiently |
| 319 | # while performing state updates to the batch. |
| 320 | return cast(list[str], self._req_ids) |
| 321 | |
| 322 | def _register_add_request(self, request: "CachedRequestState") -> int: |
| 323 | """Track add-request operations for logits processors. |
| 324 | Not applicable to pooling models. |
| 325 | """ |
| 326 | # Fill the next empty index if there is one. |
| 327 | if (new_req_index := self.batch_update_builder.pop_removed()) is None: |
| 328 | # Append to end otherwise. |
| 329 | new_req_index = self.num_reqs |
| 330 | |
| 331 | assert new_req_index < self.max_num_reqs |
| 332 | self.batch_update_builder.batch_changed = True |
| 333 | if request.sampling_params: |
| 334 | # Detailed added request metadata is only required for non-pooling |
| 335 | # models, to support logitsprocs. |
| 336 | self.batch_update_builder.added.append( |
| 337 | ( |
| 338 | new_req_index, |
| 339 | request.sampling_params, |
| 340 | request.prompt_token_ids, |
| 341 | request.output_token_ids, |
| 342 | ) |
| 343 | ) |
| 344 | |
| 345 | return new_req_index |
| 346 | |
| 347 | def add_request( |
| 348 | self, |
| 349 | request: "CachedRequestState", |
| 350 | ) -> int: |
| 351 | req_index = self._register_add_request(request) |
| 352 | |
| 353 | req_id = request.req_id |
| 354 | if req_index == len(self._req_ids): |
| 355 | self._req_ids.append(req_id) |
| 356 | self.req_output_token_ids.append(request.output_token_ids) |
| 357 | self.spec_token_ids.append([]) |
| 358 | else: |
| 359 | self._req_ids[req_index] = req_id |
| 360 | self.req_output_token_ids[req_index] = request.output_token_ids |
| 361 | self.spec_token_ids[req_index].clear() |
| 362 | |
| 363 | self.req_id_to_index[req_id] = req_index |
| 364 | |
| 365 | # Copy the prompt token ids and output token ids. |
| 366 | num_prompt_tokens = length_from_prompt_token_ids_or_embeds( |
| 367 | request.prompt_token_ids, request.prompt_embeds |
| 368 | ) |
| 369 | self.num_prompt_tokens[req_index] = num_prompt_tokens |
| 370 | start_idx = num_prompt_tokens |
| 371 | end_idx = start_idx + len(request.output_token_ids) |
| 372 | if request.prompt_token_ids is not None: |
| 373 | self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids |
| 374 | if request.prompt_is_token_ids is not None: |
| 375 | self.is_token_ids[req_index, :num_prompt_tokens] = ( |
| 376 | request.prompt_is_token_ids |
| 377 | ) |
| 378 | else: |
| 379 | self.is_token_ids[req_index, :num_prompt_tokens] = True |
| 380 | else: |
| 381 | self.is_token_ids[req_index, :num_prompt_tokens] = False |
| 382 | if request.prompt_embeds is not None: |
| 383 | self.req_prompt_embeds[req_index] = request.prompt_embeds |
| 384 | self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids |
| 385 | self.is_token_ids[req_index, start_idx:end_idx] = True |
| 386 | # Number of tokens without spec decode tokens. |
| 387 | self.num_tokens_no_spec[req_index] = request.num_tokens |
| 388 | |
| 389 | if self.use_replayssm: |
| 390 | # Ring origin = full context at (re)admission (prompt + any resumed |
| 391 | # output), so a resumed request re-anchors past the prompt. |
| 392 | self.replayssm_decode_base[req_index] = request.num_tokens |
| 393 | |
| 394 | self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens |
| 395 | self.block_table.add_row(request.block_ids, req_index) |
| 396 | |
| 397 | if sampling_params := request.sampling_params: |
| 398 | if sampling_params.sampling_type == SamplingType.GREEDY: |
| 399 | # Should avoid division by zero later when apply_temperature. |
| 400 | self.temperature_cpu[req_index] = 0.0 |
| 401 | self.greedy_reqs.add(req_id) |
| 402 | else: |
| 403 | self.temperature_cpu[req_index] = sampling_params.temperature |
| 404 | self.random_reqs.add(req_id) |
| 405 | |
| 406 | self.top_p_cpu[req_index] = sampling_params.top_p |
| 407 | if sampling_params.top_p < 1: |
| 408 | self.top_p_reqs.add(req_id) |
| 409 | top_k = sampling_params.top_k |
| 410 | if 0 < top_k < self.vocab_size: |
| 411 | self.top_k_reqs.add(req_id) |
| 412 | else: |
| 413 | top_k = self.vocab_size |
| 414 | self.top_k_cpu[req_index] = top_k |
| 415 | self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty |
| 416 | if sampling_params.frequency_penalty != 0.0: |
| 417 | self.frequency_penalties_reqs.add(req_id) |
| 418 | self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty |
| 419 | if sampling_params.presence_penalty != 0.0: |
| 420 | self.presence_penalties_reqs.add(req_id) |
| 421 | self.repetition_penalties_cpu[req_index] = ( |
| 422 | sampling_params.repetition_penalty |
| 423 | ) |
| 424 | if sampling_params.repetition_penalty != 1.0: |
| 425 | self.repetition_penalties_reqs.add(req_id) |
| 426 | |
| 427 | # NOTE(woosuk): self.generators should not include the requests that |
| 428 | # do not have their own generator. |
| 429 | if request.generator is not None: |
| 430 | self.generators[req_index] = request.generator |
| 431 | |
| 432 | if sampling_params.logprobs is not None: |
| 433 | self.num_logprobs[req_id] = ( |
| 434 | self.vocab_size |
| 435 | if sampling_params.logprobs == -1 |
| 436 | else sampling_params.logprobs |
| 437 | ) |
| 438 | |
| 439 | # Store specific token IDs to compute logprobs for (more efficient) |
| 440 | if sampling_params.logprob_token_ids is not None: |
| 441 | self.logprob_token_ids[req_id] = sampling_params.logprob_token_ids |
| 442 | |
| 443 | if sampling_params.allowed_token_ids: |
| 444 | self.has_allowed_token_ids.add(req_id) |
| 445 | if self.allowed_token_ids_mask_cpu_tensor is None: |
| 446 | # Lazy allocation for this tensor, which can be large. |
| 447 | # False means we don't fill with -inf. |
| 448 | self.allowed_token_ids_mask = torch.zeros( |
| 449 | self.max_num_reqs, |
| 450 | self.vocab_size, |
| 451 | dtype=torch.bool, |
| 452 | device=self.device, |
| 453 | ) |
| 454 | self.allowed_token_ids_mask_cpu_tensor = torch.zeros( |
| 455 | self.max_num_reqs, |
| 456 | self.vocab_size, |
| 457 | dtype=torch.bool, |
| 458 | device="cpu", |
| 459 | ) |
| 460 | self.allowed_token_ids_mask_cpu_tensor[req_index] = True |
| 461 | # False means we don't fill with -inf. |
| 462 | self.allowed_token_ids_mask_cpu_tensor[req_index][ |
| 463 | sampling_params.allowed_token_ids |
| 464 | ] = False |
| 465 | |
| 466 | if sampling_params.bad_words_token_ids: |
| 467 | self.bad_words_token_ids[req_index] = ( |
| 468 | sampling_params.bad_words_token_ids |
| 469 | ) |
| 470 | elif pooling_params := request.pooling_params: |
| 471 | pooling_states = request.pooling_states |
| 472 | assert pooling_states is not None |
| 473 | |
| 474 | self.pooling_params[req_id] = pooling_params |
| 475 | self.pooling_states[req_id] = pooling_states |
| 476 | self.logits_processing_needs_token_ids[req_index] = False |
| 477 | else: |
| 478 | raise NotImplementedError("Unrecognized request type") |
| 479 | |
| 480 | # Speculative decoding: by default 1 token is generated. |
| 481 | self.num_accepted_tokens_cpu[req_index] = 1 |
| 482 | |
| 483 | # Add request lora ID |
| 484 | if request.lora_request: |
| 485 | lora_id = request.lora_request.lora_int_id |
| 486 | if lora_id not in self.lora_id_to_request_ids: |
| 487 | self.lora_id_to_request_ids[lora_id] = set() |
| 488 | |
| 489 | self.request_lora_mapping[req_index] = lora_id |
| 490 | self.lora_id_to_request_ids[lora_id].add(request.req_id) |
| 491 | self.lora_id_to_lora_request[lora_id] = request.lora_request |
| 492 | else: |
| 493 | # No LoRA |
| 494 | self.request_lora_mapping[req_index] = 0 |
| 495 | |
| 496 | return req_index |
| 497 | |
| 498 | def update_req_spec_token_ids( |
| 499 | self, request: CachedRequestState, scheduled_spec_tokens: dict[str, list[int]] |
| 500 | ) -> None: |
| 501 | req_id = request.req_id |
| 502 | req_index = self.req_id_to_index[req_id] |
| 503 | cur_spec_token_ids = self.spec_token_ids[req_index] |
| 504 | # When speculative decoding is used with structured output, |
| 505 | # the scheduler can drop draft tokens that do not |
| 506 | # conform to the schema. This can result in |
| 507 | # scheduler_output.scheduled_spec_decode_tokens being empty, |
| 508 | # even when speculative decoding is enabled. |
| 509 | cur_spec_token_ids.clear() |
| 510 | spec_token_ids = scheduled_spec_tokens.get(req_id, ()) |
| 511 | num_spec_tokens = len(spec_token_ids) |
| 512 | request.prev_num_draft_len = num_spec_tokens |
| 513 | if not spec_token_ids: |
| 514 | return |
| 515 | |
| 516 | # For async scheduling, token_ids_cpu assigned from |
| 517 | # spec_token_ids are placeholders and will be overwritten in |
| 518 | # _prepare_input_ids. |
| 519 | start_index = self.num_tokens_no_spec[req_index] |
| 520 | end_token_index = start_index + num_spec_tokens |
| 521 | self.token_ids_cpu[req_index, start_index:end_token_index] = spec_token_ids |
| 522 | self.is_token_ids[req_index, start_index:end_token_index] = True |
| 523 | cur_spec_token_ids.extend(spec_token_ids) |
| 524 | |
| 525 | def remove_request(self, req_id: str) -> int | None: |
| 526 | """This method must always be followed by a call to condense(). |
| 527 | |
| 528 | Args: |
| 529 | req_id: request to remove |
| 530 | |
| 531 | Returns: |
| 532 | Removed request index, or `None` if `req_id` not recognized |
| 533 | |
| 534 | """ |
| 535 | req_index = self.req_id_to_index.pop(req_id, None) |
| 536 | if req_index is None: |
| 537 | return None |
| 538 | |
| 539 | self.batch_update_builder.removed_append(req_index) |
| 540 | self._req_ids[req_index] = None |
| 541 | self.req_output_token_ids[req_index] = None |
| 542 | self.spec_token_ids[req_index].clear() |
| 543 | self.block_table.clear_row(req_index) |
| 544 | self.req_prompt_embeds.pop(req_index, None) |
| 545 | |
| 546 | # LoRA |
| 547 | lora_id = self.request_lora_mapping[req_index] |
| 548 | if lora_id != 0: |
| 549 | lora_req_ids = self.lora_id_to_request_ids[lora_id] |
| 550 | lora_req_ids.discard(req_id) |
| 551 | if not lora_req_ids: |
| 552 | del self.lora_id_to_request_ids[lora_id] |
| 553 | del self.lora_id_to_lora_request[lora_id] |
| 554 | self.request_lora_mapping[req_index] = 0 |
| 555 | |
| 556 | if self.is_pooling_model: |
| 557 | self.pooling_params.pop(req_id, None) |
| 558 | self.pooling_states.pop(req_id, None) |
| 559 | return req_index |
| 560 | |
| 561 | self.greedy_reqs.discard(req_id) |
| 562 | self.random_reqs.discard(req_id) |
| 563 | self.top_p_reqs.discard(req_id) |
| 564 | self.top_k_reqs.discard(req_id) |
| 565 | self.frequency_penalties_reqs.discard(req_id) |
| 566 | self.presence_penalties_reqs.discard(req_id) |
| 567 | self.repetition_penalties_reqs.discard(req_id) |
| 568 | self.generators.pop(req_index, None) |
| 569 | self.num_logprobs.pop(req_id, None) |
| 570 | self.logprob_token_ids.pop(req_id, None) |
| 571 | if self.prev_req_id_to_index is not None: |
| 572 | self.prev_req_id_to_index.pop(req_id, None) |
| 573 | |
| 574 | self.has_allowed_token_ids.discard(req_id) |
| 575 | if self.allowed_token_ids_mask_cpu_tensor is not None: |
| 576 | # False means we don't fill with -inf. |
| 577 | self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False) |
| 578 | self.bad_words_token_ids.pop(req_index, None) |
| 579 | self.thinking_token_budget_reqs.discard(req_id) |
| 580 | return req_index |
| 581 | |
| 582 | def swap_states(self, i1: int, i2: int) -> None: |
| 583 | old_id_i1 = self._req_ids[i1] |
| 584 | old_id_i2 = self._req_ids[i2] |
| 585 | # Only swap the active token prefix for each request. Copying full |
| 586 | # max_model_len rows is expensive and unnecessary during reordering. |
| 587 | i1_active_token_count = self._get_active_token_count(i1) |
| 588 | i2_active_token_count = self._get_active_token_count(i2) |
| 589 | max_active_token_count = max(i1_active_token_count, i2_active_token_count) |
| 590 | |
| 591 | self._req_ids[i1], self._req_ids[i2] = self._req_ids[i2], self._req_ids[i1] # noqa |
| 592 | self.req_output_token_ids[i1], self.req_output_token_ids[i2] = ( |
| 593 | self.req_output_token_ids[i2], |
| 594 | self.req_output_token_ids[i1], |
| 595 | ) |
| 596 | self.spec_token_ids[i1], self.spec_token_ids[i2] = ( |
| 597 | self.spec_token_ids[i2], |
| 598 | self.spec_token_ids[i1], |
| 599 | ) |
| 600 | assert old_id_i1 is not None and old_id_i2 is not None |
| 601 | self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] = ( |
| 602 | self.req_id_to_index[old_id_i2], |
| 603 | self.req_id_to_index[old_id_i1], |
| 604 | ) |
| 605 | self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = ( |
| 606 | self.num_tokens_no_spec[i2], |
| 607 | self.num_tokens_no_spec[i1], |
| 608 | ) |
| 609 | self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] = ( |
| 610 | self.num_prompt_tokens[i2], |
| 611 | self.num_prompt_tokens[i1], |
| 612 | ) |
| 613 | if self.use_replayssm: |
| 614 | self.replayssm_decode_base[i1], self.replayssm_decode_base[i2] = ( |
| 615 | self.replayssm_decode_base[i2], |
| 616 | self.replayssm_decode_base[i1], |
| 617 | ) |
| 618 | self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = ( |
| 619 | self.num_computed_tokens_cpu[i2], |
| 620 | self.num_computed_tokens_cpu[i1], |
| 621 | ) |
| 622 | |
| 623 | # NOTE: the following is unsafe |
| 624 | # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\ |
| 625 | # self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...] |
| 626 | # instead, we need to temporarily copy the data for one of the indices |
| 627 | tmp_token_ids = self.token_ids_cpu[i1, :max_active_token_count].copy() |
| 628 | self.token_ids_cpu[i1, :max_active_token_count] = self.token_ids_cpu[ |
| 629 | i2, :max_active_token_count |
| 630 | ] |
| 631 | self.token_ids_cpu[i2, :max_active_token_count] = tmp_token_ids |
| 632 | |
| 633 | self.is_token_ids[[i1, i2], :max_active_token_count] = self.is_token_ids[ |
| 634 | [i2, i1], :max_active_token_count |
| 635 | ] |
| 636 | |
| 637 | # Swap prompt embeddings if they exist |
| 638 | embeds_i1 = self.req_prompt_embeds.get(i1) |
| 639 | embeds_i2 = self.req_prompt_embeds.get(i2) |
| 640 | if embeds_i1 is not None: |
| 641 | self.req_prompt_embeds[i2] = embeds_i1 |
| 642 | else: |
| 643 | self.req_prompt_embeds.pop(i2, None) |
| 644 | if embeds_i2 is not None: |
| 645 | self.req_prompt_embeds[i1] = embeds_i2 |
| 646 | else: |
| 647 | self.req_prompt_embeds.pop(i1, None) |
| 648 | |
| 649 | self.block_table.swap_row(i1, i2) |
| 650 | |
| 651 | self.request_lora_mapping[i1], self.request_lora_mapping[i2] = ( |
| 652 | self.request_lora_mapping[i2], |
| 653 | self.request_lora_mapping[i1], |
| 654 | ) |
| 655 | |
| 656 | if self.is_pooling_model: |
| 657 | # Sampling and logits parameters don't apply to pooling models. |
| 658 | return |
| 659 | |
| 660 | # For autoregressive models, track detailed request reordering info |
| 661 | # to support logitsprocs. |
| 662 | self.batch_update_builder.moved.append((i1, i2, MoveDirectionality.SWAP)) |
| 663 | |
| 664 | self.temperature_cpu[i1], self.temperature_cpu[i2] = ( |
| 665 | self.temperature_cpu[i2], |
| 666 | self.temperature_cpu[i1], |
| 667 | ) |
| 668 | self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1] |
| 669 | self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1] |
| 670 | self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = ( |
| 671 | self.frequency_penalties_cpu[i2], |
| 672 | self.frequency_penalties_cpu[i1], |
| 673 | ) |
| 674 | self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] = ( |
| 675 | self.presence_penalties_cpu[i2], |
| 676 | self.presence_penalties_cpu[i1], |
| 677 | ) |
| 678 | self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = ( |
| 679 | self.repetition_penalties_cpu[i2], |
| 680 | self.repetition_penalties_cpu[i1], |
| 681 | ) |
| 682 | self.num_accepted_tokens_cpu[i1], self.num_accepted_tokens_cpu[i2] = ( |
| 683 | self.num_accepted_tokens_cpu[i2], |
| 684 | self.num_accepted_tokens_cpu[i1], |
| 685 | ) |
| 686 | |
| 687 | swap_dict_values(self.generators, i1, i2) |
| 688 | swap_dict_values(self.bad_words_token_ids, i1, i2) |
| 689 | |
| 690 | if self.allowed_token_ids_mask_cpu_tensor is not None: |
| 691 | self.allowed_token_ids_mask_cpu_tensor[[i1, i2]] = ( |
| 692 | self.allowed_token_ids_mask_cpu_tensor[[i2, i1]] |
| 693 | ) |
| 694 | |
| 695 | def _get_active_token_count(self, req_index: int) -> int: |
| 696 | return int(self.num_tokens_no_spec[req_index]) + len( |
| 697 | self.spec_token_ids[req_index] |
| 698 | ) |
| 699 | |
| 700 | def condense(self) -> None: |
| 701 | """Slide non-empty requests down into lower, empty indices. |
| 702 | |
| 703 | Any consecutive empty indices at the very end of the list are not |
| 704 | filled. |
| 705 | |
| 706 | Returns: |
| 707 | swaps: list of (from,to) swap tuples for moved requests |
| 708 | empty_req_indices: indices not filled by condensation |
| 709 | |
| 710 | """ |
| 711 | num_reqs = self.num_reqs |
| 712 | |
| 713 | if not (empty_req_indices := self.batch_update_builder.removed): |
| 714 | # All removed requests were replaced by added requests, or else no |
| 715 | # requests were removed at all. No condense() needed |
| 716 | return |
| 717 | if num_reqs == 0: |
| 718 | # The batched states are empty. |
| 719 | self._req_ids.clear() |
| 720 | self.req_output_token_ids.clear() |
| 721 | self.spec_token_ids.clear() |
| 722 | return |
| 723 | |
| 724 | # NOTE(woosuk): This function assumes that the empty_req_indices |
| 725 | # is sorted in descending order. |
| 726 | last_req_index = num_reqs + len(empty_req_indices) - 1 |
| 727 | while empty_req_indices: |
| 728 | # Find the largest non-empty index. |
| 729 | while last_req_index in empty_req_indices: |
| 730 | last_req_index -= 1 |
| 731 | |
| 732 | # Find the smallest empty index. |
| 733 | empty_index = self.batch_update_builder.peek_removed() |
| 734 | assert empty_index is not None |
| 735 | if empty_index >= last_req_index: |
| 736 | break |
| 737 | |
| 738 | # Move active request down into empty request |
| 739 | # index. |
| 740 | self.batch_update_builder.pop_removed() |
| 741 | req_id = self._req_ids[last_req_index] |
| 742 | output_token_ids = self.req_output_token_ids[last_req_index] |
| 743 | assert req_id is not None |
| 744 | self._req_ids[empty_index] = req_id |
| 745 | self._req_ids[last_req_index] = None |
| 746 | self.req_output_token_ids[empty_index] = output_token_ids |
| 747 | self.req_output_token_ids[last_req_index] = None |
| 748 | self.req_id_to_index[req_id] = empty_index |
| 749 | |
| 750 | num_tokens = self._get_active_token_count(last_req_index) |
| 751 | |
| 752 | (self.spec_token_ids[last_req_index], self.spec_token_ids[empty_index]) = ( |
| 753 | self.spec_token_ids[empty_index], |
| 754 | self.spec_token_ids[last_req_index], |
| 755 | ) |
| 756 | self.spec_token_ids[last_req_index].clear() |
| 757 | |
| 758 | self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[ |
| 759 | last_req_index, :num_tokens |
| 760 | ] |
| 761 | self.is_token_ids[empty_index, :num_tokens] = self.is_token_ids[ |
| 762 | last_req_index, :num_tokens |
| 763 | ] |
| 764 | if last_req_index in self.req_prompt_embeds: |
| 765 | self.req_prompt_embeds[empty_index] = self.req_prompt_embeds.pop( |
| 766 | last_req_index |
| 767 | ) |
| 768 | self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[ |
| 769 | last_req_index |
| 770 | ] |
| 771 | self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index] |
| 772 | if self.use_replayssm: |
| 773 | self.replayssm_decode_base[empty_index] = self.replayssm_decode_base[ |
| 774 | last_req_index |
| 775 | ] |
| 776 | self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[ |
| 777 | last_req_index |
| 778 | ] |
| 779 | self.block_table.move_row(last_req_index, empty_index) |
| 780 | |
| 781 | self.request_lora_mapping[empty_index] = self.request_lora_mapping[ |
| 782 | last_req_index |
| 783 | ] |
| 784 | |
| 785 | if self.is_pooling_model: |
| 786 | last_req_index -= 1 |
| 787 | # Sampling state not used by pooling models. |
| 788 | continue |
| 789 | |
| 790 | # Autoregressive models require detailed tracking of condense |
| 791 | # operations to support logitsprocs |
| 792 | self.batch_update_builder.moved.append( |
| 793 | (last_req_index, empty_index, MoveDirectionality.UNIDIRECTIONAL) |
| 794 | ) |
| 795 | |
| 796 | self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index] |
| 797 | self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index] |
| 798 | self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index] |
| 799 | self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[ |
| 800 | last_req_index |
| 801 | ] |
| 802 | self.presence_penalties_cpu[empty_index] = self.presence_penalties_cpu[ |
| 803 | last_req_index |
| 804 | ] |
| 805 | self.repetition_penalties_cpu[empty_index] = self.repetition_penalties_cpu[ |
| 806 | last_req_index |
| 807 | ] |
| 808 | self.num_accepted_tokens_cpu[empty_index] = self.num_accepted_tokens_cpu[ |
| 809 | last_req_index |
| 810 | ] |
| 811 | generator = self.generators.pop(last_req_index, None) |
| 812 | if generator is not None: |
| 813 | self.generators[empty_index] = generator |
| 814 | |
| 815 | # TODO convert these to LogitsProcessors |
| 816 | if self.allowed_token_ids_mask_cpu_tensor is not None: |
| 817 | self.allowed_token_ids_mask_cpu_tensor[empty_index] = ( |
| 818 | self.allowed_token_ids_mask_cpu_tensor[last_req_index] |
| 819 | ) |
| 820 | self.allowed_token_ids_mask_cpu_tensor[last_req_index].fill_(False) |
| 821 | |
| 822 | bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None) |
| 823 | if bad_words_token_ids is not None: |
| 824 | self.bad_words_token_ids[empty_index] = bad_words_token_ids |
| 825 | |
| 826 | # Decrement last_req_index since it is now empty. |
| 827 | last_req_index -= 1 |
| 828 | |
| 829 | # Trim lists to the batch size. |
| 830 | del self._req_ids[num_reqs:] |
| 831 | del self.req_output_token_ids[num_reqs:] |
| 832 | del self.spec_token_ids[num_reqs:] |
| 833 | |
| 834 | def refresh_metadata(self): |
| 835 | """Apply any batch updates to sampling metadata.""" |
| 836 | if self.is_pooling_model: |
| 837 | batch_changed = self.batch_update_builder.reset() |
| 838 | if batch_changed: |
| 839 | self.sampling_metadata = self._make_sampling_metadata() |
| 840 | return |
| 841 | |
| 842 | # For non-pooling models - generate and apply logitsprocs update; |
| 843 | # reset batch update tracking. |
| 844 | # Update sampling metadata if batch state is changed. |
| 845 | batch_update = self.batch_update_builder.get_and_reset(self.num_reqs) |
| 846 | if self.thinking_budget_state_holder is not None and batch_update: |
| 847 | self.thinking_budget_state_holder.sync_batch(batch_update) |
| 848 | for logit_proc in self.logitsprocs.all: |
| 849 | logit_proc.update_state(batch_update) |
| 850 | if batch_update: |
| 851 | self.sampling_metadata = self._make_sampling_metadata() |
| 852 | |
| 853 | def _make_sampling_metadata(self) -> SamplingMetadata: |
| 854 | num_reqs = self.num_reqs |
| 855 | if not self.all_greedy: |
| 856 | temperature = copy_slice( |
| 857 | self.temperature_cpu_tensor, self.temperature, num_reqs |
| 858 | ) |
| 859 | else: |
| 860 | temperature = None |
| 861 | if not self.no_top_p: |
| 862 | copy_slice(self.top_p_cpu_tensor, self.top_p, num_reqs) |
| 863 | if not self.no_top_k: |
| 864 | copy_slice(self.top_k_cpu_tensor, self.top_k, num_reqs) |
| 865 | |
| 866 | if not self.no_penalties: |
| 867 | # Since syncing these tensors is expensive only copy them |
| 868 | # if necessary i.e. if there are requests which require |
| 869 | # penalties to be applied during sampling. |
| 870 | copy_slice( |
| 871 | self.frequency_penalties_cpu_tensor, self.frequency_penalties, num_reqs |
| 872 | ) |
| 873 | copy_slice( |
| 874 | self.presence_penalties_cpu_tensor, self.presence_penalties, num_reqs |
| 875 | ) |
| 876 | copy_slice( |
| 877 | self.repetition_penalties_cpu_tensor, |
| 878 | self.repetition_penalties, |
| 879 | num_reqs, |
| 880 | ) |
| 881 | |
| 882 | needs_prompt_token_ids = ( |
| 883 | not self.no_penalties |
| 884 | or self.logits_processing_needs_token_ids[:num_reqs].any() |
| 885 | ) |
| 886 | # The device prompt tokens are used only for applying penalties or |
| 887 | # pooling methods that explicitly request GPU token IDs. |
| 888 | # Hence copy these tensors only when there are requests which |
| 889 | # need penalties/step_pooler to be applied. |
| 890 | prompt_token_ids_cpu = ( |
| 891 | self._make_prompt_token_ids_cpu_tensor() if needs_prompt_token_ids else None |
| 892 | ) |
| 893 | prompt_token_ids = ( |
| 894 | prompt_token_ids_cpu.to(device=self.device, non_blocking=True) |
| 895 | if prompt_token_ids_cpu is not None |
| 896 | else None |
| 897 | ) |
| 898 | |
| 899 | # Only set output_token_ids if required by the current requests' |
| 900 | # sampling parameters. |
| 901 | holder = self.thinking_budget_state_holder |
| 902 | thinking_budget_tracks_reqs = ( |
| 903 | holder is not None and holder.has_tracked_requests() |
| 904 | ) |
| 905 | needs_output_token_ids = ( |
| 906 | not self.no_penalties |
| 907 | or bool(self.bad_words_token_ids) |
| 908 | or self.logitsprocs_need_output_token_ids |
| 909 | or thinking_budget_tracks_reqs |
| 910 | ) |
| 911 | output_token_ids = ( |
| 912 | cast(list[list[int]], self.req_output_token_ids) |
| 913 | if needs_output_token_ids |
| 914 | else [] |
| 915 | ) |
| 916 | |
| 917 | allowed_token_ids_mask: torch.Tensor | None = None |
| 918 | if not self.no_allowed_token_ids: |
| 919 | assert self.allowed_token_ids_mask is not None |
| 920 | copy_slice( |
| 921 | self.allowed_token_ids_mask_cpu_tensor, |
| 922 | self.allowed_token_ids_mask, |
| 923 | num_reqs, |
| 924 | ) |
| 925 | allowed_token_ids_mask = self.allowed_token_ids_mask[:num_reqs] |
| 926 | |
| 927 | # Build per-request logprob_token_ids mapping: req_index -> token_ids |
| 928 | logprob_token_ids_by_index: dict[int, list[int]] | None = None |
| 929 | if self.logprob_token_ids: |
| 930 | logprob_token_ids_by_index = {} |
| 931 | for req_id, token_ids in self.logprob_token_ids.items(): |
| 932 | if req_id in self.req_id_to_index: |
| 933 | req_index = self.req_id_to_index[req_id] |
| 934 | logprob_token_ids_by_index[req_index] = token_ids |
| 935 | |
| 936 | return SamplingMetadata( |
| 937 | temperature=temperature, |
| 938 | all_greedy=self.all_greedy, |
| 939 | all_random=self.all_random, |
| 940 | top_p=None if self.no_top_p else self.top_p[:num_reqs], |
| 941 | top_k=None if self.no_top_k else self.top_k[:num_reqs], |
| 942 | generators=self.generators, |
| 943 | max_num_logprobs=self.max_num_logprobs, |
| 944 | logprob_token_ids=logprob_token_ids_by_index, |
| 945 | prompt_token_ids=prompt_token_ids, |
| 946 | frequency_penalties=self.frequency_penalties[:num_reqs], |
| 947 | presence_penalties=self.presence_penalties[:num_reqs], |
| 948 | repetition_penalties=self.repetition_penalties[:num_reqs], |
| 949 | output_token_ids=output_token_ids, |
| 950 | spec_token_ids=self.spec_token_ids, |
| 951 | no_penalties=self.no_penalties, |
| 952 | allowed_token_ids_mask=allowed_token_ids_mask, |
| 953 | bad_words_token_ids=self.bad_words_token_ids, |
| 954 | logitsprocs=self.logitsprocs, |
| 955 | thinking_budget_state_holder=self.thinking_budget_state_holder, |
| 956 | ) |
| 957 | |
| 958 | def get_pooling_params(self) -> list[PoolingParams]: |
| 959 | assert len(self.req_ids) == len(self.pooling_params) |
| 960 | return [self.pooling_params[req_id] for req_id in self.req_ids] |
| 961 | |
| 962 | def get_pooling_states(self) -> list[PoolingStates]: |
| 963 | assert len(self.req_ids) == len(self.pooling_states) |
| 964 | return [self.pooling_states[req_id] for req_id in self.req_ids] |
| 965 | |
| 966 | def get_pooling_metadata(self) -> PoolingMetadata: |
| 967 | pooling_params = self.get_pooling_params() |
| 968 | pooling_states = self.get_pooling_states() |
| 969 | prompt_token_ids_cpu = None |
| 970 | if any(p.requires_token_ids for p in pooling_params): |
| 971 | prompt_token_ids_cpu = self._make_prompt_token_ids_cpu_tensor() |
| 972 | |
| 973 | return PoolingMetadata( |
| 974 | prompt_lens=self.num_prompt_tokens_cpu_tensor[: self.num_reqs].clone(), |
| 975 | prompt_token_ids=None, |
| 976 | prompt_token_ids_cpu=prompt_token_ids_cpu, |
| 977 | pooling_params=pooling_params, |
| 978 | pooling_states=pooling_states, |
| 979 | ) |
| 980 | |
| 981 | def _make_prompt_token_ids_cpu_tensor(self) -> torch.Tensor: |
| 982 | num_reqs = self.num_reqs |
| 983 | max_prompt_len = self.num_prompt_tokens[:num_reqs].max() |
| 984 | prompt_token_ids_cpu_tensor = torch.empty( |
| 985 | (self.num_reqs, max_prompt_len), |
| 986 | device="cpu", |
| 987 | dtype=torch.int64, |
| 988 | pin_memory=PIN_MEMORY, |
| 989 | ) |
| 990 | prompt_token_ids = prompt_token_ids_cpu_tensor.numpy() |
| 991 | prompt_token_ids[:] = self.token_ids_cpu[:num_reqs, :max_prompt_len] |
| 992 | # Use the value of vocab_size as a pad since we don't have a |
| 993 | # token_id of this value. |
| 994 | for i in range(num_reqs): |
| 995 | prompt_token_ids[i, self.num_prompt_tokens[i] :] = self.vocab_size |
| 996 | return prompt_token_ids_cpu_tensor |
| 997 | |
| 998 | def make_lora_inputs( |
| 999 | self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray |
| 1000 | ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: |
| 1001 | """Given the num_scheduled_tokens for each request in the batch, return |
| 1002 | datastructures used to activate the current LoRAs. |
| 1003 | |
| 1004 | Returns: |
| 1005 | 1. prompt_lora_mapping: A tuple of size np.sum(num_sampled_tokens) |
| 1006 | where, prompt_lora_mapping[i] is the LoRA id to use for the ith |
| 1007 | sampled token. |
| 1008 | 2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens) |
| 1009 | where, token_lora_mapping[i] is the LoRA id to use for ith |
| 1010 | token. |
| 1011 | 3. lora_requests: Set of relevant LoRA requests. |
| 1012 | |
| 1013 | """ |
| 1014 | req_lora_mapping = self.request_lora_mapping[: self.num_reqs] |
| 1015 | prompt_lora_mapping = tuple(req_lora_mapping.repeat(num_sampled_tokens)) |
| 1016 | token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens)) |
| 1017 | |
| 1018 | active_lora_requests: set[LoRARequest] = set( |
| 1019 | self.lora_id_to_lora_request.values() |
| 1020 | ) |
| 1021 | |
| 1022 | return prompt_lora_mapping, token_lora_mapping, active_lora_requests |
| 1023 | |
| 1024 | def set_async_sampled_token_ids( |
| 1025 | self, |
| 1026 | sampled_token_ids_cpu: torch.Tensor, |
| 1027 | async_copy_ready_event: torch.Event, |
| 1028 | ) -> None: |
| 1029 | """In async scheduling case, store ref to sampled_token_ids_cpu |
| 1030 | tensor and corresponding copy-ready event. Used to repair |
| 1031 | output_token_ids prior to sampling, if needed by logits processors. |
| 1032 | """ |
| 1033 | if self.sampling_metadata.output_token_ids: |
| 1034 | self.sampled_token_ids_cpu = sampled_token_ids_cpu |
| 1035 | self.async_copy_ready_event = async_copy_ready_event |
| 1036 | else: |
| 1037 | self.sampled_token_ids_cpu = None |
| 1038 | self.async_copy_ready_event = None |
| 1039 | |
| 1040 | def update_async_output_token_ids(self) -> None: |
| 1041 | """In async scheduling case, update output_token_ids in sampling metadata |
| 1042 | from prior steps sampled token ids once they've finished copying to CPU. |
| 1043 | This is called right before they are needed by the logits processors. |
| 1044 | """ |
| 1045 | output_token_ids = self.sampling_metadata.output_token_ids |
| 1046 | if self.sampled_token_ids_cpu is None or not output_token_ids: |
| 1047 | # Output token ids not needed or not async scheduling. |
| 1048 | return |
| 1049 | |
| 1050 | assert self.prev_req_id_to_index is not None |
| 1051 | sampled_token_ids = None |
| 1052 | for index, req_id in enumerate(self.req_ids): |
| 1053 | prev_index = self.prev_req_id_to_index.get(req_id) |
| 1054 | if prev_index is None: |
| 1055 | continue |
| 1056 | req_output_token_ids = output_token_ids[index] |
| 1057 | if not req_output_token_ids or req_output_token_ids[-1] != -1: |
| 1058 | # Final output id is not a placeholder, some tokens must have |
| 1059 | # been discarded after a kv-load failure. |
| 1060 | continue |
| 1061 | if sampled_token_ids is None: |
| 1062 | assert self.async_copy_ready_event is not None |
| 1063 | self.async_copy_ready_event.synchronize() |
| 1064 | sampled_token_ids = self.sampled_token_ids_cpu.tolist() |
| 1065 | # Replace placeholder token id(s) with actual sampled id(s). |
| 1066 | new_ids: list[int] = sampled_token_ids[prev_index] |
| 1067 | if not new_ids: |
| 1068 | continue |
| 1069 | num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1) |
| 1070 | # Also account for case where there may be a smaller number of |
| 1071 | # output placeholders (tokens can be discarded after kv-load |
| 1072 | # failure) or a larger number (async spec decode adds optimistic |
| 1073 | # placeholders that may exceed the actual acceptance count). |
| 1074 | first_placeholder = len(req_output_token_ids) |
| 1075 | while ( |
| 1076 | first_placeholder > 0 |
| 1077 | and req_output_token_ids[first_placeholder - 1] == -1 |
| 1078 | ): |
| 1079 | first_placeholder -= 1 |
| 1080 | num_placeholders = len(req_output_token_ids) - first_placeholder |
| 1081 | num_to_replace = min(num_sampled_ids, num_placeholders) |
| 1082 | del new_ids[num_to_replace:] |
| 1083 | req_output_token_ids[first_placeholder:] = new_ids |
| 1084 | # ^ Implicitly resizes to (first_placeholder + num_to_replace) |
| 1085 | |
| 1086 | def update_async_spec_token_ids(self, draft_token_ids: list[list[int]]) -> None: |
| 1087 | """In async scheduling case, update spec_token_ids in sampling metadata with |
| 1088 | real draft token ids from prior step. This is called right before they are |
| 1089 | needed by the rejection sampler for penalty/bad_words computation. |
| 1090 | """ |
| 1091 | if not draft_token_ids or not self.prev_req_id_to_index: |
| 1092 | return |
| 1093 | |
| 1094 | if (spec_token_ids := self.sampling_metadata.spec_token_ids) is not None: |
| 1095 | for req_id, spec_ids in zip(self.req_ids, spec_token_ids): |
| 1096 | if spec_ids: |
| 1097 | prev_index = self.prev_req_id_to_index.get(req_id) |
| 1098 | if prev_index is not None: |
| 1099 | draft_ids = draft_token_ids[prev_index] |
| 1100 | if draft_ids: |
| 1101 | del draft_ids[len(spec_ids) :] |
| 1102 | spec_ids.clear() |
| 1103 | spec_ids.extend(draft_ids) |
| 1104 | |
| 1105 | @property |
| 1106 | def num_reqs(self) -> int: |
| 1107 | return len(self.req_id_to_index) |
| 1108 | |
| 1109 | @property |
| 1110 | def all_greedy(self) -> bool: |
| 1111 | return len(self.random_reqs) == 0 |
| 1112 | |
| 1113 | @property |
| 1114 | def all_random(self) -> bool: |
| 1115 | return len(self.greedy_reqs) == 0 |
| 1116 | |
| 1117 | @property |
| 1118 | def no_top_p(self) -> bool: |
| 1119 | return len(self.top_p_reqs) == 0 |
| 1120 | |
| 1121 | @property |
| 1122 | def no_top_k(self) -> bool: |
| 1123 | return len(self.top_k_reqs) == 0 |
| 1124 | |
| 1125 | @property |
| 1126 | def no_penalties(self) -> bool: |
| 1127 | return ( |
| 1128 | len(self.presence_penalties_reqs) == 0 |
| 1129 | and len(self.frequency_penalties_reqs) == 0 |
| 1130 | and len(self.repetition_penalties_reqs) == 0 |
| 1131 | ) |
| 1132 | |
| 1133 | @property |
| 1134 | def no_thinking_budget(self) -> bool: |
| 1135 | return ( |
| 1136 | self.thinking_budget_state_holder is None |
| 1137 | or len(self.thinking_token_budget_reqs) == 0 |
| 1138 | ) |
| 1139 | |
| 1140 | @property |
| 1141 | def max_num_logprobs(self) -> int | None: |
| 1142 | return max(self.num_logprobs.values()) if self.num_logprobs else None |
| 1143 | |
| 1144 | @property |
| 1145 | def no_allowed_token_ids(self) -> bool: |
| 1146 | return len(self.has_allowed_token_ids) == 0 |