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/cache/pool.py

@@ -406,22 +406,6 @@ class TaskCacheManager:
         """True if the last bind was a steady-state increment (same tasks, +1 seq_lens)."""
         return self._bind_was_steady
 
-    def last_task_signature_matches(self, task_ids: List[str]) -> bool:
-        """Check if task_ids match the previous bind's signature.
-
-        Used by Executor to detect steady-state decode for device-to-device
-        token copy optimization.
-        """
-        if self._bind_state is None:
-            return False
-        prev_sig = self._bind_state.sig
-        # sig is tuple of req_indices, need to map task_ids to req_indices
-        try:
-            current_sig = tuple(self._states[tid].req_idx for tid in task_ids)
-            return prev_sig == current_sig
-        except KeyError:
-            return False
-
     # -- internals --
 
     def _rollback(self, state: TaskCacheState, task_id: str):

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
+        # task ids are globally unique, so equality alone proves the cached
-        # slots are reallocated to new tasks, so a fresh batch whose prefill
+        # tokens were sampled for exactly this ordered batch.  Req-index
-        # re-bind coincides with a stale signature would otherwise replay a
+        # signatures in the cache manager are deliberately NOT consulted —
-        # previous generation's tokens into ``input_ids``.
+        # they are recycled when freed slots are reallocated, which once
-        task_sig_match = self.task_cache.last_task_signature_matches(task_ids)
+        # 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)
+        # 1. The cached decode state belongs to THIS task set (task_sig)
-        # 2. Task signature matches (same task_ids in same order)
+        # 2. KV bind detected steady increment (same req_indices, seq_lens +1)
-        # 3. We have a valid cached decode state for THIS task set
+        reuse_decode_state = cache_valid and self.task_cache.bind_was_steady
-        reuse_decode_state = (
-            cache_valid and self.task_cache.bind_was_steady and task_sig_match
-        )
         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

astrai/inference/scheduler.py

