From fd7743e0e1725f9f24844e6f3446fcf0f9ade815 Mon Sep 17 00:00:00 2001 From: Byron Hsu Date: Thu, 10 Sep 2026 14:52:44 -0700 Subject: [PATCH] [Sampling] Support sampling masks with overlap scheduling (#36631) Co-authored-by: ByronHsu Co-authored-by: root Co-authored-by: Byron Hsu --- python/sglang/srt/arg_groups/field_order.py | 1 + python/sglang/srt/arg_groups/fields/exec_.py | 7 + python/sglang/srt/arg_groups/pipeline.py | 2 + .../sglang/srt/arg_groups/validation_hook.py | 17 + python/sglang/srt/disaggregation/prefill.py | 25 + python/sglang/srt/disaggregation/utils.py | 62 +- python/sglang/srt/environ.py | 5 +- python/sglang/srt/layers/logits_processor.py | 31 +- python/sglang/srt/layers/sampler.py | 200 +++---- python/sglang/srt/managers/scheduler.py | 65 ++- .../batch_result_processor.py | 158 +++++- .../sglang/srt/managers/scheduler_pp_mixin.py | 27 +- python/sglang/srt/managers/utils.py | 41 +- .../srt/sampling/sampling_batch_info.py | 44 +- .../registered/sampling/test_sampling_mask.py | 537 +++++++++++------- .../test_disaggregation_wire.py | 59 +- .../test_prefill_abort_result_cleanup.py | 54 ++ .../test_specv2_kvcache_offloading.py | 47 ++ ...st_batch_result_processor_hidden_states.py | 160 +++++- ...t_batch_result_processor_mamba_boundary.py | 4 +- .../test_generation_auxiliary_output.py | 60 +- ...test_scheduler_sampling_mask_validation.py | 76 +++ .../unit/sampling/test_sampling_batch_info.py | 48 ++ .../unit/server_args/test_server_args.py | 32 ++ 24 files changed, 1365 insertions(+), 397 deletions(-) create mode 100644 test/registered/unit/managers/test_scheduler_sampling_mask_validation.py diff --git a/python/sglang/srt/arg_groups/field_order.py b/python/sglang/srt/arg_groups/field_order.py index 69656c606..992a91b17 100644 --- a/python/sglang/srt/arg_groups/field_order.py +++ b/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", diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index 62dcd74be..e9f7e43ab 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/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.", diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index ef35ba1da..b51a627c2 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/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. diff --git a/python/sglang/srt/arg_groups/validation_hook.py b/python/sglang/srt/arg_groups/validation_hook.py index 44241c633..2a5e41dbd 100644 --- a/python/sglang/srt/arg_groups/validation_hook.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index b91ea8161..a797d375a 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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: diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 9c4acca59..ca0919bfa 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5e7d4086a..8e4125df5 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 0b7a96a69..b9835654a 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/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] diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 7781f608b..0b9323adf 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/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, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 25619df09..ef046fe91 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: 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 94c710388..9fa4e9b85 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 267410e6c..0d28ee302 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/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() diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 2695df7d7..88ee08f5d 100644 --- a/python/sglang/srt/managers/utils.py +++ b/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"], diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index d25305203..fad02f235 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/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) diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py index be27da24d..187e805dc 100644 --- a/test/registered/sampling/test_sampling_mask.py +++ b/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() diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index a3a123276..3ba0f3cb3 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/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), diff --git a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py index cc36276bc..af280b2bb 100644 --- a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py +++ b/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"])) diff --git a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py index 4c143ac3b..a1c9c8b16 100644 --- a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py +++ b/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() 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 479c7ed0f..363515e10 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 @@ -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, ) diff --git a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py index 53fa3b497..97eb071fb 100644 --- a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py +++ b/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, ) diff --git a/test/registered/unit/managers/test_generation_auxiliary_output.py b/test/registered/unit/managers/test_generation_auxiliary_output.py index 9a4df9f2c..21e6d449f 100644 --- a/test/registered/unit/managers/test_generation_auxiliary_output.py +++ b/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( diff --git a/test/registered/unit/managers/test_scheduler_sampling_mask_validation.py b/test/registered/unit/managers/test_scheduler_sampling_mask_validation.py new file mode 100644 index 000000000..ebc47b313 --- /dev/null +++ b/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() diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 886223c74..91cefe30a 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/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.""" diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 5e9f91bc6..d4f23657c 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/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)