From a0a6b07b8e45a5461137d9a71abcafeee076ff45 Mon Sep 17 00:00:00 2001 From: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:25:34 -0700 Subject: [PATCH] [TRTLLM-16186][feat] Add KV cache lending API for transfer backends Introduce a mechanism-independent lending API with stable regions, leases, resume bounds and content identities for transfer backends. Provide the staging implementation while keeping its integration obligations separate from public guarantees. Centralize V2 operations and runtime metadata in a checked lifecycle facade, reject configured lossy cold-page quantization before attachment, and add lending and integration regression coverage. Signed-off-by: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 108 +- .../pyexecutor/kv_cache/sharing/__init__.py | 191 + .../kv_cache/sharing/_fetch_ledger.py | 395 ++ .../pyexecutor/kv_cache/sharing/_identity.py | 97 + .../pyexecutor/kv_cache/sharing/_layout.py | 576 ++ .../pyexecutor/kv_cache/sharing/_lender.py | 1582 +++++ .../pyexecutor/kv_cache/sharing/_manager.py | 416 ++ .../pyexecutor/kv_cache/sharing/_ranges.py | 86 + .../pyexecutor/kv_cache/sharing/_slots.py | 229 + .../pyexecutor/kv_cache/sharing/_types.py | 522 ++ tensorrt_llm/_torch/pyexecutor/llm_request.py | 8 +- .../pyexecutor/scheduler/scheduler_v2.py | 5 +- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../executor/kv_cache/sharing/conftest.py | 978 ++++ .../kv_cache/sharing/test_host_tier.py | 938 +++ .../kv_cache/sharing/test_layout_identity.py | 1197 ++++ .../test_layout_matches_native_page_table.py | 227 + .../executor/kv_cache/sharing/test_names.py | 178 + .../kv_cache/sharing/test_public_surface.py | 895 +++ .../executor/kv_cache/sharing/test_slots.py | 1453 +++++ .../kv_cache/sharing/test_staging_lender.py | 5156 +++++++++++++++++ ...st_kv_cache_manager_v2_helix_superblock.py | 2 + .../kv_cache/test_kv_cache_v2_scheduler.py | 2 +- .../kv_cache/test_kvcm2_integration.py | 36 +- .../kv_cache/test_spec_recompute_tail.py | 1 + .../executor/test_disagg_receive_ordering.py | 1 + .../executor/test_kv_connector_v2_prefix.py | 18 +- ...est_kv_connector_v2_prefix_real_manager.py | 6 +- .../test_kv_cache_stats_behavior.py | 2 +- 29 files changed, 15255 insertions(+), 51 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/__init__.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_fetch_ledger.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_identity.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_lender.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_ranges.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_slots.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_types.py create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/conftest.py create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/test_host_tier.py create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/test_names.py create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/test_slots.py create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 1937d4890116..5f3d0eb64f1c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1187,12 +1187,13 @@ def _settle_context_cursor(req: LlmRequest, reuse: int, tokens_per_block: int) - scheduler overwrites it, and the non-chunked path wants exactly this value. """ # ``reuse`` is ``num_committed_tokens``: the locally matched prefix alone. A - # connector-served prefix is held -- pages, and the cache's ``history_length`` - # -- but never committed, because committing is for tokens this worker - # computed. Settling a re-entering first chunk on the commit depth alone - # therefore points it below KV the cache holds, and ``history_length`` cannot - # be lowered to agree. - reuse = max(reuse, req.py_connector_served_position) + # prefix served from outside the manager (a connector's load, or a staging + # fetch the request resumes from) is held -- pages, and the cache's + # ``history_length`` -- but not yet committed when the first chunk settles: + # the commit after a context step covers it. Settling a re-entering first + # chunk on the commit depth alone therefore points it below KV the cache + # holds, and ``history_length`` cannot be lowered to agree. + reuse = max(reuse, req.py_served_prefix_end) _rewind_context_cursor(req, reuse) req.set_prepopulated_prompt_len(reuse, tokens_per_block) req.context_chunk_size = req.context_remaining_length @@ -1255,6 +1256,9 @@ class KVCacheManagerV2(BaseResourceManager): # symptom is a small perturbation that flips an argmax a couple of steps into # decode, which reads as ordinary sampling variation rather than as a bug. _supports_partial_reuse = True + # The lender sharing.attach_lender installs, one per manager for its life; None without one. + # On the class so it exists without running __init__. + _sharing = None def __init__( self, @@ -1297,8 +1301,12 @@ def __init__( self._validate_speculative_config(spec_config) self.is_disagg = is_disagg self.kv_connector_manager = kv_connector_manager + codec_config = getattr(cold_page_codec_provider, "config", None) + self._has_lossy_cold_page_codec = ( + getattr(codec_config, "algorithm", None) == "quantization_for_cold_page" + ) # Filled on first use; the layer grouping does not change after init. - self._connector_life_cycle_by_group: Optional[List[AttnLifeCycle]] = None + self._layer_group_life_cycles: Optional[List[AttnLifeCycle]] = None self.kv_cache_type = kv_cache_type self.pp_layers, self.num_layers = get_pp_layers( @@ -1415,8 +1423,12 @@ def __init__( self.enable_joint_kv_cache_reuse = joint_kv_cache_reuse # A draft manager only publishes to its radix tree when it is paired. self._can_publish_block_reuse = joint_kv_cache_reuse or not self.is_draft + # Prompt tokens past a position that one-model draft layers read. The lender's refusal + # check needs this raw value even where a subclass opts out of reuse-match backoff: + # a block's reuse key covers no token past its end. + self._draft_prompt_lookahead = draft_prompt_lookahead(spec_config) or 0 # Unsupported adapters keep their main-like reuse endpoint. - self.reuse_match_backoff = draft_prompt_lookahead(spec_config) or 0 + self.reuse_match_backoff = self._draft_prompt_lookahead if not self._supports_reuse_match_backoff: self.reuse_match_backoff = 0 self.enable_partial_reuse = _effective_partial_reuse( @@ -2162,6 +2174,7 @@ def _fill_fresh_kv_pages(self, request_id: int) -> None: num_blocks = int(kv_cache.num_blocks) state = self._fresh_pages_filled.setdefault(request_id, {}) filled = 0 + fill_stream = None unfillable_dtype = False for pool_id in range(self.num_pools): base = np.frombuffer( @@ -2202,6 +2215,12 @@ def _fill_fresh_kv_pages(self, request_id: int) -> None: state[pool_id] = base.copy() if not fresh: continue + # Staging and the previous owner's forward use the manager stream; native tier + # copies hand their dependencies back to it before pages are handed out. + # Order the diagnostic fill after that work without waiting for unrelated streams. + if fill_stream is None: + fill_stream = torch.cuda.current_stream(device=self._stream.device) + fill_stream.wait_stream(self._stream) for layer_idx in self.pp_layers: if self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] != pool_id: continue @@ -2228,10 +2247,9 @@ def _fill_fresh_kv_pages(self, request_id: int) -> None: ) if filled: # The forward pass reads these pages in the same iteration, so the - # fill has to be complete before it starts. A stream-ordered write - # would be enough on one stream; the synchronise makes it true - # whichever stream the manager and the model are using. - torch.cuda.synchronize() + # fill has to be complete before it starts, whichever stream the model uses. + # Waiting for the fill stream also covers its dependency on the manager stream. + fill_stream.synchronize() if not self._fresh_fill_announced: # Once per process, at warning level, for the same reason the # guard page announces itself: a run needs proof from its own @@ -3191,6 +3209,11 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() + @property + def draft_prompt_lookahead(self) -> int: + """Prompt tokens the configured draft reads ahead, before reuse-backoff opt-outs.""" + return self._draft_prompt_lookahead + @property def blocks_in_primary_pool(self) -> int: """ @@ -3606,6 +3629,7 @@ def revert_allocate_generation(self, req: LlmRequest) -> None: f"{req.py_request_id} from {kv_cache.capacity} to " f"{reverted_cap}" ) + self._after_shrink(request_id, kv_cache) def revert_allocate_context(self, req: LlmRequest) -> bool: """Undo this iteration's context resize. False means the cache was dropped, @@ -3636,6 +3660,7 @@ def revert_allocate_context(self, req: LlmRequest) -> bool: f"request {req.py_request_id} from {kv_cache.capacity} " f"to {pre_cap}" ) + self._after_shrink(req.py_request_id, kv_cache) if pre_cap > 0: kv_cache.suspend() return True @@ -4179,9 +4204,9 @@ def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: the group. A group with no window -- full attention, or a group whose layers carry no window at all -- yields an empty stale range. """ - if self._connector_life_cycle_by_group is None: + if self._layer_group_life_cycles is None: layers = self.kv_cache_manager_py_config.layers - self._connector_life_cycle_by_group = [ + self._layer_group_life_cycles = [ AttnLifeCycle.make( getattr(layer, "window_size", None), getattr(layer, "num_sink_tokens", None), @@ -4189,16 +4214,18 @@ def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: ) for layer in (layers[int(layer_ids[0])] for layer_ids in self.impl.layer_grouping) ] - return self._connector_life_cycle_by_group + return self._layer_group_life_cycles def _stale_block_range(self, layer_group_id: int, history_length: int) -> Tuple[int, int]: """Block ordinals ``[start, end)`` that ``layer_group_id`` no longer reads. - The cache's own life cycle decides this, so the connector's view and the - allocator's cannot drift. ``start`` is not 0 whenever attention sinks are - configured: the first ``num_sink_tokens`` stay live however far below the - window they fall, and masking them would hide from the connector pages - attention still reads. + The cache's own life cycle decides this, so its readers and the allocator + cannot drift: ``get_page_indices_by_layer_group`` masks these blocks in the + page indices it gives a KV cache connector, and a staging lender's window + checks read them through ``sharing._manager``. ``start`` is not 0 whenever + attention sinks are configured: the first ``num_sink_tokens`` stay live + however far below the window they fall, and masking them would hide from + the connector pages attention still reads. """ stale = self._life_cycle_by_layer_group()[layer_group_id].get_stale_range( history_length, self.tokens_per_block @@ -4417,7 +4444,7 @@ def _accept_connector_prefix_reservations(self, scheduled_batch: ScheduledReques self.kv_connector_manager.accept_prefix_load( req, reservation.start, reservation.end, by_group ) - req.py_connector_served_position = reservation.end + req.py_served_prefix_end = reservation.end self.kv_connector_manager.update_state_after_alloc(req, flat, by_group) req.py_connector_allocation_reported = True @@ -4521,14 +4548,16 @@ def _reserve_connector_prefix( for position, end in candidates: if position <= local_end or end <= position: continue - if self._resize_for_connector_prefix(req, kv_cache, position, end): + if self._resize_for_served_prefix(req, kv_cache, position, end): return position, end - position return None - def _resize_for_connector_prefix( - self, req: LlmRequest, kv_cache, position: int, end: int - ) -> bool: - """Cover ``[0, end)`` and mark ``[0, position)`` as already valid. + def _resize_for_served_prefix(self, req: LlmRequest, kv_cache, position: int, end: int) -> bool: + """Cover ``[0, end)`` and set the history to ``position``, marking + ``[0, position)`` as already valid, for KV served from outside the + manager: a KV cache connector's load passes the served end, and a + staging lender's fetch (``sharing._manager.grow``) passes the cache's + history, or ``end`` when a layer group has a sliding window. ``position`` becomes the cache's ``history_length``, which is the sole input to the sliding-window stale-range computation -- passing the @@ -4543,7 +4572,7 @@ def _resize_for_connector_prefix( capacity = max(pre_cap, target) if not kv_cache.resize(capacity, position): logger.debug( - "req %s: could not cover a connector prefix through %d (capacity %d -> %d)", + "req %s: could not cover a served prefix through %d (capacity %d -> %d)", req.py_request_id, position, pre_cap, @@ -4580,7 +4609,7 @@ def _mark_connector_prefix_populated( # The request can leave the batch from here -- parked for an asynchronous # load, or dropped with the whole batch -- and be settled again on the way # back. Carry the served end across that gap. - req.py_connector_served_position = position + req.py_served_prefix_end = position def _mirror_draft_kv_cache(self, req: LlmRequest): """The draft manager's entry for ``req``, created on first sight. @@ -4688,6 +4717,7 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): f"{req.py_request_id}: could not resize to {capacity} tokens" f"{self._draft_pool_diagnostic()}" ) + self._after_shrink(req.py_request_id, kv_cache) for req in scheduled_batch.generation_requests: kv_cache = self._mirror_draft_kv_cache(req) @@ -5662,7 +5692,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): request.py_connector_allocation_reported = False # Same allocation, same lifetime. The pages the served position vouches # for are gone, so a replay recomputes from its own reuse match. - request.py_connector_served_position = 0 + request.py_served_prefix_end = 0 if self.kv_connector_manager is not None and not self.is_draft: # The scheduler-output deltas are keyed to the allocation too. Left # behind, the replay lands in `cached_requests` carrying a block-id @@ -5682,6 +5712,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self.impl.clear_stats_excluded(request.py_request_id) return kv_cache.discard_pending_stats() + if self._sharing is not None: + self._sharing._on_free(request.py_request_id, kv_cache) kv_cache.close() self.impl.clear_stats_excluded(request.py_request_id) if request.py_request_id in self._early_freed_index_requests: @@ -5689,6 +5721,13 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): else: self.index_mapper.remove_sequence(request.py_request_id) + def _after_shrink(self, request_id: int, kv_cache: _KVCache) -> None: + """Tell the lender, if one is attached, that ``kv_cache`` may have shrunk in place: blocks + past its capacity lost their pages, and growing again does not bring their contents back. + Every in-place shrink the manager makes calls this right after it.""" + if self._sharing is not None: + self._sharing._on_shrink(request_id, kv_cache) + def get_layer_page_index_scale(self, layer_idx: int) -> int: """Page-index scale of this layer's KV buffer. Layers in one pool can have different scales (e.g. different head_dim), so per-layer callers @@ -5933,6 +5972,8 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool return bool(has_invalid_values) def shutdown(self): + if self._sharing is not None: + self._sharing._on_shutdown() if self.sparse_metadata_batch is not None: self.sparse_metadata_batch.record_read(self._stream.cuda_stream) self.sparse_metadata_batch.close() @@ -5948,6 +5989,10 @@ def shutdown(self): if self.conversation_manager is not None: self.conversation_manager.clear() self.impl.shutdown() + if self._sharing is not None: + # After the runtime's shutdown, which raises while any cache is open, in the map or not; + # a poisoned runtime returns at once, and a close then does nothing. + self._sharing._on_caches_closed() # Shut the streaming event manager down last so removals emitted during # cache / impl teardown (via the radix tree's own event-manager # reference) are still flushed before the publisher stops. Do not null @@ -6429,6 +6474,9 @@ def update_resources( f"to capacity {new_capacity} and history length " f"{history_length} tokens at generation update" ) + # Inline check: without a lender this per-request hot path makes no extra call. + if new_capacity is not None and self._sharing is not None: + self._after_shrink(req.py_request_id, kv_cache) self._allocated_draft_lens.pop(req.py_request_id, None) def copy_batch_block_offsets( @@ -6659,3 +6707,5 @@ def reset_reuse_state(self): self.impl.clear_reusable_blocks() if self.conversation_manager is not None: self.conversation_manager.clear() + if self._sharing is not None: + self._sharing._on_reset() diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/__init__.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/__init__.py new file mode 100644 index 000000000000..0a10fed3261a --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/__init__.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Lending a KV cache manager v2's blocks to transfer backends. + +This module is the whole API: callers import only its names, the package's other modules are +private, and ``attach_lender`` selects the implementation. The docstrings of these names state every +caller obligation. + +Roles: the integrator attaches one lender to a manager; the holder lends, polls, marks arrivals and +releases; the backend moves the bytes; the waiter asks ``Lender.resume_bounds``. Lender, +lease and hold methods run only on the manager's thread, one call at a time: the builder, then the +executor loop, then the shutdown thread. Positions are tokens. Ready views describe stable, +accessible lent memory; leases and region holds delimit backend access. Cache ownership remains +with the manager. +""" + +import typing as _typing + +from ._types import ( + LayerGroupBlocks, + Lease, + LeaseView, + Lender, + MemoryRegion, + RegionHold, + ResumeBounds, + StagingOptions, +) + +if _typing.TYPE_CHECKING: + from ..kv_cache_manager_v2 import KVCacheManagerV2 + +__all__ = [ + "LayerGroupBlocks", + "Lease", + "LeaseView", + "Lender", + "MemoryRegion", + "RegionHold", + "ResumeBounds", + "StagingOptions", + "attach_lender", +] + + +def attach_lender(manager: "KVCacheManagerV2", *, scope: bytes, staging: StagingOptions) -> Lender: + """Attach the manager's one lender, which relays whole blocks through host staging. + + The lender copies blocks between a request's device pages and host slots on the manager's + execution stream, for backends that can only read and write host memory, and says where the + request may resume after a fetch. ``Lender`` states what its holder and backends keep. A + name covers ``scope``, the committed tokens up to its block's end, their multimodal digests and + the manager's reuse scope, nothing else a request carries, so the decoder self-attention pool of + an encoder-decoder model is not lent by name. + + Caller obligations: + - This implementation reserves the request's device pages before waiting for host regions. + Keep room for running requests in each pool group while fetches are parked: their locked + pages, other unreclaimable pages, the largest generation request's full resume window and + its next step must fit within ``max_util_for_resume``. V2 cannot evict or preempt a parked + fetch. Reserving space for parked fetches alone does not prevent scheduler deadlock. + - Copies use the manager's execution stream, ordered with its forward work. Any writer of + recycled pages on another stream must wait for that stream first. The diagnostic fresh- + page fill (``TRTLLM_KV_FRESH_PAGE_FILL``) already does so and then synchronizes its own + stream; a fetch that grows pages with this diagnostic enabled can wait on the CPU. + Pageable memory under confidential computing can also make a call wait for its stream. + - A failed copy whose completion cannot be established quarantines its slots permanently. + Capacity can erode until a pending lease cannot fit and blocks those behind it. Release + leases that wait too long or free their requests. Quarantined slots, unreleased leases + or region holds, and cleanup errors can keep host memory until process exit. + - From ``lend_write`` until the request resumes or its cache is dropped, keep it among the + executor's active requests and unscheduled until the resume state is applied. To resume at + ``p``, set ``py_served_prefix_end=p``, raise the cache history to ``p`` if necessary and, + back in context state with a chunk spanning the remaining prompt, call + ``set_prepopulated_prompt_len(p, tokens_per_block)``. It moves the prepopulated length + and context position together; moving either alone is not sufficient. A resume at 0 + from a later context position requires dropping the cache and recomputing. + - With routed-expert capture, fetch only before the first context step, not after a + recompute pause; otherwise skipped positions can have no recorded routes. A later + fetch also does not change the response's original ``cached_tokens`` count. + - Under attention data parallelism without a transceiver, do not park every active + request on a rank unless the resulting all-rank forward stall is acceptable. + - SWA scratch reuse must be off for fetches. A window must not abandon an uncomputed block + still held for prefix reuse; with extra speculative KV tokens, keep later consecutive + segments shorter than ``window + tokens_per_block - 1`` or compute the remainder. + Requests requiring context logits, prompt logprobs or additional per-token model + outputs cannot be fetched, because those outputs are not restored by this implementation. + - These device-page reservations, execution-stream ordering, diagnostics, slot quarantine + and executor-specific constraints belong to the current implementation and may change + when its backing mechanism changes; the ``Lender`` contract does not expose them. + - Attach on the thread that builds the executor, to the manager the executor serves. A + manager built only to estimate the KV cache size shuts down first, after which a lender + attached there lends nothing; that shutdown frees the staging memory, or keeps it until + the process exits where a lease or staging hold is still open then, a slot was lost to a + failed copy, an error in a free stopped the lender (``Lease``), or an error arose in the + lender's own work within that shutdown. + - Shut the manager down last, with the device's CUDA context used for the attach current + on the calling thread, after the shutdown steps ``Lender`` lists. A manager + that never shuts down keeps the staging memory and its own page-index host buffer until + the process exits, and a shutdown that raises keeps that buffer until a retried shutdown + returns, or else until the process exits: the runtime's shutdown raises while any cache is + open, also one whose close raised in a free, unless the runtime is poisoned. Its leases, + released ones too, and the lender's records also keep the requests' caches, and with them + the manager's device pools, until the leases and the lender are dropped; drop the last + reference on a thread that has used the device's CUDA context. + - For a request whose one-model draft has a joint-reuse pool of its own, attach a lender to + both managers, the draft pool's under a ``scope`` of its own (laid out alike, both pools + would name blocks alike), fetch through both, since a fetch fills only this manager's + blocks, and resume as ``ResumeBounds`` states, keeping what ``Lender`` states for + such a pair. + - A caller serving requests the input processor may treat differently keeps them apart with + per-item ``multi_modal_uuids``, which names and the multimodal encoder cache both cover, + since a name trusts the multimodal digests the input pipeline computes, as the manager's + own prefix reuse does. A digest covers an item's bytes, not how the input processor treats + them, so blocks with different KV bytes can share a name: a NumPy array and a + ``torch.Tensor`` with the same dtype, shape and bytes share a digest though some input + processors rescale only one, and a request's ``mm_processor_kwargs`` reach the processor + but not the digest. A ``cache_salt`` covers names but not that cache, which some models + keep by default and which keys an item's embedding by its digest and + ``mm_processor_kwargs``: salted requests can still share an embedding, so a ``cache_salt`` + alone keeps them apart only with that cache off + (``multimodal_config.encoder_cache_max_bytes=0``). + - Under ``SaveHiddenStatesDecodingConfig`` a fetch restores only this manager's KV blocks, + so the hidden states saved for a resumed request pair its whole prompt with the rows of + the tokens it computed, then stale rows, misaligned as after the manager's own prefix + reuse: the caller fetches into no such request, or accepts that. + - An attach that raises for a reason its docstring does not list may keep the staging memory + it allocated until the process exits. + + Args: + manager: A ``KVCacheManagerV2`` with a ``mapping`` that commits blocks to its prefix-reuse + tree: block reuse on, and joint reuse for a draft manager; not built for a one-model + draft that reads prompt tokens ahead, and without a KV cache connector. Subclasses are + accepted; one whose ``try_commit_blocks`` keys blocks from other tokens or another reuse + scope than the base manager's changes names silently, at worst putting wrong bytes under + a correct name. + scope: Equal exactly where KV bytes mean the same; at most 65535 bytes. The lender only puts + it, unread, into every name's namespace, so different scopes never share a name. It must + cover every configuration that changes a block's bytes beyond what the manager declares: + the model, its weights and LoRA adapters (the adapter an id maps to), the one-model + draft whose pool it is (its weights, mode and the target layers it reads), element types + and quantization, the attention (windows such as ``max_attention_window``, sinks, sparse + attention), the attention backend's K/V arrangement (head-major or token-major; the base + manager declares head-major whichever backend writes its pages), and a release that + changes what KV bytes mean. An adapter enters a name only through the request's + ``lora_task_id``: instances sharing a ``scope`` map each id to the same adapter, or else + lend no LoRA request by name or take different scopes. The framework's assembly builds + it from its configuration, so no backend guesses which settings it needs. + staging: The staging size, in whole fetches. + + Returns: + The lender, installed on ``manager`` for the manager's life. + + Raises: + TypeError: ``manager`` is not a ``KVCacheManagerV2`` or has no mapping, ``scope`` is not + ``bytes``, or ``staging`` is not ``StagingOptions``. + ValueError: A lender is attached already; the manager is context-parallel (Helix included, + ``mapping.cp_size > 1``) or pipeline-parallel (``mapping.pp_size > 1``), holds recurrent + state, commits no blocks (block reuse off, or a draft manager without joint reuse), + since a publish lends only committed blocks, has a KV cache connector or a lossy + cold-page codec (``quantization_for_cold_page``), or is built for a + one-model draft that reads prompt tokens ahead; it has a layer group of sparse buffers + (``is_sparse``), whose read-only pages a cache can lock in host memory, though sparse + attention whose buffers are not marked sparse is not refused by this; its layer groups, + layer configs and declarations do not agree; it has no pages to stage; ``scope`` is + longer than 65535 bytes; ``staging.max_bytes`` is below one fetch. One-model Eagle and + MTP-Eagle read 1 prompt token ahead, vanilla MTP its draft length, and the check reads + the manager, not ``is_draft``, so it refuses draft layers sharing the target's manager, + a joint-reuse draft pool, and the target's manager of such a draft whose layers live + elsewhere. Drafts reading no token ahead (PARD, DFlash, DSpark) and managers without + speculative decoding are accepted. A draft whose read-ahead upstream has not established + counts as reading none: for one-model DraftTarget, staging accepts the target's manager + and refuses its unpaired draft pool, which commits no blocks, so a fetch fills only the + target's blocks, as the manager's own prefix reuse does. + MemoryError: Page-locking the staging memory failed. + """ + from ._lender import _attach_staging as _attach + + return _attach(manager, scope=scope, staging=staging) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_fetch_ledger.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_fetch_ledger.py new file mode 100644 index 000000000000..ec41ad58d56e --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_fetch_ledger.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A staging lender's fetch ledger: per request, the latest fetch, what fetches delivered, the +prefix a fetch kept and a cache whose committed tail a failed copy may have mixed, and the +resume bounds they give.""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable + +import numpy as np + +from . import _manager, _ranges +from ._types import ResumeBounds + +if TYPE_CHECKING: + from tensorrt_llm.runtime.kv_cache_manager_v2 import _KVCache + + from ...llm_request import LlmRequest + from ..kv_cache_manager_v2 import KVCacheManagerV2 + from ._layout import ManagerLayout + from ._lender import _LentBlockEntries, _QueuedCopy, _StagingLease + + +@dataclass(eq=False) +class _FetchRecord: + """One write lease's fetch into one cache: ``COPY_PENDING`` once the copy of its marked entries + is queued without error, ``COPY_DONE`` once that copy has completed, or at once where its marks + queue no copy and the lease has none. ``_FetchLedger._fetch_state`` derives its ``_FetchState`` + from ``marks_queued``, ``copy`` and whether ``_FetchLedger._latest_fetches`` holds it.""" + + start: int + kv: object + marks_queued: bool = False + copy: _QueuedCopy | None = None + + +class _FetchState(enum.Enum): + """A fetch's state, which ``_FetchLedger._fetch_state`` derives. + ``_FetchLedger._latest_fetches`` holds each request's latest fetch until it is abandoned, the + request is freed, or resume_bounds, while the lender serves the manager and is not stopped, + finds the request holding another cache. ``Staging._apply_marks`` takes it out + (``_take_out_fetch``) before queuing the copy and puts it back (``_put_back_fetch``) with + ``marks_queued`` set once the copy is queued without error.""" + + # In _latest_fetches, marks not queued: resume_bounds is None and a new fetch into the cache fails + # at the call. Staging._apply_marks sets marks_queued, or abandons it where the copy fails or + # the call raises; _abandon_fetch abandons it too. + AWAITING_MARKS = "awaiting_marks" + # In _latest_fetches, its marked entries' copy queued without error and pending: as + # AWAITING_MARKS for resume_bounds and a new fetch. The copy's completion makes it COPY_DONE. + COPY_PENDING = "copy_pending" + # In _latest_fetches, its copy complete or none queued: resume_bounds answers, and a new fetch into + # the cache may start. The request's next fetch replaces it. + COPY_DONE = "copy_done" + # Out of _latest_fetches before its marks were queued: none of its entries count. + ABANDONED = "abandoned" + # Out of _latest_fetches after its marks were queued: the request's next fetch replaced it, the + # request's free dropped it, or resume_bounds dropped it as a record of a cache the request no + # longer holds. Replacing it drops none of its entries from the cache's deliveries. + RETIRED = "retired" + + +@dataclass(eq=False) +class _Deliveries: + """What the fetches into one cache delivered: per layer group, by block ordinal, the entries + whose last copy from staging was queued without error and that no shrink freed since; + ``lowest_start`` is the lowest fetch start.""" + + kv: object + lowest_start: int + # Index by local layer group, then logical block ordinal (not LeaseView tuple positions). + # Recorded when a copy queues successfully; resume bounds also wait for its completion. + delivered_block_masks: list[np.ndarray] + # (tokens computed before, max_position, the end of the last run below it), last computed + resume_bounds_memo: tuple[int, int, int] | None = None + + +# TODO: each per-request record keeps its own reference to the cache it describes, so every +# reader compares it again, and a record of a replaced cache keeps that cache until the record is +# dropped or replaced. +class _FetchLedger: + """A staging lender's per-request fetch records and the resume_bounds they give. ``Staging`` + changes the records through these methods alone and reads them through these methods, apart + from the fields of the ``_Deliveries`` that ``_deliveries_of`` returns.""" + + def __init__( + self, layout: ManagerLayout, copy_done: Callable[[_QueuedCopy | None], bool] + ) -> None: + self._layout = layout + # Whether a copy has completed, as the lender answers: each event asked at most once per + # round of its progress. + self._copy_done = copy_done + # Per-request fetch records, each with the cache it describes: the request's free drops all + # four, and resume_bounds drops one of a cache the request no longer holds. + self._latest_fetches: dict[int, _FetchRecord] = {} # the latest fetch into each request + self._deliveries: dict[int, _Deliveries] = {} + # Requests whose history was past the committed tokens after a fetch's grow: the cache it + # grew, what the latest such fetch kept of the prefix computed before it, and the history + # its grow left. + self._kept_prefixes: dict[int, tuple[object, int, int]] = {} + # Requests whose committed tail block a failed copy may have left with the fetch's bytes in + # some pools and the bytes it held before in others: the cache. ResumeBounds is empty for it. + self._mixed_tail_caches: dict[int, object] = {} + + def _record_fetch( + self, request_id: int, kv: _KVCache, start: int, kept: int, position: int, committed: int + ) -> _FetchRecord: + """Record the request's latest fetch, from ``start`` into ``kv``, and, where its grow left + the history at ``position`` past the ``committed`` tokens, ``kept``: what the fetch keeps of + the prefix computed before it. Returns the fetch's record.""" + if position > committed: + self._kept_prefixes[request_id] = (kv, kept, position) + fetch = _FetchRecord(start, kv) + self._latest_fetches[request_id] = fetch + return fetch + + def _abandon_fetch(self, lease: _StagingLease) -> None: + """Drop the write lease's fetch record: it delivered nothing, and resume_bounds counts what + earlier fetches into the cache delivered.""" + if lease._fetch is not None and self._latest_fetches.get(lease._request_id) is lease._fetch: + del self._latest_fetches[lease._request_id] + + def _unsettled_fetch_refusal(self, request_id: int, kv: _KVCache) -> str | None: + """Why the request's earlier fetch into ``kv`` keeps a new one from starting, or ``None``: + it has not settled on this rank, where the copy its marks queued may still be pending.""" + previous = self._latest_fetches.get(request_id) + if previous is not None and previous.kv is kv and not self._fetch_settled(previous): + return f"request {request_id} already has an unsettled fetch" + return None + + def _fetch_settled(self, fetch: _FetchRecord) -> bool: + """The fetch's arrived entries are marked and their copy into the request's pages is + done.""" + return fetch.marks_queued and self._copy_done(fetch.copy) + + def _fetch_state(self, request_id: int, fetch: _FetchRecord) -> _FetchState: + """The state of ``fetch``, a fetch into request ``request_id``, from ``_latest_fetches`` and + ``_fetch_settled``; see ``_FetchState``.""" + if self._latest_fetches.get(request_id) is not fetch: + return _FetchState.RETIRED if fetch.marks_queued else _FetchState.ABANDONED + if self._fetch_settled(fetch): + return _FetchState.COPY_DONE + return _FetchState.COPY_PENDING if fetch.marks_queued else _FetchState.AWAITING_MARKS + + def _take_out_fetch(self, lease: _StagingLease, fetch: _FetchRecord | None) -> bool: + """Take ``fetch``, the write lease's, out of the latest fetches before its marks' copy is + queued, if it is still the request's latest; returns whether it was.""" + current = fetch is not None and self._latest_fetches.get(lease._request_id) is fetch + if current: + del self._latest_fetches[lease._request_id] + return current + + def _invalidate_delivered_blocks(self, lease: _StagingLease, copied: list[np.ndarray]) -> None: + """Stop counting the delivered entries ``copied`` selects, before their copy is queued: one + that fails partway can leave a row holding each fetch's bytes in different pools. + ``_record_delivered_blocks`` counts them again once the copy is queued without error.""" + delivered = self._deliveries.get(lease._request_id) + if delivered is None or delivered.kv is not lease._kv: + return + for lg, block_ordinals, mask in zip( + lease._block_entries.layer_group_indices, lease._block_entries.block_ordinals, copied + ): + blocks = delivered.delivered_block_masks[lg] + entries = block_ordinals[mask] + entries = entries[entries < len(blocks)] + if blocks[entries].any(): + blocks[entries] = False + delivered.resume_bounds_memo = None + + def _record_mixed_tail(self, request_id: int, kv: _KVCache) -> None: + """Record ``kv`` as the request's cache whose committed tail block a failed copy may have + left with the fetch's bytes in some pools and the bytes it held before in others.""" + self._mixed_tail_caches[request_id] = kv + + def _put_back_fetch( + self, lease: _StagingLease, fetch: _FetchRecord, copied: list[np.ndarray] + ) -> None: + """Put ``fetch`` back as the request's latest once its marks' copy is queued without error, + with ``marks_queued`` set and the lease's copy, and add the entries ``copied`` selects to + what fetches delivered into the cache.""" + fetch.copy = lease._copy + fetch.marks_queued = True + self._latest_fetches[lease._request_id] = fetch + self._record_delivered_blocks(lease._request_id, fetch, lease._block_entries, copied) + + def _record_delivered_blocks( + self, + request_id: int, + fetch: _FetchRecord, + entries: _LentBlockEntries, + copied: list[np.ndarray], + ) -> None: + """Add a fetch's copied entries to what earlier fetches into the same cache delivered, so a + fetch split into consecutive leases counts as one.""" + delivered = self._deliveries.get(request_id) + if delivered is None or delivered.kv is not fetch.kv: + empty = [np.zeros(0, dtype=bool) for _ in range(self._layout.num_layer_groups)] + delivered = _Deliveries(fetch.kv, fetch.start, empty) + self._deliveries[request_id] = delivered + delivered.lowest_start = min(delivered.lowest_start, fetch.start) + for lg, block_ordinals, mask in zip( + entries.layer_group_indices, entries.block_ordinals, copied + ): + got = block_ordinals[mask] + if not len(got): + continue + blocks = delivered.delivered_block_masks[lg] + if int(got.max()) >= len(blocks): + blocks = np.concatenate([blocks, np.zeros(int(got.max()) + 1 - len(blocks), bool)]) + blocks[got] = True + delivered.delivered_block_masks[lg] = blocks + delivered.resume_bounds_memo = None + + def _deliveries_of(self, request_id: int) -> _Deliveries | None: + """What fetches delivered into the request's cache, or ``None``; the record may describe + a cache the request no longer holds.""" + return self._deliveries.get(request_id) + + def _drop_deliveries(self, request_id: int) -> None: + """Drop what fetches delivered into the request's cache.""" + self._deliveries.pop(request_id, None) + + def _invalidate_delivered_blocks_past_capacity( + self, delivered: _Deliveries, kv: _KVCache + ) -> None: + """Forget delivered entries at or past the cache's block count: a shrink freed their pages, + and a regrow brings pages without their contents.""" + kept = _manager.num_blocks(kv) + for blocks in delivered.delivered_block_masks: + if blocks[kept:].any(): + blocks[kept:] = False + delivered.resume_bounds_memo = None + + def _drop_records(self, request_id: int) -> None: + """Drop the request's four records, at its free.""" + for records in ( + self._latest_fetches, + self._deliveries, + self._kept_prefixes, + self._mixed_tail_caches, + ): + records.pop(int(request_id), None) + + def _computed_tokens(self, request_id: int, kv: _KVCache) -> int: + """What resume_bounds counts as computed besides delivered entries: the tokens up to the + history, but only the committed ones and what the latest fetch kept while the history stays + where that fetch's grow left it past them (a windowed fetch moves it to its end).""" + state = _manager.cache_state(kv) + advanced = self._kept_prefixes.get(request_id) + # Past where that fetch left it, the history moved as the request resumed and computed on. + if advanced is None or advanced[0] is not kv or state.history > advanced[2]: + return max(state.committed, state.history) + return max(state.committed, advanced[1]) + + def _prefix_kept_by_fetch(self, request_id: int, kv: _KVCache, start: int) -> int: + """What a fetch from ``start`` keeps of what was computed before it, since it may overwrite + every block from ``start`` on. It never passes the history, which no shrink goes below; + delivered entries past it keep their own record.""" + return min(self._computed_tokens(request_id, kv), start) + + def _resume_bounds( + self, + request: LlmRequest, + request_id: int, + manager: KVCacheManagerV2, + kv: _KVCache, + history: int, + ) -> ResumeBounds | None: + """``Staging.resume_bounds`` of a lender that serves and has not stopped, for the request's + cache ``kv`` at a history of ``history`` tokens.""" + # A record of a replaced cache (a restart) is moot. + fetch = self._latest_fetches.get(request_id) + if fetch is not None and fetch.kv is not kv: + del self._latest_fetches[request_id] + fetch = None + delivered = self._deliveries.get(request_id) + if delivered is not None and delivered.kv is not kv: + del self._deliveries[request_id] + delivered = None + advanced = self._kept_prefixes.get(request_id) + if advanced is not None and advanced[0] is not kv: + del self._kept_prefixes[request_id] + tainted = self._mixed_tail_caches.get(request_id) + if tainted is not None and tainted is not kv: + del self._mixed_tail_caches[request_id] + tainted = None + if fetch is not None and not self._fetch_settled(fetch): + return None + if tainted is not None: + # Empty: the request drops its cache and computes from 0. + return ResumeBounds(max_position=0, min_position=max(history, 1)) + known = self._computed_tokens(request_id, kv) + if delivered is None: + # Nothing delivered: what was computed and kept counts, resumed no lower than the + # history, which a windowed fetch moved to its end. + usable, floor = self._clip_to_bidirectional_runs(request, known) + return ResumeBounds(max_position=usable, min_position=max(floor, history)) + # A shrink the manager did not report still shows as blocks past the capacity. + self._invalidate_delivered_blocks_past_capacity(delivered, kv) + if delivered.resume_bounds_memo is None or delivered.resume_bounds_memo[0] != known: + usable = self._max_resume_position(manager, delivered, known) + delivered.resume_bounds_memo = ( + known, + *self._clip_to_bidirectional_runs(request, usable), + ) + _, usable, floor = delivered.resume_bounds_memo + return ResumeBounds( + max_position=usable, + min_position=max( + floor, int(self._min_resume_position(manager, history, delivered.lowest_start)) + ), + ) + + def _clip_to_bidirectional_runs(self, request: LlmRequest, usable: int) -> tuple[int, int]: + """The end lowered to the start of a run of multimodal tokens it falls strictly inside, and + the lowest floor that then leaves no position strictly inside a run below that end: the end + of the last run at or below it. Runs are those the scheduler keeps within one chunk.""" + runs = _manager.bidirectional_runs(request) + for b, e in runs: + if b < usable < e: + usable = b + return int(usable), max((e for _, e in runs if e <= usable), default=0) + + def _min_resume_position(self, manager: KVCacheManagerV2, history: int, origin: int) -> int: + """The lowest start that needs no restart: the history where a chunk resumed below it may + end below it and raise, or a window has released blocks at it (they have no pages), else + the smaller of the lowest start of a delivered fetch and the history.""" + if _ranges._floor_follows_history(manager): + return history + for lg in range(self._layout.num_layer_groups): + stale_beg, stale_end = _ranges._stale_range(manager, self._layout, lg, history) + if stale_end > stale_beg: + return history + return min(origin, history) + + def _max_resume_position( + self, manager: KVCacheManagerV2, delivered: _Deliveries, known: int + ) -> int: + """The largest start ``P >= known`` where every layer group has what it reads among the + blocks below ``known``, computed before the fetches, and delivered entries: full attention + every block below ``P``, a window its sinks and in-window blocks.""" + # TODO: a fetch none of whose entries a window reads delivers none, so the usable end stays at + # the blocks computed before the fetches, and the interval can come out empty although the + # fetch covered all that a resume at its end reads. + tpb = int(self._layout.tokens_per_block) + first = known // tpb # the blocks below were computed + last = max((len(blocks) for blocks in delivered.delivered_block_masks), default=0) + if last <= first: + return known + # Non-monotonic in P under windows, so each is checked. Blocks below known behind a window's + # history have no pages, but no start at or above the floor reads them. + top = last # no start past a full-attention group's first missing block passes that group + missing = [] # per layer group: missing[j], the undelivered blocks of [first, first + j) + for lg, blocks in enumerate(delivered.delivered_block_masks): + have = np.zeros(last - first, dtype=bool) + mine = blocks[first:last] + have[: len(mine)] = mine + missing.append(np.concatenate([[0], np.cumsum(~have)])) + if self._layout.windows[lg] is None and not have.all(): + top = min(top, first + int(np.argmin(have))) + + def gap(lg: int, a: int, b: int) -> bool: + a, b = max(a, first), min(b, last) + return b > a and missing[lg][b - first] - missing[lg][a - first] > 0 + + # The largest start first: a fetch whose leases all landed is usable at its end at once. + for end_block in range(top, first, -1): + for lg in range(len(delivered.delivered_block_masks)): + stale_beg, stale_end = _ranges._stale_range( + manager, self._layout, lg, end_block * tpb + ) + if stale_end > stale_beg: + bad = gap(lg, first, min(stale_beg, end_block)) or gap(lg, stale_end, end_block) + else: + bad = gap(lg, first, end_block) + if bad: + break + else: + return end_block * tpb + return known diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_identity.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_identity.py new file mode 100644 index 000000000000..f89dd9571840 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_identity.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Object keys and staging region labels instances exchange; equal keys mean interchangeable bytes +given the ``scope``, the ``manager`` and the caller obligations ``attach_lender`` states. The +format is a contract between instances: changing it bumps ``KEY_FORMAT_VERSION``.""" + +from __future__ import annotations + +import hashlib +from typing import Iterable, Sequence + +import numpy as np + +from ._layout import ShardDesc, canonical_order + +KEY_FORMAT_VERSION = 2 +REUSE_KEY_BYTES = 32 +_NAMESPACE_BYTES = 16 +_GROUP_BYTES = 6 # the layer group's canonical index, shard count and shard index, u16 each +OBJECT_KEY_BYTES = _NAMESPACE_BYTES + REUSE_KEY_BYTES + _GROUP_BYTES +_REGION_LAYOUT_HEX = 16 + + +def _u16(value: int) -> bytes: + return int(value).to_bytes(2, "big") + + +def namespace(scope: bytes, layout_id: bytes) -> bytes: + """The 16 bytes every object key starts with, from the caller's ``scope`` and the block layout's + ``layout_id`` (each length-prefixed). ``ValueError`` for an empty ``layout_id`` or a ``scope`` + longer than 65535 bytes.""" + if not layout_id: + raise ValueError("a namespace without layout_id would mix incompatible bytes") + if len(scope) > 0xFFFF: + raise ValueError(f"scope is {len(scope)} bytes; at most 65535 fit its length prefix") + digest = hashlib.sha256(b"trtllm-kv-object" + _u16(KEY_FORMAT_VERSION)) + digest.update(_u16(len(scope)) + scope) + digest.update(_u16(len(layout_id)) + layout_id) + return digest.digest()[:_NAMESPACE_BYTES] + + +class Identity: + """Object keys and region labels for one manager. ``layers[g]`` holds layer group ``g``'s global + layer ids and ``shards[g]`` the ``(count, index)`` share of its content this rank holds.""" + + namespace: bytes + layout_id: bytes + canonical: list[int] + + def __init__( + self, + scope: bytes, + layout_id: bytes, + layers: Sequence[Sequence[int]], + shards: Sequence[tuple[int, int]], + ) -> None: + if len(shards) != len(layers): + raise ValueError(f"{len(layers)} layer groups but {len(shards)} shards") + self.layout_id = bytes(layout_id) + self.namespace = namespace(bytes(scope), self.layout_id) + self.canonical = canonical_order(layers) + self._namespace = np.frombuffer(self.namespace, dtype=np.uint8) + self._groups = [] + for index, (count, share) in zip(self.canonical, shards): + ShardDesc(count, share) # ValueError for a share outside the count + suffix = _u16(index) + _u16(count) + _u16(share) + self._groups.append(np.frombuffer(suffix, dtype=np.uint8)) + + def object_keys(self, layer_group_index: int, reuse_keys: np.ndarray) -> np.ndarray: + """``uint8 (n, OBJECT_KEY_BYTES)``: the namespace, each entry's 32-byte reuse key, then + the group's canonical index, shard count and shard index as big-endian ``u16``.""" + if not 0 <= layer_group_index < len(self._groups): + raise ValueError(f"unknown layer group {layer_group_index}") + reuse_keys = np.ascontiguousarray(reuse_keys, dtype=np.uint8).reshape(-1, REUSE_KEY_BYTES) + out = np.empty((reuse_keys.shape[0], OBJECT_KEY_BYTES), dtype=np.uint8) + out[:, :_NAMESPACE_BYTES] = self._namespace + out[:, _NAMESPACE_BYTES : _NAMESPACE_BYTES + REUSE_KEY_BYTES] = reuse_keys + out[:, _NAMESPACE_BYTES + REUSE_KEY_BYTES :] = self._groups[layer_group_index] + return out + + def region_label(self, layer_group_indices: Iterable[int]) -> str: + """The layout-derived label for a staging region serving ``layer_group_indices``: + equal on instances laid out alike, independent of addresses, capacity and scope.""" + groups = sorted({self.canonical[int(g)] for g in layer_group_indices}) + return f"{self.layout_id.hex()[:_REGION_LAYOUT_HEX]}:lg{'+'.join(str(g) for g in groups)}" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py new file mode 100644 index 000000000000..26d98f8a6f68 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py @@ -0,0 +1,576 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The layout of a manager's blocks as the lender needs it, and ``layout_id``, the digest of how a +block's bytes are laid out. Read from the manager's public layout surface and declarations.""" + +from __future__ import annotations + +import enum +import hashlib +import json +import struct +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, Mapping, Sequence + +import numpy as np + +from . import _manager + +if TYPE_CHECKING: + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig, SsmLayerConfig + + from ..kv_cache_manager_v2 import KVCacheManagerV2 + +# A new field left out at its default keeps the version; changing the encoding or the meaning of +# an existing field bumps it. +LAYOUT_ID_VERSION = 2 + +# The fallback entry of the manager's per-role declarations (``Role.ALL``). +_ALL_ROLE = "all" + + +class BufferMapper(enum.IntEnum): + """How a buffer's bytes split across tensor-parallel ranks; the values of the mapper kinds + managers declare per role, read by value.""" + + INDEXED = 0 # head-major K/V, the default: a head range is one contiguous range per buffer + REPLICATED = 1 # bytes identical on every TP rank + NHD = 2 # token-major [token, head, dim] K/V: a head range is a slice inside every token + SECTIONED = 3 # [Sec0|Sec1|...], each section sharded independently (recurrent conv state) + + +@dataclass(frozen=True) +class BufferGeometry: + """Resharding geometry of one buffer; all ``None`` means no head axis. ``section_bytes`` sum to + the buffer's size; ``bytes_per_head * num_heads`` equals it.""" + + section_bytes: tuple[int, ...] | None = None + bytes_per_head: int | None = None + num_heads: int | None = None + + def validate(self) -> BufferGeometry: + """A copy with ``section_bytes`` as ints; ``ValueError`` unless every size given is + positive.""" + sections = self.section_bytes + if sections is not None: + sections = tuple(int(s) for s in sections) + if any(s <= 0 for s in sections): + raise ValueError(f"section sizes must be positive, got {sections}") + for name in ("bytes_per_head", "num_heads"): + value = getattr(self, name) + if value is not None and int(value) <= 0: + raise ValueError(f"{name} must be positive, got {value}") + return BufferGeometry(sections, self.bytes_per_head, self.num_heads) + + def check_size(self, size: int) -> None: + """``ValueError`` unless this geometry describes a buffer of ``size`` bytes.""" + if self.section_bytes is not None and sum(self.section_bytes) != size: + raise ValueError(f"sections {self.section_bytes} do not sum to {size} bytes") + if self.bytes_per_head is not None and size % self.bytes_per_head: + raise ValueError( + f"{size} bytes is not a whole number of {self.bytes_per_head}-byte heads" + ) + if ( + self.bytes_per_head is not None + and self.num_heads is not None + and self.bytes_per_head * self.num_heads != size + ): + raise ValueError( + f"{self.num_heads} heads of {self.bytes_per_head} bytes is not {size} bytes" + ) + + +@dataclass(frozen=True) +class ShardDesc: + """Which share (``index`` of ``count``) of a layer group's content this rank holds; ``(1, 0)`` + is the whole content, which every rank names alike. ``ValueError`` for an index outside the + count.""" + + count: int = 1 + index: int = 0 + + def __post_init__(self) -> None: + if self.count < 1 or not 0 <= self.index < self.count: + raise ValueError(f"shard {self.index} of {self.count}") + + +@dataclass(frozen=True) +class LayerGroupDesc: + """One layer group: its kind, window and sink blocks (which blocks exist), its device pool + group, its global layer ids and the share of its content this rank holds.""" + + kind: Literal["attention", "state"] + window: int | None + sink_blocks: int + pool_group: int + layers: tuple[int, ...] + shard: ShardDesc = ShardDesc() + + +@dataclass(frozen=True) +class BufferDesc: + """One buffer of one layer (global id) in a page: ``size`` bytes at ``offset`` in the slot of + ``pool`` = (pool group, pool index), how its bytes split across ranks, whether peers transfer + it, and its sub-pages per block. ``ValueError`` for a non-positive size or expansion.""" + + layer: int + role: str + pool: tuple[int, int] + offset: int + size: int + mapper: BufferMapper = BufferMapper.INDEXED + geometry: BufferGeometry = BufferGeometry() + # False for local-only buffers peers skip; they still ride along wherever whole slots are + # copied. + transfer: bool = True + # The buffer's own tokens per block is tokens_per_block / expansion; a consumer reading inside + # the buffer that does not support the value must reject the layout. + expansion: int = 1 + + def __post_init__(self) -> None: + object.__setattr__(self, "role", str(self.role)) + object.__setattr__(self, "mapper", BufferMapper(self.mapper)) + if not isinstance(self.geometry, BufferGeometry): + raise ValueError(f"geometry must be a buffer geometry, got {self.geometry!r}") + if self.size <= 0 or self.offset < 0: + raise ValueError(f"buffer at offset {self.offset} of {self.size} bytes") + if int(self.expansion) < 1: + raise ValueError(f"expansion must be a positive integer, got {self.expansion}") + self.geometry.check_size(self.size) + + +@dataclass(frozen=True) +class Layout: + """How pages are built, without addresses or slot counts: ``pool_groups[g]`` holds the slot + widths of device pool group ``g`` in pool order. ``ValueError`` for an expansion that does not + divide ``tokens_per_block``.""" + + tokens_per_block: int + pool_groups: Mapping[int, tuple[int, ...]] + layer_groups: tuple[LayerGroupDesc, ...] + buffers: tuple[BufferDesc, ...] + + def __post_init__(self) -> None: + pool_groups = { + int(g): tuple(int(w) for w in widths) for g, widths in dict(self.pool_groups).items() + } + object.__setattr__(self, "pool_groups", pool_groups) + for b in self.buffers: + if self.tokens_per_block % int(b.expansion): + raise ValueError( + f"buffer ({b.layer}, {b.role}): expansion {b.expansion} does not divide " + f"{self.tokens_per_block} tokens per block" + ) + + +# Enums encode as the member name, floats as "f64:" plus the 16 hex digits of their big-endian +# IEEE-754 binary64 bits; a None value drops an object's key and is null in an array. +def _canonical(value: object) -> object: + if isinstance(value, enum.Enum): + return value.name + if value is None or isinstance(value, (bool, str)): + return value + if isinstance(value, np.bool_): + return bool(value) + if isinstance(value, (int, np.integer)): + return int(value) + if isinstance(value, (float, np.floating)): + return "f64:" + struct.pack(">d", float(value)).hex() + if isinstance(value, Mapping): + out = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"canonical objects have string keys, got {key!r}") + if item is not None: + out[key] = _canonical(item) + return out + if isinstance(value, (list, tuple)): + return [_canonical(item) for item in value] + raise TypeError(f"{type(value).__name__} has no canonical encoding") + + +def canonical_bytes(document: Mapping) -> bytes: + """``document`` as ASCII JSON, keys sorted, no whitespace; enums by name, floats as ``f64:`` and + their big-endian bits in hex, ``None`` values dropped from objects.""" + return json.dumps( + _canonical(document), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("ascii") + + +def canonical_order(layers: Sequence[Sequence[int]]) -> list[int]: + """Each local layer group's canonical index: groups ordered by their smallest global layer id, + which pipeline parallelism does not change; ties keep the local order.""" + order = sorted(range(len(layers)), key=lambda g: (min(layers[g], default=1 << 62), g)) + canonical = [0] * len(layers) + for index, local in enumerate(order): + canonical[local] = index + return canonical + + +def _geometry_document(geometry: BufferGeometry) -> dict: + return { + "section_bytes": None if geometry.section_bytes is None else list(geometry.section_bytes), + "bytes_per_head": geometry.bytes_per_head, + "num_heads": geometry.num_heads, + } + + +def _buffer_document(layout: Layout, b: BufferDesc) -> dict: + # A buffer's position in the page: the slot widths of the pools before it in its group, plus + # its offset inside the slot. + widths = layout.pool_groups.get(b.pool[0]) + if widths is None: + raise ValueError(f"buffer ({b.layer}, {b.role}) lies in pool {b.pool}, not in the layout") + doc = { + "layer": int(b.layer), + "role": b.role, + "offset": sum(widths[: b.pool[1]]) + int(b.offset), + "size": int(b.size), + } + # Fields at their defaults are left out, so a new defaulted field keeps existing ids. + if b.mapper is not BufferMapper.INDEXED: + doc["mapper"] = b.mapper.name + geometry = {k: v for k, v in _geometry_document(b.geometry).items() if v is not None} + if geometry: + doc["geometry"] = geometry + if not b.transfer: + doc["transfer"] = False + if int(b.expansion) != 1: + doc["expansion"] = int(b.expansion) + return doc + + +def layout_document(layout: Layout) -> dict: + """The document ``layout_id`` hashes: tokens per block, groups in canonical order (layers, slot + widths, kind), buffers by (layer, role). Addresses, windows, sinks, shards, local numbering and + element types stay out.""" + canonical = canonical_order([desc.layers for desc in layout.layer_groups]) + groups = [None] * len(layout.layer_groups) + for local, desc in enumerate(layout.layer_groups): + widths = layout.pool_groups.get(desc.pool_group) + if widths is None: + raise ValueError( + f"layer group {local} maps to pool group {desc.pool_group}, which is not in the " + "layout" + ) + entry = { + "layers": sorted(int(layer) for layer in desc.layers), + "slot_bytes": [int(w) for w in widths], + } + if desc.kind != "attention": + entry["kind"] = desc.kind + groups[canonical[local]] = entry + buffers = sorted( + (_buffer_document(layout, b) for b in layout.buffers), + key=lambda d: (d["layer"], d["role"]), + ) + for before, after in zip(buffers, buffers[1:]): + if (before["layer"], before["role"]) == (after["layer"], after["role"]): + raise ValueError(f"buffer ({before['layer']}, {before['role']}) appears twice") + return { + "v": LAYOUT_ID_VERSION, + "tokens_per_block": int(layout.tokens_per_block), + "groups": groups, + "buffers": buffers, + } + + +def layout_id(layout: Layout) -> bytes: + """SHA-256 of ``canonical_bytes(layout_document(layout))``.""" + return hashlib.sha256(canonical_bytes(layout_document(layout))).digest() + + +@dataclass(frozen=True) +class DevicePool: + """One device pool of pool group ``group``: slot ``s`` lives at ``base + s * slot_bytes``.""" + + group: int + index: int + base: int + slot_bytes: int + num_slots: int + + +@dataclass(frozen=True, eq=False) +class ManagerLayout: + """What the lender reads of one manager, as plain values that hold no manager: the address-free + ``layout``, facts per layer group (tuples indexed by the local layer group) and per device pool + group (mappings keyed by the pool group).""" + + layout: Layout + tokens_per_block: int + pool_group_of: tuple[int, ...] + windows: tuple[int | None, ...] + sink_blocks: tuple[int, ...] + recurrent: tuple[bool, ...] + sparse: tuple[bool, ...] + layers: tuple[tuple[int, ...], ...] + shards: tuple[tuple[int, int], ...] + page_bytes: Mapping[int, int] + device_pools: Mapping[int, tuple[DevicePool, ...]] + + @property + def num_layer_groups(self) -> int: + """The manager's layer groups.""" + return len(self.pool_group_of) + + @property + def pool_groups(self) -> tuple[int, ...]: + """The device pool groups, sorted; staging regions follow this order.""" + return tuple(sorted(self.device_pools)) + + +def global_layer_ids(manager: KVCacheManagerV2, internal_ids: Sequence[int]) -> list[int]: + """The global id of each internal layer id, the same under any pipeline split: the declared id; + else for virtual layers ``model_layer * number_of_attention_types + attention_type``; else the + model layer.""" + # One declared id per internal layer of each layer group, in order (FP4 MLA's tail layers). + declared = _manager.declared_global_layer_ids(manager) + if declared is not None: + return [declared[int(lid)] for lid in internal_ids] + virtual = _manager.virtual_layers(manager) + if virtual is None: + return [int(manager.pp_layers[int(lid)]) for lid in internal_ids] + # The number of types counts every member of the enum, so stages holding different attention + # types agree. + inverse, num_types = virtual + return [inverse[int(lid)][0] * num_types + inverse[int(lid)][1] for lid in internal_ids] + + +def _sparse(layer_config: AttentionLayerConfig | SsmLayerConfig) -> bool: + """Whether a layer config of the runtime has sparse buffers (``is_sparse``), whose read-only + pages a cache can lock in host memory.""" + return any(getattr(b, "is_sparse", False) for b in getattr(layer_config, "buffers", ()) or ()) + + +def _expansions( + layer_config: AttentionLayerConfig | SsmLayerConfig, tokens_per_block: int +) -> dict[str, int]: + """Role -> sub-pages per block, from the layer's ``tokens_per_block_override``.""" + out = {} + for buffer in getattr(layer_config, "buffers", ()) or (): + override = getattr(buffer, "tokens_per_block_override", None) + if override is not None: + if tokens_per_block % int(override): + raise ValueError( + f"buffer {buffer.role}: {override} tokens per block does not divide " + f"{tokens_per_block}" + ) + out[str(buffer.role)] = tokens_per_block // int(override) + return out + + +class _Heads: + """Head counts of attention buffers: per rank from the manager, in total from its inputs.""" + + def __init__(self, manager: KVCacheManagerV2) -> None: + mapping = getattr(manager, "mapping", None) + dp = bool(getattr(mapping, "enable_attention_dp", False)) + self.tp_size = 1 if mapping is None or dp else max(int(getattr(mapping, "tp_size", 1)), 1) + self.tp_rank = ( + 0 if self.tp_size == 1 else int(getattr(mapping, "tp_rank", 0)) % self.tp_size + ) + self._per_rank = list(getattr(manager, "num_kv_heads_per_layer", ()) or ()) + self._total = getattr(manager, "num_kv_heads", None) + self._virtual = _manager.virtual_layers(manager) is not None + self._pp_layers = manager.pp_layers + + def per_rank(self, internal: int) -> int: + if not self._per_rank: + return 0 + index = internal if internal < len(self._per_rank) else 0 + return int(self._per_rank[index] or 0) + + def total(self, internal: int) -> int | None: + if isinstance(self._total, int): + return int(self._total) + if self._total is None or self._virtual: + return None + if internal >= len(self._pp_layers): + return None # an extra internal layer (FP4 MLA's tail) is no model layer + value = self._total[int(self._pp_layers[internal])] + return None if value is None else int(value) + + def attention_shard(self, internal: int) -> ShardDesc: + """The share of a head-sharded attention buffer this rank holds: ``T`` heads over ``tp`` + ranks, rank ``r`` holds ``ceil(T / tp)`` from ``r * T // tp``. Where any rank's range starts + inside a share, each rank's buffer is a share of its own.""" + if self.tp_size == 1: + return ShardDesc() + heads, total = self.per_rank(internal), self.total(internal) + if heads <= 0 or total is None or total <= 0: + return ShardDesc(self.tp_size, self.tp_rank) + if any(rank * total // self.tp_size % heads for rank in range(self.tp_size)): + # Some rank's heads straddle two shares, so each rank's buffer is a share of its own. + return ShardDesc(self.tp_size, self.tp_rank) + # total // heads distinct shares; fewer heads than ranks are repeated, so a single head is + # one share that every rank holds. + count = max(total // heads, 1) + index = min((self.tp_rank * total // self.tp_size) // heads, count - 1) + return ShardDesc(count, index) + + def state_shard(self) -> ShardDesc: + return ShardDesc(self.tp_size, self.tp_rank) if self.tp_size > 1 else ShardDesc() + + +def _buffer_geometry( + declared: object, mapper: BufferMapper, size: int, heads: int, state: bool +) -> BufferGeometry: + """A buffer's head geometry: declared sections or head size, completed per buffer. + ``ValueError`` if the declared or completed geometry has a size that is not positive.""" + declared = BufferGeometry( + getattr(declared, "section_bytes", None), getattr(declared, "bytes_per_head", None) + ).validate() + if mapper in (BufferMapper.REPLICATED, BufferMapper.SECTIONED): + return declared + bytes_per_head = declared.bytes_per_head + if bytes_per_head is None and not state and heads > 0 and size % heads == 0: + bytes_per_head = size // heads + if bytes_per_head is None or size % bytes_per_head: + return declared + return BufferGeometry(declared.section_bytes, bytes_per_head, size // bytes_per_head).validate() + + +# A layer group is whole when every transferred buffer holds the same bytes on every rank +# (replicated, or one KV head shared by all ranks); else it is the share this rank holds. +def _group_shard(shards: Sequence[ShardDesc], fallback: ShardDesc) -> ShardDesc: + parts = {s for s in shards if s.count != 1} + if not parts: + return ShardDesc() + return parts.pop() if len(parts) == 1 else fallback + + +def derive_layout(manager: KVCacheManagerV2) -> ManagerLayout: + """The layout of a ``KVCacheManagerV2`` and its device pools. ``ValueError`` if its layer + groups, layer configs and declarations do not agree.""" + tokens_per_block = int(manager.tokens_per_block) + layer_cfg = _manager.layer_configs(manager) + grouping = _manager.layer_group_layer_ids(manager) + pg_of_lg = _manager.layer_group_pool_indices(manager) + if len(pg_of_lg) != len(grouping): + raise ValueError(f"{len(grouping)} layer groups but {len(pg_of_lg)} pool-group indices") + internal = sorted({lid for group in grouping for lid in group}) + global_of = dict(zip(internal, global_layer_ids(manager, internal))) + + state_groups = set() + for lg, layer_ids in enumerate(grouping): + if not layer_ids or layer_ids[0] not in layer_cfg: + raise ValueError(f"layer group {lg} has no layer config (layers {layer_ids})") + if _manager.is_recurrent_layer(layer_cfg[layer_ids[0]]): + state_groups.add(lg) + + # TODO: the base manager declares head-major K/V whichever attention backend writes its pages, + # so a manager paired with a token-major backend gets the head-major layout_id unless it + # declares its layout itself. + mappers = { + role: BufferMapper(int(kind)) for role, kind in _manager.role_mapper_kinds(manager).items() + } + default_mapper = mappers.get(_ALL_ROLE, BufferMapper.INDEXED) + role_layouts = _manager.role_layouts(manager) + ignored = _manager.ignored_roles(manager) + heads = _Heads(manager) + expansions = {lid: _expansions(cfg, tokens_per_block) for lid, cfg in layer_cfg.items()} + + device_pools: dict[int, tuple[DevicePool, ...]] = {} + buffers: list[BufferDesc] = [] + shards: dict[int, list[ShardDesc]] = {lg: [] for lg in range(len(grouping))} + for pg in _manager.pool_group_descriptors(manager): + g = int(pg.pool_group_index) + pools = tuple( + DevicePool( + g, int(p.pool_index), int(p.base_address), int(p.slot_bytes), int(pg.num_slots) + ) + for p in pg.pools + ) + device_pools[g] = tuple(sorted(pools, key=lambda p: p.index)) + # One variant per layer group drawing from this pool group; pool ``i`` of a slot holds the + # ``i``-th coalesced buffer, whose members sit back to back. + for variant in pg.slot_desc.variants: + lg = int(variant.layer_group_id) + state = lg in state_groups + for pool_index, coalesced in enumerate(variant.coalesced_buffers): + size = int(coalesced.single_buffer_size) + for j, buffer_id in enumerate(coalesced.buffer_ids): + lid = int(buffer_id.layer_id) + role = str(buffer_id.role) + mapper = mappers.get(role, default_mapper) + transfer = role not in ignored + buffers.append( + BufferDesc( + layer=global_of.get(lid, lid), + role=role, + pool=(g, pool_index), + offset=j * size, + size=size, + mapper=mapper, + geometry=_buffer_geometry( + role_layouts.get(role), mapper, size, heads.per_rank(lid), state + ), + transfer=transfer, + expansion=expansions.get(lid, {}).get(role, 1), + ) + ) + if not transfer or mapper is BufferMapper.REPLICATED: + continue + shards[lg].append(heads.state_shard() if state else heads.attention_shard(lid)) + + layer_groups: list[LayerGroupDesc] = [] + for lg, layer_ids in enumerate(grouping): + first = layer_cfg[layer_ids[0]] + is_state = lg in state_groups + window = None + sink_tokens = 0 + if not is_state: + window = getattr(first, "window_size", None) + window = None if window is None else int(window) + sink_tokens = int(getattr(first, "num_sink_tokens", None) or 0) + # Layer groups with equal slot sizes may share a pool group, each drawing its own slots. + g = pg_of_lg[lg] + if g not in device_pools: + raise ValueError(f"layer group {lg} maps to pool group {g}, which has no pools") + layer_groups.append( + LayerGroupDesc( + kind="state" if is_state else "attention", + window=window, + sink_blocks=-(-sink_tokens // tokens_per_block), + pool_group=g, + layers=tuple(sorted(global_of[lid] for lid in layer_ids)), + shard=_group_shard(shards[lg], ShardDesc(heads.tp_size, heads.tp_rank)), + ) + ) + + layout = Layout( + tokens_per_block=tokens_per_block, + pool_groups={g: tuple(p.slot_bytes for p in pools) for g, pools in device_pools.items()}, + layer_groups=tuple(layer_groups), + buffers=tuple(buffers), + ) + return ManagerLayout( + layout=layout, + tokens_per_block=tokens_per_block, + pool_group_of=tuple(pg_of_lg), + windows=tuple(desc.window for desc in layer_groups), + sink_blocks=tuple(desc.sink_blocks for desc in layer_groups), + recurrent=tuple(desc.kind == "state" for desc in layer_groups), + sparse=tuple(_sparse(layer_cfg[layer_ids[0]]) for layer_ids in grouping), + layers=tuple(desc.layers for desc in layer_groups), + shards=tuple((desc.shard.count, desc.shard.index) for desc in layer_groups), + page_bytes={g: sum(p.slot_bytes for p in pools) for g, pools in device_pools.items()}, + device_pools=device_pools, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_lender.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_lender.py new file mode 100644 index 000000000000..afeb89f5b4c4 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_lender.py @@ -0,0 +1,1582 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The host-staging implementation behind the mechanism-independent lending API. + +Relays whole blocks between a request's device pages and host staging slots. + +Leases needing slots are granted first come, first served; a lease with no block entries is +ready at its first poll, with no place in line. Lender and lease methods run only on the +manager's thread. Backends never lend or call a lease; their own threads read a view, access the +memory it points to, call no lease or lender method, and tell the holder through their own +channel. + +Copies queue on the manager's execution stream, serial with the forward on the GPU: a step +queuing copies takes about their time longer. With page-locked staging and the fresh-page fill +off, no lend, poll, mark or resume_bounds call waits for them on the CPU; the manager's shutdown +does, and so does a process's first cache close, which waits for the whole device. Under +confidential computing staging is pageable, and a copy can hold the call queuing it until the +stream reaches it. + +The executor's pool rebalance suspends the caches of its active requests, parked ones among +them, and moves their pages, lent ones included: queued copies complete first, a read still +waiting for slots fails at its grant where its pages moved, and a write marked after the +rebalance copies only the block entries whose pages stayed, all that ``resume_bounds`` counts. + +Caller obligations: + - A staging backend, on the manager's thread, takes a hold with ``hold_regions()`` and + registers ``regions`` after the attach and before it touches a lease; it deregisters + before the manager shuts down, and releases the hold once deregistration is confirmed. + - Each iteration of the executor loop the holder polls every open lease and the waiter asks + ``resume_bounds`` for every waiting request, also when no new request arrives. + - At shutdown the executor loop stops; staging backends stop, deregister and release their + holds; the holder marks and releases what they give back; the manager shuts down last. + - A request whose one-model draft has a joint-reuse pool of its own shares its context + cursor with that pool: the caller fetches the same range into both managers, each + through its own lender, resumes within both intervals and publishes both. A fetch into + one alone leaves the other without the prefix while the shared cursor moves past it. + - While requests are parked in a state other than a disaggregated transfer state, the caller + keeps room for the running requests in each pool group: the pages parked requests' caches + lock, the other pages the scheduler cannot reclaim, and, for the generation request that + needs the most, the pages of its cache within its windows, which a resume locks all at + once, and the pages its next step adds stay within ``max_util_for_resume`` of the group's + GPU pages. The pages a fetch grows are outside the V2 scheduler's reach: it neither + evicts, pauses nor preempts a parked request. Its deadlock check counts no pass while any + request is in a disaggregated transfer state, as one parked in the transfer-in-progress + state is, and otherwise raises "V2 scheduler deadlock" after 1000 passes in a row that + schedule and reclaim nothing while a context or generation request waits, so without that + room a request that cannot be admitted, resume or grow can end in that deadlock. A share + of each pool group for the parked fetches alone does not ensure it, with or without a + cache tier below the GPU. The manager's ``get_page_indices_by_layer_group`` lists a + request's pages per layer group, one beam's (under beam search each further beam also + locks a page of its own for every block not wholly inside the prompt), and the runtime's + ``impl.pool_group_descs`` give each pool group's number of GPU pages and its layer groups. + - Under attention data parallelism without a cache transceiver, the caller keeps a rank from + parking all its active requests, or accepts the stall: the executor counts parked requests + as schedulable, so a rank whose active requests are all parked, at its cap of active + requests or without pages for a padding dummy, schedules nothing, and its empty batch + holds every rank's forward until a fetch settles. + - An integrator adding a writer of recycled pages off the manager's stream makes it wait on + that stream first, as the disaggregated receive does: such a writer is not ordered after + staging copies, which precede only later work on that stream, the work a page's new owner + waits for. With a KV cache connector refused, the manager's only such writer is the + fresh-page fill (``TRTLLM_KV_FRESH_PAGE_FILL``, a diagnostic off by default), which + waits for the manager stream before filling and synchronizes the fill stream afterward: + it overwrites no page a queued copy still reads, and a ``lend_write`` that grows the cache + then waits on the CPU for the queued staging copies. + - The caller releases, or frees the request of, a lease that waits too long: a slot lost to + a failed copy is never reused, and a slot is lost when a copy into or out of it may have + been queued without its completion recorded, because recording that completion failed or a + read's grant failed while queuing the copy. Such failures erode staging capacity: a lease + needing, in a pool group, a longer run of slots than any there without a lost slot waits + in line for good without failing, holding up every lease behind it until it is released or + its request is freed. A lost slot also keeps the staging memory until process exit. + +""" + +from __future__ import annotations + +import collections +import enum +import traceback +import weakref +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Deque, Mapping, Sequence + +import numpy as np +import torch +from cuda.bindings import driver as drv +from cuda.bindings import runtime as cudart + +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.logger import logger + +from . import _manager, _ranges +from ._fetch_ledger import _FetchLedger +from ._identity import Identity +from ._layout import ManagerLayout, derive_layout, layout_id +from ._slots import SlotQueue, SlotRuns, slot_counts +from ._types import LayerGroupBlocks, LeaseView, MemoryRegion, ResumeBounds, StagingOptions + +if TYPE_CHECKING: + from tensorrt_llm.runtime.kv_cache_manager_v2 import _KVCache + + from ...llm_request import LlmRequest + from ..kv_cache_manager_v2 import KVCacheManagerV2 + from ._fetch_ledger import _FetchRecord + +# Staging memory and page-index buffers kept until the process exits, by identity, oldest first; +# only a clean shutdown removes an entry. Each change is one dict operation, atomic under the GIL, +# so lenders of managers on different threads need no lock. +_retained_until_exit: dict[int, object] = {} + +_REASON_SHUT_DOWN = "the KV cache manager shut down" +_REASON_REUSE_RESET = ( + "the KV cache manager reset its reuse state, after which nothing is lent by name" +) +_REASON_STOPPED = ( + "the lender stopped after an error it logged while the KV cache manager freed a request" +) +_REASON_STOPPED_AT_SHUTDOWN = ( + "the lender stopped after an error it logged while the KV cache manager shut down" +) +_REASON_SUSPENDED = "the request's cache is suspended" +_REASON_PLACEHOLDERS = ( + "the request carries multimodal data without digests, which its object_keys do not cover" +) +_REASON_ENCODER = "the request has encoder input, which its object_keys do not cover" +_REASON_SCRATCH = ( + "the request's cache has SWA scratch reuse on: each capacity change keeps its history within " + "the scratch rewind of the old capacity, which a fetch's grow breaks" +) +_REASON_CONTEXT_OUTPUTS = ( + "the request returns context logits or asks for additional model outputs, which the executor " + "gives only for computed positions and keeps across a rewind and a recompute pause" +) + + +def _retained() -> tuple[object, ...]: + """What is kept until exit, oldest first; for tests.""" + return tuple(_retained_until_exit.values()) + + +def _retain_until_exit(owner: object) -> None: + """Keep ``owner`` until the process exits, or until ``_stop_retaining``.""" + _retained_until_exit[id(owner)] = owner + + +def _stop_retaining(owner: object) -> None: + """Drop ``owner`` from the keep list, compared by identity.""" + _retained_until_exit.pop(id(owner), None) + + +class _HostMemory: + """Host memory of exactly ``nbytes``, page-locked where pinning pays off. Pinned memory goes + only through ``free``, once; pageable memory goes with the object, which the keep list holds + until exit, so memory kept until exit stays mapped either way.""" + + def __init__(self, nbytes: int) -> None: + self.nbytes = nbytes + self._pageable: np.ndarray | None = None + if prefer_pinned(): + # The driver pins the size asked; torch's pinned allocator rounds it to a power of two. + error, address = cudart.cudaHostAlloc(nbytes, cudart.cudaHostAllocDefault) + if error != cudart.cudaError_t.cudaSuccess: + raise MemoryError(f"pinning {nbytes} bytes of staging memory failed: {error}") + self.address = int(address) + else: + self._pageable = np.empty(nbytes, dtype=np.uint8) + self.address = int(self._pageable.ctypes.data) + + def free(self) -> None: + """Release the memory; later calls do nothing.""" + address, self.address = self.address, 0 + if not address or self._pageable is not None: + self._pageable = None + return + (error,) = cudart.cudaFreeHost(address) + if error != cudart.cudaError_t.cudaSuccess: + logger.warning(f"KV cache lender: freeing the staging memory failed: {error}") + + +def _allocate_retained(nbytes: int) -> _HostMemory: + """One host allocation of ``nbytes`` for the staging regions, put on the keep list.""" + # TODO: staging duplicates the manager's host tier, which backends cannot use in its place: its + # host pools may move when they resize (mremap), the manager lends none of their pages, and its + # cold-page codec may hold a row in a form other than the row's device bytes. + memory = _HostMemory(max(int(nbytes), 1)) + _retain_until_exit(memory) + return memory + + +def _checked_layout(manager: KVCacheManagerV2) -> ManagerLayout: + """The attach's checks of the manager: a v2 manager, no context or pipeline parallelism, no + lender yet, no recurrent state, no sparse buffer.""" + _manager.require_v2(manager) + mapping = _manager.mapping(manager) + # TODO: the layout derives only tensor-parallel shards, so context-parallel ranks, Helix's + # included, would give different pages the same object_keys. + cp_size = int(mapping.cp_size) + if cp_size > 1: + raise ValueError( + f"context parallelism (cp_size={cp_size}) is not supported: its ranks would give " + "different pages the same object_keys" + ) + # TODO: each pipeline stage's staging slots and windows are its own, so one call can raise on + # one stage and lend on another. + pp_size = int(mapping.pp_size) + if pp_size > 1: + raise ValueError( + f"pipeline parallelism (pp_size={pp_size}) is not supported by staging: each stage's " + "slots and windows are its own, so one call could raise on one stage and lend on " + "another" + ) + if _manager.attached_lender(manager) is not None: + raise ValueError("a lender is already attached to this KV cache manager") + if _manager.has_lossy_cold_page_codec(manager): + raise ValueError( + "the KV cache lender cannot publish original KV bytes by name with a lossy " + "cold-page codec (quantization_for_cold_page)" + ) + layout = derive_layout(manager) + recurrent = [lg for lg, state in enumerate(layout.recurrent) if state] + if recurrent: + raise ValueError( + f"layer groups {recurrent} hold recurrent state, which the lender does not lend" + ) + # TODO: a sparse layer group's page index can name a host slot, and its locked pages move + # between GPU and host while a lease holds their indices: a decoding cache offloads its + # history, and a prefill lock promotes a host page back to the GPU. + sparse = [lg for lg, flag in enumerate(layout.sparse) if flag] + if sparse: + raise ValueError( + f"layer groups {sparse} hold sparse buffers, whose read-only pages a cache can lock " + "in host memory" + ) + return layout + + +def _check_commits(manager: KVCacheManagerV2) -> None: + """``ValueError`` for a manager that commits no blocks: a publish lends only committed ones.""" + if not _manager.commits_blocks(manager): + raise ValueError( + "staging needs a manager that commits blocks: block reuse on, and joint reuse for a " + "draft manager; with block reuse off nothing could ever be published" + ) + + +def _check_lookahead(manager: KVCacheManagerV2) -> None: + """``ValueError`` for a manager whose blocks may depend on tokens past their end: a name covers + the tokens up to the block's end, so two requests sharing them could hold different bytes. A + read-ahead upstream has not established, as one-model DraftTarget's, counts as none.""" + # TODO: staging serves no manager built for a draft that reads ahead (Eagle, MTP): a block's + # name covers only the tokens up to its end, so requests sharing them could hold different + # draft bytes under one name. + lookahead = _manager.prompt_lookahead(manager) + if lookahead > 0: + # The target manager of a draft whose layers live elsewhere is refused too: that draft's own + # pool is, so a fetch could fill the target's blocks but never the draft's. + raise ValueError( + "staging needs blocks that depend on no token past their end; this manager is built " + f"for a one-model draft that reads {lookahead} prompt tokens ahead (Eagle or MTP), so " + "requests sharing a block's tokens could hold different bytes under one name" + ) + + +def _check_connector(manager: KVCacheManagerV2) -> None: + """``ValueError`` for a manager with a KV cache connector.""" + # TODO: a KV cache connector serves a prefix at the first context chunk, from the committed + # tokens, so it cannot lower a history a windowed fetch moved, and it loads off the manager's + # stream, unordered with staging copies. + if _manager.kv_connector(manager) is not None: + raise ValueError( + "staging does not serve a manager with a KV cache connector: the connector serves a " + "request's prefix at its first context chunk, measured from the committed tokens, and " + "cannot lower a history a fetch moved" + ) + + +def _attach_staging( + manager: KVCacheManagerV2, *, scope: bytes, staging: StagingOptions, cls: type | None = None +) -> Staging: + """``attach_lender`` with the lender class as a parameter (``Staging`` when ``None``), so a + test can attach a subclass that breaks one rule.""" + layout = _checked_layout(manager) + _check_commits(manager) + _check_lookahead(manager) + _check_connector(manager) + if not isinstance(scope, bytes): + raise TypeError(f"scope must be bytes, got {type(scope).__name__}") + if not isinstance(staging, StagingOptions): + raise TypeError(f"staging must be StagingOptions, got {type(staging).__name__}") + counts = slot_counts(layout, staging) + identity = Identity(scope, layout_id(layout.layout), layout.layers, layout.shards) + slots = {g: int(counts.get(g, 0)) for g in layout.pool_groups} + sizes = {g: slots[g] * int(layout.page_bytes[g]) for g in layout.pool_groups} + # TODO: an exception after the allocation keeps the staging memory, and the page-index buffer + # once kept, until the process exits. + memory = _allocate_retained(sum(sizes.values())) + base = memory.address + regions = [] + offset = 0 + for g in layout.pool_groups: + groups = tuple(lg for lg, pg in enumerate(layout.pool_group_of) if pg == g) + layout_label = identity.region_label(groups) + regions.append( + MemoryRegion( + layout_label, base + offset, sizes[g], int(layout.page_bytes[g]), slots[g], groups + ) + ) + offset += sizes[g] + lender = (cls or Staging)( + weakref.ref(manager), + layout, + identity, + tuple(regions), + SlotQueue(slots), + weakref.ref(memory), + ) + lender._retain_index_buffer(manager) + logger.info( + f"KV cache lender: namespace {identity.namespace.hex()}, staging {offset >> 20} MiB " + f"in {len(regions)} regions" + ) + # Installed last, so a failure above leaves nothing attached. + _manager.install_lender(manager, lender) + return lender + + +class _QueuedCopy: + """Copies queued together, complete once their event says so, and never without one. ``done`` + asks the event at most once per round of the lender's progress, and never again once it + reported completion.""" + + def __init__(self, event: torch.cuda.Event | None) -> None: + self._event = event + self._done = False + self._asked_in_round: int | None = None + + def done(self, round_: int | None = None) -> bool: + """Whether the copies have completed, without waiting; ``round_=None`` always asks the + event, if there is one.""" + if self._event is None: + return False + if not self._done and (round_ is None or round_ != self._asked_in_round): + self._asked_in_round = round_ + self._done = bool(self._event.query()) + return self._done + + def wait(self) -> None: + """Block the host until the copies have completed; without an event, return at once.""" + if not self._done and self._event is not None: + self._event.synchronize() + self._done = True + + +# Copies that may be queued with no record of their completion: a read's from the grant of its +# slots until its copy's record, and a quarantined lease's. +_COPY_WITHOUT_RECORD = _QueuedCopy(None) + + +def _no_cache_reason(request_id: int) -> str: + return f"request {request_id} has no KV cache" + + +@dataclass(eq=False) +class _LentBlockEntries: + """Block entries grouped by local layer group, with aligned ordinals and slot arrays. + + Outer position selects ``layer_group_indices``; each inner position selects one block's + ordinal, device page index and, once assigned, staging slot index. + """ + + layer_group_indices: list[int] + block_ordinals: list[np.ndarray] + device_pages: list[np.ndarray] + staging_slots: list[np.ndarray] = field(default_factory=list) + + @property + def num_block_entries(self) -> int: + return sum(len(o) for o in self.block_ordinals) + + +class Staging: + """``Lender`` over one manager, referenced weakly. What its calls do depends on its mode. + + - Serving the manager: a lend past its range check, a poll and a resume_bounds call return + recyclable leases' slots and grant waiting leases in order before their own work, a mark and a + release after it. + - After the manager's reuse reset: as serving, but a lend that raises no ``ValueError`` fails + at the call. + - Stopped for good by an error its free or shutdown hook logs: a lend that raises no + ``ValueError`` fails at the call, a poll fails a lease whose view no poll returned, a mark + copies nothing, and the shutdown keeps the staging memory; while the manager is served, + progress fails the waiting leases and returns no slot, and resume_bounds gives a request with + a cache an empty interval. + - The manager shut down or gone: a lend that raises no ``ValueError`` fails at the call, a poll + runs no progress and returns the view only once ready, a mark copies nothing and resume_bounds + raises ``ValueError``; a stopped lender's poll still fails a lease whose view no poll + returned. + """ + + def __init__( + self, + manager: weakref.ref, + layout: ManagerLayout, + identity: Identity, + regions: tuple[MemoryRegion, ...], + slots: SlotQueue, + memory: weakref.ref, + ) -> None: + self._manager_ref = manager + self._layout = layout + self._identity = identity + self._regions = tuple(regions) + self._slot_queue = slots + # Only the keep list holds the staging memory strongly; this finds it there at shutdown. + self._memory_ref = memory + self._region_index_of_pool_group = {g: i for i, g in enumerate(layout.pool_groups)} + self._any_window = any(w is not None for w in layout.windows) + # The fetch records and the resume_bounds they give. The ledger asks this lender whether a + # copy completed through a weak reference, so no reference cycle keeps the lender, or a + # cache its records hold, past the lender's last reference. + lender = weakref.ref(self) + self._ledger = _FetchLedger(layout, lambda copy: lender()._copy_done(copy)) + # Leases waiting for slots, in order. + self._waiting_for_slots: Deque[_StagingLease] = collections.deque() + self._leases_holding_slots: list[_StagingLease] = [] # granted, slots not yet returned + # Open leases, held strongly: an open lease keeps the staging memory at shutdown even + # once its holder has dropped it. + self._open_leases: set[_StagingLease] = set() + # Open staging holds, held strongly too: a dropped hold still keeps the memory. + self._open_holds: set[_StagingHold] = set() + self._quarantined: list[SlotRuns] = [] # slots a failed copy may still touch; never reused + self._shut_down = False + # Whether the manager reset its reuse state; once it has, nothing is lent by name again. + self._reuse_was_reset = False + self._stop_reason: str | None = None # why lending stopped for good; the first reason stays + self._index_buffer: object | None = None # kept until the runtime's shutdown has returned + self._progress_round = 0 # rounds of progress, so each asks a copy's event at most once + + @property + def regions(self) -> tuple[MemoryRegion, ...]: + """See ``Lender.regions``.""" + return self._regions + + def hold_regions(self) -> _StagingHold: + """See ``Lender.hold_regions``.""" + if self._serving_manager() is None: + # The memory was freed or kept at the shutdown; a later hold changes neither. + return _StagingHold(None) + hold = _StagingHold(self) + self._open_holds.add(hold) + return hold + + def lend_read(self, request: LlmRequest, start: int, end: int) -> _StagingLease: + """See ``Lender.lend_read``.""" + start, end = self._whole_blocks(start, end) + request_id = int(request.py_request_id) + manager = self._serving_manager() + if manager is None: + return _StagingLease._failed_at_call(self, "read", request_id, _REASON_SHUT_DOWN) + self._progress() + if self._stop_reason is not None: + return _StagingLease._failed_at_call(self, "read", request_id, self._stop_reason) + if self._reuse_was_reset: + return _StagingLease._failed_at_call(self, "read", request_id, _REASON_REUSE_RESET) + kv = _manager.kv_of(manager, request_id) + if kv is None: + return _StagingLease._failed_at_call( + self, "read", request_id, _no_cache_reason(request_id) + ) + unnamed = self._unnamed_input_refusal(request) + if unnamed is not None: + return _StagingLease._failed_at_call(self, "read", request_id, unnamed) + state = _manager.cache_state(kv) + if end > state.committed: + raise ValueError( + f"the range ends at {end}, past the {state.committed} committed tokens" + ) + if not state.active: + return _StagingLease._failed_at_call(self, "read", request_id, _REASON_SUSPENDED) + # TODO: a publish leaves out window blocks the request's own window has passed, although + # the prefix tree may still hold their committed pages, so a fetch whose window still keeps + # such a block finds it missing. + entries = self._lendable_block_entries( + manager, state.history, self._block_entries_of_range(manager, kv, start, end) + ) + counts = self._required_slots_per_pool_group( + [len(block_ordinals) for block_ordinals in entries.block_ordinals] + ) + self._check_fits(counts) + reuse_keys = self._reuse_keys_for(manager, request, kv, entries.block_ordinals) + lease = _StagingLease(self, "read", request_id, kv, entries, reuse_keys) + self._admit(lease, counts) + return lease + + def lend_write(self, request: LlmRequest, start: int, end: int) -> _StagingLease: + """See ``Lender.lend_write``.""" + start, end = self._whole_blocks(start, end) + request_id = int(request.py_request_id) + manager = self._serving_manager() + if manager is None: + return _StagingLease._failed_at_call(self, "write", request_id, _REASON_SHUT_DOWN) + self._progress() + layout = self._layout + tpb = int(layout.tokens_per_block) + # The request computes its last prompt token itself, for its logits, so a fetch ends at + # the whole blocks before it. + prompt = _manager.prompt_length(request) + if end > (prompt - 1) // tpb * tpb: + raise ValueError( + f"the range ends at {end}, past the whole blocks before the request's last prompt " + f"token ({prompt} prompt tokens)" + ) + # Block entries and their keys come from the layout for a history of ``end``, so every + # ValueError is raised before the cache changes. The entries are counted from their ranges + # and checked to fit before any array is built. + runs = [ + _ranges._needed_ranges(manager, layout, lg, start // tpb, end // tpb, end) + for lg in range(layout.num_layer_groups) + ] + counts = self._required_slots_per_pool_group( + [sum(e - b for b, e in lg_runs) for lg_runs in runs] + ) + self._check_fits(counts) + block_ordinals = [_ranges._ordinals(lg_runs) for lg_runs in runs] + # The checks above read only the call, the request's prompt and the layout, so on a + # manager not shut down a wrong call raises whatever this rank holds for the request. + if self._stop_reason is not None: + return _StagingLease._failed_at_call(self, "write", request_id, self._stop_reason) + if self._reuse_was_reset: + return _StagingLease._failed_at_call(self, "write", request_id, _REASON_REUSE_RESET) + kv = _manager.kv_of(manager, request_id) + if kv is None: + return _StagingLease._failed_at_call( + self, "write", request_id, _no_cache_reason(request_id) + ) + unnamed = self._unnamed_input_refusal(request) + if unnamed is not None: + return _StagingLease._failed_at_call(self, "write", request_id, unnamed) + state = _manager.cache_state(kv) + if start < (state.committed // tpb) * tpb: + raise ValueError( + f"the range starts at {start}, inside the committed whole blocks of " + f"{state.committed} tokens" + ) + reuse_keys = self._reuse_keys_for(manager, request, kv, block_ordinals) + if self._any_window and end < state.history: + raise ValueError( + f"the range ends at {end}, below the history of {state.history} tokens its " + "windows keep" + ) + split = self._bidirectional_split_refusal(request, end) + if split is not None: + return _StagingLease._failed_at_call(self, "write", request_id, split) + if not state.active: + return _StagingLease._failed_at_call(self, "write", request_id, _REASON_SUSPENDED) + if self._scratch_reuse_on(kv): + return _StagingLease._failed_at_call(self, "write", request_id, _REASON_SCRATCH) + if self._returns_context_outputs(request): + return _StagingLease._failed_at_call(self, "write", request_id, _REASON_CONTEXT_OUTPUTS) + # After every ValueError: whether an earlier fetch settled is this rank's own timing. + unsettled = self._ledger._unsettled_fetch_refusal(request_id, kv) + if unsettled is not None: + return _StagingLease._failed_at_call(self, "write", request_id, unsettled) + # With a window the history moves to ``end``, so windows need pages only for the blocks a + # history of that length reads. + position = end if self._any_window else state.history + refusal = self._history_refusal(manager, state.history, position, end) + if refusal is None: + refusal = self._unwritten_pages_refusal(manager, request_id, kv, position) + if refusal is not None: + return _StagingLease._failed_at_call(self, "write", request_id, refusal) + # Read before the grow, which moves a windowed cache's history to ``end``. + kept = self._ledger._prefix_kept_by_fetch(request_id, kv, start) + # TODO: the V2 scheduler cannot reclaim the pages a parked fetch grew. + # TODO: route capture stops reading a request's prepopulated length once it holds the routes + # below it, until the request finishes, so a later resume leaves positions without routes. + # TODO: with extra KV tokens the grow locks a page past the history it leaves, which nothing + # writes, so a later all-reusable windowed lease that leaves that block behind fails. + # TODO: an exception once the grow resized the cache, in its fill or a step below, can leave + # no record of a windowed cache's moved history, so resume_bounds counts tokens never fetched as + # computed, or a fetch record nothing settles, so resume_bounds stays None. + if not _manager.grow(manager, request, kv, position, end): + return _StagingLease._failed_at_call( + self, "write", request_id, f"no free pages to grow the cache to {end} tokens" + ) + # The cache has grown, and resume_bounds accounts for it whatever this lease's outcome. + fetch = self._ledger._record_fetch(request_id, kv, start, kept, position, state.committed) + entries = self._block_entries_of_range(manager, kv, start, end) + doomed = None + for lg, lg_ordinals, lg_slots in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + if doomed is None and np.any(lg_slots < 0): + missing = lg_ordinals[lg_slots < 0].tolist() + doomed = f"layer group {lg}: blocks {missing[:8]} have no page" + lease = _StagingLease(self, "write", request_id, kv, entries, reuse_keys, fetch) + if doomed is not None: + # It fails at its first poll, which abandons the fetch. + lease._first_poll_failure = doomed + self._open_leases.add(lease) + return lease + self._admit(lease, counts) + return lease + + def resume_bounds(self, request: LlmRequest) -> ResumeBounds | None: + """See ``Lender.resume_bounds``.""" + request_id = int(request.py_request_id) + manager = self._serving_manager() + if manager is None: + raise ValueError(f"{_no_cache_reason(request_id)}: {_REASON_SHUT_DOWN}") + self._progress() + kv = _manager.kv_of(manager, request_id) + if kv is None: + raise ValueError(_no_cache_reason(request_id)) + history = _manager.cache_state(kv).history + if self._stop_reason is not None: + # Empty: no fetch through a stopped lender counts. + return ResumeBounds(max_position=0, min_position=max(history, 1)) + return self._ledger._resume_bounds(request, request_id, manager, kv, history) + + def _on_free(self, request_id: int, kv_cache: _KVCache) -> None: + """Manager hook, after the request's cache left the map: its waiting leases fail and its + fetch records go. An error in that work, which it logs, stops the lender for good.""" + try: + if self._serving_manager() is None: + return + # Granted leases go on: a read's copy is queued already, and a write's marks copy + # nothing into pages other than the ones lent. + self._fail_waiting(int(request_id), "the request was freed") + self._ledger._drop_records(request_id) + self._progress() + except Exception: + # Logged, not raised into the manager's free, which then closes the cache as usual. The + # error can leave the records disagreeing with the slots where no caller sees it, so the + # lender stops; the reason is set first, by an assignment that cannot raise. + if self._stop_reason is None: + self._stop_reason = _REASON_STOPPED + logger.error(f"KV cache lender: freeing request {request_id}: {traceback.format_exc()}") + + def _on_shrink(self, request_id: int, kv_cache: _KVCache) -> None: + """Manager hook, right after the request's cache may have shrunk in place: delivered entries + past its blocks stop counting, since the shrink freed their pages. Logs its own errors.""" + try: + delivered = self._ledger._deliveries_of(int(request_id)) + if ( + self._serving_manager() is not None + and delivered is not None + and delivered.kv is kv_cache + ): + self._ledger._invalidate_delivered_blocks_past_capacity(delivered, kv_cache) + except Exception: + # Logged, not raised into the manager's resize. All the request's delivered entries go: + # the cache may grow back before resume_bounds finds the freed ones past its blocks. + self._ledger._drop_deliveries(int(request_id)) + logger.error( + f"KV cache lender: shrink of request {request_id}: {traceback.format_exc()}" + ) + + def _on_reset(self) -> None: + """Manager hook after a reuse reset, with every cache closed as an in-place weight update + ends: object_keys stop meaning the bytes computed, so none is lent by name again; waiting + leases + and fetches went at their requests' free; granted publishes finish. Never raises.""" + # TODO: after the manager resets its reuse state the lender lends nothing by name, since + # ``scope`` is fixed at the attach and nothing identifies the weights it computes with from + # then on. + self._reuse_was_reset = True + + def _on_shutdown(self) -> None: + """Manager hook, first in its shutdown, acting once: waits for the copies whose completion + it recorded, fails waiting leases, and frees the staging memory unless a lease or a hold is + open, a slot was lost to a failed copy or the lender stopped. An error in that work, which + it logs, stops the lender for good.""" + if self._shut_down: + return + try: + # With page-locked staging and the fresh-page fill off, the lender's only host + # wait: no more work comes on the stream. + for lease in self._leases_holding_slots: + if lease._copy is not None: + lease._copy.wait() + self._progress_round += 1 + if self._stop_reason is None: # a stopped lender's records may disagree with its slots + self._recycle() + for lease in list(self._waiting_for_slots): + self._fail(lease, _REASON_SHUT_DOWN) + for lease in [ + lease for lease in self._open_leases if lease._first_poll_failure is not None + ]: + self._fail(lease, lease._first_poll_failure) + self._shut_down = True + if not self._must_retain_memory(): + # Every copy on it has completed above, and no lease or hold reaches it any more. + memory = self._memory_ref() + _stop_retaining(memory) + memory.free() + else: + logger.warning( + f"KV cache lender: keeping {sum(p.nbytes for p in self._regions) >> 20} MiB of " + f"staging memory until exit: {len(self._open_leases)} leases open, " + f"{len(self._open_holds)} staging holds open, " + f"{len(self._quarantined)} slot runs lost to failed copies" + f"{'' if self._stop_reason is None else ', and ' + self._stop_reason}" + ) + except Exception: + # Logged, not raised into the manager's shutdown; the staging memory then stays in + # the keep list until exit. As in the free hook, the records may disagree with the + # slots, so the lender stops before the log line, keeping any earlier stop's reason. + self._shut_down = True + if self._stop_reason is None: + self._stop_reason = _REASON_STOPPED_AT_SHUTDOWN + logger.error(f"KV cache lender: shutting down: {traceback.format_exc()}") + + def _on_caches_closed(self) -> None: + """Manager hook, once the runtime's shutdown has returned, after which no cache writes into + the page-index buffer, so the lender lets it go. A shutdown that raises skips this call, and + the buffer stays until a retried one returns, or until exit.""" + if self._index_buffer is None: + return + self._stop_retaining_index_buffer() + + # One rule per method, so a test subclass that breaks exactly one rule overrides one method. + + def _progress(self) -> None: + """Return the slots of recyclable leases, then grant waiting leases in order; a stopped + lender does neither and fails the waiting leases.""" + if self._serving_manager() is None: + return + if self._stop_reason is not None: + for lease in list(self._waiting_for_slots): + self._fail(lease, self._stop_reason) + return + self._progress_round += 1 + self._recycle() + self._grant_waiting() + + def _copy_done(self, copy: _QueuedCopy | None) -> bool: + """``copy`` has completed (or there is none), its event asked at most once this round.""" + return copy is None or copy.done(self._progress_round) + + def _read_copy_done(self, lease: _StagingLease) -> bool: + """A read lease's copy into its slots has completed.""" + return self._copy_done(lease._copy) + + def _recyclable(self, lease: _StagingLease) -> bool: + """The lease's slots may return: no backend access possible and no copy on them pending. + The copy is asked last, so a lease still lent costs no event query.""" + # A failed lease was never ready, so no backend has seen its slots. + if lease._holder_state() is not _HolderState.FAILED: + if not lease._released or lease._marks_due(): + return False + # TODO: every lender call asks each released lease's pending copy again, so N calls while P + # copies pend cost N*P event queries. + return self._copy_done(lease._copy) + + def _locked_block_masks(self, kv: _KVCache | None, lease: _StagingLease) -> list[np.ndarray]: + """Per layer group, entries still in the same active cache's window and still locking the + GPU page lent; a page only held may sit on another tier under the same number.""" + entries = lease._block_entries + manager = self._manager_ref() + state = _manager.cache_state(kv) if kv is not None and kv is lease._kv else None + masks = [] + for lg, block_ordinals, slots in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + same = np.zeros(len(block_ordinals), dtype=bool) + if state is not None and state.active: + pages = _manager.locked_pages(kv, lg) + inside = block_ordinals < len(pages) + same[inside] = pages[block_ordinals[inside]] == slots[inside] + stale_beg, stale_end = _ranges._stale_range( + manager, self._layout, lg, state.history + ) + same &= (block_ordinals < stale_beg) | (block_ordinals >= stale_end) + masks.append(same) + return masks + + def _unnamed_input_refusal(self, request: LlmRequest) -> str | None: + """Why the request's KV depends on more than its object_keys cover, if it does: multimodal + data the manager keys by placeholder tokens alone, or encoder input.""" + if _manager.keyed_by_placeholders(request): + return _REASON_PLACEHOLDERS + if _manager.has_encoder_input(request): + return _REASON_ENCODER + return None + + # TODO: the lender cannot see the model's sliding window, so it also keeps fetch ends and + # resumes out of runs that window would cover whole. + def _bidirectional_split_refusal(self, request: LlmRequest, end: int) -> str | None: + """Why the fetch ends strictly inside a run of multimodal tokens the scheduler keeps within + one context chunk, if it does, whatever the run's length.""" + for b, e in _manager.bidirectional_runs(request): + if b < end < e: + return ( + f"the range ends at {end}, inside the multimodal tokens [{b}, {e}), a run the " + "scheduler keeps within one context chunk" + ) + return None + + def _scratch_reuse_on(self, kv: _KVCache) -> bool: + """A write target with SWA scratch reuse on, windowed or not: each capacity change keeps its + history within the scratch rewind of the old capacity, which a fetch's grow breaks, and a + window's next chunk overwrites scratch slots.""" + return _manager.has_scratch_reuse(kv) + + # TODO: the executor gives context outputs only for the positions a request computes, prompt + # logprobs pair the logits with the prompt from its second token wherever they start, and a + # rewind or a recompute pause keeps those held, so a request returning them takes no fetch. + def _returns_context_outputs(self, request: LlmRequest) -> bool: + """Whether the request returns context logits or asks for additional model outputs, which a + fetch leaves without the fetched positions, or after a context step gapped or repeated.""" + return _manager.returns_context_outputs(request) + + def _history_refusal( + self, manager: KVCacheManagerV2, history: int, position: int, end: int + ) -> str | None: + """Why the history the fetch leaves, ``position``, admits no resume at the positions + resume_bounds would count, or ``None``: in a cache without a window whose floor follows the + history, the history stands past ``end``.""" + if self._any_window: + return None # the fetch moves the history to ``end`` + if not _ranges._floor_follows_history(manager): + return None # the floor is a delivered fetch's lowest start, or the history if lower + if position > end: + return f"the cache's history of {history} tokens stands past the fetch's end, {end}" + return None + + def _unwritten_pages_refusal( + self, manager: KVCacheManagerV2, request_id: int, kv: _KVCache, position: int + ) -> str | None: + """Why moving the history to ``position`` would have the commit store bytes the request + never wrote, or ``None``: a window leaving behind a block whose page holds tokens past those + the request computed keeps that page for the commit (a partial match's copy, a page grown + early, a row an earlier lease missed), unless a fetch into the cache delivered its row.""" + if not _manager.keeps_stale_pages(manager): + return None + computed = self._ledger._computed_tokens(request_id, kv) + first = computed // int(self._layout.tokens_per_block) # the first block past them + delivered = self._ledger._deliveries_of(request_id) + fetched = ( + delivered.delivered_block_masks + if delivered is not None and delivered.kv is kv + else None + ) + for lg in range(self._layout.num_layer_groups): + stale_beg, stale_end = _ranges._stale_range(manager, self._layout, lg, position) + beg = max(first, stale_beg) + if beg >= stale_end: + continue + unwritten = _manager.locked_pages(kv, lg)[beg:stale_end] >= 0 + if fetched is not None: + # A delivered row's page holds the bytes the fetch copied in. + got = fetched[lg][beg : beg + len(unwritten)] + unwritten[: len(got)] &= ~got + blocks = (beg + np.nonzero(unwritten)[0]).tolist() + if blocks: + # TODO: no runtime call drops one block's page in one layer group, so the fetch + # fails where the commit could store no page for those blocks instead. + return ( + f"layer group {lg} leaves blocks {blocks[:8]} behind its window at {position} " + f"tokens; past the {computed} tokens the request computed, their pages hold " + "bytes it never wrote and no fetch delivered, which the commit would store" + ) + return None + + def _fail_waiting(self, request_id: int, reason: str) -> None: + """Fail the request's leases still waiting for slots.""" + for lease in [ + lease for lease in self._waiting_for_slots if lease._request_id == request_id + ]: + self._fail(lease, reason) + + def _object_keys(self, layer_group_index: int, reuse_keys: np.ndarray) -> np.ndarray: + """The entries' object_keys: ``uint8 (n, 54)`` for ``reuse_keys`` ``uint8 (n, 32)``.""" + return self._identity.object_keys(layer_group_index, reuse_keys) + + def _must_retain_memory(self) -> bool: + """A lease or a hold is unreleased, a slot is lost to a failed copy, or the lender stopped: + the memory stays.""" + return ( + bool(self._open_leases) + or bool(self._open_holds) + or self._slots_quarantined() + or self._stop_reason is not None + ) + + def _slots_quarantined(self) -> bool: + """Whether slots are out of use for good: runs ``_quarantine`` took, or those of a holding + lease whose copy has no record of its completion, which ``_recycle`` never returns.""" + return bool(self._quarantined) or any( + lease._copy_state() is _CopyState.WITHOUT_RECORD for lease in self._leases_holding_slots + ) + + def _retain_index_buffer(self, manager: KVCacheManagerV2) -> None: + """Keep the manager's page-index buffer until the runtime's shutdown has returned: a cache a + lease or a record holds writes its page indices there as it closes, also once the manager is + gone.""" + self._index_buffer = _manager.index_buffer(manager) + _retain_until_exit(self._index_buffer) + + def _stop_retaining_index_buffer(self) -> None: + """Once the runtime's shutdown has returned: no cache writes into the page-index buffer any + more.""" + _stop_retaining(self._index_buffer) + self._index_buffer = None + + def _on_hold_released(self, hold: _StagingHold) -> None: + """The hold's holder let go; ``KeyError`` for a hold not open, which its guard prevents.""" + self._open_holds.remove(hold) + + def _check_fits(self, counts: Mapping[int, int]) -> None: + """``ValueError`` if a lease needing ``counts`` entries per pool group can never be + granted.""" + try: + self._slot_queue.check_fits_capacity(counts) + except ValueError as error: + raise ValueError( + f"the range needs more staging slots than there are ({error}); ranges of at most " + "tokens_per_fetch tokens always fit" + ) from None + + def _free_slots(self, group: int) -> int: + """Free slots of a pool group; for tests.""" + return self._slot_queue.free_slots(group) + + def _open_lease_count(self) -> int: + """Leases not yet released; for tests.""" + return len(self._open_leases) + + # Lease records, slots and grants. + + def _serving_manager(self) -> KVCacheManagerV2 | None: + """The manager while the lender serves it; ``None`` once it shut down or is gone.""" + if self._shut_down: + return None + return self._manager_ref() + + def _whole_blocks(self, start: int, end: int) -> tuple[int, int]: + start, end = int(start), int(end) + tpb = int(self._layout.tokens_per_block) + if start < 0 or end < 0 or start > end: + raise ValueError(f"bad token range [{start}, {end})") + if start % tpb or end % tpb: + raise ValueError( + f"a staging lease covers whole blocks of {tpb} tokens, got [{start}, {end})" + ) + return start, end + + def _required_slots_per_pool_group(self, entries: Sequence[int]) -> dict[int, int]: + """Required slots per pool group, summed from each layer group's block entry count.""" + counts: dict[int, int] = {} + for lg, num_entries in enumerate(entries): + g = int(self._layout.pool_group_of[lg]) + counts[g] = counts.get(g, 0) + int(num_entries) + return counts + + def _admit(self, lease: _StagingLease, counts: Mapping[int, int]) -> None: + """Record a new lease and grant it now when no lease waits and its slots are free, or at + once, out of line, when it has no entries.""" + self._open_leases.add(lease) + if lease._block_entries.num_block_entries == 0: + # Nothing to stage: ready at the first poll, without slots or a place in line. + self._grant(lease, None) + return + lease._ticket = self._slot_queue.enqueue(counts) + self._waiting_for_slots.append(lease) + self._grant_waiting(just_admitted=lease) + + def _grant_waiting(self, just_admitted: _StagingLease | None = None) -> None: + """Grant the leases in line, strictly in order; ``just_admitted`` was looked up in this + call.""" + while self._waiting_for_slots: + head = self._waiting_for_slots[0] + # TODO: an exception while slots are taken, a MemoryError say, can leave slots taken for + # no lease, or drop the head's ticket with the head still in line, after which every + # round of progress raises "not waiting" until the head fails. + runs = self._slot_queue.try_grant(head._ticket) + if runs is None: + return + self._waiting_for_slots.popleft() + head._ticket = None + try: + self._grant(head, runs, recheck=head is not just_admitted) + except Exception as error: + # One broken grant fails only its own lease: it neither stalls the line nor raises + # out of another lease's call. Slots a copy may have been queued on stay unused. + if head._kind == "read" and head._slot_runs is not None: + self._quarantine(head) + else: + head._view = head._slot_runs = None + self._slot_queue.give_back(runs) + reason = f"granting staging slots failed: {error!r}" + if head is just_admitted and head._kind == "write": + # The cache grew for it, so it fails at its first poll, which abandons the + # fetch, as a lease missing pages does. + head._first_poll_failure = reason + else: + self._fail(head, reason) + logger.warning(f"KV cache lender: {reason}") + except BaseException: + # An interrupt can leave a copy queued on the slots without its completion. + head._slot_runs = runs + self._quarantine(head) + raise + + def _grant(self, lease: _StagingLease, runs: SlotRuns | None, recheck: bool = False) -> None: + """Give ``lease`` its slots and view; a read then queues its copy into them.""" + if recheck and lease._kind == "read": + problem = self._changed_source_failure(lease) + if problem is not None: + self._slot_queue.give_back(runs) + self._fail(lease, problem) + return + self._assign_staging_slots(lease._block_entries, runs) + if lease._kind == "read" and runs is not None: + # Lost until its copy is recorded: no raise from here on leaves it ready. + lease._copy = _COPY_WITHOUT_RECORD + lease._view = self._build_view(lease._block_entries, lease._reuse_keys) + if runs is None: + return + lease._slot_runs = runs + self._leases_holding_slots.append(lease) + if lease._kind != "read": + return + copy, error = self._queue_copy(self._copy_segments(lease._block_entries), to_staging=True) + if copy is None: + self._quarantine(lease) + if error is not None: + self._fail(lease, f"the copy into staging failed: {error}") + if copy is not None: + lease._copy = copy + + def _changed_source_failure(self, lease: _StagingLease) -> str | None: + """Why a read granted after waiting cannot copy the pages it looked up any more, if so.""" + manager = self._manager_ref() + kv = _manager.kv_of(manager, lease._request_id) + if kv is not lease._kv: + return "the request's cache was freed while the lease waited" + state = _manager.cache_state(kv) + if not state.active: + return "the request's cache was suspended while the lease waited" + entries = lease._block_entries + for lg, block_ordinals, slots in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + pages = _manager.locked_pages(kv, lg) + if np.any(block_ordinals >= len(pages)) or np.any(pages[block_ordinals] != slots): + return f"layer group {lg}: pages changed while the lease waited" + stale_beg, stale_end = _ranges._stale_range(manager, self._layout, lg, state.history) + if np.any((block_ordinals >= stale_beg) & (block_ordinals < stale_end)): + return f"layer group {lg}: blocks left the window while the lease waited" + return None + + def _fail(self, lease: _StagingLease, reason: str) -> None: + """Fail a lease never seen ready: it leaves the line, and a write abandons its fetch.""" + if lease.failure is not None: + return + lease._set_failure(reason) + if lease._ticket is not None: + self._slot_queue.cancel(lease._ticket) + lease._ticket = None + self._waiting_for_slots.remove(lease) + if lease._kind == "write": + self._ledger._abandon_fetch(lease) + + def _quarantine(self, lease: _StagingLease) -> None: + """Keep the lease's slots from reuse for good: a copy on them may still be running.""" + # The copy is lost and the runs recorded before the lease leaves the holding list, so a step + # that raises leaves its slots out of use and the staging memory kept at the shutdown. + lease._copy = _COPY_WITHOUT_RECORD + # TODO: quarantined slots are never reused, so failed copies erode staging capacity: a lease + # needing a longer run than any left between lost slots waits in line for good, holding up + # every lease behind it. + if lease._slot_runs is not None: + self._quarantined.append(lease._slot_runs) + self._leases_holding_slots = [ + held for held in self._leases_holding_slots if held is not lease + ] + logger.warning( + f"KV cache lender: a failed copy took staging slots of request {lease._request_id}" + ) + + # TODO: every lender call checks every lease holding slots, even with no copy pending, so + # polling H such leases once each costs about H*H checks. + def _recycle(self) -> None: + """Return the slots of every lease ``_recyclable`` allows.""" + # Every lease is judged before any slot returns: an event query that raises returns none + # and leaves every lease held for the next call. + recyclable = [self._recyclable(lease) for lease in self._leases_holding_slots] + # TODO: an exception while slots return, a MemoryError say, can leave a lease held after + # some or all of its slots returned, so every later round of progress raises "freed twice". + holding = [] + for lease, done in zip(self._leases_holding_slots, recyclable): + if done: + self._slot_queue.give_back(lease._slot_runs) + else: + holding.append(lease) + self._leases_holding_slots = holding + + def _on_lease_released(self, lease: _StagingLease) -> None: + """The lease's holder let go. The record changes first and progress runs last, so an + unexpected error leaves the fetch abandoned rather than unsettled.""" + self._open_leases.discard(lease) + if self._serving_manager() is None: + return + if lease._slot_state() is _SlotState.WAITING: + self._fail(lease, "released while waiting for staging slots") + elif lease._kind == "write" and lease._holder_state() is _HolderState.VIEW_PENDING: + # Released before anyone saw it ready: no backend wrote, and no marks are due. + lease._first_poll_failure = None + self._ledger._abandon_fetch(lease) + self._progress() + + def _apply_marks(self, lease: _StagingLease, masks: list[np.ndarray]) -> None: + """Copy the marked entries whose page is still the one lent into the request's pages. The + fetch stays abandoned until that copy is queued without error; then its entries add to the + cache's deliveries.""" + fetch = lease._fetch + current = self._ledger._take_out_fetch(lease, fetch) + kv = _manager.kv_of(self._manager_ref(), lease._request_id) + lent = self._locked_block_masks(kv, lease) + copy = [mask & still for mask, still in zip(masks, lent)] + error = None + if any(c.any() for c in copy): + segments = self._copy_segments(lease._block_entries, copy) + self._ledger._invalidate_delivered_blocks(lease, copy) + tail = self._covers_committed_tail(kv, lease._block_entries, copy) + queued = None + try: + queued, error = self._queue_copy(segments, to_staging=False) + lease._copy = queued + finally: + # A copy may be queued without its completion on the lease, also when the call + # raised or was interrupted before the lease took it. + if queued is None or lease._copy is not queued: + lease._copy = _COPY_WITHOUT_RECORD + self._quarantine(lease) + if tail and (queued is None or error is not None): + self._ledger._record_mixed_tail(lease._request_id, kv) + if current and error is None: + self._ledger._put_back_fetch(lease, fetch, copy) + self._progress() + + def _covers_committed_tail( + self, kv: _KVCache, entries: _LentBlockEntries, copied: list[np.ndarray] + ) -> bool: + """The copy writes the block the committed tokens end inside, whose committed tokens + resume_bounds counts without a delivery.""" + tpb = int(self._layout.tokens_per_block) + committed = _manager.cache_state(kv).committed + if committed % tpb == 0: + return False + tail = committed // tpb + return any(bool(np.any(o[m] == tail)) for o, m in zip(entries.block_ordinals, copied)) + + def _block_entries_of_range( + self, manager: KVCacheManagerV2, kv: _KVCache, start: int, end: int + ) -> _LentBlockEntries: + """The blocks of ``[start, end)`` a history of ``end`` reads, per layer group, with their + device slots (-1 where a block has no page).""" + layout = self._layout + tpb = int(layout.tokens_per_block) + entries = _LentBlockEntries([], [], []) + for lg in range(layout.num_layer_groups): + block_ordinals = _ranges._needed_ordinals( + manager, layout, lg, start // tpb, end // tpb, end + ) + pages = _manager.pages(kv, lg) + slots = np.full(len(block_ordinals), -1, dtype=np.int64) + inside = block_ordinals < len(pages) + slots[inside] = pages[block_ordinals[inside]] + entries.layer_group_indices.append(lg) + entries.block_ordinals.append(block_ordinals) + entries.device_pages.append(slots) + return entries + + def _lendable_block_entries( + self, manager: KVCacheManagerV2, history: int, entries: _LentBlockEntries + ) -> _LentBlockEntries: + """``entries`` without blocks that have no page or that the request's own window has passed + (their page may hold something else).""" + out = _LentBlockEntries([], [], []) + for lg, block_ordinals, slots in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + stale_beg, stale_end = _ranges._stale_range(manager, self._layout, lg, history) + ok = (slots >= 0) & ~((block_ordinals >= stale_beg) & (block_ordinals < stale_end)) + out.layer_group_indices.append(lg) + out.block_ordinals.append(block_ordinals[ok]) + out.device_pages.append(slots[ok]) + return out + + def _reuse_keys_for( + self, + manager: KVCacheManagerV2, + request: LlmRequest, + kv: _KVCache, + block_ordinals: Sequence[np.ndarray], + ) -> list[np.ndarray]: + """Per layer group, its block entries' reuse keys, ``uint8 (n, 32)``.""" + # TODO: every lease hashes the request's whole prefix again from block 0, so a lease late in + # a long prompt costs time in proportion to the prompt. + top = max((int(o.max()) + 1 for o in block_ordinals if len(o)), default=0) + reuse_keys = _manager.block_keys(manager, request, kv, top) + columns = [ + b"".join(reuse_keys[int(o)] for o in lg_ordinals) for lg_ordinals in block_ordinals + ] + return [np.frombuffer(column, dtype=np.uint8).reshape(-1, 32) for column in columns] + + def _assign_staging_slots(self, entries: _LentBlockEntries, runs: SlotRuns | None) -> None: + """Each pool group's run goes to its layer groups in order, so the entries of one pool group + occupy consecutive slots.""" + cursor: dict[int, int] = {} + entries.staging_slots = [] + for lg, block_ordinals in zip(entries.layer_group_indices, entries.block_ordinals): + g = int(self._layout.pool_group_of[lg]) + start = runs.by_pool_group.get(g, (0, 0))[0] if runs is not None else 0 + offset = cursor.get(g, 0) + entries.staging_slots.append( + np.arange(start + offset, start + offset + len(block_ordinals), dtype=np.int64) + ) + cursor[g] = offset + len(block_ordinals) + + def _build_view(self, entries: _LentBlockEntries, reuse_keys: list[np.ndarray]) -> LeaseView: + groups = [] + for lg, block_ordinals, lg_keys, slots in zip( + entries.layer_group_indices, entries.block_ordinals, reuse_keys, entries.staging_slots + ): + index = self._region_index_of_pool_group[int(self._layout.pool_group_of[lg])] + region = self._regions[index] + groups.append( + LayerGroupBlocks( + lg, + block_ordinals, + object_keys=self._object_keys(lg, lg_keys), + addresses=region.address + slots * region.slot_bytes, + region_index=index, + ) + ) + return LeaseView(tuple(groups)) + + def _copy_segments( + self, entries: _LentBlockEntries, mask: list[np.ndarray] | None = None + ) -> list[list[int]]: + """``[staging address, device address, bytes]`` for every pool of every row ``mask`` keeps + (one boolean array per layer group; ``None``: every row), by device pool group. Segments + that continue each other on both sides are merged.""" + by_group: dict[int, list[int]] = {} + for i, lg in enumerate(entries.layer_group_indices): + by_group.setdefault(int(self._layout.pool_group_of[lg]), []).append(i) + out: list[list[int]] = [] + for g, members in by_group.items(): + dev = np.concatenate([entries.device_pages[i] for i in members]) + stg = np.concatenate([entries.staging_slots[i] for i in members]) + if mask is not None: + keep = np.concatenate([np.asarray(mask[i], dtype=bool) for i in members]) + dev, stg = dev[keep], stg[keep] + region = self._regions[self._region_index_of_pool_group[g]] + # A staging slot holds the device pools' slots back to back in pool order. + offset = 0 + for pool in self._layout.device_pools[g]: + width = int(pool.slot_bytes) + for d, t in zip(dev.tolist(), stg.tolist()): + host = region.address + t * region.slot_bytes + offset + device = int(pool.base) + d * width + last = out[-1] if out else None + if last and last[0] + last[2] == host and last[1] + last[2] == device: + last[2] += width + else: + out.append([host, device, width]) + offset += width + return out + + # TODO: copies run serially with the forward passes on the execution stream, and a copy call per + # row and pool, which entries share only where they continue each other in a pool group of one + # pool, keeps small entries and entries spanning several pools below the host-to-device bandwidth. + def _queue_copy( + self, segments: Sequence[Sequence[int]], to_staging: bool + ) -> tuple[_QueuedCopy | None, str | None]: + """Queue async copies on the manager's execution stream, then record an event covering + every copy queued, even after one failed; ``(None, error)`` if it could record none. With + page-locked staging the CPU does not wait for them; pageable staging may hold the call.""" + # The execution stream orders a copy after the forward passes that wrote its pages and + # before later work on it, which a page's new owner waits for; the lender does not rely on a + # release waiting for the copies, and a writer off that stream is not ordered after them. + stream = _manager.stream(self._manager_ref()) + handle = drv.CUstream(stream.cuda_stream) + error = None + for host, device, nbytes in segments: + dst, src = (host, device) if to_staging else (device, host) + (result,) = drv.cuMemcpyAsync( + drv.CUdeviceptr(dst), drv.CUdeviceptr(src), nbytes, handle + ) + if result != drv.CUresult.CUDA_SUCCESS: + error = f"cuMemcpyAsync of {nbytes} bytes failed: {result}" + break + try: + event = torch.cuda.Event() + event.record(stream) + except RuntimeError as record_error: + error = error or f"recording the copy's event failed: {record_error}" + logger.warning(f"KV cache lender: {error}") + return None, error + if error is not None: + logger.warning(f"KV cache lender: {error}") + return _QueuedCopy(event), error + + +class _LeaseBase: + """A lease's failure, and the checks of a write's one mark.""" + + # Set by each lease class's ``__init__``; the methods below use them. + _failure: str | None + _kind: str + _view_returned: bool + _marked: bool + _view: LeaseView | None + + @property + def failure(self) -> str | None: + """See ``Lease.failure``.""" + return self._failure + + def _accept_marks(self, masks: Sequence[np.ndarray]) -> list[np.ndarray]: + """Copies of ``masks`` for the one mark of a write whose view ``poll()`` returned; + ``ValueError`` unless one bool mask of shape ``(len(blocks),)`` per group collection.""" + if self._kind != "write": + raise RuntimeError("mark_arrived is for write leases") + if not self._view_returned: + raise RuntimeError("mark_arrived before poll() returned the view") + if self._marked: + raise RuntimeError("mark_arrived called twice") + masks, groups = list(masks), self._view.groups + if len(masks) != len(groups): + raise ValueError(f"{len(masks)} masks for {len(groups)} groups") + out = [] + for blocks, mask in zip(groups, masks): + mask = np.asarray(mask) + if mask.dtype != np.bool_ or mask.shape != (len(blocks),): + raise ValueError( + f"layer group {blocks.layer_group_index}: the mask must be bool of shape " + f"({len(blocks)},), got {mask.dtype} {mask.shape}" + ) + out.append(mask.copy()) + self._marked = True + return out + + +class _LeaseState(enum.Enum): + """A staging lease's state, which ``_StagingLease._state`` derives from its fields. ``release`` + moves any state to ``RELEASED``. Its slots and its copy vary apart from the state. The slots: + none, waiting in order in ``Staging._waiting_for_slots``, or held in + ``Staging._leases_holding_slots`` until ``_recycle`` returns them or ``_quarantine`` keeps them + from reuse. The copy (``_copy``): none, ``_COPY_WITHOUT_RECORD``, or a ``_QueuedCopy``, pending + or done. Three axes each derive from their own fields alone, released or not: ``_HolderState``, + whether the lease failed and, if not, whether a poll returned its view; ``_SlotState``, whether + it waits in line for slots; ``_CopyState``, whether its copy is ``_COPY_WITHOUT_RECORD``. + ``Staging._recyclable`` reads the first, ``Staging._on_lease_released`` the first two and + ``Staging._slots_quarantined`` the third.""" + + # In line for slots: _grant_waiting grants it in turn, or dooms a fresh write whose grant + # raised; _fail fails it. + WAITING = "waiting" + # Without slots or a view: its first poll fails it, as does the shutdown. + FAILS_AT_FIRST_POLL = "fails_at_first_poll" + # Its view and any slots: poll makes it ready, a read once its copy into them has landed, and + # fails a read whose copy is lost. + GRANTED = "granted" + # Poll returned its view: a backend may use its slots until release; mark_arrived marks a write. + READY = "ready" + # A write mark_arrived marked: its marked entries still lent go to the request's pages. + MARKED = "marked" + # Failed at the call, in line, at its grant, at a poll or at the shutdown. + FAILED = "failed" + # Let go: _recycle returns its slots once neither a backend nor a copy can reach them. + RELEASED = "released" + + +class _HolderState(enum.Enum): + """Whether a staging lease failed and, if not, whether a poll returned its view, which + ``_StagingLease._holder_state`` derives from ``_failure`` and ``_view_returned`` alone, released + or not.""" + + # Its failure is set. + FAILED = "failed" + # No failure, and no poll has returned its view. + VIEW_PENDING = "view_pending" + # No failure, and a poll returned its view. + VIEW_RETURNED = "view_returned" + + +class _SlotState(enum.Enum): + """Whether a staging lease waits in line for slots, which ``_StagingLease._slot_state`` derives + from ``_ticket`` alone.""" + + # A ticket: in line for slots. + WAITING = "waiting" + # No ticket. + NOT_WAITING = "not_waiting" + + +class _CopyState(enum.Enum): + """Whether a staging lease's copy is ``_COPY_WITHOUT_RECORD``, which + ``_StagingLease._copy_state`` derives from ``_copy`` alone.""" + + # _COPY_WITHOUT_RECORD: a copy that may be queued with no record of its completion. + WITHOUT_RECORD = "without_record" + # No copy, or one queued with the record of its completion. + NONE_OR_RECORDED = "none_or_recorded" + + +class _StagingLease(_LeaseBase): + """A staging lease; holds its lender weakly and has no finalizer. Backends may read its view's + arrays on their own threads until release, and never call it.""" + + def __init__( + self, + lender: Staging, + kind: str, + request_id: int, + kv: _KVCache | None = None, + entries: _LentBlockEntries | None = None, + reuse_keys: list[np.ndarray] | None = None, + fetch: _FetchRecord | None = None, + ) -> None: + self._lender = weakref.ref(lender) + self._kind = kind + self._request_id = request_id + self._kv = kv + self._block_entries = entries + self._reuse_keys = reuse_keys + self._fetch = fetch + self._first_poll_failure: str | None = None # why it fails at its first poll + self._ticket: int | None = None # its place in line while it waits for slots + # The granted slots; _recycle returns them and leaves this field set. + self._slot_runs: SlotRuns | None = None + self._view: LeaseView | None = None + self._copy: _QueuedCopy | None = None + self._view_returned = False + self._released = False + self._marked = False + self._failure: str | None = None + + @classmethod + def _failed_at_call( + cls, lender: Staging, kind: str, request_id: int, reason: str + ) -> _StagingLease: + """A lease failed at the call; open until released, like any other.""" + lease = cls(lender, kind, request_id) + lease._set_failure(reason) + if not lender._shut_down: + lender._open_leases.add(lease) + return lease + + def _set_failure(self, reason: str) -> None: + self._failure = reason + self._first_poll_failure = None + + def _state(self) -> _LeaseState: + """Its state, from the first of these fields that is set; see ``_LeaseState``.""" + if self._released: + return _LeaseState.RELEASED + if self._failure is not None: + return _LeaseState.FAILED + if self._first_poll_failure is not None: + return _LeaseState.FAILS_AT_FIRST_POLL + if self._ticket is not None: + return _LeaseState.WAITING + if self._marked: + return _LeaseState.MARKED + return _LeaseState.READY if self._view_returned else _LeaseState.GRANTED + + def _holder_state(self) -> _HolderState: + """See ``_HolderState``.""" + if self._failure is not None: + return _HolderState.FAILED + return _HolderState.VIEW_RETURNED if self._view_returned else _HolderState.VIEW_PENDING + + def _slot_state(self) -> _SlotState: + """See ``_SlotState``.""" + return _SlotState.WAITING if self._ticket is not None else _SlotState.NOT_WAITING + + def _copy_state(self) -> _CopyState: + """See ``_CopyState``.""" + if self._copy is _COPY_WITHOUT_RECORD: + return _CopyState.WITHOUT_RECORD + return _CopyState.NONE_OR_RECORDED + + def _marks_due(self) -> bool: + """A write whose view a poll returned and that has no mark: ``mark_arrived`` may still come + after release, and the copy it queues reads the lease's slots.""" + return self._kind == "write" and self._view_returned and not self._marked + + def __repr__(self) -> str: + return f"" + + def poll(self) -> LeaseView | None: + """See ``Lease.poll``.""" + if self._released: + raise RuntimeError("poll after release") + if self._kind == "read" and self._copy is _COPY_WITHOUT_RECORD and self._failure is None: + # A grant that raised left its copy unrecorded and the lease not failed. + self._set_failure("the copy into staging has no record of its completion") + lender = self._lender() + if lender is not None and lender._stop_reason is not None and not self._view_returned: + # A stopped lender returns only views a poll returned before, also after the shutdown. + lender._fail(self, lender._stop_reason) + return None + if lender is None or lender._serving_manager() is None: + return self._poll_not_serving() + lender._progress() + if self._failure is not None: + return None + if self._first_poll_failure is not None: + lender._fail(self, self._first_poll_failure) + return None + if self._view is None: + return None + if self._kind == "read" and not lender._read_copy_done(self): + return None + self._view_returned = True + return self._view + + def mark_arrived(self, masks: Sequence[np.ndarray]) -> None: + """See ``Lease.mark_arrived``.""" + masks = self._accept_marks(masks) + lender = self._lender() + if ( + lender is not None + and lender._serving_manager() is not None + and lender._stop_reason is None + ): + lender._apply_marks(self, masks) + + def release(self) -> None: + """See ``Lease.release``.""" + if self._released: + return + self._released = True + lender = self._lender() + if lender is not None: + lender._on_lease_released(self) + + def _poll_not_serving(self) -> LeaseView | None: + """``poll`` after the lender stopped serving: no grant or copy, just what already landed.""" + if self._failure is not None or self._view is None: + return None + if self._kind == "read" and self._copy is not None and not self._copy.done(): + return None + self._view_returned = True + return self._view + + +class _StagingHold: + """A hold on the staging memory; holds its lender weakly and has no finalizer. Without a lender + it is inert.""" + + def __init__(self, lender: Staging | None) -> None: + self._lender = weakref.ref(lender) if lender is not None else None + self._released = False + + def release(self) -> None: + """See ``RegionHold.release``.""" + if self._released: + return + self._released = True + lender = self._lender() if self._lender is not None else None + if lender is not None: + lender._on_hold_released(self) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py new file mode 100644 index 000000000000..aa1f46cd2813 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py @@ -0,0 +1,416 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The boundary for sharing's V2 operations, runtime metadata and private lifecycle hooks. + +Manager and runtime operations, including layout declarations, prefix resize and fresh-page fill, +go through this facade, as do reads of runtime layer configs, layer grouping and pool descriptors. +The layout builder reads the manager's head counts and other structural fields directly. Manager +and runtime imports stay inside the functions; the private hook protocol is checked before a +lender is installed. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple, Protocol, Sequence, runtime_checkable + +import numpy as np + +if TYPE_CHECKING: + import torch + + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + PoolGroupDesc, + SsmLayerConfig, + _KVCache, + ) + + from ...llm_request import LlmRequest + from ..kv_cache_manager_v2 import KVCacheManagerV2 + +# The largest 32-bit count: the manager's window ranges compute in 32-bit token counts. +_MAX_TOKENS = 2**31 - 1 + + +class CacheState(NamedTuple): + """A request's runtime cache as the lender reads it: tokens committed, the history length its + sliding windows keep, and whether it is active (not suspended).""" + + committed: int + history: int + active: bool + + +@runtime_checkable +class _ManagerHooks(Protocol): + """Callbacks a lender supplies for its manager's cache lifecycle.""" + + def _on_free(self, request_id: int, kv_cache: _KVCache) -> None: ... + + def _on_shrink(self, request_id: int, kv_cache: _KVCache) -> None: ... + + def _on_reset(self) -> None: ... + + def _on_shutdown(self) -> None: ... + + def _on_caches_closed(self) -> None: ... + + +def mapping(manager: KVCacheManagerV2) -> Mapping: + """The manager's mapping; ``TypeError`` for a manager without one.""" + value = getattr(manager, "mapping", None) + if value is None: + raise TypeError("the KV cache lender needs a manager with a mapping") + return value + + +def has_lossy_cold_page_codec(manager: KVCacheManagerV2) -> bool: + """Whether the configured cold-page codec changes the original KV bytes.""" + return bool(getattr(manager, "_has_lossy_cold_page_codec", False)) + + +def require_v2(manager: object) -> None: + """Raise ``TypeError`` unless ``manager`` is a ``KVCacheManagerV2`` (subclasses included).""" + from ..kv_cache_manager_v2 import KVCacheManagerV2 + + if not isinstance(manager, KVCacheManagerV2): + raise TypeError( + f"the KV cache lender needs a KVCacheManagerV2, got {type(manager).__name__}" + ) + + +def commits_blocks(manager: KVCacheManagerV2) -> bool: + """Whether the manager commits blocks to its prefix-reuse tree: block reuse on and, for a draft + manager, joint reuse with its target.""" + return bool(manager.enable_block_reuse) and bool(manager._can_publish_block_reuse) + + +def kv_connector(manager: KVCacheManagerV2) -> object | None: + """The manager's KV cache connector, or ``None`` without one.""" + return manager.kv_connector_manager + + +def context_moves_history(manager: KVCacheManagerV2) -> bool: + """Whether the manager's context update moves a request's history to the end of each chunk, + raising where that would move it back: under a block reuse policy other than all-reusable, or + when the manager commits no blocks; otherwise the update only commits.""" + from ..kv_cache_manager_v2 import BlockReusePolicy + + if not commits_blocks(manager): + return True + return manager.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE + + +def resize_ends_at_chunk(manager: KVCacheManagerV2) -> bool: + """Whether the manager's context resize sets a cache's capacity from the chunk it runs, as a + draft manager's does, so a chunk resumed below the history may leave the capacity below it and + raise there.""" + return bool(manager.is_draft) + + +def keeps_stale_pages(manager: KVCacheManagerV2) -> bool: + """Whether a block a window leaves behind before the manager commits it keeps its page, which + the commit then stores whole under the block's key: the manager commits blocks without the + minimal snapshot, as under the all-reusable policy.""" + return commits_blocks(manager) and not bool(manager.impl.commit_min_snapshot) + + +def prompt_lookahead(manager: KVCacheManagerV2) -> int: + """Prompt tokens past a position the one-model draft the manager was built for reads (Eagle 1, + vanilla MTP its draft length), wherever its layers live, also where reuse matches are not backed + off by them; 0 for a draft reading none or whose read-ahead upstream has not established.""" + return max(int(manager.draft_prompt_lookahead), int(manager.reuse_match_backoff)) + + +def attached_lender(manager: KVCacheManagerV2) -> object | None: + """The lender attached to ``manager``, or ``None``.""" + return getattr(manager, "_sharing", None) + + +def install_lender(manager: KVCacheManagerV2, lender: _ManagerHooks) -> None: + """Attach ``lender`` for the manager's life: its free, in-place shrinks, reuse reset and + shutdown then call the lender's ``_on_free``, ``_on_shrink``, ``_on_reset``, ``_on_shutdown`` + and, once the runtime's shutdown has returned, ``_on_caches_closed``.""" + hooks = ("_on_free", "_on_shrink", "_on_reset", "_on_shutdown", "_on_caches_closed") + if not isinstance(lender, _ManagerHooks) or any( + not callable(getattr(lender, name, None)) for name in hooks + ): + raise TypeError("the KV cache lender must implement all five manager lifecycle hooks") + manager._sharing = lender + + +def layer_configs(manager: KVCacheManagerV2) -> dict[int, object]: + """Internal layer id -> layer config, preferring the manager's config to the runtime's. + A layer without an explicit id uses its position in the config.""" + config = getattr(manager, "kv_cache_manager_py_config", None) + if config is None: + config = manager.impl.init_config + out: dict[int, object] = {} + for position, layer in enumerate(config.layers): + out[int(getattr(layer, "layer_id", position))] = layer + return out + + +def layer_group_layer_ids(manager: KVCacheManagerV2) -> list[list[int]]: + """Internal layer ids per runtime layer group, in layer-group order.""" + return [[int(layer_id) for layer_id in group] for group in manager.impl.layer_grouping] + + +def pool_group_descriptors(manager: KVCacheManagerV2) -> Sequence[PoolGroupDesc]: + """Runtime pool-group descriptors, in the runtime's order.""" + return manager.impl.pool_group_descs + + +def declared_global_layer_ids(manager: KVCacheManagerV2) -> dict[int, int] | None: + """Internal-to-global layer ids, or ``None`` without the manager's declaration.""" + getter = getattr(manager, "get_disagg_global_layer_ids", None) + if getter is None: + return None + table: dict[int, int] = {} + for layer_group, layer_ids in enumerate(manager.impl.layer_grouping): + ids = [int(value) for value in getter(layer_group)] + if len(ids) != len(layer_ids): + raise ValueError( + f"layer group {layer_group}: {len(ids)} declared global ids for {len(layer_ids)} layers" + ) + table.update(zip((int(layer_id) for layer_id in layer_ids), ids)) + return table + + +def disagg_declaration(manager: KVCacheManagerV2, name: str) -> dict[str, object]: + """Role-keyed layout declarations from the named manager getter.""" + getter = getattr(manager, name, None) + if getter is None: + return {} + return {str(role): value for role, value in dict(getter()).items()} + + +def role_mapper_kinds(manager: KVCacheManagerV2) -> dict[str, object]: + """The manager's per-role disaggregation mapper declarations.""" + return disagg_declaration(manager, "get_disagg_role_mapper_kinds") + + +def role_layouts(manager: KVCacheManagerV2) -> dict[str, object]: + """The manager's per-role disaggregation layout declarations.""" + return disagg_declaration(manager, "get_disagg_role_layouts") + + +def ignored_roles(manager: KVCacheManagerV2) -> frozenset[str]: + """Roles omitted by the manager's disaggregation declaration.""" + getter = getattr(manager, "get_disagg_ignored_roles", None) + return frozenset(str(role) for role in (getter() if getter else ())) + + +def layer_group_pool_indices(manager: KVCacheManagerV2) -> list[int]: + """The pool-group index of each runtime layer group, in layer-group order.""" + return [int(value) for value in manager.impl.get_life_cycle_pool_group_indices()] + + +def index_buffer(manager: KVCacheManagerV2) -> torch.Tensor: + """The host buffer the manager's caches write their page indices into, through a raw pointer + each keeps until it is detached or closed, whether or not the manager still exists.""" + return manager.host_kv_cache_block_offsets + + +def virtual_layers(manager: KVCacheManagerV2) -> tuple[dict[int, tuple[int, int]], int] | None: + """Virtual layers (DeepSeek-V4): internal layer id -> (model layer, attention type value), and + the number of attention types the enum defines. ``None`` for a manager without them.""" + virtual = getattr(manager, "_layer_attn_to_layer_id", None) + if not virtual: + return None + inverse: dict[int, tuple[int, int]] = {} + for (model_layer, attn_type), layer_id in virtual.items(): + inverse[int(layer_id)] = (int(model_layer), int(attn_type.value)) + # Every member of the enum counts, so stages holding different attention types agree. + attn_type_class = type(next(iter(virtual))[1]) + return inverse, max(int(member.value) for member in attn_type_class) + 1 + + +def stream(manager: KVCacheManagerV2) -> torch.cuda.Stream: + """The manager's execution stream: a copy queued on it runs after the forward passes that wrote + the pages and before later work on the stream, which a page's new owner waits for; a writer + off the stream is not ordered after it.""" + return manager._stream + + +def kv_of(manager: KVCacheManagerV2, request_id: int) -> _KVCache | None: + """The request's runtime cache, or ``None`` when it has none (never had one, freed, shut down). + Caches are compared by identity to see a free or a replaced cache.""" + return manager.kv_cache_map.get(request_id) + + +def cache_state(kv: _KVCache) -> CacheState: + """Committed tokens, history length and activity of runtime cache ``kv``, as plain values.""" + return CacheState(int(kv.num_committed_tokens), int(kv.history_length), bool(kv.is_active)) + + +def has_scratch_reuse(kv: _KVCache) -> bool: + """Whether runtime cache ``kv`` has SWA scratch reuse on: each capacity change then keeps its + history within the scratch rewind of the old capacity, and its window blocks may sit in scratch + slots, which the next chunk overwrites.""" + return bool(kv.enable_swa_scratch_reuse) + + +def pages(kv: _KVCache, layer_group: int) -> np.ndarray: + """``int64`` page (pool-group slot) of each block ordinal of ``layer_group``, -1 where the block + has no page.""" + return np.fromiter( + kv.get_aggregated_page_indices(layer_group, valid_only=False), dtype=np.int64 + ) + + +def locked_pages(kv: _KVCache, layer_group: int) -> np.ndarray: + """``int64`` page of each block of an active cache from its locked base page indices, -1 where + the block has none: a window block behind the history keeps at most a held page, which may be + on another tier.""" + return np.array(kv.get_base_page_indices(layer_group)[: kv.num_blocks], dtype=np.int64) + + +def num_blocks(kv: _KVCache) -> int: + """Block ordinals of runtime cache ``kv``: a shrink drops those past its capacity, with their + pages.""" + return int(kv.num_blocks) + + +def is_recurrent_layer(layer_config: AttentionLayerConfig | SsmLayerConfig) -> bool: + """Whether a layer config of the runtime describes a recurrent-state layer.""" + from tensorrt_llm.runtime import kv_cache_manager_v2 as runtime + + return isinstance(layer_config, runtime.SsmLayerConfig) + + +def stale_range(manager: KVCacheManagerV2, layer_group: int, history: int) -> tuple[int, int]: + """Block ordinals ``[beg, end)`` behind a windowed group's window at ``history``, sinks + excepted. The range adds a block's tokens to the history in 32 bits, so a longer history is read + as ``_MAX_TOKENS`` less one block, which lowers only ``end``.""" + # TODO: with a window near ``_MAX_TOKENS``, a history past that bound lowers ``end`` below + # blocks the window has passed, which then count as inside it. + longest = _MAX_TOKENS - int(manager.tokens_per_block) + beg, end = manager._stale_block_range(layer_group, min(int(history), longest)) + return int(beg), int(end) + + +def keyed_by_placeholders(request: LlmRequest) -> bool: + """Whether the request carries multimodal data that the manager keys by its placeholder tokens + alone, as ``_augment_tokens_for_block_reuse`` does without a digest, position or length.""" + return request.py_multimodal_data is not None and ( + request.multimodal_hashes is None + or request.multimodal_positions is None + or request.multimodal_lengths is None + ) + + +def has_encoder_input(request: LlmRequest) -> bool: + """Whether the request has encoder input, which its decoder's self-attention bytes depend on.""" + return request.try_get_encoder_output_len() is not None + + +def bidirectional_runs(request: LlmRequest) -> list[tuple[int, int]]: + """The runs ``[b, e)`` of multimodal tokens the scheduler keeps within one context chunk: none + unless the request's multimodal data sets ``mm_bidirectional_blocks`` and holds a + ``multimodal_embed_mask_cumsum``, the gate the scheduler uses.""" + data = request.py_multimodal_data + if not isinstance(data, dict) or not data.get("mm_bidirectional_blocks", False): + return [] + cumsum = data.get("multimodal_embed_mask_cumsum") + if cumsum is None: + return [] + multimodal = np.diff(np.asarray(cumsum, dtype=np.int64), prepend=0) == 1 + edges = np.flatnonzero(np.diff(np.concatenate([[False], multimodal, [False]]).astype(np.int8))) + return [(int(b), int(e)) for b, e in zip(edges[::2], edges[1::2])] + + +def returns_context_outputs(request: LlmRequest) -> bool: + """Whether the request returns context logits, as prompt logprobs make it do, or asks for + additional model outputs, which a model may give per context token.""" + return bool(request.py_return_context_logits) or bool(request.py_additional_outputs) + + +def prompt_length(request: LlmRequest) -> int: + """The request's prompt length.""" + return int(request.prompt_len) + + +def block_keys( + manager: KVCacheManagerV2, request: LlmRequest, kv: _KVCache, count: int +) -> list[bytes]: + """The 32-byte reuse keys of the request's first ``count`` whole blocks, over the tokens + and reuse scope the manager commits with. ``ValueError`` if the request has fewer.""" + # TODO: reuse keys miss a request's encoder input, multimodal embeddings without digests and + # how the input processor treats a digested item, so requests that differ only there share keys. + if count <= 0: + return [] + from tensorrt_llm.runtime.kv_cache_manager_v2 import sequence_to_blockchain_keys + + tpb = int(manager.tokens_per_block) + source = manager._reuse_token_source(request) + need = count * tpb + if len(source) < need: + raise ValueError( + f"request {request.py_request_id} has {len(source) // tpb} whole blocks of tokens, " + f"{count} asked" + ) + # Multimodal digests take the place of placeholder tokens, as when the manager commits. + tokens = manager._augment_tokens_for_block_reuse(source, request, 0, need) + if isinstance(tokens, np.ndarray): + tokens = tokens.tolist() + keys: list[bytes] = [] + for i, (block_tokens, key) in enumerate( + sequence_to_blockchain_keys(tpb, kv.reuse_scope, tokens) + ): + if i == 0: + continue # the root: the reuse scope's own key + if len(block_tokens) < tpb: + break + keys.append(bytes(key)) + if len(keys) == count: + break + if len(keys) < count: + raise ValueError( + f"request {request.py_request_id} has {len(keys)} whole blocks of tokens, {count} asked" + ) + return keys + + +def grow( + manager: KVCacheManagerV2, request: LlmRequest, kv: _KVCache, position: int, end: int +) -> bool: + """Cover ``[0, end)`` and move the history to ``position`` with the manager's own resize, then + run its fresh-page fill (a diagnostic, off unless set) on the pages the resize added alone. + ``False``, the cache unchanged, when pages run out.""" + before = _pages_by_pool(manager, kv) + if not manager._resize_for_served_prefix(request, kv, position, end): + return False + # TODO: grow writes the fresh-page fill's private record (``_fresh_pages_filled``) in the + # fill's own keying: per request, one int32 page array per pool. + if before is not None: + # The fill takes as fresh the pages its record lacks, and a joint draft pool's own resizes + # keep no record: the pages before the grow become it, so the fill leaves them alone. + manager._fresh_pages_filled[int(request.py_request_id)] = before + manager._fill_fresh_kv_pages(request.py_request_id) + return True + + +def _pages_by_pool(manager: KVCacheManagerV2, kv: _KVCache) -> dict[int, np.ndarray] | None: + """Runtime cache ``kv``'s base page of each block, per pool, as the fresh-page fill records + them; ``None`` with the fill off.""" + if manager._fresh_page_fill is None: + return None + count = int(kv.num_blocks) + return { + pool: np.frombuffer(kv.get_base_page_indices(pool), dtype=np.int32, count=count).copy() + for pool in range(int(manager.num_pools)) + } diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_ranges.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_ranges.py new file mode 100644 index 000000000000..82cf26e03168 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_ranges.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Which block ordinals of a layer group a history still reads under its window, and whether a +resume must stay at or past the history. Other modules call these through the module object +(``_ranges._stale_range(...)``), so each has one binding, which every caller reads.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence + +import numpy as np + +from . import _manager + +if TYPE_CHECKING: + from ..kv_cache_manager_v2 import KVCacheManagerV2 + from ._layout import ManagerLayout + + +def _stale_range( + manager: KVCacheManagerV2, layout: ManagerLayout, lg: int, history: int +) -> tuple[int, int]: + """Block ordinals ``[beg, end)`` behind layer group ``lg``'s window at ``history``.""" + if layout.windows[lg] is None: + return 0, 0 + return _manager.stale_range(manager, lg, history) + + +def _needed_ranges( + manager: KVCacheManagerV2, + layout: ManagerLayout, + lg: int, + start_block: int, + end_block: int, + history: int, +) -> list[tuple[int, int]]: + """Ordinal ranges ``[beg, end)`` of ``lg`` in ``[start_block, end_block)`` that a history of + ``history`` tokens still reads, in order and none empty: the whole range for full attention; + the sinks and the window otherwise. Arithmetic only, so a range of any length costs nothing.""" + # TODO: DeepSeek-V4 widens every window by the draft length ``D`` under any speculative + # decoding, so at 128 tokens per block and ``D >= 2`` a fetch asks for margin rows that a + # publish from a history at least ``D - 1`` tokens past the fetch's end leaves out. + stale_beg, stale_end = _stale_range(manager, layout, lg, history) + runs = ((start_block, min(end_block, stale_beg)), (max(start_block, stale_end), end_block)) + return [(beg, end) for beg, end in runs if end > beg] + + +def _ordinals(runs: Sequence[tuple[int, int]]) -> np.ndarray: + """``int64`` ordinals of ``runs``, in order.""" + if not runs: + return np.zeros(0, dtype=np.int64) + return np.concatenate([np.arange(beg, end, dtype=np.int64) for beg, end in runs]) + + +def _needed_ordinals( + manager: KVCacheManagerV2, + layout: ManagerLayout, + lg: int, + start_block: int, + end_block: int, + history: int, +) -> np.ndarray: + """Ordinals of ``lg`` in ``[start_block, end_block)`` that a history of ``history`` tokens + still reads: all of them for full attention; the sinks and the window otherwise.""" + return _ordinals(_needed_ranges(manager, layout, lg, start_block, end_block, history)) + + +# TODO: a draft pool also gives up a resume below its history whose chunk would reach it, as +# every unchunked prefill's does. +def _floor_follows_history(manager: KVCacheManagerV2) -> bool: + """Whether a request may resume only at or past its history: a chunk resumed below it may end + below it, where the manager's context update raises, or leave a draft pool's capacity below it, + where the pool's context resize raises.""" + return _manager.context_moves_history(manager) or _manager.resize_ends_at_chunk(manager) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_slots.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_slots.py new file mode 100644 index 000000000000..02bac5be3341 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_slots.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Staging sizing and slot bookkeeping, without memory. A lease takes all its slots at once, one +contiguous run per pool group, or none; leases waiting in line are served strictly in order.""" + +from __future__ import annotations + +import bisect +from collections import OrderedDict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Mapping + +import numpy as np + +from ._types import StagingOptions, _integer + +if TYPE_CHECKING: + from ._layout import ManagerLayout + + +def slots_per_fetch(layout: ManagerLayout, tokens_per_fetch: int) -> dict[int, int]: + """Slots per device pool group that any whole-block range of at most ``tokens_per_fetch`` tokens can + need: ``ceil(tokens_per_fetch / tpb)`` per layer group, with a window at most its sink blocks plus + ``ceil(window / tpb)``. ``ValueError`` if ``tokens_per_fetch < 1``, ``TypeError`` if not an int.""" + tpb = int(layout.tokens_per_block) + target = -(-_integer("tokens_per_fetch", tokens_per_fetch) // tpb) + slots: dict[int, int] = {} + for lg, group in enumerate(layout.pool_group_of): + blocks = target + window = layout.windows[lg] + # At a history on a block boundary, where every staging range ends, a window of W tokens + # reads at most ceil((W - 1) / tpb) <= ceil(W / tpb) blocks besides its sinks. + if window is not None: + blocks = min(target, int(layout.sink_blocks[lg]) + -(-int(window) // tpb)) + slots[int(group)] = slots.get(int(group), 0) + blocks + return slots + + +def slot_counts(layout: ManagerLayout, options: StagingOptions) -> dict[int, int]: + """Slots per device pool group: ``planned_fetches`` fetches, or ``max_bytes`` split by each group's + share of one fetch and rounded down to whole slots. ``ValueError`` if ``max_bytes`` is below + one fetch, so every group holds at least one fetch's slots.""" + slots = slots_per_fetch(layout, options.tokens_per_fetch) + fetch_bytes_by_pool_group = {g: slots[g] * int(layout.page_bytes[g]) for g in slots} + bytes_per_fetch = sum(fetch_bytes_by_pool_group.values()) + if bytes_per_fetch < 1: + raise ValueError("the KV cache manager has no pages to stage") + nbytes = options.planned_fetches * bytes_per_fetch + if options.max_bytes is not None: + if options.max_bytes < bytes_per_fetch: + raise ValueError( + f"max_bytes of {options.max_bytes} is below one fetch of {options.tokens_per_fetch} " + f"tokens, which needs {bytes_per_fetch} bytes" + ) + nbytes = min(nbytes, options.max_bytes) + # Each group's share of the bytes is its share of one fetch, so a cap keeps every group's + # slots of one fetch; uncapped, a group gets exactly planned_fetches times its slots. + return { + g: (nbytes * fetch_bytes_by_pool_group[g] // bytes_per_fetch) // int(layout.page_bytes[g]) + for g in slots + } + + +@dataclass(frozen=True) +class SlotRuns: + """The slots one lease holds: per pool group, ``count`` slots from ``start``.""" + + by_pool_group: Mapping[int, tuple[int, int]] + + def slots(self, group: int) -> np.ndarray: + """``int64`` slot indices of ``group``'s run, empty when the lease holds none there.""" + start, count = self.by_pool_group.get(group, (0, 0)) + return np.arange(start, start + count, dtype=np.int64) + + +class _RunAllocator: + """First-fit allocator of contiguous slot runs with coalescing frees.""" + + def __init__(self, num_slots: int) -> None: + self.num_slots = num_slots + self._free: list[tuple[int, int]] = [(0, num_slots)] if num_slots else [] + + def find(self, count: int) -> int | None: + for start, length in self._free: + if length >= count: + return start + return None + + def take(self, start: int, count: int) -> None: + for i, (s, length) in enumerate(self._free): + if s <= start and start + count <= s + length: + pieces = [] + if start > s: + pieces.append((s, start - s)) + if start + count < s + length: + pieces.append((start + count, s + length - start - count)) + self._free[i : i + 1] = pieces + return + raise ValueError(f"slots [{start}, {start + count}) are not free") + + def overlaps_free(self, start: int, count: int) -> bool: + """Some slot of ``[start, start + count)`` is free already.""" + return any(s < start + count and start < s + length for s, length in self._free) + + def give(self, start: int, count: int) -> None: + """Return a run ``SlotQueue.give_back`` checked: inside the pool, and no slot of it free.""" + if count == 0: + return + i = bisect.bisect_left(self._free, (start, 0)) + self._free.insert(i, (start, count)) + merged: list[tuple[int, int]] = [] + for s, length in self._free: + if merged and merged[-1][0] + merged[-1][1] == s: + merged[-1] = (merged[-1][0], merged[-1][1] + length) + else: + merged.append((s, length)) + self._free = merged + + @property + def free_slots(self) -> int: + return sum(length for _, length in self._free) + + +class SlotQueue: + """First-come-first-served slot queue over every pool group. Not thread-safe: like the manager, + one thread at a time uses it.""" + + def __init__(self, num_slots: Mapping[int, int]) -> None: + self._num_slots = {int(g): int(count) for g, count in num_slots.items()} + self._allocators = {g: _RunAllocator(c) for g, c in self._num_slots.items()} + self._waiting: OrderedDict[int, dict[int, int]] = OrderedDict() + self._next_ticket = 0 + + def num_slots(self, group: int) -> int: + """Slots ``group`` holds in all.""" + return self._num_slots[group] + + def check_fits_capacity(self, counts: Mapping[int, int]) -> None: + """Validate total capacity, not current free runs or queue position; queue nothing. + + ``ValueError`` for an unknown group or more slots than a group holds. + """ + self._checked_counts(counts) + + def enqueue(self, counts: Mapping[int, int]) -> int: + """Queue a request for ``counts`` slots per group and return its ticket; ``ValueError`` as + ``check_fits_capacity``.""" + wanted = self._checked_counts(counts) + ticket = self._next_ticket + self._next_ticket += 1 + self._waiting[ticket] = wanted + return ticket + + def try_grant(self, ticket: int) -> SlotRuns | None: + """Grant ``ticket`` if it is first in line and every group has a free run; else ``None``. + ``KeyError`` for a ticket that is not waiting.""" + if ticket not in self._waiting: + raise KeyError(f"staging ticket {ticket} is not waiting") + if next(iter(self._waiting)) != ticket: + return None + wanted = self._waiting[ticket] + starts = {} + for g, c in wanted.items(): + start = self._allocators[g].find(c) + if start is None: + return None + starts[g] = start + for g, c in wanted.items(): + self._allocators[g].take(starts[g], c) + del self._waiting[ticket] + return SlotRuns({g: (starts[g], c) for g, c in wanted.items()}) + + def cancel(self, ticket: int) -> None: + """Withdraw a waiting ticket; an unknown or granted ticket is ignored.""" + self._waiting.pop(ticket, None) + + def give_back(self, runs: SlotRuns) -> None: + """Return one lease's slots, all or none. ``ValueError`` for a run not held (freed twice or + never granted).""" + # Every run is checked before any is returned, so a bad one returns none. + for g, (start, count) in runs.by_pool_group.items(): + alloc = self._allocators.get(g) + if alloc is None: + raise ValueError(f"no staging slots for pool group {g}") + if count and (start < 0 or count < 0 or start + count > alloc.num_slots): + raise ValueError(f"slots [{start}, {start + count}) are outside pool group {g}") + if count and alloc.overlaps_free(start, count): + raise ValueError(f"slots [{start}, {start + count}) of pool group {g} freed twice") + for g, (start, count) in runs.by_pool_group.items(): + self._allocators[g].give(start, count) + + @property + def num_waiting(self) -> int: + """Tickets waiting in line.""" + return len(self._waiting) + + def free_slots(self, group: int) -> int: + """Slots of ``group`` not held by any lease.""" + return self._allocators[group].free_slots + + def _checked_counts(self, counts: Mapping[int, int]) -> dict[int, int]: + """``counts`` without zero entries; ``ValueError`` if it can never be granted.""" + wanted = {} + for g, c in counts.items(): + g, c = int(g), int(c) + if c < 0: + raise ValueError(f"{c} staging slots asked of pool group {g}") + if c == 0: + continue + if g not in self._num_slots: + raise ValueError(f"no staging slots for pool group {g}") + if c > self._num_slots[g]: + raise ValueError( + f"{c} staging slots asked of pool group {g}, which has {self._num_slots[g]}" + ) + wanted[g] = c + return wanted diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_types.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_types.py new file mode 100644 index 000000000000..2526e37d0b26 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_types.py @@ -0,0 +1,522 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The lender's public types and protocols. Numpy and the standard library only, so any side can +import them without a live KV cache manager.""" + +from __future__ import annotations + +import numbers +from dataclasses import dataclass +from typing import TYPE_CHECKING, NamedTuple, Optional, Protocol, Sequence, Tuple, runtime_checkable + +import numpy as np + +if TYPE_CHECKING: + from ...llm_request import LlmRequest + +# A block entry's fixed-length opaque object key. +_OBJECT_KEY_BYTES = 54 + + +def _freeze(array: np.ndarray) -> np.ndarray: + """``array`` marked read-only (a view when it is not already).""" + if array.flags.writeable: + array = array.view() + array.flags.writeable = False + return array + + +def _integer(name: str, value: object, least: int = 1) -> int: + if isinstance(value, bool) or not isinstance(value, numbers.Integral): + raise TypeError(f"{name} must be an integer, got {type(value).__name__}") + if value < least: + raise ValueError(f"{name} must be at least {least}, got {value}") + return int(value) + + +@dataclass(frozen=True) +class StagingOptions: + """Host staging capacity, planned in fetch-sized units. + + The lender computes the bytes of one fetch from the layer groups and their windows, and a range + of at most ``tokens_per_fetch`` tokens always fits. A capacity budget, not concurrency: a lease + takes one contiguous run of slots per pool group and leases that need slots are granted strictly + in order, so holes that released leases leave can make a lease wait although enough slots are + free in total; a lease with no block entries is ready at its first poll, with no place in line. + + Attributes: + tokens_per_fetch: The token count used to size one fetch unit, rounded up to whole blocks + and adjusted for each layer group's window and sinks. Not a lease's range limit: + larger ranges may fit when their required slots fit. + planned_fetches: The capacity multiplier in fetch units, before the byte cap. Not a + limit on concurrent leases or a guarantee that this many fetches run together. + max_bytes: A cap on the staging bytes, or ``None``; ``attach_lender`` raises + ``ValueError`` if it is below one fetch. + + Raises: + TypeError: A field is not an integer (``bool`` included). + ValueError: A field is not positive. + """ + + tokens_per_fetch: int + planned_fetches: int = 1 + max_bytes: Optional[int] = None + + def __post_init__(self) -> None: + object.__setattr__( + self, "tokens_per_fetch", _integer("tokens_per_fetch", self.tokens_per_fetch) + ) + object.__setattr__( + self, "planned_fetches", _integer("planned_fetches", self.planned_fetches) + ) + if self.max_bytes is not None: + object.__setattr__(self, "max_bytes", _integer("max_bytes", self.max_bytes)) + + +@dataclass(frozen=True) +class MemoryRegion: + """One memory region containing equal-sized block entries. + + A region describes where a backend may access lent memory, not the manager's cache ownership + or its allocation strategy. Its address and extent remain stable while a region hold is open. + + Attributes: + layout_label: An opaque layout-derived label, equal for regions laid out alike; not an + allocation identity. + address: The region's base virtual address. + nbytes: The region's extent in bytes. + slot_bytes: The byte extent of one block entry. + slot_count: The number of entries, back to back from ``address``. + layer_group_indices: Manager-local layer-group indices whose entries use this region. + """ + + layout_label: str + address: int + nbytes: int + slot_bytes: int + slot_count: int + layer_group_indices: Tuple[int, ...] + + +@dataclass(frozen=True, eq=False) +class LayerGroupBlocks: + """One layer group's lent KV blocks, described by aligned read-only arrays. + + Entry ``i`` describes this rank's KV for logical token block ``block_ordinals[i]`` across the + layers in this group, not a token tensor row or a physical page. Entries need not have + consecutive ordinals or addresses. Placement fields are either all given or all ``None``. + + Equal object keys mean interchangeable bytes under the scope and caller obligations of + ``attach_lender``. Keys are opaque identities, not payload checksums. Their format is not API, + but their width, 54 bytes, is. Consumers store them and compare them only for equality. + Changing the format or the layout makes previously stored objects miss, never hit wrongly. + + Attributes: + layer_group_index: The manager's local index of the layer group. + block_ordinals: ``int64 (n,)`` logical block ordinals; block ``b`` covers token positions + ``[b * tokens_per_block, (b + 1) * tokens_per_block)``. + object_keys: ``uint8 (n, 54)`` complete opaque object keys, distinct from the manager's + 32-byte reuse keys. + addresses: ``int64 (n,)`` virtual addresses; entry ``i`` occupies ``slot_bytes`` bytes at + ``addresses[i]`` in the region identified by ``region_index``. + region_index: The position of this group's region in ``Lender.regions``, not a pool-group + or layer-group identifier. + """ + + layer_group_index: int + block_ordinals: np.ndarray + object_keys: Optional[np.ndarray] = None + addresses: Optional[np.ndarray] = None + region_index: Optional[int] = None + + def __post_init__(self) -> None: + object.__setattr__( + self, "layer_group_index", _integer("layer_group_index", self.layer_group_index, 0) + ) + block_ordinals = _freeze(np.ascontiguousarray(self.block_ordinals, dtype=np.int64)) + if block_ordinals.ndim != 1: + raise ValueError( + f"block_ordinals must be one-dimensional, got shape {block_ordinals.shape}" + ) + object.__setattr__(self, "block_ordinals", block_ordinals) + placed = ( + self.object_keys is not None, + self.addresses is not None, + self.region_index is not None, + ) + if any(placed) and not all(placed): + raise ValueError("object_keys, addresses and region_index are all given or all None") + if not all(placed): + return + object_keys = np.asarray(self.object_keys) + if object_keys.dtype != np.uint8 or object_keys.shape != ( + len(block_ordinals), + _OBJECT_KEY_BYTES, + ): + raise ValueError( + f"object_keys must be uint8 ({len(block_ordinals)}, {_OBJECT_KEY_BYTES}), " + f"got {object_keys.dtype} {object_keys.shape}" + ) + addresses = np.asarray(self.addresses) + if ( + not np.issubdtype(addresses.dtype, np.integer) + or addresses.shape != block_ordinals.shape + ): + raise ValueError( + f"addresses must be integers of shape {block_ordinals.shape}, " + f"got {addresses.dtype} {addresses.shape}" + ) + object.__setattr__(self, "object_keys", _freeze(np.ascontiguousarray(object_keys))) + object.__setattr__( + self, + "addresses", + _freeze(np.ascontiguousarray(addresses, dtype=np.int64)), + ) + object.__setattr__(self, "region_index", _integer("region_index", self.region_index, 0)) + + def __len__(self) -> int: + return int(self.block_ordinals.shape[0]) + + def select(self, mask: np.ndarray) -> LayerGroupBlocks: + """The block entries where ``mask`` is true, keeping the arrays aligned. + + Args: + mask: One boolean per entry, ``True`` for each entry to keep. + + Returns: + The same layer group and region with the kept entries, in their original order. + + Raises: + ValueError: ``mask`` is not boolean with one value per block entry. + """ + mask = np.asarray(mask) + if mask.dtype != np.bool_ or mask.shape != self.block_ordinals.shape: + raise ValueError( + f"mask must be bool of shape {self.block_ordinals.shape}, got {mask.dtype} {mask.shape}" + ) + return LayerGroupBlocks( + self.layer_group_index, + self.block_ordinals[mask], + None if self.object_keys is None else self.object_keys[mask], + None if self.addresses is None else self.addresses[mask], + self.region_index, + ) + + +@dataclass(frozen=True, eq=False) +class LeaseView: + """What a ready lease lends: at most one block collection per layer group. + + When returned by ``Lease.poll``, its memory is stable and ready for the requested access. + All addresses belong to the lender's regions. Backends may read its immutable arrays on their + own threads and access only the described memory, until they relinquish access and release + the lease. + + Attributes: + groups: Block collections, at most one per layer group. Tuple positions need not equal + ``layer_group_index``. + + Raises: + TypeError: A collection is not a ``LayerGroupBlocks``. + ValueError: A layer group appears twice. + """ + + groups: Tuple[LayerGroupBlocks, ...] + + def __post_init__(self) -> None: + groups = tuple(self.groups) + seen = set() + for blocks in groups: + if not isinstance(blocks, LayerGroupBlocks): + raise TypeError(f"groups hold LayerGroupBlocks, got {type(blocks).__name__}") + if blocks.layer_group_index in seen: + raise ValueError(f"layer group {blocks.layer_group_index} appears twice") + seen.add(blocks.layer_group_index) + object.__setattr__(self, "groups", groups) + + @property + def num_block_entries(self) -> int: + """Total entries across groups; an ordinal in two groups contributes two entries.""" + return sum(len(blocks) for blocks in self.groups) + + def row_masks(self, fill_value: bool = False) -> Tuple[np.ndarray, ...]: + """Make one writable boolean mask per group, filled with ``fill_value``. + + Mask ``i`` selects entries of ``groups[i]``; its element ``j`` refers to + ``groups[i].block_ordinals[j]``, not to ordinal ``j``. Empty groups keep an empty mask. + + Args: + fill_value: What every block entry's mask value starts as. + + Returns: + One mask per collection, in ``groups`` order: the shape ``mark_arrived`` takes. + """ + return tuple(np.full(len(blocks), bool(fill_value), dtype=bool) for blocks in self.groups) + + +# TODO: the caller repeats by hand the manager's own skip past a served prefix (raising the +# history, setting py_served_prefix_end and the prepopulated length), and no public manager method +# takes a position to raise a request's history to. +class ResumeBounds(NamedTuple): + """Inclusive local bounds on token positions for resuming after a fetch. + + Resume at a position ``p`` with ``min_position <= p <= max_position``, no lower than the + request's current context position. The bounds include only previously valid cache contents + and fetched block entries that have been incorporated into the request's cache. + + The interval belongs to this rank and this manager. Across ranks, and across a target and + its joint-reuse draft pool sharing the request's cursor, use the largest minimum and smallest + maximum. Fetch into both managers before advancing their shared cursor. An empty effective + interval requires dropping the request's cache in every manager fetched into and computing + from 0, without another fetch into those caches. + + A minimum excludes positions whose required history is no longer available or which the + manager cannot rewind to. Bounds never allow resuming strictly inside a bidirectional + multimodal run. The maximum need not be a whole-block boundary; it can reach the prompt length + after the remaining prompt tokens have been computed locally. A fetch leaves the last prompt + token for local computation. + + Attributes: + max_position: The last permitted resume position. + min_position: The first permitted resume position, unless the context position is higher. + """ + + max_position: int + min_position: int + + +@runtime_checkable +class Lease(Protocol): + """One lent range: poll until the view or a final failure, mark a write once, then release. + + Methods run on the manager's thread, one call at a time. Backend threads may inspect the + ready view and access its memory, but call no lender, lease or hold method. They report + completion to the holder through their own channel. + + Caller obligations: + - Poll every open lease each executor iteration, including iterations without new requests. + Do not block waiting for a new request while a lease or backend needs progress. + - Release every lease, failed ones too, only after every backend access has ended. + Dropping an unreleased lease does not release its memory. + - Combine every rank's outcome before deciding for the distributed request. + - A failure is final. A lease failed at the lending call changed no request cache state. + - An exception not listed by the called method is not recoverable through normal lending. + Stop lending, polling, marking and querying bounds through that lender; release leases + and holds once access ends, and treat parked fetches as empty resume intervals. + A failure during manager cleanup can stop a lender permanently. A view already returned + remains held until release; otherwise the lease reports failure. + """ + + def poll(self) -> Optional[LeaseView]: + """Advance pending work and return a stable, ready view. + + Returns: + The same view on every call once ready; ``None`` while pending, or permanently on failure. + Read views contain readable bytes, and write views are ready for backend writes. + + Raises: + RuntimeError: The lease was released. + """ + ... + + @property + def failure(self) -> Optional[str]: + """Why the lease will never be ready, once that is known; for logs only.""" + ... + + def mark_arrived(self, masks: Sequence[np.ndarray]) -> None: + """Incorporate the indicated arrived block entries into the request's cache. + + Write leases only, once, after ``poll()`` returned the view, before or after release. Each + selected entry must have arrived whole and have no further backend writes in flight. + Returning does not guarantee that incorporation has completed: ``resume_bounds`` reports + when the request may use the result. Entries no longer belonging to the request are ignored; + the holder still marks and releases after the request exits. A write released before its view + was returned is abandoned and needs no mark. + + Args: + masks: One boolean array per entry of ``LeaseView.groups``, in tuple order. Positions + select entries, not logical block ordinals. ``LeaseView.row_masks()`` creates + all-false masks. + + Raises: + RuntimeError: A read lease, a second call, or a call before a view was returned. + ValueError: A wrong number, shape or dtype of masks; nothing is recorded. + """ + ... + + def release(self) -> None: + """Confirm that the backend has stopped accessing all lent memory. + + Required for every lease, failed ones too; legal in every state and idempotent. Release is + not cancellation of an access still in flight. A ready write still needs ``mark_arrived``, + which may follow release. + """ + ... + + +@runtime_checkable +class RegionHold(Protocol): + """A backend's lifetime hold on all memory regions. + + Take a hold on the manager's thread before registering ``Lender.regions``. Their addresses + and backing memory remain valid while the hold is open, including across manager shutdown. + Deregister before shutdown, and release the hold only after deregistration is confirmed. + A backend unable to confirm deregistration keeps its hold. Dropping a hold does not release + it. + """ + + def release(self) -> None: + """Confirm deregistration and that the backend cannot access the regions. + + Runs on the manager's thread, is legal in every state, and is idempotent. + """ + ... + + +# TODO: object_keys exist only in ready views, so a flow cannot look up remote hits before it reserves +# pages for them, and a publish lends a windowed group's block entries only for the window at its end: a fork +# at an earlier block or a longer next turn can miss block entries and compute from 0. +@runtime_checkable +class Lender(Protocol): + """Lend a request's KV blocks without transferring cache ownership. + + Lender, lease and hold methods run on the manager's thread, one call at a time. Backend + threads access only a ready lease's described memory. A returned view's addresses are stable + and ready, belong to ``regions``, and remain held until all backend access has ended. + + Caller obligations: + - Take a region hold before registration or the first lease's backend access. Deregister + and release the hold only after the backend can no longer access the regions. + - Poll every open lease and query bounds for every waiting request each executor iteration, + even without new requests. + - Release a lease that waits too long, or free its request; progress is not guaranteed for + a pending lease. + - At shutdown stop the executor loop and backends, confirm deregistration, mark arrived + entries and release leases and holds, then shut down the manager. + - Fetch into both a target and its joint-reuse draft pool before advancing their shared + request cursor; combine their bounds as ``ResumeBounds`` specifies. + """ + + @property + def regions(self) -> Tuple[MemoryRegion, ...]: + """The memory regions to which every ready view's addresses belong. + + Each address and extent remains stable while a region hold is open. A layout label is equal + for regions laid out alike. No ordering, adjacency or single-allocation guarantee is made. + A backend requiring one contiguous region must validate that requirement before it starts. + + Deregister before manager shutdown and release the hold after confirmation. Unreleased + leases or holds can retain memory beyond shutdown. + """ + ... + + def hold_regions(self) -> RegionHold: + """Return a new lifetime hold on every region. + + The hold does not register regions or grant a lease. Take it on the manager's thread before + registration. A hold taken after manager shutdown is inert. + + Returns: + The hold, open until its ``release``. + """ + ... + + def lend_read(self, request: LlmRequest, start: int, end: int) -> Lease: + """Lend committed whole blocks covering token positions ``[start, end)`` for reads. + + Only blocks committed to the manager's prefix-reuse tree are published by name; generated + tokens are not committed. A windowed group lends only retained blocks required at ``end``; + a block without backing memory is omitted. Access only the returned view's entries. + + The lease may fail immediately or while waiting if the request's cache is missing, inactive + or changed, its inputs cannot be named, or the manager stopped or reset reuse. After a reuse + reset the lender cannot publish again: new weights require a new manager and lender under a + scope distinguishing those weights. A previously granted publish may finish. + + Args: + request: The request whose blocks are read. + start: Inclusive first token, a multiple of ``tokens_per_block``. + end: Exclusive end token, a multiple of ``tokens_per_block``, at most the committed tokens. + + Returns: + The lease, possibly already failed. + + Raises: + ValueError: Negative, reversed or non-whole-block bounds, an end past committed tokens, + or a range exceeding the attached implementation's capacity. + """ + ... + + def lend_write(self, request: LlmRequest, start: int, end: int) -> Lease: + """Lend writable memory for fetched blocks covering token positions ``[start, end)``. + + A windowed group lends the retained blocks required at ``end``. The holder marks entries + that arrived whole and releases after backend access ends. Consecutive settled fetches can + contribute to one prefix; lend the next segment only after every rank reports settled bounds. + Abandoning a lease contributes none of its entries. + + The request must not execute while a fetch is unsettled. Previously computed tokens below + ``start`` may contribute to the prefix; retained windows may limit where it can resume. + Contents from ``start`` onward may be overwritten. Fetch only into a cache whose history is + backed by valid data. Use ``resume_bounds``, not the fetch extent, to decide where to continue. + Failure after the lending call need not undo cache changes. + + The lease may fail for the reasons in ``lend_read``, resource exhaustion, an unsettled prior + fetch, incompatible cache history or reuse, unnamed inputs, a bidirectional multimodal run + split by ``end``, or requested per-token outputs that a fetch cannot restore. + + Args: + request: The request fetched into. + start: Inclusive first token, a whole-block boundary no earlier than the committed tokens + rounded down to whole blocks. + end: Exclusive whole-block boundary, no later than the whole blocks before the last prompt + token and, for a windowed cache, no earlier than its retained history. + + Returns: + The lease, possibly already failed. + + Raises: + ValueError: Negative, reversed or non-whole-block bounds, a start within committed whole + blocks, an end past the last eligible prompt block, an end below a windowed cache's + retained history, or a range exceeding the attached implementation's capacity. + Validation precedes changes. + """ + ... + + def resume_bounds(self, request: LlmRequest) -> Optional[ResumeBounds]: + """Advance work and return local bounds once the request's fetches have settled. + + Query only while retained history not advanced by this lender is backed by valid cache + contents. The bounds cannot detect a history that runs ahead of its computed data. + + Only entries incorporated into the request's cache count as fetched. Marking arrival alone + does not make data usable. Abandoned or failed entries contribute nothing; an overwritten + partially committed block can make the interval empty. Shrinking a cache invalidates delivered + entries beyond its new capacity permanently: growing it again does not restore their contents. + + Combine rank and target/draft intervals as ``ResumeBounds`` specifies. If the effective + interval is empty, drop the request's cache in every manager fetched into and compute from 0. + + Args: + request: The request fetched into. + + Returns: + ``None`` while a fetch is unsettled, otherwise the inclusive permitted resume interval. + + Raises: + ValueError: The request has no cache, including after manager shutdown. + """ + ... diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 5b9e7f53101b..d3b49b81c09b 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1024,9 +1024,11 @@ def __init__( # allocation clears this flag so replay can report its new pages. self.py_connector_allocation_reported = False - # End of the connector prefix retained by this allocation, or 0. This - # depth survives async parking; local reuse cannot reconstruct it. - self.py_connector_served_position = 0 + # End of the prefix served from outside the manager (a KV cache + # connector's load, or a staging fetch the request resumes from) retained + # by this allocation, or 0. This depth survives async parking; local reuse + # cannot reconstruct it. + self.py_served_prefix_end = 0 self.py_result = PyResult( prompt_len=self.py_prompt_len, diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 2e213724cadf..3b7e55c8fe33 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -1166,7 +1166,10 @@ def _prepare_context_pair(self, req: LlmRequest) -> bool: self._free_kv_caches(req) return False - if draft_reuse != common_reuse: + # A served prefix is held in both pools even when their reuse policies + # commit different depths. Compare the effective prefixes before freeing it. + served_prefix = req.py_served_prefix_end + if max(draft_reuse, served_prefix) != max(common_reuse, served_prefix): # The shorter re-claim can itself fall short, so drop straight to # no reuse rather than iterate against an eviction we do not own. draft_manager.free_resources(req) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index bdeb84301f06..674ad6a72694 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -40,6 +40,7 @@ l0_a10: - unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py - unittest/_torch/executor/kv_cache/test_kv_cache_compression_manager.py - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py + - unittest/_torch/executor/kv_cache/sharing - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py - unittest/_torch/executor/test_error_classification.py diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/conftest.py b/tests/unittest/_torch/executor/kv_cache/sharing/conftest.py new file mode 100644 index 000000000000..d58a98dbfc27 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/conftest.py @@ -0,0 +1,978 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Real KV cache managers, requests and byte oracles for the lender tests. The oracles bypass the +code under test: device pages through ``impl.pool_group_descs``, staging bytes through raw host +addresses, block keys through ``sequence_to_blockchain_keys``. Nothing touches CUDA at import.""" + +from __future__ import annotations + +import ctypes +import gc +import hashlib +import inspect +import sys +import threading +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Callable, Dict, Generator, Iterator, List, Optional, Sequence, Tuple +from unittest.mock import patch + +import numpy as np +import pytest +import torch + +TPB = 32 +MAX_SEQ_LEN = 256 +SENTINEL = 0xEE +# The smallest pool the harness manager accepts: few enough pages that freed ones get reused. +POOL_TOKENS = 256 +# The host tier of a host_tier_manager, larger than any pool here. Set explicitly: the automatic +# size is capped at 80% of RLIMIT_MEMLOCK, and under a container's default 64 KiB limit it cannot +# hold what these tests move to host. +HOST_CACHE_BYTES = 64 << 20 + +_SHARING = "tensorrt_llm._torch.pyexecutor.kv_cache.sharing" +_FAILED_TEST = pytest.StashKey[bool]() + + +class _RanksAgree: + """The collectives of a TP manager built alone in this process: every rank agrees with it.""" + + local_world_size = 1 + + @staticmethod + def allreduce(value, op=None): + return value + + +@contextmanager +def _collectives_for(mapping): + """Stand-in collectives while a multi-rank manager is built in a single process.""" + if mapping is None or mapping.world_size == 1: + yield + return + from tensorrt_llm._torch.distributed.communicator import Distributed + + with patch.object(Distributed, "get", return_value=_RanksAgree()): + yield + + +def make_manager( + *, + windows: Optional[List[int]] = None, + max_tokens: int = 2048, + tokens_per_block: int = TPB, + num_layers: int = 2, + num_kv_heads: int = 4, + head_dim: int = 64, + mapping=None, + kv_cache_type=None, + dtype=None, + enable_block_reuse: bool = True, + swa_scratch_reuse: bool = False, + max_batch_size: int = 4, + max_seq_len: int = MAX_SEQ_LEN, + spec_config=None, + is_draft: bool = False, + joint_kv_cache_reuse: bool = False, + kv_connector_manager=None, + cold_page_codec_provider=None, + **kv_cache_config, +): + """A small real ``KVCacheManagerV2`` on its own execution stream, FP16 unless ``dtype``, built + for ``spec_config`` as a target or (``is_draft``) a draft pool, with ``kv_connector_manager`` if + given; further keywords go to its ``KvCacheConfig``.""" + import tensorrt_llm.bindings + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.mapping import Mapping + + batch_manager = tensorrt_llm.bindings.internal.batch_manager + config = dict(max_tokens=max_tokens, enable_block_reuse=enable_block_reuse) + if swa_scratch_reuse: + config["enable_swa_scratch_reuse"] = True + if windows is not None: + config["max_attention_window"] = windows + config.update(kv_cache_config) + mapping = mapping or Mapping(world_size=1, tp_size=1, rank=0) + with _collectives_for(mapping): + return KVCacheManagerV2( + kv_cache_config=KvCacheConfig(**config), + kv_cache_type=kv_cache_type or batch_manager.CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype or tensorrt_llm.bindings.DataType.HALF, + vocab_size=32000, + execution_stream=torch.cuda.Stream(), + spec_config=spec_config, + is_draft=is_draft, + joint_kv_cache_reuse=joint_kv_cache_reuse, + kv_connector_manager=kv_connector_manager, + cold_page_codec_provider=cold_page_codec_provider, + ) + + +def _make_deepseek_v4_manager(spec_config=None, fp8_ds_mla=False): + """A four-layer DeepSeek-V4 manager, FP8 with one KV head: one layer of each compression kind + (1, 4, 128) and a last ratio-4 layer the draft layers copy; with ``spec_config`` its one-model + draft layers follow. ``fp8_ds_mla``: the footer-scale cache, in SM100+'s 256-token blocks.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm._torch.speculative.utils import get_num_spec_layers + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig + from tensorrt_llm.mapping import Mapping + + compress_ratios = [1, 4, 128, 4] + draft_layers = get_num_spec_layers(spec_config) if spec_config is not None else 0 + draft_ratios = [compress_ratios[-1]] * draft_layers + return DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_tokens=4096, + event_buffer_max_size=0, + dtype="fp8_ds_mla" if fp8_ds_mla else "auto", + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=len(compress_ratios), + num_kv_heads=1, + head_dim=512, + tokens_per_block=256 if fp8_ds_mla else 128, + max_seq_len=1024, + max_batch_size=2, + max_input_len=1024, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.FP8, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=2 * 1025, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=compress_ratios + draft_ratios + ), + execution_stream=torch.cuda.Stream(), + spec_config=spec_config, + ) + + +def _make_hybrid_manager(): + """Nemotron-H shaped, few layers: Mamba2 state at layers 0, 2 and 4, attention at layer 3.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MambaStateConfig + from tensorrt_llm.mapping import Mapping + + pattern = "M-M*M-" + mamba_mask = [c == "M" for c in pattern] + attn_mask = [c == "*" for c in pattern] + return MambaHybridCacheManagerV2( + mamba_d_state=128, + mamba_d_conv=4, + mamba_num_heads=128, + mamba_n_groups=8, + mamba_head_dim=80, + mamba_num_layers=sum(mamba_mask), + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig( + max_tokens=2048, + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + ), + kv_cache_type=CacheType.SELF, + num_layers=sum(attn_mask), + num_kv_heads=8, + head_dim=128, + tokens_per_block=32, + max_seq_len=1024, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1), + layer_mask=attn_mask, + vocab_size=1024, + dtype=DataType.BF16, + ) + + +@contextmanager +def _managed(factory: Callable[[], object]) -> Iterator[object]: + """A manager that is shut down, with every request's cache freed, when the block exits.""" + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + mgr = factory() + try: + yield mgr + finally: + stream = getattr(mgr, "_stream", None) + if stream is not None: + stream.synchronize() + mgr.shutdown() + del mgr + gc.collect() + torch.cuda.empty_cache() + + +@pytest.fixture +def real_manager(): + """A small ``KVCacheManagerV2`` with block reuse on and its own execution stream, shut down when + the block exits: ``with real_manager(windows=None, max_tokens=2048, ...) as mgr:``.""" + + def factory(**kwargs): + return _managed(lambda: make_manager(**kwargs)) + + return factory + + +@contextmanager +def _host_tiered(factory: Callable[[], object]) -> Iterator[object]: + """``_managed``, checked to have a host tier below its GPU pool.""" + with _managed(factory) as mgr: + tiers = [str(tier) for tier in mgr.impl.cache_tier_list] + assert len(tiers) == 2 and "HOST" in tiers[1], f"no host tier below the GPU pool: {tiers}" + yield mgr + + +@pytest.fixture +def host_tier_manager(): + """``with host_tier_manager(**kwargs) as mgr:`` like ``real_manager`` on a ``POOL_TOKENS`` pool, + with a host tier of ``HOST_CACHE_BYTES`` checked present. Index slots for many one-block + requests and resume allowed up to a full pool let other requests push pages to host and back.""" + + def factory(**kwargs): + config = dict( + max_tokens=POOL_TOKENS, + max_batch_size=64, + max_util_for_resume=1.0, + host_cache_size=HOST_CACHE_BYTES, + ) + config.update(kwargs) + return _host_tiered(lambda: make_manager(**config)) + + return factory + + +@pytest.fixture +def deepseek_v4_manager(): + """``with deepseek_v4_manager(spec_config=None, fp8_ds_mla=False) as mgr:`` a small + ``DeepseekV4CacheManager``; callers skip before Blackwell.""" + + def factory(spec_config=None, fp8_ds_mla=False): + return _managed(lambda: _make_deepseek_v4_manager(spec_config, fp8_ds_mla)) + + return factory + + +@pytest.fixture +def hybrid_manager(): + """``with hybrid_manager() as mgr:`` a manager whose layer groups include recurrent state.""" + + def factory(): + return _managed(_make_hybrid_manager) + + return factory + + +# -- requests --------------------------------------------------------------------------------- + + +def make_request(request_id: int, tokens: Sequence[int], **kwargs): + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig + + return LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=list(tokens), + sampling_config=SamplingConfig(1), + is_streaming=False, + **kwargs, + ) + + +def kv(mgr, request): + """The request's runtime cache, or ``None``: read straight from the manager's map.""" + return mgr.kv_cache_map.get(request.py_request_id) + + +def closed(kv_cache) -> bool: + return kv_cache.status == kv_cache.Status.CLOSED + + +def chain_keys(kv_cache, tokens: Sequence) -> List[bytes]: + """The radix keys of every whole block of ``tokens``, computed independently of the lender.""" + from tensorrt_llm.runtime.kv_cache_manager_v2 import sequence_to_blockchain_keys + + tpb = kv_cache.tokens_per_block + keys = [] + for i, (block, key) in enumerate( + sequence_to_blockchain_keys(tpb, kv_cache.reuse_scope, list(tokens)) + ): + if i and len(block) == tpb: + keys.append(bytes(key)) + return keys + + +def fake_page(lg: int, key: bytes, nbytes: int) -> bytes: + """Deterministic page content named by (layer group, block key): equal tokens, equal bytes.""" + seed = int.from_bytes(hashlib.sha256(bytes([lg]) + key).digest()[:8], "little") + return np.random.default_rng(seed).integers(0, 256, nbytes, dtype=np.uint8).tobytes() + + +def pages(kv_cache, lg: int) -> List[int]: + return [int(p) for p in kv_cache.get_aggregated_page_indices(lg, valid_only=False)] + + +def num_layer_groups(mgr) -> int: + return len(mgr.impl.layer_grouping) + + +def pool_group_of(mgr) -> List[int]: + """Device pool group of each layer group.""" + return [int(g) for g in mgr.impl.get_life_cycle_pool_group_indices()] + + +def pool_group_ids(mgr) -> List[int]: + """Device pool groups in index order: the order of a staging lender's regions.""" + return sorted(int(pg.pool_group_index) for pg in mgr.impl.pool_group_descs) + + +def windows(mgr) -> Tuple[Optional[int], ...]: + """Each layer group's sliding window in tokens, ``None`` for full attention.""" + out = [] + for lc in mgr._life_cycle_by_layer_group(): + window = getattr(lc, "window_size", None) + out.append(None if window is None or window >= MAX_SEQ_LEN else int(window)) + return tuple(out) + + +def stale_blocks(mgr, lg: int, history: int) -> Tuple[int, int]: + beg, end = mgr._stale_block_range(lg, history) + return int(beg), int(end) + + +class DevicePages: + """Device pages addressed by (layer group, slot), straight from ``impl.pool_group_descs``; all + work runs on the manager's stream.""" + + def __init__(self, mgr): + from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor + + self._mgr = mgr + self._group_of = pool_group_of(mgr) + self._pools: Dict[int, List[torch.Tensor]] = {} + for pg in mgr.impl.pool_group_descs: + ordered = sorted(pg.pools, key=lambda p: int(p.pool_index)) + self._pools[int(pg.pool_group_index)] = [ + convert_to_torch_tensor( + TensorWrapper( + int(p.base_address), + torch.uint8, + shape=(int(pg.num_slots), int(p.slot_bytes)), + ) + ) + for p in ordered + ] + # Loading the fill kernel waits for the whole device, so load it before any stream is held. + torch.empty(1, dtype=torch.uint8, device="cuda").fill_(0) + torch.cuda.synchronize() + + def group_page_bytes(self, group: int) -> int: + return sum(t.shape[1] for t in self._pools[group]) + + def page_bytes(self, lg: int) -> int: + return self.group_page_bytes(self._group_of[lg]) + + def read(self, lg: int, slot: int) -> bytes: + self._mgr._stream.synchronize() + tensors = self._pools[self._group_of[lg]] + return b"".join(t[slot].cpu().numpy().tobytes() for t in tensors) + + def write(self, lg: int, slot: int, data: bytes) -> None: + stream = self._mgr._stream + src = torch.frombuffer(bytearray(data), dtype=torch.uint8) + offset = 0 + with torch.cuda.stream(stream): + for t in self._pools[self._group_of[lg]]: + width = t.shape[1] + t[slot].copy_(src[offset : offset + width].to(t.device)) + offset += width + stream.synchronize() + + def fill_async(self, lg: int, slot: int, byte: int) -> None: + """Fill a page on the manager's stream, not waiting: ordered after work already queued.""" + with torch.cuda.stream(self._mgr._stream): + for t in self._pools[self._group_of[lg]]: + t[slot].fill_(byte) + + +def page(mgr, request, lg: int, ordinal: int) -> bytes: + return DevicePages(mgr).read(lg, pages(kv(mgr, request), lg)[ordinal]) + + +def write_fake_kv(mgr, request, first_block: int = 0) -> None: + """Fill every page the request holds from ``first_block`` on, as a forward pass would.""" + kv_cache = kv(mgr, request) + dev = DevicePages(mgr) + keys = chain_keys(kv_cache, request.get_tokens(0)) + for lg in range(num_layer_groups(mgr)): + for ordinal, slot in enumerate(pages(kv_cache, lg)): + if ordinal < first_block or slot < 0 or ordinal >= len(keys): + continue + dev.write(lg, slot, fake_page(lg, keys[ordinal], dev.page_bytes(lg))) + + +def fill_sentinel(mgr, request, first_block: int = 0) -> None: + """Mark the request's pages from ``first_block`` on so a missing copy shows.""" + kv_cache = kv(mgr, request) + dev = DevicePages(mgr) + for lg in range(num_layer_groups(mgr)): + for ordinal, slot in enumerate(pages(kv_cache, lg)): + if ordinal >= first_block and slot >= 0: + dev.write(lg, slot, bytes([SENTINEL]) * dev.page_bytes(lg)) + + +def prefill(mgr, request) -> None: + """Run a whole prompt the way the executor does: allocate, compute (fake KV), advance, then + ``update_context_resources``, which commits the blocks.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + assert mgr.prepare_context(request) + kv_cache = kv(mgr, request) + # Real pages for every block written below; scratch slots would be overwritten. + kv_cache.enable_swa_scratch_reuse = False + first = request.context_current_position // mgr.tokens_per_block + assert mgr.resize_context(request, request.context_remaining_length) + write_fake_kv(mgr, request, first) + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + mgr.update_context_resources(batch) + assert request.context_remaining_length == 0 + assert kv_cache.num_committed_tokens == request.prompt_len + + +def published(mgr, request_id: int, tokens: Sequence[int], **kwargs): + """A request that computed ``tokens``: its blocks are committed.""" + request = make_request(request_id, tokens, **kwargs) + prefill(mgr, request) + return request + + +def admitted(mgr, request_id: int, tokens: Sequence[int]): + """A request admitted to fetch its prompt: a cache holding its local match and nothing more, + scratch reuse off as a fetch target needs. The lender grows it.""" + request = make_request(request_id, tokens) + assert mgr.prepare_context(request) + kv(mgr, request).enable_swa_scratch_reuse = False + return request + + +def host_bytes(address: int, length: int) -> bytes: + return ctypes.string_at(address, length) + + +def digest(value): + """``value`` with every byte string replaced by its length and hash, or its one repeated byte: + a failed comparison then names the entries that differ without diffing page contents.""" + if isinstance(value, (bytes, bytearray)): + if value and value.count(value[:1]) == len(value): + return f"{len(value)} x {value[0]:#04x}" + return f"{len(value)} B sha256 {hashlib.sha256(value).hexdigest()[:16]}" + if isinstance(value, dict): + return {key: digest(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [digest(item) for item in value] + return value + + +def staged_row(byte: int, layer_group: int, ordinal: int, nbytes: int) -> bytes: + """What ``stage`` writes for block ``ordinal`` of ``layer_group``: ``nbytes`` bytes drawn from a + generator seeded by ``byte`` and the row, so no two rows hold the same bytes, and a row landing + one slot over, or two pieces of a row trading places, shows, also in part of a page.""" + name = bytes([byte, layer_group]) + int(ordinal).to_bytes(8, "little") + seed = int.from_bytes(hashlib.sha256(name).digest()[:8], "little") + return np.random.default_rng(seed).integers(0, 256, nbytes, dtype=np.uint8).tobytes() + + +def stage(lender, view, byte: int) -> None: + """Fill every slot of a staging view as a backend's transfer would, each block entry with its + own bytes (``staged_row``).""" + for blocks in view.groups: + assert blocks.layer_group_index in lender.regions[blocks.region_index].layer_group_indices + length = lender.regions[blocks.region_index].slot_bytes + for ordinal, address in zip(blocks.block_ordinals.tolist(), blocks.addresses.tolist()): + ctypes.memmove( + address, staged_row(byte, blocks.layer_group_index, ordinal, length), length + ) + + +def relay(src_lender, src_view, dst_lender, dst_view, deliver=None) -> Tuple[np.ndarray, ...]: + """A backend between two staging lenders: each destination block entry takes the slot of the + source block entry of the same object key. ``deliver(group_index, entry_index)`` picks block + entries (default all); returns the masks.""" + by_object_key = {} + for blocks in src_view.groups: + assert ( + blocks.layer_group_index in src_lender.regions[blocks.region_index].layer_group_indices + ) + length = src_lender.regions[blocks.region_index].slot_bytes + for object_key, address in zip(blocks.object_keys, blocks.addresses.tolist()): + by_object_key[object_key.tobytes()] = (address, length) + masks = [] + for i, blocks in enumerate(dst_view.groups): + assert ( + blocks.layer_group_index in dst_lender.regions[blocks.region_index].layer_group_indices + ) + mask = np.zeros(len(blocks), dtype=bool) + for entry_index, (object_key, address) in enumerate( + zip(blocks.object_keys, blocks.addresses.tolist()) + ): + if deliver is not None and not deliver(i, entry_index): + continue + src, length = by_object_key[object_key.tobytes()] + assert length == dst_lender.regions[blocks.region_index].slot_bytes + ctypes.memmove(address, src, length) + mask[entry_index] = True + masks.append(mask) + return tuple(masks) + + +# -- streams ---------------------------------------------------------------------------------- + + +class _Gate: + def __init__(self, flag: torch.Tensor): + self._flag = flag + self.opened = False + self.tripped_by_watchdog = False + + def open(self) -> None: + self._flag.numpy()[0] = 1 + self.opened = True + + +def _wait_on_flag(stream: torch.cuda.Stream) -> _Gate: + from cuda.bindings import driver + + flag = torch.zeros(1, dtype=torch.int32, pin_memory=True) + (err,) = driver.cuStreamWaitValue32( + driver.CUstream(stream.cuda_stream), + driver.CUdeviceptr(flag.data_ptr()), + 1, + driver.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_GEQ, + ) + if err == driver.CUresult.CUDA_ERROR_NOT_SUPPORTED: + pytest.skip("this device cannot hold a stream on a host flag") + assert err == driver.CUresult.CUDA_SUCCESS, err + return _Gate(flag) + + +@contextmanager +def held_stream(stream: torch.cuda.Stream, strict: bool = True, watchdog_seconds: float = 120.0): + """Hold ``stream`` on a host flag until ``gate.open()``: work queued meanwhile does not start. + ``strict`` makes a torch host sync raise; a watchdog opens the gate if the test never does, and + the block then fails, so a host wait on the stream cannot pass unnoticed.""" + gate = _wait_on_flag(stream) + + def trip(): + gate.tripped_by_watchdog = True + gate.open() + + watchdog = threading.Timer(watchdog_seconds, trip) + watchdog.start() + previous = torch.cuda.get_sync_debug_mode() + if strict: + torch.cuda.set_sync_debug_mode("error") + try: + yield gate + finally: + torch.cuda.set_sync_debug_mode(previous) + gate.open() + watchdog.cancel() + watchdog.join() + stream.synchronize() + assert not gate.tripped_by_watchdog, "the stream stayed held: something waited on it" + + +@contextmanager +def gated_stream(stream: torch.cuda.Stream, open_after: float): + """Hold ``stream`` until a timer opens it after ``open_after`` seconds; a host wait on the + stream then simply lasts that long. ``gate.opened`` tells whether it has.""" + gate = _wait_on_flag(stream) + timer = threading.Timer(open_after, gate.open) + timer.start() + try: + yield gate + finally: + timer.cancel() + gate.open() + timer.join() + stream.synchronize() + + +# -- pool pressure ---------------------------------------------------------------------------- + + +def pool_pages(mgr) -> int: + (group,) = mgr.impl.pool_group_descs + return int(group.num_slots) + + +def warm_tier_moves(mgr) -> None: + """Move one committed block to host and back. The first such move in a process loads the kernels + that copy between tiers, which waits for the whole device; a held stream would stall it.""" + tokens = list(range(40_000, 40_000 + TPB + 1)) + mgr.free_resources(published(mgr, 900, tokens)) + others = Requests(mgr) + try: + while others.allocate(1, chunk=1): # until the pool is full: the block went to host + pass + finally: + others.free() + back = make_request(901, tokens) + assert mgr.prepare_context(back) + assert kv(mgr, back).num_committed_tokens == TPB, "the block did not come back from host" + mgr.free_resources(back) + torch.cuda.synchronize() + + +def tier_used(mgr, level: int) -> int: + """Slots in use at cache tier ``level`` (0 the GPU, 1 the host) over every pool group, from + the manager's storage statistics.""" + return sum(int(s.total) - int(s.free) for s in mgr.impl.get_storage_statistics(level)) + + +class Requests: + """Other requests' allocations: private pages taken straight from the pool, sentinel-filled, + a few blocks per request (a request holds at most ``MAX_SEQ_LEN`` tokens).""" + + def __init__(self, mgr, dev: Optional[DevicePages] = None): + # ``dev`` built ahead lets ``fill=False`` run while the stream is held: building the page + # views waits on the stream. + self._mgr = mgr + self._dev = dev + self._next = 100 + self.held = [] + + def allocate(self, blocks: int, fill: bool = True, chunk: Optional[int] = None) -> bool: + """Take ``blocks`` pages, ``chunk`` blocks per request. ``False`` if the pool could not give + them all; what it gave stays held. ``fill=False`` fills on the stream without waiting.""" + mgr = self._mgr + most = chunk or MAX_SEQ_LEN // TPB - 1 + while blocks: + chunk = min(blocks, most) + rid = self._next + self._next += 1 + first = 50_000 + 1_000 * rid # never matches another request's blocks + request = make_request(rid, list(range(first, first + chunk * TPB))) + if not ( + mgr.prepare_context(request) + and mgr.resize_context(request, request.context_remaining_length) + ): + mgr.free_resources(request) + return False + if fill: + fill_sentinel(mgr, request) + else: + dev = self._dev or DevicePages(mgr) + for lg in range(num_layer_groups(mgr)): + for slot in pages(kv(mgr, request), lg): + if slot >= 0: + dev.fill_async(lg, slot, SENTINEL) + self.held.append(request) + blocks -= chunk + return True + + def pages(self) -> set: + return {p for r in self.held for p in pages(kv(self._mgr, r), 0) if p >= 0} + + def free_all_but(self, keep: set) -> None: + """Free the requests holding none of the pages in ``keep``.""" + kept = [] + for request in self.held: + if keep & {p for p in pages(kv(self._mgr, request), 0) if p >= 0}: + kept.append(request) + else: + self._mgr.free_resources(request) + self.held = kept + + def free(self) -> None: + for request in self.held: + self._mgr.free_resources(request) + self.held = [] + + +def overwrite_free_pages(mgr) -> int: + """Other requests take every free page one block at a time, fill it with ``SENTINEL`` and free + it again, so whatever a freed page held is gone. Returns how many pages they took.""" + others = Requests(mgr) + try: + while others.allocate(1, chunk=1): + pass + return len(others.held) + finally: + others.free() + + +# -- memory kept until exit ------------------------------------------------------------------- + + +def retained() -> Tuple[object, ...]: + """What the lender keeps until the process exits, through its test hook.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + return _lender._retained() + + +def staging_memory(regions) -> List[object]: + """The kept host allocations holding any of ``regions``, compared by address.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + found = [] + for obj in retained(): + if not isinstance(obj, _lender._HostMemory): + continue + begin, end = obj.address, obj.address + obj.nbytes + if any(begin <= region.address < end for region in regions): + found.append(obj) + return found + + +# -- the manager's host page-index buffer ------------------------------------------------------ + +CANARY = 0x5A5A5A5A + + +class IndexRow(SimpleNamespace): + """Where a request's cache writes its base page indices of one pool in the manager's host + page-index buffer: the buffer's address, size and pinning, and the row's int32 offset in it.""" + + +def index_row(mgr, request, pool: int = 0, beam: int = 0) -> IndexRow: + buf = mgr.host_kv_cache_block_offsets + index = mgr.index_mapper.get_index(int(request.py_request_id)) + row = buf[pool, index * mgr.max_beam_width + beam, 0] + return IndexRow( + address=buf.data_ptr(), + nbytes=buf.numel() * buf.element_size(), + pinned=buf.is_pinned(), + offset=(row.data_ptr() - buf.data_ptr()) // 4, + length=row.numel(), + values=row.tolist(), + ) + + +def reclaim(row: IndexRow, tries: int = 64) -> Optional[torch.Tensor]: + """A new host int32 tensor at a freed buffer's address, every element ``CANARY``; ``None`` if + the allocator hands that address to none of ``tries`` same-size allocations.""" + held = [] + for _ in range(tries): + t = torch.empty(row.nbytes // 4, dtype=torch.int32, pin_memory=row.pinned) + if t.data_ptr() == row.address: + t.fill_(CANARY) + return t + held.append(t) + return None + + +def canary_written(canary: Optional[torch.Tensor], row: IndexRow) -> List[Tuple[int, int]]: + """``(cell, value)`` of each cell of ``row`` in ``canary`` that holds a value other than + ``CANARY``.""" + if canary is None: + return [] + cells = canary[row.offset : row.offset + row.length].numpy() + return [(int(i), int(cells[i])) for i in np.nonzero(cells != np.int32(CANARY))[0]] + + +def staging_kept(regions) -> bool: + """Whether every region lies in memory the lender keeps.""" + kept = staging_memory(regions) + return all( + any(m.address <= p.address and p.address + p.nbytes <= m.address + m.nbytes for m in kept) + for p in regions + ) + + +def pinned_range(address: int) -> Optional[Tuple[int, int]]: + """``(start, nbytes)`` of the page-locked host allocation the driver records at ``address``; + ``None`` when there is none, as after it was freed.""" + from cuda.bindings import driver + + attribute = driver.CUpointer_attribute + err, kind = driver.cuPointerGetAttribute(attribute.CU_POINTER_ATTRIBUTE_MEMORY_TYPE, address) + if err != driver.CUresult.CUDA_SUCCESS: + return None + assert kind == driver.CUmemorytype.CU_MEMORYTYPE_HOST, kind + err, start = driver.cuPointerGetAttribute( + attribute.CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, address + ) + assert err == driver.CUresult.CUDA_SUCCESS, err + err, nbytes = driver.cuPointerGetAttribute(attribute.CU_POINTER_ATTRIBUTE_RANGE_SIZE, address) + assert err == driver.CUresult.CUDA_SUCCESS, err + return int(start), int(nbytes) + + +def lender_warnings(monkeypatch) -> List[str]: + """Collects the warnings the lender module logs from now on.""" + from tensorrt_llm.logger import logger + + seen: List[str] = [] + + def spy(original): + def warn(*args, **kwargs): + caller = inspect.currentframe().f_back + if caller is not None and caller.f_globals.get("__name__", "").startswith(_SHARING): + seen.append(" ".join(str(a) for a in args)) + return original(*args, **kwargs) + + return warn + + for name in ("warning", "warning_once"): + monkeypatch.setattr(logger, name, spy(getattr(logger, name))) + return seen + + +# -- threads ---------------------------------------------------------------------------------- + + +def on_thread(call: Callable[[], object], name: str = "not-the-caller", cuda: bool = False) -> dict: + """Run ``call`` on a fresh thread, joined before returning; what it returned or raised. With + ``cuda`` the thread first synchronizes the device: like the executor's, it used CUDA before.""" + box: dict = {} + + def run(): + try: + if cuda: + torch.cuda.synchronize() + box["value"] = call() + except Exception as error: # handed back to the test + box["error"] = error + + thread = threading.Thread(target=run, name=name) + thread.start() + thread.join(120) + assert not thread.is_alive() + box["thread"] = thread.ident + return box + + +_KIT = SimpleNamespace( + SENTINEL=SENTINEL, + POOL_TOKENS=POOL_TOKENS, + CANARY=CANARY, + DevicePages=DevicePages, + Requests=Requests, + admitted=admitted, + canary_written=canary_written, + chain_keys=chain_keys, + closed=closed, + digest=digest, + fake_page=fake_page, + fill_sentinel=fill_sentinel, + gated_stream=gated_stream, + held_stream=held_stream, + host_bytes=host_bytes, + index_row=index_row, + kv=kv, + lender_warnings=lender_warnings, + make_manager=make_manager, + make_request=make_request, + managed=_managed, + num_layer_groups=num_layer_groups, + on_thread=on_thread, + overwrite_free_pages=overwrite_free_pages, + page=page, + pages=pages, + pinned_range=pinned_range, + pool_group_ids=pool_group_ids, + pool_group_of=pool_group_of, + pool_pages=pool_pages, + prefill=prefill, + published=published, + reclaim=reclaim, + relay=relay, + retained=retained, + stage=stage, + staged_row=staged_row, + stale_blocks=stale_blocks, + staging_kept=staging_kept, + staging_memory=staging_memory, + tier_used=tier_used, + warm_tier_moves=warm_tier_moves, + windows=windows, +) + + +@pytest.fixture +def kit(): + """The helpers and oracles above, for test modules that cannot import this file by name.""" + return _KIT + + +@pytest.hookimpl(wrapper=True) +def pytest_runtest_makereport( + item: pytest.Item, call: pytest.CallInfo[None] +) -> Generator[None, pytest.TestReport, pytest.TestReport]: + report = yield + if call.excinfo is not None or report.failed: + item.stash[_FAILED_TEST] = True + return report + + +@pytest.fixture(autouse=True) +def _let_go_of_what_the_test_kept(request: pytest.FixtureRequest) -> Iterator[None]: + """Free and drop the test's additions to the lender's keep list. A test that never imported + the lender kept nothing, and one that added nothing needs no collection: no finalizer adds + to it or drops from it.""" + lender = sys.modules.get(f"{_SHARING}._lender") + before = lender._retained() if lender is not None else () + yield + # Failure tracebacks can still hold native caches that write into these retained buffers. + if request.node.stash.get(_FAILED_TEST, False): + return + lender = sys.modules.get(f"{_SHARING}._lender") + if lender is None: + return + added = [owner for owner in lender._retained() if not any(owner is b for b in before)] + if not added: + return + # A lease or lender a cycle still holds closes its cache while what the cache writes to is kept. + gc.collect() + if torch.cuda.is_initialized(): + torch.cuda.synchronize() + while added: + owner = added.pop() # newest first + if isinstance(owner, lender._HostMemory): + owner.free() + lender._stop_retaining(owner) + gc.collect() + + +@pytest.fixture +def retained_cleanup() -> Callable[[bool], Iterator[None]]: + """A fresh teardown generator for tests of the autouse cleanup itself.""" + + def start(failed: bool = False) -> Iterator[None]: + request = SimpleNamespace(node=SimpleNamespace(stash={_FAILED_TEST: failed})) + return _let_go_of_what_the_test_kept.__wrapped__(request) + + return start diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_host_tier.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_host_tier.py new file mode 100644 index 000000000000..429c7e5c9571 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_host_tier.py @@ -0,0 +1,938 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The staging lender beside the KV cache's host tier. Each check first proves its move: pages to +host and back by their request's suspension or others' pressure, bystander blocks to host while lent +pages stay, lent pages to other slots by a pool rebalance, held and late-marked pages to host.""" + +import ctypes +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import StagingOptions, attach_lender + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") + +TPB = 32 +WINDOW = 64 +PROMPT = list(range(1000, 1097)) # three whole blocks and one token +END = 96 +BLOCKS = END // TPB +BYSTANDERS = [list(range(3000 + 200 * i, 3097 + 200 * i)) for i in range(2)] +WINDOWED_PROMPT = list(range(2000, 2161)) # five whole blocks and one token +WINDOWED_END = 160 +SOURCE, TARGET, SECOND = 1, 2, 3 +SCOPE = b"host-tier-suite" + + +def attach(mgr, *, tokens_per_fetch=END, planned_fetches=1): + return attach_lender( + mgr, scope=SCOPE, staging=StagingOptions(tokens_per_fetch, planned_fetches) + ) + + +def ready_view(lease, mgr, tries=3): + for _ in range(tries): + mgr._stream.synchronize() + view = lease.poll() + if view is not None: + return view + raise AssertionError(f"the lease never became ready: {lease.failure}") + + +def staged(lender, view): + """Per run, the bytes of every row's staging slot.""" + out = [] + for run in view.groups: + length = lender.regions[run.region_index].slot_bytes + out.append([ctypes.string_at(a, length) for a in run.addresses.tolist()]) + return out + + +def names(view): + return [[name.tobytes() for name in run.object_keys] for run in view.groups] + + +@contextmanager +def to_host_and_back(kit, mgr, request, dev=None): + """Suspend ``request``; one-block requests take every GPU page, its pages go to host; free the + others but those on its old pages, resume it. Yields those others and the old pages. With + ``dev`` the others fill their pages on the stream without waiting for it.""" + kv = kit.kv(mgr, request) + old = [p for p in kit.pages(kv, 0) if p >= 0] + host_before = kit.tier_used(mgr, 1) + others = kit.Requests(mgr, dev) + try: + mgr.suspend_request(request) + assert others.allocate(kit.pool_pages(mgr), fill=dev is None, chunk=1) + assert set(old) <= others.pages(), "the request's pages stayed on the GPU" + assert kit.tier_used(mgr, 1) >= host_before + len(old), "its pages did not reach host" + others.free_all_but(set(old)) + assert mgr.resume_request(request) + now = [p for p in kit.pages(kv, 0) if p >= 0] + assert len(now) == len(old) and not set(now) & set(old), "it came back to its old pages" + yield others, old + finally: + others.free() + + +# -- (a) a fetch and a later publish across the request's own trip to host --------------------- + + +def test_a_fetch_stays_usable_across_its_request_s_trip_to_host( + kit, host_tier_manager, attach=attach +): + with host_tier_manager() as mgr_a, host_tier_manager() as mgr_b: + source = kit.published(mgr_a, SOURCE, PROMPT) + lender_a = attach(mgr_a) + publish = lender_a.lend_read(source, 0, END) + publish_view = ready_view(publish, mgr_a) + target = kit.admitted(mgr_b, TARGET, PROMPT) + lender_b = attach(mgr_b) + lease = lender_b.lend_write(target, 0, END) + view = lease.poll() + assert view is not None + kit.fill_sentinel(mgr_b, target) + masks = kit.relay(lender_a, publish_view, lender_b, view) + dev = kit.DevicePages(mgr_b) + kit.warm_tier_moves(mgr_b) + with kit.held_stream(mgr_b._stream, strict=False) as gate: + # The copy into the lent pages waits behind the gate; the suspension and the move to + # host come after it on the stream, so the move carries the fetched bytes. + lease.mark_arrived(masks) + lease.release() + with to_host_and_back(kit, mgr_b, target, dev) as (others, old): + assert lender_b.resume_bounds(target) is None, "settled before its copy ran" + gate.open() + mgr_b._stream.synchronize() + assert lender_b.resume_bounds(target) == (END, 0), "not usable once its copy ran" + for ordinal in range(BLOCKS): + got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) + want = kit.digest(kit.page(mgr_a, source, 0, ordinal)) + assert got == want, "the fetched bytes are lost" + sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) + assert all(dev.read(0, slot) == sentinel for slot in old) + publish.release() + + +def test_a_publish_after_a_trip_to_host_reads_the_pages_where_they_are( + kit, host_tier_manager, attach=attach +): + with host_tier_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] + lender = attach(mgr) + first = lender.lend_read(source, 0, END) + assert kit.digest(staged(lender, ready_view(first, mgr))) == kit.digest([original]) + first.release() + with to_host_and_back(kit, mgr, source): + again = lender.lend_read(source, 0, END) + got = kit.digest(staged(lender, ready_view(again, mgr))) + assert got == kit.digest([original]), "read a stale page" + again.release() + + +# -- (b) other requests' pressure while a fetch is lent ----------------------------------------- + + +@pytest.mark.parametrize("moment", ["copy_queued", "before_marks"]) +def test_other_requests_pushing_blocks_to_host_leave_a_fetch_where_its_resume_bounds_says( + kit, host_tier_manager, moment, attach=attach +): + """``moment``: ``copy_queued``, the target stays active and its marks' copy waits on the stream + while other requests push committed blocks to host; ``before_marks``, the target itself goes to + host and back between the grant and the marks.""" + with host_tier_manager() as mgr_a, host_tier_manager() as mgr_b: + source = kit.published(mgr_a, SOURCE, PROMPT) + lender_a = attach(mgr_a) + publish = lender_a.lend_read(source, 0, END) + publish_view = ready_view(publish, mgr_a) + original = [kit.page(mgr_a, source, 0, o) for o in range(BLOCKS)] + got = kit.digest(staged(lender_a, publish_view)) + assert got == kit.digest([original]), "the relay carries rows off their blocks" + kit.warm_tier_moves(mgr_b) + for i, tokens in enumerate(BYSTANDERS): # committed blocks, then no request holds them + mgr_b.free_resources(kit.published(mgr_b, 10 + i, tokens)) + target = kit.admitted(mgr_b, TARGET, PROMPT) + lender_b = attach(mgr_b) + lease = lender_b.lend_write(target, 0, END) + view = lease.poll() + assert view is not None + kit.fill_sentinel(mgr_b, target) + masks = kit.relay(lender_a, publish_view, lender_b, view) + kv = kit.kv(mgr_b, target) + lent = kit.pages(kv, 0)[:BLOCKS] + dev = kit.DevicePages(mgr_b) + sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) + if moment == "copy_queued": + others = kit.Requests(mgr_b, dev) + host_before = kit.tier_used(mgr_b, 1) + try: + with kit.held_stream(mgr_b._stream, strict=False) as gate: + lease.mark_arrived(masks) + lease.release() + assert others.allocate(kit.pool_pages(mgr_b) - BLOCKS, fill=False, chunk=1) + moved = kit.tier_used(mgr_b, 1) - host_before + assert moved >= len(BYSTANDERS) * BLOCKS, "no bystander block reached host" + + assert kit.pages(kv, 0)[:BLOCKS] == lent, "the lent pages moved" + assert lender_b.resume_bounds(target) is None, "settled before its copy ran" + gate.open() + assert lender_b.resume_bounds(target) == (END, 0) + assert all(dev.read(0, slot) == sentinel for slot in others.pages()) + finally: + others.free() + else: + with to_host_and_back(kit, mgr_b, target) as (others, old): + assert old == lent + lease.mark_arrived(masks) + lease.release() + mgr_b._stream.synchronize() + assert all(dev.read(0, slot) == sentinel for slot in old), ( + "the marks' copy reached another request's page" + ) + assert lender_b.resume_bounds(target) == (0, 0), "counts rows that never landed" + for ordinal in range(BLOCKS): + got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) + expected = kit.page(mgr_a, source, 0, ordinal) if moment == "copy_queued" else sentinel + assert got == kit.digest(expected), "a page does not hold what its fetch says" + publish.release() + + +# -- (c) a cache whose pages are on host fails every lend at the call --------------------------- + +SERVED = "served bytes that are not its blocks'" +REACHED = "reached another request's page" + + +def use_as_a_backend(kit, mgr, lender, lease, reading, original): + """What a backend does with a lease that became ready: a read's rows are compared with the + blocks' bytes, a write's rows are filled.""" + mgr._stream.synchronize() + view = lease.poll() + if view is None: + return + if reading: + (run,) = view.groups + want = [[original[o] for o in run.block_ordinals.tolist()]] + assert kit.digest(staged(lender, view)) == kit.digest(want), SERVED + else: + kit.stage(lender, view, 0x5A) + lease.mark_arrived(view.row_masks(True)) + + +def test_a_cache_on_host_fails_every_lend_at_the_call_and_lends_again_once_resumed( + kit, host_tier_manager, to_host=True +): + """Both caches go to host before the lends. ``to_host=False`` serves only the counter-examples: + it leaves their pages on the GPU, for a lender's rule breach to show without the host tier, and + the real lender then fails the check where the resumed cache is on its old pages.""" + with host_tier_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] + target = kit.admitted(mgr, TARGET, BYSTANDERS[0]) + assert mgr.resize_context(target, target.context_remaining_length) + lender = attach(mgr, planned_fetches=2) + writer = target + kv, kv_writer = kit.kv(mgr, source), kit.kv(mgr, writer) + old = [p for p in kit.pages(kv, 0) if p >= 0] + old_writer = [p for p in kit.pages(kv_writer, 0) if p >= 0] + host_before = kit.tier_used(mgr, 1) + others = kit.Requests(mgr) + try: + mgr.suspend_request(target) + mgr.suspend_request(source) + if to_host: + assert others.allocate(kit.pool_pages(mgr), chunk=1) + assert set(old) | set(old_writer) <= others.pages() + moved = len(set(old) | set(old_writer)) + assert kit.tier_used(mgr, 1) >= host_before + moved, "its pages did not reach host" + dev = kit.DevicePages(mgr) + theirs = {slot: dev.read(0, slot) for slot in others.pages()} + failed = [lender.lend_read(source, 0, END), lender.lend_write(writer, 0, END)] + for lease, reading in zip(failed, (True, False)): + use_as_a_backend(kit, mgr, lender, lease, reading, original) + mgr._stream.synchronize() + now = {slot: dev.read(0, slot) for slot in theirs} + assert kit.digest(now) == kit.digest(theirs), REACHED + for lease in failed: + assert lease.failure is not None, "not failed at the call" + assert lease.poll() is None + groups = kit.pool_group_ids(mgr) + assert [lender._free_slots(g) for g in groups] == [p.slot_count for p in lender.regions] + for lease in failed: + lease.release() + others.free_all_but(set(old)) + assert mgr.resume_request(source) and mgr.resume_request(target) + now = [p for p in kit.pages(kv, 0) if p >= 0] + assert not set(now) & set(old), "it came back to its old pages" + read = lender.lend_read(source, 0, END) + view = ready_view(read, mgr) + assert kit.digest(staged(lender, view)) == kit.digest([original]) + write = lender.lend_write(writer, 0, END) + write_view = ready_view(write, mgr) + write.mark_arrived(write_view.row_masks()) + write.release() + read.release() + finally: + others.free() + + +# -- (d) a committed block's name does not depend on its tier ----------------------------------- + + +def test_a_committed_block_has_one_name_on_the_gpu_and_after_a_trip_to_host( + kit, host_tier_manager, attach=attach +): + with host_tier_manager() as mgr: + first = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + lease = lender.lend_read(first, 0, END) + view = ready_view(lease, mgr) + first_names, first_bytes = names(view), staged(lender, view) + lease.release() + old = kit.pages(kit.kv(mgr, first), 0)[:BLOCKS] + mgr.free_resources(first) # its whole blocks stay committed in the reuse tree + host_before = kit.tier_used(mgr, 1) + others = kit.Requests(mgr) + try: + assert others.allocate(kit.pool_pages(mgr), chunk=1) + assert set(old) <= others.pages() + assert kit.tier_used(mgr, 1) >= host_before + BLOCKS, "the blocks did not reach host" + others.free_all_but(set(old)) + second = kit.make_request(SECOND, PROMPT) + assert mgr.prepare_context(second) + kv = kit.kv(mgr, second) + assert kv.num_committed_tokens == END, "the blocks on host were not reused" + now = kit.pages(kv, 0)[:BLOCKS] + assert not set(now) & set(old), "the blocks came back to their old pages" + again = lender.lend_read(second, 0, END) + view = ready_view(again, mgr) + assert names(view) == first_names, "a block's name changed with its tier" + assert kit.digest(staged(lender, view)) == kit.digest(first_bytes) + again.release() + finally: + others.free() + + +# -- (f) a pool rebalance while a fetch is lent ------------------------------------------------- + + +class DeviceCopyGuard: + """The CUDA driver, recording every ``cuMemcpyAsync`` destination. With ``allowed``, a memcpy + into device memory outside ``allowed()`` is recorded as stray and not run: after a pool shrinks, + the slot it names may be unmapped.""" + + def __init__(self, real_driver, device_ranges=(), allowed=None): + self._real = real_driver + self._device = device_ranges # every byte a device pool ever held + self._allowed = allowed # () -> [(begin, end)] the copies may write + self.copies = [] + self.into_device = 0 + self.stray = [] + + @staticmethod + def _covered(begin, end, ranges): + for lo, hi in sorted(ranges): + if lo <= begin < hi: + begin = hi + if begin >= end: + return True + return False + + def __getattr__(self, name): + found = getattr(self._real, name) + if name != "cuMemcpyAsync": + return found + + def call(dst, src, nbytes, stream): + begin = int(dst) + self.copies.append(begin) + if self._allowed is not None and any(lo <= begin < hi for lo, hi in self._device): + self.into_device += 1 + if not self._covered(begin, begin + int(nbytes), self._allowed()): + self.stray.append((begin, int(nbytes))) + return (self._real.CUresult.CUDA_SUCCESS,) + return found(dst, src, nbytes, stream) + + return call + + +def page_ranges(kit, mgr, kv, layer_groups): + """Byte ranges of the pages ``kv`` locks in ``layer_groups``, in every pool of their groups.""" + group_of = {int(pg.pool_group_index): pg for pg in mgr.impl.pool_group_descs} + pool_group = kit.pool_group_of(mgr) + out = [] + for lg in layer_groups: + pg = group_of[pool_group[lg]] + slots = [int(s) for s in kv.get_base_page_indices(lg)[: kv.num_blocks] if int(s) >= 0] + for pool in pg.pools: + base, width = int(pool.base_address), int(pool.slot_bytes) + out.extend((base + s * width, base + (s + 1) * width) for s in slots) + return out + + +def test_a_rebalance_under_a_fetch_moves_no_copy_off_the_pages_still_lent( + kit, host_tier_manager, monkeypatch, attach=attach +): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection + + shape = dict(windows=[WINDOW, 256], num_layers=3, max_tokens=4 * kit.POOL_TOKENS) + with host_tier_manager(**shape) as mgr_a, host_tier_manager(**shape) as mgr_b: + source = kit.published(mgr_a, SOURCE, WINDOWED_PROMPT) + lender_a = attach(mgr_a, tokens_per_fetch=WINDOWED_END) + publish = lender_a.lend_read(source, 0, WINDOWED_END) + publish_view = ready_view(publish, mgr_a) + full = kit.windows(mgr_b).index(None) + pg_full = kit.pool_group_of(mgr_b)[full] + device = [] + for pg in mgr_b.impl.pool_group_descs: + for pool in pg.pools: + base = int(pool.base_address) + device.append((base, base + int(pg.num_slots) * int(pool.slot_bytes))) + # Fillers take the low slots, so the target's full-attention pages sit high in their pool. + fillers = kit.Requests(mgr_b) + stats = mgr_b.impl.get_storage_statistics + while stats(0)[pg_full].total - stats(0)[pg_full].free < 72 and fillers.allocate(7): + pass + target = kit.admitted(mgr_b, TARGET, WINDOWED_PROMPT) + lender_b = attach(mgr_b, tokens_per_fetch=WINDOWED_END) + lease = lender_b.lend_write(target, 0, WINDOWED_END) + view = lease.poll() + assert view is not None + kit.fill_sentinel(mgr_b, target) + masks = kit.relay(lender_a, publish_view, lender_b, view) + kv = kit.kv(mgr_b, target) + lent = [kit.pages(kv, lg) for lg in range(2)] + fillers.free() + # The executor's rebalance: suspend every active request, adjust the pools, resume them. + mgr_b.suspend_request(target) + _introspection.force_rebalance_precondition(mgr_b.impl, skew=0.05 if pg_full == 0 else 20) + assert mgr_b.impl.need_adjustment + mgr_b.impl.adjust() + assert mgr_b.resume_request(target) + now = [kit.pages(kv, lg) for lg in range(2)] + slots_now = { + int(pg.pool_group_index): int(pg.num_slots) for pg in mgr_b.impl.pool_group_descs + } + gone = [p for p in lent[full] if p >= slots_now[pg_full]] + assert now[full] != lent[full] and gone, "the rebalance left the lent pages where they were" + guard = DeviceCopyGuard(_lender.drv, device, lambda: page_ranges(kit, mgr_b, kv, range(2))) + monkeypatch.setattr(_lender, "drv", guard) + lease.mark_arrived(masks) + lease.release() + mgr_b._stream.synchronize() + assert guard.stray == [], "the marks' copy wrote outside the pages the target locks" + assert guard.into_device, "nothing was copied: the check saw no copy" + stayed = [ + (run.layer_group_index, o) + for run in view.groups + for o in run.block_ordinals.tolist() + if now[run.layer_group_index][o] == lent[run.layer_group_index][o] >= 0 + ] + wrong = [ + f"{lg}:{o}" + for lg, o in stayed + if kit.page(mgr_b, target, lg, o) != kit.page(mgr_a, source, lg, o) + ] + assert stayed and not wrong, ( + f"rows {wrong} of {stayed}, whose pages stayed, hold other bytes" + ) + assert lender_b.resume_bounds(target) == (0, WINDOWED_END), "counts rows that never landed" + publish.release() + + +# -- (f2) a pool rebalance with every kind of lease in flight ------------------------------------ + +REB_END = 96 # three whole blocks +REB_PROMPTS = { + "queued_read": list(range(11_000, 11_097)), # a read whose copy is queued behind the stream + "waiting_read": list(range(12_000, 12_097)), # a read waiting for slots + "unmarked_write": list(range(13_000, 13_097)), # a write granted, marked after the rebalance + "marked_write": list(range(14_000, 14_097)), # a write marked before, its copy queued +} +REB_GATE_SECONDS = 4.0 +# One pool group, its quota halved with resize(level 0): every page above the new capacity moves +# or goes to host. Two pool groups, the target ratio moved so that adjust(), the executor's own +# rebalance call, shrinks the full-attention group. +MECHANISMS = { + "resize": dict(max_tokens=2048), + "adjust": dict(windows=[WINDOW, 256], num_layers=3, max_tokens=2048), +} +STRAY = "a copy touched device memory outside the leased pages" + + +class PageCopyGuard: + """The CUDA driver, recording every ``cuMemcpyAsync``. A copy whose device side lies in a device + pool but outside ``allowed()`` (the pages the leased requests lock at the call) is stray; with + ``block`` it is recorded and not run, since its slot may be unmapped.""" + + def __init__(self, real_driver, device_ranges, allowed, block): + self._real = real_driver + self._device = device_ranges + self._allowed = allowed + self._block = block + self.copies = 0 + self.stray = [] + + def __getattr__(self, name): + found = getattr(self._real, name) + if name != "cuMemcpyAsync": + return found + + def guarded(dst, src, nbytes, stream): + self.copies += 1 + for side, address in (("dst", int(dst)), ("src", int(src))): + if not any(lo <= address < hi for lo, hi in self._device): + continue + if not DeviceCopyGuard._covered(address, address + int(nbytes), self._allowed()): + self.stray.append((side, address, int(nbytes))) + if self._block: + return (self._real.CUresult.CUDA_SUCCESS,) + return found(dst, src, nbytes, stream) + + return guarded + + +def gpu_slot(kv, lg, ordinal): + """The GPU slot the active cache ``kv`` locks for block ``ordinal`` of ``lg``, or -1.""" + if kv is None or not kv.is_active or ordinal >= kv.num_blocks: + return -1 + return int(kv.get_base_page_indices(lg)[ordinal]) + + +def published_bytes(kit, mgr, request, lg, ordinal, nbytes): + """What ``published`` wrote into block ``ordinal`` of ``lg``.""" + keys = kit.chain_keys(kit.kv(mgr, request), request.get_tokens(0)) + return kit.fake_page(lg, keys[ordinal], nbytes) + + +def staged_right(kit, mgr, lender, view, request): + """Per row of a read's view, whether its slot holds the block's published bytes.""" + out = {} + for run in view.groups: + nbytes = lender.regions[run.region_index].slot_bytes + for ordinal, address in zip(run.block_ordinals.tolist(), run.addresses.tolist()): + want = published_bytes(kit, mgr, request, run.layer_group_index, ordinal, nbytes) + out[f"{run.layer_group_index}:{ordinal}"] = kit.host_bytes(address, nbytes) == want + return out + + +def landed(kit, mgr, request, rows, pub, pub_request, dev, dev_pub): + """Per row of ``rows``, whether the request's page holds the publisher's bytes; ``None`` where + its cache locks no page.""" + kv, kv_pub = kit.kv(mgr, request), kit.kv(pub, pub_request) + out = {} + for lg, ordinals in rows: + for o in ordinals: + slot, slot_pub = gpu_slot(kv, lg, o), gpu_slot(kv_pub, lg, o) + assert slot_pub >= 0, f"the publisher holds no page for {lg}:{o}" + out[f"{lg}:{o}"] = ( + None if slot < 0 else dev.read(lg, slot) == dev_pub.read(lg, slot_pub) + ) + return out + + +def read_at(kit, mgr, lg, position): + """The block ordinals a start at ``position`` reads in layer group ``lg``.""" + blocks = list(range(position // TPB)) + if kit.windows(mgr)[lg] is None: + return blocks + beg, end = kit.stale_blocks(mgr, lg, position) + return [o for o in blocks if not beg <= o < end] + + +def locked(kit, mgr, request): + kv = kit.kv(mgr, request) + return [ + [int(s) for s in kv.get_base_page_indices(lg)[: kv.num_blocks]] + for lg in range(kit.num_layer_groups(mgr)) + ] + + +def leased_ranges(kit, mgr, caches): + """Byte ranges of the pages the active ``caches`` lock, in every pool of their groups.""" + out = [] + for kv in caches: + if kv is not None and kv.is_active: + out.extend(page_ranges(kit, mgr, kv, range(kit.num_layer_groups(mgr)))) + return out + + +def rebalance_with_leases_in_flight( + kit, host_tier_manager, monkeypatch, mechanism, attach=attach, block=False, stay=False +): + """The executor's rebalance (suspend every cache, ``resize`` or ``adjust`` the pools, resume) + with four leases in flight: a read whose copy waits behind a held stream, a read waiting for + slots, a granted write marked after the rebalance and a write marked before, its copy queued. It + moves their pages, or with ``stay`` not the waiting read's and unmarked write's. A driver guard + sees every copy to or from a device pool; bytes are checked against what ``published`` wrote.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection + + shape = MECHANISMS[mechanism] + with host_tier_manager(**shape) as mgr_a, host_tier_manager(**shape) as mgr_b: + num_lg = kit.num_layer_groups(mgr_b) + full = kit.windows(mgr_b).index(None) + moving = kit.pool_group_of(mgr_b)[full] # the pool group the rebalance shrinks + # The publisher of what the writes fetch. + pubs = { + k: kit.published(mgr_a, 30 + i, REB_PROMPTS[k]) + for i, k in enumerate(("unmarked_write", "marked_write")) + } + lender_a = attach(mgr_a, tokens_per_fetch=REB_END, planned_fetches=2) + pub_leases = {k: lender_a.lend_read(p, 0, REB_END) for k, p in pubs.items()} + pub_views = {k: ready_view(lease, mgr_a) for k, lease in pub_leases.items()} + stats = mgr_b.impl.get_storage_statistics + + def used(group): + return int(stats(0)[group].total) - int(stats(0)[group].free) + + groups = range(len(stats(0))) + totals = [int(stats(0)[g].total) for g in groups] + total = totals[moving] + goal = int(total * 0.6) + lender_b = attach(mgr_b, tokens_per_fetch=REB_END, planned_fetches=3) + leases, views, masks, early = {}, {}, {}, {} + + def lend_write(k, request): + leases[k] = lender_b.lend_write(request, 0, REB_END) + views[k] = leases[k].poll() + assert views[k] is not None, leases[k].failure + kit.fill_sentinel(mgr_b, request) # a row that never lands shows + masks[k] = kit.relay(lender_a, pub_views[k], lender_b, views[k]) + + if stay: + # Admitted and grown before the fillers: their pages sit low, below what moves. + early["waiting_read"] = kit.published(mgr_b, 22, REB_PROMPTS["waiting_read"]) + early["unmarked_write"] = kit.admitted(mgr_b, 23, REB_PROMPTS["unmarked_write"]) + lend_write("unmarked_write", early["unmarked_write"]) + # Fillers take the low slots (a fresh pool hands out slot ids in order), so the leased + # requests' pages sit above them; another pool group may fill first, so stop at 70% of it. + fillers = kit.Requests(mgr_b) + while ( + used(moving) < goal + and all(used(g) < int(totals[g] * 0.7) for g in groups if g != moving) + and fillers.allocate(min(7, goal - used(moving))) + ): + pass + reqs = { + "queued_read": kit.published(mgr_b, 21, REB_PROMPTS["queued_read"]), + "waiting_read": early.get("waiting_read") + or kit.published(mgr_b, 22, REB_PROMPTS["waiting_read"]), + "unmarked_write": early.get("unmarked_write") + or kit.admitted(mgr_b, 23, REB_PROMPTS["unmarked_write"]), + "marked_write": kit.admitted(mgr_b, 24, REB_PROMPTS["marked_write"]), + } + bystanders = [kit.published(mgr_b, 25 + i, t) for i, t in enumerate(BYSTANDERS)] + for k in ("unmarked_write", "marked_write"): + if k not in leases: + lend_write(k, reqs[k]) + rows = { + k: [(run.layer_group_index, run.block_ordinals.tolist()) for run in views[k].groups] + for k in ("unmarked_write", "marked_write") + } + fillers.free() + dev = kit.DevicePages(mgr_b) # built before the gate: building it waits for the device + before = {k: locked(kit, mgr_b, r) for k, r in reqs.items()} + for k in ("queued_read", "waiting_read"): # the oracle holds before anything moves + kv = kit.kv(mgr_b, reqs[k]) + for lg in range(num_lg): + for o in range(REB_END // TPB): + slot = gpu_slot(kv, lg, o) + if slot >= 0: + want = published_bytes(kit, mgr_b, reqs[k], lg, o, dev.page_bytes(lg)) + assert dev.read(lg, slot) == want, f"the oracle is wrong for {k} {lg}:{o}" + device = [] + for pg in mgr_b.impl.pool_group_descs: + for pool in pg.pools: + base = int(pool.base_address) + device.append((base, base + int(pg.num_slots) * int(pool.slot_bytes))) + caches = [kit.kv(mgr_b, r) for r in reqs.values()] + guard = PageCopyGuard(_lender.drv, device, lambda: leased_ranges(kit, mgr_b, caches), block) + monkeypatch.setattr(_lender, "drv", guard) + with kit.gated_stream(mgr_b._stream, open_after=REB_GATE_SECONDS) as gate: + leases["marked_write"].mark_arrived(masks["marked_write"]) # its copy queues + leases["marked_write"].release() + leases["queued_read"] = lender_b.lend_read(reqs["queued_read"], 0, REB_END) + leases["waiting_read"] = lender_b.lend_read(reqs["waiting_read"], 0, REB_END) + waiting = ( + leases["waiting_read"].poll() is None and leases["waiting_read"].failure is None + ) + granted = leases["queued_read"].failure is None + ids = list(mgr_b.kv_cache_map) + for rid in ids: # the executor's rebalance: suspend every cache, rebalance, resume + mgr_b.suspend_request(SimpleNamespace(py_request_id=rid)) + if mechanism == "adjust": + current = _introspection.current_gpu_ratio(mgr_b.impl) + assert len(current) >= 2, f"one pool group: {current}" + # Shrink the group below the lowest slot a leased request holds there, so every + # leased page in it moves. + lowest = min( + s + for k in reqs + if not (stay and k in ("waiting_read", "unmarked_write")) + for s in before[k][full] + if s >= 0 + ) + factor = min(0.45, max(lowest - 1, 1) / total) + target = list(current) + target[moving] = current[moving] * factor + scale = (1.0 - target[moving]) / (sum(target) - target[moving]) + target = [t if g == moving else t * scale for g, t in enumerate(target)] + _introspection.set_num_sampled_kv_caches(mgr_b.impl, 2001) + _introspection.set_last_adjustment_time(mgr_b.impl, 0.0) + _introspection.set_target_ratio_list_gpu(mgr_b.impl, target) + assert mgr_b.impl.need_adjustment + mgr_b.impl.adjust() + rebalanced = True + else: + quota = int(mgr_b.impl.get_quota(0)) + rebalanced = any( + bool(mgr_b.impl.resize(0, int(quota * fraction))) for fraction in (0.5, 0.55) + ) + resumed = [ + bool(mgr_b.resume_request(SimpleNamespace(py_request_id=rid))) for rid in ids + ] + if not gate.opened: + gate.open() + mgr_b._stream.synchronize() + after = {k: locked(kit, mgr_b, r) for k, r in reqs.items()} + moved = { + k: [sorted(set(b) - {-1}) != sorted(set(a) - {-1}) for b, a in zip(before[k], after[k])] + for k in reqs + } + dev = kit.DevicePages(mgr_b) # the pools changed size + dev_a = kit.DevicePages(mgr_a) + # The queued read's copy was queued before the rebalance. + view_q = ready_view(leases["queued_read"], mgr_b) if granted else None + staged_q = ( + staged_right(kit, mgr_b, lender_b, view_q, reqs["queued_read"]) if view_q else None + ) + # The unmarked write is marked after the rebalance. + leases["unmarked_write"].mark_arrived(masks["unmarked_write"]) + leases["unmarked_write"].release() + mgr_b._stream.synchronize() + # The waiting read is granted once the queued read's slots return. + leases["queued_read"].release() + view_w = None + for _ in range(3): + mgr_b._stream.synchronize() + view_w = leases["waiting_read"].poll() + if view_w is not None or leases["waiting_read"].failure is not None: + break + staged_w = ( + staged_right(kit, mgr_b, lender_b, view_w, reqs["waiting_read"]) if view_w else None + ) + mgr_b._stream.synchronize() + resume_bounds = {} + for k in ("unmarked_write", "marked_write"): + got = lender_b.resume_bounds(reqs[k]) + resume_bounds[k] = None if got is None else tuple(got) + arrived = { + k: landed(kit, mgr_b, reqs[k], rows[k], mgr_a, pubs[k], dev, dev_a) + for k in ("unmarked_write", "marked_write") + } + # Bystanders keep their own bytes wherever their pages went. + bystanders_kept = [] + for b in bystanders: + kv = kit.kv(mgr_b, b) + for lg in range(num_lg): + for o in range(len(BYSTANDERS[0]) // TPB): + slot = gpu_slot(kv, lg, o) + if slot >= 0: + want = published_bytes(kit, mgr_b, b, lg, o, dev.page_bytes(lg)) + bystanders_kept.append(dev.read(lg, slot) == want) + # ResumeBounds counts only what landed: every block a start at max_position reads. + unbacked = {} + for k in ("unmarked_write", "marked_write"): + if resume_bounds[k] is None: + unbacked[k] = ["unsettled"] + continue + unbacked[k] = [ + f"{lg}:{o}" + for lg in range(num_lg) + for o in read_at(kit, mgr_b, lg, resume_bounds[k][0]) + if arrived[k].get(f"{lg}:{o}") is not True + ] + for lease in list(leases.values()) + list(pub_leases.values()): + lease.release() + # The rebalance ran and moved what the scenario needs. + assert rebalanced, "the rebalance did not run" + assert all(resumed), f"a cache did not resume: {resumed}" + assert waiting and granted, "the leases were not in the states the scenario needs" + if stay: + assert any(moved["queued_read"]) and any(moved["marked_write"]), moved + assert not any(moved["waiting_read"]) and not any(moved["unmarked_write"]), moved + else: + assert all(any(m) for m in moved.values()), f"leased pages stayed: {moved}" + # What the docstrings promise. + assert guard.stray == [], f"{STRAY}: {guard.stray[:4]}" + # A read of [0, REB_END) lends every block a start at REB_END reads, in every layer group. + read_rows = {f"{lg}:{o}" for lg in range(num_lg) for o in read_at(kit, mgr_b, lg, REB_END)} + assert staged_q is not None and set(staged_q) == read_rows, ( + f"the queued read lent rows {sorted(staged_q or ())}, not {sorted(read_rows)}" + ) + assert all(staged_q.values()), f"the queued read staged wrong bytes: {staged_q}" + # Both writes cover [0, REB_END) at a history of REB_END, so they lend the same rows. + for k in ("unmarked_write", "marked_write"): + assert set(arrived[k]) == read_rows, ( + f"the {k} lent rows {sorted(arrived[k])}, not {sorted(read_rows)}" + ) + assert view_w is not None or leases["waiting_read"].failure is not None, ( + "the waiting read never settled" + ) + if view_w is not None: + assert set(staged_w) == read_rows, ( + f"the waiting read lent rows {sorted(staged_w)}, not {sorted(read_rows)}" + ) + assert all(staged_w.values()), f"the waiting read staged wrong bytes: {staged_w}" + if stay: + assert view_w is not None, ( + f"the waiting read failed with its pages in place: {leases['waiting_read'].failure}" + ) + assert all(arrived["unmarked_write"].values()), ( + f"the unmarked write lost rows that stayed: {arrived['unmarked_write']}" + ) + assert resume_bounds["unmarked_write"][0] == REB_END, ( + f"resume_bounds misses rows that landed: {resume_bounds['unmarked_write']}" + ) + assert all(not v for v in unbacked.values()), ( + f"resume_bounds counts rows that never landed: {unbacked}" + ) + assert all(arrived["marked_write"].values()), ( + f"the queued write's rows were lost: {arrived['marked_write']}" + ) + assert all(bystanders_kept), "a bystander's pages changed" + + +@pytest.mark.parametrize("stay", [False, True], ids=["pages_move", "pages_stay"]) +@pytest.mark.parametrize("mechanism", sorted(MECHANISMS)) +def test_a_rebalance_with_every_kind_of_lease_in_flight( + kit, host_tier_manager, monkeypatch, mechanism, stay +): + rebalance_with_leases_in_flight(kit, host_tier_manager, monkeypatch, mechanism, stay=stay) + + +# -- (g) a mark arriving after the window passed the lent blocks -------------------------------- + +LATE_PROMPT = list(range(2000, 2161)) # five whole blocks and one token +LATE_END = 2 * TPB # the lease covers blocks 0 and 1 +LATE_HISTORY = 5 * TPB # a served prefix moves the history here: blocks 0 to 2 leave the window +FETCHED = 0xAB + + +def device_slot(mgr, address): + """``(pool group, slot)`` of a device address, or ``None`` outside every device pool.""" + for pg in mgr.impl.pool_group_descs: + for pool in pg.pools: + base, width = int(pool.base_address), int(pool.slot_bytes) + if base <= address < base + int(pg.num_slots) * width: + return int(pg.pool_group_index), (address - base) // width + return None + + +def locked_by(kit, mgr, requests): + """``{(pool group, slot): request id}`` of every page ``requests`` lock, every layer group.""" + group_of = kit.pool_group_of(mgr) + out = {} + for request in requests: + kv = kit.kv(mgr, request) + for lg in range(kit.num_layer_groups(mgr)): + for slot in kv.get_base_page_indices(lg)[: kv.num_blocks]: + if int(slot) >= 0: + out[(group_of[lg], int(slot))] = request.py_request_id + return out + + +def late_mark(kit, host_tier_manager, attach, windows, fillers, keep): + """One run, ``fillers`` one-block requests allocated first and kept or freed by ``keep``, which + moves the slots the target gets. What the marks' copy wrote, or ``None`` when no lent block went + to host under its lent GPU slot's number.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + with host_tier_manager(windows=windows) as mgr: + sliding = kit.windows(mgr).index(WINDOW) + group = kit.pool_group_of(mgr)[sliding] + first, others = kit.Requests(mgr), kit.Requests(mgr) + others._next = 500 + try: + assert first.allocate(fillers, chunk=1) if fillers else True + if not keep: + first.free() + target = kit.admitted(mgr, TARGET, LATE_PROMPT) + lender = attach(mgr, tokens_per_fetch=LATE_END) + lease = lender.lend_write(target, 0, LATE_END) + view = lease.poll() + assert view is not None, lease.failure + kv = kit.kv(mgr, target) + lent = kit.pages(kv, sliding)[:2] + assert [int(s) for s in kv.get_base_page_indices(sliding)[:2]] == lent + kit.stage(lender, view, FETCHED) + # Before the marks, a served prefix moves the history past the lent blocks of the + # sliding group, as a KV connector's reservation does: they keep their pages only held. + assert mgr._resize_for_served_prefix(target, kv, LATE_HISTORY, LATE_HISTORY + 1) + stale = kit.stale_blocks(mgr, sliding, kv.history_length) + assert stale[0] <= 0 and stale[1] >= 2, "the window did not pass the lent blocks" + assert all(int(s) < 0 for s in kv.get_base_page_indices(sliding)[:2]), "still locked" + assert kit.pages(kv, sliding)[:2] == lent, "the lent blocks do not hold their pages" + # Other requests take every GPU page left: the held pages go to host. + host_before = kit.tier_used(mgr, 1) + while others.allocate(4, chunk=4): + pass + while others.allocate(1, chunk=1): + pass + now = kit.pages(kv, sliding)[:2] + owners = locked_by(kit, mgr, others.held + first.held) + moved = [o for o in range(2) if (group, lent[o]) in owners and now[o] >= 0] + assert kit.tier_used(mgr, 1) >= host_before + len(moved), "nothing reached host" + aliased = [o for o in moved if now[o] == lent[o]] + if not aliased: + lease.mark_arrived(view.row_masks(False)) + lease.release() + return None + dev = kit.DevicePages(mgr) + watched = {lent[o]: owners[(group, lent[o])] for o in aliased} + before = {s: dev.read(sliding, s) for s in watched} + log = DeviceCopyGuard(_lender.drv) + with pytest.MonkeyPatch.context() as patch: + patch.setattr(_lender, "drv", log) + lease.mark_arrived(view.row_masks(True)) + lease.release() + mgr._stream.synchronize() + written = [device_slot(mgr, dst) for dst in log.copies] + hit = sorted({w for w in written if w is not None and w in owners}) + changed = sorted(watched[s] for s in watched if dev.read(sliding, s) != before[s]) + return dict(hit=hit, changed=changed) + finally: + others.free() + first.free() + + +def test_a_late_mark_writes_no_slot_its_window_left(kit, host_tier_manager, attach=attach): + # Slots go out lowest first, so a short search finds a lent block on host under its lent GPU + # slot's number while another request locks that GPU slot. + runs = [(w, n, k) for w in ([WINDOW], [WINDOW, 256]) for k in (True, False) for n in range(4)] + for windows, fillers, keep in runs: + result = late_mark(kit, host_tier_manager, attach, windows, fillers, keep) + if result is not None: + break + else: + raise RuntimeError("no run put a lent block on host under its lent GPU slot's number") + assert not result["hit"], f"the marks' copy wrote slots other requests lock: {result['hit']}" + assert not result["changed"], f"requests {result['changed']} now hold the fetch" diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py new file mode 100644 index 000000000000..972fb1ee126f --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py @@ -0,0 +1,1197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``_layout``: ``layout_id`` over hand-built layouts (CPU), ``derive_layout`` over a stand-in +manager scripting TP ranks, pipeline stages and recurrent groups (CPU) and real ``KVCacheManagerV2`` +instances (GPU): each description's source and what moves the id. Part names: ``test_names.py``.""" + +import dataclasses +import enum +import hashlib +import json +from types import SimpleNamespace +from typing import Dict, List, Optional, Sequence + +import numpy as np +import pytest +import torch +from utils.util import skip_pre_blackwell, skip_pre_hopper + +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import Identity +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import ( + LAYOUT_ID_VERSION, + BufferDesc, + BufferGeometry, + BufferMapper, + LayerGroupDesc, + Layout, + ShardDesc, + _Heads, + canonical_bytes, + canonical_order, + derive_layout, + global_layer_ids, + layout_document, + layout_id, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + SsmLayerConfig, +) + +gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") + +# -- hand-built layouts -------------------------------------------------------------------------- + +# Device pool group 0 has two pools (a 100-byte and a 40-byte slot); pool group 1 is a single-pool +# group for a second layer group. + + +def make_layout( + *, + tokens_per_block=32, + windows=(None, 64), + layers=((0, 2), (1,)), + shard=ShardDesc(), + group_ids=(0, 1), + slot_bytes=(100, 40), +) -> Layout: + g0, g1 = group_ids + pool_groups = {g0: tuple(slot_bytes), g1: (64,)} + layer_groups = ( + LayerGroupDesc( + kind="attention", + window=windows[0], + sink_blocks=0, + pool_group=g0, + layers=layers[0], + shard=shard, + ), + LayerGroupDesc( + kind="attention", + window=windows[1], + sink_blocks=0, + pool_group=g1, + layers=layers[1], + shard=shard, + ), + ) + buffers = [] + for layer in layers[0]: + buffers += [ + BufferDesc(layer, "key", (g0, 0), 50 * (layer // 2), 50), + BufferDesc(layer, "scale", (g0, 1), 20 * (layer // 2), 20), + ] + for layer in layers[1]: + buffers.append(BufferDesc(layer, "key", (g1, 0), 0, 64)) + return Layout(tokens_per_block, pool_groups, layer_groups, tuple(buffers)) + + +def _replace_buffers(layout, **change): + return dataclasses.replace( + layout, + buffers=tuple( + dataclasses.replace(b, **change) if b.role == "scale" else b for b in layout.buffers + ), + ) + + +def _first_group(layout, **change): + first = dataclasses.replace(layout.layer_groups[0], **change) + return dataclasses.replace(layout, layer_groups=(first,) + layout.layer_groups[1:]) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "variant", + [ + pytest.param(lambda: make_layout(group_ids=(7, 3)), id="pool_group_numbering"), + pytest.param( + lambda: dataclasses.replace( + make_layout(), layer_groups=make_layout().layer_groups[::-1] + ), + id="layer_group_order", + ), + pytest.param(lambda: _first_group(make_layout(), sink_blocks=3), id="sinks"), + pytest.param(lambda: make_layout(windows=(None, 128)), id="window"), + pytest.param(lambda: make_layout(windows=(64, 64)), id="full_to_window"), + pytest.param(lambda: make_layout(shard=ShardDesc(2, 1)), id="shard"), + pytest.param( + lambda: dataclasses.replace(make_layout(), buffers=make_layout().buffers[::-1]), + id="buffer_order", + ), + ], +) +def test_layout_id_leaves_out_what_does_not_decide_how_bytes_read(variant): + """Local numbering and order, sinks and windows (which blocks exist, not how their bytes read), + and the shard position, which rides in the name so every rank computes the same ``layout_id``. + Buffers are described by global layer, so their listing order does not matter either.""" + assert layout_id(variant()) == layout_id(make_layout()) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "change", + [ + lambda: make_layout(tokens_per_block=64), + lambda: make_layout(layers=((0, 1), (2,))), + lambda: make_layout(slot_bytes=(104, 40)), + lambda: _replace_buffers(make_layout(), mapper=BufferMapper.REPLICATED), + lambda: _replace_buffers(make_layout(), geometry=BufferGeometry(bytes_per_head=10)), + lambda: _replace_buffers(make_layout(), geometry=BufferGeometry(section_bytes=(15, 5))), + lambda: _replace_buffers(make_layout(), transfer=False), + lambda: _replace_buffers(make_layout(), expansion=2), + lambda: _replace_buffers(make_layout(), role="scale2"), + lambda: _first_group(make_layout(), kind="state"), + lambda: dataclasses.replace( + make_layout(), + buffers=tuple( + dataclasses.replace(b, offset=b.offset + 10) if b.role == "scale" else b + for b in make_layout().buffers + ), + ), + ], + ids=[ + "tokens_per_block", + "layer_membership", + "slot_width", + "mapper", + "head_geometry", + "sections", + "transfer", + "expansion", + "role", + "kind", + "position_in_the_transfer_bytes", + ], +) +def test_layout_id_changes_with_every_input_that_decides_how_bytes_read(change): + assert layout_id(change()) != layout_id(make_layout()) + + +@pytest.mark.cpu_only +def test_canonical_order_sorts_groups_by_smallest_layer(): + layout = make_layout(layers=((3, 5), (1, 4))) + assert canonical_order([d.layers for d in layout.layer_groups]) == [1, 0] + assert canonical_order([d.layers for d in layout.layer_groups[::-1]]) == [0, 1] + # Ties and empty groups keep the local order; empty groups go last. + assert canonical_order([(), (2,), (2,), (0, 9)]) == [3, 1, 2, 0] + assert canonical_order([(9,), (2, 11), (5,)]) == [2, 0, 1] + + +@pytest.mark.cpu_only +def test_layout_id_is_the_hash_of_the_documented_version_2_encoding(): + """The canonical encoding, spelled out here byte for byte so a change to it cannot pass + unnoticed.""" + layout = make_layout(layers=((0,), (1,))) + expected = ( + '{"buffers":[' + '{"layer":0,"offset":0,"role":"key","size":50},' + '{"layer":0,"offset":100,"role":"scale","size":20},' + '{"layer":1,"offset":0,"role":"key","size":64}],' + '"groups":[{"layers":[0],"slot_bytes":[100,40]},{"layers":[1],"slot_bytes":[64]}],' + '"tokens_per_block":32,"v":2}' + ).encode() + assert LAYOUT_ID_VERSION == 2 + assert canonical_bytes(layout_document(layout)) == expected + assert layout_id(layout) == hashlib.sha256(expected).digest() + + +@pytest.mark.cpu_only +def test_named_keys_at_their_defaults_are_left_out_and_others_are_written(): + layout = make_layout(layers=((0,), (1,))) + plain = layout_document(layout)["buffers"][1] + assert set(plain) == {"layer", "role", "offset", "size"} + changed = layout_document( + _replace_buffers( + layout, + mapper=BufferMapper.SECTIONED, + geometry=BufferGeometry(section_bytes=(15, 5)), + transfer=False, + expansion=4, + ) + )["buffers"][1] + assert changed["mapper"] == "SECTIONED" + assert changed["geometry"] == {"section_bytes": [15, 5]} + assert changed["transfer"] is False + assert changed["expansion"] == 4 + assert layout_document(_first_group(layout, kind="state"))["groups"][0]["kind"] == "state" + assert "kind" not in layout_document(layout)["groups"][0] + + +@pytest.mark.cpu_only +def test_enums_encode_by_name_and_the_hnd_alias_is_indexed(): + hnd = _replace_buffers(make_layout(), mapper=MapperKind.HND) + indexed = _replace_buffers(make_layout(), mapper=BufferMapper.INDEXED) + assert layout_id(hnd) == layout_id(indexed) == layout_id(make_layout()) + + class Color(enum.IntEnum): + RED = 3 + + assert json.loads(canonical_bytes({"c": Color.RED, "n": None, "l": [None, 1]})) == { + "c": "RED", + "l": [None, 1], + } + + +@pytest.mark.cpu_only +def test_floats_encode_as_their_binary64_bits(): + assert canonical_bytes({"x": 1.0}) == b'{"x":"f64:3ff0000000000000"}' + assert canonical_bytes({"x": -0.0}) != canonical_bytes({"x": 0.0}) + + +@pytest.mark.cpu_only +def test_a_document_refuses_what_the_layout_does_not_hold(): + layout = make_layout() + with pytest.raises(ValueError, match="maps to pool group 9"): + layout_document(_first_group(layout, pool_group=9)) + outside = layout.buffers + (BufferDesc(7, "key", (9, 0), 0, 8),) + with pytest.raises(ValueError, match="not in the layout"): + layout_document(dataclasses.replace(layout, buffers=outside)) + twice = layout.buffers + (layout.buffers[0],) + with pytest.raises(ValueError, match="appears twice"): + layout_document(dataclasses.replace(layout, buffers=twice)) + + +# -- geometry and buffers ------------------------------------------------------------------------ + + +@pytest.mark.cpu_only +def test_a_geometry_that_does_not_fit_its_buffer_is_refused(): + with pytest.raises(ValueError, match="sum"): + BufferDesc(0, "conv", (0, 0), 0, 30, geometry=BufferGeometry(section_bytes=(10, 10))) + with pytest.raises(ValueError, match="heads"): + BufferDesc(0, "k", (0, 0), 0, 30, geometry=BufferGeometry(bytes_per_head=8)) + with pytest.raises(ValueError, match="heads"): + BufferDesc(0, "k", (0, 0), 0, 32, geometry=BufferGeometry(bytes_per_head=8, num_heads=2)) + + +@pytest.mark.cpu_only +def test_an_expansion_must_divide_the_block(): + with pytest.raises(ValueError, match="positive"): + BufferDesc(0, "k", (0, 0), 0, 32, expansion=0) + with pytest.raises(ValueError, match="does not divide"): + _replace_buffers(make_layout(), expansion=3) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("shard", [(2, 2), (0, 0), (2, -1)]) +def test_a_shard_is_a_share_of_its_count(shard): + with pytest.raises(ValueError, match="shard"): + ShardDesc(*shard) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "sections", + [(np.int64(4), np.int64(4)), np.array([4, 4]), [4, 4], (4, 4)], + ids=["numpy-ints", "numpy-array", "list", "tuple"], +) +def test_a_validated_geometry_holds_its_sections_as_a_tuple_of_ints(sections): + checked = BufferGeometry(section_bytes=sections, bytes_per_head=2).validate() + assert checked == BufferGeometry(section_bytes=(4, 4), bytes_per_head=2) + assert type(checked.section_bytes) is tuple + assert all(type(s) is int for s in checked.section_bytes) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "fields", [{"section_bytes": (4, 0)}, {"bytes_per_head": 0}, {"num_heads": -1}] +) +def test_validation_refuses_sizes_that_are_not_positive(fields): + with pytest.raises(ValueError, match="positive"): + BufferGeometry(**fields).validate() + + +# -- a stand-in manager -------------------------------------------------------------------------- + +TPB = 4 +HEAD_BYTES = TPB * 8 * 2 # 8 fp16 values per token and head + + +def coalesced(size, members): + return SimpleNamespace( + single_buffer_size=size, + buffer_ids=[SimpleNamespace(layer_id=layer, role=role) for layer, role in members], + ) + + +class Manager: + """One layer group per entry of ``groups``: ``(layer config class, [internal layer ids], + [(role, size, override)])``; group ``i`` draws from pool group ``pool_group_of[i]`` (its own by + default), one pool per role, and has window ``windows[i]`` and ``sinks[i]`` sink tokens.""" + + def __init__( + self, + groups, + *, + tp_size=1, + tp_rank=0, + attention_dp=False, + num_kv_heads=4, + pp_layers: Optional[Sequence[int]] = None, + mappers: Optional[Dict[str, MapperKind]] = None, + layouts: Optional[Dict[str, RoleLayout]] = None, + ignored=(), + pool_group_of: Optional[Sequence[int]] = None, + windows: Optional[Sequence[Optional[int]]] = None, + sinks: Optional[Sequence[Optional[int]]] = None, + ): + heads_per_rank = -(-num_kv_heads // (1 if attention_dp else tp_size)) + internal = sorted({lid for _, lids, _ in groups for lid in lids}) + pool_group_of = list(pool_group_of if pool_group_of is not None else range(len(groups))) + windows = list(windows or [None] * len(groups)) + sinks = list(sinks or [None] * len(groups)) + self.tokens_per_block = TPB + self.mapping = SimpleNamespace( + tp_size=tp_size, tp_rank=tp_rank, enable_attention_dp=attention_dp + ) + self.num_kv_heads = num_kv_heads + self.num_kv_heads_per_layer = [heads_per_rank] * len(internal) + self.pp_layers = list(pp_layers if pp_layers is not None else internal) + self.mappers = {"all": MapperKind.INDEXED, **(mappers or {})} + self.layouts = layouts or {} + self.ignored = frozenset(ignored) + layers: List[object] = [] + descs = {} + for lg, (cls, lids, roles) in enumerate(groups): + for lid in lids: + buffers = [ + BufferConfig(role=role, size=size, tokens_per_block_override=override) + for role, size, override in roles + ] + if cls is SsmLayerConfig: + layers.append(SsmLayerConfig(layer_id=lid, buffers=buffers)) + continue + layers.append( + cls( + layer_id=lid, + buffers=buffers, + sliding_window_size=windows[lg], + num_sink_tokens=sinks[lg], + ) + ) + pools = [coalesced(size, [(lid, role) for lid in lids]) for role, size, _ in roles] + variant = SimpleNamespace(layer_group_id=lg, coalesced_buffers=pools) + g = pool_group_of[lg] + if g in descs: + descs[g].slot_desc.variants.append(variant) + continue + descs[g] = SimpleNamespace( + pool_group_index=g, + num_slots=8, + pools=[ + SimpleNamespace( + pool_index=i, + base_address=0x1000 * (g + 1) + 0x100 * i, + slot_bytes=c.single_buffer_size * len(lids), + ) + for i, c in enumerate(pools) + ], + slot_desc=SimpleNamespace(variants=[variant]), + ) + self.impl = SimpleNamespace( + layer_grouping=[list(lids) for _, lids, _ in groups], + pool_group_descs=list(descs.values()), + get_life_cycle_pool_group_indices=lambda: list(pool_group_of), + init_config=SimpleNamespace(layers=layers), + ) + + def get_disagg_role_mapper_kinds(self): + return self.mappers + + def get_disagg_role_layouts(self): + return self.layouts + + def get_disagg_ignored_roles(self): + return self.ignored + + +def kv_group(lids=(0, 1), heads_per_rank=1, extra=()): + size = HEAD_BYTES * heads_per_rank + return (AttentionLayerConfig, list(lids), [("key", size, None), ("value", size, None), *extra]) + + +def attention(tp_size, tp_rank, num_kv_heads, **kwargs): + heads = -(-num_kv_heads // tp_size) + return Manager( + [kv_group(heads_per_rank=heads)], + tp_size=tp_size, + tp_rank=tp_rank, + num_kv_heads=num_kv_heads, + **kwargs, + ) + + +def recurrent_state(): + return Manager( + [(SsmLayerConfig, [0, 1], [("conv_state", 96, None), ("ssm_state", 128, None)])], + tp_size=2, + tp_rank=1, + mappers={"conv_state": MapperKind.SECTIONED}, + layouts={ + "conv_state": RoleLayout(section_bytes=(32, 32, 32)), + "ssm_state": RoleLayout(bytes_per_head=32), + }, + ) + + +TOKEN_MAJOR_SHARES = "head-major and token-major K/V share a layout_id" + + +def check_token_major_kv_has_its_own_layout_id(): + """A manager declaring token-major K/V (NHD) and one declaring the default head-major K/V (HND) + lay a block's bytes out differently, so they never share a ``layout_id``.""" + head_major = derive_layout(attention(1, 0, 4)).layout + token_major = derive_layout(attention(1, 0, 4, mappers={"all": MapperKind.NHD})).layout + assert layout_id(head_major) != layout_id(token_major), TOKEN_MAJOR_SHARES + assert {b.mapper for b in token_major.buffers} == {BufferMapper.NHD} + assert {b.mapper for b in head_major.buffers} == {BufferMapper.INDEXED} + hnd = derive_layout(attention(1, 0, 4, mappers={"all": MapperKind.HND})).layout + assert layout_id(hnd) == layout_id(head_major), "HND is the default head-major layout" + + +@pytest.mark.cpu_only +def test_token_major_kv_has_its_own_layout_id(): + check_token_major_kv_has_its_own_layout_id() + + +# (TP size, KV heads, attention DP) -> the share each rank holds, rank by rank. +KV_HEAD_SHARES = { + # DeepSeek-V4 at TP8 without attention DP: one KV head, repeated on every rank, so every rank + # describes the group as whole and all of them compute one namespace. + "single_kv_head": ((8, 1, False), [ShardDesc()] * 8), + "head_sharded": ((4, 8, False), [ShardDesc(4, r) for r in range(4)]), + "unaligned_heads": ((4, 6, False), [ShardDesc(4, r) for r in range(4)]), + # Two KV heads over four ranks: ranks 0 and 1 hold head 0, ranks 2 and 3 head 1, so there are + # two shares, not four. + "repeated_heads": ((4, 2, False), [ShardDesc(2, r // 2) for r in range(4)]), + "attention_dp_holds_whole_heads": ((4, 8, True), [ShardDesc()] * 4), +} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("case", list(KV_HEAD_SHARES)) +def test_kv_heads_name_the_share_each_rank_holds(case): + """The byte layout is the same on every rank; the share is part of the object name.""" + (tp_size, num_kv_heads, attention_dp), shares = KV_HEAD_SHARES[case] + geometries = [ + derive_layout(attention(tp_size, r, num_kv_heads, attention_dp=attention_dp)) + for r in range(tp_size) + ] + assert [geo.layout.layer_groups[0].shard for geo in geometries] == shares + assert [geo.shards for geo in geometries] == [((s.count, s.index),) for s in shares] + assert len({layout_id(geo.layout) for geo in geometries}) == 1 + keys = np.full((1, 32), ord("k"), dtype=np.uint8) + names = [ + Identity(b"head-shares", layout_id(geo.layout), geo.layers, geo.shards) + .object_keys(0, keys)[0] + .tobytes() + for geo in geometries + ] + for rank, share in enumerate(shares): + for other, other_share in enumerate(shares): + assert (names[rank] == names[other]) == (share == other_share) + + +def heads_on(rank, tp_size, total, per_rank): + """``_Heads`` of a rank that holds ``per_rank`` of ``total`` KV heads over ``tp_size`` ranks.""" + mapping = SimpleNamespace(tp_size=tp_size, tp_rank=rank, enable_attention_dp=False) + manager = SimpleNamespace( + mapping=mapping, num_kv_heads_per_layer=[per_rank], num_kv_heads=total, pp_layers=[0] + ) + return _Heads(manager) + + +# (KV heads T, TP size, heads per rank) -> each rank's share. Rank r holds its heads from +# r * T // tp; where any rank's range starts inside a share, every rank is a share of its own. +UNALIGNED_HEADS = { + (6, 4, 2): [ShardDesc(4, r) for r in range(4)], + (12, 8, 2): [ShardDesc(8, r) for r in range(8)], + (8, 4, 2): [ShardDesc(4, r) for r in range(4)], + (2, 4, 1): [ShardDesc(2, 0), ShardDesc(2, 0), ShardDesc(2, 1), ShardDesc(2, 1)], + (1, 4, 1): [ShardDesc(1, 0)] * 4, +} + + +@pytest.mark.cpu_only +def test_unaligned_head_ranges_give_every_rank_its_own_share(): + for (total, tp_size, per_rank), shares in UNALIGNED_HEADS.items(): + got = [heads_on(r, tp_size, total, per_rank).attention_shard(0) for r in range(tp_size)] + assert got == shares, f"{total} heads over {tp_size} ranks: {got}" + + +@pytest.mark.cpu_only +def test_a_group_of_replicated_buffers_is_whole_and_one_head_sharded_buffer_makes_it_a_share(): + index = ("index_key", 64, None) + replicated_only = Manager( + [(AttentionLayerConfig, [0], [index])], + tp_size=2, + tp_rank=1, + num_kv_heads=8, + mappers={"index_key": MapperKind.REPLICATED}, + ) + assert derive_layout(replicated_only).layout.layer_groups[0].shard == ShardDesc() + mixed = Manager( + [kv_group(heads_per_rank=4, extra=[index])], + tp_size=2, + tp_rank=1, + num_kv_heads=8, + mappers={"index_key": MapperKind.REPLICATED}, + ) + assert derive_layout(mixed).layout.layer_groups[0].shard == ShardDesc(2, 1) + + +@pytest.mark.cpu_only +def test_local_only_buffers_do_not_decide_the_identity(): + scratch = ("scratch", 32, None) + manager = Manager( + [(AttentionLayerConfig, [0], [("index_key", 64, None), scratch])], + tp_size=2, + tp_rank=1, + mappers={"index_key": MapperKind.REPLICATED}, + ignored=("scratch",), + ) + layout = derive_layout(manager).layout + assert layout.layer_groups[0].shard == ShardDesc() + assert {b.role: b.transfer for b in layout.buffers} == {"index_key": True, "scratch": False} + + +@pytest.mark.cpu_only +def test_groups_of_one_manager_are_judged_one_by_one(): + manager = Manager( + [ + kv_group(lids=(0,), heads_per_rank=4), + (AttentionLayerConfig, [1], [("index_key", 64, None)]), + ], + tp_size=2, + tp_rank=0, + num_kv_heads=8, + mappers={"index_key": MapperKind.REPLICATED}, + ) + groups = derive_layout(manager).layout.layer_groups + assert [g.shard for g in groups] == [ShardDesc(2, 0), ShardDesc()] + + +def mixed_shares(tp_size, tp_rank, heads): + """``heads`` KV heads in layers 0 and 1 over ``tp_size`` ranks, one head per rank in each, one + group.""" + manager = Manager([kv_group(lids=(0, 1))], tp_size=tp_size, tp_rank=tp_rank, num_kv_heads=8) + manager.num_kv_heads = list(heads) + manager.num_kv_heads_per_layer = [1, 1] + return manager + + +# name: (ranks, heads per layer). The group's first buffer is layer 0's. At TP8 the narrowest +# share equals the fallback, and the first buffer's too with eight heads first; at TP6 none does. +MIXED_SHARES = { + "eight_first": (8, (8, 4)), + "four_first": (8, (4, 8)), + "three_first": (6, (3, 2)), + "two_first": (6, (2, 3)), +} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("case", list(MIXED_SHARES)) +def test_a_group_whose_buffers_hold_different_shares_is_each_rank_alone(case): + """The two layers' buffers hold shares of different counts, so the group falls back to each + rank's own share of the ranks, and no two ranks name a block alike.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import Identity + + tp_size, heads = MIXED_SHARES[case] + keys = np.arange(3 * 32, dtype=np.uint8).reshape(3, 32) + names = set() + for rank in range(tp_size): + geometry = derive_layout(mixed_shares(tp_size, rank, heads)) + assert geometry.layout.layer_groups[0].shard == ShardDesc(tp_size, rank), ( + f"rank {rank}: the group's share is not the rank's own" + ) + identity = Identity(b"scope", layout_id(geometry.layout), geometry.layers, geometry.shards) + names |= {row.tobytes() for row in identity.object_keys(0, keys)} + assert len(names) == tp_size * len(keys), "two ranks name a block alike" + + +@pytest.mark.cpu_only +def test_recurrent_state_is_described_and_sharded_by_rank(): + geometry = derive_layout(recurrent_state()) + layout = geometry.layout + (group,) = layout.layer_groups + assert group.kind == "state" and group.shard == ShardDesc(2, 1) + assert geometry.recurrent == (True,) and geometry.windows == (None,) + by_role = {b.role: b for b in layout.buffers} + assert by_role["conv_state"].mapper is BufferMapper.SECTIONED + assert by_role["conv_state"].geometry == BufferGeometry(section_bytes=(32, 32, 32)) + assert by_role["ssm_state"].geometry == BufferGeometry(bytes_per_head=32, num_heads=4) + + +@pytest.mark.cpu_only +def test_head_geometry_is_described_per_buffer(): + layout = derive_layout(attention(2, 0, num_kv_heads=8)).layout + for b in layout.buffers: + assert b.mapper is BufferMapper.INDEXED + assert b.geometry == BufferGeometry(bytes_per_head=HEAD_BYTES, num_heads=4) + + +@pytest.mark.cpu_only +def test_layers_are_global_ids_stable_across_pipeline_stages(): + """Stage 2 of a pipeline holds model layers 10 and 11 as internal layers 0 and 1.""" + manager = Manager([kv_group()], pp_layers=[10, 11]) + geometry = derive_layout(manager) + assert geometry.layout.layer_groups[0].layers == (10, 11) + assert geometry.layers == ((10, 11),) + assert {b.layer for b in geometry.layout.buffers} == {10, 11} + assert global_layer_ids(manager, [1, 0]) == [11, 10] + # A manager without the map is refused, not named by its stage-local ids, which another + # stage reuses for other layers. + del manager.pp_layers + with pytest.raises(AttributeError, match="pp_layers"): + derive_layout(manager) + + +@pytest.mark.cpu_only +def test_virtual_layers_are_numbered_by_model_layer_and_attention_type(): + """Two attention types over three enum members: model layer ``m`` of type ``t`` is + ``3 * m + t``, whichever types this stage holds.""" + + class AttentionType(enum.Enum): + SLIDING = 0 + COMPRESSED = 1 + INDEXER = 2 + + manager = Manager([kv_group(lids=(0, 1, 2))]) + manager._layer_attn_to_layer_id = { + (4, AttentionType.SLIDING): 0, + (4, AttentionType.COMPRESSED): 1, + (5, AttentionType.SLIDING): 2, + } + assert global_layer_ids(manager, [0, 1, 2]) == [12, 13, 15] + assert derive_layout(manager).layers == ((12, 13, 15),) + + +@pytest.mark.cpu_only +def test_declared_global_layer_ids_name_extra_internal_layers(): + """A manager whose internal layers outnumber its model layers (FP4 MLA keeps its tail on + extra ones) declares the ids per layer group; they are used as declared, in group order.""" + manager = Manager( + [kv_group(lids=(1, 0)), (AttentionLayerConfig, [2, 3], [("tail", 64, None)])], + pp_layers=[0, 1], + mappers={"tail": MapperKind.REPLICATED}, + ) + declared = {0: [2, 0], 1: [1, 3]} + manager.get_disagg_global_layer_ids = lambda lg: declared[lg] + assert global_layer_ids(manager, [0, 1, 2, 3]) == [0, 2, 1, 3] + layout = derive_layout(manager).layout + assert [g.layers for g in layout.layer_groups] == [(0, 2), (1, 3)] + assert {(b.layer, b.role) for b in layout.buffers if b.role == "tail"} == { + (1, "tail"), + (3, "tail"), + } + manager.get_disagg_global_layer_ids = lambda lg: declared[lg][:1] + with pytest.raises(ValueError, match="declared global ids"): + global_layer_ids(manager, [0]) + + +@pytest.mark.cpu_only +def test_a_buffer_with_its_own_tokens_per_block_is_expanded(): + manager = Manager( + [kv_group(extra=[("index_key", 64, 2)])], mappers={"index_key": MapperKind.REPLICATED} + ) + layout = derive_layout(manager).layout + assert {b.role: b.expansion for b in layout.buffers} == {"key": 1, "value": 1, "index_key": 2} + plain = derive_layout( + Manager( + [kv_group(extra=[("index_key", 64, None)])], + mappers={"index_key": MapperKind.REPLICATED}, + ) + ).layout + assert layout_id(layout) != layout_id(plain) + + +@pytest.mark.cpu_only +def test_the_manager_layout_holds_per_group_and_per_pool_group_facts(): + """Two attention groups share pool group 3, each with its own window and sinks; a state group + has pool group 1 and neither window nor sinks.""" + manager = Manager( + [ + kv_group(lids=(0,)), + kv_group(lids=(1,)), + (SsmLayerConfig, [2], [("conv_state", 96, None)]), + ], + pool_group_of=[3, 3, 1], + windows=[None, 8, 16], + sinks=[None, 5, 4], + ) + geometry = derive_layout(manager) + assert geometry.num_layer_groups == 3 + assert geometry.tokens_per_block == TPB + assert geometry.pool_group_of == (3, 3, 1) + assert [d.pool_group for d in geometry.layout.layer_groups] == [3, 3, 1] + assert geometry.pool_groups == (1, 3) + assert geometry.windows == (None, 8, None) + assert geometry.sink_blocks == (0, 2, 0) # 5 sink tokens take two 4-token blocks + assert geometry.recurrent == (False, False, True) + assert geometry.layers == ((0,), (1,), (2,)) + assert geometry.shards == ((1, 0),) * 3 + assert dict(geometry.page_bytes) == {3: 2 * HEAD_BYTES, 1: 96} + assert dict(geometry.layout.pool_groups) == {3: (HEAD_BYTES, HEAD_BYTES), 1: (96,)} + pools = { + g: [(p.group, p.index, p.base, p.slot_bytes, p.num_slots) for p in ps] + for g, ps in geometry.device_pools.items() + } + assert pools == { + 3: [(3, 0, 0x4000, HEAD_BYTES, 8), (3, 1, 0x4100, HEAD_BYTES, 8)], + 1: [(1, 0, 0x2000, 96, 8)], + } + + +@pytest.mark.cpu_only +def test_a_state_group_is_told_by_its_config_class_not_its_name(): + """Only the runtime's state layer config makes a group a state group, whatever another layer + config class is named.""" + + class SsmNamedAttentionConfig(SimpleNamespace): + """An attention layer config whose class name reads like the state one.""" + + _, lids, roles = kv_group(lids=(0,)) + manager = Manager( + [(SsmNamedAttentionConfig, lids, roles), (SsmLayerConfig, [1], [("conv_state", 96, None)])] + ) + assert derive_layout(manager).recurrent == (False, True) + + +@pytest.mark.cpu_only +def test_a_manager_whose_parts_disagree_is_refused(): + no_config = Manager([kv_group()]) + no_config.impl.init_config.layers = [] + with pytest.raises(ValueError, match="has no layer config"): + derive_layout(no_config) + miscounted = Manager([kv_group()]) + miscounted.impl.get_life_cycle_pool_group_indices = lambda: [0, 0] + with pytest.raises(ValueError, match="pool-group indices"): + derive_layout(miscounted) + poolless = Manager([kv_group()]) + poolless.impl.get_life_cycle_pool_group_indices = lambda: [5] + with pytest.raises(ValueError, match="which has no pools"): + derive_layout(poolless) + with pytest.raises(ValueError, match="does not divide"): + derive_layout(Manager([kv_group(extra=[("index_key", 64, 3)])])) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "build,expected", + [ + ( + lambda: Manager( + [kv_group(heads_per_rank=4, extra=[("index_key", 64, 2)])], + tp_size=2, + tp_rank=1, + num_kv_heads=8, + mappers={"index_key": MapperKind.REPLICATED}, + ), + "069d747522257b726c6b28d917cb2c260350e26c55550bb687bea7350a76f51c", + ), + ( + recurrent_state, + "3eba9af0efca5fe1e8701398e86b43ada86f07ba249f483be435925ffeaab6fa", + ), + ], + ids=["head_sharded_with_expanded_replicated_role", "recurrent_state"], +) +def test_the_derived_layout_id_is_pinned(build, expected): + """The derivation and the encoding together: a change here moves every name built on it.""" + assert layout_id(derive_layout(build()).layout).hex() == expected + + +# -- real managers ------------------------------------------------------------------------------- + +REAL_TPB = 32 +WINDOW = 64 + + +def native_pool_groups(mgr): + """``{g: [(base, slot bytes)...]}`` and slot counts, read straight off the manager.""" + pools, counts = {}, {} + for pg in mgr.impl.pool_group_descs: + g = int(pg.pool_group_index) + pools[g] = [ + (int(p.base_address), int(p.slot_bytes)) + for p in sorted(pg.pools, key=lambda p: int(p.pool_index)) + ] + counts[g] = int(pg.num_slots) + return pools, counts + + +def check_pool_groups(mgr, geometry): + layout = geometry.layout + pools, counts = native_pool_groups(mgr) + assert set(layout.pool_groups) == set(pools) + assert geometry.pool_groups == tuple(sorted(pools)) + for g, native in pools.items(): + widths = tuple(w for _, w in native) + assert layout.pool_groups[g] == widths + assert geometry.page_bytes[g] == sum(widths) + # Slot i at base + i * the layout's slot width: a pool is its slots back to back. + assert [(p.base, p.slot_bytes * p.num_slots) for p in geometry.device_pools[g]] == [ + (base, width * counts[g]) for base, width in native + ] + + +def check_buffers(mgr, layout): + """Every (layer, role) of the manager appears once, inside its pool's slot.""" + seen = {} + for b in layout.buffers: + assert (b.layer, b.role) not in seen + seen[(b.layer, b.role)] = b + width = layout.pool_groups[b.pool[0]][b.pool[1]] + assert 0 <= b.offset and b.offset + b.size <= width + internal = sorted({int(bid.layer_id) for bid in mgr.impl.all_buffer_ids}) + global_of = dict(zip(internal, global_layer_ids(mgr, internal))) + native = {(global_of[int(bid.layer_id)], str(bid.role)) for bid in mgr.impl.all_buffer_ids} + assert set(seen) == native + + +@gpu +def test_a_full_attention_cache_is_one_layer_group_over_one_pool_group(real_manager): + with real_manager() as mgr: + geometry = derive_layout(mgr) + layout = geometry.layout + + assert layout.tokens_per_block == geometry.tokens_per_block == REAL_TPB + (group,) = layout.layer_groups + assert group.kind == "attention" + assert group.window is None + assert group.layers == (0, 1) + g = int(mgr.impl.get_life_cycle_pool_group_indices()[0]) + assert group.pool_group == g + assert geometry.pool_group_of == (g,) + assert geometry.windows == (None,) + assert geometry.sink_blocks == (0,) + assert geometry.recurrent == (False,) + assert geometry.shards == ((1, 0),) + check_pool_groups(mgr, geometry) + check_buffers(mgr, layout) + # K and V are head-sharded (INDEXED): four heads of ``tokens_per_block * head_dim`` fp16 + # values on this rank. + for b in layout.buffers: + assert b.mapper is BufferMapper.INDEXED + assert b.transfer and b.expansion == 1 + assert b.geometry == BufferGeometry(bytes_per_head=REAL_TPB * 64 * 2, num_heads=4) + assert group.shard == ShardDesc() + + +@gpu +def test_two_windows_are_two_layer_groups_with_their_own_windows(real_manager): + with real_manager(windows=[WINDOW, 256]) as mgr: + geometry = derive_layout(mgr) + layout = geometry.layout + + assert len(layout.layer_groups) == 2 + native = [int(g) for g in mgr.impl.get_life_cycle_pool_group_indices()] + assert list(geometry.pool_group_of) == native + by_layer = {} + for lg, group in enumerate(layout.layer_groups): + assert group.pool_group == native[lg] + assert geometry.windows[lg] == group.window + for layer in group.layers: + by_layer[layer] = group.window + # Layer 0 slides; layer 1's window reaches max_seq_len and normalizes to full attention. + assert mgr.max_seq_len <= 256 + assert by_layer == {0: WINDOW, 1: None} + check_pool_groups(mgr, geometry) + check_buffers(mgr, layout) + + +@gpu +def test_layout_id_ignores_slot_counts_and_addresses(real_manager): + with real_manager(max_tokens=2048) as small: + small_id = layout_id(derive_layout(small).layout) + small_pools = native_pool_groups(small) + with real_manager(max_tokens=8192) as large: + large_geometry = derive_layout(large) + large_pools = native_pool_groups(large) + assert large_pools[1] != small_pools[1], "the two caches must differ in slot counts" + assert layout_id(large_geometry.layout) == small_id + + +@gpu +def test_layout_id_leaves_the_window_out(real_manager): + """One window for every layer: the same single layer group, only its window differs. The + window says which blocks exist, not how a block's bytes read, so the id stays.""" + with real_manager(windows=[WINDOW]) as mgr: + before = derive_layout(mgr).layout + with real_manager(windows=[2 * WINDOW]) as mgr: + after = derive_layout(mgr).layout + assert [g.window for g in before.layer_groups] != [g.window for g in after.layer_groups] + assert layout_id(after) == layout_id(before) + + +@gpu +def test_layout_id_follows_the_tokens_per_block(real_manager): + with real_manager() as mgr: + before = derive_layout(mgr).layout + with real_manager(tokens_per_block=64) as mgr: + after = derive_layout(mgr).layout + assert len(after.layer_groups) == len(before.layer_groups) + assert layout_id(after) != layout_id(before) + + +@gpu +def test_layout_id_leaves_the_element_type_out(real_manager): + """FP16 and BF16 caches differ in no size or offset, so they share a ``layout_id``; what the + bytes mean is for the caller's scope to say.""" + from tensorrt_llm.bindings import DataType + + with real_manager(dtype=DataType.HALF) as mgr: + fp16, fp16_type = derive_layout(mgr).layout, mgr.dtype + with real_manager(dtype=DataType.BF16) as mgr: + bf16, bf16_type = derive_layout(mgr).layout, mgr.dtype + assert fp16_type != bf16_type, "the two caches hold the same element type: proves nothing" + assert layout_id(bf16) == layout_id(fp16) + + +def fp4_mla(kit, monkeypatch, backend): + """A three-layer FP4 MLA manager on ``backend``, shut down when the block exits.""" + from tensorrt_llm._torch.attention.backends.fp4_mla import ( + FP4_MLA_ATTENTION_BACKEND_ENV, + FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, + ) + from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 + from tensorrt_llm._torch.kimi_k3_cache_policy import KIMI_K3_BF16_KV_LAYERS_ENV + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.mapping import Mapping + + monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, backend) + monkeypatch.setenv(FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, "0") + monkeypatch.delenv(KIMI_K3_BF16_KV_LAYERS_ENV, raising=False) + return kit.managed( + lambda: Fp4MlaKVCacheManagerV2( + KvCacheConfig( + max_tokens=512, dtype="nvfp4", enable_block_reuse=True, host_cache_size=0 + ), + CacheType.SELFKONLY, + num_layers=3, + num_kv_heads=1, + head_dim=576, + tokens_per_block=128, + max_seq_len=512, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1), + dtype=DataType.NVFP4, + max_num_tokens=512, + pretrained_config=SimpleNamespace(kv_lora_rank=512), + ) + ) + + +@gpu +@skip_pre_hopper +def test_fp4_mla_buffers_carry_the_declared_global_layer_ids(kit, monkeypatch): + """The tail lives on internal layers past the model's; in the layout its buffers take the ids + the manager declares (``2 * layer + 1``), next to the cache's ``2 * layer``.""" + with fp4_mla(kit, monkeypatch, "triton") as mgr: + buffers = {(b.layer, b.role) for b in derive_layout(mgr).layout.buffers} + assert {layer for layer, role in buffers if role == "key"} == {0, 2, 4} + assert {layer for layer, role in buffers if role == "mla_hp_tail"} == {1, 3, 5} + + +M3_TPB = 128 +M3_KV_HEADS = 2 +M3_HEAD_DIM = 128 +M3_INDEX_DIM = 128 +M3_SPARSE_LAYER = 3 +M3_KV_DTYPES = {"nvfp4": "NVFP4", "fp8": "FP8"} + + +def minimax_m3(kit, kv_dtype, max_tokens=1024): + """A four-layer MiniMax-M3 manager: layers 0 to 2 dense, layer 3 sparse with a bf16 index-K. + With NVFP4 KV only the sparse layer is NVFP4 (packed K/V and block scales); dense layers FP8. + Shut down when the block exits.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3KVCacheManagerV2 + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.mapping import Mapping + + return kit.managed( + lambda: MiniMaxM3KVCacheManagerV2( + kv_cache_config=KvCacheConfig( + max_tokens=max_tokens, dtype=kv_dtype, enable_block_reuse=True, host_cache_size=0 + ), + kv_cache_type=CacheType.SELF, + num_layers=4, + num_kv_heads=M3_KV_HEADS, + head_dim=M3_HEAD_DIM, + tokens_per_block=M3_TPB, + max_seq_len=512, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1), + dtype=getattr(DataType, M3_KV_DTYPES[kv_dtype]), + vocab_size=32000, + max_num_tokens=512, + sparse_layer_ids=[M3_SPARSE_LAYER], + disable_index_value_layer_ids=[M3_SPARSE_LAYER], + sparse_index_dim=M3_INDEX_DIM, + sparse_attention_config=SimpleNamespace(implementation="msa", indexer_kv_dtype="bf16"), + ) + ) + + +def check_minimax_m3_nvfp4(mgr, layout): + """Every buffer once, with the size, mapper and heads its storage has: FP8 K/V on dense layers, + half-width K/V plus one scale byte per 16 elements on the sparse layer, replicated index-K.""" + check_buffers(mgr, layout) + fp8 = (M3_TPB * M3_KV_HEADS * M3_HEAD_DIM, BufferMapper.INDEXED, M3_KV_HEADS) + packed = (fp8[0] // 2, BufferMapper.INDEXED, M3_KV_HEADS) + scale = (fp8[0] // 16, BufferMapper.INDEXED, M3_KV_HEADS) + expected = {(layer, role): fp8 for layer in range(4) for role in ("key", "value")} + expected.update({(M3_SPARSE_LAYER, role): packed for role in ("key", "value")}) + expected.update( + {(M3_SPARSE_LAYER, role): scale for role in ("key_block_scale", "value_block_scale")} + ) + expected[(M3_SPARSE_LAYER, "index_key")] = ( + M3_TPB * M3_INDEX_DIM * 2, + BufferMapper.REPLICATED, + None, + ) + described = { + (b.layer, b.role): (b.size, b.mapper, b.geometry.num_heads) for b in layout.buffers + } + assert described == expected + assert all(b.transfer and b.expansion == 1 for b in layout.buffers) + assert sorted(layer for group in layout.layer_groups for layer in group.layers) == [0, 1, 2, 3] + assert all(group.shard == ShardDesc() for group in layout.layer_groups) + + +def minimax_m3_nvfp4_lies(layout): + """Layouts that each misstate one NVFP4 fact: a lost scale, a scale on a dense layer, a scale + as wide as its data and a head-sharded index-K.""" + scale = next(b for b in layout.buffers if b.role == "key_block_scale") + index = next(b for b in layout.buffers if b.role == "index_key") + + def buffers(bufs): + return dataclasses.replace(layout, buffers=tuple(bufs)) + + def swap(old, new): + return buffers(new if b is old else b for b in layout.buffers) + + return { + "scale dropped": buffers(b for b in layout.buffers if b is not scale), + "scale on a dense layer": buffers((*layout.buffers, dataclasses.replace(scale, layer=0))), + "scale as wide as its data": swap( + scale, dataclasses.replace(scale, size=scale.size * 8, geometry=BufferGeometry()) + ), + "index-K head-sharded": swap(index, dataclasses.replace(index, mapper=BufferMapper.NHD)), + } + + +@gpu +@skip_pre_hopper +def test_minimax_m3_nvfp4_describes_its_packed_kv_and_block_scales(kit): + with minimax_m3(kit, "nvfp4") as mgr: + geometry = derive_layout(mgr) + check_pool_groups(mgr, geometry) + check_minimax_m3_nvfp4(mgr, geometry.layout) + for name, lie in minimax_m3_nvfp4_lies(geometry.layout).items(): + with pytest.raises(AssertionError): + check_minimax_m3_nvfp4(mgr, lie) + pytest.fail(f"the {name} lie went unnoticed") + + +def check_minimax_m3_ids(builds, id_of): + """``id_of(layout, slot_counts)`` is one id for every NVFP4 build and another for FP8.""" + ids = {} + for kv_dtype, layout, counts in builds: + ids.setdefault(kv_dtype, set()).add(id_of(layout, counts)) + assert len(ids["nvfp4"]) == 1, "the NVFP4 builds disagree" + assert ids["nvfp4"] != ids["fp8"], "NVFP4 and FP8 share an id" + + +@gpu +@skip_pre_hopper +def test_minimax_m3_nvfp4_and_fp8_caches_have_their_own_layout_ids(kit): + """NVFP4 builds that differ only in slot count agree; an FP8 cache of the same model reads its + bytes differently.""" + builds = [] + for kv_dtype, max_tokens in (("nvfp4", 1024), ("fp8", 1024), ("nvfp4", 16384)): + with minimax_m3(kit, kv_dtype, max_tokens) as mgr: + builds.append((kv_dtype, derive_layout(mgr).layout, native_pool_groups(mgr)[1])) + assert builds[0][2] != builds[2][2], "the two NVFP4 caches must differ in slot counts" + check_minimax_m3_ids(builds, lambda layout, counts: layout_id(layout)) + lies = { + "one id for every cache": ( + lambda layout, counts: bytes(32), + "NVFP4 and FP8 share an id", + ), + "an id over slot counts": ( + lambda layout, counts: layout_id(layout) + repr(sorted(counts.items())).encode(), + "the NVFP4 builds disagree", + ), + } + for name, (lie, caught_by) in lies.items(): + with pytest.raises(AssertionError, match=caught_by): + check_minimax_m3_ids(builds, lie) + pytest.fail(f"{name} went unnoticed") + + +@gpu +@skip_pre_blackwell +@pytest.mark.parametrize("fp8_ds_mla", [False, True], ids=["fp8", "fp8_ds_mla"]) +def test_the_deepseek_v4_layout_covers_every_virtual_layer_and_pool( + deepseek_v4_manager, fp8_ds_mla +): + with deepseek_v4_manager(fp8_ds_mla=fp8_ds_mla) as mgr: + assert mgr.use_fp8_ds_mla == fp8_ds_mla + geometry = derive_layout(mgr) + layout = geometry.layout + + life_cycles = mgr._life_cycle_by_layer_group() + assert len(layout.layer_groups) == len(life_cycles) >= 2 + assert list(geometry.windows) == [lc.window_size for lc in life_cycles] + assert None in geometry.windows and any(w is not None for w in geometry.windows) + + # Several pools per group: a page is their slots back to back. + native = {int(pg.pool_group_index): pg for pg in mgr.impl.pool_group_descs} + assert max(len(pg.pools) for pg in native.values()) >= 2 + for g, pg in native.items(): + widths = tuple( + int(p.slot_bytes) for p in sorted(pg.pools, key=lambda p: int(p.pool_index)) + ) + assert layout.pool_groups[g] == widths + assert geometry.page_bytes[g] == sum(widths) + + # Buffers and layer groups name layers by global layer id: for virtual layers, model layer + # times the number of attention types plus a type the internal layer holds. + internal = sorted(int(i) for group in mgr.impl.layer_grouping for i in group) + global_of = dict(zip(internal, global_layer_ids(mgr, internal))) + assert len(set(global_of.values())) == len(internal) + virtual = mgr._layer_attn_to_layer_id + num_types = max(t.value for t in type(next(iter(virtual))[1])) + 1 + for (model_layer, attn_type), layer_id in virtual.items(): + assert global_of[layer_id] // num_types == model_layer + held = {t.value for (m, t), lid in virtual.items() if lid == layer_id} + assert global_of[layer_id] % num_types in held + described = {(b.layer, b.role) for b in layout.buffers} + assert described == { + (global_of[int(b.layer_id)], str(b.role)) for b in mgr.impl.all_buffer_ids + } + layers = sorted(layer for group in layout.layer_groups for layer in group.layers) + assert layers == sorted(global_of.values()) + # One KV head: every layer group is the same content on every rank. + assert all(group.shard == ShardDesc() for group in layout.layer_groups) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py new file mode 100644 index 000000000000..e5f47987b30d --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``derive_layout`` against the page table native builds from the same live manager, for three +model shapes with small caches and no checkpoint: each transferred buffer's layer, pool, offset, +size and mapper, and each layer group's kind, window, heads, pools and view roles; and each +buffer's offset and size against the runtime's address and page size of its layer and role.""" + +from __future__ import annotations + +import dataclasses +from collections import defaultdict + +import pytest +import torch + +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager +from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, MapperKind +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import ( + BufferMapper, + ShardDesc, + derive_layout, +) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") + +# Both dense shapes: bf16, a 4096-token pool, sequences of up to 2048 tokens. +DENSE = dict(dtype=DataType.BF16, max_tokens=4096, max_seq_len=2048) +# TinyLlama-1.1B: 22 layers, 4 KV heads of 64. Rank 1 of TP2 holds heads 2 and 3. +TINYLLAMA_TP2_RANK1 = dict( + num_layers=22, num_kv_heads=4, head_dim=64, mapping=Mapping(world_size=2, rank=1, tp_size=2) +) +# DeepSeek-V3-Lite: 30 MLA layers, one 576-wide latent per token (512 + 64 rope), K only. +DSV3_LITE_MLA = dict( + kv_cache_type=CacheType.SELFKONLY, + num_layers=30, + num_kv_heads=1, + head_dim=512 + 64, + tokens_per_block=64, +) +# shape -> (keywords for conftest's ``real_manager``, or ``None`` for its ``hybrid_manager``: the +# Nemotron-Nano-9B-v2 shape with fewer layers; the shard every layer group names) +SHAPES = { + "tinyllama_tp2_rank1": ({**DENSE, **TINYLLAMA_TP2_RANK1}, ShardDesc(2, 1)), + "dsv3_lite_mla": ({**DENSE, **DSV3_LITE_MLA}, ShardDesc()), + "nemotron_h_hybrid": (None, ShardDesc()), +} + + +def _as_list(values): + return None if values is None else [int(v) for v in values] + + +def native_groups(page_table): + """What native's page table says per layer group. Pools are named by base address, layers by + the global ids of ``local_layers``; order-free parts are sorted.""" + out = [] + for lg in page_table.layer_groups: + pools = [ + (int(p.base_address), int(p.slot_bytes), int(p.num_slots)) + for p in page_table.pool_groups[lg.pool_group_idx].pools + ] + global_of = {int(ll.local_layer_id): int(ll.global_layer_id) for ll in lg.local_layers} + state = lg.kind == CacheKind.STATE + rows, roles, geometry = [], {}, {} + for pv in lg.pool_views: + key = (pools[pv.pool_idx][0], MapperKind(int(pv.mapper_kind)).name) + assert key not in roles, f"two views of pool and mapper {key}" + roles[key] = sorted(pv.pool_role) + if state: + geometry[key] = (_as_list(pv.section_bytes), pv.bytes_per_head) + for e in pv.buffer_entries: + layer = global_of[int(e["local_layer_id"])] + rows.append((layer, key[0], int(e["offset"]), int(e["size"]), key[1])) + out.append( + { + "kind": "state" if state else "attention", + "window": None if state else lg.sliding_window_size, + "heads": None if state else int(lg.kv_head_num_per_rank), + "layers": sorted(global_of.values()), + "pools": pools, + "rows": sorted(rows), + "roles": roles, + "state_geometry": geometry, + } + ) + return out + + +def layout_groups(geometry): + """The same per layer group from ``derive_layout``'s result, without local-only buffers (native + skips ignored roles); transferred buffers no group claims are listed last. Attention heads are + the head-sharded buffers' ``num_heads``: native's attention views have no per-head size.""" + layout = geometry.layout + out, claimed = [], set() + for desc in layout.layer_groups: + g = desc.pool_group + pools = [(p.base, p.slot_bytes, p.num_slots) for p in geometry.device_pools[g]] + assert layout.pool_groups[g] == tuple(w for _, w, _ in pools) + base_of = {p.index: p.base for p in geometry.device_pools[g]} + state = desc.kind == "state" + rows, roles, geometries, heads = [], defaultdict(list), defaultdict(set), set() + for i, b in enumerate(layout.buffers): + if b.layer not in desc.layers or b.pool[0] != g or not b.transfer: + continue + assert b.expansion == 1, "native reads no expansion" + claimed.add(i) + key = (base_of[b.pool[1]], b.mapper.name) + rows.append((b.layer, key[0], b.offset, b.size, key[1])) + roles[key].append(b.role) + if state: + section = b.geometry.section_bytes + section = None if section is None else tuple(section) + geometries[key].add((section, b.geometry.bytes_per_head)) + elif b.mapper in (BufferMapper.INDEXED, BufferMapper.NHD): + heads.add(b.geometry.num_heads) + assert all(len(v) == 1 for v in geometries.values()), f"roles disagree: {dict(geometries)}" + assert len(heads) <= 1, heads + out.append( + { + "kind": desc.kind, + "window": desc.window, + "heads": None if state else (heads.pop() if heads else 0), + "layers": sorted(desc.layers), + "pools": pools, + "rows": sorted(rows), + "roles": {k: sorted(set(v)) for k, v in roles.items()}, + "state_geometry": {k: (_as_list(s), h) for k, ((s, h),) in geometries.items()}, + } + ) + stray = [ + (b.layer, b.role) for i, b in enumerate(layout.buffers) if b.transfer and i not in claimed + ] + if stray: + out.append({"unclaimed buffers": stray}) + return out + + +def lies(geometry): + """Layout fields that each misstate one thing native reads.""" + layout = geometry.layout + first = next(b for b in layout.buffers if b.transfer) + + def buffers(**change): + return tuple(dataclasses.replace(b, **change) if b is first else b for b in layout.buffers) + + def first_group(**change): + return tuple( + dataclasses.replace(d, **change) if i == 0 else d + for i, d in enumerate(layout.layer_groups) + ) + + other = ( + BufferMapper.INDEXED if first.mapper is BufferMapper.REPLICATED else BufferMapper.REPLICATED + ) + out = { + "offset": {"buffers": buffers(offset=first.offset + first.size)}, + "layer": {"buffers": buffers(layer=first.layer + 1000)}, + "mapper": {"buffers": buffers(mapper=other)}, + "extra buffer": { + "buffers": layout.buffers + (dataclasses.replace(first, layer=first.layer + 1000),) + }, + "window": {"layer_groups": first_group(window=64)}, + } + # The first layer group pointed at another device pool group, when the layout has one. + mine = layout.layer_groups[0].pool_group + others = sorted(g for g in layout.pool_groups if g != mine) + if others: + out["pool group"] = {"layer_groups": first_group(pool_group=others[0])} + return out + + +@pytest.mark.parametrize("shape", list(SHAPES)) +def test_layout_matches_the_page_table_native_builds(real_manager, hybrid_manager, shape): + kwargs, shard = SHAPES[shape] + with real_manager(**kwargs) if kwargs else hybrid_manager() as mgr: + geometry = derive_layout(mgr) + page_table = build_page_table_from_manager(mgr) + native = native_groups(page_table) + derived = layout_groups(geometry) + + layout = geometry.layout + assert layout.tokens_per_block == page_table.tokens_per_block + assert len(derived) == len(native), derived[len(native) :] + assert all(group["rows"] for group in native) + for lg, (ours, theirs) in enumerate(zip(derived, native)): + for key in theirs: + assert ours[key] == theirs[key], f"layer group {lg}: {key}" + assert [d.shard for d in layout.layer_groups] == [shard] * len(native) + assert geometry.shards == ((shard.count, shard.index),) * len(native) + # The lender finds a group's pools through either; they must agree. + assert tuple(d.pool_group for d in layout.layer_groups) == geometry.pool_group_of + assert geometry.recurrent == tuple(g["kind"] == "state" for g in native) + + # The comparison is not vacuous: each misstated field shows. + for name, change in lies(geometry).items(): + lie = dataclasses.replace(geometry, layout=dataclasses.replace(layout, **change)) + assert layout_groups(lie) != native, f"the {name} lie went unnoticed" + + # Every buffer sits where the runtime keeps its layer and role. Native's views list roles + # per pool and mapper, so a layer's K and V trading places shows only here. + internal = { + int(ll.global_layer_id): int(ll.local_layer_id) + for lg in page_table.layer_groups + for ll in lg.local_layers + } + base = {(p.group, p.index): p.base for ps in geometry.device_pools.values() for p in ps} + for b in layout.buffers: + lid = internal[b.layer] + address = int(mgr.impl.get_mem_pool_base_address(lid, b.role, PageIndexMode.SHARED)) + page = int(mgr.impl.get_page_stride(lid, b.role)) * b.expansion + assert (address - base[b.pool], page) == (b.offset, b.size), (b.layer, b.role) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_names.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_names.py new file mode 100644 index 000000000000..5373a69c9cca --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_names.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Row and part names: who shares a name and who never does, including the ranks of a +tensor-parallel group. One full name and the part names of one layout are pinned: stored objects +stay reachable only while the bytes stay the same.""" + +import hashlib + +import numpy as np +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import ( + KEY_FORMAT_VERSION, + OBJECT_KEY_BYTES, + REUSE_KEY_BYTES, + Identity, + namespace, +) + +pytestmark = pytest.mark.cpu_only + +WHOLE = (1, 0) + + +def _reuse_keys(n, salt=0): + return np.array( + [[(salt * 31 + i * 7 + b) % 256 for b in range(REUSE_KEY_BYTES)] for i in range(n)], + dtype=np.uint8, + ) + + +def _identity(scope=b"", layout_id=b"layout", layers=((0,),), shards=None): + return Identity(scope, layout_id, layers, shards or [WHOLE] * len(layers)) + + +def _object_keys(identity, layer_group=0, n=3, salt=0): + return [row.tobytes() for row in identity.object_keys(layer_group, _reuse_keys(n, salt))] + + +def _golden(): + return Identity( + hashlib.sha256(b"golden-compute").digest(), + hashlib.sha256(b"golden-layout").digest(), + [(4, 5), (0, 1, 2, 3)], + [(4, 3), WHOLE], + ) + + +def test_a_name_is_namespace_then_key_then_the_layer_group_s_identity(): + # Local group 0 holds the largest layer, so it is canonical group 5 of six. + layers = [(50,), (0,), (10,), (20,), (30,), (40,)] + identity = _identity(layers=layers, shards=[(4, 3)] + [WHOLE] * 5) + rows = identity.object_keys(0, _reuse_keys(3)) + assert rows.shape == (3, OBJECT_KEY_BYTES) and rows.dtype == np.uint8 + names = [row.tobytes() for row in rows] + keys = _reuse_keys(3) + for i, name in enumerate(names): + assert len(name) == OBJECT_KEY_BYTES == 54 + assert name[: len(identity.namespace)] == identity.namespace + assert name[len(identity.namespace) : -6] == keys[i].tobytes() + # canonical group 5, shard 3 of 4: big-endian uint16 each + assert name[-6:] == bytes([0, 5, 0, 4, 0, 3]) + assert len(set(names)) == 3 + + +def test_one_full_name_is_pinned(): + """Any change to these bytes must bump ``KEY_FORMAT_VERSION``, so old objects become misses.""" + key = np.frombuffer(bytes(range(REUSE_KEY_BYTES)), dtype=np.uint8).reshape(1, REUSE_KEY_BYTES) + assert KEY_FORMAT_VERSION == 2 + assert _golden().object_keys(0, key)[0].tobytes().hex() == ( + "ac0cbcea5a240bde17f6bac29386876a" + "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f" + "000100040003" + ) + + +def test_region_labels_are_pinned(): + """A part is named by the layout's first 8 bytes and the canonical groups it serves.""" + identity = _golden() + assert identity.region_label([0]) == "94ede4f3bb6ba356:lg1" + assert identity.region_label([1]) == "94ede4f3bb6ba356:lg0" + assert identity.region_label([1, 0, 0]) == "94ede4f3bb6ba356:lg0+1" + + +def test_region_labels_do_not_depend_on_the_scope(): + layers = [(0,), (1,)] + a = _identity(b"model", layers=layers) + b = _identity(b"other model", layers=layers) + assert [a.region_label([g]) for g in range(2)] == [b.region_label([g]) for g in range(2)] + assert a.namespace != b.namespace + + +def test_local_group_numbering_does_not_leak_into_names(): + # The same canonical groups, numbered the other way round locally, name the same objects. + a = _identity(layers=[(0,), (1,)]) + b = _identity(layers=[(1,), (0,)]) + assert _object_keys(a, 0) == _object_keys(b, 1) and _object_keys(a, 1) == _object_keys(b, 0) + assert _object_keys(a, 0) != _object_keys(a, 1) + assert a.region_label([0]) == b.region_label([1]) + + +@pytest.mark.parametrize("other", [(b"model", b"other-layout"), (b"other-model", b"layout")]) +def test_a_different_scope_or_layout_never_collides(other): + a = _identity(b"model", b"layout") + b = _identity(*other) + assert not set(_object_keys(a)) & set(_object_keys(b)) + + +def test_namespace_fields_are_length_prefixed(): + assert namespace(b"a", b"bc") != namespace(b"ab", b"c") + + +def test_an_unknown_layer_group_or_an_empty_layout_is_rejected(): + with pytest.raises(ValueError, match="unknown layer group"): + _identity().object_keys(1, _reuse_keys(1)) + with pytest.raises(ValueError, match="layout_id"): + _identity(b"model", b"") + with pytest.raises(ValueError, match="layout_id"): + namespace(b"model", b"") + + +def test_a_scope_longer_than_its_length_prefix_is_rejected(): + assert len(namespace(b"s" * 0xFFFF, b"layout")) == 16 + with pytest.raises(ValueError, match="at most 65535"): + namespace(b"s" * 0x10000, b"layout") + + +# -- shards -------------------------------------------------------------------------------------- + + +def test_a_replicated_group_is_named_alike_on_every_rank(): + """Every buffer of the group holds the same bytes on every rank (a single KV head, a + replicated indexer): the group is whole on each rank, so ranks share its names.""" + ranks = [_identity(b"model", layers=[(0, 1), (2, 3)]) for _ in range(4)] + assert len({tuple(_object_keys(identity)) for identity in ranks}) == 1 + assert _object_keys(ranks[0])[0][-4:] == bytes([0, 1, 0, 0]), "not one share of one" + + +def test_a_head_sharded_group_is_named_per_share(): + layers = [(0, 1), (2, 3)] + ranks = [_identity(b"model", layers=layers, shards=[(4, r), (4, r)]) for r in range(4)] + names = [set(_object_keys(identity)) for identity in ranks] + for i in range(4): + for j in range(i + 1, 4): + assert not names[i] & names[j] + + +def test_groups_of_one_rank_are_judged_one_by_one(): + """A replicated group stays shared across ranks even when another group of the same rank is + head-sharded.""" + layers = [(0, 1), (2, 3)] + a = _identity(b"m", layers=layers, shards=[(2, 0), WHOLE]) + b = _identity(b"m", layers=layers, shards=[(2, 1), WHOLE]) + assert _object_keys(a, 1) == _object_keys(b, 1) + assert not set(_object_keys(a, 0)) & set(_object_keys(b, 0)) + + +@pytest.mark.parametrize("shard", [(2, 2), (0, 0), (2, -1)]) +def test_a_shard_names_a_share_of_its_count(shard): + with pytest.raises(ValueError, match="shard"): + _identity(shards=[shard]) + + +def test_every_layer_group_has_a_shard(): + with pytest.raises(ValueError, match="shards"): + Identity(b"", b"layout", [(0,), (1,)], [WHOLE]) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py new file mode 100644 index 000000000000..d08fa1ea6cd5 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py @@ -0,0 +1,895 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The lender's public surface and boundaries: the nine names, their members, what importing the +package loads, and who may import its private modules. No GPU.""" + +import ast +import dataclasses +import inspect +import json +import os +import re +import subprocess +import sys +import weakref +from pathlib import Path +from typing import List, Protocol, Set + +import numpy as np +import pytest + +import tensorrt_llm +from tensorrt_llm._torch.pyexecutor.kv_cache import sharing +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + LayerGroupBlocks, + Lease, + LeaseView, + Lender, + MemoryRegion, + RegionHold, + ResumeBounds, + StagingOptions, + attach_lender, +) + +SH = "tensorrt_llm._torch.pyexecutor.kv_cache.sharing" +MGR = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" +PUBLIC = sorted( + "LayerGroupBlocks Lease MemoryRegion RegionHold ResumeBounds LeaseView Lender StagingOptions " + "attach_lender".split() +) +MODULES = [ + "__init__.py", + "_fetch_ledger.py", + "_identity.py", + "_layout.py", + "_lender.py", + "_manager.py", + "_ranges.py", + "_slots.py", +] +MODULES = sorted(MODULES + ["_types.py"]) +PKG_DIR = Path(sharing.__file__).resolve().parent +ROOT = Path(tensorrt_llm.__file__).resolve().parent +TESTS_DIR = Path(__file__).resolve().parent +# Manager and runtime members only the manager facade may touch. +INTERNALS = { + "impl", + "init_config", + "layer_grouping", + "pool_group_descs", + "_has_lossy_cold_page_codec", + "get_disagg_global_layer_ids", + "get_disagg_role_mapper_kinds", + "get_disagg_role_layouts", + "get_disagg_ignored_roles", + "get_life_cycle_pool_group_indices", + "_reuse_token_source", + "_augment_tokens_for_block_reuse", + "_stale_block_range", + "_resize_for_served_prefix", + "_fill_fresh_kv_pages", + "_fresh_page_fill", + "_fresh_pages_filled", + "_can_publish_block_reuse", + "block_reuse_policy", + "_draft_prompt_lookahead", + "kv_connector_manager", + "commit_min_snapshot", + "reuse_match_backoff", + "_stream", + "_layer_attn_to_layer_id", + "_sharing", + "host_kv_cache_block_offsets", + "kv_cache_map", + "py_multimodal_data", + "multimodal_hashes", + "multimodal_positions", + "multimodal_lengths", + "try_get_encoder_output_len", + "py_return_context_logits", + "py_additional_outputs", + "py_result", + "additional_context_outputs", + "prompt_len", + "is_draft", +} + + +FIRST_YEAR = 2026 +COPYRIGHT = re.compile( + r"# SPDX-FileCopyrightText: Copyright \(c\) (?:(20\d\d)-)?(20\d\d) NVIDIA CORPORATION & " + r"AFFILIATES\. All rights reserved\.$" +) +LICENSE = [ + "# SPDX-License-Identifier: Apache-2.0", + "#", + '# Licensed under the Apache License, Version 2.0 (the "License");', + "# you may not use this file except in compliance with the License.", + "# You may obtain a copy of the License at", + "#", + "# http://www.apache.org/licenses/LICENSE-2.0", + "#", + "# Unless required by applicable law or agreed to in writing, software", + '# distributed under the License is distributed on an "AS IS" BASIS,', + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.", + "# See the License for the specific language governing permissions and", + "# limitations under the License.", +] + + +def has_header(source: str) -> bool: + """The NVIDIA header: a copyright year, or a range of years, ending in 2026 or later, then the + Apache license text.""" + lines = source.splitlines() + match = COPYRIGHT.match(lines[0]) if lines else None + if match is None or lines[1 : 1 + len(LICENSE)] != LICENSE: + return False + first, last = int(match.group(1) or match.group(2)), int(match.group(2)) + return first <= last and last >= FIRST_YEAR + + +def package_files() -> List[Path]: + return sorted(PKG_DIR.glob("*.py")) + + +def module_name(path: Path, root: Path = ROOT) -> str: + relative = path.resolve().relative_to(root.parent).with_suffix("") + parts = list(relative.parts) + if parts[-1] == "__init__": + parts.pop() + return ".".join(parts) + + +def package_of(path: Path, root: Path = ROOT) -> str: + """The package a file's relative imports resolve against.""" + name = module_name(path, root) + return name if path.name == "__init__.py" else name.rpartition(".")[0] + + +def imported_modules(source: str, package: str) -> Set[str]: + """Every module (and ``module.name``) a source imports, lazy imports and + ``importlib.import_module`` with a literal included; relative imports resolved.""" + found: Set[str] = set() + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Import): + found.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + if node.level: + base = package.split(".") + base = base[: len(base) - (node.level - 1)] + module = ".".join(base + ([node.module] if node.module else [])) + else: + module = node.module or "" + found.add(module) + found.update(f"{module}.{alias.name}" for alias in node.names) + elif isinstance(node, ast.Call): + func = node.func + name = func.attr if isinstance(func, ast.Attribute) else getattr(func, "id", "") + first = node.args[0] if node.args else None + if name in ("import_module", "__import__") and isinstance(first, ast.Constant): + found.add(str(first.value)) + return found + + +PRIVATE_IMPORT = re.compile(r"^" + re.escape(SH) + r"\._") + + +def private_imports(source: str, package: str) -> List[str]: + """What a source outside the package imports of the package's private modules.""" + hits = sorted(m for m in imported_modules(source, package) if PRIVATE_IMPORT.match(m)) + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Constant) and isinstance(node.value, str): + if "kv_cache.sharing._" in node.value: + hits.append(node.value) + return hits + + +def private_importers(root: Path, pkg_dir: Path) -> dict: + """Modules under ``root`` outside ``pkg_dir`` that import the package's private modules.""" + offenders = {} + for path in root.rglob("*.py"): + if pkg_dir in path.resolve().parents: + continue + source = path.read_text(encoding="utf-8", errors="replace") + if "sharing" not in source: + continue + try: + hits = private_imports(source, package_of(path, root)) + except SyntaxError: + continue # not importable, so it imports nothing + if hits: + offenders[path.relative_to(root).as_posix()] = hits + return offenders + + +def disaggregation_imports(source: str, package: str) -> List[str]: + prefix = "tensorrt_llm._torch.disaggregation" + return sorted(m for m in imported_modules(source, package) if m.startswith(prefix)) + + +def internal_reads(source: str) -> List[str]: + """Attribute names (or ``getattr`` literals) of manager internals a source uses.""" + hits = [] + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Attribute) and node.attr in INTERNALS: + hits.append(node.attr) + elif isinstance(node, ast.Constant) and node.value in INTERNALS: + hits.append(str(node.value)) + return sorted(hits) + + +def threads_or_finalizers(source: str) -> List[str]: + hits = [] + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Import): + hits += [a.name for a in node.names if a.name.split(".")[0] in _THREADING] + elif isinstance(node, ast.ImportFrom) and (node.module or "").split(".")[0] in _THREADING: + hits.append(node.module) + elif isinstance(node, ast.ImportFrom) and any(a.name == "finalize" for a in node.names): + hits.append("finalize") + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "__del__": + hits.append("__del__") + elif isinstance(node, ast.Attribute) and node.attr == "finalize": + hits.append("finalize") + return hits + + +_THREADING = {"threading", "_thread", "concurrent", "multiprocessing"} + + +def identifiers(source: str) -> Set[str]: + """Every name the code defines or uses: variables, attributes, functions, classes, arguments + and imported names; comments and docstrings aside.""" + found: Set[str] = set() + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Name): + found.add(node.id) + elif isinstance(node, ast.Attribute): + found.add(node.attr) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + found.add(node.name) + elif isinstance(node, ast.arg): + found.add(node.arg) + elif isinstance(node, ast.alias): + found.add((node.asname or node.name).split(".")[-1]) + return found + + +def compute_identities(source: str) -> List[str]: + return sorted(n for n in identifiers(source) if "compute_id" in n) + + +# What no file of the package holds. Which instance computed a block is the caller's: the lender +# names blocks by layout, scope and reuse key alone. +PACKAGE_SCANS = { + "disaggregation_imports": lambda path: disaggregation_imports( + path.read_text(), package_of(path) + ), + "threads_locks_or_finalizers": lambda path: threads_or_finalizers(path.read_text()), + "compute_identities": lambda path: compute_identities(path.read_text()), +} + + +def run_python(code: str) -> dict: + """Run ``code`` in a new interpreter importing this tree's ``tensorrt_llm``; it prints JSON.""" + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join(filter(None, [str(ROOT.parent), env.get("PYTHONPATH")])) + done = subprocess.run( + [sys.executable, "-c", code], + cwd=str(ROOT.parent), + env=env, + capture_output=True, + text=True, + timeout=600, + ) + assert done.returncode == 0, done.stderr[-4000:] + result = json.loads(done.stdout.strip().splitlines()[-1]) + # Proof the child imported the package under test, not another installed copy. + assert Path(result["file"]).resolve() == Path(sharing.__file__).resolve() + return result + + +# -- the names -------------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("raw,backoff", [(0, 0), (1, 0), (1, 2)]) +def test_the_manager_exposes_read_only_raw_lookahead_and_the_lender_keeps_both_guards(raw, backoff): + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + manager = object.__new__(KVCacheManagerV2) + manager._draft_prompt_lookahead = raw + manager.reuse_match_backoff = backoff + assert manager.draft_prompt_lookahead == raw + with pytest.raises(AttributeError): + manager.draft_prompt_lookahead = 0 + assert _manager.prompt_lookahead(manager) == max(raw, backoff) + + +@pytest.mark.cpu_only +def test_importing_the_package_loads_its_types_alone_and_exposes_nothing_else(): + result = run_python( + "import json, sys\n" + f"import {SH} as s\n" + "print(json.dumps({'file': s.__file__, 'all': s.__all__,\n" + " 'public': sorted(n for n in vars(s) if not n.startswith('_')),\n" + f" 'loaded': sorted(m for m in sys.modules if m.startswith('{SH}'))}}))\n" + ) + assert result["public"] == sorted(result["all"]) == PUBLIC + assert result["loaded"] == [SH, f"{SH}._types"] + + +def subpackages(pkg_dir: Path) -> List[str]: + """The directories under ``pkg_dir`` that could hold modules.""" + return [p.name for p in pkg_dir.iterdir() if p.is_dir() and p.name != "__pycache__"] + + +@pytest.mark.cpu_only +def test_the_package_is_its_init_and_eight_private_modules(tmp_path): + assert [p.name for p in package_files()] == MODULES + assert subpackages(PKG_DIR) == [] + (tmp_path / "__pycache__").mkdir() + (tmp_path / "planted").mkdir() + assert subpackages(tmp_path) == ["planted"], "the listing sees no subpackage" + + +@pytest.mark.cpu_only +def test_the_protocols_keep_their_members_and_stay_unrelated(): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + def members(cls): + return {n for n in vars(cls) if not n.startswith("_")} + + assert members(Lease) == {"poll", "failure", "mark_arrived", "release"} + assert members(Lender) == { + "regions", + "hold_regions", + "lend_read", + "lend_write", + "resume_bounds", + } + assert members(RegionHold) == {"release"} + assert Lease not in RegionHold.__mro__ and RegionHold not in Lease.__mro__ + for protocol in (Lease, Lender, RegionHold): + assert Protocol in protocol.__mro__ + assert not isinstance(object(), protocol), "checkable at run time" + assert isinstance(Lease.failure, property) and isinstance(Lender.regions, property) + # The checks are structural: a lease passes as a hold. + lease = object.__new__(_lender._StagingLease) + assert isinstance(lease, RegionHold) + + def params(func): + return list(inspect.signature(func).parameters) + + assert params(Lease.poll) == ["self"] + assert params(Lease.mark_arrived) == ["self", "masks"] + assert params(Lease.release) == ["self"] + assert params(RegionHold.release) == ["self"] + assert params(Lender.hold_regions) == ["self"] + assert params(Lender.lend_read) == ["self", "request", "start", "end"] + assert params(Lender.lend_write) == ["self", "request", "start", "end"] + assert params(Lender.resume_bounds) == ["self", "request"] + + +def public_names(obj) -> Set[str]: + return {n for n in dir(obj) if not n.startswith("_")} + + +LEASE_MEMBERS = {"poll", "failure", "mark_arrived", "release"} +STAGING_MEMBERS = {"regions", "hold_regions", "lend_read", "lend_write", "resume_bounds"} + + +@pytest.mark.cpu_only +def test_lenders_leases_and_holds_show_only_their_protocol(): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + class Owner: # anything a weak reference can point to + pass + + owner = Owner() + ref = weakref.ref(owner) + layout = SimpleNamespace(pool_groups=(), windows=()) # all the constructors read + shown = { + "staging lender": (_lender.Staging(ref, layout, None, (), None, ref), STAGING_MEMBERS), + "staging lease": (_lender._StagingLease(owner, "read", 1), LEASE_MEMBERS), + "staging hold": (_lender._StagingHold(owner), {"release"}), + } + assert {what: public_names(obj) for what, (obj, _) in shown.items()} == { + what: members for what, (_, members) in shown.items() + } + # A lease showing its view as an attribute, and a lender showing a manager hook, are found. + leaky = _lender._StagingLease(owner, "read", 1) + leaky.view = None + assert public_names(leaky) - LEASE_MEMBERS == {"view"} + hooked = type("Hooked", (_lender.Staging,), {"on_free": _lender.Staging._on_free}) + assert public_names(hooked(ref, layout, None, (), None, ref)) - STAGING_MEMBERS == {"on_free"} + + +@pytest.mark.cpu_only +def test_the_types_keep_their_fields(): + def fields(cls): + """``name`` or ``name=default`` per field, in order.""" + missing = dataclasses.MISSING + return " ".join( + f.name if f.default is missing else f"{f.name}={f.default!r}" + for f in dataclasses.fields(cls) + ) + + assert fields(StagingOptions) == "tokens_per_fetch planned_fetches=1 max_bytes=None" + assert ( + fields(MemoryRegion) + == "layout_label address nbytes slot_bytes slot_count layer_group_indices" + ) + assert ( + fields(LayerGroupBlocks) + == "layer_group_index block_ordinals object_keys=None addresses=None region_index=None" + ) + assert fields(LeaseView) == "groups" + for cls in (StagingOptions, MemoryRegion, LayerGroupBlocks, LeaseView): + assert cls.__dataclass_params__.frozen, cls + assert MemoryRegion("p", 4096, 8192, 4096, 2, (0,)) == MemoryRegion( + "p", 4096, 8192, 4096, 2, (0,) + ) + assert ResumeBounds._fields == ("max_position", "min_position") + assert issubclass(ResumeBounds, tuple) + + def beyond_fields(cls): + named = {f.name for f in dataclasses.fields(cls)} + return {n for n in dir(cls) if not n.startswith("_")} - named + + assert beyond_fields(LayerGroupBlocks) == {"select"} + assert beyond_fields(LeaseView) == {"num_block_entries", "row_masks"} + assert beyond_fields(MemoryRegion) == beyond_fields(StagingOptions) == set() + assert {n for n in dir(ResumeBounds) if not n.startswith("_")} == { + "max_position", + "min_position", + "count", + "index", + } + + +@pytest.mark.cpu_only +def test_public_names_are_independent_of_the_lending_mechanism(): + assert {name for name in sharing.__all__ if "Staging" in name} == {"StagingOptions"} + for name in ("StagingLender", "StagingRegion", "StagingHold", "attach_staging"): + assert not hasattr(sharing, name), name + assert not hasattr(Lender, "hold_staging") + assert not hasattr(LeaseView, "block_masks") + assert not hasattr(LayerGroupBlocks, "staging_addresses") + + +@pytest.mark.cpu_only +def test_views_and_regions_match_the_backend_contract_member_names(): + """Exercise the structural surface without importing the unmerged backend contract.""" + region = MemoryRegion("layout", 4096, 128, 64, 2, (3,)) + group = LayerGroupBlocks( + layer_group_index=3, + block_ordinals=np.array([0, 7]), + object_keys=np.zeros((2, 54), dtype=np.uint8), + addresses=np.array([4096, 4160]), + region_index=0, + ) + view = LeaseView((group,)) + assert all( + hasattr(region, name) + for name in ("layout_label", "address", "nbytes", "slot_bytes", "slot_count") + ) + assert all( + hasattr(group, name) + for name in ("layer_group_index", "object_keys", "addresses", "region_index", "__len__") + ) + assert all(hasattr(view, name) for name in ("groups", "row_masks")) + assert list(inspect.signature(LeaseView.row_masks).parameters) == ["self", "fill_value"] + false_masks = view.row_masks() + true_masks = view.row_masks(fill_value=True) + assert len(group) == 2 and len(false_masks) == len(true_masks) == len(view.groups) == 1 + for masks, fill in ((false_masks, False), (true_masks, True)): + assert masks[0].shape == (2,) and masks[0].dtype == np.bool_ + assert masks[0].flags.writeable and np.all(masks[0] == fill) + assert region.nbytes == region.slot_bytes * region.slot_count + assert np.all(group.addresses >= region.address) + assert np.all(group.addresses + region.slot_bytes <= region.address + region.nbytes) + + +@pytest.mark.cpu_only +def test_interface_docstrings_state_guarantees_without_staging_mechanisms(): + forbidden = re.compile(r"\b(staging|stream|copy|copies|host|device|gpu)\b", re.IGNORECASE) + for cls in (MemoryRegion, LayerGroupBlocks, LeaseView, ResumeBounds, Lease, RegionHold, Lender): + entries = [cls] + [member for name, member in vars(cls).items() if not name.startswith("_")] + for member in entries: + doc = inspect.getdoc(member) or "" + assert not forbidden.search(doc), (cls.__name__, doc) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "invalid_hook", ["_on_free", "_on_shrink", "_on_reset", "_on_shutdown", "_on_caches_closed"] +) +@pytest.mark.parametrize("replacement", [None, 1, "absent"]) +def test_install_checks_every_private_lifecycle_hook_before_attaching(invalid_hook, replacement): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + hooks = { + name: lambda *args: None + for name in ("_on_free", "_on_shrink", "_on_reset", "_on_shutdown", "_on_caches_closed") + } + if replacement == "absent": + hooks.pop(invalid_hook) + else: + hooks[invalid_hook] = replacement + manager = SimpleNamespace() + with pytest.raises(TypeError, match="all five manager lifecycle hooks"): + _manager.install_lender(manager, SimpleNamespace(**hooks)) + assert "_sharing" not in vars(manager) + + +@pytest.mark.cpu_only +def test_install_accepts_the_complete_private_lifecycle_protocol(): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + lender = SimpleNamespace( + **{ + name: lambda *args: None + for name in ("_on_free", "_on_shrink", "_on_reset", "_on_shutdown", "_on_caches_closed") + } + ) + manager = SimpleNamespace() + assert isinstance(lender, _manager._ManagerHooks) + _manager.install_lender(manager, lender) + assert manager._sharing is lender + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("config_source", ["manager", "runtime", "runtime_none"]) +def test_manager_facade_reads_layer_configs_by_id_with_the_existing_fallback(config_source): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + layers = [ + SimpleNamespace(layer_id="7"), + SimpleNamespace(layer_id=4), + SimpleNamespace(), + SimpleNamespace(layer_id=0), + ] + config = SimpleNamespace(layers=layers) + if config_source == "manager": + manager = SimpleNamespace(kv_cache_manager_py_config=config) + else: + manager = SimpleNamespace(impl=SimpleNamespace(init_config=config)) + if config_source == "runtime_none": + manager.kv_cache_manager_py_config = None + actual = _manager.layer_configs(manager) + assert list(actual) == [7, 4, 2, 0] + assert all(actual[layer_id] is layer for layer_id, layer in zip([7, 4, 2, 0], layers)) + + +@pytest.mark.cpu_only +def test_manager_facade_preserves_layer_group_order_and_normalizes_ids(): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + grouping = [[np.int64(5), "7"], [2]] + manager = SimpleNamespace(impl=SimpleNamespace(layer_grouping=grouping)) + actual = _manager.layer_group_layer_ids(manager) + assert actual == [[5, 7], [2]] + assert all(type(layer_id) is int for group in actual for layer_id in group) + assert actual is not grouping and all(a is not b for a, b in zip(actual, grouping)) + + +@pytest.mark.cpu_only +def test_manager_facade_returns_the_runtime_pool_group_descriptors_unchanged(): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + descriptors = [object(), object()] + manager = SimpleNamespace(impl=SimpleNamespace(pool_group_descs=descriptors)) + assert _manager.pool_group_descriptors(manager) is descriptors + + +@pytest.mark.cpu_only +def test_the_attach_functions_keep_their_signatures(): + kind = inspect.Parameter + + def shape(func): + return [(p.name, p.kind, p.default) for p in inspect.signature(func).parameters.values()] + + assert shape(attach_lender) == [ + ("manager", kind.POSITIONAL_OR_KEYWORD, kind.empty), + ("scope", kind.KEYWORD_ONLY, kind.empty), + ("staging", kind.KEYWORD_ONLY, kind.empty), + ] + + +@pytest.mark.cpu_only +def test_a_name_is_as_long_as_the_identity_makes_it(): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _identity, _types + + assert _types._OBJECT_KEY_BYTES == _identity.OBJECT_KEY_BYTES == 54 + + +def public_lending_names(cls) -> List[str]: + """The public names of ``cls`` about lending.""" + words = re.compile(r"lend|lent|sharing|staging|retain") + return [n for n in dir(cls) if not n.startswith("_") and words.search(n)] + + +@pytest.mark.cpu_only +def test_the_manager_has_no_public_name_about_lending(): + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + + assert "_sharing" in vars(KVCacheManagerV2) and KVCacheManagerV2._sharing is None + assert public_lending_names(KVCacheManagerV2) == [] + planted = type("Planted", (KVCacheManagerV2,), {"lend_read": None, "staging_part": None}) + assert public_lending_names(planted) == ["lend_read", "staging_part"], "the scan sees nothing" + + +# -- boundaries ------------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_nothing_outside_the_package_imports_its_private_modules(): + assert private_importers(ROOT, PKG_DIR) == {} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("scan", list(PACKAGE_SCANS)) +def test_the_package_holds_nothing_a_scan_looks_for(scan): + hits = {path.name: PACKAGE_SCANS[scan](path) for path in package_files()} + assert {name: found for name, found in hits.items() if found} == {} + + +@pytest.mark.cpu_only +def test_the_manager_never_imports_the_package(): + path = ROOT / "_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py" + source = path.read_text() + loaded = imported_modules(source, package_of(path)) + assert [m for m in loaded if m.startswith(SH)] == [] + planted = source + "\nfrom .sharing import attach_lender\n" + assert [m for m in imported_modules(planted, package_of(path)) if m.startswith(SH)] + + +@pytest.mark.cpu_only +def test_the_native_path_loads_nothing_of_the_package(): + result = run_python( + "import json, sys\n" + "import tensorrt_llm._torch.disaggregation.transceiver\n" + f"import {MGR}\n" + f"loaded = sorted(m for m in sys.modules if m.startswith('{SH}'))\n" + f"import {SH} as s\n" + "print(json.dumps({'file': s.__file__, 'loaded': loaded}))\n" + ) + assert result["loaded"] == [] + + +@pytest.mark.cpu_only +def test_only_the_manager_facade_touches_manager_internals(): + offenders = {} + for path in package_files(): + if path.name == "_manager.py": + continue + hits = internal_reads(path.read_text()) + if hits: + offenders[path.name] = hits + assert offenders == {} + + +@pytest.mark.cpu_only +def test_the_context_output_check_reads_the_request_alone(): + """The fetch's context-output check reads two request attributes and never the executor's + result, whose additional outputs concatenate their chunks at each read.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + class Unread: + def __getattr__(self, name): + raise AssertionError(f"read py_result.{name}") + + class Request: + py_result = Unread() + + def __init__(self, logits, outputs): + self.py_return_context_logits, self.py_additional_outputs = logits, outputs + + cases = [(False, None, False), (False, [], False), (True, None, True), (False, ["x"], True)] + for logits, outputs, returns in cases: + assert _manager.returns_context_outputs(Request(logits, outputs)) is returns + + +# -- headers ---------------------------------------------------------------------------------- + + +def written_files() -> List[Path]: + return package_files() + sorted(TESTS_DIR.glob("*.py")) + + +@pytest.mark.cpu_only +def test_every_new_file_carries_the_license_header(): + assert [p.name for p in written_files() if not has_header(p.read_text())] == [] + body = "\n".join(LICENSE) + + def header(years: str) -> str: + owner = "NVIDIA CORPORATION & AFFILIATES. All rights reserved." + return f"# SPDX-FileCopyrightText: Copyright (c) {years} {owner}\n{body}" + + for years in ("2026", "2027", "2025-2026", "2026-2027", "2031"): + assert has_header(header(years)), years + for years in ("2025", "2024-2025", "2027-2026"): + assert not has_header(header(years)), years + assert not has_header("# SPDX-FileCopyrightText: Copyright (c) 2026 Someone else.\n" + body) + assert not has_header(header("2026").splitlines()[0] + "\n# SPDX-License-Identifier: MIT\n") + + +# -- the public types' own rules -------------------------------------------------------------- + + +def staging_blocks(n=3, layer_group_index=0, region_index=0): + return LayerGroupBlocks( + layer_group_index, + np.arange(n), + np.repeat(np.arange(n, dtype=np.uint8)[:, None], 54, axis=1), + np.arange(n, dtype=np.int64) * 4096, + region_index, + ) + + +@pytest.mark.cpu_only +def test_layer_group_blocks_carry_all_three_placements_or_none_in_their_shapes(): + block_ordinals = np.arange(3) + object_keys = np.zeros((3, 54), np.uint8) + addresses = np.arange(3, dtype=np.int64) + refused = [ + (ValueError, (0, block_ordinals, object_keys, None, None)), + (ValueError, (0, block_ordinals, None, addresses, 0)), + (ValueError, (0, block_ordinals, object_keys, addresses, None)), + (ValueError, (0, np.zeros((3, 1)))), + (ValueError, (0, block_ordinals, np.zeros((3, 53), np.uint8), addresses, 0)), + (ValueError, (0, block_ordinals, np.zeros((3, 54), np.int8), addresses, 0)), + (ValueError, (0, block_ordinals, object_keys, addresses[:2], 0)), + (ValueError, (-1, block_ordinals)), + (ValueError, (0, block_ordinals, object_keys, addresses, -1)), + (TypeError, (True, block_ordinals)), + (TypeError, (0, block_ordinals, object_keys, addresses, 1.0)), + ] + for error, args in refused: + with pytest.raises(error): + LayerGroupBlocks(*args) + unplaced = LayerGroupBlocks(1, block_ordinals) + assert (unplaced.object_keys, unplaced.addresses, unplaced.region_index) == ( + None, + None, + None, + ) + assert len(unplaced) == 3 and unplaced.block_ordinals.dtype == np.int64 + + +@pytest.mark.cpu_only +def test_layer_group_blocks_arrays_are_read_only_views(): + block_ordinals = np.arange(3, dtype=np.int64) + blocks = LayerGroupBlocks( + 0, block_ordinals, np.zeros((3, 54), np.uint8), np.arange(3, dtype=np.int64), 0 + ) + for array in (blocks.block_ordinals, blocks.object_keys, blocks.addresses): + assert not array.flags.writeable + with pytest.raises(ValueError): + array[0] = 1 + assert block_ordinals.flags.writeable, "the caller's own array is left as it was" + + +@pytest.mark.cpu_only +def test_select_keeps_the_marked_block_entries_in_order(): + blocks = staging_blocks(4, layer_group_index=2, region_index=1) + picked = blocks.select(np.array([True, False, True, True])) + assert picked.block_ordinals.tolist() == [0, 2, 3] + assert picked.addresses.tolist() == [0, 8192, 12288] + np.testing.assert_array_equal(picked.object_keys, blocks.object_keys[[0, 2, 3]]) + assert picked.object_keys.shape == (3, 54) and ( + picked.layer_group_index, + picked.region_index, + ) == (2, 1) + assert LayerGroupBlocks(0, np.arange(2)).select(np.array([False, True])).object_keys is None + with pytest.raises(ValueError): + blocks.select(np.array([True, False])) + with pytest.raises(ValueError): + blocks.select(np.array([1, 0, 1, 1])) + + +@pytest.mark.cpu_only +def test_a_lease_view_holds_one_block_collection_per_layer_group(): + view = LeaseView([staging_blocks(3, 0), staging_blocks(2, 1, region_index=1)]) + assert isinstance(view.groups, tuple) and view.num_block_entries == 5 + masks = view.row_masks() + assert [m.shape for m in masks] == [(3,), (2,)] + assert all(m.dtype == np.bool_ and not m.any() and m.flags.writeable for m in masks) + assert all(m.all() for m in view.row_masks(True)) + assert LeaseView(()).num_block_entries == 0 and LeaseView(()).row_masks() == () + with pytest.raises(ValueError): + LeaseView((staging_blocks(3, 0), staging_blocks(1, 0))) + with pytest.raises(TypeError): + LeaseView((object(),)) + + +@pytest.mark.cpu_only +def test_sparse_block_entries_keep_group_position_masks_and_placement_alignment(): + keys = np.arange(3 * 54, dtype=np.uint8).reshape(3, 54) + blocks = LayerGroupBlocks( + layer_group_index=7, + block_ordinals=np.array([0, 3, 9]), + object_keys=keys, + addresses=np.array([12288, 4096, 8192]), + region_index=2, + ) + empty = LayerGroupBlocks(layer_group_index=2, block_ordinals=np.array([], dtype=np.int64)) + repeated = LayerGroupBlocks(layer_group_index=4, block_ordinals=np.array([3, 9])) + view = LeaseView(groups=(blocks, empty, repeated)) + # Five group/block entries, but only three distinct token block ordinals. + assert view.num_block_entries == 5 + masks = view.row_masks() + assert [mask.shape for mask in masks] == [(3,), (0,), (2,)] + masks[0][1] = True + picked = view.groups[0].select(masks[0]) + assert picked.block_ordinals.tolist() == [3] + assert picked.addresses.tolist() == [4096] + np.testing.assert_array_equal(picked.object_keys, keys[[1]]) + assert (picked.layer_group_index, picked.region_index) == (7, 2) + assert not masks[2].any() + + +@pytest.mark.cpu_only +def test_renamed_bounds_keep_maximum_first_and_region_fields_accept_keywords(): + bounds = ResumeBounds(max_position=96, min_position=32) + assert bounds._fields == ("max_position", "min_position") + assert tuple(bounds) == (96, 32) + assert ResumeBounds(0, 1).min_position > ResumeBounds(0, 1).max_position + region = MemoryRegion( + layout_label="layout:lg0+1", + address=4096, + nbytes=8192, + slot_bytes=4096, + slot_count=2, + layer_group_indices=(7, 2), + ) + assert region.layout_label == "layout:lg0+1" + assert region.layer_group_indices == (7, 2) + + +@pytest.mark.cpu_only +def test_staging_options_take_positive_integers(): + options = StagingOptions(256) + assert (options.tokens_per_fetch, options.planned_fetches, options.max_bytes) == (256, 1, None) + assert StagingOptions(256, planned_fetches=4, max_bytes=1 << 30).max_bytes == 1 << 30 + for bad in (dict(tokens_per_fetch=0), dict(tokens_per_fetch=8, planned_fetches=0)): + with pytest.raises(ValueError): + StagingOptions(**bad) + with pytest.raises(ValueError): + StagingOptions(8, max_bytes=0) + for bad in ( + dict(tokens_per_fetch=True), + dict(tokens_per_fetch=8.0), + dict(tokens_per_fetch=8, max_bytes=1.5), + ): + with pytest.raises(TypeError): + StagingOptions(**bad) + with pytest.raises(dataclasses.FrozenInstanceError): + options.planned_fetches = 2 diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_slots.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_slots.py new file mode 100644 index 000000000000..b13004823378 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_slots.py @@ -0,0 +1,1453 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Staging sizing and the slot queue alone, over plain layouts, the usable end resume_bounds finds and +how it and fetch ends keep out of runs of multimodal tokens, with the runtime's own life cycles +behind a stand-in manager, the delivered rows a shrink hook drops when its own work raises, the +slots a read's grant or a write's marks keep when a copy's completion cannot be recorded, also when +the quarantine itself raises, a write whose grant and failure both raise, a lender stopped by an +error its free or shutdown hook logs, the records of a replaced cache resume_bounds drops, the states +a fetch walks through, and the bytes the test kit stages: no real manager, no memory, no GPU.""" + +import ast +import gc +import inspect +import itertools +import random +import weakref +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + StagingOptions, + _fetch_ledger, + _lender, + _ranges, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._slots import ( + SlotQueue, + SlotRuns, + slot_counts, + slots_per_fetch, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttnLifeCycle + +pytestmark = pytest.mark.cpu_only + +MiB = 1 << 20 +GiB = 1 << 30 + + +def layout(tokens_per_block, groups): + """What the sizing reads of a manager's layout, for layer groups ``(window, sink blocks, pool + group, page bytes)``.""" + return SimpleNamespace( + tokens_per_block=tokens_per_block, + windows=tuple(window for window, _, _, _ in groups), + sink_blocks=tuple(sinks for _, sinks, _, _ in groups), + pool_group_of=tuple(g for _, _, g, _ in groups), + page_bytes={g: page for _, _, g, page in groups}, + ) + + +def one_fetch(lay, tokens_per_fetch): + return sum( + rows * lay.page_bytes[g] for g, rows in slots_per_fetch(lay, tokens_per_fetch).items() + ) + + +# TinyLlama-1.1B, TP=1, 32 tokens per block: 22 layers of K and V, 4 heads of 64 bf16 each. +TINYLLAMA_PAGE = 22 * 2 * 4 * 64 * 2 * 32 +TINYLLAMA = layout(32, [(None, 0, 0, TINYLLAMA_PAGE)]) + +# DeepSeek-V4-Pro as its cache manager lays it out on every rank, 128 tokens per block: a +# 128-token window, compressed full attention, and the compressors' state with an 8-token window. +DSV4_PAGES = {1: 20250624, 0: 572672, 2: 39321600} +DSV4_PRO = layout( + 128, [(128, 0, 1, DSV4_PAGES[1]), (None, 0, 0, DSV4_PAGES[0]), (8, 0, 2, DSV4_PAGES[2])] +) + + +# -- sizing ----------------------------------------------------------------------------------- + + +def test_tinyllama_stages_eight_whole_prompts(): + assert slots_per_fetch(TINYLLAMA, 2048) == {0: 64} + assert one_fetch(TINYLLAMA, 2048) == 44 * MiB + assert slot_counts(TINYLLAMA, StagingOptions(2048, planned_fetches=8)) == {0: 8 * 64} + + +def test_deepseek_v4_windows_take_their_window_and_the_full_group_the_range(): + assert slots_per_fetch(DSV4_PRO, 4096) == {1: 1, 0: 32, 2: 1} + assert 74 * MiB < one_fetch(DSV4_PRO, 4096) < 75 * MiB + assert slot_counts(DSV4_PRO, StagingOptions(4096, planned_fetches=8)) == {1: 8, 0: 256, 2: 8} + + +def test_deepseek_v4_long_ranges_capped_at_one_fetch_keep_one_fetch(): + one = one_fetch(DSV4_PRO, 1 << 20) + assert slots_per_fetch(DSV4_PRO, 1 << 20) == {1: 1, 0: 8192, 2: 1} + assert 4 * GiB < one < 5 * GiB + options = StagingOptions(1 << 20, planned_fetches=8, max_bytes=one) + assert slot_counts(DSV4_PRO, options) == {1: 1, 0: 8192, 2: 1} + # A cap between one and eight fetches: every group keeps its rows of one fetch and more. + cap = 960 * GiB // 8 // 4 + counts = slot_counts(DSV4_PRO, StagingOptions(1 << 20, planned_fetches=8, max_bytes=cap)) + assert counts[0] >= 8192 and counts[1] >= 1 and counts[2] >= 1 + assert sum(counts[g] * DSV4_PAGES[g] for g in counts) <= cap + + +def test_short_ranges_take_no_more_than_their_blocks_even_in_a_window(): + assert slots_per_fetch(DSV4_PRO, 100) == {1: 1, 0: 1, 2: 1} + assert slot_counts(DSV4_PRO, StagingOptions(100)) == {1: 1, 0: 1, 2: 1} + + +def test_a_window_wider_than_the_range_takes_the_range_blocks(): + lay = layout(32, [(1024, 0, 0, 100)]) # a window of 32 blocks over a range of 8 + assert slots_per_fetch(lay, 256) == {0: 8} + assert slot_counts(lay, StagingOptions(256, planned_fetches=2)) == {0: 16} + + +def test_sink_blocks_count_with_the_window(): + lay = layout(32, [(64, 1, 0, 100), (None, 0, 1, 10)]) + assert slots_per_fetch(lay, 320) == {0: 1 + 2, 1: 10} + assert one_fetch(lay, 320) == 3 * 100 + 10 * 10 + assert slot_counts(lay, StagingOptions(320, planned_fetches=3)) == {0: 9, 1: 30} + + +def test_layer_groups_sharing_a_pool_group_add_up(): + lay = layout(16, [(None, 0, 0, 8), (32, 0, 0, 8)]) + assert slots_per_fetch(lay, 160) == {0: 10 + 2} + + +def test_a_partial_block_counts_as_a_whole_one(): + assert slots_per_fetch(TINYLLAMA, 33) == {0: 2} + assert slots_per_fetch(TINYLLAMA, 1) == {0: 1} + + +@pytest.mark.parametrize("tokens_per_fetch", [0, -32]) +def test_tokens_per_fetch_must_be_positive(tokens_per_fetch): + with pytest.raises(ValueError): + slots_per_fetch(TINYLLAMA, tokens_per_fetch) + + +def test_max_bytes_below_one_fetch_is_rejected(): + one = one_fetch(DSV4_PRO, 4096) + with pytest.raises(ValueError, match="one fetch"): + slot_counts(DSV4_PRO, StagingOptions(4096, max_bytes=one - 1)) + assert slot_counts(DSV4_PRO, StagingOptions(4096, max_bytes=one)) == slots_per_fetch( + DSV4_PRO, 4096 + ) + + +def test_max_bytes_above_planned_fetches_changes_nothing(): + options = StagingOptions(4096, planned_fetches=2, max_bytes=100 * one_fetch(DSV4_PRO, 4096)) + assert slot_counts(DSV4_PRO, options) == {1: 2, 0: 64, 2: 2} + + +def random_layout(rng): + """A layout with up to four layer groups over up to three pool groups, and its life cycles.""" + tpb = rng.choice([4, 8, 16, 32]) + pages = {g: rng.randint(1, 1000) for g in range(rng.randint(1, 3))} + groups, life_cycles = [], [] + for _ in range(rng.randint(1, 4)): + window = rng.choice([None, rng.randint(1, 6 * tpb)]) + sink_tokens = 0 if window is None else rng.choice([0, 0, rng.randint(1, 3 * tpb)]) + life_cycle = AttnLifeCycle.make(window, sink_tokens, tpb) + g = rng.choice(sorted(pages)) + groups.append((window, int(life_cycle.num_sink_blocks), g, pages[g])) + life_cycles.append(life_cycle) + return layout(tpb, groups), life_cycles + + +@pytest.mark.parametrize("seed", range(40)) +def test_any_whole_block_range_of_at_most_tokens_per_fetch_fits_one_fetch(seed): + # Rows a range [start, end) needs: its blocks that a history of end still reads, as the + # runtime's own life cycle says; the longest range ending at end needs the most. + rng = random.Random(seed) + lay, life_cycles = random_layout(rng) + tpb = lay.tokens_per_block + tokens_per_fetch = rng.randint(1, 12 * tpb) + bound = slots_per_fetch(lay, tokens_per_fetch) + span = tokens_per_fetch // tpb + for end_block in range(1, 40): + end = end_block * tpb + need = dict.fromkeys(bound, 0) + for lg, life_cycle in enumerate(life_cycles): + stale = life_cycle.get_stale_range(end, tpb) + blocks = range(max(0, end_block - span), end_block) + need[lay.pool_group_of[lg]] += sum(not stale.beg <= b < stale.end for b in blocks) + assert all(need[g] <= bound[g] for g in bound), (end, need, bound) + + +@pytest.mark.parametrize("seed", range(40)) +def test_counts_are_planned_fetches_fetches_or_the_cap_rounded_down(seed): + rng = random.Random(seed) + lay, _ = random_layout(rng) + tokens_per_fetch = rng.randint(1, 12 * lay.tokens_per_block) + rows = slots_per_fetch(lay, tokens_per_fetch) + one = one_fetch(lay, tokens_per_fetch) + planned_fetches = rng.randint(1, 9) + assert slot_counts(lay, StagingOptions(tokens_per_fetch, planned_fetches)) == { + g: planned_fetches * rows[g] for g in rows + } + cap = rng.randint(one, planned_fetches * one) + counts = slot_counts(lay, StagingOptions(tokens_per_fetch, planned_fetches, max_bytes=cap)) + assert sum(counts[g] * lay.page_bytes[g] for g in counts) <= cap + # As many whole fetches as the cap holds hold slots at once. + whole = cap // one + assert all(whole * rows[g] <= counts[g] <= planned_fetches * rows[g] for g in rows) + assert all(counts[g] * one <= cap * rows[g] < (counts[g] + 1) * one for g in rows) + + +def test_a_fractional_fetch_budget_keeps_slots_for_a_shorter_range() -> None: + lay = layout(32, [(None, 0, 0, 8)]) + slots = SlotQueue(slot_counts(lay, StagingOptions(320, planned_fetches=3, max_bytes=200))) + assert slots.num_slots(0) == 25 + for rows in (10, 10, 5): + assert slots.try_grant(slots.enqueue({0: rows})) is not None + assert slots.try_grant(slots.enqueue({0: 1})) is None + + +@pytest.mark.parametrize("max_bytes_in_fetches", [None, 1, 2.5]) +def test_planned_fetches_ranges_hold_slots_at_once_and_one_more_waits(max_bytes_in_fetches): + rows = slots_per_fetch(DSV4_PRO, 4096) + max_bytes = None + at_once = 3 + if max_bytes_in_fetches is not None: + max_bytes = int(max_bytes_in_fetches * one_fetch(DSV4_PRO, 4096)) + at_once = int(max_bytes_in_fetches) + slots = SlotQueue( + slot_counts(DSV4_PRO, StagingOptions(4096, planned_fetches=3, max_bytes=max_bytes)) + ) + for _ in range(at_once): + assert slots.try_grant(slots.enqueue(rows)) is not None + assert slots.try_grant(slots.enqueue(rows)) is None + assert slots.num_waiting == 1 + + +# -- the slot queue --------------------------------------------------------------------------- + + +def take_now(slots, counts): + ticket = slots.enqueue(counts) + runs = slots.try_grant(ticket) + assert runs is not None, f"{counts} did not fit" + return runs + + +def test_slots_of_one_lease_are_contiguous_per_group(): + slots = SlotQueue({0: 8, 1: 8}) + assert (slots.num_slots(0), slots.free_slots(1)) == (8, 8) + runs = take_now(slots, {0: 3, 1: 5}) + assert runs.slots(0).tolist() == [0, 1, 2] + assert runs.slots(1).tolist() == [0, 1, 2, 3, 4] + assert runs.slots(7).tolist() == [] + assert runs.slots(0).dtype.name == "int64" + + +def test_a_lease_is_all_or_nothing_across_groups(): + slots = SlotQueue({0: 8, 1: 4}) + held = take_now(slots, {1: 3}) + ticket = slots.enqueue({0: 2, 1: 2}) + + assert slots.try_grant(ticket) is None + # Group 0 had room, but nothing was taken from it. + assert slots.free_slots(0) == 8 + assert slots.free_slots(1) == 1 + + slots.give_back(held) + assert slots.try_grant(ticket) is not None + assert (slots.free_slots(0), slots.free_slots(1)) == (6, 2) + + +def test_contiguity_waits_for_a_run_even_when_enough_slots_are_free(): + slots = SlotQueue({0: 6}) + a, b, c = (take_now(slots, {0: 2}) for _ in range(3)) + slots.give_back(a) + slots.give_back(c) + assert slots.free_slots(0) == 4 + + ticket = slots.enqueue({0: 3}) + assert slots.try_grant(ticket) is None, "two free runs of 2 are not a run of 3" + + slots.give_back(b) + assert slots.try_grant(ticket).slots(0).tolist() == [0, 1, 2] + + +def test_waiting_leases_are_served_strictly_in_order(): + slots = SlotQueue({0: 4}) + held = take_now(slots, {0: 3}) + first = slots.enqueue({0: 2}) + second = slots.enqueue({0: 1}) + + # One slot is free and ``second`` would fit, but ``first`` asked earlier. + assert slots.try_grant(second) is None + assert slots.try_grant(first) is None + assert slots.num_waiting == 2 + + slots.give_back(held) + assert slots.try_grant(second) is None, "still behind the first" + assert slots.try_grant(first).slots(0).tolist() == [0, 1] + assert slots.try_grant(second).slots(0).tolist() == [2] + assert slots.num_waiting == 0 + + +def test_cancelling_the_head_lets_the_next_lease_through(): + slots = SlotQueue({0: 4}) + held = take_now(slots, {0: 3}) + first = slots.enqueue({0: 4}) + second = slots.enqueue({0: 1}) + assert slots.try_grant(second) is None + + slots.cancel(first) + assert slots.try_grant(second).slots(0).tolist() == [3] + slots.give_back(held) + # Cancelling an unknown or granted ticket does nothing. + slots.cancel(first) + slots.cancel(12345) + assert (slots.free_slots(0), slots.num_waiting) == (3, 0) + + +def test_a_ticket_that_is_not_waiting_is_a_key_error(): + slots = SlotQueue({0: 8}) + with pytest.raises(KeyError): + slots.try_grant(12345) + ticket = slots.enqueue({0: 1}) + assert slots.try_grant(ticket) is not None + with pytest.raises(KeyError): + slots.try_grant(ticket) + + +def test_frees_coalesce_into_one_run(): + slots = SlotQueue({0: 8}) + held = [take_now(slots, {0: 2}) for _ in range(4)] + for runs in (held[1], held[3], held[0], held[2]): + slots.give_back(runs) + assert slots.free_slots(0) == 8 + # Only a single coalesced run can serve the whole group at once. + assert take_now(slots, {0: 8}).slots(0).tolist() == list(range(8)) + + +def test_a_double_free_is_rejected_and_changes_nothing(): + slots = SlotQueue({0: 4, 1: 4}) + runs = take_now(slots, {0: 2, 1: 2}) + slots.give_back(SlotRuns({1: (0, 2)})) + assert (slots.free_slots(0), slots.free_slots(1)) == (2, 4) + + with pytest.raises(ValueError, match="freed twice"): + slots.give_back(runs) + # Group 0's run was not returned on the way to finding group 1's double free. + assert (slots.free_slots(0), slots.free_slots(1)) == (2, 4) + + +@pytest.mark.parametrize("run", [(3, 2), (-1, 1)], ids=["past_the_end", "negative"]) +def test_freeing_slots_outside_the_group_is_rejected(run): + slots = SlotQueue({0: 4}) + take_now(slots, {0: 4}) + with pytest.raises(ValueError, match="outside"): + slots.give_back(SlotRuns({0: run})) + assert slots.free_slots(0) == 0 + + +def test_freeing_slots_of_an_unknown_group_is_rejected(): + slots = SlotQueue({0: 4}) + take_now(slots, {0: 4}) + with pytest.raises(ValueError, match="pool group 5"): + slots.give_back(SlotRuns({0: (0, 4), 5: (0, 1)})) + assert slots.free_slots(0) == 0 + + +@pytest.mark.parametrize( + "counts", [{0: 5}, {3: 1}, {0: -1}], ids=["too_many", "unknown", "negative"] +) +def test_a_request_that_can_never_be_granted_is_rejected_at_once(counts): + slots = SlotQueue({0: 4}) + with pytest.raises(ValueError): + slots.check_fits_capacity(counts) + with pytest.raises(ValueError): + slots.enqueue(counts) + assert slots.num_waiting == 0 + assert slots.free_slots(0) == 4 + + +def test_check_queues_nothing(): + slots = SlotQueue({0: 4}) + slots.check_fits_capacity({0: 4}) + assert slots.num_waiting == 0 + # A later ask is first in line. + assert take_now(slots, {0: 4}).slots(0).tolist() == [0, 1, 2, 3] + + +def test_zero_counts_need_no_slots(): + slots = SlotQueue({0: 1}) + take_now(slots, {0: 1}) + assert take_now(slots, {0: 0, 9: 0}).by_pool_group == {} + + +class _Owner: + """What a stand-in lender's weak references point at.""" + + +_OWNER = _Owner() + + +def staging_over(windows, sinks, tpb): + """A staging lender and a manager stand-in answering what resume_bounds asks: the tokens per block + and each layer group's stale range, from the runtime's own life cycle.""" + cycles = [AttnLifeCycle.make(w, s * tpb, tpb) for w, s in zip(windows, sinks)] + + def stale_block_range(lg, history): + stale = cycles[lg].get_stale_range(history, tpb) + return stale.beg, stale.end + + manager = SimpleNamespace(tokens_per_block=tpb, _stale_block_range=stale_block_range) + lay = SimpleNamespace( + tokens_per_block=tpb, windows=tuple(windows), pool_groups=(), num_layer_groups=len(windows) + ) + ref = weakref.ref(_OWNER) + return _lender.Staging(ref, lay, None, (), None, ref), manager + + +def reads_all(manager, windows, blocks, first, end_block, tpb): + """Every layer group has, among the blocks from ``first`` and the delivered ones, every block a + start at ``end_block * tpb`` reads: all of them for full attention, else all but the stale.""" + for lg, delivered in enumerate(blocks): + beg, end = (0, 0) + if windows[lg] is not None: + beg, end = manager._stale_block_range(lg, end_block * tpb) + for ordinal in range(first, end_block): + if not beg <= ordinal < end and not (ordinal < len(delivered) and delivered[ordinal]): + return False + return True + + +def test_max_resume_position_is_the_largest_start_every_layer_group_reads(): + """Random layouts, delivered rows and computed prefixes: the usable end is the largest + whole-block start past the computed prefix where every layer group has what it reads, whatever + the lowest fetch start, else the computed prefix.""" + rng = random.Random(0) + for _ in range(3000): + tpb = rng.choice([16, 32]) + windows = [None if rng.random() < 0.25 else rng.randint(1, 200) for _ in range(3)] + windows = windows[: rng.randint(1, 3)] + sinks = [0 if w is None else rng.choice([0, 0, 1, 2]) for w in windows] + known = rng.choice([0, rng.randint(0, 12) * tpb, rng.randint(0, 12 * tpb)]) + origin = rng.choice([known // tpb * tpb, rng.randint(0, 14) * tpb]) + length = rng.randint(0, 20) + blocks = [] + for _ in windows: + row = np.zeros(max(0, length + rng.choice([0, -1, 1])), dtype=bool) + row[rng.choice([0, origin // tpb, rng.randint(0, 20)]) :] = True + for _ in range(rng.choice([0, 0, 1, 2])): + if len(row): + row[rng.randrange(len(row))] = False + row[: origin // tpb] = False # fetches start at origin or later + blocks.append(row) + lender, manager = staging_over(windows, sinks, tpb) + first = known // tpb + last = max((len(b) for b in blocks), default=0) + ends = range(first + 1, last + 1) + starts = [e for e in ends if reads_all(manager, windows, blocks, first, e, tpb)] + want = starts[-1] * tpb if starts else known + delivered = _fetch_ledger._Deliveries(None, origin, [b.copy() for b in blocks]) + got = lender._ledger._max_resume_position(manager, delivered, known) + assert got == want, ( + f"windows {windows}, sinks {sinks}, tpb {tpb}, origin {origin}, known {known}, " + f"delivered {[np.nonzero(b)[0].tolist() for b in blocks]}: usable until {got}, " + f"not {want}" + ) + + +def test_max_resume_position_after_a_whole_fetch_checks_its_end_alone(monkeypatch): + """Two windowed layer groups, nothing computed and 4096 blocks delivered in each: the end of the + delivered rows is usable, found with one stale-range lookup per layer group, not one per block + below it, so resume_bounds after each lease of a split fetch costs no more as the fetch grows.""" + lender, manager = staging_over([4096, 8192], [0, 0], 32) + blocks = [np.ones(4096, dtype=bool), np.ones(4096, dtype=bool)] + real_stale, calls = _ranges._stale_range, [] + + def counted(*args): + calls.append(args) + return real_stale(*args) + + monkeypatch.setattr(_ranges, "_stale_range", counted) + got = lender._ledger._max_resume_position( + manager, _fetch_ledger._Deliveries(None, 0, blocks), 0 + ) + assert got == 4096 * 32, f"usable until {got}" + assert len(calls) == 2, f"{len(calls)} stale-range lookups for one usable end" + + +def test_max_resume_position_checks_no_start_past_a_missing_full_attention_block(monkeypatch): + """One full-attention layer group, nothing computed and 8192 blocks delivered but one: no start + past the missing block has every block it reads, so none is checked, and resume_bounds after each + later lease of a split fetch that lost an early block checks no more starts as it grows.""" + lender, manager = staging_over([None], [0], 32) + real_stale, calls = _ranges._stale_range, [] + + def counted(*args): + calls.append(args) + return real_stale(*args) + + monkeypatch.setattr(_ranges, "_stale_range", counted) + for hole, usable, checked in ((0, 0, 0), (4000, 4000 * 32, 1)): + blocks = np.ones(8192, dtype=bool) + blocks[hole] = False + calls.clear() + got = lender._ledger._max_resume_position( + manager, _fetch_ledger._Deliveries(None, 0, [blocks]), 0 + ) + assert got == usable, f"usable until {got} with block {hole} missing" + assert len(calls) == checked, f"{len(calls)} starts checked with block {hole} missing" + + +def test_a_shrink_hook_that_raises_forgets_what_the_request_was_delivered(monkeypatch): + """A shrink hook whose own work raises, which it logs: the request's delivered rows all go, + since the cache may grow back before resume_bounds sees the freed ones past its blocks.""" + lender, _ = staging_over([None], [0], 32) + kv = object() + lender._ledger._deliveries[7] = _fetch_ledger._Deliveries(kv, 0, [np.ones(3, dtype=bool)]) + lender._ledger._mixed_tail_caches[7] = kv + raised = [] + + def num_blocks(cache): + raised.append(cache) + raise MemoryError("planted") + + monkeypatch.setattr(_lender._manager, "num_blocks", num_blocks) + lender._on_shrink(7, kv) + assert raised == [kv], "the hook's work did not raise: the check proves nothing" + assert 7 not in lender._ledger._deliveries, "the rows the shrink freed stay delivered" + assert lender._ledger._mixed_tail_caches.get(7) is kv, "the hook dropped the request's taint" + + +# -- runs of multimodal tokens ---------------------------------------------------------------- + + +def multimodal(runs, length, flagged=True): + """A request stand-in whose multimodal data marks ``runs`` as tokens that attend both ways.""" + mask = np.zeros(length, dtype=np.int64) + for beg, end in runs: + mask[beg:end] = 1 + data = {"mm_bidirectional_blocks": flagged, "multimodal_embed_mask_cumsum": np.cumsum(mask)} + return SimpleNamespace(py_multimodal_data=data) + + +def strictly_inside(runs, position): + return any(beg < position < end for beg, end in runs) + + +def test_resume_bounds_and_fetch_ends_keep_out_of_bidirectional_runs(): + """Random bidirectional runs of multimodal tokens: a fetch end strictly inside a run is refused, + others not; resume_bounds lowers its end to the largest position at or below it not strictly inside + one, with a floor at the end of the last run at or below that, so none from floor to end is.""" + rng = random.Random(0) + # With a window too, which the refusal must not read: the model's window can be narrower. + lender, windowed = (staging_over(w, [0] * len(w), 32)[0] for w in ([None], [64, None])) + for _ in range(3000): + length = rng.randint(1, 400) + runs, pos = [], rng.randint(0, 20) + while pos < length - 1 and rng.random() < 0.7: + stop = min(length, pos + rng.randint(1, 120)) + runs.append((pos, stop)) + pos = stop + rng.randint(1, 40) + request = multimodal(runs, length) + usable = rng.randint(0, length) + got, floor = lender._ledger._clip_to_bidirectional_runs(request, usable) + want = max(p for p in range(usable + 1) if not strictly_inside(runs, p)) + case = f"runs {runs}, usable {usable}" + assert got == want, f"{case}: the end lowered to {got}, not {want}" + assert floor == max((stop for _, stop in runs if stop <= want), default=0), ( + f"{case}: a floor of {floor}" + ) + assert not any(strictly_inside(runs, p) for p in range(floor, got + 1)), ( + f"{case}: [{floor}, {got}] holds a position strictly inside a run" + ) + end = rng.randint(0, length) + for each in (lender, windowed): + refused = each._bidirectional_split_refusal(request, end) is not None + assert refused == strictly_inside(runs, end), ( + f"{case}: an end at {end} refused: {refused}" + ) + + +def test_a_request_without_bidirectional_runs_keeps_its_interval_and_ends(): + """No flag, no multimodal data, or no mask: nothing lowers, rises or fails.""" + lender, _ = staging_over([64, None], [0, 0], 32) + for request in ( + multimodal([(1, 300)], 400, flagged=False), + SimpleNamespace(py_multimodal_data=None), + SimpleNamespace(py_multimodal_data={"mm_bidirectional_blocks": True}), + ): + assert lender._ledger._clip_to_bidirectional_runs(request, 200) == (200, 0) + assert lender._bidirectional_split_refusal(request, 200) is None + + +# -- copies whose completion cannot be recorded ---------------------------------------------- + + +def raising_once(call, error, fired): + """``call``, except that its first call raises ``error`` and notes it in ``fired``.""" + + def once(*args, **kwargs): + if not fired: + fired.append(True) + raise error + return call(*args, **kwargs) + + return once + + +class _AppendRaisingOnce(list): + """A record list whose first append raises ``MemoryError`` and notes it in ``fired``.""" + + def __init__(self, fired): + super().__init__() + self.fired = fired + + def append(self, item): + if not self.fired: + self.fired.append(True) + raise MemoryError("planted") + super().append(item) + + +def _stand_in_view(rows, keys): + return "view" + + +@pytest.mark.parametrize("fault", ["none", "before", "record", "after"]) +def test_a_read_grant_keeps_the_slots_its_queued_copy_may_write_out_of_use(monkeypatch, fault): + """A read granted 2 of 4 slots queues into them a copy whose completion cannot be recorded: the + read fails and its slots stay out of use, also with a second fault in the quarantine's record or + the read's failure after it. A fault in its view, before it takes slots, returns them.""" + lay = SimpleNamespace( + pool_groups=(0,), + pool_group_of=(0,), + windows=(None,), + num_layer_groups=1, + tokens_per_block=1, + ) + ref = weakref.ref(_OWNER) + lender = _lender.Staging(ref, lay, None, (), SlotQueue({0: 4}), ref) + fired = [] + view = _stand_in_view + if fault == "before": + view = raising_once(view, MemoryError("planted"), fired) + monkeypatch.setattr(lender, "_build_view", view) + monkeypatch.setattr(lender, "_copy_segments", lambda rows: []) + monkeypatch.setattr( + lender, "_queue_copy", lambda segments, to_staging: (None, "planted: no event") + ) + if fault == "record": + lender._quarantined = _AppendRaisingOnce(fired) + elif fault == "after": + failing = raising_once(lender._fail, MemoryError("planted"), fired) + monkeypatch.setattr(lender, "_fail", failing) + rows = _lender._LentBlockEntries([0], [np.arange(2)], [np.arange(2)]) + read = _lender._StagingLease(lender, "read", 7, None, rows, [np.zeros((2, 32), np.uint8)]) + lender._admit(read, {0: 2}) + assert bool(fired) == (fault != "none"), "the planted fault did not strike: proves nothing" + assert read.failure is not None and "planted" in read.failure, read.failure + if fault == "before": + assert lender._free_slots(0) == 4, "slots no copy reached stayed out of use" + assert not lender._quarantined, "slots no copy reached were recorded as lost" + return + assert lender._free_slots(0) == 2, "slots a queued copy may still write came back" + assert any(r is read._slot_runs for r in lender._quarantined), "the lost slots are not recorded" + + +class _IterRaisingOnce(list): + """A holding list whose first iteration raises ``MemoryError`` and notes it in ``fired``.""" + + def __init__(self, fired): + super().__init__() + self.fired = fired + + def __iter__(self): + if not self.fired: + self.fired.append(True) + raise MemoryError("planted") + return super().__iter__() + + +def _unrecorded_copy(segments, to_staging): + raise MemoryError("planted: the copy is queued, its record is not") + + +def _one_pool_group(): + return SimpleNamespace( + pool_groups=(0,), + pool_group_of=(0,), + windows=(None,), + num_layer_groups=1, + tokens_per_block=1, + ) + + +def _write_target(): + """A stand-in manager whose request 7 holds blocks 0 and 1 on pages 0 and 1.""" + kv = SimpleNamespace( + num_committed_tokens=0, + history_length=0, + is_active=True, + num_blocks=2, + get_base_page_indices=lambda lg: [0, 1], + ) + manager = _Owner() + manager.kv_cache_map = {7: kv} + return kv, manager + + +@pytest.mark.parametrize("fault", ["none", "call", "record", "holding", "log", "append", "partial"]) +def test_a_read_whose_grant_and_quarantine_both_raise_is_never_ready(monkeypatch, fault): + """A read granted 2 of 4 slots: its grant raises once its view is set (recording its copy, + entering the holding list or failing a partly failed copy), then its quarantine raises too (at + its call, the slots' record, the holding list or its log line). As after a manager hook logged + the error, it fails at its next poll, returning no view, its slots stay out of use after + release, and where a copy may be queued the shutdown would keep the staging memory. With no + fault it is ready once its copy landed.""" + ref = weakref.ref(_OWNER) + lender = _lender.Staging(ref, _one_pool_group(), None, (), SlotQueue({0: 4}), ref) + planted = [[], []] # the grant's fault, then the quarantine's + landed = SimpleNamespace(query=lambda: True) + monkeypatch.setattr(lender, "_build_view", _stand_in_view) + monkeypatch.setattr(lender, "_copy_segments", lambda rows: []) + if fault == "none": + copied = (_lender._QueuedCopy(landed), None) + monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging: copied) + elif fault == "partial": + copied = (_lender._QueuedCopy(landed), "planted: a copy call failed") + monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging: copied) + failing = raising_once(lender._fail, MemoryError("planted"), planted[0]) + monkeypatch.setattr(lender, "_fail", failing) + else: + monkeypatch.setattr(lender, "_queue_copy", _unrecorded_copy) + planted[0].append(True) + if fault == "append": + planted[0].clear() + lender._leases_holding_slots = _AppendRaisingOnce(planted[0]) + if fault in ("call", "append", "partial"): + failing = raising_once(lender._quarantine, MemoryError("planted"), planted[1]) + monkeypatch.setattr(lender, "_quarantine", failing) + elif fault == "record": + lender._quarantined = _AppendRaisingOnce(planted[1]) + elif fault == "holding": + lender._leases_holding_slots = _IterRaisingOnce(planted[1]) + elif fault == "log": + failing = raising_once(_lender.logger.warning, MemoryError("planted"), planted[1]) + monkeypatch.setattr(_lender.logger, "warning", failing) + rows = _lender._LentBlockEntries([0], [np.arange(2)], [np.arange(2)]) + read = _lender._StagingLease(lender, "read", 7, None, rows, [np.zeros((2, 32), np.uint8)]) + if fault == "none": + lender._admit(read, {0: 2}) + else: + with pytest.raises(MemoryError, match="planted"): # what a manager hook logs and goes on + lender._admit(read, {0: 2}) + assert all(planted), "a planted fault did not strike: proves nothing" + view = read.poll() + read.release() + if fault == "none": + assert view == "view" and read.failure is None, "a read whose copy landed was not ready" + assert lender._free_slots(0) == 4, "the slots of a read whose copy landed stayed out of use" + assert not lender._must_retain_memory(), "nothing open, yet the staging memory would stay" + return + assert view is None, "a read with no record of its copy was ready" + assert read.failure is not None, "a read with no record of its copy did not fail" + assert lender._free_slots(0) == 2, "slots a queued copy may still write came back" + if fault != "append": # past the append, a copy may be queued + assert lender._must_retain_memory(), ( + "the shutdown would free staging a queued copy may write" + ) + + +@pytest.mark.parametrize("fault", ["call", "record"]) +@pytest.mark.parametrize("copy", ["raised", "unrecorded"]) +def test_marks_whose_copy_and_quarantine_both_fail_keep_the_slots_out_of_use( + monkeypatch, copy, fault +): + """A write granted 2 of 4 slots is marked; its marks' copy raises once queued or returns with no + record of its completion, then its quarantine raises too, at its call or the slots' record: + after the release the slots stay out of use and the shutdown would keep the staging memory.""" + kv, manager = _write_target() + ref = weakref.ref(manager) + lender = _lender.Staging(ref, _one_pool_group(), None, (), SlotQueue({0: 4}), ref) + fired = [] + monkeypatch.setattr( + lender, "_build_view", lambda rows, keys: SimpleNamespace(groups=rows.block_ordinals) + ) + monkeypatch.setattr(lender, "_copy_segments", lambda rows, mask: []) + if copy == "raised": + monkeypatch.setattr(lender, "_queue_copy", _unrecorded_copy) + else: + monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging: (None, "no event")) + if fault == "call": + failing = raising_once(lender._quarantine, MemoryError("planted"), fired) + monkeypatch.setattr(lender, "_quarantine", failing) + else: + lender._quarantined = _AppendRaisingOnce(fired) + rows = _lender._LentBlockEntries([0], [np.arange(2)], [np.arange(2)]) + write = _lender._StagingLease(lender, "write", 7, kv, rows) + lender._admit(write, {0: 2}) + view = write.poll() + with pytest.raises(MemoryError, match="planted"): + write.mark_arrived([np.ones(len(run), dtype=bool) for run in view.groups]) + assert fired, "the quarantine did not raise: proves nothing" + write.release() + assert lender._free_slots(0) == 2, "slots a queued copy may still read came back" + assert lender._must_retain_memory(), ( + "the shutdown would free staging a queued copy may still read" + ) + + +class _HoldingAppendRaisingOnce(list): + """A holding list, with the leases of ``held``, whose first append raises ``MemoryError`` and + notes it in ``fired``.""" + + def __init__(self, held, fired): + super().__init__(held) + self.fired = fired + + def append(self, item): + if not self.fired: + self.fired.append(True) + raise MemoryError("planted") + super().append(item) + + +def test_a_write_whose_grant_and_failure_both_raise_is_never_ready(monkeypatch): + """A write waits for the slots a read holds, and the read's release grants it: entering the + holding list raises, then failing the write raises too. As after a manager hook logged that + error, the write never returns a view over the slots it gave back for the next grant.""" + kv, manager = _write_target() + ref = weakref.ref(manager) + lender = _lender.Staging(ref, _one_pool_group(), None, (), SlotQueue({0: 2}), ref) + landed = SimpleNamespace(query=lambda: True) + monkeypatch.setattr( + lender, "_build_view", lambda rows, keys: SimpleNamespace(groups=rows.block_ordinals) + ) + monkeypatch.setattr(lender, "_copy_segments", lambda rows, mask=None: []) + copied = (_lender._QueuedCopy(landed), None) + monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging: copied) + read = _lender._StagingLease( + lender, + "read", + 8, + None, + _lender._LentBlockEntries([0], [np.arange(2)], [np.arange(2)]), + None, + ) + lender._admit(read, {0: 2}) + assert read.poll() is not None, "the read was not ready: proves nothing" + write = _lender._StagingLease( + lender, "write", 7, kv, _lender._LentBlockEntries([0], [np.arange(2)], [np.arange(2)]) + ) + lender._admit(write, {0: 2}) + assert write.poll() is None and write.failure is None, "the write did not wait" + appended, failed = [], [] + recycle = lender._recycle + + def recycling(): # the holding list, once rebuilt, raises at the write's grant + recycle() + lender._leases_holding_slots = _HoldingAppendRaisingOnce( + lender._leases_holding_slots, appended + ) + + monkeypatch.setattr(lender, "_recycle", recycling) + monkeypatch.setattr(lender, "_fail", raising_once(lender._fail, MemoryError("planted"), failed)) + with pytest.raises(MemoryError, match="planted"): # what a manager hook logs and goes on + read.release() + assert appended and failed, "a planted fault did not strike: proves nothing" + assert write.poll() is None, "a write was ready over slots it gave back" + + +# -- a lender stopped by an error its free or shutdown hook logs ------------------------------ + + +class _Memory: + """The staging memory a stand-in lender finds at its shutdown; records each free.""" + + def __init__(self): + self.freed = [] + + def free(self): + self.freed.append(True) + + +def _cache(blocks=2, committed=0): + """A stand-in cache whose blocks ``0..blocks-1`` lock pages ``0..blocks-1`` in every layer + group, with ``committed`` committed tokens, which its history also has.""" + pages = list(range(blocks)) + return SimpleNamespace( + num_committed_tokens=committed, + history_length=committed, + is_active=True, + num_blocks=blocks, + get_base_page_indices=lambda lg: pages, + get_aggregated_page_indices=lambda lg, valid_only=False: pages, + ) + + +def _rows_of(count): + """Blocks ``0..count-1`` on pages ``0..count-1``, in each of two layer groups.""" + ordinals = [np.arange(count), np.arange(count)] + return _lender._LentBlockEntries([0, 1], ordinals, [o.copy() for o in ordinals]) + + +def _request(request_id, prompt_len=4): + return SimpleNamespace( + py_request_id=request_id, + prompt_len=prompt_len, + py_multimodal_data=None, + try_get_encoder_output_len=lambda: None, + ) + + +def _a_write_waits(monkeypatch, fault, error=MemoryError): + """Two layer groups on pool groups 0 and 1 of four slots each: a read holds slots 0 and 1 of + both with its copy pending, a write into request 8 whose view a poll returned holds slot 2, a + write whose view no poll returned holds slot 3, and a write into request 7 waits for 2 + 2; + request 10's cache has 8 committed blocks. The read is released and its copy lands. ``error`` + is planted in pool group 1, changed second, for its next take or return of slots as ``fault`` + says: the take of the waiting write's slots or the return of the read's.""" + manager = _Owner() + manager.kv_cache_map = {7: _cache(), 8: _cache(), 10: _cache(blocks=8, committed=8)} + layout = SimpleNamespace( + pool_groups=(0, 1), + pool_group_of=(0, 1), + windows=(None, None), + num_layer_groups=2, + tokens_per_block=1, + ) + memory = _Memory() + lender = _lender.Staging( + weakref.ref(manager), layout, None, (), SlotQueue({0: 4, 1: 4}), weakref.ref(memory) + ) + s = SimpleNamespace( + lender=lender, manager=manager, memory=memory, copies=[], errors=[], fired=[] + ) + landed = [] + event = SimpleNamespace(query=lambda: bool(landed), synchronize=lambda: None) + + def memcpy(segments, to_staging): + s.copies.append(to_staging) + return _lender._QueuedCopy(event), None + + monkeypatch.setattr( + lender, "_build_view", lambda rows, keys: SimpleNamespace(groups=rows.block_ordinals) + ) + monkeypatch.setattr(lender, "_copy_segments", lambda rows, mask=None: []) + monkeypatch.setattr(lender, "_queue_copy", memcpy) + monkeypatch.setattr(_lender.logger, "error", s.errors.append) + + def lent(kind, request_id, count): + kv = manager.kv_cache_map.get(request_id) + fetch = None + if kind == "write" and kv is not None: # recorded as lend_write leaves it + fetch = lender._ledger._latest_fetches[request_id] = _fetch_ledger._FetchRecord(0, kv) + lease = _lender._StagingLease(lender, kind, request_id, kv, _rows_of(count), None, fetch) + lender._admit(lease, {0: count, 1: count}) + return lease + + s.read = lent("read", 11, 2) + s.seen = lent("write", 8, 1) + s.seen_view = s.seen.poll() + s.granted = lent("write", 12, 1) + s.waiting = lent("write", 7, 2) + assert s.seen_view is not None and s.granted._view is not None + assert s.waiting.poll() is None and s.waiting.failure is None, "the write did not wait" + s.read.release() # its copy pends, so its slots stay out of use + landed.append(True) + if fault != "none": + alloc = lender._slot_queue._allocators[1] + planted = raising_once(getattr(alloc, fault), error("planted"), s.fired) + monkeypatch.setattr(alloc, fault, planted) + return s + + +def _freed_while_a_write_waits(monkeypatch, fault, error=MemoryError): + """``_a_write_waits``, then request 9's free runs the manager's hook.""" + s = _a_write_waits(monkeypatch, fault, error) + s.lender._on_free(9, object()) # the manager's hook: it raises nothing into the free + assert bool(s.fired) == (fault != "none"), "the planted fault did not strike: proves nothing" + return s + + +@pytest.mark.parametrize( + "fault, error", + [ + ("none", None), + ("take", MemoryError), + ("give", MemoryError), + ("take", KeyError), + ("give", KeyError), + ], + ids=["none", "take", "give", "take_key_error", "give_key_error"], +) +def test_an_error_the_free_hook_logs_stops_the_lender(monkeypatch, fault, error): + """Another request's free returns the slots of a released read whose copy landed and grants + them to the write waiting for them, and a planted MemoryError or KeyError strikes between the + changes of the two pool groups: at the take of the write's slots or at the return of the read's. + The hook logs it, raises nothing into the free, and stops the lender for good: the next lender + call fails the waiting write, a poll fails a lease no poll returned the view of and returns a + view one did, lends fail at the call after their range checks, resume_bounds is an empty interval, + a mark copies nothing, no slot returns or is granted, and the shutdown keeps the staging memory + with nothing open. ResumeBounds of a request without a cache still raises. A range needing more + slots than a part has still raises in ``lend_write``, which checks the stop after its slot + check, and fails at the call in ``lend_read``, which checks the stop before any check of the + cache, each where it checks a reuse reset. Without the fault the free grants the waiting write, + and that range raises in both lends.""" + s = _freed_while_a_write_waits(monkeypatch, fault, error) + lender, waiting = s.lender, s.waiting + big = _request(10, prompt_len=9) # [0, 8) needs 8 slots of each pool group, which has 4 + if fault == "none": + assert not s.errors, s.errors + assert waiting.poll() is not None and waiting.failure is None, "the free granted nothing" + assert s.granted.poll() is not None and s.seen.poll() is s.seen_view + for lend in (lender.lend_read, lender.lend_write): + with pytest.raises(ValueError, match="more staging slots"): + lend(big, 0, 8) + for lease in (s.seen, s.granted, waiting): + lease.mark_arrived([np.full(len(run), lease is s.seen) for run in lease._view.groups]) + lease.release() + assert s.copies == [True, False], "no copy of the marked rows into the request's pages" + lender._on_shutdown() + assert s.memory.freed == [True], "the shutdown kept staging that nothing holds" + return + assert len(s.errors) == 1 and "planted" in s.errors[0], s.errors + free = (0, 2) if fault == "take" else (2, 0) # pool group 0 changed, pool group 1 not + assert (lender._free_slots(0), lender._free_slots(1)) == free, "no fault between the groups" + requests = [_request(7), _request(8)] + intervals = [lender.resume_bounds(request) for request in requests] # progress comes first + assert waiting.failure is not None, "the waiting write neither became ready nor failed" + stop = waiting.failure + assert stop == _lender._REASON_STOPPED, stop + assert waiting.poll() is None + for request, interval in zip(requests, intervals): + assert interval is not None and interval.max_position < interval.min_position, ( + f"request {request.py_request_id}: resume_bounds {interval}, not an empty interval" + ) + no_cache = _request(9) + with pytest.raises(ValueError, match="has no KV cache"): + lender.resume_bounds(no_cache) + assert s.granted.poll() is None and s.granted.failure == stop, "a view went out after the stop" + assert s.seen.poll() is s.seen_view, "a view a poll returned before the stop was taken back" + for lend in (lender.lend_read, lender.lend_write): + lease = lend(no_cache, 0, 2) + assert lease.failure == stop, f"{lend.__name__} did not fail at the call: {lease.failure}" + lease.release() + with pytest.raises(ValueError): + lender.lend_read(no_cache, 2, 0) + with pytest.raises(ValueError): + lender.lend_write(no_cache, 0, 4) # past the whole blocks before the last prompt token + with pytest.raises(ValueError, match="more staging slots"): + lender.lend_write(big, 0, 8) + read = lender.lend_read(big, 0, 8) + assert read.failure == stop, f"lend_read of the range did not fail at the call: {read.failure}" + read.release() + marks = [np.ones(len(run), dtype=bool) for run in s.seen_view.groups] + s.seen.mark_arrived(marks) + assert s.copies == [True], "a mark queued a copy into the request's pages" + with pytest.raises(RuntimeError, match="twice"): + s.seen.mark_arrived(marks) + for lease in (s.seen, s.granted, waiting): + lease.release() + assert lender._open_lease_count() == 0 + assert (lender._free_slots(0), lender._free_slots(1)) == free, "slots returned or granted" + assert lender._must_retain_memory(), "the shutdown would free the staging memory" + warnings = [] + monkeypatch.setattr(_lender.logger, "warning", warnings.append) + lender._on_shutdown() + assert not s.memory.freed, "the shutdown freed staging whose slot records may be wrong" + assert any(stop in warning for warning in warnings), f"the stop is not logged: {warnings}" + assert (lender._free_slots(0), lender._free_slots(1)) == free, "the shutdown returned slots" + assert len(s.errors) == 1, f"the shutdown logged an error: {s.errors[1:]}" + + +@pytest.mark.parametrize("fault", ["take", "give"]) +def test_a_shutdown_right_after_the_free_hook_stopped_the_lender(monkeypatch, fault): + """As above, with the manager's shutdown the first call after the free: the shutdown itself + fails the waiting write, before any poll, a later poll fails a lease no poll returned the view + of and returns a view one did, the staging memory stays, and the shutdown logs no error of its + own.""" + s = _freed_while_a_write_waits(monkeypatch, fault) + s.lender._on_shutdown() + assert s.waiting.failure == _lender._REASON_SHUT_DOWN, ( + "the shutdown left the waiting write in line" + ) + assert s.waiting.poll() is None + assert s.granted.poll() is None and s.granted.failure == _lender._REASON_STOPPED, ( + "a view went out" + ) + assert s.seen.poll() is s.seen_view, "a view a poll returned before the stop was taken back" + assert not s.memory.freed, "the shutdown freed staging whose slot records may be wrong" + assert len(s.errors) == 1, f"the shutdown logged an error: {s.errors[1:]}" + + +@pytest.mark.parametrize("before", ["live", "stopped"]) +def test_an_error_the_shutdown_logs_stops_the_lender(monkeypatch, before): + """The manager's shutdown meets an error in its own work and logs it, raising nothing into the + shutdown: on a live lender a planted MemoryError between the returns of the two pool groups of + a released read whose copy landed; on a lender the free hook stopped, a planted MemoryError in + the warning that it keeps the staging memory. The lender stops, its first reason staying: the + waiting write fails by its next poll, a poll fails a lease no poll returned the view of rather + than handing the view out, a view a poll returned keeps returning, and the staging memory + stays.""" + if before == "live": + s = _a_write_waits(monkeypatch, "give") + struck = s.fired + else: + s = _freed_while_a_write_waits(monkeypatch, "take") + struck = [] + warned = raising_once(_lender.logger.warning, MemoryError("planted"), struck) + monkeypatch.setattr(_lender.logger, "warning", warned) + s.lender._on_shutdown() # the manager's hook: it raises nothing into the shutdown + assert struck, "the planted fault did not strike: proves nothing" + assert len(s.errors) == 1 + (before == "stopped") and "planted" in s.errors[-1], s.errors + assert s.waiting.poll() is None and s.waiting.failure is not None, "the waiting write pends" + assert s.granted.poll() is None and s.granted.failure is not None, "a view went out" + assert s.seen.poll() is s.seen_view, "a view a poll returned was taken back" + assert not s.memory.freed, "the shutdown freed staging whose slot records may be wrong" + with pytest.raises(ValueError, match="shut down"): + s.lender.resume_bounds(_request(7)) + first = _lender._REASON_STOPPED_AT_SHUTDOWN if before == "live" else _lender._REASON_STOPPED + assert s.granted.failure == first, f"not the first reason: {s.granted.failure}" + + +def test_a_clean_shutdown_does_not_stop_the_lender(monkeypatch): + """Without a fault the shutdown fails the waiting write with its own reason and keeps the + staging memory for the open leases, and a write no poll returned the view of returns its view + at its first poll after it: only an error the shutdown logs stops the lender.""" + s = _a_write_waits(monkeypatch, "none") + warnings = [] + monkeypatch.setattr(_lender.logger, "warning", warnings.append) + s.lender._on_shutdown() + assert not s.errors, s.errors + assert s.waiting.failure == _lender._REASON_SHUT_DOWN and s.waiting.poll() is None + assert s.granted.poll() is not None and s.granted.failure is None, "a clean shutdown stopped it" + assert s.seen.poll() is s.seen_view + assert warnings and not s.memory.freed, "leases are open: the memory stays" + + +@pytest.mark.parametrize("hook", ["free", "shutdown"]) +def test_a_hook_whose_error_log_raises_still_stops_the_lender(monkeypatch, hook): + """A planted MemoryError strikes in a hook's own work, at the take in the free or at the return + in the shutdown as above, and the hook's error log line raises a KeyError, which goes out into + the manager's free or shutdown: the lender has stopped all the same, so the waiting write and a + lease no poll returned the view of fail at their next poll, and a view a poll returned keeps + returning.""" + s = _a_write_waits(monkeypatch, "take" if hook == "free" else "give") + logged = [] + log = raising_once(s.errors.append, KeyError("log line"), logged) + monkeypatch.setattr(_lender.logger, "error", log) + with pytest.raises(KeyError, match="log line"): + if hook == "free": + s.lender._on_free(9, object()) + else: + s.lender._on_shutdown() + assert s.fired and logged, "a planted fault did not strike: proves nothing" + assert s.waiting.poll() is None and s.waiting.failure is not None, "the waiting write pends" + assert s.granted.poll() is None and s.granted.failure is not None, "a view went out" + assert s.seen.poll() is s.seen_view, "a view a poll returned was taken back" + stop = _lender._REASON_STOPPED if hook == "free" else _lender._REASON_STOPPED_AT_SHUTDOWN + assert s.waiting.failure == s.granted.failure == stop, (s.waiting.failure, s.granted.failure) + + +# -- records of a cache the request no longer holds ------------------------------------------- + +# ResumeBounds of request 7 in ``_records_of``: without records, and with one record of each kind +# describing the request's current cache. +NO_RECORDS = (4, 4) +CURRENT_RECORDS = {"fetch": None, "delivered": (6, 4), "advanced": (2, 4), "tainted": (0, 4)} + + +def _records_of(monkeypatch, kind, owner): + """A lender over a stand-in manager whose request 7 holds a cache of 8 one-token blocks in two + full-attention layer groups, with 2 committed tokens, a history of 4 and the floor at the + history. One record of ``kind`` describes that cache, or with ``owner`` set to ``"replaced"`` + the cache the request held before: an undelivered fetch, blocks 2 to 5 delivered, a fetch from + 2 that kept the 2 tokens below it at a history of 4, or a taint.""" + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: True) + manager = _Owner() + kv, old = _cache(blocks=8, committed=2), _cache(blocks=8, committed=2) + kv.history_length = old.history_length = 4 + manager.kv_cache_map = {7: kv} + layout = SimpleNamespace( + pool_groups=(0, 1), + pool_group_of=(0, 1), + windows=(None, None), + num_layer_groups=2, + tokens_per_block=1, + ) + ref = weakref.ref(manager) + lender = _lender.Staging(ref, layout, None, (), SlotQueue({0: 4, 1: 4}), ref) + cache = kv if owner == "current" else old + rows = np.zeros(6, dtype=bool) + rows[2:] = True + records, record = { + "fetch": (lender._ledger._latest_fetches, _fetch_ledger._FetchRecord(2, cache)), + "delivered": ( + lender._ledger._deliveries, + _fetch_ledger._Deliveries(cache, 2, [rows, rows.copy()]), + ), + "advanced": (lender._ledger._kept_prefixes, (cache, 2, 4)), + "tainted": (lender._ledger._mixed_tail_caches, cache), + }[kind] + records[7] = record + return SimpleNamespace(lender=lender, manager=manager, record=record, records=records) + + +@pytest.mark.parametrize("owner", ["replaced", "current"]) +@pytest.mark.parametrize("kind", list(CURRENT_RECORDS)) +def test_resume_bounds_drops_the_records_of_a_cache_the_request_no_longer_holds( + monkeypatch, kind, owner +): + """The manager's free takes a request's cache out of its map, then calls + ``discard_pending_stats`` and then the lender's free hook, so where ``discard_pending_stats`` + raises, the lender keeps the request's records of a cache the request no longer holds. + ResumeBounds of its next cache ignores and drops such a record of each kind the lender keeps per + request, answering as without records; the same record of the current cache changes resume_bounds + and stays.""" + s = _records_of(monkeypatch, kind, owner) + resume_bounds = s.lender.resume_bounds(_request(7)) + if owner == "current": + assert resume_bounds == CURRENT_RECORDS[kind] != NO_RECORDS, ( + f"resume_bounds {resume_bounds}" + ) + assert s.records.get(7) is s.record, "resume_bounds dropped a record of the current cache" + return + assert resume_bounds == NO_RECORDS, ( + f"resume_bounds {resume_bounds} reads a record of a replaced cache" + ) + assert 7 not in s.records, "resume_bounds kept a record of a replaced cache" + + +# -- a fetch's states ------------------------------------------------------------------------- + + +def test_a_fetch_walks_through_the_states_its_enum_names(monkeypatch): + """Write leases into a stand-in cache, ``_fetch_state`` read after each call: the first fetch + open, delivered with its copy pending, settled once the copy completed, then replaced by the + next, which is abandoned when its lease is released before a poll returned the view; a third is + abandoned when the copy of its marks fails. A new fetch is refused while one is unsettled.""" + states = _fetch_ledger._FetchState + lay = SimpleNamespace( + pool_groups=(0,), + pool_group_of=(0,), + windows=(None,), + num_layer_groups=1, + tokens_per_block=1, + ) + kv = SimpleNamespace( + num_committed_tokens=0, + history_length=0, + is_active=True, + num_blocks=2, + get_base_page_indices=lambda lg: [0, 1], + ) + manager = _Owner() + manager.kv_cache_map = {7: kv} + ref = weakref.ref(manager) + lender = _lender.Staging(ref, lay, None, (), SlotQueue({0: 4}), ref) + landed, failure = [], {} + event = SimpleNamespace(query=lambda: bool(landed)) + monkeypatch.setattr( + lender, "_build_view", lambda rows, keys: SimpleNamespace(groups=rows.block_ordinals) + ) + monkeypatch.setattr(lender, "_copy_segments", lambda rows, mask: []) + monkeypatch.setattr( + lender, + "_queue_copy", + lambda segments, to_staging: (_lender._QueuedCopy(event), failure.get("error")), + ) + seen = set() + + def state(lease, expected): + got = lender._ledger._fetch_state(7, lease._fetch) + seen.add(got) + assert got is expected, f"{got.name}, not {expected.name}" + + def lent(): + """A write lease as ``lend_write`` leaves it once the cache grew: recorded, granted.""" + fetch = _fetch_ledger._FetchRecord(0, kv) + lender._ledger._latest_fetches[7] = fetch + rows = _lender._LentBlockEntries([0], [np.arange(2)], [np.arange(2)]) + lease = _lender._StagingLease(lender, "write", 7, kv, rows, None, fetch) + lender._admit(lease, {0: 2}) + return lease + + def marked(lease): + lease.mark_arrived([np.ones(len(run), dtype=bool) for run in lease.poll().groups]) + + first = lent() + state(first, states.AWAITING_MARKS) + assert lender._ledger._unsettled_fetch_refusal(7, kv), "a fetch started while another was open" + marked(first) + state(first, states.COPY_PENDING) + assert lender._ledger._unsettled_fetch_refusal(7, kv), ( + "a fetch started while another's copy was pending" + ) + first.release() + landed.append(True) + lender._progress() # as a lend, a poll or a resume_bounds call does first; it asks the copy again + state(first, states.COPY_DONE) + assert lender._ledger._unsettled_fetch_refusal(7, kv) is None, ( + "a settled fetch kept a new one from starting" + ) + second = lent() + state(first, states.RETIRED) + state(second, states.AWAITING_MARKS) + second.release() # before a poll returned its view + state(second, states.ABANDONED) + third = lent() + failure["error"] = "planted" + marked(third) + state(third, states.ABANDONED) + third.release() + assert seen == set(states), f"states never seen: {set(states) - seen}" + + +# -- the hooks the manager calls -------------------------------------------------------------- + +# The manager's module, beside the sharing package. +MANAGER_SOURCE = Path(_lender.__file__).resolve().parents[1] / "kv_cache_manager_v2.py" +# Each lender method the manager calls on its ``_sharing``: positional arguments, keywords. +HOOKS = { + "_on_free": (2, ()), + "_on_shrink": (2, ()), + "_on_shutdown": (0, ()), + "_on_caches_closed": (0, ()), + "_on_reset": (0, ()), +} + + +def hook_calls(source): + """Each method ``source`` calls on an attribute named ``_sharing``, with the positional + arguments and the keywords of each of its calls.""" + calls = {} + for node in ast.walk(ast.parse(source)): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Attribute) + and node.func.value.attr == "_sharing" + ): + shape = (len(node.args), tuple(k.arg for k in node.keywords)) + calls.setdefault(node.func.attr, set()).add(shape) + return calls + + +def unbound_hooks(cls, calls): + """The methods of ``calls`` that ``cls`` lacks, or whose signature cannot bind one of their + calls.""" + unbound = [] + for name, shapes in sorted(calls.items()): + try: + signature = inspect.signature(getattr(cls, name)) + for args, keywords in shapes: + signature.bind(None, *[None] * args, **dict.fromkeys(keywords)) + except (AttributeError, TypeError, ValueError): + unbound.append(name) + return unbound + + +def test_the_lender_has_every_hook_the_manager_calls(monkeypatch): + """The manager calls five lender methods on its ``_sharing``, and ``Staging`` has each, taking + the arguments the manager passes. The check finds a ``Staging`` planted without one of them, + and one planted with a hook that takes fewer arguments.""" + calls = hook_calls(MANAGER_SOURCE.read_text()) + assert calls == {name: {shape} for name, shape in HOOKS.items()}, calls + assert unbound_hooks(_lender.Staging, calls) == [] + with monkeypatch.context() as patched: + patched.delattr(_lender.Staging, "_on_reset") + assert unbound_hooks(_lender.Staging, calls) == ["_on_reset"], "a missing hook passed" + with monkeypatch.context() as patched: + patched.setattr(_lender.Staging, "_on_free", lambda self, request_id: None) + assert unbound_hooks(_lender.Staging, calls) == ["_on_free"], "a short hook passed" + + +# -- the bytes the test kit stages ------------------------------------------------------------ + + +def test_the_pieces_of_a_staged_row_differ(kit): + """The 32-byte pieces of a staged row are pairwise distinct, so two of them trading places + shows in the row's bytes.""" + row = kit.staged_row(0x5A, 1, 3, 4096) + pieces = [row[i : i + 32] for i in range(0, len(row), 32)] + assert len(set(pieces)) == len(pieces), "a staged row repeats a piece, so a swap would not show" + + +# -- a lease's axes ------------------------------------------------------------------------------ + + +def test_each_axis_of_a_lease_follows_its_own_fields_alone(): + """Each axis of a staging lease follows its own fields alone: the holder axis the failure and + whether a poll returned the view, the slot axis the ticket, the copy axis the copy. Checked + over every combination of those four fields, ticket 0 (the first a slot queue issues) and an + empty failure among them, and of four lease fields no axis reads.""" + lender, _ = staging_over([None], [0], 32) + holder, slots, copies = _lender._HolderState, _lender._SlotState, _lender._CopyState + queued = _lender._QueuedCopy(SimpleNamespace(query=lambda: False)) + recorded = {None: True, _lender._COPY_WITHOUT_RECORD: False, queued: True} + for failure, seen, ticket, copy, released, marked, runs, doomed in itertools.product( + [None, "failed", ""], + [False, True], + [None, 0, 3], + recorded, + [False, True], + [False, True], + [None, SlotRuns({0: (0, 1)})], + [None, "doomed"], + ): + lease = _lender._StagingLease(lender, "write", 7) + lease._failure, lease._view_returned, lease._ticket = failure, seen, ticket + lease._copy, lease._released, lease._marked = copy, released, marked + lease._slot_runs, lease._first_poll_failure = runs, doomed + failed = failure is not None + want = ( + holder.FAILED if failed else holder.VIEW_RETURNED if seen else holder.VIEW_PENDING, + slots.WAITING if ticket is not None else slots.NOT_WAITING, + copies.NONE_OR_RECORDED if recorded[copy] else copies.WITHOUT_RECORD, + ) + got = (lease._holder_state(), lease._slot_state(), lease._copy_state()) + assert got == want, ( + f"failure {failure}, seen {seen}, ticket {ticket}, copy {copy}, released {released}, " + f"marked {marked}, runs {runs}, doomed {doomed}: {got}" + ) + + +class _Cache: + """A stand-in cache a weak reference can follow.""" + + +def test_a_lender_and_its_records_go_with_its_last_reference(): + """With the cyclic collector off, a lender and a cache its fetch records hold go as soon as the + last reference to the lender does: its ledger reaches it only through a weak reference, so no + reference cycle keeps either until a collection.""" + lender, _ = staging_over([None], [0], 32) + cache = _Cache() + lender._ledger._kept_prefixes[7] = (cache, 0, 1) + lender_ref, cache_ref = weakref.ref(lender), weakref.ref(cache) + collecting = gc.isenabled() + gc.disable() + try: + del lender, cache + assert lender_ref() is None, "a reference cycle keeps the lender after its last reference" + assert cache_ref() is None, "the lender's records keep a cache after its last reference" + finally: + if collecting: + gc.enable() diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py new file mode 100644 index 000000000000..82d5bb7fe711 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py @@ -0,0 +1,5156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The staging lender's contract over real KV cache managers, driven through the public API; a few +oracles read lender internals the API does not show. Allocates device pools.""" + +import gc +import mmap +import threading +import weakref +from types import SimpleNamespace +from typing import List, Set + +import numpy as np +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + Lease, + LeaseView, + Lender, + RegionHold, + ResumeBounds, + StagingOptions, + attach_lender, +) +from tensorrt_llm._utils import prefer_pinned + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") + +TPB = 32 +WINDOW = 64 +PROMPT = list(range(1000, 1097)) # three whole blocks and one token +OTHER_PROMPT = list(range(5000, 5097)) +END = 96 +BLOCKS = END // TPB +WINDOWED_PROMPT = list(range(2000, 2161)) # five whole blocks and one token +OTHER_WINDOWED_PROMPT = list(range(6000, 6161)) +WINDOWED_END = 160 +SOURCE, TARGET, OTHER = 1, 2, 3 +SCOPE = b"lender-contract-suite" + + +def attach(mgr, *, tokens_per_fetch, planned_fetches=1, max_bytes=None, scope=SCOPE): + """The public attach, as an integrator makes it.""" + options = StagingOptions(tokens_per_fetch, planned_fetches, max_bytes) + return attach_lender(mgr, scope=scope, staging=options) + + +def attach_breaking(rules): + """An attach installing a lender whose ``rules`` (method name -> function, or a function of the + manager returning them) replace the real ones, on the lender or on its fetch ledger, whichever + defines the method: a lender breaking exactly those rules.""" + + def attach_rule_breaker( + mgr, *, tokens_per_fetch, planned_fetches=1, max_bytes=None, scope=SCOPE + ): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _fetch_ledger, _lender + + chosen = dict(rules(mgr) if callable(rules) else rules) + on_ledger = {n: f for n, f in chosen.items() if _on_ledger(n)} + lender_rules = {n: f for n, f in chosen.items() if n not in on_ledger} + cls = type("RuleBreaker", (_lender.Staging,), lender_rules) + options = StagingOptions(tokens_per_fetch, planned_fetches, max_bytes) + lender = _lender._attach_staging(mgr, scope=scope, staging=options, cls=cls) + if on_ledger: + # A ledger of the rule-breaking class in place of the real one, before it records + # anything; its rules reach the lender through _lender_ref. + ledger = type("LedgerRuleBreaker", (_fetch_ledger._FetchLedger,), on_ledger) + lender._ledger = ledger(lender._layout, lender._ledger._copy_done) + lender._ledger._lender_ref = weakref.ref(lender) + return lender + + return attach_rule_breaker + + +def _on_ledger(name): + """Whether rule ``name`` is a method of the lender's fetch ledger rather than of the lender.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _fetch_ledger + + return not name.startswith("__") and name in vars(_fetch_ledger._FetchLedger) + + +def ready_view(lease, mgr, tries=3): + """The lease's view once the copies queued so far have run; fails if it never comes.""" + for _ in range(tries): + mgr._stream.synchronize() + view = lease.poll() + if view is not None: + return view + raise AssertionError(f"the lease never became ready: {lease.failure}") + + +def by_group(view): + return {run.layer_group_index: run.block_ordinals.tolist() for run in view.groups} + + +def one_fetch_bytes(kit, mgr, tokens_per_fetch): + """Bytes one fetch of ``tokens_per_fetch`` tokens needs, from the device pages: every block for full + attention, a window's blocks for a window (no sinks here).""" + dev = kit.DevicePages(mgr) + total = 0 + for lg, window in enumerate(kit.windows(mgr)): + blocks = -(-tokens_per_fetch // TPB) + if window is not None: + blocks = min(blocks, -(-window // TPB)) + total += blocks * dev.page_bytes(lg) + return total + + +def check_staged(kit, mgr, lender, request, view): + """Every row's slot, inside its region, holds the request's device page of that block.""" + kv = kit.kv(mgr, request) + dev = kit.DevicePages(mgr) + rows = 0 + for run in view.groups: + region = lender.regions[run.region_index] + assert run.layer_group_index in lender.regions[run.region_index].layer_group_indices + assert region.slot_bytes == dev.page_bytes(run.layer_group_index) + slots = kit.pages(kv, run.layer_group_index) + for address, ordinal in zip(run.addresses.tolist(), run.block_ordinals.tolist()): + assert region.address <= address < region.address + region.nbytes + assert (address - region.address) % region.slot_bytes == 0 + device = kit.digest(dev.read(run.layer_group_index, slots[ordinal])) + assert kit.digest(kit.host_bytes(address, region.slot_bytes)) == device + rows += 1 + return rows + + +def state_of(kit, mgr, lender, request): + """Capacity, history, committed tokens and resume_bounds of the request, or what stands for none.""" + kv = kit.kv(mgr, request) + try: + resume_bounds = lender.resume_bounds(request) + except ValueError: + resume_bounds = "no cache" + if kv is None: + return None, resume_bounds + return (kv.capacity, kv.history_length, kv.num_committed_tokens), resume_bounds + + +# -- attach ----------------------------------------------------------------------------------- + + +def test_a_manager_takes_one_lender_for_its_life(real_manager): + with real_manager() as mgr: + for scope in ("text", 1, bytearray(SCOPE)): # bytes() would take the last two + with pytest.raises(TypeError, match="scope must be bytes"): + attach_lender(mgr, scope=scope, staging=StagingOptions(END)) + with pytest.raises(TypeError, match="staging must be StagingOptions"): + attach_lender(mgr, scope=SCOPE, staging=END) + lender = attach(mgr, tokens_per_fetch=END) # the refused calls attached nothing + assert isinstance(lender, Lender) + with pytest.raises(ValueError, match="already attached"): + attach(mgr, tokens_per_fetch=END) + + +def test_recurrent_state_is_refused(kit, hybrid_manager): + before = len(kit.retained()) + with hybrid_manager() as mgr: + with pytest.raises(ValueError, match="recurrent"): + attach(mgr, tokens_per_fetch=256) + assert len(kit.retained()) == before, "a refused attach allocated staging" + + +def sparse_layer(monkeypatch, index): + """Managers built from now on mark the buffers of their layer ``index`` sparse, which no model + here does: the runtime may then lock that layer group's read-only pages in host memory.""" + from dataclasses import replace + + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig, BufferConfig + + build = KVCacheManagerV2._build_cache_config + + def with_a_sparse_layer(self, config): + layers = list(config.layers) + layer = layers[index] + layers[index] = AttentionLayerConfig( + layer_id=layer.layer_id, + buffers=[ + BufferConfig(b.role, b.size, b.tokens_per_block_override, is_sparse=True) + for b in layer.buffers + ], + sliding_window_size=layer.sliding_window_size, + num_sink_tokens=layer.num_sink_tokens, + ) + return build(self, replace(config, layers=layers)) + + monkeypatch.setattr(KVCacheManagerV2, "_build_cache_config", with_a_sparse_layer) + + +# case: (manager keywords, the layer whose buffers are sparse); "later" puts them in a layer group +# after a dense one +SPARSE_CASES = {"first": ({}, 0), "later": ({"windows": [WINDOW, 256]}, 1)} + + +@pytest.mark.parametrize("case", list(SPARSE_CASES)) +def test_a_layer_group_of_sparse_buffers_is_refused(kit, host_tier_manager, monkeypatch, case): + """A cache can lock a sparse layer group's read-only pages in host memory, where a page index + names a host slot that staging would address as a device slot: the attach refuses such a + manager, naming the group and keeping nothing, also where it follows a dense one. The runtime + builds sparse buffers only with host memory as the cache level below the GPU, as here.""" + keywords, index = SPARSE_CASES[case] + sparse_layer(monkeypatch, index) + before = len(kit.retained()) + with host_tier_manager(**keywords) as mgr: + layers = {int(layer.layer_id): layer for layer in mgr.impl.init_config.layers} + flags = [ + any(b.is_sparse for b in layers[int(ids[0])].buffers) for ids in mgr.impl.layer_grouping + ] + sparse = [lg for lg, flag in enumerate(flags) if flag] + assert sparse and (0 in sparse) == (case == "first"), ( + f"sparse buffers in layer groups {sparse}: the check proves nothing" + ) + message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) + assert f"layer groups {sparse} hold sparse buffers" in message + assert "_sharing" not in vars(mgr) + assert len(kit.retained()) == before, "a refused attach kept something" + + +def stand_in(cp_type=None): + """The smallest manager the attach path reads before any other work: an uninitialised + ``KVCacheManagerV2`` whose mapping is rank 0 of two ``cp_type`` ranks, or has no mapping.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.mapping import Mapping + + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + if cp_type is not None: + manager.mapping = Mapping(world_size=2, cp_size=2, rank=0, cp_config={"cp_type": cp_type}) + return manager + + +def refusal(call, kind=ValueError) -> str: + """The message of the ``kind`` error ``call`` raises; the check fails on anything else.""" + try: + call() + except kind as error: + return str(error) + except AssertionError: + raise + except Exception as error: # handed to the check as its failure + raise AssertionError(f"expected a {kind.__name__}, got {error!r}") from error + raise AssertionError(NOT_REFUSED) + + +PAST_THE_CHECKS = "the attach got past its checks" +NOT_REFUSED = "the call went through" + + +def _without_parallelism(manager=None): + """A mapping read as one rank: what a lender misreading the manager's parallelism sees.""" + return SimpleNamespace(cp_size=1, pp_size=1) + + +def _defaulting_a_missing_mapping(manager): + return getattr(manager, "mapping", None) or _without_parallelism() + + +# Managers the attach refuses before any other work, as (manager, error, words), and a wrong +# reading of the manager's mapping that would let one of them through. +REFUSED = { + "context_parallel": ( + [ + (lambda: stand_in("HELIX"), ValueError, "context parallelism"), + (lambda: stand_in("ULYSSES"), ValueError, "context parallelism"), + ], + _without_parallelism, + ), + "not_v2_or_unmapped": ( + [(SimpleNamespace, TypeError, "KVCacheManagerV2"), (stand_in, TypeError, "mapping")], + _defaulting_a_missing_mapping, + ), +} + + +@pytest.mark.parametrize("case", list(REFUSED)) +def test_a_manager_no_lender_serves_is_refused_before_any_work(monkeypatch, case): + """Each ``(manager, kind, words)`` of the case is refused by the attach with a ``kind`` error + naming ``words``, before the layout is derived: a tripwire there fails the check.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + def tripwire(manager): + raise AssertionError(PAST_THE_CHECKS) + + monkeypatch.setattr(_lender, "derive_layout", tripwire) + before = len(_lender._retained()) + for make, kind, words in REFUSED[case][0]: + staging = make() + assert words in refusal(lambda: attach(staging, tokens_per_fetch=END), kind) + assert "_sharing" not in vars(staging) + assert len(_lender._retained()) == before, "a refused attach allocated staging" + + +def test_an_attach_that_raises_leaves_nothing_attached(real_manager, monkeypatch, attach=attach): + """An attach whose last step before the install raises, its log line here, leaves the manager + without a lender, so a second attach goes through.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + def raising(*args, **kwargs): + raise RuntimeError("planted") + + with real_manager() as mgr: + with monkeypatch.context() as patched: + patched.setattr(_lender, "logger", SimpleNamespace(info=raising)) + with pytest.raises(RuntimeError, match="planted"): + attach(mgr, tokens_per_fetch=END) + assert "_sharing" not in vars(mgr), "a failed attach left its lender installed" + assert isinstance(attach(mgr, tokens_per_fetch=END), Lender) + + +def pipeline_stage(rank): + """Stage ``rank`` of a two-stage pipeline.""" + from tensorrt_llm.mapping import Mapping + + return Mapping(world_size=2, pp_size=2, rank=rank) + + +@pytest.mark.parametrize("rank", [0, 1], ids=["first_stage", "last_stage"]) +def test_staging_refuses_a_pipeline_stage(kit, real_manager, rank, attach=attach): + """Each stage's staging slots and windows are its own, so one call could raise on one stage and + lend on another: staging refuses every stage at the attach, which then changed nothing.""" + with real_manager(mapping=pipeline_stage(rank), num_layers=4, windows=[256, WINDOW]) as mgr: + assert mgr.mapping.pp_size == 2 and len(mgr.pp_layers) == 2 + before = len(kit.retained()) + message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) + assert "pipeline parallelism" in message, message + assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + + +def unpaired_draft(mgr) -> None: + """The state of a draft manager without joint reuse: it never publishes blocks for reuse.""" + mgr._can_publish_block_reuse = False + + +COMMITS_NOTHING = { + "block_reuse_off": ({"enable_block_reuse": False}, None), + "unpaired_draft": ({}, unpaired_draft), +} + + +@pytest.mark.parametrize("case", list(COMMITS_NOTHING)) +def test_staging_refuses_a_manager_that_commits_no_blocks(kit, real_manager, case, attach=attach): + """Staging publishes committed blocks only, so a manager that commits none is refused at the + attach, which then changed nothing.""" + manager_kwargs, adjust = COMMITS_NOTHING[case] + with real_manager(**manager_kwargs) as mgr: + if adjust is not None: + adjust(mgr) + before = len(kit.retained()) + message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) + assert "block reuse" in message, message + assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + + +@pytest.mark.parametrize("capacity_only", [False, True]) +def test_staging_refuses_a_manager_with_a_kv_cache_connector( + kit, real_manager, capacity_only, attach=attach +): + """A KV cache connector serves a request's prefix at its first context chunk, measured from the + committed tokens, so it cannot lower a history a windowed fetch moved, and its loads run off the + manager's stream. Staging refuses either transfer or capacity-only connectors at the attach, + which then changed nothing. A stand-in exposing the connector mode suffices here.""" + connector = SimpleNamespace(capacity_only=capacity_only) + with real_manager(windows=[WINDOW, 256], kv_connector_manager=connector) as mgr: + assert mgr.kv_connector_manager is connector + before = len(kit.retained()) + message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) + assert "KV cache connector" in message, message + assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + + +def test_lender_refuses_configured_lossy_cold_page_quantization_before_allocating( + kit, real_manager, monkeypatch, attach=attach +): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + from tensorrt_llm.llmapi.llm_args import ColdPageQuantizationCompressionConfig + from tensorrt_llm.runtime.kv_cache_manager_v2 import create_default_kv_cache_cold_page_codec + + # The stand-in supplies a native codec so the real constructor records the actual config. + provider = SimpleNamespace( + config=ColdPageQuantizationCompressionConfig(), + create_cold_page_codec=lambda *args, **kwargs: create_default_kv_cache_cold_page_codec(), + ) + with real_manager(cold_page_codec_provider=provider) as mgr: + assert mgr._has_lossy_cold_page_codec + before = len(kit.retained()) + + def tripwire(*args, **kwargs): + raise AssertionError("a refused codec must not derive layouts or allocate memory") + + monkeypatch.setattr(_lender, "derive_layout", tripwire) + monkeypatch.setattr(_lender, "_allocate_retained", tripwire) + with pytest.raises(ValueError, match="lossy.*cold-page codec"): + attach(mgr, tokens_per_fetch=END) + assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + + +@pytest.mark.parametrize("custom_lossless", [False, True]) +def test_lender_accepts_default_and_custom_lossless_cold_page_codecs( + real_manager, custom_lossless, attach=attach +): + from tensorrt_llm.runtime.kv_cache_manager_v2 import create_default_kv_cache_cold_page_codec + + provider = None + if custom_lossless: + provider = SimpleNamespace( + create_cold_page_codec=lambda *args, **kwargs: create_default_kv_cache_cold_page_codec() + ) + with real_manager(cold_page_codec_provider=provider) as mgr: + assert not mgr._has_lossy_cold_page_codec + assert isinstance(attach(mgr, tokens_per_fetch=END), Lender) + + +def eagle3(): + """One-model Eagle 3: its draft at position ``i`` reads prompt token ``i + 1``.""" + from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig + + return Eagle3DecodingConfig(max_draft_len=1, speculative_model="draft-model") + + +def mtp(draft_len=1, vanilla=False): + """One-model MTP: MTP-Eagle reads one prompt token ahead, vanilla MTP ``draft_len``.""" + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + config = MTPDecodingConfig(max_draft_len=draft_len, use_mtp_vanilla=vanilla) + config.num_nextn_predict_layers = draft_len # the checkpoint sets it at model load + return config + + +def dflash(): + """A standalone DFlash drafter: its draft KV at ``i`` depends on prompt tokens ``[0, i]``.""" + from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig + + return DFlashDecodingConfig(max_draft_len=2, speculative_model="draft-model") + + +def pard(): + """PARD: its drafter reads the prompt as the target does.""" + from tensorrt_llm.llmapi.llm_args import PARDDecodingConfig + + return PARDDecodingConfig(max_draft_len=2, speculative_model="draft-model") + + +def ngram(): + """A drafter outside the one-model engine, over its own request view, as a two-model draft.""" + from tensorrt_llm.llmapi.llm_args import NGramDecodingConfig + + return NGramDecodingConfig(max_draft_len=2, max_matching_ngram_size=2) + + +# Managers built for a one-model draft reading prompt tokens past a position: case -> (manager +# keywords, tokens read ahead). The draft layers share the target's manager, or have a joint-reuse +# draft pool of their own. +READS_AHEAD = { + "eagle3_layers_in_the_target": (lambda: dict(spec_config=eagle3()), 1), + "mtp_eagle_layers_in_the_target": (lambda: dict(spec_config=mtp()), 1), + "vanilla_mtp_layers_in_the_target": (lambda: dict(spec_config=mtp(2, vanilla=True)), 2), + "eagle3_joint_reuse_draft_pool": ( + lambda: dict(spec_config=eagle3(), is_draft=True, joint_kv_cache_reuse=True), + 1, + ), +} +# Managers whose blocks depend on no token past their end, all committing blocks. +READS_NONE = { + "no_speculation": lambda: dict(), + "pard_target": lambda: dict(spec_config=pard()), + "dflash_joint_reuse_draft_pool": lambda: dict( + spec_config=dflash(), is_draft=True, joint_kv_cache_reuse=True + ), + "own_view_draft_pool": lambda: dict( + spec_config=ngram(), is_draft=True, joint_kv_cache_reuse=True + ), +} +READ_AHEAD_WORDS = "past their end" +OVER_REFUSED = "the attach refused it" + + +def accepted(call): + """What ``call`` returns; the check fails on the ``ValueError`` of a refusal.""" + try: + return call() + except ValueError as error: + raise AssertionError(f"{OVER_REFUSED}: {error}") from error + + +def check_a_draft_reading_ahead_is_refused(kit, mgr, attach, ahead): + """Staging refuses ``mgr``, built for a draft reading ``ahead`` prompt tokens past a position, + and the refused attach changed nothing.""" + before = len(kit.retained()) + message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) + assert READ_AHEAD_WORDS in message and str(ahead) in message, message + assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + + +@pytest.mark.parametrize("case", list(READS_AHEAD)) +def test_staging_refuses_a_manager_built_for_a_draft_reading_past_a_block( + kit, real_manager, case, attach=attach +): + make, ahead = READS_AHEAD[case] + with real_manager(**make()) as mgr: + assert mgr.draft_prompt_lookahead == ahead + assert mgr.reuse_match_backoff == ahead + check_a_draft_reading_ahead_is_refused(kit, mgr, attach, ahead) + + +@skip_pre_blackwell +@pytest.mark.parametrize("fp8_ds_mla", [False, True], ids=["fp8", "fp8_ds_mla"]) +def test_staging_refuses_a_deepseek_v4_manager_whose_mtp_layers_read_ahead( + kit, deepseek_v4_manager, fp8_ds_mla +): + from tensorrt_llm._torch.speculative import draft_prompt_lookahead + + # The class opts out of the reuse backoff, so the backoff alone would not see the draft. + spec_config = mtp() + assert draft_prompt_lookahead(spec_config) == 1 + with deepseek_v4_manager(spec_config=spec_config, fp8_ds_mla=fp8_ds_mla) as mgr: + assert mgr.draft_prompt_lookahead == 1 + assert mgr.reuse_match_backoff == 0 + check_a_draft_reading_ahead_is_refused(kit, mgr, attach, 1) + + +@pytest.mark.parametrize("case", list(READS_NONE)) +def test_staging_accepts_a_manager_whose_blocks_read_no_token_past_their_end( + kit, real_manager, case, attach=attach +): + with real_manager(**READS_NONE[case]()) as mgr: + assert mgr.draft_prompt_lookahead == 0 + assert mgr.reuse_match_backoff == 0 + assert isinstance(accepted(lambda: attach(mgr, tokens_per_fetch=END)), Lender) + + +# -- sizing ----------------------------------------------------------------------------------- + + +# planned_fetches, max_bytes in fetches (None for no cap), and the fetches staging then holds at once +SIZING = { + "planned_fetches": (2, None, 2), + "max_bytes_of_one_fetch": (4, 1, 1), + "max_bytes_of_two_and_a_half_fetches": (4, 2.5, 2), +} + + +@pytest.mark.parametrize("case", list(SIZING)) +def test_staging_holds_planned_fetches_fetches_at_once_capped_by_max_bytes(kit, real_manager, case): + planned_fetches, cap, holding = SIZING[case] + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + one = one_fetch_bytes(kit, mgr, END) + with pytest.raises(ValueError, match="one fetch"): + attach(mgr, tokens_per_fetch=END, planned_fetches=4, max_bytes=one - 1) + max_bytes = None if cap is None else int(cap * one) + lender = attach( + mgr, tokens_per_fetch=END, planned_fetches=planned_fetches, max_bytes=max_bytes + ) + (region,) = lender.regions + assert region.slot_bytes == kit.DevicePages(mgr).page_bytes(0) + assert region.nbytes == region.slot_count * region.slot_bytes + assert holding * one <= region.nbytes <= (max_bytes or holding * one) + leases = [lender.lend_read(source, 0, END) for _ in range(holding + 1)] + for lease in leases[:holding]: + ready_view(lease, mgr) + assert leases[-1].poll() is None and leases[-1].failure is None, "one fetch too many" + leases[0].release() + check_staged(kit, mgr, lender, source, ready_view(leases[-1], mgr)) + for lease in leases[1:]: + lease.release() + + +@pytest.mark.parametrize("num_layers", [2, 3], ids=["one_pool_group", "two_pool_groups"]) +def test_a_windowed_range_of_tokens_per_fetch_always_fits(kit, real_manager, num_layers): + with real_manager(windows=[WINDOW, 256], num_layers=num_layers) as mgr: + source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_WINDOWED_PROMPT) + lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + dev = kit.DevicePages(mgr) + rows = {} + for lg, (window, g) in enumerate(zip(kit.windows(mgr), kit.pool_group_of(mgr))): + rows[g] = rows.get(g, 0) + (WINDOW // TPB if window else WINDOWED_END // TPB) + groups = kit.pool_group_ids(mgr) + assert len(lender.regions) == len(groups) == num_layers - 1 + group_of = kit.pool_group_of(mgr) + for g, region in zip(groups, lender.regions): + assert (region.slot_count, region.slot_bytes) == (rows[g], dev.group_page_bytes(g)) + assert region.nbytes == region.slot_count * region.slot_bytes + assert region.layer_group_indices == tuple( + lg for lg, pg in enumerate(group_of) if pg == g + ) + # The regions' layer groups are disjoint and cover every layer group. + held = sorted(lg for region in lender.regions for lg in region.layer_group_indices) + assert held == list(range(len(group_of))) + read = lender.lend_read(source, 0, WINDOWED_END) # takes every slot + assert check_staged(kit, mgr, lender, source, ready_view(read, mgr)) == sum(rows.values()) + read.release() + write = lender.lend_write(target, 0, WINDOWED_END) + view = ready_view(write, mgr) + assert view.num_block_entries == sum(rows.values()) + write.mark_arrived(view.row_masks()) + write.release() + + +def test_the_staging_memory_is_the_parts_and_nothing_more(kit, real_manager, attach=attach): + # Three layers: two pool groups, so two regions. + with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: + lender = attach(mgr, tokens_per_fetch=WINDOWED_END, planned_fetches=2) + regions = lender.regions + assert len(regions) == len(kit.pool_group_ids(mgr)) == 2 + (memory,) = kit.staging_memory(regions) + assert memory.address <= regions[0].address < memory.address + memory.nbytes + assert memory.nbytes == sum(p.nbytes for p in regions), ( + "staging is not the sum of its regions" + ) + + +class FreeSpy: + """The CUDA runtime, recording the address of every page-locked host allocation freed.""" + + def __init__(self, runtime): + self._runtime = runtime + self.freed = [] + + def __getattr__(self, name): + found = getattr(self._runtime, name) + if name != "cudaFreeHost": + return found + + def free(address): + self.freed.append(int(address)) + return found(address) + + return free + + +def test_the_staging_memory_is_pinned_at_its_size_and_freed_once(kit, real_manager, monkeypatch): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + if not prefer_pinned(): + pytest.skip("staging is pageable where pinning does not pay off") + spy = FreeSpy(_lender.cudart) + monkeypatch.setattr(_lender, "cudart", spy) + with real_manager() as mgr: + regions = attach(mgr, tokens_per_fetch=END, planned_fetches=3).regions + base, total = regions[0].address, sum(p.nbytes for p in regions) + whole_pages = -(-total // mmap.PAGESIZE) * mmap.PAGESIZE + assert 1 << (total - 1).bit_length() > whole_pages, "a power of two would pass unseen" + start, pinned = kit.pinned_range(base) + assert start == base and total <= pinned <= whole_pages, ( + f"{pinned} bytes pinned for {total}" + ) + mgr.shutdown() + assert spy.freed == [base] and kit.pinned_range(base) is None, "not freed at the shutdown" + mgr.shutdown() + gc.collect() + assert spy.freed == [base], "freed more than once" + + +@pytest.mark.parametrize("pinned", [False, True], ids=["pageable", "pinned"]) +@pytest.mark.parametrize("failed", [False, True], ids=["passed", "failed"]) +def test_test_teardown_frees_only_added_staging( + kit, retained_cleanup, monkeypatch, pinned: bool, failed: bool +) -> None: + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + if pinned and not prefer_pinned(): + pytest.skip("staging is pageable where pinning does not pay off") + torch.cuda.init() + monkeypatch.setattr(_lender, "prefer_pinned", lambda: pinned) + before = _lender._allocate_retained(mmap.PAGESIZE) + cleanup = retained_cleanup(failed) + next(cleanup) + added = _lender._allocate_retained(mmap.PAGESIZE) + address = added.address + index_owner = np.empty(1, dtype=np.int32) + _lender._retain_until_exit(index_owner) + try: + if pinned: + assert kit.pinned_range(address) is not None + with pytest.raises(StopIteration): + next(cleanup) + if failed: + assert added.address == address, "a failed test's staging was freed" + assert any(owner is added for owner in kit.retained()) + assert any(owner is index_owner for owner in kit.retained()) + return + assert added.address == 0 and added._pageable is None, "teardown did not free staging" + assert not any(owner is added or owner is index_owner for owner in kit.retained()) + assert before.address != 0 and any(owner is before for owner in kit.retained()) + if pinned: + assert kit.pinned_range(address) is None + assert kit.pinned_range(before.address) is not None + finally: + for owner in (added, before): + owner.free() + _lender._stop_retaining(owner) + _lender._stop_retaining(index_owner) + + +def test_pageable_staging_relays_bytes_and_frees_nothing_page_locked( + kit, real_manager, monkeypatch +): + """Where pinning does not pay off, as under confidential computing, staging is pageable: it is + not page-locked, each lender's stays its own while both live, a publish relayed into a fetch + lands byte for byte, and the shutdown releases it without freeing any page-locked memory.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + spy = FreeSpy(_lender.cudart) + monkeypatch.setattr(_lender, "cudart", spy) + monkeypatch.setattr(_lender, "prefer_pinned", lambda: False) + with real_manager() as mgr_a, real_manager() as mgr_b: + source = kit.published(mgr_a, SOURCE, PROMPT) + target = kit.admitted(mgr_b, TARGET, PROMPT) + lender_a, lender_b = ( + attach(mgr_a, tokens_per_fetch=END), + attach(mgr_b, tokens_per_fetch=END), + ) + regions = lender_a.regions + lender_b.regions + locked = [kit.pinned_range(region.address) for region in regions] + assert locked == [None] * len(regions), f"pageable staging was page-locked: {locked}" + # A pageable buffer freed early can come back to the next lender at the same address. + spans = sorted((region.address, region.address + region.nbytes) for region in regions) + assert all(end <= begin for (_, end), (begin, _) in zip(spans, spans[1:])), ( + f"two lenders' staging overlaps: {spans}" + ) + memories = kit.staging_memory(regions) + assert len(memories) == 2, f"{len(memories)} kept allocations hold two lenders' staging" + for memory in memories: + kept = memory._pageable + assert kept is not None and kept.ctypes.data == memory.address, ( + "the pageable staging is not an array the lender keeps" + ) + publish = lender_a.lend_read(source, 0, END) + publish_view = ready_view(publish, mgr_a) + fetch = lender_b.lend_write(target, 0, END) + view = fetch.poll() + assert view is not None, f"the write lease was not ready: {fetch.failure}" + kit.fill_sentinel(mgr_b, target) + fetch.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) + mgr_b._stream.synchronize() + assert lender_b.resume_bounds(target) == (END, 0) + for ordinal in range(BLOCKS): + got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) + assert got == kit.digest(kit.page(mgr_a, source, 0, ordinal)), "a fetched page differs" + for lease in (fetch, publish): + lease.release() + for mgr in (mgr_a, mgr_b): + mgr.shutdown() + assert spy.freed == [], "pageable staging was freed as page-locked memory" + assert not kit.staging_kept(regions), "the staging memory was kept past the shutdown" + + +# -- publish ---------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "tokens", [PROMPT, PROMPT[:END]], ids=["with_a_partial_block", "on_a_block_boundary"] +) +def test_a_publish_lends_the_committed_pages_under_their_reuse_keys(kit, real_manager, tokens): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, tokens) + lender = attach(mgr, tokens_per_fetch=END) + lease = lender.lend_read(source, 0, END) + assert isinstance(lease, Lease) + view = ready_view(lease, mgr) + assert isinstance(view, LeaseView) and lease.poll() is view, "one view, every poll" + (run,) = view.groups + assert run.block_ordinals.tolist() == list(range(BLOCKS)) and run.region_index == 0 + assert check_staged(kit, mgr, lender, source, view) == BLOCKS + kv = kit.kv(mgr, source) + assert [bytes(name[16:48]) for name in run.object_keys] == kit.chain_keys(kv, tokens)[ + :BLOCKS + ] + assert len({bytes(name[:16]) for name in run.object_keys}) == 1, "one namespace" + assert len({bytes(name[48:]) for name in run.object_keys}) == 1, ( + "one layer group, one shard" + ) + # The manager committed under those keys: its tree finds the block after them by them. + extra = list(range(7000, 7000 + TPB)) + probed = mgr.impl.probe_first_new_block_key(kv.reuse_scope, PROMPT[:END] + extra) + assert probed == kit.chain_keys(kv, PROMPT[:END] + extra)[BLOCKS] + lease.release() + with pytest.raises(RuntimeError): + lease.poll() + + +def test_a_multimodal_publish_is_named_by_the_keys_the_manager_commits(kit, real_manager): + image = dict(multimodal_positions=[40], multimodal_lengths=[16]) + with real_manager() as mgr: + first = kit.published(mgr, SOURCE, PROMPT, multimodal_hashes=[list(range(1, 9))], **image) + second = kit.published( + mgr, OTHER, PROMPT, multimodal_hashes=[list(range(8, 0, -1))], **image + ) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + leases = [lender.lend_read(r, 0, END) for r in (first, second)] + names = [ready_view(lease, mgr).groups[0].object_keys for lease in leases] + keys = [[bytes(n[16:48]) for n in run_names] for run_names in names] + kv = kit.kv(mgr, first) + augmented = list(mgr._augment_tokens_for_block_reuse(PROMPT, first)) + assert keys[0] == kit.chain_keys(kv, augmented)[:BLOCKS] + extra = list(range(7000, 7000 + TPB)) + probed = mgr.impl.probe_first_new_block_key(kv.reuse_scope, augmented[:END] + extra) + assert probed == kit.chain_keys(kv, augmented[:END] + extra)[BLOCKS] + # The image sits in block 1: the block before it is shared, the rest is not. + assert keys[0][0] == keys[1][0] and not set(keys[0][1:]) & set(keys[1][1:]) + assert keys[0][1:] != kit.chain_keys(kv, PROMPT)[1:BLOCKS], "the image changes the keys" + for lease in leases: + lease.release() + + +@pytest.mark.parametrize( + "scope", [dict(cache_salt="tenant-a"), dict(lora_task_id=7)], ids=["cache_salt", "lora_task_id"] +) +def test_a_publish_is_named_under_the_request_s_reuse_scope(kit, real_manager, scope): + with real_manager() as mgr: + plain = kit.published(mgr, SOURCE, PROMPT) + scoped = kit.published(mgr, OTHER, PROMPT, **scope) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + leases = [lender.lend_read(r, 0, END) for r in (plain, scoped)] + keys = [ + [bytes(n[16:48]) for n in ready_view(lease, mgr).groups[0].object_keys] + for lease in leases + ] + kv = kit.kv(mgr, scoped) + assert keys[1] == kit.chain_keys(kv, PROMPT)[:BLOCKS], "a name ignores the reuse scope" + # The manager committed under those keys: its tree finds the block after them by them. + extra = list(range(7000, 7000 + TPB)) + probed = mgr.impl.probe_first_new_block_key(kv.reuse_scope, PROMPT[:END] + extra) + assert probed == kit.chain_keys(kv, PROMPT[:END] + extra)[BLOCKS] + assert not set(keys[0]) & set(keys[1]), "the reuse scope changes no key" + for lease in leases: + lease.release() + + +def test_a_windowed_publish_lends_only_what_the_window_at_its_end_reads(kit, real_manager): + with real_manager(windows=[WINDOW, 256]) as mgr: + source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) + lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + windows = kit.windows(mgr) + sliding, full = windows.index(WINDOW), windows.index(None) + lease = lender.lend_read(source, 0, WINDOWED_END) + view = ready_view(lease, mgr) + beg, end = kit.stale_blocks(mgr, sliding, WINDOWED_END) + assert end > beg, "the window must leave early blocks behind" + blocks = range(WINDOWED_END // TPB) + in_window = [o for o in blocks if not beg <= o < end] + assert by_group(view) == {full: list(blocks), sliding: in_window} + addresses = np.concatenate([run.addresses for run in view.groups]) + assert len(set(addresses.tolist())) == len(addresses), "every row has its own slot" + assert check_staged(kit, mgr, lender, source, view) == len(addresses) + lease.release() + + # The request's own window has since passed what a window at 96 reads: none is lent. + early = lender.lend_read(source, 0, END) + rows = by_group(ready_view(early, mgr)) + assert rows[full] == list(range(BLOCKS)) + assert rows.get(sliding, []) == [] + early.release() + + +def test_a_range_without_rows_is_ready_at_once_and_waits_for_no_one(kit, real_manager): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + holder = lender.lend_read(source, 0, END) # every slot + ready_view(holder, mgr) + empty = lender.lend_read(source, END, END) + view = empty.poll() + assert view is not None and view.num_block_entries == 0 + assert len(view.groups) == kit.num_layer_groups(mgr) + empty.release() + holder.release() + + +# -- ready means copied ----------------------------------------------------------------------- + + +def test_a_publish_is_ready_only_once_its_copy_ran(kit, real_manager, attach=attach): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + with kit.held_stream(mgr._stream) as gate: + lease = lender.lend_read(source, 0, END) + assert lease.poll() is None, "ready before its copy into staging ran" + assert lease.poll() is None and lease.failure is None + gate.open() + check_staged(kit, mgr, lender, source, ready_view(lease, mgr)) + lease.release() + + +# -- no slot reused while a copy or a backend may touch it ----------------------------------- + + +def test_a_slot_returns_only_after_the_copy_into_it_ran(kit, real_manager, attach=attach): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + # Grown ahead, so the held section only lends. + assert mgr._resize_for_served_prefix(target, kit.kv(mgr, target), 0, END) + lender = attach(mgr, tokens_per_fetch=END) # one fetch: three slots + with kit.held_stream(mgr._stream) as gate: + read = lender.lend_read(source, 0, END) # every slot, its copy held + read.release() + write = lender.lend_write(target, 0, END) + assert write.poll() is None, "slots handed out while a copy into them was queued" + gate.open() + view = ready_view(write, mgr) + write.mark_arrived(view.row_masks()) + write.release() + + +def test_leases_wait_for_slots_strictly_in_order(kit, real_manager, attach=attach): + """Waiting leases are granted in order, by the lender's own calls, without new lends.""" + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) # three slots + first = lender.lend_read(source, 0, 2 * TPB) + ready_view(first, mgr) + second = lender.lend_read(source, 0, END) + third = lender.lend_read(source, 0, TPB) + assert second.poll() is None + assert third.poll() is None, "a later lease that fits jumped the one waiting ahead" + first.release() + ready_view(second, mgr) + assert third.poll() is None + second.release() + assert by_group(ready_view(third, mgr)) == {0: [0]} + third.release() + + +def test_planned_fetches_is_a_budget_and_holes_can_make_a_lease_wait(kit, real_manager): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + (group,) = kit.pool_group_ids(mgr) + (region,) = lender.regions + assert region.slot_count == 2 * BLOCKS + first = lender.lend_read(source, 0, 2 * TPB) # slots 0 and 1 + second = lender.lend_read(source, 0, TPB) # slot 2 + third = lender.lend_read(source, 0, TPB) # slot 3 + for lease in (first, second, third): + ready_view(lease, mgr) + first.release() + assert lender._free_slots(group) == 4 # slots 0, 1, 4 and 5 + fetch = lender.lend_read(source, 0, END) + assert fetch.poll() is None and fetch.failure is None, "no three free slots in a row" + second.release() # slots 0 to 2 are free in a row + (run,) = ready_view(fetch, mgr).groups + assert ((run.addresses - region.address) // region.slot_bytes).tolist() == [0, 1, 2] + for lease in (third, fetch): + lease.release() + + +# -- fetch ------------------------------------------------------------------------------------ + + +@pytest.mark.parametrize("local_blocks", [0, 2], ids=["cold", "local_prefix"]) +def test_a_fetch_is_usable_once_its_copy_lands(kit, real_manager, local_blocks, attach=attach): + with real_manager() as mgr_a, real_manager() as mgr_b: + source = kit.published(mgr_a, SOURCE, PROMPT) + lender_a = attach(mgr_a, tokens_per_fetch=END) + publish = lender_a.lend_read(source, 0, END) + publish_view = ready_view(publish, mgr_a) + if local_blocks: + kit.published(mgr_b, 7, PROMPT[: local_blocks * TPB] + list(range(9000, 9033))) + target = kit.admitted(mgr_b, TARGET, PROMPT) + kv = kit.kv(mgr_b, target) + local, history = kv.num_committed_tokens, kv.history_length + assert local == local_blocks * TPB + lender_b = attach(mgr_b, tokens_per_fetch=END) + assert lender_b.resume_bounds(target) == (max(local, history), history) + + lease = lender_b.lend_write(target, local, END) + assert lender_b.resume_bounds(target) is None, "a fetch counts from its call" + assert kv.capacity >= END and kv.history_length == history + view = lease.poll() + assert view is not None, "a write is ready at its first poll after the grant" + assert by_group(view) == {0: list(range(local_blocks, BLOCKS))} + kit.fill_sentinel(mgr_b, target, first_block=local_blocks) + masks = kit.relay(lender_a, publish_view, lender_b, view) + with kit.held_stream(mgr_b._stream) as gate: + lease.mark_arrived(masks) + assert lender_b.resume_bounds(target) is None, "settled before the copy into pages ran" + gate.open() + resume_bounds = lender_b.resume_bounds(target) + assert isinstance(resume_bounds, ResumeBounds) and all( + type(v) is int for v in resume_bounds + ) + assert resume_bounds == (END, min(local, history)) + for ordinal in range(BLOCKS): + got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) + assert got == kit.digest(kit.page(mgr_a, source, 0, ordinal)), "a fetched page differs" + lease.release() + publish.release() + + +@pytest.mark.parametrize( + "rows,usable", + [((0, 1), 64), ((0, 2), 32), ((), 0), ((1, 2), 0)], + ids=["prefix", "gap", "nothing", "no_first_block"], +) +def test_a_partial_arrival_is_usable_up_to_its_first_gap(kit, real_manager, rows, usable): + with real_manager() as mgr_a, real_manager() as mgr_b: + source = kit.published(mgr_a, SOURCE, PROMPT) + lender_a = attach(mgr_a, tokens_per_fetch=END) + publish = lender_a.lend_read(source, 0, END) + publish_view = ready_view(publish, mgr_a) + target = kit.admitted(mgr_b, TARGET, PROMPT) + lender_b = attach(mgr_b, tokens_per_fetch=END) + lease = lender_b.lend_write(target, 0, END) + view = lease.poll() + kit.fill_sentinel(mgr_b, target) + kit.relay(lender_a, publish_view, lender_b, view, deliver=lambda run, row: row in rows) + lease.mark_arrived((np.isin(np.arange(BLOCKS), rows),)) + mgr_b._stream.synchronize() + assert lender_b.resume_bounds(target) == (usable, 0) + sentinel = bytes([kit.SENTINEL]) * kit.DevicePages(mgr_b).page_bytes(0) + for ordinal in range(BLOCKS): + expected = kit.page(mgr_a, source, 0, ordinal) if ordinal in rows else sentinel + got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) + assert got == kit.digest(expected), "an unmarked row was copied" + lease.release() + publish.release() + + +@pytest.mark.parametrize("partial", [False, True], ids=["complete", "partial"]) +def test_a_windowed_fetch_is_usable_only_from_its_floor(kit, real_manager, partial): + windows = [WINDOW, 256] + with real_manager(windows=windows) as mgr_a, real_manager(windows=windows) as mgr_b: + source = kit.published(mgr_a, SOURCE, WINDOWED_PROMPT) + lender_a = attach(mgr_a, tokens_per_fetch=WINDOWED_END) + publish = lender_a.lend_read(source, 0, WINDOWED_END) + publish_view = ready_view(publish, mgr_a) + target = kit.admitted(mgr_b, TARGET, WINDOWED_PROMPT) + lender_b = attach(mgr_b, tokens_per_fetch=WINDOWED_END) + lease = lender_b.lend_write(target, 0, WINDOWED_END) + assert kit.kv(mgr_b, target).history_length == WINDOWED_END, "the history moves to the end" + view = lease.poll() + assert by_group(view) == by_group(publish_view) + kit.fill_sentinel(mgr_b, target) + deliver = (lambda run, row: row < 2) if partial else None + masks = kit.relay(lender_a, publish_view, lender_b, view, deliver) + lease.mark_arrived(masks) + mgr_b._stream.synchronize() + resume_bounds = lender_b.resume_bounds(target) + assert resume_bounds.min_position == WINDOWED_END, "the window released the early blocks" + if partial: + assert resume_bounds.max_position < resume_bounds.min_position, "empty: compute from 0" + for run in view.groups: # the two rows of each run that arrived hold the source's bytes + for ordinal in run.block_ordinals.tolist()[:2]: + lg = run.layer_group_index + got = kit.digest(kit.page(mgr_b, target, lg, ordinal)) + assert got == kit.digest(kit.page(mgr_a, source, lg, ordinal)) + else: + assert resume_bounds.max_position == WINDOWED_END + for run in view.groups: + for ordinal in run.block_ordinals.tolist(): + lg = run.layer_group_index + got = kit.digest(kit.page(mgr_b, target, lg, ordinal)) + assert got == kit.digest(kit.page(mgr_a, source, lg, ordinal)) + lease.release() + publish.release() + + +def test_mark_arrived_takes_one_mask_per_run_once_on_a_write_seen_ready(kit, real_manager): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + read = lender.lend_read(source, 0, END) + read_view = ready_view(read, mgr) + with pytest.raises(RuntimeError): + read.mark_arrived(read_view.row_masks()) + write = lender.lend_write(target, 0, END) + with pytest.raises(RuntimeError): + write.mark_arrived(()) # before poll() gave the view + view = write.poll() + bad = ( + (), + (np.ones(BLOCKS - 1, bool),), + (np.ones(BLOCKS, bool),) * 2, + (np.ones(BLOCKS, np.int64),), # one per row, but not bool + ) + for masks in bad: + with pytest.raises(ValueError): + write.mark_arrived(masks) + assert lender.resume_bounds(target) is None, "a refused mask recorded something" + write.mark_arrived(view.row_masks()) + with pytest.raises(RuntimeError): + write.mark_arrived(view.row_masks()) + mgr._stream.synchronize() + kv = kit.kv(mgr, target) + assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + write.release() + read.release() + + +@pytest.mark.parametrize( + "mark_first", [True, False], ids=["mark_then_release", "release_then_mark"] +) +def test_either_order_of_marks_and_release_copies_the_same(kit, real_manager, mark_first): + with real_manager() as mgr: + first = kit.admitted(mgr, TARGET, PROMPT) + second = kit.admitted(mgr, OTHER, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END) # the second waits for the first's slots + lease = lender.lend_write(first, 0, END) + view = lease.poll() + kit.fill_sentinel(mgr, first) + kit.stage(lender, view, 0x5A) + waiting = lender.lend_write(second, 0, END) + if not mark_first: + lease.release() + assert waiting.poll() is None, "slots returned before the marks" + with kit.held_stream(mgr._stream) as gate: + lease.mark_arrived(view.row_masks(True)) + lease.release() + assert waiting.poll() is None, "slots returned before the copy out of them ran" + gate.open() + waiting_view = ready_view(waiting, mgr) + nbytes = kit.DevicePages(mgr).page_bytes(0) + got = kit.digest([kit.page(mgr, first, 0, o) for o in range(BLOCKS)]) + assert got == kit.digest([kit.staged_row(0x5A, 0, o, nbytes) for o in range(BLOCKS)]) + assert lender.resume_bounds(first) == (END, 0) + waiting.mark_arrived(waiting_view.row_masks()) + waiting.release() + + +# -- the request exits during a lease --------------------------------------------------------- + + +def test_arrivals_for_a_freed_request_copy_nothing(kit, real_manager, attach=attach): + with real_manager(max_tokens=kit.POOL_TOKENS) as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + lease = lender.lend_write(target, 0, END) + view = lease.poll() + assert view is not None + kv = kit.kv(mgr, target) + lent = {kit.pages(kv, 0)[o] for o in view.groups[0].block_ordinals.tolist()} + kit.stage(lender, view, 0x5A) + mgr.free_resources(target) + with pytest.raises(ValueError): + lender.resume_bounds(target) + others = kit.Requests(mgr) + try: + assert others.allocate(kit.pool_pages(mgr)) + assert lent <= others.pages(), "no lent page was reused: the check proves nothing" + lease.mark_arrived(view.row_masks(True)) + mgr._stream.synchronize() + dev = kit.DevicePages(mgr) + sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) + for slot in sorted(lent): + got = kit.digest(dev.read(0, slot)) + assert got == kit.digest(sentinel), "an arrival overwrote another request's page" + finally: + lease.release() + others.free() + + +def commit_fetched(mgr, request, end): + """The manager commits ``request``'s blocks up to ``end``, as after a context chunk.""" + request.context_chunk_size = end - request.context_current_position + request.move_to_next_context_chunk() + mgr.try_commit_blocks(request) + + +def test_a_commit_during_a_fetch_copies_nothing_into_pages_it_returned( + kit, real_manager, attach=attach +): + """A commit while a write is open rebases the target onto blocks another request committed and + returns the lent pages to the pool: arrivals reach neither those pages nor the shared ones.""" + with real_manager(max_tokens=kit.POOL_TOKENS, max_batch_size=16) as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + lease = lender.lend_write(target, 0, END) + view = lease.poll() + assert view is not None + kv = kit.kv(mgr, target) + lent = {kit.pages(kv, 0)[o] for o in view.groups[0].block_ordinals.tolist()} + kit.stage(lender, view, 0x5A) + source = kit.published(mgr, SOURCE, PROMPT) + shared = kit.pages(kit.kv(mgr, source), 0)[:BLOCKS] + committed_bytes = kit.digest([kit.page(mgr, source, 0, o) for o in range(BLOCKS)]) + commit_fetched(mgr, target, END) + assert kv.num_committed_tokens == END + assert kit.pages(kv, 0)[:BLOCKS] == shared, "no rebase: the check proves nothing" + others = kit.Requests(mgr) + try: + while others.allocate(1, chunk=1): + pass + assert lent <= others.pages(), "no returned page was reused: the check proves nothing" + lease.mark_arrived(view.row_masks(True)) + mgr._stream.synchronize() + dev = kit.DevicePages(mgr) + sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) + for slot in sorted(lent): + got = kit.digest(dev.read(0, slot)) + assert got == kit.digest(sentinel), "an arrival overwrote another request's page" + got = kit.digest([kit.page(mgr, source, 0, o) for o in range(BLOCKS)]) + assert got == committed_bytes, "an arrival overwrote the blocks the target rebased onto" + finally: + lease.release() + others.free() + + +@pytest.mark.parametrize("change", ["freed", "suspended"]) +def test_a_waiting_read_fails_when_its_cache_changes(kit, real_manager, change, attach=attach): + """A read waiting for slots fails at once when its request is freed, and at its grant when its + cache was suspended; either way its slots come back.""" + with real_manager() as mgr: + holder = kit.published(mgr, SOURCE, PROMPT) + other = kit.published(mgr, OTHER, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + head = lender.lend_read(holder, 0, END) + ready_view(head, mgr) + waiting = lender.lend_read(other, 0, 2 * TPB) + assert waiting.poll() is None and waiting.failure is None + if change == "freed": + mgr.free_resources(other) + assert waiting.failure is not None, "a lease waiting for slots outlived its request" + else: + mgr.suspend_request(other) + head.release() + assert waiting.poll() is None and waiting.failure is not None, "copied a moved cache" + waiting.release() + again = lender.lend_read(holder, 0, END) # its slots came back + ready_view(again, mgr) + again.release() + + +def test_a_publish_granted_before_its_request_is_freed_copies_what_it_lent(kit, real_manager): + # No host tier: the freed pages go straight to the other requests. With one, they would move + # to host first, and the first such move in a process waits for the whole device: behind the + # held stream, that is until the watchdog opens it. + with real_manager(max_tokens=kit.POOL_TOKENS, host_cache_size=0) as mgr: + # A process's first cache close waits for the whole device, the gated stream included: + # close one before the gate. + mgr.free_resources(kit.published(mgr, 900, list(range(40_000, 40_000 + TPB + 1)))) + torch.cuda.synchronize() + source = kit.published(mgr, SOURCE, PROMPT) + original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] + lent = set(kit.pages(kit.kv(mgr, source), 0)[:BLOCKS]) + lender = attach(mgr, tokens_per_fetch=END) + dev = kit.DevicePages(mgr) # built ahead: building it waits for the device + others = kit.Requests(mgr, dev) + with kit.held_stream(mgr._stream, strict=False) as gate: + lease = lender.lend_read(source, 0, END) # granted, its copy queued behind the gate + mgr.free_resources(source) + assert lease.failure is None and kit.kv(mgr, source) is None + # Other requests take every page and overwrite it, on the stream behind the copy. + assert others.allocate(kit.pool_pages(mgr), fill=False) + assert lent <= others.pages(), "no freed page was taken: the check proves nothing" + gate.open() + sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) + assert all(dev.read(0, slot) == sentinel for slot in lent), "the pages were not overwritten" + (run,) = ready_view(lease, mgr).groups + length = lender.regions[run.region_index].slot_bytes + got = kit.digest([kit.host_bytes(a, length) for a in run.addresses.tolist()]) + assert got == kit.digest(original), "the copy read a later owner's writes" + lease.release() + others.free() + + +# -- abandoned fetches ------------------------------------------------------------------------ + + +@pytest.mark.parametrize("waiting", [False, True], ids=["granted", "waiting_in_line"]) +@pytest.mark.parametrize("windowed", [False, True], ids=["full", "windowed"]) +def test_an_abandoned_fetch_leaves_only_its_growth( + kit, real_manager, windowed, waiting, attach=attach +): + tokens, other, end = ( + (WINDOWED_PROMPT, OTHER_WINDOWED_PROMPT, WINDOWED_END) + if windowed + else (PROMPT, OTHER_PROMPT, END) + ) + with real_manager(windows=[WINDOW, 256] if windowed else None) as mgr: + target = kit.admitted(mgr, TARGET, tokens) + second = kit.admitted(mgr, OTHER, other) + lender = attach(mgr, tokens_per_fetch=end) + head = None + if waiting: # a publish takes every slot first, so the write waits in line + holder = kit.published(mgr, SOURCE, list(range(8000, 8000 + len(tokens)))) + head = lender.lend_read(holder, 0, end) + ready_view(head, mgr) + lease = lender.lend_write(target, 0, end) + assert lender.resume_bounds(target) is None + if waiting: + assert lease.poll() is None and lease.failure is None, "the write did not wait" + lease.release() # before any grant or view: no backend wrote, the fetch is abandoned + if head is not None: + head.release() # the slots come back; the released write takes none of them + kv = kit.kv(mgr, target) + own = (kv.num_committed_tokens, kv.history_length) + assert lender.resume_bounds(target) == own, "the abandoned fetch still counts" + if windowed: + assert kv.history_length == end, "the growth moved the window's history" + before = state_of(kit, mgr, lender, target) + with pytest.raises(ValueError, match="below the history"): + lender.lend_write(target, 0, END) + assert state_of(kit, mgr, lender, target) == before + again = lender.lend_write(second, 0, end) + view = again.poll() + assert view is not None, "the abandoned fetch kept its slots" + again.mark_arrived(view.row_masks()) + again.release() + + +# -- a fetch split into leases, and a shrink of the cache it went into ------------------------- + +SPLIT_PROMPT = list(range(3000, 3129)) # four whole blocks and one token +SMALL_POOL = dict(max_tokens=256, max_batch_size=64) + + +def stale_anywhere(kit, mgr, history): + """Some windowed layer group has released blocks behind its window at ``history``.""" + for lg, window in enumerate(kit.windows(mgr)): + if window is not None: + beg, end = kit.stale_blocks(mgr, lg, history) + if end > beg: + return True + return False + + +def fetch_segment(kit, mgr, lender, target, start, end, byte, arrived=None): + """One write lease over ``[start, end)``: every slot staged from ``byte``, each row with its + own bytes (``kit.staged_row``), the rows of block ordinals ``arrived`` (default all) marked, + released, the copy into pages run.""" + lease = lender.lend_write(target, start, end) + view = lease.poll() + assert view is not None, f"the write lease was not ready: {lease.failure}" + kit.stage(lender, view, byte) + masks = tuple( + np.ones(len(run), dtype=bool) if arrived is None else np.isin(run.block_ordinals, arrived) + for run in view.groups + ) + lease.mark_arrived(masks) + lease.release() + mgr._stream.synchronize() + + +SEGMENTS = { + # name: (manager kwargs, segments [(start, end, arrived)], max_position after each segment) + "two_halves": ({}, [(0, 64, None), (64, 128, None)], [64, 128]), + "three_segments": ({}, [(0, 32, None), (32, 96, None), (96, 128, None)], [32, 96, 128]), + "windowed": ({"windows": [WINDOW, 256]}, [(0, 64, None), (64, 128, None)], [64, 128]), + # A gap stays a gap across segments. + "first_landed_nothing": ({}, [(0, 64, []), (64, 128, None)], [0, 0]), + "first_landed_one_block": ({}, [(0, 64, [0]), (64, 128, None)], [32, 32]), +} + + +@pytest.mark.parametrize("case", list(SEGMENTS)) +def test_consecutive_leases_add_up_to_one_fetch(kit, real_manager, case, attach=attach): + """Consecutive leases into one cache, nothing committed in between: resume_bounds covers every + segment that landed, and nothing more.""" + manager_kwargs, segments, expected = SEGMENTS[case] + with real_manager(**manager_kwargs) as mgr: + target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) + kv = kit.kv(mgr, target) + assert kv.num_committed_tokens == 0, "no local prefix: every usable token was fetched" + lender = attach(mgr, tokens_per_fetch=64) + bytes_of = {} + for i, ((start, end, arrived), usable) in enumerate(zip(segments, expected)): + byte = 0x11 * (i + 1) + fetch_segment(kit, mgr, lender, target, start, end, byte, arrived) + for ordinal in range(start // TPB, end // TPB): + if arrived is None or ordinal in arrived: + bytes_of[ordinal] = byte + assert kv.num_committed_tokens == 0, "a fetch committed" + resume_bounds = lender.resume_bounds(target) + floor = kv.history_length if stale_anywhere(kit, mgr, kv.history_length) else 0 + assert resume_bounds is not None and resume_bounds.min_position == floor + assert resume_bounds.max_position >= usable, ( + f"after segment {i} [{start}, {end}): resume_bounds {tuple(resume_bounds)}; an earlier " + "segment's delivered prefix was dropped" + ) + assert resume_bounds.max_position == usable, ( + f"after segment {i}: resume_bounds {tuple(resume_bounds)} counts what never landed" + ) + # The bytes are in the pages: only the bookkeeping decides whether they count. A window's + # blocks behind the final history are released, and no resume reads them. + dev = kit.DevicePages(mgr) + for lg in range(kit.num_layer_groups(mgr)): + stale_beg, stale_end = kit.stale_blocks(mgr, lg, kv.history_length) + slots = kit.pages(kv, lg) + for ordinal, byte in bytes_of.items(): + if stale_beg <= ordinal < stale_end: + continue + assert slots[ordinal] >= 0, f"layer group {lg} lost block {ordinal}'s page" + got = kit.digest(dev.read(lg, slots[ordinal])) + staged = kit.digest(kit.staged_row(byte, lg, ordinal, dev.page_bytes(lg))) + assert got == staged, f"layer group {lg} block {ordinal} holds other bytes" + + +# name: the span of the lease after the first, two blocks long +EXTRA_TOKEN_SPLITS = {"window_and_a_block": 3 * TPB, "within_the_window": 2 * TPB} + + +@pytest.mark.parametrize("case", list(EXTRA_TOKEN_SPLITS)) +def test_a_later_lease_past_the_window_and_a_block_fails_with_extra_kv_tokens( + kit, real_manager, case, attach=attach +): + """All-reusable, a window of ``WINDOW`` tokens and PARD's extra KV token: the first lease's grow + gives block 2, which holds that token, a page past the history the lease leaves, and nothing + writes it. A consecutive lease spanning at least ``WINDOW + TPB - 1`` tokens leaves that block + behind, so it fails at the call and changes nothing in the cache; a shorter one is lent.""" + span = EXTRA_TOKEN_SPLITS[case] + with real_manager(spec_config=pard(), windows=[WINDOW, 256]) as mgr: + assert mgr.num_extra_kv_tokens == 1, "no extra KV token: nothing checked" + target = kit.admitted(mgr, TARGET, WINDOWED_PROMPT) + lender = attach(mgr, tokens_per_fetch=3 * TPB) + fetch_segment(kit, mgr, lender, target, 0, 2 * TPB, 0x5A) + before = state_of(kit, mgr, lender, target) + assert before[1].max_position == 2 * TPB, f"resume_bounds after the first lease {before[1]}" + sliding = kit.windows(mgr).index(WINDOW) + assert kit.pages(kit.kv(mgr, target), sliding)[2] >= 0, ( + "block 2 has no page: nothing checked" + ) + lease = lender.lend_write(target, 2 * TPB, 2 * TPB + span) + if span >= WINDOW + TPB - 1: + assert lease.failure is not None and "never wrote" in lease.failure, ( + f"a later lease of {span} tokens was lent past an unwritten page: {lease.failure}" + ) + assert state_of(kit, mgr, lender, target) == before, ( + "the failed lease changed the cache" + ) + else: + view = ready_view(lease, mgr) + kit.stage(lender, view, 0x6B) + lease.mark_arrived(view.row_masks(True)) + mgr._stream.synchronize() + for lg in range( + kit.num_layer_groups(mgr) + ): # what resume_bounds counts, the window's alone + fetched = [([2, 3], 0x6B)] if lg == sliding else [([0, 1], 0x5A), ([2, 3], 0x6B)] + for ordinals, byte in fetched: + assert holds(kit, mgr, target, lg, ordinals, staged_bytes(kit, byte)), ( + f"layer group {lg}: fetched blocks {ordinals} lack their staged bytes" + ) + end = 2 * TPB + span + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == (end, end), ( + f"resume_bounds {tuple(resume_bounds)} after the later lease" + ) + lease.release() + + +def test_an_abandoned_segment_keeps_what_earlier_ones_delivered(kit, real_manager, attach=attach): + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) + lender = attach(mgr, tokens_per_fetch=64) + fetch_segment(kit, mgr, lender, target, 0, 64, 0x11) + assert lender.resume_bounds(target) == (64, 0) + second = lender.lend_write(target, 64, 128) + assert lender.resume_bounds(target) is None, "a fetch counts from its call" + second.release() # before anyone saw it ready: abandoned, nothing written + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == (64, 0), ( + f"resume_bounds {tuple(resume_bounds)} after an abandoned segment; the first segment's " + "delivered prefix was dropped" + ) + assert holds(kit, mgr, target, 0, [0, 1], staged_bytes(kit, 0x11)), ( + "the first segment's rows do not hold the bytes staged for them" + ) + + +LOST_SEGMENTS = { + # name: (the first lease's end, the later lease's end, resume_bounds expected; None: empty) + "window_keeps_the_blocks": (TPB, 2 * TPB, (TPB, 0)), + "window_released_blocks": (2 * TPB, 4 * TPB, None), +} + + +@pytest.mark.parametrize("lost", ["abandoned", "nothing_arrived"]) +@pytest.mark.parametrize("case", list(LOST_SEGMENTS)) +def test_a_lost_later_segment_empties_the_interval_only_where_a_window_released_blocks( + kit, real_manager, case, lost, attach=attach +): + """All-reusable, a window of 64 tokens: a fetch in two leases, the first delivered whole, the + later one lost after moving the history to its end. Where no window released blocks there, the + request may resume below that end, up to where the first lease's rows reach, and its next chunk + runs from there; where one did, the floor is that end, which no row reaches.""" + first_end, later_end, expected = LOST_SEGMENTS[case] + with real_manager(windows=[WINDOW, 256]) as mgr: + target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) + kv = kit.kv(mgr, target) + lender = attach(mgr, tokens_per_fetch=2 * TPB) + fetch_segment(kit, mgr, lender, target, 0, first_end, 0x11) + for lg in range(kit.num_layer_groups(mgr)): + assert holds(kit, mgr, target, lg, range(first_end // TPB), staged_bytes(kit, 0x11)) + later = lender.lend_write(target, first_end, later_end) + assert kv.history_length == later_end, "the later lease did not move the history" + if lost == "abandoned": + later.release() # before its first poll + else: + later.mark_arrived(later.poll().row_masks(False)) + later.release() + mgr._stream.synchronize() + resume_bounds = lender.resume_bounds(target) + assert stale_anywhere(kit, mgr, later_end) == (expected is None) + if expected is None: + assert resume_bounds.min_position == later_end > resume_bounds.max_position, ( + f"resume_bounds {tuple(resume_bounds)} claims a start whose window blocks were released" + ) + return + assert resume_bounds == expected, ( + f"resume_bounds {tuple(resume_bounds)}: no resume below the later lease's end {later_end}, " + "where no window released blocks" + ) + resumed(kit, mgr, target, resume_bounds.max_position) + assert next_chunk(kit, mgr, target) is None + + +def tp_rank(rank): + """Rank ``rank`` of two tensor-parallel ranks.""" + from tensorrt_llm.mapping import Mapping + + return Mapping(world_size=2, tp_size=2, rank=rank) + + +RAISED_ON_ONE_RANK = "a ValueError on one rank only" + + +def test_whether_an_earlier_fetch_settled_is_each_rank_s_own(kit, real_manager, attach=attach): + """Two managers stand for the two tensor-parallel ranks of one request, which a caller fetches + in two leases. Both ranks mark the first; rank 1's copy into its pages still waits behind + earlier work on its stream when both lend the second. Rank 0 gets the lease and rank 1 a lease + failed at the call, which changed nothing: an outcome the caller combines, where a ValueError + on one rank alone would leave that rank's loop while the other waits for it. Lent only once + every rank's resume_bounds is not None, the second lease is granted on both.""" + with ( + real_manager(num_kv_heads=4, mapping=tp_rank(0)) as rank0, + real_manager(num_kv_heads=4, mapping=tp_rank(1)) as rank1, + ): + ranks = (rank0, rank1) + targets = [kit.admitted(mgr, TARGET, SPLIT_PROMPT) for mgr in ranks] + lenders = [attach(mgr, tokens_per_fetch=2 * TPB, planned_fetches=2) for mgr in ranks] + firsts = [lender.lend_write(t, 0, 2 * TPB) for lender, t in zip(lenders, targets)] + views = [ready_view(lease, mgr) for lease, mgr in zip(firsts, ranks)] + for lender, view in zip(lenders, views): + kit.stage(lender, view, 0x5A) # the backend filled the first segment on both ranks + firsts[0].mark_arrived(views[0].row_masks(True)) + rank0._stream.synchronize() # rank 0's copy into its pages has run + seconds, outcomes = [], [] + with kit.held_stream(rank1._stream, strict=False): + firsts[1].mark_arrived(views[1].row_masks(True)) # rank 1's copy stays queued + before = [ + state_of(kit, mgr, lender, t) for mgr, lender, t in zip(ranks, lenders, targets) + ] + for lender, target in zip(lenders, targets): + try: + lease = lender.lend_write(target, 2 * TPB, 4 * TPB) + except ValueError as error: + outcomes.append(f"ValueError: {error}") + continue + seconds.append(lease) + outcomes.append("lease" if lease.failure is None else "failed at the call") + after = state_of(kit, rank1, lenders[1], targets[1]) + try: + assert before[0][1] is not None and before[1][1] is None, ( + f"resume_bounds {before}: rank 1's copy was not held, so this proves nothing" + ) + raised = [outcome.startswith("ValueError") for outcome in outcomes] + assert not any(raised), f"{RAISED_ON_ONE_RANK}: {outcomes}" + assert outcomes == ["lease", "failed at the call"], outcomes + assert after == before[1], "the lease failed at the call changed rank 1's cache" + for lease in seconds: + lease.release() + seconds = [] + # The caller lends the next segment once every rank's resume_bounds is not None. + for mgr in ranks: + mgr._stream.synchronize() + assert all(lender.resume_bounds(t) is not None for lender, t in zip(lenders, targets)) + for rank, (mgr, target) in enumerate(zip(ranks, targets)): + assert holds(kit, mgr, target, 0, [0, 1], staged_bytes(kit, 0x5A)), ( + f"rank {rank}: the first segment's rows landed off their blocks" + ) + seconds = [ + lender.lend_write(t, 2 * TPB, 4 * TPB) for lender, t in zip(lenders, targets) + ] + assert [lease.failure for lease in seconds] == [None, None] + for lease, mgr in zip(seconds, ranks): + lease.mark_arrived(ready_view(lease, mgr).row_masks()) + finally: + for mgr in ranks: + mgr._stream.synchronize() + for lease in [*firsts, *seconds]: + lease.release() + + +FRESH_FILL = 2.5 # every fp16 element 0x4100: no staged byte pattern reads as it + + +def test_a_resume_with_the_fresh_page_fill_on_keeps_the_fetched_blocks( + kit, real_manager, monkeypatch +): + """The manager's fresh-page fill writes pages as a request first gets them. The pages a fetch + grew are the request's own once the fetch lands: resuming fills only pages past them.""" + monkeypatch.setenv("TRTLLM_KV_FRESH_PAGE_FILL", str(FRESH_FILL)) + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + fetch_segment(kit, mgr, lender, target, 0, 64, 0x5A) + assert lender.resume_bounds(target).max_position == 64 + # The executor resumes the request where the fetch left it usable. + target.context_chunk_size = target.prompt_len - target.context_current_position + target.set_prepopulated_prompt_len(64, TPB) + target.context_chunk_size = target.prompt_len - 64 + assert mgr.resize_context(target, target.context_remaining_length) + mgr._stream.synchronize() + dev = kit.DevicePages(mgr) + fetched = [kit.digest(kit.staged_row(0x5A, 0, o, dev.page_bytes(0))) for o in range(2)] + filled = np.full(dev.page_bytes(0) // 2, FRESH_FILL, dtype=np.float16).tobytes() + got = [kit.digest(kit.page(mgr, target, 0, ordinal)) for ordinal in range(3)] + assert got[2] == kit.digest(filled), "the fill did not run on the resume" + assert got[:2] == fetched, "the resume filled the fetched blocks as fresh pages" + mgr.free_resources(target) + + +def test_a_draft_pool_fetch_fills_only_the_pages_its_grow_added(kit, real_manager, monkeypatch): + """A joint-reuse draft pool's own context resizes run no fresh-page fill. Under per_request the + draft computed ``[0, LOCAL)`` and committed nothing; a fetch of the rest through the draft + pool's lender fills the pages its grow added and leaves the blocks the draft computed.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + monkeypatch.setenv("TRTLLM_KV_FRESH_PAGE_FILL", str(FRESH_FILL)) + spec = dflash() + with real_manager( + spec_config=spec, is_draft=True, joint_kv_cache_reuse=True, **PER_REQUEST + ) as draft: + request = kit.admitted(draft, TARGET, SPLIT_PROMPT) + kv = kit.kv(draft, request) + dev = kit.DevicePages(draft) + # The draft pool's first chunk as the executor runs it: admission, forward, update. + request.context_chunk_size = LOCAL + assert draft.try_allocate_draft_context(request, LOCAL) + for lg in range(kit.num_layer_groups(draft)): + for ordinal, slot in enumerate(kit.pages(kv, lg)): + if ordinal < LOCAL // TPB and slot >= 0: + dev.write(lg, slot, chunk_bytes(dev, lg, ordinal)) + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + draft.update_context_resources(batch) + assert (kv.history_length, kv.num_committed_tokens) == (LOCAL, 0), "the chunk committed" + assert request.py_request_id not in draft._fresh_pages_filled, ( + "the draft pool's own resizes ran the fill: the check proves nothing" + ) + lender = attach(draft, tokens_per_fetch=2 * TPB, scope=SCOPE + b"/draft-pool") + blocks_before = kv.num_blocks + lease = lender.lend_write(request, LOCAL, 4 * TPB) + view = ready_view(lease, draft) + torch.cuda.synchronize() # the fill runs on the current stream + filled = kit.digest(np.full(dev.page_bytes(0) // 2, FRESH_FILL, np.float16).tobytes()) + added = range(blocks_before, 4) + assert len(added) and all( + kit.digest(kit.page(draft, request, 0, o)) == filled for o in added + ), "the fill did not run on the pages the fetch grew" + assert holds(kit, draft, request, 0, range(LOCAL // TPB), chunk_bytes), ( + "the fetch's fill overwrote the blocks the draft computed" + ) + kit.stage(lender, view, 0x5A) + lease.mark_arrived(view.row_masks(True)) + lease.release() + draft._stream.synchronize() + assert holds(kit, draft, request, 0, [2, 3], staged_bytes(kit, 0x5A)) + assert lender.resume_bounds(request) == (4 * TPB, LOCAL) + + +def test_a_second_draft_pool_fetch_fills_only_the_pages_its_own_grow_added( + kit, real_manager, monkeypatch +): + """Under per_request, a first fetch through a joint-reuse draft pool's lender leaves the + fresh-page fill a record of the pages it grew. The draft then computes a chunk through its own + context resize, which runs no fill and leaves that record stale. A second fetch fills only the + pages its own grow added and leaves the chunk the draft computed.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + monkeypatch.setenv("TRTLLM_KV_FRESH_PAGE_FILL", str(FRESH_FILL)) + spec = dflash() + with real_manager( + spec_config=spec, is_draft=True, joint_kv_cache_reuse=True, **PER_REQUEST + ) as draft: + request = kit.admitted(draft, TARGET, LONG_PROMPT) + kv = kit.kv(draft, request) + dev = kit.DevicePages(draft) + lender = attach(draft, tokens_per_fetch=2 * TPB, scope=SCOPE + b"/draft-pool") + fetch_segment(kit, draft, lender, request, 0, LOCAL, 0x5A) + assert lender.resume_bounds(request).max_position == LOCAL + # The draft pool computes blocks 2 to 4 as the executor runs it after the resume. + chunk = range(LOCAL // TPB, 5) + resumed(kit, draft, request, LOCAL) + request.context_chunk_size = len(chunk) * TPB + assert draft.try_allocate_draft_context(request, len(chunk) * TPB) + for lg in range(kit.num_layer_groups(draft)): + slots = kit.pages(kv, lg) + for ordinal in chunk: + if slots[ordinal] >= 0: + dev.write(lg, slots[ordinal], chunk_bytes(dev, lg, ordinal)) + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + draft.update_context_resources(batch) + assert (kv.history_length, kv.num_committed_tokens) == (5 * TPB, 0), "the chunk committed" + record = draft._fresh_pages_filled.get(request.py_request_id, {}) + known = min((len(pages) for pages in record.values()), default=kv.num_blocks) + assert known < chunk[-1] + 1, ( + f"the fill's record covers {known} blocks, the draft's whole chunk: nothing to check" + ) + blocks_before = kv.num_blocks + lease = lender.lend_write(request, 5 * TPB, 7 * TPB) + view = ready_view(lease, draft) + torch.cuda.synchronize() # the fill runs on the current stream + assert holds(kit, draft, request, 0, chunk, chunk_bytes), ( + "the second fetch's fill overwrote the chunk the draft computed" + ) + filled = kit.digest(np.full(dev.page_bytes(0) // 2, FRESH_FILL, np.float16).tobytes()) + grown = range(blocks_before, kv.num_blocks) + assert len(grown) and all( + kit.digest(kit.page(draft, request, 0, o)) == filled for o in grown + ), "the fill did not run on the pages the second fetch grew" + kit.stage(lender, view, 0x6B) + lease.mark_arrived(view.row_masks(True)) + lease.release() + draft._stream.synchronize() + assert holds(kit, draft, request, 0, [0, 1], staged_bytes(kit, 0x5A)) + assert holds(kit, draft, request, 0, [5, 6], staged_bytes(kit, 0x6B)) + assert lender.resume_bounds(request) == (7 * TPB, 5 * TPB) + + +@pytest.mark.parametrize("taker", ["context", "fetch"]) +def test_the_fresh_page_fill_lets_a_freed_publisher_s_copy_run_first( + kit, real_manager, monkeypatch, taker, fill_liar=None +): + """A publish queues its copy into staging on the manager's stream, then its request may exit. + With the fresh-page fill on, a request taking its pages by the context path or a fetch's grow + fills them on the current stream after that copy, so staging holds the publisher's bytes.""" + monkeypatch.setenv("TRTLLM_KV_FRESH_PAGE_FILL", str(FRESH_FILL)) + # No host tier: the freed pages go straight to the next request. + with real_manager(max_tokens=kit.POOL_TOKENS, host_cache_size=0, max_batch_size=16) as mgr: + # A process's first cache close waits for the whole device, the gated stream included: + # close one before the gate. + mgr.free_resources(kit.published(mgr, 900, list(range(40_000, 40_000 + TPB + 1)))) + torch.cuda.synchronize() + source = kit.published(mgr, SOURCE, PROMPT) + original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] + lent = set(kit.pages(kit.kv(mgr, source), 0)[:BLOCKS]) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + fetching = kit.admitted(mgr, TARGET, OTHER_PROMPT) if taker == "fetch" else None + others = kit.Requests(mgr) + # Other requests take every free page. Their fills load the fill's kernels before the + # gate: a first load waits for the whole device. + while others.allocate(1, chunk=1): + pass + fill = mgr._fill_fresh_kv_pages + if fill_liar is not None: + fill = fill_liar(fill) + began = [] # per fill: whether the gate had opened when it began + + def reject_device_sync(*args, **kwargs): + raise AssertionError("the diagnostic fill waited for the whole device") + + def recording_fill(request_id): + unrelated_stream = torch.cuda.Stream() + with kit.held_stream(unrelated_stream, strict=False) as unrelated_gate: + with monkeypatch.context() as guard: + guard.setattr(torch.cuda, "synchronize", reject_device_sync) + began.append(gate.opened) + fill(request_id) + assert not unrelated_gate.opened, "the fill waited for an unrelated stream" + + mgr._fill_fresh_kv_pages = recording_fill # dropped below: it holds the manager + write = None + try: + with kit.gated_stream(mgr._stream, open_after=10.0) as gate: + read = lender.lend_read(source, 0, END) # its copy waits behind the gate + mgr.free_resources(source) + if taker == "context": + taken = kit.make_request(TARGET, OTHER_PROMPT) + assert mgr.prepare_context(taken) + assert mgr.resize_context(taken, taken.context_remaining_length) + else: + taken = fetching + write = lender.lend_write(taken, 0, END) + assert write.failure is None, write.failure + finally: + del mgr._fill_fresh_kv_pages + assert began and not began[0], "the copy ran before the fill: the check proves nothing" + assert lent & set(kit.pages(kit.kv(mgr, taken), 0)), "no freed page was taken" + (run,) = ready_view(read, mgr).groups + length = lender.regions[run.region_index].slot_bytes + got = kit.digest([kit.host_bytes(a, length) for a in run.addresses.tolist()]) + assert got == kit.digest(original), "the fill overwrote a page before its copy ran" + read.release() + if write is not None: + write.release() + mgr.free_resources(taken) + others.free() + + +@pytest.mark.parametrize("ask_between", [False, True], ids=["regrown_unasked", "asked_between"]) +def test_a_rollback_voids_what_the_fetch_delivered(kit, real_manager, ask_between, attach=attach): + """The manager's own context rollback undoes the growth ``lend_write`` made: the same cache + shrinks and frees the fetched pages, which others overwrite before the cache grows back. + ``ask_between``: resume_bounds is asked between the rollback and the regrow.""" + with real_manager() as mgr_a, real_manager(**SMALL_POOL) as mgr_b: + source = kit.published(mgr_a, SOURCE, PROMPT) + lender_a = attach(mgr_a, tokens_per_fetch=END) + publish = lender_a.lend_read(source, 0, END) + publish_view = ready_view(publish, mgr_a) + source_pages = [kit.digest(kit.page(mgr_a, source, 0, b)) for b in range(BLOCKS)] + target = kit.admitted(mgr_b, TARGET, PROMPT) + lender_b = attach(mgr_b, tokens_per_fetch=END) + kv = kit.kv(mgr_b, target) + lease = lender_b.lend_write(target, 0, END) + assert target.py_ctx_pre_resize_cap == 0, "the rollback would not undo the growth" + view = lease.poll() + kit.fill_sentinel(mgr_b, target) + lease.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) + lease.release() + publish.release() + mgr_b._stream.synchronize() + assert [kit.digest(kit.page(mgr_b, target, 0, b)) for b in range(BLOCKS)] == source_pages + assert lender_b.resume_bounds(target) == (END, 0) + + assert mgr_b.revert_allocate_context(target) is True + assert kit.kv(mgr_b, target) is kv, "the rollback replaced the cache" + assert kv.capacity == 0 and kit.pages(kv, 0) == [] and kv.num_committed_tokens == 0 + if ask_between: + resume_bounds = lender_b.resume_bounds(target) + assert resume_bounds == (0, 0), ( + f"resume_bounds {tuple(resume_bounds)} counts freed blocks" + ) + assert kit.overwrite_free_pages(mgr_b) >= BLOCKS + assert mgr_b.resize_context(target, END) + regrown = [kit.digest(kit.page(mgr_b, target, 0, b)) for b in range(BLOCKS)] + assert regrown != source_pages, "the regrown pages still hold the fetched bytes" + resume_bounds = lender_b.resume_bounds(target) + assert resume_bounds == (0, 0), ( + f"resume_bounds {tuple(resume_bounds)} counts blocks whose pages now hold {regrown}" + ) + mgr_b.free_resources(target) + + +def test_resume_bounds_voids_the_rows_of_a_shrink_no_hook_reported(kit, real_manager): + unreported = attach_breaking({"_on_shrink": lambda self, request_id, kv_cache: None}) + test_a_rollback_voids_what_the_fetch_delivered( + kit, real_manager, ask_between=True, attach=unreported + ) + + +def test_a_rollback_to_a_local_prefix_keeps_only_the_prefix(kit, real_manager, attach=attach): + """The rollback's other branch: a local match of two blocks, a fetch of the third; the rollback + shrinks to the prefix and suspends the cache.""" + local = 2 * TPB + with real_manager() as mgr_a, real_manager(**SMALL_POOL) as mgr_b: + source = kit.published(mgr_a, SOURCE, PROMPT) + lender_a = attach(mgr_a, tokens_per_fetch=END) + publish = lender_a.lend_read(source, 0, END) + publish_view = ready_view(publish, mgr_a) + kit.published(mgr_b, 7, PROMPT[:local] + list(range(9000, 9033))) + target = kit.admitted(mgr_b, TARGET, PROMPT) + lender_b = attach(mgr_b, tokens_per_fetch=END) + kv = kit.kv(mgr_b, target) + assert kv.num_committed_tokens == local + lease = lender_b.lend_write(target, local, END) + view = lease.poll() + kit.fill_sentinel(mgr_b, target, first_block=2) + lease.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) + lease.release() + publish.release() + mgr_b._stream.synchronize() + assert lender_b.resume_bounds(target) == (END, local) + got = kit.digest(kit.page(mgr_b, target, 0, 2)) + assert got == kit.digest(kit.page(mgr_a, source, 0, 2)), ( + "the fetched block 2 holds bytes other than the publisher's" + ) + assert mgr_b.revert_allocate_context(target) is True + assert kit.kv(mgr_b, target) is kv + assert (kv.capacity, kv.is_active, kv.num_committed_tokens) == (local, False, local) + resume_bounds = lender_b.resume_bounds(target) + assert resume_bounds == (local, local), ( + f"resume_bounds {tuple(resume_bounds)} counts block 2, which the rollback freed" + ) + mgr_b.free_resources(target) + + +GENERATION_SHRINKS = ["revert_allocate_generation", "update_resources"] +GENERATION_DRAFT = 2 * TPB - 1 # draft tokens the generation path grows the cache for + + +@pytest.mark.parametrize("path", GENERATION_SHRINKS) +def test_a_generation_shrink_voids_the_fetched_rows_past_the_capacity( + kit, real_manager, path, attach=attach +): + """The generation path grows the cache for a token and ``GENERATION_DRAFT`` draft tokens, and a + fetch of ``[0, END)`` grows it on to ``END``. ``revert_allocate_generation``, or the generation + update reclaiming every draft slot, shrinks the same cache and leaves fetched blocks past its + capacity; others overwrite the freed pages before the generation path grows the cache back.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + from tensorrt_llm.bindings import LlmRequestState + + with real_manager() as mgr_a, real_manager(**SMALL_POOL) as mgr_b: + source = kit.published(mgr_a, SOURCE, PROMPT) + lender_a = attach(mgr_a, tokens_per_fetch=END) + publish = lender_a.lend_read(source, 0, END) + publish_view = ready_view(publish, mgr_a) + source_pages = [kit.digest(kit.page(mgr_a, source, 0, b)) for b in range(BLOCKS)] + target = kit.admitted(mgr_b, TARGET, PROMPT) + target.py_draft_tokens = [0] * GENERATION_DRAFT + assert mgr_b.try_allocate_generation(target) + kv = kit.kv(mgr_b, target) + assert kv.capacity == 1 + GENERATION_DRAFT + lender_b = attach(mgr_b, tokens_per_fetch=END) + lease = lender_b.lend_write(target, 0, END) + view = lease.poll() + kit.fill_sentinel(mgr_b, target) + lease.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) + lease.release() + publish.release() + mgr_b._stream.synchronize() + assert [kit.digest(kit.page(mgr_b, target, 0, b)) for b in range(BLOCKS)] == source_pages + assert lender_b.resume_bounds(target) == (END, 0) + + if path == "revert_allocate_generation": + mgr_b.revert_allocate_generation(target) + else: + # The update moves the history to the request's last token, which no capacity goes + # below, unless KV compression manages the history. + mgr_b.kv_compression_manages_history = True + target.state = LlmRequestState.GENERATION_IN_PROGRESS + batch = ScheduledRequests() + batch.append_generation_request(target) + mgr_b.update_resources(batch) + assert kit.kv(mgr_b, target) is kv, "the shrink replaced the cache" + kept = kv.num_blocks + assert 0 < kept < BLOCKS, f"{kept} blocks left: the shrink freed no fetched block, or all" + assert kit.overwrite_free_pages(mgr_b) >= BLOCKS - kept + assert mgr_b.try_allocate_generation(target) + assert kv.num_blocks >= BLOCKS + regrown = [kit.digest(kit.page(mgr_b, target, 0, b)) for b in range(BLOCKS)] + assert regrown[:kept] == source_pages[:kept], "a block the shrink kept lost its bytes" + assert all(r != s for r, s in zip(regrown[kept:], source_pages[kept:])), ( + "a regrown page still holds the fetched bytes" + ) + resume_bounds = lender_b.resume_bounds(target) + assert resume_bounds == (kept * TPB, 0), ( + f"resume_bounds {tuple(resume_bounds)} counts blocks whose pages now hold {regrown[kept:]}" + ) + mgr_b.free_resources(target) + + +@pytest.mark.parametrize("path", GENERATION_SHRINKS) +def test_a_generation_shrink_no_hook_reported_is_caught(kit, real_manager, path): + """A lender that ignores the shrink hook, as under a manager that never calls it, counts the + rows the generation shrink freed once the cache grows back.""" + unreported = attach_breaking({"_on_shrink": lambda self, request_id, kv_cache: None}) + with pytest.raises(AssertionError, match="counts blocks whose pages now hold"): + test_a_generation_shrink_voids_the_fetched_rows_past_the_capacity( + kit, real_manager, path, attach=unreported + ) + pytest.fail("a lender that ignores the shrink hook went unnoticed") + + +# -- a prefix computed before a fetch and not committed --------------------------------------- + +LOCAL = 2 * TPB # the context chunk computed before the fetch +PER_REQUEST = {"block_reuse_config": {"policy": "per_request"}} + + +def computed(kit, mgr, request_id, tokens, upto=LOCAL, conversation=None, **inputs): + """A request carrying ``inputs`` whose first context chunk, ``[0, upto)``, ran and committed + nothing, as under a reuse policy other than all-reusable; block ``b`` of layer group ``lg`` + holds ``chunk_bytes(dev, lg, b)``.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + from tensorrt_llm.conversation_params import ConversationParams + + request = kit.make_request(request_id, tokens, **inputs) + if conversation is not None: + request.py_conversation_params = ConversationParams(conversation_id=conversation) + assert mgr.prepare_context(request) + kv = kit.kv(mgr, request) + kv.enable_swa_scratch_reuse = False # a fetch target needs it off + request.context_chunk_size = upto + assert mgr.resize_context(request, upto) + dev = kit.DevicePages(mgr) + for lg in range(kit.num_layer_groups(mgr)): + for ordinal, slot in enumerate(kit.pages(kv, lg)): + if slot >= 0: + dev.write(lg, slot, chunk_bytes(dev, lg, ordinal)) + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + mgr.update_context_resources(batch) + assert (kv.history_length, kv.num_committed_tokens) == (upto, 0), "the chunk was committed" + return request + + +def chunk_bytes(dev, lg, ordinal): + """What ``computed`` wrote into block ``ordinal`` of layer group ``lg``: bytes drawn from a + generator seeded by both, so the same block of another layer group, or a page whose pools or + pieces trade places, holds other bytes.""" + rng = np.random.default_rng([0x40, lg, ordinal]) + return rng.integers(0, 256, dev.page_bytes(lg), dtype=np.uint8).tobytes() + + +def staged_bytes(kit, byte): + """What ``fetch_segment`` staged from ``byte`` into block ``ordinal`` of ``lg``.""" + return lambda dev, lg, ordinal: kit.staged_row(byte, lg, ordinal, dev.page_bytes(lg)) + + +def holds(kit, mgr, request, lg, ordinals, expected): + """Whether blocks ``ordinals`` of ``lg`` hold ``expected(dev, lg, ordinal)`` in the request's + pages.""" + dev = kit.DevicePages(mgr) + got = [kit.digest(kit.page(mgr, request, lg, o)) for o in ordinals] + return got == [kit.digest(expected(dev, lg, o)) for o in ordinals] + + +OUTCOMES = ["abandoned", "delivered"] + + +@pytest.mark.parametrize("outcome", OUTCOMES) +@pytest.mark.parametrize("policy", ["per_request", "per_conversation"]) +def test_a_fetch_keeps_the_tokens_computed_before_it( + kit, real_manager, policy, outcome, attach=attach +): + """A chunk computed ``[0, LOCAL)`` and committed nothing; a fetch of the rest of the prompt is + abandoned before its first poll, or delivered whole. The chunk's tokens count either way.""" + conversation = "kept-chunk" if policy == "per_conversation" else None + with real_manager(block_reuse_config={"policy": policy}) as mgr: + target = computed(kit, mgr, TARGET, SPLIT_PROMPT, conversation=conversation) + lender = attach(mgr, tokens_per_fetch=2 * TPB) + assert lender.resume_bounds(target) == (LOCAL, LOCAL) + if outcome == "abandoned": + lease = lender.lend_write(target, LOCAL, 4 * TPB) + assert lender.resume_bounds(target) is None, "a fetch counts from its call" + lease.release() # before its first poll: no backend wrote + expected = (LOCAL, LOCAL) + else: + fetch_segment(kit, mgr, lender, target, LOCAL, 4 * TPB, 0x5A) + assert holds(kit, mgr, target, 0, [2, 3], staged_bytes(kit, 0x5A)) + expected = (4 * TPB, LOCAL) + assert holds(kit, mgr, target, 0, [0, 1], chunk_bytes), "the fetch changed the chunk" + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == expected, ( + f"resume_bounds {tuple(resume_bounds)}: the tokens computed before the fetch were dropped" + ) + + +def test_a_windowed_fetch_counts_what_its_window_still_holds(kit, real_manager, attach=attach): + """The chunk under a window of 96 tokens. A fetch of ``[LOCAL, 128)`` moves the history to 128: + the window releases block 0 but keeps block 1, which it reads at 128 beside the fetched blocks 2 + and 3. Delivered whole, the request resumes at 128.""" + with real_manager(windows=[96, 256], **PER_REQUEST) as mgr: + target = computed(kit, mgr, TARGET, SPLIT_PROMPT) + kv = kit.kv(mgr, target) + sliding = kit.windows(mgr).index(96) + lender = attach(mgr, tokens_per_fetch=2 * TPB) + fetch_segment(kit, mgr, lender, target, LOCAL, 4 * TPB, 0x5A) + assert kv.history_length == 4 * TPB, "the fetch did not move the history to its end" + assert kit.stale_blocks(mgr, sliding, 4 * TPB) == (0, 1) + assert kit.pages(kv, sliding)[0] < 0, "the window kept block 0" + assert holds(kit, mgr, target, sliding, [1], chunk_bytes), "the window lost block 1" + assert holds(kit, mgr, target, sliding, [2, 3], staged_bytes(kit, 0x5A)) + assert holds(kit, mgr, target, 1 - sliding, [0, 1], chunk_bytes) + assert holds(kit, mgr, target, 1 - sliding, [2, 3], staged_bytes(kit, 0x5A)), ( + "the full-attention group's fetched blocks do not hold the bytes staged for them" + ) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == (4 * TPB, 4 * TPB), ( + f"resume_bounds {tuple(resume_bounds)}: the tokens computed before the fetch were dropped" + ) + + +# name: (windows, the fetch's start, resume_bounds, or None for an empty interval) +WINDOW_STARTS = { + "from_the_computed_prefix": ([WINDOW], 0, (4 * TPB, 4 * TPB)), + "past_it": ([WINDOW], 2 * TPB, (4 * TPB, 4 * TPB)), + "past_it_with_a_wider_window": ([WINDOW, 256], 2 * TPB, None), +} + + +@pytest.mark.parametrize("case", list(WINDOW_STARTS)) +def test_a_fetch_counts_the_rows_each_window_reads_whatever_its_start( + kit, real_manager, case, attach=attach +): + """All-reusable, nothing computed, a fetch ending at 128: resume_bounds counts what each layer group + reads at 128 among the delivered rows, whatever the fetch's start. The window of ``WINDOW`` + tokens reads blocks 2 and 3 alone; a group whose window covers the prompt reads 0 and 1 too.""" + windows, start, expected = WINDOW_STARTS[case] + with real_manager(windows=windows) as mgr: + target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) + lender = attach(mgr, tokens_per_fetch=4 * TPB) + fetch_segment(kit, mgr, lender, target, start, 4 * TPB, 0x5A) + sliding = kit.windows(mgr).index(WINDOW) + assert holds(kit, mgr, target, sliding, [2, 3], staged_bytes(kit, 0x5A)), ( + "the window's fetched blocks lack their staged bytes" + ) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.min_position == 4 * TPB, f"resume_bounds {tuple(resume_bounds)}" + if expected is None: + assert resume_bounds.max_position < resume_bounds.min_position, ( + f"resume_bounds {tuple(resume_bounds)} counts blocks 0 and 1, which the wider window reads " + "and nothing delivered" + ) + else: + assert resume_bounds == expected, ( + f"resume_bounds {tuple(resume_bounds)}: the window at {4 * TPB} reads blocks 2 and 3 " + "alone, which landed" + ) + + +def test_an_abandoned_windowed_fetch_claims_nothing(kit, real_manager, attach=attach): + """The chunk under a window of 64 tokens. A fetch of ``[LOCAL, 128)`` moves the history to 128, + and the window releases blocks 0 and 1, which a start of 64 reads; abandoned, the fetch wrote + neither block 2 nor block 3, which a start of 128 reads. No start is usable.""" + with real_manager(windows=[WINDOW, 256], **PER_REQUEST) as mgr: + target = computed(kit, mgr, TARGET, SPLIT_PROMPT) + kv = kit.kv(mgr, target) + sliding = kit.windows(mgr).index(WINDOW) + lender = attach(mgr, tokens_per_fetch=2 * TPB) + lender.lend_write(target, LOCAL, 4 * TPB).release() # before its first poll + assert kv.history_length == 4 * TPB, "the fetch did not move the history to its end" + assert all(p < 0 for p in kit.pages(kv, sliding)[:2]), "the window kept the chunk" + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.min_position == 4 * TPB + assert resume_bounds.max_position < resume_bounds.min_position, ( + f"resume_bounds {tuple(resume_bounds)} claims a start whose blocks were released or never " + "written" + ) + + +def test_a_fetch_from_below_the_history_keeps_only_what_lies_below_it( + kit, real_manager, monkeypatch, attach=attach +): + """The chunk ``[0, LOCAL)``, then a fetch from block 1: its copy into pages writes block 1 and + fails at block 3, so the fetch is abandoned with block 1 overwritten. Only block 0 counts.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + with real_manager(**PER_REQUEST) as mgr: + target = computed(kit, mgr, TARGET, SPLIT_PROMPT) + lender = attach(mgr, tokens_per_fetch=3 * TPB) + lease = lender.lend_write(target, TPB, 4 * TPB) + view = lease.poll() + assert by_group(view) == {0: [1, 2, 3]} + kit.stage(lender, view, 0x5A) + driver = FailingDriver(_lender.drv, fail_on=2) + with monkeypatch.context() as patched: + patched.setattr(_lender, "drv", driver) + # Blocks 1 and 3 arrived: their staging slots are apart, so each takes its own copy. + lease.mark_arrived(tuple(np.isin(run.block_ordinals, [1, 3]) for run in view.groups)) + lease.release() + assert driver.calls >= 2, "a single copy: nothing failed partway" + assert holds(kit, mgr, target, 0, [0], chunk_bytes) + assert holds(kit, mgr, target, 0, [1], staged_bytes(kit, 0x5A)), ( + "the copy did not bring block 1 its own row" + ) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position <= TPB, ( + f"resume_bounds {tuple(resume_bounds)} counts block 1, which the abandoned fetch overwrote" + ) + # Nothing was delivered, so the request may resume only at its history, past TPB. + assert resume_bounds == (TPB, LOCAL), ( + f"resume_bounds {tuple(resume_bounds)}: block 0, computed before the fetch, was dropped" + ) + + +def test_a_refetch_whose_copy_fails_partway_counts_none_of_the_rows_it_was_to_overwrite( + kit, real_manager, monkeypatch, attach=attach +): + """``[0, 128)`` fetched and delivered, nothing committed; then ``[32, 128)`` fetched again with + blocks 1 and 3 marked, its copy failing after block 1's. Neither marked block counts any more, + since a copy failing partway can leave a row holding each fetch's bytes; block 0 still does.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) + lender = attach(mgr, tokens_per_fetch=4 * TPB) + fetch_segment(kit, mgr, lender, target, 0, 4 * TPB, 0x11) + assert lender.resume_bounds(target) == (4 * TPB, 0) + lease = lender.lend_write(target, TPB, 4 * TPB) + view = lease.poll() + assert by_group(view) == {0: [1, 2, 3]} + kit.stage(lender, view, 0x22) + driver = FailingDriver(_lender.drv, fail_on=2) + with monkeypatch.context() as patched: + patched.setattr(_lender, "drv", driver) + # Blocks 1 and 3: their staging slots are apart, so each takes its own copy. + lease.mark_arrived(tuple(np.isin(run.block_ordinals, [1, 3]) for run in view.groups)) + lease.release() + assert driver.calls >= 2, "a single copy: nothing failed partway" + assert holds(kit, mgr, target, 0, [1], staged_bytes(kit, 0x22)), ( + "the copy did not bring block 1 the second fetch's row" + ) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position <= TPB, ( + f"resume_bounds {tuple(resume_bounds)} counts rows the failed copy was to overwrite" + ) + assert resume_bounds == (TPB, 0), ( + f"resume_bounds {tuple(resume_bounds)}: block 0, which the second fetch left alone, was dropped" + ) + + +class StatsDiscardRaising: + """A ``kv_cache_map`` entry whose ``discard_pending_stats`` raises, so the manager's free stops + before the lender's free hook; the manager's other calls on it go to the cache.""" + + def __init__(self, kv_cache): + self.kv_cache = kv_cache + + def __getattr__(self, name): + return getattr(self.kv_cache, name) + + def discard_pending_stats(self): + raise RuntimeError("planted") + + +FAILED_COPIES = ["returned", "raised", "interrupted"] +TAIL_MIXED = "counts the committed tokens of block 1, which the failed copy may have mixed" + + +@pytest.mark.parametrize("failure", [None, *FAILED_COPIES], ids=["control", *FAILED_COPIES]) +def test_a_copy_failing_over_the_block_the_committed_tokens_end_inside_empties_the_interval( + kit, real_manager, monkeypatch, failure, attach=attach +): + """All-reusable: the target's local match ends inside block 1, and a fetch from block 1 marks + blocks 1 and 3. Block 1's copy is queued, and lands, before block 3's fails (the second memcpy + returns an error, or raises an ``Exception`` or an interruption that is none): the stricter + order for the rule, which reads the marks, not how far the copy got. Block 1 may then hold the + fetch's bytes in some pools and its old bytes in others, so the interval is empty, then and + later, though not for a cache that replaces it under the same id after a free that raised before + the lender's free hook, which keeps the records. The control counts block 1, which holds its + staged bytes. A lender tainting only the failing copy call's rows, or only a block the copy left + incomplete, keeps counting block 1.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + with real_manager() as mgr: + kit.published(mgr, SOURCE, SIBLING_PROMPT) + target = kit.admitted(mgr, TARGET, MATCHED_PROMPT) + kv = kit.kv(mgr, target) + assert kv.num_committed_tokens == len(SHARED) and len(SHARED) % TPB, ( + "no partial match: the check proves nothing" + ) + lender = attach(mgr, tokens_per_fetch=3 * TPB) + lease = lender.lend_write(target, TPB, 4 * TPB) + view = lease.poll() + assert view is not None, lease.failure + assert by_group(view) == {0: [1, 2, 3]} + kit.stage(lender, view, 0x5A) + planted = {"raised": MemoryError, "interrupted": Interrupted}.get(failure) + driver = FailingDriver(_lender.drv, fail_on=2 if failure else None, raises=planted) + # Blocks 1 and 3: their staging slots are apart, so each takes its own copy. + masks = tuple(np.isin(run.block_ordinals, [1, 3]) for run in view.groups) + with monkeypatch.context() as patched: + patched.setattr(_lender, "drv", driver) + if planted is None: + lease.mark_arrived(masks) + else: + with pytest.raises(planted, match="planted"): + lease.mark_arrived(masks) + lease.release() + mgr._stream.synchronize() + assert driver.calls >= 2, "a single copy: nothing failed partway" + for run in view.groups: # block 1, in every layer group the view has + assert holds(kit, mgr, target, run.layer_group_index, [1], staged_bytes(kit, 0x5A)), ( + f"group {run.layer_group_index}: the copy did not reach block 1: the check proves nothing" + ) + if failure is None: + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == (2 * TPB, TPB), f"resume_bounds {tuple(resume_bounds)}" + return + for _ in range(2): # and on a later call + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position < resume_bounds.min_position, ( + f"resume_bounds {tuple(resume_bounds)} {TAIL_MIXED}" + ) + mgr.kv_cache_map[TARGET] = StatsDiscardRaising(kv) + with pytest.raises(RuntimeError, match="planted"): + mgr.free_resources(target) + assert ( + TARGET not in mgr.kv_cache_map and lender._ledger._mixed_tail_caches.get(TARGET) is kv + ), "the free left the cache in the map or ran the lender's hook: the check proves nothing" + # The free's steps past the raise, so the request id is admitted again. + kv.close() + mgr.index_mapper.remove_sequence(TARGET) + target = kit.admitted(mgr, TARGET, MATCHED_PROMPT) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position >= resume_bounds.min_position, ( + f"resume_bounds {tuple(resume_bounds)}: the cache that replaced the failed one stays empty" + ) + + +def test_a_copy_failing_over_that_block_in_a_later_layer_group_empties_the_interval( + kit, real_manager, monkeypatch, attach=attach +): + """All-reusable with two layer groups: the target's local match ends inside block 1, which a + window keeps. ``[32, 128)`` is delivered whole, then fetched again with blocks 1 and 3 marked in + the second layer group alone, its copy failing after block 1's. The rule reads the marks in + every layer group, so the interval is empty.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + with real_manager(windows=[4 * TPB, 256]) as mgr: + kit.published(mgr, SOURCE, SIBLING_PROMPT) + target = kit.admitted(mgr, TARGET, MATCHED_PROMPT) + assert kit.kv(mgr, target).num_committed_tokens == len(SHARED), ( + "no partial match: the check proves nothing" + ) + lender = attach(mgr, tokens_per_fetch=3 * TPB) + fetch_segment(kit, mgr, lender, target, TPB, 4 * TPB, 0x11) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position >= resume_bounds.min_position, ( + f"resume_bounds {tuple(resume_bounds)} is empty before the refetch: the check proves nothing" + ) + lease = lender.lend_write(target, TPB, 4 * TPB) + view = lease.poll() + assert view is not None, lease.failure + assert by_group(view) == {0: [1, 2, 3], 1: [1, 2, 3]} + kit.stage(lender, view, 0x5A) + driver = FailingDriver(_lender.drv, fail_on=2) + # Blocks 1 and 3 of the second layer group: their staging slots are apart. + first, *later = view.groups + masks = ( + np.zeros(len(first), dtype=bool), + *(np.isin(r.block_ordinals, [1, 3]) for r in later), + ) + with monkeypatch.context() as patched: + patched.setattr(_lender, "drv", driver) + lease.mark_arrived(masks) + lease.release() + mgr._stream.synchronize() + assert driver.calls >= 2, "a single copy: nothing failed partway" + assert holds(kit, mgr, target, first.layer_group_index, [1], staged_bytes(kit, 0x11)), ( + "the copy reached block 1 of the first layer group: the check proves nothing" + ) + assert holds(kit, mgr, target, later[0].layer_group_index, [1], staged_bytes(kit, 0x5A)), ( + "the copy did not reach block 1 of the second layer group: the check proves nothing" + ) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position < resume_bounds.min_position, ( + f"resume_bounds {tuple(resume_bounds)} {TAIL_MIXED}" + ) + + +def test_a_second_fetch_keeps_what_the_first_kept(kit, real_manager, attach=attach): + """The chunk under a window of 64 tokens; a fetch of ``[LOCAL, 128)`` abandoned, which moved the + history to 128; then a fetch of ``[128, 160)`` delivered. A start of 160 reads block 3, which + nothing computed or fetched.""" + with real_manager(windows=[WINDOW, 256], **PER_REQUEST) as mgr: + target = computed(kit, mgr, TARGET, WINDOWED_PROMPT) + kv = kit.kv(mgr, target) + lender = attach(mgr, tokens_per_fetch=2 * TPB) + lender.lend_write(target, LOCAL, 4 * TPB).release() # before its first poll + assert kv.history_length == 4 * TPB + fetch_segment(kit, mgr, lender, target, 4 * TPB, WINDOWED_END, 0x5A) + for lg in range(kit.num_layer_groups(mgr)): + assert holds(kit, mgr, target, lg, [4], staged_bytes(kit, 0x5A)), f"layer group {lg}" + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.min_position == WINDOWED_END + assert resume_bounds.max_position < resume_bounds.min_position, ( + f"resume_bounds {tuple(resume_bounds)} counts block 3, which nothing computed or fetched" + ) + + +def test_a_rollback_keeps_the_computed_chunk_and_voids_the_fetches( + kit, real_manager, attach=attach +): + """The chunk ``[0, TPB)``, then two one-block fetches delivered; the manager's context rollback + shrinks the cache back to the chunk, freeing both fetched blocks.""" + with real_manager(**PER_REQUEST) as mgr: + target = computed(kit, mgr, TARGET, SPLIT_PROMPT, upto=TPB) + target.py_ctx_pre_resize_cap = None # a later resize grew nothing + kv = kit.kv(mgr, target) + lender = attach(mgr, tokens_per_fetch=TPB) + fetch_segment(kit, mgr, lender, target, TPB, 2 * TPB, 0x11) + assert lender.resume_bounds(target) == (2 * TPB, TPB) + fetch_segment(kit, mgr, lender, target, 2 * TPB, 3 * TPB, 0x22) + assert lender.resume_bounds(target) == (3 * TPB, TPB) + for lg in range(kit.num_layer_groups(mgr)): # the rows that resume_bounds counts + for ordinal, byte in ((1, 0x11), (2, 0x22)): + assert holds(kit, mgr, target, lg, [ordinal], staged_bytes(kit, byte)), ( + f"layer group {lg}: fetched block {ordinal} lacks its staged bytes" + ) + assert target.py_ctx_pre_resize_cap == TPB, "the rollback would not shrink to the chunk" + assert mgr.revert_allocate_context(target) is True + assert kit.kv(mgr, target) is kv, "the rollback replaced the cache" + assert (kv.capacity, kv.history_length, kv.num_blocks) == (TPB, TPB, 1) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == (TPB, TPB), ( + f"resume_bounds {tuple(resume_bounds)} counts block 1, which the rollback freed" + ) + mgr.free_resources(target) + + +# -- where a request resumes when its context update moves the history ------------------------ + +LONG_PROMPT = list(range(4000, 4225)) # seven whole blocks and one token + + +def resumed(kit, mgr, request, at): + """``documented_resume`` as far as the manager calls that follow read it: the history raised to + ``at`` if below it, the context position moved there. Through the V2 scheduler a resume past the + first chunk needs the whole step.""" + kv = kit.kv(mgr, request) + if kv.history_length < at: + kv.resize(None, at) + request.context_current_position = at + + +def next_chunk(kit, mgr, request, tokens=TPB): + """The request's next context chunk of ``tokens`` run as the executor runs it, block ``b`` + written with ``chunk_bytes``; the manager's error, or ``None``.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + first = request.context_current_position + request.context_chunk_size = min(tokens, request.prompt_len - first) + assert mgr.resize_context(request, request.context_chunk_size) + kv = kit.kv(mgr, request) + dev = kit.DevicePages(mgr) + end_block = (first + request.context_chunk_size) // TPB + for lg in range(kit.num_layer_groups(mgr)): + for ordinal, slot in enumerate(kit.pages(kv, lg)): + if first // TPB <= ordinal < end_block and slot >= 0: + dev.write(lg, slot, chunk_bytes(dev, lg, ordinal)) + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + try: + mgr.update_context_resources(batch) + except ValueError as error: + return str(error) + return None + + +RESUMES = { + # name: (manager kwargs, prompt, tokens computed first, fetch range, block ordinals arrived + # (None: all), resume_bounds expected (None: any interval whose starts all resume)) + "delivered": ({}, SPLIT_PROMPT, LOCAL, (LOCAL, 4 * TPB), None, (4 * TPB, LOCAL)), + # A window that releases nothing by 128, and a fetch that delivered no row: a remote miss. + "windowed_nothing_delivered": ( + {"windows": [4 * TPB, 256]}, + SPLIT_PROMPT, + LOCAL, + (LOCAL, 4 * TPB), + [], + None, + ), + # A window that releases nothing by 192, and a fetch whose first block alone arrived. + "windowed_partial": ( + {"windows": [6 * TPB, 256]}, + LONG_PROMPT, + LOCAL, + (LOCAL, 6 * TPB), + [2], + None, + ), + # A fetch from below the history whose first block did not arrive. + "below_the_history_partial": ({}, SPLIT_PROMPT, 4 * TPB, (TPB, 4 * TPB), [2, 3], None), + # A fetch from below the history, delivered whole. + "below_the_history_delivered": ( + {}, + LONG_PROMPT, + 4 * TPB, + (TPB, 6 * TPB), + None, + (6 * TPB, 4 * TPB), + ), + # Nothing computed first, a window that releases nothing by 192, and two blocks arrived. + "cold_windowed_partial": ( + {"windows": [6 * TPB, 256]}, + LONG_PROMPT, + 0, + (0, 6 * TPB), + [0, 1], + None, + ), +} + + +@pytest.mark.parametrize("policy", ["per_request", "per_conversation"]) +@pytest.mark.parametrize("case", list(RESUMES)) +def test_resume_bounds_claims_only_starts_the_context_update_takes( + kit, real_manager, case, policy, attach=attach +): + """Per request and per conversation, the context update moves the history to each chunk's end + and raises for a chunk ending below it: the next chunk runs from the lowest start resume_bounds + claims, and the floor is at least the history. No window released blocks, so none sets it.""" + from tensorrt_llm.conversation_params import ConversationParams + + kwargs, prompt, upto, (start, end), arrived, expected = RESUMES[case] + conversation = "claims" if policy == "per_conversation" else None + with real_manager(**kwargs, block_reuse_config={"policy": policy}) as mgr: + if upto: + target = computed(kit, mgr, TARGET, prompt, upto=upto, conversation=conversation) + else: + target = kit.make_request(TARGET, prompt) + if conversation is not None: + target.py_conversation_params = ConversationParams(conversation_id=conversation) + assert mgr.prepare_context(target) + kit.kv(mgr, target).enable_swa_scratch_reuse = False # a fetch target needs it off + kv = kit.kv(mgr, target) + lender = attach(mgr, tokens_per_fetch=end - start) + fetch_segment(kit, mgr, lender, target, start, end, 0x5A, arrived) + delivered = range(start // TPB, end // TPB) if arrived is None else arrived + for lg in range(kit.num_layer_groups(mgr)): + assert holds(kit, mgr, target, lg, delivered, staged_bytes(kit, 0x5A)), f"group {lg}" + history = kv.history_length + assert not stale_anywhere(kit, mgr, history), "a window released blocks" + resume_bounds = lender.resume_bounds(target) + floor = resume_bounds.min_position + if resume_bounds.max_position >= floor: + # The lowest start claimed: from a higher one the chunk ends no lower. + resumed(kit, mgr, target, floor) + error = next_chunk(kit, mgr, target) + assert error is None, ( + f"resume_bounds {tuple(resume_bounds)} lets the request resume at {floor}, below its " + f"history of {history}, where its next chunk fails: {error}" + ) + assert floor >= history, ( + f"resume_bounds {tuple(resume_bounds)} floors below the history {history}" + ) + if expected is not None: + assert resume_bounds == expected, ( + f"resume_bounds {tuple(resume_bounds)}: what was computed or delivered was dropped" + ) + + +def test_a_later_fetch_keeps_what_the_request_computed_after_resuming( + kit, real_manager, attach=attach +): + """Per request: the chunk ``[0, LOCAL)``, a fetch of ``[LOCAL, 128)`` delivered, a resume at 128 + as a remote tier does, ``[128, 192)`` computed, then a fetch of ``[192, 224)`` delivered. Every + block below 224 holds what the request computed or fetched, so it resumes at 224.""" + with real_manager(**PER_REQUEST) as mgr: + target = computed(kit, mgr, TARGET, LONG_PROMPT) + kv = kit.kv(mgr, target) + lender = attach(mgr, tokens_per_fetch=2 * TPB) + fetch_segment(kit, mgr, lender, target, LOCAL, 4 * TPB, 0x5A) + first = lender.resume_bounds(target) + assert first == (4 * TPB, LOCAL), f"resume_bounds {tuple(first)} after the first fetch" + resumed(kit, mgr, target, first.max_position) + assert next_chunk(kit, mgr, target, 2 * TPB) is None + assert (kv.history_length, kv.num_committed_tokens) == (6 * TPB, 0), "the chunk committed" + fetch_segment(kit, mgr, lender, target, 6 * TPB, 7 * TPB, 0x6B) + assert holds(kit, mgr, target, 0, [0, 1, 4, 5], chunk_bytes) + assert holds(kit, mgr, target, 0, [2, 3], staged_bytes(kit, 0x5A)) + assert holds(kit, mgr, target, 0, [6], staged_bytes(kit, 0x6B)) + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == (7 * TPB, 6 * TPB), ( + f"resume_bounds {tuple(resume_bounds)}: the tokens computed after resuming were dropped" + ) + + +# -- a joint-reuse draft pool, which shares the request's cursor ---------------------------- + +PAIRED_END = 128 # SPLIT_PROMPT's whole blocks before its last token + + +def schedule(scheduler, request) -> List[int]: + """The ids of the context requests one round of the V2 scheduler admits.""" + return [r.py_request_id for r in scheduler.schedule_request([request], set()).context_requests] + + +def fill_blocks(kit, mgr, request, ordinals, byte): + """Write the request's blocks ``ordinals`` as ``kit.stage`` stages them from ``byte``.""" + dev = kit.DevicePages(mgr) + for lg in range(kit.num_layer_groups(mgr)): + slots = kit.pages(kit.kv(mgr, request), lg) + for ordinal in ordinals: + dev.write(lg, slots[ordinal], kit.staged_row(byte, lg, ordinal, dev.page_bytes(lg))) + + +@pytest.mark.parametrize("draft_delivers", [True, False], ids=["draft_hit", "draft_miss"]) +def test_a_joint_draft_pool_resumes_where_both_lenders_allow( + kit, real_manager, draft_delivers, attach=attach +): + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + from tensorrt_llm.bindings import LlmRequestState + from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy + + spec, byte = dflash(), 0x5A # what any worker computing these tokens holds + with ( + real_manager(spec_config=spec, joint_kv_cache_reuse=True) as target, + real_manager(spec_config=spec, is_draft=True, joint_kv_cache_reuse=True) as draft, + ): + pools = (target, draft) + scheduler = KVCacheV2Scheduler( + max_batch_size=4, + max_num_tokens=2048, + kv_cache_manager=target, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + draft_kv_cache_manager=draft, + ) + # The draft pool under a scope of its own, as attach_lender asks of such a pair. + lenders = [ + attach(target, tokens_per_fetch=PAIRED_END), + attach(draft, tokens_per_fetch=PAIRED_END, scope=SCOPE + b"/draft-pool"), + ] + request = kit.make_request(TARGET, SPLIT_PROMPT) + assert schedule(scheduler, request) == [TARGET] + for mgr in pools: + fill_blocks(kit, mgr, request, range(5), kit.SENTINEL) # nothing computed yet + request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS # parked to fetch + intervals = [] + for mgr, lender, delivered in zip(pools, lenders, (True, draft_delivers)): + lease = lender.lend_write(request, 0, PAIRED_END) + view = lease.poll() + kit.stage(lender, view, byte) + lease.mark_arrived(view.row_masks(delivered)) + lease.release() + mgr._stream.synchronize() + intervals.append(lender.resume_bounds(request)) + # Each interval covers only its own pool's blocks: the request resumes where both allow. + usable = min(interval.max_position for interval in intervals) + assert max(interval.min_position for interval in intervals) <= usable + for mgr in pools: + kv = kit.kv(mgr, request) + if kv.history_length < usable: + kv.resize(None, usable) + request.py_served_prefix_end = usable + request.state = LlmRequestState.CONTEXT_INIT + assert schedule(scheduler, request) == [TARGET] + assert request.context_current_position == usable + batch = ScheduledRequests() + batch.append_context_request(request) + for mgr in pools: + mgr.prepare_resources(batch) + first = usable // TPB + last = (usable + request.context_chunk_size - 1) // TPB + for mgr in pools: # the forward computes the chunk's blocks + fill_blocks(kit, mgr, request, range(first, last + 1), byte) + request.move_to_next_context_chunk() + for mgr in pools: + mgr.update_context_resources(batch) + assert kit.kv(draft, request).num_committed_tokens == len(SPLIT_PROMPT) + nbytes = kit.DevicePages(draft).page_bytes(0) + blocks = range(PAIRED_END // TPB) + held = [kit.digest(kit.staged_row(byte, 0, o, nbytes)) for o in blocks] + committed = [kit.digest(kit.page(draft, request, 0, o)) for o in blocks] + assert committed == held, "the draft pool committed blocks it lacks" + for lg in range(kit.num_layer_groups(target)): + nbytes = kit.DevicePages(target).page_bytes(lg) + held = [kit.digest(kit.staged_row(byte, lg, o, nbytes)) for o in blocks] + committed = [kit.digest(kit.page(target, request, lg, o)) for o in blocks] + assert committed == held, "the target pool committed blocks it lacks" + # The whole fetch is used where both pools' rows landed, none where the draft's did not. + assert usable == (PAIRED_END if draft_delivers else 0), ( + f"resumed at {usable}, not where both pools' fetched blocks allow" + ) + + +JOINT_WINDOW = 192 # wider than PAIRED_END: no window releases a block at the fetch's end +# name: (block ordinals delivered into the target pool, into the draft pool) +JOINT_DELIVERIES = { + "short_in_both": ([0, 1], [0, 1]), + "draft_miss": ([0, 1, 2, 3], []), + "full": ([0, 1, 2, 3], [0, 1, 2, 3]), +} + + +@pytest.mark.parametrize("case", list(JOINT_DELIVERIES)) +def test_a_windowed_joint_draft_pool_resumes_only_where_its_context_resize_holds( + kit, real_manager, case, attach=attach +): + """All-reusable, both pools windowed wider than the fetch: the target pool may resume below the + history the fetch moved; the draft pool's context resize sets its capacity from the chunk, so + its floor is that history. Wherever both may resume, they prepare a chunk there smaller than the + fetch and hold the staged bytes below it; a short delivery in either pool leaves none.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + from tensorrt_llm.bindings import LlmRequestState + from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy + + spec, byte, windows = dflash(), 0x5A, [JOINT_WINDOW] + with ( + real_manager(spec_config=spec, joint_kv_cache_reuse=True, windows=windows) as target, + real_manager( + spec_config=spec, is_draft=True, joint_kv_cache_reuse=True, windows=windows + ) as draft, + ): + pools = (target, draft) + for mgr in pools: + assert JOINT_WINDOW in kit.windows(mgr), "a pool has no window: nothing checked" + beg, end = kit.stale_blocks(mgr, kit.windows(mgr).index(JOINT_WINDOW), PAIRED_END) + assert end <= beg, "a window released blocks at the fetch's end: nothing checked" + scheduler = KVCacheV2Scheduler( + max_batch_size=4, + max_num_tokens=2048, + kv_cache_manager=target, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + draft_kv_cache_manager=draft, + ) + lenders = [ + attach(target, tokens_per_fetch=PAIRED_END), + attach(draft, tokens_per_fetch=PAIRED_END, scope=SCOPE + b"/draft-pool"), + ] + request = kit.make_request(TARGET, SPLIT_PROMPT) + assert schedule(scheduler, request) == [TARGET] + request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS # parked to fetch + intervals = [] + for mgr, lender, ordinals in zip(pools, lenders, JOINT_DELIVERIES[case]): + lease = lender.lend_write(request, 0, PAIRED_END) + view = ready_view(lease, mgr) + kit.stage(lender, view, byte) + lease.mark_arrived(tuple(np.isin(run.block_ordinals, ordinals) for run in view.groups)) + lease.release() + mgr._stream.synchronize() + intervals.append(lender.resume_bounds(request)) + floor = max(interval.min_position for interval in intervals) + usable = min(interval.max_position for interval in intervals) + if floor <= usable: + # The pair resumes at the smaller usable end and runs a chunk smaller than the fetch. + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = usable + request.context_chunk_size = min(TPB, request.prompt_len - usable) + assert target.resize_context(request, request.context_chunk_size) + batch = ScheduledRequests() + batch.append_context_request(request) + try: + draft.prepare_resources(batch) + except (ValueError, RuntimeError) as error: + raise AssertionError( + f"the draft pool's context resize at {usable} raised: {error}" + ) from error + for name, mgr in zip(("target", "draft"), pools): # the resume reads every row below + for lg in range(kit.num_layer_groups(mgr)): + rows = range(usable // TPB) + assert holds(kit, mgr, request, lg, rows, staged_bytes(kit, byte)), ( + f"{name} pool, group {lg}: a row the resume reads lacks its staged bytes" + ) + if case == "full": + assert (floor, usable) == (PAIRED_END, PAIRED_END), (floor, usable) + else: + assert floor > usable, f"the pair may resume in [{floor}, {usable}] after a short fetch" + + +# -- a request resumed past its first context chunk, through the V2 scheduler ----------------- + +RESUME_END = 6 * TPB # the fetch's end in LONG_PROMPT, past the chunk the request computed first + + +def documented_resume(kit, pools, request, at): + """``ResumeBounds``'s resume at ``at``, a KV cache connector's skip: the served position set, each + cache's history raised to ``at`` where below it, then the chunk spanning to the prompt's end and + the prepopulated length and context position moved to ``at`` together.""" + first_chunk_resume(kit, pools, request, at) + request.context_chunk_size = request.prompt_len - request.context_current_position + request.set_prepopulated_prompt_len(at, TPB) + + +def first_chunk_resume(kit, pools, request, at): # the served position and the histories alone + for mgr in pools: + kv = kit.kv(mgr, request) + if kv.history_length < at: + kv.resize(None, at) + request.py_served_prefix_end = at + + +def runs_its_chunk(kit, pools, request): + """The chunk the scheduler gave the request, run as the executor runs it: each pool prepared, + every block the chunk touches written with ``chunk_bytes``, the context position moved and each + pool's context update run; the first update's error, or ``None``.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + batch = ScheduledRequests() + batch.append_context_request(request) + for mgr in pools: + mgr.prepare_resources(batch) + first = request.context_current_position + end = first + request.context_chunk_size + for mgr in pools: + kv = kit.kv(mgr, request) + dev = kit.DevicePages(mgr) + for lg in range(kit.num_layer_groups(mgr)): + for ordinal, slot in enumerate(kit.pages(kv, lg)): + if first // TPB <= ordinal < -(-end // TPB) and slot >= 0: + dev.write(lg, slot, chunk_bytes(dev, lg, ordinal)) + request.move_to_next_context_chunk() + try: + for mgr in pools: + mgr.update_context_resources(batch) + except (ValueError, RuntimeError) as error: + return str(error) + return None + + +@pytest.mark.parametrize( + "target_policy,draft_policy", + [ + ("per_request", None), + ("per_request", "per_request"), + ("per_request", "all_reusable"), + ("all_reusable", "per_request"), + ("all_reusable", "all_reusable"), + ], +) +def test_a_request_past_its_first_chunk_resumes_where_the_documented_step_puts_it( + kit, real_manager, target_policy, draft_policy, resume=documented_resume +): + """Through the V2 scheduler chunking at ``LOCAL`` tokens, alone or with a same/mixed-policy + joint-reuse draft pool: the request computes its first chunk, a fetch of ``[LOCAL, RESUME_END)`` + arrives whole, and the request resumes at ``max_position`` by the step ``ResumeBounds`` states. Its + next chunk starts there with the prepopulated length moved alike, every block below holds what + the request computed or the fetch delivered, and each pool's context update takes the chunk.""" + from contextlib import ExitStack + + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + from tensorrt_llm.bindings import LlmRequestState + from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy + + byte = 0x5A + paired = draft_policy is not None + target_config = {"block_reuse_config": {"policy": target_policy}} + with ExitStack() as stack: + if paired: + spec = dflash() + pools = ( + stack.enter_context( + real_manager(spec_config=spec, joint_kv_cache_reuse=True, **target_config) + ), + stack.enter_context( + real_manager( + spec_config=spec, + is_draft=True, + joint_kv_cache_reuse=True, + block_reuse_config={"policy": draft_policy}, + ) + ), + ) + else: + pools = (stack.enter_context(real_manager(**target_config)),) + scheduler = KVCacheV2Scheduler( + max_batch_size=4, + max_num_tokens=LOCAL, + kv_cache_manager=pools[0], + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, TPB), + draft_kv_cache_manager=pools[1] if paired else None, + ) + scopes = (SCOPE, SCOPE + b"/draft-pool") # the draft pool under a scope of its own + lenders = [ + attach(m, tokens_per_fetch=RESUME_END - LOCAL, scope=s) for m, s in zip(pools, scopes) + ] + request = kit.make_request(TARGET, LONG_PROMPT) + assert schedule(scheduler, request) == [TARGET] + for mgr in pools: + kit.kv(mgr, request).enable_swa_scratch_reuse = False # a fetch target needs it off + assert runs_its_chunk(kit, pools, request) is None + assert request.context_current_position == LOCAL, "the first chunk was not LOCAL tokens" + request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS # parked to fetch + intervals = [] + for mgr, lender in zip(pools, lenders): + fetch_segment(kit, mgr, lender, request, LOCAL, RESUME_END, byte) + intervals.append(tuple(lender.resume_bounds(request))) + assert intervals == [(RESUME_END, LOCAL)] * len(pools), f"resume_bounds {intervals}" + at = RESUME_END + request.state = LlmRequestState.CONTEXT_INIT + resume(kit, pools, request, at) + assert schedule(scheduler, request) == [TARGET] + assert request.context_current_position == at, ( + f"the next chunk starts at {request.context_current_position}, not at {at}" + ) + assert request.prepopulated_prompt_len == at, ( + f"the prepopulated length stays at {request.prepopulated_prompt_len}, not at {at}" + ) + computed_blocks, fetched_blocks = range(LOCAL // TPB), range(LOCAL // TPB, at // TPB) + for mgr in pools: + for lg in range(kit.num_layer_groups(mgr)): + assert holds(kit, mgr, request, lg, computed_blocks, chunk_bytes), f"group {lg}" + assert holds(kit, mgr, request, lg, fetched_blocks, staged_bytes(kit, byte)), ( + f"group {lg}: a block the resume reads lacks its staged bytes" + ) + error = runs_its_chunk(kit, pools, request) + assert error is None, f"the context update did not take the chunk at {at}: {error}" + for mgr in pools: + assert kit.kv(mgr, request).history_length == request.prompt_len + + +SHIFTED_AT = 3 * TPB + 24 # inside the fetched block, past the context position + + +def test_a_request_resumes_at_an_unaligned_point_past_its_context_position( + kit, real_manager, resume=documented_resume, below=False +): + """All-reusable, through the V2 scheduler: a partial local match and a first chunk leave the + context position inside block 3, and a fetch of block 3 starts the interval below it. Resumed + past it at an unaligned point, the request's next chunk starts there over the bytes it reads.""" + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + from tensorrt_llm.bindings import LlmRequestState + from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy + + byte, prompt = 0x5A, SHARED + list(range(7400, 7493)) # four whole blocks and five tokens + with real_manager() as mgr: + scheduler = KVCacheV2Scheduler( + max_batch_size=4, + max_num_tokens=LOCAL, + kv_cache_manager=mgr, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, TPB), + ) + sibling = kit.published(mgr, SOURCE, SIBLING_PROMPT) + lender = attach(mgr, tokens_per_fetch=TPB) + request = kit.make_request(TARGET, prompt) + assert schedule(scheduler, request) == [TARGET] + kv = kit.kv(mgr, request) + assert kv.num_committed_tokens == len(SHARED), "no partial match: the check proves nothing" + kv.enable_swa_scratch_reuse = False # a fetch target needs it off + assert runs_its_chunk(kit, [mgr], request) is None + position = len(SHARED) + LOCAL + assert request.context_current_position == position, "the first chunk was not LOCAL tokens" + request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS # parked to fetch + fetch_segment(kit, mgr, lender, request, 3 * TPB, 4 * TPB, byte) + resume_bounds = lender.resume_bounds(request) + assert tuple(resume_bounds) == (4 * TPB, 3 * TPB), f"resume_bounds {tuple(resume_bounds)}" + at = resume_bounds.min_position if below else SHIFTED_AT + request.state = LlmRequestState.CONTEXT_INIT + try: + resume(kit, [mgr], request, at) + except RuntimeError as error: + raise AssertionError(f"the step at {at} raised: {error}") from error + try: + scheduled = schedule(scheduler, request) + except RuntimeError as error: + raise AssertionError(f"the next scheduling pass raised: {error}") from error + assert scheduled == [TARGET] + assert request.context_current_position == at, ( + f"the next chunk starts at {request.context_current_position}, not at {at}" + ) + assert request.prepopulated_prompt_len == at, ( + f"the prepopulated length stays at {request.prepopulated_prompt_len}, not at {at}" + ) + block_0 = kit.chain_keys(kit.kv(mgr, sibling), SIBLING_PROMPT)[0] + + # what the sibling's prefill wrote, not read back from its page + def written(dev, lg, ordinal): + return kit.fake_page(lg, block_0, dev.page_bytes(lg)) + + for lg in range(kit.num_layer_groups(mgr)): + assert holds(kit, mgr, request, lg, [0], written), f"group {lg}: block 0" + assert holds(kit, mgr, request, lg, [1, 2], chunk_bytes), f"group {lg}: blocks 1-2" + assert holds(kit, mgr, request, lg, [3], staged_bytes(kit, byte)), ( + f"group {lg}: the fetched block lacks its staged bytes" + ) + error = runs_its_chunk(kit, [mgr], request) + assert error is None, f"the context update did not take the chunk at {at}: {error}" + assert kv.history_length == request.prompt_len + + +# -- failures at the call --------------------------------------------------------------------- + + +@pytest.mark.parametrize("state", ["no_cache", "suspended", "shut_down"]) +def test_a_lease_failed_at_the_call_changed_nothing(kit, real_manager, state): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + if state == "no_cache": + source, target = kit.make_request(8, PROMPT), kit.make_request(9, OTHER_PROMPT) + elif state == "suspended": + for request in (source, target): + mgr.suspend_request(request) + assert not kit.kv(mgr, request).is_active + else: + mgr.shutdown() + before = [state_of(kit, mgr, lender, r) for r in (source, target)] + leases = [lender.lend_read(source, 0, END), lender.lend_write(target, 0, END)] + for lease in leases: + assert lease.failure is not None, "not failed at the call" + assert lease.poll() is None + assert [state_of(kit, mgr, lender, r) for r in (source, target)] == before + for lease in leases: + lease.release() + lease.release() + + +def test_a_write_the_pool_cannot_grow_fails_at_the_call_and_changed_nothing(kit, real_manager): + with real_manager(max_tokens=kit.POOL_TOKENS) as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + others = kit.Requests(mgr) + assert others.allocate(kit.pool_pages(mgr)), "others could not take the whole pool" + before = state_of(kit, mgr, lender, target) + lease = lender.lend_write(target, 0, END) + assert "no free pages" in str(lease.failure) and lease.poll() is None + assert state_of(kit, mgr, lender, target) == before + lease.release() + others.free() + lease = lender.lend_write(target, 0, END) # with pages free again it goes through + view = lease.poll() + assert view is not None + lease.mark_arrived(view.row_masks()) + lease.release() + + +def test_a_write_the_pool_cannot_grow_leaves_a_computed_prefix_as_it_was( + kit, real_manager, attach=attach +): + """Per request: the chunk ``[0, LOCAL)`` computed and not committed, then a fetch of + ``[TPB, 3 * TPB)``, from below the history, that the pool cannot grow the cache for. The lease + fails at the call, and the cache and its resume_bounds stay as they were: the chunk still counts.""" + with real_manager(max_tokens=kit.POOL_TOKENS, **PER_REQUEST) as mgr: + target = computed(kit, mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=2 * TPB) + others = kit.Requests(mgr) + free = kit.pool_pages(mgr) - int(kit.kv(mgr, target).num_blocks) # past the chunk's pages + assert others.allocate(free), "others could not take the rest of the pool" + before = state_of(kit, mgr, lender, target) + assert before[1] == (LOCAL, LOCAL), f"resume_bounds before the fetch {before[1]}" + lease = lender.lend_write(target, TPB, 3 * TPB) + assert "no free pages" in str(lease.failure) and lease.poll() is None, lease.failure + after = state_of(kit, mgr, lender, target) + assert after == before, f"the failed grow changed the cache or its resume_bounds: {after}" + lease.release() + others.free() + + +def test_a_write_left_with_a_block_without_a_page_fails_at_its_first_poll( + kit, real_manager, monkeypatch +): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + kv = kit.kv(mgr, target) + + def without_block_one(read): + def pages(kv_cache, lg): + found = np.array(read(kv_cache, lg), dtype=np.int64) + if kv_cache is kv and len(found) > 1: + found[1] = -1 + return found + + return pages + + with monkeypatch.context() as patched: + for name in ("pages", "locked_pages"): + patched.setattr(_manager, name, without_block_one(getattr(_manager, name))) + lease = lender.lend_write(target, 0, END) + assert lease.failure is None, "the cache grew: nothing fails at the call after that" + assert lender.resume_bounds(target) is None + assert lease.poll() is None and lease.failure is not None + assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + lease.release() + + +GROWN_THEN_FAILED = "the write failed at the call after its cache grew" + + +def test_a_write_whose_grant_raises_within_the_call_fails_at_its_first_poll( + kit, real_manager, monkeypatch, attach=attach +): + """``lend_write`` grows the cache before taking slots, so when its grant raises, the lease it + returns has not failed and fails at its first poll, abandoning the fetch like one with a block + left without a page. A read changes no cache, so one whose grant raises fails at the call.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + view = _lender.Staging._build_view + planted = [] + + def raising_once(self, rows, keys): + if not planted: + planted.append(True) + raise MemoryError("planted") + return view(self, rows, keys) + + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + kv = kit.kv(mgr, target) + with monkeypatch.context() as patched: + patched.setattr(_lender.Staging, "_build_view", raising_once) + read = lender.lend_read(source, 0, END) + assert planted and "planted" in str(read.failure), "the read's grant did not raise" + planted.clear() + lease = lender.lend_write(target, 0, END) + assert planted, "the write's grant ran after the call: the check proves nothing" + assert kv.capacity >= END, "the write did not grow the cache" + assert lease.failure is None, f"{GROWN_THEN_FAILED}: {lease.failure}" + assert lender.resume_bounds(target) is None + assert lease.poll() is None and "planted" in str(lease.failure), lease.failure + assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + lease.release() + read.release() + + +SHARED = list(range(7000, 7040)) # one whole block and eight tokens +SIBLING_PROMPT = SHARED + list(range(7100, 7160)) # block 1 holds other tokens past the eight +MATCHED_PROMPT = SHARED + list(range(7200, 7353)) # six whole blocks and one token +# name: the fetch's end, at which the window leaves block 1 behind or keeps it +MATCHED_FETCHES = {"window_leaves_it": 6 * TPB, "window_keeps_it": 3 * TPB} + + +@pytest.mark.parametrize("case", list(MATCHED_FETCHES)) +def test_a_fetch_from_a_partially_matched_block(kit, real_manager, case, attach=attach): + """All-reusable: the target's local match ends inside block 1, which its resume copied from a + sibling, so past the match the page holds the sibling's tokens. A fetch from block 1 whose + window leaves the block behind fails at the call, changing nothing, while one whose window keeps + it fetches the block, every row of which then holds its staged bytes. Either way native reuse + of the window group's block 1 later serves the target's bytes.""" + with real_manager(windows=[WINDOW, 256]) as mgr: + window_lg = kit.windows(mgr).index(WINDOW) + dev = kit.DevicePages(mgr) + sibling = kit.published(mgr, SOURCE, SIBLING_PROMPT) + sibling_bytes = kit.digest(kit.page(mgr, sibling, window_lg, 1)) + target = kit.admitted(mgr, TARGET, MATCHED_PROMPT) + kv = kit.kv(mgr, target) + assert kv.num_committed_tokens == len(SHARED), "no partial match: the check proves nothing" + mgr._stream.synchronize() # the resume's copy of the matched page + assert kit.digest(kit.page(mgr, target, window_lg, 1)) == sibling_bytes, ( + "block 1 holds no copy of the sibling's page: the check proves nothing" + ) + start, end = TPB, MATCHED_FETCHES[case] + stale_beg, stale_end = kit.stale_blocks(mgr, window_lg, end) + leaves = stale_beg <= 1 < stale_end + assert leaves == (case == "window_leaves_it") + lender = attach(mgr, tokens_per_fetch=end - start) + before = (kv.capacity, kv.history_length, kv.num_committed_tokens) + lease = lender.lend_write(target, start, end) + if lease.failure is None: + # Granted: the fetch lands whole, and the target resumes where resume_bounds allows. + view = ready_view(lease, mgr) + kit.stage(lender, view, 0x5A) + lease.mark_arrived(view.row_masks(True)) + lease.release() + mgr._stream.synchronize() + for run in view.groups: # every row the resume reads + ordinals = run.block_ordinals.tolist() + assert holds( + kit, mgr, target, run.layer_group_index, ordinals, staged_bytes(kit, 0x5A) + ), f"layer group {run.layer_group_index}: a fetched row lacks its staged bytes" + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position >= resume_bounds.min_position, tuple(resume_bounds) + resumed(kit, mgr, target, resume_bounds.min_position) + else: + assert leaves, f"the fetch failed although the window keeps block 1: {lease.failure}" + assert "never wrote" in lease.failure, lease.failure + lease.release() + assert (kv.capacity, kv.history_length, kv.num_committed_tokens) == before, ( + "the fetch that failed at the call changed the cache" + ) + # The target computes the rest of its prompt, and the context update commits it. + assert next_chunk(kit, mgr, target, tokens=len(MATCHED_PROMPT)) is None + assert kv.num_committed_tokens == len(MATCHED_PROMPT), "the prompt was not committed" + # A later request sharing the target's first three blocks takes block 1 from the tree. + later = kit.admitted(mgr, OTHER, MATCHED_PROMPT[: 3 * TPB + 1]) + slot = kit.pages(kit.kv(mgr, later), window_lg)[1] + if slot < 0: + # Native reuse holds no window page for the block, so it serves no bytes there. + assert lease.failure is None and leaves, "no window page for block 1: nothing checked" + return + got = kit.digest(dev.read(window_lg, slot)) + assert got != sibling_bytes, ( + "native reuse serves block 1 with the sibling's bytes past the target's match" + ) + expected = staged_bytes(kit, 0x5A) if lease.failure is None else chunk_bytes + assert got == kit.digest(expected(dev, window_lg, 1)), "block 1 holds other bytes" + + +def test_a_fetch_past_pages_grown_before_it_fails_at_the_call(kit, real_manager, attach=attach): + """All-reusable: the target's first chunk was sized before the fetch, so the blocks past its + history have pages nothing wrote. A fetch whose window leaves them behind fails at the call, + changing nothing, and native reuse of those blocks later serves what the request computed.""" + with real_manager(windows=[WINDOW, 256]) as mgr: + window_lg = kit.windows(mgr).index(WINDOW) + dev = kit.DevicePages(mgr) + target = kit.admitted(mgr, TARGET, MATCHED_PROMPT) + kv = kit.kv(mgr, target) + assert (kv.num_committed_tokens, kv.history_length) == (0, 0), "a local match" + # The scheduler sizes the first chunk before the fetch. + target.context_chunk_size = 4 * TPB + assert mgr.resize_context(target, 4 * TPB) + kit.fill_sentinel(mgr, target) # what a page nothing wrote holds + unwritten = kit.digest(bytes([kit.SENTINEL]) * dev.page_bytes(window_lg)) + end = 6 * TPB + stale_beg, stale_end = kit.stale_blocks(mgr, window_lg, end) + assert stale_beg <= 1 < stale_end and kit.pages(kv, window_lg)[1] >= 0, ( + "block 1 has no page or the window keeps it: the check proves nothing" + ) + lender = attach(mgr, tokens_per_fetch=end) + before = (kv.capacity, kv.history_length, kv.num_committed_tokens) + lease = lender.lend_write(target, 0, end) + if lease.failure is None: + # Granted: the fetch lands whole, and the target resumes where resume_bounds allows. + view = ready_view(lease, mgr) + kit.stage(lender, view, 0x5A) + lease.mark_arrived(view.row_masks(True)) + lease.release() + mgr._stream.synchronize() + for run in view.groups: # every row the resume reads + ordinals = run.block_ordinals.tolist() + assert holds( + kit, mgr, target, run.layer_group_index, ordinals, staged_bytes(kit, 0x5A) + ), f"layer group {run.layer_group_index}: a fetched row lacks its staged bytes" + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position >= resume_bounds.min_position, tuple(resume_bounds) + resumed(kit, mgr, target, resume_bounds.min_position) + else: + assert "never wrote" in lease.failure, lease.failure + lease.release() + assert (kv.capacity, kv.history_length, kv.num_committed_tokens) == before, ( + "the fetch that failed at the call changed the cache" + ) + # The target computes the rest of its prompt, and the context update commits it. + assert next_chunk(kit, mgr, target, tokens=len(MATCHED_PROMPT)) is None + assert kv.num_committed_tokens == len(MATCHED_PROMPT), "the prompt was not committed" + # A later request sharing the target's first three blocks takes block 1 from the tree. + later = kit.admitted(mgr, OTHER, MATCHED_PROMPT[: 3 * TPB + 1]) + slot = kit.pages(kit.kv(mgr, later), window_lg)[1] + if slot < 0: + # Native reuse holds no window page for the block, so it serves no bytes there. + assert lease.failure is None, "no window page for block 1: nothing checked" + return + got = kit.digest(dev.read(window_lg, slot)) + assert got != unwritten, ( + "native reuse serves block 1 with bytes nothing computed or fetched" + ) + expected = staged_bytes(kit, 0x5A) if lease.failure is None else chunk_bytes + assert got == kit.digest(expected(dev, window_lg, 1)), "block 1 holds other bytes" + + +def test_a_fetch_past_pages_grown_before_it_goes_through_where_the_commit_keeps_none( + kit, real_manager, attach=attach +): + """Per request: the manager commits only the snapshot a request's window needs, so pages a + window leaves behind are dropped, not kept for the commit. A fetch past pages the target's first + chunk was sized with goes through, and the target resumes with what it fetched.""" + with real_manager(windows=[WINDOW, 256], **PER_REQUEST) as mgr: + assert mgr.impl.commit_min_snapshot, "passed pages stay: the check proves nothing" + window_lg = kit.windows(mgr).index(WINDOW) + target = kit.admitted(mgr, TARGET, MATCHED_PROMPT) + kv = kit.kv(mgr, target) + # The scheduler sizes the first chunk before the fetch. + target.context_chunk_size = 4 * TPB + assert mgr.resize_context(target, 4 * TPB) + kit.fill_sentinel(mgr, target) + end = 6 * TPB + stale_beg, stale_end = kit.stale_blocks(mgr, window_lg, end) + assert stale_beg <= 1 < stale_end and kit.pages(kv, window_lg)[1] >= 0, ( + "block 1 has no page or the window keeps it: the check proves nothing" + ) + lender = attach(mgr, tokens_per_fetch=end) + lease = lender.lend_write(target, 0, end) + assert lease.failure is None, ( + f"the fetch failed although the commit keeps no page a window passed: {lease.failure}" + ) + view = ready_view(lease, mgr) + kit.stage(lender, view, 0x5A) + lease.mark_arrived(view.row_masks(True)) + lease.release() + mgr._stream.synchronize() + passed = kit.pages(kv, window_lg)[stale_beg:stale_end] + assert all(slot < 0 for slot in passed), f"the window's passed blocks kept pages: {passed}" + resume_bounds = lender.resume_bounds(target) + assert resume_bounds.max_position == end, tuple(resume_bounds) + resumed(kit, mgr, target, resume_bounds.max_position) + kept = [o for o in range(end // TPB) if not stale_beg <= o < stale_end] + for lg in range(kit.num_layer_groups(mgr)): + ordinals = kept if lg == window_lg else range(end // TPB) + assert holds(kit, mgr, target, lg, ordinals, staged_bytes(kit, 0x5A)), ( + f"layer group {lg}: the target resumes without the bytes it fetched" + ) + assert next_chunk(kit, mgr, target, tokens=len(MATCHED_PROMPT)) is None + + +@pytest.mark.parametrize("case", ["missed", "delivered", "computed"]) +def test_a_later_lease_past_window_rows_an_earlier_one_missed_fails_at_the_call( + kit, real_manager, case, attach=attach +): + """All-reusable, a window of ``WINDOW`` tokens: the request computed block 0, and the first + lease of a fetch, ``[32, 64)``, delivers block 1's full-attention row and, in "delivered" alone, + its window row. The later lease ``[64, 128)`` leaves blocks 0 and 1 behind its window, and the + commit stores block 1's pages whole. Where its window row was missed, nothing wrote that page + and the later lease fails at the call, changing nothing; where it was delivered, or the request + computed block 1 from where resume_bounds let it resume, it is lent. Native reuse of blocks 0 and 1 + later serves what the request computed or fetched.""" + with real_manager(windows=[WINDOW, 256]) as mgr: + sliding = kit.windows(mgr).index(WINDOW) + dev = kit.DevicePages(mgr) + target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) + kv = kit.kv(mgr, target) + assert kv.num_committed_tokens == 0, "a local match: the check proves nothing" + # Block 0 first, so the check starts past it and reads block 1's delivered row there. + assert next_chunk(kit, mgr, target, tokens=TPB) is None + assert kv.num_committed_tokens == TPB, "block 0 was not committed" + lender = attach(mgr, tokens_per_fetch=2 * TPB) + first = lender.lend_write(target, TPB, 2 * TPB) + kit.fill_sentinel(mgr, target, 1) # what a page nothing wrote holds + view = ready_view(first, mgr) + kit.stage(lender, view, 0x11) + delivered = case == "delivered" + first.mark_arrived( + tuple( + np.full(len(run), delivered or run.layer_group_index != sliding) + for run in view.groups + ) + ) + first.release() + mgr._stream.synchronize() + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == ((2 if delivered else 1) * TPB, TPB), tuple(resume_bounds) + if case == "computed": # from where resume_bounds lets it resume, through the lease's end + resumed(kit, mgr, target, resume_bounds.max_position) + assert next_chunk(kit, mgr, target, tokens=TPB) is None + assert kv.num_committed_tokens == 2 * TPB, "the chunk was not committed" + assert kit.stale_blocks(mgr, sliding, 4 * TPB) == (0, 2), "nothing left behind: no check" + assert min(kit.pages(kv, sliding)[:2]) >= 0, "blocks 0 and 1 have no page: no check" + before = state_of(kit, mgr, lender, target) + later = lender.lend_write(target, 2 * TPB, 4 * TPB) + if case == "missed": + assert later.failure is not None and "never wrote" in later.failure, ( + f"a later lease was lent past window rows an earlier one missed: {later.failure}" + ) + later.release() + assert state_of(kit, mgr, lender, target) == before, ( + "the failed lease changed the cache" + ) + else: + assert later.failure is None, f"the later lease failed: {later.failure}" + view = ready_view(later, mgr) + kit.stage(lender, view, 0x22) + later.mark_arrived(view.row_masks(True)) + later.release() + mgr._stream.synchronize() + resume_bounds = lender.resume_bounds(target) + assert resume_bounds == (4 * TPB, 4 * TPB), tuple(resume_bounds) + # The target computes the rest of its prompt from where resume_bounds lets it resume, and the + # context update commits it. + resumed(kit, mgr, target, resume_bounds.max_position) + assert next_chunk(kit, mgr, target, tokens=len(SPLIT_PROMPT)) is None + assert kv.num_committed_tokens == len(SPLIT_PROMPT), "the prompt was not committed" + # A later request sharing the target's first two blocks takes them from the tree. + other = kit.admitted(mgr, OTHER, SPLIT_PROMPT[: 2 * TPB] + OTHER_PROMPT[:33]) + unwritten = kit.digest(bytes([kit.SENTINEL]) * dev.page_bytes(sliding)) + for ordinal in (0, 1): + slot = kit.pages(kit.kv(mgr, other), sliding)[ordinal] + assert slot >= 0, f"native reuse took no window page for block {ordinal}: no check" + got = kit.digest(dev.read(sliding, slot)) + assert got != unwritten, ( + f"native reuse serves block {ordinal} with bytes nothing computed or fetched" + ) + expected = staged_bytes(kit, 0x11) if delivered and ordinal == 1 else chunk_bytes + assert got == kit.digest(expected(dev, sliding, ordinal)), f"block {ordinal} differs" + + +# -- argument errors -------------------------------------------------------------------------- + + +def test_a_read_refuses_a_bad_range_at_the_call(kit, real_manager): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + kit.published(mgr, 7, OTHER_PROMPT[: 2 * TPB] + list(range(9000, 9033))) + sized = kit.admitted(mgr, TARGET, OTHER_PROMPT) # its first two blocks match locally + # Sized for the rest of its prompt, as before its next chunk: block 2 has a page, no KV. + assert mgr.resize_context(sized, sized.context_remaining_length) + assert kit.kv(mgr, sized).num_committed_tokens == 2 * TPB + lender = attach(mgr, tokens_per_fetch=4 * TPB) + for request, start, end, why in ( + (source, -TPB, TPB, "bad token range"), + (source, TPB, 0, "bad token range"), + (source, 0, 50, "whole blocks"), + (source, 16, 48, "whole blocks"), + (source, 0, 4 * TPB, "past the 97 committed tokens"), + (source, 0, 1 << 50, "past the 97 committed tokens"), # of any length + (sized, 0, END, "past the 64 committed tokens"), # inside the prompt's whole blocks + ): + with pytest.raises(ValueError, match=why): + lender.lend_read(request, start, end) + lease = lender.lend_read(source, 0, END) + ready_view(lease, mgr) + lease.release() + + +def test_a_write_refuses_a_bad_range_before_changing_anything(kit, real_manager): + with real_manager() as mgr: + kit.published(mgr, 7, PROMPT[: 2 * TPB] + list(range(9000, 9033))) + matched = kit.admitted(mgr, TARGET, PROMPT) # its first two blocks match locally + fresh = kit.admitted(mgr, OTHER, OTHER_PROMPT) + assert kit.kv(mgr, matched).num_committed_tokens == 2 * TPB + lender = attach(mgr, tokens_per_fetch=4 * TPB, planned_fetches=2) + before = [state_of(kit, mgr, lender, r) for r in (matched, fresh)] + bad = [ + (fresh, -TPB, TPB, "bad token range"), + (fresh, TPB, 0, "bad token range"), + (fresh, 0, 50, "covers whole blocks"), + (fresh, 16, 48, "covers whole blocks"), + (matched, TPB, END, "inside the committed whole blocks"), + (fresh, 0, 4 * TPB, "past the whole blocks before the request's last prompt token"), + ] + for request, start, end, why in bad: + with pytest.raises(ValueError, match=why): + lender.lend_write(request, start, end) + assert [state_of(kit, mgr, lender, r) for r in (matched, fresh)] == before + leases = [lender.lend_write(matched, 2 * TPB, END), lender.lend_write(fresh, 0, END)] + unsettled = state_of(kit, mgr, lender, fresh) + again = lender.lend_write(fresh, 0, END) # one unsettled fetch per cache, on this rank + assert again.failure is not None and "unsettled" in again.failure, again.failure + assert again.poll() is None and state_of(kit, mgr, lender, fresh) == unsettled + again.release() + for lease in leases: + view = lease.poll() + lease.mark_arrived(view.row_masks()) + lease.release() + + +ALIGNED_PROMPT = list(range(11000, 11128)) # exactly four whole blocks + + +@pytest.mark.parametrize("windowed", [False, True], ids=["full", "windowed"]) +def test_a_fetch_ends_at_the_whole_blocks_before_the_last_prompt_token(kit, real_manager, windowed): + # The request computes its last prompt token itself, for its logits: on a prompt of whole + # blocks, fetching all of them would leave it a resume at the prompt's end. A range of any + # length past them is refused alike. + with real_manager(windows=[WINDOW, 256] if windowed else None) as mgr: + target = kit.admitted(mgr, TARGET, ALIGNED_PROMPT) + end = len(ALIGNED_PROMPT) + lender = attach(mgr, tokens_per_fetch=end) + before = state_of(kit, mgr, lender, target) + for past in (end, 1 << 50): # 1 << 50: whole blocks whose ordinals would take 256 TiB + with pytest.raises(ValueError, match="past the whole blocks before"): + lender.lend_write(target, 0, past) + assert state_of(kit, mgr, lender, target) == before + lease = lender.lend_write(target, 0, end - TPB) # the block of the last token left out + view = lease.poll() + assert view is not None + lease.mark_arrived(view.row_masks()) + lease.release() + + +def test_a_replaced_cache_takes_a_new_fetch(kit, real_manager): + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + first = lender.lend_write(target, 0, END) + first_view = first.poll() # unsettled: seen ready, not marked + mgr.free_resources(target) + target = kit.admitted(mgr, TARGET, PROMPT) + second = lender.lend_write(target, 0, END) + view = second.poll() + assert view is not None and lender.resume_bounds(target) is None + second.mark_arrived(view.row_masks()) + first.mark_arrived(first_view.row_masks()) + for lease in (first, second): + lease.release() + + +def test_a_windowed_target_with_scratch_reuse_fails_at_the_call(kit, real_manager, attach=attach): + with real_manager(windows=[WINDOW, 256], swa_scratch_reuse=True) as mgr: + target = kit.make_request(TARGET, WINDOWED_PROMPT) + assert mgr.prepare_context(target) + kv = kit.kv(mgr, target) + assert kv.enable_swa_scratch_reuse, "the target must start with scratch reuse on" + lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + before = state_of(kit, mgr, lender, target) + try: + lease = lender.lend_write(target, 0, WINDOWED_END) + except RuntimeError as error: # the manager refusing the grow a lender let through + raise AssertionError(f"not failed at the call: {error}") from error + assert lease.failure is not None and "scratch" in lease.failure, "not failed at the call" + assert lease.poll() is None + assert state_of(kit, mgr, lender, target) == before + lease.release() + kv.enable_swa_scratch_reuse = False + lease = lender.lend_write(target, 0, WINDOWED_END) # with it off the write goes through + view = lease.poll() + assert view is not None + lease.mark_arrived(view.row_masks()) + lease.release() + + +def test_a_windowless_target_with_scratch_reuse_fails_at_the_call(kit, real_manager, attach=attach): + """Without a window too, a cache with SWA scratch reuse on keeps its history within the scratch + rewind of its old capacity at every capacity change, which a fetch's grow breaks: the fetch + fails at the call and changes nothing, and with scratch reuse off it goes through.""" + with real_manager(swa_scratch_reuse=True) as mgr: + assert not any(kit.windows(mgr)), "a window: the windowed check covers that" + target = kit.make_request(TARGET, SPLIT_PROMPT) + assert mgr.prepare_context(target) + kv = kit.kv(mgr, target) + assert kv.enable_swa_scratch_reuse, "the target must start with scratch reuse on" + lender = attach(mgr, tokens_per_fetch=PAIRED_END) + before = state_of(kit, mgr, lender, target) + try: + lease = lender.lend_write(target, 0, PAIRED_END) + except RuntimeError as error: # the manager refusing the grow a lender let through + raise AssertionError(f"not failed at the call: {error}") from error + try: + failure = lease.failure + assert failure is not None, "not failed at the call" + assert "scratch rewind of the old capacity" in failure, failure + assert lease.poll() is None + assert state_of(kit, mgr, lender, target) == before, "the fetch changed the cache" + finally: + lease.release() + kv.enable_swa_scratch_reuse = False + lease = lender.lend_write(target, 0, PAIRED_END) # with it off the write goes through + view = lease.poll() + assert view is not None + lease.mark_arrived(view.row_masks()) + lease.release() + + +def test_a_fetch_into_a_request_returning_context_outputs_fails_at_the_call( + kit, real_manager, attach=attach +): + """A request returning context logits, as prompt logprobs make it do, fails at the call whether + it holds some or none, as does one asking for additional model outputs, caches unchanged: the + executor gives these only for positions a request computes. One returning none is lent.""" + with real_manager() as mgr: + held = kit.make_request(TARGET, PROMPT, return_context_logits=True) + fresh = kit.make_request(OTHER, OTHER_PROMPT, return_context_logits=True) + outputs = kit.make_request(OTHER + 1, CONTROL_PROMPT, additional_outputs=["context_output"]) + plain = kit.make_request(OTHER + 2, SPLIT_PROMPT) + for request in (held, fresh, outputs, plain): + assert mgr.prepare_context(request) + held.py_result.append_context_logits(torch.zeros(TPB, 1, 8, device="cuda")) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + for request, returns in ((held, True), (fresh, True), (outputs, True), (plain, False)): + before = state_of(kit, mgr, lender, request) + lease = lender.lend_write(request, 0, END) + if returns: + assert lease.failure is not None, "not failed at the call" + assert "context" in lease.failure, lease.failure + assert lease.poll() is None + assert state_of(kit, mgr, lender, request) == before, "the lease changed the cache" + else: + assert lease.failure is None, "refused a request returning no context outputs" + lease.mark_arrived(ready_view(lease, mgr).row_masks()) + lease.release() + + +@pytest.mark.parametrize("policy", ["per_request", "per_conversation"]) +def test_a_fetch_ending_below_a_history_the_update_moved_fails_at_the_call( + kit, real_manager, policy, attach=attach +): + """Full attention under a policy whose context update moves the history: the request computed + ``[0, 128)`` without a commit, so its history stands past a fetch of ``[64, 96)``. The floor + would follow it past all the fetch brings, an empty interval dropping the computed tokens; the + fetch fails at the call instead, the cache unchanged, and the request resumes at its history.""" + conversation = "history-past" if policy == "per_conversation" else None + with real_manager(block_reuse_config={"policy": policy}) as mgr: + target = computed(kit, mgr, TARGET, SPLIT_PROMPT, upto=4 * TPB, conversation=conversation) + lender = attach(mgr, tokens_per_fetch=TPB) + before = state_of(kit, mgr, lender, target) + lease = lender.lend_write(target, 2 * TPB, 3 * TPB) + try: + assert lease.failure is not None, f"not failed at the call, from {before}" + assert "stands past" in lease.failure, lease.failure + assert lease.poll() is None + assert state_of(kit, mgr, lender, target) == before, "the fetch changed the cache" + finally: + lease.release() + assert lender.resume_bounds(target) == (4 * TPB, 4 * TPB) + + +def test_a_draft_pool_fetch_ending_below_its_history_fails_at_the_call( + kit, real_manager, attach=attach +): + """Windowless under the all-reusable policy, a joint-reuse draft pool's floor follows its + history, so a fetch of ``[64, 96)`` into a cache whose history stands at 128 fails at the call, + the cache unchanged, while a target pool in the same state lends. Each pool's context capacity + is allocated as admission does and its history moved to 128 without a commit: such a pool's + context update commits every chunk it runs, so no native path leaves the history there.""" + spec, at = dflash(), 4 * TPB + with ( + real_manager(spec_config=spec, joint_kv_cache_reuse=True) as target, + real_manager(spec_config=spec, is_draft=True, joint_kv_cache_reuse=True) as draft, + ): + failures = [] + for mgr, scope in ((target, SCOPE), (draft, SCOPE + b"/draft-pool")): + request = kit.admitted(mgr, TARGET, SPLIT_PROMPT) + request.context_chunk_size = at + if mgr.is_draft: + assert mgr.try_allocate_draft_context(request, at) + else: + assert mgr.resize_context(request, at) + kv = kit.kv(mgr, request) + kv.resize(None, at) + assert (kv.history_length, kv.num_committed_tokens) == (at, 0) + lender = attach(mgr, tokens_per_fetch=TPB, scope=scope) + before = state_of(kit, mgr, lender, request) + lease = lender.lend_write(request, 2 * TPB, 3 * TPB) + try: + failures.append(lease.failure) + if lease.failure is not None: + assert state_of(kit, mgr, lender, request) == before, "the fetch changed it" + finally: + lease.release() + assert failures[0] is None, f"the target pool's fetch failed at the call: {failures[0]}" + assert failures[1] is not None and "stands past" in failures[1], ( + f"the draft pool's fetch below its history was not failed at the call: {failures[1]}" + ) + + +def test_a_range_longer_than_staging_holds_is_refused(kit, real_manager, attach=attach): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=2 * TPB) # two slots + before = state_of(kit, mgr, lender, target) + assert "more staging slots" in refusal(lambda: lender.lend_read(source, 0, END)) + assert lender._open_lease_count() == 0, "a refused publish left a lease open" + assert "more staging slots" in refusal(lambda: lender.lend_write(target, 0, END)) + assert state_of(kit, mgr, lender, target) == before + # A range of tokens_per_fetch tokens fits exactly. + read = lender.lend_read(source, 0, 2 * TPB) + ready_view(read, mgr) + read.release() + write = lender.lend_write(target, 0, 2 * TPB) + view = write.poll() + assert view is not None + write.mark_arrived(view.row_masks()) + write.release() + + +# What this rank holds for a request that fails a call whose range is right: no KV cache, a reset +# of the manager's reuse state, and multimodal data without digests. +REQUEST_STATES = ["no_cache", "reuse_reset", "unnamed"] + + +@pytest.mark.parametrize("state", REQUEST_STATES) +def test_a_wrong_call_raises_whatever_the_request_s_state(kit, real_manager, state, attach=attach): + """A range that is not whole blocks or, for a fetch, ends past the whole blocks before the + request's last prompt token or needs more slots than a region has is the caller's error, the same + on every rank: on a manager not shut down it raises ``ValueError`` whatever this rank holds for + the request, also where a right range fails at the call, and changes nothing.""" + with real_manager() as mgr: + lender = attach(mgr, tokens_per_fetch=2 * TPB) # two slots + if state == "no_cache": + target = kit.make_request(TARGET, OTHER_PROMPT) + elif state == "reuse_reset": + mgr.reset_reuse_state() # what reset_prefix_cache does when a weight update ends + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + else: + target = carrying(kit, TARGET, OTHER_PROMPT, dict(py_multimodal_data=IMAGE)) + assert mgr.prepare_context(target) + before = state_of(kit, mgr, lender, target) + for lend, end, why in ( + (lender.lend_read, 50, "whole blocks"), + (lender.lend_write, 50, "whole blocks"), + (lender.lend_write, 4 * TPB, "past the whole blocks before the request's last prompt"), + (lender.lend_write, END, "more staging slots"), + ): + assert why in refusal(lambda: lend(target, 0, end)) + assert state_of(kit, mgr, lender, target) == before, "a wrong call changed the cache" + lease = lender.lend_write(target, 0, 2 * TPB) # a right range: the request's state decides + try: + assert lease.failure is not None and lease.poll() is None, ( + f"a right range went through for a request in state {state}: this proves nothing" + ) + finally: + lease.release() + + +# -- resume_bounds -------------------------------------------------------------------------------- + + +def test_resume_bounds_without_a_fetch_is_the_cache_s_own_interval(kit, real_manager): + with real_manager(windows=[WINDOW, 256]) as mgr: + source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) + lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + kv = kit.kv(mgr, source) + resume_bounds = lender.resume_bounds(source) + assert resume_bounds == (max(kv.num_committed_tokens, kv.history_length), kv.history_length) + assert all(type(v) is int for v in resume_bounds) + with pytest.raises(ValueError): + lender.resume_bounds(kit.make_request(9, PROMPT)) + + +def test_resume_bounds_waits_on_nothing_and_settles_by_being_asked(kit, real_manager): + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + lease = lender.lend_write(target, 0, END) + view = lease.poll() + with kit.held_stream(mgr._stream) as gate: + lease.mark_arrived(view.row_masks(True)) + assert lender.resume_bounds(target) is None + assert lender.resume_bounds(target) is None + gate.open() + assert lender.resume_bounds(target) == (END, 0) + lease.release() + + +# -- a reset of the manager's reuse state ---------------------------------------------------- + +RESET = "reset its reuse state" + + +def test_a_reuse_reset_stops_lending_by_name(kit, real_manager, attach=attach): + # After an in-place weight update the manager computes other bytes for the same tokens, whose + # names stay what they were: nothing may be lent by name any more. + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + publish = lender.lend_read(source, 0, END) # granted: its copy is queued + fetch = lender.lend_write(target, 0, END) + fetch_view = fetch.poll() # seen ready, not marked + for request in (source, target): # the manager resets only with every cache closed + mgr.free_resources(request) + mgr.reset_reuse_state() # what reset_prefix_cache does when a weight update ends + (run,) = ready_view(publish, mgr).groups # granted before the reset: it finishes + length = lender.regions[run.region_index].slot_bytes + staged = [kit.host_bytes(a, length) for a in run.addresses.tolist()] + assert kit.digest(staged) == kit.digest(original), "the granted publish lost its bytes" + fetch.mark_arrived(fetch_view.row_masks(True)) # its request is gone: nothing is copied + # The same tokens, computed after the reset, would go out under the names of the old bytes. + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + before = [state_of(kit, mgr, lender, r) for r in (source, target)] + later = [lender.lend_read(source, 0, END), lender.lend_write(target, 0, END)] + for lease in later: + assert lease.poll() is None and RESET in str(lease.failure), ( + "lent by name after the reset" + ) + assert [state_of(kit, mgr, lender, r) for r in (source, target)] == before + for lease in (publish, fetch, *later): + lease.release() + + +# -- shutdown --------------------------------------------------------------------------------- + + +def test_the_manager_s_shutdown_waits_for_the_lender_s_copies(kit, real_manager, monkeypatch): + """At the shutdown the lender waits for its copies before it frees the staging memory they + write. The manager's own teardown waits on the stream as well, so the order shows only at the + free.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + free = _lender._HostMemory.free + frees = [] # per free of staging memory: whether the stream had run the gated copy by then + + def recording_free(memory): + frees.append(mgr._stream.query()) + gate.open() + mgr._stream.synchronize() # nothing is freed under a copy, also for a lender that is wrong + free(memory) + + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + regions = lender.regions + monkeypatch.setattr(_lender._HostMemory, "free", recording_free) + with kit.held_stream(mgr._stream, strict=False) as gate: + lease = lender.lend_read(source, 0, END) + event = lease._copy._event + assert event is not None + synchronize = event.synchronize + + def finish_copy(): + gate.open() + synchronize() + + monkeypatch.setattr(event, "synchronize", finish_copy) + lease.release() # its copy waits behind the gate until shutdown waits on its event + mgr.shutdown() + assert frees == [True], "the staging memory was freed before the copy into it ran" + assert not kit.staging_kept(regions), "no lease was open: the memory is freed" + + +def test_shutdown_fails_waiting_leases_and_later_calls_only_end_records(kit, real_manager): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + held = lender.lend_read(source, 0, END) + ready_view(held, mgr) + waiting = lender.lend_read(source, 0, END) + assert waiting.poll() is None + mgr.shutdown() + assert waiting.failure is not None and waiting.poll() is None + for lease in (lender.lend_read(source, 0, END), lender.lend_write(target, 0, END)): + assert lease.failure is not None and lease.poll() is None + lease.release() + with pytest.raises(ValueError): + lender.resume_bounds(source) + for lease in (held, held, waiting): + lease.release() + mgr.shutdown() # a second shutdown frees nothing and returns + + +@pytest.mark.parametrize("failed", [False, True], ids=["ready_lease", "failed_lease"]) +def test_an_unreleased_lease_keeps_the_staging_memory_until_exit( + kit, real_manager, monkeypatch, failed, attach=attach +): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + regions = lender.regions + assert kit.staging_memory(regions) + if failed: + lease = lender.lend_read(kit.make_request(9, PROMPT), 0, END) + assert lease.failure is not None + address, content = None, None + else: + lease = lender.lend_read(source, 0, END) + address = int(ready_view(lease, mgr).groups[0].addresses[0]) + content = kit.host_bytes(address, regions[0].slot_bytes) + warnings = kit.lender_warnings(monkeypatch) + mgr.shutdown() + assert kit.staging_kept(regions), "staging freed at shutdown with a lease open" + assert warnings, "keeping the memory until exit is logged" + if address is not None: + got = kit.digest(kit.host_bytes(address, regions[0].slot_bytes)) + assert got == kit.digest(content), "the kept staging memory changed" + lease.release() + assert kit.staging_kept(regions), "a release after the shutdown freed it" + + +def test_a_shutdown_hook_that_fails_keeps_the_memory_and_raises_nothing( + kit, real_manager, monkeypatch +): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + def planted(self): + raise RuntimeError("planted") + + with real_manager() as mgr: + regions = attach(mgr, tokens_per_fetch=END).regions + monkeypatch.setattr(_lender.Staging, "_must_retain_memory", planted) + mgr.shutdown() + assert kit.staging_kept(regions) + + +# -- regions holds ------------------------------------------------------------------------------ + + +def released(hold) -> None: + """``hold.release()``; the check fails if it raises, since release is legal in every state.""" + try: + hold.release() + except Exception as error: # handed to the check as its failure + raise AssertionError(f"release raised {error!r}") from error + + +def test_a_hold_dropped_unreleased_still_keeps_the_staging_memory(kit, real_manager, attach=attach): + with real_manager() as mgr: + lender = attach(mgr, tokens_per_fetch=END) + regions = lender.regions + assert kit.staging_memory(regions) + lender.hold_regions() # dropped unreleased + gc.collect() + mgr.shutdown() + assert kit.staging_kept(regions), "a dropped hold let the memory go" + + +def test_a_hold_releases_once_in_every_state(kit, real_manager, attach=attach): + with real_manager() as mgr: + lender = attach(mgr, tokens_per_fetch=END) + regions = lender.regions + assert kit.staging_memory(regions) + first, second = lender.hold_regions(), lender.hold_regions() + released(first) + released(first) # does nothing: the second hold still holds + mgr.shutdown() + assert kit.staging_kept(regions), "a second release let the other hold go" + late = lender.hold_regions() # after the shutdown: inert + assert isinstance(late, RegionHold) + for hold in (late, late): + released(hold) + gone = weakref.ref(lender) + del mgr, lender + gc.collect() + assert gone() is None, "the lender outlived its manager" + for hold in (second, second): # its lender is gone + released(hold) + + +# What is released before the manager's shutdown, in order; the rest is released after it. +ORDERS = { + "neither": (), + "lease": ("lease",), + "hold": ("hold",), + "lease_then_hold": ("lease", "hold"), + "hold_then_lease": ("hold", "lease"), +} + + +@pytest.mark.parametrize("order", list(ORDERS.values()), ids=list(ORDERS)) +def test_a_lease_and_a_hold_keep_the_memory_until_both_go( + kit, real_manager, monkeypatch, order, attach=attach +): + """A ready lease and a regions hold released in ``order`` before the shutdown, the rest after it, + and a failed lease released at once: only with both gone does the shutdown free the memory.""" + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + regions = lender.regions + refs = [weakref.ref(m) for m in kit.staging_memory(regions)] + assert refs + lender.lend_read(kit.make_request(9, PROMPT), 0, END).release() # failed, released + hold = lender.hold_regions() + assert isinstance(hold, RegionHold) + lease = lender.lend_read(source, 0, END) + ready_view(lease, mgr) + ends = {"lease": lease.release, "hold": lambda: released(hold)} + for what in order: + ends[what]() + warnings = kit.lender_warnings(monkeypatch) + mgr.shutdown() + both = len(order) == 2 + if both: + assert not kit.staging_kept(regions) and warnings == [], "kept with both released" + else: + assert kit.staging_kept(regions), f"freed at shutdown after releasing {order}" + assert warnings, "keeping the memory until exit is logged" + for what in [w for w in ("lease", "hold") if w not in order]: + ends[what]() + mgr.shutdown() + assert kit.staging_kept(regions) == (not both), ( + "a release or a shutdown after it changed it" + ) + gc.collect() + assert all(ref() is None for ref in refs) == both, "kept memory went, or freed memory stayed" + + +@pytest.mark.parametrize("order", list(ORDERS.values()), ids=list(ORDERS)) +def test_a_failed_copy_and_a_hold_keep_the_memory_in_every_order( + kit, real_manager, monkeypatch, order, attach=attach +): + def no_event(self, stream=None): + raise RuntimeError("planted") + + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + regions = lender.regions + (group,) = kit.pool_group_ids(mgr) + assert kit.staging_memory(regions) + hold = lender.hold_regions() + with monkeypatch.context() as patched: + patched.setattr(torch.cuda.Event, "record", no_event) + lease = lender.lend_read(source, 0, END) # its copy has no event: slots lost + assert lease.failure is not None + mgr._stream.synchronize() + ends = {"lease": lease.release, "hold": lambda: released(hold)} + for what in order: + ends[what]() + lender.resume_bounds(source) # any call makes progress + assert lender._free_slots(group) == 0, "slots a copy may still touch were reused" + warnings = kit.lender_warnings(monkeypatch) + mgr.shutdown() + assert kit.staging_kept(regions), f"freed at shutdown after releasing {order}" + assert warnings, "keeping the memory until exit is logged" + for what in [w for w in ("lease", "hold") if w not in order]: + ends[what]() + assert kit.staging_kept(regions), "a release after the shutdown freed it" + + +# -- progress --------------------------------------------------------------------------------- + + +class EventQueries: + """Counts ``torch.cuda.Event.query`` per event within each lender call made through ``call``.""" + + def __init__(self, monkeypatch): + self.asked: List[int] = [] + self.last: Set[int] = set() # the events the latest call asked + query = torch.cuda.Event.query + + def counted(event): + self.asked.append(id(event)) + return query(event) + + monkeypatch.setattr(torch.cuda.Event, "query", counted) + + def call(self, method, *args): + first = len(self.asked) + result = method(*args) + asked = self.asked[first:] + assert len(asked) == len(set(asked)), "a copy's event was asked twice in one lender call" + self.last = set(asked) + return result + + +HELD_READS = 4 + + +def test_each_call_asks_each_copy_s_event_at_most_once( + kit, real_manager, monkeypatch, attach=attach +): + """Each lender call asks each copy's event at most once, a lease still lent costs no query + beyond its own poll's, and a copy reported complete is never asked again: a call asks only its + own copy and the pending copies of released leases, not one per held lease.""" + with real_manager() as mgr: + sources = [ + kit.published(mgr, 10 + i, [10000 * (i + 1) + t for t in range(END + 1)]) + for i in range(HELD_READS) + ] + target = kit.admitted(mgr, TARGET, PROMPT) + assert mgr._resize_for_served_prefix(target, kit.kv(mgr, target), 0, END) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=HELD_READS + 1) + events = EventQueries(monkeypatch) + with kit.held_stream(mgr._stream) as gate: + reads = [events.call(lender.lend_read, source, 0, END) for source in sources] + write = events.call(lender.lend_write, target, 0, END) + view = events.call(write.poll) + events.call(write.mark_arrived, view.row_masks(True)) + events.call(write.release) + released = id(write._copy._event) + for _ in range(2): + for lease in reads: + assert events.call(lease.poll) is None + own = id(lease._copy._event) + assert events.last == {own, released}, ( + f"a read's poll asked {len(events.last)} copies, not its own and the " + "released one" + ) + assert events.call(lender.resume_bounds, target) is None + assert events.last == {released}, ( + f"resume_bounds asked {len(events.last)} copies, not the released one" + ) + gate.open() + mgr._stream.synchronize() + assert events.call(lender.resume_bounds, target) is not None + for lease in reads: + assert events.call(lease.poll) is not None + first = len(events.asked) + for _ in range(3): + for lease in reads: + assert events.call(lease.poll) is not None + assert events.call(lender.resume_bounds, target) is not None + assert len(events.asked) == first, "a copy reported complete was asked again" + for lease in reads: + events.call(lease.release) + assert len(events.asked) == first, "a copy reported complete was asked again" + + +# -- a copy that fails partway ---------------------------------------------------------------- + + +class FailingDriver: + """The CUDA driver, except that its ``fail_on``-th memcpy call, of any flavour, fails: it + returns an error, or raises ``raises("planted")``.""" + + def __init__(self, real_driver, fail_on, raises=None): + self._real = real_driver + self._fail_on = fail_on + self._raises = raises + self.calls = 0 + + def __getattr__(self, name): + found = getattr(self._real, name) + if not name.startswith("cuMemcpy"): + return found + + def call(*args, **kwargs): + self.calls += 1 + if self.calls == self._fail_on: + if self._raises is not None: + raise self._raises("planted") + return (self._real.CUresult.CUDA_ERROR_INVALID_VALUE,) + return found(*args, **kwargs) + + return call + + +def test_a_publish_whose_copy_fails_partway_fails_and_frees_its_slots_after_the_rest( + kit, real_manager, monkeypatch +): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + # Three layers: two pool groups, so the copy is several memcpys. + with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: + source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) + lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + groups = kit.pool_group_ids(mgr) + driver = FailingDriver(_lender.drv, fail_on=2) + monkeypatch.setattr(_lender, "drv", driver) + with kit.held_stream(mgr._stream) as gate: + lease = lender.lend_read(source, 0, WINDOWED_END) # every slot + assert driver.calls >= 2, "a single memcpy: nothing failed partway" + assert lease.failure is not None and lease.poll() is None + lease.release() + assert [lender._free_slots(g) for g in groups] == [0] * len(groups), ( + "slots came back while a copy into them was queued" + ) + gate.open() + lender.resume_bounds(source) # any call makes progress + assert [lender._free_slots(g) for g in groups] == [p.slot_count for p in lender.regions] + + +def test_marks_whose_copy_fails_partway_abandon_the_fetch_and_keep_the_lease_ready( + kit, real_manager, monkeypatch +): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: + target = kit.admitted(mgr, TARGET, WINDOWED_PROMPT) + lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + groups = kit.pool_group_ids(mgr) + lease = lender.lend_write(target, 0, WINDOWED_END) + view = lease.poll() + kit.stage(lender, view, 0x5A) + driver = FailingDriver(_lender.drv, fail_on=2) + monkeypatch.setattr(_lender, "drv", driver) + kv = kit.kv(mgr, target) + with kit.held_stream(mgr._stream) as gate: + lease.mark_arrived(view.row_masks(True)) # raises nothing + assert driver.calls >= 2, "a single memcpy: nothing failed partway" + assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + assert lease.poll() is view, "the lease stays ready" + lease.release() + assert [lender._free_slots(g) for g in groups] == [0] * len(groups) + gate.open() + lender.resume_bounds(target) + assert [lender._free_slots(g) for g in groups] == [p.slot_count for p in lender.regions] + + +# -- errors inside the lender ----------------------------------------------------------------- + + +def test_an_event_query_that_raises_returns_no_slot_twice( + kit, real_manager, monkeypatch, attach=attach +): + """Two released publishes' copies have run, and the slot return's query of the later one's + event raises: that call raises having returned no slot, and the next one returns every slot + once.""" + with real_manager() as mgr: + first = kit.published(mgr, SOURCE, PROMPT) + second = kit.published(mgr, OTHER, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + (group,) = kit.pool_group_ids(mgr) + with kit.held_stream(mgr._stream): + earlier = lender.lend_read(first, 0, END) # granted and released first: asked first + later = lender.lend_read(second, 0, END) + earlier.release() + later.release() + assert lender._free_slots(group) == 0 + query = torch.cuda.Event.query + planted = later._copy._event + + def raises_once(event): + if event is planted: + monkeypatch.setattr(torch.cuda.Event, "query", query) + raise RuntimeError("planted") + return query(event) + + monkeypatch.setattr(torch.cuda.Event, "query", raises_once) + with pytest.raises(RuntimeError, match="planted"): + lender.resume_bounds(first) # any call makes progress + try: + lender.resume_bounds(first) + except ValueError as error: + raise AssertionError(f"a slot came back twice: {error}") from None + assert lender._free_slots(group) == lender.regions[0].slot_count + again = lender.lend_read(first, 0, END) + ready_view(again, mgr) + again.release() + + +def test_marks_that_raise_leave_the_fetch_abandoned_and_the_slots_returnable( + kit, real_manager, monkeypatch +): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + def planted(*args): + raise RuntimeError("planted") + + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + other = kit.admitted(mgr, OTHER, OTHER_PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + (group,) = kit.pool_group_ids(mgr) + lease = lender.lend_write(target, 0, END) + view = lease.poll() + kv = kit.kv(mgr, target) + with monkeypatch.context() as patched: + patched.setattr(_lender.Staging, "_locked_block_masks", planted) + with pytest.raises(RuntimeError, match="planted"): + lease.mark_arrived(view.row_masks(True)) + assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + lease.release() + waiting = lender.lend_write(other, 0, END) + waiting_view = waiting.poll() + assert waiting_view is not None, "the lease's slots never came back" + opened = lender._open_lease_count() + with monkeypatch.context() as patched: + patched.setattr(_lender.Staging, "_progress", planted) + with pytest.raises(RuntimeError, match="planted"): + waiting.release() # counted before the progress that raises + assert lender._open_lease_count() == opened - 1 + waiting.mark_arrived(waiting_view.row_masks()) + lender.resume_bounds(other) + assert lender._free_slots(group) == lender.regions[0].slot_count + + +def test_marks_whose_copy_raises_once_queued_lose_the_slots_it_reads( + kit, real_manager, monkeypatch, attach=attach, planted=MemoryError +): + """The marked rows' copy is queued, then creating its event raises: the mark raises, the fetch + stays abandoned, and the slots the copy reads are lost, not returned at the release.""" + + def no_event(*args, **kwargs): + raise planted("planted") + + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + (group,) = kit.pool_group_ids(mgr) + lease = lender.lend_write(target, 0, END) + view = lease.poll() + kv = kit.kv(mgr, target) + with kit.held_stream(mgr._stream): + with monkeypatch.context() as patched: + patched.setattr(torch.cuda, "Event", no_event) + with pytest.raises(planted, match="planted"): + lease.mark_arrived(view.row_masks(True)) # the copy is queued, then raises + assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + lease.release() + assert lender._free_slots(group) == 0, ( + "slots came back while a copy out of them was queued" + ) + lender.resume_bounds(target) # any call makes progress + assert lender._free_slots(group) == 0, "slots a copy may still read were reused" + + +class Interrupted(BaseException): + """An interruption that is no ``Exception``, as ``KeyboardInterrupt`` is.""" + + +def test_marks_whose_copy_is_interrupted_once_queued_lose_the_slots_it_reads( + kit, real_manager, monkeypatch, attach=attach +): + """``test_marks_whose_copy_raises_once_queued_lose_the_slots_it_reads`` with an interruption + that is no ``Exception`` in place of its ``MemoryError``: the slots a queued copy reads are lost + whatever the mark raises.""" + test_marks_whose_copy_raises_once_queued_lose_the_slots_it_reads( + kit, real_manager, monkeypatch, attach, planted=Interrupted + ) + + +SLOTS_BACK = "slots came back while a copy on them was queued" + + +def interrupted_event(*args, **kwargs): + raise Interrupted("planted") + + +def test_an_interrupt_in_a_read_s_grant_keeps_its_slots_out_of_use( + kit, real_manager, monkeypatch, attach=attach +): + """A read waits for slots another read holds, whose release grants it: the copy into them is + queued, then creating its event is interrupted. The release raises; after the read's release the + slots stay out of use, no grant hands them out, and the shutdown keeps the staging memory.""" + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + regions = lender.regions + (group,) = kit.pool_group_ids(mgr) + head = lender.lend_read(source, 0, END) + ready_view(head, mgr) + waiting = lender.lend_read(source, 0, END) + assert waiting.poll() is None, "the read was granted at once: the check proves nothing" + with kit.held_stream(mgr._stream): + with monkeypatch.context() as patched: + patched.setattr(torch.cuda, "Event", interrupted_event) + with pytest.raises(Interrupted, match="planted"): + head.release() # grants the waiting read: its copy is queued, then interrupted + runs = waiting._slot_runs + waiting.release() + assert lender._free_slots(group) == 0, SLOTS_BACK + lender.resume_bounds(source) # any call makes progress + assert any(r is runs for r in lender._quarantined), SLOTS_BACK + later = lender.lend_read(source, 0, END) + assert later.poll() is None and later.failure is None, "a grant handed the slots out" + later.release() + warnings = kit.lender_warnings(monkeypatch) + mgr.shutdown() + assert kit.staging_kept(regions), "staging freed at shutdown with slots out of use" + assert warnings, "keeping the memory until exit is logged" + + +def interrupting_as_it_takes_a_copy(lease): + """Turn ``lease`` into one whose first setting of ``_copy`` to a copy raises ``Interrupted``.""" + + class Interrupting(type(lease)): + def __setattr__(self, name, value): + if name == "_copy" and value is not None and not self.__dict__.get("_interrupted"): + self.__dict__["_interrupted"] = True + raise Interrupted("planted") + super().__setattr__(name, value) + + lease.__class__ = Interrupting + + +def test_an_interrupt_before_a_write_takes_its_copy_keeps_the_slots_out_of_use( + kit, real_manager, monkeypatch, attach=attach +): + """A write's marked rows are copied into its pages, and the lease taking that copy is + interrupted: the mark raises, the slots the copy reads stay out of use after the release, no + grant hands them out, and the shutdown keeps the staging memory.""" + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + regions = lender.regions + (group,) = kit.pool_group_ids(mgr) + lease = lender.lend_write(target, 0, END) + view = lease.poll() + interrupting_as_it_takes_a_copy(lease) + with kit.held_stream(mgr._stream): + with pytest.raises(Interrupted, match="planted"): + lease.mark_arrived(view.row_masks(True)) # the copy is queued, then interrupted + runs = lease._slot_runs + lease.release() + assert lender._free_slots(group) == 0, SLOTS_BACK + lender.resume_bounds(target) # any call makes progress + assert any(r is runs for r in lender._quarantined), SLOTS_BACK + later = lender.lend_write(target, 0, END) + assert later.poll() is None and later.failure is None, "a grant handed the slots out" + later.release() + warnings = kit.lender_warnings(monkeypatch) + mgr.shutdown() + assert kit.staging_kept(regions), "staging freed at shutdown with slots out of use" + assert warnings, "keeping the memory until exit is logged" + + +def test_a_grant_that_raises_fails_only_its_own_lease( + kit, real_manager, monkeypatch, attach=attach +): + """Two reads wait in line for the slots a third holds, and the first one's grant raises when + the third's release frees them: that read fails, the healthy one behind it is granted in the + same release and stages its source's bytes, resume_bounds answers and a later read goes through.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + + with real_manager() as mgr: + holder = kit.published(mgr, SOURCE, PROMPT) + broken = kit.published(mgr, OTHER, OTHER_PROMPT) + fetched = kit.admitted(mgr, TARGET, list(range(8000, 8097))) + lender = attach(mgr, tokens_per_fetch=END) + write = lender.lend_write(fetched, 0, END) + write.mark_arrived(write.poll().row_masks()) + write.release() + head = lender.lend_read(holder, 0, END) + ready_view(head, mgr) + waiting = lender.lend_read(broken, 0, 2 * TPB) + assert waiting.poll() is None + behind = lender.lend_read(holder, 0, TPB) # healthy, in line behind the broken one + assert behind.poll() is None + broken_kv = kit.kv(mgr, broken) + + real_kv_of = _manager.kv_of + + def kv_of(manager, request_id): + if request_id == OTHER: + raise RuntimeError("planted") + return real_kv_of(manager, request_id) + + def raising_for_broken(read): + def call(kv_cache, *rest): + if kv_cache is broken_kv: + raise RuntimeError("planted") + return read(kv_cache, *rest) + + return call + + with monkeypatch.context() as patched: + patched.setattr(_manager, "kv_of", kv_of) + for name in ("cache_state", "pages", "locked_pages"): + patched.setattr(_manager, name, raising_for_broken(getattr(_manager, name))) + head.release() # grants both waiting reads, and the first one's grant raises + # The lender's own record, read before any other call: a poll would grant it itself. + granted = behind._view is not None + assert waiting.poll() is None and waiting.failure is not None + assert granted and behind.failure is None, ( + f"the read behind the broken one was not granted in that release: {behind.failure}" + ) + assert isinstance(lender.resume_bounds(fetched), ResumeBounds) + view = ready_view(behind, mgr) + assert 0 < view.num_block_entries == check_staged(kit, mgr, lender, holder, view) + behind.release() + waiting.release() + again = lender.lend_read(holder, 0, END) + ready_view(again, mgr) + again.release() + + +@pytest.mark.parametrize("fault", ["none", "take"], ids=["control", "take_raises"]) +def test_an_error_in_the_free_hook_stops_the_lender(kit, real_manager, monkeypatch, fault): + """Two pool groups. A released publish's copy into its slots is pending while a fetch waits for + them; another request's free through the manager then returns the publish's slots and grants + them to the fetch, and the take of the fetch's second pool group raises a planted MemoryError + once its first pool group changed. The free returns, and the lender has stopped: the waiting + fetch and a lease no poll returned the view of fail, a view a poll returned stays, lends fail at + the call after their range checks, a mark copies nothing, resume_bounds is an empty interval, no + slot returns, and the shutdown keeps the staging memory with nothing open. Without the fault + the free grants the waiting fetch and a mark copies its rows.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: + source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_WINDOWED_PROMPT) + seen_target = kit.admitted(mgr, 4, PROMPT) + granted_target = kit.admitted(mgr, 5, list(range(8000, 8097))) + other = kit.admitted(mgr, OTHER, list(range(9000, 9097))) + lender = attach(mgr, tokens_per_fetch=WINDOWED_END, planned_fetches=3) + regions = lender.regions + groups = kit.pool_group_ids(mgr) + assert len(regions) == len(groups) == 2, "a grant changes one pool group, then the other" + seen = lender.lend_write(seen_target, 0, 2 * TPB) + seen_view = seen.poll() + granted = lender.lend_write(granted_target, 0, TPB) # granted, its view never polled + assert seen_view is not None and granted.failure is None + kit.fill_sentinel(mgr, seen_target) + kit.stage(lender, seen_view, 0x5A) + # Grown ahead, so the held section only lends. + kv = kit.kv(mgr, target) + assert mgr._resize_for_served_prefix(target, kv, WINDOWED_END, WINDOWED_END) + with kit.held_stream(mgr._stream) as gate: + publish = lender.lend_read(source, 0, WINDOWED_END) # its copy held + fetch = lender.lend_write(target, 0, WINDOWED_END) + assert fetch.poll() is None and fetch.failure is None, "the fetch did not wait" + publish.release() # its copy pends, so its slots stay out of use + gate.open() + taken = [] + + def planted(take): + def call(start, count): + taken.append(count) + if len(taken) == 2: # the fetch's second pool group, its first one changed + raise MemoryError("planted") + return take(start, count) + + return call + + if fault == "take": + for alloc in lender._slot_queue._allocators.values(): + monkeypatch.setattr(alloc, "take", planted(alloc.take)) + mgr.free_resources(other) # the manager's free runs the lender's hook + assert kit.kv(mgr, other) is None + if fault == "none": + assert fetch.poll() is not None and fetch.failure is None, "the free granted nothing" + assert granted.poll() is not None + seen.mark_arrived(seen_view.row_masks(True)) + mgr._stream.synchronize() + for run in seen_view.groups: + nbytes = lender.regions[run.region_index].slot_bytes + for ordinal in run.block_ordinals.tolist(): + got = kit.digest(kit.page(mgr, seen_target, run.layer_group_index, ordinal)) + want = kit.digest(kit.staged_row(0x5A, run.layer_group_index, ordinal, nbytes)) + assert got == want, "a mark copied nothing" + for lease in (fetch, granted): + lease.mark_arrived(lease.poll().row_masks()) + for lease in (seen, granted, fetch): + lease.release() + mgr.shutdown() + assert not kit.staging_kept(regions), "nothing open, yet the staging memory stays" + return + assert len(taken) == 2, "the planted fault did not strike: proves nothing" + free = {g: lender._free_slots(g) for g in groups} + intervals = [lender.resume_bounds(r) for r in (target, seen_target)] # progress comes first + assert fetch.failure is not None, "the waiting fetch neither became ready nor failed" + stop = fetch.failure + assert stop == _lender._REASON_STOPPED, stop + assert fetch.poll() is None + for interval in intervals: + assert interval is not None and interval.max_position < interval.min_position, ( + f"resume_bounds {interval} after the stop, not an empty interval" + ) + assert granted.poll() is None and granted.failure == stop, "a view went out after the stop" + assert seen.poll() is seen_view, "a view a poll returned before the stop was taken back" + for lend, request in ((lender.lend_read, source), (lender.lend_write, target)): + lease = lend(request, 0, WINDOWED_END) + assert lease.failure == stop, ( + f"{lend.__name__} did not fail at the call: {lease.failure}" + ) + lease.release() + with pytest.raises(ValueError): + lender.lend_read(source, 0, TPB + 1) # not whole blocks + with pytest.raises(ValueError): + lender.lend_write(target, 0, WINDOWED_END + TPB) # past the last prompt token's block + seen.mark_arrived(seen_view.row_masks(True)) + with pytest.raises(RuntimeError): + seen.mark_arrived(seen_view.row_masks(True)) + mgr._stream.synchronize() + for run in seen_view.groups: + sentinel = bytes([kit.SENTINEL]) * lender.regions[run.region_index].slot_bytes + for ordinal in run.block_ordinals.tolist(): + got = kit.digest(kit.page(mgr, seen_target, run.layer_group_index, ordinal)) + assert got == kit.digest(sentinel), "a mark copied rows after the stop" + for lease in (seen, granted, fetch): + lease.release() + assert lender._open_lease_count() == 0 + assert {g: lender._free_slots(g) for g in groups} == free, "slots returned or granted" + warnings = kit.lender_warnings(monkeypatch) + mgr.shutdown() + assert kit.staging_kept(regions), ( + "the shutdown freed staging whose slot records may be wrong" + ) + assert warnings, "keeping the staging memory until exit is logged" + + +# -- names ------------------------------------------------------------------------------------ + + +def _tp2(rank): + from tensorrt_llm.mapping import Mapping + + return Mapping(world_size=2, tp_size=2, rank=rank) + + +def names_by_rank(kit, real_manager, attach, num_kv_heads, scopes=(SCOPE, SCOPE)): + """Each rank of a TP2 pair publishes ``PROMPT``: its names, and its region names.""" + names, region_labels = [], [] + for rank, scope in enumerate(scopes): + with real_manager(num_kv_heads=num_kv_heads, mapping=_tp2(rank)) as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END, scope=scope) + lease = lender.lend_read(source, 0, END) + (run,) = ready_view(lease, mgr).groups + names.append(np.array(run.object_keys)) + region_labels.append([p.layout_label for p in lender.regions]) + lease.release() + return names, region_labels + + +def test_a_replicated_group_is_named_alike_on_every_rank(kit, real_manager, attach=attach): + (rank0, rank1), (parts0, parts1) = names_by_rank(kit, real_manager, attach, num_kv_heads=1) + assert rank0.shape == (BLOCKS, 54) and rank0.dtype == np.uint8 + assert np.array_equal(rank0, rank1), "ranks holding the same bytes must share their names" + assert parts0 == parts1 + + +def test_a_head_sharded_group_names_each_rank_s_share(kit, real_manager): + (rank0, rank1), (parts0, parts1) = names_by_rank(kit, real_manager, attach, num_kv_heads=4) + assert np.array_equal(rank0[:, :50], rank1[:, :50]), "same scope, layout and blocks" + for rank, names in enumerate((rank0, rank1)): + shard = (2).to_bytes(2, "big") + rank.to_bytes(2, "big") + assert {bytes(n[50:54]) for n in names} == {shard} + assert parts0 == parts1 + + +def test_block_names_follow_the_scope_and_region_labels_only_the_layout(kit, real_manager): + names = [] + for scope in (b"model-a", b"model-b"): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END, scope=scope) + lease = lender.lend_read(source, 0, END) + names.append( + ([bytes(n) for n in ready_view(lease, mgr).groups[0].object_keys], lender.regions) + ) + lease.release() + (first, first_regions), (second, second_regions) = names + assert [n[16:] for n in first] == [n[16:] for n in second], "same blocks, same keys" + assert not set(first) & set(second) + assert [p.layout_label for p in first_regions] == [p.layout_label for p in second_regions], ( + "laid out alike" + ) + with real_manager(tokens_per_block=2 * TPB) as mgr: + other_layout = attach(mgr, tokens_per_fetch=END, scope=b"model-a").regions + assert [p.layout_label for p in other_layout] != [p.layout_label for p in first_regions], ( + "another layout" + ) + + +CONTROL_PROMPT = list(range(9100, 9197)) +IMAGE = {"image": {"pixel_values": torch.zeros(4)}} +DIGESTS = dict( + multimodal_hashes=[list(range(1, 9))], multimodal_positions=[40], multimodal_lengths=[16] +) +# What a request carries beyond the tokens its names cover, how its lease fails, and what a control +# request whose names cover what it carries holds instead. +UNNAMED = { + "multimodal_without_digests": ( + dict(py_multimodal_data=IMAGE), + "without digests", + dict(py_multimodal_data=IMAGE, **DIGESTS), + ), + "encoder_input": (dict(encoder_input_tokens=list(range(300, 316))), "encoder input", {}), +} + + +def carrying(kit, request_id, tokens, inputs): + """A request carrying ``inputs``, at its context as the executor leaves it after any encoder.""" + from tensorrt_llm.bindings import LlmRequestState + + request = kit.make_request(request_id, tokens, **inputs) + request.state = LlmRequestState.CONTEXT_INIT + return request + + +@pytest.mark.parametrize("case", list(UNNAMED)) +def test_a_request_whose_kv_its_names_miss_is_not_lent(kit, real_manager, case, attach=attach): + inputs, why, covered = UNNAMED[case] + with real_manager() as mgr: + source = carrying(kit, SOURCE, PROMPT, inputs) + kit.prefill(mgr, source) + target = carrying(kit, TARGET, OTHER_PROMPT, inputs) + assert mgr.prepare_context(target) + control = kit.published(mgr, OTHER, CONTROL_PROMPT, **covered) + lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + before = [state_of(kit, mgr, lender, r) for r in (source, target)] + refused = [lender.lend_read(source, 0, END), lender.lend_write(target, 0, END)] + for lease in refused: + assert lease.poll() is None and why in str(lease.failure), ( + f"lent a request carrying {case.replace('_', ' ')}" + ) + assert [state_of(kit, mgr, lender, r) for r in (source, target)] == before + lent = lender.lend_read(control, 0, END) + assert check_staged(kit, mgr, lender, control, ready_view(lent, mgr)) == BLOCKS + for lease in (*refused, lent): + lease.release() + + +# name: (windows, multimodal run [b, e), the fetch's end, mm_bidirectional_blocks set, refused) +SPLIT_RUNS = { + "end_inside_a_run_longer_than_the_window": ([WINDOW, 256], (1, 130), 128, True, True), + "end_past_the_run": ([WINDOW, 256], (1, 130), 160, True, False), + "end_inside_a_run_within_the_window": ([WINDOW, 256], (1, 65), 32, True, True), + # The manager's window holds the run, or there is none: the model's may still be shorter. + "end_inside_a_run_within_a_wider_window": ([2 * WINDOW, 256], (1, 100), 96, True, True), + "end_inside_a_run_without_a_window": (None, (1, 130), 128, True, True), + "run_not_bidirectional": ([WINDOW, 256], (1, 130), 128, False, False), + "end_past_a_run_the_floor_may_fall_below": ([80, 256], (1, 87), 96, True, False), +} + + +def bidirectional(runs, flagged, length): + """Multimodal data whose tokens ``runs`` attend both ways where ``flagged``, with digests.""" + mask = np.zeros(length, dtype=np.int64) + for beg, end in runs: + mask[beg:end] = 1 + data = dict(IMAGE, mm_bidirectional_blocks=flagged) + data["multimodal_embed_mask_cumsum"] = torch.from_numpy(np.cumsum(mask)) + return dict(py_multimodal_data=data, **DIGESTS) + + +@pytest.mark.parametrize("case", list(SPLIT_RUNS)) +def test_a_fetch_ending_inside_a_bidirectional_run_fails_at_the_call( + kit, real_manager, case, attach=attach +): + """A request whose multimodal tokens attend both ways in runs the scheduler keeps within one + chunk: a fetch whose end falls strictly inside a run fails at the call, whatever the run's + length and the manager's windows, which need not be the model's; every other fetch is lent.""" + windows, run, end, flagged, refused = SPLIT_RUNS[case] + with real_manager(windows=windows) as mgr: + inputs = bidirectional([run], flagged, len(WINDOWED_PROMPT)) + target = carrying(kit, TARGET, WINDOWED_PROMPT, inputs) + assert mgr.prepare_context(target) + kit.kv(mgr, target).enable_swa_scratch_reuse = False # a fetch target needs it off + lender = attach(mgr, tokens_per_fetch=end) + before = state_of(kit, mgr, lender, target) + lease = lender.lend_write(target, 0, end) + if refused: + assert lease.poll() is None and "multimodal" in str(lease.failure), ( + f"lent a fetch ending inside the run {run}: {case}" + ) + assert state_of(kit, mgr, lender, target) == before, ( + "the failed lease changed the cache" + ) + else: + ready_view(lease, mgr) + lease.release() + + +# name: (windows, multimodal runs, the fetch's end, block ordinals that arrive (None: all), +# mm_bidirectional_blocks set, resume_bounds) +RUN_READINESS = { + "delivered_rows_ending_inside_a_run": (None, [(40, 100)], 160, [0, 1], True, (40, 0)), + "a_floor_below_a_run": (None, [(5, 30), (40, 100)], 160, None, True, (160, 100)), + "a_floor_below_a_run_no_window_releases": ([80, 256], [(1, 87)], 96, None, True, (96, 87)), + "without_the_flag": (None, [(40, 100)], 160, None, False, (160, 0)), +} + + +@pytest.mark.parametrize("case", list(RUN_READINESS)) +def test_resume_bounds_keeps_resumes_out_of_bidirectional_runs( + kit, real_manager, case, attach=attach +): + """All-reusable, nothing computed, a fetch past every run into a request whose multimodal tokens + attend both ways: resume_bounds ends at the start of a run the delivered rows end inside and starts + at the end of the last run below its end; a request without the flag keeps its interval.""" + windows, runs, end, arrived, flagged, expected = RUN_READINESS[case] + with real_manager(windows=windows) as mgr: + inputs = bidirectional(runs, flagged, len(WINDOWED_PROMPT)) + target = carrying(kit, TARGET, WINDOWED_PROMPT, inputs) + assert mgr.prepare_context(target) + kit.kv(mgr, target).enable_swa_scratch_reuse = False # a fetch target needs it off + lender = attach(mgr, tokens_per_fetch=end) + fetch_segment(kit, mgr, lender, target, 0, end, 0x5A, arrived) + for _ in range(2): # the second answer comes from the record resume_bounds keeps + resume_bounds = tuple(lender.resume_bounds(target)) + assert resume_bounds == expected, ( + f"resume_bounds {resume_bounds} lets the request resume inside a run of {runs}: {case}" + ) + + +def test_resume_bounds_keeps_a_resume_out_of_the_run_the_context_position_is_in( + kit, real_manager, attach=attach +): + """Per-request, a first chunk that ended inside a run, as a whole-block local match can leave + it, and a fetch past the run abandoned: with nothing delivered, resume_bounds ends at the run's + start, below the context position, so the request does not resume inside the run.""" + with real_manager(**PER_REQUEST) as mgr: + inputs = bidirectional([(40, 100)], True, len(SPLIT_PROMPT)) + target = computed(kit, mgr, TARGET, SPLIT_PROMPT, **inputs) + lender = attach(mgr, tokens_per_fetch=2 * TPB) + lender.lend_write(target, LOCAL, 4 * TPB).release() # before its first poll: abandoned + resume_bounds = lender.resume_bounds(target) + assert tuple(resume_bounds) == (40, LOCAL), ( + f"resume_bounds {tuple(resume_bounds)} lets the request resume inside the run [40, 100)" + ) + + +# -- threads ---------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("pinned", [False, True], ids=["pageable", "pinned"]) +def test_threads_take_turns_and_the_lender_starts_none(kit, real_manager, monkeypatch, pinned): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + monkeypatch.setattr(_lender, "prefer_pinned", lambda: pinned) + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + before = set(threading.enumerate()) + built = kit.on_thread( + lambda: attach(mgr, tokens_per_fetch=END, planned_fetches=2), "builder", cuda=True + ) + lender = built["value"] + regions = lender.regions + (memory,) = kit.staging_memory(regions) + base = regions[0].address + if pinned: + assert kit.pinned_range(base) is not None + + def executor_loop(): + read = lender.lend_read(source, 0, END) + view = ready_view(read, mgr) + # A backend reads the arrays on its own thread until the release. + seen = kit.on_thread(lambda: [bytes(n) for n in view.groups[0].object_keys], "backend") + write = lender.lend_write(target, 0, END) + write.mark_arrived(write.poll().row_masks(True)) + read.release() + write.release() + return seen["value"] + + looped = kit.on_thread(executor_loop, "executor-loop", cuda=True) + assert "error" not in looped and len(looped["value"]) == BLOCKS + assert set(threading.enumerate()) <= before, "the lender started a thread" + assert "error" not in kit.on_thread(mgr.shutdown, "shutdown", cuda=True) + assert memory.address == 0 and memory._pageable is None + assert not kit.staging_kept(regions), "the shutdown thread kept the staging memory" + if pinned: + assert kit.pinned_range(base) is None, "the shutdown thread did not free pinned memory" + with pytest.raises(ValueError): + lender.resume_bounds(source) + + +# -- the staging lease's states -------------------------------------------------------------- + +UNMARKED_SLOTS_BACK = "a write released unmarked gave its slots back while a backend may write" + + +def test_a_staging_lease_walks_through_the_states_its_enum_names( + kit, real_manager, monkeypatch, attach=attach +): + """Reads and writes through every ``_LeaseState``, ``_state()`` read after each call: a read + granted, ready and released, slots back; a waiting read granted; a write marked and released, + slots back; a doomed write failed at its first poll; a lease failed at the call; a write + released ready and unmarked, slots still held; a waiting read failed at its request's free.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + states = _lender._LeaseState + seen = set() + + def state(lease, expected): + got = lease._state() + seen.add(got) + assert got is expected, f"{lease!r}: {got.name}, not {expected.name}" + + def holding(lease): + return any(held is lease for held in lender._leases_holding_slots) + + view = _lender.Staging._build_view + planted = [] + + def raising_once(self, rows, keys): + if not planted: + planted.append(True) + raise MemoryError("planted") + return view(self, rows, keys) + + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + read = lender.lend_read(source, 0, END) + state(read, states.GRANTED) + ready_view(read, mgr) + state(read, states.READY) + waiting = lender.lend_read(source, 0, END) + state(waiting, states.WAITING) + read.release() + state(read, states.RELEASED) + assert not holding(read), "a released read kept its slots" + state(waiting, states.GRANTED) + ready_view(waiting, mgr) + waiting.release() + write = lender.lend_write(kit.admitted(mgr, TARGET, OTHER_PROMPT), 0, END) + state(write, states.GRANTED) + write_view = write.poll() + state(write, states.READY) + write.mark_arrived(write_view.row_masks()) + state(write, states.MARKED) + write.release() + state(write, states.RELEASED) + mgr._stream.synchronize() + lender.resume_bounds(source) # any call makes progress + assert not holding(write), "a marked write kept its slots once its copy had landed" + with monkeypatch.context() as patched: + patched.setattr(_lender.Staging, "_build_view", raising_once) + doomed = lender.lend_write(kit.admitted(mgr, 4, list(range(8000, 8097))), 0, END) + state(doomed, states.FAILS_AT_FIRST_POLL) + assert doomed.poll() is None + state(doomed, states.FAILED) + doomed.release() + failed = lender.lend_read(kit.make_request(9, PROMPT), 0, END) + state(failed, states.FAILED) + failed.release() + unmarked = lender.lend_write(kit.admitted(mgr, 5, list(range(9000, 9097))), 0, END) + assert unmarked.poll() is not None + state(unmarked, states.READY) + unmarked.release() + state(unmarked, states.RELEASED) + lender.resume_bounds(source) + assert holding(unmarked), UNMARKED_SLOTS_BACK + freed = kit.published(mgr, 6, list(range(10000, 10097))) + in_line = lender.lend_read(freed, 0, END) + state(in_line, states.WAITING) + mgr.free_resources(freed) + state(in_line, states.FAILED) + in_line.release() + assert seen == set(states), f"states never seen: {set(states) - seen}" + + +# -- the manager's page-index buffer, kept from the attach until the runtime's shutdown returns + + +def test_the_index_buffer_is_kept_from_the_attach_until_the_shutdown( + kit, real_manager, attach=attach +): + with real_manager() as mgr: + before = len(kit.retained()) + buffer = mgr.host_kv_cache_block_offsets + attach(mgr, tokens_per_fetch=END) + assert any(o is buffer for o in kit.retained()), "not kept from the attach" + mgr.shutdown() + assert not any(o is buffer for o in kit.retained()), "kept past the shutdown" + assert len(kit.retained()) == before + + +class ClosingSecondTime: + """A ``kv_cache_map`` entry whose first ``close()`` raises and whose later ones close the cache; + the manager's other calls on it go to the cache.""" + + def __init__(self, kv_cache): + self.kv_cache = kv_cache + self.calls = 0 + + def __getattr__(self, name): + return getattr(self.kv_cache, name) + + def close(self): + self.calls += 1 + if self.calls == 1: + raise RuntimeError("planted") + self.kv_cache.close() + + +BUFFER_GONE_EARLY = "the buffer went before every cache closed" + + +def test_the_index_buffer_stays_kept_through_a_close_that_raises(kit, real_manager, attach=attach): + """A close that raises stops the manager's shutdown with caches still open, which a released + lease or a record of the lender may hold and which write their page indices into the buffer as + they close: staging keeps the buffer until the runtime's shutdown returns, here a retry's.""" + with real_manager() as mgr: + buffer = mgr.host_kv_cache_block_offsets + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + lease = lender.lend_write(target, 0, END) + lease.mark_arrived(ready_view(lease, mgr).row_masks()) + lease.release() + kv = kit.kv(mgr, target) + request_id = next(rid for rid, cache in mgr.kv_cache_map.items() if cache is kv) + mgr.kv_cache_map[request_id] = ClosingSecondTime(kv) # the released lease's own cache + with pytest.raises(RuntimeError, match="planted"): + mgr.shutdown() + assert not kit.closed(kv), "the planted close closed the cache: the check proves nothing" + assert any(o is buffer for o in kit.retained()), BUFFER_GONE_EARLY + mgr.shutdown() + assert kit.closed(kv), "the retried shutdown left the cache open" + assert not any(o is buffer for o in kit.retained()), "kept past a shutdown that closed all" + + +def test_the_index_buffer_stays_kept_past_a_free_whose_close_raised( + kit, real_manager, wrap=None, attach=attach +): + """A close that raises in a free leaves the cache open outside the manager's map, held by a + released read, so the runtime's shutdown raises: staging keeps the buffer until a shutdown + returns, here a retry's once the cache has closed.""" + with real_manager() as mgr: + buffer = mgr.host_kv_cache_block_offsets + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, tokens_per_fetch=END) + lease = lender.lend_read(source, 0, END) + ready_view(lease, mgr) + lease.release() # a released read holds its cache until it is dropped + kv = kit.kv(mgr, source) + mgr.kv_cache_map[source.py_request_id] = ClosingSecondTime(kv) + with pytest.raises(RuntimeError, match="planted"): + mgr.free_resources(source) + assert not kit.closed(kv) and source.py_request_id not in mgr.kv_cache_map, ( + "the free closed the cache or left it in the map: the check proves nothing" + ) + try: + if wrap is not None: + wrap(mgr) + with pytest.raises(RuntimeError, match="still open"): + mgr.shutdown() + assert any(o is buffer for o in kit.retained()), BUFFER_GONE_EARLY + finally: + kv.close() + mgr.shutdown() + assert not any(o is buffer for o in kit.retained()), "kept past a shutdown that returned" + + +# A cache writes -1 for every block into its manager's host page-index buffer as it closes. A new +# tensor reclaims a freed buffer's address with a canary the cache must neither read nor write. +FREED = "touched the page-index buffer freed with its manager" + + +def fresh_device(): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + + +def check_a_lease_outliving_its_manager(kit): + """A staging read outlives its manager and lender, collected without a shutdown, holding the + request's cache; it also keeps its staging memory as it was.""" + fresh_device() + mgr = kit.make_manager() + request = kit.published(mgr, SOURCE, PROMPT) + kv = kit.kv(mgr, request) + blocks = int(kv.num_blocks) + lender = attach_lender(mgr, scope=b"index-buffer", staging=StagingOptions(END)) + lease = lender.lend_read(request, 0, END) + mgr._stream.synchronize() # a staging read is ready once its copy into staging ran + view = lease.poll() + assert view is not None + (run,) = view.groups + address, slot_bytes = int(run.addresses[0]), lender.regions[run.region_index].slot_bytes + regions, content = lender.regions, kit.host_bytes(address, slot_bytes) + row = kit.index_row(mgr, request) + own = list(kv.get_base_page_indices(0)[:blocks]) + assert row.values[:blocks] == own and -1 not in own, "the cache writes elsewhere" + watched = [weakref.ref(o) for o in (mgr.host_kv_cache_block_offsets, mgr, lender)] + mgr._stream.synchronize() + del mgr, lender + gc.collect() + assert watched[1]() is None and watched[2]() is None, "the manager or lender was not collected" + canary = None if watched[0]() is not None else kit.reclaim(row) + if watched[0]() is None: + assert canary is not None, "inconclusive: no allocation reused the freed buffer's address" + read = list(kv.get_base_page_indices(0)[:blocks]) + assert kit.staging_kept(regions), "the staging memory went with its manager" + now = kit.digest(kit.host_bytes(address, slot_bytes)) + assert now == kit.digest(content), "the kept staging memory changed" + del kv # the lease now holds the cache's last reference + lease.release() + lease.release() + del lease # a staging lease keeps its cache past its release + gc.collect() + written = kit.canary_written(canary, row) + assert not (canary is not None and read == [kit.CANARY] * blocks) and not written, ( + f"the lent cache {FREED}: it read {read} and its close wrote {written} (cell, value)" + ) + + +def test_a_cache_outliving_its_manager_touches_no_freed_index_buffer(kit): + check_a_lease_outliving_its_manager(kit) + + +# -- DeepSeek-V4 ------------------------------------------------------------------------------ + + +@skip_pre_blackwell +@pytest.mark.parametrize("fp8_ds_mla", [False, True], ids=["fp8", "fp8_ds_mla"]) +def test_a_deepseek_v4_prompt_is_fetched_through_staging_byte_for_byte( + kit, deepseek_v4_manager, fp8_ds_mla +): + """Three whole blocks published by one DeepSeek-V4 manager and fetched by another, byte for + byte, in FP8 at 128 tokens per block and in the footer-scale cache at 256.""" + with ( + deepseek_v4_manager(fp8_ds_mla=fp8_ds_mla) as mgr_a, + deepseek_v4_manager(fp8_ds_mla=fp8_ds_mla) as mgr_b, + ): + tpb = mgr_a.tokens_per_block + assert tpb == (256 if fp8_ds_mla else 128) + end = 3 * tpb + prompt = list(range(3000, 3000 + end + 1)) + source = kit.published(mgr_a, SOURCE, prompt) + lender_a = attach(mgr_a, tokens_per_fetch=end) + publish = lender_a.lend_read(source, 0, end) + publish_view = ready_view(publish, mgr_a) + assert len(publish_view.groups) == kit.num_layer_groups(mgr_a) + assert ( + check_staged(kit, mgr_a, lender_a, source, publish_view) + == publish_view.num_block_entries + ) + target = kit.admitted(mgr_b, TARGET, prompt) + lender_b = attach(mgr_b, tokens_per_fetch=end) + assert [p.layout_label for p in lender_a.regions] == [ + p.layout_label for p in lender_b.regions + ] + lease = lender_b.lend_write(target, 0, end) + view = lease.poll() + assert by_group(view) == by_group(publish_view) + kit.fill_sentinel(mgr_b, target) + lease.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) + mgr_b._stream.synchronize() + assert lender_b.resume_bounds(target) == (end, end) + for run in view.groups: + for ordinal in run.block_ordinals.tolist(): + lg = run.layer_group_index + got = kit.digest(kit.page(mgr_b, target, lg, ordinal)) + assert got == kit.digest(kit.page(mgr_a, source, lg, ordinal)) + lease.release() + publish.release() + + +def ngram_drafting(draft_len): + """NGram drafting ``draft_len`` tokens; ``ngram`` drafts 2.""" + from tensorrt_llm.llmapi.llm_args import NGramDecodingConfig + + return NGramDecodingConfig(max_draft_len=draft_len, max_matching_ngram_size=2) + + +# Drafting without read-ahead widens each DeepSeek-V4 window by the draft length, so a fetch of +# [0, 512) also asks for that margin; layer group 0 has windows of 128. case -> (spec config, +# tokens the publisher computes past 512, rows the publish lacks per layer group, resume_bounds) +DSV4_TPB = 128 +DSV4_END = 4 * DSV4_TPB +DSV4_PROMPT = list(range(60_000, 60_000 + DSV4_END + 2)) # four whole blocks and two tokens +DSV4_TARGET = DSV4_PROMPT[:DSV4_END] + [69_999, 69_998] # the same four blocks, then others +DSV4_DRAFT_MARGIN = { + # A window of 128 at 512 keeps block 3 alone, which a publisher two tokens past still keeps. + "no_speculation": (lambda: None, 2, {}, (DSV4_END, DSV4_END)), + # A draft length of 2 keeps block 2 too, which only a publisher whose history stands at the + # fetch's end still keeps; past it, the interval is empty. + "draft2_publisher_at_the_end": (ngram, 0, {}, (DSV4_END, DSV4_END)), + "draft2_publisher_one_past": (ngram, 1, {0: [2]}, (0, DSV4_END)), + # A draft length of 3 keeps block 2 for a publisher one token past, not two. + "draft3_publisher_one_past": (lambda: ngram_drafting(3), 1, {}, (DSV4_END, DSV4_END)), + "draft3_publisher_two_past": (lambda: ngram_drafting(3), 2, {0: [2]}, (0, DSV4_END)), + # PARD also grows the cache by the extra tokens embedded DSpark does. + "pard2_publisher_two_past": (pard, 2, {0: [2]}, (0, DSV4_END)), +} +NOT_COMPUTED = "holds bytes other than the ones the publisher computed" + + +@skip_pre_blackwell +@pytest.mark.parametrize("case", list(DSV4_DRAFT_MARGIN)) +def test_a_deepseek_v4_fetch_under_speculative_decoding_asks_for_the_window_margin( + kit, deepseek_v4_manager, case +): + """DeepSeek-V4 keeps every window the draft length wider under any speculative decoding, also + one reading no token ahead, which staging accepts. A fetch asks for every row the cache's + windows keep at its end, the margin's included; the rows the publish holds arrive with the bytes + the publisher computed, and resume_bounds counts a position only where a resume there finds those + bytes in every page the cache keeps below it. At 128 tokens per block and a draft length of 2 or + more, a publisher whose history stands at least the draft length less one token past the fetch's + end has passed a margin row, so the interval is empty and the request computes from 0.""" + make, tail, lacking, expected = DSV4_DRAFT_MARGIN[case] + with ( + deepseek_v4_manager(spec_config=make()) as pub, + deepseek_v4_manager(spec_config=make()) as mgr, + ): + assert kit.windows(mgr)[0] == DSV4_TPB + mgr.max_draft_len + assert mgr.reuse_match_backoff == 0 + source = kit.published(pub, SOURCE, DSV4_PROMPT[: DSV4_END + tail]) + publisher = attach(pub, tokens_per_fetch=DSV4_END) + publish = publisher.lend_read(source, 0, DSV4_END) + publish_view = ready_view(publish, pub) + assert ( + check_staged(kit, pub, publisher, source, publish_view) + == publish_view.num_block_entries + ) + held = {name.tobytes() for run in publish_view.groups for name in run.object_keys} + target = kit.admitted(mgr, TARGET, DSV4_TARGET) + lender = attach(mgr, tokens_per_fetch=DSV4_END) + lease = lender.lend_write(target, 0, DSV4_END) + view = ready_view(lease, mgr) + kit.fill_sentinel(mgr, target) + + def published_row(i, row): + return view.groups[i].object_keys[row].tobytes() in held + + lease.mark_arrived(kit.relay(publisher, publish_view, lender, view, published_row)) + lease.release() + publish.release() + mgr._stream.synchronize() + usable, floor = lender.resume_bounds(target) + kv = kit.kv(mgr, target) + keys = kit.chain_keys(kv, DSV4_TARGET) + dev = kit.DevicePages(mgr) + + def wrong(lg, ordinal): + """The block's page, digested, unless it holds what the publisher computed.""" + got = kit.digest(dev.read(lg, kit.pages(kv, lg)[ordinal])) + computed = kit.digest(kit.fake_page(lg, keys[ordinal], dev.page_bytes(lg))) + return None if got == computed else got + + for i, run in enumerate(view.groups): + lg = run.layer_group_index + for row, ordinal in enumerate(run.block_ordinals.tolist()): + got = wrong(lg, ordinal) if published_row(i, row) else None + assert got is None, f"fetched block {ordinal} of layer group {lg} {NOT_COMPUTED}" + if floor <= usable: + for lg in range(kit.num_layer_groups(mgr)): + beg, end = kit.stale_blocks(mgr, lg, usable) + for ordinal in range(-(-usable // DSV4_TPB)): + got = None if beg <= ordinal < end else wrong(lg, ordinal) + assert got is None, ( + f"block {ordinal} of layer group {lg}, which the cache keeps at the resume " + f"point {usable}, {NOT_COMPUTED}: {got}" + ) + resume_bounds = (usable, floor) + assert resume_bounds == expected, ( + f"resume_bounds {resume_bounds}, the rules give {expected}" + ) + asked = {} + for i, run in enumerate(view.groups): + rows = enumerate(run.block_ordinals.tolist()) + missing = [ordinal for row, ordinal in rows if not published_row(i, row)] + if missing: + asked[run.layer_group_index] = missing + assert asked == lacking, f"the fetch's rows the publish lacks: {asked}, expected {lacking}" diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index e1f13c50e62b..21f146d1cf88 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py @@ -231,6 +231,8 @@ def test_update_resources_leaves_history_untouched_under_helix() -> None: kv_cache_map={7: kv}, kv_compression_manages_history=False, _has_cp_helix=True, + # No lender attached, so a resize has no one to tell. + _sharing=None, ) batch = SimpleNamespace(generation_requests=[req]) KVCacheManagerV2.update_resources(mgr, batch) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index dc154cde2b1d..59b54805eaa3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -373,7 +373,7 @@ def make_ctx_request( req.context_remaining_length = context_remaining_length req.prompt_len = prompt_len or context_remaining_length req.context_current_position = 0 - req.py_connector_served_position = 0 + req.py_served_prefix_end = 0 req.expect_snapshot_points = [] req.num_draft_tokens = num_draft_tokens req.has_draft_tokens = num_draft_tokens > 0 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 123f8100eb6d..75e32a232825 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -1612,7 +1612,7 @@ class _ContextRequest: is_dummy_request: bool = False return_perf_metrics: bool = False context_current_position: int = 0 - py_connector_served_position: int = 0 + py_served_prefix_end: int = 0 py_num_connector_matched_tokens: int = 0 prepopulated_prompt: tuple[int, int] | None = None multimodal_hashes: None = None @@ -2938,10 +2938,13 @@ def _make_fresh_fill_manager( monkeypatch: pytest.MonkeyPatch, fill_value: float | None = _FRESH_FILL, ) -> KVCacheManagerV2: - # The production path synchronises so the fill lands before the forward - # pass; these tests run on CPU tensors. - monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + # These tests use CPU tensors and an explicit stream double; a device-wide wait is forbidden. + monkeypatch.setattr( + torch.cuda, "synchronize", Mock(side_effect=AssertionError("device-wide fill sync")) + ) + monkeypatch.setattr(torch.cuda, "current_stream", Mock(return_value=Mock())) manager = object.__new__(KVCacheManagerV2) + manager._stream = SimpleNamespace(device=torch.device("cuda:0")) manager._fresh_page_fill = fill_value manager._fresh_pages_filled = {} manager._fresh_fill_announced = True @@ -2962,11 +2965,29 @@ def test_fresh_fill_covers_pages_on_allocation(monkeypatch: pytest.MonkeyPatch) buffer = torch.zeros(8, 2) kv_cache = _FreshFillKVCache([2, 5]) manager = _make_fresh_fill_manager(buffer, kv_cache, monkeypatch) + other_buffer = torch.zeros_like(buffer) + manager.pp_layers = [0, 1] + manager.layer_offsets = {0: 0, 1: 1} + manager.layer_to_pool_mapping_dict = {0: 0, 1: 0} + manager.get_buffers = lambda layer_idx: (buffer, other_buffer)[layer_idx] + order = Mock() + order.attach_mock(torch.cuda.current_stream.return_value, "stream") + fill = Mock(wraps=kv_cache_v2_module._fill_kv_pages) + monkeypatch.setattr(kv_cache_v2_module, "_fill_kv_pages", fill) + order.attach_mock(fill, "fill") manager._fill_fresh_kv_pages(1) + torch.cuda.current_stream.assert_called_once_with(device=manager._stream.device) + assert order.mock_calls == [ + call.stream.wait_stream(manager._stream), + call.fill(buffer, [2, 5], _FRESH_FILL), + call.fill(other_buffer, [2, 5], _FRESH_FILL), + call.stream.synchronize(), + ] assert torch.equal(buffer[2], torch.full((2,), _FRESH_FILL)) assert torch.equal(buffer[5], torch.full((2,), _FRESH_FILL)) assert torch.equal(buffer[0], torch.zeros(2)) + assert torch.equal(other_buffer, buffer) # The owner writes its KV into page 2; growing by one page must fill only # the new page, not rewrite what the owner already stored. @@ -2975,6 +2996,12 @@ def test_fresh_fill_covers_pages_on_allocation(monkeypatch: pytest.MonkeyPatch) manager._fill_fresh_kv_pages(1) assert torch.equal(buffer[6], torch.full((2,), _FRESH_FILL)) assert torch.equal(buffer[2], torch.full((2,), 1.5)) + assert order.mock_calls[-4:] == [ + call.stream.wait_stream(manager._stream), + call.fill(buffer, [6], _FRESH_FILL), + call.fill(other_buffer, [6], _FRESH_FILL), + call.stream.synchronize(), + ] def test_fresh_fill_skips_committed_prefix(monkeypatch: pytest.MonkeyPatch) -> None: @@ -3002,6 +3029,7 @@ def test_fresh_fill_disabled_or_unknown_request_is_a_noop( manager._fresh_page_fill = _FRESH_FILL manager._fill_fresh_kv_pages(99) assert torch.equal(buffer, torch.zeros(4, 2)) + torch.cuda.current_stream.assert_not_called() def test_fresh_fill_refills_page_reassigned_after_release( diff --git a/tests/unittest/_torch/executor/kv_cache/test_spec_recompute_tail.py b/tests/unittest/_torch/executor/kv_cache/test_spec_recompute_tail.py index 20466cc9c5b2..e8f2ecbf1f62 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_spec_recompute_tail.py +++ b/tests/unittest/_torch/executor/kv_cache/test_spec_recompute_tail.py @@ -200,6 +200,7 @@ def __init__(self, prompt_len): self.is_first_context_chunk = True self.py_connector_allocation_reported = False self.py_connector_served_position = 0 + self.py_served_prefix_end = 0 self.context_current_position = 0 self.prepopulated_prompt_len = 0 self.context_chunk_size = prompt_len diff --git a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py index 1706c9013df7..e9c83dd2385c 100644 --- a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py +++ b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py @@ -50,6 +50,7 @@ def _coordinator(manager: Mock, receive: Mock) -> DisaggTransferCoordinator: def _manager() -> Mock: """Run the real admission and resource hooks with mocked cache allocation.""" manager = Mock(spec=KVCacheManagerV2) + manager._sharing = None # no lender: free_resources closes and removes the cache itself manager._stream = Mock() manager._disagg_receive_ready = {} manager._pending_preemption = {} diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py index ccae819f0f9f..fb1214565139 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py @@ -96,7 +96,7 @@ def __init__(self, request_id=0, prompt_len=PROMPT_LEN): self.is_disagg_generation_init_state = False self.py_num_connector_matched_tokens = 0 self.py_connector_allocation_reported = False - self.py_connector_served_position = 0 + self.py_served_prefix_end = 0 self.py_ctx_pre_resize_cap = None self.py_draft_tokens = [] self._context_chunk_size = prompt_len @@ -849,10 +849,10 @@ def test_a_dead_allocation_drops_the_floor(self): manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) assert manager.prepare_context(req) manager._run_kv_connector_hooks(scheduled(req)) - assert req.py_connector_served_position == 64 + assert req.py_served_prefix_end == 64 req.py_connector_allocation_reported = False # what free_resources does - req.py_connector_served_position = 0 # ... and this + req.py_served_prefix_end = 0 # ... and this assert manager.prepare_context(req) @@ -907,7 +907,7 @@ def prepare( def test_reserve_before_budget_and_accept_after_final_trim(self, load_async: bool) -> None: manager, connector, req, cache = self.prepare(load_async=load_async) assert req.context_current_position == 64 - assert req.py_connector_served_position == 0 + assert req.py_served_prefix_end == 0 assert cache.history_length == 0 assert connector.accepted == [] assert connector.allocs == [] @@ -921,7 +921,7 @@ def test_reserve_before_budget_and_accept_after_final_trim(self, load_async: boo req.context_chunk_size = 32 manager.report_batch_to_connector(batch) - assert req.py_connector_served_position == 64 + assert req.py_served_prefix_end == 64 assert req.context_chunk_size == 32 assert len(connector.allocs) == 1 reservation, groups = connector.accepted[0] @@ -945,7 +945,7 @@ def test_rejection_rewinds_even_without_capacity_growth(self, capacity: int) -> assert req.context_current_position == 0 assert req.prepopulated_prompt_len == 0 assert req.context_chunk_size == PROMPT_LEN - assert req.py_connector_served_position == 0 + assert req.py_served_prefix_end == 0 manager.release_unused_connector_reservations(set()) assert connector.releases == [(1, 0, 64)] @@ -978,7 +978,7 @@ def test_partial_local_prefix_and_last_prompt_token( ) assert req.context_current_position == expected_end assert req.context_remaining_length == PROMPT_LEN - expected_end - assert req.py_connector_served_position == 0 + assert req.py_served_prefix_end == 0 reservation = connector.get_prefix_reservation(req) if expected_end == committed: assert reservation is None @@ -1002,7 +1002,7 @@ def test_dispatched_load_cannot_lose_its_allocation(self) -> None: with pytest.raises(RuntimeError, match="while request .* is loading"): manager.free_resources(req) assert not cache.closed - assert req.py_connector_served_position == 64 + assert req.py_served_prefix_end == 64 assert manager.kv_cache_map[req.request_id] is cache def test_final_batch_releases_only_rejected_requests(self) -> None: @@ -1040,6 +1040,6 @@ def test_final_admission_rejects_an_invalid_local_allocation() -> None: manager.report_batch_to_connector(scheduled(req)) assert connector.accepted == [] assert connector.get_prefix_reservation(req) is not None - assert req.py_connector_served_position == 0 + assert req.py_served_prefix_end == 0 manager.release_unused_connector_reservations(set()) assert connector.releases == [(1, 0, 64)] diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py index 88f0ef7ac143..78ed3910f0c9 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py @@ -467,7 +467,7 @@ def test_attention_sinks_are_never_masked(vswa_manager: KVCacheManagerV2) -> Non life_cycles = list(vswa_manager._life_cycle_by_layer_group()) life_cycles[sliding] = AttnLifeCycle.make(VSWA_WINDOW, num_sink_tokens, tokens_per_block) - vswa_manager._connector_life_cycle_by_group = life_cycles + vswa_manager._layer_group_life_cycles = life_cycles request = make_request(prompt_len=VSWA_PROMPT_LEN) assert schedule(vswa_manager, request) @@ -668,7 +668,7 @@ def test_a_sliding_single_group_reports_sentinels_in_the_flat_list(): def test_a_served_prefix_survives_re_entry_under_a_sliding_window(): """Here the rewind is not only wasted work. - ``_resize_for_connector_prefix`` sets the cache's ``history_length`` to the + ``_resize_for_served_prefix`` sets the cache's ``history_length`` to the served end, and raising history unlocks the blocks the window has passed. A cursor settled below that end points the forward pass at ordinals whose pages are gone. @@ -724,7 +724,7 @@ def test_rejected_connector_candidate_releases_and_rewinds( assert request.context_current_position == 0 assert request.prepopulated_prompt_len == 0 assert request.context_chunk_size == PROMPT_LEN - assert request.py_connector_served_position == 0 + assert request.py_served_prefix_end == 0 assert schedule(manager, request) assert connector.get_prefix_reservation(request).reservation_id == 2 assert connector.queries == [(request.request_id, 0), (request.request_id, 0)] diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index 17c96cd59a6b..3d8bb0ceebb8 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py @@ -63,7 +63,7 @@ class _StatsRequest: draft_tokens: list[int] = field(default_factory=list) state: LlmRequestState = LlmRequestState.GENERATION_IN_PROGRESS context_current_position: int = 0 - py_connector_served_position: int = 0 + py_served_prefix_end: int = 0 context_chunk_size: int = 0 expect_snapshot_points: list[int] = field(default_factory=list) prepopulated_prompt: tuple[int, int] | None = None