diff --git a/tests/v1/worker/test_gpu_sampler_flags.py b/tests/v1/worker/test_gpu_sampler_flags.py new file mode 100644 index 0000000000..3fd1cb3b3c --- /dev/null +++ b/tests/v1/worker/test_gpu_sampler_flags.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("CUDA required for sampler flag tests", allow_module_level=True) + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.states import RequestState + +DEVICE = torch.device("cuda") +VOCAB_SIZE = 128 + + +class MockReasoningConfig: + reasoning_start_token_ids = [90] + reasoning_end_token_ids = [91] + natural_reasoning_end_token_ids = [91] + + +def _make_sampler() -> Sampler: + req_states = RequestState( + max_num_reqs=4, + max_model_len=64, + max_num_batched_tokens=16, + num_speculative_steps=1, + vocab_size=VOCAB_SIZE, + device=DEVICE, + ) + return Sampler( + max_num_reqs=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + req_states=req_states, + reasoning_config=MockReasoningConfig(), + ) + + +@pytest.mark.parametrize( + ("sampling_params", "expected"), + [ + pytest.param(SamplingParams(), False, id="defaults"), + pytest.param(SamplingParams(temperature=0.0), False, id="greedy"), + pytest.param( + SamplingParams(thinking_token_budget=3), True, id="thinking-budget" + ), + pytest.param(SamplingParams(logit_bias={1: 1.0}), True, id="logit-bias"), + pytest.param(SamplingParams(frequency_penalty=0.1), True, id="penalty"), + pytest.param(SamplingParams(_bad_words_token_ids=[[1]]), True, id="bad-words"), + pytest.param(SamplingParams(temperature=0.7), True, id="temperature"), + pytest.param(SamplingParams(min_p=0.1), True, id="min-p"), + pytest.param(SamplingParams(top_k=10), True, id="top-k"), + pytest.param(SamplingParams(top_p=0.9), True, id="top-p"), + pytest.param( + SamplingParams.for_sampler_warmup(), True, id="all-logits-processors" + ), + ], +) +def test_logits_processing_cache_matches_request_features( + sampling_params: SamplingParams, expected: bool +): + sampler = _make_sampler() + sampler.add_request(3, prompt_len=1, sampling_params=sampling_params) + + assert sampler.needs_logits_processing[3] == expected + + +def test_logits_processing_cache_is_overwritten_when_slot_is_reused(): + sampler = _make_sampler() + sampler.add_request(3, 1, SamplingParams.for_sampler_warmup()) + sampler.add_request(3, 1, SamplingParams()) + + assert not sampler.needs_logits_processing[3] + + +def test_logits_processing_cache_only_checks_active_requests(): + sampler = _make_sampler() + sampler.add_request(0, 1, SamplingParams(temperature=0.0)) + sampler.add_request(2, 1, SamplingParams.for_sampler_warmup()) + + sampling_only = np.array([0], dtype=np.int32) + with_processing = np.array([0, 2], dtype=np.int32) + + assert not np.any(sampler.needs_logits_processing[sampling_only]) + assert np.any(sampler.needs_logits_processing[with_processing]) diff --git a/tests/v1/worker/test_gpu_thinking_budget.py b/tests/v1/worker/test_gpu_thinking_budget.py index 0b0f19f7db..2bc697b084 100644 --- a/tests/v1/worker/test_gpu_thinking_budget.py +++ b/tests/v1/worker/test_gpu_thinking_budget.py @@ -12,6 +12,7 @@ if not torch.cuda.is_available(): ) from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState from vllm.v1.worker.gpu.states import RequestState @@ -185,6 +186,44 @@ def test_v2_thinking_budget_ignores_plain_request(): assert torch.all(out == 0) +def test_v2_greedy_sampling_applies_thinking_budget(): + """Greedy-only requests must not bypass thinking-budget processing.""" + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + sampler = Sampler( + max_num_reqs=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + req_states=req_states, + reasoning_config=MockReasoningConfig(), + ) + sampler.add_request( + req_idx=3, + prompt_len=1, + sampling_params=SamplingParams( + temperature=0.0, + thinking_token_budget=3, + ), + ) + sampler.apply_staged_writes() + + idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE) + idx_mapping_np = idx_mapping.cpu().numpy() + expanded_idx_mapping = idx_mapping.clone() + input_ids = torch.tensor([12], dtype=torch.int32, device=DEVICE) + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = sampler.apply_sampling_params( + logits, + expanded_idx_mapping, + idx_mapping, + idx_mapping_np, + torch.tensor([4], dtype=torch.int32, device=DEVICE), + input_ids, + torch.tensor([0], dtype=torch.int32, device=DEVICE), + ) + + assert out[0, END].item() == pytest.approx(1.0e9) + + def test_v2_thinking_budget_latest_prefill_end_disables_forcing(): req_states = _make_req_states( [1, START, 10, 11, 12, END, 13], diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index bb7dc40b02..26371f014e 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -53,6 +53,7 @@ class Sampler: self.bad_words_state = BadWordsState(req_states) self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device) self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config) + self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool) self.num_speculative_tokens = num_speculative_tokens self.return_sampling_mask = return_sampling_mask self.use_flashinfer = ( @@ -69,6 +70,22 @@ class Sampler: self.logprob_token_ids_state.add_request(req_idx, sampling_params) self.thinking_budget_state.add_request(req_idx, sampling_params) + states = self.sampling_states + temperature = states.temperature.np[req_idx] + self.needs_logits_processing[req_idx] = ( + self.logit_bias_state.use_logit_bias[req_idx] + or self.penalties_state.use_penalty[req_idx] + or self.bad_words_state.num_bad_words.np[req_idx] > 0 + or ( + self.thinking_budget_state.enabled + and self.thinking_budget_state.use_thinking_budget[req_idx] + ) + or (temperature != 0.0 and temperature != 1.0) + or states.min_p.np[req_idx] != 0.0 + or states.top_k.np[req_idx] != states.vocab_size + or states.top_p.np[req_idx] != 1.0 + ) + def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() self.penalties_state.apply_staged_writes() @@ -172,7 +189,7 @@ class Sampler: expanded_local_pos: torch.Tensor, skip_top_k_top_p: bool = False, ) -> torch.Tensor: - if not self._requires_logits_processing(idx_mapping_np): + if not np.any(self.needs_logits_processing[idx_mapping_np]): return logits # Copy logits to a new FP32 tensor. @@ -228,24 +245,6 @@ class Sampler: logits, expanded_idx_mapping, idx_mapping_np ) - def _requires_logits_processing(self, idx_mapping_np: np.ndarray) -> bool: - if np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]): - return True - if np.any(self.penalties_state.use_penalty[idx_mapping_np]): - return True - if np.any(self.bad_words_state.num_bad_words.np[idx_mapping_np] > 0): - return True - - states = self.sampling_states - temperatures = states.temperature.np[idx_mapping_np] - if np.any((temperatures != 0.0) & (temperatures != 1.0)): - return True - if np.any(states.min_p.np[idx_mapping_np] != 0.0): - return True - if np.any(states.top_k.np[idx_mapping_np] != states.vocab_size): - return True - return bool(np.any(states.top_p.np[idx_mapping_np] != 1.0)) - def sample( self, logits: torch.Tensor,