import logging from typing import Callable, Dict, List, Optional, Tuple import torch import torch.distributed as dist from torch import nn from sglang.srt.distributed import get_tp_group from sglang.srt.layers.dp_attention import ( get_attention_tp_group, is_dp_attention_enabled, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.utils.hash import murmur_hash32 from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logprobs from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_params import TOP_K_ALL from sglang.srt.server_args import get_global_server_args from sglang.srt.utils.common import crash_on_warnings, get_bool_env_var, is_cuda, is_npu if is_cuda(): from sgl_kernel import ( min_p_sampling_from_probs, top_k_renorm_prob, top_k_top_p_sampling_from_probs, top_p_renorm_prob, ) if is_npu(): import torch_npu logger = logging.getLogger(__name__) SYNC_TOKEN_IDS_ACROSS_TP = get_bool_env_var("SYNC_TOKEN_IDS_ACROSS_TP") SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB") _CUSTOM_SAMPLER_FACTORIES: Dict[str, Callable[[], "Sampler"]] = {} _BUILT_IN_SAMPLING_BACKENDS = {"flashinfer", "pytorch", "ascend"} class Sampler(nn.Module): def __init__(self): super().__init__() self.use_nan_detection = get_global_server_args().enable_nan_detection self.tp_sync_group = get_tp_group().device_group if is_dp_attention_enabled(): self.tp_sync_group = get_attention_tp_group().device_group self.rl_on_policy_target = get_global_server_args().rl_on_policy_target # In RL on-policy mode, deterministic inference is automatically enabled. self.enable_deterministic = ( get_global_server_args().enable_deterministic_inference ) # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer. self.use_log_softmax_logprob = self.rl_on_policy_target is not None self.use_ascend_backend = get_global_server_args().sampling_backend == "ascend" def _preprocess_logits( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo ) -> torch.Tensor: """Apply custom logit processors and handle NaN detection.""" # Apply the custom logit processors if registered in the sampling info if sampling_info.has_custom_logit_processor: apply_custom_logit_processor(logits, sampling_info) # Detect and handle NaN values in logits if self.use_nan_detection and torch.any(torch.isnan(logits)): logger.warning("Detected errors during sampling! NaN in the logits.") logits = torch.where( torch.isnan(logits), torch.full_like(logits, -1e5), logits ) if crash_on_warnings(): raise ValueError("Detected errors during sampling! NaN in the logits.") return logits def forward( self, logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo, return_logprob: bool, top_logprobs_nums: List[int], token_ids_logprobs: List[List[int]], positions: torch.Tensor, ): """Run a sampler & compute logprobs and update logits_output accordingly. Args: logits_output: The logits from the model forward sampling_info: Metadata for sampling return_logprob: If set, store the output logprob information to logits_output top_logprobs_nums: Number of top lobprobs per sequence in a batch token_ids_logprobs: Per-sequence list of specific token IDs to retrieve logprobs for. Each element is a list of token IDs (or None) for one sequence in the batch. This is used in speculative decoding. positions: The positions of the tokens in the sequence. Used for deterministic sampling to get the unique seed for each position. """ logits = logits_output.next_token_logits # Preprocess logits (custom processors and NaN handling) logits = self._preprocess_logits(logits, sampling_info) if sampling_info.is_all_greedy: # Use torch.argmax if all requests use greedy sampling batch_next_token_ids = torch.argmax(logits, -1) if return_logprob: original_logprobs = logprobs = torch.nn.functional.log_softmax( logits, dim=-1 ) else: simple_sampling_case = ( not sampling_info.need_top_p_sampling and not sampling_info.need_top_k_sampling and not sampling_info.need_min_p_sampling ) # If requested, cache original logprobs before temperature scaling. if return_logprob and SGLANG_RETURN_ORIGINAL_LOGPROB: original_logprobs = torch.log_softmax(logits, dim=-1) # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer. logprobs_via_logsoftmax_kernel = None if self.rl_on_policy_target is not None: # TODO: use more inplace ops to save memory logits_div_temperature = ( logits.bfloat16().div(sampling_info.temperatures).bfloat16() ) logprobs_via_logsoftmax_kernel = torch.log_softmax( logits_div_temperature, dim=-1 ) del logits_div_temperature if self.use_ascend_backend: # Ascend backend: sample from logits directly. batch_next_token_ids, logprobs = self._forward_ascend_backend( logits, sampling_info, simple_sampling_case, return_logprob ) elif ( self.use_log_softmax_logprob and self.enable_deterministic and simple_sampling_case ): # RL on-policy path: sample from logprobs to match the trainer. batch_next_token_ids = self._sample_from_logprobs( logprobs_via_logsoftmax_kernel, sampling_info, positions, ) if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB: logprobs = logprobs_via_logsoftmax_kernel else: # Standard path: do softmax and sample from probs. logits.div_(sampling_info.temperatures) # In-place op to save memory logits[:] = torch.softmax(logits, dim=-1) probs = logits batch_next_token_ids = self._sample_from_probs( probs, sampling_info, positions, simple_sampling_case ) if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB: logprobs = ( logprobs_via_logsoftmax_kernel if logprobs_via_logsoftmax_kernel is not None else torch.log(probs) ) del probs # Attach logprobs to logits_output (in-place modification) if return_logprob: if SGLANG_RETURN_ORIGINAL_LOGPROB: logprobs = original_logprobs self._attach_logprobs_to_output( logits_output, logprobs, top_logprobs_nums, token_ids_logprobs, sampling_info, batch_next_token_ids, ) self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info) return batch_next_token_ids def _sample_from_probs( self, probs: torch.Tensor, sampling_info: SamplingBatchInfo, positions: torch.Tensor, simple_sampling_case: bool, ) -> torch.Tensor: """Sample from probability distribution (after softmax). Used for standard sampling with flashinfer/pytorch backends. Handles both simple (direct multinomial) and complex (top-k/top-p/min-p) cases. """ if simple_sampling_case: batch_next_token_ids = sampling_from_probs_torch( probs, sampling_seed=sampling_info.sampling_seed, positions=positions, ) else: backend = get_global_server_args().sampling_backend if backend == "flashinfer": assert ( sampling_info.sampling_seed is None ), "Sampling seed is not supported for flashinfer backend" if sampling_info.need_min_p_sampling: probs = top_k_renorm_prob(probs, sampling_info.top_ks) probs = top_p_renorm_prob(probs, sampling_info.top_ps) batch_next_token_ids = min_p_sampling_from_probs( probs, sampling_info.min_ps ) else: batch_next_token_ids = top_k_top_p_sampling_from_probs( probs.contiguous(), sampling_info.top_ks, sampling_info.top_ps, filter_apply_order="joint", check_nan=self.use_nan_detection, ) elif backend == "pytorch": # A slower fallback implementation with torch native operations. batch_next_token_ids = top_k_top_p_min_p_sampling_from_probs_torch( probs, sampling_info.top_ks, sampling_info.top_ps, sampling_info.min_ps, sampling_info.need_min_p_sampling, sampling_info.sampling_seed, positions, ) else: raise ValueError(f"Invalid sampling backend: {backend}") return batch_next_token_ids def _sample_from_logprobs( self, logprobs: torch.Tensor, sampling_info: SamplingBatchInfo, positions: torch.Tensor, ) -> torch.Tensor: """Sample from log-probabilities using the Gumbel trick. Used for deterministic sampling with simple cases (no top-k/top-p/min-p). Requires sampling_seed to be set in sampling_info. """ assert ( sampling_info.sampling_seed is not None ), "sampling_seed is required for sampling from logprobs" sampled_index = multinomial_with_seed( logprobs, sampling_info.sampling_seed, positions ) return sampled_index.view(-1).to(torch.int32) def _sample_from_logits( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo, simple_sampling_case: bool, ) -> torch.Tensor: """Sample from temperature-scaled logits without softmax. Used for the Ascend NPU backend which handles softmax internally. """ if simple_sampling_case: probs = torch.softmax(logits, dim=-1) batch_next_token_ids = torch.multinomial(probs, num_samples=1).view(-1) return batch_next_token_ids.to(torch.int32) else: assert ( self.use_ascend_backend ), "Only ascend backend supports sampling from logits" batch_next_token_ids = top_k_top_p_min_p_sampling_from_logits_ascend( logits, sampling_info.top_ks, sampling_info.top_ps, sampling_info.min_ps, sampling_info.need_min_p_sampling, ) return batch_next_token_ids.to(torch.int32) def _forward_ascend_backend( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo, simple_sampling_case: bool, return_logprob: bool, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Handle the full Ascend backend sampling path. Ascend backend has fused kernels that handle softmax internally, so we sample directly from temperature-scaled logits. Returns: A tuple of (batch_next_token_ids, logprobs). logprobs is None when return_logprob is False or SGLANG_RETURN_ORIGINAL_LOGPROB is set. """ logits.div_(sampling_info.temperatures) batch_next_token_ids = self._sample_from_logits( logits, sampling_info, simple_sampling_case ) logprobs = None if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB: logprobs = torch.log_softmax(logits, dim=-1) return batch_next_token_ids, logprobs def _attach_logprobs_to_output( self, logits_output: LogitsProcessorOutput, logprobs: torch.Tensor, top_logprobs_nums: List[int], token_ids_logprobs: List[List[int]], sampling_info: SamplingBatchInfo, batch_next_token_ids: torch.Tensor, ): # clamp to avoid -inf values logprobs.clamp_(min=torch.finfo(logprobs.dtype).min) # Attach logprobs to logits_output (in-place modification) if any(x > 0 for x in top_logprobs_nums): ( logits_output.next_token_top_logprobs_val, logits_output.next_token_top_logprobs_idx, ) = get_top_logprobs(logprobs, top_logprobs_nums) if any(x is not None for x in token_ids_logprobs): ( logits_output.next_token_token_ids_logprobs_val, logits_output.next_token_token_ids_logprobs_idx, ) = get_token_ids_logprobs(logprobs, token_ids_logprobs) logits_output.next_token_logprobs = logprobs[ torch.arange(len(batch_next_token_ids), device=sampling_info.device), batch_next_token_ids, ] def _sync_token_ids_across_tp( self, batch_next_token_ids: torch.Tensor, sampling_info: SamplingBatchInfo ): if SYNC_TOKEN_IDS_ACROSS_TP or sampling_info.grammars: # For performance reasons, SGLang does not sync the final token IDs across TP ranks by default. # This saves one all-reduce, but the correctness of this approach depends on the determinism of several operators: # the last all-reduce, the last lm_head matmul, and all sampling kernels. # These kernels are deterministic in most cases, but there are some rare instances where they are not deterministic. # In such cases, enable this env variable to prevent hanging due to TP ranks becoming desynchronized. # When using xgrammar, this becomes more likely so we also do the sync when grammar is used. torch.distributed.all_reduce( batch_next_token_ids, op=dist.ReduceOp.MIN, group=self.tp_sync_group, ) def compute_logprobs_only( self, logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo, return_logprob: bool, top_logprobs_nums: List[int], token_ids_logprobs: List[List[int]], ) -> None: """ Compute logprobs for requested token IDs without performing sampling. Optimized for prefill-only scoring requests that need token probabilities but don't require next token generation. """ if logits_output.next_token_logits is None: logger.warning("No logits available for logprob computation") return # Check if any requests actually need logprobs computation needs_token_ids_logprobs = any( token_ids is not None and len(token_ids) > 0 for token_ids in token_ids_logprobs ) needs_top_logprobs = any(x > 0 for x in top_logprobs_nums) if not (needs_token_ids_logprobs or needs_top_logprobs): return # Preprocess logits (custom processors and NaN handling) logits = self._preprocess_logits(logits_output.next_token_logits, sampling_info) # Compute logprobs logprobs = torch.nn.functional.log_softmax(logits, dim=-1) # Handle top logprobs if requested if needs_top_logprobs: ( logits_output.next_token_top_logprobs_val, logits_output.next_token_top_logprobs_idx, ) = get_top_logprobs(logprobs, top_logprobs_nums) # Handle token_ids logprobs if requested if needs_token_ids_logprobs: ( logits_output.next_token_token_ids_logprobs_val, logits_output.next_token_token_ids_logprobs_idx, ) = get_token_ids_logprobs_batch_optimized(logprobs, token_ids_logprobs) def register_sampler_backend(backend: str, factory: Callable[[], "Sampler"]) -> None: """Register a custom sampler factory for a backend string.""" if not backend: raise ValueError("backend must be a non-empty string") from sglang.srt.server_args import SAMPLING_BACKEND_CHOICES if backend in _CUSTOM_SAMPLER_FACTORIES: logger.warning("Overriding existing sampler factory for backend '%s'", backend) SAMPLING_BACKEND_CHOICES.add(backend) _CUSTOM_SAMPLER_FACTORIES[backend] = factory def create_sampler(backend: Optional[str] = None) -> "Sampler": """Create a sampler honoring custom backend registrations.""" server_args = get_global_server_args() backend = backend or (server_args.sampling_backend if server_args else None) if backend in _CUSTOM_SAMPLER_FACTORIES: sampler = _CUSTOM_SAMPLER_FACTORIES[backend]() if not isinstance(sampler, Sampler): raise TypeError( f"Custom sampler factory for backend '{backend}' must return a Sampler" ) return sampler if backend is None or backend in _BUILT_IN_SAMPLING_BACKENDS: return Sampler() raise ValueError( f"Unknown sampling backend '{backend}'. Register it via register_sampler_backend()." ) def top_k_top_p_min_p_sampling_from_probs_torch( probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor, min_ps: torch.Tensor, need_min_p_sampling: bool, sampling_seed: Optional[torch.Tensor], positions: torch.Tensor, ): """ A top-k, top-p and min-p sampling implementation with native pytorch operations. When sampling_seed is not None, deterministic inference will be enabled, it will sample with the sampling_seed of each request. """ probs_sort, probs_idx = probs.sort(dim=-1, descending=True) probs_sum = torch.cumsum(probs_sort, dim=-1) probs_sort[ torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1) >= top_ks.view(-1, 1) ] = 0.0 probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0 if need_min_p_sampling: # TODO: probs_sort should be re-normalized for the use of multinomial_with_seed assert ( sampling_seed is None ), "With sampling seed, multinomial_with_seed will provide wrong results" min_p_thresholds = probs_sort[:, 0] * min_ps probs_sort[probs_sort < min_p_thresholds.view(-1, 1)] = 0.0 if sampling_seed is None: sampled_index = torch.multinomial(probs_sort, num_samples=1) else: # NOTE: when using top-k/top-p/min-p sampling, we need to modify probs before we # apply log to get logprobs. Therefore, we cannot use log_softmax directly. # For now, we use log to the modified probs to get logprobs, but for numerical # stability, we'd better come up with a solution to use log_softmax. logprobs = probs_sort.to(torch.float64) # Using float64 for numerical stability del probs_sort logprobs.log_() sampled_index = multinomial_with_seed(logprobs, sampling_seed, positions) # int32 range is enough to represent the token ids probs_idx = probs_idx.to(torch.int32) batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(-1) return batch_next_token_ids def top_k_top_p_min_p_sampling_from_logits_ascend( logits: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor, min_ps: torch.Tensor, need_min_p_sampling: bool, ): """A top-k, top-p and min-p sampling implementation for ascend npu with torch_npu interface. Takes temperature-scaled logits as input (softmax is applied internally). """ # torch_npu.npu_top_k_top_p requires top_k value range in [1, 1024] if hasattr(torch_npu, "npu_top_k_top_p") and torch.all( (top_ks <= 1024) & (top_ks >= 1) ): logits_top_k_top_p = torch_npu.npu_top_k_top_p(logits, top_ps, top_ks) probs_top_k_top_p = logits_top_k_top_p.softmax(dim=-1) if need_min_p_sampling: min_p_thresholds = probs_top_k_top_p.max(dim=-1) * min_ps min_p_mask = probs_top_k_top_p < min_p_thresholds.view(-1, 1) probs_top_k_top_p.masked_fill_(min_p_mask, 0.0) batch_next_token_ids = torch.multinomial(probs_top_k_top_p, num_samples=1) else: probs = torch.softmax(logits, dim=-1) probs_sort, probs_idx = probs.sort(dim=-1, descending=True) # when top_k is -1 (in which sglang turns it to TOP_K_ALL), make it explicitly equal to logit's size topk_all_mask = top_ks == TOP_K_ALL top_ks.masked_fill_(topk_all_mask, probs.shape[1]) top_k_mask = torch.arange(0, probs.shape[-1], device=probs.device).view( 1, -1 ) >= top_ks.view(-1, 1) probs_sort.masked_fill_(top_k_mask, 0.0) probs_sum = torch.cumsum(probs_sort, dim=-1) top_p_mask = probs_sum - probs_sort > top_ps.view(-1, 1) probs_sort.masked_fill_(top_p_mask, 0.0) if need_min_p_sampling: min_p_thresholds = probs_sort[:, 0] * min_ps min_p_mask = probs_sort < min_p_thresholds.view(-1, 1) probs_sort.masked_fill_(min_p_mask, 0.0) sampled_index = torch.multinomial(probs_sort, num_samples=1) probs_idx = probs_idx.to(torch.int32) batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index) return batch_next_token_ids.view(-1) @torch.compile(dynamic=True) def multinomial_with_seed( logprobs: torch.Tensor, seed: torch.Tensor, positions: torch.Tensor ) -> torch.Tensor: """ Samples n elements from an input tensor `inputs` of shape (n, m) using a unique random seed for each row. This is a deterministic batched alternative to `torch.multinomial`. Args: inputs: A float tensor of shape (n, m) representing n categorical distributions with m categories each. The values are treated as weights and do not need to sum to 1. seed: An integer tensor of shape (n,) containing the random seed for each corresponding row in `inputs`. positions: The positions of the tokens in the sequence. Used for deterministic sampling to get the unique seed for each position. Returns: A tensor of shape (n,) where the i-th element is an index sampled from the distribution in `inputs[i]` using `seed[i]`. """ n, m = logprobs.shape seed = seed.to(torch.uint64) col_indices = torch.arange(m, device=logprobs.device) hashed = murmur_hash32(seed, positions, col_indices) # NOTE (sehoon): it is critical to keep gumbel noise calculation in float64 to avoid numerical instability. # keeping logprobs in float64 is less critical, but we found it's still safer to keep it in float64. x = hashed.to(torch.float64) / torch.iinfo(torch.uint32).max # x is a uniform sample in [0, 1]. get gumbel noise from it. # which is equivalent to -log(-log(x)) # keep everything in in-place operations to avoid unnecessary memory allocations. x.log_().clamp_(min=torch.finfo(x.dtype).min).neg_() # -log(x) x.log_().neg_() # -log(-log(x)) == gumbel noise # add gumbel noise to logprobs x.add_(logprobs.to(torch.float64)) return torch.argmax(x, dim=1, keepdim=True) def sampling_from_probs_torch( probs: torch.Tensor, sampling_seed: Optional[torch.Tensor] = None, positions: Optional[torch.Tensor] = None, ): """A sampling implementation with native pytorch operations, without top-k, top-p, or min-p filtering. Note: For deterministic sampling from logprobs, use Sampler._sample_from_logprobs instead. """ if sampling_seed is None: sampled_index = torch.multinomial(probs, num_samples=1) else: # Deterministic sampling: convert probs to logprobs and use gumbel trick sampled_index = multinomial_with_seed( torch.log(probs), sampling_seed, positions ) batch_next_token_ids = sampled_index.view(-1).to(torch.int32) return batch_next_token_ids def top_p_normalize_probs_torch( probs: torch.Tensor, top_ps: torch.Tensor, ): # See also top_k_top_p_min_p_sampling_from_probs_torch probs_sort, probs_idx = probs.sort(dim=-1, descending=True) probs_sum = torch.cumsum(probs_sort, dim=-1) probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0 probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True)) return torch.zeros_like(probs_sort).scatter_(-1, probs_idx, probs_sort) def get_token_ids_logprobs_batch_optimized( logprobs: torch.Tensor, token_ids_logprobs: List[List[int]], ) -> Tuple[List, List]: """ Vectorized batch processing for token ID logprobs extraction. Uses a single GPU kernel call for the entire batch instead of multiple separate calls, significantly improving performance for large batches. Args: logprobs: Log probabilities tensor [batch_size, vocab_size] token_ids_logprobs: List of token IDs to extract logprobs for Example: # Input: batch_size=3, vocab_size=5 logprobs = torch.tensor([ [-1.2, -2.1, -0.8, -3.0, -1.5], # batch 0 [-0.5, -1.8, -2.2, -1.1, -2.7], # batch 1 [-2.0, -0.9, -1.4, -2.8, -1.6], # batch 2 ]) token_ids_logprobs = [[1, 3], [2], [0, 2, 4]] # Output: # values = [tensor([-2.1, -3.0]), tensor([-2.2]), tensor([-2.0, -1.4, -1.6])] # indices = [[1, 3], [2], [0, 2, 4]] """ batch_size = len(token_ids_logprobs) device = logprobs.device # Step 1: Calculate lengths for each request, treating None as empty list # Example: [[1, 3], [2], [0, 2, 4]] -> token_lengths = tensor([2, 1, 3]) token_lengths = torch.tensor( [len(token_ids or []) for token_ids in token_ids_logprobs], device=device ) total_tokens = int(token_lengths.sum().item()) # 2 + 1 + 3 = 6 # Handle edge case where no tokens are requested if total_tokens == 0: return [logprobs.new_empty(0) for _ in token_ids_logprobs], [ [] for _ in token_ids_logprobs ] # Step 2: Build flattened indices using torch operations # Example: row_indices = [0, 0, 1, 2, 2, 2] (batch indices repeated by their lengths) row_indices = torch.repeat_interleave( torch.arange(batch_size, device=device), token_lengths ) # Example: col_indices = [1, 3, 2, 0, 2, 4] (flattened token IDs from all requests) col_indices = torch.tensor( [ token_id for token_ids in token_ids_logprobs for token_id in (token_ids or []) ], device=device, dtype=torch.long, ) # Step 3: Single vectorized gather operation # Example: logprobs[row_indices, col_indices] -> [-2.1, -3.0, -2.2, -2.0, -1.4, -1.6] gathered_logprobs = logprobs[row_indices, col_indices] # Step 4: Split results back per request using torch operations # Example: split tensor [6] into chunks of sizes [2, 1, 3] -> [tensor(2), tensor(1), tensor(3)] split_logprobs = torch.split_with_sizes( gathered_logprobs, token_lengths.tolist(), dim=0 ) # Step 5: Format output to match expected return structure # Example: Convert split tensors back to list format with proper empty handling # i=0: [1,3] -> append split_logprobs[0] and [1,3] # i=1: [2] -> append split_logprobs[1] and [2] # i=2: [0,2,4] -> append split_logprobs[2] and [0,2,4] output_token_ids_logprobs_val = [] output_token_ids_logprobs_idx = [] for i, token_ids in enumerate(token_ids_logprobs): if token_ids is not None and len(token_ids) > 0: output_token_ids_logprobs_val.append(split_logprobs[i]) output_token_ids_logprobs_idx.append(token_ids) else: output_token_ids_logprobs_val.append(logprobs.new_empty(0)) output_token_ids_logprobs_idx.append([]) return output_token_ids_logprobs_val, output_token_ids_logprobs_idx def apply_custom_logit_processor( logits: torch.Tensor, sampling_batch_info: SamplingBatchInfo, num_tokens_in_batch: int = 1, ): """Apply custom logit processors to the logits. This function will modify the logits in-place. num_tokens_in_batch is needed to support spec decoding, where each batch can contain multiple tokens. By default, we assume each batch contains only 1 token. """ assert logits.shape[0] == len(sampling_batch_info) * num_tokens_in_batch, ( f"The batch size of logits ({logits.shape[0]}) does not match the batch size of " f"sampling_batch_info ({len(sampling_batch_info)}) x num_tokens_in_batch " f"({num_tokens_in_batch})" ) for _, ( processor, batch_mask, ) in sampling_batch_info.custom_logit_processor.items(): # Get the batch indices that need to be processed batch_indices = batch_mask.nonzero(as_tuple=True)[0] assert batch_mask.shape[0] == len(sampling_batch_info), ( f"The number of batch mask ({batch_mask.shape[0]}) does not match the number of " f"sampling_batch_info ({len(sampling_batch_info)})" ) batch_mask = torch.repeat_interleave(batch_mask, num_tokens_in_batch) # Apply the processor to the logits logits[batch_mask] = processor( logits[batch_mask], [sampling_batch_info.custom_params[i] for i in batch_indices], ) logger.debug( f"Custom logit processor {processor.__class__.__name__} is applied." )