[Perf][Sampler] Speed up the small-batch split top-p pipeline (#61996)

Signed-off-by: Bowen Wang <bowen@inferact.ai>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

tests/v1/sample/test_topk_topp_sampler.py

@@ -293,7 +293,7 @@ class TestTritonTopkTopp:
 
         self._compare_results(logits, k, p=None)
 
-    @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024])
+    @pytest.mark.parametrize("batch_size", [1, 8, 32, 48, 128, 512, 1024])
     @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256])
     def test_topp_only(self, batch_size: int, vocab_size: int):
         """Test top-p only (k=None)."""
@@ -307,7 +307,7 @@ class TestTritonTopkTopp:
 
         self._compare_results(logits, k=None, p=p)
 
-    @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024])
+    @pytest.mark.parametrize("batch_size", [1, 8, 32, 48, 128, 512, 1024])
     @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256])
     def test_topk_and_topp(self, batch_size: int, vocab_size: int):
         """Test combined top-k and top-p."""
@@ -328,6 +328,24 @@ class TestTritonTopkTopp:
 
         self._compare_results(logits, k, p)
 
+    @pytest.mark.parametrize("batch_size", [8, 48, 64])
+    @pytest.mark.parametrize("vocab_size", [151936, 200064, 262144])
+    def test_topk_and_topp_large_vocab(self, batch_size: int, vocab_size: int):
+        """Mixed batches at vocab sizes where the split top-p pipeline uses
+        its largest split counts (up to _SPLIT_MAX_SPLITS)."""
+        logits = torch.randn(
+            batch_size, vocab_size, generator=self.generator, dtype=torch.float32
+        )
+        k = torch.randint(1, 100, (batch_size,), generator=self.generator)
+        p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1
+        # Rows cycle through top-k + top-p, top-p only, top-k only, neither.
+        k[1::4] = vocab_size
+        p[2::4] = 1.0
+        k[3::4] = vocab_size
+        p[3::4] = 1.0
+
+        self._compare_results(logits, k, p)
+
     def test_both_disabled(self):
         """Test when both k and p are None (should be no-op)."""
         from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
@@ -788,6 +806,38 @@ class TestTritonTopkTopp:
             f"{kept_topk_topp.tolist()}"
         )
 
