| 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.platforms import current_platform |
| 10 | from vllm.utils.torch_utils import async_tensor_h2d |
| 11 | from vllm.v1.sample.logits_processor.interface import ( |
| 12 | BatchUpdate, |
| 13 | MoveDirectionality, |
| 14 | ) |
| 15 | |
| 16 | if TYPE_CHECKING: |
| 17 | from vllm.config.reasoning import ReasoningConfig |
| 18 | |
| 19 | |
| 20 | def maybe_create_thinking_budget_state_holder( |
| 21 | reasoning_config: "ReasoningConfig | None", |
| 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( |
| 30 | reasoning_config, max_num_seqs, num_spec_tokens, device, is_pin_memory |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | class ThinkingBudgetStateHolder: |
| 35 | """Tracks thinking sections and forces end tokens when budget is exceeded.""" |
| 36 | |
| 37 | think_start_token_ids: list[int] |
| 38 | think_end_token_ids: list[int] |
| 39 | |
| 40 | def __init__( |
| 41 | self, |
| 42 | reasoning_config: "ReasoningConfig | None", |
| 43 | max_num_seqs: int, |
| 44 | num_spec_tokens: int, |
| 45 | device: torch.device, |
| 46 | is_pin_memory: bool, |
| 47 | ): |
| 48 | _ = is_pin_memory # API parity with logits processors |
| 49 | max_num_reqs = max_num_seqs |
| 50 | self.in_spec_mode = num_spec_tokens > 0 |
| 51 | self.num_spec_tokens = num_spec_tokens |
| 52 | |
| 53 | # No separate enable flag: a non-``None`` ``reasoning_config`` is the switch. |
| 54 | self.is_enabled = reasoning_config is not None |
| 55 | |
| 56 | if reasoning_config is None: |
| 57 | self.think_start_token_ids = [] |
| 58 | self.think_end_token_ids = [] |
| 59 | else: |
| 60 | rs = reasoning_config.reasoning_start_token_ids |
| 61 | re = reasoning_config.reasoning_end_token_ids |
| 62 | self.think_start_token_ids = rs if rs else [] |
| 63 | self.think_end_token_ids = re if re else [] |
| 64 | |
| 65 | self.device = device |
| 66 | self._state: dict[int, dict[str, Any]] = {} |
| 67 | self.cu_num_tokens: dict[int, int] = {} |
| 68 | |
| 69 | if self.num_spec_tokens > 0: |
| 70 | self._mask_capacity = max_num_reqs * (self.num_spec_tokens + 1) |
| 71 | else: |
| 72 | self._mask_capacity = max_num_reqs |
| 73 | |
| 74 | def has_tracked_requests(self) -> bool: |
| 75 | """True when ``sync_batch`` has state for a ``thinking_token_budget`` row. |
| 76 | |
| 77 | Used to decide whether sampling needs output-token rows and spec combining; |
| 78 | distinct from merely having a holder instance (reasoning may be on with no |
| 79 | budgeted requests in this batch). |
| 80 | """ |
| 81 | return bool(self._state) |
| 82 | |
| 83 | def sync_batch(self, batch_update: BatchUpdate | None) -> None: |
| 84 | """Add/remove/move per-request state only (no _update_think_state).""" |
| 85 | if not self.is_enabled or not batch_update: |
| 86 | return |
| 87 | for index in batch_update.removed: |
| 88 | self._state.pop(index, None) |
| 89 | |
| 90 | for index, params, prompt_tok_ids, output_tok_ids in batch_update.added: |
| 91 | thinking_token_budget = params.thinking_token_budget |
| 92 | if thinking_token_budget is not None: |
| 93 | self._state[index] = self._init_state_entry( |
| 94 | prompt_tok_ids, thinking_token_budget |
| 95 | ) |
| 96 | self._state[index]["output_tok_ids"] = output_tok_ids |
| 97 | self._state[index]["spec_token_ids"] = [] |
| 98 | else: |
| 99 | self._state.pop(index, None) |
| 100 | |
| 101 | for i1, i2, direction in batch_update.moved: |
| 102 | if direction == MoveDirectionality.SWAP: |
| 103 | state1 = self._state.pop(i1, None) |
| 104 | state2 = self._state.pop(i2, None) |
| 105 | if state1 is not None: |
| 106 | self._state[i2] = state1 |
| 107 | if state2 is not None: |
| 108 | self._state[i1] = state2 |
| 109 | else: |
| 110 | state = self._state.pop(i1, None) |
| 111 | if state is not None: |
| 112 | self._state[i2] = state |
| 113 | |
| 114 | def update_state( |
| 115 | self, |
| 116 | output_token_ids: list[list[int]], |
| 117 | spec_token_ids: list[list[int]] | None, |
| 118 | repeat_indices: torch.Tensor | None = None, |
| 119 | ) -> None: |
| 120 | """Refresh output/spec from sampling rows and recompute think state.""" |
| 121 | if not self.is_enabled or not self._state: |
| 122 | return |
| 123 | |
| 124 | spec_lists = spec_token_ids or [] |
| 125 | last_row_for_req: dict[int, int] | None = None |
| 126 | if repeat_indices is not None: |
| 127 | last_row_for_req = {} |
| 128 | rpt = repeat_indices.cpu().tolist() |
| 129 | for batch_row, req_i in enumerate(rpt): |
| 130 | last_row_for_req[req_i] = batch_row |
| 131 | |
| 132 | for seq_idx, state in list(self._state.items()): |
| 133 | if last_row_for_req is not None: |
| 134 | output_row: int | None = last_row_for_req.get(seq_idx) |
| 135 | if output_row is None or output_row >= len(output_token_ids): |
| 136 | continue |
| 137 | state["output_tok_ids"] = output_token_ids[output_row] |
| 138 | elif seq_idx >= len(output_token_ids): |
| 139 | continue |
| 140 | else: |
| 141 | state["output_tok_ids"] = output_token_ids[seq_idx] |
| 142 | if seq_idx < len(spec_lists): |
| 143 | state["spec_token_ids"] = list(spec_lists[seq_idx]) |
| 144 | else: |
| 145 | state["spec_token_ids"] = [] |
| 146 | state["in_spec_mode"] = self.in_spec_mode |
| 147 | state["force_index"] = [] |
| 148 | self._update_think_state(state) |
| 149 | |
| 150 | def apply_to_logits( |
| 151 | self, |
| 152 | logits: torch.Tensor, |
| 153 | predict_bonus_token: bool, |
| 154 | spec_token_ids: list[list[int]] | None, |
| 155 | ) -> torch.Tensor: |
| 156 | """Mask and bump logits for forced end-of-thinking tokens.""" |
| 157 | if not self.is_enabled or not self._state: |
| 158 | return logits |
| 159 | spec_lists = spec_token_ids or [] |
| 160 | return self._apply_forcing_to_logits(logits, predict_bonus_token, spec_lists) |
| 161 | |
| 162 | @staticmethod |
| 163 | def _find_last_sequence_index(target_list: list[int], token_ids: list[int]) -> int: |
| 164 | if not token_ids: |
| 165 | return -1 |
| 166 | for i in range(len(target_list) - len(token_ids), -1, -1): |
| 167 | if target_list[i : i + len(token_ids)] == token_ids: |
| 168 | return i |
| 169 | return -1 |
| 170 | |
| 171 | def _init_state_entry( |
| 172 | self, prompt_tok_ids: list[int] | None, thinking_token_budget: int |
| 173 | ) -> dict[str, Any]: |
| 174 | if prompt_tok_ids is None: |
| 175 | last_start = -1 |
| 176 | last_end = -1 |
| 177 | in_think = False |
| 178 | think_count = 0 |
| 179 | start_thinking = -1 |
| 180 | countdown = thinking_token_budget |
| 181 | continue_thinking = False |
| 182 | in_end = False |
| 183 | else: |
| 184 | start_thinking = -1 |
| 185 | countdown = thinking_token_budget |
| 186 | continue_thinking = False |
| 187 | in_end = False |
| 188 | last_start = self._find_last_sequence_index( |
| 189 | prompt_tok_ids, self.think_start_token_ids |
| 190 | ) |
| 191 | last_end = self._find_last_sequence_index( |
| 192 | prompt_tok_ids, self.think_end_token_ids |
| 193 | ) |
| 194 | in_think = last_start > last_end |
| 195 | # load metrics such as think count, start thinking |
| 196 | # if request is in thinking mode, already |
| 197 | if in_think: |
| 198 | think_count = len(prompt_tok_ids) - ( |
| 199 | last_start + len(self.think_start_token_ids) |
| 200 | ) |
| 201 | start_thinking = len(prompt_tok_ids) - think_count - 1 |
| 202 | countdown -= think_count |
| 203 | continue_thinking = True |
| 204 | # check if the token is exhausted within prompt |
| 205 | token_exhausted = thinking_token_budget - think_count |
| 206 | in_end = token_exhausted <= 0 |
| 207 | else: |
| 208 | think_count = 0 |
| 209 | |
| 210 | return { |
| 211 | "in_think": in_think, |
| 212 | "in_end": in_end, |
| 213 | "check_count_down": countdown, |
| 214 | "think_count": think_count, |
| 215 | "end_count": 0, |
| 216 | "prompt_tok_ids": prompt_tok_ids, |
| 217 | "output_tok_ids": [], |
| 218 | "thinking_token_budget": thinking_token_budget, |
| 219 | "prev_output_length": 0, |
| 220 | "spec_token_ids": [], |
| 221 | "force_index": [], |
| 222 | "start_thinking": start_thinking, |
| 223 | "end_thinking": -1, |
| 224 | "in_spec_mode": False, |
| 225 | "bonus_token_forced": False, |
| 226 | "continue_thinking": continue_thinking, |
| 227 | "scan_offset": 0, |
| 228 | } |
| 229 | |
| 230 | def _update_think_state(self, state: dict[str, Any]) -> None: |
| 231 | if state.get("thinking_token_budget", -1) == -1: |
| 232 | return |
| 233 | if len(self.think_end_token_ids) == 0: |
| 234 | state["thinking_token_budget"] = -1 |
| 235 | state["in_end"] = False |
| 236 | state["force_index"] = [] |
| 237 | return |
| 238 | |
| 239 | if state["start_thinking"] == -1: |
| 240 | scan_offset = state.get("scan_offset", 0) |
| 241 | output_slice = state.get("output_tok_ids", [])[scan_offset:] |
| 242 | start_thinking = self._find_last_sequence_index( |
| 243 | output_slice, self.think_start_token_ids |
| 244 | ) |
| 245 | if start_thinking >= 0: |
| 246 | start_thinking += scan_offset |
| 247 | state["start_thinking"] = start_thinking |
| 248 | if state["end_thinking"] == -1: |
| 249 | scan_offset = state.get("scan_offset", 0) |
| 250 | output_slice = state.get("output_tok_ids", [])[scan_offset:] |
| 251 | end_thinking = self._find_last_sequence_index( |
| 252 | output_slice, self.think_end_token_ids |
| 253 | ) |
| 254 | if end_thinking >= 0: |
| 255 | end_thinking += scan_offset |
| 256 | state["end_thinking"] = end_thinking |
| 257 | |
| 258 | if ( |
| 259 | not state.get("in_end", False) |
| 260 | and state["start_thinking"] >= 0 |
| 261 | and state["end_thinking"] >= 0 |
| 262 | and state["end_thinking"] > state["start_thinking"] |
| 263 | and not state.get("continue_thinking", False) |
| 264 | ): |
| 265 | state["in_think"] = False |
| 266 | state["think_count"] = 0 |
| 267 | state["continue_thinking"] = False |
| 268 | state["start_thinking"] = -1 |
| 269 | state["end_thinking"] = -1 |
| 270 | state["scan_offset"] = len(state.get("output_tok_ids", [])) |
| 271 | state["check_count_down"] = state["thinking_token_budget"] |
| 272 | return |
| 273 | |
| 274 | if state["start_thinking"] == -1: |
| 275 | return |
| 276 | |
| 277 | if state["continue_thinking"]: |
| 278 | sampled_tokens_from_previous_step = len( |
| 279 | state.get("output_tok_ids", []) |
| 280 | ) - state.get("prev_output_length", 0) |
| 281 | else: |
| 282 | if state["prev_output_length"] == 0: |
| 283 | sampled_tokens_from_previous_step = len( |
| 284 | state.get("output_tok_ids", []) |
| 285 | ) - len(self.think_start_token_ids) |
| 286 | else: |
| 287 | sampled_tokens_from_previous_step = ( |
| 288 | len(state.get("output_tok_ids", [])) - state["prev_output_length"] |
| 289 | ) |
| 290 | current_step_countdown = ( |
| 291 | state["check_count_down"] - sampled_tokens_from_previous_step |
| 292 | ) |
| 293 | predicted_countdown = current_step_countdown - len(state["spec_token_ids"]) - 1 |
| 294 | # We only proceed further if we have counted down the thinking budget |
| 295 | # to 0 or less and when we are in the "in think" mode. |
| 296 | # Exception: when continue_thinking=True and a natural </think> is |
| 297 | # detected (end_thinking != -1), fall through to handle the exit — |
| 298 | # even if the budget hasn't expired yet. For continue_thinking=False, |
| 299 | # the early natural-end detection block above already handles it. |
| 300 | natural_end_with_continue = ( |
| 301 | state.get("continue_thinking", False) and state["end_thinking"] != -1 |
| 302 | ) |
| 303 | if ( |
| 304 | not state.get("in_end", False) |
| 305 | and predicted_countdown >= 0 |
| 306 | and state["start_thinking"] > -1 |
| 307 | and not natural_end_with_continue |
| 308 | ): |
| 309 | state["check_count_down"] = current_step_countdown |
| 310 | state["prev_output_length"] = len(state.get("output_tok_ids", [])) |
| 311 | return |
| 312 | output = state.get("output_tok_ids", []) |
| 313 | if not output: |
| 314 | # When in_end was set at init (budget=0, prompt already in think), |
| 315 | # we must force the first generated token to be the end token; |
| 316 | # otherwise apply() sees in_end=True but force_index=[] and |
| 317 | # allows an extra thinking token. |
| 318 | if state.get("in_end", False): |
| 319 | state["force_index"] = [0] |
| 320 | return |
| 321 | |
| 322 | # Track previous output length for incremental processing |
| 323 | prev_length = state.get("prev_output_length", 0) |
| 324 | current_length = len(output) |
| 325 | |
| 326 | if current_length <= prev_length: |
| 327 | if state.get("in_end", False): |
| 328 | remaining_budget = state["thinking_token_budget"] - state["think_count"] |
| 329 | spec_len = len(state["spec_token_ids"]) |
| 330 | if spec_len > 0: |
| 331 | if 0 < remaining_budget < spec_len: |
| 332 | state["force_index"] = [remaining_budget] |
| 333 | elif remaining_budget <= 0: |
| 334 | state["force_index"] = [0] |
| 335 | else: |
| 336 | state["force_index"] = [spec_len] |
| 337 | else: |
| 338 | state["force_index"] = [0] |
| 339 | return |
| 340 | |
| 341 | state["prev_output_length"] = current_length |
| 342 | |
| 343 | start_len = len(self.think_start_token_ids) |
| 344 | absolute_start_pos = state["start_thinking"] |
| 345 | |
| 346 | if state["continue_thinking"] and state["end_thinking"] > -1: |
| 347 | absolute_end_pos = state["end_thinking"] + len( |
| 348 | state.get("prompt_tok_ids") or [] |
| 349 | ) |
| 350 | else: |
| 351 | absolute_end_pos = state["end_thinking"] |
| 352 | # Update state based on recent sequences |
| 353 | # This is the case where we are in end mode, but the rejection sampler |
| 354 | # rejected a token before the end token, |
| 355 | # so we need to go back to think mode and wait for the next end token |
| 356 | # eg with 999: [2,4,5,999] -> [3,-1,-1,-1] |
| 357 | if state["in_end"] and state["end_count"] == 0: |
| 358 | new_tokens = output[prev_length:] |
| 359 | stopping_thinking = ( |
| 360 | self.think_end_token_ids[state["end_count"]] in new_tokens |
| 361 | ) |
| 362 | if not stopping_thinking: |
| 363 | state["in_think"] = True |
| 364 | state["in_end"] = False |
| 365 | state["end_count"] = 0 |
| 366 | state["bonus_token_forced"] = False |
| 367 | |
| 368 | if not state["in_end"]: |
| 369 | if absolute_start_pos >= 0 and absolute_end_pos >= 0: |
| 370 | # Case: ...<end>...<start>... - entering think mode |
| 371 | if absolute_start_pos > absolute_end_pos: |
| 372 | new_think_count = current_length - (absolute_start_pos + start_len) |
| 373 | state["in_think"] = True |
| 374 | state["think_count"] = new_think_count |
| 375 | else: |
| 376 | # Case: ...<start>...<end>... - exiting think mode |
| 377 | state["in_think"] = False |
| 378 | state["think_count"] = 0 |
| 379 | state["continue_thinking"] = False |
| 380 | state["start_thinking"] = -1 |
| 381 | state["end_thinking"] = -1 |
| 382 | state["scan_offset"] = len(state.get("output_tok_ids", [])) |
| 383 | |
| 384 | elif absolute_start_pos >= 0 and not state["continue_thinking"]: |
| 385 | # Found think start - entering think mode |
| 386 | new_think_count = current_length - (absolute_start_pos + start_len) |
| 387 | state["in_think"] = True |
| 388 | state["think_count"] = new_think_count |
| 389 | |
| 390 | elif absolute_end_pos >= 0: |
| 391 | # Found think end - exiting think mode |
| 392 | state["in_think"] = False |
| 393 | state["think_count"] = 0 |
| 394 | state["continue_thinking"] = False |
| 395 | state["start_thinking"] = -1 |
| 396 | state["end_thinking"] = -1 |
| 397 | state["scan_offset"] = len(state.get("output_tok_ids", [])) |
| 398 | |
| 399 | elif state["in_think"]: |
| 400 | # Continue thinking mode, increment count by new tokens |
| 401 | prompt_tok_ids = state.get("prompt_tok_ids") or [] |
| 402 | think_tokens_in_prompt = len(prompt_tok_ids) - ( |
| 403 | absolute_start_pos + start_len |
| 404 | ) |
| 405 | state["think_count"] = ( |
| 406 | len(state["output_tok_ids"]) + think_tokens_in_prompt |
| 407 | ) |
| 408 | if state["in_think"]: |
| 409 | remaining_budget = max( |
| 410 | 0, state["thinking_token_budget"] - state["think_count"] |
| 411 | ) |
| 412 | state["check_count_down"] = remaining_budget |
| 413 | else: |
| 414 | state["check_count_down"] = state["thinking_token_budget"] |
| 415 | |
| 416 | total_thinking_tokens = ( |
| 417 | state["think_count"] + len(state["spec_token_ids"]) + 1 |
| 418 | ) |
| 419 | # Check if need to transition to end mode |
| 420 | # If we have more thinking tokens than the budget, |
| 421 | # we need to transition to end mode |
| 422 | if ( |
| 423 | state["in_think"] |
| 424 | and total_thinking_tokens > state["thinking_token_budget"] |
| 425 | ): |
| 426 | # Calculate force_index: position within spec_token_ids where |
| 427 | # forcing starts. If we're already over budget without spec |
| 428 | # tokens, force from position 0. Force from the position |
| 429 | # where budget is exceeded. |
| 430 | state["in_think"] = False |
| 431 | state["in_end"] = True |
| 432 | state["end_count"] = 0 |
| 433 | state["check_count_down"] = state["thinking_token_budget"] |
| 434 | remaining_budget = state["thinking_token_budget"] - state["think_count"] |
| 435 | spec_len = len(state["spec_token_ids"]) |
| 436 | if 0 < remaining_budget < spec_len: |
| 437 | state["force_index"] = [remaining_budget] |
| 438 | |
| 439 | elif remaining_budget <= 0: |
| 440 | state["force_index"] = [0] |
| 441 | |
| 442 | else: |
| 443 | # remaining_budget >= spec_len: all spec tokens are within |
| 444 | # budget; force the bonus token position |
| 445 | state["force_index"] = [len(state["spec_token_ids"])] |
| 446 | |
| 447 | else: |
| 448 | state["force_index"] = [] |
| 449 | if len(state["spec_token_ids"]) > 0: |
| 450 | for i, token_id in enumerate(state["spec_token_ids"]): |
| 451 | if state["end_count"] + 1 < len(self.think_end_token_ids): |
| 452 | if token_id == self.think_end_token_ids[state["end_count"] + 1]: |
| 453 | state["end_count"] += 1 |
| 454 | else: |
| 455 | state["end_count"] += 1 |
| 456 | state["force_index"] = [i] |
| 457 | break |
| 458 | else: |
| 459 | state["end_count"] += 1 |
| 460 | if len(state["force_index"]) == 0: |
| 461 | state["end_count"] += 1 |
| 462 | state["force_index"] = [len(state["spec_token_ids"])] |
| 463 | else: |
| 464 | state["end_count"] += 1 |
| 465 | state["force_index"] = [0] |
| 466 | if state["end_count"] >= len(self.think_end_token_ids): |
| 467 | state.update( |
| 468 | { |
| 469 | "in_end": False, |
| 470 | "end_count": 0, |
| 471 | "check_count_down": state["thinking_token_budget"], |
| 472 | "start_thinking": -1, |
| 473 | "end_thinking": -1, |
| 474 | "think_count": 0, |
| 475 | "continue_thinking": False, |
| 476 | "scan_offset": len(state.get("output_tok_ids", [])), |
| 477 | } |
| 478 | ) |
| 479 | |
| 480 | def _apply_forcing_to_logits( |
| 481 | self, |
| 482 | logits: torch.Tensor, |
| 483 | predict_bonus_token: bool, |
| 484 | spec_token_ids_for_layout: list[list[int]], |
| 485 | ) -> torch.Tensor: |
| 486 | cumulative_total = 0 |
| 487 | self.cu_num_tokens.clear() |
| 488 | |
| 489 | n_layout = len(spec_token_ids_for_layout) |
| 490 | if self._state: |
| 491 | n_layout = max(n_layout, max(self._state.keys()) + 1) |
| 492 | |
| 493 | for index in range(n_layout): |
| 494 | self.cu_num_tokens[index] = cumulative_total |
| 495 | spec_tokens = ( |
| 496 | spec_token_ids_for_layout[index] |
| 497 | if index < len(spec_token_ids_for_layout) |
| 498 | else [] |
| 499 | ) |
| 500 | if self.in_spec_mode: |
| 501 | cumulative_total += len(spec_tokens) if not predict_bonus_token else 1 |
| 502 | else: |
| 503 | cumulative_total += 1 |
| 504 | |
| 505 | # Build the active index / forced-token lists entirely on CPU so we |
| 506 | # avoid per-iteration scalar sync writes to GPU tensors. |
| 507 | active_indices_cpu: list[int] = [] |
| 508 | force_tokens_cpu: list[int] = [] |
| 509 | |
| 510 | for seq_idx in sorted(self._state.keys()): |
| 511 | if seq_idx not in self.cu_num_tokens: |
| 512 | continue |
| 513 | state = self._state[seq_idx] |
| 514 | if state.get("in_end", False): |
| 515 | # logits processor in spec mode are called twice |
| 516 | # once for bonus token logits and |
| 517 | # second time for the target logits |
| 518 | # in case the force index is bonus token index |
| 519 | # we change the force index to 0 |
| 520 | if predict_bonus_token: |
| 521 | if state.get("force_index") and state["force_index"][0] < len( |
| 522 | state["spec_token_ids"] |
| 523 | ): |
| 524 | continue |
| 525 | else: |
| 526 | state["force_index"] = [0] |
| 527 | # continue enforcing the end thinking tokens |
| 528 | if state["end_count"] > 0: |
| 529 | state["bonus_token_forced"] = False |
| 530 | if state and not state["bonus_token_forced"]: |
| 531 | force_index = state.get("force_index", []) |
| 532 | if len(force_index) == 0: |
| 533 | continue |
| 534 | end_count = state.get("end_count", 0) |
| 535 | for force_idx in force_index: |
| 536 | if end_count < len(self.think_end_token_ids): |
| 537 | mask_idx = self.cu_num_tokens[seq_idx] + force_idx |
| 538 | if ( |
| 539 | mask_idx < self._mask_capacity |
| 540 | and mask_idx < logits.shape[0] |
| 541 | ): |
| 542 | active_indices_cpu.append(mask_idx) |
| 543 | force_tokens_cpu.append( |
| 544 | self.think_end_token_ids[end_count] |
| 545 | ) |
| 546 | if predict_bonus_token: |
| 547 | if state["end_count"] > 0: |
| 548 | state["bonus_token_forced"] = False |
| 549 | state["force_index"] = [] |
| 550 | else: |
| 551 | state["bonus_token_forced"] = True |
| 552 | |
| 553 | if active_indices_cpu: |
| 554 | device = logits.device |
| 555 | if current_platform.is_rocm() and logits.is_contiguous(): |
| 556 | # Flattened index_fill avoids ROCm faults seen with 2-D |
| 557 | # advanced-indexing writes on the thinking-budget path. |
| 558 | vocab_size = logits.shape[1] |
| 559 | flat_indices_cpu = [ |
| 560 | row * vocab_size + token |
| 561 | for row, token in zip(active_indices_cpu, force_tokens_cpu) |
| 562 | ] |
| 563 | flat_indices = async_tensor_h2d( |
| 564 | flat_indices_cpu, dtype=torch.long, device=device |
| 565 | ) |
| 566 | logits.view(-1).index_fill_(0, flat_indices, 1e9) |
| 567 | elif current_platform.is_rocm(): |
| 568 | fill = logits.new_tensor(1e9) |
| 569 | for row, token in zip(active_indices_cpu, force_tokens_cpu): |
| 570 | logits[row, token] = fill |
| 571 | else: |
| 572 | active_indices = async_tensor_h2d( |
| 573 | active_indices_cpu, dtype=torch.long, device=device |
| 574 | ) |
| 575 | force_tokens = async_tensor_h2d( |
| 576 | force_tokens_cpu, dtype=torch.long, device=device |
| 577 | ) |
| 578 | # Avoid CPU->GPU sync. |
| 579 | fill = logits.new_full((len(active_indices_cpu),), 1e9) |
| 580 | logits.index_put_((active_indices, force_tokens), fill) |
| 581 | |
| 582 | return logits |