refactor: harden inference cache state and attention dispatch

- split KVCache into phase-specific PrefillKVCache/DecodeKVCache types selected by start_pos
- unify steady-state detection in TaskCacheManager
- guard decode steady-state reuse with the cached task signature so recycled req slots cannot replay a prior generation's tokens and positions
- collapse attention backend fwd_decode/fwd_prefill into a single subclass-owned forward with a shared _check_fwd guard
- fix thread-safety gap in weight update and validate prefill inputs before KV allocation
- centralize magic constants in InferenceConfig and align docs with behavior

astrai/config/inference_config.py

@@ -0,0 +1,23 @@
+"""Inference engine configuration."""
+
+from astrai.config.base import BaseConfig
+
+
+class InferenceConfig(BaseConfig):
+    """Configuration for inference workspace and execution parameters.
+
+    Centralizes magic constants previously scattered across inference modules.
+
+    Args:
+        max_splits (int): Maximum number of splits for split-KV attention (decode partial results). Defaults to 32.
+        q_tile_rows (int): Number of rows per Q tile in prefill ragged batching. Defaults to 64.
+        prefill_warmup_len (int): Prompt length for prefill warmup (cuBLAS auto-tuning). Defaults to 64.
+        default_rep_window (int): Default repetition penalty window size for frequency penalty. Defaults to 64.
+        max_recent_tasks (int): Maximum number of recent tasks tracked for aggregate statistics. Defaults to 128.
+    """
+
+    max_splits: int = 32
+    q_tile_rows: int = 64
+    prefill_warmup_len: int = 64
+    default_rep_window: int = 64
+    max_recent_tasks: int = 128

astrai/extension/backend/attention.py

