[Sampling] Support sampling masks with overlap scheduling (#36631)

Co-authored-by: ByronHsu <ByronHsu@users.noreply.github.com>
Co-authored-by: root <root@slurm-h200-208-179.slurm-compute.tenant-slurm.svc.cluster.local>
Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>

python/sglang/srt/arg_groups/field_order.py

@@ -496,6 +496,7 @@ POSITIONAL_FIELD_ORDER = (
     "return_hidden_states_mode",
     "enable_return_routed_experts",
     "enable_return_indexer_topk",
+    "sampling_mask_max_tokens",
     "disable_outlines_disk_cache",
     "enable_mis",
     "weight_cache_mode",

python/sglang/srt/arg_groups/fields/exec_.py

@@ -95,6 +95,13 @@ class ExecFeatures(msgspec.Struct):
         bool,
         "Enable returning indexer topk indices of layers with indexer with responses.",
     ] = False
+    sampling_mask_max_tokens: A[
+        int,
+        "The maximum number of token IDs in a returned sampling mask. Requests "
+        "are aborted if their realized sampling support exceeds this limit. "
+        "Use the same value on disaggregated prefill and decode nodes; clients "
+        "should set top_k below the limit to leave headroom for cutoff ties.",
+    ] = 4096
     disable_outlines_disk_cache: A[
         bool,
         "Disable disk cache of outlines to avoid possible crashes related to file system or high concurrency.",

python/sglang/srt/arg_groups/pipeline.py

@@ -95,9 +95,11 @@ def run_resolution_pipeline(server_args: Any) -> None:
     from sglang.srt.arg_groups.validation_hook import (
         validate_experimental_sgl_marlin,
         validate_prefill_decode_interval,
+        validate_sampling_mask_max_tokens,
     )
 
     validate_prefill_decode_interval(server_args)
+    validate_sampling_mask_max_tokens(server_args)
 
     # Reject an explicitly enabled but incompatible hardware runtime before
     # model path resolution, downloads, or the dummy-model short circuit.

python/sglang/srt/arg_groups/validation_hook.py

@@ -17,6 +17,7 @@ from sglang.srt.arg_groups.overrides import (
 from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
     parse_ib_device_config,
 )
+from sglang.srt.environ import envs
 from sglang.srt.runtime_context import get_platform
 from sglang.srt.utils.common import torch_release
 from sglang.srt.utils.runai_utils import is_runai_obj_uri
@@ -438,6 +439,22 @@ def validate_prefill_decode_interval(server_args: Any):
         raise ValueError("--prefill-decode-interval must be non-negative.")
 
 
+def validate_sampling_mask_max_tokens(server_args: Any):
+    if envs.SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.is_set():
+        raise ValueError(
+            "SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS is no longer supported. "
+            "Unset it. To enable sampling masks for disaggregated serving, set "
+            "SGLANG_ENABLE_DISAGG_SAMPLING_MASK=1 and use the same positive "
+            "--sampling-mask-max-tokens value on both prefill and decode servers."
+        )
+    cfg = resolving_view(server_args)
+    if cfg.sampling_mask_max_tokens <= 0:
+        raise ValueError(
+            "--sampling-mask-max-tokens must be positive "
+            f"(got {cfg.sampling_mask_max_tokens})."
+        )
+
+
 def check_two_batch_overlap(server_args: Any):
     # With no EP a2a backend, two-batch-overlap is only valid on the non-EP
     # DP TP-MoE path (overlapping the DP all_gatherv / reduce_scatterv with

python/sglang/srt/disaggregation/prefill.py

@@ -754,6 +754,10 @@ class SchedulerDisaggregationPrefillMixin:
             batch=batch,
             logits_output=logits_output,
         )
+        if logits_output is not None and logits_output.sampling_mask_output is not None:
+            self.batch_result_processor.materialize_sampling_mask_output(
+                batch.reqs, logits_output
+            )
 
         def advance_logprob_pt(i: int, req: Req) -> None:
             nonlocal logprob_pt
@@ -783,6 +787,27 @@ class SchedulerDisaggregationPrefillMixin:
                     advance_logprob_pt(i, req)
                     continue
 
+                sampling_mask_finish_reason = None
+                if req.return_sampling_mask:
+                    assert logits_output is not None
+                    statuses = logits_output.next_token_sampling_mask_status
+                    status = None if statuses is None else statuses[i]
+                    sampling_mask_finish_reason = (
+                        self.batch_result_processor.get_sampling_mask_finish_reason(
+                            status=status
+                        )
+                    )
+                if sampling_mask_finish_reason is not None:
+                    req.to_finish = sampling_mask_finish_reason
+                    req.time_stats.trace_ctx.abort(
+                        abort_info={"reason": sampling_mask_finish_reason.message}
+                    )
+                    if self._retire_aborted_prefill_result(req):
+                        req.time_stats.set_completion_time()
+                        aborted_reqs.append(req)
+                    advance_logprob_pt(i, req)
+                    continue
+
                 req.output_ids.append(next_token_id)
                 if req.grammar is not None:
                     try:

python/sglang/srt/disaggregation/utils.py

@@ -329,18 +329,14 @@ class MetadataBuffers:
         size: int,
         hidden_size: int,
         hidden_states_dtype: torch.dtype,
+        max_sampling_mask_tokens: int,
         max_top_logprobs_num: int = 128,
-        max_sampling_mask_tokens: Optional[int] = None,
         custom_mem_pool: torch.cuda.MemPool = None,
         output_dsa_topk_indices_dim: int = 0,
     ):
         self.custom_mem_pool = custom_mem_pool
         self.output_dsa_topk_indices_dim = output_dsa_topk_indices_dim
-        if max_sampling_mask_tokens is None:
-            max_sampling_mask_tokens = (
-                envs.SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.get()
-            )
-        self.enable_sampling_mask = max_sampling_mask_tokens > 0
+        self.enable_sampling_mask = envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.get()
         bootstrap_room_dtype = torch.uint64
         device = "cpu"
         if is_npu():
@@ -423,38 +419,23 @@ class MetadataBuffers:
             self.output_token_logprobs_idx,
             self.output_top_logprobs_val,
             self.output_top_logprobs_idx,
+            self.output_token_sampling_mask_len,
+            self.output_token_sampling_mask_idx,
+            self.output_token_sampling_logprobs,
+            self.output_topk_p,
+            self.output_topk_index,
+            self.output_hidden_states,
         ]
-        if self.enable_sampling_mask:
-            bufs.extend(
-                [
-                    self.output_token_sampling_mask_len,
-                    self.output_token_sampling_mask_idx,
-                    self.output_token_sampling_logprobs,
-                ]
-            )
-        bufs.extend(
-            [
-                self.output_topk_p,
-                self.output_topk_index,
-                self.output_hidden_states,
-            ]
-        )
         if self.output_dsa_topk_indices is not None:
             bufs.append(self.output_dsa_topk_indices)
         bufs.append(self.bootstrap_room)
+        bufs = [buf for buf in bufs if buf is not None]
         ptrs = [buf.data_ptr() for buf in bufs]
         data_lens = [buf.nbytes for buf in bufs]
         item_lens = [buf[0].nbytes for buf in bufs]
         return ptrs, data_lens, item_lens
 
     def get_buf(self, idx: int):
-        sampling_mask_len = None
-        sampling_mask_idx = None
-        sampling_logprobs = None
-        if self.enable_sampling_mask:
-            sampling_mask_len = self.output_token_sampling_mask_len[idx].clone()
-            sampling_mask_idx = self.output_token_sampling_mask_idx[idx].clone()
-            sampling_logprobs = self.output_token_sampling_logprobs[idx].clone()
         return (
             self.output_ids[idx].clone(),
             self.cached_tokens[idx].clone(),
@@ -462,9 +443,21 @@ class MetadataBuffers:
             self.output_token_logprobs_idx[idx].clone(),
             self.output_top_logprobs_val[idx].clone(),
             self.output_top_logprobs_idx[idx].clone(),
-            sampling_mask_len,
-            sampling_mask_idx,
-            sampling_logprobs,
+            (
+                self.output_token_sampling_mask_len[idx].clone()
+                if self.enable_sampling_mask
+                else None
+            ),
+            (
+                self.output_token_sampling_mask_idx[idx].clone()
+                if self.enable_sampling_mask
+                else None
+            ),
+            (
+                self.output_token_sampling_logprobs[idx].clone()
+                if self.enable_sampling_mask
+                else None
+            ),
             self.output_topk_p[idx].clone(),
             self.output_topk_index[idx].clone(),
             self.output_hidden_states[idx].clone(),
@@ -532,11 +525,6 @@ class MetadataBuffers:
                     device="cpu",
                 )
         if req.return_sampling_mask:
-            if not self.enable_sampling_mask:
-                raise RuntimeError(
-                    "return_sampling_mask with disaggregation requires "
-                    "SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS > 0."
-                )
             # Sentinel -1: the decode side records None for this handoff token.
             self.output_token_sampling_mask_len[req.metadata_buffer_index][0] = -1
             sampling_masks = req.output_token_sampling_mask
@@ -551,7 +539,7 @@ class MetadataBuffers:
                         raise RuntimeError(
                             f"Sampling mask length {mask_len} exceeds disaggregation "
                             f"metadata capacity {max_mask_len}. Increase "
-                            "SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS."
+                            "--sampling-mask-max-tokens."
                         )
                     self.output_token_sampling_mask_len[req.metadata_buffer_index][
                         0

python/sglang/srt/environ.py

@@ -666,6 +666,10 @@ class Envs:
     # computed dynamically at runtime based on cpu_count; see disaggregation backends.
     SGLANG_DISAGGREGATION_THREAD_POOL_SIZE = EnvInt(None)
     SGLANG_DISAGGREGATION_QUEUE_SIZE = EnvInt(4)
+    # Enable on both P and D with the same --sampling-mask-max-tokens value.
+    SGLANG_ENABLE_DISAGG_SAMPLING_MASK = EnvBool(False)
+    # Retained only to reject the removed setting during startup.
+    SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS = EnvInt(None)
     SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT = EnvInt(300)
     SGLANG_DISAGGREGATION_ZMQ_SEND_TIMEOUT = EnvInt(1)
     SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0)
@@ -679,7 +683,6 @@ class Envs:
     SGLANG_DISAGGREGATION_ZMQ_MAX_SOCKETS = EnvInt(16384)
     SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False)
     SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK = EnvBool(False)
-    SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS = EnvInt(0)
     SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120)
     # Deferred decode-side KV release: on abort, hold an in-flight request's KV
     # pages/slot until the prefill acks the transfer drained, or the timeout

python/sglang/srt/layers/logits_processor.py

@@ -16,6 +16,7 @@
 import dataclasses
 import logging
 from contextlib import contextmanager
+from enum import IntEnum
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -83,6 +84,30 @@ _UNQUANTIZED_LM_HEAD_METHODS = {
 _autotune_run_lm_head: Optional[bool] = None
 
 
+class SamplingMaskStatus(IntEnum):
+    """Ordered by severity so distributed MAX reaches one policy decision."""
+
+    OK = 0
+    OVERFLOW = 1
+    INVALID = 2
+
+
+@dataclasses.dataclass
+class SamplingMaskOutput:
+    """Tensor result for opted-in rows in batch order."""
+
+    token_ids: torch.Tensor
+    lengths: torch.Tensor
+    selected_logprobs: 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)
+        self.statuses = fn(self.statuses)
+
+
 def _trace_e2e_logits(stage: str, **fields) -> None:
     if not envs.SGLANG_TRACE_LOGITS_E2E.get():
         return
@@ -196,10 +221,12 @@ class LogitsProcessorOutput:
         List[Union[List[float], torch.Tensor]]
     ] = None
     next_token_token_ids_logprobs_idx: Optional[List] = None
-    # Sparse top-k/top-p/min-p support ids and selected-token logprob after
-    # truncation/renormalization. Only populated when requested.
+    # Post-filter support IDs, bounded by server capacity, and selected-token
+    # logprob 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_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
     # The logprobs of input tokens.        shape: [#token]

python/sglang/srt/layers/sampler.py

@@ -12,7 +12,11 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import (
     is_dp_attention_enabled,
 )
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.logits_processor import (
+    LogitsProcessorOutput,
+    SamplingMaskOutput,
+    SamplingMaskStatus,
+)
 from sglang.srt.layers.logprob_processor import (
     OutputLogprobProcessor,
 )
