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