From 6389b156df4913c15670963fa84358e156fca218 Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Thu, 24 Sep 2026 18:43:33 +0000
Subject: [PATCH 1/5] [Sampling] Stream sampling masks as per-request arrays

Treat output-token sampling supports as bulk numeric data from the sampler
to the tokenizer manager:

- The scheduler copies each committed token's support and logprobs from
  numpy views of the step's host buffers into the request's
  SamplingMaskRows. Its buffers double in capacity and are recycled once no
  handed-off chunk references them, so appends stay linear in the bytes and
  steady-state appends touch no new pages.
- On emission the output streamer hands each request's queued rows off
  without copying, as a SamplingMaskChunk of int32/float32 arrays. The
  detokenizer and multi-tokenizer routing pass the per-request list through
  like any other per-request field.
- The tokenizer manager expands each chunk into the unchanged response lists.

Rows the sampler marks OK already fit the packed width, and a P/D prefill
always queues its first token's row before writing the metadata buffer,
whose width is the same --sampling-mask-max-tokens. The scheduler therefore
no longer re-validates row lengths, and the P/D handoff drops its
unavailable-row sentinel and capacity check.

This replaces a per-row .tolist() on the scheduler and three Python-object
round trips per support entry across the detokenizer and tokenizer.
---
 python/sglang/srt/disaggregation/decode.py    |  26 ++---
 python/sglang/srt/disaggregation/utils.py     |  56 ++-------
 python/sglang/srt/layers/logits_processor.py  |   7 +-
 .../srt/managers/detokenizer_manager.py       |   1 -
 python/sglang/srt/managers/io_struct.py       |  13 +--
 .../srt/managers/multi_tokenizer_mixin.py     |   6 -
 python/sglang/srt/managers/schedule_batch.py  |  11 +-
 .../batch_result_processor.py                 |  27 ++---
 .../scheduler_components/output_streamer.py   |  23 +---
 .../sglang/srt/managers/tokenizer_manager.py  |  11 +-
 python/sglang/srt/sampling/sampling_mask.py   | 106 ++++++++++++++++++
 .../registered/sampling/test_sampling_mask.py |  28 ++++-
 .../test_disaggregation_wire.py               |  17 +--
 ...st_batch_result_processor_hidden_states.py |  24 +++-
 .../managers/test_flat_raw_top_logprobs.py    |   1 -
 .../managers/test_multi_tokenizer_mixin.py    |  22 +++-
 .../test_output_streamer_customized_info.py   |  89 ++++++++++++++-
 .../test_tokenizer_manager_rid_cleanup.py     |   2 +
 18 files changed, 310 insertions(+), 160 deletions(-)
 create mode 100644 python/sglang/srt/sampling/sampling_mask.py

diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py
index 3d233172f0e7..a7f85d5873a8 100644
--- a/python/sglang/srt/disaggregation/decode.py
+++ b/python/sglang/srt/disaggregation/decode.py
@@ -2554,23 +2554,15 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest):
                 "sampling mask buffer disabled on decode side"
             )
             sampling_mask_len = int(output_token_sampling_mask_len[0].item())
-            if sampling_mask_len < 0:
-                decode_req.req.output_token_sampling_mask.append(None)
-                decode_req.req.output_token_sampling_logprobs.append(None)
-            else:
-                decode_req.req.output_token_sampling_mask.append(
-                    output_token_sampling_mask_idx[:sampling_mask_len].cpu().tolist()
-                )
-                if decode_req.req.sampling_logprobs_mode == "support":
-                    decode_req.req.output_token_sampling_logprobs.append(
-                        output_token_sampling_logprobs[:sampling_mask_len]
-                        .cpu()
-                        .tolist()
-                    )
-                else:
-                    decode_req.req.output_token_sampling_logprobs.append(
-                        float(output_token_sampling_logprobs[0].item())
-                    )
+            num_logprobs = (
+                sampling_mask_len
+                if decode_req.req.sampling_logprobs_mode == "support"
+                else 1
+            )
+            decode_req.req.sampling_mask_rows.append(
+                output_token_sampling_mask_idx[:sampling_mask_len].cpu().numpy(),
+                output_token_sampling_logprobs[:num_logprobs].cpu().numpy(),
+            )
 
         decode_req.kv_receiver.clear()
         decode_req.kv_receiver = None
diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py
index 2876f1d814a3..859f01b93af0 100644
--- a/python/sglang/srt/disaggregation/utils.py
+++ b/python/sglang/srt/disaggregation/utils.py
@@ -518,51 +518,17 @@ def set_buf(self, req: Req):
                     device="cpu",
                 )
         if req.return_sampling_mask:
-            # 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
-            sampling_logprobs = req.output_token_sampling_logprobs
-            if sampling_masks:
-                sampling_mask = sampling_masks[0]
-                sampling_logprobs_row = (
-                    sampling_logprobs[0] if sampling_logprobs else None
-                )
-                if sampling_mask is not None and sampling_logprobs_row is not None:
-                    mask_len = len(sampling_mask)
-                    max_mask_len = self.output_token_sampling_mask_idx.shape[1]
-                    if mask_len > max_mask_len:
-                        raise RuntimeError(
-                            f"Sampling mask length {mask_len} exceeds disaggregation "
-                            f"metadata capacity {max_mask_len}. Increase "
-                            "--sampling-mask-max-tokens."
-                        )
-                    self.output_token_sampling_mask_len[req.metadata_buffer_index][
-                        0
-                    ] = mask_len
-                    if mask_len:
-                        self.output_token_sampling_mask_idx[
-                            req.metadata_buffer_index, :mask_len
-                        ].copy_(
-                            torch.tensor(
-                                sampling_mask,
-                                dtype=torch.int32,
-                                device=self.output_token_sampling_mask_idx.device,
-                            )
-                        )
-                    if req.sampling_logprobs_mode == "support" and mask_len:
-                        self.output_token_sampling_logprobs[
-                            req.metadata_buffer_index, :mask_len
-                        ].copy_(
-                            torch.tensor(
-                                sampling_logprobs_row,
-                                dtype=torch.float32,
-                                device=self.output_token_sampling_logprobs.device,
-                            )
-                        )
-                    elif req.sampling_logprobs_mode == "selected":
-                        self.output_token_sampling_logprobs[
-                            req.metadata_buffer_index, 0
-                        ] = float(sampling_logprobs_row)
+            # Prefill streams a request only once its KV transfer ends or it aborts,
+            # so the first token's row is the only one queued here.
+            chunk = req.sampling_mask_rows.view()
+            mask_len = len(chunk.token_ids)
+            self.output_token_sampling_mask_len[req.metadata_buffer_index][0] = mask_len
+            self.output_token_sampling_mask_idx[
+                req.metadata_buffer_index, :mask_len
+            ].copy_(torch.from_numpy(chunk.token_ids))
+            self.output_token_sampling_logprobs[
+                req.metadata_buffer_index, : len(chunk.logprobs)
+            ].copy_(torch.from_numpy(chunk.logprobs))
         # For PD + spec decode
         if req.hidden_states_tensor is not None:
             # speculative_eagle_topk should not be greater than 16 currently
diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py
index 5280910431a9..5cb046c4e84b 100644
--- a/python/sglang/srt/layers/logits_processor.py
+++ b/python/sglang/srt/layers/logits_processor.py
@@ -19,6 +19,7 @@
 from enum import IntEnum
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+import numpy as np
 import torch
 from torch import nn
 
@@ -229,10 +230,8 @@ class LogitsProcessorOutput:
     # Post-filter support IDs and requested behavior logprobs, bounded by server
     # capacity. Logprobs are normalized over the full realized support.
     sampling_mask_output: Optional[SamplingMaskOutput] = None
-    next_token_sampling_mask_idx: Optional[List[Optional[List[int]]]] = None
-    next_token_sampling_logprobs: Optional[
-        List[Optional[Union[float, List[float]]]]
-    ] = None
+    next_token_sampling_mask_idx: Optional[List[Optional[np.ndarray]]] = None
+    next_token_sampling_logprobs: Optional[List[Optional[np.ndarray]]] = None
     next_token_sampling_mask_status: Optional[List[Optional[int]]] = None
 
     ## Part 3: Prefill-only. This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor
diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py
index a15bb1d7b136..e27726f695e3 100644
--- a/python/sglang/srt/managers/detokenizer_manager.py
+++ b/python/sglang/srt/managers/detokenizer_manager.py
@@ -495,7 +495,6 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
             output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx,
             output_token_entropy_val=recv_obj.output_token_entropy_val,
             output_token_sampling_mask=recv_obj.output_token_sampling_mask,
-            output_token_sampling_logprobs=recv_obj.output_token_sampling_logprobs,
             output_hidden_states=recv_obj.output_hidden_states,
             routed_experts=routed_experts,
             indexer_topk=indexer_topk,
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index e246f2123ec5..cdf56a2b7b85 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -63,6 +63,7 @@
     get_return_hidden_states_mode,
 )
 from sglang.srt.multimodal.mm_utils import has_valid_data
+from sglang.srt.sampling.sampling_mask import SamplingMaskChunk
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.utils import ImageData, VideoData
 from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d
@@ -1543,12 +1544,7 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True):
     output_token_ids_logprobs_val: TokenIdsLogprobValues
     output_token_ids_logprobs_idx: TokenIdsLogprobIndices
     output_token_entropy_val: Optional[List[Optional[float]]]
-    # Per-request chunks of output-token sampling supports. None when no request
-    # in the batch asks for return_sampling_mask.
-    output_token_sampling_mask: Optional[List[List[List[int]]]]
-    # Per-request chunks. Each output-token entry is a selected-token scalar or
-    # a list aligned with output_token_sampling_mask, according to the request.
-    output_token_sampling_logprobs: Optional[List[List[Union[float, List[float]]]]]
+    output_token_sampling_mask: Optional[List[Optional[SamplingMaskChunk]]]
 
     # Hidden states
     output_hidden_states: OutputHiddenStates
@@ -1643,10 +1639,7 @@ class BatchStrOutput(BaseBatchReq, kw_only=True):
     output_token_ids_logprobs_val: TokenIdsLogprobValues
     output_token_ids_logprobs_idx: TokenIdsLogprobIndices
     output_token_entropy_val: Optional[List[Optional[float]]]
-    # Detokenizer pass-through for BatchTokenIDOutput.output_token_sampling_*;
-    # support-mode logprobs are aligned elementwise with the token IDs.
-    output_token_sampling_mask: Optional[List[List[List[int]]]]
-    output_token_sampling_logprobs: Optional[List[List[Union[float, List[float]]]]]
+    output_token_sampling_mask: Optional[List[Optional[SamplingMaskChunk]]]
 
     # Hidden states
     output_hidden_states: OutputHiddenStates
diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py
index ff88f27e8d67..48cb012c13ec 100644
--- a/python/sglang/srt/managers/multi_tokenizer_mixin.py
+++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py
@@ -247,9 +247,6 @@ def _handle_output_by_index(output, i):
             output_token_sampling_mask=_extract_field_by_index(
                 output, "output_token_sampling_mask", i, check_length=False
             ),
-            output_token_sampling_logprobs=_extract_field_by_index(
-                output, "output_token_sampling_logprobs", i, check_length=False
-            ),
             output_hidden_states=_extract_field_by_index(
                 output, "output_hidden_states", i, check_length=False
             ),
@@ -368,9 +365,6 @@ def _handle_output_by_index(output, i):
             output_token_sampling_mask=_extract_field_by_index(
                 output, "output_token_sampling_mask", i, check_length=False
             ),
-            output_token_sampling_logprobs=_extract_field_by_index(
-                output, "output_token_sampling_logprobs", i, check_length=False
-            ),
             output_hidden_states=_extract_field_by_index(
                 output, "output_hidden_states", i, check_length=False
             ),
diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py
index 33a0453a19ad..06bea0e6a83f 100755
--- a/python/sglang/srt/managers/schedule_batch.py
+++ b/python/sglang/srt/managers/schedule_batch.py
@@ -141,6 +141,7 @@
     SchedulerReqTimeStats,
 )
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+from sglang.srt.sampling.sampling_mask import SamplingMaskRows
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.utils import flatten_nested_list
 from sglang.srt.utils.token_sequence_matcher import TokenSequenceMatcher
