minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: ByronHsu <ByronHsu@users.noreply.github.com>
Co-authored-by: root <root@slurm-h200-208-179.slurm-compute.tenant-slurm.svc.cluster.local>
Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>
This commit is contained in:
Byron Hsu
2026-09-10 14:52:44 -07:00
committed by GitHub
co-authored by ByronHsu root Byron Hsu
parent 55b45cb45a
commit fd7743e0e1
24 changed files with 1365 additions and 397 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/field_order.py
+1
View File
@@ -496,6 +496,7 @@ POSITIONAL_FIELD_ORDER = (
"return_hidden_states_mode",
"enable_return_routed_experts",
"enable_return_indexer_topk",
"sampling_mask_max_tokens",
"disable_outlines_disk_cache",
"enable_mis",
"weight_cache_mode",
python/sglang/srt/arg_groups/fields/exec_.py
+7
View File
@@ -95,6 +95,13 @@ class ExecFeatures(msgspec.Struct):
bool,
"Enable returning indexer topk indices of layers with indexer with responses.",
] = False
sampling_mask_max_tokens: A[
int,
"The maximum number of token IDs in a returned sampling mask. Requests "
"are aborted if their realized sampling support exceeds this limit. "
"Use the same value on disaggregated prefill and decode nodes; clients "
"should set top_k below the limit to leave headroom for cutoff ties.",
] = 4096
disable_outlines_disk_cache: A[
bool,
"Disable disk cache of outlines to avoid possible crashes related to file system or high concurrency.",
python/sglang/srt/arg_groups/pipeline.py
+2
View File
@@ -95,9 +95,11 @@ def run_resolution_pipeline(server_args: Any) -> None:
from sglang.srt.arg_groups.validation_hook import (
validate_experimental_sgl_marlin,
validate_prefill_decode_interval,
validate_sampling_mask_max_tokens,
)
validate_prefill_decode_interval(server_args)
validate_sampling_mask_max_tokens(server_args)
# Reject an explicitly enabled but incompatible hardware runtime before
# model path resolution, downloads, or the dummy-model short circuit.
python/sglang/srt/arg_groups/validation_hook.py
+17
View File
@@ -17,6 +17,7 @@ from sglang.srt.arg_groups.overrides import (
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
parse_ib_device_config,
)
from sglang.srt.environ import envs
from sglang.srt.runtime_context import get_platform
from sglang.srt.utils.common import torch_release
from sglang.srt.utils.runai_utils import is_runai_obj_uri
@@ -438,6 +439,22 @@ def validate_prefill_decode_interval(server_args: Any):
raise ValueError("--prefill-decode-interval must be non-negative.")
def validate_sampling_mask_max_tokens(server_args: Any):
if envs.SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.is_set():
raise ValueError(
"SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS is no longer supported. "
"Unset it. To enable sampling masks for disaggregated serving, set "
"SGLANG_ENABLE_DISAGG_SAMPLING_MASK=1 and use the same positive "
"--sampling-mask-max-tokens value on both prefill and decode servers."
)
cfg = resolving_view(server_args)
if cfg.sampling_mask_max_tokens <= 0:
raise ValueError(
"--sampling-mask-max-tokens must be positive "
f"(got {cfg.sampling_mask_max_tokens})."
)
def check_two_batch_overlap(server_args: Any):
# With no EP a2a backend, two-batch-overlap is only valid on the non-EP
# DP TP-MoE path (overlapping the DP all_gatherv / reduce_scatterv with
python/sglang/srt/disaggregation/prefill.py
+25
View File
@@ -754,6 +754,10 @@ class SchedulerDisaggregationPrefillMixin:
batch=batch,
logits_output=logits_output,
)
if logits_output is not None and logits_output.sampling_mask_output is not None:
self.batch_result_processor.materialize_sampling_mask_output(
batch.reqs, logits_output
)
def advance_logprob_pt(i: int, req: Req) -> None:
nonlocal logprob_pt
@@ -783,6 +787,27 @@ class SchedulerDisaggregationPrefillMixin:
advance_logprob_pt(i, req)
continue
sampling_mask_finish_reason = None
if req.return_sampling_mask:
assert logits_output is not None
statuses = logits_output.next_token_sampling_mask_status
status = None if statuses is None else statuses[i]
sampling_mask_finish_reason = (
self.batch_result_processor.get_sampling_mask_finish_reason(
status=status
)
)
if sampling_mask_finish_reason is not None:
req.to_finish = sampling_mask_finish_reason
req.time_stats.trace_ctx.abort(
abort_info={"reason": sampling_mask_finish_reason.message}
)
if self._retire_aborted_prefill_result(req):
req.time_stats.set_completion_time()
aborted_reqs.append(req)
advance_logprob_pt(i, req)
continue
req.output_ids.append(next_token_id)
if req.grammar is not None:
try:
python/sglang/srt/disaggregation/utils.py
+25 -37
View File
@@ -329,18 +329,14 @@ class MetadataBuffers:
size: int,
hidden_size: int,
hidden_states_dtype: torch.dtype,
max_sampling_mask_tokens: int,
max_top_logprobs_num: int = 128,
max_sampling_mask_tokens: Optional[int] = None,
custom_mem_pool: torch.cuda.MemPool = None,
output_dsa_topk_indices_dim: int = 0,
):
self.custom_mem_pool = custom_mem_pool
self.output_dsa_topk_indices_dim = output_dsa_topk_indices_dim
if max_sampling_mask_tokens is None:
max_sampling_mask_tokens = (
envs.SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.get()
)
self.enable_sampling_mask = max_sampling_mask_tokens > 0
self.enable_sampling_mask = envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.get()
bootstrap_room_dtype = torch.uint64
device = "cpu"
if is_npu():
@@ -423,38 +419,23 @@ class MetadataBuffers:
self.output_token_logprobs_idx,
self.output_top_logprobs_val,
self.output_top_logprobs_idx,
self.output_token_sampling_mask_len,
self.output_token_sampling_mask_idx,
self.output_token_sampling_logprobs,
self.output_topk_p,
self.output_topk_index,
self.output_hidden_states,
]
if self.enable_sampling_mask:
bufs.extend(
[
self.output_token_sampling_mask_len,
self.output_token_sampling_mask_idx,
self.output_token_sampling_logprobs,
]
)
bufs.extend(
[
self.output_topk_p,
self.output_topk_index,
self.output_hidden_states,
]
)
if self.output_dsa_topk_indices is not None:
bufs.append(self.output_dsa_topk_indices)
bufs.append(self.bootstrap_room)
bufs = [buf for buf in bufs if buf is not None]
ptrs = [buf.data_ptr() for buf in bufs]
data_lens = [buf.nbytes for buf in bufs]
item_lens = [buf[0].nbytes for buf in bufs]
return ptrs, data_lens, item_lens
def get_buf(self, idx: int):
sampling_mask_len = None
sampling_mask_idx = None
sampling_logprobs = None
if self.enable_sampling_mask:
sampling_mask_len = self.output_token_sampling_mask_len[idx].clone()
sampling_mask_idx = self.output_token_sampling_mask_idx[idx].clone()
sampling_logprobs = self.output_token_sampling_logprobs[idx].clone()
return (
self.output_ids[idx].clone(),
self.cached_tokens[idx].clone(),
@@ -462,9 +443,21 @@ class MetadataBuffers:
self.output_token_logprobs_idx[idx].clone(),
self.output_top_logprobs_val[idx].clone(),
self.output_top_logprobs_idx[idx].clone(),
sampling_mask_len,
sampling_mask_idx,
sampling_logprobs,
(
self.output_token_sampling_mask_len[idx].clone()
if self.enable_sampling_mask
else None
),
(
self.output_token_sampling_mask_idx[idx].clone()
if self.enable_sampling_mask
else None
),
(
self.output_token_sampling_logprobs[idx].clone()
if self.enable_sampling_mask
else None
),
self.output_topk_p[idx].clone(),
self.output_topk_index[idx].clone(),
self.output_hidden_states[idx].clone(),
@@ -532,11 +525,6 @@ class MetadataBuffers:
device="cpu",
)
if req.return_sampling_mask:
if not self.enable_sampling_mask:
raise RuntimeError(
"return_sampling_mask with disaggregation requires "
"SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS > 0."
)
# Sentinel -1: the decode side records None for this handoff token.
self.output_token_sampling_mask_len[req.metadata_buffer_index][0] = -1
sampling_masks = req.output_token_sampling_mask
@@ -551,7 +539,7 @@ class MetadataBuffers:
raise RuntimeError(
f"Sampling mask length {mask_len} exceeds disaggregation "
f"metadata capacity {max_mask_len}. Increase "
"SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS."
"--sampling-mask-max-tokens."
)
self.output_token_sampling_mask_len[req.metadata_buffer_index][
0
python/sglang/srt/environ.py
+4 -1
View File
@@ -666,6 +666,10 @@ class Envs:
# computed dynamically at runtime based on cpu_count; see disaggregation backends.
SGLANG_DISAGGREGATION_THREAD_POOL_SIZE = EnvInt(None)
SGLANG_DISAGGREGATION_QUEUE_SIZE = EnvInt(4)
# Enable on both P and D with the same --sampling-mask-max-tokens value.
SGLANG_ENABLE_DISAGG_SAMPLING_MASK = EnvBool(False)
# Retained only to reject the removed setting during startup.
SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS = EnvInt(None)
SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT = EnvInt(300)
SGLANG_DISAGGREGATION_ZMQ_SEND_TIMEOUT = EnvInt(1)
SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0)
@@ -679,7 +683,6 @@ class Envs:
SGLANG_DISAGGREGATION_ZMQ_MAX_SOCKETS = EnvInt(16384)
SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False)
SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK = EnvBool(False)
SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS = EnvInt(0)
SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120)
# Deferred decode-side KV release: on abort, hold an in-flight request's KV
# pages/slot until the prefill acks the transfer drained, or the timeout
python/sglang/srt/layers/logits_processor.py
+29 -2
View File
@@ -16,6 +16,7 @@
import dataclasses
import logging
from contextlib import contextmanager
from enum import IntEnum
from typing import Any, Dict, List, Optional, Tuple, Union
import torch
@@ -83,6 +84,30 @@ _UNQUANTIZED_LM_HEAD_METHODS = {
_autotune_run_lm_head: Optional[bool] = None
class SamplingMaskStatus(IntEnum):
"""Ordered by severity so distributed MAX reaches one policy decision."""
OK = 0
OVERFLOW = 1
INVALID = 2
@dataclasses.dataclass
class SamplingMaskOutput:
"""Tensor result for opted-in rows in batch order."""
token_ids: torch.Tensor
lengths: torch.Tensor
selected_logprobs: torch.Tensor
statuses: torch.Tensor
def map_device_tensors(self, fn) -> None:
self.token_ids = fn(self.token_ids)
self.lengths = fn(self.lengths)
self.selected_logprobs = fn(self.selected_logprobs)
self.statuses = fn(self.statuses)
def _trace_e2e_logits(stage: str, **fields) -> None:
if not envs.SGLANG_TRACE_LOGITS_E2E.get():
return
@@ -196,10 +221,12 @@ class LogitsProcessorOutput:
List[Union[List[float], torch.Tensor]]
] = None
next_token_token_ids_logprobs_idx: Optional[List] = None
# Sparse top-k/top-p/min-p support ids and selected-token logprob after
# truncation/renormalization. Only populated when requested.
# Post-filter support IDs, bounded by server capacity, and selected-token
# logprob over the full realized support.
sampling_mask_output: Optional[SamplingMaskOutput] = None
next_token_sampling_mask_idx: Optional[List[Optional[List[int]]]] = None
next_token_sampling_logprobs: Optional[List[Optional[float]]] = None
next_token_sampling_mask_status: Optional[List[Optional[int]]] = None
## Part 3: Prefill-only. This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor
# The logprobs of input tokens. shape: [#token]
python/sglang/srt/layers/sampler.py
+100 -100
View File
@@ -12,7 +12,11 @@ from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.logits_processor import (
LogitsProcessorOutput,
SamplingMaskOutput,
SamplingMaskStatus,
)
from sglang.srt.layers.logprob_processor import (
OutputLogprobProcessor,
)
@@ -93,12 +97,23 @@ class _SamplingMaskCapture(NamedTuple):
batch_rows: torch.Tensor
def _select_sampling_mask_rows(
tensor: torch.Tensor, batch_indices: torch.Tensor
) -> torch.Tensor:
"""Select opted-in rows, returning the input when every row opts in."""
if batch_indices.numel() == tensor.shape[0]:
return tensor
return tensor.index_select(0, batch_indices)
class Sampler(nn.Module):
def __init__(self):
super().__init__()
self.tp_sync_group = get_tp_group().device_group
self.cp_sync_group = None
if is_dp_attention_enabled():
self.tp_sync_group = get_parallel().attn_tp_group.device_group
self.cp_sync_group = get_parallel().attn_cp_group.device_group
self.rl_on_policy_target = get_exec().deterministic.rl_on_policy_target
# In RL on-policy mode, deterministic inference is automatically enabled.
@@ -108,6 +123,7 @@ class Sampler(nn.Module):
# In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer.
self.use_log_softmax_logprob = self.rl_on_policy_target is not None
self.use_ascend_backend = get_exec().kernel.sampling_backend == "ascend"
self.sampling_mask_max_tokens = get_exec().features.sampling_mask_max_tokens
self.output_logprob_processor = OutputLogprobProcessor()
@@ -157,7 +173,8 @@ class Sampler(nn.Module):
_trace_e2e_sampler("preprocess_enter")
logits = self._preprocess_logits(logits, sampling_info)
_trace_e2e_sampler("preprocess_returned")
return_sampling_mask = any(sampling_info.return_sampling_masks or [])
sampling_mask_batch_indices = sampling_info.sampling_mask_batch_indices
return_sampling_mask = sampling_mask_batch_indices is not None
sampling_mask_capture = None
if sampling_info.is_all_greedy:
@@ -249,7 +266,6 @@ class Sampler(nn.Module):
sampling_info,
positions,
simple_sampling_case,
return_sampling_mask=return_sampling_mask,
)
if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB:
logprobs = (
@@ -275,9 +291,12 @@ class Sampler(nn.Module):
_trace_e2e_sampler("token_sync_returned")
if return_sampling_mask:
assert sampling_mask_batch_indices is not None
if sampling_info.is_all_greedy:
self._attach_greedy_sampling_mask_to_output(
logits_output, sampling_info, batch_next_token_ids
logits_output.sampling_mask_output = (
self._build_greedy_sampling_mask_output(
sampling_mask_batch_indices, batch_next_token_ids
)
)
else:
assert sampling_mask_capture is not None
@@ -286,9 +305,7 @@ class Sampler(nn.Module):
sampling_mask_capture = sampling_mask_capture._replace(
selected_weight=None
)
self._attach_sampling_mask_to_output(
logits_output,
sampling_info,
logits_output.sampling_mask_output = self._build_sampling_mask_output(
batch_next_token_ids,
sampling_mask_capture,
)
@@ -302,8 +319,6 @@ class Sampler(nn.Module):
sampling_info: SamplingBatchInfo,
positions: torch.Tensor,
simple_sampling_case: bool,
*,
return_sampling_mask: bool = False,
) -> Tuple[torch.Tensor, Optional[_SamplingMaskCapture]]:
"""Sample from probability distribution (after softmax).
@@ -311,29 +326,13 @@ class Sampler(nn.Module):
Handles both simple (direct multinomial) and complex (top-k/top-p/min-p) cases.
Capture work is performed only when return_sampling_mask is enabled.
"""
capture_rows = sampling_info.sampling_mask_batch_indices
return_sampling_mask = capture_rows is not None
sampling_mask_capture = None
capture_rows = None
capture_all_rows = False
if return_sampling_mask:
capture_rows_list = [
i
for i, should_return in enumerate(
sampling_info.return_sampling_masks or []
)
if should_return
]
if not capture_rows_list:
raise RuntimeError(
"Sampling-mask capture requested without any opted-in batch rows."
)
capture_rows = torch.tensor(
capture_rows_list, device=probs.device, dtype=torch.long
select_capture_rows = partial(
_select_sampling_mask_rows, batch_indices=capture_rows
)
capture_all_rows = capture_rows_list == list(range(probs.shape[0]))
def select_capture_rows(tensor: torch.Tensor) -> torch.Tensor:
assert capture_rows is not None
return tensor if capture_all_rows else tensor.index_select(0, capture_rows)
if simple_sampling_case:
batch_next_token_ids = sampling_from_probs_torch(
@@ -457,62 +456,48 @@ class Sampler(nn.Module):
raise ValueError(f"Invalid sampling backend: {backend}")
return batch_next_token_ids, sampling_mask_capture
def _attach_greedy_sampling_mask_to_output(
def _build_greedy_sampling_mask_output(
self,
logits_output: LogitsProcessorOutput,
sampling_info: SamplingBatchInfo,
batch_indices: torch.Tensor,
batch_next_token_ids: torch.Tensor,
) -> None:
tokens = batch_next_token_ids.to(torch.int32).cpu().tolist()
masks = []
logprobs = []
for i, should_return in enumerate(sampling_info.return_sampling_masks or []):
if should_return:
masks.append([int(tokens[i])])
logprobs.append(0.0)
else:
masks.append(None)
logprobs.append(None)
logits_output.next_token_sampling_mask_idx = masks
logits_output.next_token_sampling_logprobs = logprobs
) -> SamplingMaskOutput:
token_ids = batch_next_token_ids.index_select(0, batch_indices).to(torch.int32)
num_requests = batch_indices.numel()
return SamplingMaskOutput(
token_ids=token_ids.view(-1, 1),
lengths=torch.ones(
num_requests, dtype=torch.int32, device=token_ids.device
),
selected_logprobs=torch.zeros(
num_requests, dtype=torch.float32, device=token_ids.device
),
statuses=torch.full(
(num_requests,),
SamplingMaskStatus.OK,
dtype=torch.int32,
device=token_ids.device,
),
)
def _attach_sampling_mask_to_output(
def _build_sampling_mask_output(
self,
logits_output: LogitsProcessorOutput,
sampling_info: SamplingBatchInfo,
batch_next_token_ids: torch.Tensor,
sampling_mask_capture: _SamplingMaskCapture,
) -> None:
return_sampling_masks = sampling_info.return_sampling_masks or []
if not return_sampling_masks:
logits_output.next_token_sampling_mask_idx = []
logits_output.next_token_sampling_logprobs = []
return
requested_rows_list = [
i for i, should_return in enumerate(return_sampling_masks) if should_return
]
requested_rows = sampling_mask_capture.batch_rows
) -> SamplingMaskOutput:
"""Pack captured positive support into the fixed-cap device result."""
batch_indices = sampling_mask_capture.batch_rows
weights = sampling_mask_capture.weights
token_ids = sampling_mask_capture.token_ids
selected_weight = sampling_mask_capture.selected_weight
sampled_tokens = batch_next_token_ids.index_select(0, requested_rows).view(
-1, 1
)
sampled_tokens = batch_next_token_ids.index_select(0, batch_indices).view(-1, 1)
sampled_tokens_int32 = sampled_tokens.to(torch.int32)
if token_ids is None:
support_token_ids = (
torch.arange(
weights.shape[-1], device=weights.device, dtype=torch.int32
)
.view(1, -1)
.expand_as(weights)
)
selected_from_weights = torch.gather(
weights, 1, sampled_tokens.long()
).squeeze(1)
sampled_in_capture = selected_from_weights > 0
else:
sampled_matches = token_ids == sampled_tokens.to(token_ids.dtype)
sampled_matches = token_ids == sampled_tokens_int32
sampled_in_capture = sampled_matches.any(dim=-1)
selected_positions = sampled_matches.to(torch.int32).argmax(
dim=-1, keepdim=True
@@ -520,48 +505,63 @@ class Sampler(nn.Module):
selected_from_weights = torch.gather(
weights, 1, selected_positions
).squeeze(1)
support_token_ids = token_ids
if selected_weight is None:
selected_weight = selected_from_weights
support = weights > 0
support_mass = weights.sum(dim=-1, dtype=torch.float32)
support_mass = torch.where(support, weights, torch.zeros_like(weights)).sum(
dim=-1, dtype=torch.float32
)
selected_weight = selected_weight.float()
selected_logprobs = torch.log(selected_weight / support_mass)
valid = (
invalid = ~(
sampled_in_capture
& (selected_weight > 0)
& (support_mass > 0)
& torch.isfinite(selected_logprobs)
)
if not bool(torch.all(valid).item()):
invalid_rows = (~valid).nonzero(as_tuple=True)[0].cpu().tolist()
raise RuntimeError(
"Sampled token is outside captured positive sampling support "
f"for batch rows {invalid_rows}."
realized_lengths = support.sum(dim=-1, dtype=torch.int32)
overflow = realized_lengths > self.sampling_mask_max_tokens
statuses = torch.where(
invalid,
SamplingMaskStatus.INVALID,
torch.where(
overflow,
SamplingMaskStatus.OVERFLOW,
SamplingMaskStatus.OK,
),
).to(torch.int32)
# All replicas must make the same request-abort decision.
if dist.is_initialized():
if dist.get_world_size(self.tp_sync_group) > 1:
dist.all_reduce(
statuses, op=dist.ReduceOp.MAX, group=self.tp_sync_group
)
if (
self.cp_sync_group is not None
and dist.get_world_size(self.cp_sync_group) > 1
):
dist.all_reduce(
statuses, op=dist.ReduceOp.MAX, group=self.cp_sync_group
)
packed_size = min(self.sampling_mask_max_tokens, weights.shape[-1])
if token_ids is None:
_, packed_positions = torch.topk(
weights, k=packed_size, dim=-1, largest=True, sorted=True
)
packed_token_ids = packed_positions.to(torch.int32)
else:
# The PyTorch producer already sorts weights and IDs together.
packed_token_ids = token_ids[:, :packed_size].contiguous()
flat_rows, flat_cols = support.nonzero(as_tuple=True)
flat_ids = support_token_ids[flat_rows, flat_cols].to(torch.int32)
mask_lengths = support.sum(dim=-1, dtype=torch.int32)
flat_ids_cpu = flat_ids.cpu().tolist()
mask_lengths_cpu = mask_lengths.cpu().tolist()
selected_logprobs_cpu = selected_logprobs.cpu().tolist()
masks = [None] * len(return_sampling_masks)
logprobs = [None] * len(return_sampling_masks)
cursor = 0
for capture_row, batch_row in enumerate(requested_rows_list):
mask_len = int(mask_lengths_cpu[capture_row])
row_ids = flat_ids_cpu[cursor : cursor + mask_len]
cursor += mask_len
masks[batch_row] = row_ids
logprobs[batch_row] = float(selected_logprobs_cpu[capture_row])
logits_output.next_token_sampling_mask_idx = masks
logits_output.next_token_sampling_logprobs = logprobs
return SamplingMaskOutput(
token_ids=packed_token_ids,
lengths=realized_lengths.clamp(max=packed_size),
selected_logprobs=selected_logprobs,
statuses=statuses,
)
def _sample_from_logprobs(
self,
python/sglang/srt/managers/scheduler.py
+34 -31
View File
@@ -309,7 +309,6 @@ from sglang.srt.platforms import current_platform
from sglang.srt.plugins import load_plugins
from sglang.srt.rust_server.server import RustServer
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.sampling.sampling_params import TOP_K_ALL
from sglang.srt.server_args import PortArgs, ServerArgs, compute_world_size
from sglang.srt.session.session_controller import SessionController
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
@@ -1503,6 +1502,7 @@ class Scheduler(
buffer_size,
hidden_size=disagg_hidden_size,
hidden_states_dtype=disagg_hidden_states_dtype,
max_sampling_mask_tokens=self.server_args.sampling_mask_max_tokens,
custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
)
@@ -1549,6 +1549,7 @@ class Scheduler(
buffer_size,
hidden_size=disagg_hidden_size,
hidden_states_dtype=disagg_hidden_states_dtype,
max_sampling_mask_tokens=self.server_args.sampling_mask_max_tokens,
custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
)
@@ -2884,30 +2885,29 @@ class Scheduler(
self._add_request_to_queue(req)
return
if (
req.return_sampling_mask
and self.disaggregation_mode != DisaggregationMode.NULL
and not self.disagg_metadata_buffers.enable_sampling_mask
):
error_msg = (
"return_sampling_mask with disaggregation requires "
"SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS > 0."
)
req.set_finish_with_abort(error_msg)
self.init_req_max_new_tokens(req)
self._add_request_to_queue(req)
return
if req.return_sampling_mask and req.sampling_params.top_k == TOP_K_ALL:
error_msg = (
"return_sampling_mask requires finite top_k; top_p-only sampling "
"is valid but can return huge masks in the tail, blowing up "
"metadata, so we need a safety cap."
)
req.set_finish_with_abort(error_msg)
self.init_req_max_new_tokens(req)
self._add_request_to_queue(req)
return
if req.return_sampling_mask:
if (
self.disaggregation_mode != DisaggregationMode.NULL
and not self.disagg_metadata_buffers.enable_sampling_mask
):
self._reject_sampling_mask_request(
req,
"return_sampling_mask requires "
"SGLANG_ENABLE_DISAGG_SAMPLING_MASK=1 on both prefill and "
"decode servers when using disaggregated serving.",
)
return
top_k = req.sampling_params.top_k
sampling_mask_cap = self.server_args.sampling_mask_max_tokens
if top_k != 1 and not (1 < top_k <= sampling_mask_cap):
error_msg = (
"return_sampling_mask requires top_k=1 for greedy sampling "
f"or finite 1 < top_k <= {sampling_mask_cap}; got top_k="
f"{top_k}. Lower top_k or increase "
"--sampling-mask-max-tokens."
)
self._reject_sampling_mask_request(req, error_msg)
return
if req.return_sampling_mask and not self.spec_algorithm.is_none():
# Spec workers do not emit one sampling support per accepted token, so
@@ -2916,9 +2916,7 @@ class Scheduler(
error_msg = (
"return_sampling_mask is not supported with speculative decoding."
)
req.set_finish_with_abort(error_msg)
self.init_req_max_new_tokens(req)
self._add_request_to_queue(req)
self._reject_sampling_mask_request(req, error_msg)
return
if req.return_sampling_mask and get_exec().kernel.sampling_backend == "ascend":
@@ -2928,9 +2926,7 @@ class Scheduler(
"return_sampling_mask is not supported with the ascend "
"sampling backend."
)
req.set_finish_with_abort(error_msg)
self.init_req_max_new_tokens(req)
self._add_request_to_queue(req)
self._reject_sampling_mask_request(req, error_msg)
return
# Handle multimodal inputs
@@ -3170,6 +3166,13 @@ class Scheduler(
else:
raise ValueError(f"Invalid {self.disaggregation_mode=}")
def _reject_sampling_mask_request(self, req: Req, error_msg: str) -> None:
"""Return a sampling-mask validation error without running the model."""
logger.error(f"{error_msg}, {req.rid=}")
req.time_stats.trace_ctx.abort(abort_info={"reason": error_msg})
prepare_abort(req, error_msg, status_code=HTTPStatus.BAD_REQUEST)
self.output_streamer.stream_output([req], req.return_logprob)
def _set_or_validate_priority(self, req: Req) -> bool:
"""Set the default priority value, or abort the request based on the priority scheduling mode."""
if self.enable_priority_scheduling and req.priority is None:
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+144 -14
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import logging
from dataclasses import dataclass
from http import HTTPStatus
from typing import (
TYPE_CHECKING,
Callable,
@@ -15,7 +16,10 @@ import torch
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.environ import envs
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.logits_processor import (
LogitsProcessorOutput,
SamplingMaskStatus,
)
from sglang.srt.managers.schedule_batch import (
FINISH_ABORT,
FINISH_MATCHED_TOKEN,
@@ -287,6 +291,7 @@ class SchedulerBatchResultProcessor:
# Move next_token_ids and logprobs to cpu
next_token_ids = next_token_ids.tolist()
self.move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
self.materialize_sampling_mask_output(batch.reqs, logits_output)
self._validate_pp_skip_output_comm(batch, result)
@@ -299,6 +304,19 @@ class SchedulerBatchResultProcessor:
logprob_pt = 0
for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
should_commit_output = (
not req.finished()
and not req.is_retracted
and req.inflight_middle_chunks <= 0
)
sampling_mask_finish_reason = None
if should_commit_output and req.return_sampling_mask:
assert logits_output is not None
statuses = logits_output.next_token_sampling_mask_status
status = None if statuses is None else statuses[i]
sampling_mask_finish_reason = self.get_sampling_mask_finish_reason(
status=status
)
if (
batch.return_hidden_states
and logits_output.hidden_states is not None
@@ -311,9 +329,7 @@ class SchedulerBatchResultProcessor:
capture_hidden_mode=prefill_hidden_capture_mode,
extend_input_len=extend_input_len_per_req[i],
store=(
not req.finished()
and not req.is_retracted
and req.inflight_middle_chunks <= 0
should_commit_output and sampling_mask_finish_reason is None
),
)
@@ -328,7 +344,10 @@ class SchedulerBatchResultProcessor:
if req.inflight_middle_chunks <= 0:
req.time_stats.set_prefill_finished_time()
if req.beam_group is not None:
if sampling_mask_finish_reason is not None:
req.to_finish = sampling_mask_finish_reason
req.update_finish_state(0)
elif req.beam_group is not None:
# The relay point already replaced the sampled-token
# append; the group owns all finish semantics.
self.beam_coordinator.commit_prefill(
@@ -351,16 +370,22 @@ class SchedulerBatchResultProcessor:
):
req.kv.kv_committed_len += 1
if req.finished():
self._maybe_collect_routed_experts(req)
self._maybe_collect_indexer_topk(req)
release_kv_cache(req, self.tree_cache)
if sampling_mask_finish_reason is None:
self._maybe_collect_routed_experts(req)
self._maybe_collect_indexer_topk(req)
release_kv_cache(
req,
self.tree_cache,
is_insert=sampling_mask_finish_reason is None,
)
req.time_stats.set_completion_time()
elif not batch.decoding_reqs or req not in batch.decoding_reqs:
maybe_cache_unfinished_req(req, self.tree_cache)
if get_memory().enable_hisparse:
self.hisparse_coordinator.admit_request_into_staging(req)
self._maybe_collect_customized_info(i, req, logits_output)
if sampling_mask_finish_reason is None:
self._maybe_collect_customized_info(i, req, logits_output)
if batch.return_logprob:
logprob_pt = self._apply_prefill_logprobs(
@@ -371,8 +396,12 @@ class SchedulerBatchResultProcessor:
extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
next_token_ids=next_token_ids,
logprob_pt=logprob_pt,
store=sampling_mask_finish_reason is None,
)
if sampling_mask_finish_reason is not None:
continue
if req.return_sampling_mask:
self.add_sampling_mask_return_values(i, req, logits_output)
@@ -520,7 +549,9 @@ class SchedulerBatchResultProcessor:
extend_logprob_start_len_per_req: Optional[List[int]],
next_token_ids: List[int],
logprob_pt: int,
store: bool = True,
) -> int:
"""Advance the logprob cursor, optionally storing this request's values."""
assert extend_logprob_start_len_per_req is not None
assert extend_input_len_per_req is not None
extend_logprob_start_len = extend_logprob_start_len_per_req[i]
@@ -532,7 +563,7 @@ class SchedulerBatchResultProcessor:
extend_logprob_start_len,
)
if req.return_logprob:
if store and req.return_logprob:
self.logprob_result_processor.add_logprob_return_values(
i,
req,
@@ -907,6 +938,7 @@ class SchedulerBatchResultProcessor:
result.next_token_ids,
result.can_run_cuda_graph,
)
self.materialize_sampling_mask_output(batch.reqs, logits_output)
next_token_ids, next_token_logprobs = self._normalize_decode_outputs(
batch=batch,
@@ -967,13 +999,27 @@ class SchedulerBatchResultProcessor:
next_token_id = next_token_ids[i]
is_spec = not batch.spec_algorithm.is_none()
req.output_ids.extend(next_token_id)
new_accept_len = len(next_token_id)
self._maybe_update_reasoning_tokens(req, next_token_id)
req.time_stats.set_last_decode_finish_time()
sampling_mask_finish_reason = None
if req.return_sampling_mask:
statuses = logits_output.next_token_sampling_mask_status
status = None if statuses is None else statuses[i]
sampling_mask_finish_reason = self.get_sampling_mask_finish_reason(
status=status
)
if sampling_mask_finish_reason is not None:
req.to_finish = sampling_mask_finish_reason
new_accept_len = 0
else:
req.output_ids.extend(next_token_id)
new_accept_len = len(next_token_id)
self._maybe_update_reasoning_tokens(req, next_token_id)
req.update_finish_state(new_accept_len)
if sampling_mask_finish_reason is not None:
self._handle_sampling_mask_abort(req)
continue
self._handle_finish_state_updated_req(req, batch, result, i, logits_output)
if req.return_logprob:
@@ -1123,6 +1169,90 @@ class SchedulerBatchResultProcessor:
None if logprobs is None else logprobs[i]
)
@staticmethod
def materialize_sampling_mask_output(
reqs: List[Req],
output: Optional[LogitsProcessorOutput],
) -> None:
"""Convert opted-in tensor rows to batch-aligned Python results."""
if output is None or output.sampling_mask_output is None:
return
sampling_output = output.sampling_mask_output
batch_indices = [i for i, req in enumerate(reqs) if req.return_sampling_mask]
lengths = sampling_output.lengths.tolist()
selected_logprobs = sampling_output.selected_logprobs.tolist()
statuses = sampling_output.statuses.tolist()
assert len(batch_indices) == len(lengths)
batch_size = len(reqs)
masks = [None] * batch_size
logprobs = [None] * batch_size
status_by_batch = [None] * batch_size
token_ids = sampling_output.token_ids.cpu()
packed_width = token_ids.shape[1]
for row, batch_index in enumerate(batch_indices):
status = int(statuses[row])
length = int(lengths[row])
if status == SamplingMaskStatus.OK and not (0 <= length <= packed_width):
status = SamplingMaskStatus.INVALID
status_by_batch[batch_index] = status
if status == SamplingMaskStatus.OK:
masks[batch_index] = token_ids[row, :length].tolist()
logprobs[batch_index] = float(selected_logprobs[row])
output.next_token_sampling_mask_idx = masks
output.next_token_sampling_logprobs = logprobs
output.next_token_sampling_mask_status = status_by_batch
output.sampling_mask_output = None
def get_sampling_mask_finish_reason(
self,
*,
status: Optional[int],
) -> Optional[FINISH_ABORT]:
if status == SamplingMaskStatus.OK:
return None
if status == SamplingMaskStatus.OVERFLOW:
return FINISH_ABORT(
"Sampling support exceeds --sampling-mask-max-tokens="
f"{get_exec().features.sampling_mask_max_tokens}, commonly because "
"of cutoff ties. Lower top_k to leave more headroom or increase "
"the server limit.",
HTTPStatus.BAD_REQUEST,
"BadRequestError",
)
if status is None:
return FINISH_ABORT(
"The sampling backend did not return captured sampling support.",
HTTPStatus.INTERNAL_SERVER_ERROR,
"InternalServerError",
)
assert status == SamplingMaskStatus.INVALID
return FINISH_ABORT(
"The sampling backend selected a token outside its captured sampling "
"support.",
HTTPStatus.INTERNAL_SERVER_ERROR,
"InternalServerError",
)
def _handle_sampling_mask_abort(self, req: Req) -> None:
"""Release a request whose sampled token must not be committed."""
if req.multimodal_inputs is not None and req.session is None:
req.multimodal_inputs.release_features()
if get_disagg().disaggregation_decode_enable_offload_kvcache:
self.decode_offload_manager.finalize_release_on_finish(req)
else:
if get_memory().enable_hisparse:
self.hisparse_coordinator.request_finished(req)
prepare_release = getattr(
self.model_worker, "prepare_for_kv_cache_release", None
)
if callable(prepare_release):
prepare_release(req)
release_kv_cache(req, self.tree_cache, is_insert=False)
req.time_stats.set_completion_time()
def _handle_finish_state_updated_req(
self,
req: Req,
python/sglang/srt/managers/scheduler_pp_mixin.py
+22 -5
View File
@@ -780,11 +780,15 @@ class SchedulerPPMixin:
tensor_dict["draft_topk_index"] = draft_input.topk_index.contiguous()
tensor_dict["draft_hidden_states"] = draft_input.hidden_states.contiguous()
if batch.return_logprob:
logprob_dict = get_logprob_dict_from_result(result)
has_sampling_mask_output = (
result.logits_output is not None
and result.logits_output.sampling_mask_output is not None
)
if batch.return_logprob or has_sampling_mask_output:
logits_output_dict = get_logprob_dict_from_result(result)
tensor_dict = {
**tensor_dict,
**logprob_dict,
**logits_output_dict,
}
auxiliary_output = (
result.logits_output.auxiliary_device_output
@@ -904,7 +908,10 @@ class SchedulerPPMixin:
extend_input_len_per_req = None
extend_logprob_start_len_per_req = None
if batch.return_logprob:
if (
batch.return_logprob
or pp_outputs.tensors.get("sampling_mask_token_ids") is not None
):
(
logits_output,
extend_input_len_per_req,
@@ -1079,7 +1086,17 @@ class SchedulerPPMixin:
target, mb_metadata[next_mb_id], next_pp_outputs
)
d2h_event = self.device_module.Event()
d2h_event.record(self.device_module.current_stream())
if (
batch_result.logits_output is not None
and batch_result.logits_output.sampling_mask_output is not None
):
batch_result.copy_done = d2h_event
batch_result.copy_to_cpu(
return_logprob=target.return_logprob,
return_hidden_states=False,
)
else:
d2h_event.record(self.device_module.current_stream())
if send_first:
send_output_work = _do_send()
python/sglang/srt/managers/utils.py
+36 -5
View File
@@ -11,7 +11,10 @@ import torch
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.logits_processor import (
LogitsProcessorOutput,
SamplingMaskOutput,
)
from sglang.srt.managers import io_struct
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
@@ -174,7 +177,13 @@ class GenerationBatchResult:
# Sub-objects only declare their device fields; the single copy+safety
# primitive (_async_d2h: pinned D2H + record_stream) is injected here so
# all device->host copying and lifetime safety lives in one place.
sampling_mask_output = (
self.logits_output.sampling_mask_output
if self.logits_output is not None
else None
)
for holder in (
sampling_mask_output,
self.routed_experts_output,
self.indexer_topk_output,
self.expert_distribution_metrics,
@@ -246,9 +255,11 @@ def validate_input_length(
def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict:
"""Build the tensor payload needed to reconstruct PP output processing state."""
logits_output = result.logits_output
assert logits_output is not None
sampling_mask_output = logits_output.sampling_mask_output
return {
"extend_input_len_per_req": result.extend_input_len_per_req,
@@ -258,8 +269,20 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict:
"next_token_top_logprobs_idx": result.logits_output.next_token_top_logprobs_idx,
"next_token_token_ids_logprobs_val": result.logits_output.next_token_token_ids_logprobs_val,
"next_token_token_ids_logprobs_idx": result.logits_output.next_token_token_ids_logprobs_idx,
"next_token_sampling_mask_idx": result.logits_output.next_token_sampling_mask_idx,
"next_token_sampling_logprobs": result.logits_output.next_token_sampling_logprobs,
"sampling_mask_token_ids": (
None if sampling_mask_output is None else sampling_mask_output.token_ids
),
"sampling_mask_lengths": (
None if sampling_mask_output is None else sampling_mask_output.lengths
),
"sampling_mask_selected_logprobs": (
None
if sampling_mask_output is None
else sampling_mask_output.selected_logprobs
),
"sampling_mask_statuses": (
None if sampling_mask_output is None else sampling_mask_output.statuses
),
"input_token_logprobs": result.logits_output.input_token_logprobs,
"input_top_logprobs_val": result.logits_output.input_top_logprobs_val,
"input_top_logprobs_idx": result.logits_output.input_top_logprobs_idx,
@@ -271,6 +294,15 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict:
def get_logprob_from_pp_outputs(
next_pp_outputs: PPProxyTensors,
) -> tuple[LogitsProcessorOutput, list[int], list[int]]:
"""Reconstruct output processing state received from the last PP stage."""
sampling_mask_output = None
if next_pp_outputs["sampling_mask_token_ids"] is not None:
sampling_mask_output = SamplingMaskOutput(
token_ids=next_pp_outputs["sampling_mask_token_ids"],
lengths=next_pp_outputs["sampling_mask_lengths"],
selected_logprobs=next_pp_outputs["sampling_mask_selected_logprobs"],
statuses=next_pp_outputs["sampling_mask_statuses"],
)
logits_output = LogitsProcessorOutput(
# Do not send logits and hidden states because they are large
next_token_logits=None,
@@ -284,8 +316,7 @@ def get_logprob_from_pp_outputs(
next_token_token_ids_logprobs_idx=next_pp_outputs[
"next_token_token_ids_logprobs_idx"
],
next_token_sampling_mask_idx=next_pp_outputs["next_token_sampling_mask_idx"],
next_token_sampling_logprobs=next_pp_outputs["next_token_sampling_logprobs"],
sampling_mask_output=sampling_mask_output,
input_token_logprobs=next_pp_outputs["input_token_logprobs"],
input_top_logprobs_val=next_pp_outputs["input_top_logprobs_val"],
input_top_logprobs_idx=next_pp_outputs["input_top_logprobs_idx"],
python/sglang/srt/sampling/sampling_batch_info.py
+43 -1
View File
@@ -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)