+    @pytest.mark.parametrize("batch_size", [8, 48])
+    def test_skip_split_without_topp_only_rows(
+        self, batch_size: int, monkeypatch: pytest.MonkeyPatch
+    ):
+        """has_topp_only_rows=False skips the split pipeline; the result must
+        not change when every top-p row also has a top-k, grammar rows
+        (k >= #finite logits) included."""
+        import vllm.v1.sample.ops.topk_topp_triton as tt
+
+        vocab_size = 32000
+        logits = torch.randn(
+            batch_size, vocab_size, generator=self.generator, dtype=torch.float32
+        )
+        logits[1] = float("-inf")
+        idx = torch.randperm(vocab_size, generator=self.generator)[:30]
+        logits[1, idx] = torch.randn(30, generator=self.generator)
+        k = torch.randint(
+            50, 100, (batch_size,), generator=self.generator, dtype=torch.int32
+        )
+        p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1
+        k[::4] = vocab_size
+        p[::4] = 1.0
+
+        expected = tt.apply_top_k_top_p_triton(logits.clone(), k, p)
+        monkeypatch.setattr(
+            tt, "_apply_topp_split", lambda *args: pytest.fail("split launched")
+        )
+        actual = tt.apply_top_k_top_p_triton(
+            logits.clone(), k, p, has_topp_only_rows=False
+        )
+        assert torch.equal(actual, expected)
+
     def test_topp_tie_break_deterministic(self):
         """Boundary-duplicate trimming must be deterministic across calls.
 

tests/v1/worker/test_gpu_sampler_flags.py

@@ -109,6 +109,30 @@ def test_custom_processor_add_request_gates_pipeline_flag(admitted: bool):
     assert sampler.needs_logits_processing[3] == admitted
 
 
+def test_has_topp_only_rows():
+    """A row is top-p-only when it has top-p and no effective top-k; top_k
+    of 0 (the default) or at least the logits width disables top-k."""
+    sampler = _make_sampler()
+    sampler.add_request(0, SamplingParams(top_k=10, top_p=0.9))
+    sampler.add_request(1, SamplingParams(top_k=10))
+    sampler.add_request(2, SamplingParams(top_p=0.9))
+    sampler.add_request(3, SamplingParams(top_k=VOCAB_SIZE + 1, top_p=0.9))
+    states = sampler.sampling_states
+    states.apply_staged_writes()
+
+    def has_topp_only_rows(idx: list[int], vocab_size: int = VOCAB_SIZE) -> bool:
+        idx_np = np.array(idx, dtype=np.int32)
+        idx_gpu = torch.from_numpy(idx_np).to(DEVICE)
+        return states.get_top_k_top_p(idx_gpu, idx_np, vocab_size)[2]
+
+    assert not has_topp_only_rows([0, 1])
+    assert has_topp_only_rows([0, 2])
+    assert has_topp_only_rows([3])
+    # The kernel compares top_k against the logits width, so a top_k at or
+    # above a narrower logits width counts as no top-k.
+    assert has_topp_only_rows([0], vocab_size=10)
+
+
 def test_logits_processing_cache_only_checks_active_requests():
     sampler = _make_sampler()
     sampler.add_request(0, SamplingParams(temperature=0.0))

tests/watermarking/test_watermarking.py

@@ -368,6 +368,7 @@ def test_gpu_sampler_respects_mixed_request_watermarking(
         None,
         None,
         False,
+        False,
     )
 
     assert torch.equal(sampled, torch.tensor([3, 7]))
@@ -383,7 +384,7 @@ def test_gpu_sampler_filters_top_k_top_p_before_watermarking(
     # reference path; the assertion is about the call order, not the kernel.
     monkeypatch.setattr(
         "vllm.v1.watermarking.gpu_sampler.apply_top_k_top_p",
-        lambda logits, k, p: apply_top_k_top_p_pytorch(
+        lambda logits, k, p, **kwargs: apply_top_k_top_p_pytorch(
             logits, k, p, allow_cpu_sync=True
         ),
     )
@@ -415,6 +416,7 @@ def test_gpu_sampler_filters_top_k_top_p_before_watermarking(
         torch.tensor([2]),
         torch.tensor([0.8]),
         False,
+        False,
     )
 
     assert watermarker.captured_logits is not None
@@ -482,6 +484,7 @@ def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch):
         None,
         None,
         False,
+        False,
     )
 
     assert torch.equal(sampled, torch.tensor([3, 7]))
@@ -517,6 +520,7 @@ def test_gpu_sampler_can_disable_context_deduplication(monkeypatch):
         None,
         None,
         False,
+        False,
     )
 
     assert torch.equal(sampled, torch.tensor([7, 7]))
@@ -956,6 +960,7 @@ def test_gpu_sampler_uses_fused_gumbel_for_repeated_contexts():
         None,
         None,
         False,
+        False,
     )
 
     assert torch.equal(repeated, torch.tensor([True, False], device=device))
@@ -992,6 +997,7 @@ def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch):
         None,
         None,
         False,
+        False,
     )
 
     assert actual is expected

vllm/model_executor/models/diffusion_gemma.py

@@ -1404,14 +1404,17 @@ class DiffusionSampler:
         # unaffected; only the canvas exploration is constrained. Applied
         # before canvas padding so phantom positions stay uniform.
         if num_decode > 0:
-            top_k, top_p = self.sampling_states.get_top_k_top_p(
+            top_k, top_p, has_topp_only_rows = self.sampling_states.get_top_k_top_p(
                 decode_slots.repeat_interleave(
                     valid_canvas_len, output_size=int(valid_canvas_len_np.sum())
                 ),
                 decode_slots_np,
+                logits.shape[-1],
             )
             if top_k is not None or top_p is not None:
-                logits = apply_top_k_top_p(logits.float(), top_k, top_p)
+                logits = apply_top_k_top_p(
+                    logits.float(), top_k, top_p, has_topp_only_rows=has_topp_only_rows
+                )
 
         # Where each decode request's rows start in the flat logits. Tiles
         # below gather and pad a request to the tile's width.

vllm/v1/sample/ops/topk_topp_sampler.py

@@ -484,13 +484,23 @@ class TopKTopPSampler(nn.Module):
 
 
 def apply_top_k_top_p(
-    logits: torch.Tensor, k: torch.Tensor | None, p: torch.Tensor | None
+    logits: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor | None,
+    has_topp_only_rows: bool = True,
 ) -> torch.Tensor:
+    """Apply top-k and/or top-p masking to logits.
+
+    has_topp_only_rows must be True if any row has p < 1 and
+    k >= logits.shape[-1] (no active top-k).
+    """
     if p is None and k is None:
         return logits
 
     if HAS_TRITON:
-        return apply_top_k_top_p_triton(logits, k, p)
+        return apply_top_k_top_p_triton(
+            logits, k, p, has_topp_only_rows=has_topp_only_rows
+        )
 
     is_cpu = current_platform.is_cpu()
     return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=is_cpu)

vllm/v1/sample/ops/topk_topp_triton.py

@@ -987,13 +987,19 @@ _SPLIT_MAX_BATCH = 64
 _SPLIT_MAX_SPLITS = 32  # power of two
 _SPLIT_FANOUT = 8
 _SPLIT_ROUNDS = 5
+# The step kernels are latency-bound at one program per SM, so keep splitting
+# past that while slices span at least this many entries; shorter slices no
+# longer amortize each program's row stats and combine replay.
+_SPLIT_MIN_SLICE = 7168
 
 _TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {}
 
 
-def _topp_split_count(batch_size: int, num_sm: int) -> int:
+def _topp_split_count(batch_size: int, vocab_size: int, num_sm: int) -> int:
     splits = 1
-    while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:
+    while splits < _SPLIT_MAX_SPLITS and (
+        splits * 2 * batch_size <= num_sm or splits * 2 * _SPLIT_MIN_SLICE <= vocab_size
+    ):
         splits *= 2
     return splits
 
@@ -1079,7 +1085,7 @@ def _topp_split_stats(
     num_sm: int,
 ) -> DispatchSpec:
     batch_size, vocab_size = logits.shape
-    splits = _topp_split_count(batch_size, num_sm)
+    splits = _topp_split_count(batch_size, vocab_size, num_sm)
     return (logits.shape[0] * splits,), dict(
         K=k if k is not None else logits,
         HAS_K=k is not None,
@@ -1373,7 +1379,7 @@ def _topp_split_step(
     num_sm: int,
 ) -> DispatchSpec:
     batch_size, vocab_size = logits.shape
-    splits = _topp_split_count(batch_size, num_sm)
+    splits = _topp_split_count(batch_size, vocab_size, num_sm)
     return (logits.shape[0] * splits,), dict(
         K=k if k is not None else logits,
         HAS_K=k is not None,
@@ -1555,7 +1561,7 @@ def _topp_split_mask(
     num_sm: int,
 ) -> DispatchSpec:
     batch_size, vocab_size = logits.shape
-    splits = _topp_split_count(batch_size, num_sm)
+    splits = _topp_split_count(batch_size, vocab_size, num_sm)
     return (logits.shape[0] * splits,), dict(
         K=k if k is not None else logits,
         HAS_K=k is not None,
@@ -1619,6 +1625,7 @@ def apply_top_k_top_p_triton(
     k: torch.Tensor | None,
     p: torch.Tensor | None,
     mask_value: float = float("-inf"),
+    has_topp_only_rows: bool = True,
 ) -> torch.Tensor:
     """Apply combined top-k and top-p masking using Triton.
 
