diff --git a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py index 142161093f12..e259aa202f2a 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py +++ b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py @@ -74,6 +74,20 @@ def get_block_ids( must translate before returning. """ + #: V1 hands back a pre-eviction chain the legacy trim already understands. + exposes_block_ordinals: bool = False + + @abstractmethod + def get_block_ordinals( + self, + req: LlmRequest, + group_idx: int, + lg: AttentionLayerGroup, + ) -> np.ndarray: + """Entry ``i`` is the slot for tokens ``[i * tokens_per_block, ...)``, + ``-1`` where no page is bound. Position must not come from a length. + """ + @abstractmethod def commit_blocks_for_reuse(self, req: LlmRequest) -> None: """Commit KV blocks to radix tree for future prefix reuse. @@ -123,6 +137,10 @@ def get_block_ids(self, req, group_idx, lg): # noqa: ARG002 ) return np.asarray(pool_indices, dtype=np.int64) + def get_block_ordinals(self, req, group_idx, lg): + # V1 exposes the whole chain in order, so ids are already positional. + return self.get_block_ids(req, group_idx, lg) + def commit_blocks_for_reuse(self, req: LlmRequest) -> None: if not self.enable_block_reuse: return @@ -165,6 +183,16 @@ def get_block_ids(self, req, group_idx, lg): # noqa: ARG002 dtype=np.int64, ) + exposes_block_ordinals = True + + def get_block_ordinals(self, req, group_idx, lg): # noqa: ARG002 + return np.fromiter( + self._mgr.kv_cache_map[req.py_request_id].get_aggregated_page_indices( + group_idx, valid_only=False + ), + dtype=np.int64, + ) + def commit_blocks_for_reuse(self, req: LlmRequest) -> None: self._mgr.try_commit_blocks(req) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 9067443c16c9..9b5a77a2bdb1 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -50,7 +50,7 @@ GenTransferStatus, KvCacheTransceiver, ) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_kv_capacity_tokens from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( MambaHybridCacheManager, MambaHybridCacheManagerV2, @@ -318,9 +318,20 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: window_size = lg.sliding_window_size if window_size is not None: - allocated_blocks = ( - req.prompt_len + self._kv_cache_manager.num_extra_kv_tokens + tpb - 1 - ) // tpb + stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb) + if req.py_beam_width == 1 and adapter.exposes_block_ordinals: + # Position comes from the index, never from a length. + ordinals = adapter.get_block_ordinals(req, idx, lg) + in_window = ordinals[stale_end:total_blocks] + groups.append(in_window[in_window >= 0].astype(np.int64)) + continue + + capacity_tokens = get_kv_capacity_tokens( + req, + self._kv_cache_manager.num_extra_kv_tokens, + include_drafts=is_gen_only, + ) + allocated_blocks = (capacity_tokens + tpb - 1) // tpb beam0_block_ids, tail_block_ids = self._split_packed_beam_block_ids( block_ids, req.py_beam_width, @@ -333,24 +344,27 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: if tail_block_ids.size > 0 else beam0_block_ids ) - # Current PyExecutor cache managers disable KV-cache token sinks, - # so SWA block lists contain an evictable prompt prefix followed - # by the speculative scratch tail. If token sinks are enabled, - # this must use block-ordinal metadata to preserve the sink prefix. - # Remove scratch before trimming stale prompt blocks; otherwise a - # boundary-crossing allocation can displace initialized prompt KV. - scratch_blocks = max(0, allocated_blocks - total_blocks) + # Beam search still walks the packed-beam path below. + expected_valid = max(0, total_blocks - stale_end) + surplus_blocks = max(0, block_ids.size - expected_valid) + scratch_blocks = min(max(0, allocated_blocks - total_blocks), surplus_blocks) if scratch_blocks > 0: if req.py_beam_width != 1: raise ValueError("speculative scratch blocks require beam_width == 1") - block_ids = ( - block_ids[:-scratch_blocks] - if scratch_blocks < block_ids.size - else np.array([], dtype=np.int64) + block_ids = block_ids[:-scratch_blocks] + # Only two sizes are explainable here: the manager kept the + # stale head (total_blocks) or pruned it (expected_valid). + size_is_unexplained = expected_valid < block_ids.size < total_blocks + if req.py_beam_width == 1 and size_is_unexplained: + logger.warning_once( + "disagg KV slice: windowed layer group has " + f"{block_ids.size} pruned blocks but only {expected_valid} " + f"are in window (prompt_len={req.prompt_len}, tpb={tpb}, " + f"window={window_size}, allocated={allocated_blocks}). " + "Some producer disagrees with get_kv_capacity_tokens; " + "the trim is about to drop live KV.", + key="disagg_kv_slice_capacity_divergence", ) - # Drop stale blocks the manager may still expose (V1 pre-eviction). - stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb) - expected_valid = max(0, total_blocks - stale_end) # Stale prefix already pruned above; skip reuse-hit blocks that # land inside the window. Clamp to 0: ctx side has cached_per_lg # synthetically 0, and a reuse hit may fall entirely inside the diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index e3bfd3c2db4c..9cd9a7b046ff 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -92,7 +92,13 @@ KVCacheV2SsmLifeCycleIterationStats, KVCacheV2SsmSnapshotIterationStats, ) -from .llm_request import LlmRequest, LlmRequestState, SamplingConfig, get_draft_token_length +from .llm_request import ( + LlmRequest, + LlmRequestState, + SamplingConfig, + get_draft_token_length, + get_kv_capacity_tokens, +) from .resource_manager import ( BaseResourceManager, CacheTypeCpp, @@ -2661,7 +2667,7 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: # Helix requests carry the rank-local strided slice in prompt_len; # the global ledger sizes off the full prompt instead. prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len - target = prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens + target = get_kv_capacity_tokens(req, self.num_extra_kv_tokens, prompt_len=prompt_len) capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 3a2f977f6084..6167c3c9b084 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1702,3 +1702,17 @@ def get_draft_token_length(request: LlmRequest) -> int: if request.py_draft_tokens is not None: return len(request.py_draft_tokens) return 0 + + +def get_kv_capacity_tokens(request: LlmRequest, + num_extra_kv_tokens: int, + *, + prompt_len: Optional[int] = None, + include_drafts: bool = True) -> int: + """KV tokens a request occupies, defined once. + + Consumers that each spell the sum out drift apart silently. + """ + base = request.prompt_len if prompt_len is None else prompt_len + draft_len = get_draft_token_length(request) if include_drafts else 0 + return base + draft_len + num_extra_kv_tokens diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index eaca868014cf..aab9a8fc7de6 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -28,7 +28,9 @@ ) from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, LocalLayer from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm._torch.pyexecutor.llm_request import get_kv_capacity_tokens from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.logger import logger pytestmark = pytest.mark.cpu_only @@ -385,6 +387,8 @@ def _build_transceiver_for_kv_slice( cached_tokens: int = 0, is_generation_only: bool = False, beam_width: int = 1, + draft_len: int = 0, + exposes_block_ordinals: bool = True, ): """Stub a KvCacheTransceiverV2 so _create_kv_slice runs without dist setup. @@ -392,23 +396,48 @@ def _build_transceiver_for_kv_slice( - reuse adapter: tokens_per_block, per-layer-group cached count, block ids - page table: layer groups - cache manager: num_extra_kv_tokens (read in this code path) + - request: py_draft_tokens (read via get_draft_token_length) + + draft_len is the number of draft tokens the request currently carries. The + manager allocates for them on top of num_extra_kv_tokens, so the default + block list has to cover both. """ layer_group = AttentionLayerGroup( pool_group_idx=0, kv_head_num_per_rank=1, sliding_window_size=sliding_window_size, ) - total_blocks = (prompt_len + num_extra_kv_tokens + tokens_per_block - 1) // tokens_per_block + # Sized through the shared helper so the stub tracks the allocator's + # definition instead of drifting with the consumer under test. + total_blocks = ( + get_kv_capacity_tokens( + SimpleNamespace(prompt_len=prompt_len, py_draft_tokens=[0] * draft_len), + num_extra_kv_tokens, + ) + + tokens_per_block + - 1 + ) // tokens_per_block if block_ids is None: block_ids = np.arange(total_blocks, dtype=np.int64) else: block_ids = np.asarray(block_ids, dtype=np.int64) + # A manager exposes one entry per block ordinal, -1 where no page is bound. + # The stale prefix the caller trimmed out of block_ids is exactly that. + stale_end = 0 + if sliding_window_size is not None: + stale_end = max(0, (prompt_len + 1 - sliding_window_size) // tokens_per_block) + ordinals = np.concatenate( + [np.full(stale_end, -1, dtype=np.int64), np.asarray(block_ids, dtype=np.int64)] + ) + reuse_adapter = SimpleNamespace( tokens_per_block=tokens_per_block, get_cached_token_count_per_layer_group=lambda req, layer_groups: [cached_tokens] * len(layer_groups), get_block_ids=lambda req, idx, lg: block_ids, + get_block_ordinals=lambda req, idx, lg: ordinals, + exposes_block_ordinals=exposes_block_ordinals, ) page_table = SimpleNamespace(layer_groups=[layer_group]) cache_manager = SimpleNamespace(num_extra_kv_tokens=num_extra_kv_tokens) @@ -422,6 +451,7 @@ def _build_transceiver_for_kv_slice( prompt_len=prompt_len, py_request_id=0, py_beam_width=beam_width, + py_draft_tokens=list(range(draft_len)), is_generation_only_request=lambda: is_generation_only, ) return transceiver, req @@ -480,6 +510,7 @@ def test_swa_caps_oversized_non_speculative_v1_list_before_window_trim(self): prompt_len=32, block_ids=[100, 101, 102, 103, 104], sliding_window_size=16, + exposes_block_ordinals=False, ) kv_slice = transceiver._create_kv_slice(req) @@ -505,14 +536,19 @@ def test_swa_allocation_cap_preserves_packed_beam_tails(self): np.array([102, 103, 200, 201, 202], dtype=np.int64), ) + # A manager either hands back the whole chain (ordinals unavailable, legacy + # count-based trim) or the pruned valid list (ordinals available). Pairing + # the two spellings keeps each case honest about which it simulates. @pytest.mark.parametrize( - "block_ids", + "block_ids,exposes_block_ordinals", ( - pytest.param([100, 101, 102, 103, 104], id="v1-pre-eviction"), - pytest.param([102, 103, 104], id="v2-valid-only"), + pytest.param([100, 101, 102, 103, 104], False, id="v1-pre-eviction"), + pytest.param([102, 103, 104], True, id="v2-valid-only"), ), ) - def test_swa_trims_speculative_tail_before_stale_prompt_blocks(self, block_ids): + def test_swa_trims_speculative_tail_before_stale_prompt_blocks( + self, block_ids, exposes_block_ordinals + ): transceiver, req = _build_transceiver_for_kv_slice( num_extra_kv_tokens=2, prompt_len=32, @@ -520,6 +556,7 @@ def test_swa_trims_speculative_tail_before_stale_prompt_blocks(self, block_ids): sliding_window_size=16, cached_tokens=16, is_generation_only=True, + exposes_block_ordinals=exposes_block_ordinals, ) kv_slice = transceiver._create_kv_slice(req) @@ -573,6 +610,94 @@ def test_dspark_disagg_boundary_keeps_only_initialized_swa(self, prompt_len): ) +class TestSwaDraftTokenBlockAccounting: + """The SWA slice must cover the prompt window, whatever the request drafts. + + A generation-side windowed list is trimmed to its last expected_valid + entries, because for a sliding window the live blocks are a suffix. That is + only safe while the list holds prompt blocks alone: a speculative block left + in it makes the list longer for a reason the trim reads as "extra window", + so it drops real prompt blocks off the FRONT and every survivor lands one + block late. + + _create_kv_slice sizes the speculative tail from num_extra_kv_tokens plus + the request's own draft tokens. Dropping either term undercounts the tail, + which leaves a block behind exactly when the missing tokens cross a block + boundary -- a few percent of prompt lengths, so it costs accuracy without + ever looking like a crash. + + Expected values here come from prompt_len, the window and tokens_per_block + only, never from an allocation size, so they cannot drift with the code + under test. + """ + + TPB = 8 + WINDOW = 16 + TAG = 1000 # block id TAG+i holds tokens [i*TPB, (i+1)*TPB) + + def _blocks(self, prompt_len, max_draft_len): + """Position-tagged list a V2 manager exposes: valid blocks, no stale prefix.""" + num_extra = max(0, max_draft_len - 1) + allocated = (prompt_len + max_draft_len + num_extra + self.TPB - 1) // self.TPB + stale_end = max(0, (prompt_len + 1 - self.WINDOW) // self.TPB) + return np.arange(self.TAG + stale_end, self.TAG + allocated, dtype=np.int64) + + def _expected(self, prompt_len): + total_blocks = (prompt_len + self.TPB - 1) // self.TPB + stale_end = max(0, (prompt_len + 1 - self.WINDOW) // self.TPB) + return np.arange(self.TAG + stale_end, self.TAG + total_blocks, dtype=np.int64) + + # A contiguous span of more than one block period, so the prompt lengths + # whose draft tokens cross a boundary are covered by construction rather + # than by picking them out by hand. + _PROMPT_LENS = tuple(range(24, 41)) + + @pytest.mark.parametrize("max_draft_len", (2, 3, 5)) + @pytest.mark.parametrize("prompt_len", _PROMPT_LENS) + def test_window_blocks_are_not_shifted_by_draft_tokens(self, prompt_len, max_draft_len): + transceiver, req = _build_transceiver_for_kv_slice( + num_extra_kv_tokens=max_draft_len - 1, + prompt_len=prompt_len, + tokens_per_block=self.TPB, + block_ids=self._blocks(prompt_len, max_draft_len), + sliding_window_size=self.WINDOW, + is_generation_only=True, + draft_len=max_draft_len, + ) + + produced = transceiver._create_kv_slice(req).block_ids_per_layer_groups[0] + expected = self._expected(prompt_len) + + shift = int(produced[0] - expected[0]) if produced.size and expected.size else 0 + assert shift == 0, ( + f"prompt_len={prompt_len} max_draft_len={max_draft_len}: the window's " + f"first block moved by {shift}. Block {expected[0]} holds tokens " + f"[{(expected[0] - self.TAG) * self.TPB}, ...) but the slice starts at " + f"block {produced[0]}, so the peer writes every block one position late." + ) + np.testing.assert_array_equal(produced, expected) + + @pytest.mark.parametrize("max_draft_len", (2, 3, 5)) + def test_span_contains_a_boundary_crossing(self, max_draft_len): + """Guards the guard: the span above must contain the case that bites. + + Asserted on the inputs, not on _create_kv_slice, so it keeps holding + once the accounting is correct -- it pins the test data, not the bug. + """ + num_extra = max_draft_len - 1 + crossing = [ + prompt_len + for prompt_len in self._PROMPT_LENS + if (prompt_len + max_draft_len + num_extra + self.TPB - 1) // self.TPB + > (prompt_len + num_extra + self.TPB - 1) // self.TPB + ] + assert crossing, ( + f"no prompt_len in {self._PROMPT_LENS} makes the draft tokens cross a " + f"block boundary at max_draft_len={max_draft_len}; the test above would " + "pass without ever exercising the undercount" + ) + + # --------------------------------------------------------------------------- # CacheReuseAdapter.get_cached_token_count_per_layer_group: SWA clamp. # --------------------------------------------------------------------------- @@ -910,3 +1035,228 @@ def test_shutdown_is_idempotent(self): tc.shutdown() tc.shutdown() # second call short-circuits on the _shutdown guard. tc._transfer_worker.shutdown.assert_called_once() + + +# --------------------------------------------------------------------------- +# The capacity contract: one derivation, several consumers. +# --------------------------------------------------------------------------- + +_DIVERGENCE_KEY = "disagg_kv_slice_capacity_divergence" + + +class TestKvCapacityContract: + """Asserts positions, not counts. + + The defect kept every count correct while dropping a live block. + """ + + TPB = 8 + WINDOW = 16 + TAG = 1000 # block id TAG+i holds tokens [i*TPB, (i+1)*TPB) + + @staticmethod + def _req(prompt_len, draft_len): + return SimpleNamespace(prompt_len=prompt_len, py_draft_tokens=[0] * draft_len) + + def test_helper_is_the_sum_of_its_three_terms(self): + req = self._req(100, 3) + assert get_kv_capacity_tokens(req, 2) == 100 + 3 + 2 + assert get_kv_capacity_tokens(req, 0) == 100 + 3 + assert get_kv_capacity_tokens(self._req(100, 0), 2) == 100 + 2 + # The override exists for context-parallel Helix requests, whose + # rank-local prompt_len is not the length the ledger sizes off. + assert get_kv_capacity_tokens(req, 2, prompt_len=64) == 64 + 3 + 2 + + def _manager_blocks(self, prompt_len, draft_len, num_extra): + """Manager's list, sized through the shared helper the allocator uses.""" + allocated = ( + get_kv_capacity_tokens(self._req(prompt_len, draft_len), num_extra) + self.TPB - 1 + ) // self.TPB + stale_end = max(0, (prompt_len + 1 - self.WINDOW) // self.TPB) + return np.arange(self.TAG + stale_end, self.TAG + allocated, dtype=np.int64) + + def _blocks_in_window(self, prompt_len): + """Expected output, from prompt/window/tpb alone -- never from a size.""" + total_blocks = (prompt_len + self.TPB - 1) // self.TPB + stale_end = max(0, (prompt_len + 1 - self.WINDOW) // self.TPB) + return np.arange(self.TAG + stale_end, self.TAG + total_blocks, dtype=np.int64) + + # Longer than one block period, so boundary-crossing lengths are covered. + _PROMPT_LENS = tuple(range(24, 41)) + + @pytest.mark.parametrize("num_extra", (0, 1, 2)) + @pytest.mark.parametrize("draft_len", (0, 1, 3, 5)) + @pytest.mark.parametrize("prompt_len", _PROMPT_LENS) + def test_slice_follows_the_shared_capacity_definition(self, prompt_len, draft_len, num_extra): + transceiver, req = _build_transceiver_for_kv_slice( + num_extra_kv_tokens=num_extra, + prompt_len=prompt_len, + tokens_per_block=self.TPB, + block_ids=self._manager_blocks(prompt_len, draft_len, num_extra), + sliding_window_size=self.WINDOW, + is_generation_only=True, + draft_len=draft_len, + ) + + produced = transceiver._create_kv_slice(req).block_ids_per_layer_groups[0] + expected = self._blocks_in_window(prompt_len) + + shift = int(produced[0] - expected[0]) if produced.size and expected.size else 0 + assert shift == 0, ( + f"prompt_len={prompt_len} draft_len={draft_len} num_extra={num_extra}: " + f"the window's first block moved by {shift}. Block {expected[0]} holds " + f"tokens [{(expected[0] - self.TAG) * self.TPB}, ...) but the slice " + f"starts at block {produced[0]}, so the receiver writes every block " + "one position late." + ) + np.testing.assert_array_equal(produced, expected) + + def test_span_contains_a_boundary_crossing(self): + """Guards the guard: the sweep must contain a boundary-crossing case. + + Asserted on inputs, so it pins the test data rather than the bug. + """ + crossing = [ + (p, d, e) + for p in self._PROMPT_LENS + for d in (1, 3, 5) + for e in (0, 1, 2) + if (get_kv_capacity_tokens(self._req(p, d), e) + self.TPB - 1) // self.TPB + > (get_kv_capacity_tokens(self._req(p, 0), e) + self.TPB - 1) // self.TPB + ] + assert crossing, ( + "no parameter combination makes the draft tokens cross a block " + "boundary; the sweep above would pass without ever exercising the " + "case that produced the disagg speculative-decoding KV shift" + ) + + +class TestCapacityDivergenceIsReported: + """A warning that also fires on healthy traffic is one nobody reads.""" + + TPB = 8 + WINDOW = 16 + + @pytest.fixture(autouse=True) + def _reset_log_once(self): + logger._appeared_keys.discard(_DIVERGENCE_KEY) + yield + logger._appeared_keys.discard(_DIVERGENCE_KEY) + + def _run(self, block_ids, prompt_len=36, draft_len=3, num_extra=2): + transceiver, req = _build_transceiver_for_kv_slice( + num_extra_kv_tokens=num_extra, + prompt_len=prompt_len, + tokens_per_block=self.TPB, + block_ids=block_ids, + sliding_window_size=self.WINDOW, + is_generation_only=True, + draft_len=draft_len, + exposes_block_ordinals=False, + ) + transceiver._create_kv_slice(req) + return _DIVERGENCE_KEY in logger._appeared_keys + + def test_pruned_list_longer_than_the_window_is_reported(self): + """One block more than the shared definition allows: the 6676406 shape.""" + prompt_len, draft_len, num_extra = 36, 3, 2 + allocated = ( + get_kv_capacity_tokens( + SimpleNamespace(prompt_len=prompt_len, py_draft_tokens=[0] * draft_len), + num_extra, + ) + + self.TPB + - 1 + ) // self.TPB + stale_end = max(0, (prompt_len + 1 - self.WINDOW) // self.TPB) + # One block MORE than the definition allows; arange(stale_end, + # allocated) would be the correct list and would test nothing. + pruned = np.arange(stale_end, allocated + 1, dtype=np.int64) + total_blocks = (prompt_len + self.TPB - 1) // self.TPB + scratch = max(0, allocated - total_blocks) + assert total_blocks > pruned.size - scratch > total_blocks - stale_end, ( + "input must land strictly between the two healthy post-scratch " + "sizes, or it is not a divergence at all" + ) + + assert self._run(pruned), ( + "a pruned windowed list longer than the in-window count went " + "unreported; the trim silently dropped live KV, which is exactly " + "how the disagg speculative-decoding KV shift stayed invisible" + ) + + def test_full_pre_eviction_list_is_not_reported(self): + """Stay quiet when the surplus really is the stale head. + + A V1 manager may return every block, and then keeping the tail is right. + """ + prompt_len, draft_len, num_extra = 36, 3, 2 + allocated = ( + get_kv_capacity_tokens( + SimpleNamespace(prompt_len=prompt_len, py_draft_tokens=[0] * draft_len), + num_extra, + ) + + self.TPB + - 1 + ) // self.TPB + full = np.arange(allocated, dtype=np.int64) + + assert not self._run(full), ( + "the detector fired on a full pre-eviction list, where trimming to " + "the last expected_valid entries is the documented, correct " + "behaviour; a warning that fires on healthy traffic gets ignored" + ) + + +class TestWindowSelectionIsPositional: + """A hole must not shift its neighbours. + + Count-based selection passes the happy path and fails exactly here. + """ + + TPB = 8 + WINDOW = 16 + TAG = 1000 + + def _slice(self, prompt_len, ordinals): + transceiver, req = _build_transceiver_for_kv_slice( + num_extra_kv_tokens=2, + prompt_len=prompt_len, + tokens_per_block=self.TPB, + block_ids=np.asarray(ordinals, dtype=np.int64), + sliding_window_size=self.WINDOW, + is_generation_only=True, + draft_len=3, + ) + transceiver._reuse_adapter.get_block_ordinals = lambda req_, idx, lg, _o=np.asarray( + ordinals, dtype=np.int64 + ): _o + return transceiver._create_kv_slice(req).block_ids_per_layer_groups[0] + + def _window(self, prompt_len): + total = (prompt_len + self.TPB - 1) // self.TPB + stale = max(0, (prompt_len + 1 - self.WINDOW) // self.TPB) + return stale, total + + def test_scratch_past_the_prompt_is_not_selected(self): + """Ordinals beyond the prompt are speculative scratch, never transferred.""" + prompt_len = 33 + stale, total = self._window(prompt_len) + ordinals = np.arange(self.TAG, self.TAG + total + 2, dtype=np.int64) + + produced = self._slice(prompt_len, ordinals) + np.testing.assert_array_equal(produced, ordinals[stale:total]) + + def test_a_hole_does_not_shift_its_neighbours(self): + """Selecting by count would shift everything after the hole.""" + prompt_len = 33 + stale, total = self._window(prompt_len) + ordinals = np.arange(self.TAG, self.TAG + total + 2, dtype=np.int64) + hole = total - 1 + assert stale <= hole < total, "the hole must land inside the window" + ordinals[hole] = -1 + + produced = self._slice(prompt_len, ordinals) + expected = np.array([o for o in ordinals[stale:total] if o >= 0], dtype=np.int64) + np.testing.assert_array_equal(produced, expected) + assert self.TAG + hole not in produced.tolist()