[Refactor] Move output logprob processing into the logprob_processor layer (#31624)
This commit is contained in:
Liangsheng Yin
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user