diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py index e8bbacfd4..d78651707 100644 --- a/python/sglang/srt/layers/logprob_processor.py +++ b/python/sglang/srt/layers/logprob_processor.py @@ -1,6 +1,7 @@ from __future__ import annotations import dataclasses +import logging from enum import Enum, auto from typing import TYPE_CHECKING, Callable, List, Optional, Tuple @@ -9,8 +10,11 @@ import torch from sglang.srt.environ import envs if TYPE_CHECKING: - from sglang.srt.layers.logits_processor import LogitsMetadata + from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding + from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo + +logger = logging.getLogger(__name__) class LogprobStage(Enum): @@ -614,3 +618,212 @@ class InputLogprobProcessor: ), sampled_logits, ) + + +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 + + +@dataclasses.dataclass +class OutputLogprobsResult: + """Output-side counterpart of InputLogprobsResult. + + Built by OutputLogprobProcessor; write_to() flushes the populated fields + onto LogitsProcessorOutput, so the IPC / D2H wire format stays unchanged. + """ + + token_logprobs: Optional[torch.Tensor] = None + top_logprobs_val: Optional[List] = None + top_logprobs_idx: Optional[List] = None + token_ids_logprobs_val: Optional[List] = None + token_ids_logprobs_idx: Optional[List] = None + + def write_to(self, logits_output: LogitsProcessorOutput) -> None: + if self.token_logprobs is not None: + logits_output.next_token_logprobs = self.token_logprobs + if self.top_logprobs_val is not None: + logits_output.next_token_top_logprobs_val = self.top_logprobs_val + logits_output.next_token_top_logprobs_idx = self.top_logprobs_idx + if self.token_ids_logprobs_val is not None: + logits_output.next_token_token_ids_logprobs_val = ( + self.token_ids_logprobs_val + ) + logits_output.next_token_token_ids_logprobs_idx = ( + self.token_ids_logprobs_idx + ) + + +class OutputLogprobProcessor: + """Output (decode) logprob processing: logprobs -> topk / token-ids / + sampled-token gather, attached onto LogitsProcessorOutput. + + Only logits/logprobs are needed here; sampler-side concerns (custom + logit processors, NaN sanitizing) are injected via ``preprocess_fn``. + """ + + def attach_logprobs_to_output( + self, + logits_output: LogitsProcessorOutput, + logprobs: torch.Tensor, + top_logprobs_nums: List[int], + token_ids_logprobs: List[List[int]], + batch_next_token_ids: torch.Tensor, + ): + # clamp to avoid -inf values + logprobs.clamp_(min=torch.finfo(logprobs.dtype).min) + + result = OutputLogprobsResult() + if any(x > 0 for x in top_logprobs_nums): + ( + result.top_logprobs_val, + result.top_logprobs_idx, + ) = get_top_logprobs(logprobs, top_logprobs_nums, no_copy_to_cpu=True) + + if any(x is not None for x in token_ids_logprobs): + ( + result.token_ids_logprobs_val, + result.token_ids_logprobs_idx, + ) = get_token_ids_logprobs( + logprobs, token_ids_logprobs, no_copy_to_cpu=True + ) + + result.token_logprobs = logprobs[ + torch.arange(len(batch_next_token_ids), device=batch_next_token_ids.device), + batch_next_token_ids, + ] + result.write_to(logits_output) + + def compute_logprobs_only( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + top_logprobs_nums: List[int], + token_ids_logprobs: List[List[int]], + preprocess_fn: Callable, + ) -> 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 = preprocess_fn(logits_output.next_token_logits, sampling_info) + + # Compute logprobs + logprobs = torch.nn.functional.log_softmax(logits, dim=-1) + + result = OutputLogprobsResult() + # Handle top logprobs if requested + if needs_top_logprobs: + ( + result.top_logprobs_val, + result.top_logprobs_idx, + ) = get_top_logprobs(logprobs, top_logprobs_nums, no_copy_to_cpu=True) + + # Handle token_ids logprobs if requested + if needs_token_ids_logprobs: + ( + result.token_ids_logprobs_val, + result.token_ids_logprobs_idx, + ) = get_token_ids_logprobs_batch_optimized(logprobs, token_ids_logprobs) + result.write_to(logits_output) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 68280bfd3..b208eb69c 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -11,7 +11,9 @@ from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.layers.logprob_processor import get_token_ids_logprobs, get_top_logprobs +from sglang.srt.layers.logprob_processor import ( + OutputLogprobProcessor, +) from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_params import TOP_K_ALL @@ -78,6 +80,8 @@ class Sampler(nn.Module): self.use_log_softmax_logprob = self.rl_on_policy_target is not None self.use_ascend_backend = get_server_args().sampling_backend == "ascend" + self.output_logprob_processor = OutputLogprobProcessor() + def _preprocess_logits( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo ) -> torch.Tensor: @@ -210,12 +214,11 @@ class Sampler(nn.Module): if return_logprob: if SGLANG_RETURN_ORIGINAL_LOGPROB: logprobs = original_logprobs - self._attach_logprobs_to_output( + self.output_logprob_processor.attach_logprobs_to_output( logits_output, logprobs, top_logprobs_nums, token_ids_logprobs, - sampling_info, batch_next_token_ids, ) @@ -470,38 +473,6 @@ class Sampler(nn.Module): 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, no_copy_to_cpu=True) - - 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, no_copy_to_cpu=True - ) - - 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 ): @@ -527,46 +498,13 @@ class Sampler(nn.Module): 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 + self.output_logprob_processor.compute_logprobs_only( + logits_output=logits_output, + sampling_info=sampling_info, + top_logprobs_nums=top_logprobs_nums, + token_ids_logprobs=token_ids_logprobs, + preprocess_fn=self._preprocess_logits, ) - 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, no_copy_to_cpu=True) - - # 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: @@ -801,94 +739,6 @@ def top_p_normalize_probs_torch( 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,