| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | pytest.importorskip("triton") |
| 11 | if not torch.cuda.is_available(): |
| 12 | pytest.skip( |
| 13 | "CUDA required for Model Runner V2 thinking budget tests", |
| 14 | allow_module_level=True, |
| 15 | ) |
| 16 | |
| 17 | from vllm.sampling_params import SamplingParams |
| 18 | from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext |
| 19 | from vllm.v1.worker.gpu.sample.sampler import Sampler |
| 20 | from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState |
| 21 | from vllm.v1.worker.gpu.states import RequestState |
| 22 | |
| 23 | DEVICE = torch.device("cuda") |
| 24 | START = 90 |
| 25 | END = 91 |
| 26 | END_A = 92 |
| 27 | END_B = 93 |
| 28 | VOCAB_SIZE = 128 |
| 29 | |
| 30 | |
| 31 | class MockReasoningConfig: |
| 32 | reasoning_start_token_ids = [START] |
| 33 | reasoning_end_token_ids = [END] |
| 34 | natural_reasoning_end_token_ids = [END] |
| 35 | |
| 36 | |
| 37 | class MockMultiTokenEndReasoningConfig: |
| 38 | reasoning_start_token_ids = [START] |
| 39 | reasoning_end_token_ids = [END_A, END_B] |
| 40 | natural_reasoning_end_token_ids = [END_A, END_B] |
| 41 | |
| 42 | |
| 43 | class MockDistinctEndReasoningConfig: |
| 44 | reasoning_start_token_ids = [START] |
| 45 | reasoning_end_token_ids = [END_A, END_B] |
| 46 | natural_reasoning_end_token_ids = [END] |
| 47 | |
| 48 | |
| 49 | def _make_req_states(tokens: list[int], prompt_len: int = 1) -> RequestState: |
| 50 | req_states = RequestState( |
| 51 | max_num_reqs=4, |
| 52 | max_model_len=max(64, len(tokens) + 1), |
| 53 | max_num_batched_tokens=16, |
| 54 | num_speculative_steps=4, |
| 55 | vocab_size=VOCAB_SIZE, |
| 56 | device=DEVICE, |
| 57 | ) |
| 58 | req_states.add_request( |
| 59 | req_id="req", |
| 60 | prompt_len=prompt_len, |
| 61 | all_token_ids=tokens, |
| 62 | num_computed_tokens=len(tokens), |
| 63 | max_tokens=32, |
| 64 | ) |
| 65 | req_states.apply_staged_writes() |
| 66 | return req_states |
| 67 | |
| 68 | |
| 69 | def _apply( |
| 70 | state: ThinkingBudgetState, |
| 71 | logits: torch.Tensor, |
| 72 | input_ids: list[int], |
| 73 | local_pos: list[int], |
| 74 | ) -> torch.Tensor: |
| 75 | idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE) |
| 76 | expanded_idx_mapping = torch.tensor( |
| 77 | [3] * len(input_ids), dtype=torch.int32, device=DEVICE |
| 78 | ) |
| 79 | idx_mapping_np = idx_mapping.cpu().numpy() |
| 80 | state.apply( |
| 81 | logits, |
| 82 | LogitsContext( |
| 83 | expanded_idx_mapping=expanded_idx_mapping, |
| 84 | idx_mapping=idx_mapping, |
| 85 | idx_mapping_np=idx_mapping_np, |
| 86 | expanded_local_pos=torch.tensor( |
| 87 | local_pos, dtype=torch.int32, device=DEVICE |
| 88 | ), |
| 89 | input_ids=torch.tensor(input_ids, dtype=torch.int32, device=DEVICE), |
| 90 | pos=torch.zeros(len(input_ids), dtype=torch.int32, device=DEVICE), |
| 91 | seq_lens_upper_bound_np=np.full(1, len(input_ids), dtype=np.int64), |
| 92 | ), |
| 93 | ) |
| 94 | return logits.cpu() |
| 95 | |
| 96 | |
| 97 | def test_v2_thinking_budget_forces_end_after_budget_reached(): |
| 98 | req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) |
| 99 | state = ThinkingBudgetState(req_states, MockReasoningConfig()) |
| 100 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 101 | state.apply_staged_writes() |
| 102 | |
| 103 | logits = torch.arange(VOCAB_SIZE, dtype=torch.float32, device=DEVICE).view(1, -1) |
| 104 | expected = logits.cpu() |
| 105 | out = _apply(state, logits, input_ids=[12], local_pos=[0]) |
| 106 | |
| 107 | expected[0, END] = 1.0e9 |
| 108 | torch.testing.assert_close(out, expected) |
| 109 | |
| 110 | |
| 111 | def test_v2_thinking_budget_restores_masked_end_token(): |
| 112 | req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) |
| 113 | state = ThinkingBudgetState(req_states, MockReasoningConfig()) |
| 114 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 115 | state.apply_staged_writes() |
| 116 | |
| 117 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 118 | logits[0, END] = -float("inf") |
| 119 | out = _apply(state, logits, input_ids=[12], local_pos=[0]) |
| 120 | |
| 121 | assert out[0, END] == pytest.approx(1.0e9) |
| 122 | |
| 123 | |
| 124 | def test_v2_thinking_budget_allows_tokens_before_budget(): |
| 125 | req_states = _make_req_states([1, START, 10, 11], prompt_len=1) |
| 126 | state = ThinkingBudgetState(req_states, MockReasoningConfig()) |
| 127 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 128 | state.apply_staged_writes() |
| 129 | |
| 130 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 131 | out = _apply(state, logits, input_ids=[11], local_pos=[0]) |
| 132 | |
| 133 | assert torch.all(out == 0) |
| 134 | |
| 135 | |
| 136 | def test_v2_thinking_budget_continues_multi_token_end_marker(): |
| 137 | req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) |
| 138 | state = ThinkingBudgetState(req_states, MockMultiTokenEndReasoningConfig()) |
| 139 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 140 | state.apply_staged_writes() |
| 141 | |
| 142 | logits = torch.zeros((2, VOCAB_SIZE), device=DEVICE) |
| 143 | out = _apply( |
| 144 | state, |
| 145 | logits, |
| 146 | input_ids=[12, END_A], |
| 147 | local_pos=[0, 1], |
| 148 | ) |
| 149 | |
| 150 | assert out[0, END_A] == pytest.approx(1.0e9) |
| 151 | assert out[1, END_B] == pytest.approx(1.0e9) |
| 152 | |
| 153 | |
| 154 | def test_v2_thinking_budget_uses_distinct_forced_end_marker(): |
| 155 | req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) |
| 156 | state = ThinkingBudgetState(req_states, MockDistinctEndReasoningConfig()) |
| 157 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 158 | state.apply_staged_writes() |
| 159 | |
| 160 | logits = torch.zeros((2, VOCAB_SIZE), device=DEVICE) |
| 161 | out = _apply( |
| 162 | state, |
| 163 | logits, |
| 164 | input_ids=[12, END_A], |
| 165 | local_pos=[0, 1], |
| 166 | ) |
| 167 | |
| 168 | assert out[0, END_A] == pytest.approx(1.0e9) |
| 169 | assert out[1, END_B] == pytest.approx(1.0e9) |
| 170 | |
| 171 | |
| 172 | def test_v2_thinking_budget_stops_after_natural_end_marker(): |
| 173 | req_states = _make_req_states( |
| 174 | [1, START, 10, END, 20, 21, 22], |
| 175 | prompt_len=1, |
| 176 | ) |
| 177 | state = ThinkingBudgetState(req_states, MockDistinctEndReasoningConfig()) |
| 178 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 179 | state.apply_staged_writes() |
| 180 | |
| 181 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 182 | out = _apply(state, logits, input_ids=[22], local_pos=[0]) |
| 183 | |
| 184 | assert torch.all(out == 0) |
| 185 | |
| 186 | |
| 187 | def test_v2_thinking_budget_ignores_plain_request(): |
| 188 | req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) |
| 189 | state = ThinkingBudgetState(req_states, MockReasoningConfig()) |
| 190 | state.add_request(3, SamplingParams()) |
| 191 | state.apply_staged_writes() |
| 192 | |
| 193 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 194 | out = _apply(state, logits, input_ids=[12], local_pos=[0]) |
| 195 | |
| 196 | assert torch.all(out == 0) |
| 197 | |
| 198 | |
| 199 | def test_v2_greedy_sampling_applies_thinking_budget(): |
| 200 | """Greedy-only requests must not bypass thinking-budget processing.""" |
| 201 | req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) |
| 202 | sampler = Sampler( |
| 203 | vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()), |
| 204 | max_num_reqs=4, |
| 205 | vocab_size=VOCAB_SIZE, |
| 206 | device=DEVICE, |
| 207 | req_states=req_states, |
| 208 | ) |
| 209 | sampler.add_request( |
| 210 | req_idx=3, |
| 211 | sampling_params=SamplingParams( |
| 212 | temperature=0.0, |
| 213 | thinking_token_budget=3, |
| 214 | ), |
| 215 | ) |
| 216 | sampler.apply_staged_writes() |
| 217 | |
| 218 | idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE) |
| 219 | idx_mapping_np = idx_mapping.cpu().numpy() |
| 220 | expanded_idx_mapping = idx_mapping.clone() |
| 221 | input_ids = torch.tensor([12], dtype=torch.int32, device=DEVICE) |
| 222 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 223 | out = sampler.apply_sampling_params( |
| 224 | logits, |
| 225 | expanded_idx_mapping, |
| 226 | idx_mapping, |
| 227 | idx_mapping_np, |
| 228 | torch.tensor([4], dtype=torch.int32, device=DEVICE), |
| 229 | input_ids, |
| 230 | torch.tensor([0], dtype=torch.int32, device=DEVICE), |
| 231 | np.full(1, 4, dtype=np.int64), |
| 232 | ) |
| 233 | |
| 234 | assert out[0, END].item() == pytest.approx(1.0e9) |
| 235 | |
| 236 | |
| 237 | def test_v2_thinking_budget_latest_prefill_end_disables_forcing(): |
| 238 | req_states = _make_req_states( |
| 239 | [1, START, 10, 11, 12, END, 13], |
| 240 | prompt_len=1, |
| 241 | ) |
| 242 | state = ThinkingBudgetState(req_states, MockReasoningConfig()) |
| 243 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 244 | state.apply_staged_writes() |
| 245 | |
| 246 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 247 | out = _apply(state, logits, input_ids=[13], local_pos=[0]) |
| 248 | |
| 249 | assert torch.all(out == 0) |
| 250 | |
| 251 | |
| 252 | def test_v2_thinking_budget_uses_latest_prefill_start_boundary(): |
| 253 | req_states = _make_req_states( |
| 254 | [1, START, 10, 11, 12, END, 13, START, 14, 15, 16], |
| 255 | prompt_len=1, |
| 256 | ) |
| 257 | state = ThinkingBudgetState(req_states, MockReasoningConfig()) |
| 258 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 259 | state.apply_staged_writes() |
| 260 | |
| 261 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 262 | out = _apply(state, logits, input_ids=[16], local_pos=[0]) |
| 263 | |
| 264 | assert out[0, END] == pytest.approx(1.0e9) |
| 265 | |
| 266 | |
| 267 | def test_v2_thinking_budget_incrementally_scans_long_generation(): |
| 268 | """Guard against rescanning the full token history on every decode step.""" |
| 269 | tokens = [1, START, *([10] * 16382)] |
| 270 | req_states = _make_req_states(tokens) |
| 271 | state = ThinkingBudgetState(req_states, MockReasoningConfig()) |
| 272 | state.add_request(3, SamplingParams(thinking_token_budget=32768)) |
| 273 | state.apply_staged_writes() |
| 274 | |
| 275 | _apply(state, torch.zeros((1, VOCAB_SIZE), device=DEVICE), [10], [0]) |
| 276 | assert state.cached_scan_pos[3].item() == len(tokens) |
| 277 | |
| 278 | req_states.all_token_ids.stage_write(3, len(tokens), [10]) |
| 279 | req_states.total_len.stage_write_elem(3, len(tokens) + 1) |
| 280 | req_states.apply_staged_writes() |
| 281 | _apply(state, torch.zeros((1, VOCAB_SIZE), device=DEVICE), [10], [0]) |
| 282 | |
| 283 | assert state.cached_scan_pos[3].item() == len(tokens) + 1 |
| 284 | |
| 285 | |
| 286 | def test_v2_thinking_budget_clamps_oversized_budget(): |
| 287 | """Budgets beyond int32 must not crash and behave as unlimited.""" |
| 288 | req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) |
| 289 | state = ThinkingBudgetState(req_states, MockReasoningConfig()) |
| 290 | state.add_request(3, SamplingParams(thinking_token_budget=2**40)) |
| 291 | state.apply_staged_writes() |
| 292 | |
| 293 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 294 | out = _apply(state, logits, input_ids=[12], local_pos=[0]) |
| 295 | |
| 296 | assert torch.all(out == 0) |
| 297 | |
| 298 | |
| 299 | def test_v2_thinking_budget_continues_end_prefix_from_prompt(): |
| 300 | """A resumed prompt ending with a partial forced-end marker must not |
| 301 | restart the marker sequence and duplicate its first token.""" |
| 302 | req_states = _make_req_states([1, START, 10, 11, END_A], prompt_len=5) |
| 303 | state = ThinkingBudgetState(req_states, MockMultiTokenEndReasoningConfig()) |
| 304 | state.add_request(3, SamplingParams(thinking_token_budget=3)) |
| 305 | state.apply_staged_writes() |
| 306 | |
| 307 | logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) |
| 308 | out = _apply(state, logits, input_ids=[END_A], local_pos=[0]) |
| 309 | |
| 310 | assert out[0, END_B] == pytest.approx(1.0e9) |
| 311 | assert out[0, END_A] == 0 |