[Performance][MRV2] Cache logits-processing request state (#52329)
Signed-off-by: Rui Rui4 Cheng <chengrui4@lenovo.com> Co-authored-by: Rui Rui4 Cheng <chengrui4@lenovo.com>
Cheng Rui committed
Aug 16, 2026 at 19:36 UTC
6664d397bf091cb9371cba481d4efb8233436fe6
3 files changed
+147
-19
tests/v1/worker/test_gpu_sampler_flags.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import numpy as np
5
+import pytest
6
+import torch
7
+
8
+pytest.importorskip("triton")
9
+if not torch.cuda.is_available():
10
+ pytest.skip("CUDA required for sampler flag tests", allow_module_level=True)
11
+
12
+from vllm.sampling_params import SamplingParams
13
+from vllm.v1.worker.gpu.sample.sampler import Sampler
14
+from vllm.v1.worker.gpu.states import RequestState
15
+
16
+DEVICE = torch.device("cuda")
17
+VOCAB_SIZE = 128
18
+
19
+
20
+class MockReasoningConfig:
21
+ reasoning_start_token_ids = [90]
22
+ reasoning_end_token_ids = [91]
23
+ natural_reasoning_end_token_ids = [91]
24
+
25
+
26
+def _make_sampler() -> Sampler:
27
+ req_states = RequestState(
28
+ max_num_reqs=4,
29
+ max_model_len=64,
30
+ max_num_batched_tokens=16,
31
+ num_speculative_steps=1,
32
+ vocab_size=VOCAB_SIZE,
33
+ device=DEVICE,
34
+ )
35
+ return Sampler(
36
+ max_num_reqs=4,
37
+ vocab_size=VOCAB_SIZE,
38
+ device=DEVICE,
39
+ req_states=req_states,
40
+ reasoning_config=MockReasoningConfig(),
41
+ )
42
+
43
+
44
+@pytest.mark.parametrize(
45
+ ("sampling_params", "expected"),
46
+ [
47
+ pytest.param(SamplingParams(), False, id="defaults"),
48
+ pytest.param(SamplingParams(temperature=0.0), False, id="greedy"),
49
+ pytest.param(
50
+ SamplingParams(thinking_token_budget=3), True, id="thinking-budget"
51
+ ),
52
+ pytest.param(SamplingParams(logit_bias={1: 1.0}), True, id="logit-bias"),
53
+ pytest.param(SamplingParams(frequency_penalty=0.1), True, id="penalty"),
54
+ pytest.param(SamplingParams(_bad_words_token_ids=[[1]]), True, id="bad-words"),
55
+ pytest.param(SamplingParams(temperature=0.7), True, id="temperature"),
56
+ pytest.param(SamplingParams(min_p=0.1), True, id="min-p"),
57
+ pytest.param(SamplingParams(top_k=10), True, id="top-k"),
58
+ pytest.param(SamplingParams(top_p=0.9), True, id="top-p"),
59
+ pytest.param(
60
+ SamplingParams.for_sampler_warmup(), True, id="all-logits-processors"
61
+ ),
62
+ ],
63
+)
64
+def test_logits_processing_cache_matches_request_features(
65
+ sampling_params: SamplingParams, expected: bool
66
+):
67
+ sampler = _make_sampler()
68
+ sampler.add_request(3, prompt_len=1, sampling_params=sampling_params)
69
+
70
+ assert sampler.needs_logits_processing[3] == expected
71
+
72
+
73
+def test_logits_processing_cache_is_overwritten_when_slot_is_reused():
74
+ sampler = _make_sampler()
75
+ sampler.add_request(3, 1, SamplingParams.for_sampler_warmup())
76
+ sampler.add_request(3, 1, SamplingParams())
77
+
78
+ assert not sampler.needs_logits_processing[3]
79
+
80
+
81
+def test_logits_processing_cache_only_checks_active_requests():
82
+ sampler = _make_sampler()
83
+ sampler.add_request(0, 1, SamplingParams(temperature=0.0))
84
+ sampler.add_request(2, 1, SamplingParams.for_sampler_warmup())
85
+
86
+ sampling_only = np.array([0], dtype=np.int32)
87
+ with_processing = np.array([0, 2], dtype=np.int32)
88
+
89
+ assert not np.any(sampler.needs_logits_processing[sampling_only])
90
+ assert np.any(sampler.needs_logits_processing[with_processing])
tests/v1/worker/test_gpu_thinking_budget.py
+39
@@ -12,6 +12,7 @@ if not torch.cuda.is_available():
12
)
13
14
from vllm.sampling_params import SamplingParams
15
+from vllm.v1.worker.gpu.sample.sampler import Sampler
16
from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
17
from vllm.v1.worker.gpu.states import RequestState
18
@@ -185,6 +186,44 @@ def test_v2_thinking_budget_ignores_plain_request():
186
assert torch.all(out == 0)
187
188
189
+def test_v2_greedy_sampling_applies_thinking_budget():
190
+ """Greedy-only requests must not bypass thinking-budget processing."""
191
+ req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
192
+ sampler = Sampler(
193
+ max_num_reqs=4,
194
+ vocab_size=VOCAB_SIZE,
195
+ device=DEVICE,
196
+ req_states=req_states,
197
+ reasoning_config=MockReasoningConfig(),
198
+ )
199
+ sampler.add_request(
200
+ req_idx=3,
201
+ prompt_len=1,
202
+ sampling_params=SamplingParams(
203
+ temperature=0.0,
204
+ thinking_token_budget=3,
205
+ ),
206
+ )
207
+ sampler.apply_staged_writes()
208
+
209
+ idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE)
210
+ idx_mapping_np = idx_mapping.cpu().numpy()
211
+ expanded_idx_mapping = idx_mapping.clone()
212
+ input_ids = torch.tensor([12], dtype=torch.int32, device=DEVICE)
213
+ logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
214
+ out = sampler.apply_sampling_params(
215
+ logits,
216
+ expanded_idx_mapping,
217
+ idx_mapping,
218
+ idx_mapping_np,
219
+ torch.tensor([4], dtype=torch.int32, device=DEVICE),
220
+ input_ids,
221
+ torch.tensor([0], dtype=torch.int32, device=DEVICE),
222
+ )
223
+
224
+ assert out[0, END].item() == pytest.approx(1.0e9)
225
+
226
+
227
def test_v2_thinking_budget_latest_prefill_end_disables_forcing():
228
req_states = _make_req_states(
229
[1, START, 10, 11, 12, END, 13],
vllm/v1/worker/gpu/sample/sampler.py
+18
-19
@@ -53,6 +53,7 @@ class Sampler:
53
self.bad_words_state = BadWordsState(req_states)
54
self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
55
self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config)
56
+ self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool)
57
self.num_speculative_tokens = num_speculative_tokens
58
self.return_sampling_mask = return_sampling_mask
59
self.use_flashinfer = (
@@ -69,6 +70,22 @@ class Sampler:
70
self.logprob_token_ids_state.add_request(req_idx, sampling_params)
71
self.thinking_budget_state.add_request(req_idx, sampling_params)
72
73
+ states = self.sampling_states
74
+ temperature = states.temperature.np[req_idx]
75
+ self.needs_logits_processing[req_idx] = (
76
+ self.logit_bias_state.use_logit_bias[req_idx]
77
+ or self.penalties_state.use_penalty[req_idx]
78
+ or self.bad_words_state.num_bad_words.np[req_idx] > 0
79
+ or (
80
+ self.thinking_budget_state.enabled
81
+ and self.thinking_budget_state.use_thinking_budget[req_idx]
82
+ )
83
+ or (temperature != 0.0 and temperature != 1.0)
84
+ or states.min_p.np[req_idx] != 0.0
85
+ or states.top_k.np[req_idx] != states.vocab_size
86
+ or states.top_p.np[req_idx] != 1.0
87
+ )
88
+
89
def apply_staged_writes(self) -> None:
90
self.sampling_states.apply_staged_writes()
91
self.penalties_state.apply_staged_writes()
@@ -172,7 +189,7 @@ class Sampler:
189
expanded_local_pos: torch.Tensor,
190
skip_top_k_top_p: bool = False,
191
) -> torch.Tensor:
175
- if not self._requires_logits_processing(idx_mapping_np):
192
+ if not np.any(self.needs_logits_processing[idx_mapping_np]):
193
return logits
194
195
# Copy logits to a new FP32 tensor.
@@ -228,24 +245,6 @@ class Sampler:
245
logits, expanded_idx_mapping, idx_mapping_np
246
)
247
231
- def _requires_logits_processing(self, idx_mapping_np: np.ndarray) -> bool:
232
- if np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]):
233
- return True
234
- if np.any(self.penalties_state.use_penalty[idx_mapping_np]):
235
- return True
236
- if np.any(self.bad_words_state.num_bad_words.np[idx_mapping_np] > 0):
237
- return True
238
-
239
- states = self.sampling_states
240
- temperatures = states.temperature.np[idx_mapping_np]
241
- if np.any((temperatures != 0.0) & (temperatures != 1.0)):
242
- return True
243
- if np.any(states.min_p.np[idx_mapping_np] != 0.0):
244
- return True
245
- if np.any(states.top_k.np[idx_mapping_np] != states.vocab_size):
246
- return True
247
- return bool(np.any(states.top_p.np[idx_mapping_np] != 1.0))
248
-
248
def sample(
249
self,
250
logits: torch.Tensor,