Return top-p/top-k sampling mask/nucleas (#27408)

python/sglang/srt/layers/sampler.py

@@ -114,6 +114,7 @@ class Sampler(nn.Module):
 
         # Preprocess logits (custom processors and NaN handling)
         logits = self._preprocess_logits(logits, sampling_info)
+        return_sampling_mask = any(sampling_info.return_sampling_masks or [])
 
         if sampling_info.is_all_greedy:
             if _use_aiter and not _disable_aiter_greedy_sample:
@@ -123,6 +124,10 @@ class Sampler(nn.Module):
                 _aiter_greedy_sample(batch_next_token_ids, logits)
             else:
                 batch_next_token_ids = torch.argmax(logits, -1)
+            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
@@ -183,6 +188,16 @@ class Sampler(nn.Module):
                 batch_next_token_ids = self._sample_from_probs(
                     probs, sampling_info, positions, simple_sampling_case
                 )
+                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
@@ -260,6 +275,120 @@ class Sampler(nn.Module):
                 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
+
+    def _attach_greedy_sampling_mask_to_output(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        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
+
+    def _attach_sampling_mask_to_output(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        batch_next_token_ids: torch.Tensor,
+        sampling_mask_data: Tuple[
+            torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+        ],
+    ) -> 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)
+        )
+
+        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)
+
+        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 = []
+        cursor = 0
+        for i, should_return in enumerate(return_sampling_masks):
+            mask_len = int(mask_lengths_cpu[i])
+            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)
+
+        logits_output.next_token_sampling_mask_idx = masks
+        logits_output.next_token_sampling_logprobs = logprobs
+
     def _sample_from_logprobs(
         self,
         logprobs: torch.Tensor,