@@ -93,12 +97,23 @@ class _SamplingMaskCapture(NamedTuple):
     batch_rows: torch.Tensor
 
 
+def _select_sampling_mask_rows(
+    tensor: torch.Tensor, batch_indices: torch.Tensor
+) -> torch.Tensor:
+    """Select opted-in rows, returning the input when every row opts in."""
+    if batch_indices.numel() == tensor.shape[0]:
+        return tensor
+    return tensor.index_select(0, batch_indices)
+
+
 class Sampler(nn.Module):
     def __init__(self):
         super().__init__()
         self.tp_sync_group = get_tp_group().device_group
+        self.cp_sync_group = None
         if is_dp_attention_enabled():
             self.tp_sync_group = get_parallel().attn_tp_group.device_group
+            self.cp_sync_group = get_parallel().attn_cp_group.device_group
 
         self.rl_on_policy_target = get_exec().deterministic.rl_on_policy_target
         # In RL on-policy mode, deterministic inference is automatically enabled.
@@ -108,6 +123,7 @@ class Sampler(nn.Module):
         # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer.
         self.use_log_softmax_logprob = self.rl_on_policy_target is not None
         self.use_ascend_backend = get_exec().kernel.sampling_backend == "ascend"
+        self.sampling_mask_max_tokens = get_exec().features.sampling_mask_max_tokens
 
         self.output_logprob_processor = OutputLogprobProcessor()
 
@@ -157,7 +173,8 @@ class Sampler(nn.Module):
         _trace_e2e_sampler("preprocess_enter")
         logits = self._preprocess_logits(logits, sampling_info)
         _trace_e2e_sampler("preprocess_returned")
-        return_sampling_mask = any(sampling_info.return_sampling_masks or [])
+        sampling_mask_batch_indices = sampling_info.sampling_mask_batch_indices
+        return_sampling_mask = sampling_mask_batch_indices is not None
         sampling_mask_capture = None
 
         if sampling_info.is_all_greedy:
@@ -249,7 +266,6 @@ class Sampler(nn.Module):
                     sampling_info,
                     positions,
                     simple_sampling_case,
