refactor: extract shared steady-state increment detection

- add _BindState dataclass and _is_steady_increment() to cache.py
- replace _bind_sig/_bind_seq_lens dual fields with single _bind_state
- replace DecodeSteadyState bare tuple with named dataclass
- use _is_steady_increment() in both PagePool.bind_tasks and Executor.execute_decode

astrai/inference/core/cache.py

@@ -13,9 +13,8 @@ Two modes:
 """
 
 import threading
-from collections import OrderedDict
 from dataclasses import dataclass
-from typing import Callable, Dict, List, Optional
+from typing import Callable, Dict, List, Optional, OrderedDict
 
 import torch
 from torch import Tensor
@@ -23,6 +22,30 @@ from torch import Tensor
 from astrai.inference.core.workspace import InferenceWorkspace
 
 
+@dataclass
+class _BindState:
+    """Cached bind metadata for steady-state decode increment detection."""
+
+    sig: tuple
+    seq_lens: List[int]
+
+
+def _is_steady_increment(
+    prev_sig: Optional[tuple],
+    prev_vals: Optional[List[int]],
+    cur_sig: tuple,
+    cur_vals: List[int],
+) -> bool:
+    """True when the same ordered set has every value +1 from the previous step."""
+    return (
+        prev_sig is not None
+        and prev_vals is not None
+        and prev_sig == cur_sig
+        and len(prev_vals) == len(cur_vals)
+        and all(c == p + 1 for c, p in zip(cur_vals, prev_vals))
+    )
+
+
 def page_hash(
     token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
 ) -> int:
@@ -350,13 +373,10 @@ class PagePool:
         self._task_pages: Dict[str, List[int]] = {}
         self._lock = threading.Lock()
 
-        # Steady-state decode validation state: the ordered task set and its
+        # Steady-state decode validation: when the same ordered task set
-        # Python seq_lens mirror.  When the same set advances every sequence
+        # advances every sequence by exactly one token per step, bind_tasks
-        # by exactly one token per step, bind_tasks updates the stable
+        # updates the stable buffers in-place instead of re-cumsumming.
-        # buffers in-place (+=1 / +=inc) instead of re-cumsumming.  Any
+        self._bind_state: Optional[_BindState] = None
-        # task-set change is a miss and rebuilds.
-        self._bind_sig: Optional[tuple] = None
-        self._bind_seq_lens: Optional[List[int]] = None
 
     # ---- task lifecycle ----
 
@@ -529,13 +549,11 @@ class PagePool:
         inc_buf = workspace.inc
         ocl_buf = workspace.out_cache_loc
 
+        prev = self._bind_state
         incremental = (
             start_pos is None
-            and self._bind_sig is not None
+            and prev is not None
-            and self._bind_sig == sig
+            and _is_steady_increment(prev.sig, prev.seq_lens, sig, seq_lens)
-            and self._bind_seq_lens is not None
-            and len(self._bind_seq_lens) == b
-            and all(s == p + 1 for s, p in zip(seq_lens, self._bind_seq_lens))
         )
         if incremental:
             # Steady-state decode: advance the stable buffers in-place.
@@ -557,8 +575,7 @@ class PagePool:
             req_pool_indices = rpi_buf[:b]
             seq_lens_t = sl_buf[:b]
             kv_indptr = kvp_buf[: b + 1]
-        self._bind_sig = sig
+        self._bind_state = _BindState(sig, list(seq_lens))
-        self._bind_seq_lens = list(seq_lens)
 
         if start_pos is not None:
             seq_len = seq_lens[0]

astrai/inference/core/executor.py

@@ -14,7 +14,7 @@ from astrai.extension.attention_backend import (
     attn_backend,
     get_backend,
 )
-from astrai.inference.core.cache import PagePool
+from astrai.inference.core.cache import PagePool, _is_steady_increment
 from astrai.inference.core.graph import CudaGraphContext
 from astrai.inference.core.task import Task
 from astrai.inference.core.workspace import InferenceWorkspace
@@ -54,6 +54,19 @@ class SamplingBatchInfo:
     has_freq: bool  # any frequency_penalty != 0 (avoids per-step GPU .any())
 
 
+@dataclass
+class DecodeSteadyState:
+    """Cached decode metadata for the steady-state case.
+
+    When the same ordered task set decodes one token per step, sampling
+    params and task signature are reused; only positions advance by 1.
+    """
+
+    task_sig: tuple
+    positions: list[int]
+    sampling_info: SamplingBatchInfo
+
+
 def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
     pin = str(device).startswith("cuda")
     freq_penalties = torch.tensor(
@@ -165,11 +178,10 @@ class Executor:
         self.device = device or next(model.parameters()).device
         self.dtype = dtype or next(model.parameters()).dtype
 
-        # Per-step decode cache for the steady-state case where the same
+        # Per-step decode cache for the steady-state case (same ordered
-        # ordered task set decodes one token per step.  Sampling params are
+        # task set decodes one token per step).  Sampling params stay
-        # constant across steps; position_ids grows by exactly 1.  Single-slot:
+        # constant; only positions advance.
-        # any task-set change is a cache miss.
+        self._decode_cache: Optional[DecodeSteadyState] = None
-        self._decode_cache: Optional[tuple] = None
 
         # Pre-allocated fixed-shape buffers for the decode hot path
         # (input_ids, decode mask, KV bind metadata).  Eagerly sized at init
@@ -340,20 +352,18 @@ class Executor:
 
         sig = tuple(task_ids)
         cached = self._decode_cache
-        if (
+        prev_sig = cached.task_sig if cached is not None else None
-            cached is not None
+        prev_pos = cached.positions if cached is not None else None
-            and cached[0] == sig
+        if _is_steady_increment(prev_sig, prev_pos, sig, cur_positions):
-            and cur_positions == [p + 1 for p in cached[1]]
+            info = cached.sampling_info
-        ):
-            info = cached[2]
             ws.position_ids[:b] += 1
-            self._decode_cache = (sig, cur_positions, info)
+            self._decode_cache = DecodeSteadyState(sig, cur_positions, info)
         else:
             info = _build_sampling_batch_info(tasks, self.device)
             ws.position_ids[:b].copy_(
                 torch.tensor(cur_positions, dtype=torch.long, device=self.device)
             )
-            self._decode_cache = (sig, cur_positions, info)
+            self._decode_cache = DecodeSteadyState(sig, cur_positions, info)
 
         total_len = max(cur_positions) + 1
         input_mask = ws.decode_mask(ws.position_ids[:b], total_len)