[Sampling] Add selected/support sampling logprob modes

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

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,8 +519,10 @@ 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)
                     max_mask_len = self.output_token_sampling_mask_idx.shape[1]
                     if mask_len > max_mask_len:
@@ -540,9 +544,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

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: Optional[Literal["selected", "support"]] = None
     reasoning_effort: ReasoningEffortType = Field(
         default=None,
         description="Constrains effort on reasoning for reasoning models. "

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,

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

python/sglang/srt/layers/sampler.py

@@ -176,6 +176,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,7 +299,9 @@ 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:
@@ -309,6 +314,7 @@ def forward(
                 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")
@@ -461,6 +467,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 +479,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,6 +500,7 @@ 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
@@ -549,18 +566,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,
         )
 

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,14 @@ 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. 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
@@ -768,6 +775,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, None, "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 +958,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 +1132,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 +1528,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 +1626,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

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)

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,

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -1191,16 +1191,32 @@ 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 = 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):
                 status = SamplingMaskStatus.INVALID
             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

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.