minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Refactor] Move output logprob processing into the logprob_processor layer (#31624)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-17 22:33:52 -07:00
committed by GitHub
parent 19c53c44a0
commit 639261f7b2
2 changed files with 226 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/logprob_processor.py
+214 -1
View File
@@ -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)
python/sglang/srt/layers/sampler.py
+12 -162
View File
@@ -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,