diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4930ba563..943d94506 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -834,9 +834,13 @@ class Envs: # Sparse Embeddings SGLANG_EMBEDDINGS_SPARSE_HEAD = EnvStr(None) - # Logits processor - SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK = EnvBool(True) - SGLANG_LOGITS_PROCESSER_CHUNK_SIZE = EnvInt(2048) + # Logprob processor + SGLANG_ENABLE_LOGPROB_CHUNK = EnvBoolWithAlias( + True, deprecated_name="SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK" + ) + SGLANG_LOGPROB_CHUNK_SIZE = EnvIntWithAlias( + 2048, deprecated_name="SGLANG_LOGITS_PROCESSER_CHUNK_SIZE" + ) # Tool-Call behavior SGLANG_TOOL_STRICT_LEVEL = EnvInt(ToolStrictLevel.OFF) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 8fb06959f..cf55b97cf 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -473,16 +473,13 @@ class LogitsProcessor(nn.Module): skip_chunking_for_dp_attn=self.do_tensor_parallel_all_gather_dp_attn, ) - return LogitsProcessorOutput( + logits_output = LogitsProcessorOutput( next_token_logits=sampled_logits, 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, ) + logprobs_result.write_input_to(logits_output) + return logits_output def _get_pruned_states( self, diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py index b23723e9a..5ea86a6a5 100644 --- a/python/sglang/srt/layers/logprob_processor.py +++ b/python/sglang/srt/layers/logprob_processor.py @@ -12,7 +12,6 @@ from sglang.srt.environ import envs if TYPE_CHECKING: 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__) @@ -23,12 +22,43 @@ class LogprobStage(Enum): @dataclasses.dataclass -class InputLogprobsResult: - input_token_logprobs: torch.Tensor - input_top_logprobs_val: Optional[List] = None - input_top_logprobs_idx: Optional[List] = None - input_token_ids_logprobs_val: Optional[List] = None - input_token_ids_logprobs_idx: Optional[List] = None +class LogprobResult: + """Logprob fields produced by Input/OutputLogprobProcessor. + + Input (prefill) always fills token_logprobs; output (decode / scoring) + fills on demand. write_input_to / write_output_to flush 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_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( @@ -135,29 +165,27 @@ def get_token_ids_logprobs(logprobs, token_ids_logprobs, no_copy_to_cpu=False): def get_top_logprobs_chunk( logprobs: torch.Tensor, - logits_metadata: LogitsMetadata, top_k_nums: List[int], pruned_lens: List[int], - input_top_logprobs_val: List, - input_top_logprobs_idx: List, + top_logprobs_val: List, + top_logprobs_idx: List, split_pruned_len: int, ) -> int: """Get top-k logprobs for each sequence in the chunk. Args: 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 pruned_lens: List of pruned lengths for each sequence - input_top_logprobs_val: List to store top-k logprob values - input_top_logprobs_idx: List to store top-k token indices + top_logprobs_val: List to store top-k logprob values + top_logprobs_idx: List to store top-k token indices split_pruned_len: Length of pruned tokens from previous chunk Returns: int: Number of remaining tokens to process in next chunk """ # 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) values = ret.values.tolist() indices = ret.indices.tolist() @@ -175,8 +203,8 @@ def get_top_logprobs_chunk( if pruned_len <= 0: # if pruned length is less than or equal to 0, # there is no top-k logprobs to process - input_top_logprobs_val.append([]) - input_top_logprobs_idx.append([]) + top_logprobs_val.append([]) + top_logprobs_idx.append([]) continue # Get the top-k logprobs @@ -195,11 +223,11 @@ def get_top_logprobs_chunk( # Split-sequence continuations extend; everyone else owns a fresh # (possibly empty) entry. if split_pruned_len > 0: - input_top_logprobs_val[-1].extend(val) - input_top_logprobs_idx[-1].extend(idx) + top_logprobs_val[-1].extend(val) + top_logprobs_idx[-1].extend(idx) else: - input_top_logprobs_val.append(val) - input_top_logprobs_idx.append(idx) + top_logprobs_val.append(val) + top_logprobs_idx.append(idx) pt += pruned_len return next_split_pruned_len @@ -209,19 +237,18 @@ def get_token_ids_logprobs_chunk( logprobs: torch.Tensor, token_ids_logprobs: List[int], pruned_lens: List[int], - input_token_ids_logprobs_val: List, - input_token_ids_logprobs_idx: List, + token_ids_logprobs_val: List, + token_ids_logprobs_idx: List, split_pruned_len: int = 0, ): """Get token_ids logprobs for each sequence in the chunk. Args: 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 pruned_lens: List of pruned lengths for each sequence - input_token_ids_logprobs_val: List to store token logprob values - input_token_ids_logprobs_idx: List to store token indices + token_ids_logprobs_val: List to store token logprob values + token_ids_logprobs_idx: List to store token indices split_pruned_len: Length of pruned tokens from previous chunk Returns: @@ -245,8 +272,8 @@ def get_token_ids_logprobs_chunk( if pruned_len <= 0: # if pruned length is less than or equal to 0, # there is no token ids logprobs to process - input_token_ids_logprobs_val.append([]) - input_token_ids_logprobs_idx.append([]) + token_ids_logprobs_val.append([]) + token_ids_logprobs_idx.append([]) continue # Get the token ids logprobs @@ -264,11 +291,11 @@ def get_token_ids_logprobs_chunk( # Split-sequence continuations extend; everyone else owns a fresh # (possibly empty) entry. if split_pruned_len > 0: - input_token_ids_logprobs_val[-1].extend(val) - input_token_ids_logprobs_idx[-1].extend(idx) + token_ids_logprobs_val[-1].extend(val) + token_ids_logprobs_idx[-1].extend(idx) else: - input_token_ids_logprobs_val.append(val) - input_token_ids_logprobs_idx.append(idx) + token_ids_logprobs_val.append(val) + token_ids_logprobs_idx.append(idx) pt += pruned_len return next_split_pruned_len @@ -349,9 +376,9 @@ class InputLogprobProcessor: def __init__(self): # 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 - self.logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get() + self.logprobs_chunk_size = envs.SGLANG_LOGPROB_CHUNK_SIZE.get() def forward( self, @@ -363,7 +390,7 @@ class InputLogprobProcessor: get_logits_fn: Callable, logits_metadata: LogitsMetadata, 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 # single-chunk: the collective schedule cannot depend on per-rank rows. if ( @@ -396,24 +423,24 @@ class InputLogprobProcessor: get_logits_fn: Callable, logits_metadata: LogitsMetadata, chunk_size: int, - ) -> Tuple[InputLogprobsResult, torch.Tensor]: + ) -> Tuple[LogprobResult, torch.Tensor]: """Compute input logprobs chunk by chunk to cap peak memory.""" total_size = pruned_states.shape[0] num_chunks = (total_size + chunk_size - 1) // chunk_size - input_token_logprobs = [] + token_logprobs = [] if logits_metadata.extend_return_top_logprob: - input_top_logprobs_val = [] - input_top_logprobs_idx = [] + top_logprobs_val = [] + top_logprobs_idx = [] else: - input_top_logprobs_val = None - input_top_logprobs_idx = None + top_logprobs_val = None + top_logprobs_idx = None if logits_metadata.extend_token_ids_logprob: - input_token_ids_logprobs_val = [] - input_token_ids_logprobs_idx = [] + token_ids_logprobs_val = [] + token_ids_logprobs_idx = [] else: - input_token_ids_logprobs_val = None - input_token_ids_logprobs_idx = None + token_ids_logprobs_val = None + token_ids_logprobs_idx = None # 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. @@ -472,11 +499,9 @@ class InputLogprobProcessor: # 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, # 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 - chunk_input_logprobs = torch.nn.functional.log_softmax( - chunk_input_logprobs, dim=-1 - ) + chunk_logprobs = torch.nn.functional.log_softmax(chunk_logprobs, dim=-1) # 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. @@ -491,12 +516,11 @@ class InputLogprobProcessor: chunk_slice ] split_len_topk = get_top_logprobs_chunk( - chunk_input_logprobs, - logits_metadata, + chunk_logprobs, top_k_nums, pruned_lens, - input_top_logprobs_val, - input_top_logprobs_idx, + top_logprobs_val, + top_logprobs_idx, split_len_topk, ) @@ -507,24 +531,22 @@ class InputLogprobProcessor: chunk_slice ] split_len_token_ids = get_token_ids_logprobs_chunk( - chunk_input_logprobs, + chunk_logprobs, token_ids_logprobs, pruned_lens, - input_token_ids_logprobs_val, - input_token_ids_logprobs_idx, + token_ids_logprobs_val, + token_ids_logprobs_idx, split_len_token_ids, ) # Get the logprob of the requested token ids - chunk_input_token_logprobs = chunk_input_logprobs[ - torch.arange( - chunk_input_logprobs.shape[0], device=chunk_input_logprobs.device - ), + chunk_token_logprobs = chunk_logprobs[ + torch.arange(chunk_logprobs.shape[0], device=chunk_logprobs.device), 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. - del chunk_input_logprobs + del chunk_logprobs # Restore the full-pruned lm_head batch_info after chunk iteration. if num_chunks > 1 and hasattr(lm_head, "reset_lm_head_pass"): @@ -534,15 +556,15 @@ class InputLogprobProcessor: lm_head.reset_lm_head_pass() # Concatenate the results - input_token_logprobs = torch.cat(input_token_logprobs, dim=0) + token_logprobs = torch.cat(token_logprobs, dim=0) return ( - InputLogprobsResult( - input_token_logprobs=input_token_logprobs, - input_top_logprobs_val=input_top_logprobs_val, - input_top_logprobs_idx=input_top_logprobs_idx, - input_token_ids_logprobs_val=input_token_ids_logprobs_val, - input_token_ids_logprobs_idx=input_token_ids_logprobs_idx, + LogprobResult( + token_logprobs=token_logprobs, + top_logprobs_val=top_logprobs_val, + top_logprobs_idx=top_logprobs_idx, + token_ids_logprobs_val=token_ids_logprobs_val, + token_ids_logprobs_idx=token_ids_logprobs_idx, ), sampled_logits, ) @@ -636,55 +658,26 @@ def get_token_ids_logprobs_batch_optimized( 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. + 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 logit processors, NaN sanitizing) are injected via ``preprocess_fn``. """ - def attach_logprobs_to_output( + def compute_logprobs( self, - logits_output: LogitsProcessorOutput, logprobs: torch.Tensor, top_logprobs_nums: List[int], token_ids_logprobs: List[List[int]], batch_next_token_ids: torch.Tensor, - ): + ) -> LogprobResult: # clamp to avoid -inf values logprobs.clamp_(min=torch.finfo(logprobs.dtype).min) - result = OutputLogprobsResult() + result = LogprobResult() if any(x > 0 for x in top_logprobs_nums): ( result.top_logprobs_val, @@ -703,16 +696,15 @@ class OutputLogprobProcessor: torch.arange(len(batch_next_token_ids), device=batch_next_token_ids.device), batch_next_token_ids, ] - result.write_to(logits_output) + return result def compute_logprobs_only( self, - logits_output: LogitsProcessorOutput, - sampling_info: SamplingBatchInfo, + next_token_logits: Optional[torch.Tensor], top_logprobs_nums: List[int], token_ids_logprobs: List[List[int]], - preprocess_fn: Callable, - ) -> None: + preprocess_fn: Callable[[torch.Tensor], torch.Tensor], + ) -> Optional[LogprobResult]: """ Compute logprobs for requested token IDs without performing sampling. @@ -720,9 +712,9 @@ class OutputLogprobProcessor: 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") - return + return None # Check if any requests actually need logprobs computation needs_token_ids_logprobs = any( @@ -732,15 +724,15 @@ class OutputLogprobProcessor: needs_top_logprobs = any(x > 0 for x in top_logprobs_nums) if not (needs_token_ids_logprobs or needs_top_logprobs): - return + return None # 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 logprobs = torch.nn.functional.log_softmax(logits, dim=-1) - result = OutputLogprobsResult() + result = LogprobResult() # Handle top logprobs if requested if needs_top_logprobs: ( @@ -754,4 +746,4 @@ class OutputLogprobProcessor: 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) + return result diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index b208eb69c..de587ec5f 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -1,4 +1,5 @@ import logging +from functools import partial from typing import Callable, Dict, List, Optional, Tuple import torch @@ -210,17 +211,16 @@ class Sampler(nn.Module): ) del probs - # Attach logprobs to logits_output (in-place modification) if return_logprob: if SGLANG_RETURN_ORIGINAL_LOGPROB: logprobs = original_logprobs - self.output_logprob_processor.attach_logprobs_to_output( - logits_output, + logprob_result = self.output_logprob_processor.compute_logprobs( logprobs, top_logprobs_nums, token_ids_logprobs, batch_next_token_ids, ) + logprob_result.write_output_to(logits_output) self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info) @@ -494,17 +494,17 @@ class Sampler(nn.Module): self, logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo, - return_logprob: bool, top_logprobs_nums: List[int], token_ids_logprobs: List[List[int]], ) -> None: - self.output_logprob_processor.compute_logprobs_only( - logits_output=logits_output, - sampling_info=sampling_info, + logprob_result = self.output_logprob_processor.compute_logprobs_only( + next_token_logits=logits_output.next_token_logits, top_logprobs_nums=top_logprobs_nums, 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: diff --git a/python/sglang/srt/lora/backend/lmhead_mixing.py b/python/sglang/srt/lora/backend/lmhead_mixing.py index e7ed98176..96e5b9129 100644 --- a/python/sglang/srt/lora/backend/lmhead_mixing.py +++ b/python/sglang/srt/lora/backend/lmhead_mixing.py @@ -32,8 +32,8 @@ class LoRABackendLmHeadMixing: Returns None if logprobs chunking is disabled or the pruned token count does not exceed the logprobs chunk size. """ - logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get() - enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get() + logprobs_chunk_size = envs.SGLANG_LOGPROB_CHUNK_SIZE.get() + enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGPROB_CHUNK.get() pruned_total = sum(pruned_lens) if not enable_logprobs_chunk or pruned_total <= logprobs_chunk_size: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 0ca99255e..547494b48 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1552,7 +1552,6 @@ class ModelRunner: self.sampler.compute_logprobs_only( logits_output, forward_batch.sampling_info, - forward_batch.return_logprob, forward_batch.top_logprobs_nums, forward_batch.token_ids_logprobs, ) diff --git a/test/manual/test_logprobs.py b/test/manual/test_logprobs.py index 28b3e2723..9256d1b4b 100644 --- a/test/manual/test_logprobs.py +++ b/test/manual/test_logprobs.py @@ -242,10 +242,10 @@ class TestLogprobsDense(unittest.TestCase): chunk_size = kwargs.pop("chunk_size", None) if chunk_size is not None: print(f"Setting chunk size to {chunk_size}") - os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "True" - os.environ["SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"] = str(chunk_size) + os.environ["SGLANG_ENABLE_LOGPROB_CHUNK"] = "True" + os.environ["SGLANG_LOGPROB_CHUNK_SIZE"] = str(chunk_size) else: - os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "False" + os.environ["SGLANG_ENABLE_LOGPROB_CHUNK"] = "False" # Create engine with merged configuration engine_config = {**DEFAULT_ENGINE_CONFIG, **kwargs} diff --git a/test/registered/core/test_srt_endpoint.py b/test/registered/core/test_srt_endpoint.py index 04c2fa2ed..81f3fa7ba 100644 --- a/test/registered/core/test_srt_endpoint.py +++ b/test/registered/core/test_srt_endpoint.py @@ -46,7 +46,7 @@ class TestSRTEndpoint(CustomTestCase): # The tiny logprob chunk size routes this file's logprob tests # through the multi-chunk stitching path (requests at or below 64 # 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=( "--enable-custom-logit-processor", "--mem-fraction-static", diff --git a/test/registered/lora/test_lora_hf_sgl_logprob_diff.py b/test/registered/lora/test_lora_hf_sgl_logprob_diff.py index efebc2e1b..4c713e186 100644 --- a/test/registered/lora/test_lora_hf_sgl_logprob_diff.py +++ b/test/registered/lora/test_lora_hf_sgl_logprob_diff.py @@ -536,8 +536,8 @@ class TestLoRAHFSGLLogprobDifference(CustomTestCase): """ saved = {} env_overrides = { - "SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK": "true", - "SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "4", + "SGLANG_ENABLE_LOGPROB_CHUNK": "true", + "SGLANG_LOGPROB_CHUNK_SIZE": "4", } for key, val in env_overrides.items(): saved[key] = os.environ.get(key) diff --git a/test/registered/lora/test_lora_moe_tp_logprob_diff.py b/test/registered/lora/test_lora_moe_tp_logprob_diff.py index 6ee668fe2..20b42e77e 100644 --- a/test/registered/lora/test_lora_moe_tp_logprob_diff.py +++ b/test/registered/lora/test_lora_moe_tp_logprob_diff.py @@ -161,7 +161,7 @@ class TestMoELoRATP2Logprobs(CustomTestCase): prompts = MOE_LORA_TEST_PROMPTS[:3] baseline = _run_sglang_moe_lora(tp_size=2, prompts=prompts) 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) for i in range(len(prompts)): diff --git a/test/registered/unit/layers/test_logprob_chunk_stitching.py b/test/registered/unit/layers/test_logprob_chunk_stitching.py index 48489fdc5..9ba2a2146 100644 --- a/test/registered/unit/layers/test_logprob_chunk_stitching.py +++ b/test/registered/unit/layers/test_logprob_chunk_stitching.py @@ -108,25 +108,21 @@ class TestLogprobChunkStitching(CustomTestCase): ref, ref_sampled = _run(proc, batch, False, 10**9) got, got_sampled = _run(proc, batch, True, chunk_size) label = f"specs={list(combo)} chunk={chunk_size}" - self.assertEqual( - ref.input_top_logprobs_val, got.input_top_logprobs_val, label - ) - self.assertEqual( - ref.input_top_logprobs_idx, got.input_top_logprobs_idx, label - ) + self.assertEqual(ref.top_logprobs_val, got.top_logprobs_val, label) + self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label) if with_token_ids: self.assertEqual( - ref.input_token_ids_logprobs_val, - got.input_token_ids_logprobs_val, + ref.token_ids_logprobs_val, + got.token_ids_logprobs_val, label, ) self.assertEqual( - ref.input_token_ids_logprobs_idx, - got.input_token_ids_logprobs_idx, + ref.token_ids_logprobs_idx, + got.token_ids_logprobs_idx, label, ) 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) self.assertGreater(tried, 1000)