@@ -16,6 +16,7 @@ from astrai.extension import (
 from astrai.inference.cache import PagePool, TaskCacheManager
 from astrai.inference.metrics import MetricsCollector
 from astrai.inference.runtime.executor import Executor
+from astrai.inference.runtime.stepper import Stepper
 from astrai.inference.task import (
     STOP,
     GenerationResult,
@@ -23,6 +24,7 @@ from astrai.inference.task import (
     TaskManager,
     TaskStatus,
 )
+from astrai.inference.versioning import PolicyVersionGuard
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
 
@@ -119,49 +121,32 @@ class InferenceScheduler:
                 enable_cuda_graph=enable_cuda_graph,
             )
 
+        self._stepper = Stepper(
+            self._cache, self._task_cache, self._executor, self._metrics
+        )
+
         self._stop_event = threading.Event()
         self._loop_thread: Optional[threading.Thread] = None
-        self._weight_lock = threading.RLock()
+        self._policy_guard = PolicyVersionGuard(
-        self._policy_version = policy_version
+            policy_version,
+            ensure_ready=self._ensure_weight_update_ready,
+            on_commit=self._task_cache.invalidate_cache,
+        )
+        # Synchronous generation shares the guard's generation/weight mutex.
+        self._weight_lock = self._policy_guard.lock
 
     @property
     def policy_version(self) -> int:
         """Version of the model weights used for subsequent generations."""
-        return self._policy_version
+        return self._policy_guard.policy_version
-
-    def _validate_weight_version(
-        self, policy_version: int, *, require_advance: bool = False
-    ) -> None:
-        if (
-            isinstance(policy_version, bool)
-            or not isinstance(policy_version, int)
-            or policy_version < 0
-        ):
-            raise ValueError("policy_version must be a non-negative integer")
-        if policy_version < self._policy_version:
-            raise ValueError(
-                f"policy_version cannot move backwards from "
-                f"{self._policy_version} to {policy_version}"
-            )
-        if require_advance and policy_version == self._policy_version:
-            raise ValueError(
-                f"policy_version must advance beyond {self._policy_version} "
-                "when model weights are mutated"
-            )
 
     def _ensure_weight_update_ready(self) -> None:
-        """Check weight update preconditions. Must be called under _weight_lock."""
+        """Check weight update preconditions. Must be called under the lock."""
         if self._loop_thread is not None and self._loop_thread.is_alive():
             raise RuntimeError("Stop the scheduler before updating model weights")
         if self._task_mgr.get_active_tasks() or self._task_mgr.get_waiting_tasks():
             raise RuntimeError("Cannot update model weights while tasks are queued")
 
-    def _commit_weight_version(self, policy_version: int) -> int:
-        self._task_cache.invalidate_cache()
-        self._policy_version = policy_version
-        return self._policy_version
-
-    @_with_weight_lock
     def update_weights(self, policy_version: int) -> int:
         """Acknowledge an in-place weight update and invalidate stale KV state.
 
@@ -170,13 +155,8 @@ class InferenceScheduler:
         version update makes that lifecycle visible and prevents prefix KV
         entries produced by older weights from being reused.
         """
-        self._validate_weight_version(policy_version)
+        return self._policy_guard.update_weights(policy_version)
-        if policy_version == self._policy_version:
-            return self._policy_version
-        self._ensure_weight_update_ready()
-        return self._commit_weight_version(policy_version)
 
-    @_with_weight_lock
     def apply_weight_update(
         self, policy_version: Optional[int], update: Callable[[], T]
     ) -> T:
@@ -186,24 +166,11 @@ class InferenceScheduler:
         callers that only need "advance by one" (e.g. ``optimizer.step()``)
         without a read-compute-write race on the current version.
         """
-        if not callable(update):
+        return self._policy_guard.apply_weight_update(policy_version, update)
-            raise TypeError("update must be callable")
-        if policy_version is None:
-            policy_version = self._policy_version + 1
-        else:
-            self._validate_weight_version(policy_version, require_advance=True)
-        self._ensure_weight_update_ready()
 
-        result = update()
-        self._commit_weight_version(policy_version)
-        return result
-
-    @_with_weight_lock
     def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T:
         """Inspect state while the scheduler's policy version remains stable."""
-        if not callable(inspect):
+        return self._policy_guard.with_policy_snapshot(inspect)
-            raise TypeError("inspect must be callable")
-        return inspect(self._policy_version)
 
     def add_task(self, prompt: str, **kwargs) -> str:
         return self._task_mgr.add_task(prompt, **kwargs)
@@ -226,7 +193,7 @@ class InferenceScheduler:
     def get_stats(self) -> Dict[str, Any]:
         stats = self._task_mgr.get_stats()
         stats["kv_cache_tasks"] = self._task_cache.task_count
-        stats["policy_version"] = self._policy_version
+        stats["policy_version"] = self._policy_guard.policy_version
         return stats
 
     @property
@@ -242,105 +209,11 @@ class InferenceScheduler:
             return nullcontext()
         return attn_backend(self._backend)
 
-    @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 every active task by one token (prefill + decode).
+        """Advance every active task by one token; see :class:`Stepper`."""
-
+        return self._stepper.step(tasks, return_logprobs=return_logprobs)
-        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.
-
-        Args:
-            tasks: Active tasks to advance by one token.
-            return_logprobs: Forwarded to ``execute_decode``; 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._cache.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
 
     def _run_generation_loop(self):
         stop_ids = self._task_mgr.tokenizer.stop_ids
@@ -389,7 +262,7 @@ class InferenceScheduler:
                         if task.status != TaskStatus.ABORTED
                     ]
 
-                    decoded, aborted = self._step(active)
+                    decoded, aborted = self._stepper.step(active)
 
                     for t in aborted:
                         self._task_mgr.invoke_callback(t.task_id, STOP)
@@ -533,7 +406,9 @@ class InferenceScheduler:
 
             with self._backend_context():
                 while live:
-                    decoded, aborted = self._step(live, return_logprobs=return_logprobs)
+                    decoded, aborted = self._stepper.step(
+                        live, return_logprobs=return_logprobs
+                    )
                     for task in aborted:
                         runtime_errors[task.task_id] = "kv_cache_extension_failed"
                     live = [t for t in decoded if not t.is_finished(stop_ids)]

astrai/inference/versioning.py

@@ -0,0 +1,129 @@
+"""Policy version protocol for weights shared between trainer and server.
+
+The guard owns the monotonic version counter, the RLock that serializes
+weight publication against generation (``run_batch`` acquires the same
+lock), and the validation/commit rules around both.  Scheduler-specific
+preconditions (no in-flight generation, no queued tasks) and side effects
+(dropping stale KV entries) are injected as callables so the guard stays
+free of inference-subsystem knowledge.
+"""
+
+import threading
+from functools import wraps
+from typing import Callable, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+def _locked(method):
+    @wraps(method)
+    def synchronized(self, *args, **kwargs):
+        with self._lock:
+            return method(self, *args, **kwargs)
+
+    return synchronized
+
+
+class PolicyVersionGuard:
+    """Monotonic policy-version protocol over shared in-place weights.
+
+    The scheduler and the in-process trainer mutate the same model object;
+    every weight mutation must publish a new version atomically under the
+    generation lock.  Versions never move backwards and every mutation
+    advances (``apply_weight_update``) or repeats the live version
+    (``update_weights`` idempotently).
+    """
+
+    def __init__(
+        self,
+        policy_version: int,
+        ensure_ready: Callable[[], None],
+        on_commit: Callable[[], None],
+    ):
+        if (
+            isinstance(policy_version, bool)
+            or not isinstance(policy_version, int)
+            or policy_version < 0
+        ):
+            raise ValueError("policy_version must be a non-negative integer")
+        self._lock = threading.RLock()
+        self._policy_version = policy_version
+        self._ensure_ready = ensure_ready
+        self._on_commit = on_commit
+
+    @property
+    def lock(self) -> threading.RLock:
+        """Generation/weight mutex; synchronous generation acquires it too."""
+        return self._lock
+
+    @property
+    def policy_version(self) -> int:
+        """Version of the model weights used for subsequent generations."""
+        return self._policy_version
+
+    def _validate(self, policy_version: int, *, require_advance: bool = False) -> None:
+        if (
+            isinstance(policy_version, bool)
+            or not isinstance(policy_version, int)
+            or policy_version < 0
+        ):
+            raise ValueError("policy_version must be a non-negative integer")
+        if policy_version < self._policy_version:
+            raise ValueError(
+                f"policy_version cannot move backwards from "
+                f"{self._policy_version} to {policy_version}"
+            )
+        if require_advance and policy_version == self._policy_version:
+            raise ValueError(
+                f"policy_version must advance beyond {self._policy_version} "
+                "when model weights are mutated"
+            )
+
+    @_locked
+    def update_weights(self, policy_version: int) -> int:
+        """Acknowledge an in-place weight update and invalidate stale KV state.
+
+        The scheduler owns the same model object as the in-process trainer,
+        so weights have already changed when this method is called. The
+        explicit version update makes that lifecycle visible and prevents
+        prefix KV entries produced by older weights from being reused.
+        """
+        self._validate(policy_version)
+        if policy_version == self._policy_version:
+            return self._policy_version
+        self._ensure_ready()
+        return self._commit(policy_version)
+
+    @_locked
+    def apply_weight_update(
+        self, policy_version: Optional[int], update: Callable[[], T]
+    ) -> T:
+        """Mutate shared weights and publish their version without generation.
+
+        ``policy_version=None`` derives ``live + 1`` under the same lock, for
+        callers that only need "advance by one" (e.g. ``optimizer.step()``)
+        without a read-compute-write race on the current version.
+        """
+        if not callable(update):
+            raise TypeError("update must be callable")
+        if policy_version is None:
+            policy_version = self._policy_version + 1
+        else:
+            self._validate(policy_version, require_advance=True)
+        self._ensure_ready()
+
+        result = update()
+        self._commit(policy_version)
+        return result
+
+    @_locked
+    def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T:
+        """Inspect state while the policy version remains stable."""
+        if not callable(inspect):
+            raise TypeError("inspect must be callable")
+        return inspect(self._policy_version)
+
+    def _commit(self, policy_version: int) -> int:
+        self._on_commit()
+        self._policy_version = policy_version
+        return self._policy_version

