refactor: centralize batch state and split scheduler duties

- make decode steady state self-validating: gate the token fill and sampling reuse on the decode cache's own task signature, drop TaskCacheManager.last_task_signature_matches whose req-index signature recycles with slot reuse
- extract PolicyVersionGuard (version protocol plus generation/weight mutex) and Stepper (shared one-token advancement) out of the scheduler, keeping its public API unchanged
- pin the input staging buffer only on CUDA devices so CPU-only workspaces allocate

astrai/inference/runtime/executor.py

@@ -415,14 +415,14 @@ class Executor:
         # inference-mode context.
         #
         # ``cache_valid`` checks the decode cache's own task signature:
-        # req-index signatures in the cache manager are recycled when freed
-        # slots are reallocated to new tasks, so a fresh batch whose prefill
-        # re-bind coincides with a stale signature would otherwise replay a
-        # previous generation's tokens into ``input_ids``.
-        task_sig_match = self.task_cache.last_task_signature_matches(task_ids)
+        # task ids are globally unique, so equality alone proves the cached
+        # tokens were sampled for exactly this ordered batch.  Req-index
+        # signatures in the cache manager are deliberately NOT consulted —
+        # they are recycled when freed slots are reallocated, which once
+        # let a fresh batch replay a previous generation's tokens.
         cached = self._decode_cache
         cache_valid = cached is not None and cached.task_sig == task_sig
-        if task_sig_match and cache_valid and cached.last_tokens is not None:
+        if cache_valid and cached.last_tokens is not None:
             with torch.inference_mode():
                 input_ids = ws.fill_input_ids_from_device(cached.last_tokens)
         else:
@@ -433,12 +433,9 @@ class Executor:
         kv_cache = self.task_cache.bind(task_ids, ws)
 
         # Reuse sampling state only if all conditions hold:
-        # 1. KV bind detected steady increment (same req_indices, seq_lens +1)
-        # 2. Task signature matches (same task_ids in same order)
-        # 3. We have a valid cached decode state for THIS task set
-        reuse_decode_state = (
-            cache_valid and self.task_cache.bind_was_steady and task_sig_match
-        )
+        # 1. The cached decode state belongs to THIS task set (task_sig)
+        # 2. KV bind detected steady increment (same req_indices, seq_lens +1)
+        reuse_decode_state = cache_valid and self.task_cache.bind_was_steady
         if reuse_decode_state:
             info = cached.sampling_info
             ws.position_ids[:b] += 1

astrai/inference/runtime/stepper.py

@@ -0,0 +1,127 @@
+"""One-token advancement primitive shared by every scheduling mode."""
+
+from contextlib import nullcontext
+from typing import Dict, List, Optional, Tuple
+
+from astrai.extension import AttentionBackend, attn_backend
+from astrai.inference.cache import PagePool, TaskCacheManager
+from astrai.inference.metrics import MetricsCollector
+from astrai.inference.runtime.executor import Executor
+from astrai.inference.task import Task, TaskStatus
+
+
+class Stepper:
+    """Advance every active task by one token (prefill + decode).
+
+    Single shared primitive for both the continuous-batching loop and the
+    synchronous ``run_batch`` path, so the two cannot drift.
+
+    Tasks must already be allocated in the KV cache. Tasks without output
+    are prefilled first and sample their first token from the final prompt
+    position. Tasks with output extend the cache by one position and decode
+    from their latest generated token.
+    """
+
+    def __init__(
+        self,
+        pool: PagePool,
+        task_cache: TaskCacheManager,
+        executor: Executor,
+        metrics: MetricsCollector,
+    ):
+        self._pool = pool
+        self._task_cache = task_cache
+        self._executor = executor
+        self._metrics = metrics
+
+    @staticmethod
+    def _task_backend_groups(tasks: List[Task]):
+        groups = {}
+        for task in tasks:
+            groups.setdefault(task.backend, (task.backend, []))[1].append(task)
+        return groups.values()
+
+    def step(
+        self, tasks: List[Task], return_logprobs: bool = False
+    ) -> Tuple[List[Task], List[Task]]:
+        """Advance ``tasks`` by one token.
+
+        Args:
+            tasks: Active tasks to advance by one token.
+            return_logprobs: Forwarded to the executor; per-token logprobs
+                are recorded on each task's ``output_logprobs``.
+
+        Returns:
+            ``(decoded, aborted)``: tasks that produced a new token (its ID
+            already appended to ``output_ids``) and tasks that hit the
+            sequence cap and were marked ``ABORTED``.
+        """
+        to_prefill = [t for t in tasks if not t.prefill_done and t.prompt_ids]
+        prefilled_ids = set()
+        produced: List[Task] = []
+        if to_prefill:
+            for t in to_prefill:
+                t.input_tokens = len(t.prompt_ids)
+
+            groups: Dict[Tuple[int, Optional[AttentionBackend]], List[Task]] = {}
+            for t in to_prefill:
+                start_pos = min(
+                    self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1
+                )
+                groups.setdefault((start_pos, t.backend), []).append(t)
+
+            for (start_pos, _), group in groups.items():
+                backend = group[0].backend
+                backend_context = (
+                    attn_backend(backend) if backend is not None else nullcontext()
+                )
+                with (
+                    backend_context,
+                    self._metrics.record([t.task_id for t in group], "prefill"),
+                ):
+                    prefilled, step_out = self._executor.execute_prefill(
+                        group, start_pos=start_pos, return_logprobs=return_logprobs
+                    )
+
+                for t, out in zip(prefilled, step_out):
+                    t.output_ids.append(out[0] if return_logprobs else out)
+                    t.output_tokens += 1
+                    t.mark_prefill_done()
+                    prefilled_ids.add(t.task_id)
+                    produced.append(t)
+
+                start_logical_page = start_pos // self._pool.page_size
+                for t in group:
+                    self._task_cache.task_record_hashes(
+                        t.task_id, t.prompt_ids, start_logical_page
+                    )
+
+        decoded: List[Task] = []
+        aborted: List[Task] = []
+        for t in tasks:
+            if t.task_id in prefilled_ids:
+                continue
+            if self._task_cache.task_extend(t.task_id, t.next_pos):
+                decoded.append(t)
+            else:
+                t.status = TaskStatus.ABORTED
+                aborted.append(t)
+
+        for backend, group in self._task_backend_groups(decoded):
+            backend_context = (
+                attn_backend(backend) if backend is not None else nullcontext()
+            )
+            with (
+                backend_context,
+                self._metrics.record([t.task_id for t in group], "decode"),
+            ):
+                step_out = self._executor.execute_decode(
+                    group, return_logprobs=return_logprobs
+                )
+            for t, out in zip(group, step_out):
+                t.output_ids.append(out[0] if return_logprobs else out)
+                t.output_tokens += 1
+                t.advance_kv()
+                produced.append(t)
+
+        return produced, aborted