diff --git a/benchmarks/kernels/benchmark_thinking_budget.py b/benchmarks/kernels/benchmark_thinking_budget.py new file mode 100644 index 0000000000..7a0a1b4d59 --- /dev/null +++ b/benchmarks/kernels/benchmark_thinking_budget.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import dataclasses +from collections.abc import Callable +from typing import Literal + +import numpy as np +import torch + +from vllm.sampling_params import SamplingParams +from vllm.triton_utils import triton +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState +from vllm.v1.worker.gpu.states import RequestState + +START_TOKEN_ID = 90 +END_TOKEN_ID = 91 +TOKEN_ID = 10 +VOCAB_SIZE = 128 +QUANTILES = [0.5, 0.2, 0.8] + + +class ReasoningConfig: + reasoning_start_token_ids = [START_TOKEN_ID] + reasoning_end_token_ids = [END_TOKEN_ID] + natural_reasoning_end_token_ids = [END_TOKEN_ID] + + +@dataclasses.dataclass +class BenchmarkCase: + req_states: RequestState + state: ThinkingBudgetState + run: Callable[[], None] + + +def create_case( + history_len: int, + device: torch.device, + extra_tokens: int = 0, + history_type: Literal["reasoning", "prefill"] = "reasoning", + batch_size: int = 1, + budget_type: Literal["active", "forced", "mixed"] = "active", +) -> BenchmarkCase: + assert history_len >= 2 + if history_type == "prefill": + tokens = [*([TOKEN_ID] * (history_len - 1)), START_TOKEN_ID] + else: + tokens = [TOKEN_ID, START_TOKEN_ID, *([TOKEN_ID] * (history_len - 2))] + req_states = RequestState( + max_num_reqs=batch_size, + max_model_len=history_len + extra_tokens + 1, + max_num_batched_tokens=batch_size, + num_speculative_steps=1, + vocab_size=VOCAB_SIZE, + device=device, + ) + req_indices = [] + for req_num in range(batch_size): + req_id = f"benchmark-{req_num}" + req_states.add_request( + req_id=req_id, + prompt_len=1, + all_token_ids=tokens, + num_computed_tokens=history_len, + max_tokens=extra_tokens + 1, + ) + req_indices.append(req_states.req_id_to_index[req_id]) + req_states.apply_staged_writes() + + state = ThinkingBudgetState(req_states, ReasoningConfig()) + active_budget = history_len + extra_tokens + 1 + forced_budget = max(1, history_len - 2) + for req_num, req_idx in enumerate(req_indices): + if budget_type == "forced": + budget = forced_budget + elif budget_type == "mixed" and req_num % 2 == 1: + budget = None + else: + budget = active_budget + state.add_request( + req_idx, + SamplingParams(thinking_token_budget=budget), + ) + state.apply_staged_writes() + + idx_mapping = torch.tensor(req_indices, dtype=torch.int32, device=device) + logits = torch.zeros((batch_size, VOCAB_SIZE), device=device) + idx_mapping_np = idx_mapping.cpu().numpy() + input_ids = torch.full((batch_size,), TOKEN_ID, dtype=torch.int32, device=device) + local_pos = torch.zeros(batch_size, dtype=torch.int32, device=device) + + def run() -> None: + state.apply( + logits, + idx_mapping, + idx_mapping, + idx_mapping_np, + input_ids, + local_pos, + ) + + run() + torch.accelerator.synchronize(device) + active_req_indices = [ + req_idx + for req_num, req_idx in enumerate(req_indices) + if budget_type != "mixed" or req_num % 2 == 0 + ] + assert torch.all(state.cached_scan_pos[active_req_indices] == history_len).item() + if budget_type == "forced": + assert torch.all(logits[:, END_TOKEN_ID] == 1.0e9).item() + return BenchmarkCase(req_states, state, run) + + +def benchmark_cached( + case: BenchmarkCase, warmup_ms: int, rep_ms: int +) -> tuple[float, float, float]: + median_ms, min_ms, max_ms = triton.testing.do_bench( + case.run, + warmup=warmup_ms, + rep=rep_ms, + quantiles=QUANTILES, + ) + return median_ms, min_ms, max_ms + + +def summarize(timings_ms: list[float]) -> tuple[float, float, float]: + median_ms, min_ms, max_ms = np.quantile(timings_ms, QUANTILES) + return float(median_ms), float(min_ms), float(max_ms) + + +def time_cuda_call(run: Callable[[], None]) -> float: + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) + + +def benchmark_incremental_decode( + history_len: int, + device: torch.device, + warmup_steps: int, + iterations: int, +) -> tuple[float, float, float]: + case = create_case(history_len, device, warmup_steps + iterations) + timings_ms: list[float] = [] + total_len = history_len + for step in range(warmup_steps + iterations): + case.req_states.all_token_ids.stage_write(0, total_len, [TOKEN_ID]) + total_len += 1 + case.req_states.total_len.stage_write_elem(0, total_len) + case.req_states.apply_staged_writes() + torch.accelerator.synchronize(device) + + elapsed_ms = time_cuda_call(case.run) + if step >= warmup_steps: + timings_ms.append(elapsed_ms) + + assert case.state.cached_scan_pos[0].item() == total_len + return summarize(timings_ms) + + +def benchmark_cold_scan( + history_len: int, + device: torch.device, + warmup_steps: int, + iterations: int, + history_type: Literal["reasoning", "prefill"], +) -> tuple[float, float, float]: + case = create_case(history_len, device, history_type=history_type) + timings_ms: list[float] = [] + for step in range(warmup_steps + iterations): + case.state.cached_last_start.fill_(-1) + case.state.cached_last_end.fill_(-1) + case.state.cached_scan_pos.zero_() + torch.accelerator.synchronize(device) + + elapsed_ms = time_cuda_call(case.run) + if step >= warmup_steps: + timings_ms.append(elapsed_ms) + + return summarize(timings_ms) + + +def main() -> None: + parser = FlexibleArgumentParser( + description=( + "Benchmark thinking-budget scan, forced-end, and batched request " + "overhead across token history lengths, including worst-case cold " + "resume reconstruction." + ) + ) + parser.add_argument( + "--history-lengths", + type=int, + nargs="+", + default=[128, 16384, 32768], + ) + parser.add_argument( + "--modes", + nargs="+", + choices=[ + "cached", + "incremental-decode", + "cold-prefill", + "cold-resume-worst-case", + "forced-end", + "batched-budgeted", + "batched-mixed", + ], + default=[ + "cached", + "incremental-decode", + "cold-prefill", + "cold-resume-worst-case", + "forced-end", + "batched-budgeted", + "batched-mixed", + ], + ) + parser.add_argument( + "--batch-sizes", + type=int, + nargs="+", + default=[1, 8, 32], + help="Batch sizes used by batched-budgeted and batched-mixed modes.", + ) + parser.add_argument("--device", type=str, default="cuda:0") + parser.add_argument( + "--warmup-ms", + type=int, + default=100, + help="Warmup duration passed to triton.testing.do_bench.", + ) + parser.add_argument( + "--rep-ms", + type=int, + default=500, + help="Measurement duration passed to triton.testing.do_bench.", + ) + parser.add_argument( + "--warmup-steps", + type=int, + default=10, + help="Warmup steps for incremental-decode and cold modes.", + ) + parser.add_argument( + "--iterations", + type=int, + default=100, + help="Measured steps for incremental-decode and cold modes.", + ) + parser.add_argument( + "--max-slowdown", + type=float, + default=None, + help=( + "Fail if cached or incremental-decode median latency exceeds this " + "multiple of its shortest-history latency. Cold modes are excluded " + "because worst-case cold-resume cost scales with scanned history." + ), + ) + args = parser.parse_args() + + history_lengths = sorted(set(args.history_lengths)) + if not history_lengths or history_lengths[0] < 2: + parser.error("history lengths must be at least 2") + if args.warmup_steps < 0 or args.iterations <= 0: + parser.error("warmup steps must be non-negative and iterations positive") + batch_sizes = sorted(set(args.batch_sizes)) + if not batch_sizes or batch_sizes[0] <= 0: + parser.error("batch sizes must be positive") + + device = torch.device(args.device) + results: dict[str, dict[int, tuple[float, float, float]]] = {} + for mode in args.modes: + mode_batch_sizes = batch_sizes if mode.startswith("batched-") else [1] + if mode == "batched-mixed": + mode_batch_sizes = [size for size in mode_batch_sizes if size >= 2] + if not mode_batch_sizes: + parser.error("batched-mixed requires a batch size of at least 2") + for batch_size in mode_batch_sizes: + label = f"{mode}-b{batch_size}" if mode.startswith("batched-") else mode + mode_results: dict[int, tuple[float, float, float]] = {} + for history_len in history_lengths: + if mode == "cached": + case = create_case(history_len, device) + result = benchmark_cached(case, args.warmup_ms, args.rep_ms) + elif mode == "incremental-decode": + result = benchmark_incremental_decode( + history_len, + device, + args.warmup_steps, + args.iterations, + ) + elif mode == "cold-prefill": + result = benchmark_cold_scan( + history_len, + device, + args.warmup_steps, + args.iterations, + "prefill", + ) + elif mode == "cold-resume-worst-case": + result = benchmark_cold_scan( + history_len, + device, + args.warmup_steps, + args.iterations, + "reasoning", + ) + else: + budget_type = ( + "forced" + if mode == "forced-end" + else "mixed" + if mode == "batched-mixed" + else "active" + ) + case = create_case( + history_len, + device, + batch_size=batch_size, + budget_type=budget_type, + ) + result = benchmark_cached(case, args.warmup_ms, args.rep_ms) + mode_results[history_len] = result + results[label] = mode_results + + print("mode history_len median_us p20_us p80_us slowdown") + for label, mode_results in results.items(): + baseline_ms = mode_results[history_lengths[0]][0] + for history_len in history_lengths: + median_ms, min_ms, max_ms = mode_results[history_len] + slowdown = median_ms / baseline_ms + print( + f"{label:<26} {history_len:>11} {median_ms * 1000:>9.3f} " + f"{min_ms * 1000:>6.3f} {max_ms * 1000:>6.3f} " + f"{slowdown:>7.3f}x" + ) + + if args.max_slowdown is not None: + for mode in ("cached", "incremental-decode"): + if mode not in results: + continue + mode_results = results[mode] + baseline_ms = mode_results[history_lengths[0]][0] + worst_history_len = max( + mode_results, key=lambda length: mode_results[length][0] + ) + worst_slowdown = mode_results[worst_history_len][0] / baseline_ms + if worst_slowdown > args.max_slowdown: + raise SystemExit( + f"thinking-budget {mode} slowdown {worst_slowdown:.3f}x " + f"at history length {worst_history_len} exceeds " + f"{args.max_slowdown:.3f}x" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index 4c57467381..ec7a977f1a 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -72,9 +72,7 @@ def server(): "0.4", "--no-async-scheduling", ] - # thinking_token_budget is not yet supported by the V2 model runner. - env_dict = {"VLLM_USE_V2_MODEL_RUNNER": "0"} - with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server: + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server @@ -90,9 +88,7 @@ def server_with_auto_reasoning_config(): "0.4", "--no-async-scheduling", ] - # thinking_token_budget is not yet supported by the V2 model runner. - env_dict = {"VLLM_USE_V2_MODEL_RUNNER": "0"} - with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server: + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server @@ -126,8 +122,7 @@ def server_qwen35_fp8_mtp_tp2(): } ), ] - # thinking_token_budget is not yet supported by the V2 model runner. - env_dict: dict[str, str] = {"VLLM_USE_V2_MODEL_RUNNER": "0"} + env_dict: dict[str, str] = {} # With 4+ GPUs, run TP=2 on physical devices 2,3 so module-scoped 0.6B servers # on 0,1 do not exhaust memory on the same devices as this worker. if current_platform.device_count() >= 4: diff --git a/tests/v1/worker/test_gpu_thinking_budget.py b/tests/v1/worker/test_gpu_thinking_budget.py new file mode 100644 index 0000000000..0b0f19f7db --- /dev/null +++ b/tests/v1/worker/test_gpu_thinking_budget.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for Model Runner V2 thinking budget tests", + allow_module_level=True, + ) + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState +from vllm.v1.worker.gpu.states import RequestState + +DEVICE = torch.device("cuda") +START = 90 +END = 91 +END_A = 92 +END_B = 93 +VOCAB_SIZE = 128 + + +class MockReasoningConfig: + reasoning_start_token_ids = [START] + reasoning_end_token_ids = [END] + natural_reasoning_end_token_ids = [END] + + +class MockMultiTokenEndReasoningConfig: + reasoning_start_token_ids = [START] + reasoning_end_token_ids = [END_A, END_B] + natural_reasoning_end_token_ids = [END_A, END_B] + + +class MockDistinctEndReasoningConfig: + reasoning_start_token_ids = [START] + reasoning_end_token_ids = [END_A, END_B] + natural_reasoning_end_token_ids = [END] + + +def _make_req_states(tokens: list[int], prompt_len: int = 1) -> RequestState: + req_states = RequestState( + max_num_reqs=4, + max_model_len=max(64, len(tokens) + 1), + max_num_batched_tokens=16, + num_speculative_steps=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + ) + req_states.add_request( + req_id="req", + prompt_len=prompt_len, + all_token_ids=tokens, + num_computed_tokens=len(tokens), + max_tokens=32, + ) + req_states.apply_staged_writes() + return req_states + + +def _apply( + state: ThinkingBudgetState, + logits: torch.Tensor, + input_ids: list[int], + local_pos: list[int], +) -> torch.Tensor: + idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE) + expanded_idx_mapping = torch.tensor( + [3] * len(input_ids), dtype=torch.int32, device=DEVICE + ) + idx_mapping_np = idx_mapping.cpu().numpy() + state.apply( + logits, + expanded_idx_mapping, + idx_mapping, + idx_mapping_np, + torch.tensor(input_ids, dtype=torch.int32, device=DEVICE), + torch.tensor(local_pos, dtype=torch.int32, device=DEVICE), + ) + return logits.cpu() + + +def test_v2_thinking_budget_forces_end_after_budget_reached(): + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + state = ThinkingBudgetState(req_states, MockReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.arange(VOCAB_SIZE, dtype=torch.float32, device=DEVICE).view(1, -1) + expected = logits.cpu() + out = _apply(state, logits, input_ids=[12], local_pos=[0]) + + expected[0, END] = 1.0e9 + torch.testing.assert_close(out, expected) + + +def test_v2_thinking_budget_restores_masked_end_token(): + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + state = ThinkingBudgetState(req_states, MockReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + logits[0, END] = -float("inf") + out = _apply(state, logits, input_ids=[12], local_pos=[0]) + + assert out[0, END] == pytest.approx(1.0e9) + + +def test_v2_thinking_budget_allows_tokens_before_budget(): + req_states = _make_req_states([1, START, 10, 11], prompt_len=1) + state = ThinkingBudgetState(req_states, MockReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = _apply(state, logits, input_ids=[11], local_pos=[0]) + + assert torch.all(out == 0) + + +def test_v2_thinking_budget_continues_multi_token_end_marker(): + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + state = ThinkingBudgetState(req_states, MockMultiTokenEndReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.zeros((2, VOCAB_SIZE), device=DEVICE) + out = _apply( + state, + logits, + input_ids=[12, END_A], + local_pos=[0, 1], + ) + + assert out[0, END_A] == pytest.approx(1.0e9) + assert out[1, END_B] == pytest.approx(1.0e9) + + +def test_v2_thinking_budget_uses_distinct_forced_end_marker(): + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + state = ThinkingBudgetState(req_states, MockDistinctEndReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.zeros((2, VOCAB_SIZE), device=DEVICE) + out = _apply( + state, + logits, + input_ids=[12, END_A], + local_pos=[0, 1], + ) + + assert out[0, END_A] == pytest.approx(1.0e9) + assert out[1, END_B] == pytest.approx(1.0e9) + + +def test_v2_thinking_budget_stops_after_natural_end_marker(): + req_states = _make_req_states( + [1, START, 10, END, 20, 21, 22], + prompt_len=1, + ) + state = ThinkingBudgetState(req_states, MockDistinctEndReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = _apply(state, logits, input_ids=[22], local_pos=[0]) + + assert torch.all(out == 0) + + +def test_v2_thinking_budget_ignores_plain_request(): + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + state = ThinkingBudgetState(req_states, MockReasoningConfig()) + state.add_request(3, SamplingParams()) + state.apply_staged_writes() + + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = _apply(state, logits, input_ids=[12], local_pos=[0]) + + assert torch.all(out == 0) + + +def test_v2_thinking_budget_latest_prefill_end_disables_forcing(): + req_states = _make_req_states( + [1, START, 10, 11, 12, END, 13], + prompt_len=1, + ) + state = ThinkingBudgetState(req_states, MockReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = _apply(state, logits, input_ids=[13], local_pos=[0]) + + assert torch.all(out == 0) + + +def test_v2_thinking_budget_uses_latest_prefill_start_boundary(): + req_states = _make_req_states( + [1, START, 10, 11, 12, END, 13, START, 14, 15, 16], + prompt_len=1, + ) + state = ThinkingBudgetState(req_states, MockReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = _apply(state, logits, input_ids=[16], local_pos=[0]) + + assert out[0, END] == pytest.approx(1.0e9) + + +def test_v2_thinking_budget_incrementally_scans_long_generation(): + """Guard against rescanning the full token history on every decode step.""" + tokens = [1, START, *([10] * 16382)] + req_states = _make_req_states(tokens) + state = ThinkingBudgetState(req_states, MockReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=32768)) + state.apply_staged_writes() + + _apply(state, torch.zeros((1, VOCAB_SIZE), device=DEVICE), [10], [0]) + assert state.cached_scan_pos[3].item() == len(tokens) + + req_states.all_token_ids.stage_write(3, len(tokens), [10]) + req_states.total_len.stage_write_elem(3, len(tokens) + 1) + req_states.apply_staged_writes() + _apply(state, torch.zeros((1, VOCAB_SIZE), device=DEVICE), [10], [0]) + + assert state.cached_scan_pos[3].item() == len(tokens) + 1 + + +def test_v2_thinking_budget_clamps_oversized_budget(): + """Budgets beyond int32 must not crash and behave as unlimited.""" + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + state = ThinkingBudgetState(req_states, MockReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=2**40)) + state.apply_staged_writes() + + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = _apply(state, logits, input_ids=[12], local_pos=[0]) + + assert torch.all(out == 0) + + +def test_v2_thinking_budget_continues_end_prefix_from_prompt(): + """A resumed prompt ending with a partial forced-end marker must not + restart the marker sequence and duplicate its first token.""" + req_states = _make_req_states([1, START, 10, 11, END_A], prompt_len=5) + state = ThinkingBudgetState(req_states, MockMultiTokenEndReasoningConfig()) + state.add_request(3, SamplingParams(thinking_token_budget=3)) + state.apply_staged_writes() + + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = _apply(state, logits, input_ids=[END_A], local_pos=[0]) + + assert out[0, END_B] == pytest.approx(1.0e9) + assert out[0, END_A] == 0 diff --git a/vllm/config/reasoning.py b/vllm/config/reasoning.py index ff5546e05e..0ac93be082 100644 --- a/vllm/config/reasoning.py +++ b/vllm/config/reasoning.py @@ -13,10 +13,10 @@ from vllm.tokenizers import cached_tokenizer_from_config class ReasoningConfig: """Configuration for reasoning models. - Set `reasoning_start_str` and `reasoning_end_str` to the strings that delimit - the reasoning block (e.g. `""` and `""`). The - corresponding token IDs are derived automatically via - `initialize_token_ids` and are not intended to be set directly. + Set `reasoning_start_str` and `reasoning_end_str` to the strings used to + enter and forcibly terminate reasoning. The end string may include a + transition phrase before the parser's natural reasoning end marker. Token + IDs are derived automatically by `initialize_token_ids`. """ reasoning_parser: str = "" @@ -24,7 +24,7 @@ class ReasoningConfig: reasoning_start_str: str = "" """String that indicates the start of reasoning.""" reasoning_end_str: str = "" - """String that indicates the end of reasoning content.""" + """String forced when the thinking budget is exhausted.""" _reasoning_start_token_ids: list[int] | None = field( default=None, init=False, repr=False @@ -34,8 +34,11 @@ class ReasoningConfig: _reasoning_end_token_ids: list[int] | None = field( default=None, init=False, repr=False ) - """Private backing field for `reasoning_end_token_ids`. Set by - `initialize_token_ids`. Not intended to be configured directly.""" + """Private backing field for forced reasoning end token IDs.""" + _natural_reasoning_end_token_ids: list[int] | None = field( + default=None, init=False, repr=False + ) + """Token IDs that naturally terminate reasoning, as defined by the parser.""" _enabled: bool = field(default=False, init=False, repr=False) """Private field indicating whether reasoning token IDs have been initialized. @@ -55,15 +58,20 @@ class ReasoningConfig: @property def reasoning_end_token_ids(self) -> list[int] | None: - """Token IDs derived from `reasoning_end_str`. Set automatically by - `initialize_token_ids`. Not intended to be configured directly.""" + """Token IDs forced when the thinking budget is exhausted.""" return self._reasoning_end_token_ids + @property + def natural_reasoning_end_token_ids(self) -> list[int] | None: + """Token IDs that indicate the model naturally ended reasoning.""" + return self._natural_reasoning_end_token_ids + def initialize_token_ids(self, model_config: ModelConfig) -> None: """Initialize reasoning token IDs from strings using the tokenizer.""" if ( self._reasoning_start_token_ids is not None and self._reasoning_end_token_ids is not None + and self._natural_reasoning_end_token_ids is not None ): self._enabled = True return # Already initialized @@ -71,9 +79,8 @@ class ReasoningConfig: tokenizer = cached_tokenizer_from_config(model_config=model_config) reasoning_start_str = self.reasoning_start_str reasoning_end_str = self.reasoning_end_str - if self.reasoning_parser is not None and ( - not reasoning_start_str or not reasoning_end_str - ): + natural_reasoning_end_str = "" + if self.reasoning_parser: parser_cls = ReasoningParserManager.get_reasoning_parser( self.reasoning_parser ) @@ -85,6 +92,10 @@ class ReasoningConfig: end_token = reasoning_parser.reasoning_end_str if end_token and not reasoning_end_str: reasoning_end_str = end_token + natural_reasoning_end_str = end_token or "" + + if not natural_reasoning_end_str: + natural_reasoning_end_str = reasoning_end_str if not reasoning_start_str or not reasoning_end_str: # If we don't have valid strings to tokenize, @@ -96,8 +107,15 @@ class ReasoningConfig: self._reasoning_end_token_ids = tokenizer.encode( reasoning_end_str, add_special_tokens=False ) + self._natural_reasoning_end_token_ids = tokenizer.encode( + natural_reasoning_end_str, add_special_tokens=False + ) - if not self._reasoning_start_token_ids or not self._reasoning_end_token_ids: + if ( + not self._reasoning_start_token_ids + or not self._reasoning_end_token_ids + or not self._natural_reasoning_end_token_ids + ): raise ValueError( f"ReasoningConfig: failed to tokenize reasoning strings: " f"reasoning_start_str='{self.reasoning_start_str}', " diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index bb9963a03c..edac3441c7 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2377,12 +2377,6 @@ class VllmConfig: f"Model Runner V2 does not yet support: {', '.join(unsupported)}" ) - if self.reasoning_config is not None: - logger.warning_once( - "Model Runner V2 does not yet support the thinking_token_budget " - "request parameter. Set VLLM_USE_V2_MODEL_RUNNER=0 if this is required." - ) - def validate_block_size(self) -> None: """Validate block_size against DCP and mamba constraints. diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 359d1952a1..0e5995ee6f 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -108,22 +108,15 @@ class InputProcessor: self.tokenizer, ) - if params.thinking_token_budget is not None: - if ( - self.vllm_config.reasoning_config is None - or not self.vllm_config.reasoning_config.enabled - ): - raise VLLMValidationError( - "thinking_token_budget is set but reasoning_config is " - "not configured. Please set --reasoning-parser " - "and/or --reasoning-config to use thinking_token_budget." - ) - if self.use_v2_model_runner: - raise VLLMValidationError( - "thinking_token_budget is not yet supported by the V2 " - "model runner. Run vLLM with VLLM_USE_V2_MODEL_RUNNER=0 " - "to use thinking_token_budget." - ) + if params.thinking_token_budget is not None and ( + self.vllm_config.reasoning_config is None + or not self.vllm_config.reasoning_config.enabled + ): + raise VLLMValidationError( + "thinking_token_budget is set but reasoning_config is " + "not configured. Please set --reasoning-parser " + "and/or --reasoning-config to use thinking_token_budget." + ) elif isinstance(params, PoolingParams): supported_pooling_tasks = [ task for task in supported_tasks if task in POOLING_TASKS diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5fa39436b1..3a68fb49fb 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -371,6 +371,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): logprobs_mode=self.model_config.logprobs_mode, num_speculative_tokens=self.decode_query_len, use_fp64_gumbel=self.model_config.use_fp64_gumbel, + reasoning_config=self.vllm_config.reasoning_config, ) custom = self.model_state.custom_sampler(self.sampler) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 3503b699f9..2f03355e5d 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -6,6 +6,7 @@ import torch import vllm.envs as envs from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode +from vllm.config.reasoning import ReasoningConfig from vllm.sampling_params import SamplingParams from vllm.v1.sample.ops.topk_topp_sampler import ( apply_top_k_top_p, @@ -24,6 +25,7 @@ from vllm.v1.worker.gpu.sample.logprob import ( from vllm.v1.worker.gpu.sample.output import SamplerOutput from vllm.v1.worker.gpu.sample.penalties import PenaltiesState from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates +from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState from vllm.v1.worker.gpu.states import RequestState @@ -37,6 +39,7 @@ class Sampler: logprobs_mode: LogprobsMode = "raw_logprobs", num_speculative_tokens: int = 1, use_fp64_gumbel: bool = False, + reasoning_config: ReasoningConfig | None = None, ): self.logprobs_mode = logprobs_mode self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default. @@ -48,6 +51,7 @@ class Sampler: self.logit_bias_state = LogitBiasState(max_num_reqs, device) 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.num_speculative_tokens = num_speculative_tokens self.use_flashinfer = flashinfer_sampler_supported() @@ -59,6 +63,7 @@ class Sampler: self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params) self.bad_words_state.add_request(req_idx, sampling_params) self.logprob_token_ids_state.add_request(req_idx, sampling_params) + self.thinking_budget_state.add_request(req_idx, sampling_params) def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() @@ -66,6 +71,7 @@ class Sampler: self.logit_bias_state.apply_staged_writes() self.bad_words_state.apply_staged_writes() self.logprob_token_ids_state.apply_staged_writes() + self.thinking_budget_state.apply_staged_writes() def __call__( self, @@ -73,6 +79,7 @@ class Sampler: input_batch: InputBatch, ) -> SamplerOutput: expanded_idx_mapping = input_batch.expanded_idx_mapping + idx_mapping = input_batch.idx_mapping idx_mapping_np = input_batch.idx_mapping_np cu_num_logits_np = input_batch.cu_num_logits_np expanded_local_pos = input_batch.expanded_local_pos @@ -92,6 +99,7 @@ class Sampler: sampled, processed_logits = self.sample( logits, expanded_idx_mapping, + idx_mapping, idx_mapping_np, pos, input_ids, @@ -146,6 +154,7 @@ class Sampler: self, logits: torch.Tensor, expanded_idx_mapping: torch.Tensor, + idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray, pos: torch.Tensor, input_ids: torch.Tensor, @@ -181,6 +190,17 @@ class Sampler: expanded_local_pos, ) + # Force the reasoning end marker once a request's thinking budget is + # reached; applied before temperature so the forced token is always kept. + self.thinking_budget_state.apply( + logits, + expanded_idx_mapping, + idx_mapping, + idx_mapping_np, + input_ids, + expanded_local_pos, + ) + # Apply temperature in place. self.sampling_states.apply_temperature( logits, expanded_idx_mapping, idx_mapping_np @@ -219,6 +239,7 @@ class Sampler: self, logits: torch.Tensor, expanded_idx_mapping: torch.Tensor, + idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray, pos: torch.Tensor, input_ids: torch.Tensor, @@ -228,6 +249,7 @@ class Sampler: processed_logits = self.apply_sampling_params( logits, expanded_idx_mapping, + idx_mapping, idx_mapping_np, pos, input_ids, diff --git a/vllm/v1/worker/gpu/sample/thinking_budget.py b/vllm/v1/worker/gpu/sample/thinking_budget.py new file mode 100644 index 0000000000..eca7929860 --- /dev/null +++ b/vllm/v1/worker/gpu/sample/thinking_budget.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + +import numpy as np +import torch + +from vllm.sampling_params import SamplingParams +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor +from vllm.v1.worker.gpu.states import RequestState + +if TYPE_CHECKING: + from vllm.config.reasoning import ReasoningConfig + +_INT32_MAX = np.iinfo(np.int32).max +_COLD_SCAN_BLOCK = 1024 + + +class ThinkingBudgetState: + """Model Runner V2 state for per-request thinking token budgets.""" + + def __init__( + self, + req_states: RequestState, + reasoning_config: "ReasoningConfig | None", + ): + self.req_states = req_states + self.max_num_reqs = req_states.max_num_reqs + self.device = req_states.device + + start_ids = ( + [] + if reasoning_config is None + else reasoning_config.reasoning_start_token_ids or [] + ) + end_ids = ( + [] + if reasoning_config is None + else reasoning_config.reasoning_end_token_ids or [] + ) + natural_end_ids = ( + [] + if reasoning_config is None + else reasoning_config.natural_reasoning_end_token_ids or [] + ) + self.enabled = bool(start_ids and end_ids and natural_end_ids) + if not self.enabled: + return + + self.thinking_token_budget = UvaBackedTensor( + self.max_num_reqs, dtype=torch.int32 + ) + self.thinking_token_budget.np.fill(-1) + self.thinking_token_budget.copy_to_uva() + self.use_thinking_budget = np.zeros(self.max_num_reqs, dtype=bool) + + self.cached_last_start = torch.full( + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device + ) + self.cached_last_end = torch.full( + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device + ) + self.cached_scan_pos = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self._reset_reqs: list[int] = [] + self._budget_dirty = False + + self.reasoning_start_token_ids = torch.tensor( + start_ids, dtype=torch.int32, device=self.device + ) + self.natural_reasoning_end_token_ids = torch.tensor( + natural_end_ids, dtype=torch.int32, device=self.device + ) + self.reasoning_end_token_ids = torch.tensor( + end_ids, dtype=torch.int32, device=self.device + ) + + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: + if not self.enabled: + return + budget = sampling_params.thinking_token_budget + self.use_thinking_budget[req_idx] = budget is not None + if budget is None: + budget = -1 + else: + budget = min(budget, _INT32_MAX) + self._reset_reqs.append(req_idx) + if self.thinking_token_budget.np[req_idx] != budget: + self.thinking_token_budget.np[req_idx] = budget + self._budget_dirty = True + + def apply_staged_writes(self) -> None: + if not self.enabled: + return + if self._reset_reqs: + idx = async_tensor_h2d( + self._reset_reqs, dtype=torch.int64, device=self.device + ) + self.cached_last_start.index_fill_(0, idx, -1) + self.cached_last_end.index_fill_(0, idx, -1) + self.cached_scan_pos.index_fill_(0, idx, 0) + self._reset_reqs.clear() + if self._budget_dirty: + self.thinking_token_budget.copy_to_uva() + self._budget_dirty = False + + def apply( + self, + logits: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + idx_mapping: torch.Tensor, + idx_mapping_np: np.ndarray, + input_ids: torch.Tensor, + expanded_local_pos: torch.Tensor, + ) -> None: + if not self.enabled or not np.any(self.use_thinking_budget[idx_mapping_np]): + return + + apply_thinking_budget( + logits, + idx_mapping, + expanded_idx_mapping, + self.thinking_token_budget.gpu, + self.req_states.all_token_ids.gpu, + self.req_states.total_len.gpu, + input_ids, + expanded_local_pos, + self.cached_last_start, + self.cached_last_end, + self.cached_scan_pos, + self.reasoning_start_token_ids, + self.natural_reasoning_end_token_ids, + self.reasoning_end_token_ids, + ) + + +@triton.jit +def _load_effective_token( + all_token_ids_ptr, + all_token_ids_stride, + input_ids_ptr, + cur_req_first_pos, + req_state_idx, + total_len, + pos, +): + if pos < total_len: + return tl.load(all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos) + # In decode/spec-decode, input_ids at local position 0 is the already + # committed last sampled token. Effective draft-prefix positions start at + # local position 1. + input_pos = cur_req_first_pos + pos - total_len + 1 + return tl.load(input_ids_ptr + input_pos) + + +@triton.jit +def _update_committed_marker_cache_kernel( + req_ids_ptr, + thinking_token_budget_ptr, + all_token_ids_ptr, + all_token_ids_stride, + total_len_ptr, + cached_last_start_ptr, + cached_last_end_ptr, + cached_scan_pos_ptr, + reasoning_start_token_ids_ptr, + natural_reasoning_end_token_ids_ptr, + START_LEN: tl.constexpr, + NATURAL_END_LEN: tl.constexpr, + MAX_LEN: tl.constexpr, + BLOCK: tl.constexpr, +): + req_state_idx = tl.load(req_ids_ptr + tl.program_id(0)) + budget = tl.load(thinking_token_budget_ptr + req_state_idx) + if budget < 0: + return + + total_len = tl.load(total_len_ptr + req_state_idx) + scan_pos = tl.load(cached_scan_pos_ptr + req_state_idx) + last_start = tl.load(cached_last_start_ptr + req_state_idx) + last_end = tl.load(cached_last_end_ptr + req_state_idx) + + if scan_pos > total_len: + scan_pos = 0 + last_start = -1 + last_end = -1 + + if scan_pos == 0 and last_start < 0 and last_end < 0: + # Cold scan: walk backward in vectorized blocks, stopping at the first + # block with a marker; only the relative order of the two positions + # found matters below. + block_hi = total_len + while block_hi > 0 and last_start < 0 and last_end < 0: + block_lo = block_hi - BLOCK + if block_lo < 0: + block_lo = 0 + offs = block_lo + tl.arange(0, BLOCK) + + start_match = (offs < block_hi) & (offs + START_LEN <= total_len) + for j in tl.static_range(0, START_LEN): + expected = tl.load(reasoning_start_token_ids_ptr + j) + actual = tl.load( + all_token_ids_ptr + req_state_idx * all_token_ids_stride + offs + j, + mask=offs + j < total_len, + other=-1, + ) + start_match = start_match & (actual == expected) + + end_match = (offs < block_hi) & (offs + NATURAL_END_LEN <= total_len) + for j in tl.static_range(0, NATURAL_END_LEN): + expected = tl.load(natural_reasoning_end_token_ids_ptr + j) + actual = tl.load( + all_token_ids_ptr + req_state_idx * all_token_ids_stride + offs + j, + mask=offs + j < total_len, + other=-1, + ) + end_match = end_match & (actual == expected) + + last_start = tl.max(tl.where(start_match, offs, -1), axis=0) + last_end = tl.max(tl.where(end_match, offs, -1), axis=0) + block_hi = block_lo + else: + for i in tl.range(scan_pos, total_len): + if i + START_LEN <= total_len: + start_match = True + for j in tl.static_range(0, START_LEN): + expected = tl.load(reasoning_start_token_ids_ptr + j) + actual = tl.load( + all_token_ids_ptr + req_state_idx * all_token_ids_stride + i + j + ) + start_match = start_match & (actual == expected) + if start_match: + last_start = i + + if i + NATURAL_END_LEN <= total_len: + end_match = True + for j in tl.static_range(0, NATURAL_END_LEN): + expected = tl.load(natural_reasoning_end_token_ids_ptr + j) + actual = tl.load( + all_token_ids_ptr + req_state_idx * all_token_ids_stride + i + j + ) + end_match = end_match & (actual == expected) + if end_match: + last_end = i + + tl.store(cached_last_start_ptr + req_state_idx, last_start) + tl.store(cached_last_end_ptr + req_state_idx, last_end) + new_scan_pos = total_len - (MAX_LEN - 1) + if new_scan_pos < 0: + new_scan_pos = 0 + tl.store(cached_scan_pos_ptr + req_state_idx, new_scan_pos) + + +@triton.jit +def _thinking_budget_kernel( + logits_ptr, + logits_stride, + expanded_idx_mapping_ptr, + thinking_token_budget_ptr, + all_token_ids_ptr, + all_token_ids_stride, + total_len_ptr, + input_ids_ptr, + expanded_local_pos_ptr, + cached_last_start_ptr, + cached_last_end_ptr, + reasoning_start_token_ids_ptr, + natural_reasoning_end_token_ids_ptr, + reasoning_end_token_ids_ptr, + START_LEN: tl.constexpr, + NATURAL_END_LEN: tl.constexpr, + END_LEN: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) + budget = tl.load(thinking_token_budget_ptr + req_state_idx) + if budget < 0: + return + + local_pos = tl.load(expanded_local_pos_ptr + token_idx) + cur_req_first_pos = token_idx - local_pos + total_len = tl.load(total_len_ptr + req_state_idx) + effective_len = total_len + local_pos + + last_start = tl.load(cached_last_start_ptr + req_state_idx) + last_end = tl.load(cached_last_end_ptr + req_state_idx) + + start_lo = total_len - START_LEN + 1 + if start_lo < 0: + start_lo = 0 + for i in tl.range(start_lo, effective_len - START_LEN + 1): + start_match = True + for j in tl.static_range(0, START_LEN): + expected = tl.load(reasoning_start_token_ids_ptr + j) + actual = _load_effective_token( + all_token_ids_ptr, + all_token_ids_stride, + input_ids_ptr, + cur_req_first_pos, + req_state_idx, + total_len, + i + j, + ) + start_match = start_match & (actual == expected) + if start_match: + last_start = i + + end_lo = total_len - NATURAL_END_LEN + 1 + if end_lo < 0: + end_lo = 0 + for i in tl.range(end_lo, effective_len - NATURAL_END_LEN + 1): + end_match = True + for j in tl.static_range(0, NATURAL_END_LEN): + expected = tl.load(natural_reasoning_end_token_ids_ptr + j) + actual = _load_effective_token( + all_token_ids_ptr, + all_token_ids_stride, + input_ids_ptr, + cur_req_first_pos, + req_state_idx, + total_len, + i + j, + ) + end_match = end_match & (actual == expected) + if end_match: + last_end = i + + if last_start < 0 or last_start <= last_end: + return + + reasoning_start = last_start + START_LEN + # If the request resumes from a prompt that already contains generated + # reasoning content, count it against the remaining budget. + num_reasoning_tokens = effective_len - reasoning_start + if num_reasoning_tokens < budget: + return + + # If the tail already ends with a prefix of the forced end sequence + # (even from a resumed prompt), continue from the next marker token. + end_prefix_len = 0 + max_prefix_len = END_LEN - 1 + if effective_len < max_prefix_len: + max_prefix_len = effective_len + + for prefix_len in tl.static_range(1, END_LEN): + if prefix_len <= max_prefix_len: + prefix_match = True + suffix_start = effective_len - prefix_len + for j in tl.static_range(0, END_LEN): + if j < prefix_len: + expected = tl.load(reasoning_end_token_ids_ptr + j) + actual = _load_effective_token( + all_token_ids_ptr, + all_token_ids_stride, + input_ids_ptr, + cur_req_first_pos, + req_state_idx, + total_len, + suffix_start + j, + ) + prefix_match = prefix_match & (actual == expected) + if prefix_match: + end_prefix_len = prefix_len + + force_token_id = tl.load(reasoning_end_token_ids_ptr + end_prefix_len) + tl.store(logits_ptr + token_idx * logits_stride + force_token_id, 1.0e9) + + +def apply_thinking_budget( + logits: torch.Tensor, + req_ids: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + thinking_token_budget: torch.Tensor, + all_token_ids: torch.Tensor, + total_len: torch.Tensor, + input_ids: torch.Tensor, + expanded_local_pos: torch.Tensor, + cached_last_start: torch.Tensor, + cached_last_end: torch.Tensor, + cached_scan_pos: torch.Tensor, + reasoning_start_token_ids: torch.Tensor, + natural_reasoning_end_token_ids: torch.Tensor, + reasoning_end_token_ids: torch.Tensor, +) -> None: + num_tokens = logits.shape[0] + start_len = reasoning_start_token_ids.shape[0] + natural_end_len = natural_reasoning_end_token_ids.shape[0] + end_len = reasoning_end_token_ids.shape[0] + + _update_committed_marker_cache_kernel[(req_ids.shape[0],)]( + req_ids, + thinking_token_budget, + all_token_ids, + all_token_ids.stride(0), + total_len, + cached_last_start, + cached_last_end, + cached_scan_pos, + reasoning_start_token_ids, + natural_reasoning_end_token_ids, + START_LEN=start_len, + NATURAL_END_LEN=natural_end_len, + MAX_LEN=max(start_len, natural_end_len), + BLOCK=_COLD_SCAN_BLOCK, + ) + + _thinking_budget_kernel[(num_tokens,)]( + logits, + logits.stride(0), + expanded_idx_mapping, + thinking_token_budget, + all_token_ids, + all_token_ids.stride(0), + total_len, + input_ids, + expanded_local_pos, + cached_last_start, + cached_last_end, + reasoning_start_token_ids, + natural_reasoning_end_token_ids, + reasoning_end_token_ids, + START_LEN=start_len, + NATURAL_END_LEN=natural_end_len, + END_LEN=end_len, + ) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 63d479efa0..a4be95be00 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -141,6 +141,7 @@ class RejectionSampler: processed_logits = self.sampler.apply_sampling_params( logits, expanded_idx_mapping, + idx_mapping, idx_mapping_np, pos, draft_sampled,