@@ -1633,6 +1640,9 @@ def apply_top_k_top_p_triton(
         p: [batch_size] float32 tensor of top-p values per row (0 to 1),
             or None to disable top-p
         mask_value: Value for masked positions (default: -inf)
+        has_topp_only_rows: Whether any row has p < 1 and k >= vocab_size.
+            False lets small batches skip the split top-p pipeline; pass it
+            only if no such row exists, as such a row would get no top-p.
 
     Returns:
         The masked logits tensor. It may or may not be modified in-place.
@@ -1715,7 +1725,7 @@ def apply_top_k_top_p_triton(
         mask_value,
         num_sm,
     )
-    if use_split:
+    if use_split and has_topp_only_rows:
         _apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm)
 
     return logits

vllm/v1/watermarking/gpu_sampler.py

@@ -58,6 +58,7 @@ class GPUWatermarkSampler(Sampler):
         pos: torch.Tensor,
         top_k: torch.Tensor | None,
         top_p: torch.Tensor | None,
+        has_topp_only_rows: bool,
         use_fused_sampler: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         enabled = self.watermarking.np[idx_mapping_np] & (
@@ -71,10 +72,13 @@ class GPUWatermarkSampler(Sampler):
                 pos,
                 top_k,
                 top_p,
+                has_topp_only_rows,
                 use_fused_sampler,
             )
 
