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

[Refactor] Unify logprob results into a single LogprobResult and rename chunk env vars (#31733)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-20 12:44:27 -07:00
committed by GitHub
parent e149cdb337
commit ff6c755952
11 changed files with 139 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+7 -3
View File
@@ -834,9 +834,13 @@ class Envs:
# Sparse Embeddings # Sparse Embeddings
SGLANG_EMBEDDINGS_SPARSE_HEAD = EnvStr(None) SGLANG_EMBEDDINGS_SPARSE_HEAD = EnvStr(None)
# Logits processor # Logprob processor
SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK = EnvBool(True) SGLANG_ENABLE_LOGPROB_CHUNK = EnvBoolWithAlias(
SGLANG_LOGITS_PROCESSER_CHUNK_SIZE = EnvInt(2048) True, deprecated_name="SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"
)
SGLANG_LOGPROB_CHUNK_SIZE = EnvIntWithAlias(
2048, deprecated_name="SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"
)
# Tool-Call behavior # Tool-Call behavior
SGLANG_TOOL_STRICT_LEVEL = EnvInt(ToolStrictLevel.OFF) SGLANG_TOOL_STRICT_LEVEL = EnvInt(ToolStrictLevel.OFF)
python/sglang/srt/layers/logits_processor.py
+3 -6
View File
@@ -473,16 +473,13 @@ class LogitsProcessor(nn.Module):
skip_chunking_for_dp_attn=self.do_tensor_parallel_all_gather_dp_attn, skip_chunking_for_dp_attn=self.do_tensor_parallel_all_gather_dp_attn,
) )
return LogitsProcessorOutput( logits_output = LogitsProcessorOutput(
next_token_logits=sampled_logits, next_token_logits=sampled_logits,
hidden_states=hidden_states_to_store, hidden_states=hidden_states_to_store,
input_token_logprobs=logprobs_result.input_token_logprobs,
input_top_logprobs_val=logprobs_result.input_top_logprobs_val,
input_top_logprobs_idx=logprobs_result.input_top_logprobs_idx,
input_token_ids_logprobs_val=logprobs_result.input_token_ids_logprobs_val,
input_token_ids_logprobs_idx=logprobs_result.input_token_ids_logprobs_idx,
mm_input_embeds=logits_metadata.mm_input_embeds, mm_input_embeds=logits_metadata.mm_input_embeds,
) )
logprobs_result.write_input_to(logits_output)
return logits_output
def _get_pruned_states( def _get_pruned_states(
self, self,
python/sglang/srt/layers/logprob_processor.py
+105 -113
View File
@@ -12,7 +12,6 @@ from sglang.srt.environ import envs
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -23,12 +22,43 @@ class LogprobStage(Enum):
@dataclasses.dataclass @dataclasses.dataclass
class InputLogprobsResult: class LogprobResult:
input_token_logprobs: torch.Tensor """Logprob fields produced by Input/OutputLogprobProcessor.
input_top_logprobs_val: Optional[List] = None
input_top_logprobs_idx: Optional[List] = None Input (prefill) always fills token_logprobs; output (decode / scoring)
input_token_ids_logprobs_val: Optional[List] = None fills on demand. write_input_to / write_output_to flush populated fields
input_token_ids_logprobs_idx: Optional[List] = None 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_input_to(self, logits_output: LogitsProcessorOutput) -> None:
if self.token_logprobs is not None:
logits_output.input_token_logprobs = self.token_logprobs
if self.top_logprobs_val is not None:
logits_output.input_top_logprobs_val = self.top_logprobs_val
logits_output.input_top_logprobs_idx = self.top_logprobs_idx
if self.token_ids_logprobs_val is not None:
logits_output.input_token_ids_logprobs_val = self.token_ids_logprobs_val
logits_output.input_token_ids_logprobs_idx = self.token_ids_logprobs_idx
def write_output_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
)
def get_top_logprobs_raw( def get_top_logprobs_raw(
@@ -135,29 +165,27 @@ def get_token_ids_logprobs(logprobs, token_ids_logprobs, no_copy_to_cpu=False):
def get_top_logprobs_chunk( def get_top_logprobs_chunk(
logprobs: torch.Tensor, logprobs: torch.Tensor,
logits_metadata: LogitsMetadata,
top_k_nums: List[int], top_k_nums: List[int],
pruned_lens: List[int], pruned_lens: List[int],
input_top_logprobs_val: List, top_logprobs_val: List,
input_top_logprobs_idx: List, top_logprobs_idx: List,
split_pruned_len: int, split_pruned_len: int,
) -> int: ) -> int:
"""Get top-k logprobs for each sequence in the chunk. """Get top-k logprobs for each sequence in the chunk.
Args: Args:
logprobs: Log probabilities tensor of shape [seq_len, vocab_size] logprobs: Log probabilities tensor of shape [seq_len, vocab_size]
logits_metadata: Metadata containing top-k and pruned length info
top_k_nums: List of top-k numbers for each sequence top_k_nums: List of top-k numbers for each sequence
pruned_lens: List of pruned lengths for each sequence pruned_lens: List of pruned lengths for each sequence
input_top_logprobs_val: List to store top-k logprob values top_logprobs_val: List to store top-k logprob values
input_top_logprobs_idx: List to store top-k token indices top_logprobs_idx: List to store top-k token indices
split_pruned_len: Length of pruned tokens from previous chunk split_pruned_len: Length of pruned tokens from previous chunk
Returns: Returns:
int: Number of remaining tokens to process in next chunk int: Number of remaining tokens to process in next chunk
""" """
# Empty chunks still walk the slice to emit placeholder entries. # Empty chunks still walk the slice to emit placeholder entries.
max_k = max(logits_metadata.top_logprobs_nums) max_k = max(top_k_nums)
ret = logprobs.topk(max_k, dim=1) ret = logprobs.topk(max_k, dim=1)
values = ret.values.tolist() values = ret.values.tolist()
indices = ret.indices.tolist() indices = ret.indices.tolist()
@@ -175,8 +203,8 @@ def get_top_logprobs_chunk(
if pruned_len <= 0: if pruned_len <= 0:
# if pruned length is less than or equal to 0, # if pruned length is less than or equal to 0,
# there is no top-k logprobs to process # there is no top-k logprobs to process
input_top_logprobs_val.append([]) top_logprobs_val.append([])
input_top_logprobs_idx.append([]) top_logprobs_idx.append([])
continue continue
# Get the top-k logprobs # Get the top-k logprobs
@@ -195,11 +223,11 @@ def get_top_logprobs_chunk(
# Split-sequence continuations extend; everyone else owns a fresh # Split-sequence continuations extend; everyone else owns a fresh
# (possibly empty) entry. # (possibly empty) entry.
if split_pruned_len > 0: if split_pruned_len > 0:
input_top_logprobs_val[-1].extend(val) top_logprobs_val[-1].extend(val)
input_top_logprobs_idx[-1].extend(idx) top_logprobs_idx[-1].extend(idx)
else: else:
input_top_logprobs_val.append(val) top_logprobs_val.append(val)
input_top_logprobs_idx.append(idx) top_logprobs_idx.append(idx)
pt += pruned_len pt += pruned_len
return next_split_pruned_len return next_split_pruned_len
@@ -209,19 +237,18 @@ def get_token_ids_logprobs_chunk(
logprobs: torch.Tensor, logprobs: torch.Tensor,
token_ids_logprobs: List[int], token_ids_logprobs: List[int],
pruned_lens: List[int], pruned_lens: List[int],
input_token_ids_logprobs_val: List, token_ids_logprobs_val: List,
input_token_ids_logprobs_idx: List, token_ids_logprobs_idx: List,
split_pruned_len: int = 0, split_pruned_len: int = 0,
): ):
"""Get token_ids logprobs for each sequence in the chunk. """Get token_ids logprobs for each sequence in the chunk.
Args: Args:
logprobs: Log probabilities tensor of shape [seq_len, vocab_size] logprobs: Log probabilities tensor of shape [seq_len, vocab_size]
logits_metadata: Metadata containing token IDs and pruned length info
token_ids_logprobs: List of token IDs for each sequence token_ids_logprobs: List of token IDs for each sequence
pruned_lens: List of pruned lengths for each sequence pruned_lens: List of pruned lengths for each sequence
input_token_ids_logprobs_val: List to store token logprob values token_ids_logprobs_val: List to store token logprob values
input_token_ids_logprobs_idx: List to store token indices token_ids_logprobs_idx: List to store token indices
split_pruned_len: Length of pruned tokens from previous chunk split_pruned_len: Length of pruned tokens from previous chunk
Returns: Returns:
@@ -245,8 +272,8 @@ def get_token_ids_logprobs_chunk(
if pruned_len <= 0: if pruned_len <= 0:
# if pruned length is less than or equal to 0, # if pruned length is less than or equal to 0,
# there is no token ids logprobs to process # there is no token ids logprobs to process
input_token_ids_logprobs_val.append([]) token_ids_logprobs_val.append([])
input_token_ids_logprobs_idx.append([]) token_ids_logprobs_idx.append([])
continue continue
# Get the token ids logprobs # Get the token ids logprobs
@@ -264,11 +291,11 @@ def get_token_ids_logprobs_chunk(
# Split-sequence continuations extend; everyone else owns a fresh # Split-sequence continuations extend; everyone else owns a fresh
# (possibly empty) entry. # (possibly empty) entry.
if split_pruned_len > 0: if split_pruned_len > 0:
input_token_ids_logprobs_val[-1].extend(val) token_ids_logprobs_val[-1].extend(val)
input_token_ids_logprobs_idx[-1].extend(idx) token_ids_logprobs_idx[-1].extend(idx)
else: else:
input_token_ids_logprobs_val.append(val) token_ids_logprobs_val.append(val)
input_token_ids_logprobs_idx.append(idx) token_ids_logprobs_idx.append(idx)
pt += pruned_len pt += pruned_len
return next_split_pruned_len return next_split_pruned_len
@@ -349,9 +376,9 @@ class InputLogprobProcessor:
def __init__(self): def __init__(self):
# enable chunked logprobs processing # enable chunked logprobs processing
self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get() self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGPROB_CHUNK.get()
# chunk size for logprobs processing # chunk size for logprobs processing
self.logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get() self.logprobs_chunk_size = envs.SGLANG_LOGPROB_CHUNK_SIZE.get()
def forward( def forward(
self, self,
@@ -363,7 +390,7 @@ class InputLogprobProcessor:
get_logits_fn: Callable, get_logits_fn: Callable,
logits_metadata: LogitsMetadata, logits_metadata: LogitsMetadata,
skip_chunking_for_dp_attn: bool = False, skip_chunking_for_dp_attn: bool = False,
) -> Tuple[InputLogprobsResult, torch.Tensor]: ) -> Tuple[LogprobResult, torch.Tensor]:
# Non-chunked = one chunk covering every row. DP-attention must stay # Non-chunked = one chunk covering every row. DP-attention must stay
# single-chunk: the collective schedule cannot depend on per-rank rows. # single-chunk: the collective schedule cannot depend on per-rank rows.
if ( if (
@@ -396,24 +423,24 @@ class InputLogprobProcessor:
get_logits_fn: Callable, get_logits_fn: Callable,
logits_metadata: LogitsMetadata, logits_metadata: LogitsMetadata,
chunk_size: int, chunk_size: int,
) -> Tuple[InputLogprobsResult, torch.Tensor]: ) -> Tuple[LogprobResult, torch.Tensor]:
"""Compute input logprobs chunk by chunk to cap peak memory.""" """Compute input logprobs chunk by chunk to cap peak memory."""
total_size = pruned_states.shape[0] total_size = pruned_states.shape[0]
num_chunks = (total_size + chunk_size - 1) // chunk_size num_chunks = (total_size + chunk_size - 1) // chunk_size
input_token_logprobs = [] token_logprobs = []
if logits_metadata.extend_return_top_logprob: if logits_metadata.extend_return_top_logprob:
input_top_logprobs_val = [] top_logprobs_val = []
input_top_logprobs_idx = [] top_logprobs_idx = []
else: else:
input_top_logprobs_val = None top_logprobs_val = None
input_top_logprobs_idx = None top_logprobs_idx = None
if logits_metadata.extend_token_ids_logprob: if logits_metadata.extend_token_ids_logprob:
input_token_ids_logprobs_val = [] token_ids_logprobs_val = []
input_token_ids_logprobs_idx = [] token_ids_logprobs_idx = []
else: else:
input_token_ids_logprobs_val = None token_ids_logprobs_val = None
input_token_ids_logprobs_idx = None token_ids_logprobs_idx = None
# If a single sequence is split into multiple chunks, we need to keep track # If a single sequence is split into multiple chunks, we need to keep track
# of the pruned length of the sequences in the previous chunks. # of the pruned length of the sequences in the previous chunks.
@@ -472,11 +499,9 @@ class InputLogprobProcessor:
# Compute the logprobs of the chunk. Free the raw logits before the # Compute the logprobs of the chunk. Free the raw logits before the
# out-of-place log_softmax: keeping all three alive is a 3x peak, # out-of-place log_softmax: keeping all three alive is a 3x peak,
# which OOMs when the single chunk covers a large batch. # which OOMs when the single chunk covers a large batch.
chunk_input_logprobs = chunk_logits[chunk_indices] chunk_logprobs = chunk_logits[chunk_indices]
del chunk_logits del chunk_logits
chunk_input_logprobs = torch.nn.functional.log_softmax( chunk_logprobs = torch.nn.functional.log_softmax(chunk_logprobs, dim=-1)
chunk_input_logprobs, dim=-1
)
# End at the last row inside the chunk; token_to_seq_idx[end_idx] # End at the last row inside the chunk; token_to_seq_idx[end_idx]
# belongs to the next chunk and would emit its sequence twice. # belongs to the next chunk and would emit its sequence twice.
@@ -491,12 +516,11 @@ class InputLogprobProcessor:
chunk_slice chunk_slice
] ]
split_len_topk = get_top_logprobs_chunk( split_len_topk = get_top_logprobs_chunk(
chunk_input_logprobs, chunk_logprobs,
logits_metadata,
top_k_nums, top_k_nums,
pruned_lens, pruned_lens,
input_top_logprobs_val, top_logprobs_val,
input_top_logprobs_idx, top_logprobs_idx,
split_len_topk, split_len_topk,
) )
@@ -507,24 +531,22 @@ class InputLogprobProcessor:
chunk_slice chunk_slice
] ]
split_len_token_ids = get_token_ids_logprobs_chunk( split_len_token_ids = get_token_ids_logprobs_chunk(
chunk_input_logprobs, chunk_logprobs,
token_ids_logprobs, token_ids_logprobs,
pruned_lens, pruned_lens,
input_token_ids_logprobs_val, token_ids_logprobs_val,
input_token_ids_logprobs_idx, token_ids_logprobs_idx,
split_len_token_ids, split_len_token_ids,
) )
# Get the logprob of the requested token ids # Get the logprob of the requested token ids
chunk_input_token_logprobs = chunk_input_logprobs[ chunk_token_logprobs = chunk_logprobs[
torch.arange( torch.arange(chunk_logprobs.shape[0], device=chunk_logprobs.device),
chunk_input_logprobs.shape[0], device=chunk_input_logprobs.device
),
logits_metadata.extend_input_logprob_token_ids_gpu[mask_indices], logits_metadata.extend_input_logprob_token_ids_gpu[mask_indices],
] ]
input_token_logprobs.append(chunk_input_token_logprobs) token_logprobs.append(chunk_token_logprobs)
# Free before the next chunk's logits (bf16 + fp32) materialize. # Free before the next chunk's logits (bf16 + fp32) materialize.
del chunk_input_logprobs del chunk_logprobs
# Restore the full-pruned lm_head batch_info after chunk iteration. # Restore the full-pruned lm_head batch_info after chunk iteration.
if num_chunks > 1 and hasattr(lm_head, "reset_lm_head_pass"): if num_chunks > 1 and hasattr(lm_head, "reset_lm_head_pass"):
@@ -534,15 +556,15 @@ class InputLogprobProcessor:
lm_head.reset_lm_head_pass() lm_head.reset_lm_head_pass()
# Concatenate the results # Concatenate the results
input_token_logprobs = torch.cat(input_token_logprobs, dim=0) token_logprobs = torch.cat(token_logprobs, dim=0)
return ( return (
InputLogprobsResult( LogprobResult(
input_token_logprobs=input_token_logprobs, token_logprobs=token_logprobs,
input_top_logprobs_val=input_top_logprobs_val, top_logprobs_val=top_logprobs_val,
input_top_logprobs_idx=input_top_logprobs_idx, top_logprobs_idx=top_logprobs_idx,
input_token_ids_logprobs_val=input_token_ids_logprobs_val, token_ids_logprobs_val=token_ids_logprobs_val,
input_token_ids_logprobs_idx=input_token_ids_logprobs_idx, token_ids_logprobs_idx=token_ids_logprobs_idx,
), ),
sampled_logits, sampled_logits,
) )
@@ -636,55 +658,26 @@ def get_token_ids_logprobs_batch_optimized(
return output_token_ids_logprobs_val, output_token_ids_logprobs_idx 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: class OutputLogprobProcessor:
"""Output (decode) logprob processing: logprobs -> topk / token-ids / """Output (decode) logprob processing: logprobs -> topk / token-ids /
sampled-token gather, attached onto LogitsProcessorOutput. sampled-token gather, returned as a LogprobResult for the caller to
write back onto LogitsProcessorOutput.
Only logits/logprobs are needed here; sampler-side concerns (custom Only logits/logprobs are needed here; sampler-side concerns (custom
logit processors, NaN sanitizing) are injected via ``preprocess_fn``. logit processors, NaN sanitizing) are injected via ``preprocess_fn``.
""" """
def attach_logprobs_to_output( def compute_logprobs(
self, self,
logits_output: LogitsProcessorOutput,
logprobs: torch.Tensor, logprobs: torch.Tensor,
top_logprobs_nums: List[int], top_logprobs_nums: List[int],
token_ids_logprobs: List[List[int]], token_ids_logprobs: List[List[int]],
batch_next_token_ids: torch.Tensor, batch_next_token_ids: torch.Tensor,
): ) -> LogprobResult:
# clamp to avoid -inf values # clamp to avoid -inf values
logprobs.clamp_(min=torch.finfo(logprobs.dtype).min) logprobs.clamp_(min=torch.finfo(logprobs.dtype).min)
result = OutputLogprobsResult() result = LogprobResult()
if any(x > 0 for x in top_logprobs_nums): if any(x > 0 for x in top_logprobs_nums):
( (
result.top_logprobs_val, result.top_logprobs_val,
@@ -703,16 +696,15 @@ class OutputLogprobProcessor:
torch.arange(len(batch_next_token_ids), device=batch_next_token_ids.device), torch.arange(len(batch_next_token_ids), device=batch_next_token_ids.device),
batch_next_token_ids, batch_next_token_ids,
] ]
result.write_to(logits_output) return result
def compute_logprobs_only( def compute_logprobs_only(
self, self,
logits_output: LogitsProcessorOutput, next_token_logits: Optional[torch.Tensor],
sampling_info: SamplingBatchInfo,
top_logprobs_nums: List[int], top_logprobs_nums: List[int],
token_ids_logprobs: List[List[int]], token_ids_logprobs: List[List[int]],
preprocess_fn: Callable, preprocess_fn: Callable[[torch.Tensor], torch.Tensor],
) -> None: ) -> Optional[LogprobResult]:
""" """
Compute logprobs for requested token IDs without performing sampling. Compute logprobs for requested token IDs without performing sampling.
@@ -720,9 +712,9 @@ class OutputLogprobProcessor:
but don't require next token generation. but don't require next token generation.
""" """
if logits_output.next_token_logits is None: if next_token_logits is None:
logger.warning("No logits available for logprob computation") logger.warning("No logits available for logprob computation")
return return None
# Check if any requests actually need logprobs computation # Check if any requests actually need logprobs computation
needs_token_ids_logprobs = any( needs_token_ids_logprobs = any(
@@ -732,15 +724,15 @@ class OutputLogprobProcessor:
needs_top_logprobs = any(x > 0 for x in top_logprobs_nums) needs_top_logprobs = any(x > 0 for x in top_logprobs_nums)
if not (needs_token_ids_logprobs or needs_top_logprobs): if not (needs_token_ids_logprobs or needs_top_logprobs):
return return None
# Preprocess logits (custom processors and NaN handling) # Preprocess logits (custom processors and NaN handling)
logits = preprocess_fn(logits_output.next_token_logits, sampling_info) logits = preprocess_fn(next_token_logits)
# Compute logprobs # Compute logprobs
logprobs = torch.nn.functional.log_softmax(logits, dim=-1) logprobs = torch.nn.functional.log_softmax(logits, dim=-1)
result = OutputLogprobsResult() result = LogprobResult()
# Handle top logprobs if requested # Handle top logprobs if requested
if needs_top_logprobs: if needs_top_logprobs:
( (
@@ -754,4 +746,4 @@ class OutputLogprobProcessor:
result.token_ids_logprobs_val, result.token_ids_logprobs_val,
result.token_ids_logprobs_idx, result.token_ids_logprobs_idx,
) = get_token_ids_logprobs_batch_optimized(logprobs, token_ids_logprobs) ) = get_token_ids_logprobs_batch_optimized(logprobs, token_ids_logprobs)
result.write_to(logits_output) return result
python/sglang/srt/layers/sampler.py
+8 -8
View File
@@ -1,4 +1,5 @@
import logging import logging
from functools import partial
from typing import Callable, Dict, List, Optional, Tuple from typing import Callable, Dict, List, Optional, Tuple
import torch import torch
@@ -210,17 +211,16 @@ class Sampler(nn.Module):
) )
del probs del probs
# Attach logprobs to logits_output (in-place modification)
if return_logprob: if return_logprob:
if SGLANG_RETURN_ORIGINAL_LOGPROB: if SGLANG_RETURN_ORIGINAL_LOGPROB:
logprobs = original_logprobs logprobs = original_logprobs
self.output_logprob_processor.attach_logprobs_to_output( logprob_result = self.output_logprob_processor.compute_logprobs(
logits_output,
logprobs, logprobs,
top_logprobs_nums, top_logprobs_nums,
token_ids_logprobs, token_ids_logprobs,
batch_next_token_ids, batch_next_token_ids,
) )
logprob_result.write_output_to(logits_output)
self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info) self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info)
@@ -494,17 +494,17 @@ class Sampler(nn.Module):
self, self,
logits_output: LogitsProcessorOutput, logits_output: LogitsProcessorOutput,
sampling_info: SamplingBatchInfo, sampling_info: SamplingBatchInfo,
return_logprob: bool,
top_logprobs_nums: List[int], top_logprobs_nums: List[int],
token_ids_logprobs: List[List[int]], token_ids_logprobs: List[List[int]],
) -> None: ) -> None:
self.output_logprob_processor.compute_logprobs_only( logprob_result = self.output_logprob_processor.compute_logprobs_only(
logits_output=logits_output, next_token_logits=logits_output.next_token_logits,
sampling_info=sampling_info,
top_logprobs_nums=top_logprobs_nums, top_logprobs_nums=top_logprobs_nums,
token_ids_logprobs=token_ids_logprobs, token_ids_logprobs=token_ids_logprobs,
preprocess_fn=self._preprocess_logits, preprocess_fn=partial(self._preprocess_logits, sampling_info=sampling_info),
) )
if logprob_result is not None:
logprob_result.write_output_to(logits_output)
def register_sampler_backend(backend: str, factory: Callable[[], "Sampler"]) -> None: def register_sampler_backend(backend: str, factory: Callable[[], "Sampler"]) -> None:
python/sglang/srt/lora/backend/lmhead_mixing.py
+2 -2
View File
@@ -32,8 +32,8 @@ class LoRABackendLmHeadMixing:
Returns None if logprobs chunking is disabled or the pruned token Returns None if logprobs chunking is disabled or the pruned token
count does not exceed the logprobs chunk size. count does not exceed the logprobs chunk size.
""" """
logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get() logprobs_chunk_size = envs.SGLANG_LOGPROB_CHUNK_SIZE.get()
enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get() enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGPROB_CHUNK.get()
pruned_total = sum(pruned_lens) pruned_total = sum(pruned_lens)
if not enable_logprobs_chunk or pruned_total <= logprobs_chunk_size: if not enable_logprobs_chunk or pruned_total <= logprobs_chunk_size:
python/sglang/srt/model_executor/model_runner.py
-1
View File
@@ -1552,7 +1552,6 @@ class ModelRunner:
self.sampler.compute_logprobs_only( self.sampler.compute_logprobs_only(
logits_output, logits_output,
forward_batch.sampling_info, forward_batch.sampling_info,
forward_batch.return_logprob,
forward_batch.top_logprobs_nums, forward_batch.top_logprobs_nums,
forward_batch.token_ids_logprobs, forward_batch.token_ids_logprobs,
) )
test/manual/test_logprobs.py
+3 -3
View File
@@ -242,10 +242,10 @@ class TestLogprobsDense(unittest.TestCase):
chunk_size = kwargs.pop("chunk_size", None) chunk_size = kwargs.pop("chunk_size", None)
if chunk_size is not None: if chunk_size is not None:
print(f"Setting chunk size to {chunk_size}") print(f"Setting chunk size to {chunk_size}")
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "True" os.environ["SGLANG_ENABLE_LOGPROB_CHUNK"] = "True"
os.environ["SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"] = str(chunk_size) os.environ["SGLANG_LOGPROB_CHUNK_SIZE"] = str(chunk_size)
else: else:
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "False" os.environ["SGLANG_ENABLE_LOGPROB_CHUNK"] = "False"
# Create engine with merged configuration # Create engine with merged configuration
engine_config = {**DEFAULT_ENGINE_CONFIG, **kwargs} engine_config = {**DEFAULT_ENGINE_CONFIG, **kwargs}
test/registered/core/test_srt_endpoint.py
+1 -1
View File
@@ -46,7 +46,7 @@ class TestSRTEndpoint(CustomTestCase):
# The tiny logprob chunk size routes this file's logprob tests # The tiny logprob chunk size routes this file's logprob tests
# through the multi-chunk stitching path (requests at or below 64 # through the multi-chunk stitching path (requests at or below 64
# rows still cover the non-chunked path). # rows still cover the non-chunked path).
env={**SERVER_ENV, "SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "64"}, env={**SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"},
other_args=( other_args=(
"--enable-custom-logit-processor", "--enable-custom-logit-processor",
"--mem-fraction-static", "--mem-fraction-static",
test/registered/lora/test_lora_hf_sgl_logprob_diff.py
+2 -2
View File
@@ -536,8 +536,8 @@ class TestLoRAHFSGLLogprobDifference(CustomTestCase):
""" """
saved = {} saved = {}
env_overrides = { env_overrides = {
"SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK": "true", "SGLANG_ENABLE_LOGPROB_CHUNK": "true",
"SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "4", "SGLANG_LOGPROB_CHUNK_SIZE": "4",
} }
for key, val in env_overrides.items(): for key, val in env_overrides.items():
saved[key] = os.environ.get(key) saved[key] = os.environ.get(key)
test/registered/lora/test_lora_moe_tp_logprob_diff.py
+1 -1
View File
@@ -161,7 +161,7 @@ class TestMoELoRATP2Logprobs(CustomTestCase):
prompts = MOE_LORA_TEST_PROMPTS[:3] prompts = MOE_LORA_TEST_PROMPTS[:3]
baseline = _run_sglang_moe_lora(tp_size=2, prompts=prompts) baseline = _run_sglang_moe_lora(tp_size=2, prompts=prompts)
torch.cuda.empty_cache() torch.cuda.empty_cache()
with envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.override(16): with envs.SGLANG_LOGPROB_CHUNK_SIZE.override(16):
chunked = _run_sglang_moe_lora(tp_size=2, prompts=prompts) chunked = _run_sglang_moe_lora(tp_size=2, prompts=prompts)
for i in range(len(prompts)): for i in range(len(prompts)):
test/registered/unit/layers/test_logprob_chunk_stitching.py
+7 -11
View File
@@ -108,25 +108,21 @@ class TestLogprobChunkStitching(CustomTestCase):
ref, ref_sampled = _run(proc, batch, False, 10**9) ref, ref_sampled = _run(proc, batch, False, 10**9)
got, got_sampled = _run(proc, batch, True, chunk_size) got, got_sampled = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size}" label = f"specs={list(combo)} chunk={chunk_size}"
self.assertEqual( self.assertEqual(ref.top_logprobs_val, got.top_logprobs_val, label)
ref.input_top_logprobs_val, got.input_top_logprobs_val, label self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
)
self.assertEqual(
ref.input_top_logprobs_idx, got.input_top_logprobs_idx, label
)
if with_token_ids: if with_token_ids:
self.assertEqual( self.assertEqual(
ref.input_token_ids_logprobs_val, ref.token_ids_logprobs_val,
got.input_token_ids_logprobs_val, got.token_ids_logprobs_val,
label, label,
) )
self.assertEqual( self.assertEqual(
ref.input_token_ids_logprobs_idx, ref.token_ids_logprobs_idx,
got.input_token_ids_logprobs_idx, got.token_ids_logprobs_idx,
label, label,
) )
torch.testing.assert_close( torch.testing.assert_close(
ref.input_token_logprobs, got.input_token_logprobs, msg=label ref.token_logprobs, got.token_logprobs, msg=label
) )
torch.testing.assert_close(ref_sampled, got_sampled, msg=label) torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
self.assertGreater(tried, 1000) self.assertGreater(tried, 1000)