@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 68dd7db810

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 13:45 UTC

[Reasoning] Support for speculative decoding with thinking budget (#34668)

Signed-off-by: rishitdholakia13 <rishit+github@cohere.com> Signed-off-by: rishitdholakia13 <123388671+rishitdholakia13@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

rishitdholakia13 committed Apr 29, 2026 at 02:14 UTC 68dd7db81001267c846907769adc14bb32566190

12 files changed +1143 -365

tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py

+181 -4

@@ -1,18 +1,62 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -"""E2E tests for thinking_token_budget with reasoning models."""

4 +"""E2E tests for ``thinking_token_budget`` with reasoning models.

5 +

6 +Covers Qwen3-0.6B and Qwen3.5 FP8 + MTP.

7 +"""

8 +

9 +import asyncio

10 +import json

11 +from typing import Literal

13 import openai

14 import pytest

15 import pytest_asyncio

10 -from tests.utils import RemoteOpenAIServer

17 +from tests.utils import RemoteOpenAIServer, multi_gpu_only, requires_fp8

18 +from vllm.platforms import current_platform

19 +from vllm.tokenizers import get_tokenizer

21 MODEL_NAME = "Qwen/Qwen3-0.6B"

22 +QWEN35_FP8_MTP_MODEL = "Qwen/Qwen3.5-35B-A3B-FP8"

23 MESSAGES = [{"role": "user", "content": "What is 1+1? Be concise."}]

24 THINK_BUDGET = 5

26 +REASONING_START_STR = "<think>"

27 +REASONING_END_STR = "</think>"

28 +

29 +

30 +def _count_reasoning_decode_token_ids_between_markers(

31 + full_token_ids: list[int],

32 + reasoning_start_ids: list[int],

33 + reasoning_end_ids: list[int],

34 +) -> int | None:

35 + """Count decode tokens in the thinking span (after last start, before first end)."""

36 +

37 + if not reasoning_start_ids or not reasoning_end_ids:

38 + raise ValueError("reasoning marker token id lists must be non-empty")

39 +

40 + def _last_subseq_index(haystack: list[int], needle: list[int]) -> int:

41 + n = len(needle)

42 + if n > len(haystack):

43 + return -1

44 + for i in range(len(haystack) - n, -1, -1):

45 + if haystack[i : i + n] == needle:

46 + return i

47 + return -1

48 +

49 + last_start = _last_subseq_index(full_token_ids, reasoning_start_ids)

50 + if last_start < 0:

51 + return None

52 +

53 + pos_after_start = last_start + len(reasoning_start_ids)

54 + end_n = len(reasoning_end_ids)

55 + for j in range(pos_after_start, len(full_token_ids) - end_n + 1):

56 + if full_token_ids[j : j + end_n] == reasoning_end_ids:

57 + return j - pos_after_start

58 + return len(full_token_ids) - pos_after_start

59 +

61 @pytest.fixture(scope="module")

62 def server():

@@ -48,6 +92,51 @@ def server_with_auto_reasoning_config():

92 yield remote_server

95 +@pytest.fixture(scope="module")

96 +def server_qwen35_fp8_mtp_tp2():

97 + """Qwen3.5-35B FP8 with MTP speculative decoding and tensor parallel size 2."""

98 + if current_platform.device_count() < 2:

99 + pytest.skip("Need at least 2 GPUs for --tensor-parallel-size 2")

100 + if not current_platform.supports_fp8():

101 + pytest.skip("FP8 is not supported on this platform")

102 +

103 + spec_cfg = {

104 + "method": "mtp",

105 + "num_speculative_tokens": 2,

106 + "max_model_len": 32768,

107 + }

108 + args = [

109 + "--tensor-parallel-size",

110 + "2",

111 + "--max-model-len",

112 + "32768",

113 + "--speculative-config",

114 + json.dumps(spec_cfg),

115 + "--reasoning-parser",

116 + "qwen3",

117 + "--reasoning-config",

118 + json.dumps(

119 + {

120 + "reasoning_start_str": REASONING_START_STR,

121 + "reasoning_end_str": REASONING_END_STR,

122 + }

123 + ),

124 + ]

125 + # With 4+ GPUs, run TP=2 on physical devices 2,3 so module-scoped 0.6B servers

126 + # on 0,1 do not exhaust memory on the same devices as this worker.

127 + env_dict = None

128 + if current_platform.device_count() >= 4:

129 + env_dict = {"CUDA_VISIBLE_DEVICES": "2,3"}

130 +

131 + with RemoteOpenAIServer(

132 + QWEN35_FP8_MTP_MODEL,

133 + args,

134 + max_wait_seconds=3000,

135 + env_dict=env_dict,

136 + ) as remote_server:

137 + yield remote_server

138 +

139 +

140 @pytest_asyncio.fixture

141 async def client(request, server, server_with_auto_reasoning_config):

142 server_map = {

@@ -89,8 +178,10 @@ async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):

178 async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI):

179 """Test that thinking_token_budget limits the number of reasoning tokens.

180

92 - In streaming mode each reasoning delta corresponds to one token, so

93 - counting non-empty reasoning_content chunks gives the exact token count.

181 + Counts non-empty streaming ``delta.reasoning`` chunks (coarse proxy; each

182 + chunk may represent multiple decode tokens — see

183 + ``_count_reasoning_decode_token_ids_between_markers`` and the Qwen3.5 MTP

184 + test for id-based checks).

185 """

186

187 reasoning_token_count = 0

@@ -110,3 +201,89 @@ async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI

201 f"reasoning tokens ({reasoning_token_count}) exceeded "

202 f"thinking_token_budget ({THINK_BUDGET})"

203 )

204 +

205 +

206 +@pytest.mark.asyncio

207 +@multi_gpu_only(num_gpus=2)

208 +@requires_fp8

209 +async def test_thinking_token_budget_qwen35_fp8_mtp_concurrent_mixed_budget_and_plain(

210 + server_qwen35_fp8_mtp_tp2,

211 +):

212 + """Concurrent chat requests: some with ``thinking_token_budget``, some without.

213 +

214 + Exercises the scheduler / input processor under a mixed batch on the same

215 + Qwen3.5 FP8 + MTP (TP=2) server. Budgeted calls are checked with

216 + ``_count_reasoning_decode_token_ids_between_markers`` on full token ids.

217 + """

218 +

219 + _batch_spec: list[tuple[Literal["budget"], int] | tuple[Literal["plain"], None]] = [

220 + ("budget", 1),

221 + ("budget", 12),

222 + ("plain", None),

223 + ("budget", 20),

224 + ("budget", 14),

225 + ("plain", None),

226 + ("plain", None),

227 + ("budget", 12),

228 + ("plain", None),

229 + ]

230 +

231 + tokenizer = get_tokenizer(tokenizer_name=QWEN35_FP8_MTP_MODEL)

232 + start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False))

233 + end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False))

234 +

235 + async with server_qwen35_fp8_mtp_tp2.get_async_client() as client:

236 +

237 + async def budgeted_call(expected_budget: int):

238 + return await client.chat.completions.create(

239 + model=QWEN35_FP8_MTP_MODEL,

240 + messages=MESSAGES,

241 + max_tokens=256,

242 + stream=False,

243 + extra_body={

244 + "thinking_token_budget": expected_budget,

245 + "return_token_ids": True,

246 + },

247 + )

248 +

249 + async def plain_call():

250 + return await client.chat.completions.create(

251 + model=QWEN35_FP8_MTP_MODEL,

252 + messages=MESSAGES,

253 + max_tokens=256,

254 + stream=False,

255 + )

256 +

257 + coros = []

258 + for row in _batch_spec:

259 + if row[0] == "budget":

260 + b = row[1]

261 + assert isinstance(b, int)

262 + coros.append(budgeted_call(b))

263 + else:

264 + coros.append(plain_call())

265 + results = await asyncio.gather(*coros)

266 +

267 + for i, (response, (kind, expected_budget)) in enumerate(

268 + zip(results, _batch_spec, strict=True)

269 + ):

270 + msg = response.choices[0].message

271 + assert msg.content or getattr(msg, "reasoning", None), (

272 + f"index {i} ({kind}): empty message"

273 + )

274 +

275 + if kind == "budget":

276 + assert expected_budget is not None

277 + assert response.prompt_token_ids is not None

278 + assert response.choices[0].token_ids is not None

279 + full_ids = list(response.prompt_token_ids) + list(

280 + response.choices[0].token_ids

281 + )

282 + n_reason = _count_reasoning_decode_token_ids_between_markers(

283 + full_ids, start_ids, end_ids

284 + )

285 + assert n_reason is not None, f"index {i}: missing reasoning start in ids"

286 + assert n_reason == expected_budget, (

287 + f"index {i}: reasoning decode token ids ({n_reason}) != "

288 + f"thinking_token_budget ({expected_budget})"

289 + )

tests/v1/logits_processors/test_correctness.py

+338 -85

@@ -30,10 +30,13 @@ from vllm.v1.sample.logits_processor import (

30 MinPLogitsProcessor,

31 MinTokensLogitsProcessor,

32 MoveDirectionality,

33 - ThinkingTokenBudgetLogitsProcessor,

33 build_logitsprocs,

34 )

35 from vllm.v1.sample.metadata import SamplingMetadata