@@ -1220,7 +1221,6 @@ def __init__(
         # TODO (Byron): send_output_token_logprobs_offset and send_decode_id_offset can be different in disaggregation mode
         # because the decode server does not have the first output token logprobs
         self.send_output_token_logprobs_offset: int = 0
-        self.send_output_sampling_mask_offset: int = 0
 
         # Logprobs (arguments)
         self.return_logprob = return_logprob
@@ -1256,12 +1256,9 @@ def __init__(
             # Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring)
             self.logprob.output_token_ids_logprobs_val = []
             self.logprob.output_token_ids_logprobs_idx = []
-        if return_sampling_mask:
-            self.output_token_sampling_mask = []
-            self.output_token_sampling_logprobs = []
-        else:
-            self.output_token_sampling_mask = None
-            self.output_token_sampling_logprobs = None
+        self.sampling_mask_rows: Optional[SamplingMaskRows] = (
+            SamplingMaskRows() if return_sampling_mask else None
+        )
         self.hidden_states: List[List[float]] = []
         self.hidden_states_tensor = None  # Note: use tensor instead of list to transfer hidden_states when PD + MTP
         self.output_topk_p = None
diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py
index efaf0e45b1f3..e0f922fe8df6 100644
--- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py
@@ -1163,11 +1163,9 @@ def add_sampling_mask_return_values(
         output: LogitsProcessorOutput,
     ) -> None:
         """Attach sparse sampling support metadata to the return values."""
-        mask = output.next_token_sampling_mask_idx
-        logprobs = output.next_token_sampling_logprobs
-        req.output_token_sampling_mask.append(None if mask is None else mask[i])
-        req.output_token_sampling_logprobs.append(
-            None if logprobs is None else logprobs[i]
+        req.sampling_mask_rows.append(
+            output.next_token_sampling_mask_idx[i],
+            output.next_token_sampling_logprobs[i],
         )
 
     @staticmethod
@@ -1175,14 +1173,13 @@ def materialize_sampling_mask_output(
         reqs: List[Req],
         output: Optional[LogitsProcessorOutput],
     ) -> None:
-        """Convert opted-in tensor rows to batch-aligned Python results."""
+        """Convert opted-in tensor rows to batch-aligned host rows."""
         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)
 
@@ -1190,13 +1187,13 @@ def materialize_sampling_mask_output(
         masks = [None] * batch_size
         logprobs = [None] * batch_size
         status_by_batch = [None] * batch_size
-        token_ids = sampling_output.token_ids.cpu()
+        token_ids = sampling_output.token_ids.cpu().numpy()
+        selected_logprobs = sampling_output.selected_logprobs.cpu().numpy()
         support_logprobs = (
             None
             if sampling_output.support_logprobs is None
-            else sampling_output.support_logprobs.cpu()
+            else sampling_output.support_logprobs.cpu().numpy()
         )
-        packed_width = token_ids.shape[1]
         support_row = 0
         for row, batch_index in enumerate(batch_indices):
             returns_support_logprobs = (
@@ -1204,17 +1201,13 @@ def materialize_sampling_mask_output(
             )
             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()
+                masks[batch_index] = token_ids[row, :length]
                 if returns_support_logprobs:
-                    logprobs[batch_index] = support_logprobs[
-                        support_row, :length
-                    ].tolist()
+                    logprobs[batch_index] = support_logprobs[support_row, :length]
                 else:
-                    logprobs[batch_index] = float(selected_logprobs[row])
+                    logprobs[batch_index] = selected_logprobs[row : row + 1]
             if returns_support_logprobs:
                 support_row += 1
 
diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py
index 42382f3024ab..89cf4e53d18f 100644
--- a/python/sglang/srt/managers/scheduler_components/output_streamer.py
+++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py
@@ -32,6 +32,7 @@
 )
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.runtime_context import get_observability, get_parallel, get_serving
+from sglang.srt.sampling.sampling_mask import SamplingMaskChunk
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils.weight_versions import compute_weight_version_spans
@@ -374,8 +375,7 @@ class _GenerationStreamAccumulator:
     input_token_ids_logprobs_idx: Optional[list] = None
     output_token_ids_logprobs_val: Optional[list] = None
     output_token_ids_logprobs_idx: Optional[list] = None
-    output_token_sampling_mask: Optional[list[list[list[int]]]] = None
-    output_token_sampling_logprobs: Optional[list[list[float | list[float]]]] = None
+    output_token_sampling_mask: Optional[list[Optional[SamplingMaskChunk]]] = None
     # Rust server mode: the Rust detokenizer reconstructs text/ids from the raw
     # output tokens itself and never consumes the scheduler's incremental-detok
     # offsets (decode_ids / read_offset), so that per-step bookkeeping is skipped.
@@ -407,7 +407,6 @@ def __post_init__(self) -> None:
             self.output_token_ids_logprobs_idx = []
         if self.return_sampling_mask:
             self.output_token_sampling_mask = []
-            self.output_token_sampling_logprobs = []
 
     def _beam_admits(self, *, req: Req) -> bool:
         # Only the leader is ever streamed, and only at group finish.
@@ -624,22 +623,9 @@ def accept(self, *, req: Req) -> None:
 
         if self.return_sampling_mask:
             if req.return_sampling_mask:
-                send_output_sampling_mask_offset = req.send_output_sampling_mask_offset
-                sampling_mask_end = len(req.output_token_sampling_mask)
-                self.output_token_sampling_mask.append(
-                    req.output_token_sampling_mask[
-                        send_output_sampling_mask_offset:sampling_mask_end
-                    ]
-                )
-                self.output_token_sampling_logprobs.append(
-                    req.output_token_sampling_logprobs[
-                        send_output_sampling_mask_offset:sampling_mask_end
-                    ]
-                )
-                req.send_output_sampling_mask_offset = sampling_mask_end
+                self.output_token_sampling_mask.append(req.sampling_mask_rows.take())
             else:
-                self.output_token_sampling_mask.append([])
-                self.output_token_sampling_logprobs.append([])
+                self.output_token_sampling_mask.append(None)
 
         if self.return_hidden_states:
             if req.return_hidden_states:
@@ -748,7 +734,6 @@ def to_payload(
             output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx,
             output_token_entropy_val=None,
             output_token_sampling_mask=self.output_token_sampling_mask,
-            output_token_sampling_logprobs=self.output_token_sampling_logprobs,
             output_hidden_states=self.output_hidden_states,
             routed_experts=self.routed_experts,
             indexer_topk=self.indexer_topk,
diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
index dff813f3b804..7140fbbd5c64 100644
--- a/python/sglang/srt/managers/tokenizer_manager.py
+++ b/python/sglang/srt/managers/tokenizer_manager.py
@@ -2397,12 +2397,11 @@ async def _handle_batch_output(
             ):
                 output_sampling_mask = recv_obj.output_token_sampling_mask
                 if output_sampling_mask is not None:
-                    state.output_token_sampling_mask.extend(output_sampling_mask[i])
-                    output_sampling_logprobs = recv_obj.output_token_sampling_logprobs
-                    if output_sampling_logprobs is not None:
-                        state.output_token_sampling_logprobs.extend(
-                            output_sampling_logprobs[i]
-                        )
+                    masks, logprobs = output_sampling_mask[i].to_lists(
+                        support_logprobs=state.obj.sampling_logprobs_mode == "support"
+                    )
+                    state.output_token_sampling_mask.extend(masks)
+                    state.output_token_sampling_logprobs.extend(logprobs)
                     meta_info["output_token_sampling_mask"] = (
                         state.output_token_sampling_mask
                     )
diff --git a/python/sglang/srt/sampling/sampling_mask.py b/python/sglang/srt/sampling/sampling_mask.py
new file mode 100644
index 000000000000..d9c4e4bde35a
--- /dev/null
+++ b/python/sglang/srt/sampling/sampling_mask.py
@@ -0,0 +1,106 @@
+"""Host-side transport of output-token sampling masks."""
+
+import sys
+from array import array
+from typing import List, Optional, Tuple, Union
+
+import msgspec
+import numpy as np
+
+
+class SamplingMaskChunk(msgspec.Struct, frozen=True, kw_only=True, array_like=True):
+    """One request's sampling masks in one output batch, one row per output token."""
+
+    # int32 [num_rows]
+    lengths: np.ndarray
+    # int32 [sum(lengths)], row after row
+    token_ids: np.ndarray
+    # float32 [sum(lengths)] in support mode, [num_rows] in selected mode
+    logprobs: np.ndarray
+
+    def to_lists(
+        self, support_logprobs: bool
+    ) -> Tuple[List[List[int]], List[Union[float, List[float]]]]:
+        masks: List[List[int]] = []
+        values: List[Union[float, List[float]]] = (
+            [] if support_logprobs else self.logprobs.tolist()
+        )
+        start = 0
+        for length in self.lengths.tolist():
+            end = start + length
+            masks.append(self.token_ids[start:end].tolist())
+            if support_logprobs:
+                values.append(self.logprobs[start:end].tolist())
+            start = end
+        return masks, values
+
+
+class _GrowableArray(msgspec.Struct):
+    """Append-only 1-D buffer; doubling capacity copies each value O(1) times.
+
+    ``reset`` keeps the buffer as a spare, reused once no view of it is alive,
+    so steady-state appends touch no new pages.
+    """
+
+    buffer: np.ndarray
+    size: int = 0
+    spare: Optional[np.ndarray] = None
+
+    def extend(self, values: np.ndarray) -> None:
+        end = self.size + len(values)
+        if end > len(self.buffer):
+            self.buffer = self._grow(end)
+        self.buffer[self.size : end] = values
+        self.size = end
+
+    def view(self) -> np.ndarray:
+        return self.buffer[: self.size]
+
+    def reset(self) -> None:
+        self.spare = self.buffer
+        self.buffer = np.empty(0, dtype=self.buffer.dtype)
+        self.size = 0
+
+    def _grow(self, end: int) -> np.ndarray:
+        spare, self.spare = self.spare, None
+        # Two references, the local and the argument, mean no view is alive.
+        if spare is not None and len(spare) >= end and sys.getrefcount(spare) == 2:
+            grown = spare
+        else:
+            grown = np.empty(max(end, 2 * len(self.buffer)), dtype=self.buffer.dtype)
+        grown[: self.size] = self.buffer[: self.size]
+        return grown
+
+
+class SamplingMaskRows(msgspec.Struct):
+    """One request's SamplingMaskChunk rows, queued until they are streamed."""
+
+    _lengths: array = msgspec.field(default_factory=lambda: array("i"))
+    _token_ids: _GrowableArray = msgspec.field(
+        default_factory=lambda: _GrowableArray(np.empty(0, dtype=np.int32))
+    )
+    _logprobs: _GrowableArray = msgspec.field(
+        default_factory=lambda: _GrowableArray(np.empty(0, dtype=np.float32))
+    )
+
+    def append(self, token_ids: np.ndarray, logprobs: np.ndarray) -> None:
+        """Queue one token's int32 support and float32 logprobs."""
+        self._lengths.append(len(token_ids))
+        self._token_ids.extend(token_ids)
+        self._logprobs.extend(logprobs)
+
+    def view(self) -> SamplingMaskChunk:
+        """View the queued token IDs and logprobs without copying them."""
+        return SamplingMaskChunk(
+            lengths=np.array(self._lengths, dtype=np.int32),
+            token_ids=self._token_ids.view(),
+            logprobs=self._logprobs.view(),
+        )
+
+    def take(self) -> SamplingMaskChunk:
+        """Hand off the queued rows without copying them, leaving the queue empty."""
+        chunk = self.view()
+        del self._lengths[:]
+        self._token_ids.reset()
+        self._logprobs.reset()
+        return chunk
diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py
index 6094a9e9e788..5bb347332f83 100644
--- a/test/registered/sampling/test_sampling_mask.py
+++ b/test/registered/sampling/test_sampling_mask.py
@@ -203,8 +203,8 @@ def test_hard_exclusion_replay_in_mixed_batch(self):
                         simple_sampling_case=False,
                     )
                 output = self._materialize(sampled, capture, requested_rows=[0, 1])
-                support = output.next_token_sampling_mask_idx[0]
-                sampling_logprobs = output.next_token_sampling_logprobs[0]
+                support = output.next_token_sampling_mask_idx[0].tolist()
+                sampling_logprobs = output.next_token_sampling_logprobs[0].tolist()
                 self.assertEqual(set(support), {0, 1, 3})
                 self.assertIn(int(sampled[0]), support)
                 expected = original[0, sampled[0]] - original[0, support].logsumexp(0)
@@ -280,9 +280,13 @@ def test_flashinfer_capture_only_materializes_requested_rows(self):
 
         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.assertEqual(
+            set(output.next_token_sampling_mask_idx[1].tolist()), {0, 1, 2}
+        )
         self.assertIsNone(output.next_token_sampling_mask_idx[2])
-        self.assertEqual(set(output.next_token_sampling_mask_idx[3]), {0, 1, 2})
+        self.assertEqual(
+            set(output.next_token_sampling_mask_idx[3].tolist()), {0, 1, 2}
+        )
         self.assertIsNone(output.next_token_sampling_logprobs[0])
         self.assertIsNotNone(output.next_token_sampling_logprobs[1])
         self.assertIsNone(output.next_token_sampling_logprobs[2])
@@ -714,8 +718,20 @@ def test_greedy_device_output_survives_async_copy(self):
             ],
             output,
         )
-        self.assertEqual(output.next_token_sampling_mask_idx, [[3], None, [5]])
-        self.assertEqual(output.next_token_sampling_logprobs, [[0.0], None, [0.0]])
+        self.assertEqual(
+            [
+                None if row is None else row.tolist()
+                for row in output.next_token_sampling_mask_idx
+            ],
+            [[3], None, [5]],
+        )
+        self.assertEqual(
+            [
+                None if row is None else row.tolist()
+                for row in 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
diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py
index 22e1ec977dcf..9e324f15ee91 100644
--- a/test/registered/unit/disaggregation/test_disaggregation_wire.py
+++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py
@@ -59,6 +59,7 @@
     QSATokenToKVPool,
 )
 from sglang.srt.runtime_context import get_context
+from sglang.srt.sampling.sampling_mask import SamplingMaskRows
 from sglang.srt.speculative.eagle_disaggregation import (
     build_eagle_disagg_draft_input,
 )
@@ -713,6 +714,13 @@ def _make_req(
         sampling_logprobs=None,
         sampling_logprobs_mode="support",
     ):
+        sampling_mask_rows = None
+        if sampling_mask is not None:
+            sampling_mask_rows = SamplingMaskRows()
+            sampling_mask_rows.append(
+                np.array(sampling_mask, np.int32),
+                np.array(sampling_logprobs, np.float32),
+            )
         return SimpleNamespace(
             metadata_buffer_index=metadata_buffer_index,
             output_ids=[101],
@@ -724,12 +732,7 @@ def _make_req(
             return_logprob=False,
             return_sampling_mask=sampling_mask is not None,
             sampling_logprobs_mode=sampling_logprobs_mode,
-            output_token_sampling_mask=(
-                None if sampling_mask is None else [sampling_mask]
-            ),
-            output_token_sampling_logprobs=(
-                None if sampling_logprobs is None else [sampling_logprobs]
-            ),
+            sampling_mask_rows=sampling_mask_rows,
             hidden_states_tensor=torch.tensor([1.0, 2.0]),
             output_topk_p=torch.tensor([1.0]),
             output_topk_index=torch.tensor([7]),
@@ -813,7 +816,7 @@ def test_sampling_mask_selected_logprob_uses_first_metadata_slot(self):
                 self._make_req(
                     None,
                     sampling_mask=[7, 8, 9],
-                    sampling_logprobs=-0.5,
+                    sampling_logprobs=[-0.5],
                     sampling_logprobs_mode="selected",
                 )
             )
diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py
index 477979c4adf7..103432afbd62 100644
--- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py
+++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py
@@ -76,10 +76,13 @@ def test_selected_and_support_modes_share_one_batch(self):
             ],
             output=output,
         )
-        self.assertEqual(output.next_token_sampling_mask_idx, [[7, 8], [9, 10]])
         self.assertEqual(
-            output.next_token_sampling_logprobs,
-            [-0.5, [-0.25, -1.5]],
+            [row.tolist() for row in output.next_token_sampling_mask_idx],
+            [[7, 8], [9, 10]],
+        )
+        self.assertEqual(
+            [row.tolist() for row in output.next_token_sampling_logprobs],
+            [[-0.5], [-0.25, -1.5]],
         )
 
     def test_selected_mode_does_not_require_support_tensor(self):
@@ -102,7 +105,7 @@ def test_selected_mode_does_not_require_support_tensor(self):
             ],
             output=output,
         )
