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
This commit is contained in:
2026-09-04 14:28:04 +08:00
parent e13fe53475
commit ae7fc3059a
12 changed files with 244 additions and 152 deletions
+7
View File
@@ -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