astrai/inference/workspace.py

@@ -80,7 +80,10 @@ class InferenceWorkspace:
                 (max_batch_size,), dtype=torch.long, device=device
             )
             self._pin = torch.empty(
-                (max_batch_size,), dtype=torch.long, pin_memory=True
+                (max_batch_size,),
+                dtype=torch.long,
+                pin_memory=torch.cuda.is_available()
+                and torch.device(device).type == "cuda",
             )
 
             # KV-cache bind metadata (fixed shape, written by

tests/inference/test_scheduler.py

@@ -12,6 +12,7 @@ from astrai.extension import CudaBackend, TorchNativeBackend, get_backend
 from astrai.inference import GenerationResult, InferenceScheduler
 from astrai.inference.metrics import MetricsCollector
 from astrai.inference.runtime.executor import DecodeSteadyState, Executor
+from astrai.inference.runtime.stepper import Stepper
 from astrai.inference.task import Task
 from astrai.model.transformer import AutoRegressiveLM
 from tests.helpers import FakeTokenizer, make_rollout_config
@@ -119,10 +120,14 @@ def test_generation_loop_activates_backend_in_worker_thread():
 
 def test_step_splits_decode_batch_by_request_backend():
     scheduler = object.__new__(InferenceScheduler)
+    scheduler._cache = SimpleNamespace(page_size=1)
     scheduler._task_cache = MagicMock()
     scheduler._task_cache.task_extend.return_value = True
     scheduler._metrics = MetricsCollector()
     scheduler._executor = MagicMock()
+    scheduler._stepper = Stepper(
+        scheduler._cache, scheduler._task_cache, scheduler._executor, scheduler._metrics
+    )
 
     observed = []
 
@@ -157,6 +162,9 @@ def test_step_batches_ragged_prefill_with_shared_cache_start():
     scheduler._task_cache.task_cached.return_value = 0
     scheduler._metrics = MetricsCollector()
     scheduler._executor = MagicMock()
+    scheduler._stepper = Stepper(
+        scheduler._cache, scheduler._task_cache, scheduler._executor, scheduler._metrics
+    )
 
     short = Task("short", [1, 2, 3])
     long = Task("long", [4, 5, 6, 7, 8])
@@ -701,8 +709,8 @@ def test_run_batch_details_report_extension_failure_and_cleanup(device):
     scheduler, _tok, _model = _make_real_scheduler(device)
     try:
         with patch.object(
-            scheduler,
+            scheduler._stepper,
-            "_step",
+            "step",
             side_effect=lambda tasks, **_kwargs: ([], list(tasks)),
         ):
             result = scheduler.run_batch([[10, 20]], max_tokens=2, return_details=True)[
@@ -722,9 +730,6 @@ def test_decode_does_not_reuse_previous_batch_state():
     executor.device = torch.device("cpu")
     executor.task_cache = MagicMock()
     executor.task_cache.bind_was_steady = True
-    executor.task_cache.last_task_signature_matches.return_value = (
-        False  # Different task
-    )
     executor.task_cache.bind.return_value = MagicMock()
     executor._graph_supported = False
     executor._graph_ctx = SimpleNamespace(enabled=False)
@@ -771,7 +776,6 @@ def test_decode_fills_input_ids_from_device_on_matching_signature():
     executor.device = torch.device("cpu")
     executor.task_cache = MagicMock()
     executor.task_cache.bind_was_steady = True
-    executor.task_cache.last_task_signature_matches.return_value = True  # Same task
     executor.task_cache.bind.return_value = MagicMock()
     executor._graph_supported = False
     executor._graph_ctx = SimpleNamespace(enabled=False)