@@ -343,9 +343,11 @@ def attention(
 class AttentionBackend(ABC):
     """Abstract base for attention computation strategies.
 
-    Subclasses implement ``fwd_decode`` (q_len == 1, with cache) and
+    Subclasses implement a single ``forward`` and branch on ``fwd``
-    ``fwd_prefill`` (q_len > 1, with or without cache). The public
+    ("decode" / "prefill", or None for training) wherever their kernels
-    ``forward`` method dispatches based on q_len.
+    split — the mode taxonomy is the caller's, not the base class's, so
+    it lives in the implementations. ``_check_fwd`` is the shared guard
+    against unknown mode strings.
 
     Capability contract — every backend declares:
 
@@ -365,7 +367,6 @@ class AttentionBackend(ABC):
         with attn_backend(TorchNativeBackend):          # class
             ...
         with TorchNativeBackend():                       # instance
-            ...
     """
 
     def __enter__(self) -> "AttentionBackend":
@@ -398,7 +399,15 @@ class AttentionBackend(ABC):
         Called on the canonical singleton instance (or a caller-provided
         one); must be side-effect free.
         """
+        return True
 
+    @staticmethod
+    def _check_fwd(fwd: Optional[str]) -> None:
+        """Reject unknown forward modes loudly."""
+        if fwd not in (None, "prefill", "decode"):
+            raise ValueError(f"unsupported attention forward mode: {fwd}")
+
+    @abstractmethod
     def forward(
         self,
         q: Tensor,
@@ -410,7 +419,7 @@ class AttentionBackend(ABC):
         is_causal: bool = False,
         fwd: Optional[str] = None,
     ) -> Tensor:
-        """Dispatch to decode or extend based on q_len.
+        """Run one attention call; ``fwd`` selects the mode.
 
         Args:
             q: [batch, q_len, n_heads, head_dim]
@@ -420,41 +429,11 @@ class AttentionBackend(ABC):
             layer_id: transformer layer index for buffer access.
             attn_mask: pre-built attention mask compatible with SDPA.
             is_causal: whether to apply causal masking.
+            fwd: "prefill" / "decode" for inference, None for training.
 
         Returns:
             [batch, q_len, n_heads * head_dim]
         """
-        if fwd == "decode":
-            return self.fwd_decode(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
-        if fwd == "prefill" or fwd is None:
-            return self.fwd_prefill(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
-        raise ValueError(f"unsupported attention forward mode: {fwd}")
-
-    @abstractmethod
-    def fwd_decode(
-        self,
-        q: Tensor,
-        k: Tensor,
-        v: Tensor,
-        kv_cache: Optional["KVCache"],
-        layer_id: int,
-        attn_mask: Optional[Tensor] = None,
-        is_causal: bool = False,
-    ) -> Tensor:
-        """Single-token decode with KV cache."""
-
-    @abstractmethod
-    def fwd_prefill(
-        self,
-        q: Tensor,
-        k: Tensor,
-        v: Tensor,
-        kv_cache: Optional["KVCache"],
-        layer_id: int,
-        attn_mask: Optional[Tensor] = None,
-        is_causal: bool = False,
-    ) -> Tensor:
-        """Multi-token prefill or training forward."""
 
     @staticmethod
     def supports_graph() -> bool:
@@ -498,31 +477,7 @@ class TorchNativeBackend(AttentionBackend):
     ) -> bool:
         return True
 
-    def fwd_decode(
+    def forward(
-        self,
-        q: Tensor,
-        k: Tensor,
-        v: Tensor,
-        kv_cache: Optional["KVCache"],
-        layer_id: int,
-        attn_mask: Optional[Tensor] = None,
-        is_causal: bool = False,
-    ) -> Tensor:
-        return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
-
-    def fwd_prefill(
-        self,
-        q: Tensor,
-        k: Tensor,
-        v: Tensor,
-        kv_cache: Optional["KVCache"],
-        layer_id: int,
-        attn_mask: Optional[Tensor] = None,
-        is_causal: bool = False,
-    ) -> Tensor:
-        return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
-
-    def _forward(
         self,
         q: Tensor,
         k: Tensor,
@@ -531,7 +486,9 @@ class TorchNativeBackend(AttentionBackend):
         layer_id: int,
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> Tensor:
+        self._check_fwd(fwd)
         if q.ndim == 4:
             n_rep = q.size(2) // k.size(2)
             if n_rep > 1:
@@ -633,7 +590,7 @@ class CudaBackend(AttentionBackend):
     def supports_graph() -> bool:
         return True
 
-    def fwd_decode(
+    def forward(
         self,
         q: Tensor,
         k: Tensor,
@@ -642,10 +599,23 @@ class CudaBackend(AttentionBackend):
         layer_id: int,
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> Tensor:
+        self._check_fwd(fwd)
         if kv_cache is None:
             raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
+        if fwd == "decode":
+            return self._decode(q, k, v, kv_cache, layer_id)
+        return self._prefill(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
 
+    def _decode(
+        self,
+        q: Tensor,
+        k: Tensor,
+        v: Tensor,
+        kv_cache: "KVCache",
+        layer_id: int,
+    ) -> Tensor:
         kv_indptr = kv_cache.kv_indptr
 
         out = attn_paged_decode(
@@ -664,19 +634,16 @@ class CudaBackend(AttentionBackend):
         )
         return out
 
-    def fwd_prefill(
+    def _prefill(
         self,
         q: Tensor,
         k: Tensor,
         v: Tensor,
-        kv_cache: Optional["KVCache"],
+        kv_cache: "KVCache",
         layer_id: int,
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
     ) -> Tensor:
-        if kv_cache is None:
-            raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
-
         loc = kv_cache.out_cache_loc
         kv_cache.k_buffer[layer_id, loc] = k
         kv_cache.v_buffer[layer_id, loc] = v
@@ -734,7 +701,7 @@ class FlashAttnBackend(AttentionBackend):
         # back to TorchNativeBackend instead of silently ignoring the mask.
         return attn_mask is None
 
-    def fwd_decode(
+    def forward(
         self,
         q: Tensor,
         k: Tensor,
@@ -743,20 +710,11 @@ class FlashAttnBackend(AttentionBackend):
         layer_id: int,
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> Tensor:
-        return self._forward_packed(q, k, v, kv_cache, layer_id)
+        self._check_fwd(fwd)
-
+        # Decode is always packed; prefill/training split by layout.
-    def fwd_prefill(
+        if fwd == "decode" or q.ndim == 3:
-        self,
-        q: Tensor,
-        k: Tensor,
-        v: Tensor,
-        kv_cache: Optional["KVCache"],
-        layer_id: int,
-        attn_mask: Optional[Tensor] = None,
-        is_causal: bool = False,
-    ) -> Tensor:
-        if q.ndim == 3:
             return self._forward_packed(q, k, v, kv_cache, layer_id)
         return self._forward_dense(q, k, v, attn_mask, is_causal)
 

astrai/inference/cache/buffer.py

@@ -2,7 +2,10 @@
 
 Layer 1 — ``KVStorage``:   flat token-level K/V GPU buffers [n_layers, size, n_kv_heads, head_dim]
 Layer 2 — ``ReqToTokenPool``: index table [req_idx, pos] → physical token slot
-Layer 3 — ``KVCache``:       pure dataclass passed to the model for direct buffer access
+Layer 3 — ``BaseKVCache``:    shared fields for all cache modes
+         ``PrefillKVCache``:  prefill-specific layout
+         ``DecodeKVCache``:   decode-specific layout
+         ``KVCache``:         union type for backward compatibility
 
 These classes have no knowledge of tasks, allocation policies, or scheduling.
 They are the "dumb" physical storage layer.
@@ -10,7 +13,7 @@ They are the "dumb" physical storage layer.
 
 import threading
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import torch
 from torch import Tensor
@@ -74,8 +77,8 @@ class KVStorage:
 
 
 @dataclass
-class KVCache:
+class BaseKVCache:
-    """Pure data struct passed to model for KV cache I/O.
+    """Shared fields for all KV cache modes.
 
     The attention layer does raw buffer indexing — no methods, no abstraction.
     """
@@ -85,12 +88,36 @@ class KVCache:
     req_to_token: Tensor
     req_pool_indices: Tensor
     seq_lens: Tensor
-    out_cache_loc: Tensor
+    max_len: int
-    max_len: int = 0
+    kv_indptr: Tensor  # Always present in both modes
-    kv_indptr: Optional[Tensor] = None
+
-    qo_indptr: Optional[Tensor] = None
+
-    q_tile_to_batch: Optional[Tensor] = None
+@dataclass
-    q_tile_to_index: Optional[Tensor] = None
+class PrefillKVCache(BaseKVCache):
-    decode_o_part: Optional[Tensor] = None
+    """Prefill-specific KV cache layout.
-    decode_ml_part: Optional[Tensor] = None
+
-    decode_out: Optional[Tensor] = None
+    Handles packed ragged batching where prompts have variable lengths.
+    """
+
+    out_cache_loc: Tensor  # [total_q_tokens] - flattened write locations
+    qo_indptr: Tensor  # [B+1] - prefix sum of q_lens for unpacking
+    q_tile_to_batch: Tensor  # [num_q_tiles] - maps Q tiles to batch indices
+    q_tile_to_index: Tensor  # [num_q_tiles] - maps Q tiles to local indices
+
+
+@dataclass
+class DecodeKVCache(BaseKVCache):
+    """Decode-specific KV cache layout.
+
+    Single-token incremental generation with split-KV partial results.
+    """
+
+    out_cache_loc: Tensor  # [B] - one write position per request
+    qo_indptr: Tensor  # [B+1] - sequential [0, 1, 2, ..., B]
+    decode_o_part: Tensor  # [B, max_q_heads, MAX_SPLITS, head_dim]
+    decode_ml_part: Tensor  # [B, max_q_heads, MAX_SPLITS, 2]
+    decode_out: Tensor  # [B, max_q_heads, head_dim]
+
+
+# Backward compatibility: union type for existing code
+KVCache = Union[PrefillKVCache, DecodeKVCache]

astrai/inference/cache/pool.py

@@ -16,7 +16,13 @@ from typing import Dict, List, Optional
 
 import torch
 
-from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
+from astrai.inference.cache.buffer import (
+    DecodeKVCache,
+    KVCache,
+    KVStorage,
+    PrefillKVCache,
+    ReqToTokenPool,
+)
 from astrai.inference.cache.strategy import (
     AllocationStrategy,
     Allocator,
@@ -31,6 +37,8 @@ from astrai.inference.workspace import Q_TILE_ROWS, InferenceWorkspace
 # continues to work unchanged after the file split.
 __all__ = [
     "KVCache",
+    "PrefillKVCache",
+    "DecodeKVCache",
     "KVStorage",
     "ReqToTokenPool",
     "Allocator",
@@ -232,7 +240,20 @@ class PagePool:
             )
             q_tile_to_batch = workspace.q_tile_to_batch[:n_tiles]
             q_tile_to_index = workspace.q_tile_to_index[:n_tiles]
-            decode_o_part = decode_ml_part = decode_out = None
+
+            return PrefillKVCache(
+                k_buffer=self._storage.k_buffer,
+                v_buffer=self._storage.v_buffer,
+                req_to_token=self._req_pool.req_to_token,
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens_t,
+                max_len=max(seq_lens),
+                kv_indptr=kv_indptr,
+                out_cache_loc=out_cache_loc,
+                qo_indptr=qo_indptr,
+                q_tile_to_batch=q_tile_to_batch,
+                q_tile_to_index=q_tile_to_index,
+            )
         else:
             # ---- decode: out_cache_loc is a single column (last position) ----
             write_pos = seq_lens_t - 1
@@ -241,23 +262,20 @@ class PagePool:
             out_cache_loc = ocl_buf[:b].reshape(-1)
             workspace.qo_indptr[: b + 1].copy_(inc_buf[: b + 1])
             qo_indptr = workspace.qo_indptr[: b + 1]
-            q_tile_to_batch = q_tile_to_index = None
             decode_o_part = getattr(workspace, "decode_o_part", None)
             decode_ml_part = getattr(workspace, "decode_ml_part", None)
             decode_out = getattr(workspace, "decode_out", None)
 
-        return KVCache(
+            return DecodeKVCache(
                 k_buffer=self._storage.k_buffer,
                 v_buffer=self._storage.v_buffer,
                 req_to_token=self._req_pool.req_to_token,
                 req_pool_indices=req_pool_indices,
                 seq_lens=seq_lens_t,
-            out_cache_loc=out_cache_loc,
                 max_len=max(seq_lens),
                 kv_indptr=kv_indptr,
+                out_cache_loc=out_cache_loc,
                 qo_indptr=qo_indptr,
-            q_tile_to_batch=q_tile_to_batch,
-            q_tile_to_index=q_tile_to_index,
                 decode_o_part=decode_o_part,
                 decode_ml_part=decode_ml_part,
                 decode_out=decode_out,
@@ -385,8 +403,25 @@ class TaskCacheManager:
 
     @property
     def bind_was_steady(self) -> bool:
+        """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/metrics.py

@@ -7,6 +7,10 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Any, Deque, Dict, Generator, List, Literal, Optional
 
+from astrai.config.inference_config import InferenceConfig
+
+_config = InferenceConfig()
+
 
 @dataclass
 class TaskTiming:
@@ -123,7 +127,7 @@ class MetricsCollector:
         stats = metrics.get_stats()
     """
 
-    def __init__(self, max_recent: int = 128):
+    def __init__(self, max_recent: int = _config.max_recent_tasks):
         self._timings: Dict[str, TaskTiming] = {}
         self._completed: Deque[TaskTiming] = deque(maxlen=max_recent)
         self._lock = threading.Lock()

astrai/inference/runtime/executor.py

@@ -7,6 +7,7 @@ from typing import List, Optional
 import torch
 from torch import Tensor
 
+from astrai.config.inference_config import InferenceConfig
 from astrai.extension.backend.attention import (
     CudaBackend,
     get_backend,
@@ -19,6 +20,7 @@ from astrai.inference.workspace import InferenceWorkspace
 from astrai.model.automodel import AutoModel
 
 logger = logging.getLogger(__name__)
+_config = InferenceConfig()
 
 
 @contextmanager
@@ -114,7 +116,7 @@ def _warmup_cuda_graphs(
     # shapes on first call (F.linear is the dominant cost).  This also warms
     # up the CUDA context (driver init) and compiles the graph-capture trace
     # that follows.  Custom .so kernels do NOT need this — they are pre-built.
-    warmup_len = 64
+    warmup_len = _config.prefill_warmup_len
     tid = "_warmup_prefill"
     if task_cache.task_alloc(tid, list(range(warmup_len))):
         with (
@@ -312,9 +314,20 @@ class Executor:
     ):
         tasks = sorted(tasks, key=lambda t: t.task_id)
         batch_sz = len(tasks)
+
+        # Validate batch size bounds
+        if batch_sz > self._workspace.max_batch_size:
+            raise ValueError(
+                f"Batch size {batch_sz} exceeds max_batch_size "
+                f"{self._workspace.max_batch_size}"
+            )
+
         prompt_lens = [len(t.prompt_ids) for t in tasks]
+
+        # Validate inputs before any resource allocation
         if any(start_pos >= prompt_len for prompt_len in prompt_lens):
             raise ValueError("prefill start_pos must precede every prompt end")
+
         q_lens = [prompt_len - start_pos for prompt_len in prompt_lens]
 
         input_ids = torch.tensor(
@@ -380,10 +393,18 @@ class Executor:
             return []
 
         b = len(tasks)
+
+        # Validate batch size bounds
+        if b > self._workspace.max_batch_size:
+            raise ValueError(
+                f"Batch size {b} exceeds max_batch_size "
+                f"{self._workspace.max_batch_size}"
+            )
+
         ws = self._workspace
         task_ids = [t.task_id for t in tasks]
-        cur_positions = [t.next_pos for t in tasks]
         task_sig = tuple(task_ids)
+        cur_positions = [t.next_pos for t in tasks]
 
         # ---- pre-replay: update input buffers in-place ----
 
@@ -392,9 +413,16 @@ class Executor:
         # slots — fill input ids device-to-device.  inference_mode guards
         # the read because the source was produced under sampling's
         # 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)
         cached = self._decode_cache
-        sig_match = cached is not None and cached.task_sig == task_sig
+        cache_valid = cached is not None and cached.task_sig == task_sig
-        if sig_match and cached.last_tokens is not None:
+        if task_sig_match and 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:
@@ -404,9 +432,15 @@ class Executor:
 
         kv_cache = self.task_cache.bind(task_ids, ws)
 
-        reuse_decode_state = self.task_cache.bind_was_steady and sig_match
+        # 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
+        )
         if reuse_decode_state:
-            info = self._decode_cache.sampling_info
+            info = cached.sampling_info
             ws.position_ids[:b] += 1
         else:
             info = _build_sampling_batch_info(tasks, self.device)

astrai/inference/scheduler.py

@@ -150,6 +150,7 @@ class InferenceScheduler:
             )
 
     def _ensure_weight_update_ready(self) -> None:
+        """Check weight update preconditions. Must be called under _weight_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():

astrai/inference/task.py

@@ -18,6 +18,7 @@ from typing import (
 
 from tokenizers.decoders import DecodeStream
 
+from astrai.config.inference_config import InferenceConfig
 from astrai.inference.metrics import MetricsCollector
 from astrai.tokenize.tokenizer import AutoTokenizer
 
@@ -25,6 +26,7 @@ if TYPE_CHECKING:
     from astrai.extension import AttentionBackend
 
 STOP = object()
+_config = InferenceConfig()
 
 
 @dataclass(frozen=True)
@@ -85,7 +87,7 @@ class Task:
         top_p: float = 1.0,
         top_k: int = 50,
         frequency_penalty: float = 0.0,
-        rep_window: int = 64,
+        rep_window: int = _config.default_rep_window,
         backend: Optional["AttentionBackend"] = None,
     ):
         self.task_id = task_id

astrai/inference/workspace.py

@@ -9,8 +9,11 @@ the hot loop — a prerequisite for CUDA-graph capture.
 import torch
 from torch import Tensor
 
-_MAX_SPLITS = 32
+from astrai.config.inference_config import InferenceConfig
-Q_TILE_ROWS = 64
+
+_CONFIG = InferenceConfig()
+MAX_SPLITS = _CONFIG.max_splits
+Q_TILE_ROWS = _CONFIG.q_tile_rows
 
 
 class InferenceWorkspace:
@@ -22,8 +25,9 @@ class InferenceWorkspace:
     - ``decode_mask``: a ``[B, 1, total_len]`` validity mask, the RHS
       ``arange`` pre-computed so only a single ``torch.ge(out=)`` runs per
       step.
-    - ``input_ids``: per-step token IDs filled from host (pinned, double-
+    - ``input_ids``: per-step token IDs filled from host — values are
-      buffered so an in-flight async H2D copy never races the next fill).
+      staged through a pinned buffer and bulk-copied into the stable
+      device buffer (fixed address for CUDA-graph capture).
     - KV-cache bind metadata (``req_pool_indices``, ``seq_lens``,
       ``kv_indptr``, ``inc``, ``out_cache_loc``), written by
       ``PagePool.bind_tasks`` when the Executor passes this workspace.
@@ -71,17 +75,13 @@ class InferenceWorkspace:
 
             # Per-step token IDs.  Values come from host Python lists every
             # step, so the device buffer is pre-allocated (stable address for
-            # CUDA-graph capture) and filled via a host staging buffer.  A
+            # CUDA-graph capture) and filled via a host staging buffer.
-            # double buffer keeps a copy in flight from being overwritten by
-            # the next fill.
             self.input_ids = torch.empty(
                 (max_batch_size,), dtype=torch.long, device=device
             )
-            self._pin = [
+            self._pin = torch.empty(
-                torch.empty((max_batch_size,), dtype=torch.long),
+                (max_batch_size,), dtype=torch.long, pin_memory=True
-                torch.empty((max_batch_size,), dtype=torch.long),
+            )
-            ]
-            self._pin_idx = 0
 
             # KV-cache bind metadata (fixed shape, written by
             # ``PagePool.bind_tasks`` when the Executor passes this
@@ -124,15 +124,15 @@ class InferenceWorkspace:
             # Split-KV partial-result buffers for decode (persistent, one
             # global alloc per process — mirrors FlashInfer's workspace
             # pattern).  Shape:
-            # [max_batch_size, max_q_heads, _MAX_SPLITS, head_dim] (o_part)
+            # [max_batch_size, max_q_heads, MAX_SPLITS, head_dim] (o_part)
-            # [max_batch_size, max_q_heads, _MAX_SPLITS, 2]     (ml_part)
+            # [max_batch_size, max_q_heads, MAX_SPLITS, 2]     (ml_part)
             self.decode_o_part = torch.empty(
-                (max_batch_size, max_q_heads, _MAX_SPLITS, head_dim),
+                (max_batch_size, max_q_heads, MAX_SPLITS, head_dim),
                 dtype=torch.float32,
                 device=device,
             )
             self.decode_ml_part = torch.empty(
-                (max_batch_size, max_q_heads, _MAX_SPLITS, 2),
+                (max_batch_size, max_q_heads, MAX_SPLITS, 2),
                 dtype=torch.float32,
                 device=device,
             )
@@ -148,16 +148,13 @@ class InferenceWorkspace:
     def fill_input_ids(self, ids: "list[int]") -> Tensor:
         """Write ``ids`` into the device buffer and return ``[B]``.
 
-        Host values are staged through the double buffer and copied into the
+        Host values are staged through a pinned buffer and copied synchronously
-        stable device buffer (``copy_`` without pinning is synchronous, so
+        into the stable device buffer.
-        the alternating buffers guard against an in-flight transfer).
         """
         b = len(ids)
-        pin = self._pin[self._pin_idx]
-        self._pin_idx ^= 1
         for i, v in enumerate(ids):
-            pin[i] = v
+            self._pin[i] = v
-        self.input_ids[:b].copy_(pin[:b])
+        self.input_ids[:b].copy_(self._pin[:b])
         return self.input_ids[:b]
 
     def fill_input_ids_from_device(self, tokens: Tensor) -> Tensor:

docs/developer/cuda_kernels.md

@@ -417,7 +417,7 @@ cycle belong under `TYPE_CHECKING`.
 
 `astrai/extension/backend/attention.py` provides the backend abstraction:
 
-- **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
+- **`AttentionBackend`** (ABC): single abstract `forward`; each subclass branches on `fwd` ("decode" / "prefill" / None) internally, `_check_fwd` guards unknown modes
 - **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU.
 - **`FlashAttnBackend`**: Optional flash-attn dispatch via `flash_attn_varlen_func` over gathered flat K/V.
 - **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback)

tests/extension/test_backend.py

@@ -154,14 +154,18 @@ class _DummyBackend(AttentionBackend):
     def supports_call(self, q, kv_cache, attn_mask, is_causal, fwd) -> bool:
         return True
 
-    def fwd_decode(
+    def forward(
-        self, q, k, v, kv_cache=None, layer_id=0, attn_mask=None, is_causal=False
+        self,
-    ):
+        q,
-        return q
+        k,
-
+        v,
-    def fwd_prefill(
+        kv_cache=None,
-        self, q, k, v, kv_cache=None, layer_id=0, attn_mask=None, is_causal=False
+        layer_id=0,
+        attn_mask=None,
+        is_causal=False,
+        fwd=None,
     ):
+        self._check_fwd(fwd)
         return q
 
 

tests/inference/test_scheduler.py

@@ -185,6 +185,7 @@ def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits():
     executor.task_cache = MagicMock()
     executor.task_cache.bind.return_value = MagicMock()
     executor._workspace = MagicMock()
+    executor._workspace.max_batch_size = 16  # Add max_batch_size for validation
     all_logits = torch.arange(42, dtype=torch.float32).reshape(6, 7)
     executor.model = MagicMock(return_value={"logits": all_logits})
     executor._sample_logits = MagicMock(
@@ -721,11 +722,15 @@ 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)
 
     workspace = MagicMock()
+    workspace.max_batch_size = 16
     workspace.position_ids = torch.tensor([2], dtype=torch.long)
     workspace.fill_input_ids.return_value = torch.tensor([7], dtype=torch.long)
     workspace.decode_mask.return_value = torch.ones(1, 1, 9, dtype=torch.bool)
@@ -766,11 +771,13 @@ 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)
 
     workspace = MagicMock()
+    workspace.max_batch_size = 16
     workspace.position_ids = torch.tensor([2], dtype=torch.long)
     workspace.fill_input_ids_from_device.return_value = torch.tensor(
         [9], dtype=torch.long