-                    return_sampling_mask=return_sampling_mask,
                 )
                 if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB:
                     logprobs = (
@@ -275,9 +291,12 @@ class Sampler(nn.Module):
         _trace_e2e_sampler("token_sync_returned")
 
         if return_sampling_mask:
+            assert sampling_mask_batch_indices is not None
             if sampling_info.is_all_greedy:
-                self._attach_greedy_sampling_mask_to_output(
-                    logits_output, sampling_info, batch_next_token_ids
+                logits_output.sampling_mask_output = (
+                    self._build_greedy_sampling_mask_output(
+                        sampling_mask_batch_indices, batch_next_token_ids
+                    )
                 )
             else:
                 assert sampling_mask_capture is not None
@@ -286,9 +305,7 @@ class Sampler(nn.Module):
                     sampling_mask_capture = sampling_mask_capture._replace(
                         selected_weight=None
                     )
-                self._attach_sampling_mask_to_output(
-                    logits_output,
-                    sampling_info,
+                logits_output.sampling_mask_output = self._build_sampling_mask_output(
                     batch_next_token_ids,
                     sampling_mask_capture,
                 )
@@ -302,8 +319,6 @@ class Sampler(nn.Module):
         sampling_info: SamplingBatchInfo,
         positions: torch.Tensor,
         simple_sampling_case: bool,
-        *,
-        return_sampling_mask: bool = False,
     ) -> Tuple[torch.Tensor, Optional[_SamplingMaskCapture]]:
         """Sample from probability distribution (after softmax).
 
@@ -311,29 +326,13 @@ class Sampler(nn.Module):
         Handles both simple (direct multinomial) and complex (top-k/top-p/min-p) cases.
         Capture work is performed only when return_sampling_mask is enabled.
         """
+        capture_rows = sampling_info.sampling_mask_batch_indices
+        return_sampling_mask = capture_rows is not None
         sampling_mask_capture = None
-        capture_rows = None
-        capture_all_rows = False
         if return_sampling_mask:
-            capture_rows_list = [
-                i
-                for i, should_return in enumerate(
-                    sampling_info.return_sampling_masks or []
-                )
-                if should_return
-            ]
-            if not capture_rows_list:
-                raise RuntimeError(
-                    "Sampling-mask capture requested without any opted-in batch rows."
-                )
-            capture_rows = torch.tensor(
-                capture_rows_list, device=probs.device, dtype=torch.long
+            select_capture_rows = partial(
+                _select_sampling_mask_rows, batch_indices=capture_rows
             )
-            capture_all_rows = capture_rows_list == list(range(probs.shape[0]))
-
-        def select_capture_rows(tensor: torch.Tensor) -> torch.Tensor:
-            assert capture_rows is not None
-            return tensor if capture_all_rows else tensor.index_select(0, capture_rows)
 
         if simple_sampling_case:
             batch_next_token_ids = sampling_from_probs_torch(
@@ -457,62 +456,48 @@ class Sampler(nn.Module):
                 raise ValueError(f"Invalid sampling backend: {backend}")
         return batch_next_token_ids, sampling_mask_capture
 
-    def _attach_greedy_sampling_mask_to_output(
+    def _build_greedy_sampling_mask_output(
         self,
-        logits_output: LogitsProcessorOutput,
-        sampling_info: SamplingBatchInfo,
+        batch_indices: torch.Tensor,
         batch_next_token_ids: torch.Tensor,
-    ) -> None:
-        tokens = batch_next_token_ids.to(torch.int32).cpu().tolist()
-        masks = []
-        logprobs = []
-        for i, should_return in enumerate(sampling_info.return_sampling_masks or []):
-            if should_return:
-                masks.append([int(tokens[i])])
-                logprobs.append(0.0)
-            else:
-                masks.append(None)
-                logprobs.append(None)
-        logits_output.next_token_sampling_mask_idx = masks
-        logits_output.next_token_sampling_logprobs = logprobs
+    ) -> SamplingMaskOutput:
+        token_ids = batch_next_token_ids.index_select(0, batch_indices).to(torch.int32)
+        num_requests = batch_indices.numel()
+        return SamplingMaskOutput(
+            token_ids=token_ids.view(-1, 1),
+            lengths=torch.ones(
+                num_requests, dtype=torch.int32, device=token_ids.device
+            ),
+            selected_logprobs=torch.zeros(
+                num_requests, dtype=torch.float32, device=token_ids.device
+            ),
+            statuses=torch.full(
+                (num_requests,),
+                SamplingMaskStatus.OK,
+                dtype=torch.int32,
+                device=token_ids.device,
+            ),
+        )
 
-    def _attach_sampling_mask_to_output(
+    def _build_sampling_mask_output(
         self,
-        logits_output: LogitsProcessorOutput,
-        sampling_info: SamplingBatchInfo,
         batch_next_token_ids: torch.Tensor,
         sampling_mask_capture: _SamplingMaskCapture,
-    ) -> None:
-        return_sampling_masks = sampling_info.return_sampling_masks or []
-        if not return_sampling_masks:
-            logits_output.next_token_sampling_mask_idx = []
-            logits_output.next_token_sampling_logprobs = []
-            return
-
-        requested_rows_list = [
-            i for i, should_return in enumerate(return_sampling_masks) if should_return
-        ]
-        requested_rows = sampling_mask_capture.batch_rows
+    ) -> 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, requested_rows).view(
-            -1, 1
-        )
+        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:
-            support_token_ids = (
-                torch.arange(
-                    weights.shape[-1], device=weights.device, dtype=torch.int32
-                )
-                .view(1, -1)
-                .expand_as(weights)
-            )
             selected_from_weights = torch.gather(
                 weights, 1, sampled_tokens.long()
             ).squeeze(1)
             sampled_in_capture = selected_from_weights > 0
         else:
-            sampled_matches = token_ids == sampled_tokens.to(token_ids.dtype)
+            sampled_matches = token_ids == sampled_tokens_int32
             sampled_in_capture = sampled_matches.any(dim=-1)
             selected_positions = sampled_matches.to(torch.int32).argmax(
                 dim=-1, keepdim=True
@@ -520,48 +505,63 @@ class Sampler(nn.Module):
             selected_from_weights = torch.gather(
                 weights, 1, selected_positions
             ).squeeze(1)
-            support_token_ids = token_ids
-
         if selected_weight is None:
             selected_weight = selected_from_weights
 
         support = weights > 0
-        support_mass = weights.sum(dim=-1, dtype=torch.float32)
+        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)
-        valid = (
+        invalid = ~(
             sampled_in_capture
             & (selected_weight > 0)
             & (support_mass > 0)
             & torch.isfinite(selected_logprobs)
         )
-        if not bool(torch.all(valid).item()):
-            invalid_rows = (~valid).nonzero(as_tuple=True)[0].cpu().tolist()
-            raise RuntimeError(
-                "Sampled token is outside captured positive sampling support "
-                f"for batch rows {invalid_rows}."
+        realized_lengths = support.sum(dim=-1, dtype=torch.int32)
+        overflow = realized_lengths > self.sampling_mask_max_tokens
+        statuses = torch.where(
+            invalid,
+            SamplingMaskStatus.INVALID,
+            torch.where(
+                overflow,
+                SamplingMaskStatus.OVERFLOW,
+                SamplingMaskStatus.OK,
+            ),
+        ).to(torch.int32)
+
+        # All replicas must make the same request-abort decision.
+        if dist.is_initialized():
+            if dist.get_world_size(self.tp_sync_group) > 1:
+                dist.all_reduce(
+                    statuses, op=dist.ReduceOp.MAX, group=self.tp_sync_group
+                )
+            if (
+                self.cp_sync_group is not None
+                and dist.get_world_size(self.cp_sync_group) > 1
+            ):
+                dist.all_reduce(
+                    statuses, op=dist.ReduceOp.MAX, group=self.cp_sync_group
+                )
+
+        packed_size = min(self.sampling_mask_max_tokens, weights.shape[-1])
+        if token_ids is None:
+            _, 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()
 
-        flat_rows, flat_cols = support.nonzero(as_tuple=True)
-        flat_ids = support_token_ids[flat_rows, flat_cols].to(torch.int32)
-        mask_lengths = support.sum(dim=-1, dtype=torch.int32)
-
-        flat_ids_cpu = flat_ids.cpu().tolist()
-        mask_lengths_cpu = mask_lengths.cpu().tolist()
-        selected_logprobs_cpu = selected_logprobs.cpu().tolist()
-
-        masks = [None] * len(return_sampling_masks)
-        logprobs = [None] * len(return_sampling_masks)
-        cursor = 0
-        for capture_row, batch_row in enumerate(requested_rows_list):
-            mask_len = int(mask_lengths_cpu[capture_row])
-            row_ids = flat_ids_cpu[cursor : cursor + mask_len]
-            cursor += mask_len
-            masks[batch_row] = row_ids
-            logprobs[batch_row] = float(selected_logprobs_cpu[capture_row])
-
-        logits_output.next_token_sampling_mask_idx = masks
-        logits_output.next_token_sampling_logprobs = logprobs
+        return SamplingMaskOutput(
+            token_ids=packed_token_ids,
+            lengths=realized_lengths.clamp(max=packed_size),
+            selected_logprobs=selected_logprobs,
+            statuses=statuses,
+        )
 
     def _sample_from_logprobs(
         self,

python/sglang/srt/managers/scheduler.py

@@ -309,7 +309,6 @@ from sglang.srt.platforms import current_platform
 from sglang.srt.plugins import load_plugins
 from sglang.srt.rust_server.server import RustServer
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
-from sglang.srt.sampling.sampling_params import TOP_K_ALL
 from sglang.srt.server_args import PortArgs, ServerArgs, compute_world_size
 from sglang.srt.session.session_controller import SessionController
 from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
@@ -1503,6 +1502,7 @@ class Scheduler(
                 buffer_size,
                 hidden_size=disagg_hidden_size,
                 hidden_states_dtype=disagg_hidden_states_dtype,
+                max_sampling_mask_tokens=self.server_args.sampling_mask_max_tokens,
                 custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
                 output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
             )
@@ -1549,6 +1549,7 @@ class Scheduler(
                 buffer_size,
                 hidden_size=disagg_hidden_size,
                 hidden_states_dtype=disagg_hidden_states_dtype,
+                max_sampling_mask_tokens=self.server_args.sampling_mask_max_tokens,
                 custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
                 output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
             )
@@ -2884,30 +2885,29 @@ class Scheduler(
                 self._add_request_to_queue(req)
                 return
 
-        if (
-            req.return_sampling_mask
-            and self.disaggregation_mode != DisaggregationMode.NULL
-            and not self.disagg_metadata_buffers.enable_sampling_mask
-        ):
-            error_msg = (
-                "return_sampling_mask with disaggregation requires "
-                "SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS > 0."
-            )
-            req.set_finish_with_abort(error_msg)
-            self.init_req_max_new_tokens(req)
-            self._add_request_to_queue(req)
-            return
-
-        if req.return_sampling_mask and req.sampling_params.top_k == TOP_K_ALL:
-            error_msg = (
-                "return_sampling_mask requires finite top_k; top_p-only sampling "
-                "is valid but can return huge masks in the tail, blowing up "
-                "metadata, so we need a safety cap."
-            )
-            req.set_finish_with_abort(error_msg)
-            self.init_req_max_new_tokens(req)
-            self._add_request_to_queue(req)
-            return
+        if req.return_sampling_mask:
+            if (
+                self.disaggregation_mode != DisaggregationMode.NULL
+                and not self.disagg_metadata_buffers.enable_sampling_mask
+            ):
+                self._reject_sampling_mask_request(
+                    req,
+                    "return_sampling_mask requires "
+                    "SGLANG_ENABLE_DISAGG_SAMPLING_MASK=1 on both prefill and "
+                    "decode servers when using disaggregated serving.",
+                )
+                return
+            top_k = req.sampling_params.top_k
+            sampling_mask_cap = self.server_args.sampling_mask_max_tokens
+            if top_k != 1 and not (1 < top_k <= sampling_mask_cap):
+                error_msg = (
+                    "return_sampling_mask requires top_k=1 for greedy sampling "
+                    f"or finite 1 < top_k <= {sampling_mask_cap}; got top_k="
+                    f"{top_k}. Lower top_k or increase "
+                    "--sampling-mask-max-tokens."
+                )
+                self._reject_sampling_mask_request(req, error_msg)
+                return
 
         if req.return_sampling_mask and not self.spec_algorithm.is_none():
             # Spec workers do not emit one sampling support per accepted token, so
@@ -2916,9 +2916,7 @@ class Scheduler(
             error_msg = (
                 "return_sampling_mask is not supported with speculative decoding."
             )
-            req.set_finish_with_abort(error_msg)
-            self.init_req_max_new_tokens(req)
-            self._add_request_to_queue(req)
+            self._reject_sampling_mask_request(req, error_msg)
             return
 
         if req.return_sampling_mask and get_exec().kernel.sampling_backend == "ascend":
@@ -2928,9 +2926,7 @@ class Scheduler(
                 "return_sampling_mask is not supported with the ascend "
                 "sampling backend."
             )
-            req.set_finish_with_abort(error_msg)
-            self.init_req_max_new_tokens(req)
-            self._add_request_to_queue(req)
+            self._reject_sampling_mask_request(req, error_msg)
             return
 
         # Handle multimodal inputs
@@ -3170,6 +3166,13 @@ class Scheduler(
         else:
             raise ValueError(f"Invalid {self.disaggregation_mode=}")
 
+    def _reject_sampling_mask_request(self, req: Req, error_msg: str) -> None:
+        """Return a sampling-mask validation error without running the model."""
+        logger.error(f"{error_msg}, {req.rid=}")
+        req.time_stats.trace_ctx.abort(abort_info={"reason": error_msg})
+        prepare_abort(req, error_msg, status_code=HTTPStatus.BAD_REQUEST)
+        self.output_streamer.stream_output([req], req.return_logprob)
+
     def _set_or_validate_priority(self, req: Req) -> bool:
         """Set the default priority value, or abort the request based on the priority scheduling mode."""
         if self.enable_priority_scheduling and req.priority is None:

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

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass
+from http import HTTPStatus
 from typing import (
     TYPE_CHECKING,
     Callable,
@@ -15,7 +16,10 @@ import torch
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.environ import envs
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.logits_processor import (
+    LogitsProcessorOutput,
+    SamplingMaskStatus,
+)
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
     FINISH_MATCHED_TOKEN,
@@ -287,6 +291,7 @@ class SchedulerBatchResultProcessor:
             # Move next_token_ids and logprobs to cpu
             next_token_ids = next_token_ids.tolist()
             self.move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
+            self.materialize_sampling_mask_output(batch.reqs, logits_output)
 
             self._validate_pp_skip_output_comm(batch, result)
 
@@ -299,6 +304,19 @@ class SchedulerBatchResultProcessor:
             logprob_pt = 0
 
             for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
+                should_commit_output = (
+                    not req.finished()
+                    and not req.is_retracted
+                    and req.inflight_middle_chunks <= 0
+                )
+                sampling_mask_finish_reason = None
+                if should_commit_output and req.return_sampling_mask:
+                    assert logits_output is not None
+                    statuses = logits_output.next_token_sampling_mask_status
+                    status = None if statuses is None else statuses[i]
+                    sampling_mask_finish_reason = self.get_sampling_mask_finish_reason(
+                        status=status
+                    )
                 if (
                     batch.return_hidden_states
                     and logits_output.hidden_states is not None
@@ -311,9 +329,7 @@ class SchedulerBatchResultProcessor:
                         capture_hidden_mode=prefill_hidden_capture_mode,
                         extend_input_len=extend_input_len_per_req[i],
                         store=(
-                            not req.finished()
-                            and not req.is_retracted
-                            and req.inflight_middle_chunks <= 0
+                            should_commit_output and sampling_mask_finish_reason is None
                         ),
                     )
 
@@ -328,7 +344,10 @@ class SchedulerBatchResultProcessor:
                 if req.inflight_middle_chunks <= 0:
                     req.time_stats.set_prefill_finished_time()
 
-                    if req.beam_group is not None:
+                    if sampling_mask_finish_reason is not None:
+                        req.to_finish = sampling_mask_finish_reason
+                        req.update_finish_state(0)
+                    elif req.beam_group is not None:
                         # The relay point already replaced the sampled-token
                         # append; the group owns all finish semantics.
                         self.beam_coordinator.commit_prefill(
@@ -351,16 +370,22 @@ class SchedulerBatchResultProcessor:
                     ):
                         req.kv.kv_committed_len += 1
                     if req.finished():
-                        self._maybe_collect_routed_experts(req)
-                        self._maybe_collect_indexer_topk(req)
-                        release_kv_cache(req, self.tree_cache)
+                        if sampling_mask_finish_reason is None:
+                            self._maybe_collect_routed_experts(req)
+                            self._maybe_collect_indexer_topk(req)
+                        release_kv_cache(
+                            req,
+                            self.tree_cache,
+                            is_insert=sampling_mask_finish_reason is None,
+                        )
                         req.time_stats.set_completion_time()
                     elif not batch.decoding_reqs or req not in batch.decoding_reqs:
                         maybe_cache_unfinished_req(req, self.tree_cache)
                         if get_memory().enable_hisparse:
                             self.hisparse_coordinator.admit_request_into_staging(req)
 
-                    self._maybe_collect_customized_info(i, req, logits_output)
+                    if sampling_mask_finish_reason is None:
+                        self._maybe_collect_customized_info(i, req, logits_output)
 
                     if batch.return_logprob:
                         logprob_pt = self._apply_prefill_logprobs(
@@ -371,8 +396,12 @@ class SchedulerBatchResultProcessor:
                             extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
                             next_token_ids=next_token_ids,
                             logprob_pt=logprob_pt,
+                            store=sampling_mask_finish_reason is None,
                         )
 
+                    if sampling_mask_finish_reason is not None:
+                        continue
+
                     if req.return_sampling_mask:
                         self.add_sampling_mask_return_values(i, req, logits_output)
 
@@ -520,7 +549,9 @@ class SchedulerBatchResultProcessor:
         extend_logprob_start_len_per_req: Optional[List[int]],
         next_token_ids: List[int],
         logprob_pt: int,
+        store: bool = True,
     ) -> int:
+        """Advance the logprob cursor, optionally storing this request's values."""
         assert extend_logprob_start_len_per_req is not None
         assert extend_input_len_per_req is not None
         extend_logprob_start_len = extend_logprob_start_len_per_req[i]
@@ -532,7 +563,7 @@ class SchedulerBatchResultProcessor:
             extend_logprob_start_len,
         )
 
-        if req.return_logprob:
+        if store and req.return_logprob:
             self.logprob_result_processor.add_logprob_return_values(
                 i,
                 req,
@@ -907,6 +938,7 @@ class SchedulerBatchResultProcessor:
             result.next_token_ids,
             result.can_run_cuda_graph,
         )
+        self.materialize_sampling_mask_output(batch.reqs, logits_output)
 
         next_token_ids, next_token_logprobs = self._normalize_decode_outputs(
             batch=batch,
@@ -967,13 +999,27 @@ class SchedulerBatchResultProcessor:
             next_token_id = next_token_ids[i]
             is_spec = not batch.spec_algorithm.is_none()
 
-            req.output_ids.extend(next_token_id)
-            new_accept_len = len(next_token_id)
-
-            self._maybe_update_reasoning_tokens(req, next_token_id)
             req.time_stats.set_last_decode_finish_time()
+            sampling_mask_finish_reason = None
+            if req.return_sampling_mask:
+                statuses = logits_output.next_token_sampling_mask_status
+                status = None if statuses is None else statuses[i]
+                sampling_mask_finish_reason = self.get_sampling_mask_finish_reason(
+                    status=status
+                )
+            if sampling_mask_finish_reason is not None:
+                req.to_finish = sampling_mask_finish_reason
+                new_accept_len = 0
+            else:
+                req.output_ids.extend(next_token_id)
+                new_accept_len = len(next_token_id)
+                self._maybe_update_reasoning_tokens(req, next_token_id)
             req.update_finish_state(new_accept_len)
 
+            if sampling_mask_finish_reason is not None:
+                self._handle_sampling_mask_abort(req)
+                continue
+
             self._handle_finish_state_updated_req(req, batch, result, i, logits_output)
 
             if req.return_logprob:
@@ -1123,6 +1169,90 @@ class SchedulerBatchResultProcessor:
             None if logprobs is None else logprobs[i]
         )
 
+    @staticmethod
+    def materialize_sampling_mask_output(
+        reqs: List[Req],
+        output: Optional[LogitsProcessorOutput],
+    ) -> None:
+        """Convert opted-in tensor rows to batch-aligned Python results."""
+        if output is None or output.sampling_mask_output is None:
+            return
+
+        sampling_output = output.sampling_mask_output
+        batch_indices = [i for i, req in enumerate(reqs) if req.return_sampling_mask]
+        lengths = sampling_output.lengths.tolist()
+        selected_logprobs = sampling_output.selected_logprobs.tolist()
+        statuses = sampling_output.statuses.tolist()
+        assert len(batch_indices) == len(lengths)
+
+        batch_size = len(reqs)
+        masks = [None] * batch_size
+        logprobs = [None] * batch_size
+        status_by_batch = [None] * batch_size
+        token_ids = sampling_output.token_ids.cpu()
+        packed_width = token_ids.shape[1]
+        for row, batch_index in enumerate(batch_indices):
+            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])
+
+        output.next_token_sampling_mask_idx = masks
+        output.next_token_sampling_logprobs = logprobs
+        output.next_token_sampling_mask_status = status_by_batch
+        output.sampling_mask_output = None
+
+    def get_sampling_mask_finish_reason(
+        self,
+        *,
+        status: Optional[int],
+    ) -> Optional[FINISH_ABORT]:
+        if status == SamplingMaskStatus.OK:
+            return None
+        if status == SamplingMaskStatus.OVERFLOW:
+            return FINISH_ABORT(
+                "Sampling support exceeds --sampling-mask-max-tokens="
+                f"{get_exec().features.sampling_mask_max_tokens}, commonly because "
+                "of cutoff ties. Lower top_k to leave more headroom or increase "
+                "the server limit.",
+                HTTPStatus.BAD_REQUEST,
+                "BadRequestError",
+            )
+        if status is None:
+            return FINISH_ABORT(
+                "The sampling backend did not return captured sampling support.",
+                HTTPStatus.INTERNAL_SERVER_ERROR,
+                "InternalServerError",
+            )
+        assert status == SamplingMaskStatus.INVALID
+        return FINISH_ABORT(
+            "The sampling backend selected a token outside its captured sampling "
+            "support.",
+            HTTPStatus.INTERNAL_SERVER_ERROR,
+            "InternalServerError",
+        )
+
+    def _handle_sampling_mask_abort(self, req: Req) -> None:
+        """Release a request whose sampled token must not be committed."""
+        if req.multimodal_inputs is not None and req.session is None:
+            req.multimodal_inputs.release_features()
+        if get_disagg().disaggregation_decode_enable_offload_kvcache:
+            self.decode_offload_manager.finalize_release_on_finish(req)
+        else:
+            if get_memory().enable_hisparse:
+                self.hisparse_coordinator.request_finished(req)
+            prepare_release = getattr(
+                self.model_worker, "prepare_for_kv_cache_release", None
+            )
+            if callable(prepare_release):
+                prepare_release(req)
+            release_kv_cache(req, self.tree_cache, is_insert=False)
+        req.time_stats.set_completion_time()
+
     def _handle_finish_state_updated_req(
         self,
         req: Req,

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -780,11 +780,15 @@ class SchedulerPPMixin:
             tensor_dict["draft_topk_index"] = draft_input.topk_index.contiguous()
             tensor_dict["draft_hidden_states"] = draft_input.hidden_states.contiguous()
 
-        if batch.return_logprob:
-            logprob_dict = get_logprob_dict_from_result(result)
+        has_sampling_mask_output = (
+            result.logits_output is not None
+            and result.logits_output.sampling_mask_output is not None
+        )
+        if batch.return_logprob or has_sampling_mask_output:
+            logits_output_dict = get_logprob_dict_from_result(result)
             tensor_dict = {
                 **tensor_dict,
-                **logprob_dict,
+                **logits_output_dict,
             }
         auxiliary_output = (
             result.logits_output.auxiliary_device_output
@@ -904,7 +908,10 @@ class SchedulerPPMixin:
         extend_input_len_per_req = None
         extend_logprob_start_len_per_req = None
 
-        if batch.return_logprob:
+        if (
+            batch.return_logprob
+            or pp_outputs.tensors.get("sampling_mask_token_ids") is not None
+        ):
             (
                 logits_output,
                 extend_input_len_per_req,
@@ -1079,7 +1086,17 @@ class SchedulerPPMixin:
                     target, mb_metadata[next_mb_id], next_pp_outputs
                 )
                 d2h_event = self.device_module.Event()
-                d2h_event.record(self.device_module.current_stream())
+                if (
+                    batch_result.logits_output is not None
+                    and batch_result.logits_output.sampling_mask_output is not None
+                ):
+                    batch_result.copy_done = d2h_event
+                    batch_result.copy_to_cpu(
+                        return_logprob=target.return_logprob,
+                        return_hidden_states=False,
+                    )
+                else:
+                    d2h_event.record(self.device_module.current_stream())
 
         if send_first:
             send_output_work = _do_send()

python/sglang/srt/managers/utils.py

@@ -11,7 +11,10 @@ import torch
 
 from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
 from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.logits_processor import (
+    LogitsProcessorOutput,
+    SamplingMaskOutput,
+)
 from sglang.srt.managers import io_struct
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
@@ -174,7 +177,13 @@ class GenerationBatchResult:
         # Sub-objects only declare their device fields; the single copy+safety
         # primitive (_async_d2h: pinned D2H + record_stream) is injected here so
         # all device->host copying and lifetime safety lives in one place.
+        sampling_mask_output = (
+            self.logits_output.sampling_mask_output
+            if self.logits_output is not None
+            else None
+        )
         for holder in (
+            sampling_mask_output,
             self.routed_experts_output,
             self.indexer_topk_output,
             self.expert_distribution_metrics,
@@ -246,9 +255,11 @@ def validate_input_length(
 
 
 def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict:
+    """Build the tensor payload needed to reconstruct PP output processing state."""
 
     logits_output = result.logits_output
     assert logits_output is not None
+    sampling_mask_output = logits_output.sampling_mask_output
 
     return {
         "extend_input_len_per_req": result.extend_input_len_per_req,
@@ -258,8 +269,20 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict:
         "next_token_top_logprobs_idx": result.logits_output.next_token_top_logprobs_idx,
         "next_token_token_ids_logprobs_val": result.logits_output.next_token_token_ids_logprobs_val,
         "next_token_token_ids_logprobs_idx": result.logits_output.next_token_token_ids_logprobs_idx,
-        "next_token_sampling_mask_idx": result.logits_output.next_token_sampling_mask_idx,
-        "next_token_sampling_logprobs": result.logits_output.next_token_sampling_logprobs,
+        "sampling_mask_token_ids": (
+            None if sampling_mask_output is None else sampling_mask_output.token_ids
+        ),
+        "sampling_mask_lengths": (
+            None if sampling_mask_output is None else sampling_mask_output.lengths
+        ),
+        "sampling_mask_selected_logprobs": (
+            None
+            if sampling_mask_output is None
+            else sampling_mask_output.selected_logprobs
+        ),
+        "sampling_mask_statuses": (
+            None if sampling_mask_output is None else sampling_mask_output.statuses
+        ),
         "input_token_logprobs": result.logits_output.input_token_logprobs,
         "input_top_logprobs_val": result.logits_output.input_top_logprobs_val,
         "input_top_logprobs_idx": result.logits_output.input_top_logprobs_idx,
@@ -271,6 +294,15 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict:
 def get_logprob_from_pp_outputs(
     next_pp_outputs: PPProxyTensors,
 ) -> tuple[LogitsProcessorOutput, list[int], list[int]]:
+    """Reconstruct output processing state received from the last PP stage."""
+    sampling_mask_output = None
+    if next_pp_outputs["sampling_mask_token_ids"] is not None:
+        sampling_mask_output = SamplingMaskOutput(
+            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"],
+            statuses=next_pp_outputs["sampling_mask_statuses"],
+        )
     logits_output = LogitsProcessorOutput(
         # Do not send logits and hidden states because they are large
         next_token_logits=None,
@@ -284,8 +316,7 @@ def get_logprob_from_pp_outputs(
         next_token_token_ids_logprobs_idx=next_pp_outputs[
             "next_token_token_ids_logprobs_idx"
         ],
-        next_token_sampling_mask_idx=next_pp_outputs["next_token_sampling_mask_idx"],
-        next_token_sampling_logprobs=next_pp_outputs["next_token_sampling_logprobs"],
+        sampling_mask_output=sampling_mask_output,
         input_token_logprobs=next_pp_outputs["input_token_logprobs"],
         input_top_logprobs_val=next_pp_outputs["input_top_logprobs_val"],
         input_top_logprobs_idx=next_pp_outputs["input_top_logprobs_idx"],

python/sglang/srt/sampling/sampling_batch_info.py

@@ -74,8 +74,10 @@ class SamplingBatchInfo:
     # Used for deterministic sampling
     sampling_seed: Optional[torch.Tensor] = None
 
-    # Per-request flag for returning sparse sampling support metadata.
+    # 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
+    sampling_mask_batch_indices: Optional[torch.Tensor] = None
 
     # Device
     device: str = "cuda"
@@ -145,6 +147,9 @@ class SamplingBatchInfo:
             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]
+        sampling_mask_batch_indices = cls._make_sampling_mask_batch_indices(
+            return_sampling_masks, device
+        )
 
         if has_custom_logit_processor:
             # Merge the same type of custom logit processors together
@@ -210,6 +215,7 @@ class SamplingBatchInfo:
             device=device,
             logit_bias=logit_bias,
             return_sampling_masks=return_sampling_masks,
+            sampling_mask_batch_indices=sampling_mask_batch_indices,
         )
         ret.adjusted_from_schedule_batch(batch, vocab_size)
         return ret
@@ -228,6 +234,22 @@ class SamplingBatchInfo:
     ):
         pass
 
+    @staticmethod
+    def _make_sampling_mask_batch_indices(
+        return_sampling_masks: List[bool], device: str
+    ) -> Optional[torch.Tensor]:
+        """Build device row indices for requests that return sampling masks."""
+        indices = [
+            i for i, should_return in enumerate(return_sampling_masks) if should_return
+        ]
+        if not indices:
+            return None
+        return torch.tensor(
+            indices,
+            dtype=torch.long,
+            pin_memory=is_pin_memory_available(device),
+        ).to(device, non_blocking=True)
+
     def __len__(self):
         return len(self.temperatures)
 
@@ -338,6 +360,12 @@ class SamplingBatchInfo:
             self.return_sampling_masks = [
                 self.return_sampling_masks[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.adjusted_filter_batch(keep_indices, keep_indices_device)
 
@@ -436,6 +464,20 @@ class SamplingBatchInfo:
             self.return_sampling_masks is not None
             or other.return_sampling_masks is not None
         ):
+            if other.sampling_mask_batch_indices is not None:
+                other_sampling_mask_batch_indices = (
+                    other.sampling_mask_batch_indices + self_len
+                )
+                self.sampling_mask_batch_indices = (
+                    other_sampling_mask_batch_indices
+                    if self.sampling_mask_batch_indices is None
+                    else torch.cat(
+                        [
+                            self.sampling_mask_batch_indices,
+                            other_sampling_mask_batch_indices,
+                        ]
+                    )
+                )
             self.return_sampling_masks = (
                 self.return_sampling_masks or [False] * self_len
             ) + (other.return_sampling_masks or [False] * other_len)

test/registered/sampling/test_sampling_mask.py

@@ -1,3 +1,4 @@
+import json
 import math
 import unittest
 from types import SimpleNamespace
@@ -7,8 +8,14 @@ import requests
 import torch
 
 from sglang.srt.layers import sampler as sampler_module
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.layers.sampler import Sampler
+from sglang.srt.layers.logits_processor import (
+    LogitsProcessorOutput,
+    SamplingMaskStatus,
+)
+from sglang.srt.layers.sampler import Sampler, _SamplingMaskCapture
+from sglang.srt.managers.scheduler_components.batch_result_processor import (
+    SchedulerBatchResultProcessor,
+)
 from sglang.srt.sampling.custom_logit_processor import (
     DisallowedTokensLogitsProcessor,
     Qwen3ThinkingBudgetLogitProcessor,
@@ -24,7 +31,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=139, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=250, stage="base-b", runner_config="2-gpu-large")
 register_amd_ci(est_time=320, suite="stage-b-test-1-gpu-small-amd")
 
 _MAX_NEW_TOKENS = 4
@@ -36,9 +43,11 @@ _SERVER_ARGS = (
     "--mem-fraction-static",
     "0.7",
     "--enable-custom-logit-processor",
+    "--sampling-mask-max-tokens",
+    "64",
 )
 _INVALID_SAMPLING_MASK_ERROR = (
-    "top_p-only sampling is valid but can return huge masks in the tail"
+    "return_sampling_mask requires top_k=1 for greedy sampling"
 )
 
 
@@ -46,6 +55,88 @@ class TestSamplingMaskCapture(CustomTestCase):
     def setUp(self):
         self.sampler = Sampler.__new__(Sampler)
         torch.nn.Module.__init__(self.sampler)
+        self.sampler.sampling_mask_max_tokens = 4096
+        self.sampler.tp_sync_group = None
+        self.sampler.cp_sync_group = None
+
+    def test_default_sampling_does_not_construct_capture_helpers(self):
+        """Requests without masks must bypass capture-only allocations."""
+        probs = torch.tensor([[0.6, 0.4]])
+        info = SimpleNamespace(sampling_mask_batch_indices=None, sampling_seed=None)
+        with patch.object(
+            sampler_module, "partial", side_effect=AssertionError("capture helper")
+        ):
+            _, capture = self.sampler._sample_from_probs(
+                probs=probs,
+                sampling_info=info,
+                positions=torch.tensor([0]),
+                simple_sampling_case=True,
+            )
+        self.assertIsNone(capture)
+
+    def _sample(
+        self, probs, backend, *, top_k=2, top_p=0.45, min_p=0.0, requested_rows=None
+    ):
+        batch_size = len(probs)
+        if requested_rows is None:
+            requested_rows = range(batch_size)
+        sampling_info = SimpleNamespace(
+            sampling_seed=None,
+            need_top_k_sampling=True,
+            need_top_p_sampling=top_p < 1.0,
+            need_min_p_sampling=min_p > 0.0,
+            top_ks=torch.full((batch_size,), top_k, dtype=torch.int32, device="cuda"),
+            top_ps=torch.full((batch_size,), top_p, device="cuda"),
+            min_ps=torch.full((batch_size,), min_p, device="cuda"),
+            sampling_mask_batch_indices=torch.tensor(requested_rows, device="cuda"),
+        )
+        with patch.object(
+            sampler_module,
+            "get_exec",
+            return_value=SimpleNamespace(
+                kernel=SimpleNamespace(sampling_backend=backend)
+            ),
+        ):
+            return self.sampler._sample_from_probs(
+                probs,
+                sampling_info,
+                positions=torch.zeros(batch_size, dtype=torch.int64, device="cuda"),
+                simple_sampling_case=False,
+            )
+
+    def _materialize(self, sampled, capture, requested_rows):
+        output = LogitsProcessorOutput(
+            next_token_logits=None,
+            sampling_mask_output=self.sampler._build_sampling_mask_output(
+                sampled, capture
+            ),
+        )
+        output.sampling_mask_output.map_device_tensors(lambda tensor: tensor.cpu())
+        SchedulerBatchResultProcessor.materialize_sampling_mask_output(
+            [
+                SimpleNamespace(return_sampling_mask=i in requested_rows)
+                for i in range(len(sampled))
+            ],
+            output,
+        )
+        return output
+
+    def test_min_p_capture_matches_filtered_support_and_logprob(self):
+        backends = ("pytorch",) if is_hip() else ("pytorch", "flashinfer")
+        for backend in backends:
+            with self.subTest(backend=backend):
+                probs = torch.tensor([[0.4, 0.3, 0.2, 0.1]], device="cuda")
+                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)
+                self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK])
+                self.assertEqual(output.lengths.tolist(), [2])
+                self.assertEqual(set(output.token_ids[0, :2].tolist()), {0, 1})
+                expected = (0.4 if sampled.item() == 0 else 0.3) / 0.7
+                self.assertAlmostEqual(
+                    output.selected_logprobs.item(), math.log(expected), places=6
+                )
 
     def test_hard_exclusion_replay_in_mixed_batch(self):
         backends = ["pytorch"] if is_hip() else ["pytorch", "flashinfer"]
@@ -77,6 +168,7 @@ class TestSamplingMaskCapture(CustomTestCase):
                         )
                     },
                     return_sampling_masks=[True, True],
+                    sampling_mask_batch_indices=torch.tensor([0, 1], device="cuda"),
                 )
                 logits = self.sampler._preprocess_logits(logits, info)
                 with patch(
@@ -90,12 +182,8 @@ class TestSamplingMaskCapture(CustomTestCase):
                         info,
                         positions=torch.zeros(2, dtype=torch.int64, device="cuda"),
                         simple_sampling_case=False,
-                        return_sampling_mask=True,
                     )
-                output = LogitsProcessorOutput(next_token_logits=None)
-                self.sampler._attach_sampling_mask_to_output(
-                    output, info, sampled, capture
-                )
+                output = self._materialize(sampled, capture, requested_rows=[0, 1])
                 support = output.next_token_sampling_mask_idx[0]
                 self.assertEqual(set(support), {0, 1, 3})
                 self.assertIn(int(sampled[0]), support)
@@ -126,29 +214,7 @@ class TestSamplingMaskCapture(CustomTestCase):
         expected_ids = expected_support.nonzero(as_tuple=True)[0].tolist()
         self.assertEqual(expected_ids, [0, 1, 2])
 
-        sampling_info = SimpleNamespace(
-            sampling_seed=None,
-            need_top_k_sampling=True,
-            need_top_p_sampling=True,
-            need_min_p_sampling=False,
-            top_ks=torch.full((batch_size,), top_k, dtype=torch.int32, device="cuda"),
-            top_ps=torch.full((batch_size,), top_p, device="cuda"),
-            min_ps=torch.zeros(batch_size, device="cuda"),
-            return_sampling_masks=[True] * batch_size,
-        )
-        with patch(
-            "sglang.srt.layers.sampler.get_exec",
-            return_value=SimpleNamespace(
-                kernel=SimpleNamespace(sampling_backend="flashinfer")
-            ),
-        ):
-            sampled, capture = self.sampler._sample_from_probs(
-                probs,
-                sampling_info,
-                positions=torch.zeros(batch_size, dtype=torch.int64, device="cuda"),
-                simple_sampling_case=False,
-                return_sampling_mask=True,
-            )
+        sampled, capture = self._sample(probs, "flashinfer", top_k=top_k, top_p=top_p)
 
         self.assertIsNotNone(capture)
         self.assertEqual(capture.batch_rows.cpu().tolist(), list(range(batch_size)))
@@ -164,46 +230,24 @@ class TestSamplingMaskCapture(CustomTestCase):
     @unittest.skipIf(is_hip(), "FlashInfer is not available on ROCm")
     def test_flashinfer_capture_only_materializes_requested_rows(self):
         batch_size = 4
-        top_k = 2
-        top_p = 0.45
         requested_rows = [1, 3]
         probs = torch.tensor([[0.4, 0.2, 0.2, 0.1, 0.1]], device="cuda").repeat(
             batch_size, 1
         )
-        sampling_info = SimpleNamespace(
-            sampling_seed=None,
-            need_top_k_sampling=True,
-            need_top_p_sampling=True,
-            need_min_p_sampling=False,
-            top_ks=torch.full((batch_size,), top_k, dtype=torch.int32, device="cuda"),
-            top_ps=torch.full((batch_size,), top_p, device="cuda"),
-            min_ps=torch.zeros(batch_size, device="cuda"),
-            return_sampling_masks=[False, True, False, True],
-        )
-        top_k_renorm = sampler_module.top_k_renorm_prob
-        top_p_renorm = sampler_module.top_p_renorm_prob
         with (
-            patch(
-                "sglang.srt.layers.sampler.get_exec",
-                return_value=SimpleNamespace(
-                    kernel=SimpleNamespace(sampling_backend="flashinfer")
-                ),
-            ),
-            patch(
-                "sglang.srt.layers.sampler.top_k_renorm_prob",
-                wraps=top_k_renorm,
+            patch.object(
+                sampler_module,
+                "top_k_renorm_prob",
+                wraps=sampler_module.top_k_renorm_prob,
             ) as top_k_mock,
-            patch(
-                "sglang.srt.layers.sampler.top_p_renorm_prob",
-                wraps=top_p_renorm,
+            patch.object(
+                sampler_module,
+                "top_p_renorm_prob",
+                wraps=sampler_module.top_p_renorm_prob,
             ) as top_p_mock,
         ):
-            sampled, capture = self.sampler._sample_from_probs(
-                probs,
-                sampling_info,
-                positions=torch.zeros(batch_size, dtype=torch.int64, device="cuda"),
-                simple_sampling_case=False,
-                return_sampling_mask=True,
+            sampled, capture = self._sample(
+                probs, "flashinfer", requested_rows=requested_rows
             )
 
         self.assertIsNotNone(capture)
@@ -212,10 +256,7 @@ class TestSamplingMaskCapture(CustomTestCase):
         self.assertEqual(tuple(top_k_mock.call_args.args[0].shape), (2, 5))
         self.assertEqual(tuple(top_p_mock.call_args.args[0].shape), (2, 5))
 
-        output = LogitsProcessorOutput(next_token_logits=None)
-        self.sampler._attach_sampling_mask_to_output(
-            output, sampling_info, sampled, capture
-        )
+        output = self._materialize(sampled, capture, requested_rows)
         self.assertIsNone(output.next_token_sampling_mask_idx[0])
         self.assertEqual(set(output.next_token_sampling_mask_idx[1]), {0, 1, 2})
         self.assertIsNone(output.next_token_sampling_mask_idx[2])
@@ -231,39 +272,14 @@ class TestSamplingMaskCapture(CustomTestCase):
         probs = torch.tensor([[0.4, 0.2, 0.2, 0.1, 0.1]], device="cuda").repeat(
             batch_size, 1
         )
-        sampling_info = SimpleNamespace(
-            sampling_seed=None,
-            need_top_k_sampling=True,
-            need_top_p_sampling=True,
-            need_min_p_sampling=False,
-            top_ks=torch.full((batch_size,), 2, dtype=torch.int32, device="cuda"),
-            top_ps=torch.full((batch_size,), 0.45, device="cuda"),
-            min_ps=torch.zeros(batch_size, device="cuda"),
-            return_sampling_masks=[False, True, False, True],
-        )
-        with patch(
-            "sglang.srt.layers.sampler.get_exec",
-            return_value=SimpleNamespace(
-                kernel=SimpleNamespace(sampling_backend="pytorch")
-            ),
-        ):
-            sampled, capture = self.sampler._sample_from_probs(
-                probs,
-                sampling_info,
-                positions=torch.zeros(batch_size, dtype=torch.int64, device="cuda"),
-                simple_sampling_case=False,
-                return_sampling_mask=True,
-            )
+        sampled, capture = self._sample(probs, "pytorch", requested_rows=requested_rows)
 
         self.assertIsNotNone(capture)
         self.assertEqual(capture.batch_rows.cpu().tolist(), requested_rows)
         self.assertEqual(tuple(capture.weights.shape), (len(requested_rows), 5))
         self.assertEqual(tuple(capture.token_ids.shape), (len(requested_rows), 5))
 
-        output = LogitsProcessorOutput(next_token_logits=None)
-        self.sampler._attach_sampling_mask_to_output(
-            output, sampling_info, sampled, capture
-        )
+        output = self._materialize(sampled, capture, requested_rows)
         for batch_row in requested_rows:
             self.assertIn(
                 int(sampled[batch_row]),
@@ -297,18 +313,38 @@ class SamplingMaskTestMixin:
         return_logprob=False,
         top_logprobs_num=0,
         custom_logit_processor=None,
+        stream=False,
     ):
         payload = {
             "text": "The capital of France is",
-            "sampling_params": sampling_params,
+            "sampling_params": {
+                "temperature": 1.0,
+                "max_new_tokens": _MAX_NEW_TOKENS,
+                "ignore_eos": True,
+                **sampling_params,
+            },
             "return_sampling_mask": return_sampling_mask,
+            "stream": stream,
         }
         if custom_logit_processor is not None:
             payload["custom_logit_processor"] = custom_logit_processor
         if return_logprob:
             payload["return_logprob"] = True
             payload["top_logprobs_num"] = top_logprobs_num
-        return requests.post(self.base_url + "/generate", json=payload, timeout=60)
+        return requests.post(
+            self.base_url + "/generate", json=payload, stream=stream, timeout=60
+        )
+
+    def _assert_sampling_masks(self, output_ids, meta_info):
+        masks = meta_info["output_token_sampling_mask"]
+        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.assertIn(token_id, mask)
+            self.assertEqual(len(mask), len(set(mask)))
+        return masks
 
     def _generate_sampling_masks(self, sampling_params):
         response = self._post_generate(sampling_params)
@@ -317,23 +353,13 @@ class SamplingMaskTestMixin:
         output = response.json()
         meta_info = output["meta_info"]
         output_ids = output["output_ids"]
-        sampling_masks = meta_info["output_token_sampling_mask"]
 
         self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
         self.assertEqual(meta_info["completion_tokens"], len(output_ids))
         self.assertEqual(
             meta_info["output_token_sampling_mask_length"], len(output_ids)
         )
-        self.assertEqual(len(sampling_masks), len(output_ids))
-        for output_id, sampling_mask in zip(output_ids, sampling_masks):
-            self.assertIn(output_id, sampling_mask)
-            self.assertEqual(len(sampling_mask), len(set(sampling_mask)))
-        return sampling_masks
-
-    def _assert_rejects_unbounded_sampling_mask(self, sampling_params):
-        response = self._post_generate(sampling_params)
-        self.assertEqual(response.status_code, 400, response.text)
-        self.assertIn(_INVALID_SAMPLING_MASK_ERROR, response.text)
+        return self._assert_sampling_masks(output_ids, meta_info)
 
 
 class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
@@ -393,40 +419,19 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
                 self.assertEqual(recovery.status_code, 200, recovery.text)
 
     def test_generate_returns_sampling_mask(self):
-        top_p_sampling_masks = self._generate_sampling_masks(
-            {
-                "temperature": 1.0,
-                "top_k": _TOP_K,
-                "top_p": _TOP_P,
-                "max_new_tokens": _MAX_NEW_TOKENS,
-                "ignore_eos": True,
-            }
-        )
-        for sampling_mask in top_p_sampling_masks:
-            self.assertGreater(len(sampling_mask), 0)
+        for params, min_size in (
+            ({"top_p": _TOP_P}, 1),
+            ({}, _TOP_K),
+            ({"top_p": 1.0}, _TOP_K),
+        ):
+            with self.subTest(sampling_params=params):
+                masks = self._generate_sampling_masks({"top_k": _TOP_K, **params})
+                for mask in masks:
+                    self.assertGreaterEqual(len(mask), min_size)
 
-        top_k_sampling_masks = self._generate_sampling_masks(
-            {
-                "temperature": 1.0,
-                "top_k": _TOP_K,
-                "max_new_tokens": _MAX_NEW_TOKENS,
-                "ignore_eos": True,
-            }
-        )
-        for sampling_mask in top_k_sampling_masks:
-            self.assertGreaterEqual(len(sampling_mask), _TOP_K)
-
-        top_k_top_p_one_sampling_masks = self._generate_sampling_masks(
-            {
-                "temperature": 1.0,
-                "top_k": _TOP_K,
-                "top_p": 1.0,
-                "max_new_tokens": _MAX_NEW_TOKENS,
-                "ignore_eos": True,
-            }
-        )
-        for sampling_mask in top_k_top_p_one_sampling_masks:
-            self.assertGreaterEqual(len(sampling_mask), _TOP_K)
+    def test_generate_returns_greedy_singleton_mask(self):
+        masks = self._generate_sampling_masks({"temperature": 0.0})
+        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.
@@ -444,13 +449,7 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
         """
         top_k, top_p = _TOP_K, _TOP_P
         response = self._post_generate(
-            {
-                "temperature": 1.0,
-                "top_k": top_k,
-                "top_p": top_p,
-                "max_new_tokens": _MAX_NEW_TOKENS,
-                "ignore_eos": True,
-            },
+            {"top_k": top_k, "top_p": top_p},
             return_logprob=True,
             top_logprobs_num=_TOP_LOGPROBS_NUM,
         )
@@ -459,12 +458,10 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
         output = response.json()
         meta_info = output["meta_info"]
         output_ids = output["output_ids"]
-        sampling_masks = meta_info["output_token_sampling_mask"]
+        sampling_masks = self._assert_sampling_masks(output_ids, meta_info)
         sampling_logprobs = meta_info["output_token_sampling_logprobs"]
         top_logprobs = meta_info["output_top_logprobs"]  # [logprob, id, text] per token
 
-        self.assertEqual(len(sampling_masks), len(output_ids))
-        self.assertEqual(len(sampling_logprobs), len(output_ids))
         self.assertEqual(len(top_logprobs), len(output_ids))
 
         for output_id, mask, mask_logprob, step_top_logprobs in zip(
@@ -476,7 +473,6 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
 
             mask_set = set(mask)
 
-            self.assertIn(output_id, mask_set)
             self.assertTrue(mask_set.issubset(probs))
             top_k_cutoff = sorted(probs.values(), reverse=True)[top_k - 1]
             for token_id in mask_set:
@@ -508,33 +504,115 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
 
         choice = response.json()["choices"][0]
         output_ids = choice["response_token_ids"]
-        meta_info = choice["meta_info"]
-        sampling_masks = meta_info["output_token_sampling_mask"]
-        sampling_logprobs = meta_info["output_token_sampling_logprobs"]
+        self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
+        self._assert_sampling_masks(output_ids, choice["meta_info"])
+
+    def test_generate_streams_aligned_sampling_masks(self):
+        response = self._post_generate({"top_k": _TOP_K, "top_p": _TOP_P}, stream=True)
+        self.assertEqual(response.status_code, 200, response.text)
+
+        output_ids = []
+        for line in response.iter_lines():
+            if not line.startswith(b"data: ") or line[6:] == b"[DONE]":
+                continue
+            chunk = json.loads(line[6:])
+            output_ids = chunk["output_ids"]
+            self._assert_sampling_masks(output_ids, chunk["meta_info"])
 
         self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
-        self.assertEqual(len(sampling_masks), len(output_ids))
-        self.assertEqual(len(sampling_logprobs), len(output_ids))
-        for output_id, sampling_mask in zip(output_ids, sampling_masks):
-            self.assertIn(output_id, sampling_mask)
 
     def test_generate_rejects_unbounded_sampling_mask(self):
-        self._assert_rejects_unbounded_sampling_mask(
-            {
-                "temperature": 1.0,
-                "top_p": _TOP_P,
-                "max_new_tokens": _MAX_NEW_TOKENS,
-                "ignore_eos": True,
-            }
+        for params in ({"top_p": _TOP_P}, {"top_k": 65}, {"top_p": 1.0}):
+            with self.subTest(sampling_params=params):
+                response = self._post_generate(params)
+                self.assertEqual(response.status_code, 400, response.text)
+                self.assertIn(_INVALID_SAMPLING_MASK_ERROR, response.text)
+
+
+class TestSamplingMaskPacking(CustomTestCase):
+    def setUp(self):
+        self.sampler = Sampler.__new__(Sampler)
+        self.sampler.sampling_mask_max_tokens = 3
+        self.sampler.tp_sync_group = None
+        self.sampler.cp_sync_group = None
+
+    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(selected, capture)
+                self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.INVALID])
+
+    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,
         )
-        self._assert_rejects_unbounded_sampling_mask(
-            {
-                "temperature": 1.0,
-                "top_p": 1.0,
-                "max_new_tokens": _MAX_NEW_TOKENS,
-                "ignore_eos": True,
-            }
+        output = self.sampler._build_sampling_mask_output(torch.tensor([1]), capture)
+        self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK])
+        self.assertAlmostEqual(output.selected_logprobs.item(), math.log(0.25))
+
+    def test_greedy_device_output_survives_async_copy(self):
+        from sglang.srt.managers.utils import GenerationBatchResult
+
+        tokens = torch.tensor([3, 4, 5], device="cuda")
+        output = LogitsProcessorOutput(
+            next_token_logits=None,
+            sampling_mask_output=self.sampler._build_greedy_sampling_mask_output(
+                torch.tensor([0, 2], device="cuda"), tokens
+            ),
         )
+        result = GenerationBatchResult(
+            logits_output=output, next_token_ids=tokens, copy_done=torch.cuda.Event()
+        )
+        result.copy_to_cpu(return_logprob=False)
+        result.copy_done.synchronize()
+        self.assertEqual(output.sampling_mask_output.token_ids.device.type, "cpu")
+        SchedulerBatchResultProcessor.materialize_sampling_mask_output(
+            [
+                SimpleNamespace(return_sampling_mask=flag)
+                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])
+
+    def test_overflow_never_materializes_a_partial_mask(self):
+        # Simulate a top-k cutoff tie: a nominal top_k below the cap can still
+        # produce more positive weights than the fixed transport can hold.
+        capture = _SamplingMaskCapture(
+            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
+        )
+
+        output = LogitsProcessorOutput(
+            next_token_logits=None,
+            sampling_mask_output=sampling_output,
+        )
+        SchedulerBatchResultProcessor.materialize_sampling_mask_output(
+            [SimpleNamespace(return_sampling_mask=True)], output
+        )
+        self.assertEqual(
+            output.next_token_sampling_mask_status,
+            [SamplingMaskStatus.OVERFLOW],
+        )
+        self.assertEqual(output.next_token_sampling_mask_idx, [None])
+        self.assertEqual(output.next_token_sampling_logprobs, [None])
 
 
 class TestSamplingMaskDeterministic(SamplingMaskTestMixin, CustomTestCase):
@@ -548,32 +626,20 @@ class TestSamplingMaskDeterministic(SamplingMaskTestMixin, CustomTestCase):
 
     def test_return_sampling_mask_preserves_deterministic_sampling(self):
         sampling_params = {
-            "temperature": 1.0,
             "top_k": _TOP_K,
             "top_p": 1.0,
             "sampling_seed": _SAMPLING_SEED,
-            "max_new_tokens": _MAX_NEW_TOKENS,
-            "ignore_eos": True,
         }
 
-        with_mask_response = self._post_generate(
-            sampling_params, return_sampling_mask=True
-        )
-        self.assertEqual(with_mask_response.status_code, 200, with_mask_response.text)
-
-        without_mask_response = self._post_generate(
-            sampling_params, return_sampling_mask=False
-        )
-        self.assertEqual(
-            without_mask_response.status_code, 200, without_mask_response.text
-        )
-
-        with_mask_output = with_mask_response.json()
-        without_mask_output = without_mask_response.json()
-        self.assertEqual(
-            with_mask_output["output_ids"], without_mask_output["output_ids"]
-        )
-        self.assertEqual(with_mask_output["text"], without_mask_output["text"])
+        outputs = []
+        for return_mask in (False, True):
+            response = self._post_generate(
+                sampling_params, return_sampling_mask=return_mask
+            )
+            self.assertEqual(response.status_code, 200, response.text)
+            output = response.json()
+            outputs.append((output["output_ids"], output["text"]))
+        self.assertEqual(outputs[0], outputs[1])
 
 
 class TestSamplingMaskPytorch(TestSamplingMask):
@@ -584,5 +650,88 @@ class TestSamplingMaskPytorch(TestSamplingMask):
         cls._launch_server(("--sampling-backend", "pytorch"))
 
 
+@unittest.skipIf(is_hip(), "The AMD sampling-mask CI suite provides only one GPU.")
+class TestDistributedSamplingMask(CustomTestCase):
+    def _check_parallel_config(self, *, tp_size, pp_size):
+        process = None
+        try:
+            process = popen_launch_server(
+                "Qwen/Qwen2.5-0.5B-Instruct",
+                DEFAULT_URL_FOR_TEST,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--tp-size",
+                    str(tp_size),
+                    "--pp-size",
+                    str(pp_size),
+                    "--sampling-mask-max-tokens",
+                    "64",
+                    "--mem-fraction-static",
+                    "0.5",
+                    "--max-running-requests",
+                    "8",
+                    "--cuda-graph-max-bs-decode",
+                    "8",
+                ],
+            )
+            for return_logprob in (False, True):
+                with self.subTest(return_logprob=return_logprob):
+                    output = self._generate(
+                        return_sampling_mask=True, return_logprob=return_logprob
+                    )
+                    token_ids = output["output_ids"]
+                    meta = output["meta_info"]
+                    masks = meta["output_token_sampling_mask"]
+                    logprobs = meta["output_token_sampling_logprobs"]
+                    self.assertEqual(len(token_ids), 4)
+                    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):
+                        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)
+                    if return_logprob:
+                        self.assertEqual(len(meta["output_token_logprobs"]), 4)
+
+            ordinary = self._generate(return_sampling_mask=False, return_logprob=False)
+            self.assertEqual(len(ordinary["output_ids"]), 4)
+            self.assertNotIn("output_token_sampling_mask", ordinary["meta_info"])
+        finally:
+            if process is not None:
+                kill_process_tree(process.pid)
+                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,
+                "return_logprob": return_logprob,
+            },
+            timeout=120,
+        )
+        self.assertEqual(response.status_code, 200, response.text)
+        return response.json()
+
+    def test_tp2_sampling_mask(self):
+        """Exercise status synchronization across two tensor-parallel ranks."""
+        self._check_parallel_config(tp_size=2, pp_size=1)
+
+    def test_pp2_sampling_mask(self):
+        """Exercise mask transport between two live pipeline stages."""
+        self._check_parallel_config(tp_size=1, pp_size=2)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -46,6 +46,7 @@ from sglang.srt.speculative.eagle_disaggregation import (
     build_eagle_disagg_draft_input,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=11, suite="base-a-test-cpu")
 
@@ -335,9 +336,14 @@ class TestMooncakePPStaging(unittest.TestCase):
         )
 
 
-class TestEagleDsaSeedTransfer(unittest.TestCase):
+class TestEagleDsaSeedTransfer(CustomTestCase):
     @staticmethod
-    def _make_req(seed, metadata_buffer_index=0):
+    def _make_req(
+        seed,
+        metadata_buffer_index=0,
+        sampling_mask=None,
+        sampling_logprob=None,
+    ):
         return SimpleNamespace(
             metadata_buffer_index=metadata_buffer_index,
             output_ids=[101],
@@ -347,7 +353,13 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
             cached_tokens_storage=0,
             multimodal_inputs=None,
             return_logprob=False,
-            return_sampling_mask=False,
+            return_sampling_mask=sampling_mask is not None,
+            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]
+            ),
             hidden_states_tensor=torch.tensor([1.0, 2.0]),
             output_topk_p=torch.tensor([1.0]),
             output_topk_index=torch.tensor([7]),
@@ -360,6 +372,7 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
             size=2,
             hidden_size=2,
             hidden_states_dtype=torch.float32,
+            max_sampling_mask_tokens=16,
             output_dsa_topk_indices_dim=3,
         )
         seed = torch.tensor([4, 5, 6], dtype=torch.int32)
@@ -378,6 +391,46 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
         self.assertEqual(data_lens[-2], buffers.output_dsa_topk_indices.nbytes)
         self.assertEqual(item_lens[-2], buffers.output_dsa_topk_indices[0].nbytes)
 
+    def test_sampling_mask_metadata_is_opt_in(self):
+        """Disabled masks stay off the wire; enabled masks round-trip at capacity."""
+        schemas = []
+        for enabled in (False, True):
+            with (
+                self.subTest(enabled=enabled),
+                envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(enabled),
+            ):
+                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] if enabled else None,
+                        sampling_logprob=-1.25 if enabled else None,
+                    )
+                )
+                schemas.append(buffers.get_buf_infos())
+                if enabled:
+                    self.assertEqual(
+                        buffers.output_token_sampling_mask_idx.shape, (1, 3)
+                    )
+                    length, mask, logprob = 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)
+                else:
+                    self.assertIsNone(buffers.output_token_sampling_mask_len)
+                    self.assertIsNone(buffers.output_token_sampling_mask_idx)
+                    self.assertIsNone(buffers.output_token_sampling_logprobs)
+                    self.assertEqual(buffers.get_buf(0)[6:9], (None, None, None))
+        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)
+
     def test_decode_input_requires_valid_seed_for_every_request(self):
         seeds = (
             torch.tensor([1, 2, 3], dtype=torch.int32),

test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py

@@ -5,8 +5,13 @@ import pytest
 import torch
 
 from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput, SamplingMaskStatus
 from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo
+from sglang.srt.managers.scheduler_components.batch_result_processor import (
+    SchedulerBatchResultProcessor,
+)
 from sglang.srt.managers.utils import GenerationBatchResult
+from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=1, suite="base-a-test-cpu")
@@ -205,5 +210,54 @@ def test_aborted_result_releases_mamba_allocated_before_kv():
     scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
 
 
+@pytest.mark.parametrize("transport_error", [False, True])
+@pytest.mark.parametrize(
+    "status,http_status,err_type",
+    [
+        (SamplingMaskStatus.OVERFLOW, 400, "BadRequestError"),
+        (SamplingMaskStatus.INVALID, 500, "InternalServerError"),
+    ],
+)
+@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req)
+def test_sampling_mask_abort_preserves_error_and_releases_once(
+    release_kv_cache, status, http_status, err_type, transport_error
+):
+    """A failed sender notification must not leak ownership or lose the API error."""
+    scheduler = _Scheduler()
+    scheduler.batch_result_processor.get_sampling_mask_finish_reason = lambda **kwargs: (
+        SchedulerBatchResultProcessor.get_sampling_mask_finish_reason(None, **kwargs)
+    )
+    req = _Req(inflight_middle_chunks=0)
+    req.to_finish = None
+    req.return_sampling_mask = True
+    req.time_stats.trace_ctx = Mock()
+    if transport_error:
+        req.disagg_kv_sender.abort.side_effect = RuntimeError("transport is down")
+    result = GenerationBatchResult(
+        next_token_ids=torch.tensor([11]),
+        logits_output=LogitsProcessorOutput(
+            next_token_logits=None, next_token_sampling_mask_status=[status]
+        ),
+    )
+
+    with get_context().override_server_args(sampling_mask_max_tokens=64):
+        scheduler.process_batch_result_disagg_prefill(_batch(req), result)
+        scheduler.process_batch_result_disagg_prefill(_batch(req), result)
+
+    assert req.finished_reason.status_code == http_status
+    assert req.finished_reason.err_type == err_type
+    assert req.output_ids == []
+    assert not req.kv.holds_kv and not req.kv.holds_mamba
+    assert req.metadata_buffer_index == -1
+    assert not req.pending_bootstrap
+    assert req.rid not in scheduler.disagg_prefill_pending_chunk_rids
+    release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False)
+    req.disagg_kv_sender.abort.assert_called_once_with()
+    scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(7)
+    scheduler.tree_cache.release_aborted_request.assert_called_once_with(req.rid)
+    scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
+    scheduler.send_kv_chunk.assert_not_called()
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__, "-v"]))

test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py

@@ -9,6 +9,7 @@ Requires: torch, sglang (run in an environment with sglang installed)
 
 import gc
 import unittest
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 from weakref import WeakKeyDictionary as WeakKeyDict
 
@@ -20,9 +21,14 @@ from sglang.srt.disaggregation.decode_kvcache_offload_manager import (
 from sglang.srt.disaggregation.kv_events import OffloadedState
 from sglang.srt.managers.cache_controller import HiCacheAck
 from sglang.srt.managers.schedule_batch import ReqKvInfo
+from sglang.srt.managers.scheduler_components.batch_result_processor import (
+    SchedulerBatchResultProcessor,
+)
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
@@ -440,5 +446,46 @@ class TestReleaseFinishedReq(unittest.TestCase):
         self.assertEqual(len(manager.offload_inflight), 0)
 
 
+class TestSamplingMaskAbortOffload(CustomTestCase):
+    def test_abort_waits_for_existing_offload_before_reusing_slots(self):
+        """An abort must not recycle slots while a previous D2H copy reads them."""
+        for inflight in (False, True):
+            with self.subTest(inflight=inflight):
+                manager, freed = _make_manager(pool_size=32)
+                req = _make_mock_req(0, 20, 20)
+                req.multimodal_inputs = None
+                req.finished.return_value = True
+                manager.req_to_token_pool.free.side_effect = lambda req: setattr(
+                    req.kv, "req_pool_idx", None
+                )
+                processor = SimpleNamespace(decode_offload_manager=manager)
+                if inflight:
+                    manager.offload_inflight[req] = 1
+                    manager.ongoing_offload[1] = (req, torch.arange(4), [1], 0.0)
+                    manager.cache_controller = MagicMock()
+                    manager.cache_controller.ack_write_queue = [
+                        HiCacheAck(None, _FinishedEvent(), [1])
+                    ]
+                    manager._trigger_backup = MagicMock(return_value="hash")
+
+                with get_context().override_server_args(
+                    disaggregation_decode_enable_offload_kvcache=True,
+                    enable_hisparse=False,
+                ):
+                    SchedulerBatchResultProcessor._handle_sampling_mask_abort(
+                        processor, req
+                    )
+
+                if inflight:
+                    self.assertEqual(freed, [])
+                    self.assertEqual(req.kv.req_pool_idx, 0)
+                    manager._check_offload_progress(1)
+                self.assertEqual(len(freed), 1)
+                self.assertTrue(torch.equal(freed[0], torch.arange(20)))
+                self.assertIsNone(req.kv.req_pool_idx)
+                manager.finalize_release_on_finish(req)
+                self.assertEqual(len(freed), 1)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/managers/test_batch_result_processor_hidden_states.py

@@ -4,6 +4,8 @@ from unittest.mock import Mock, patch
 
 import torch
 
+from sglang.srt.environ import envs
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput, SamplingMaskStatus
 from sglang.srt.managers.scheduler_components.batch_result_processor import (
     SchedulerBatchResultProcessor,
 )
@@ -49,6 +51,31 @@ def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProc
     )
 
 
+class TestSamplingMaskMaterialization(CustomTestCase):
+    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()
+        packed_ids.shape = (2, 3)
+        packed_ids.cpu.return_value = torch.tensor([[7, 8, 0], [9, 0, 0]])
+        output = LogitsProcessorOutput(
+            next_token_logits=None,
+            sampling_mask_output=SimpleNamespace(
+                token_ids=packed_ids,
+                lengths=torch.tensor([2, 1]),
+                selected_logprobs=torch.tensor([-0.5, -0.25]),
+                statuses=torch.tensor([SamplingMaskStatus.OK, SamplingMaskStatus.OK]),
+            ),
+        )
+        SchedulerBatchResultProcessor.materialize_sampling_mask_output(
+            reqs=[SimpleNamespace(return_sampling_mask=x) 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.assertIsNone(output.sampling_mask_output)
+
+
 class _PrefillReq:
     def __init__(self, *, rid: str, inflight_middle_chunks: int, return_hidden_states):
         self.rid = rid
@@ -138,6 +165,7 @@ class TestPrefillHiddenStateOffsets(CustomTestCase):
                     logits_output=SimpleNamespace(
                         hidden_states=hidden_states,
                         customized_info=None,
+                        sampling_mask_output=None,
                     ),
                     next_token_ids=torch.tensor([0, 1]),
                     extend_input_len_per_req=[2, 3],
@@ -165,6 +193,134 @@ class TestPrefillHiddenStateOffsets(CustomTestCase):
                 self.assertEqual(last.hidden_states, [[22.0]])
 
 
+class TestPrefillSkippedOutput(CustomTestCase):
+    def test_sampling_mask_middle_chunk_does_not_require_logits_output(self):
+        """A non-token-producing PP chunk may omit its logits output."""
+        req = _PrefillReq(
+            rid="middle",
+            inflight_middle_chunks=1,
+            return_hidden_states=False,
+        )
+        req.return_sampling_mask = True
+        batch = SimpleNamespace(
+            reqs=[req],
+            return_logprob=False,
+            return_hidden_states=False,
+            return_hidden_states_mode=CaptureHiddenMode.NULL,
+            spec_info=None,
+            prefill_stats=None,
+            dp_cooperation_info=None,
+        )
+        result = SimpleNamespace(
+            copy_done=None,
+            auxiliary_host_output=None,
+            routed_experts_output=None,
+            indexer_topk_output=None,
+            logits_output=None,
+            next_token_ids=torch.zeros(1, dtype=torch.int64),
+            extend_input_len_per_req=None,
+            extend_logprob_start_len_per_req=None,
+            grammar_advanced=False,
+            can_run_cuda_graph=False,
+            skipped_output_comm=True,
+        )
+        processor = _make_processor(self)
+
+        with patch.object(
+            envs.SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM,
+            "get",
+            return_value=True,
+        ):
+            processor.process_batch_result_prefill(batch, result)
+
+        self.assertEqual(req.inflight_middle_chunks, 0)
+        self.assertEqual(req.output_ids, [])
+        processor.output_streamer.stream_output.assert_called_once_with(
+            [req], False, req
+        )
+
+
+class TestDecodeWithoutLogits(CustomTestCase):
+    def test_pipeline_result_commits_token_without_sampling_metadata(self):
+        processor = _make_processor(self)
+        req = _DecodeReq()
+        req.return_hidden_states = False
+        batch = SimpleNamespace(
+            reqs=[req],
+            return_logprob=False,
+            spec_algorithm=SimpleNamespace(is_none=lambda: True),
+            batch_size=lambda: 1,
+        )
+        result = GenerationBatchResult(
+            logits_output=None,
+            next_token_ids=torch.tensor([8]),
+        )
+
+        with (
+            patch.object(
+                SchedulerBatchResultProcessor, "_maybe_update_reasoning_tokens"
+            ),
+            patch.object(
+                SchedulerBatchResultProcessor, "_handle_finish_state_updated_req"
+            ),
+        ):
+            processor.process_batch_result_decode(batch, result)
+
+        self.assertEqual(req.output_ids, [8])
+        self.assertEqual(processor.metrics_reporter.num_generated_tokens, 1)
+        processor.output_streamer.stream_output.assert_called_once_with([req], False)
+
+
+class TestSamplingMaskStatusErrors(CustomTestCase):
+    def test_decode_abort_releases_cache_without_committing_token(self):
+        processor = _make_processor(self)
+        req = _DecodeReq()
+        req.output_ids = [7]
+        req.return_sampling_mask = True
+        req.multimodal_inputs = None
+        req.update_finish_state = Mock()
+        batch = SimpleNamespace(
+            reqs=[req],
+            return_logprob=False,
+            spec_algorithm=SimpleNamespace(is_none=lambda: True),
+            batch_size=lambda: 1,
+        )
+        result = GenerationBatchResult(
+            logits_output=LogitsProcessorOutput(
+                next_token_logits=None,
+                next_token_sampling_mask_status=[SamplingMaskStatus.OVERFLOW],
+            ),
+            next_token_ids=torch.tensor([8]),
+        )
+        with patch(
+            "sglang.srt.managers.scheduler_components.batch_result_processor.release_kv_cache"
+        ) as release:
+            processor.process_batch_result_decode(batch, result)
+
+        self.assertEqual(req.output_ids, [7])
+        self.assertEqual(req.to_finish.status_code, 400)
+        req.update_finish_state.assert_called_once_with(0)
+        processor.model_worker.prepare_for_kv_cache_release.assert_called_once_with(req)
+        release.assert_called_once_with(req, processor.tree_cache, is_insert=False)
+        processor.output_streamer.stream_output.assert_called_once_with([req], False)
+
+    def test_overflow_and_invalid_have_distinct_http_errors(self):
+        processor = _make_processor(self)
+
+        overflow = processor.get_sampling_mask_finish_reason(
+            status=SamplingMaskStatus.OVERFLOW
+        )
+        self.assertEqual(overflow.status_code, 400)
+        self.assertEqual(overflow.err_type, "BadRequestError")
+        self.assertIn("cutoff ties", overflow.message)
+
+        invalid = processor.get_sampling_mask_finish_reason(
+            status=SamplingMaskStatus.INVALID
+        )
+        self.assertEqual(invalid.status_code, 500)
+        self.assertEqual(invalid.err_type, "InternalServerError")
+
+
 class TestDecodeHiddenStateRetention(CustomTestCase):
     def test_last_mode_multi_step_storage_stays_bounded(self):
         processor = _make_processor(self)
@@ -180,7 +336,9 @@ class TestDecodeHiddenStateRetention(CustomTestCase):
 
         def result(hidden_states):
             return GenerationBatchResult(
-                logits_output=SimpleNamespace(hidden_states=hidden_states),
+                logits_output=SimpleNamespace(
+                    hidden_states=hidden_states, sampling_mask_output=None
+                ),
                 speculative_num_draft_tokens=4,
             )
 

test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py

@@ -80,7 +80,9 @@ def _make_processor() -> SchedulerBatchResultProcessor:
 
 def _make_result():
     return GenerationBatchResult(
-        logits_output=SimpleNamespace(hidden_states=None, customized_info=None),
+        logits_output=SimpleNamespace(
+            hidden_states=None, customized_info=None, sampling_mask_output=None
+        ),
         next_token_ids=[4],
         speculative_num_draft_tokens=0,
     )

test/registered/unit/managers/test_generation_auxiliary_output.py

@@ -6,13 +6,17 @@ import pytest
 import torch
 
 from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.logits_processor import (
+    LogitsProcessorOutput,
+    SamplingMaskOutput,
+    SamplingMaskStatus,
+)
 from sglang.srt.managers.scheduler import Scheduler
 from sglang.srt.managers.scheduler_components.batch_result_processor import (
     SchedulerBatchResultProcessor,
 )
 from sglang.srt.managers.scheduler_pp_mixin import PPBatchMetadata
-from sglang.srt.managers.utils import GenerationBatchResult
+from sglang.srt.managers.utils import GenerationBatchResult, get_logprob_from_pp_outputs
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.runtime_context import publish, reset_context
@@ -112,6 +116,58 @@ def test_auxiliary_output_releases_device_holder_after_copy():
     assert result.copy_done.record_count == 1
 
 
+def test_sampling_mask_output_uses_generation_result_copy_path():
+    sampling_output = SamplingMaskOutput(
+        token_ids=torch.tensor([[3, 5]], dtype=torch.int32),
+        lengths=torch.tensor([2], dtype=torch.int32),
+        selected_logprobs=torch.tensor([-0.5]),
+        statuses=torch.tensor([SamplingMaskStatus.OK], dtype=torch.int32),
+    )
+    result = GenerationBatchResult(
+        logits_output=LogitsProcessorOutput(
+            next_token_logits=None,
+            sampling_mask_output=sampling_output,
+        ),
+        next_token_ids=torch.tensor([3]),
+        copy_done=CopyDone(),
+    )
+
+    with patch(
+        "sglang.srt.managers.utils._async_d2h",
+        side_effect=lambda tensor: tensor.clone(),
+    ) as copy_tensor:
+        result.copy_to_cpu(return_logprob=False)
+
+    assert copy_tensor.call_count == 5
+    assert sampling_output.token_ids.tolist() == [[3, 5]]
+    assert sampling_output.lengths.tolist() == [2]
+    assert sampling_output.statuses.tolist() == [SamplingMaskStatus.OK]
+    assert result.copy_done.record_count == 1
+
+
+def test_pipeline_sampling_mask_round_trip_without_logprobs():
+    sampling_output = SamplingMaskOutput(
+        token_ids=torch.tensor([[3, 5]], dtype=torch.int32),
+        lengths=torch.tensor([2], dtype=torch.int32),
+        selected_logprobs=torch.tensor([-0.5]),
+        statuses=torch.tensor([SamplingMaskStatus.OK], dtype=torch.int32),
+    )
+    result = GenerationBatchResult(
+        logits_output=LogitsProcessorOutput(
+            next_token_logits=None, sampling_mask_output=sampling_output
+        ),
+        next_token_ids=torch.tensor([3]),
+    )
+    payload = Scheduler._pp_prepare_tensor_dict(
+        SimpleNamespace(), result, SimpleNamespace(return_logprob=False)
+    )
+    output, _, _ = get_logprob_from_pp_outputs(PPProxyTensors(payload))
+    for name in ("token_ids", "lengths", "selected_logprobs", "statuses"):
+        torch.testing.assert_close(
+            getattr(output.sampling_mask_output, name), getattr(sampling_output, name)
+        )
+
+
 def test_non_pp_auxiliary_output_only_requires_host_copy_support():
     device_output = HostOnlyDeviceOutput(torch.tensor([1.0, 2.0]))
     result = GenerationBatchResult(

test/registered/unit/managers/test_scheduler_sampling_mask_validation.py

@@ -0,0 +1,76 @@
+"""Reject PD sampling-mask requests before entering transfer queues."""
+
+import unittest
+from http import HTTPStatus
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.scheduler import Scheduler
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+class TestSchedulerSamplingMaskValidation(CustomTestCase):
+    def test_disabled_pd_masks_return_bad_request_before_admission(self):
+        """Neither PD role may queue mask requests without metadata buffers."""
+        for mode in (DisaggregationMode.PREFILL, DisaggregationMode.DECODE):
+            with self.subTest(mode=mode):
+                scheduler = Scheduler.__new__(Scheduler)
+                scheduler.enable_session_radix_cache = False
+                scheduler.model_config = SimpleNamespace(
+                    hf_eos_token_id={1}, vocab_size=128
+                )
+                scheduler.disaggregation_mode = mode
+                scheduler.disagg_metadata_buffers = SimpleNamespace(
+                    enable_sampling_mask=False
+                )
+                scheduler.metrics_reporter = SimpleNamespace(enable_metrics=False)
+                scheduler.tokenizer = None
+                scheduler.dllm_config = None
+                scheduler._maybe_namespace_elastic_radix_cache = MagicMock()
+                scheduler.spec_algorithm = SimpleNamespace(
+                    is_dflash_family=lambda: False,
+                    is_uno=lambda: False,
+                )
+                scheduler._add_request_to_queue = MagicMock()
+                scheduler.output_streamer = MagicMock()
+                recv_req = MagicMock(
+                    session_params=None,
+                    session_id=None,
+                    input_embeds=None,
+                    bootstrap_port=1,
+                    bootstrap_room=9,
+                )
+                req = MagicMock(return_sampling_mask=True, return_logprob=False)
+                with (
+                    patch(
+                        "sglang.srt.managers.scheduler.BeamCoordinator.request_beam_width",
+                        return_value=1,
+                    ),
+                    patch("sglang.srt.managers.scheduler.Req", return_value=req),
+                    patch("sglang.srt.managers.scheduler.prepare_abort") as abort,
+                ):
+                    scheduler.handle_generate_request(recv_req)
+                abort.assert_called_once()
+                self.assertIs(abort.call_args.args[0], req)
+                self.assertIn(
+                    "SGLANG_ENABLE_DISAGG_SAMPLING_MASK=1",
+                    abort.call_args.args[1],
+                )
+                self.assertEqual(
+                    abort.call_args.kwargs["status_code"], HTTPStatus.BAD_REQUEST
+                )
+                scheduler.output_streamer.stream_output.assert_called_once_with(
+                    [req], False
+                )
+                scheduler._add_request_to_queue.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/sampling/test_sampling_batch_info.py

@@ -100,6 +100,54 @@ class TestSamplingBatchInfoLen(CustomTestCase):
         self.assertEqual(len(info), 5)
 
 
+class TestSamplingMaskBatchIndices(CustomTestCase):
+    def test_filter_removes_last_opted_in_row_then_merge_restores_capture(self):
+        info = _make_info(
+            batch_size=2,
+            return_sampling_masks=[False, True],
+            sampling_mask_batch_indices=torch.tensor([1]),
+        )
+        info.filter_batch([0], torch.tensor([0]))
+        self.assertIsNone(info.sampling_mask_batch_indices)
+        other = _make_info(
+            batch_size=1,
+            return_sampling_masks=[True],
+            sampling_mask_batch_indices=torch.tensor([0]),
+        )
+        info.merge_batch(other)
+        self.assertEqual(info.return_sampling_masks, [False, True])
+        self.assertEqual(info.sampling_mask_batch_indices.tolist(), [1])
+
+    def test_filter_rebuilds_row_indices(self):
+        info = _make_info(
+            batch_size=4,
+            return_sampling_masks=[False, True, False, True],
+            sampling_mask_batch_indices=torch.tensor([1, 3]),
+        )
+
+        info.filter_batch([1, 2, 3], torch.tensor([1, 2, 3]))
+
+        self.assertEqual(info.return_sampling_masks, [True, False, True])
+        self.assertEqual(info.sampling_mask_batch_indices.tolist(), [0, 2])
+
+    def test_merge_offsets_rhs_row_indices(self):
+        lhs = _make_info(
+            batch_size=2,
+            return_sampling_masks=[False, True],
+            sampling_mask_batch_indices=torch.tensor([1]),
+        )
+        rhs = _make_info(
+            batch_size=3,
+            return_sampling_masks=[True, False, True],
+            sampling_mask_batch_indices=torch.tensor([0, 2]),
+        )
+
+        lhs.merge_batch(rhs)
+
+        self.assertEqual(lhs.return_sampling_masks, [False, True, True, False, True])
+        self.assertEqual(lhs.sampling_mask_batch_indices.tolist(), [1, 2, 4])
+
+
 class TestMergeCustomLogitProcessor(CustomTestCase):
     def test_both_none_returns_none(self):
         """Test that merging two None processor dicts returns None."""

test/registered/unit/server_args/test_server_args.py

@@ -186,6 +186,38 @@ class TestPrepareServerArgs(CustomTestCase):
         ):
             ServerArgs(model_path="dummy", prefill_decode_interval=-1).resolve_once()
 
+    def test_sampling_mask_max_tokens(self):
+        self.assertEqual(ServerArgs(model_path="dummy").sampling_mask_max_tokens, 4096)
+        self.assertEqual(
+            ServerArgs(
+                model_path="dummy", sampling_mask_max_tokens=8192
+            ).sampling_mask_max_tokens,
+            8192,
+        )
+        with self.assertRaisesRegex(
+            ValueError, "--sampling-mask-max-tokens must be positive"
+        ):
+            prepare_server_args(
+                ["--model-path", "dummy", "--sampling-mask-max-tokens", "0"]
+            ).resolve_once()
+
+    def test_legacy_sampling_mask_env_requires_migration(self):
+        """Legacy configuration must not silently disable PD sampling masks."""
+        for value in ("0", "128", "invalid", ""):
+            for enabled in (False, True):
+                with (
+                    self.subTest(value=value, enabled=enabled),
+                    envs.SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.override(value),
+                    envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(enabled),
+                    self.assertRaisesRegex(
+                        ValueError,
+                        "SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.*"
+                        "Unset it.*SGLANG_ENABLE_DISAGG_SAMPLING_MASK=1.*"
+                        "--sampling-mask-max-tokens.*prefill and decode",
+                    ),
+                ):
+                    ServerArgs(model_path="dummy").resolve_once()
+
     def test_dsv4_prefill_backend_cli_choices(self):
         parser = server_args_module.argparse.ArgumentParser()
         ServerArgs.add_cli_args(parser)