36 +from vllm.v1.sample.thinking_budget_state import (

37 + ThinkingBudgetStateHolder,

38 + maybe_create_thinking_budget_state_holder,

39 +)

41 PIN_MEMORY_AVAILABLE = is_pin_memory_available()

42 MAX_NUM_REQS = 256

@@ -48,8 +51,10 @@ MAX_NUM_PROMPT_TOKENS = 64

51 MIN_TOKENS_LEN_THRESHOLD = 5

52 REQS_PER_LOGITPROC = 50

53 STR_NO_LOGITPROC = "none"

54 +# Thinking budget uses ``ThinkingBudgetStateHolder`` (not a logits processor).

55 +STR_THINKING_BUDGET = "thinking_budget"

52 -# ThinkingTokenBudgetLogitsProcessor testing constants

57 +# Thinking token budget testing constants

58 THINKING_TOKEN_BUDGET = 5

59 THINK_START_TOKEN_ID = 999

60 THINK_END_TOKEN_ID = 998

@@ -80,15 +85,8 @@ class LogitsProcsRequestParams:

85 if num_tokens > 0:

86 # Use diverse random tokens

87 self.out_tokens = [random.randint(1, 950) for _ in range(num_tokens)]

83 - # Set first token for ThinkingTokenBudget testing

84 - is_thinking_processor = (

85 - logitproc_type is ThinkingTokenBudgetLogitsProcessor

86 - or (

87 - hasattr(logitproc_type, "__name__")

88 - and logitproc_type.__name__ == "ThinkingTokenBudgetLogitsProcessor"

89 - )

90 - )

91 - if is_thinking_processor:

88 + # Think-start seed for ``STR_THINKING_BUDGET`` rows.

89 + if logitproc_type == STR_THINKING_BUDGET:

90 self.out_tokens[0] = THINK_START_TOKEN_ID

91 else:

92 self.out_tokens = []

@@ -102,7 +100,7 @@ class LogitsProcsRequestParams:

100

101

102 class MockReasoningConfig:

105 - """Mock reasoning config for testing ThinkingTokenBudgetLogitsProcessor."""

103 + """Minimal reasoning config for ``ThinkingBudgetStateHolder`` tests."""

104

105 reasoning_start_token_ids = [THINK_START_TOKEN_ID]

106 reasoning_end_token_ids = [THINK_END_TOKEN_ID]

@@ -137,6 +135,18 @@ def _generate_fake_sampling_metadata(

135 is_pin_memory=PIN_MEMORY_AVAILABLE,

136 is_pooling_model=False,

137 )

138 + num_spec = (

139 + vllm_config.speculative_config.num_speculative_tokens

140 + if vllm_config.speculative_config

141 + else 0

142 + )

143 + thinking_holder = maybe_create_thinking_budget_state_holder(

144 + vllm_config.reasoning_config,

145 + vllm_config.scheduler_config.max_num_seqs,

146 + num_spec,

147 + device,

148 + PIN_MEMORY_AVAILABLE,

149 + )

