diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index f78f48b3e8..0cd44fe213 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/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. diff --git a/tests/v1/worker/test_gpu_sampler_flags.py b/tests/v1/worker/test_gpu_sampler_flags.py index 3f05a2cfb5..7982be1938 100644 --- a/tests/v1/worker/test_gpu_sampler_flags.py +++ b/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)) diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index eec43845db..3379e385a0 100644 --- a/tests/watermarking/test_watermarking.py +++ b/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 diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 7f4f9ffba9..541ae4e2f1 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/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. diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 085f092341..02d62298cd 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/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) diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index d5849a8d67..6945adc662 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/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 diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py index 776dece208..48547057fa 100644 --- a/vllm/v1/watermarking/gpu_sampler.py +++ b/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": diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 7ffc5bf493..4ad56292e4 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/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, diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index 47cf5de856..9e598b5a91 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/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))