# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Datastructures defining a GPU input batch from dataclasses import dataclass from typing import cast import numpy as np import torch from vllm.config.reasoning import ReasoningConfig from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams, SamplingType from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.collection_utils import swap_dict_values from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.outputs import LogprobsTensors from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates from vllm.v1.sample.logits_processor import ( BatchUpdateBuilder, LogitsProcessors, MoveDirectionality, ) from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.thinking_budget_state import ( maybe_create_thinking_budget_state_holder, ) from vllm.v1.utils import copy_slice from vllm.v1.worker.block_table import MultiGroupBlockTable, SlotMappingMode @dataclass class CachedRequestState: req_id: str prompt_token_ids: list[int] | None mm_features: list[MultiModalFeatureSpec] sampling_params: SamplingParams | None generator: torch.Generator | None block_ids: tuple[list[int], ...] num_computed_tokens: int output_token_ids: list[int] mrope_positions: torch.Tensor | None = None mrope_position_delta: int | None = None xdrope_positions: torch.Tensor | None = None lora_request: LoRARequest | None = None prompt_embeds: torch.Tensor | None = None # To accumulate prompt logprobs tensor chunks across prefill steps. in_progress_prompt_logprobs_cpu: LogprobsTensors | None = None # Per-position mask for mixed-mode inputs (e.g chat completion with # prompt_embeds content parts). See `Request.prompt_is_token_ids`. prompt_is_token_ids: list[bool] | None = None # Used when both async_scheduling and spec_decode are enabled. prev_num_draft_len: int = 0 # for pooling models pooling_params: PoolingParams | None = None pooling_states: PoolingStates | None = None def __post_init__(self): self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds( self.prompt_token_ids, self.prompt_embeds ) if self.pooling_params is not None: self.pooling_states = PoolingStates() @property def num_tokens(self) -> int: return self.num_prompt_tokens + len(self.output_token_ids) def get_token_id(self, idx: int) -> int: if idx < self.num_prompt_tokens: if self.prompt_token_ids is None: raise ValueError( f"Tried to access token index {idx}, but that token was " "provided via prompt_embeds, and its ID is unknown." ) return self.prompt_token_ids[idx] if idx - self.num_prompt_tokens < len(self.output_token_ids): return self.output_token_ids[idx - self.num_prompt_tokens] return -1 class InputBatch: def __init__( self, max_num_reqs: int, max_model_len: int, max_num_batched_tokens: int, device: torch.device, vocab_size: int, block_sizes: list[int], # The block_size of each kv cache group kernel_block_sizes: list[int], max_num_blocks_per_req: list[int], logitsprocs: LogitsProcessors | None = None, logitsprocs_need_output_token_ids: bool = False, num_spec_tokens: int = 0, is_pooling_model: bool = False, cp_kv_cache_interleave_size: int = 1, reasoning_config: ReasoningConfig | None = None, slot_mapping_modes: list[SlotMappingMode] | None = None, ): self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder( reasoning_config, max_num_reqs, num_spec_tokens, device, PIN_MEMORY, ) self.thinking_token_budget_reqs: set[str] = set() self.is_pooling_model = is_pooling_model self.max_num_reqs = max_num_reqs self.max_model_len = max_model_len self.max_num_batched_tokens = max_num_batched_tokens self.device = device self.vocab_size = vocab_size self._req_ids: list[str | None] = [] self.req_id_to_index: dict[str, int] = {} # TODO(woosuk): This buffer could be too large if max_model_len is big. # Find a way to reduce the CPU memory usage. # This buffer is not directly transferred to the GPU, so it does not # need to be pinned. self.token_ids_cpu_tensor = torch.zeros( (max_num_reqs, max_model_len), device="cpu", dtype=torch.int32, pin_memory=False, ) self.token_ids_cpu = self.token_ids_cpu_tensor.numpy() self.is_token_ids_tensor = torch.zeros( (max_num_reqs, max_model_len), device="cpu", dtype=bool, pin_memory=False, ) self.is_token_ids = self.is_token_ids_tensor.numpy() # Store prompt embeddings per request to avoid OOM from large upfront # allocation if max_model_len is big. # Maps req_index -> tensor of shape (num_prompt_tokens, hidden_size) self.req_prompt_embeds: dict[int, torch.Tensor] = {} self.num_tokens_no_spec_cpu_tensor = torch.zeros( (max_num_reqs,), device="cpu", dtype=torch.int32, pin_memory=PIN_MEMORY, ) self.num_tokens_no_spec = self.num_tokens_no_spec_cpu_tensor.numpy() self.num_prompt_tokens_cpu_tensor = torch.zeros( (max_num_reqs,), device="cpu", dtype=torch.int32, pin_memory=PIN_MEMORY, ) self.num_prompt_tokens = self.num_prompt_tokens_cpu_tensor.numpy() self.num_computed_tokens_cpu_tensor = torch.zeros( (max_num_reqs,), device="cpu", dtype=torch.int32, pin_memory=PIN_MEMORY, ) self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy() # Block table. self.block_table = MultiGroupBlockTable( max_num_reqs=max_num_reqs, max_num_batched_tokens=max_num_batched_tokens, pin_memory=PIN_MEMORY, device=device, block_sizes=block_sizes, kernel_block_sizes=kernel_block_sizes, max_num_blocks=max_num_blocks_per_req, cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, slot_mapping_modes=slot_mapping_modes, ) # Sampling-related. self.temperature = torch.empty( (max_num_reqs,), dtype=torch.float32, device=device ) self.temperature_cpu_tensor = torch.empty( (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=PIN_MEMORY ) self.temperature_cpu = self.temperature_cpu_tensor.numpy() self.greedy_reqs: set[str] = set() self.random_reqs: set[str] = set() self.top_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) self.top_p_cpu_tensor = torch.empty( (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=PIN_MEMORY ) self.top_p_cpu = self.top_p_cpu_tensor.numpy() self.top_p_reqs: set[str] = set() self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device) self.top_k_cpu_tensor = torch.empty( (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY ) self.top_k_cpu = self.top_k_cpu_tensor.numpy() self.top_k_reqs: set[str] = set() # Frequency penalty related data structures self.frequency_penalties = torch.empty( (max_num_reqs,), dtype=torch.float, device=device ) self.frequency_penalties_cpu_tensor = torch.empty( (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY ) self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy() self.frequency_penalties_reqs: set[str] = set() # Presence penalty related data structures self.presence_penalties = torch.empty( (max_num_reqs,), dtype=torch.float, device=device ) self.presence_penalties_cpu_tensor = torch.empty( (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY ) self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy() self.presence_penalties_reqs: set[str] = set() # Repetition penalty related data structures self.repetition_penalties = torch.empty( (max_num_reqs,), dtype=torch.float, device=device ) self.repetition_penalties_cpu_tensor = torch.empty( (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY ) self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy() self.repetition_penalties_reqs: set[str] = set() # Speculative decoding self.num_accepted_tokens_cpu_tensor = torch.ones( (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY ) self.num_accepted_tokens_cpu = self.num_accepted_tokens_cpu_tensor.numpy() # lora related self.request_lora_mapping = np.zeros((self.max_num_reqs,), dtype=np.int64) self.lora_id_to_request_ids: dict[int, set[str]] = {} self.lora_id_to_lora_request: dict[int, LoRARequest] = {} # req_index -> generator # NOTE(woosuk): The indices of the requests that do not have their own # generator should not be included in the dictionary. self.generators: dict[int, torch.Generator] = {} self.num_logprobs: dict[str, int] = {} # req_id -> list of specific token IDs to compute logprobs for # More efficient than num_logprobs=-1 when only a few tokens are needed self.logprob_token_ids: dict[str, list[int]] = {} # Internal representation of per-step batch state changes, used for # reordering persistent batch and generating logitsprocs batch state # updates. Should reset each step. self.batch_update_builder = BatchUpdateBuilder() # TODO convert this to LogitsProcessor self.has_allowed_token_ids: set[str] = set() # NOTE(lufang): In the mask tensor, if the corresponding token allowed, # the value is False. Since we use masked_fill_ to set -inf. self.allowed_token_ids_mask: torch.Tensor | None = None self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None # req_index -> bad_words_token_ids self.bad_words_token_ids: dict[int, list[list[int]]] = {} self.logits_processing_needs_token_ids = np.zeros(max_num_reqs, dtype=bool) self.req_output_token_ids: list[list[int] | None] = [] # Store provided logitsprocs. If none are provided, initialize empty # data structure self.logitsprocs = logitsprocs or LogitsProcessors() self.logitsprocs_need_output_token_ids = logitsprocs_need_output_token_ids # Store last speculative tokens for sampler. self.spec_token_ids: list[list[int]] = [[] for _ in range(max_num_reqs)] # This is updated each time the batch constituents change. self.sampling_metadata = self._make_sampling_metadata() # for pooling models self.pooling_params: dict[str, PoolingParams] = {} self.pooling_states: dict[str, PoolingStates] = {} # Cached reference to the GPU tensor of previously sampled tokens self.prev_sampled_token_ids: torch.Tensor | None = None self.prev_req_id_to_index: dict[str, int] | None = None # These are used to update output_token_ids with real sampled # ids from prior step, if required by current sampling params # (e.g. penalties). self.sampled_token_ids_cpu: torch.Tensor | None = None self.async_copy_ready_event: torch.Event | None = None @property def req_ids(self) -> list[str]: # None elements should only be present transiently # while performing state updates to the batch. return cast(list[str], self._req_ids) def _register_add_request(self, request: "CachedRequestState") -> int: """Track add-request operations for logits processors. Not applicable to pooling models. """ # Fill the next empty index if there is one. if (new_req_index := self.batch_update_builder.pop_removed()) is None: # Append to end otherwise. new_req_index = self.num_reqs assert new_req_index < self.max_num_reqs self.batch_update_builder.batch_changed = True if request.sampling_params: # Detailed added request metadata is only required for non-pooling # models, to support logitsprocs. self.batch_update_builder.added.append( ( new_req_index, request.sampling_params, request.prompt_token_ids, request.output_token_ids, ) ) return new_req_index def add_request( self, request: "CachedRequestState", ) -> int: req_index = self._register_add_request(request) req_id = request.req_id if req_index == len(self._req_ids): self._req_ids.append(req_id) self.req_output_token_ids.append(request.output_token_ids) self.spec_token_ids.append([]) else: self._req_ids[req_index] = req_id self.req_output_token_ids[req_index] = request.output_token_ids self.spec_token_ids[req_index].clear() self.req_id_to_index[req_id] = req_index # Copy the prompt token ids and output token ids. num_prompt_tokens = length_from_prompt_token_ids_or_embeds( request.prompt_token_ids, request.prompt_embeds ) self.num_prompt_tokens[req_index] = num_prompt_tokens start_idx = num_prompt_tokens end_idx = start_idx + len(request.output_token_ids) if request.prompt_token_ids is not None: self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids if request.prompt_is_token_ids is not None: self.is_token_ids[req_index, :num_prompt_tokens] = ( request.prompt_is_token_ids ) else: self.is_token_ids[req_index, :num_prompt_tokens] = True else: self.is_token_ids[req_index, :num_prompt_tokens] = False if request.prompt_embeds is not None: self.req_prompt_embeds[req_index] = request.prompt_embeds self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids self.is_token_ids[req_index, start_idx:end_idx] = True # Number of tokens without spec decode tokens. self.num_tokens_no_spec[req_index] = request.num_tokens self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens self.block_table.add_row(request.block_ids, req_index) if sampling_params := request.sampling_params: if sampling_params.sampling_type == SamplingType.GREEDY: # Should avoid division by zero later when apply_temperature. self.temperature_cpu[req_index] = 0.0 self.greedy_reqs.add(req_id) else: self.temperature_cpu[req_index] = sampling_params.temperature self.random_reqs.add(req_id) self.top_p_cpu[req_index] = sampling_params.top_p if sampling_params.top_p < 1: self.top_p_reqs.add(req_id) top_k = sampling_params.top_k if 0 < top_k < self.vocab_size: self.top_k_reqs.add(req_id) else: top_k = self.vocab_size self.top_k_cpu[req_index] = top_k self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty if sampling_params.frequency_penalty != 0.0: self.frequency_penalties_reqs.add(req_id) self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty if sampling_params.presence_penalty != 0.0: self.presence_penalties_reqs.add(req_id) self.repetition_penalties_cpu[req_index] = ( sampling_params.repetition_penalty ) if sampling_params.repetition_penalty != 1.0: self.repetition_penalties_reqs.add(req_id) # NOTE(woosuk): self.generators should not include the requests that # do not have their own generator. if request.generator is not None: self.generators[req_index] = request.generator if sampling_params.logprobs is not None: self.num_logprobs[req_id] = ( self.vocab_size if sampling_params.logprobs == -1 else sampling_params.logprobs ) # Store specific token IDs to compute logprobs for (more efficient) if sampling_params.logprob_token_ids is not None: self.logprob_token_ids[req_id] = sampling_params.logprob_token_ids if sampling_params.allowed_token_ids: self.has_allowed_token_ids.add(req_id) if self.allowed_token_ids_mask_cpu_tensor is None: # Lazy allocation for this tensor, which can be large. # False means we don't fill with -inf. self.allowed_token_ids_mask = torch.zeros( self.max_num_reqs, self.vocab_size, dtype=torch.bool, device=self.device, ) self.allowed_token_ids_mask_cpu_tensor = torch.zeros( self.max_num_reqs, self.vocab_size, dtype=torch.bool, device="cpu", ) self.allowed_token_ids_mask_cpu_tensor[req_index] = True # False means we don't fill with -inf. self.allowed_token_ids_mask_cpu_tensor[req_index][ sampling_params.allowed_token_ids ] = False if sampling_params.bad_words_token_ids: self.bad_words_token_ids[req_index] = ( sampling_params.bad_words_token_ids ) elif pooling_params := request.pooling_params: pooling_states = request.pooling_states assert pooling_states is not None self.pooling_params[req_id] = pooling_params self.pooling_states[req_id] = pooling_states self.logits_processing_needs_token_ids[req_index] = ( pooling_params.requires_token_ids ) else: raise NotImplementedError("Unrecognized request type") # Speculative decoding: by default 1 token is generated. self.num_accepted_tokens_cpu[req_index] = 1 # Add request lora ID if request.lora_request: lora_id = request.lora_request.lora_int_id if lora_id not in self.lora_id_to_request_ids: self.lora_id_to_request_ids[lora_id] = set() self.request_lora_mapping[req_index] = lora_id self.lora_id_to_request_ids[lora_id].add(request.req_id) self.lora_id_to_lora_request[lora_id] = request.lora_request else: # No LoRA self.request_lora_mapping[req_index] = 0 return req_index def update_req_spec_token_ids( self, request: CachedRequestState, scheduled_spec_tokens: dict[str, list[int]] ) -> None: req_id = request.req_id req_index = self.req_id_to_index[req_id] cur_spec_token_ids = self.spec_token_ids[req_index] # When speculative decoding is used with structured output, # the scheduler can drop draft tokens that do not # conform to the schema. This can result in # scheduler_output.scheduled_spec_decode_tokens being empty, # even when speculative decoding is enabled. cur_spec_token_ids.clear() spec_token_ids = scheduled_spec_tokens.get(req_id, ()) num_spec_tokens = len(spec_token_ids) request.prev_num_draft_len = num_spec_tokens if not spec_token_ids: return # For async scheduling, token_ids_cpu assigned from # spec_token_ids are placeholders and will be overwritten in # _prepare_input_ids. start_index = self.num_tokens_no_spec[req_index] end_token_index = start_index + num_spec_tokens self.token_ids_cpu[req_index, start_index:end_token_index] = spec_token_ids self.is_token_ids[req_index, start_index:end_token_index] = True cur_spec_token_ids.extend(spec_token_ids) def remove_request(self, req_id: str) -> int | None: """This method must always be followed by a call to condense(). Args: req_id: request to remove Returns: Removed request index, or `None` if `req_id` not recognized """ req_index = self.req_id_to_index.pop(req_id, None) if req_index is None: return None self.batch_update_builder.removed_append(req_index) self._req_ids[req_index] = None self.req_output_token_ids[req_index] = None self.spec_token_ids[req_index].clear() self.block_table.clear_row(req_index) # LoRA lora_id = self.request_lora_mapping[req_index] if lora_id != 0: lora_req_ids = self.lora_id_to_request_ids[lora_id] lora_req_ids.discard(req_id) if not lora_req_ids: del self.lora_id_to_request_ids[lora_id] del self.lora_id_to_lora_request[lora_id] self.request_lora_mapping[req_index] = 0 if self.is_pooling_model: self.pooling_params.pop(req_id, None) self.pooling_states.pop(req_id, None) return req_index self.greedy_reqs.discard(req_id) self.random_reqs.discard(req_id) self.top_p_reqs.discard(req_id) self.top_k_reqs.discard(req_id) self.frequency_penalties_reqs.discard(req_id) self.presence_penalties_reqs.discard(req_id) self.repetition_penalties_reqs.discard(req_id) self.generators.pop(req_index, None) self.num_logprobs.pop(req_id, None) self.logprob_token_ids.pop(req_id, None) if self.prev_req_id_to_index is not None: self.prev_req_id_to_index.pop(req_id, None) self.has_allowed_token_ids.discard(req_id) if self.allowed_token_ids_mask_cpu_tensor is not None: # False means we don't fill with -inf. self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False) self.bad_words_token_ids.pop(req_index, None) self.thinking_token_budget_reqs.discard(req_id) return req_index def swap_states(self, i1: int, i2: int) -> None: old_id_i1 = self._req_ids[i1] old_id_i2 = self._req_ids[i2] # Only swap the active token prefix for each request. Copying full # max_model_len rows is expensive and unnecessary during reordering. i1_active_token_count = self._get_active_token_count(i1) i2_active_token_count = self._get_active_token_count(i2) max_active_token_count = max(i1_active_token_count, i2_active_token_count) self._req_ids[i1], self._req_ids[i2] = self._req_ids[i2], self._req_ids[i1] # noqa self.req_output_token_ids[i1], self.req_output_token_ids[i2] = ( self.req_output_token_ids[i2], self.req_output_token_ids[i1], ) self.spec_token_ids[i1], self.spec_token_ids[i2] = ( self.spec_token_ids[i2], self.spec_token_ids[i1], ) assert old_id_i1 is not None and old_id_i2 is not None self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] = ( self.req_id_to_index[old_id_i2], self.req_id_to_index[old_id_i1], ) self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = ( self.num_tokens_no_spec[i2], self.num_tokens_no_spec[i1], ) self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] = ( self.num_prompt_tokens[i2], self.num_prompt_tokens[i1], ) self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = ( self.num_computed_tokens_cpu[i2], self.num_computed_tokens_cpu[i1], ) # NOTE: the following is unsafe # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\ # self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...] # instead, we need to temporarily copy the data for one of the indices tmp_token_ids = self.token_ids_cpu[i1, :max_active_token_count].copy() self.token_ids_cpu[i1, :max_active_token_count] = self.token_ids_cpu[ i2, :max_active_token_count ] self.token_ids_cpu[i2, :max_active_token_count] = tmp_token_ids self.is_token_ids[[i1, i2], :max_active_token_count] = self.is_token_ids[ [i2, i1], :max_active_token_count ] # Swap prompt embeddings if they exist embeds_i1 = self.req_prompt_embeds.get(i1) embeds_i2 = self.req_prompt_embeds.get(i2) if embeds_i1 is not None: self.req_prompt_embeds[i2] = embeds_i1 else: self.req_prompt_embeds.pop(i2, None) if embeds_i2 is not None: self.req_prompt_embeds[i1] = embeds_i2 else: self.req_prompt_embeds.pop(i1, None) self.block_table.swap_row(i1, i2) self.request_lora_mapping[i1], self.request_lora_mapping[i2] = ( self.request_lora_mapping[i2], self.request_lora_mapping[i1], ) if self.is_pooling_model: # Sampling and logits parameters don't apply to pooling models. return # For autoregressive models, track detailed request reordering info # to support logitsprocs. self.batch_update_builder.moved.append((i1, i2, MoveDirectionality.SWAP)) self.temperature_cpu[i1], self.temperature_cpu[i2] = ( self.temperature_cpu[i2], self.temperature_cpu[i1], ) self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1] self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1] self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = ( self.frequency_penalties_cpu[i2], self.frequency_penalties_cpu[i1], ) self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] = ( self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1], ) self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = ( self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1], ) self.num_accepted_tokens_cpu[i1], self.num_accepted_tokens_cpu[i2] = ( self.num_accepted_tokens_cpu[i2], self.num_accepted_tokens_cpu[i1], ) swap_dict_values(self.generators, i1, i2) swap_dict_values(self.bad_words_token_ids, i1, i2) if self.allowed_token_ids_mask_cpu_tensor is not None: ( self.allowed_token_ids_mask_cpu_tensor[i1], self.allowed_token_ids_mask_cpu_tensor[i2], ) = ( self.allowed_token_ids_mask_cpu_tensor[i2], self.allowed_token_ids_mask_cpu_tensor[i1], ) def _get_active_token_count(self, req_index: int) -> int: return int(self.num_tokens_no_spec[req_index]) + len( self.spec_token_ids[req_index] ) def condense(self) -> None: """Slide non-empty requests down into lower, empty indices. Any consecutive empty indices at the very end of the list are not filled. Returns: swaps: list of (from,to) swap tuples for moved requests empty_req_indices: indices not filled by condensation """ num_reqs = self.num_reqs if not (empty_req_indices := self.batch_update_builder.removed): # All removed requests were replaced by added requests, or else no # requests were removed at all. No condense() needed return if num_reqs == 0: # The batched states are empty. self._req_ids.clear() self.req_output_token_ids.clear() self.spec_token_ids.clear() return # NOTE(woosuk): This function assumes that the empty_req_indices # is sorted in descending order. last_req_index = num_reqs + len(empty_req_indices) - 1 while empty_req_indices: # Find the largest non-empty index. while last_req_index in empty_req_indices: last_req_index -= 1 # Find the smallest empty index. empty_index = self.batch_update_builder.peek_removed() assert empty_index is not None if empty_index >= last_req_index: break # Move active request down into empty request # index. self.batch_update_builder.pop_removed() req_id = self._req_ids[last_req_index] output_token_ids = self.req_output_token_ids[last_req_index] assert req_id is not None self._req_ids[empty_index] = req_id self._req_ids[last_req_index] = None self.req_output_token_ids[empty_index] = output_token_ids self.req_output_token_ids[last_req_index] = None self.req_id_to_index[req_id] = empty_index num_tokens = self._get_active_token_count(last_req_index) (self.spec_token_ids[last_req_index], self.spec_token_ids[empty_index]) = ( self.spec_token_ids[empty_index], self.spec_token_ids[last_req_index], ) self.spec_token_ids[last_req_index].clear() self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[ last_req_index, :num_tokens ] self.is_token_ids[empty_index, :num_tokens] = self.is_token_ids[ last_req_index, :num_tokens ] if last_req_index in self.req_prompt_embeds: self.req_prompt_embeds[empty_index] = self.req_prompt_embeds.pop( last_req_index ) self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[ last_req_index ] self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index] self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[ last_req_index ] self.block_table.move_row(last_req_index, empty_index) self.request_lora_mapping[empty_index] = self.request_lora_mapping[ last_req_index ] if self.is_pooling_model: last_req_index -= 1 # Sampling state not used by pooling models. continue # Autoregressive models require detailed tracking of condense # operations to support logitsprocs self.batch_update_builder.moved.append( (last_req_index, empty_index, MoveDirectionality.UNIDIRECTIONAL) ) self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index] self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index] self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index] self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[ last_req_index ] self.presence_penalties_cpu[empty_index] = self.presence_penalties_cpu[ last_req_index ] self.repetition_penalties_cpu[empty_index] = self.repetition_penalties_cpu[ last_req_index ] self.num_accepted_tokens_cpu[empty_index] = self.num_accepted_tokens_cpu[ last_req_index ] generator = self.generators.pop(last_req_index, None) if generator is not None: self.generators[empty_index] = generator # TODO convert these to LogitsProcessors if self.allowed_token_ids_mask_cpu_tensor is not None: self.allowed_token_ids_mask_cpu_tensor[empty_index] = ( self.allowed_token_ids_mask_cpu_tensor[last_req_index] ) bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None) if bad_words_token_ids is not None: self.bad_words_token_ids[empty_index] = bad_words_token_ids # Decrement last_req_index since it is now empty. last_req_index -= 1 # Trim lists to the batch size. del self._req_ids[num_reqs:] del self.req_output_token_ids[num_reqs:] del self.spec_token_ids[num_reqs:] def refresh_metadata(self): """Apply any batch updates to sampling metadata.""" if self.is_pooling_model: batch_changed = self.batch_update_builder.reset() if batch_changed: self.sampling_metadata = self._make_sampling_metadata() return # For non-pooling models - generate and apply logitsprocs update; # reset batch update tracking. # Update sampling metadata if batch state is changed. batch_update = self.batch_update_builder.get_and_reset(self.num_reqs) if self.thinking_budget_state_holder is not None and batch_update: self.thinking_budget_state_holder.sync_batch(batch_update) for logit_proc in self.logitsprocs.all: logit_proc.update_state(batch_update) if batch_update: self.sampling_metadata = self._make_sampling_metadata() def _make_sampling_metadata(self) -> SamplingMetadata: num_reqs = self.num_reqs if not self.all_greedy: temperature = copy_slice( self.temperature_cpu_tensor, self.temperature, num_reqs ) else: temperature = None if not self.no_top_p: copy_slice(self.top_p_cpu_tensor, self.top_p, num_reqs) if not self.no_top_k: copy_slice(self.top_k_cpu_tensor, self.top_k, num_reqs) if not self.no_penalties: # Since syncing these tensors is expensive only copy them # if necessary i.e. if there are requests which require # penalties to be applied during sampling. copy_slice( self.frequency_penalties_cpu_tensor, self.frequency_penalties, num_reqs ) copy_slice( self.presence_penalties_cpu_tensor, self.presence_penalties, num_reqs ) copy_slice( self.repetition_penalties_cpu_tensor, self.repetition_penalties, num_reqs, ) needs_prompt_token_ids = ( not self.no_penalties or self.logits_processing_needs_token_ids[:num_reqs].any() ) # The prompt tokens are used only for applying penalties or # step pooling during the sampling/pooling process. # Hence copy these tensors only when there are requests which # need penalties/step_pooler to be applied. prompt_token_ids_cpu = ( self._make_prompt_token_ids_cpu_tensor() if needs_prompt_token_ids else None ) prompt_token_ids = ( prompt_token_ids_cpu.to(device=self.device, non_blocking=True) if prompt_token_ids_cpu is not None else None ) # Only set output_token_ids if required by the current requests' # sampling parameters. holder = self.thinking_budget_state_holder thinking_budget_tracks_reqs = ( holder is not None and holder.has_tracked_requests() ) needs_output_token_ids = ( not self.no_penalties or bool(self.bad_words_token_ids) or self.logitsprocs_need_output_token_ids or thinking_budget_tracks_reqs ) output_token_ids = ( cast(list[list[int]], self.req_output_token_ids) if needs_output_token_ids else [] ) allowed_token_ids_mask: torch.Tensor | None = None if not self.no_allowed_token_ids: assert self.allowed_token_ids_mask is not None copy_slice( self.allowed_token_ids_mask_cpu_tensor, self.allowed_token_ids_mask, num_reqs, ) allowed_token_ids_mask = self.allowed_token_ids_mask[:num_reqs] # Build per-request logprob_token_ids mapping: req_index -> token_ids logprob_token_ids_by_index: dict[int, list[int]] | None = None if self.logprob_token_ids: logprob_token_ids_by_index = {} for req_id, token_ids in self.logprob_token_ids.items(): if req_id in self.req_id_to_index: req_index = self.req_id_to_index[req_id] logprob_token_ids_by_index[req_index] = token_ids return SamplingMetadata( temperature=temperature, all_greedy=self.all_greedy, all_random=self.all_random, top_p=None if self.no_top_p else self.top_p[:num_reqs], top_k=None if self.no_top_k else self.top_k[:num_reqs], generators=self.generators, max_num_logprobs=self.max_num_logprobs, logprob_token_ids=logprob_token_ids_by_index, prompt_token_ids=prompt_token_ids, frequency_penalties=self.frequency_penalties[:num_reqs], presence_penalties=self.presence_penalties[:num_reqs], repetition_penalties=self.repetition_penalties[:num_reqs], output_token_ids=output_token_ids, spec_token_ids=self.spec_token_ids, no_penalties=self.no_penalties, allowed_token_ids_mask=allowed_token_ids_mask, bad_words_token_ids=self.bad_words_token_ids, logitsprocs=self.logitsprocs, thinking_budget_state_holder=self.thinking_budget_state_holder, ) def get_pooling_params(self) -> list[PoolingParams]: assert len(self.req_ids) == len(self.pooling_params) return [self.pooling_params[req_id] for req_id in self.req_ids] def get_pooling_states(self) -> list[PoolingStates]: assert len(self.req_ids) == len(self.pooling_states) return [self.pooling_states[req_id] for req_id in self.req_ids] def get_pooling_metadata(self) -> PoolingMetadata: pooling_params = self.get_pooling_params() pooling_states = self.get_pooling_states() prompt_token_ids_cpu = None if any(p.requires_token_ids for p in pooling_params): prompt_token_ids_cpu = self._make_prompt_token_ids_cpu_tensor() return PoolingMetadata( prompt_lens=self.num_prompt_tokens_cpu_tensor[: self.num_reqs].clone(), prompt_token_ids=self.sampling_metadata.prompt_token_ids, prompt_token_ids_cpu=prompt_token_ids_cpu, pooling_params=pooling_params, pooling_states=pooling_states, ) def _make_prompt_token_ids_cpu_tensor(self) -> torch.Tensor: num_reqs = self.num_reqs max_prompt_len = self.num_prompt_tokens[:num_reqs].max() prompt_token_ids_cpu_tensor = torch.empty( (self.num_reqs, max_prompt_len), device="cpu", dtype=torch.int64, pin_memory=PIN_MEMORY, ) prompt_token_ids = prompt_token_ids_cpu_tensor.numpy() prompt_token_ids[:] = self.token_ids_cpu[:num_reqs, :max_prompt_len] # Use the value of vocab_size as a pad since we don't have a # token_id of this value. for i in range(num_reqs): prompt_token_ids[i, self.num_prompt_tokens[i] :] = self.vocab_size return prompt_token_ids_cpu_tensor def make_lora_inputs( self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: """ Given the num_scheduled_tokens for each request in the batch, return datastructures used to activate the current LoRAs. Returns: 1. prompt_lora_mapping: A tuple of size np.sum(num_sampled_tokens) where, prompt_lora_mapping[i] is the LoRA id to use for the ith sampled token. 2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens) where, token_lora_mapping[i] is the LoRA id to use for ith token. 3. lora_requests: Set of relevant LoRA requests. """ req_lora_mapping = self.request_lora_mapping[: self.num_reqs] prompt_lora_mapping = tuple(req_lora_mapping.repeat(num_sampled_tokens)) token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens)) active_lora_requests: set[LoRARequest] = set( self.lora_id_to_lora_request.values() ) return prompt_lora_mapping, token_lora_mapping, active_lora_requests def set_async_sampled_token_ids( self, sampled_token_ids_cpu: torch.Tensor, async_copy_ready_event: torch.Event, ) -> None: """ In async scheduling case, store ref to sampled_token_ids_cpu tensor and corresponding copy-ready event. Used to repair output_token_ids prior to sampling, if needed by logits processors. """ if self.sampling_metadata.output_token_ids: self.sampled_token_ids_cpu = sampled_token_ids_cpu self.async_copy_ready_event = async_copy_ready_event else: self.sampled_token_ids_cpu = None self.async_copy_ready_event = None def update_async_output_token_ids(self) -> None: """ In async scheduling case, update output_token_ids in sampling metadata from prior steps sampled token ids once they've finished copying to CPU. This is called right before they are needed by the logits processors. """ output_token_ids = self.sampling_metadata.output_token_ids if self.sampled_token_ids_cpu is None or not output_token_ids: # Output token ids not needed or not async scheduling. return assert self.prev_req_id_to_index is not None sampled_token_ids = None for index, req_id in enumerate(self.req_ids): prev_index = self.prev_req_id_to_index.get(req_id) if prev_index is None: continue req_output_token_ids = output_token_ids[index] if not req_output_token_ids or req_output_token_ids[-1] != -1: # Final output id is not a placeholder, some tokens must have # been discarded after a kv-load failure. continue if sampled_token_ids is None: assert self.async_copy_ready_event is not None self.async_copy_ready_event.synchronize() sampled_token_ids = self.sampled_token_ids_cpu.tolist() # Replace placeholder token id(s) with actual sampled id(s). new_ids: list[int] = sampled_token_ids[prev_index] if not new_ids: continue num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1) # Also account for case where there may be a smaller number of # output placeholders (tokens can be discarded after kv-load # failure) or a larger number (async spec decode adds optimistic # placeholders that may exceed the actual acceptance count). first_placeholder = len(req_output_token_ids) while ( first_placeholder > 0 and req_output_token_ids[first_placeholder - 1] == -1 ): first_placeholder -= 1 num_placeholders = len(req_output_token_ids) - first_placeholder num_to_replace = min(num_sampled_ids, num_placeholders) del new_ids[num_to_replace:] req_output_token_ids[first_placeholder:] = new_ids # ^ Implicitly resizes to (first_placeholder + num_to_replace) def update_async_spec_token_ids(self, draft_token_ids: list[list[int]]) -> None: """ In async scheduling case, update spec_token_ids in sampling metadata with real draft token ids from prior step. This is called right before they are needed by the rejection sampler for penalty/bad_words computation. """ if not draft_token_ids or not self.prev_req_id_to_index: return if (spec_token_ids := self.sampling_metadata.spec_token_ids) is not None: for req_id, spec_ids in zip(self.req_ids, spec_token_ids): if spec_ids: prev_index = self.prev_req_id_to_index.get(req_id) if prev_index is not None: draft_ids = draft_token_ids[prev_index] if draft_ids: del draft_ids[len(spec_ids) :] spec_ids.clear() spec_ids.extend(draft_ids) @property def num_reqs(self) -> int: return len(self.req_id_to_index) @property def all_greedy(self) -> bool: return len(self.random_reqs) == 0 @property def all_random(self) -> bool: return len(self.greedy_reqs) == 0 @property def no_top_p(self) -> bool: return len(self.top_p_reqs) == 0 @property def no_top_k(self) -> bool: return len(self.top_k_reqs) == 0 @property def no_penalties(self) -> bool: return ( len(self.presence_penalties_reqs) == 0 and len(self.frequency_penalties_reqs) == 0 and len(self.repetition_penalties_reqs) == 0 ) @property def no_thinking_budget(self) -> bool: return ( self.thinking_budget_state_holder is None or len(self.thinking_token_budget_reqs) == 0 ) @property def max_num_logprobs(self) -> int | None: return max(self.num_logprobs.values()) if self.num_logprobs else None @property def no_allowed_token_ids(self) -> bool: return len(self.has_allowed_token_ids) == 0