diff --git a/daser/connector/scheduler/lifecycle.py b/daser/connector/scheduler/lifecycle.py index ddcc69c..7c43e0b 100644 --- a/daser/connector/scheduler/lifecycle.py +++ b/daser/connector/scheduler/lifecycle.py @@ -132,10 +132,6 @@ def get_num_new_matched_tokens( logger.debug("[CONNECTOR] cache miss req=%s", request.request_id[:8]) return 0, False - extra_tokens = _contiguous_prefix_tokens(chunks, num_computed_tokens) - if extra_tokens <= 0: - return 0, False - pending_store = ( None if skip_save @@ -148,6 +144,10 @@ def get_num_new_matched_tokens( if pending_store is not None: self._pending_alloc[request.request_id] = pending_store + extra_tokens = _contiguous_prefix_tokens(chunks, num_computed_tokens) + if extra_tokens <= 0: + return 0, False + available = len(tokens) - num_computed_tokens if extra_tokens >= available: extra_tokens = available - 1 diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index d33ea9d..6ed81ed 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -2163,7 +2163,7 @@ class Output: def test_prefix_mode_hit_tracks_store_from_first_missing_slot(): - """Warm prefix hits should not rewrite loaded slots, only later misses.""" + """Warm prefix hits track suffix stores even when GPU coverage is longer.""" class MockIPCClient: def lookup(self, tokens, model_id): @@ -2211,6 +2211,10 @@ class MockRequest: assert pending.rolling_key == "hit-0" assert pending.rolling_slot_index == 1 + gpu_hit_connector = MockConnector() + assert gpu_hit_connector.get_num_new_matched_tokens(MockRequest(), 8) == (0, False) + assert gpu_hit_connector.pending_alloc["req"].start_slot_index == 1 + def test_prefix_mode_hit_still_allocates_missing_slot_stores_after_load(): """Prefix-hit requests still store newly computed suffix slots."""