diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index d9b290c20..8f10c1abf 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -1,6 +1,6 @@ import logging from functools import partial -from typing import Callable, Dict, List, Optional, Tuple +from typing import Callable, Dict, List, NamedTuple, Optional, Tuple import torch import torch.distributed as dist @@ -84,6 +84,15 @@ def _trace_e2e_sampler(stage: str, **fields) -> None: print(f"SGLANG_TRACE_SAMPLER_E2E {rank} stage={stage} {details}", flush=True) +class _SamplingMaskCapture(NamedTuple): + """Compact post-filter weights and their original batch-row mapping.""" + + weights: torch.Tensor + token_ids: Optional[torch.Tensor] + selected_weight: Optional[torch.Tensor] + batch_rows: torch.Tensor + + class Sampler(nn.Module): def __init__(self): super().__init__() @@ -149,6 +158,7 @@ class Sampler(nn.Module): logits = self._preprocess_logits(logits, sampling_info) _trace_e2e_sampler("preprocess_returned") return_sampling_mask = any(sampling_info.return_sampling_masks or []) + sampling_mask_capture = None if sampling_info.is_all_greedy: _trace_e2e_sampler("greedy_enter") @@ -162,10 +172,6 @@ class Sampler(nn.Module): _trace_e2e_sampler( "greedy_returned", output_shape=tuple(batch_next_token_ids.shape) ) - if return_sampling_mask: - self._attach_greedy_sampling_mask_to_output( - logits_output, sampling_info, batch_next_token_ids - ) if return_logprob: original_logprobs = logprobs = torch.nn.functional.log_softmax( logits, dim=-1 @@ -238,19 +244,13 @@ class Sampler(nn.Module): logits[:] = torch.softmax(logits, dim=-1) probs = logits - batch_next_token_ids = self._sample_from_probs( - probs, sampling_info, positions, simple_sampling_case + batch_next_token_ids, sampling_mask_capture = self._sample_from_probs( + probs, + sampling_info, + positions, + simple_sampling_case, + return_sampling_mask=return_sampling_mask, ) - if return_sampling_mask: - sampling_mask_data = self._compute_sampling_mask_from_probs( - probs, sampling_info - ) - self._attach_sampling_mask_to_output( - logits_output, - sampling_info, - batch_next_token_ids, - sampling_mask_data, - ) if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB: logprobs = ( logprobs_via_logsoftmax_kernel @@ -274,6 +274,25 @@ class Sampler(nn.Module): self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info) _trace_e2e_sampler("token_sync_returned") + if return_sampling_mask: + if sampling_info.is_all_greedy: + self._attach_greedy_sampling_mask_to_output( + logits_output, sampling_info, batch_next_token_ids + ) + else: + assert sampling_mask_capture is not None + if SYNC_TOKEN_IDS_ACROSS_TP or sampling_info.grammars: + # Token synchronization can replace the producer-selected token. + sampling_mask_capture = sampling_mask_capture._replace( + selected_weight=None + ) + self._attach_sampling_mask_to_output( + logits_output, + sampling_info, + batch_next_token_ids, + sampling_mask_capture, + ) + _trace_e2e_sampler("forward_returned") return batch_next_token_ids @@ -283,18 +302,57 @@ class Sampler(nn.Module): sampling_info: SamplingBatchInfo, positions: torch.Tensor, simple_sampling_case: bool, - ) -> torch.Tensor: + *, + return_sampling_mask: bool = False, + ) -> Tuple[torch.Tensor, Optional[_SamplingMaskCapture]]: """Sample from probability distribution (after softmax). Used for standard sampling with flashinfer/pytorch backends. 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. """ + 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 + ) + 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( probs, sampling_seed=sampling_info.sampling_seed, positions=positions, ) + if return_sampling_mask: + capture_probs = select_capture_rows(probs) + capture_tokens = select_capture_rows(batch_next_token_ids) + selected_weight = torch.gather( + capture_probs, 1, capture_tokens.long().view(-1, 1) + ).squeeze(1) + sampling_mask_capture = _SamplingMaskCapture( + weights=capture_probs, + token_ids=None, + selected_weight=selected_weight, + batch_rows=capture_rows, + ) else: backend = get_exec().kernel.sampling_backend if backend == "flashinfer": @@ -307,6 +365,27 @@ class Sampler(nn.Module): batch_next_token_ids = min_p_sampling_from_probs( probs, sampling_info.min_ps ) + if return_sampling_mask: + capture_probs = select_capture_rows(probs) + capture_min_ps = select_capture_rows(sampling_info.min_ps) + capture_tokens = select_capture_rows(batch_next_token_ids) + min_p_thresholds = ( + capture_probs.max(dim=-1).values * capture_min_ps + ) + filtered_probs = capture_probs.masked_fill( + capture_probs < min_p_thresholds.view(-1, 1), 0 + ) + selected_weight = torch.gather( + filtered_probs, + 1, + capture_tokens.long().view(-1, 1), + ).squeeze(1) + sampling_mask_capture = _SamplingMaskCapture( + weights=filtered_probs, + token_ids=None, + selected_weight=selected_weight, + batch_rows=capture_rows, + ) else: batch_next_token_ids = top_k_top_p_sampling_from_probs( probs.contiguous(), @@ -314,9 +393,42 @@ class Sampler(nn.Module): sampling_info.top_ps, filter_apply_order="joint", ) + if return_sampling_mask: + # Correctness invariant: the fused joint sampler and these + # separate renormalization primitives must share cutoff, + # tie, and joint-support semantics so captured positive + # support exactly describes the sampler's action space. + capture_probs = select_capture_rows(probs) + capture_top_ks = select_capture_rows(sampling_info.top_ks) + capture_top_ps = select_capture_rows(sampling_info.top_ps) + capture_tokens = select_capture_rows(batch_next_token_ids) + filtered_probs = capture_probs + if sampling_info.need_top_k_sampling: + filtered_probs = top_k_renorm_prob( + capture_probs, capture_top_ks + ) + if sampling_info.need_top_p_sampling: + top_p_probs = top_p_renorm_prob( + capture_probs, capture_top_ps + ) + if filtered_probs is capture_probs: + filtered_probs = top_p_probs + else: + filtered_probs.masked_fill_(top_p_probs <= 0, 0) + selected_weight = torch.gather( + filtered_probs, + 1, + capture_tokens.long().view(-1, 1), + ).squeeze(1) + sampling_mask_capture = _SamplingMaskCapture( + weights=filtered_probs, + token_ids=None, + selected_weight=selected_weight, + batch_rows=capture_rows, + ) elif backend == "pytorch": # A slower fallback implementation with torch native operations. - batch_next_token_ids = top_k_top_p_min_p_sampling_from_probs_torch( + sample_result = top_k_top_p_min_p_sampling_from_probs_torch( probs, sampling_info.top_ks, sampling_info.top_ps, @@ -324,39 +436,26 @@ class Sampler(nn.Module): sampling_info.need_min_p_sampling, sampling_info.sampling_seed, positions, + return_filtered_probs=return_sampling_mask, ) + if return_sampling_mask: + ( + batch_next_token_ids, + filtered_probs, + token_ids, + selected_weight, + ) = sample_result + sampling_mask_capture = _SamplingMaskCapture( + weights=select_capture_rows(filtered_probs), + token_ids=select_capture_rows(token_ids), + selected_weight=select_capture_rows(selected_weight), + batch_rows=capture_rows, + ) + else: + batch_next_token_ids = sample_result else: raise ValueError(f"Invalid sampling backend: {backend}") - return batch_next_token_ids - - def _compute_sampling_mask_from_probs( - self, probs: torch.Tensor, sampling_info: SamplingBatchInfo - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Return sorted token ids, sorted probs, keep mask, and raw probs.""" - vocab_size = probs.shape[-1] - max_top_k = sampling_info.sampling_mask_max_top_k - if 0 < max_top_k < vocab_size: - probs_sort, probs_idx = torch.topk( - probs, - k=max_top_k, - dim=-1, - largest=True, - sorted=True, - ) - positions = torch.arange(max_top_k, device=probs.device).view(1, -1) - else: - probs_sort, probs_idx = probs.sort(dim=-1, descending=True) - positions = torch.arange(vocab_size, device=probs.device).view(1, -1) - probs_sum = torch.cumsum(probs_sort, dim=-1) - - keep_mask = positions < sampling_info.top_ks.view(-1, 1) - keep_mask &= (probs_sum - probs_sort) <= sampling_info.top_ps.view(-1, 1) - - if sampling_info.need_min_p_sampling: - min_p_thresholds = probs_sort[:, 0] * sampling_info.min_ps - keep_mask &= probs_sort >= min_p_thresholds.view(-1, 1) - - return probs_idx, probs_sort, keep_mask, probs + return batch_next_token_ids, sampling_mask_capture def _attach_greedy_sampling_mask_to_output( self, @@ -382,63 +481,84 @@ class Sampler(nn.Module): logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo, batch_next_token_ids: torch.Tensor, - sampling_mask_data: Tuple[ - torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor - ], + sampling_mask_capture: _SamplingMaskCapture, ) -> None: - probs_idx, probs_sort, keep_mask, probs = sampling_mask_data 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 - sampled_tokens = batch_next_token_ids.view(-1, 1) - sampled_matches_all = probs_idx == sampled_tokens - sampled_in_idx = sampled_matches_all.any(dim=-1) - - # The sampler is the source of truth for the rollout action space. If a - # backend/numeric edge chooses a token just outside the reconstructed - # prefix, include that sampled token so training can replay a support - # that contained the rollout action. - effective_keep_mask = keep_mask | sampled_matches_all - selected_raw_probs = torch.gather(probs, 1, sampled_tokens).squeeze(1) - support_mass = torch.where( - effective_keep_mask, probs_sort, torch.zeros_like(probs_sort) - ).sum(dim=-1) - support_mass = support_mass + torch.where( - sampled_in_idx, torch.zeros_like(selected_raw_probs), selected_raw_probs - ) - selected_logprobs = torch.log( - selected_raw_probs.float() - / support_mass.float().clamp_min(torch.finfo(torch.float32).tiny) + requested_rows_list = [ + i for i, should_return in enumerate(return_sampling_masks) if should_return + ] + requested_rows = 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 ) + 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_in_capture = sampled_matches.any(dim=-1) + selected_positions = sampled_matches.to(torch.int32).argmax( + dim=-1, keepdim=True + ) + selected_from_weights = torch.gather( + weights, 1, selected_positions + ).squeeze(1) + support_token_ids = token_ids - flat_rows, flat_cols = effective_keep_mask.nonzero(as_tuple=True) - flat_ids = probs_idx[flat_rows, flat_cols].to(torch.int32) - mask_lengths = effective_keep_mask.sum(dim=-1, dtype=torch.int32) + if selected_weight is None: + selected_weight = selected_from_weights + + support = weights > 0 + support_mass = weights.sum(dim=-1, dtype=torch.float32) + selected_weight = selected_weight.float() + selected_logprobs = torch.log(selected_weight / support_mass) + valid = ( + 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}." + ) + + 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() - sampled_in_idx_cpu = sampled_in_idx.cpu().tolist() - sampled_tokens_cpu = batch_next_token_ids.to(torch.int32).cpu().tolist() selected_logprobs_cpu = selected_logprobs.cpu().tolist() - masks = [] - logprobs = [] + masks = [None] * len(return_sampling_masks) + logprobs = [None] * len(return_sampling_masks) cursor = 0 - for i, should_return in enumerate(return_sampling_masks): - mask_len = int(mask_lengths_cpu[i]) + 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 - if not sampled_in_idx_cpu[i]: - row_ids.append(int(sampled_tokens_cpu[i])) - if should_return: - masks.append(row_ids) - logprobs.append(float(selected_logprobs_cpu[i])) - else: - masks.append(None) - logprobs.append(None) + 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 @@ -602,11 +722,17 @@ def top_k_top_p_min_p_sampling_from_probs_torch( need_min_p_sampling: bool, sampling_seed: Optional[torch.Tensor], positions: torch.Tensor, + *, + return_filtered_probs: bool = False, ): """ A top-k, top-p and min-p sampling implementation with native pytorch operations. When sampling_seed is not None, deterministic inference will be enabled, it will sample with the sampling_seed of each request. + + By default, returns only sampled token IDs. With return_filtered_probs=True, + also returns the actual filtered weights, their token-ID permutation, and + the selected weights. """ probs_sort, probs_idx = probs.sort(dim=-1, descending=True) probs_sum = torch.cumsum(probs_sort, dim=-1) @@ -631,14 +757,20 @@ def top_k_top_p_min_p_sampling_from_probs_torch( # apply log to get logprobs. Therefore, we cannot use log_softmax directly. # For now, we use log to the modified probs to get logprobs, but for numerical # stability, we'd better come up with a solution to use log_softmax. - logprobs = probs_sort.to(torch.float64) # Using float64 for numerical stability - del probs_sort + logprobs = probs_sort.to(torch.float64, copy=return_filtered_probs) + if not return_filtered_probs: + del probs_sort logprobs.log_() sampled_index = multinomial_with_seed(logprobs, sampling_seed, positions) + if return_filtered_probs: + selected_weight = torch.gather(probs_sort, 1, sampled_index).view(-1) + # int32 range is enough to represent the token ids probs_idx = probs_idx.to(torch.int32) batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(-1) + if return_filtered_probs: + return batch_next_token_ids, probs_sort, probs_idx, selected_weight return batch_next_token_ids diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 9bf2c9683..d25305203 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -76,7 +76,6 @@ class SamplingBatchInfo: # Per-request flag for returning sparse sampling support metadata. return_sampling_masks: Optional[List[bool]] = None - sampling_mask_max_top_k: int = 0 # Device device: str = "cuda" @@ -146,10 +145,6 @@ 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_max_top_k = max( - (r.sampling_params.top_k for r in reqs if r.return_sampling_mask), - default=0, - ) if has_custom_logit_processor: # Merge the same type of custom logit processors together @@ -215,7 +210,6 @@ class SamplingBatchInfo: device=device, logit_bias=logit_bias, return_sampling_masks=return_sampling_masks, - sampling_mask_max_top_k=sampling_mask_max_top_k, ) ret.adjusted_from_schedule_batch(batch, vocab_size) return ret @@ -445,9 +439,6 @@ class SamplingBatchInfo: self.return_sampling_masks = ( self.return_sampling_masks or [False] * self_len ) + (other.return_sampling_masks or [False] * other_len) - self.sampling_mask_max_top_k = max( - self.sampling_mask_max_top_k, other.sampling_mask_max_top_k - ) # Note: because the __len()__ operator is defined on the temperatures tensor, # please make sure any merge operation with len(self) or len(other) is done before diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py index b86c0d35a..7d3f4e5dc 100644 --- a/test/registered/sampling/test_sampling_mask.py +++ b/test/registered/sampling/test_sampling_mask.py @@ -1,9 +1,15 @@ import math import unittest +from types import SimpleNamespace +from unittest.mock import patch import requests +import torch -from sglang.srt.utils import kill_process_tree +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.utils import is_hip, kill_process_tree from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import ( DEFAULT_SMALL_MODEL_NAME_FOR_TEST, @@ -19,6 +25,7 @@ register_amd_ci(est_time=320, suite="stage-b-test-1-gpu-small-amd") _MAX_NEW_TOKENS = 4 _TOP_P = 0.99 _TOP_K = 10 +_TOP_LOGPROBS_NUM = 128 _SAMPLING_SEED = 1234 _SERVER_ARGS = ( "--mem-fraction-static", @@ -29,6 +36,180 @@ _INVALID_SAMPLING_MASK_ERROR = ( ) +class TestSamplingMaskCapture(CustomTestCase): + def setUp(self): + self.sampler = Sampler.__new__(Sampler) + torch.nn.Module.__init__(self.sampler) + + @unittest.skipIf(is_hip(), "FlashInfer is not available on ROCm") + def test_flashinfer_joint_cutoff_ties_match_capture(self): + batch_size = 256 + top_k = 2 + top_p = 0.45 + base_probs = torch.tensor([[0.4, 0.2, 0.2, 0.1, 0.1]], device="cuda") + probs = base_probs.repeat(batch_size, 1) + + # Derive the threshold-based joint support independently. Both filters + # cut at 0.2, so the tied entries must survive even though this yields + # more support entries than top_k. + sorted_probs = base_probs[0].sort(descending=True).values + top_k_cutoff = sorted_probs[top_k - 1] + mass_before = sorted_probs.cumsum(dim=-1) - sorted_probs + top_p_cutoff = sorted_probs[mass_before <= top_p][-1] + expected_support = (base_probs[0] >= top_k_cutoff) & ( + base_probs[0] >= top_p_cutoff + ) + 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, + ) + + self.assertIsNotNone(capture) + self.assertEqual(capture.batch_rows.cpu().tolist(), list(range(batch_size))) + actual_support = capture.weights > 0 + self.assertTrue( + torch.equal(actual_support, expected_support.expand_as(actual_support)) + ) + self.assertGreater(int(actual_support[0].sum().item()), top_k) + self.assertTrue( + bool(actual_support.gather(1, sampled.view(-1, 1)).all().item()) + ) + + @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, + ) as top_k_mock, + patch( + "sglang.srt.layers.sampler.top_p_renorm_prob", + wraps=top_p_renorm, + ) 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, + ) + + 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(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 + ) + 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]) + self.assertEqual(set(output.next_token_sampling_mask_idx[3]), {0, 1, 2}) + self.assertIsNone(output.next_token_sampling_logprobs[0]) + self.assertIsNotNone(output.next_token_sampling_logprobs[1]) + self.assertIsNone(output.next_token_sampling_logprobs[2]) + self.assertIsNotNone(output.next_token_sampling_logprobs[3]) + + def test_pytorch_capture_compacts_requested_rows(self): + batch_size = 4 + 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,), 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, + ) + + 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 + ) + for batch_row in requested_rows: + self.assertIn( + int(sampled[batch_row]), + output.next_token_sampling_mask_idx[batch_row], + ) + self.assertIsNotNone(output.next_token_sampling_logprobs[batch_row]) + self.assertIsNone(output.next_token_sampling_mask_idx[0]) + self.assertIsNone(output.next_token_sampling_mask_idx[2]) + + class SamplingMaskTestMixin: @classmethod def _launch_server(cls, other_args=()): @@ -79,6 +260,7 @@ class SamplingMaskTestMixin: 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): @@ -88,6 +270,8 @@ class SamplingMaskTestMixin: class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase): + _sampling_backend = "flashinfer" + @classmethod def setUpClass(cls): cls._launch_server() @@ -102,12 +286,8 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase): "ignore_eos": True, } ) - # The mask keeps at most top_k tokens, plus possibly the actually - # sampled token when the sampling kernel picks one just outside the - # mask's topk reconstruction (fp cumsum divergence); see - # Sampler._attach_sampling_mask_to_output. for sampling_mask in top_p_sampling_masks: - self.assertLessEqual(len(sampling_mask), _TOP_K + 1) + self.assertGreater(len(sampling_mask), 0) top_k_sampling_masks = self._generate_sampling_masks( { @@ -118,7 +298,7 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase): } ) for sampling_mask in top_k_sampling_masks: - self.assertIn(len(sampling_mask), (_TOP_K, _TOP_K + 1)) + self.assertGreaterEqual(len(sampling_mask), _TOP_K) top_k_top_p_one_sampling_masks = self._generate_sampling_masks( { @@ -130,18 +310,19 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase): } ) for sampling_mask in top_k_top_p_one_sampling_masks: - self.assertIn(len(sampling_mask), (_TOP_K, _TOP_K + 1)) + self.assertGreaterEqual(len(sampling_mask), _TOP_K) def test_sampling_mask_matches_topk_logprobs(self): """Check the returned mask and its renormalized logprobs. - We get the per-token full-vocab logprobs via ``return_logprob`` with - ``top_logprobs_num == top_k``, which covers every token the mask can - contain. With ``temperature=1.0`` these are the sampler's distribution, - so ``p = exp(logprob)`` are the exact probabilities. For each token, we check: + We get a wide prefix of full-vocab logprobs via ``return_logprob`` so + cutoff ties that extend beyond ``top_k`` are visible. With + ``temperature=1.0`` these are the sampler's distribution, so + ``p = exp(logprob)`` are the exact probabilities. For each token, we check: - 1. the returned mask matches the nucleus reconstructed from those probs, - 2. sampling_logprob == log(p[sampled] / sum(p[t] for t in mask)). + 1. the sampled token is in the returned top-k-bounded mask, + 2. every mask token is present in the returned top logprobs, + 3. sampling_logprob == log(p[sampled] / sum(p[t] for t in mask)). """ top_k, top_p = _TOP_K, _TOP_P response = self._post_generate( @@ -153,7 +334,7 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase): "ignore_eos": True, }, return_logprob=True, - top_logprobs_num=top_k, + top_logprobs_num=_TOP_LOGPROBS_NUM, ) self.assertEqual(response.status_code, 200, response.text) @@ -175,19 +356,13 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase): int(tid): math.exp(logprob) for logprob, tid, _ in step_top_logprobs } - reconstructed = [] - mass_before = 0.0 - for logprob, tid, _ in step_top_logprobs: - if mass_before <= top_p: - reconstructed.append(int(tid)) - mass_before += math.exp(logprob) - if output_id not in reconstructed: - reconstructed.append(output_id) - # ``<= 1``: fp32 (server) and fp64 (here) cumsums may split on the - # single token straddling the top_p cut. - self.assertLessEqual(len(set(mask) ^ set(reconstructed)), 1) + mask_set = set(mask) - support_mass = sum(probs[tid] for tid in mask) + self.assertIn(output_id, mask_set) + self.assertLessEqual(len(mask_set), top_k) + self.assertTrue(mask_set.issubset(probs)) + + support_mass = sum(probs[token_id] for token_id in mask_set) expected_logprob = math.log(probs[output_id] / support_mass) self.assertAlmostEqual(mask_logprob, expected_logprob, delta=1e-2) @@ -280,5 +455,13 @@ class TestSamplingMaskDeterministic(SamplingMaskTestMixin, CustomTestCase): self.assertEqual(with_mask_output["text"], without_mask_output["text"]) +class TestSamplingMaskPytorch(TestSamplingMask): + _sampling_backend = "pytorch" + + @classmethod + def setUpClass(cls): + cls._launch_server(("--sampling-backend", "pytorch")) + + if __name__ == "__main__": unittest.main()