-        self.assertEqual(output.next_token_sampling_logprobs, [-0.5])
+        self.assertEqual(output.next_token_sampling_logprobs[0].tolist(), [-0.5])
 
     def test_packed_ids_are_copied_before_per_request_slicing(self):
         """Non-overlap capture must not perform one device copy per request."""
@@ -130,9 +133,18 @@ def test_packed_ids_are_copied_before_per_request_slicing(self):
             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,
+            [
+                None if row is None else row.tolist()
+                for row in output.next_token_sampling_mask_idx
+            ],
+            [[7, 8], None, [9]],
+        )
+        self.assertEqual(
+            [
+                None if row is None else row.tolist()
+                for row in output.next_token_sampling_logprobs
+            ],
             [[-0.5, -0.75], None, [-0.25]],
         )
         self.assertIsNone(output.sampling_mask_output)
diff --git a/test/registered/unit/managers/test_flat_raw_top_logprobs.py b/test/registered/unit/managers/test_flat_raw_top_logprobs.py
index 646ccb9d2823..9535d8e1bd19 100644
--- a/test/registered/unit/managers/test_flat_raw_top_logprobs.py
+++ b/test/registered/unit/managers/test_flat_raw_top_logprobs.py
@@ -636,7 +636,6 @@ def _make_batch_token_id_output(**overrides) -> BatchTokenIDOutput:
         output_token_ids_logprobs_idx=[[], []],
         output_token_entropy_val=None,
         output_token_sampling_mask=None,