150 fake_sampling_metadata = SamplingMetadata(

151 temperature=torch.full((batch_size,), 0.0),

152 all_greedy=True,

@@ -156,6 +166,7 @@ def _generate_fake_sampling_metadata(

166 allowed_token_ids_mask=None,

167 bad_words_token_ids={},

168 logitsprocs=logitsprocs,

169 + thinking_budget_state_holder=thinking_holder,

170 )

171 return fake_sampling_metadata

172

@@ -187,7 +198,7 @@ def _sampling_params_from_logitproc(logitproc_type: LogitprocType) -> SamplingPa

198

199 def _generate_mixed_logitsprocs_batch_params(

200 reqs_per_logitproc: int,

190 - logitsprocs_types: list[str],

201 + logitsprocs_types: list[LogitprocType],

202 ) -> list[LogitsProcsRequestParams]:

203 """Define key params for a batch of requests with a different

204 logitproc enabled per request.

@@ -450,23 +461,21 @@ def _thinking_budget_validate(

461 request_params: LogitsProcsRequestParams,

462 step_idx: int,

463 ) -> None:

453 - """Validate thinking token budget processor behavior"""

454 - # Get the ThinkingTokenBudgetLogitsProcessor instance

455 - tb_processor: ThinkingTokenBudgetLogitsProcessor = next(

456 - test_fakes.get_logitsprocs_by_cls(ThinkingTokenBudgetLogitsProcessor)

457 - )

464 + """Validate ``ThinkingBudgetStateHolder`` thinking-budget behavior.

465

459 - # Get current request state

460 - state = tb_processor._state.get(batch_index)

466 + State is keyed by **batch slot** (same index space as logits rows), matching

467 + ``sync_batch`` / sampler integration (see PR #34668 discussion).

468 + """

469 + holder = test_fakes.sampling_metadata.thinking_budget_state_holder

470 + assert holder is not None

471 + state = holder._state.get(batch_index)

472 params = request_params.params

473

463 - # Validate thinking token budget configuration

474 if hasattr(params, "thinking_token_budget") and params.thinking_token_budget:

465 - # State should exist for requests with thinking_token_budget

475 if state is None:

476 _raise_error_invalid(

477 msg_suffix=(

469 - f"Expected state for batch {batch_index} "

478 + f"Expected holder state for batch slot {batch_index} "

479 f"with thinking_token_budget={params.thinking_token_budget}"

480 ),

481 batch_index=batch_index,

@@ -474,10 +483,8 @@ def _thinking_budget_validate(

483 step_idx=step_idx,

484 )

485

477 - # Validate budget matches what was set

486 expected_budget = params.thinking_token_budget

487 actual_budget = state["thinking_token_budget"]

480 -

488 if actual_budget != expected_budget:

489 _raise_error_invalid(

490 msg_suffix=(

@@ -488,13 +495,9 @@ def _thinking_budget_validate(

495 step_idx=step_idx,

496 )

497

491 - # Check if we're in thinking mode and validate token counting

498 output_tokens = request_params.out_tokens

493 -

494 - # Find if thinking has started in output tokens

499 + start_tokens = holder.think_start_token_ids

500 thinking_started = False

496 - start_tokens = tb_processor.reasoning_start_token_ids

497 -

501 if len(start_tokens) > 0:

502 for i in range(len(output_tokens) - len(start_tokens) + 1):

503 if output_tokens[i : i + len(start_tokens)] == start_tokens:

@@ -502,61 +505,42 @@ def _thinking_budget_validate(

505 break

506

507 if thinking_started:

505 - # If budget is exceeded, validate end token forcing

508 think_count = state["think_count"]

509 budget = state["thinking_token_budget"]

510 + if think_count >= budget and not state["in_end"]:

511 + _raise_error_invalid(

512 + msg_suffix=(

513 + f"Budget exceeded ({think_count} >= {budget}) but "

514 + "in_end is false"

515 + ),

516 + batch_index=batch_index,

517 + request_params=request_params,

518 + step_idx=step_idx,

519 + )

520

509 - if think_count >= budget:

510 - if not state["in_end"]:

521 + end_tokens = holder.think_end_token_ids

522 + if (

523 + think_count >= budget

524 + and state["in_end"]

525 + and len(end_tokens) > 0

526 + and holder.has_tracked_requests()

527 + ):

528 + expected_end_token_id = end_tokens[

529 + min(state["end_count"], len(end_tokens) - 1)

530 + ]

531 + # Holder bumps forced vocab positions to 1e9 (does not -inf others).

532 + forced_logit = float(logits_new[batch_index, expected_end_token_id])

533 + if forced_logit < 1.0e8:

534 _raise_error_invalid(

535 msg_suffix=(

513 - f"Budget exceeded ({think_count} >= "

514 - f"{budget}) but not "

515 - "forcing end tokens"

536 + f"Expected forced end token {expected_end_token_id} "

537 + f"with large logit, got {forced_logit}"

538 ),

539 batch_index=batch_index,

540 request_params=request_params,

541 step_idx=step_idx,

542 )

543

522 - # Validate that only end tokens are allowed

523 - end_tokens = tb_processor.reasoning_end_token_ids

524 - if len(end_tokens) > 0:

525 - expected_end_token_id = end_tokens[

526 - min(state["end_count"], len(end_tokens) - 1)

527 - ]

528 -

529 - # Check logits masking

530 - batch_logits = logits_new[batch_index]

531 - for token_id in range(len(batch_logits)):

532 - logit_value = batch_logits[token_id]

533 -

534 - if token_id == expected_end_token_id:

535 - # End token should not be masked

536 - if logit_value == -float("inf"):

537 - _raise_error_invalid(

538 - msg_suffix=(

539 - f"End token {token_id} should not be "

540 - "masked but is"

541 - ),

542 - batch_index=batch_index,

543 - request_params=request_params,

544 - step_idx=step_idx,

545 - )

546 - else:

547 - # All other tokens should be masked when forcing end

548 - if logit_value != -float("inf"):

549 - _raise_error_invalid(

550 - msg_suffix=(

551 - f"Token {token_id} should be masked "

552 - f"when forcing end tokens, but "

553 - f"logit={logit_value}"

554 - ),

555 - batch_index=batch_index,

556 - request_params=request_params,

557 - step_idx=step_idx,

558 - )

559 -

544

545 def _none_validate(

546 test_fakes: LogitsprocsTestFakes,

@@ -604,7 +588,7 @@ logitsprocs_test_mapping = {

588 MinTokensLogitsProcessor: LogitsprocTestHelpers(

589 gen_request_fxn=_min_tokens_params, eval_fxn=_min_tokens_validate

590 ),

607 - ThinkingTokenBudgetLogitsProcessor: LogitsprocTestHelpers(

591 + STR_THINKING_BUDGET: LogitsprocTestHelpers(

592 gen_request_fxn=_thinking_budget_params, eval_fxn=_thinking_budget_validate

593 ),

594 }

@@ -614,20 +598,17 @@ def _get_test_cases() -> list[list[str]]:

598 """Each test case is a set of logitsprocs"""

599 logitsprocs_types = list(logitsprocs_test_mapping.keys())

600

617 - # Isolate ThinkingTokenBudgetLogitsProcessor from all other processors

618 - # to avoid unexpected modification of logits interference

619 - thinking_processor = ThinkingTokenBudgetLogitsProcessor

601 + # Isolate thinking-budget handling from other processors to avoid cross-talk.

602 + thinking_id: LogitprocType = STR_THINKING_BUDGET

603 other_processors = [

621 - p

622 - for p in logitsprocs_types

623 - if p != STR_NO_LOGITPROC and p != thinking_processor

604 + p for p in logitsprocs_types if p != STR_NO_LOGITPROC and p != thinking_id

605 ]

606

607 return (

608 [[STR_NO_LOGITPROC]]

609 + [[logitproc_type, STR_NO_LOGITPROC] for logitproc_type in other_processors]

610 + [other_processors]

630 - + [[thinking_processor]]

611 + + [[thinking_id]]

612 )

613

614

@@ -802,12 +783,23 @@ def _assert_valid(

783 )

784

785

786 +def _slot_outputs_for_metadata(

787 + persistent_batch: list[LogitsProcsRequestParams], pad_len: int

788 +) -> list[list[int]]:

789 + """Per-batch-slot output token ids aligned with ``SamplingMetadata`` rows."""

790 + rows: list[list[int]] = [[] for _ in range(pad_len)]

791 + for i, req in enumerate(persistent_batch):

792 + if i < pad_len:

793 + rows[i] = list(req.out_tokens)

794 + return rows

795 +

796 +

797 @create_new_process_for_each_test()

798 @pytest.mark.parametrize("device", DEVICES)

799 @pytest.mark.parametrize("reqs_per_logitproc", [REQS_PER_LOGITPROC])

800 @pytest.mark.parametrize("logitsprocs_under_test", _get_test_cases())

801 def test_logitsprocs(

810 - device: str, reqs_per_logitproc: int, logitsprocs_under_test: list[str]

802 + device: str, reqs_per_logitproc: int, logitsprocs_under_test: list[LogitprocType]

803 ):

804 random.seed(40)

805 torch.set_default_device(device)

@@ -855,9 +847,10 @@ def test_logitsprocs(

847 # Apply fake batch update to logitsprocs

848 fake_update_logitsprocs_state(test_fakes, batch_update)

849

858 - # Emulate application of logits processors in engine

850 + # Emulate application of logits processors + thinking holder (sampler order).

851 slice_idxs = [req.workload_index for req in persistent_batch]

860 - logits_w_lp = fake_apply_logitsprocs(test_fakes, slice_idxs).cpu()

852 + slot_rows = _slot_outputs_for_metadata(persistent_batch, workload_size)

853 + logits_w_lp = fake_apply_logitsprocs(test_fakes, slice_idxs, slot_rows).cpu()

854

855 _assert_valid(

856 batch_size=batch_size,

@@ -869,3 +862,263 @@ def test_logitsprocs(

862 )

863

864 step_idx += 1

865 +

866 +

867 +class MockReasoningNoEndTokens:

868 + """Reasoning config with no end token ids (disables enforcement in holder)."""

869 +

870 + reasoning_start_token_ids = [THINK_START_TOKEN_ID]

871 + reasoning_end_token_ids: list[int] = []

872 +

873 +

874 +def test_maybe_create_thinking_budget_holder_without_reasoning():

875 + cfg = VllmConfig()

876 + assert cfg.reasoning_config is None

877 + assert (

878 + maybe_create_thinking_budget_state_holder(

879 + None,

880 + cfg.scheduler_config.max_num_seqs,

881 + 0,

882 + torch.device("cpu"),

883 + False,

884 + )

885 + is None

886 + )

887 +

888 +

889 +def test_thinking_budget_holder_has_tracked_after_sync_add():

890 + vc = VllmConfig()

891 + vc.reasoning_config = MockReasoningConfig()

892 + h = ThinkingBudgetStateHolder(

893 + vc.reasoning_config,

894 + vc.scheduler_config.max_num_seqs,

895 + 0,

896 + torch.device("cpu"),

897 + False,

898 + )

899 + assert not h.has_tracked_requests()

900 + h.sync_batch(

901 + BatchUpdate(

902 + batch_size=1,

903 + removed=(),

904 + added=[

905 + (

906 + 0,

907 + SamplingParams(thinking_token_budget=3),

908 + None,

909 + [THINK_START_TOKEN_ID],

910 + )

911 + ],

912 + moved=(),

913 + )

914 + )

915 + assert h.has_tracked_requests()

916 + assert h._state[0]["thinking_token_budget"] == 3

917 +

918 +

919 +def test_thinking_budget_holder_sync_remove_clears_state():

920 + vc = VllmConfig()

921 + vc.reasoning_config = MockReasoningConfig()

922 + h = ThinkingBudgetStateHolder(

923 + vc.reasoning_config,

924 + vc.scheduler_config.max_num_seqs,

925 + 0,

926 + torch.device("cpu"),

927 + False,

928 + )

929 + h.sync_batch(

930 + BatchUpdate(

931 + batch_size=1,

932 + removed=(),

933 + added=[

934 + (

935 + 0,

936 + SamplingParams(thinking_token_budget=3),

937 + None,

938 + [],

939 + )

940 + ],

941 + moved=(),

942 + )

943 + )

944 + assert h.has_tracked_requests()

945 + h.sync_batch(BatchUpdate(batch_size=0, removed=(0,), added=(), moved=()))

946 + assert not h.has_tracked_requests()

947 +

948 +

949 +def test_thinking_budget_holder_sync_add_without_budget_drops_row():

950 + vc = VllmConfig()

951 + vc.reasoning_config = MockReasoningConfig()

952 + h = ThinkingBudgetStateHolder(

953 + vc.reasoning_config,

954 + vc.scheduler_config.max_num_seqs,

955 + 0,

956 + torch.device("cpu"),

957 + False,

958 + )

959 + h.sync_batch(

960 + BatchUpdate(

961 + batch_size=1,

962 + removed=(),

963 + added=[(0, SamplingParams(), None, [])],

964 + moved=(),

965 + )

966 + )

967 + assert not h.has_tracked_requests()

968 +

969 +

970 +def test_thinking_budget_holder_swap_exchanges_state():

971 + vc = VllmConfig()

972 + vc.reasoning_config = MockReasoningConfig()

973 + h = ThinkingBudgetStateHolder(

974 + vc.reasoning_config,

975 + vc.scheduler_config.max_num_seqs,

976 + 0,

977 + torch.device("cpu"),

978 + False,

979 + )

980 + h.sync_batch(

981 + BatchUpdate(

982 + batch_size=2,

983 + removed=(),

984 + added=[

985 + (

986 + 0,

987 + SamplingParams(thinking_token_budget=3),

988 + None,

989 + [],

990 + ),

991 + (

992 + 1,

993 + SamplingParams(thinking_token_budget=7),

994 + None,

995 + [],

996 + ),

997 + ],

998 + moved=(),

999 + )

1000 + )

1001 + b0, b1 = h._state[0]["thinking_token_budget"], h._state[1]["thinking_token_budget"]

1002 + h.sync_batch(

1003 + BatchUpdate(

1004 + batch_size=2,

1005 + removed=(),

1006 + added=(),

1007 + moved=[(0, 1, MoveDirectionality.SWAP)],

1008 + )

1009 + )

1010 + assert h._state[0]["thinking_token_budget"] == b1

1011 + assert h._state[1]["thinking_token_budget"] == b0

1012 +

1013 +

1014 +def test_thinking_budget_holder_unidirectional_move():

1015 + vc = VllmConfig()

1016 + vc.reasoning_config = MockReasoningConfig()

1017 + h = ThinkingBudgetStateHolder(

1018 + vc.reasoning_config,

1019 + vc.scheduler_config.max_num_seqs,

1020 + 0,

1021 + torch.device("cpu"),

1022 + False,

1023 + )

1024 + h.sync_batch(

1025 + BatchUpdate(

1026 + batch_size=2,

1027 + removed=(),

1028 + added=[

1029 + (

1030 + 1,

1031 + SamplingParams(thinking_token_budget=4),

1032 + None,

1033 + [],

1034 + ),

1035 + ],

1036 + moved=(),

1037 + )

1038 + )

1039 + assert 1 in h._state and 0 not in h._state

1040 + h.sync_batch(

1041 + BatchUpdate(

1042 + batch_size=2,

1043 + removed=(),

1044 + added=(),

1045 + moved=[(1, 0, MoveDirectionality.UNIDIRECTIONAL)],

1046 + )

1047 + )

1048 + assert 0 in h._state and 1 not in h._state

1049 + assert h._state[0]["thinking_token_budget"] == 4

1050 +

1051 +

1052 +def test_thinking_budget_holder_update_state_repeat_indices_last_row_wins():

1053 + vc = VllmConfig()

1054 + vc.reasoning_config = MockReasoningConfig()

1055 + h = ThinkingBudgetStateHolder(

1056 + vc.reasoning_config,

1057 + vc.scheduler_config.max_num_seqs,

1058 + 0,

1059 + torch.device("cpu"),

1060 + False,

1061 + )

1062 + h.sync_batch(

1063 + BatchUpdate(

1064 + batch_size=1,

1065 + removed=(),

1066 + added=[

1067 + (

1068 + 0,

1069 + SamplingParams(thinking_token_budget=5),

1070 + None,

1071 + [THINK_START_TOKEN_ID],

1072 + )

1073 + ],

1074 + moved=(),

1075 + )

1076 + )

1077 + out_lists = [[THINK_START_TOKEN_ID], [THINK_START_TOKEN_ID, 10, 11, 12, 13, 14]]

1078 + h.update_state(

1079 + out_lists,

1080 + None,

1081 + torch.tensor([0, 0], dtype=torch.long),

1082 + )

1083 + assert h._state[0]["output_tok_ids"] == out_lists[1]

1084 +

1085 +

1086 +def test_thinking_budget_holder_spec_mode_tensor_layout():

1087 + h = ThinkingBudgetStateHolder(

1088 + MockReasoningConfig(),

1089 + 8,

1090 + 2,

1091 + torch.device("cpu"),

1092 + False,

1093 + )

1094 + assert h.in_spec_mode

1095 + assert h.mask.shape[0] == 8 * (2 + 1)

1096 +

1097 +

1098 +def test_thinking_budget_holder_empty_end_tokens_disables_row():

1099 + vc = VllmConfig()

1100 + vc.reasoning_config = MockReasoningNoEndTokens()

1101 + h = ThinkingBudgetStateHolder(

1102 + vc.reasoning_config,

1103 + vc.scheduler_config.max_num_seqs,

1104 + 0,

1105 + torch.device("cpu"),

1106 + False,

1107 + )

1108 + h.sync_batch(

1109 + BatchUpdate(

1110 + batch_size=1,

1111 + removed=(),

1112 + added=[

1113 + (

1114 + 0,

1115 + SamplingParams(thinking_token_budget=5),

1116 + None,

1117 + [THINK_START_TOKEN_ID],

1118 + )

1119 + ],

1120 + moved=(),

1121 + )

1122 + )

1123 + h.update_state([[THINK_START_TOKEN_ID, 1]], None, None)

1124 + assert h._state[0]["thinking_token_budget"] == -1

tests/v1/sample/utils.py

+23 -2

@@ -198,22 +198,43 @@ class LogitsprocsTestFakes(NamedTuple):

198

199 def fake_update_logitsprocs_state(

200 test_fakes: LogitsprocsTestFakes,

201 - batch_update: BatchUpdate,

201 + batch_update: BatchUpdate | None,

202 ) -> None:

203 """Imitate logits processors persistent batch state update

204 in engine core"""

205 for logitproc in test_fakes.get_logitsprocs():

206 logitproc.update_state(batch_update)

207 + holder = test_fakes.sampling_metadata.thinking_budget_state_holder

208 + if holder is not None:

209 + holder.sync_batch(batch_update)

210

211

212 def fake_apply_logitsprocs(

213 test_fakes: LogitsprocsTestFakes,

214 slice_indices: list[int],

215 + slot_output_token_ids: list[list[int]] | None = None,

216 ) -> torch.Tensor:

213 - """Imitate application of logits processors in engine core"""

217 + """Imitate application of logits processors in engine core.

218 +

219 + When ``thinking_budget_state_holder`` has tracked requests, this mirrors

220 + :meth:`Sampler.apply_logits_processors` by refreshing per-slot

221 + ``output_token_ids`` (if ``slot_output_token_ids`` is provided), then

222 + ``update_state`` + ``apply_to_logits`` on the holder after built-in logits

223 + processors.

224 + """

225 logits = test_fakes.logits[torch.tensor(slice_indices, dtype=torch.long)].clone()

226 for processor in test_fakes.get_logitsprocs():

227 logits = processor.apply(logits)

228 +

229 + md = test_fakes.sampling_metadata

230 + holder = md.thinking_budget_state_holder

231 + if holder is not None and holder.has_tracked_requests():

232 + if slot_output_token_ids is not None:

233 + for i, toks in enumerate(slot_output_token_ids):

234 + if i < len(md.output_token_ids):

235 + md.output_token_ids[i] = list(toks)

236 + holder.update_state(md.output_token_ids, md.spec_token_ids, None)

237 + logits = holder.apply_to_logits(logits, False, md.spec_token_ids)

238 return logits

239

240

tests/v1/streaming_input/test_gpu_model_runner_streaming.py

-1

@@ -39,7 +39,6 @@ def mock_model_runner_with_input_batch():

39 vocab_size=32000,

40 block_sizes=[16],

41 kernel_block_sizes=[16],

42 - is_spec_decode=False,

42 logitsprocs=None,

43 is_pooling_model=False,

44 )

vllm/v1/sample/logits_processor/__init__.py

-3

@@ -18,7 +18,6 @@ from vllm.v1.sample.logits_processor.builtin import (

18 LogitBiasLogitsProcessor,

19 MinPLogitsProcessor,

20 MinTokensLogitsProcessor,

21 - ThinkingTokenBudgetLogitsProcessor,

21 process_dict_updates,

22 )

23 from vllm.v1.sample.logits_processor.interface import (

@@ -51,7 +50,6 @@ BUILTIN_LOGITS_PROCESSORS: list[type[LogitsProcessor]] = [

50 MinTokensLogitsProcessor,

51 LogitBiasLogitsProcessor,

52 MinPLogitsProcessor,

54 - ThinkingTokenBudgetLogitsProcessor,

53 ]

@@ -356,5 +354,4 @@ __all__ = [

354 "STR_POOLING_REJECTS_LOGITSPROCS",

355 "LOGITSPROCS_GROUP",

356 "AdapterLogitsProcessor",

359 - "ThinkingTokenBudgetLogitsProcessor",

357 ]

vllm/v1/sample/logits_processor/builtin.py

+1 -262

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Callable, Sequence

4 -from typing import TYPE_CHECKING, Any, TypeVar

4 +from typing import TYPE_CHECKING, TypeVar

6 import numpy as np

7 import torch

@@ -291,267 +291,6 @@ class MinTokensLogitsProcessor(LogitsProcessor):

291 return logits

292

293

294 -class ThinkingTokenBudgetLogitsProcessor(LogitsProcessor):

295 - """Limits the number of tokens allowed inside a 'thinking' section."""

296 -

297 - def __init__(

298 - self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool

299 - ):

300 - reasoning_config = vllm_config.reasoning_config

301 - max_num_reqs = vllm_config.scheduler_config.max_num_seqs

302 -

303 - # Check if thinking is enabled

304 - self.is_enabled = reasoning_config is not None and reasoning_config.enabled

305 -

306 - self.reasoning_start_token_ids = getattr(

307 - reasoning_config, "reasoning_start_token_ids", []

308 - )

309 - self.reasoning_end_token_ids = getattr(

310 - reasoning_config, "reasoning_end_token_ids", []

311 - )

312 -

313 - self.pin_memory = is_pin_memory

314 - self.device = device

315 - # Per-request state tracking for thinking token management

316 - # Key: request_index, Value: state dict containing:

317 - # "in_think": bool - currently in thinking mode

318 - # "in_end": bool - currently forcing end tokens output

319 - # "check_count_down": int - steps remaining until next think

320 - # start/end token parsing

321 - # "think_count": int - number of thinking tokens generated

322 - # "end_count": int - number of end tokens forced so far

323 - # "thinking_token_budget": int - max allowed thinking tokens

324 - # "output_tok_ids": list[int] - generated output tokens

325 - # "prev_output_length": int - previous output length for

326 - # incremental processing

327 - self._state: dict[int, dict[str, Any]] = {}

328 -

329 - # Preallocate reusable tensors

330 - self.mask = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)

331 - self.force_token_ids = torch.full(

332 - (max_num_reqs,), -1, dtype=torch.long, device=device

333 - )

334 -

335 - @staticmethod

336 - def _find_last_sequence_index(target_list: list[int], token_ids: list[int]) -> int:

337 - """

338 - Returns the index of the last occurrence of token_ids in target_list.

339 -

340 - Args:

341 - target_list (list[int]): The list of token IDs.

342 - token_ids (list[int]): The sequence of token IDs to find.

343 - """

344 - if not token_ids:

345 - return -1

346 - for i in range(len(target_list) - len(token_ids), -1, -1):

347 - if target_list[i : i + len(token_ids)] == token_ids:

348 - return i

349 - return -1

350 -

351 - def _init_state_entry(

352 - self, prompt_tok_ids: list[int] | None, thinking_token_budget: int

353 - ) -> dict[str, Any]:

354 - """Initializes the tracking state for a given sequence index."""

355 - if prompt_tok_ids is None:

356 - last_start = -1

357 - last_end = -1

358 - in_think = False

359 - think_count = 0

360 - else:

361 - last_start = self._find_last_sequence_index(

362 - prompt_tok_ids, self.reasoning_start_token_ids

363 - )

364 - last_end = self._find_last_sequence_index(

365 - prompt_tok_ids, self.reasoning_end_token_ids

366 - )

367 - in_think = last_start > last_end

368 - if in_think:

369 - think_count = len(prompt_tok_ids) - (

370 - last_start + len(self.reasoning_start_token_ids)

371 - )

372 - else:

373 - think_count = 0

374 -

375 - return {

376 - "in_think": in_think, # Currently in thinking mode

377 - "in_end": in_think and thinking_token_budget == 0,

378 - "check_count_down": thinking_token_budget,

379 - "think_count": think_count, # Number of tokens in thinking section

380 - "end_count": 0, # Number of end tokens forced so far

381 - "prompt_tok_ids": prompt_tok_ids,

382 - "output_tok_ids": [],

383 - "thinking_token_budget": thinking_token_budget,

384 - "prev_output_length": 0,

385 - # Track previous output length for incremental updates

386 - }

387 -

388 - def _update_think_state(self, state: dict[str, Any]):

389 - """Updates the state based on newly generated output tokens."""

390 - if not state.get("in_end", False) and state.get("check_count_down", 0) > 0:

391 - state["check_count_down"] -= 1

392 - return

393 -

394 - output = state.get("output_tok_ids", [])

395 - if not output:

396 - return

397 -

398 - # Track previous output length for incremental processing

399 - prev_length = state.get("prev_output_length", 0)

400 - current_length = len(output)

401 -

402 - if current_length <= prev_length:

403 - return

404 -

405 - # Process only newly added tokens

406 - new_tokens = output[prev_length:]

407 - state["prev_output_length"] = current_length

408 -

409 - # Check if new tokens contain think start or end sequences

410 - start_len = len(self.reasoning_start_token_ids)

411 - end_len = len(self.reasoning_end_token_ids)

412 -

413 - # Look for think sequences in recent tokens (including boundary)

414 - # Check overlapping regions where sequences might span boundaries

415 - check_start_idx = max(0, prev_length - max(start_len, end_len) + 1)

416 - recent_tokens = output[check_start_idx:]

417 -

418 - # Find any think start/end sequences in recent tokens

419 - recent_start_pos = self._find_last_sequence_index(

420 - recent_tokens, self.reasoning_start_token_ids

421 - )

422 - recent_end_pos = self._find_last_sequence_index(

423 - recent_tokens, self.reasoning_end_token_ids

424 - )

425 -

426 - # Update state based on recent sequences

427 - if not state["in_end"]:

428 - if recent_start_pos >= 0 and recent_end_pos >= 0:

429 - if recent_start_pos > recent_end_pos:

430 - # Case: ...<end>...<start>... - entering think mode

431 - absolute_start_pos = check_start_idx + recent_start_pos

432 - new_think_count = current_length - (absolute_start_pos + start_len)

433 - state["in_think"] = True

434 - state["think_count"] = new_think_count

435 - else:

436 - # Case: ...<start>...<end>... - exiting think mode

437 - state["in_think"] = False

438 - state["think_count"] = 0

439 - elif recent_start_pos >= 0:

440 - # Found think start - entering think mode

441 - absolute_start_pos = check_start_idx + recent_start_pos

442 - new_think_count = current_length - (absolute_start_pos + start_len)

443 - state["in_think"] = True

444 - state["think_count"] = new_think_count

445 - elif recent_end_pos >= 0:

446 - # Found think end - exiting think mode

447 - state["in_think"] = False

448 - state["think_count"] = 0

449 - elif state["in_think"]:

450 - # Continue thinking mode, increment count by new tokens

451 - state["think_count"] += len(new_tokens)

452 -

453 - # Set countdown based on current state

454 - if state["in_think"]:

455 - remaining_budget = max(

456 - 0, state["thinking_token_budget"] - state["think_count"]

457 - )

458 - state["check_count_down"] = max(0, remaining_budget - 1)

459 - else:

460 - state["check_count_down"] = state["thinking_token_budget"]

461 -

462 - # Check if need to transition to end mode

463 - if (

464 - state["in_think"]

465 - and state["think_count"] >= state["thinking_token_budget"]

466 - ):

467 - state["in_think"] = False

468 - state["in_end"] = True

469 - state["end_count"] = 0

470 - state["check_count_down"] = state["thinking_token_budget"]

471 - else:

472 - # In end mode

473 - state["end_count"] += 1

474 - if state["end_count"] >= len(self.reasoning_end_token_ids):

475 - state.update(

476 - {

477 - "in_end": False,

478 - "end_count": 0,

479 - "check_count_down": state["thinking_token_budget"],

480 - }

481 - )

482 -

483 - def is_argmax_invariant(self) -> bool:

484 - """This logits processor can change the outcome of

485 - greedy sampling by forcing that the thinking section

486 - ends after a certain number of tokens."""

487 - return False

488 -

489 - def update_state(self, batch_update: BatchUpdate | None):

490 - if not self.is_enabled:

491 - return

492 - if batch_update:

493 - for index, params, prompt_tok_ids, output_tok_ids in batch_update.added:

494 - thinking_token_budget = params.thinking_token_budget

495 -

496 - if thinking_token_budget is not None:

497 - self._state[index] = self._init_state_entry(

498 - prompt_tok_ids, thinking_token_budget

499 - )

500 - self._state[index]["output_tok_ids"] = output_tok_ids

501 - else:

502 - # Remove state if no thinking budget

503 - self._state.pop(index, None)

504 -

505 - for index in batch_update.removed:

506 - self._state.pop(index, {})

507 -

508 - for i1, i2, direction in batch_update.moved:

509 - if direction == MoveDirectionality.SWAP:

510 - state1 = self._state.pop(i1, None)

511 - state2 = self._state.pop(i2, None)

512 - if state1 is not None:

513 - self._state[i2] = state1

514 - if state2 is not None:

515 - self._state[i1] = state2

516 - else:

517 - state = self._state.pop(i1, None)

518 - if state is not None:

519 - self._state[i2] = state

520 -

521 - for state in self._state.values():

522 - self._update_think_state(state)

523 -

524 - def apply(self, logits: torch.Tensor) -> torch.Tensor:

525 - if not self.is_enabled or not self._state:

526 - return logits

527 -

528 - batch_size = logits.size(0)

529 - self.mask[:batch_size] = False

530 -

531 - for i in range(batch_size):

532 - state = self._state.get(i)

533 - if state and state["in_end"]:

534 - self.mask[i] = True

535 - self.force_token_ids[i] = self.reasoning_end_token_ids[

536 - state["end_count"]

537 - ]

538 -

539 - # Check in CPU first not to sync with GPU

540 - has_active_thinking = any(

541 - state.get("in_end", False) for state in self._state.values()

542 - )

543 -

544 - if has_active_thinking:

545 - current_mask = self.mask[:batch_size]

546 - active_indices = current_mask.nonzero(as_tuple=False).view(-1)

547 - if len(active_indices) > 0:

548 - force_tokens = self.force_token_ids[active_indices]

549 - # Apply a large value for the end thinking token id index

550 - logits[active_indices, force_tokens] = 1e9

551 -

552 - return logits

553 -

554 -

294 def process_dict_updates(

295 req_entries: dict[int, T],

296 batch_update: BatchUpdate | None,

vllm/v1/sample/metadata.py

@@ -1,11 +1,14 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +from __future__ import annotations

5 +

6 from dataclasses import dataclass

8 import torch

10 from vllm.v1.sample.logits_processor import LogitsProcessors

11 +from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder

14 @dataclass

@@ -47,3 +50,6 @@ class SamplingMetadata:

51 # Speculative token ids

52 spec_token_ids: list[list[int]] | None = None

53 + # When non-None, use ``holder.has_tracked_requests()`` to see if this batch applies

54 + # thinking-token-budget logits (holder may exist with an empty tracking set).

55 + thinking_budget_state_holder: ThinkingBudgetStateHolder | None = None

vllm/v1/sample/rejection_sampler.py

+16 -3

@@ -290,16 +290,24 @@ class RejectionSampler(nn.Module):

290 any_penalties_or_bad_words = (

291 sampling_metadata.bad_words_token_ids or has_penalties

292 )

293 + holder = sampling_metadata.thinking_budget_state_holder

294 + needs_thinking = holder is not None and holder.has_tracked_requests()

295

296 output_token_ids = sampling_metadata.output_token_ids

295 - if any_penalties_or_bad_words:

297 + if any_penalties_or_bad_words or needs_thinking:

298 output_token_ids = self._combine_outputs_with_spec_tokens(

299 output_token_ids,

300 sampling_metadata.spec_token_ids,

301 )

302

303 # Calculate indices of target logits.

302 - if sampling_metadata.allowed_token_ids_mask is not None or has_penalties:

304 + repeat_indices: torch.Tensor | None = None

305 + need_repeat_indices = (

306 + sampling_metadata.allowed_token_ids_mask is not None

307 + or has_penalties

308 + or needs_thinking

309 + )

310 + if need_repeat_indices:

311 num_requests = len(metadata.num_draft_tokens)

312 num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu")

313 original_indices = torch.arange(num_requests, device="cpu")

@@ -327,7 +335,12 @@ class RejectionSampler(nn.Module):

335 logits = processor.apply_with_spec_decode(

336 logits, metadata.num_draft_tokens

337 )

330 -

338 + if holder is not None and holder.has_tracked_requests():

339 + logits = holder.apply_to_logits(

340 + logits,

341 + predict_bonus_token=False,

342 + spec_token_ids=sampling_metadata.spec_token_ids,

343 + )

344 return logits

345

346 @staticmethod

vllm/v1/sample/sampler.py

+16 -1

@@ -364,9 +364,13 @@ class Sampler(nn.Module):

364 any_penalties_or_bad_words = (

365 bool(bad_words_token_ids) or not sampling_metadata.no_penalties

366 )

367 + holder = sampling_metadata.thinking_budget_state_holder

368 + needs_thinking_combine = holder is not None and holder.has_tracked_requests()

369

370 output_token_ids = sampling_metadata.output_token_ids

369 - if predict_bonus_token and any_penalties_or_bad_words:

371 + if predict_bonus_token and (

372 + any_penalties_or_bad_words or needs_thinking_combine

373 + ):

374 # Combine base outputs with spec tokens when speculative decoding

375 # is enabled.

376 output_token_ids = self._combine_outputs_with_spec_tokens(

@@ -388,6 +392,17 @@ class Sampler(nn.Module):

392

393 # Apply penalties (e.g., freq_penalties).

394 logits = self.apply_penalties(logits, sampling_metadata, output_token_ids)

395 + if holder is not None and holder.has_tracked_requests():

396 + holder.update_state(

397 + output_token_ids,

398 + sampling_metadata.spec_token_ids,

399 + repeat_indices=None,

400 + )

401 + logits = holder.apply_to_logits(

402 + logits,

403 + predict_bonus_token,

404 + sampling_metadata.spec_token_ids,

405 + )

406 return logits

407

408 @staticmethod

vllm/v1/sample/thinking_budget_state.py new

+528

@@ -0,0 +1,528 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Per-batch thinking token budget state; applied after penalties at sample time."""

4 +

5 +from typing import TYPE_CHECKING, Any

6 +

7 +import torch

8 +

9 +from vllm.v1.sample.logits_processor.interface import (

10 + BatchUpdate,

11 + MoveDirectionality,

12 +)

13 +

14 +if TYPE_CHECKING:

15 + from vllm.config.reasoning import ReasoningConfig

16 +

17 +

18 +def maybe_create_thinking_budget_state_holder(

19 + reasoning_config: "ReasoningConfig | None",

20 + max_num_seqs: int,

21 + num_spec_tokens: int,

22 + device: torch.device,

23 + is_pin_memory: bool,

24 +) -> "ThinkingBudgetStateHolder | None":

25 + if reasoning_config is None:

26 + return None

27 + return ThinkingBudgetStateHolder(

28 + reasoning_config, max_num_seqs, num_spec_tokens, device, is_pin_memory

29 + )

30 +

31 +

32 +class ThinkingBudgetStateHolder:

33 + """Tracks thinking sections and forces end tokens when budget is exceeded."""

34 +

35 + think_start_token_ids: list[int]

36 + think_end_token_ids: list[int]

37 +

38 + def __init__(

39 + self,

40 + reasoning_config: "ReasoningConfig | None",

41 + max_num_seqs: int,

42 + num_spec_tokens: int,

43 + device: torch.device,

44 + is_pin_memory: bool,

45 + ):

46 + _ = is_pin_memory # API parity with logits processors

47 + max_num_reqs = max_num_seqs

48 + self.in_spec_mode = num_spec_tokens > 0

49 + self.num_spec_tokens = num_spec_tokens

50 +

51 + # No separate enable flag: a non-``None`` ``reasoning_config`` is the switch.

52 + self.is_enabled = reasoning_config is not None

53 +

54 + if reasoning_config is None:

55 + self.think_start_token_ids = []

56 + self.think_end_token_ids = []

57 + else:

58 + rs = reasoning_config.reasoning_start_token_ids

59 + re = reasoning_config.reasoning_end_token_ids

60 + self.think_start_token_ids = rs if rs else []

61 + self.think_end_token_ids = re if re else []

62 +

63 + self.device = device

64 + self._state: dict[int, dict[str, Any]] = {}

65 + self.cu_num_tokens: dict[int, int] = {}

66 +

67 + if self.num_spec_tokens > 0:

68 + self.mask = torch.zeros(

69 + max_num_reqs * (self.num_spec_tokens + 1),

70 + dtype=torch.bool,

71 + device=device,

72 + )

73 + self.force_token_ids = torch.full(

74 + (max_num_reqs * (self.num_spec_tokens + 1),),

75 + -1,

76 + dtype=torch.long,

77 + device=device,

78 + )

79 + else:

80 + self.mask = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)

81 + self.force_token_ids = torch.full(

82 + (max_num_reqs,), -1, dtype=torch.long, device=device

83 + )

84 +

85 + def has_tracked_requests(self) -> bool:

86 + """True when ``sync_batch`` has state for a ``thinking_token_budget`` row.

87 +

88 + Used to decide whether sampling needs output-token rows and spec combining;

89 + distinct from merely having a holder instance (reasoning may be on with no

90 + budgeted requests in this batch).

91 + """

92 + return bool(self._state)

93 +

94 + def sync_batch(self, batch_update: BatchUpdate | None) -> None:

95 + """Add/remove/move per-request state only (no _update_think_state)."""

96 + if not self.is_enabled or not batch_update:

97 + return

98 + for index in batch_update.removed:

99 + self._state.pop(index, None)

100 +

101 + for index, params, prompt_tok_ids, output_tok_ids in batch_update.added:

102 + thinking_token_budget = params.thinking_token_budget

103 + if thinking_token_budget is not None:

104 + self._state[index] = self._init_state_entry(

105 + prompt_tok_ids, thinking_token_budget

106 + )

107 + self._state[index]["output_tok_ids"] = output_tok_ids

108 + self._state[index]["spec_token_ids"] = []

109 + else:

110 + self._state.pop(index, None)

111 +

112 + for i1, i2, direction in batch_update.moved:

113 + if direction == MoveDirectionality.SWAP:

114 + state1 = self._state.get(i1)

115 + state2 = self._state.get(i2)

116 + if state1 is not None:

117 + self._state[i2] = state1

118 + if state2 is not None:

119 + self._state[i1] = state2

120 + else:

121 + state = self._state.pop(i1, None)

122 + if state is not None:

123 + self._state[i2] = state

124 +

125 + def update_state(

126 + self,

127 + output_token_ids: list[list[int]],

128 + spec_token_ids: list[list[int]] | None,

129 + repeat_indices: torch.Tensor | None = None,

130 + ) -> None:

131 + """Refresh output/spec from sampling rows and recompute think state."""

132 + if not self.is_enabled or not self._state:

133 + return

134 +

135 + spec_lists = spec_token_ids or []

136 + last_row_for_req: dict[int, int] | None = None

137 + if repeat_indices is not None:

138 + last_row_for_req = {}

139 + rpt = repeat_indices.cpu().tolist()

140 + for batch_row, req_i in enumerate(rpt):

141 + last_row_for_req[req_i] = batch_row

142 +

143 + for seq_idx, state in list(self._state.items()):

144 + if last_row_for_req is not None:

145 + output_row: int | None = last_row_for_req.get(seq_idx)

146 + if output_row is None or output_row >= len(output_token_ids):

147 + continue

148 + state["output_tok_ids"] = output_token_ids[output_row]

149 + elif seq_idx >= len(output_token_ids):

150 + continue

151 + else:

152 + state["output_tok_ids"] = output_token_ids[seq_idx]

153 + if seq_idx < len(spec_lists):

154 + state["spec_token_ids"] = list(spec_lists[seq_idx])

155 + else:

156 + state["spec_token_ids"] = []

157 + state["in_spec_mode"] = self.in_spec_mode

158 + state["force_index"] = []

159 + if len(state["output_tok_ids"]) > 0:

160 + spec_len = len(state["spec_token_ids"])

161 + # Only strip draft suffix when there are spec tokens; ``[:-0]`` would

162 + # clear the whole list (Python treats stop index 0 as "up to empty").

163 + if spec_len > 0 and len(state["output_tok_ids"]) >= spec_len:

164 + state["output_tok_ids"] = state["output_tok_ids"][:-spec_len]

165 + self._update_think_state(state)

166 +

167 + def apply_to_logits(

168 + self,

169 + logits: torch.Tensor,

170 + predict_bonus_token: bool,

171 + spec_token_ids: list[list[int]] | None,

172 + ) -> torch.Tensor:

173 + """Mask and bump logits for forced end-of-thinking tokens."""

174 + if not self.is_enabled or not self._state:

175 + return logits

176 + spec_lists = spec_token_ids or []

177 + return self._apply_forcing_to_logits(logits, predict_bonus_token, spec_lists)

178 +

179 + @staticmethod

180 + def _find_last_sequence_index(target_list: list[int], token_ids: list[int]) -> int:

181 + if not token_ids:

182 + return -1

183 + for i in range(len(target_list) - len(token_ids), -1, -1):

184 + if target_list[i : i + len(token_ids)] == token_ids:

185 + return i

186 + return -1

187 +

188 + def _init_state_entry(

189 + self, prompt_tok_ids: list[int] | None, thinking_token_budget: int

190 + ) -> dict[str, Any]:

191 + if prompt_tok_ids is None:

192 + last_start = -1

193 + last_end = -1

194 + in_think = False

195 + think_count = 0

196 + start_thinking = -1

197 + countdown = thinking_token_budget

198 + continue_thinking = False

199 + in_end = False

200 + else:

201 + start_thinking = -1

202 + countdown = thinking_token_budget

203 + continue_thinking = False

204 + in_end = False

205 + last_start = self._find_last_sequence_index(

206 + prompt_tok_ids, self.think_start_token_ids

207 + )

208 + last_end = self._find_last_sequence_index(

209 + prompt_tok_ids, self.think_end_token_ids

210 + )

211 + in_think = last_start > last_end

212 + # load metrics such as think count, start thinking

213 + # if request is in thinking mode, already

214 + if in_think:

215 + think_count = len(prompt_tok_ids) - (

216 + last_start + len(self.think_start_token_ids)

217 + )

218 + start_thinking = len(prompt_tok_ids) - think_count - 1

219 + countdown -= think_count

220 + continue_thinking = True

221 + # check if the token is exhausted within prompt

222 + token_exhausted = thinking_token_budget - think_count

223 + in_end = token_exhausted <= 0

224 + else:

225 + think_count = 0

226 +

227 + return {

228 + "in_think": in_think,

229 + "in_end": in_end,

230 + "check_count_down": countdown,

231 + "think_count": think_count,

232 + "end_count": 0,

233 + "prompt_tok_ids": prompt_tok_ids,

234 + "output_tok_ids": [],

235 + "thinking_token_budget": thinking_token_budget,

236 + "prev_output_length": 0,

237 + "spec_token_ids": [],

238 + "force_index": [],

239 + "start_thinking": start_thinking,

240 + "end_thinking": -1,

241 + "in_spec_mode": False,

242 + "bonus_token_forced": False,

243 + "continue_thinking": continue_thinking,

244 + }

245 +

246 + def _update_think_state(self, state: dict[str, Any]) -> None:

247 + if state.get("thinking_token_budget", -1) == -1:

248 + return

249 + if len(self.think_end_token_ids) == 0:

250 + state["thinking_token_budget"] = -1

251 + state["in_end"] = False

252 + state["force_index"] = []

253 + return

254 +

255 + if state["start_thinking"] == -1:

256 + start_thinking = self._find_last_sequence_index(

257 + state.get("output_tok_ids", []), self.think_start_token_ids

258 + )

259 + state["start_thinking"] = start_thinking

260 + if state["end_thinking"] == -1:

261 + end_thinking = self._find_last_sequence_index(

262 + state.get("output_tok_ids", []), self.think_end_token_ids

263 + )

264 + state["end_thinking"] = end_thinking

265 +

266 + if state["start_thinking"] == -1:

267 + return

268 +

269 + if state["continue_thinking"]:

270 + sampled_tokens_from_previous_step = len(

271 + state.get("output_tok_ids", [])

272 + ) - state.get("prev_output_length", 0)

273 + else:

274 + if state["prev_output_length"] == 0:

275 + sampled_tokens_from_previous_step = len(

276 + state.get("output_tok_ids", [])

277 + ) - len(self.think_start_token_ids)

278 + else:

279 + sampled_tokens_from_previous_step = (

280 + len(state.get("output_tok_ids", [])) - state["prev_output_length"]

281 + )

282 + current_step_countdown = (

283 + state["check_count_down"] - sampled_tokens_from_previous_step

284 + )

285 + predicted_countdown = current_step_countdown - len(state["spec_token_ids"]) - 1

286 + # We only proceed further if we have counted down the thinking budget

287 + # to 0 or less and when we are in the "in think" mode.

288 + if (

289 + not state.get("in_end", False)

290 + and predicted_countdown >= 0

291 + and state["start_thinking"] > -1

292 + ):

293 + state["check_count_down"] = current_step_countdown

294 + state["prev_output_length"] = len(state.get("output_tok_ids", []))

295 + return

296 + output = state.get("output_tok_ids", [])

297 + if not output:

298 + # When in_end was set at init (budget=0, prompt already in think),

299 + # we must force the first generated token to be the end token;

300 + # otherwise apply() sees in_end=True but force_index=[] and

301 + # allows an extra thinking token.

302 + if state.get("in_end", False):

303 + state["force_index"] = [0]

304 + return

305 +

306 + # Track previous output length for incremental processing

307 + prev_length = state.get("prev_output_length", 0)

308 + current_length = len(output)

309 +

310 + if current_length <= prev_length:

311 + if state.get("in_end", False):

312 + remaining_budget = state["thinking_token_budget"] - state["think_count"]

313 + spec_len = len(state["spec_token_ids"])

314 + if spec_len > 0:

315 + if 0 < remaining_budget < spec_len:

316 + state["force_index"] = [remaining_budget]

317 + elif remaining_budget <= 0:

318 + state["force_index"] = [0]

319 + else:

320 + state["force_index"] = [spec_len]

321 + else:

322 + state["force_index"] = [0]

323 + return

324 +

325 + state["prev_output_length"] = current_length

326 +

327 + start_len = len(self.think_start_token_ids)

328 + absolute_start_pos = state["start_thinking"]

329 +

330 + if state["continue_thinking"] and state["end_thinking"] > -1:

331 + absolute_end_pos = state["end_thinking"] + len(

332 + state.get("prompt_tok_ids") or []

333 + )

334 + else:

335 + absolute_end_pos = state["end_thinking"]

336 + # Update state based on recent sequences

337 + # This is the case where we are in end mode, but the rejection sampler

338 + # rejected a token before the end token,

339 + # so we need to go back to think mode and wait for the next end token

340 + # eg with 999: [2,4,5,999] -> [3,-1,-1,-1]

341 + if state["in_end"] and state["end_count"] == 0:

342 + new_tokens = output[prev_length:]

343 + stopping_thinking = (

344 + self.think_end_token_ids[state["end_count"]] in new_tokens

345 + )

346 + if not stopping_thinking:

347 + state["in_think"] = True

348 + state["in_end"] = False

349 + state["end_count"] = 0

350 + state["bonus_token_forced"] = False

351 +

352 + if not state["in_end"]:

353 + if absolute_start_pos >= 0 and absolute_end_pos >= 0:

354 + # Case: ...<end>...<start>... - entering think mode

355 + if absolute_start_pos > absolute_end_pos:

356 + new_think_count = current_length - (absolute_start_pos + start_len)

357 + state["in_think"] = True

358 + state["think_count"] = new_think_count

359 + else:

360 + # Case: ...<start>...<end>... - exiting think mode

361 + state["in_think"] = False

362 + state["think_count"] = 0

363 +

364 + elif absolute_start_pos >= 0 and not state["continue_thinking"]:

365 + # Found think start - entering think mode

366 + new_think_count = current_length - (absolute_start_pos + start_len)

367 + state["in_think"] = True

368 + state["think_count"] = new_think_count

369 +

370 + elif absolute_end_pos >= 0:

371 + # Found think end - exiting think mode

372 + state["in_think"] = False

373 + state["think_count"] = 0

374 +

375 + elif state["in_think"]:

376 + # Continue thinking mode, increment count by new tokens

377 + prompt_tok_ids = state.get("prompt_tok_ids") or []

378 + think_tokens_in_prompt = len(prompt_tok_ids) - (

379 + absolute_start_pos + start_len

380 + )

381 + state["think_count"] = (

382 + len(state["output_tok_ids"]) + think_tokens_in_prompt

383 + )

384 + if state["in_think"]:

385 + remaining_budget = max(

386 + 0, state["thinking_token_budget"] - state["think_count"]

387 + )

388 + state["check_count_down"] = remaining_budget

389 + else:

390 + state["check_count_down"] = state["thinking_token_budget"]

391 +

392 + total_thinking_tokens = (

393 + state["think_count"] + len(state["spec_token_ids"]) + 1

394 + )

395 + # Check if need to transition to end mode

396 + # If we have more thinking tokens than the budget,

397 + # we need to transition to end mode

398 + if (

399 + state["in_think"]

400 + and total_thinking_tokens > state["thinking_token_budget"]

401 + ):

402 + # Calculate force_index: position within spec_token_ids where

403 + # forcing starts. If we're already over budget without spec

404 + # tokens, force from position 0. Force from the position

405 + # where budget is exceeded.

406 + state["in_think"] = False

407 + state["in_end"] = True

408 + state["end_count"] = 0

409 + state["check_count_down"] = state["thinking_token_budget"]

410 + remaining_budget = state["thinking_token_budget"] - state["think_count"]

411 + spec_len = len(state["spec_token_ids"])

412 + if 0 < remaining_budget < spec_len:

413 + state["force_index"] = [remaining_budget]

414 +

415 + elif remaining_budget <= 0:

416 + state["force_index"] = [0]

417 +

418 + else:

419 + # remaining_budget >= spec_len: all spec tokens are within

420 + # budget; force the bonus token position

421 + state["force_index"] = [len(state["spec_token_ids"])]

422 +

423 + else:

424 + state["force_index"] = []

425 + if len(state["spec_token_ids"]) > 0:

426 + for i, token_id in enumerate(state["spec_token_ids"]):

427 + if state["end_count"] + 1 < len(self.think_end_token_ids):

428 + if token_id == self.think_end_token_ids[state["end_count"] + 1]:

429 + state["end_count"] += 1

430 + else:

431 + state["end_count"] += 1

432 + state["force_index"] = [i]

433 + break

434 + else:

435 + state["end_count"] += 1

436 + if len(state["force_index"]) == 0:

437 + state["end_count"] += 1

438 + state["force_index"] = [len(state["spec_token_ids"])]

439 + else:

440 + state["end_count"] += 1

441 + state["force_index"] = [0]

442 + if state["end_count"] >= len(self.think_end_token_ids):

443 + state.update(

444 + {

445 + "in_end": False,

446 + "end_count": 0,

447 + "check_count_down": state["thinking_token_budget"],

448 + }

449 + )

450 +

451 + def _apply_forcing_to_logits(

452 + self,

453 + logits: torch.Tensor,

454 + predict_bonus_token: bool,

455 + spec_token_ids_for_layout: list[list[int]],

456 + ) -> torch.Tensor:

457 + self.mask[:] = False

458 + cumulative_total = 0

459 + self.cu_num_tokens.clear()

460 +

461 + n_layout = len(spec_token_ids_for_layout)

462 + if self._state:

463 + n_layout = max(n_layout, max(self._state.keys()) + 1)

464 +

465 + for index in range(n_layout):

466 + self.cu_num_tokens[index] = cumulative_total

467 + spec_tokens = (

468 + spec_token_ids_for_layout[index]

469 + if index < len(spec_token_ids_for_layout)

470 + else []

471 + )

472 + if self.in_spec_mode:

473 + cumulative_total += len(spec_tokens) if not predict_bonus_token else 1

474 + else:

475 + cumulative_total += 1

476 +

477 + for seq_idx in sorted(self._state.keys()):

478 + if seq_idx not in self.cu_num_tokens:

479 + continue

480 + state = self._state[seq_idx]

481 + if state.get("in_end", False):

482 + # logits processor in spec mode are called twice

483 + # once for bonus token logits and

484 + # second time for the target logits

485 + # in case the force index is bonus token index

486 + # we change the force index to 0

487 + if predict_bonus_token:

488 + if state.get("force_index") and state["force_index"][0] < len(

489 + state["spec_token_ids"]

490 + ):

491 + continue

492 + else:

493 + state["force_index"] = [0]

494 + # continue enforcing the end thinking tokens

495 + if state["end_count"] > 0:

496 + state["bonus_token_forced"] = False

497 + if state and not state["bonus_token_forced"]:

498 + force_index = state.get("force_index", [])

499 + if len(force_index) == 0:

500 + continue

501 + end_count = state.get("end_count", 0)

502 + for force_idx in force_index:

503 + if end_count < len(self.think_end_token_ids):

504 + mask_idx = self.cu_num_tokens[seq_idx] + force_idx

505 + if mask_idx < len(self.mask) and mask_idx < logits.shape[0]:

506 + self.mask[mask_idx] = True

507 + self.force_token_ids[mask_idx] = (

508 + self.think_end_token_ids[end_count]

509 + )

510 + if predict_bonus_token:

511 + if state["end_count"] > 0:

512 + state["bonus_token_forced"] = False

513 + state["force_index"] = []

514 + else:

515 + state["bonus_token_forced"] = True

516 +

517 + has_active_thinking = any(

518 + state.get("in_end", False) for state in self._state.values()

519 + )

520 +

521 + if has_active_thinking:

522 + active_indices = self.mask.nonzero(as_tuple=False).view(-1)

523 +

524 + if len(active_indices) > 0:

525 + force_tokens = self.force_token_ids[active_indices]

526 + logits[active_indices, force_tokens] = 1e9

527 +

528 + return logits

vllm/v1/worker/gpu_input_batch.py

+30 -2

@@ -8,6 +8,7 @@ from typing import cast

8 import numpy as np

9 import torch

11 +from vllm.config.reasoning import ReasoningConfig

12 from vllm.lora.request import LoRARequest

13 from vllm.multimodal.inputs import MultiModalFeatureSpec

14 from vllm.pooling_params import PoolingParams

@@ -22,6 +23,9 @@ from vllm.v1.sample.logits_processor import (

23 MoveDirectionality,

24 )

25 from vllm.v1.sample.metadata import SamplingMetadata

26 +from vllm.v1.sample.thinking_budget_state import (

27 + maybe_create_thinking_budget_state_holder,

28 +)

29 from vllm.v1.utils import copy_slice

30 from vllm.v1.worker.block_table import MultiGroupBlockTable

@@ -92,12 +96,20 @@ class InputBatch:

96 max_num_blocks_per_req: list[int] | None = None,

97 logitsprocs: LogitsProcessors | None = None,

98 logitsprocs_need_output_token_ids: bool = False,

95 - is_spec_decode: bool = False,

99 + num_spec_tokens: int = 0,

100 is_pooling_model: bool = False,

101 cp_kv_cache_interleave_size: int = 1,

102 + reasoning_config: ReasoningConfig | None = None,

103 ):

