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

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(
-        self, q, k, v, kv_cache=None, layer_id=0, attn_mask=None, is_causal=False
-    ):
-        return q
-
-    def fwd_prefill(
-        self, q, k, v, kv_cache=None, layer_id=0, attn_mask=None, is_causal=False
+    def forward(
+        self,
+        q,
+        k,
+        v,
+        kv_cache=None,
+        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