-        processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
+        processed_logits = apply_top_k_top_p(
+            processed_logits, top_k, top_p, has_topp_only_rows=has_topp_only_rows
+        )
         contexts = self._get_contexts(expanded_idx_mapping)
         repeated_contexts = None
         if self.deduplicate_contexts != "none":

vllm/v1/worker/gpu/sample/sampler.py

@@ -298,8 +298,8 @@ class Sampler:
             seq_lens_upper_bound_np,
             skip_top_k_top_p=True,
         )
-        top_k, top_p = self.sampling_states.get_top_k_top_p(
-            expanded_idx_mapping, idx_mapping_np
+        top_k, top_p, has_topp_only_rows = self.sampling_states.get_top_k_top_p(
+            expanded_idx_mapping, idx_mapping_np, logits.shape[-1]
         )
         # Don't use a fused sampler if no requests use top_k/top_p, if there are
         # any greedy requests or per-request seeds, or if post-processed
@@ -321,6 +321,7 @@ class Sampler:
             pos,
             top_k,
             top_p,
+            has_topp_only_rows,
             use_fused_sampler,
         )
 
@@ -332,6 +333,7 @@ class Sampler:
         pos: torch.Tensor,
         top_k: torch.Tensor | None,
         top_p: torch.Tensor | None,
+        has_topp_only_rows: bool,
         use_fused_sampler: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if use_fused_sampler:
@@ -342,7 +344,9 @@ class Sampler:
             else:  # Use XPU sampler
                 sampled, _ = xpu_sample(processed_logits, top_k, top_p)
         else:
-            processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
+            processed_logits = apply_top_k_top_p(
+                processed_logits, top_k, top_p, has_topp_only_rows=has_topp_only_rows
+            )
             sampled = gumbel_sample(
                 processed_logits,
                 expanded_idx_mapping,

vllm/v1/worker/gpu/sample/states.py

@@ -103,13 +103,22 @@ class SamplingStates:
         apply_min_p(logits, expanded_idx_mapping, self.min_p.gpu)
 
     def get_top_k_top_p(
-        self, expanded_idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-        do_top_k = np.any(self.top_k.np[idx_mapping_np] != self.vocab_size)
-        do_top_p = np.any(self.top_p.np[idx_mapping_np] != 1.0)
+        self,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        vocab_size: int,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, bool]:
+        """Returns (top_k, top_p, has_topp_only_rows)."""
+        top_k_np = self.top_k.np[idx_mapping_np]
+        top_p_np = self.top_p.np[idx_mapping_np]
+        do_top_k = np.any(top_k_np != self.vocab_size)
+        do_top_p = np.any(top_p_np != 1.0)
         top_k = self.top_k.gpu[expanded_idx_mapping] if do_top_k else None
         top_p = self.top_p.gpu[expanded_idx_mapping] if do_top_p else None
-        return top_k, top_p
+        has_topp_only_rows = bool(
+            do_top_p and np.any((top_k_np >= vocab_size) & (top_p_np < 1.0))
+        )
+        return top_k, top_p, has_topp_only_rows
 
     def apply_top_k_top_p(
         self,
@@ -117,10 +126,14 @@ class SamplingStates:
         expanded_idx_mapping: torch.Tensor,
         idx_mapping_np: np.ndarray,
     ) -> torch.Tensor:
-        top_k, top_p = self.get_top_k_top_p(expanded_idx_mapping, idx_mapping_np)
+        top_k, top_p, has_topp_only_rows = self.get_top_k_top_p(
+            expanded_idx_mapping, idx_mapping_np, logits.shape[-1]
+        )
         if top_k is None and top_p is None:
             return logits
-        return apply_top_k_top_p(logits, top_k, top_p)
+        return apply_top_k_top_p(
+            logits, top_k, top_p, has_topp_only_rows=has_topp_only_rows
+        )
 
     def any_greedy(self, idx_mapping_np: np.ndarray) -> bool:
         return bool(np.any(self.temperature.np[idx_mapping_np] == 0.0))