@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ed908cf0a0

[Bugfix] Fix thinking_token_budget not enforced after natural </think> re-entry (#45984)

Signed-off-by: Ashwin Giridharan <girida@amazon.com>

Ashwin Giridharan committed Jul 10, 2026 at 15:47 UTC ed908cf0a08af20839b82dc0638dae91b1ad630b
3 files changed +453 -47
tests/v1/logits_processors/test_correctness.py
+407 -3
@@ -145,6 +145,7 @@ def _generate_fake_sampling_metadata(
145 vllm_config.scheduler_config.max_num_seqs,
146 num_spec,
147 device,
148 + is_pin_memory=False,
149 )
150 fake_sampling_metadata = SamplingMetadata(
151 temperature=torch.full((batch_size,), 0.0),
@@ -879,6 +880,7 @@ def test_maybe_create_thinking_budget_holder_without_reasoning():
880 cfg.scheduler_config.max_num_seqs,
881 0,
882 torch.device("cpu"),
883 + is_pin_memory=False,
884 )
885 is None
886 )
@@ -1251,12 +1253,20 @@ def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index
1253 out.append(tok)
1254 h.update_state([out], None, None)
1255 assert h._state[0]["end_thinking"] == -1 # not present yet
1254 - expected_end_idx = len(out) # marker appended next
1256 out.extend(end)
1257 h.update_state([out], None, None)
1258
1258 - assert h._state[0]["start_thinking"] == 0
1259 - assert h._state[0]["end_thinking"] == expected_end_idx
1259 + # After a natural exit, start_thinking and end_thinking are reset to -1
1260 + # (state machine prepared for next block). Verify the exit happened at the
1261 + # correct position by checking scan_offset (set to len(output) after exit).
1262 + assert not h._state[0]["in_think"], "Should have exited think mode after </think>"
1263 + assert h._state[0]["start_thinking"] == -1, (
1264 + "start_thinking should be reset after natural exit"
1265 + )
1266 + assert h._state[0]["scan_offset"] == len(out), (
1267 + f"scan_offset should point past the end marker; "
1268 + f"expected {len(out)}, got {h._state[0]['scan_offset']}"
1269 + )
1270
1271
1272 # --- Thinking budget re-entry tests (issue #43708) ---
@@ -1381,3 +1391,397 @@ class TestThinkingBudgetReentry:
1391
1392 assert not holder._state[0]["in_end"]
1393 assert not holder._state[0]["in_think"]
1394 +
1395 +
1396 +# --- Thinking budget natural-end re-entry tests (issue #45974) ---
1397 +# After a model naturally emits </think> before exhausting its budget,
1398 +# subsequent <think> blocks must still be tracked and budget-enforced.
1399 +
1400 +
1401 +class TestThinkingBudgetNaturalEndReentry:
1402 + """Tests for thinking budget enforcement across multiple think blocks.
1403 +
1404 + Covers both natural-end re-entry (issue #45974) and forced-end re-entry
1405 + (issue #43708).
1406 + """
1407 +
1408 + THINK_START = 100
1409 + THINK_END_SINGLE = [200]
1410 + THINK_END_MULTI = [200, 201, 202]
1411 + BUDGET = 10
1412 + CONTENT_TOKEN = 50
1413 + THINK_TOKEN = 60
1414 +
1415 + @staticmethod
1416 + def _make_holder(end_token_ids):
1417 + from dataclasses import dataclass
1418 +
1419 + from vllm.v1.sample.thinking_budget_state import (
1420 + ThinkingBudgetStateHolder,
1421 + )
1422 +
1423 + @dataclass
1424 + class FakeReasoningConfig:
1425 + reasoning_start_token_ids: list[int]
1426 + reasoning_end_token_ids: list[int]
1427 + enabled: bool = True
1428 +
1429 + cfg = FakeReasoningConfig(
1430 + reasoning_start_token_ids=[TestThinkingBudgetNaturalEndReentry.THINK_START],
1431 + reasoning_end_token_ids=end_token_ids,
1432 + )
1433 + return ThinkingBudgetStateHolder(
1434 + reasoning_config=cfg,
1435 + max_num_seqs=8,
1436 + num_spec_tokens=0,
1437 + device=torch.device("cpu"),
1438 + is_pin_memory=False,
1439 + )
1440 +
1441 + @staticmethod
1442 + def _sync_batch(holder, budget):
1443 + from unittest.mock import MagicMock
1444 +
1445 + params = MagicMock()
1446 + params.thinking_token_budget = budget
1447 + batch_update = MagicMock(
1448 + removed=[],
1449 + added=[(0, params, None, [])],
1450 + moved=[],
1451 + )
1452 + holder.sync_batch(batch_update)
1453 +
1454 + @staticmethod
1455 + def _step(holder, output_tok_ids):
1456 + holder.update_state(
1457 + output_token_ids=[output_tok_ids],
1458 + spec_token_ids=None,
1459 + repeat_indices=None,
1460 + )
1461 +
1462 + # --- Natural-end re-entry tests ---
1463 +
1464 + def test_natural_end_reentry_single_token(self):
1465 + """After natural </think>, a new <think> block must be enforced."""
1466 + holder = self._make_holder(self.THINK_END_SINGLE)
1467 + self._sync_batch(holder, self.BUDGET)
1468 +
1469 + output = []
1470 +
1471 + # Block 1: 6 tokens (under budget) + natural </think>
1472 + output.append(self.THINK_START)
1473 + self._step(holder, list(output))
1474 + for _ in range(6):
1475 + output.append(self.THINK_TOKEN)
1476 + self._step(holder, list(output))
1477 + output.append(self.THINK_END_SINGLE[0])
1478 + self._step(holder, list(output))
1479 +
1480 + assert not holder._state[0]["in_end"]
1481 + assert not holder._state[0]["in_think"]
1482 +
1483 + # Some content
1484 + for _ in range(3):
1485 + output.append(self.CONTENT_TOKEN)
1486 + self._step(holder, list(output))
1487 +
1488 + # Block 2: re-entry — should be budget-tracked
1489 + output.append(self.THINK_START)
1490 + self._step(holder, list(output))
1491 + for _ in range(self.BUDGET + 1):
1492 + output.append(self.THINK_TOKEN)
1493 + self._step(holder, list(output))
1494 +
1495 + assert holder._state[0]["in_end"], (
1496 + "Second thinking block after natural end must be budget-enforced"
1497 + )
1498 +
1499 + def test_natural_end_reentry_multi_token(self):
1500 + """Multi-token </think> natural end still allows Block 2 enforcement."""
1501 + holder = self._make_holder(self.THINK_END_MULTI)
1502 + self._sync_batch(holder, self.BUDGET)
1503 +
1504 + output = []
1505 +
1506 + # Block 1: 4 tokens + natural multi-token </think>
1507 + output.append(self.THINK_START)
1508 + self._step(holder, list(output))
1509 + for _ in range(4):
1510 + output.append(self.THINK_TOKEN)
1511 + self._step(holder, list(output))
1512 + for tok in self.THINK_END_MULTI:
1513 + output.append(tok)
1514 + self._step(holder, list(output))
1515 +
1516 + assert not holder._state[0]["in_think"]
1517 +
1518 + # Content
1519 + for _ in range(5):
1520 + output.append(self.CONTENT_TOKEN)
1521 + self._step(holder, list(output))
1522 +
1523 + # Block 2: re-entry
1524 + output.append(self.THINK_START)
1525 + self._step(holder, list(output))
1526 + for _ in range(self.BUDGET + 1):
1527 + output.append(self.THINK_TOKEN)
1528 + self._step(holder, list(output))
1529 +
1530 + assert holder._state[0]["in_end"], (
1531 + "Second block after multi-token natural end must be enforced"
1532 + )
1533 +
1534 + def test_natural_end_immediate_reentry(self):
1535 + """Natural </think> followed immediately by <think> with no content."""
1536 + holder = self._make_holder(self.THINK_END_SINGLE)
1537 + self._sync_batch(holder, self.BUDGET)
1538 +
1539 + output = []
1540 +
1541 + # Block 1: 5 tokens + natural end
1542 + output.append(self.THINK_START)
1543 + self._step(holder, list(output))
1544 + for _ in range(5):
1545 + output.append(self.THINK_TOKEN)
1546 + self._step(holder, list(output))
1547 + output.append(self.THINK_END_SINGLE[0])
1548 + self._step(holder, list(output))
1549 +
1550 + # Immediate re-entry — NO content between blocks
1551 + output.append(self.THINK_START)
1552 + self._step(holder, list(output))
1553 + for _ in range(self.BUDGET + 1):
1554 + output.append(self.THINK_TOKEN)
1555 + self._step(holder, list(output))
1556 +
1557 + assert holder._state[0]["in_end"], (
1558 + "Immediate re-entry after natural end must be budget-enforced"
1559 + )
1560 +
1561 + def test_multiple_natural_reentries(self):
1562 + """Block 1 natural -> Block 2 natural -> Block 3 must be enforced."""
1563 + holder = self._make_holder(self.THINK_END_SINGLE)
1564 + self._sync_batch(holder, self.BUDGET)
1565 +
1566 + output = []
1567 +
1568 + # Block 1: 4 tokens + natural end
1569 + output.append(self.THINK_START)
1570 + self._step(holder, list(output))
1571 + for _ in range(4):
1572 + output.append(self.THINK_TOKEN)
1573 + self._step(holder, list(output))
1574 + output.append(self.THINK_END_SINGLE[0])
1575 + self._step(holder, list(output))
1576 +
1577 + # Content
1578 + for _ in range(2):
1579 + output.append(self.CONTENT_TOKEN)
1580 + self._step(holder, list(output))
1581 +
1582 + # Block 2: 3 tokens + natural end
1583 + output.append(self.THINK_START)
1584 + self._step(holder, list(output))
1585 + for _ in range(3):
1586 + output.append(self.THINK_TOKEN)
1587 + self._step(holder, list(output))
1588 + output.append(self.THINK_END_SINGLE[0])
1589 + self._step(holder, list(output))
1590 +
1591 + # Content
1592 + for _ in range(2):
1593 + output.append(self.CONTENT_TOKEN)
1594 + self._step(holder, list(output))
1595 +
1596 + # Block 3: exceed budget
1597 + output.append(self.THINK_START)
1598 + self._step(holder, list(output))
1599 + for _ in range(self.BUDGET + 1):
1600 + output.append(self.THINK_TOKEN)
1601 + self._step(holder, list(output))
1602 +
1603 + assert holder._state[0]["in_end"], (
1604 + "Third block after two natural ends must be budget-enforced"
1605 + )
1606 +
1607 + def test_natural_then_forced_in_block_2(self):
1608 + """Block 1 ends naturally, Block 2 exceeds budget -> forced end."""
1609 + budget = 5
1610 + holder = self._make_holder(self.THINK_END_SINGLE)
1611 + self._sync_batch(holder, budget)
1612 +
1613 + output = []
1614 +
1615 + # Block 1: 3 tokens + natural end (under budget)
1616 + output.append(self.THINK_START)
1617 + self._step(holder, list(output))
1618 + for _ in range(3):
1619 + output.append(self.THINK_TOKEN)
1620 + self._step(holder, list(output))
1621 + output.append(self.THINK_END_SINGLE[0])
1622 + self._step(holder, list(output))
1623 +
1624 + assert not holder._state[0]["in_end"]
1625 +
1626 + # Content
1627 + for _ in range(3):
1628 + output.append(self.CONTENT_TOKEN)
1629 + self._step(holder, list(output))
1630 +
1631 + # Block 2: exceed budget
1632 + output.append(self.THINK_START)
1633 + self._step(holder, list(output))
1634 + for _ in range(budget + 1):
1635 + output.append(self.THINK_TOKEN)
1636 + self._step(holder, list(output))
1637 +
1638 + assert holder._state[0]["in_end"], (
1639 + "Block 2 must trigger forced end after natural end in Block 1"
1640 + )
1641 +
1642 + def test_budget_one_natural_end_reentry(self):
1643 + """Budget = 1: natural end with 0 think tokens, re-entry enforced."""
1644 + budget = 1
1645 + holder = self._make_holder(self.THINK_END_SINGLE)
1646 + self._sync_batch(holder, budget)
1647 +
1648 + output = []
1649 +
1650 + # Block 1: immediate natural end (0 think tokens)
1651 + output.append(self.THINK_START)
1652 + self._step(holder, list(output))
1653 + output.append(self.THINK_END_SINGLE[0])
1654 + self._step(holder, list(output))
1655 +
1656 + assert not holder._state[0]["in_end"]
1657 +
1658 + # Content
1659 + for _ in range(3):
1660 + output.append(self.CONTENT_TOKEN)
1661 + self._step(holder, list(output))
1662 +
1663 + # Block 2: 2 think tokens (should be enforced at budget=1)
1664 + output.append(self.THINK_START)
1665 + self._step(holder, list(output))
1666 + output.append(self.THINK_TOKEN)
1667 + self._step(holder, list(output))
1668 + output.append(self.THINK_TOKEN)
1669 + self._step(holder, list(output))
1670 +
1671 + assert holder._state[0]["in_end"], (
1672 + "Budget=1 must enforce after 1 token in re-entry block"
1673 + )
1674 +
1675 + def test_partial_end_sequence_in_content(self):
1676 + """Partial end sequence in content shouldn't trigger false exit."""
1677 + holder = self._make_holder(self.THINK_END_MULTI)
1678 + self._sync_batch(holder, self.BUDGET)
1679 +
1680 + output = []
1681 +
1682 + # Block 1: natural end with full sequence [200, 201, 202]
1683 + output.append(self.THINK_START)
1684 + self._step(holder, list(output))
1685 + for _ in range(3):
1686 + output.append(self.THINK_TOKEN)
1687 + self._step(holder, list(output))
1688 + for tok in self.THINK_END_MULTI:
1689 + output.append(tok)
1690 + self._step(holder, list(output))
1691 +
1692 + # Content with partial end sequence [200, 201] but not [200, 201, 202]
1693 + output.append(self.CONTENT_TOKEN)
1694 + self._step(holder, list(output))
1695 + output.append(200)
1696 + self._step(holder, list(output))
1697 + output.append(201)
1698 + self._step(holder, list(output))
1699 + output.append(self.CONTENT_TOKEN)
1700 + self._step(holder, list(output))
1701 +
1702 + # Block 2: re-entry must still be enforced
1703 + output.append(self.THINK_START)
1704 + self._step(holder, list(output))
1705 + for _ in range(self.BUDGET + 1):
1706 + output.append(self.THINK_TOKEN)
1707 + self._step(holder, list(output))
1708 +
1709 + assert holder._state[0]["in_end"], (
1710 + "Partial end tokens in content must not prevent Block 2 enforcement"
1711 + )
1712 +
1713 + def test_continue_thinking_natural_end_reentry(self):
1714 + """Prompt-prefill inside <think>: after natural end, Block 2 must be enforced.
1715 +
1716 + When a request arrives with the prompt already inside a <think> block,
1717 + continue_thinking is set to True. This test verifies that after a
1718 + natural </think> exit, continue_thinking is cleared and Block 2
1719 + enforcement works correctly.
1720 + """
1721 + from dataclasses import dataclass
1722 + from unittest.mock import MagicMock
1723 +
1724 + from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder
1725 +
1726 + @dataclass
1727 + class FakeReasoningConfig:
1728 + reasoning_start_token_ids: list[int]
1729 + reasoning_end_token_ids: list[int]
1730 + enabled: bool = True
1731 +
1732 + cfg = FakeReasoningConfig(
1733 + reasoning_start_token_ids=[self.THINK_START],
1734 + reasoning_end_token_ids=self.THINK_END_SINGLE,
1735 + )
1736 + holder = ThinkingBudgetStateHolder(
1737 + reasoning_config=cfg,
1738 + max_num_seqs=8,
1739 + num_spec_tokens=0,
1740 + device=torch.device("cpu"),
1741 + is_pin_memory=False,
1742 + )
1743 +
1744 + # Simulate: prompt already inside <think> (sets continue_thinking=True)
1745 + prompt_tok_ids = [self.THINK_START]
1746 + params = MagicMock()
1747 + params.thinking_token_budget = self.BUDGET
1748 + batch_update = MagicMock(
1749 + removed=[],
1750 + added=[(0, params, prompt_tok_ids, [])],
1751 + moved=[],
1752 + )
1753 + holder.sync_batch(batch_update)
1754 + assert holder._state[0]["continue_thinking"], (
1755 + "Prompt inside <think> must set continue_thinking=True"
1756 + )
1757 +
1758 + # Block 1: emit some think tokens then natural </think>
1759 + output = list(prompt_tok_ids)
1760 + for _ in range(4):
1761 + output.append(self.THINK_TOKEN)
1762 + holder.update_state([output], None, None)
1763 + output.append(self.THINK_END_SINGLE[0])
1764 + holder.update_state([output], None, None)
1765 +
1766 + assert not holder._state[0]["in_think"], "Should have exited think mode"
1767 + assert not holder._state[0]["continue_thinking"], (
1768 + "continue_thinking must be cleared after natural end"
1769 + )
1770 +
1771 + # Content tokens
1772 + for _ in range(3):
1773 + output.append(self.CONTENT_TOKEN)
1774 + holder.update_state([output], None, None)
1775 +
1776 + # Block 2: re-entry must be detected and budget enforced
1777 + output.append(self.THINK_START)
1778 + holder.update_state([output], None, None)
1779 + for _ in range(self.BUDGET + 1):
1780 + output.append(self.THINK_TOKEN)
1781 + holder.update_state([output], None, None)
1782 +
1783 + assert holder._state[0]["in_end"], (
1784 + "Block 2 after natural end of prompt-prefill session must be enforced"
1785 + )
1786 +
1787 + # --- Forced-end re-entry tests ---
vllm/v1/sample/thinking_budget_state.py
+45 -44
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any
7 import torch
8
9 from vllm.platforms import current_platform
10 -from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
10 +from vllm.utils.torch_utils import async_tensor_h2d
11 from vllm.v1.sample.logits_processor.interface import (
12 BatchUpdate,
13 MoveDirectionality,
@@ -22,11 +22,12 @@ def maybe_create_thinking_budget_state_holder(
22 max_num_seqs: int,
23 num_spec_tokens: int,
24 device: torch.device,
25 + is_pin_memory: bool,
26 ) -> "ThinkingBudgetStateHolder | None":
27 if reasoning_config is None:
28 return None
29 return ThinkingBudgetStateHolder(
29 - reasoning_config, max_num_seqs, num_spec_tokens, device, PIN_MEMORY
30 + reasoning_config, max_num_seqs, num_spec_tokens, device, is_pin_memory
31 )
32
33
@@ -173,19 +174,6 @@ class ThinkingBudgetStateHolder:
174 return i
175 return -1
176
176 - @staticmethod
177 - def _find_last_sequence_index_from(
178 - target_list: list[int], token_ids: list[int], search_start: int
179 - ) -> int:
180 - """Last occurrence of ``token_ids`` at or after ``search_start``."""
181 - if not token_ids:
182 - return -1
183 - lo = max(0, search_start)
184 - for i in range(len(target_list) - len(token_ids), lo - 1, -1):
185 - if target_list[i : i + len(token_ids)] == token_ids:
186 - return i
187 - return -1
188 -
177 def _init_state_entry(
178 self, prompt_tok_ids: list[int] | None, thinking_token_budget: int
179 ) -> dict[str, Any]:
@@ -239,8 +227,6 @@ class ThinkingBudgetStateHolder:
227 "force_index": [],
228 "start_thinking": start_thinking,
229 "end_thinking": -1,
242 - "start_search_pos": 0,
243 - "end_search_pos": 0,
230 "in_spec_mode": False,
231 "bonus_token_forced": False,
232 "continue_thinking": continue_thinking,
@@ -256,41 +242,40 @@ class ThinkingBudgetStateHolder:
242 state["force_index"] = []
243 return
244
259 - output_tok_ids = state.get("output_tok_ids", [])
245 if state["start_thinking"] == -1:
261 - seq_len = len(self.think_start_token_ids)
246 scan_offset = state.get("scan_offset", 0)
263 - start_thinking = self._find_last_sequence_index_from(
264 - output_tok_ids,
265 - self.think_start_token_ids,
266 - max(scan_offset, state["start_search_pos"] - (seq_len - 1)),
247 + output_slice = state.get("output_tok_ids", [])[scan_offset:]
248 + start_thinking = self._find_last_sequence_index(
249 + output_slice, self.think_start_token_ids
250 )
268 - if start_thinking >= 0 and scan_offset > 0:
269 - # Re-entry after a forced end: budget was already exhausted
270 - # in a prior block, so immediately force-close this one.
271 - # scan_offset > 0 is only set after forced-end completion
272 - # (never after natural end), so this won't block legitimate
273 - # re-entries where budget remains.
274 - state["start_thinking"] = start_thinking
275 - state["in_think"] = False
276 - state["in_end"] = True
277 - state["end_count"] = 0
278 - state["force_index"] = [0]
279 - return
251 + if start_thinking >= 0:
252 + start_thinking += scan_offset
253 state["start_thinking"] = start_thinking
281 - if start_thinking == -1:
282 - state["start_search_pos"] = len(output_tok_ids)
254 if state["end_thinking"] == -1:
284 - seq_len = len(self.think_end_token_ids)
255 scan_offset = state.get("scan_offset", 0)
286 - end_thinking = self._find_last_sequence_index_from(
287 - output_tok_ids,
288 - self.think_end_token_ids,
289 - max(scan_offset, state["end_search_pos"] - (seq_len - 1)),
256 + output_slice = state.get("output_tok_ids", [])[scan_offset:]
257 + end_thinking = self._find_last_sequence_index(
258 + output_slice, self.think_end_token_ids
259 )
260 + if end_thinking >= 0:
261 + end_thinking += scan_offset
262 state["end_thinking"] = end_thinking
292 - if end_thinking == -1:
293 - state["end_search_pos"] = len(output_tok_ids)
263 +
264 + if (
265 + not state.get("in_end", False)
266 + and state["start_thinking"] >= 0
267 + and state["end_thinking"] >= 0
268 + and state["end_thinking"] > state["start_thinking"]
269 + and not state.get("continue_thinking", False)
270 + ):
271 + state["in_think"] = False
272 + state["think_count"] = 0
273 + state["continue_thinking"] = False
274 + state["start_thinking"] = -1
275 + state["end_thinking"] = -1
276 + state["scan_offset"] = len(state.get("output_tok_ids", []))
277 + state["check_count_down"] = state["thinking_token_budget"]
278 + return
279
280 if state["start_thinking"] == -1:
281 return
@@ -314,10 +299,18 @@ class ThinkingBudgetStateHolder:
299 predicted_countdown = current_step_countdown - len(state["spec_token_ids"]) - 1
300 # We only proceed further if we have counted down the thinking budget
301 # to 0 or less and when we are in the "in think" mode.
302 + # Exception: when continue_thinking=True and a natural </think> is
303 + # detected (end_thinking != -1), fall through to handle the exit —
304 + # even if the budget hasn't expired yet. For continue_thinking=False,
305 + # the early natural-end detection block above already handles it.
306 + natural_end_with_continue = (
307 + state.get("continue_thinking", False) and state["end_thinking"] != -1
308 + )
309 if (
310 not state.get("in_end", False)
311 and predicted_countdown >= 0
312 and state["start_thinking"] > -1
313 + and not natural_end_with_continue
314 ):
315 state["check_count_down"] = current_step_countdown
316 state["prev_output_length"] = len(state.get("output_tok_ids", []))
@@ -389,6 +382,10 @@ class ThinkingBudgetStateHolder:
382 # Case: ...<start>...<end>... - exiting think mode
383 state["in_think"] = False
384 state["think_count"] = 0
385 + state["continue_thinking"] = False
386 + state["start_thinking"] = -1
387 + state["end_thinking"] = -1
388 + state["scan_offset"] = len(state.get("output_tok_ids", []))
389
390 elif absolute_start_pos >= 0 and not state["continue_thinking"]:
391 # Found think start - entering think mode
@@ -400,6 +397,10 @@ class ThinkingBudgetStateHolder:
397 # Found think end - exiting think mode
398 state["in_think"] = False
399 state["think_count"] = 0
400 + state["continue_thinking"] = False
401 + state["start_thinking"] = -1
402 + state["end_thinking"] = -1
403 + state["scan_offset"] = len(state.get("output_tok_ids", []))
404
405 elif state["in_think"]:
406 # Continue thinking mode, increment count by new tokens
vllm/v1/worker/gpu_input_batch.py
+1
@@ -113,6 +113,7 @@ class InputBatch:
113 max_num_reqs,
114 num_spec_tokens,
115 device,
116 + PIN_MEMORY,
117 )
118 self.thinking_token_budget_reqs: set[str] = set()
119 self.is_pooling_model = is_pooling_model