update states only in update_state method

Signed-off-by: Sungjae Lee <33976427+llsj14@users.noreply.github.com>

Author: llsj14

vllm/v1/sample/logits_processor.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass, field
 from enum import Enum
 from itertools import chain
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 from torch._prims_common import DeviceLikeType
@@ -502,9 +502,9 @@ def __init__(self, reasoning_config: ReasoningConfig, pin_memory: bool, device:
         self.think_end_token_id = reasoning_config.think_end_token_id
         self.pin_memory = pin_memory
         self.device = device
-        self._state = {}
+        self._state: dict[int, dict[str, Any]] = {}
 
-    def _find_last_token_index(self, tokens, token_id):
+    def _find_last_token_index(self, tokens: list[int], token_id: int) -> int:
         try:
             return len(tokens) - tokens[::-1].index(token_id) - 1
         except ValueError:
@@ -516,71 +516,61 @@ def is_argmax_invariant(self) -> bool:
         return False
 
     def update_state(self, batch_update: Optional[BatchUpdate]):
-        if batch_update is None:
-            return
-
-        for index, params, prompt_tok_ids, output_tok_ids in batch_update.added:
-            max_think_tokens = params.max_think_tokens if isinstance(params, SamplingParams) else None
-
-            if max_think_tokens is None:
-                continue
-
-            last_think_start_idx = self._find_last_token_index(prompt_tok_ids, self.think_start_token_id)
-            last_think_end_idx = self._find_last_token_index(prompt_tok_ids, self.think_end_token_id)
-
-            in_think = False
-            count = 0
+        if batch_update:
+            for index, params, prompt_tok_ids, output_tok_ids in batch_update.added:
+                max_think_tokens = params.max_think_tokens if isinstance(params, SamplingParams) else None
+                if max_think_tokens is not None:
+                    last_start = self._find_last_token_index(prompt_tok_ids, self.think_start_token_id)
+                    last_end = self._find_last_token_index(prompt_tok_ids, self.think_end_token_id)
+                    in_think = last_start > last_end
+                    count = len(prompt_tok_ids) - (last_start + 1) if in_think else 0
+
+                    self._state[index] = {
+                        "in_think": in_think,
+                        "count": count,
+                        "prompt_tok_ids": prompt_tok_ids,
+                        "output_tok_ids": output_tok_ids,
+                        "max_think_tokens": max_think_tokens,
+                    }
 
-            if last_think_start_idx > last_think_end_idx:
-                in_think = True
-                count = len(prompt_tok_ids) - (last_think_start_idx + 1)
+            for index in batch_update.removed:
+                self._state.pop(index, None)
 
-            self._state[index] = {
-                "in_think": in_think,
-                "count": count,
-                "prompt_tok_ids": prompt_tok_ids,
-                "output_tok_ids": output_tok_ids,
-                "max_think_tokens": max_think_tokens,
-            }
+            for i1, i2, direction in batch_update.moved:
+                if direction == MoveDirectionality.SWAP:
+                    self._state[i1], self._state[i2] = self._state[i2], self._state[i1]
+                else:
+                    self._state[i2] = self._state.pop(i1, None)
 
-        for index in batch_update.removed:
-            self._state.pop(index, None)
+        # Update in_think and count for all active requests
+        for state in self._state.values():
+            output = state["output_tok_ids"]
+            if not output:
+                continue
 
-        for i1, i2, direction in batch_update.moved:
-            if direction == MoveDirectionality.SWAP:
-                self._state[i1], self._state[i2] = self._state[i2], self._state[i1]
-            else:
-                self._state[i2] = self._state.pop(i1, None)
+            last_tok = output[-1]
+            if last_tok == self.think_start_token_id:
+                state["in_think"] = True
+                state["count"] = 0
+            elif last_tok == self.think_end_token_id:
+                state["in_think"] = False
+                state["count"] = 0
+            elif state["in_think"]:
+                state["count"] += 1
 
     def apply(self, logits: torch.Tensor) -> torch.Tensor:
         batch_size = logits.size(0)
-        if batch_size == 0:
+        if not self._state:
             return logits
 
         mask = torch.zeros(batch_size, dtype=torch.bool, device=logits.device)
         end_token_id = self.think_end_token_id
 
         for index in range(batch_size):
-            state = self._state.get(index, None)
-            if not state or not state.get("output_tok_ids"):
+            state = self._state.get(index)
+            if not state:
                 continue
 
-            last_tok = state["output_tok_ids"][-1]
-            in_think = state["in_think"]
-            count = state["count"]
-
-            if last_tok == self.think_start_token_id:
-                in_think = True
-                count = 0
-            elif last_tok == self.think_end_token_id:
-                in_think = False
-                count = 0
-            elif in_think:
-                count += 1
-
-            state["in_think"] = in_think
-            state["count"] = count
-
             if state["in_think"] and state["count"] >= state["max_think_tokens"]:
                 mask[index] = True