104 + self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder(

105 + reasoning_config,

106 + max_num_reqs,

107 + num_spec_tokens,

108 + device,

109 + pin_memory,

110 + )

111 + self.thinking_token_budget_reqs: set[str] = set()

112 self.is_pooling_model = is_pooling_model

100 - self.is_spec_decode = is_spec_decode

113 self.max_num_reqs = max_num_reqs

114 self.max_model_len = max_model_len

115 self.max_num_batched_tokens = max_num_batched_tokens

@@ -540,6 +552,7 @@ class InputBatch:

552 # False means we don't fill with -inf.

553 self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)

554 self.bad_words_token_ids.pop(req_index, None)

555 + self.thinking_token_budget_reqs.discard(req_id)

556 return req_index

557

558 def swap_states(self, i1: int, i2: int) -> None:

@@ -800,6 +813,8 @@ class InputBatch:

813 # reset batch update tracking.

814 # Update sampling metadata if batch state is changed.

815 batch_update = self.batch_update_builder.get_and_reset(self.num_reqs)

816 + if self.thinking_budget_state_holder is not None and batch_update:

817 + self.thinking_budget_state_holder.sync_batch(batch_update)

818 for logit_proc in self.logitsprocs.all:

819 logit_proc.update_state(batch_update)

820 if batch_update:

