iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Feat] Support thinking_token_budget in Model Runner V2 (#46727)

Browse Source 
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Chauncey
2026-08-07 03:36:18 +00:00
committed by GitHub
co-authored by Nick Hill
parent 5ac2684976
commit 72c0d67657
10 changed files with 1122 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/kernels/benchmark_thinking_budget.py
+365
View File
@@ -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()
tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
+3 -8
View File
@@ -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:
tests/v1/worker/test_gpu_thinking_budget.py
+262
View File
@@ -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
vllm/config/reasoning.py
+31 -13
View File
@@ -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. `"<think>"` and `"</think>"`). 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}', "
vllm/config/vllm.py
-6
View File
@@ -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.
vllm/v1/engine/input_processor.py
+9 -16
View File
@@ -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
vllm/v1/worker/gpu/model_runner.py
+1
View File
@@ -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)
vllm/v1/worker/gpu/sample/sampler.py
+22
View File
@@ -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,
vllm/v1/worker/gpu/sample/thinking_budget.py
+428
View File
@@ -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,
)
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+1
View File
@@ -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,