| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import random |
| 5 | from collections.abc import Callable |
| 6 | from typing import NamedTuple, NoReturn, TypeAlias |
| 7 | |
| 8 | import numpy as np |
| 9 | import pytest |
| 10 | import torch |
| 11 | |
| 12 | from tests.utils import create_new_process_for_each_test |
| 13 | from tests.v1.sample.utils import ( |
| 14 | LogitsprocsTestFakes, |
| 15 | create_fake_logits, |
| 16 | create_mock_reasoning_config, |
| 17 | create_penalty_tensor, |
| 18 | create_prompt_tokens_tensor, |
| 19 | fake_apply_logitsprocs, |
| 20 | fake_update_logitsprocs_state, |
| 21 | ) |
| 22 | from vllm.config import VllmConfig |
| 23 | from vllm.platforms import current_platform |
| 24 | from vllm.sampling_params import ( |
| 25 | SamplingParams, |
| 26 | StructuredOutputsParams, |
| 27 | validate_thinking_token_budget, |
| 28 | ) |
| 29 | from vllm.utils.platform_utils import is_pin_memory_available |
| 30 | from vllm.v1.sample.logits_processor import ( |
| 31 | BatchUpdate, |
| 32 | BatchUpdateBuilder, |
| 33 | LogitBiasLogitsProcessor, |
| 34 | LogitsProcessor, |
| 35 | MinPLogitsProcessor, |
| 36 | MinTokensLogitsProcessor, |
| 37 | MoveDirectionality, |
| 38 | build_logitsprocs, |
| 39 | ) |
| 40 | from vllm.v1.sample.metadata import SamplingMetadata |
| 41 | from vllm.v1.sample.thinking_budget_state import ( |
| 42 | ThinkingBudgetStateHolder, |
| 43 | maybe_create_thinking_budget_state_holder, |
| 44 | ) |
| 45 | |
| 46 | PIN_MEMORY_AVAILABLE = is_pin_memory_available() |
| 47 | MAX_NUM_REQS = 256 |
| 48 | VOCAB_SIZE = 1024 |
| 49 | NUM_OUTPUT_TOKENS = 20 |
| 50 | DEVICE_TYPE = current_platform.device_type |
| 51 | DEVICES = [ |
| 52 | f"{DEVICE_TYPE}:{i}" |
| 53 | for i in range(1 if current_platform.device_count() == 1 else 2) |
| 54 | ] |
| 55 | MAX_NUM_PROMPT_TOKENS = 64 |
| 56 | MIN_TOKENS_LEN_THRESHOLD = 5 |
| 57 | REQS_PER_LOGITPROC = 50 |
| 58 | STR_NO_LOGITPROC = "none" |
| 59 | # Thinking budget uses ``ThinkingBudgetStateHolder`` (not a logits processor). |
| 60 | STR_THINKING_BUDGET = "thinking_budget" |
| 61 | |
| 62 | # Thinking token budget testing constants |
| 63 | THINKING_TOKEN_BUDGET = 5 |
| 64 | THINK_START_TOKEN_ID = 999 |
| 65 | THINK_END_TOKEN_ID = 998 |
| 66 | |
| 67 | # LogitsProcessor subclass or "none" |
| 68 | LogitprocType: TypeAlias = type[LogitsProcessor] | str |
| 69 | |
| 70 | |
| 71 | class LogitsProcsRequestParams: |
| 72 | """Encapsulates key params for a single request in a batch. |
| 73 | |
| 74 | Params can be customized based on the enabled logitproc |
| 75 | """ |
| 76 | |
| 77 | workload_index: int |
| 78 | logitproc_type: LogitprocType # Logitproc enabled, specified by str id |
| 79 | out_tokens: list[int] # Output tokens required for min tokens test |
| 80 | prompt_tokens: list[int] # Dummy prompt tokens placeholder |
| 81 | params: SamplingParams # Settings customized for logitproc |
| 82 | |
| 83 | def __init__(self, workload_index: int, logitproc_type: LogitprocType): |
| 84 | self.workload_index = workload_index |
| 85 | self.logitproc_type = logitproc_type |
| 86 | # Number of output tokens is randomly 0 or twice the min-tokens |
| 87 | # threshold which will be used in testing. |
| 88 | # Generate diverse random tokens for all processors (more realistic) |
| 89 | num_tokens = MIN_TOKENS_LEN_THRESHOLD * random.randint(0, 2) |
| 90 | if num_tokens > 0: |
| 91 | # Use diverse random tokens |
| 92 | self.out_tokens = [random.randint(1, 950) for _ in range(num_tokens)] |
| 93 | # Think-start seed for ``STR_THINKING_BUDGET`` rows. |
| 94 | if logitproc_type == STR_THINKING_BUDGET: |
| 95 | self.out_tokens[0] = THINK_START_TOKEN_ID |
| 96 | else: |
| 97 | self.out_tokens = [] |
| 98 | self.prompt_tokens = [] |
| 99 | self.params = _sampling_params_from_logitproc(logitproc_type) |
| 100 | |
| 101 | def __str__(self): |
| 102 | """For debugging.""" |
| 103 | summ = ", ".join(f"{k}={v}" for k, v in vars(self).items()) |
| 104 | return f"MyClass({summ})" |
| 105 | |
| 106 | |
| 107 | def _generate_fake_sampling_metadata( |
| 108 | num_output_tokens: int, |
| 109 | batch_size: int, |
| 110 | vocab_size: int, |
| 111 | device: torch.device, |
| 112 | ) -> SamplingMetadata: |
| 113 | """Generate fake sampling metadata with fake logitsprocs.""" |
| 114 | output_token_ids: list[list[int]] = [] |
| 115 | prompt_token_ids: list[list[int]] = [] |
| 116 | for _ in range(batch_size): |
| 117 | output_token_ids.append( |
| 118 | np.random.randint(0, vocab_size, size=num_output_tokens).tolist() |
| 119 | ) |
| 120 | prompt_token_ids.append( |
| 121 | np.random.randint( |
| 122 | 0, vocab_size, size=np.random.randint(1, MAX_NUM_PROMPT_TOKENS) |
| 123 | ).tolist() |
| 124 | ) |
| 125 | |
| 126 | vllm_config = VllmConfig() |
| 127 | vllm_config.reasoning_config = create_mock_reasoning_config( |
| 128 | [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID] |
| 129 | ) |
| 130 | |
| 131 | logitsprocs = build_logitsprocs( |
| 132 | vllm_config=vllm_config, |
| 133 | device=device, |
| 134 | is_pin_memory=PIN_MEMORY_AVAILABLE, |
| 135 | is_pooling_model=False, |
| 136 | ) |
| 137 | num_spec = ( |
| 138 | vllm_config.speculative_config.num_speculative_tokens |
| 139 | if vllm_config.speculative_config |
| 140 | else 0 |
| 141 | ) |
| 142 | thinking_holder = maybe_create_thinking_budget_state_holder( |
| 143 | vllm_config.reasoning_config, |
| 144 | vllm_config.scheduler_config.max_num_seqs, |
| 145 | num_spec, |
| 146 | device, |
| 147 | is_pin_memory=False, |
| 148 | ) |
| 149 | fake_sampling_metadata = SamplingMetadata( |
| 150 | temperature=torch.full((batch_size,), 0.0), |
| 151 | all_greedy=True, |
| 152 | all_random=False, |
| 153 | top_p=None, |
| 154 | top_k=None, |
| 155 | generators={}, |
| 156 | max_num_logprobs=0, |
| 157 | prompt_token_ids=create_prompt_tokens_tensor( |
| 158 | prompt_token_ids, vocab_size, device |
| 159 | ), |
| 160 | output_token_ids=output_token_ids, |
| 161 | frequency_penalties=create_penalty_tensor(batch_size, 0.0, device), |
| 162 | presence_penalties=create_penalty_tensor(batch_size, 0.0, device), |
| 163 | repetition_penalties=create_penalty_tensor(batch_size, 1.0, device), |
| 164 | no_penalties=True, |
| 165 | allowed_token_ids_mask=None, |
| 166 | bad_words_token_ids={}, |
| 167 | logitsprocs=logitsprocs, |
| 168 | thinking_budget_state_holder=thinking_holder, |
| 169 | ) |
| 170 | return fake_sampling_metadata |
| 171 | |
| 172 | |
| 173 | def _generate_test_fakes(batch_size: int, device: str) -> LogitsprocsTestFakes: |
| 174 | """Generate fake logits and sampling metadata.""" |
| 175 | fake_logits = create_fake_logits(batch_size, VOCAB_SIZE) |
| 176 | # Create one dominant token per batch, to support min-p test |
| 177 | for i in range(batch_size): |
| 178 | fake_logits[i, 0] = 10.0 # High logit for first token |
| 179 | fake_logits[i, 1:] = 1e-2 # Others remain low |
| 180 | sampling_metadata = _generate_fake_sampling_metadata( |
| 181 | NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device) |
| 182 | ) |
| 183 | return LogitsprocsTestFakes( |
| 184 | logits=fake_logits, |
| 185 | sampling_metadata=sampling_metadata, |
| 186 | ) |
| 187 | |
| 188 | |
| 189 | def _sampling_params_from_logitproc(logitproc_type: LogitprocType) -> SamplingParams: |
| 190 | """Customize request SamplingParams for a specified logitproc.""" |
| 191 | # SamplingParams for req with no logitproc |
| 192 | kwargs = {"min_p": 0.0, "logit_bias": None, "min_tokens": 0} |
| 193 | if fxn := logitsprocs_test_mapping[logitproc_type].gen_request_fxn: |
| 194 | fxn(kwargs) |
| 195 | return SamplingParams(**kwargs) |
| 196 | |
| 197 | |
| 198 | def _generate_mixed_logitsprocs_batch_params( |
| 199 | reqs_per_logitproc: int, |
| 200 | logitsprocs_types: list[LogitprocType], |
| 201 | ) -> list[LogitsProcsRequestParams]: |
| 202 | """Define key params for a batch of requests with a different |
| 203 | logitproc enabled per request. |
| 204 | |
| 205 | The batch will have `reqs_per_logitproc` repeats for all |
| 206 | `logitsprocs_types` under test, including the case where |
| 207 | no logitsproc is enabled. The batch is randomly shuffled. The |
| 208 | size of the batch is `reqs_per_logitproc` times |
| 209 | `n = len(logitsprocs_types)` |
| 210 | |
| 211 | Args: |
| 212 | reqs_per_logitproc: number of requests using each logitproc |
| 213 | logitsprocs_types: logitsprocs under test |
| 214 | |
| 215 | Returns: |
| 216 | List of per-request params which configure the engine for that request's |
| 217 | enabled logitproc |
| 218 | |
| 219 | """ |
| 220 | batch_size = len(logitsprocs_types) * reqs_per_logitproc |
| 221 | # Generate multiple repeats of key params for each logitproc; |
| 222 | # apply random inverse permutation to the iteration |
| 223 | # over logitsprocs, such that logitsprocs are shuffled. |
| 224 | batch_perm = random.sample(range(batch_size), k=batch_size) |
| 225 | return [ |
| 226 | LogitsProcsRequestParams( |
| 227 | workload_index=idx, |
| 228 | logitproc_type=logitsprocs_types[pdx // reqs_per_logitproc], |
| 229 | ) |
| 230 | for idx, pdx in enumerate(batch_perm) |
| 231 | ] |
| 232 | |
| 233 | |
| 234 | def _raise_error_invalid( |
| 235 | msg_suffix: str, |
| 236 | batch_index: int, |
| 237 | request_params: LogitsProcsRequestParams, |
| 238 | step_idx: int, |
| 239 | err_cls: type[Exception] = ValueError, |
| 240 | ) -> NoReturn: |
| 241 | raise err_cls( |
| 242 | f"Validation failed for step={step_idx}, " |
| 243 | f"batch_index={batch_index}, " |
| 244 | f"workload_index={request_params.workload_index}, " |
| 245 | f"req_params={request_params}. Reason: {msg_suffix}" |
| 246 | ) |
| 247 | |
| 248 | |
| 249 | def _logit_bias_params(kwargs: dict) -> None: |
| 250 | """Logit bias config.""" |
| 251 | kwargs["logit_bias"] = { |
| 252 | random.randint(0, VOCAB_SIZE - 1): random.choice([-0.1, 0.2]) |
| 253 | } |
| 254 | |
| 255 | |
| 256 | def _logit_bias_validate( |
| 257 | test_fakes: LogitsprocsTestFakes, |
| 258 | persistent_batch: list[LogitsProcsRequestParams], |
| 259 | logits_new: torch.Tensor, |
| 260 | batch_index: int, |
| 261 | request_params: LogitsProcsRequestParams, |
| 262 | step_idx: int, |
| 263 | ) -> None: |
| 264 | """Validate logit bias logitproc applied correctly.""" |
| 265 | logit_bias = request_params.params.logit_bias |
| 266 | assert logit_bias is not None |
| 267 | logits_old = test_fakes.logits[persistent_batch[batch_index].workload_index].cpu() |
| 268 | logits_new = logits_new[batch_index].cpu() |
| 269 | for token_id in range(VOCAB_SIZE): |
| 270 | logit_old_value = logits_old[token_id] |
| 271 | logit_new_value = logits_new[token_id] |
| 272 | if token_id in logit_bias: |
| 273 | bias_value = logit_bias[token_id] |
| 274 | exp_value = bias_value + logit_old_value |
| 275 | if logit_new_value != pytest.approx(exp_value): |
| 276 | _raise_error_invalid( |
| 277 | msg_suffix=( |
| 278 | f"Biased token {token_id} logit value {logit_new_value} " |
| 279 | f"does not match expected value {exp_value} " |
| 280 | f"given bias {bias_value}" |
| 281 | ), |
| 282 | batch_index=batch_index, |
| 283 | request_params=request_params, |
| 284 | step_idx=step_idx, |
| 285 | ) |
| 286 | |
| 287 | else: |
| 288 | if logit_new_value != pytest.approx(logit_old_value): |
| 289 | _raise_error_invalid( |
| 290 | msg_suffix=( |
| 291 | f"Unbiased token {token_id} logit value {logit_new_value} " |
| 292 | f"does not match expected value {logit_old_value}" |
| 293 | ), |
| 294 | batch_index=batch_index, |
| 295 | request_params=request_params, |
| 296 | step_idx=step_idx, |
| 297 | ) |
| 298 | |
| 299 | |
| 300 | def _min_p_params(kwargs: dict) -> None: |
| 301 | """Min-p logitproc config.""" |
| 302 | kwargs["min_p"] = 0.1 |
| 303 | |
| 304 | |
| 305 | def _min_p_validate( |
| 306 | test_fakes: LogitsprocsTestFakes, |
| 307 | persistent_batch: list[LogitsProcsRequestParams], |
| 308 | logits_new: torch.Tensor, |
| 309 | batch_index: int, |
| 310 | request_params: LogitsProcsRequestParams, |
| 311 | step_idx: int, |
| 312 | ) -> None: |
| 313 | """Validate min-p logitproc applied correctly.""" |
| 314 | for token_id in range(VOCAB_SIZE): |
| 315 | logits_for_token = logits_new[batch_index][token_id] |
| 316 | if token_id == 0: |
| 317 | # Dominant token should always be unmasked |
| 318 | if logits_for_token == -float("inf"): |
| 319 | _raise_error_invalid( |
| 320 | msg_suffix="Invalid: dominant token 0 masked (-inf)", |
| 321 | batch_index=batch_index, |
| 322 | request_params=request_params, |
| 323 | step_idx=step_idx, |
| 324 | ) |
| 325 | else: |
| 326 | if request_params.params.min_p > 0.0: |
| 327 | # Non-dominant tokens should be masked when min_p > 0 |
| 328 | if logits_for_token != -float("inf"): |
| 329 | _raise_error_invalid( |
| 330 | msg_suffix=f"Invalid: non-dominant token {token_id} not masked", |
| 331 | batch_index=batch_index, |
| 332 | request_params=request_params, |
| 333 | step_idx=step_idx, |
| 334 | ) |
| 335 | else: |
| 336 | # No masking when min_p is 0 |
| 337 | if logits_for_token == -float("inf"): |
| 338 | _raise_error_invalid( |
| 339 | msg_suffix=f"Invalid: token {token_id} masked when min_p=0.0", |
| 340 | batch_index=batch_index, |
| 341 | request_params=request_params, |
| 342 | step_idx=step_idx, |
| 343 | ) |
| 344 | |
| 345 | |
| 346 | def _min_tokens_params(kwargs: dict) -> None: |
| 347 | """Min-tokens logitproc config.""" |
| 348 | kwargs["min_tokens"] = MIN_TOKENS_LEN_THRESHOLD |
| 349 | kwargs["stop_token_ids"] = [ |
| 350 | np.random.randint(0, VOCAB_SIZE - 1) |
| 351 | for _ in range(np.random.randint(0, VOCAB_SIZE)) |
| 352 | ] |
| 353 | |
| 354 | |
| 355 | def _min_tokens_validate( |
| 356 | test_fakes: LogitsprocsTestFakes, |
| 357 | persistent_batch: list[LogitsProcsRequestParams], |
| 358 | logits_new: torch.Tensor, |
| 359 | batch_index: int, |
| 360 | request_params: LogitsProcsRequestParams, |
| 361 | step_idx: int, |
| 362 | ) -> None: |
| 363 | """Validate min-tokens logitsproc applied correctly.""" |
| 364 | ref_num_out_tokens = len(request_params.out_tokens) |
| 365 | min_reached = ref_num_out_tokens >= MIN_TOKENS_LEN_THRESHOLD |
| 366 | ref_all_stop_token_ids = request_params.params.all_stop_token_ids |
| 367 | mt_lp = next(test_fakes.get_logitsprocs_by_cls(MinTokensLogitsProcessor)) |
| 368 | assert isinstance(mt_lp, MinTokensLogitsProcessor) |
| 369 | min_tok = mt_lp.min_toks.get(batch_index, None) |
| 370 | |
| 371 | # Validate min-token logits processor state |
| 372 | if min_tok: |
| 373 | (_, out_tok, all_stop_token_ids, _) = min_tok |
| 374 | num_out_tokens = len(out_tok) |
| 375 | if num_out_tokens != ref_num_out_tokens: |
| 376 | _raise_error_invalid( |
| 377 | msg_suffix=( |
| 378 | "Number of output tokens in min-token logit processor " |
| 379 | f"request metadata ({num_out_tokens}) does not match " |
| 380 | f"reference ({ref_num_out_tokens})." |
| 381 | ), |
| 382 | batch_index=batch_index, |
| 383 | request_params=request_params, |
| 384 | step_idx=step_idx, |
| 385 | ) |
| 386 | if ref_all_stop_token_ids != all_stop_token_ids: |
| 387 | _raise_error_invalid( |
| 388 | msg_suffix=( |
| 389 | "Stop token ids do not match reference; all_stop_token_ids: " |
| 390 | f"{sorted(all_stop_token_ids)}, ref_all_stop_token_ids: " |
| 391 | f"{sorted(ref_all_stop_token_ids)}" |
| 392 | ), |
| 393 | batch_index=batch_index, |
| 394 | request_params=request_params, |
| 395 | step_idx=step_idx, |
| 396 | ) |
| 397 | if min_reached: |
| 398 | _raise_error_invalid( |
| 399 | msg_suffix=( |
| 400 | "Expected min-tokens request with min reached, but batch " |
| 401 | "index is recognized by min-tokens logits processor." |
| 402 | ), |
| 403 | batch_index=batch_index, |
| 404 | request_params=request_params, |
| 405 | step_idx=step_idx, |
| 406 | err_cls=RuntimeError, |
| 407 | ) |
| 408 | |
| 409 | elif not min_reached: |
| 410 | _raise_error_invalid( |
| 411 | msg_suffix=( |
| 412 | "Expected min-tokens request with min not reached, but batch " |
| 413 | "index is not recognized by min-tokens logits processor." |
| 414 | ), |
| 415 | batch_index=batch_index, |
| 416 | request_params=request_params, |
| 417 | step_idx=step_idx, |
| 418 | err_cls=RuntimeError, |
| 419 | ) |
| 420 | |
| 421 | # Validate min-token logits |
| 422 | for token_id in range(VOCAB_SIZE): |
| 423 | logits_for_token = logits_new[batch_index][token_id] |
| 424 | if token_id in ref_all_stop_token_ids and not min_reached: |
| 425 | if logits_for_token != -float("inf"): |
| 426 | _raise_error_invalid( |
| 427 | msg_suffix=( |
| 428 | f"Token {token_id} is a stop token and " |
| 429 | "the sequence has not reached min length, " |
| 430 | "but the token is not masked " |
| 431 | f"(logit={logits_for_token})" |
| 432 | ), |
| 433 | batch_index=batch_index, |
| 434 | request_params=request_params, |
| 435 | step_idx=step_idx, |
| 436 | ) |
| 437 | else: |
| 438 | if logits_for_token == -float("inf"): |
| 439 | _raise_error_invalid( |
| 440 | msg_suffix=( |
| 441 | f"Token {token_id} should not be masked but " |
| 442 | f"is (output len={ref_num_out_tokens})" |
| 443 | ), |
| 444 | batch_index=batch_index, |
| 445 | request_params=request_params, |
| 446 | step_idx=step_idx, |
| 447 | ) |
| 448 | |
| 449 | |
| 450 | def _make_min_tokens_processor() -> MinTokensLogitsProcessor: |
| 451 | return MinTokensLogitsProcessor(VllmConfig(), torch.device("cpu"), False) |
| 452 | |
| 453 | |
| 454 | # Constructing the processor allocates pinned host memory, which creates a |
| 455 | # CUDA driver context even though the tensors stay on CPU. Run these in their |
| 456 | # own process so the pytest process can still fork for `test_logitsprocs`. |
| 457 | @create_new_process_for_each_test() |
| 458 | def test_min_tokens_keeps_all_masked_behavior_without_structured_output(): |
| 459 | processor = _make_min_tokens_processor() |
| 460 | processor.update_state( |
| 461 | BatchUpdate( |
| 462 | batch_size=1, |
| 463 | removed=(), |
| 464 | added=[ |
| 465 | ( |
| 466 | 0, |
| 467 | SamplingParams(min_tokens=2, stop_token_ids=[0]), |
| 468 | None, |
| 469 | [], |
| 470 | ) |
| 471 | ], |
| 472 | moved=(), |
| 473 | ) |
| 474 | ) |
| 475 | |
| 476 | logits = torch.full((1, 3), -float("inf")) |
| 477 | logits[0, 0] = 1.0 |
| 478 | |
| 479 | processor.apply(logits) |
| 480 | |
| 481 | assert torch.isneginf(logits).all() |
| 482 | |
| 483 | |
| 484 | @create_new_process_for_each_test() |
| 485 | def test_min_tokens_restores_all_masked_structured_output_stop_token(): |
| 486 | processor = _make_min_tokens_processor() |
| 487 | processor.update_state( |
| 488 | BatchUpdate( |
| 489 | batch_size=1, |
| 490 | removed=(), |
| 491 | added=[ |
| 492 | ( |
| 493 | 0, |
| 494 | SamplingParams( |
| 495 | min_tokens=2, |
| 496 | stop_token_ids=[0], |
| 497 | structured_outputs=StructuredOutputsParams( |
| 498 | json={"type": "boolean"} |
| 499 | ), |
| 500 | ), |
| 501 | None, |
| 502 | [], |
| 503 | ) |
| 504 | ], |
| 505 | moved=(), |
| 506 | ) |
| 507 | ) |
| 508 | |
| 509 | logits = torch.full((1, 3), -float("inf")) |
| 510 | logits[0, 0] = 1.0 |
| 511 | |
| 512 | processor.apply(logits) |
| 513 | |
| 514 | assert logits[0, 0] == 1.0 |
| 515 | assert torch.isneginf(logits[0, 1:]).all() |
| 516 | |
| 517 | |
| 518 | @create_new_process_for_each_test() |
| 519 | def test_min_tokens_restores_all_masked_structured_output_stop_token_spec_decode(): |
| 520 | processor = _make_min_tokens_processor() |
| 521 | processor.update_state( |
| 522 | BatchUpdate( |
| 523 | batch_size=1, |
| 524 | removed=(), |
| 525 | added=[ |
| 526 | ( |
| 527 | 0, |
| 528 | SamplingParams( |
| 529 | min_tokens=2, |
| 530 | stop_token_ids=[0, 1], |
| 531 | structured_outputs=StructuredOutputsParams( |
| 532 | json={"type": "boolean"} |
| 533 | ), |
| 534 | ), |
| 535 | None, |
| 536 | [], |
| 537 | ) |
| 538 | ], |
| 539 | moved=(), |
| 540 | ) |
| 541 | ) |
| 542 | |
| 543 | logits = torch.full((2, 4), -float("inf")) |
| 544 | logits[0, 0] = 1.0 |
| 545 | logits[0, 1] = 2.0 |
| 546 | logits[1, 0] = 3.0 |
| 547 | |
| 548 | processor.apply_with_spec_decode(logits, [2]) |
| 549 | |
| 550 | assert logits[0, 0] == 1.0 |
| 551 | assert logits[0, 1] == 2.0 |
| 552 | assert logits[1, 0] == 3.0 |
| 553 | assert torch.isneginf(logits[1, 1:]).all() |
| 554 | assert torch.isneginf(logits[:, 2:]).all() |
| 555 | |
| 556 | |
| 557 | def _thinking_budget_params(kwargs: dict) -> None: |
| 558 | """Set SamplingParams kwargs for thinking token budget tests.""" |
| 559 | kwargs["thinking_token_budget"] = THINKING_TOKEN_BUDGET |
| 560 | |
| 561 | |
| 562 | def _thinking_budget_validate( |
| 563 | test_fakes: LogitsprocsTestFakes, |
| 564 | persistent_batch: list[LogitsProcsRequestParams], |
| 565 | logits_new: torch.Tensor, |
| 566 | batch_index: int, |
| 567 | request_params: LogitsProcsRequestParams, |
| 568 | step_idx: int, |
| 569 | ) -> None: |
| 570 | """Validate ``ThinkingBudgetStateHolder`` thinking-budget behavior. |
| 571 | |
| 572 | State is keyed by **batch slot** (same index space as logits rows), matching |
| 573 | ``sync_batch`` / sampler integration (see PR #34668 discussion). |
| 574 | """ |
| 575 | holder = test_fakes.sampling_metadata.thinking_budget_state_holder |
| 576 | assert holder is not None |
| 577 | state = holder._state.get(batch_index) |
| 578 | params = request_params.params |
| 579 | |
| 580 | if hasattr(params, "thinking_token_budget") and params.thinking_token_budget: |
| 581 | if state is None: |
| 582 | _raise_error_invalid( |
| 583 | msg_suffix=( |
| 584 | f"Expected holder state for batch slot {batch_index} " |
| 585 | f"with thinking_token_budget={params.thinking_token_budget}" |
| 586 | ), |
| 587 | batch_index=batch_index, |
| 588 | request_params=request_params, |
| 589 | step_idx=step_idx, |
| 590 | ) |
| 591 | |
| 592 | expected_budget = params.thinking_token_budget |
| 593 | actual_budget = state["thinking_token_budget"] |
| 594 | if actual_budget != expected_budget: |
| 595 | _raise_error_invalid( |
| 596 | msg_suffix=( |
| 597 | f"Budget mismatch: expected {expected_budget}, got {actual_budget}" |
| 598 | ), |
| 599 | batch_index=batch_index, |
| 600 | request_params=request_params, |
| 601 | step_idx=step_idx, |
| 602 | ) |
| 603 | |
| 604 | output_tokens = request_params.out_tokens |
| 605 | start_tokens = holder.think_start_token_ids |
| 606 | thinking_started = False |
| 607 | if len(start_tokens) > 0: |
| 608 | for i in range(len(output_tokens) - len(start_tokens) + 1): |
| 609 | if output_tokens[i : i + len(start_tokens)] == start_tokens: |
| 610 | thinking_started = True |
| 611 | break |
| 612 | |
| 613 | if thinking_started: |
| 614 | think_count = state["think_count"] |
| 615 | budget = state["thinking_token_budget"] |
| 616 | if think_count >= budget and not state["in_end"]: |
| 617 | _raise_error_invalid( |
| 618 | msg_suffix=( |
| 619 | f"Budget exceeded ({think_count} >= {budget}) but " |
| 620 | "in_end is false" |
| 621 | ), |
| 622 | batch_index=batch_index, |
| 623 | request_params=request_params, |
| 624 | step_idx=step_idx, |
| 625 | ) |
| 626 | |
| 627 | end_tokens = holder.think_end_token_ids |
| 628 | if ( |
| 629 | think_count >= budget |
| 630 | and state["in_end"] |
| 631 | and len(end_tokens) > 0 |
| 632 | and holder.has_tracked_requests() |
| 633 | ): |
| 634 | expected_end_token_id = end_tokens[ |
| 635 | min(state["end_count"], len(end_tokens) - 1) |
| 636 | ] |
| 637 | # Holder bumps forced vocab positions to 1e9 (does not -inf others). |
| 638 | forced_logit = float(logits_new[batch_index, expected_end_token_id]) |
| 639 | if forced_logit < 1.0e8: |
| 640 | _raise_error_invalid( |
| 641 | msg_suffix=( |
| 642 | f"Expected forced end token {expected_end_token_id} " |
| 643 | f"with large logit, got {forced_logit}" |
| 644 | ), |
| 645 | batch_index=batch_index, |
| 646 | request_params=request_params, |
| 647 | step_idx=step_idx, |
| 648 | ) |
| 649 | |
| 650 | |
| 651 | def _none_validate( |
| 652 | test_fakes: LogitsprocsTestFakes, |
| 653 | persistent_batch: list[LogitsProcsRequestParams], |
| 654 | logits_new: torch.Tensor, |
| 655 | batch_index: int, |
| 656 | request_params: LogitsProcsRequestParams, |
| 657 | step_idx: int, |
| 658 | ) -> None: |
| 659 | """Validate that no logits processors are applied.""" |
| 660 | logits = test_fakes.logits[persistent_batch[batch_index].workload_index].cpu() |
| 661 | ref_logits = logits_new[batch_index] |
| 662 | if not torch.all(ref_logits == logits): |
| 663 | mismatch_toks = (ref_logits != logits).nonzero(as_tuple=True)[0].tolist() |
| 664 | mismatch_strs = [] |
| 665 | for token in mismatch_toks: |
| 666 | val = float(logits[token]) |
| 667 | ref_val = float(ref_logits[token]) |
| 668 | mismatch_strs.append(f"({token=},{val=},{ref_val=})") |
| 669 | _raise_error_invalid( |
| 670 | msg_suffix=( |
| 671 | f"Unexpected modification of logits: {','.join(mismatch_strs)}" |
| 672 | ), |
| 673 | batch_index=batch_index, |
| 674 | request_params=request_params, |
| 675 | step_idx=step_idx, |
| 676 | ) |
| 677 | |
| 678 | |
| 679 | class LogitsprocTestHelpers(NamedTuple): |
| 680 | """Supports setting up and validating logitsprocs unit tests.""" |
| 681 | |
| 682 | eval_fxn: Callable |
| 683 | gen_request_fxn: Callable | None = None |
| 684 | |
| 685 | |
| 686 | logitsprocs_test_mapping: dict[LogitprocType, LogitsprocTestHelpers] = { |
| 687 | STR_NO_LOGITPROC: LogitsprocTestHelpers(eval_fxn=_none_validate), |
| 688 | LogitBiasLogitsProcessor: LogitsprocTestHelpers( |
| 689 | gen_request_fxn=_logit_bias_params, eval_fxn=_logit_bias_validate |
| 690 | ), |
| 691 | MinPLogitsProcessor: LogitsprocTestHelpers( |
| 692 | gen_request_fxn=_min_p_params, eval_fxn=_min_p_validate |
| 693 | ), |
| 694 | MinTokensLogitsProcessor: LogitsprocTestHelpers( |
| 695 | gen_request_fxn=_min_tokens_params, eval_fxn=_min_tokens_validate |
| 696 | ), |
| 697 | STR_THINKING_BUDGET: LogitsprocTestHelpers( |
| 698 | gen_request_fxn=_thinking_budget_params, eval_fxn=_thinking_budget_validate |
| 699 | ), |
| 700 | } |
| 701 | |
| 702 | |
| 703 | def _get_test_cases() -> list[list[LogitprocType]]: |
| 704 | """Each test case is a set of logitsprocs.""" |
| 705 | logitsprocs_types: list[LogitprocType] = list(logitsprocs_test_mapping.keys()) |
| 706 | |
| 707 | # Isolate thinking-budget handling from other processors to avoid cross-talk. |
| 708 | thinking_id: LogitprocType = STR_THINKING_BUDGET |
| 709 | other_processors: list[LogitprocType] = [ |
| 710 | p for p in logitsprocs_types if p != STR_NO_LOGITPROC and p != thinking_id |
| 711 | ] |
| 712 | |
| 713 | test_cases: list[list[LogitprocType]] = [[STR_NO_LOGITPROC]] |
| 714 | test_cases.extend( |
| 715 | [logitproc_type, STR_NO_LOGITPROC] for logitproc_type in other_processors |
| 716 | ) |
| 717 | test_cases.extend([other_processors, [thinking_id]]) |
| 718 | return test_cases |
| 719 | |
| 720 | |
| 721 | def _generate_fake_step_update( |
| 722 | persistent_batch: list[LogitsProcsRequestParams], |
| 723 | workload_params: list[LogitsProcsRequestParams], |
| 724 | wdx: int, |
| 725 | batch_update_builder: BatchUpdateBuilder, |
| 726 | ) -> tuple[BatchUpdate | None, int, int]: |
| 727 | batch_size = len(persistent_batch) |
| 728 | workload_size = len(workload_params) |
| 729 | workload_reqs_remaining = workload_size - wdx |
| 730 | max_add_remove_per_step = max(1, int(0.2 * workload_size)) |
| 731 | |
| 732 | # 50% of steps: add no reqs |
| 733 | # Other 50%: add a limited number of reqs (less than the number |
| 734 | # of workload reqs remaining, less than an arbitrary max) |
| 735 | # If no workload reqs remain: 100% of steps have 0 adds |
| 736 | num_step_add = ( |
| 737 | random.choice( |
| 738 | [ |
| 739 | 0, |
| 740 | random.randint( |
| 741 | 1, min(max_add_remove_per_step, workload_reqs_remaining) |
| 742 | ), |
| 743 | ] |
| 744 | ) |
| 745 | if workload_reqs_remaining |
| 746 | else 0 |
| 747 | ) |
| 748 | |
| 749 | # 50% of steps: remove no requests |
| 750 | # Other 50%: remove a limited number of reqs (less than the number |
| 751 | # persistent batch reqs remaining, less than an arbitrary max) |
| 752 | # If persistent batch is empty: 100% of steps have 0 removals until |
| 753 | # more requests are added. Assume that removed requests are always |
| 754 | # drawn from the current batch, before new adds |
| 755 | num_step_remove = ( |
| 756 | random.choice([0, random.randint(1, min(max_add_remove_per_step, batch_size))]) |
| 757 | if batch_size |
| 758 | else 0 |
| 759 | ) |
| 760 | |
| 761 | num_step_add_replace = min(num_step_add, num_step_remove) |
| 762 | |
| 763 | # Generate fake removed request indices drawn from persistent batch indices |
| 764 | for removal in random.sample(range(batch_size), num_step_remove): |
| 765 | batch_update_builder.removed_append(removal) |
| 766 | |
| 767 | # Get added requests from workload |
| 768 | for add_req_params in workload_params[wdx : (wdx + num_step_add_replace)]: |
| 769 | # Replace as many removed requests as possible with added requests |
| 770 | add_remove_idx = batch_update_builder.pop_removed() |
| 771 | assert add_remove_idx is not None |
| 772 | batch_update_builder.added.append( |
| 773 | ( |
| 774 | add_remove_idx, |
| 775 | add_req_params.params, |
| 776 | add_req_params.prompt_tokens, |
| 777 | add_req_params.out_tokens, |
| 778 | ) |
| 779 | ) |
| 780 | persistent_batch[add_remove_idx] = add_req_params |
| 781 | |
| 782 | # Append remaining added requests to end of batch |
| 783 | add_reqs_append = workload_params[ |
| 784 | (wdx + num_step_add_replace) : (wdx + num_step_add) |
| 785 | ] |
| 786 | batch_update_builder.added.extend( |
| 787 | [ |
| 788 | ( |
| 789 | adx + batch_size, |
| 790 | add_req_params.params, |
| 791 | add_req_params.prompt_tokens, |
| 792 | add_req_params.out_tokens, |
| 793 | ) |
| 794 | for adx, add_req_params in enumerate(add_reqs_append) |
| 795 | ] |
| 796 | ) |
| 797 | persistent_batch.extend(add_reqs_append) |
| 798 | pre_condense_batch_size = len(persistent_batch) |
| 799 | wdx += num_step_add # Update workload offset |
| 800 | |
| 801 | # Simulate condensing persistent batch |
| 802 | last_nonempty_index = pre_condense_batch_size - 1 |
| 803 | condensed_to_idxs = set() |
| 804 | while batch_update_builder.removed: |
| 805 | if ( |
| 806 | last_nonempty_index in batch_update_builder.removed |
| 807 | or last_nonempty_index in condensed_to_idxs |
| 808 | ): |
| 809 | last_nonempty_index -= 1 |
| 810 | continue |
| 811 | # last_nonempty_index is the highest persistent batch index that was |
| 812 | # not removed |
| 813 | first_empty_index = batch_update_builder.peek_removed() |
| 814 | assert first_empty_index is not None |
| 815 | if first_empty_index > last_nonempty_index: |
| 816 | break |
| 817 | # first_empty_index is the lowest removed persistent batch index |
| 818 | # that is less than last_nonempty_index |
| 819 | # |
| 820 | # move last_nonempty_index -> first_empty_index |
| 821 | batch_update_builder.pop_removed() |
| 822 | condensed_to_idxs.add(first_empty_index) |
| 823 | persistent_batch[first_empty_index] = persistent_batch[last_nonempty_index] |
| 824 | batch_update_builder.moved.append( |
| 825 | (last_nonempty_index, first_empty_index, MoveDirectionality.UNIDIRECTIONAL) |
| 826 | ) |
| 827 | |
| 828 | last_nonempty_index -= 1 |
| 829 | |
| 830 | # Now removed requests & gaps left by non-removed requests that got |
| 831 | # moved downward are grouped consecutively in the upper indices of |
| 832 | # the persistent batch. Truncate them to get condensed persistent batch |
| 833 | condensed_batch_size = batch_size + num_step_add - num_step_remove |
| 834 | persistent_batch[:] = persistent_batch[0:condensed_batch_size] |
| 835 | |
| 836 | if condensed_batch_size > 1: |
| 837 | # Simulate arbitrary batch ordering in the kernel backend |
| 838 | # Generate a random number k of non-overlapping swap tuples |
| 839 | k = random.randint(0, condensed_batch_size // 2) |
| 840 | idxs = list(range(condensed_batch_size)) |
| 841 | random.shuffle(idxs) |
| 842 | swaps = [tuple(sorted([idxs[2 * i], idxs[2 * i + 1]])) for i in range(k)] |
| 843 | batch_update_builder.moved.extend( |
| 844 | [(sw[0], sw[1], MoveDirectionality.SWAP) for sw in swaps] |
| 845 | ) |
| 846 | for adx, bdx in swaps: |
| 847 | persistent_batch[adx], persistent_batch[bdx] = ( |
| 848 | persistent_batch[bdx], |
| 849 | persistent_batch[adx], |
| 850 | ) |
| 851 | |
| 852 | return ( |
| 853 | batch_update_builder.get_and_reset(condensed_batch_size), |
| 854 | wdx, |
| 855 | workload_size - wdx, |
| 856 | ) |
| 857 | |
| 858 | |
| 859 | def _assert_valid( |
| 860 | batch_size: int, |
| 861 | persistent_batch: list[LogitsProcsRequestParams], |
| 862 | test_fakes: LogitsprocsTestFakes, |
| 863 | slice_idxs: list[int], |
| 864 | logits_w_lp: torch.Tensor, |
| 865 | step_idx: int, |
| 866 | ) -> None: |
| 867 | if not slice_idxs: |
| 868 | # Trivial case of empty persistent batch |
| 869 | assert len(persistent_batch) == 0 |
| 870 | if logits_w_lp.shape[0] != 0: |
| 871 | raise ValueError( |
| 872 | "Fake persistent batch is empty but logitsprocs " |
| 873 | f"output batch has shape {logits_w_lp.shape}" |
| 874 | ) |
| 875 | return |
| 876 | |
| 877 | # Validate logits for each fake request |
| 878 | for batch_index in range(batch_size): |
| 879 | request_params = persistent_batch[batch_index] |
| 880 | # Invoke the appropriate validation function for |
| 881 | # the logitproc employed by this request |
| 882 | fxn = logitsprocs_test_mapping[request_params.logitproc_type].eval_fxn |
| 883 | fxn( |
| 884 | test_fakes=test_fakes, |
| 885 | persistent_batch=persistent_batch, |
| 886 | logits_new=logits_w_lp, |
| 887 | batch_index=batch_index, |
| 888 | request_params=request_params, |
| 889 | step_idx=step_idx, |
| 890 | ) |
| 891 | |
| 892 | |
| 893 | def _slot_outputs_for_metadata( |
| 894 | persistent_batch: list[LogitsProcsRequestParams], pad_len: int |
| 895 | ) -> list[list[int]]: |
| 896 | """Per-batch-slot output token ids aligned with ``SamplingMetadata`` rows.""" |
| 897 | rows: list[list[int]] = [[] for _ in range(pad_len)] |
| 898 | for i, req in enumerate(persistent_batch): |
| 899 | if i < pad_len: |
| 900 | rows[i] = list(req.out_tokens) |
| 901 | return rows |
| 902 | |
| 903 | |
| 904 | @create_new_process_for_each_test() |
| 905 | @pytest.mark.parametrize("device", DEVICES) |
| 906 | @pytest.mark.parametrize("reqs_per_logitproc", [REQS_PER_LOGITPROC]) |
| 907 | @pytest.mark.parametrize("logitsprocs_under_test", _get_test_cases()) |
| 908 | def test_logitsprocs( |
| 909 | device: str, reqs_per_logitproc: int, logitsprocs_under_test: list[LogitprocType] |
| 910 | ): |
| 911 | random.seed(40) |
| 912 | torch.set_default_device(device) |
| 913 | |
| 914 | # Define a shuffled batch of requests which individually use a different |
| 915 | # logitproc, or no logitproc at all |
| 916 | workload_params = _generate_mixed_logitsprocs_batch_params( |
| 917 | reqs_per_logitproc=reqs_per_logitproc, logitsprocs_types=logitsprocs_under_test |
| 918 | ) |
| 919 | workload_size = len(workload_params) |
| 920 | |
| 921 | # Create fake test data structures for testing. |
| 922 | test_fakes = _generate_test_fakes(workload_size, device) |
| 923 | |
| 924 | wdx = 0 # Next request index in workload to add |
| 925 | persistent_batch: list[ |
| 926 | LogitsProcsRequestParams |
| 927 | ] = [] # Persistent batch state, as list of workload indices |
| 928 | |
| 929 | # Generate fake removed request indices from current persistent |
| 930 | # batch before adds |
| 931 | batch_update_builder = BatchUpdateBuilder() |
| 932 | |
| 933 | # Break when entire workload has been added previously and persistent |
| 934 | # batch is empty |
| 935 | workload_reqs_remaining = workload_size |
| 936 | batch_size = 0 |
| 937 | step_idx = 0 |
| 938 | while True: |
| 939 | if not (workload_reqs_remaining or batch_size): |
| 940 | break |
| 941 | |
| 942 | ( |
| 943 | batch_update, |
| 944 | wdx, |
| 945 | workload_reqs_remaining, |
| 946 | ) = _generate_fake_step_update( |
| 947 | persistent_batch=persistent_batch, |
| 948 | workload_params=workload_params, |
| 949 | wdx=wdx, |
| 950 | batch_update_builder=batch_update_builder, |
| 951 | ) |
| 952 | batch_size = len(persistent_batch) |
| 953 | |
| 954 | # Apply fake batch update to logitsprocs |
| 955 | fake_update_logitsprocs_state(test_fakes, batch_update) |
| 956 | |
| 957 | # Emulate application of logits processors + thinking holder (sampler order). |
| 958 | slice_idxs = [req.workload_index for req in persistent_batch] |
| 959 | slot_rows = _slot_outputs_for_metadata(persistent_batch, workload_size) |
| 960 | logits_w_lp = fake_apply_logitsprocs(test_fakes, slice_idxs, slot_rows).cpu() |
| 961 | |
| 962 | _assert_valid( |
| 963 | batch_size=batch_size, |
| 964 | persistent_batch=persistent_batch, |
| 965 | test_fakes=test_fakes, |
| 966 | slice_idxs=slice_idxs, |
| 967 | logits_w_lp=logits_w_lp, |
| 968 | step_idx=step_idx, |
| 969 | ) |
| 970 | |
| 971 | step_idx += 1 |
| 972 | |
| 973 | |
| 974 | def test_maybe_create_thinking_budget_holder_without_reasoning(): |
| 975 | cfg = VllmConfig() |
| 976 | assert cfg.reasoning_config is None |
| 977 | assert ( |
| 978 | maybe_create_thinking_budget_state_holder( |
| 979 | None, |
| 980 | cfg.scheduler_config.max_num_seqs, |
| 981 | 0, |
| 982 | torch.device("cpu"), |
| 983 | is_pin_memory=False, |
| 984 | ) |
| 985 | is None |
| 986 | ) |
| 987 | |
| 988 | |
| 989 | def test_thinking_budget_holder_has_tracked_after_sync_add(): |
| 990 | vc = VllmConfig() |
| 991 | vc.reasoning_config = create_mock_reasoning_config( |
| 992 | [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID] |
| 993 | ) |
| 994 | h = ThinkingBudgetStateHolder( |
| 995 | vc.reasoning_config, |
| 996 | vc.scheduler_config.max_num_seqs, |
| 997 | 0, |
| 998 | torch.device("cpu"), |
| 999 | False, |
| 1000 | ) |
| 1001 | assert not h.has_tracked_requests() |
| 1002 | h.sync_batch( |
| 1003 | BatchUpdate( |
| 1004 | batch_size=1, |
| 1005 | removed=(), |
| 1006 | added=[ |
| 1007 | ( |
| 1008 | 0, |
| 1009 | SamplingParams(thinking_token_budget=3), |
| 1010 | None, |
| 1011 | [THINK_START_TOKEN_ID], |
| 1012 | ) |
| 1013 | ], |
| 1014 | moved=(), |
| 1015 | ) |
| 1016 | ) |
| 1017 | assert h.has_tracked_requests() |
| 1018 | assert h._state[0]["thinking_token_budget"] == 3 |
| 1019 | |
| 1020 | |
| 1021 | def test_thinking_budget_holder_sync_remove_clears_state(): |
| 1022 | vc = VllmConfig() |
| 1023 | vc.reasoning_config = create_mock_reasoning_config( |
| 1024 | [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID] |
| 1025 | ) |
| 1026 | h = ThinkingBudgetStateHolder( |
| 1027 | vc.reasoning_config, |
| 1028 | vc.scheduler_config.max_num_seqs, |
| 1029 | 0, |
| 1030 | torch.device("cpu"), |
| 1031 | False, |
| 1032 | ) |
| 1033 | h.sync_batch( |
| 1034 | BatchUpdate( |
| 1035 | batch_size=1, |
| 1036 | removed=(), |
| 1037 | added=[ |
| 1038 | ( |
| 1039 | 0, |
| 1040 | SamplingParams(thinking_token_budget=3), |
| 1041 | None, |
| 1042 | [], |
| 1043 | ) |
| 1044 | ], |
| 1045 | moved=(), |
| 1046 | ) |
| 1047 | ) |
| 1048 | assert h.has_tracked_requests() |
| 1049 | h.sync_batch(BatchUpdate(batch_size=0, removed=(0,), added=(), moved=())) |
| 1050 | assert not h.has_tracked_requests() |
| 1051 | |
| 1052 | |
| 1053 | def test_thinking_budget_holder_sync_add_without_budget_drops_row(): |
| 1054 | vc = VllmConfig() |
| 1055 | vc.reasoning_config = create_mock_reasoning_config( |
| 1056 | [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID] |
| 1057 | ) |
| 1058 | h = ThinkingBudgetStateHolder( |
| 1059 | vc.reasoning_config, |
| 1060 | vc.scheduler_config.max_num_seqs, |
| 1061 | 0, |
| 1062 | torch.device("cpu"), |
| 1063 | False, |
| 1064 | ) |
| 1065 | h.sync_batch( |
| 1066 | BatchUpdate( |
| 1067 | batch_size=1, |
| 1068 | removed=(), |
| 1069 | added=[(0, SamplingParams(), None, [])], |
| 1070 | moved=(), |
| 1071 | ) |
| 1072 | ) |
| 1073 | assert not h.has_tracked_requests() |
| 1074 | |
| 1075 | |
| 1076 | def test_thinking_budget_holder_swap_exchanges_state(): |
| 1077 | vc = VllmConfig() |
| 1078 | vc.reasoning_config = create_mock_reasoning_config( |
| 1079 | [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID] |
| 1080 | ) |
| 1081 | h = ThinkingBudgetStateHolder( |
| 1082 | vc.reasoning_config, |
| 1083 | vc.scheduler_config.max_num_seqs, |
| 1084 | 0, |
| 1085 | torch.device("cpu"), |
| 1086 | False, |
| 1087 | ) |
| 1088 | h.sync_batch( |
| 1089 | BatchUpdate( |
| 1090 | batch_size=2, |
| 1091 | removed=(), |
| 1092 | added=[ |
| 1093 | ( |
| 1094 | 0, |
| 1095 | SamplingParams(thinking_token_budget=3), |
| 1096 | None, |
| 1097 | [], |
| 1098 | ), |
| 1099 | ( |
| 1100 | 1, |
| 1101 | SamplingParams(thinking_token_budget=7), |
| 1102 | None, |
| 1103 | [], |
| 1104 | ), |
| 1105 | ], |
| 1106 | moved=(), |
| 1107 | ) |
| 1108 | ) |
| 1109 | b0, b1 = h._state[0]["thinking_token_budget"], h._state[1]["thinking_token_budget"] |
| 1110 | h.sync_batch( |
| 1111 | BatchUpdate( |
| 1112 | batch_size=2, |
| 1113 | removed=(), |
| 1114 | added=(), |
| 1115 | moved=[(0, 1, MoveDirectionality.SWAP)], |
| 1116 | ) |
| 1117 | ) |
| 1118 | assert h._state[0]["thinking_token_budget"] == b1 |
| 1119 | assert h._state[1]["thinking_token_budget"] == b0 |
| 1120 | |
| 1121 | |
| 1122 | def test_thinking_budget_holder_unidirectional_move(): |
| 1123 | vc = VllmConfig() |
| 1124 | vc.reasoning_config = create_mock_reasoning_config( |
| 1125 | [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID] |
| 1126 | ) |
| 1127 | h = ThinkingBudgetStateHolder( |
| 1128 | vc.reasoning_config, |
| 1129 | vc.scheduler_config.max_num_seqs, |
| 1130 | 0, |
| 1131 | torch.device("cpu"), |
| 1132 | False, |
| 1133 | ) |
| 1134 | h.sync_batch( |
| 1135 | BatchUpdate( |
| 1136 | batch_size=2, |
| 1137 | removed=(), |
| 1138 | added=[ |
| 1139 | ( |
| 1140 | 1, |
| 1141 | SamplingParams(thinking_token_budget=4), |
| 1142 | None, |
| 1143 | [], |
| 1144 | ), |
| 1145 | ], |
| 1146 | moved=(), |
| 1147 | ) |
| 1148 | ) |
| 1149 | assert 1 in h._state and 0 not in h._state |
| 1150 | h.sync_batch( |
| 1151 | BatchUpdate( |
| 1152 | batch_size=2, |
| 1153 | removed=(), |
| 1154 | added=(), |
| 1155 | moved=[(1, 0, MoveDirectionality.UNIDIRECTIONAL)], |
| 1156 | ) |
| 1157 | ) |
| 1158 | assert 0 in h._state and 1 not in h._state |
| 1159 | assert h._state[0]["thinking_token_budget"] == 4 |
| 1160 | |
| 1161 | |
| 1162 | def test_thinking_budget_holder_update_state_repeat_indices_last_row_wins(): |
| 1163 | vc = VllmConfig() |
| 1164 | vc.reasoning_config = create_mock_reasoning_config( |
| 1165 | [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID] |
| 1166 | ) |
| 1167 | h = ThinkingBudgetStateHolder( |
| 1168 | vc.reasoning_config, |
| 1169 | vc.scheduler_config.max_num_seqs, |
| 1170 | 0, |
| 1171 | torch.device("cpu"), |
| 1172 | False, |
| 1173 | ) |
| 1174 | h.sync_batch( |
| 1175 | BatchUpdate( |
| 1176 | batch_size=1, |
| 1177 | removed=(), |
| 1178 | added=[ |
| 1179 | ( |
| 1180 | 0, |
| 1181 | SamplingParams(thinking_token_budget=5), |
| 1182 | None, |
| 1183 | [THINK_START_TOKEN_ID], |
| 1184 | ) |
| 1185 | ], |
| 1186 | moved=(), |
| 1187 | ) |
| 1188 | ) |
| 1189 | out_lists = [[THINK_START_TOKEN_ID], [THINK_START_TOKEN_ID, 10, 11, 12, 13, 14]] |
| 1190 | h.update_state( |
| 1191 | out_lists, |
| 1192 | None, |
| 1193 | torch.tensor([0, 0], dtype=torch.long), |
| 1194 | ) |
| 1195 | assert h._state[0]["output_tok_ids"] == out_lists[1] |
| 1196 | |
| 1197 | |
| 1198 | def test_thinking_budget_holder_spec_mode_tensor_layout(): |
| 1199 | h = ThinkingBudgetStateHolder( |
| 1200 | create_mock_reasoning_config([THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]), |
| 1201 | 8, |
| 1202 | 2, |
| 1203 | torch.device("cpu"), |
| 1204 | False, |
| 1205 | ) |
| 1206 | assert h.in_spec_mode |
| 1207 | assert h._mask_capacity == 8 * (2 + 1) |
| 1208 | |
| 1209 | |
| 1210 | def test_thinking_budget_holder_empty_end_tokens_disables_row(): |
| 1211 | vc = VllmConfig() |
| 1212 | vc.reasoning_config = create_mock_reasoning_config([THINK_START_TOKEN_ID], []) |
| 1213 | h = ThinkingBudgetStateHolder( |
| 1214 | vc.reasoning_config, |
| 1215 | vc.scheduler_config.max_num_seqs, |
| 1216 | 0, |
| 1217 | torch.device("cpu"), |
| 1218 | False, |
| 1219 | ) |
| 1220 | h.sync_batch( |
| 1221 | BatchUpdate( |
| 1222 | batch_size=1, |
| 1223 | removed=(), |
| 1224 | added=[ |
| 1225 | ( |
| 1226 | 0, |
| 1227 | SamplingParams(thinking_token_budget=5), |
| 1228 | None, |
| 1229 | [THINK_START_TOKEN_ID], |
| 1230 | ) |
| 1231 | ], |
| 1232 | moved=(), |
| 1233 | ) |
| 1234 | ) |
| 1235 | h.update_state([[THINK_START_TOKEN_ID, 1]], None, None) |
| 1236 | assert h._state[0]["thinking_token_budget"] == -1 |
| 1237 | |
| 1238 | |
| 1239 | def test_thinking_budget_enforced_without_penalties(): |
| 1240 | """Regression test for gpu_input_batch.py bug. |
| 1241 | |
| 1242 | When thinking_budget_tracks_reqs=True and no penalties/bad_words are set, |
| 1243 | the old code computed needs_output_token_ids=False (inverted condition: |
| 1244 | ``or not thinking_budget_tracks_reqs``), causing update_state to receive |
| 1245 | an empty list and skip _update_think_state for every request. |
| 1246 | |
| 1247 | Fix: changed ``or not thinking_budget_tracks_reqs`` to |
| 1248 | ``or thinking_budget_tracks_reqs`` so that output_token_ids is populated |
| 1249 | whenever the thinking budget state holder has tracked requests. |
| 1250 | |
| 1251 | This test verifies that update_state correctly calls _update_think_state |
| 1252 | (setting in_end=True) when given the real output_token_ids, and that |
| 1253 | passing an empty list (the pre-fix behavior) prevents budget enforcement. |
| 1254 | """ |
| 1255 | vc = VllmConfig() |
| 1256 | vc.reasoning_config = create_mock_reasoning_config( |
| 1257 | [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID] |
| 1258 | ) |
| 1259 | budget = 3 # allow 3 thinking tokens |
| 1260 | |
| 1261 | h = ThinkingBudgetStateHolder( |
| 1262 | vc.reasoning_config, |
| 1263 | vc.scheduler_config.max_num_seqs, |
| 1264 | 0, |
| 1265 | torch.device("cpu"), |
| 1266 | False, |
| 1267 | ) |
| 1268 | output_token_ids: list[int] = [] |
| 1269 | h.sync_batch( |
| 1270 | BatchUpdate( |
| 1271 | batch_size=1, |
| 1272 | removed=(), |
| 1273 | added=[ |
| 1274 | ( |
| 1275 | 0, |
| 1276 | SamplingParams(thinking_token_budget=budget), |
| 1277 | None, |
| 1278 | output_token_ids, |
| 1279 | ) |
| 1280 | ], |
| 1281 | moved=(), |
| 1282 | ) |
| 1283 | ) |
| 1284 | assert h.has_tracked_requests() |
| 1285 | |
| 1286 | # Simulate the buggy behavior: update_state receives empty list. |
| 1287 | # _update_think_state is skipped → in_end stays False → no budget enforcement. |
| 1288 | h.update_state([], None, None) |
| 1289 | assert not h._state[0].get("in_end", False), ( |
| 1290 | "With empty output_token_ids, in_end should stay False (budget not yet tracked)" |
| 1291 | ) |
| 1292 | |
| 1293 | # Simulate the correct behavior: output_token_ids is the live list. |
| 1294 | # Step 1: think-start token appears. |
| 1295 | output_token_ids.append(THINK_START_TOKEN_ID) |
| 1296 | h.update_state([output_token_ids], None, None) |
| 1297 | assert not h._state[0].get("in_end", False), ( |
| 1298 | "Still within budget after 0 think tokens" |
| 1299 | ) |
| 1300 | |
| 1301 | # Steps 2–4: 3 thinking tokens (hits the budget exactly). |
| 1302 | for tok in [1, 2, 3]: |
| 1303 | output_token_ids.append(tok) |
| 1304 | h.update_state([output_token_ids], None, None) |
| 1305 | |
| 1306 | # After exactly `budget` thinking tokens the holder should force end token. |
| 1307 | assert h._state[0].get("in_end", False), ( |
| 1308 | "Budget exceeded: in_end should be True so that apply_to_logits " |
| 1309 | "forces the end token" |
| 1310 | ) |
| 1311 | |
| 1312 | |
| 1313 | @pytest.mark.parametrize( |
| 1314 | ("raw_value", "expected"), |
| 1315 | [ |
| 1316 | (None, None), |
| 1317 | (-1, None), |
| 1318 | (10, 10), |
| 1319 | (0, 0), |
| 1320 | ], |
| 1321 | ) |
| 1322 | def test_validate_thinking_token_budget(raw_value, expected): |
| 1323 | assert validate_thinking_token_budget(raw_value) == expected |
| 1324 | |
| 1325 | |
| 1326 | def test_sampling_params_minus_one_normalizes_to_none(): |
| 1327 | params = SamplingParams(thinking_token_budget=-1) |
| 1328 | assert params.thinking_token_budget is None |
| 1329 | |
| 1330 | |
| 1331 | @pytest.mark.parametrize("invalid_budget", [-2, 0.6, 10.5, True]) |
| 1332 | def test_validate_thinking_token_budget_rejects_invalid(invalid_budget): |
| 1333 | from vllm.exceptions import VLLMValidationError |
| 1334 | |
| 1335 | with pytest.raises(VLLMValidationError, match="thinking_token_budget"): |
| 1336 | validate_thinking_token_budget(invalid_budget) |
| 1337 | |
| 1338 | |
| 1339 | @pytest.mark.parametrize("invalid_budget", [-2, 0.6, 10.5]) |
| 1340 | def test_thinking_budget_invalid_budget_rejected(invalid_budget): |
| 1341 | from vllm.exceptions import VLLMValidationError |
| 1342 | |
| 1343 | with pytest.raises(VLLMValidationError, match="thinking_token_budget"): |
| 1344 | SamplingParams(thinking_token_budget=invalid_budget) |
| 1345 | |
| 1346 | |
| 1347 | def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index(): |
| 1348 | """Test thinking budget enforced for a long thinking run, |
| 1349 | then a natural end marker.""" |
| 1350 | h = ThinkingBudgetStateHolder( |
| 1351 | create_mock_reasoning_config([THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]), |
| 1352 | 8, |
| 1353 | 0, |
| 1354 | torch.device("cpu"), |
| 1355 | False, |
| 1356 | ) |
| 1357 | h.sync_batch( |
| 1358 | BatchUpdate( |
| 1359 | batch_size=1, |
| 1360 | removed=(), |
| 1361 | added=[(0, SamplingParams(thinking_token_budget=10_000), None, [])], |
| 1362 | moved=(), |
| 1363 | ) |
| 1364 | ) |
| 1365 | start = [THINK_START_TOKEN_ID] |
| 1366 | end = [THINK_END_TOKEN_ID] |
| 1367 | |
| 1368 | out: list[int] = list(start) |
| 1369 | h.update_state([out], None, None) |
| 1370 | for tok in range(500): # 500 filler thinking tokens, one decode step each |
| 1371 | out.append(tok) |
| 1372 | h.update_state([out], None, None) |
| 1373 | assert h._state[0]["end_thinking"] == -1 # not present yet |
| 1374 | out.extend(end) |
| 1375 | h.update_state([out], None, None) |
| 1376 | |
| 1377 | # After a natural exit, start_thinking and end_thinking are reset to -1 |
| 1378 | # (state machine prepared for next block). Verify the exit happened at the |
| 1379 | # correct position by checking scan_offset (set to len(output) after exit). |
| 1380 | assert not h._state[0]["in_think"], "Should have exited think mode after </think>" |
| 1381 | assert h._state[0]["start_thinking"] == -1, ( |
| 1382 | "start_thinking should be reset after natural exit" |
| 1383 | ) |
| 1384 | assert h._state[0]["scan_offset"] == len(out), ( |
| 1385 | f"scan_offset should point past the end marker; " |
| 1386 | f"expected {len(out)}, got {h._state[0]['scan_offset']}" |
| 1387 | ) |
| 1388 | |
| 1389 | |
| 1390 | # --- Thinking budget re-entry tests (issue #43708) --- |
| 1391 | # Regression tests: after budget forces end-of-thinking token sequence, |
| 1392 | # the state machine must detect and enforce budget on subsequent blocks. |
| 1393 | |
| 1394 | |
| 1395 | class TestThinkingBudgetReentry: |
| 1396 | THINK_START = 100 |
| 1397 | THINK_END_SINGLE = [200] |
| 1398 | THINK_END_MULTI = [200, 201, 202] |
| 1399 | BUDGET = 5 |
| 1400 | CONTENT_TOKEN = 50 |
| 1401 | THINK_TOKEN = 60 |
| 1402 | |
| 1403 | @staticmethod |
| 1404 | def _make_holder(end_token_ids: list[int]) -> ThinkingBudgetStateHolder: |
| 1405 | cfg = create_mock_reasoning_config( |
| 1406 | [TestThinkingBudgetReentry.THINK_START], end_token_ids |
| 1407 | ) |
| 1408 | return ThinkingBudgetStateHolder( |
| 1409 | reasoning_config=cfg, |
| 1410 | max_num_seqs=8, |
| 1411 | num_spec_tokens=0, |
| 1412 | device=torch.device("cpu"), |
| 1413 | is_pin_memory=False, |
| 1414 | ) |
| 1415 | |
| 1416 | @staticmethod |
| 1417 | def _sync_batch(holder: ThinkingBudgetStateHolder, budget: int) -> None: |
| 1418 | holder.sync_batch( |
| 1419 | BatchUpdate( |
| 1420 | batch_size=1, |
| 1421 | removed=(), |
| 1422 | added=[(0, SamplingParams(thinking_token_budget=budget), None, [])], |
| 1423 | moved=(), |
| 1424 | ) |
| 1425 | ) |
| 1426 | |
| 1427 | @staticmethod |
| 1428 | def _step(holder: ThinkingBudgetStateHolder, output_tok_ids: list[int]) -> None: |
| 1429 | holder.update_state( |
| 1430 | output_token_ids=[output_tok_ids], |
| 1431 | spec_token_ids=None, |
| 1432 | repeat_indices=None, |
| 1433 | ) |
| 1434 | |
| 1435 | def _exhaust_budget(self, holder: ThinkingBudgetStateHolder) -> list[int]: |
| 1436 | output = [self.THINK_START] |
| 1437 | self._step(holder, list(output)) |
| 1438 | for _ in range(self.BUDGET): |
| 1439 | output.append(self.THINK_TOKEN) |
| 1440 | self._step(holder, list(output)) |
| 1441 | |
| 1442 | assert holder._state[0]["in_end"] |
| 1443 | return output |
| 1444 | |
| 1445 | def _accept_end_tokens( |
| 1446 | self, |
| 1447 | holder: ThinkingBudgetStateHolder, |
| 1448 | output: list[int], |
| 1449 | end_token_ids: list[int], |
| 1450 | ) -> None: |
| 1451 | for tok in end_token_ids: |
| 1452 | output.append(tok) |
| 1453 | self._step(holder, list(output)) |
| 1454 | |
| 1455 | def test_single_token_end_reentry(self): |
| 1456 | holder = self._make_holder(self.THINK_END_SINGLE) |
| 1457 | self._sync_batch(holder, self.BUDGET) |
| 1458 | |
| 1459 | output = self._exhaust_budget(holder) |
| 1460 | self._accept_end_tokens(holder, output, self.THINK_END_SINGLE) |
| 1461 | |
| 1462 | for _ in range(3): |
| 1463 | output.append(self.CONTENT_TOKEN) |
| 1464 | self._step(holder, list(output)) |
| 1465 | |
| 1466 | output.append(self.THINK_START) |
| 1467 | self._step(holder, list(output)) |
| 1468 | for _ in range(self.BUDGET): |
| 1469 | output.append(self.THINK_TOKEN) |
| 1470 | self._step(holder, list(output)) |
| 1471 | |
| 1472 | assert holder._state[0]["in_end"], ( |
| 1473 | "Second thinking block must also be budget-enforced" |
| 1474 | ) |
| 1475 | |
| 1476 | def test_multi_token_end_reentry(self): |
| 1477 | holder = self._make_holder(self.THINK_END_MULTI) |
| 1478 | self._sync_batch(holder, self.BUDGET) |
| 1479 | |
| 1480 | output = self._exhaust_budget(holder) |
| 1481 | self._accept_end_tokens(holder, output, self.THINK_END_MULTI) |
| 1482 | |
| 1483 | assert not holder._state[0]["in_end"] |
| 1484 | |
| 1485 | output.append(self.THINK_START) |
| 1486 | self._step(holder, list(output)) |
| 1487 | for _ in range(self.BUDGET): |
| 1488 | output.append(self.THINK_TOKEN) |
| 1489 | self._step(holder, list(output)) |
| 1490 | |
| 1491 | assert holder._state[0]["in_end"], ( |
| 1492 | "Immediate re-entry after multi-token end must be enforced" |
| 1493 | ) |
| 1494 | |
| 1495 | def test_single_block_not_broken(self): |
| 1496 | holder = self._make_holder(self.THINK_END_SINGLE) |
| 1497 | self._sync_batch(holder, self.BUDGET) |
| 1498 | |
| 1499 | output = self._exhaust_budget(holder) |
| 1500 | self._accept_end_tokens(holder, output, self.THINK_END_SINGLE) |
| 1501 | |
| 1502 | for _ in range(20): |
| 1503 | output.append(self.CONTENT_TOKEN) |
| 1504 | self._step(holder, list(output)) |
| 1505 | |
| 1506 | assert not holder._state[0]["in_end"] |
| 1507 | assert not holder._state[0]["in_think"] |
| 1508 | |
| 1509 | |
| 1510 | # --- Thinking budget natural-end re-entry tests (issue #45974) --- |
| 1511 | # After a model naturally emits </think> before exhausting its budget, |
| 1512 | # subsequent <think> blocks must still be tracked and budget-enforced. |
| 1513 | |
| 1514 | |
| 1515 | class TestThinkingBudgetNaturalEndReentry: |
| 1516 | """Tests for thinking budget enforcement across multiple think blocks. |
| 1517 | |
| 1518 | Covers both natural-end re-entry (issue #45974) and forced-end re-entry |
| 1519 | (issue #43708). |
| 1520 | """ |
| 1521 | |
| 1522 | THINK_START = 100 |
| 1523 | THINK_END_SINGLE = [200] |
| 1524 | THINK_END_MULTI = [200, 201, 202] |
| 1525 | BUDGET = 10 |
| 1526 | CONTENT_TOKEN = 50 |
| 1527 | THINK_TOKEN = 60 |
| 1528 | |
| 1529 | @staticmethod |
| 1530 | def _make_holder(end_token_ids): |
| 1531 | from vllm.v1.sample.thinking_budget_state import ( |
| 1532 | ThinkingBudgetStateHolder, |
| 1533 | ) |
| 1534 | |
| 1535 | cfg = create_mock_reasoning_config( |
| 1536 | [TestThinkingBudgetNaturalEndReentry.THINK_START], end_token_ids |
| 1537 | ) |
| 1538 | return ThinkingBudgetStateHolder( |
| 1539 | reasoning_config=cfg, |
| 1540 | max_num_seqs=8, |
| 1541 | num_spec_tokens=0, |
| 1542 | device=torch.device("cpu"), |
| 1543 | is_pin_memory=False, |
| 1544 | ) |
| 1545 | |
| 1546 | @staticmethod |
| 1547 | def _sync_batch(holder, budget): |
| 1548 | from unittest.mock import MagicMock |
| 1549 | |
| 1550 | params = MagicMock() |
| 1551 | params.thinking_token_budget = budget |
| 1552 | batch_update = MagicMock( |
| 1553 | removed=[], |
| 1554 | added=[(0, params, None, [])], |
| 1555 | moved=[], |
| 1556 | ) |
| 1557 | holder.sync_batch(batch_update) |
| 1558 | |
| 1559 | @staticmethod |
| 1560 | def _step(holder, output_tok_ids): |
| 1561 | holder.update_state( |
| 1562 | output_token_ids=[output_tok_ids], |
| 1563 | spec_token_ids=None, |
| 1564 | repeat_indices=None, |
| 1565 | ) |
| 1566 | |
| 1567 | # --- Natural-end re-entry tests --- |
| 1568 | |
| 1569 | def test_natural_end_reentry_single_token(self): |
| 1570 | """After natural </think>, a new <think> block must be enforced.""" |
| 1571 | holder = self._make_holder(self.THINK_END_SINGLE) |
| 1572 | self._sync_batch(holder, self.BUDGET) |
| 1573 | |
| 1574 | output = [] |
| 1575 | |
| 1576 | # Block 1: 6 tokens (under budget) + natural </think> |
| 1577 | output.append(self.THINK_START) |
| 1578 | self._step(holder, list(output)) |
| 1579 | for _ in range(6): |
| 1580 | output.append(self.THINK_TOKEN) |
| 1581 | self._step(holder, list(output)) |
| 1582 | output.append(self.THINK_END_SINGLE[0]) |
| 1583 | self._step(holder, list(output)) |
| 1584 | |
| 1585 | assert not holder._state[0]["in_end"] |
| 1586 | assert not holder._state[0]["in_think"] |
| 1587 | |
| 1588 | # Some content |
| 1589 | for _ in range(3): |
| 1590 | output.append(self.CONTENT_TOKEN) |
| 1591 | self._step(holder, list(output)) |
| 1592 | |
| 1593 | # Block 2: re-entry — should be budget-tracked |
| 1594 | output.append(self.THINK_START) |
| 1595 | self._step(holder, list(output)) |
| 1596 | for _ in range(self.BUDGET + 1): |
| 1597 | output.append(self.THINK_TOKEN) |
| 1598 | self._step(holder, list(output)) |
| 1599 | |
| 1600 | assert holder._state[0]["in_end"], ( |
| 1601 | "Second thinking block after natural end must be budget-enforced" |
| 1602 | ) |
| 1603 | |
| 1604 | def test_natural_end_reentry_multi_token(self): |
| 1605 | """Multi-token </think> natural end still allows Block 2 enforcement.""" |
| 1606 | holder = self._make_holder(self.THINK_END_MULTI) |
| 1607 | self._sync_batch(holder, self.BUDGET) |
| 1608 | |
| 1609 | output = [] |
| 1610 | |
| 1611 | # Block 1: 4 tokens + natural multi-token </think> |
| 1612 | output.append(self.THINK_START) |
| 1613 | self._step(holder, list(output)) |
| 1614 | for _ in range(4): |
| 1615 | output.append(self.THINK_TOKEN) |
| 1616 | self._step(holder, list(output)) |
| 1617 | for tok in self.THINK_END_MULTI: |
| 1618 | output.append(tok) |
| 1619 | self._step(holder, list(output)) |
| 1620 | |
| 1621 | assert not holder._state[0]["in_think"] |
| 1622 | |
| 1623 | # Content |
| 1624 | for _ in range(5): |
| 1625 | output.append(self.CONTENT_TOKEN) |
| 1626 | self._step(holder, list(output)) |
| 1627 | |
| 1628 | # Block 2: re-entry |
| 1629 | output.append(self.THINK_START) |
| 1630 | self._step(holder, list(output)) |
| 1631 | for _ in range(self.BUDGET + 1): |
| 1632 | output.append(self.THINK_TOKEN) |
| 1633 | self._step(holder, list(output)) |
| 1634 | |
| 1635 | assert holder._state[0]["in_end"], ( |
| 1636 | "Second block after multi-token natural end must be enforced" |
| 1637 | ) |
| 1638 | |
| 1639 | def test_natural_end_immediate_reentry(self): |
| 1640 | """Natural </think> followed immediately by <think> with no content.""" |
| 1641 | holder = self._make_holder(self.THINK_END_SINGLE) |
| 1642 | self._sync_batch(holder, self.BUDGET) |
| 1643 | |
| 1644 | output = [] |
| 1645 | |
| 1646 | # Block 1: 5 tokens + natural end |
| 1647 | output.append(self.THINK_START) |
| 1648 | self._step(holder, list(output)) |
| 1649 | for _ in range(5): |
| 1650 | output.append(self.THINK_TOKEN) |
| 1651 | self._step(holder, list(output)) |
| 1652 | output.append(self.THINK_END_SINGLE[0]) |
| 1653 | self._step(holder, list(output)) |
| 1654 | |
| 1655 | # Immediate re-entry — NO content between blocks |
| 1656 | output.append(self.THINK_START) |
| 1657 | self._step(holder, list(output)) |
| 1658 | for _ in range(self.BUDGET + 1): |
| 1659 | output.append(self.THINK_TOKEN) |
| 1660 | self._step(holder, list(output)) |
| 1661 | |
| 1662 | assert holder._state[0]["in_end"], ( |
| 1663 | "Immediate re-entry after natural end must be budget-enforced" |
| 1664 | ) |
| 1665 | |
| 1666 | def test_multiple_natural_reentries(self): |
| 1667 | """Block 1 natural -> Block 2 natural -> Block 3 must be enforced.""" |
| 1668 | holder = self._make_holder(self.THINK_END_SINGLE) |
| 1669 | self._sync_batch(holder, self.BUDGET) |
| 1670 | |
| 1671 | output = [] |
| 1672 | |
| 1673 | # Block 1: 4 tokens + natural end |
| 1674 | output.append(self.THINK_START) |
| 1675 | self._step(holder, list(output)) |
| 1676 | for _ in range(4): |
| 1677 | output.append(self.THINK_TOKEN) |
| 1678 | self._step(holder, list(output)) |
| 1679 | output.append(self.THINK_END_SINGLE[0]) |
| 1680 | self._step(holder, list(output)) |
| 1681 | |
| 1682 | # Content |
| 1683 | for _ in range(2): |
| 1684 | output.append(self.CONTENT_TOKEN) |
| 1685 | self._step(holder, list(output)) |
| 1686 | |
| 1687 | # Block 2: 3 tokens + natural end |
| 1688 | output.append(self.THINK_START) |
| 1689 | self._step(holder, list(output)) |
| 1690 | for _ in range(3): |
| 1691 | output.append(self.THINK_TOKEN) |
| 1692 | self._step(holder, list(output)) |
| 1693 | output.append(self.THINK_END_SINGLE[0]) |
| 1694 | self._step(holder, list(output)) |
| 1695 | |
| 1696 | # Content |
| 1697 | for _ in range(2): |
| 1698 | output.append(self.CONTENT_TOKEN) |
| 1699 | self._step(holder, list(output)) |
| 1700 | |
| 1701 | # Block 3: exceed budget |
| 1702 | output.append(self.THINK_START) |
| 1703 | self._step(holder, list(output)) |
| 1704 | for _ in range(self.BUDGET + 1): |
| 1705 | output.append(self.THINK_TOKEN) |
| 1706 | self._step(holder, list(output)) |
| 1707 | |
| 1708 | assert holder._state[0]["in_end"], ( |
| 1709 | "Third block after two natural ends must be budget-enforced" |
| 1710 | ) |
| 1711 | |
| 1712 | def test_natural_then_forced_in_block_2(self): |
| 1713 | """Block 1 ends naturally, Block 2 exceeds budget -> forced end.""" |
| 1714 | budget = 5 |
| 1715 | holder = self._make_holder(self.THINK_END_SINGLE) |
| 1716 | self._sync_batch(holder, budget) |
| 1717 | |
| 1718 | output = [] |
| 1719 | |
| 1720 | # Block 1: 3 tokens + natural end (under budget) |
| 1721 | output.append(self.THINK_START) |
| 1722 | self._step(holder, list(output)) |
| 1723 | for _ in range(3): |
| 1724 | output.append(self.THINK_TOKEN) |
| 1725 | self._step(holder, list(output)) |
| 1726 | output.append(self.THINK_END_SINGLE[0]) |
| 1727 | self._step(holder, list(output)) |
| 1728 | |
| 1729 | assert not holder._state[0]["in_end"] |
| 1730 | |
| 1731 | # Content |
| 1732 | for _ in range(3): |
| 1733 | output.append(self.CONTENT_TOKEN) |
| 1734 | self._step(holder, list(output)) |
| 1735 | |
| 1736 | # Block 2: exceed budget |
| 1737 | output.append(self.THINK_START) |
| 1738 | self._step(holder, list(output)) |
| 1739 | for _ in range(budget + 1): |
| 1740 | output.append(self.THINK_TOKEN) |
| 1741 | self._step(holder, list(output)) |
| 1742 | |
| 1743 | assert holder._state[0]["in_end"], ( |
| 1744 | "Block 2 must trigger forced end after natural end in Block 1" |
| 1745 | ) |
| 1746 | |
| 1747 | def test_budget_one_natural_end_reentry(self): |
| 1748 | """Budget = 1: natural end with 0 think tokens, re-entry enforced.""" |
| 1749 | budget = 1 |
| 1750 | holder = self._make_holder(self.THINK_END_SINGLE) |
| 1751 | self._sync_batch(holder, budget) |
| 1752 | |
| 1753 | output = [] |
| 1754 | |
| 1755 | # Block 1: immediate natural end (0 think tokens) |
| 1756 | output.append(self.THINK_START) |
| 1757 | self._step(holder, list(output)) |
| 1758 | output.append(self.THINK_END_SINGLE[0]) |
| 1759 | self._step(holder, list(output)) |
| 1760 | |
| 1761 | assert not holder._state[0]["in_end"] |
| 1762 | |
| 1763 | # Content |
| 1764 | for _ in range(3): |
| 1765 | output.append(self.CONTENT_TOKEN) |
| 1766 | self._step(holder, list(output)) |
| 1767 | |
| 1768 | # Block 2: 2 think tokens (should be enforced at budget=1) |
| 1769 | output.append(self.THINK_START) |
| 1770 | self._step(holder, list(output)) |
| 1771 | output.append(self.THINK_TOKEN) |
| 1772 | self._step(holder, list(output)) |
| 1773 | output.append(self.THINK_TOKEN) |
| 1774 | self._step(holder, list(output)) |
| 1775 | |
| 1776 | assert holder._state[0]["in_end"], ( |
| 1777 | "Budget=1 must enforce after 1 token in re-entry block" |
| 1778 | ) |
| 1779 | |
| 1780 | def test_partial_end_sequence_in_content(self): |
| 1781 | """Partial end sequence in content shouldn't trigger false exit.""" |
| 1782 | holder = self._make_holder(self.THINK_END_MULTI) |
| 1783 | self._sync_batch(holder, self.BUDGET) |
| 1784 | |
| 1785 | output = [] |
| 1786 | |
| 1787 | # Block 1: natural end with full sequence [200, 201, 202] |
| 1788 | output.append(self.THINK_START) |
| 1789 | self._step(holder, list(output)) |
| 1790 | for _ in range(3): |
| 1791 | output.append(self.THINK_TOKEN) |
| 1792 | self._step(holder, list(output)) |
| 1793 | for tok in self.THINK_END_MULTI: |
| 1794 | output.append(tok) |
| 1795 | self._step(holder, list(output)) |
| 1796 | |
| 1797 | # Content with partial end sequence [200, 201] but not [200, 201, 202] |
| 1798 | output.append(self.CONTENT_TOKEN) |
| 1799 | self._step(holder, list(output)) |
| 1800 | output.append(200) |
| 1801 | self._step(holder, list(output)) |
| 1802 | output.append(201) |
| 1803 | self._step(holder, list(output)) |
| 1804 | output.append(self.CONTENT_TOKEN) |
| 1805 | self._step(holder, list(output)) |
| 1806 | |
| 1807 | # Block 2: re-entry must still be enforced |
| 1808 | output.append(self.THINK_START) |
| 1809 | self._step(holder, list(output)) |
| 1810 | for _ in range(self.BUDGET + 1): |
| 1811 | output.append(self.THINK_TOKEN) |
| 1812 | self._step(holder, list(output)) |
| 1813 | |
| 1814 | assert holder._state[0]["in_end"], ( |
| 1815 | "Partial end tokens in content must not prevent Block 2 enforcement" |
| 1816 | ) |
| 1817 | |
| 1818 | def test_continue_thinking_natural_end_reentry(self): |
| 1819 | """Prompt-prefill inside <think>: after natural end, Block 2 must be enforced. |
| 1820 | |
| 1821 | When a request arrives with the prompt already inside a <think> block, |
| 1822 | continue_thinking is set to True. This test verifies that after a |
| 1823 | natural </think> exit, continue_thinking is cleared and Block 2 |
| 1824 | enforcement works correctly. |
| 1825 | """ |
| 1826 | from unittest.mock import MagicMock |
| 1827 | |
| 1828 | from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder |
| 1829 | |
| 1830 | cfg = create_mock_reasoning_config([self.THINK_START], self.THINK_END_SINGLE) |
| 1831 | holder = ThinkingBudgetStateHolder( |
| 1832 | reasoning_config=cfg, |
| 1833 | max_num_seqs=8, |
| 1834 | num_spec_tokens=0, |
| 1835 | device=torch.device("cpu"), |
| 1836 | is_pin_memory=False, |
| 1837 | ) |
| 1838 | |
| 1839 | # Simulate: prompt already inside <think> (sets continue_thinking=True) |
| 1840 | prompt_tok_ids = [self.THINK_START] |
| 1841 | params = MagicMock() |
| 1842 | params.thinking_token_budget = self.BUDGET |
| 1843 | batch_update = MagicMock( |
| 1844 | removed=[], |
| 1845 | added=[(0, params, prompt_tok_ids, [])], |
| 1846 | moved=[], |
| 1847 | ) |
| 1848 | holder.sync_batch(batch_update) |
| 1849 | assert holder._state[0]["continue_thinking"], ( |
| 1850 | "Prompt inside <think> must set continue_thinking=True" |
| 1851 | ) |
| 1852 | |
| 1853 | # Block 1: emit some think tokens then natural </think> |
| 1854 | output = list(prompt_tok_ids) |
| 1855 | for _ in range(4): |
| 1856 | output.append(self.THINK_TOKEN) |
| 1857 | holder.update_state([output], None, None) |
| 1858 | output.append(self.THINK_END_SINGLE[0]) |
| 1859 | holder.update_state([output], None, None) |
| 1860 | |
| 1861 | assert not holder._state[0]["in_think"], "Should have exited think mode" |
| 1862 | assert not holder._state[0]["continue_thinking"], ( |
| 1863 | "continue_thinking must be cleared after natural end" |
| 1864 | ) |
| 1865 | |
| 1866 | # Content tokens |
| 1867 | for _ in range(3): |
| 1868 | output.append(self.CONTENT_TOKEN) |
| 1869 | holder.update_state([output], None, None) |
| 1870 | |
| 1871 | # Block 2: re-entry must be detected and budget enforced |
| 1872 | output.append(self.THINK_START) |
| 1873 | holder.update_state([output], None, None) |
| 1874 | for _ in range(self.BUDGET + 1): |
| 1875 | output.append(self.THINK_TOKEN) |
| 1876 | holder.update_state([output], None, None) |
| 1877 | |
| 1878 | assert holder._state[0]["in_end"], ( |
| 1879 | "Block 2 after natural end of prompt-prefill session must be enforced" |
| 1880 | ) |
| 1881 | |
| 1882 | # --- Forced-end re-entry tests --- |