-        output_token_sampling_logprobs=None,
         output_hidden_states=None,
         routed_experts=None,
         indexer_topk=None,
diff --git a/test/registered/unit/managers/test_multi_tokenizer_mixin.py b/test/registered/unit/managers/test_multi_tokenizer_mixin.py
index 7796146b7126..02a05ca209e5 100644
--- a/test/registered/unit/managers/test_multi_tokenizer_mixin.py
+++ b/test/registered/unit/managers/test_multi_tokenizer_mixin.py
@@ -1,5 +1,8 @@
 import unittest
 
+import numpy as np
+
+from sglang.srt.sampling.sampling_mask import SamplingMaskChunk
 from sglang.srt.utils.weight_versions import WeightVersionSpan
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import maybe_stub_sgl_kernel
@@ -69,8 +72,18 @@ def _make_batch_str_output() -> BatchStrOutput:
         output_token_ids_logprobs_val=[[], []],
         output_token_ids_logprobs_idx=[[], []],
         output_token_entropy_val=[0.0, 0.0],
-        output_token_sampling_mask=[[[7, 8]], [[9]]],
-        output_token_sampling_logprobs=[[[-0.5, -1.0]], [[0.0]]],
+        output_token_sampling_mask=[
+            SamplingMaskChunk(
+                lengths=np.array([2], np.int32),
+                token_ids=np.array([7, 8], np.int32),
+                logprobs=np.array([-0.5, -1.0], np.float32),
+            ),
+            SamplingMaskChunk(
+                lengths=np.array([1], np.int32),
+                token_ids=np.array([9], np.int32),
+                logprobs=np.array([0.0], np.float32),
+            ),
+        ],
         output_hidden_states=[None, None],
         routed_experts=[None, None],
         indexer_topk=[None, None],
@@ -123,10 +136,9 @@ def test_batch_str_output_keeps_sampling_distribution_aligned(self):
 
         single_output = _handle_output_by_index(output, 0)
 
-        self.assertEqual(single_output.output_token_sampling_mask, [[[7, 8]]])
+        (chunk,) = single_output.output_token_sampling_mask
         self.assertEqual(
-            single_output.output_token_sampling_logprobs,
-            [[[-0.5, -1.0]]],
+            chunk.to_lists(support_logprobs=True), ([[7, 8]], [[-0.5, -1.0]])
         )
 
     def test_batch_str_output_without_weight_versions_stays_none(self):
diff --git a/test/registered/unit/managers/test_output_streamer_customized_info.py b/test/registered/unit/managers/test_output_streamer_customized_info.py
index d6b438f2ee79..d06d2ab532b1 100644
--- a/test/registered/unit/managers/test_output_streamer_customized_info.py
+++ b/test/registered/unit/managers/test_output_streamer_customized_info.py
@@ -1,13 +1,23 @@
+import pickle
 import unittest
+from array import array
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
+import numpy as np
+
 from sglang.srt.disaggregation.utils import DisaggregationMode
-from sglang.srt.managers.io_struct import unwrap_from_pickle
+from sglang.srt.managers import io_struct
+from sglang.srt.managers.io_struct import (
+    msgpack_decode,
+    msgpack_encode,
+    unwrap_from_pickle,
+)
 from sglang.srt.managers.scheduler_components.output_streamer import (
     SchedulerOutputStreamer,
     _GenerationStreamAccumulator,
 )
