From b53f9ebd00c78cf50b63f00a8e2976bbb7968882 Mon Sep 17 00:00:00 2001 From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:29:34 +0000 Subject: [PATCH 1/5] feat: add sampling logprob output modes --- python/sglang/srt/disaggregation/decode.py | 13 +- python/sglang/srt/disaggregation/utils.py | 35 ++- .../sglang/srt/entrypoints/openai/protocol.py | 1 + .../srt/entrypoints/openai/serving_chat.py | 1 + python/sglang/srt/layers/logits_processor.py | 13 +- python/sglang/srt/layers/sampler.py | 76 +++---- python/sglang/srt/managers/io_struct.py | 33 ++- python/sglang/srt/managers/schedule_batch.py | 3 + python/sglang/srt/managers/scheduler.py | 1 + .../batch_result_processor.py | 36 +++- .../scheduler_components/output_streamer.py | 4 +- .../sglang/srt/managers/tokenizer_manager.py | 14 +- python/sglang/srt/managers/utils.py | 6 + .../srt/sampling/sampling_batch_info.py | 63 ++++++ .../sglang/srt/session/session_controller.py | 1 + .../registered/sampling/test_sampling_mask.py | 201 +++++++++++++++--- .../test_disaggregation_wire.py | 52 ++++- .../unit/entrypoints/openai/test_protocol.py | 1 + .../entrypoints/openai/test_serving_chat.py | 2 + .../unit/managers/test_auxiliary_output.py | 14 +- ...st_batch_result_processor_hidden_states.py | 89 +++++++- .../unit/managers/test_io_struct.py | 14 ++ .../managers/test_multi_tokenizer_mixin.py | 15 +- .../managers/test_sampling_mask_validation.py | 18 +- .../unit/sampling/test_sampling_batch_info.py | 26 +++ 25 files changed, 615 insertions(+), 117 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 1f6dfabcc2a1..3a2a8eacfe8d 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -2396,9 +2396,16 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest): decode_req.req.output_token_sampling_mask.append( output_token_sampling_mask_idx[:sampling_mask_len].cpu().tolist() ) - decode_req.req.output_token_sampling_logprobs.append( - float(output_token_sampling_logprobs[0].item()) - ) + if decode_req.req.sampling_logprobs_mode == "support": + decode_req.req.output_token_sampling_logprobs.append( + output_token_sampling_logprobs[:sampling_mask_len] + .cpu() + .tolist() + ) + else: + decode_req.req.output_token_sampling_logprobs.append( + float(output_token_sampling_logprobs[0].item()) + ) decode_req.kv_receiver.clear() decode_req.kv_receiver = None diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 7f2b8f45b9fd..2a8c2aea9ab3 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -349,7 +349,9 @@ def __init__( (size, max_sampling_mask_tokens), dtype=torch.int32, device=device ) self.output_token_sampling_logprobs = torch.zeros( - (size, 16), dtype=torch.float32, device=device + (size, max_sampling_mask_tokens), + dtype=torch.float32, + device=device, ) # For PD + spec decode self.output_topk_p = torch.zeros( @@ -517,9 +519,19 @@ def set_buf(self, req: Req): sampling_logprobs = req.output_token_sampling_logprobs if sampling_masks: sampling_mask = sampling_masks[0] - sampling_logprob = sampling_logprobs[0] if sampling_logprobs else None - if sampling_mask is not None and sampling_logprob is not None: + sampling_logprobs_row = ( + sampling_logprobs[0] if sampling_logprobs else None + ) + if sampling_mask is not None and sampling_logprobs_row is not None: mask_len = len(sampling_mask) + if ( + req.sampling_logprobs_mode == "support" + and len(sampling_logprobs_row) != mask_len + ): + raise RuntimeError( + "Sampling mask IDs and logprobs must have equal lengths; " + f"got {mask_len} and {len(sampling_logprobs_row)}." + ) max_mask_len = self.output_token_sampling_mask_idx.shape[1] if mask_len > max_mask_len: raise RuntimeError( @@ -540,9 +552,20 @@ def set_buf(self, req: Req): device=self.output_token_sampling_mask_idx.device, ) ) - self.output_token_sampling_logprobs[req.metadata_buffer_index][ - 0 - ] = float(sampling_logprob) + if req.sampling_logprobs_mode == "support" and mask_len: + self.output_token_sampling_logprobs[ + req.metadata_buffer_index, :mask_len + ].copy_( + torch.tensor( + sampling_logprobs_row, + dtype=torch.float32, + device=self.output_token_sampling_logprobs.device, + ) + ) + elif req.sampling_logprobs_mode == "selected": + self.output_token_sampling_logprobs[ + req.metadata_buffer_index, 0 + ] = float(sampling_logprobs_row) # For PD + spec decode if req.hidden_states_tensor is not None: # speculative_eagle_topk should not be greater than 16 currently diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index d4579b24b820..19410851057c 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -894,6 +894,7 @@ class ChatCompletionRequest(BaseModel): return_input_ids_in_sglext: bool = False return_output_ids_in_sglext: bool = False return_sampling_mask: bool = False + sampling_logprobs_mode: Literal["selected", "support"] = "selected" reasoning_effort: ReasoningEffortType = Field( default=None, description="Constrains effort on reasoning for reasoning models. " diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index b737c83acdd7..92842acff2b7 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -1239,6 +1239,7 @@ def _convert_to_internal_request( logprob_start_len=-1, top_logprobs_num=request.top_logprobs or 0, return_sampling_mask=request.return_sampling_mask, + sampling_logprobs_mode=request.sampling_logprobs_mode, stream=request.stream, return_text_in_logprobs=True, modalities=processed_messages.modalities, diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 84dd8d0a95c4..5280910431a9 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -93,17 +93,20 @@ class SamplingMaskStatus(IntEnum): @dataclasses.dataclass class SamplingMaskOutput: - """Tensor result for opted-in rows in batch order.""" + """Sampling-support IDs and optional full-support behavior logprobs.""" token_ids: torch.Tensor lengths: torch.Tensor selected_logprobs: torch.Tensor + support_logprobs: Optional[torch.Tensor] statuses: torch.Tensor def map_device_tensors(self, fn) -> None: self.token_ids = fn(self.token_ids) self.lengths = fn(self.lengths) self.selected_logprobs = fn(self.selected_logprobs) + if self.support_logprobs is not None: + self.support_logprobs = fn(self.support_logprobs) self.statuses = fn(self.statuses) @@ -223,11 +226,13 @@ class LogitsProcessorOutput: List[Union[List[float], torch.Tensor]] ] = None next_token_token_ids_logprobs_idx: Optional[List] = None - # Post-filter support IDs, bounded by server capacity, and selected-token - # logprob over the full realized support. + # Post-filter support IDs and requested behavior logprobs, bounded by server + # capacity. Logprobs are normalized over the full realized support. sampling_mask_output: Optional[SamplingMaskOutput] = None next_token_sampling_mask_idx: Optional[List[Optional[List[int]]]] = None - next_token_sampling_logprobs: Optional[List[Optional[float]]] = None + next_token_sampling_logprobs: Optional[ + List[Optional[Union[float, List[float]]]] + ] = None next_token_sampling_mask_status: Optional[List[Optional[int]]] = None ## Part 3: Prefill-only. This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index acd35db2cff2..ed8d3d4760b4 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -88,11 +88,10 @@ def _trace_e2e_sampler(stage: str, **fields) -> None: class _SamplingMaskCapture(NamedTuple): - """Compact post-filter weights and their original batch-row mapping.""" + """Post-filter weights and their original batch-row mapping.""" weights: torch.Tensor token_ids: Optional[torch.Tensor] - selected_weight: Optional[torch.Tensor] batch_rows: torch.Tensor @@ -176,6 +175,9 @@ def forward( _trace_e2e_sampler("preprocess_returned") sampling_mask_batch_indices = sampling_info.sampling_mask_batch_indices return_sampling_mask = sampling_mask_batch_indices is not None + sampling_support_logprobs_capture_indices = ( + sampling_info.sampling_support_logprobs_capture_indices + ) sampling_mask_capture = None if sampling_info.is_all_greedy: @@ -296,19 +298,17 @@ def forward( if sampling_info.is_all_greedy: logits_output.sampling_mask_output = ( self._build_greedy_sampling_mask_output( - sampling_mask_batch_indices, batch_next_token_ids + sampling_mask_batch_indices, + batch_next_token_ids, + sampling_support_logprobs_capture_indices, ) ) else: assert sampling_mask_capture is not None - if SYNC_TOKEN_IDS_ACROSS_TP or sampling_info.grammars: - # Token synchronization can replace the producer-selected token. - sampling_mask_capture = sampling_mask_capture._replace( - selected_weight=None - ) logits_output.sampling_mask_output = self._build_sampling_mask_output( batch_next_token_ids, sampling_mask_capture, + sampling_support_logprobs_capture_indices, ) _trace_e2e_sampler("forward_returned") @@ -343,14 +343,9 @@ def _sample_from_probs( ) if return_sampling_mask: capture_probs = select_capture_rows(probs) - capture_tokens = select_capture_rows(batch_next_token_ids) - selected_weight = torch.gather( - capture_probs, 1, capture_tokens.long().view(-1, 1) - ).squeeze(1) sampling_mask_capture = _SamplingMaskCapture( weights=capture_probs, token_ids=None, - selected_weight=selected_weight, batch_rows=capture_rows, ) else: @@ -368,22 +363,15 @@ def _sample_from_probs( if return_sampling_mask: capture_probs = select_capture_rows(probs) capture_min_ps = select_capture_rows(sampling_info.min_ps) - capture_tokens = select_capture_rows(batch_next_token_ids) min_p_thresholds = ( capture_probs.max(dim=-1).values * capture_min_ps ) filtered_probs = capture_probs.masked_fill( capture_probs < min_p_thresholds.view(-1, 1), 0 ) - selected_weight = torch.gather( - filtered_probs, - 1, - capture_tokens.long().view(-1, 1), - ).squeeze(1) sampling_mask_capture = _SamplingMaskCapture( weights=filtered_probs, token_ids=None, - selected_weight=selected_weight, batch_rows=capture_rows, ) else: @@ -401,7 +389,6 @@ def _sample_from_probs( capture_probs = select_capture_rows(probs) capture_top_ks = select_capture_rows(sampling_info.top_ks) capture_top_ps = select_capture_rows(sampling_info.top_ps) - capture_tokens = select_capture_rows(batch_next_token_ids) filtered_probs = capture_probs if sampling_info.need_top_k_sampling: filtered_probs = top_k_renorm_prob( @@ -415,15 +402,9 @@ def _sample_from_probs( filtered_probs = top_p_probs else: filtered_probs.masked_fill_(top_p_probs <= 0, 0) - selected_weight = torch.gather( - filtered_probs, - 1, - capture_tokens.long().view(-1, 1), - ).squeeze(1) sampling_mask_capture = _SamplingMaskCapture( weights=filtered_probs, token_ids=None, - selected_weight=selected_weight, batch_rows=capture_rows, ) elif backend == "pytorch": @@ -443,12 +424,10 @@ def _sample_from_probs( batch_next_token_ids, filtered_probs, token_ids, - selected_weight, ) = sample_result sampling_mask_capture = _SamplingMaskCapture( weights=select_capture_rows(filtered_probs), token_ids=select_capture_rows(token_ids), - selected_weight=select_capture_rows(selected_weight), batch_rows=capture_rows, ) else: @@ -461,6 +440,7 @@ def _build_greedy_sampling_mask_output( self, batch_indices: torch.Tensor, batch_next_token_ids: torch.Tensor, + support_capture_indices: Optional[torch.Tensor], ) -> SamplingMaskOutput: token_ids = batch_next_token_ids.index_select(0, batch_indices).to(torch.int32) num_requests = batch_indices.numel() @@ -472,6 +452,15 @@ def _build_greedy_sampling_mask_output( selected_logprobs=torch.zeros( num_requests, dtype=torch.float32, device=token_ids.device ), + support_logprobs=( + torch.zeros( + (support_capture_indices.numel(), 1), + dtype=torch.float32, + device=token_ids.device, + ) + if support_capture_indices is not None + else None + ), statuses=torch.full( (num_requests,), SamplingMaskStatus.OK, @@ -484,12 +473,12 @@ def _build_sampling_mask_output( self, batch_next_token_ids: torch.Tensor, sampling_mask_capture: _SamplingMaskCapture, + support_capture_indices: Optional[torch.Tensor], ) -> SamplingMaskOutput: """Pack captured positive support into the fixed-cap device result.""" batch_indices = sampling_mask_capture.batch_rows weights = sampling_mask_capture.weights token_ids = sampling_mask_capture.token_ids - selected_weight = sampling_mask_capture.selected_weight sampled_tokens = batch_next_token_ids.index_select(0, batch_indices).view(-1, 1) sampled_tokens_int32 = sampled_tokens.to(torch.int32) if token_ids is None: @@ -506,18 +495,14 @@ def _build_sampling_mask_output( selected_from_weights = torch.gather( weights, 1, selected_positions ).squeeze(1) - if selected_weight is None: - selected_weight = selected_from_weights - support = weights > 0 support_mass = torch.where(support, weights, torch.zeros_like(weights)).sum( dim=-1, dtype=torch.float32 ) - selected_weight = selected_weight.float() - selected_logprobs = torch.log(selected_weight / support_mass) + selected_logprobs = torch.log(selected_from_weights.float() / support_mass) invalid = ~( sampled_in_capture - & (selected_weight > 0) + & (selected_from_weights > 0) & (support_mass > 0) & torch.isfinite(selected_logprobs) ) @@ -549,18 +534,27 @@ def _build_sampling_mask_output( packed_size = min(self.sampling_mask_max_tokens, weights.shape[-1]) if token_ids is None: - _, packed_positions = torch.topk( + packed_weights, packed_positions = torch.topk( weights, k=packed_size, dim=-1, largest=True, sorted=True ) packed_token_ids = packed_positions.to(torch.int32) else: # The PyTorch producer already sorts weights and IDs together. packed_token_ids = token_ids[:, :packed_size].contiguous() + packed_weights = weights[:, :packed_size] + + support_logprobs = None + if support_capture_indices is not None: + support_logprobs = torch.log( + packed_weights.index_select(0, support_capture_indices).float() + / support_mass.index_select(0, support_capture_indices).unsqueeze(-1) + ) return SamplingMaskOutput( token_ids=packed_token_ids, lengths=realized_lengths.clamp(max=packed_size), selected_logprobs=selected_logprobs, + support_logprobs=support_logprobs, statuses=statuses, ) @@ -733,8 +727,7 @@ def top_k_top_p_min_p_sampling_from_probs_torch( with the sampling_seed of each request. By default, returns only sampled token IDs. With return_filtered_probs=True, - also returns the actual filtered weights, their token-ID permutation, and - the selected weights. + also returns the actual filtered weights and their token-ID permutation. """ probs_sort, probs_idx = probs.sort(dim=-1, descending=True) probs_sum = torch.cumsum(probs_sort, dim=-1) @@ -765,14 +758,11 @@ def top_k_top_p_min_p_sampling_from_probs_torch( logprobs.log_() sampled_index = multinomial_with_seed(logprobs, sampling_seed, positions) - if return_filtered_probs: - selected_weight = torch.gather(probs_sort, 1, sampled_index).view(-1) - # int32 range is enough to represent the token ids probs_idx = probs_idx.to(torch.int32) batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(-1) if return_filtered_probs: - return batch_next_token_ids, probs_sort, probs_idx, selected_weight + return batch_next_token_ids, probs_sort, probs_idx return batch_next_token_ids diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 8f5309db2c52..c69de94c66f3 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -59,6 +59,7 @@ Modality, MultimodalProcessorOutput, ReturnHiddenStatesMode, + SamplingLogprobsMode, get_return_hidden_states_mode, ) from sglang.srt.multimodal.mm_utils import has_valid_data @@ -232,8 +233,13 @@ class GenerateReqInput: top_logprobs_num: Optional[Union[List[int], int]] = None # If return logprobs, the token ids to return logprob for. token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None - # Whether to return output-token sampling support and renormalized logprobs. + # Whether to return each output token's sampling support and behavior logprob. return_sampling_mask: Optional[Union[List[bool], bool]] = None + # Return either the selected token's behavior logprob or behavior logprobs for + # the full sampling support. Only used when return_sampling_mask is enabled. + sampling_logprobs_mode: Union[List[SamplingLogprobsMode], SamplingLogprobsMode] = ( + "selected" + ) # Whether to detokenize tokens in text in the returned logprobs. return_text_in_logprobs: bool = False # Return prompt top logprobs as flat arrays plus shape metadata instead of @@ -768,6 +774,9 @@ def normalize_param(param, default_value, param_name): self.return_sampling_mask = normalize_param( self.return_sampling_mask, False, "return_sampling_mask" ) + self.sampling_logprobs_mode = normalize_param( + self.sampling_logprobs_mode, "selected", "sampling_logprobs_mode" + ) # Handle token_ids_logprob specially due to its nested structure if not self.token_ids_logprob: # covers both None and [] @@ -948,6 +957,7 @@ def __getitem__(self, i): top_logprobs_num=self.top_logprobs_num[i], token_ids_logprob=self.token_ids_logprob[i], return_sampling_mask=self.return_sampling_mask[i], + sampling_logprobs_mode=self.sampling_logprobs_mode[i], return_text_in_logprobs=self.return_text_in_logprobs, return_flat_raw_top_logprobs=self.return_flat_raw_top_logprobs, return_flat_raw_top_logprobs_b64=self.return_flat_raw_top_logprobs_b64, @@ -1121,8 +1131,11 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): kv_hints: Optional[KvHintsEnvelope] = None # Internal PP control bit, set by PP0 before forwarding the request. - # Keep at the end to preserve the positional Rust wire schema. + # Keep tail fields append-only to preserve the positional Rust wire schema. pp_prefetch_ticketed: bool = False + # Shape of output_token_sampling_logprobs for each output token. This is a + # defaulted tail field so older IPC senders decode as selected mode. + sampling_logprobs_mode: SamplingLogprobsMode = "selected" def wrap_pickle_fields(self): self.time_stats = wrap_as_pickle(self.time_stats) @@ -1514,10 +1527,10 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): output_token_entropy_val: Optional[List[Optional[float]]] # Per-request chunks of output-token sampling supports. None when no request # in the batch asks for return_sampling_mask. - output_token_sampling_mask: Optional[List[List]] - # Per-request chunks of selected-token logprobs renormalized over the - # corresponding sampling supports. None when sampling masks are not returned. - output_token_sampling_logprobs: Optional[List[List]] + output_token_sampling_mask: Optional[List[List[List[int]]]] + # Per-request chunks. Each output-token entry is a selected-token scalar or + # a list aligned with output_token_sampling_mask, according to the request. + output_token_sampling_logprobs: Optional[List[List[Union[float, List[float]]]]] # Hidden states output_hidden_states: OutputHiddenStates @@ -1612,10 +1625,10 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): output_token_ids_logprobs_val: TokenIdsLogprobValues output_token_ids_logprobs_idx: TokenIdsLogprobIndices output_token_entropy_val: Optional[List[Optional[float]]] - # Detokenizer pass-through for BatchTokenIDOutput.output_token_sampling_*. - # None when sampling masks are not returned. - output_token_sampling_mask: Optional[List[List]] - output_token_sampling_logprobs: Optional[List[List]] + # Detokenizer pass-through for BatchTokenIDOutput.output_token_sampling_*; + # support-mode logprobs are aligned elementwise with the token IDs. + output_token_sampling_mask: Optional[List[List[List[int]]]] + output_token_sampling_logprobs: Optional[List[List[Union[float, List[float]]]]] # Hidden states output_hidden_states: OutputHiddenStates diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 66fa6fbe3fa9..840080d36cc8 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -165,6 +165,7 @@ ReturnHiddenStatesMode = Union[bool, Literal["last"]] +SamplingLogprobsMode = Literal["selected", "support"] def get_return_hidden_states_mode( @@ -989,6 +990,7 @@ def __init__( dllm_config: Optional[DllmConfig] = None, token_ids_logprob: List[int] = None, return_sampling_mask: bool = False, + sampling_logprobs_mode: SamplingLogprobsMode = "selected", return_flat_raw_top_logprobs: bool = False, stream: bool = False, origin_input_ids_unpadded: Optional[array[int]] = None, @@ -1226,6 +1228,7 @@ def __init__( self.temp_scaled_logprobs = False self.top_p_normalized_logprobs = False self.return_sampling_mask = return_sampling_mask + self.sampling_logprobs_mode = sampling_logprobs_mode self.return_flat_raw_top_logprobs = return_flat_raw_top_logprobs # Logprobs (return values) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f0d8533e03ff..95a5e2fbcaa0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2752,6 +2752,7 @@ def handle_generate_request( top_logprobs_num=recv_req.top_logprobs_num, token_ids_logprob=recv_req.token_ids_logprob, return_sampling_mask=recv_req.return_sampling_mask, + sampling_logprobs_mode=recv_req.sampling_logprobs_mode, return_flat_raw_top_logprobs=recv_req.return_flat_raw_top_logprobs, stream=recv_req.stream, lora_id=recv_req.lora_id, diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 7ce4572f594a..c8694982f5db 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -1191,8 +1191,35 @@ def materialize_sampling_mask_output( logprobs = [None] * batch_size status_by_batch = [None] * batch_size token_ids = sampling_output.token_ids.cpu() + support_logprobs = ( + None + if sampling_output.support_logprobs is None + else sampling_output.support_logprobs.cpu() + ) packed_width = token_ids.shape[1] + support_row_count = sum( + reqs[batch_index].sampling_logprobs_mode == "support" + for batch_index in batch_indices + ) + if support_row_count and support_logprobs is None: + raise RuntimeError( + "Support sampling logprobs were requested but not returned by " + "the sampler." + ) + expected_support_shape = (support_row_count, packed_width) + if support_logprobs is not None and ( + tuple(support_logprobs.shape) != expected_support_shape + ): + raise RuntimeError( + "Sampling support IDs and logprobs have incompatible packed " + f"shapes; expected {expected_support_shape}, got " + f"{tuple(support_logprobs.shape)}." + ) + support_row = 0 for row, batch_index in enumerate(batch_indices): + returns_support_logprobs = ( + reqs[batch_index].sampling_logprobs_mode == "support" + ) status = int(statuses[row]) length = int(lengths[row]) if status == SamplingMaskStatus.OK and not (0 <= length <= packed_width): @@ -1200,7 +1227,14 @@ def materialize_sampling_mask_output( status_by_batch[batch_index] = status if status == SamplingMaskStatus.OK: masks[batch_index] = token_ids[row, :length].tolist() - logprobs[batch_index] = float(selected_logprobs[row]) + if returns_support_logprobs: + logprobs[batch_index] = support_logprobs[ + support_row, :length + ].tolist() + else: + logprobs[batch_index] = float(selected_logprobs[row]) + if returns_support_logprobs: + support_row += 1 output.next_token_sampling_mask_idx = masks output.next_token_sampling_logprobs = logprobs diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 197644589244..42382f3024ab 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -374,8 +374,8 @@ class _GenerationStreamAccumulator: input_token_ids_logprobs_idx: Optional[list] = None output_token_ids_logprobs_val: Optional[list] = None output_token_ids_logprobs_idx: Optional[list] = None - output_token_sampling_mask: Optional[list] = None - output_token_sampling_logprobs: Optional[list] = None + output_token_sampling_mask: Optional[list[list[list[int]]]] = None + output_token_sampling_logprobs: Optional[list[list[float | list[float]]]] = None # Rust server mode: the Rust detokenizer reconstructs text/ids from the raw # output tokens itself and never consumes the scheduler's incremental-detok # offsets (decode_ids / read_offset), so that per-step bookkeeping is skipped. diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 254e0ff11195..3f80a686267c 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -289,8 +289,12 @@ def get_crash_dump_output(self) -> Dict[Any, Any]: input_token_ids_logprobs_idx: List = dataclasses.field(default_factory=list) output_token_ids_logprobs_val: List = dataclasses.field(default_factory=list) output_token_ids_logprobs_idx: List = dataclasses.field(default_factory=list) - output_token_sampling_mask: List = dataclasses.field(default_factory=list) - output_token_sampling_logprobs: List = dataclasses.field(default_factory=list) + output_token_sampling_mask: List[List[int]] = dataclasses.field( + default_factory=list + ) + output_token_sampling_logprobs: List[Union[float, List[float]]] = dataclasses.field( + default_factory=list + ) # Cached flat-format prompt top logprob fields; rebuilt only when more # prefill chunks arrive, so streaming decode chunks reuse the payload. @@ -1377,6 +1381,11 @@ def _validate_one_request( "return_sampling_mask only supports DisallowedTokensLogitsProcessor " "among custom logit processors." ) + if obj.sampling_logprobs_mode == "support" and not obj.return_sampling_mask: + raise ValueError( + "sampling_logprobs_mode='support' requires " + "return_sampling_mask=true." + ) def _validate_mm_limits( self, obj: Union[GenerateReqInput, EmbeddingReqInput] @@ -1494,6 +1503,7 @@ def _create_tokenized_object( top_logprobs_num=obj.top_logprobs_num, token_ids_logprob=obj.token_ids_logprob, return_sampling_mask=obj.return_sampling_mask, + sampling_logprobs_mode=obj.sampling_logprobs_mode, return_flat_raw_top_logprobs=obj.return_flat_raw_top_logprobs, stream=obj.stream, rid=obj.rid, diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 322dba9efeac..ed40ce3694a5 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -304,6 +304,11 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict: if sampling_mask_output is None else sampling_mask_output.selected_logprobs ), + "sampling_mask_support_logprobs": ( + None + if sampling_mask_output is None + else sampling_mask_output.support_logprobs + ), "sampling_mask_statuses": ( None if sampling_mask_output is None else sampling_mask_output.statuses ), @@ -325,6 +330,7 @@ def get_logprob_from_pp_outputs( token_ids=next_pp_outputs["sampling_mask_token_ids"], lengths=next_pp_outputs["sampling_mask_lengths"], selected_logprobs=next_pp_outputs["sampling_mask_selected_logprobs"], + support_logprobs=next_pp_outputs["sampling_mask_support_logprobs"], statuses=next_pp_outputs["sampling_mask_statuses"], ) logits_output = LogitsProcessorOutput( diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 8440c426fcb2..a37efcd31809 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -95,7 +95,9 @@ class SamplingBatchInfo: # CPU request flags and their derived device indices. The flags make batch # filtering cheap; the indices keep sampler work limited to opted-in rows. return_sampling_masks: Optional[List[bool]] = None + return_sampling_support_logprobs: Optional[List[bool]] = None sampling_mask_batch_indices: Optional[torch.Tensor] = None + sampling_support_logprobs_capture_indices: Optional[torch.Tensor] = None # Device device: str = "cuda" @@ -185,9 +187,21 @@ def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int): and any(r.custom_logit_processor for r in reqs) # check the flag first. ) # then check the requests. return_sampling_masks = [r.return_sampling_mask for r in reqs] + return_sampling_support_logprobs = [ + r.return_sampling_mask + and getattr(r, "sampling_logprobs_mode", "selected") == "support" + for r in reqs + ] sampling_mask_batch_indices = cls._make_sampling_mask_batch_indices( return_sampling_masks, device ) + sampling_support_logprobs_capture_indices = ( + cls._make_sampling_support_logprobs_capture_indices( + return_sampling_masks, + return_sampling_support_logprobs, + device, + ) + ) if has_custom_logit_processor: # Merge the same type of custom logit processors together @@ -253,7 +267,11 @@ def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int): device=device, logit_bias=logit_bias, return_sampling_masks=return_sampling_masks, + return_sampling_support_logprobs=return_sampling_support_logprobs, sampling_mask_batch_indices=sampling_mask_batch_indices, + sampling_support_logprobs_capture_indices=( + sampling_support_logprobs_capture_indices + ), ) ret.adjusted_from_schedule_batch(batch, vocab_size) return ret @@ -284,6 +302,30 @@ def _make_sampling_mask_batch_indices( return None return _rows_to_device_indices(indices, device) + @staticmethod + def _make_sampling_support_logprobs_capture_indices( + return_sampling_masks: List[bool], + return_sampling_support_logprobs: Optional[List[bool]], + device: str, + ) -> Optional[torch.Tensor]: + """Map support-mode requests to rows in the compact mask capture.""" + if return_sampling_support_logprobs is None: + return None + capture_index = -1 + indices = [] + for return_mask, return_support in zip( + return_sampling_masks, + return_sampling_support_logprobs, + strict=True, + ): + if return_mask: + capture_index += 1 + if return_support: + indices.append(capture_index) + if not indices: + return None + return _rows_to_device_indices(indices, device) + def __len__(self): return len(self.temperatures) @@ -394,12 +436,23 @@ def filter_batch(self, keep_indices: List[int], keep_indices_device: torch.Tenso self.return_sampling_masks = [ self.return_sampling_masks[i] for i in keep_indices ] + if self.return_sampling_support_logprobs is not None: + self.return_sampling_support_logprobs = [ + self.return_sampling_support_logprobs[i] for i in keep_indices + ] if self.sampling_mask_batch_indices is not None: self.sampling_mask_batch_indices = ( self._make_sampling_mask_batch_indices( self.return_sampling_masks, self.device ) ) + self.sampling_support_logprobs_capture_indices = ( + self._make_sampling_support_logprobs_capture_indices( + self.return_sampling_masks, + self.return_sampling_support_logprobs, + self.device, + ) + ) self.adjusted_filter_batch(keep_indices, keep_indices_device) @@ -475,6 +528,16 @@ def merge_batch(self, other: SamplingBatchInfo): self.return_sampling_masks = ( self.return_sampling_masks or [False] * self_len ) + (other.return_sampling_masks or [False] * other_len) + self.return_sampling_support_logprobs = ( + self.return_sampling_support_logprobs or [False] * self_len + ) + (other.return_sampling_support_logprobs or [False] * other_len) + self.sampling_support_logprobs_capture_indices = ( + self._make_sampling_support_logprobs_capture_indices( + self.return_sampling_masks, + self.return_sampling_support_logprobs, + self.device, + ) + ) # Note: because the __len()__ operator is defined on the temperatures tensor, # please make sure any merge operation with len(self) or len(other) is done before diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index c23c9653c6e3..7f28b8594a70 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -310,6 +310,7 @@ def create_req( top_logprobs_num=req.top_logprobs_num, token_ids_logprob=req.token_ids_logprob, return_sampling_mask=req.return_sampling_mask, + sampling_logprobs_mode=req.sampling_logprobs_mode, vocab_size=vocab_size, eos_token_ids=eos_token_ids, require_reasoning=req.require_reasoning, diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py index f43faa1c8e85..248f3e100954 100644 --- a/test/registered/sampling/test_sampling_mask.py +++ b/test/registered/sampling/test_sampling_mask.py @@ -111,13 +111,20 @@ def _materialize(self, sampled, capture, requested_rows): output = LogitsProcessorOutput( next_token_logits=None, sampling_mask_output=self.sampler._build_sampling_mask_output( - sampled, capture + sampled, + capture, + support_capture_indices=torch.arange( + len(requested_rows), device=sampled.device + ), ), ) output.sampling_mask_output.map_device_tensors(lambda tensor: tensor.cpu()) SchedulerBatchResultProcessor.materialize_sampling_mask_output( [ - SimpleNamespace(return_sampling_mask=i in requested_rows) + SimpleNamespace( + return_sampling_mask=i in requested_rows, + sampling_logprobs_mode="support", + ) for i in range(len(sampled)) ], output, @@ -132,13 +139,21 @@ def test_min_p_capture_matches_filtered_support_and_logprob(self): sampled, capture = self._sample( probs, backend, top_k=3, top_p=1.0, min_p=0.6 ) - output = self.sampler._build_sampling_mask_output(sampled, capture) + output = self.sampler._build_sampling_mask_output( + sampled, + capture, + support_capture_indices=torch.tensor([0], device="cuda"), + ) self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK]) self.assertEqual(output.lengths.tolist(), [2]) - self.assertEqual(set(output.token_ids[0, :2].tolist()), {0, 1}) + support = output.token_ids[0, :2].tolist() + self.assertEqual(set(support), {0, 1}) + sampled_index = support.index(sampled.item()) expected = (0.4 if sampled.item() == 0 else 0.3) / 0.7 self.assertAlmostEqual( - output.selected_logprobs.item(), math.log(expected), places=6 + output.support_logprobs[0, sampled_index].item(), + math.log(expected), + places=6, ) def test_hard_exclusion_replay_in_mixed_batch(self): @@ -189,11 +204,14 @@ def test_hard_exclusion_replay_in_mixed_batch(self): ) output = self._materialize(sampled, capture, requested_rows=[0, 1]) support = output.next_token_sampling_mask_idx[0] + sampling_logprobs = output.next_token_sampling_logprobs[0] self.assertEqual(set(support), {0, 1, 3}) self.assertIn(int(sampled[0]), support) expected = original[0, sampled[0]] - original[0, support].logsumexp(0) self.assertAlmostEqual( - output.next_token_sampling_logprobs[0], expected.item(), places=5 + sampling_logprobs[support.index(int(sampled[0]))], + expected.item(), + places=5, ) self.assertIn(2, output.next_token_sampling_mask_idx[1]) @@ -318,6 +336,7 @@ def _post_generate( top_logprobs_num=0, custom_logit_processor=None, stream=False, + sampling_logprobs_mode="support", ): payload = { "text": "The capital of France is", @@ -330,6 +349,8 @@ def _post_generate( "return_sampling_mask": return_sampling_mask, "stream": stream, } + if return_sampling_mask and sampling_logprobs_mode is not None: + payload["sampling_logprobs_mode"] = sampling_logprobs_mode if custom_logit_processor is not None: payload["custom_logit_processor"] = custom_logit_processor if return_logprob: @@ -341,13 +362,19 @@ def _post_generate( def _assert_sampling_masks(self, output_ids, meta_info): masks = meta_info["output_token_sampling_mask"] + sampling_logprobs = meta_info["output_token_sampling_logprobs"] self.assertEqual(len(masks), len(output_ids)) - self.assertEqual( - len(meta_info["output_token_sampling_logprobs"]), len(output_ids) - ) - for token_id, mask in zip(output_ids, masks): + self.assertEqual(len(sampling_logprobs), len(output_ids)) + for token_id, mask, logprobs in zip( + output_ids, masks, sampling_logprobs, strict=True + ): self.assertIn(token_id, mask) self.assertEqual(len(mask), len(set(mask))) + self.assertEqual(len(mask), len(logprobs)) + self.assertTrue(all(math.isfinite(logprob) for logprob in logprobs)) + self.assertAlmostEqual( + sum(math.exp(logprob) for logprob in logprobs), 1.0, delta=1e-5 + ) return masks def _generate_sampling_masks(self, sampling_params): @@ -397,6 +424,7 @@ def test_disallowed_tokens_with_replay(self): meta = output["meta_info"] token = output["output_ids"][0] mask = meta["output_token_sampling_mask"][0] + sampling_logprobs = meta["output_token_sampling_logprobs"][0] self.assertTrue(set(mask).isdisjoint(blocked)) self.assertIn(token, mask) probs = { @@ -404,7 +432,22 @@ def test_disallowed_tokens_with_replay(self): } expected = math.log(probs[token] / sum(probs[tid] for tid in mask)) self.assertAlmostEqual( - meta["output_token_sampling_logprobs"][0], expected, delta=1e-2 + sampling_logprobs[mask.index(token)], expected, delta=1e-2 + ) + + def test_sampling_logprobs_default_to_selected_token(self): + response = self._post_generate( + {"top_k": _TOP_K, "top_p": _TOP_P}, + sampling_logprobs_mode=None, + ) + self.assertEqual(response.status_code, 200, response.text) + meta = response.json()["meta_info"] + self.assertTrue(meta["output_token_sampling_logprobs"]) + self.assertTrue( + all( + isinstance(logprob, float) + for logprob in meta["output_token_sampling_logprobs"] + ) ) def test_rejected_processors_do_not_break_generation(self): @@ -438,7 +481,7 @@ def test_generate_returns_greedy_singleton_mask(self): self.assertTrue(all(len(mask) == 1 for mask in masks)) def test_sampling_mask_matches_topk_logprobs(self): - """Check the returned mask and its renormalized logprobs. + """Check the returned mask and its aligned behavior logprobs. We get a wide prefix of full-vocab logprobs via ``return_logprob`` so cutoff ties that extend beyond ``top_k`` are visible. With @@ -449,7 +492,7 @@ def test_sampling_mask_matches_topk_logprobs(self): 2. every mask token is in the returned top logprobs and at or above the top-k cutoff (ties at the cutoff survive, so the mask may exceed ``top_k``), - 3. sampling_logprob == log(p[sampled] / sum(p[t] for t in mask)). + 3. each sampling logprob equals log(p[token] / sum(p[t] for t in mask)). """ top_k, top_p = _TOP_K, _TOP_P response = self._post_generate( @@ -468,8 +511,8 @@ def test_sampling_mask_matches_topk_logprobs(self): self.assertEqual(len(top_logprobs), len(output_ids)) - for output_id, mask, mask_logprob, step_top_logprobs in zip( - output_ids, sampling_masks, sampling_logprobs, top_logprobs + for output_id, mask, mask_logprobs, step_top_logprobs in zip( + output_ids, sampling_masks, sampling_logprobs, top_logprobs, strict=True ): probs = { int(tid): math.exp(logprob) for logprob, tid, _ in step_top_logprobs @@ -484,8 +527,9 @@ def test_sampling_mask_matches_topk_logprobs(self): self.assertGreaterEqual(probs[token_id], top_k_cutoff * (1 - 1e-3)) support_mass = sum(probs[token_id] for token_id in mask_set) - expected_logprob = math.log(probs[output_id] / support_mass) - self.assertAlmostEqual(mask_logprob, expected_logprob, delta=1e-2) + for token_id, sampling_logprob in zip(mask, mask_logprobs, strict=True): + expected_logprob = math.log(probs[token_id] / support_mass) + self.assertAlmostEqual(sampling_logprob, expected_logprob, delta=1e-2) def test_chat_completions_returns_sampling_mask(self): response = requests.post( @@ -499,6 +543,7 @@ def test_chat_completions_returns_sampling_mask(self): "max_tokens": _MAX_NEW_TOKENS, "ignore_eos": True, "return_sampling_mask": True, + "sampling_logprobs_mode": "support", "return_meta_info": True, "return_token_ids": True, }, @@ -540,29 +585,99 @@ def setUp(self): self.sampler.tp_sync_group = None self.sampler.cp_sync_group = None - def test_selected_token_must_have_positive_captured_weight(self): + def test_sampled_token_must_have_positive_captured_weight(self): for token_ids in (None, torch.tensor([[2, 1, 0]], dtype=torch.int32)): with self.subTest(sorted_capture=token_ids is not None): capture = _SamplingMaskCapture( batch_rows=torch.tensor([0]), weights=torch.tensor([[0.7, 0.3, 0.0]]), token_ids=token_ids, - selected_weight=None, ) selected = torch.tensor([2 if token_ids is None else 0]) - output = self.sampler._build_sampling_mask_output(selected, capture) + output = self.sampler._build_sampling_mask_output( + selected, capture, support_capture_indices=None + ) self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.INVALID]) - def test_synced_token_logprob_is_recomputed_from_capture(self): + def test_selected_mode_does_not_build_support_logprobs(self): + capture = _SamplingMaskCapture( + batch_rows=torch.tensor([0]), + weights=torch.tensor([[0.6, 0.4, 0.0]]), + token_ids=None, + ) + output = self.sampler._build_sampling_mask_output( + torch.tensor([0]), capture, support_capture_indices=None + ) + self.assertIsNone(output.support_logprobs) + self.assertAlmostEqual(output.selected_logprobs.item(), math.log(0.6)) + + def test_support_logprobs_align_with_packed_token_ids(self): + captures = ( + _SamplingMaskCapture( + batch_rows=torch.tensor([0]), + weights=torch.tensor([[0.1, 0.6, 0.0, 0.3]]), + token_ids=None, + ), + _SamplingMaskCapture( + batch_rows=torch.tensor([0]), + weights=torch.tensor([[0.6, 0.3, 0.1, 0.0]]), + token_ids=torch.tensor([[1, 3, 0, 2]], dtype=torch.int32), + ), + ) + for capture in captures: + with self.subTest(sorted_capture=capture.token_ids is not None): + output = self.sampler._build_sampling_mask_output( + torch.tensor([3]), + capture, + support_capture_indices=torch.tensor([0]), + ) + self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK]) + self.assertEqual(output.lengths.tolist(), [3]) + self.assertEqual(output.token_ids[0, :3].tolist(), [1, 3, 0]) + torch.testing.assert_close( + output.support_logprobs[0, :3].exp(), + torch.tensor([0.6, 0.3, 0.1]), + ) + + def test_support_logprobs_only_pack_support_mode_rows(self): + capture = _SamplingMaskCapture( + batch_rows=torch.tensor([0, 1]), + weights=torch.tensor( + [ + [0.7, 0.3, 0.0], + [0.1, 0.6, 0.3], + ] + ), + token_ids=None, + ) + output = self.sampler._build_sampling_mask_output( + torch.tensor([0, 1]), + capture, + support_capture_indices=torch.tensor([1]), + ) + self.assertEqual(tuple(output.selected_logprobs.shape), (2,)) + self.assertEqual(tuple(output.support_logprobs.shape), (1, 3)) + torch.testing.assert_close( + output.support_logprobs[0].exp(), + torch.tensor([0.6, 0.3, 0.1]), + ) + + def test_synced_token_is_validated_against_capture(self): capture = _SamplingMaskCapture( batch_rows=torch.tensor([0]), weights=torch.tensor([[0.6, 0.2, 0.0]]), token_ids=torch.tensor([[2, 1, 0]], dtype=torch.int32), - selected_weight=None, ) - output = self.sampler._build_sampling_mask_output(torch.tensor([1]), capture) + output = self.sampler._build_sampling_mask_output( + torch.tensor([1]), + capture, + support_capture_indices=torch.tensor([0]), + ) self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK]) - self.assertAlmostEqual(output.selected_logprobs.item(), math.log(0.25)) + support = output.token_ids[0, :2].tolist() + self.assertAlmostEqual( + output.support_logprobs[0, support.index(1)].item(), math.log(0.25) + ) def test_greedy_device_output_survives_async_copy(self): from sglang.srt.managers.utils import GenerationBatchResult @@ -571,7 +686,9 @@ def test_greedy_device_output_survives_async_copy(self): output = LogitsProcessorOutput( next_token_logits=None, sampling_mask_output=self.sampler._build_greedy_sampling_mask_output( - torch.tensor([0, 2], device="cuda"), tokens + torch.tensor([0, 2], device="cuda"), + tokens, + support_capture_indices=torch.tensor([0, 1], device="cuda"), ), ) result = GenerationBatchResult( @@ -582,13 +699,16 @@ def test_greedy_device_output_survives_async_copy(self): self.assertEqual(output.sampling_mask_output.token_ids.device.type, "cpu") SchedulerBatchResultProcessor.materialize_sampling_mask_output( [ - SimpleNamespace(return_sampling_mask=flag) + SimpleNamespace( + return_sampling_mask=flag, + sampling_logprobs_mode="support", + ) for flag in (True, False, True) ], output, ) self.assertEqual(output.next_token_sampling_mask_idx, [[3], None, [5]]) - self.assertEqual(output.next_token_sampling_logprobs, [0.0, None, 0.0]) + self.assertEqual(output.next_token_sampling_logprobs, [[0.0], None, [0.0]]) def test_overflow_never_materializes_a_partial_mask(self): # Simulate a top-k cutoff tie: a nominal top_k below the cap can still @@ -597,11 +717,12 @@ def test_overflow_never_materializes_a_partial_mask(self): batch_rows=torch.tensor([0]), weights=torch.tensor([[0.2, 0.2, 0.2, 0.2, 0.2]]), token_ids=None, - selected_weight=torch.tensor([0.2]), ) sampling_output = self.sampler._build_sampling_mask_output( - torch.tensor([0]), capture + torch.tensor([0]), + capture, + support_capture_indices=torch.tensor([0]), ) output = LogitsProcessorOutput( @@ -609,7 +730,13 @@ def test_overflow_never_materializes_a_partial_mask(self): sampling_mask_output=sampling_output, ) SchedulerBatchResultProcessor.materialize_sampling_mask_output( - [SimpleNamespace(return_sampling_mask=True)], output + [ + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="support", + ) + ], + output, ) self.assertEqual( output.next_token_sampling_mask_status, @@ -691,12 +818,19 @@ def _check_parallel_config(self, *, tp_size, pp_size): self.assertEqual(meta["output_token_sampling_mask_length"], 4) self.assertEqual(len(masks), 4) self.assertEqual(len(logprobs), 4) - for token_id, mask, logprob in zip(token_ids, masks, logprobs): + for token_id, mask, step_logprobs in zip( + token_ids, masks, logprobs, strict=True + ): self.assertIn(token_id, mask) self.assertEqual(len(mask), len(set(mask))) self.assertLessEqual(len(mask), 64) - self.assertTrue(math.isfinite(logprob)) - self.assertLessEqual(logprob, 0.0) + self.assertEqual(len(mask), len(step_logprobs)) + self.assertTrue( + all(math.isfinite(logprob) for logprob in step_logprobs) + ) + self.assertTrue( + all(logprob <= 0.0 for logprob in step_logprobs) + ) if return_logprob: self.assertEqual(len(meta["output_token_logprobs"]), 4) @@ -721,6 +855,7 @@ def _generate(self, *, return_sampling_mask, return_logprob): "ignore_eos": True, }, "return_sampling_mask": return_sampling_mask, + "sampling_logprobs_mode": "support", "return_logprob": return_logprob, }, timeout=120, diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 95bffc6f74cb..680147ebb767 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -710,7 +710,8 @@ def _make_req( seed, metadata_buffer_index=0, sampling_mask=None, - sampling_logprob=None, + sampling_logprobs=None, + sampling_logprobs_mode="support", ): return SimpleNamespace( metadata_buffer_index=metadata_buffer_index, @@ -722,11 +723,12 @@ def _make_req( multimodal_inputs=None, return_logprob=False, return_sampling_mask=sampling_mask is not None, + sampling_logprobs_mode=sampling_logprobs_mode, output_token_sampling_mask=( None if sampling_mask is None else [sampling_mask] ), output_token_sampling_logprobs=( - None if sampling_logprob is None else [sampling_logprob] + None if sampling_logprobs is None else [sampling_logprobs] ), hidden_states_tensor=torch.tensor([1.0, 2.0]), output_topk_p=torch.tensor([1.0]), @@ -777,7 +779,7 @@ def test_sampling_mask_metadata_is_opt_in(self): self._make_req( None, sampling_mask=[7, 8, 9] if enabled else None, - sampling_logprob=-1.25 if enabled else None, + sampling_logprobs=[-1.25, -1.5, -2.0] if enabled else None, ) ) schemas.append(buffers.get_buf_infos()) @@ -785,10 +787,10 @@ def test_sampling_mask_metadata_is_opt_in(self): self.assertEqual( buffers.output_token_sampling_mask_idx.shape, (1, 3) ) - length, mask, logprob = buffers.get_buf(0)[6:9] + length, mask, logprobs = buffers.get_buf(0)[6:9] self.assertEqual(length[0].item(), 3) self.assertEqual(mask.tolist(), [7, 8, 9]) - self.assertAlmostEqual(logprob[0].item(), -1.25) + self.assertEqual(logprobs.tolist(), [-1.25, -1.5, -2.0]) else: self.assertIsNone(buffers.output_token_sampling_mask_len) self.assertIsNone(buffers.output_token_sampling_mask_idx) @@ -797,7 +799,45 @@ def test_sampling_mask_metadata_is_opt_in(self): disabled_ptrs, _, disabled_sizes = schemas[0] enabled_ptrs, _, enabled_sizes = schemas[1] self.assertEqual(len(enabled_ptrs) - len(disabled_ptrs), 3) - self.assertEqual(sum(enabled_sizes) - sum(disabled_sizes), 3 * 4 + 128) + self.assertEqual(sum(enabled_sizes) - sum(disabled_sizes), 2 * 3 * 4 + 64) + + def test_sampling_mask_metadata_requires_aligned_logprobs(self): + with envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(True): + buffers = MetadataBuffers( + size=1, + hidden_size=2, + hidden_states_dtype=torch.float32, + max_sampling_mask_tokens=3, + ) + with self.assertRaisesRegex(RuntimeError, "must have equal lengths"): + buffers.set_buf( + self._make_req( + None, + sampling_mask=[7, 8, 9], + sampling_logprobs=[-0.5, -1.0], + ) + ) + + def test_sampling_mask_selected_logprob_uses_first_metadata_slot(self): + with envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(True): + buffers = MetadataBuffers( + size=1, + hidden_size=2, + hidden_states_dtype=torch.float32, + max_sampling_mask_tokens=3, + ) + buffers.set_buf( + self._make_req( + None, + sampling_mask=[7, 8, 9], + sampling_logprobs=-0.5, + sampling_logprobs_mode="selected", + ) + ) + length, mask, logprobs = buffers.get_buf(0)[6:9] + self.assertEqual(length[0].item(), 3) + self.assertEqual(mask.tolist(), [7, 8, 9]) + self.assertEqual(logprobs[0].item(), -0.5) def test_decode_input_requires_valid_seed_for_every_request(self): seeds = ( diff --git a/test/registered/unit/entrypoints/openai/test_protocol.py b/test/registered/unit/entrypoints/openai/test_protocol.py index c580e691f409..2fefdf865626 100644 --- a/test/registered/unit/entrypoints/openai/test_protocol.py +++ b/test/registered/unit/entrypoints/openai/test_protocol.py @@ -189,6 +189,7 @@ def test_basic_chat_completion_request(self): self.assertEqual(request.temperature, None) # default self.assertFalse(request.stream) # default self.assertFalse(request.return_sampling_mask) + self.assertEqual(request.sampling_logprobs_mode, "selected") self.assertEqual(request.tool_choice, "none") # default when no tools def test_image_content_hash_validation(self): diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 343f16354959..42115543cc56 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -610,11 +610,13 @@ def test_convert_to_internal_request_single(self): ) self.basic_req.return_sampling_mask = True + self.basic_req.sampling_logprobs_mode = "support" self.basic_req.return_meta_info = True adapted, processed = self.chat._convert_to_internal_request(self.basic_req) self.assertIsInstance(adapted, GenerateReqInput) self.assertFalse(adapted.stream) self.assertTrue(adapted.return_sampling_mask) + self.assertEqual(adapted.sampling_logprobs_mode, "support") self.assertEqual(adapted.session_id, "session-1") self.assertEqual(processed, self.basic_req) diff --git a/test/registered/unit/managers/test_auxiliary_output.py b/test/registered/unit/managers/test_auxiliary_output.py index 4ef9b389214b..c6c12a272e77 100644 --- a/test/registered/unit/managers/test_auxiliary_output.py +++ b/test/registered/unit/managers/test_auxiliary_output.py @@ -131,6 +131,7 @@ def test_sampling_mask_output_uses_generation_result_copy_path(): token_ids=torch.tensor([[3, 5]], dtype=torch.int32), lengths=torch.tensor([2], dtype=torch.int32), selected_logprobs=torch.tensor([-0.5]), + support_logprobs=torch.tensor([[-0.5, -1.0]]), statuses=torch.tensor([SamplingMaskStatus.OK], dtype=torch.int32), ) result = GenerationBatchResult( @@ -148,8 +149,10 @@ def test_sampling_mask_output_uses_generation_result_copy_path(): ) as copy_tensor: result.copy_to_cpu(return_logprob=False) - assert copy_tensor.call_count == 5 + assert copy_tensor.call_count == 6 assert sampling_output.token_ids.tolist() == [[3, 5]] + assert sampling_output.selected_logprobs.tolist() == [-0.5] + assert sampling_output.support_logprobs.tolist() == [[-0.5, -1.0]] assert sampling_output.lengths.tolist() == [2] assert sampling_output.statuses.tolist() == [SamplingMaskStatus.OK] assert result.copy_done.record_count == 1 @@ -160,6 +163,7 @@ def test_pipeline_sampling_mask_round_trip_without_logprobs(): token_ids=torch.tensor([[3, 5]], dtype=torch.int32), lengths=torch.tensor([2], dtype=torch.int32), selected_logprobs=torch.tensor([-0.5]), + support_logprobs=torch.tensor([[-0.5, -1.0]]), statuses=torch.tensor([SamplingMaskStatus.OK], dtype=torch.int32), ) result = GenerationBatchResult( @@ -172,7 +176,13 @@ def test_pipeline_sampling_mask_round_trip_without_logprobs(): object.__new__(Scheduler), result, SimpleNamespace(return_logprob=False) ) output, _, _ = get_logprob_from_pp_outputs(PPProxyTensors(payload)) - for name in ("token_ids", "lengths", "selected_logprobs", "statuses"): + for name in ( + "token_ids", + "lengths", + "selected_logprobs", + "support_logprobs", + "statuses", + ): torch.testing.assert_close( getattr(output.sampling_mask_output, name), getattr(sampling_output, name) ) diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py index 363515e10343..acb9f5a9cfdc 100644 --- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py +++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py @@ -52,6 +52,58 @@ def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProc class TestSamplingMaskMaterialization(CustomTestCase): + def test_selected_and_support_modes_share_one_batch(self): + output = LogitsProcessorOutput( + next_token_logits=None, + sampling_mask_output=SimpleNamespace( + token_ids=torch.tensor([[7, 8], [9, 10]], dtype=torch.int32), + lengths=torch.tensor([2, 2]), + selected_logprobs=torch.tensor([-0.5, -0.25]), + support_logprobs=torch.tensor([[-0.25, -1.5]]), + statuses=torch.tensor([SamplingMaskStatus.OK, SamplingMaskStatus.OK]), + ), + ) + SchedulerBatchResultProcessor.materialize_sampling_mask_output( + reqs=[ + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="selected", + ), + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="support", + ), + ], + output=output, + ) + self.assertEqual(output.next_token_sampling_mask_idx, [[7, 8], [9, 10]]) + self.assertEqual( + output.next_token_sampling_logprobs, + [-0.5, [-0.25, -1.5]], + ) + + def test_selected_mode_does_not_require_support_tensor(self): + output = LogitsProcessorOutput( + next_token_logits=None, + sampling_mask_output=SimpleNamespace( + token_ids=torch.tensor([[7, 8]], dtype=torch.int32), + lengths=torch.tensor([2]), + selected_logprobs=torch.tensor([-0.5]), + support_logprobs=None, + statuses=torch.tensor([SamplingMaskStatus.OK]), + ), + ) + SchedulerBatchResultProcessor.materialize_sampling_mask_output( + reqs=[ + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="selected", + ) + ], + output=output, + ) + self.assertEqual(output.next_token_sampling_logprobs, [-0.5]) + def test_packed_ids_are_copied_before_per_request_slicing(self): """Non-overlap capture must not perform one device copy per request.""" packed_ids = Mock() @@ -63,18 +115,51 @@ def test_packed_ids_are_copied_before_per_request_slicing(self): token_ids=packed_ids, lengths=torch.tensor([2, 1]), selected_logprobs=torch.tensor([-0.5, -0.25]), + support_logprobs=torch.tensor([[-0.5, -0.75, 0.0], [-0.25, 0.0, 0.0]]), statuses=torch.tensor([SamplingMaskStatus.OK, SamplingMaskStatus.OK]), ), ) SchedulerBatchResultProcessor.materialize_sampling_mask_output( - reqs=[SimpleNamespace(return_sampling_mask=x) for x in (True, False, True)], + reqs=[ + SimpleNamespace( + return_sampling_mask=x, + sampling_logprobs_mode="support", + ) + for x in (True, False, True) + ], output=output, ) packed_ids.cpu.assert_called_once_with() self.assertEqual(output.next_token_sampling_mask_idx, [[7, 8], None, [9]]) - self.assertEqual(output.next_token_sampling_logprobs, [-0.5, None, -0.25]) + self.assertEqual( + output.next_token_sampling_logprobs, + [[-0.5, -0.75], None, [-0.25]], + ) self.assertIsNone(output.sampling_mask_output) + def test_sampling_mask_ids_and_logprobs_must_have_identical_shapes(self): + output = LogitsProcessorOutput( + next_token_logits=None, + sampling_mask_output=SimpleNamespace( + token_ids=torch.tensor([[7, 8]], dtype=torch.int32), + lengths=torch.tensor([2]), + selected_logprobs=torch.tensor([-0.5]), + support_logprobs=torch.tensor([[-0.5]]), + statuses=torch.tensor([SamplingMaskStatus.OK]), + ), + ) + + with self.assertRaisesRegex(RuntimeError, "incompatible packed shapes"): + SchedulerBatchResultProcessor.materialize_sampling_mask_output( + reqs=[ + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="support", + ) + ], + output=output, + ) + class _PrefillReq: def __init__(self, *, rid: str, inflight_middle_chunks: int, return_hidden_states): diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index bd198317a634..1ef9aa68e29c 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -412,6 +412,20 @@ def setUp(self): rid=["id1", "id2"], ) + def test_sampling_logprobs_mode_normalization(self): + req = copy.deepcopy(self.base_req) + req.return_sampling_mask = [True, True] + req.sampling_logprobs_mode = ["selected", "support"] + + req.normalize_batch_and_arguments() + + self.assertEqual( + req.sampling_logprobs_mode, + ["selected", "support"], + ) + self.assertEqual(req[0].sampling_logprobs_mode, "selected") + self.assertEqual(req[1].sampling_logprobs_mode, "support") + def test_single_image_to_list_of_lists(self): """Test that a single image is converted to a list of single-image lists.""" req = copy.deepcopy(self.base_req) diff --git a/test/registered/unit/managers/test_multi_tokenizer_mixin.py b/test/registered/unit/managers/test_multi_tokenizer_mixin.py index 7468d1fc5042..7796146b7126 100644 --- a/test/registered/unit/managers/test_multi_tokenizer_mixin.py +++ b/test/registered/unit/managers/test_multi_tokenizer_mixin.py @@ -69,8 +69,8 @@ def _make_batch_str_output() -> BatchStrOutput: output_token_ids_logprobs_val=[[], []], output_token_ids_logprobs_idx=[[], []], output_token_entropy_val=[0.0, 0.0], - output_token_sampling_mask=[[], []], - output_token_sampling_logprobs=[[], []], + output_token_sampling_mask=[[[7, 8]], [[9]]], + output_token_sampling_logprobs=[[[-0.5, -1.0]], [[0.0]]], output_hidden_states=[None, None], routed_experts=[None, None], indexer_topk=[None, None], @@ -118,6 +118,17 @@ def test_batch_str_output_keeps_weight_versions_nested_per_request(self): [[WeightVersionSpan(version="v2", start=0, end=2)]], ) + def test_batch_str_output_keeps_sampling_distribution_aligned(self): + output = _make_batch_str_output() + + single_output = _handle_output_by_index(output, 0) + + self.assertEqual(single_output.output_token_sampling_mask, [[[7, 8]]]) + self.assertEqual( + single_output.output_token_sampling_logprobs, + [[[-0.5, -1.0]]], + ) + def test_batch_str_output_without_weight_versions_stays_none(self): """An output from an older server without the field splits into None.""" output = _make_batch_str_output() diff --git a/test/registered/unit/managers/test_sampling_mask_validation.py b/test/registered/unit/managers/test_sampling_mask_validation.py index a41dfe548595..5a530b91246b 100644 --- a/test/registered/unit/managers/test_sampling_mask_validation.py +++ b/test/registered/unit/managers/test_sampling_mask_validation.py @@ -29,12 +29,18 @@ def setUp(self): self.manager.validate_total_tokens = False self.manager.is_generation = True - def _validate(self, processor, return_sampling_mask=True): + def _validate( + self, + processor, + return_sampling_mask=True, + sampling_logprobs_mode="selected", + ): req = GenerateReqInput( input_ids=[1, 2, 3], sampling_params={"top_k": 10}, custom_logit_processor=processor, return_sampling_mask=return_sampling_mask, + sampling_logprobs_mode=sampling_logprobs_mode, ) self.manager._validate_one_request(req, req.input_ids) @@ -78,6 +84,16 @@ def test_other_processors_still_work_without_sampling_masks(self): Qwen3ThinkingBudgetLogitProcessor.to_str(), return_sampling_mask=False ) + def test_support_logprobs_require_sampling_mask(self): + with self.assertRaisesRegex(ValueError, "requires return_sampling_mask=true"): + self._validate( + None, + return_sampling_mask=False, + sampling_logprobs_mode="support", + ) + + self._validate(None, sampling_logprobs_mode="support") + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 016f6d2341e0..bef594b0250a 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/test/registered/unit/sampling/test_sampling_batch_info.py @@ -103,6 +103,32 @@ def test_len_matches_batch_size(self): class TestSamplingMaskBatchIndices(CustomTestCase): + def test_filter_and_merge_preserve_support_logprob_modes(self): + info = _make_info( + batch_size=3, + return_sampling_masks=[True, True, False], + return_sampling_support_logprobs=[False, True, False], + sampling_mask_batch_indices=torch.tensor([0, 1]), + ) + info.filter_batch([1, 2], torch.tensor([1, 2])) + self.assertEqual(info.return_sampling_support_logprobs, [True, False]) + self.assertEqual(info.sampling_support_logprobs_capture_indices.tolist(), [0]) + + other = _make_info( + batch_size=2, + return_sampling_masks=[True, True], + return_sampling_support_logprobs=[False, True], + sampling_mask_batch_indices=torch.tensor([0, 1]), + ) + info.merge_batch(other) + self.assertEqual( + info.return_sampling_support_logprobs, + [True, False, False, True], + ) + self.assertEqual( + info.sampling_support_logprobs_capture_indices.tolist(), [0, 2] + ) + def test_filter_removes_last_opted_in_row_then_merge_restores_capture(self): info = _make_info( batch_size=2, From e9225f7550970d81dff4499e93bfeeec61fe05d3 Mon Sep 17 00:00:00 2001 From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:55:54 +0000 Subject: [PATCH 2/5] Preserve selected sampling logprob path --- python/sglang/srt/layers/sampler.py | 46 +++++++++++++++++-- .../srt/sampling/sampling_batch_info.py | 3 +- .../registered/sampling/test_sampling_mask.py | 12 ++++- 3 files changed, 52 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index ed8d3d4760b4..d32d2e9757ad 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -88,10 +88,11 @@ def _trace_e2e_sampler(stage: str, **fields) -> None: class _SamplingMaskCapture(NamedTuple): - """Post-filter weights and their original batch-row mapping.""" + """Compact post-filter weights and their original batch-row mapping.""" weights: torch.Tensor token_ids: Optional[torch.Tensor] + selected_weight: Optional[torch.Tensor] batch_rows: torch.Tensor @@ -305,6 +306,11 @@ def forward( ) else: assert sampling_mask_capture is not None + if SYNC_TOKEN_IDS_ACROSS_TP or sampling_info.grammars: + # Token synchronization can replace the producer-selected token. + sampling_mask_capture = sampling_mask_capture._replace( + selected_weight=None + ) logits_output.sampling_mask_output = self._build_sampling_mask_output( batch_next_token_ids, sampling_mask_capture, @@ -343,9 +349,14 @@ def _sample_from_probs( ) if return_sampling_mask: capture_probs = select_capture_rows(probs) + capture_tokens = select_capture_rows(batch_next_token_ids) + selected_weight = torch.gather( + capture_probs, 1, capture_tokens.long().view(-1, 1) + ).squeeze(1) sampling_mask_capture = _SamplingMaskCapture( weights=capture_probs, token_ids=None, + selected_weight=selected_weight, batch_rows=capture_rows, ) else: @@ -363,15 +374,22 @@ def _sample_from_probs( if return_sampling_mask: capture_probs = select_capture_rows(probs) capture_min_ps = select_capture_rows(sampling_info.min_ps) + capture_tokens = select_capture_rows(batch_next_token_ids) min_p_thresholds = ( capture_probs.max(dim=-1).values * capture_min_ps ) filtered_probs = capture_probs.masked_fill( capture_probs < min_p_thresholds.view(-1, 1), 0 ) + selected_weight = torch.gather( + filtered_probs, + 1, + capture_tokens.long().view(-1, 1), + ).squeeze(1) sampling_mask_capture = _SamplingMaskCapture( weights=filtered_probs, token_ids=None, + selected_weight=selected_weight, batch_rows=capture_rows, ) else: @@ -389,6 +407,7 @@ def _sample_from_probs( capture_probs = select_capture_rows(probs) capture_top_ks = select_capture_rows(sampling_info.top_ks) capture_top_ps = select_capture_rows(sampling_info.top_ps) + capture_tokens = select_capture_rows(batch_next_token_ids) filtered_probs = capture_probs if sampling_info.need_top_k_sampling: filtered_probs = top_k_renorm_prob( @@ -402,9 +421,15 @@ def _sample_from_probs( filtered_probs = top_p_probs else: filtered_probs.masked_fill_(top_p_probs <= 0, 0) + selected_weight = torch.gather( + filtered_probs, + 1, + capture_tokens.long().view(-1, 1), + ).squeeze(1) sampling_mask_capture = _SamplingMaskCapture( weights=filtered_probs, token_ids=None, + selected_weight=selected_weight, batch_rows=capture_rows, ) elif backend == "pytorch": @@ -424,10 +449,12 @@ def _sample_from_probs( batch_next_token_ids, filtered_probs, token_ids, + selected_weight, ) = sample_result sampling_mask_capture = _SamplingMaskCapture( weights=select_capture_rows(filtered_probs), token_ids=select_capture_rows(token_ids), + selected_weight=select_capture_rows(selected_weight), batch_rows=capture_rows, ) else: @@ -479,6 +506,7 @@ def _build_sampling_mask_output( batch_indices = sampling_mask_capture.batch_rows weights = sampling_mask_capture.weights token_ids = sampling_mask_capture.token_ids + selected_weight = sampling_mask_capture.selected_weight sampled_tokens = batch_next_token_ids.index_select(0, batch_indices).view(-1, 1) sampled_tokens_int32 = sampled_tokens.to(torch.int32) if token_ids is None: @@ -495,14 +523,18 @@ def _build_sampling_mask_output( selected_from_weights = torch.gather( weights, 1, selected_positions ).squeeze(1) + if selected_weight is None: + selected_weight = selected_from_weights + support = weights > 0 support_mass = torch.where(support, weights, torch.zeros_like(weights)).sum( dim=-1, dtype=torch.float32 ) - selected_logprobs = torch.log(selected_from_weights.float() / support_mass) + selected_weight = selected_weight.float() + selected_logprobs = torch.log(selected_weight / support_mass) invalid = ~( sampled_in_capture - & (selected_from_weights > 0) + & (selected_weight > 0) & (support_mass > 0) & torch.isfinite(selected_logprobs) ) @@ -727,7 +759,8 @@ def top_k_top_p_min_p_sampling_from_probs_torch( with the sampling_seed of each request. By default, returns only sampled token IDs. With return_filtered_probs=True, - also returns the actual filtered weights and their token-ID permutation. + also returns the actual filtered weights, their token-ID permutation, and + the selected weights. """ probs_sort, probs_idx = probs.sort(dim=-1, descending=True) probs_sum = torch.cumsum(probs_sort, dim=-1) @@ -758,11 +791,14 @@ def top_k_top_p_min_p_sampling_from_probs_torch( logprobs.log_() sampled_index = multinomial_with_seed(logprobs, sampling_seed, positions) + if return_filtered_probs: + selected_weight = torch.gather(probs_sort, 1, sampled_index).view(-1) + # int32 range is enough to represent the token ids probs_idx = probs_idx.to(torch.int32) batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(-1) if return_filtered_probs: - return batch_next_token_ids, probs_sort, probs_idx + return batch_next_token_ids, probs_sort, probs_idx, selected_weight return batch_next_token_ids diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index a37efcd31809..2472f76a810c 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -188,8 +188,7 @@ def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int): ) # then check the requests. return_sampling_masks = [r.return_sampling_mask for r in reqs] return_sampling_support_logprobs = [ - r.return_sampling_mask - and getattr(r, "sampling_logprobs_mode", "selected") == "support" + r.return_sampling_mask and r.sampling_logprobs_mode == "support" for r in reqs ] sampling_mask_batch_indices = cls._make_sampling_mask_batch_indices( diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py index 248f3e100954..99676c575996 100644 --- a/test/registered/sampling/test_sampling_mask.py +++ b/test/registered/sampling/test_sampling_mask.py @@ -585,13 +585,14 @@ def setUp(self): self.sampler.tp_sync_group = None self.sampler.cp_sync_group = None - def test_sampled_token_must_have_positive_captured_weight(self): + def test_selected_token_must_have_positive_captured_weight(self): for token_ids in (None, torch.tensor([[2, 1, 0]], dtype=torch.int32)): with self.subTest(sorted_capture=token_ids is not None): capture = _SamplingMaskCapture( batch_rows=torch.tensor([0]), weights=torch.tensor([[0.7, 0.3, 0.0]]), token_ids=token_ids, + selected_weight=None, ) selected = torch.tensor([2 if token_ids is None else 0]) output = self.sampler._build_sampling_mask_output( @@ -604,6 +605,7 @@ def test_selected_mode_does_not_build_support_logprobs(self): batch_rows=torch.tensor([0]), weights=torch.tensor([[0.6, 0.4, 0.0]]), token_ids=None, + selected_weight=torch.tensor([0.6]), ) output = self.sampler._build_sampling_mask_output( torch.tensor([0]), capture, support_capture_indices=None @@ -617,11 +619,13 @@ def test_support_logprobs_align_with_packed_token_ids(self): batch_rows=torch.tensor([0]), weights=torch.tensor([[0.1, 0.6, 0.0, 0.3]]), token_ids=None, + selected_weight=torch.tensor([0.3]), ), _SamplingMaskCapture( batch_rows=torch.tensor([0]), weights=torch.tensor([[0.6, 0.3, 0.1, 0.0]]), token_ids=torch.tensor([[1, 3, 0, 2]], dtype=torch.int32), + selected_weight=torch.tensor([0.3]), ), ) for capture in captures: @@ -649,6 +653,7 @@ def test_support_logprobs_only_pack_support_mode_rows(self): ] ), token_ids=None, + selected_weight=torch.tensor([0.7, 0.6]), ) output = self.sampler._build_sampling_mask_output( torch.tensor([0, 1]), @@ -662,11 +667,12 @@ def test_support_logprobs_only_pack_support_mode_rows(self): torch.tensor([0.6, 0.3, 0.1]), ) - def test_synced_token_is_validated_against_capture(self): + def test_synced_token_logprob_is_recomputed_from_capture(self): capture = _SamplingMaskCapture( batch_rows=torch.tensor([0]), weights=torch.tensor([[0.6, 0.2, 0.0]]), token_ids=torch.tensor([[2, 1, 0]], dtype=torch.int32), + selected_weight=None, ) output = self.sampler._build_sampling_mask_output( torch.tensor([1]), @@ -674,6 +680,7 @@ def test_synced_token_is_validated_against_capture(self): support_capture_indices=torch.tensor([0]), ) self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK]) + self.assertAlmostEqual(output.selected_logprobs.item(), math.log(0.25)) support = output.token_ids[0, :2].tolist() self.assertAlmostEqual( output.support_logprobs[0, support.index(1)].item(), math.log(0.25) @@ -717,6 +724,7 @@ def test_overflow_never_materializes_a_partial_mask(self): batch_rows=torch.tensor([0]), weights=torch.tensor([[0.2, 0.2, 0.2, 0.2, 0.2]]), token_ids=None, + selected_weight=torch.tensor([0.2]), ) sampling_output = self.sampler._build_sampling_mask_output( From b309fe1335712e456921048daf98b48e10096f1b Mon Sep 17 00:00:00 2001 From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:22:42 +0000 Subject: [PATCH 3/5] Test session sampling logprob mode propagation --- test/registered/unit/mem_cache/test_session_token_share_unit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/registered/unit/mem_cache/test_session_token_share_unit.py b/test/registered/unit/mem_cache/test_session_token_share_unit.py index cf12e063c689..b762f5026a2e 100644 --- a/test/registered/unit/mem_cache/test_session_token_share_unit.py +++ b/test/registered/unit/mem_cache/test_session_token_share_unit.py @@ -39,6 +39,7 @@ def _recv(rid, input_ids, max_new_tokens=8): top_logprobs_num=0, token_ids_logprob=None, return_sampling_mask=False, + sampling_logprobs_mode="selected", require_reasoning=False, return_hidden_states=False, return_routed_experts=False, From 76443cab7f6da691f0094156e87627881476f328 Mon Sep 17 00:00:00 2001 From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com> Date: Thu, 24 Sep 2026 01:03:40 +0000 Subject: [PATCH 4/5] Address sampling logprob mode review --- python/sglang/srt/disaggregation/utils.py | 8 ------ .../sglang/srt/entrypoints/openai/protocol.py | 2 +- python/sglang/srt/managers/io_struct.py | 11 ++++---- .../batch_result_processor.py | 18 ------------- .../sglang/srt/managers/tokenizer_manager.py | 6 ++--- .../test_disaggregation_wire.py | 17 ------------- .../unit/entrypoints/openai/test_protocol.py | 2 +- ...st_batch_result_processor_hidden_states.py | 23 ----------------- .../unit/managers/test_io_struct.py | 4 +++ .../managers/test_sampling_mask_validation.py | 25 ++++++++++++------- 10 files changed, 31 insertions(+), 85 deletions(-) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 2a8c2aea9ab3..a8b9115ed91a 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -524,14 +524,6 @@ def set_buf(self, req: Req): ) if sampling_mask is not None and sampling_logprobs_row is not None: mask_len = len(sampling_mask) - if ( - req.sampling_logprobs_mode == "support" - and len(sampling_logprobs_row) != mask_len - ): - raise RuntimeError( - "Sampling mask IDs and logprobs must have equal lengths; " - f"got {mask_len} and {len(sampling_logprobs_row)}." - ) max_mask_len = self.output_token_sampling_mask_idx.shape[1] if mask_len > max_mask_len: raise RuntimeError( diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 19410851057c..88b33753a898 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -894,7 +894,7 @@ class ChatCompletionRequest(BaseModel): return_input_ids_in_sglext: bool = False return_output_ids_in_sglext: bool = False return_sampling_mask: bool = False - sampling_logprobs_mode: Literal["selected", "support"] = "selected" + sampling_logprobs_mode: Optional[Literal["selected", "support"]] = None reasoning_effort: ReasoningEffortType = Field( default=None, description="Constrains effort on reasoning for reasoning models. " diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index c69de94c66f3..6c24e7896115 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -236,10 +236,11 @@ class GenerateReqInput: # Whether to return each output token's sampling support and behavior logprob. return_sampling_mask: Optional[Union[List[bool], bool]] = None # Return either the selected token's behavior logprob or behavior logprobs for - # the full sampling support. Only used when return_sampling_mask is enabled. - sampling_logprobs_mode: Union[List[SamplingLogprobsMode], SamplingLogprobsMode] = ( - "selected" - ) + # the full sampling support. When omitted, selected mode is used if + # return_sampling_mask is enabled. + sampling_logprobs_mode: Optional[ + Union[List[Optional[SamplingLogprobsMode]], SamplingLogprobsMode] + ] = None # Whether to detokenize tokens in text in the returned logprobs. return_text_in_logprobs: bool = False # Return prompt top logprobs as flat arrays plus shape metadata instead of @@ -775,7 +776,7 @@ def normalize_param(param, default_value, param_name): self.return_sampling_mask, False, "return_sampling_mask" ) self.sampling_logprobs_mode = normalize_param( - self.sampling_logprobs_mode, "selected", "sampling_logprobs_mode" + self.sampling_logprobs_mode, None, "sampling_logprobs_mode" ) # Handle token_ids_logprob specially due to its nested structure diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index c8694982f5db..efaf0e45b1f3 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -1197,24 +1197,6 @@ def materialize_sampling_mask_output( else sampling_output.support_logprobs.cpu() ) packed_width = token_ids.shape[1] - support_row_count = sum( - reqs[batch_index].sampling_logprobs_mode == "support" - for batch_index in batch_indices - ) - if support_row_count and support_logprobs is None: - raise RuntimeError( - "Support sampling logprobs were requested but not returned by " - "the sampler." - ) - expected_support_shape = (support_row_count, packed_width) - if support_logprobs is not None and ( - tuple(support_logprobs.shape) != expected_support_shape - ): - raise RuntimeError( - "Sampling support IDs and logprobs have incompatible packed " - f"shapes; expected {expected_support_shape}, got " - f"{tuple(support_logprobs.shape)}." - ) support_row = 0 for row, batch_index in enumerate(batch_indices): returns_support_logprobs = ( diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 3f80a686267c..f6164c6d0172 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1381,9 +1381,9 @@ def _validate_one_request( "return_sampling_mask only supports DisallowedTokensLogitsProcessor " "among custom logit processors." ) - if obj.sampling_logprobs_mode == "support" and not obj.return_sampling_mask: + if obj.sampling_logprobs_mode is not None and not obj.return_sampling_mask: raise ValueError( - "sampling_logprobs_mode='support' requires " + "sampling_logprobs_mode can only be set when " "return_sampling_mask=true." ) @@ -1503,7 +1503,7 @@ def _create_tokenized_object( top_logprobs_num=obj.top_logprobs_num, token_ids_logprob=obj.token_ids_logprob, return_sampling_mask=obj.return_sampling_mask, - sampling_logprobs_mode=obj.sampling_logprobs_mode, + sampling_logprobs_mode=obj.sampling_logprobs_mode or "selected", return_flat_raw_top_logprobs=obj.return_flat_raw_top_logprobs, stream=obj.stream, rid=obj.rid, diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 680147ebb767..22e1ec977dcf 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -801,23 +801,6 @@ def test_sampling_mask_metadata_is_opt_in(self): self.assertEqual(len(enabled_ptrs) - len(disabled_ptrs), 3) self.assertEqual(sum(enabled_sizes) - sum(disabled_sizes), 2 * 3 * 4 + 64) - def test_sampling_mask_metadata_requires_aligned_logprobs(self): - with envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(True): - buffers = MetadataBuffers( - size=1, - hidden_size=2, - hidden_states_dtype=torch.float32, - max_sampling_mask_tokens=3, - ) - with self.assertRaisesRegex(RuntimeError, "must have equal lengths"): - buffers.set_buf( - self._make_req( - None, - sampling_mask=[7, 8, 9], - sampling_logprobs=[-0.5, -1.0], - ) - ) - def test_sampling_mask_selected_logprob_uses_first_metadata_slot(self): with envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(True): buffers = MetadataBuffers( diff --git a/test/registered/unit/entrypoints/openai/test_protocol.py b/test/registered/unit/entrypoints/openai/test_protocol.py index 2fefdf865626..fe312945c370 100644 --- a/test/registered/unit/entrypoints/openai/test_protocol.py +++ b/test/registered/unit/entrypoints/openai/test_protocol.py @@ -189,7 +189,7 @@ def test_basic_chat_completion_request(self): self.assertEqual(request.temperature, None) # default self.assertFalse(request.stream) # default self.assertFalse(request.return_sampling_mask) - self.assertEqual(request.sampling_logprobs_mode, "selected") + self.assertIsNone(request.sampling_logprobs_mode) self.assertEqual(request.tool_choice, "none") # default when no tools def test_image_content_hash_validation(self): diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py index acb9f5a9cfdc..477979c4adf7 100644 --- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py +++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py @@ -137,29 +137,6 @@ def test_packed_ids_are_copied_before_per_request_slicing(self): ) self.assertIsNone(output.sampling_mask_output) - def test_sampling_mask_ids_and_logprobs_must_have_identical_shapes(self): - output = LogitsProcessorOutput( - next_token_logits=None, - sampling_mask_output=SimpleNamespace( - token_ids=torch.tensor([[7, 8]], dtype=torch.int32), - lengths=torch.tensor([2]), - selected_logprobs=torch.tensor([-0.5]), - support_logprobs=torch.tensor([[-0.5]]), - statuses=torch.tensor([SamplingMaskStatus.OK]), - ), - ) - - with self.assertRaisesRegex(RuntimeError, "incompatible packed shapes"): - SchedulerBatchResultProcessor.materialize_sampling_mask_output( - reqs=[ - SimpleNamespace( - return_sampling_mask=True, - sampling_logprobs_mode="support", - ) - ], - output=output, - ) - class _PrefillReq: def __init__(self, *, rid: str, inflight_middle_chunks: int, return_hidden_states): diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index 1ef9aa68e29c..5444ff5bcc60 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -413,6 +413,10 @@ def setUp(self): ) def test_sampling_logprobs_mode_normalization(self): + default_req = copy.deepcopy(self.base_req) + default_req.normalize_batch_and_arguments() + self.assertEqual(default_req.sampling_logprobs_mode, [None, None]) + req = copy.deepcopy(self.base_req) req.return_sampling_mask = [True, True] req.sampling_logprobs_mode = ["selected", "support"] diff --git a/test/registered/unit/managers/test_sampling_mask_validation.py b/test/registered/unit/managers/test_sampling_mask_validation.py index 5a530b91246b..963bd94d2327 100644 --- a/test/registered/unit/managers/test_sampling_mask_validation.py +++ b/test/registered/unit/managers/test_sampling_mask_validation.py @@ -33,7 +33,7 @@ def _validate( self, processor, return_sampling_mask=True, - sampling_logprobs_mode="selected", + sampling_logprobs_mode=None, ): req = GenerateReqInput( input_ids=[1, 2, 3], @@ -84,15 +84,22 @@ def test_other_processors_still_work_without_sampling_masks(self): Qwen3ThinkingBudgetLogitProcessor.to_str(), return_sampling_mask=False ) - def test_support_logprobs_require_sampling_mask(self): - with self.assertRaisesRegex(ValueError, "requires return_sampling_mask=true"): - self._validate( - None, - return_sampling_mask=False, - sampling_logprobs_mode="support", - ) + def test_sampling_logprobs_mode_requires_sampling_mask(self): + self._validate(None, return_sampling_mask=False) + for mode in ("selected", "support"): + with self.subTest(mode=mode): + with self.assertRaisesRegex( + ValueError, "can only be set when return_sampling_mask=true" + ): + self._validate( + None, + return_sampling_mask=False, + sampling_logprobs_mode=mode, + ) - self._validate(None, sampling_logprobs_mode="support") + for mode in (None, "selected", "support"): + with self.subTest(mode=mode): + self._validate(None, sampling_logprobs_mode=mode) if __name__ == "__main__": From 6260d8d556732249fc92d56a121617f57f74a240 Mon Sep 17 00:00:00 2001 From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com> Date: Thu, 24 Sep 2026 03:25:55 +0000 Subject: [PATCH 5/5] Fix distributed sampling mask control request --- .../registered/sampling/test_sampling_mask.py | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py index 99676c575996..6094a9e9e788 100644 --- a/test/registered/sampling/test_sampling_mask.py +++ b/test/registered/sampling/test_sampling_mask.py @@ -851,22 +851,22 @@ def _check_parallel_config(self, *, tp_size, pp_size): process.wait(timeout=30) def _generate(self, *, return_sampling_mask, return_logprob): - response = requests.post( - DEFAULT_URL_FOR_TEST + "/generate", - json={ - "text": "The capital of France is", - "sampling_params": { - "temperature": 0.8, - "top_k": 8, - "top_p": 0.9, - "max_new_tokens": 4, - "ignore_eos": True, - }, - "return_sampling_mask": return_sampling_mask, - "sampling_logprobs_mode": "support", - "return_logprob": return_logprob, + payload = { + "text": "The capital of France is", + "sampling_params": { + "temperature": 0.8, + "top_k": 8, + "top_p": 0.9, + "max_new_tokens": 4, + "ignore_eos": True, }, - timeout=120, + "return_sampling_mask": return_sampling_mask, + "return_logprob": return_logprob, + } + if return_sampling_mask: + payload["sampling_logprobs_mode"] = "support" + response = requests.post( + DEFAULT_URL_FOR_TEST + "/generate", json=payload, timeout=120 ) self.assertEqual(response.status_code, 200, response.text) return response.json()