@@ -853,10 +868,15 @@ class InputBatch:

868

869 # Only set output_token_ids if required by the current requests'

870 # sampling parameters.

871 + holder = self.thinking_budget_state_holder

872 + thinking_budget_tracks_reqs = (

873 + holder is not None and holder.has_tracked_requests()

874 + )

875 needs_output_token_ids = (

876 not self.no_penalties

877 or bool(self.bad_words_token_ids)

878 or self.logitsprocs_need_output_token_ids

879 + or not thinking_budget_tracks_reqs

880 )

881 output_token_ids = (

882 cast(list[list[int]], self.req_output_token_ids)

@@ -902,6 +922,7 @@ class InputBatch:

922 allowed_token_ids_mask=allowed_token_ids_mask,

923 bad_words_token_ids=self.bad_words_token_ids,

924 logitsprocs=self.logitsprocs,

925 + thinking_budget_state_holder=self.thinking_budget_state_holder,

926 )

927

928 def get_pooling_params(self) -> list[PoolingParams]:

@@ -1076,6 +1097,13 @@ class InputBatch:

1097 and len(self.repetition_penalties_reqs) == 0

1098 )

1099

1100 + @property

1101 + def no_thinking_budget(self) -> bool:

1102 + return (

1103 + self.thinking_budget_state_holder is None

1104 + or len(self.thinking_token_budget_reqs) == 0

1105 + )

