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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user