+from sglang.srt.sampling.sampling_mask import SamplingMaskRows
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils.weight_versions import (
     WeightVersionSpan,
@@ -27,6 +37,7 @@ def __init__(
         customized_info=None,
         *,
         finished=False,
+        sampling_mask_rows=None,
     ):
         self.rid = rid
         self.http_worker_ipc = None
@@ -58,7 +69,8 @@ def __init__(
         self.return_hidden_states = False
         self.return_routed_experts = False
         self.return_indexer_topk = False
-        self.return_sampling_mask = False
+        self.return_sampling_mask = sampling_mask_rows is not None
+        self.sampling_mask_rows = sampling_mask_rows
         self.mm_image_tokens = 0
         self.mm_audio_tokens = 0
         self.mm_video_tokens = 0
@@ -76,12 +88,13 @@ def check_match_stop_str_prefix(self):
         return False
 
 
-def _accumulator(current_weight_version="default"):
+def _accumulator(current_weight_version="default", return_sampling_mask=False):
     return _GenerationStreamAccumulator(
         return_logprob=False,
         return_hidden_states=False,
         return_routed_experts=False,
         return_indexer_topk=False,
+        return_sampling_mask=return_sampling_mask,
         spec_algorithm=SpeculativeAlgorithm.NONE,
         disaggregation_mode=DisaggregationMode.NULL,
         default_stream_interval=1,
@@ -369,5 +382,75 @@ def test_payload_omits_spans_while_all_requests_stream(self):
         self.assertIsNone(payload.weight_versions)
 
 
+class TestOutputStreamerSamplingMasks(unittest.TestCase):
+    def test_rows_stream_once_and_expand_to_response_lists(self):
+        """Queued rows cross either IPC codec once, in batch order, as per-token lists;
+        rows queued later never overwrite a chunk that is still referenced."""
+        codecs = {
+            "pickle": lambda payload: pickle.loads(
+                pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)
+            ),
+            "msgpack": lambda payload: msgpack_decode(msgpack_encode(payload)),
+        }
+        for codec, round_trip in codecs.items():
+            with (
+                self.subTest(codec=codec),
+                patch.object(io_struct, "_USE_PICKLE_IPC", codec == "pickle"),
+            ):
+                selected_rows = SamplingMaskRows()
+                selected_rows.append(
+                    np.array([5, 3], np.int32), np.array([-0.5], np.float32)
+                )
+                selected_rows.append(
+                    np.array([7], np.int32), np.array([0.0], np.float32)
+                )
+                support_rows = SamplingMaskRows()
+                support_rows.append(
+                    np.array([9, 8], np.int32), np.array([-0.25, -1.5], np.float32)
+                )
+                support_rows.append(
+                    np.array([6], np.int32), np.array([0.0], np.float32)
+                )
+                selected = _FakeReq(
+                    "selected", array("q", [5, 7]), sampling_mask_rows=selected_rows
+                )
+                support = _FakeReq(
+                    "support", array("q", [9, 6]), sampling_mask_rows=support_rows
+                )
+                accumulator = _accumulator(return_sampling_mask=True)
+                for req in (selected, _FakeReq("plain", array("q", [1])), support):
+                    accumulator.accept(req=req)
+                payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False)
+                masks = round_trip(payload).output_token_sampling_mask
+                self.assertEqual(
+                    masks[0].to_lists(support_logprobs=False),
+                    ([[5, 3], [7]], [-0.5, 0.0]),
+                )
+                self.assertIsNone(masks[1])
+                self.assertEqual(
+                    masks[2].to_lists(support_logprobs=True),
+                    ([[9, 8], [6]], [[-0.25, -1.5], [0.0]]),
+                )
+
+                selected.output_ids.append(4)
+                selected.sampling_mask_rows.append(
+                    np.array([4], np.int32), np.array([-2.0], np.float32)
+                )
+                self.assertEqual(
+                    payload.output_token_sampling_mask[0].to_lists(
+                        support_logprobs=False
+                    ),
+                    ([[5, 3], [7]], [-0.5, 0.0]),
+                )
+                accumulator = _accumulator(return_sampling_mask=True)
+                accumulator.accept(req=selected)
+                masks = round_trip(
+                    accumulator.to_payload(dp_rank=0, is_idle_batch=False)
+                ).output_token_sampling_mask
+                self.assertEqual(
+                    masks[0].to_lists(support_logprobs=False), ([[4]], [-2.0])
+                )
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
index 2c09c2efc90a..0fbf721f7f6e 100644
--- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
+++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
@@ -201,6 +201,8 @@ def _make_batch_str_output(rid: str, finished_reason=None) -> BatchStrOutput:
             kwargs[f.name] = [[]]
         elif f.name in _PER_REQUEST_OPTIONAL_FIELDS:
             kwargs[f.name] = [None]
+        elif f.name == "output_token_sampling_mask":
+            kwargs[f.name] = None
         # Fields with class defaults — skip, let the default be used
         elif (
             f.default is not msgspec.NODEFAULT

From f1da9579c02802a917a3161797f1f2d58e55a27b Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Thu, 24 Sep 2026 21:54:58 +0000
Subject: [PATCH 2/5] [Sampling] Test that clients receive sampling masks
 unchanged

The host transport must hand clients exactly what the sampler produced. Add
two checks of that contract:

- A CPU test drives sampler-shaped outputs through materialization, the
  per-request queues, stream emissions at arbitrary steps and both IPC
  codecs, and compares every request's lists with the direct conversion of
  the sampler's rows: token_ids[row, :length] with selected_logprobs[row] or
  support_logprobs[row, :length].
- A GPU test with deterministic inference checks that, for the same seeded
  tokens, selected and support mode return the same supports, each selected
  logprob equals the sampled token's support logprob, and a per-token stream
  matches the response emitted every 50 tokens.
---
 .../registered/sampling/test_sampling_mask.py | 55 +++++++++++
 .../test_output_streamer_customized_info.py   | 99 +++++++++++++++++--
 2 files changed, 147 insertions(+), 7 deletions(-)

diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py
index 5bb347332f83..7aeec97cc1ab 100644
--- a/test/registered/sampling/test_sampling_mask.py
+++ b/test/registered/sampling/test_sampling_mask.py
@@ -796,6 +796,61 @@ def test_return_sampling_mask_preserves_deterministic_sampling(self):
             outputs.append((output["output_ids"], output["text"]))
         self.assertEqual(outputs[0], outputs[1])
 
+    def test_modes_and_streaming_return_identical_sampling_masks(self):
+        """For the same seeded tokens, selected and support mode report the same
+        supports, each selected logprob is the sampled token's support logprob, and a
+        stream that emits every token matches the response emitted every 50 tokens."""
+        sampling_params = {
+            "top_k": _TOP_K,
+            "top_p": 1.0,
+            "sampling_seed": _SAMPLING_SEED,
+            "max_new_tokens": 64,
+        }
+        responses = {}
+        for mode, stream in (
+            ("support", False),
+            ("selected", False),
+            ("support", True),
+        ):
+            response = self._post_generate(
+                sampling_params, sampling_logprobs_mode=mode, stream=stream
+            )
+            self.assertEqual(response.status_code, 200, response.text)
+            if stream:
+                chunks = [
+                    json.loads(line[6:])
+                    for line in response.iter_lines()
+                    if line.startswith(b"data: ") and line[6:] != b"[DONE]"
+                ]
+                responses[mode, stream] = chunks[-1]
+            else:
+                responses[mode, stream] = response.json()
+
+        support = responses["support", False]
+        output_ids = support["output_ids"]
+        masks = self._assert_sampling_masks(output_ids, support["meta_info"])
+        support_logprobs = support["meta_info"]["output_token_sampling_logprobs"]
+
+        selected = responses["selected", False]
+        self.assertEqual(selected["output_ids"], output_ids)
+        self.assertEqual(selected["meta_info"]["output_token_sampling_mask"], masks)
+        self.assertEqual(
+            selected["meta_info"]["output_token_sampling_logprobs"],
+            [
+                logprobs[mask.index(token)]
+                for token, mask, logprobs in zip(output_ids, masks, support_logprobs)
+            ],
+        )
+
+        streamed = responses["support", True]
+        self.assertEqual(streamed["output_ids"], output_ids)
+        for key in (
+            "output_token_sampling_mask",
+            "output_token_sampling_logprobs",
+            "output_token_sampling_mask_length",
+        ):
+            self.assertEqual(streamed["meta_info"][key], support["meta_info"][key])
+
 
 class TestSamplingMaskPytorch(TestSamplingMask):
     _sampling_backend = "pytorch"
diff --git a/test/registered/unit/managers/test_output_streamer_customized_info.py b/test/registered/unit/managers/test_output_streamer_customized_info.py
index d06d2ab532b1..8817dc762561 100644
--- a/test/registered/unit/managers/test_output_streamer_customized_info.py
+++ b/test/registered/unit/managers/test_output_streamer_customized_info.py
@@ -1,3 +1,4 @@
+import dataclasses
 import pickle
 import unittest
 from array import array
@@ -5,14 +6,19 @@
 from unittest.mock import Mock, patch
 
 import numpy as np
+import torch
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput, SamplingMaskStatus
 from sglang.srt.managers import io_struct
 from sglang.srt.managers.io_struct import (
     msgpack_decode,
     msgpack_encode,
     unwrap_from_pickle,
 )
+from sglang.srt.managers.scheduler_components.batch_result_processor import (
+    SchedulerBatchResultProcessor,
+)
 from sglang.srt.managers.scheduler_components.output_streamer import (
     SchedulerOutputStreamer,
     _GenerationStreamAccumulator,
@@ -382,17 +388,19 @@ def test_payload_omits_spans_while_all_requests_stream(self):
         self.assertIsNone(payload.weight_versions)
 
 
+_IPC_ROUND_TRIPS = {
+    "pickle": lambda payload: pickle.loads(
+        pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)
+    ),
+    "msgpack": lambda payload: msgpack_decode(msgpack_encode(payload)),
+}
+
+
 class TestOutputStreamerSamplingMasks(unittest.TestCase):
     def test_rows_stream_once_and_expand_to_response_lists(self):
         """Queued rows cross either IPC codec once, in batch order, as per-token lists;
         rows queued later never overwrite a chunk that is still referenced."""
-        codecs = {
-            "pickle": lambda payload: pickle.loads(
-                pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)
-            ),
-            "msgpack": lambda payload: msgpack_decode(msgpack_encode(payload)),
-        }
-        for codec, round_trip in codecs.items():
+        for codec, round_trip in _IPC_ROUND_TRIPS.items():
             with (
                 self.subTest(codec=codec),
                 patch.object(io_struct, "_USE_PICKLE_IPC", codec == "pickle"),
@@ -451,6 +459,83 @@ def test_rows_stream_once_and_expand_to_response_lists(self):
                     masks[0].to_lists(support_logprobs=False), ([[4]], [-2.0])
                 )
 
+    def test_clients_receive_the_sampler_rows_unchanged(self):
+        """Across decode steps, stream emissions and either IPC codec, each request gets
+        the sampler's rows exactly as the tensor-to-list conversion defines them."""
+        processor = SchedulerBatchResultProcessor(
+            **{f.name: None for f in dataclasses.fields(SchedulerBatchResultProcessor)}
+        )
+        modes = ["selected", "support", None, "support", "selected"]
+        mask_reqs = [i for i, mode in enumerate(modes) if mode is not None]
+        support_reqs = [i for i in mask_reqs if modes[i] == "support"]
+        for codec, round_trip in _IPC_ROUND_TRIPS.items():
+            with (
+                self.subTest(codec=codec),
+                patch.object(io_struct, "_USE_PICKLE_IPC", codec == "pickle"),
+            ):
+                generator = torch.Generator().manual_seed(0)
+                reqs = []
+                for i, mode in enumerate(modes):
+                    rows = None if mode is None else SamplingMaskRows()
+                    req = _FakeReq(f"r{i}", array("q"), sampling_mask_rows=rows)
+                    req.sampling_logprobs_mode = mode
+                    reqs.append(req)
+                expected = [([], []) for _ in modes]
+                received = [([], []) for _ in modes]
+                for step in range(12):
+                    width = (1, 7, 64)[step % 3]
+                    lengths = torch.randint(
+                        1, width + 1, (len(mask_reqs),), generator=generator
+                    )
+                    token_ids = torch.randint(
+                        0, 1 << 20, (len(mask_reqs), width), generator=generator
+                    ).int()
+                    selected = torch.randn(len(mask_reqs), generator=generator)
+                    support = torch.randn(len(support_reqs), width, generator=generator)
+                    output = LogitsProcessorOutput(
+                        next_token_logits=None,
+                        sampling_mask_output=SimpleNamespace(
+                            token_ids=token_ids,
+                            lengths=lengths,
+                            selected_logprobs=selected,
+                            support_logprobs=support,
+                            statuses=torch.full(
+                                (len(mask_reqs),), SamplingMaskStatus.OK
+                            ),
+                        ),
+                    )
+                    SchedulerBatchResultProcessor.materialize_sampling_mask_output(
+                        reqs, output
+                    )
+                    for row, i in enumerate(mask_reqs):
+                        processor.add_sampling_mask_return_values(i, reqs[i], output)
+                        length = int(lengths[row])
+                        expected[i][0].append(token_ids[row, :length].tolist())
+                        if modes[i] == "support":
+                            support_row = support_reqs.index(i)
+                            expected[i][1].append(
+                                support[support_row, :length].tolist()
+                            )
+                        else:
+                            expected[i][1].append(float(selected[row]))
+                    if step in (0, 1, 5, 6, 11):
+                        accumulator = _accumulator(return_sampling_mask=True)
+                        for req in reqs:
+                            accumulator.accept(req=req)
+                        chunks = round_trip(
+                            accumulator.to_payload(dp_rank=0, is_idle_batch=False)
+                        ).output_token_sampling_mask
+                        for i, chunk in enumerate(chunks):
+                            if modes[i] is None:
+                                self.assertIsNone(chunk)
+                                continue
+                            masks, logprobs = chunk.to_lists(
+                                support_logprobs=modes[i] == "support"
+                            )
+                            received[i][0].extend(masks)
+                            received[i][1].extend(logprobs)
+                self.assertEqual(received, expected)
+
 
 if __name__ == "__main__":
     unittest.main()

From 60349679d0f42dc1a915f0b76d82a15d731e7d1e Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Thu, 24 Sep 2026 22:47:00 +0000
Subject: [PATCH 3/5] [Sampling] Pickle sampling-mask chunks as raw buffers

Unpickling a chunk through numpy's array reducer leaves about nine
GC-tracked objects per request, against three for main's nested lists.
When every message carries one short chunk per request, as with per-token
streaming, those objects trigger extra full collections of the tokenizer
manager's heap. Streaming at top_k=1000 x 256 in support mode spent 19% more
tokenizer CPU than main.

Pickle a chunk as its three raw buffers and rebuild it with np.frombuffer,
which leaves no tracked objects; the struct itself is untracked. The bytes
are unchanged, and protocols below 5 keep the default reduction.
---
 python/sglang/srt/sampling/sampling_mask.py | 27 ++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/sampling/sampling_mask.py b/python/sglang/srt/sampling/sampling_mask.py
index d9c4e4bde35a..aeebac6b97c1 100644
--- a/python/sglang/srt/sampling/sampling_mask.py
+++ b/python/sglang/srt/sampling/sampling_mask.py
@@ -1,5 +1,6 @@
 """Host-side transport of output-token sampling masks."""
 
+import pickle
 import sys
 from array import array
 from typing import List, Optional, Tuple, Union
@@ -8,7 +9,9 @@
 import numpy as np
 
 
-class SamplingMaskChunk(msgspec.Struct, frozen=True, kw_only=True, array_like=True):
+class SamplingMaskChunk(
+    msgspec.Struct, frozen=True, kw_only=True, array_like=True, gc=False
+):
     """One request's sampling masks in one output batch, one row per output token."""
 
     # int32 [num_rows]
@@ -18,6 +21,20 @@ class SamplingMaskChunk(msgspec.Struct, frozen=True, kw_only=True, array_like=Tr
     # float32 [sum(lengths)] in support mode, [num_rows] in selected mode
     logprobs: np.ndarray
 
+    def __reduce_ex__(self, protocol: int):
+        # Unpickling with numpy's reducer leaves GC-tracked objects per array; with a
+        # chunk per request per message they trigger full collections downstream.
+        if protocol < 5:
+            return super().__reduce_ex__(protocol)
+        return (
+            _chunk_from_buffers,
+            (
+                pickle.PickleBuffer(self.lengths),
+                pickle.PickleBuffer(self.token_ids),
+                pickle.PickleBuffer(self.logprobs),
+            ),
+        )
+
     def to_lists(
         self, support_logprobs: bool
     ) -> Tuple[List[List[int]], List[Union[float, List[float]]]]:
@@ -35,6 +52,14 @@ def to_lists(
         return masks, values
 
 
+def _chunk_from_buffers(lengths, token_ids, logprobs) -> SamplingMaskChunk:
+    return SamplingMaskChunk(
+        lengths=np.frombuffer(lengths, dtype=np.int32),
+        token_ids=np.frombuffer(token_ids, dtype=np.int32),
+        logprobs=np.frombuffer(logprobs, dtype=np.float32),
+    )
+
+
 class _GrowableArray(msgspec.Struct):
     """Append-only 1-D buffer; doubling capacity copies each value O(1) times.
 

From 3e796ca8b219d166e9d9adc25f5cb59578e5e2dd Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Thu, 24 Sep 2026 23:14:53 +0000
Subject: [PATCH 4/5] [Sampling] Pickle sampling-mask chunk buffers as
 read-only bytes

Writable buffers unpickle as bytearrays, and np.frombuffer keeps a GC-tracked
memoryview as the base of an array over a bytearray. That left three tracked
objects per request per message and still doubled the tokenizer manager's
young-generation collections during streaming, which promoted its streamed
lists into more full collections.

Pickle read-only views instead: they unpickle as bytes, which np.frombuffer
wraps without a tracked base. Received chunks are only read, so read-only
arrays suffice.
---
 python/sglang/srt/sampling/sampling_mask.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/sampling/sampling_mask.py b/python/sglang/srt/sampling/sampling_mask.py
index aeebac6b97c1..cb3c319b07c2 100644
--- a/python/sglang/srt/sampling/sampling_mask.py
+++ b/python/sglang/srt/sampling/sampling_mask.py
@@ -22,16 +22,15 @@ class SamplingMaskChunk(
     logprobs: np.ndarray
 
     def __reduce_ex__(self, protocol: int):
-        # Unpickling with numpy's reducer leaves GC-tracked objects per array; with a
-        # chunk per request per message they trigger full collections downstream.
+        # Read-only buffers unpickle as bytes, which np.frombuffer wraps without the
+        # GC-tracked objects that numpy's reducer or a bytearray would leave per array.
         if protocol < 5:
             return super().__reduce_ex__(protocol)
         return (
             _chunk_from_buffers,
-            (
-                pickle.PickleBuffer(self.lengths),
-                pickle.PickleBuffer(self.token_ids),
-                pickle.PickleBuffer(self.logprobs),
+            tuple(
+                pickle.PickleBuffer(memoryview(values).toreadonly())
+                for values in (self.lengths, self.token_ids, self.logprobs)
             ),
         )
 

From abdee233fa11cc1090c3ea83d14403d0a7bf6e2d Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Fri, 25 Sep 2026 08:39:29 +0000
Subject: [PATCH 5/5] [Sampling] Test the PD handoff of the first sampling-mask
 row

The prefill worker writes the first token's row into the metadata buffers,
and the decode worker reads it back into its queue. Check that the rows the
decode worker streams equal the row the prefill worker queued, in selected
and support mode.
---
 .../test_disaggregation_wire.py               | 62 +++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py
index 9e324f15ee91..ccd765278fa3 100644
--- a/test/registered/unit/disaggregation/test_disaggregation_wire.py
+++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py
@@ -24,6 +24,7 @@
     unpack_int_lists,
     unpack_list_of_buffers,
 )
+from sglang.srt.disaggregation.decode import DecodeRequest, DecodeTransferQueue
 from sglang.srt.disaggregation.decode_schedule_batch_mixin import (
     ScheduleBatchDisaggregationDecodeMixin,
 )
@@ -825,6 +826,67 @@ def test_sampling_mask_selected_logprob_uses_first_metadata_slot(self):
             self.assertEqual(mask.tolist(), [7, 8, 9])
             self.assertEqual(logprobs[0].item(), -0.5)
 
+    def test_sampling_mask_row_reaches_decode_unchanged(self):
+        """The prefill worker's first-token row is the row the decode worker streams."""
+        for mode, logprobs, expected in (
+            ("support", [-1.25, -1.5, -2.0], [[-1.25, -1.5, -2.0]]),
+            ("selected", [-0.5], [-0.5]),
+        ):
+            with (
+                self.subTest(sampling_logprobs_mode=mode),
+                envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(True),
+            ):
+                buffers = MetadataBuffers(
+                    size=1,
+                    hidden_size=2,
+                    hidden_states_dtype=torch.float32,
+                    max_sampling_mask_tokens=4,
+                )
+                buffers.set_buf(
+                    self._make_req(
+                        None,
+                        sampling_mask=[7, 8, 9],
+                        sampling_logprobs=logprobs,
+                        sampling_logprobs_mode=mode,
+                    )
+                )
+                queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
+                queue.scheduler = SimpleNamespace(
+                    kv_checksum_computer=None,
+                    batch_result_processor=SimpleNamespace(
+                        _maybe_update_reasoning_tokens=lambda req, token_id: None
+                    ),
+                )
+                queue.spec_algorithm = SimpleNamespace(is_none=lambda: True)
+                queue.metadata_buffers = buffers
+                req = SimpleNamespace(
+                    rid="r0",
+                    bootstrap_host="127.0.0.1",
+                    bootstrap_room=9,
+                    output_ids=[],
+                    return_logprob=False,
+                    return_sampling_mask=True,
+                    sampling_logprobs_mode=mode,
+                    sampling_mask_rows=SamplingMaskRows(),
+                    time_stats=SimpleNamespace(set_wait_queue_entry_time=lambda: None),
+                )
+
+                queue._commit_transfer_to_req(
+                    DecodeRequest(
+                        req=req,
+                        kv_receiver=SimpleNamespace(clear=lambda: None),
+                        metadata_buffer_index=0,
+                    )
+                )
+
+                self.assertEqual(req.output_ids, [101])
+                self.assertEqual(
+                    req.sampling_mask_rows.take().to_lists(
+                        support_logprobs=mode == "support"
+                    ),
+                    ([[7, 8, 9]], expected),
+                )
+
     def test_decode_input_requires_valid_seed_for_every_request(self):
         seeds = (
             torch.tensor([1, 2, 3], dtype=torch.int32),