1106 +

1107 @property

1108 def max_num_logprobs(self) -> int | None:

1109 return max(self.num_logprobs.values()) if self.num_logprobs else None

vllm/v1/worker/gpu_model_runner.py

+4 -2

@@ -629,7 +629,7 @@ class GPUModelRunner(

629 vocab_size=self.model_config.get_vocab_size(),

630 block_sizes=[placeholder_block_size],

631 kernel_block_sizes=[placeholder_block_size],

632 - is_spec_decode=bool(self.vllm_config.speculative_config),

632 + num_spec_tokens=self.num_spec_tokens,

633 logitsprocs=build_logitsprocs(

634 self.vllm_config,

635 self.device,

@@ -645,6 +645,7 @@ class GPUModelRunner(

645 or self.vllm_config.reasoning_config is not None,

646 is_pooling_model=self.is_pooling_model,

647 cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,

648 + reasoning_config=self.vllm_config.reasoning_config,

649 )

650

651 # Separate cuda stream for overlapping transfer of sampled token ids from

@@ -6504,10 +6505,11 @@ class GPUModelRunner(

6505 block_sizes=block_sizes,

6506 kernel_block_sizes=kernel_block_sizes,

6507 max_num_blocks_per_req=max_num_blocks,

6507 - is_spec_decode=bool(self.vllm_config.speculative_config),

6508 + num_spec_tokens=self.num_spec_tokens,

6509 logitsprocs=self.input_batch.logitsprocs,

6510 logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids,

6511 is_pooling_model=self.is_pooling_model,

6512 + reasoning_config=self.vllm_config.reasoning_config,

6513 )

6514

6515 assert self._init_block_sizes == block_sizes, (