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/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
-Q_TILE_ROWS = 64
+from astrai.config.inference_config import InferenceConfig
+
+_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-
-      buffered so an in-flight async H2D copy never races the next fill).
+    - ``input_ids``: per-step token IDs filled from host — values are
+      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
-            # double buffer keeps a copy in flight from being overwritten by
-            # the next fill.
+            # CUDA-graph capture) and filled via a host staging buffer.
             self.input_ids = torch.empty(
                 (max_batch_size,), dtype=torch.long, device=device
             )
-            self._pin = [
-                torch.empty((max_batch_size,), dtype=torch.long),
-                torch.empty((max_batch_size,), dtype=torch.long),
-            ]
-            self._pin_idx = 0
+            self._pin = torch.empty(
+                (max_batch_size,), dtype=torch.long, pin_memory=True
+            )
 
             # 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, 2]     (ml_part)
+            # [max_batch_size, max_q_heads, MAX_SPLITS, head_dim] (o_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
-        stable device buffer (``copy_`` without pinning is synchronous, so
-        the alternating buffers guard against an in-flight transfer).
+        Host values are staged through a pinned buffer and copied synchronously
+        into the stable device buffer.
         """
         b = len(ids)
-        pin = self._pin[self._pin_idx]
-        self._pin_idx ^= 1
         for i, v in enumerate(ids):
-            pin[i] = v
-        self.input_ids[:b].copy_(pin[:b])
+            self._pin[i] = v
+        self.input_ids[:b].copy_(self._pin[:b])
         return self.input_ids[:b]
 
     def fill_input_ids_from_device(self, tokens: Tensor) -> Tensor: