diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index dc16209c31ad..89ae82287b88 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -42,6 +42,19 @@ These methods run on the leader process and drive the connector's behavior. * **`update_state_after_alloc(self, request: LlmRequest, block_ids: list[int])`** * **Description**: a callback to update internal state after KV cache blocks have been allocated for the prefill. + * **Note**: on `KVCacheManagerV2` with chunked prefill, `block_ids` covers only the blocks allocated for the first chunk, because V2 allocates per chunk rather than for the whole prompt. The remaining blocks arrive as append-deltas in `RequestData.new_block_ids` on subsequent chunks. A connector that treats this callback as its only source of block ids will under-plan; drive off `build_connector_meta` instead. + * **Note**: on `KVCacheManagerV2` under sliding-window attention, a block that the window has already passed holds no page, and is reported as `-1` (`BAD_PAGE_INDEX`) **in place** rather than being dropped from the list. This keeps each entry aligned with its block ordinal, so entry `i` always describes prompt tokens `[i * tokens_per_block, (i+1) * tokens_per_block)` and an append-delta over successive calls stays valid. Connectors must skip `-1` entries rather than treating them as page slots. The same applies to `RequestData.new_block_ids` and to `cache_block_ids` in `request_finished`. + +* **`cancel_load(self, request: LlmRequest, start: int, end: int)`** + * **Description**: Optional, with a no-op default. Tells the connector that the runtime will not consume KV it offered from `get_num_new_matched_tokens` for prompt tokens `[start, end)`, so any ownership taken for that range can be released. Offsets are absolute prompt positions, on the same scale as `num_computed_tokens`. + * **When it fires**: only on `KVCacheManagerV2`, which asks during a speculative scheduling pass and resolves the answer later. Two things can happen in between, and both are reported here: the runtime may fail to allocate pages to cover the offer, in which case the request falls back to computing the prefix locally; or the request may be cancelled, time out or fail before it ever reaches a batch, in which case the whole offer is released. A third case -- the local cache overtaking part of the offer because another request committed the same prefix -- is handled by the same callback but cannot arise today, since a request's local match is fixed when its cache is created and only its own completed forward passes extend it. + * **Caveat**: best-effort. For a synchronous load nothing has been transferred yet, so cancelling is exact. For `is_async=True` the transfer necessarily started inside `get_num_new_matched_tokens`, so it may already be in flight. + +##### Serving a prefix on `KVCacheManagerV2` + +V1 answers `get_num_new_matched_tokens` from C++ while the block manager holds its radix-tree mutex, so the local match and the query are atomic and the answer is consumed immediately. V2 has no such mutex, and its scheduling pass is speculative -- a prepared request can still be dropped at the token budget, at resize, at multimodal alignment or at cross attention, and retried in a later iteration. + +The contract for connectors is unchanged, and in particular `get_num_new_matched_tokens` is still called **exactly once per request** on both managers -- a request that is asked and then deferred is not asked again when it comes back. What differs is that on V2 the runtime may resolve the answer in a later iteration than the one it asked in, and may by then be unable to honour part or all of it. That is what `cancel_load` reports. #### 2. Worker Interface (`KvCacheConnectorWorker`) diff --git a/examples/llm-api/llm_kv_cache_connector.py b/examples/llm-api/llm_kv_cache_connector.py index 882478993e5a..f5d9ba0ba51a 100644 --- a/examples/llm-api/llm_kv_cache_connector.py +++ b/examples/llm-api/llm_kv_cache_connector.py @@ -116,6 +116,21 @@ def register_kv_caches(self, kv_cache_tensor: torch.Tensor): assert self.kv_cache_tensor is None, "KV cache tensor already registered" self.kv_cache_tensor = kv_cache_tensor + def register_kv_cache_layout(self, layout): + # KVCacheManagerV2 describes its pools rather than handing over one + # tensor. This example targets the uniform case, where a layer group's + # buffers coalesce into a single whole-slot region; indexing that + # region by page slot is the direct analogue of indexing a pool by + # block id, so the load/save paths below need no change. + assert self.kv_cache_tensor is None, "KV cache tensor already registered" + if len(layout.groups) != 1 or len(layout.groups[0].regions) != 1: + raise NotImplementedError( + "This example connector handles a single layer group with a " + "single coalesced region (uniform attention, one pool). Got " + f"{len(layout.groups)} layer group(s) and " + f"{sum(len(g.regions) for g in layout.groups)} region(s).") + self.kv_cache_tensor = layout.groups[0].regions[0].as_tensor() + def start_load_kv(self, stream: torch.cuda.Stream): # Do all loads synchronously, and blockwise. for path, block_id in self._metadata.load: diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 45bc4b5e5a56..821987150900 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -269,7 +269,7 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: assert self._page_table is not None layer_groups = self._page_table.layer_groups - is_gen_only = req.is_generation_only_request() + is_gen_only = req.is_generation_only_request cached_per_lg = ( adapter.get_cached_token_count_per_layer_group(req, layer_groups) if is_gen_only diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 95c92473efa5..a074c2d9b1d4 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -653,9 +653,9 @@ def _validate_or_fallback_kv_cache_manager_v2( # also go through the V2-incompatible-feature gate below. if issubclass(kv_cache_manager_cls, KVCacheManagerV2): sparse_attn_config = model_config.sparse_attention_config + # The KV connector is supported on V2 via the pool layout + # registration path, so it no longer forces a fallback to V1. incompat: List[str] = [] - if self._kv_connector_manager is not None: - incompat.append("kv_connector_manager") if self._max_beam_width is not None and self._max_beam_width > 1: incompat.append("max_beam_width > 1") if incompat: @@ -682,8 +682,8 @@ def _validate_or_fallback_kv_cache_manager_v2( raise NotImplementedError( "Hybrid Mamba cache managers do not support " f"{incompat_str}; CppMambaHybridCacheManager does not " - "provide a compatible fallback. Use max_beam_width=1 " - "and disable the KV connector.") + "provide a compatible fallback. Disable the listed " + "features to run hybrid linear models.") # Plain V2 (explicitly enabled or selected by a model preference): # V2 was a preference, not a structural requirement, so we can # safely fall back to V1. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 99da2a42265c..49f6f27e9998 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -54,6 +54,7 @@ if TYPE_CHECKING: from ..resource_manager import KVCacheManager + from .kv_cache_layout import KvCacheLayout # Used to store data for a single inflight request. @@ -85,6 +86,11 @@ class RequestData: # remote object id) MUST mix cache_salt into their identifiers, # otherwise blocks from a different salt could be incorrectly reused. cache_salt: Optional[str] = None + # New page slot indices keyed by layer group, populated under + # KVCacheManagerV2. Blocks with no page in a group -- the sliding-window + # case -- appear as BAD_PAGE_INDEX in place, keeping ordinals stable. + # Empty under the V1 manager, whose block IDs are a single flat space. + new_block_ids_by_layer_group: Dict[int, List[int]] = field(default_factory=dict) # A class to store some basic data regarding all inflight requests. @@ -135,6 +141,30 @@ def register_kv_caches(self, kv_cache_tensor: torch.Tensor): kv_cache_tensor: The contiguous KV cache tensor. """ + def register_kv_cache_layout(self, layout: "KvCacheLayout") -> None: + """ + Register the KV cache pools described by ``layout``. + + Called instead of ``register_kv_caches`` when the KV cache manager is + ``KVCacheManagerV2``, whose memory cannot be expressed as a single + tensor: there is one slot address space per pool and one page-index + space per layer group. + + The bytes for page slot ``i`` of a region live at + ``region.base + region.stride * i`` for ``region.size`` bytes, or + equivalently at ``region.as_tensor()[i]``. Page indices arriving in + ``RequestData.new_block_ids`` are scoped to a layer group and index the + regions of that group. + + Args: + layout: Description of the KV cache pools; see ``KvCacheLayout``. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement register_kv_cache_layout, so it " + "cannot run against KVCacheManagerV2. Implement it, or select the V1 KV " + "cache manager with kv_cache_config.use_kv_cache_manager_v2=False." + ) + @abstractmethod def start_load_kv(self, stream: torch.cuda.Stream): """ @@ -255,6 +285,30 @@ def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): block_ids: The KV cacheblock IDs that were allocated. """ + def cancel_load(self, request: LlmRequest, start: int, end: int): + """ + Called when the runtime will not consume KV that was offered by + ``get_num_new_matched_tokens`` for prompt tokens ``[start, end)``. + + Offsets are absolute prompt positions, on the same scale as + ``num_computed_tokens``. This is raised when the local cache overtook + part or all of an offer while the request was waiting to be scheduled, + or when the runtime could not allocate pages to cover it. + + Best-effort: release any ownership taken in + ``get_num_new_matched_tokens`` for that range. For a synchronous load + nothing has been transferred yet, so this is exact. For an + asynchronous load the transfer necessarily began inside + ``get_num_new_matched_tokens`` (see ``take_scheduled_requests_pending_load``), + so it may already be in flight and cancelling is genuinely lossy. + + Args: + request: The request whose offer is being handed back. + start: First prompt position that will not be consumed. + end: One past the last prompt position that will not be consumed. + """ + return + def wait_for_initialization(self): """ Some connectors need to wait for some resources to be initialized. @@ -316,21 +370,49 @@ def loading_ids(self) -> Set[int]: class KvCacheConnectorSchedulerOutputRequest: def __init__(self): self.block_ids = [] + self.block_ids_by_layer_group: Dict[int, List[int]] = {} self.tokens = [] def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManager"): - block_ids = kv_cache_manager.get_cache_indices(req) + from ..kv_cache_manager_v2 import KVCacheManagerV2 + + is_v2 = isinstance(kv_cache_manager, KVCacheManagerV2) tokens = req.get_tokens(0) - # Commit hashes for any blocks that have become full since the last call - # and read back the full cumulative chain. The C++ side sets each block's - # mBlockKey/mHash on first call, so subsequent calls become pure lookups. - block_hashes = kv_cache_manager.commit_and_get_block_hashes(req) + new_block_ids_by_layer_group: Dict[int, List[int]] = {} + if is_v2: + # Block hashes and retention priorities have no V2 accessor yet, so + # they are reported empty rather than guessed at. + block_hashes = [] + indices_by_group = kv_cache_manager.get_page_indices_by_layer_group(req) + for layer_group_id, indices in indices_by_group.items(): + seen = self.block_ids_by_layer_group.setdefault(layer_group_id, []) + new_ids = indices[len(seen) :] + seen.extend(new_ids) + new_block_ids_by_layer_group[layer_group_id] = new_ids + # With a single layer group -- every non-VSWA, non-hybrid model -- + # ``new_block_ids`` carries that group's indices so connectors that + # do not reason about layer groups keep working. With several groups + # it is left empty and ``new_block_ids_by_layer_group`` is the only + # correct source. + new_block_ids = ( + next(iter(new_block_ids_by_layer_group.values())) + if len(new_block_ids_by_layer_group) == 1 + else [] + ) + self.block_ids.extend(new_block_ids) + else: + block_ids = kv_cache_manager.get_cache_indices(req) - new_block_ids = block_ids[len(self.block_ids) :] - new_tokens = tokens[len(self.tokens) :] + # Commit hashes for any blocks that have become full since the last call + # and read back the full cumulative chain. The C++ side sets each block's + # mBlockKey/mHash on first call, so subsequent calls become pure lookups. + block_hashes = kv_cache_manager.commit_and_get_block_hashes(req) + + new_block_ids = block_ids[len(self.block_ids) :] + self.block_ids.extend(new_block_ids) - self.block_ids.extend(new_block_ids) + new_tokens = tokens[len(self.tokens) :] self.tokens.extend(new_tokens) if req.state in ( @@ -346,9 +428,13 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag ) # Specdec with draft tokens is not supported yet. # Get retention priority for each new block only if retention config is provided - # (for priority-based offload filtering) + # (for priority-based offload filtering). Priorities stay None on + # KVCacheManagerV2: it does not implement `KvCacheRetentionConfig` at + # all -- its per-page priority comes from `custom_priority_callback`, + # which KVCacheManagerV2 never overrides -- so every page carries the + # default and reporting it would misdescribe what the user asked for. priorities = None - if req.kv_cache_retention_config is not None: + if not is_v2 and req.kv_cache_retention_config is not None: priorities = [ kv_cache_manager.get_priority_by_block_id(block_id) for block_id in new_block_ids ] @@ -362,6 +448,7 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag block_hashes=block_hashes, priorities=priorities, cache_salt=req.cache_salt, + new_block_ids_by_layer_group=new_block_ids_by_layer_group, ) @@ -463,7 +550,18 @@ def _run_on_leader(self, f: Callable[[], Any]) -> Any: res = None return mpi_broadcast(res, root=0) - def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> int: + def query_num_new_matched_tokens( + self, request: LlmRequest, num_computed_tokens: int + ) -> Tuple[int, bool]: + """Ask the connector how much of the prompt it can serve. No side effects. + + This is the half of ``get_num_new_matched_tokens`` that must run at most + once per request, because the connector ABC promises exactly one query + per request and connectors take ownership of remote blocks in it. + Callers that cannot commit to consuming the answer in the same iteration + (KVCacheManagerV2, whose scheduling pass is speculative) query here and + commit later via ``commit_new_matched_tokens``. + """ if request.is_generation_only_request: raise RuntimeError("Connector API is not supported for generation-only requests!") @@ -474,6 +572,18 @@ def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: i if num_tokens == 0 and load_kv_async: raise RuntimeError("load_kv_async must be False when num_tokens is 0!") + return num_tokens, load_kv_async + + def commit_new_matched_tokens( + self, request: LlmRequest, num_tokens: int, load_kv_async: bool + ) -> None: + """Register the runtime's side of an answered query. + + Must run in the iteration the request is actually scheduled: + ``external_loads`` is consumed and cleared by every + ``build_scheduler_output``, and ``new_async_requests.loading`` is read + by that same call to suppress the request's ``RequestData``. + """ # TODO(jthomson04): This part is a bit ugly. # When the connector indicates that a request will be loaded # asynchronously, we need to suspend its execution. This is @@ -488,8 +598,21 @@ def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: i request.py_num_connector_matched_tokens = num_tokens + def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> int: + """Query and commit in one step. + + This is the V1 entry point, called from C++ while the block manager + holds the radix-tree mutex, so the local match and the query are atomic + with respect to the tree and the answer can be committed immediately. + """ + num_tokens, load_kv_async = self.query_num_new_matched_tokens(request, num_computed_tokens) + self.commit_new_matched_tokens(request, num_tokens, load_kv_async) return num_tokens + def cancel_load(self, request: LlmRequest, start: int, end: int) -> None: + if self.scheduler is not None: + self.scheduler.cancel_load(request, start, end) + def should_add_sequence(self, request: LlmRequest) -> bool: req_id = request.request_id return req_id not in self.finished_async_loading_requests diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py new file mode 100644 index 000000000000..8c4b1a77198e --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py @@ -0,0 +1,245 @@ +# 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. +"""KV cache layout description handed to a KV connector under KVCacheManagerV2. + +V2 allocates *pool groups*, each holding ``num_slots`` slots. A slot holds the +coalesced buffers of one *layer group* (a life cycle), and a buffer is keyed by +``BufferId(layer_id, role)``. A connector therefore cannot be handed a single +pool tensor: there is one slot address space per pool, and one index space per +layer group. + +Instead it is handed :class:`KvCacheLayout` -- a description of the byte ranges +that repeat per page slot. The addressing contract is V2's own, taken verbatim +from ``AggregatedPageDesc``:: + + (base + stride * i + Range(0, size) for i in aggregated_page_indices) + +where ``i`` comes from ``_KVCache.get_aggregated_page_indices(layer_group_id)``. + +Because ranges are described rather than implied, this covers MLA (a pool simply +has no VALUE buffer), sliding-window attention and hybrid models (one layer group +per window size) without any of them being special cases. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +import torch + +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor + +if TYPE_CHECKING: + from ..kv_cache_manager_v2 import KVCacheManagerV2 + +__all__ = [ + "KvCacheBufferRef", + "KvCacheLayerGroupLayout", + "KvCacheLayout", + "KvCacheRegion", + "build_kv_cache_layout_v2", +] + + +@dataclass(frozen=True) +class KvCacheBufferRef: + """One ``(layer, role)`` buffer covered by a region, in memory order.""" + + #: Global model layer index -- the same index space the per-layer connector + #: hooks (``wait_for_layer_load`` / ``save_kv_layer``) receive. + layer_id: int + #: Native role name as the cache manager spells it, e.g. "key" / "value". + #: Deliberately not an enum: roles are an open vocabulary on the manager + #: side, and a connector should not need updating when one is added. + role: str + #: Page expansion factor for heterogeneous tokens-per-block layers. + expansion: int = 1 + + +@dataclass(frozen=True) +class KvCacheRegion: + """A contiguous byte range that repeats once per page slot. + + The data for page slot ``i`` lives at ``base + stride * i``, for ``size`` + bytes. ``size`` is not necessarily equal to ``stride``: a region covers one + run of adjacent buffers within a slot, and a slot may hold several runs. + """ + + base: int + size: int + stride: int + num_slots: int + buffers: Tuple[KvCacheBufferRef, ...] + + def address_of(self, slot_id: int) -> int: + """Device address of this region for ``slot_id``.""" + if not 0 <= slot_id < self.num_slots: + raise IndexError(f"slot_id {slot_id} out of range [0, {self.num_slots})") + return self.base + self.stride * slot_id + + def as_tensor(self, dtype: torch.dtype = torch.uint8) -> torch.Tensor: + """A strided ``[num_slots, size // itemsize]`` view; row ``i`` is slot ``i``. + + Defaults to ``uint8``. A region may span several roles whose element + types differ, and a connector that only moves bytes should not have to + care; callers that want a typed view can pass ``dtype`` explicitly. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + if self.size % itemsize or self.stride % itemsize: + raise ValueError( + f"region size {self.size} and stride {self.stride} must both be " + f"multiples of {dtype} itemsize {itemsize}" + ) + return convert_to_torch_tensor( + TensorWrapper( + self.base, + dtype, + shape=(self.num_slots, self.size // itemsize), + strides=(self.stride // itemsize, 1), + ) + ) + + +@dataclass(frozen=True) +class KvCacheLayerGroupLayout: + """One layer group -- the unit that page indices are scoped to.""" + + layer_group_id: int + #: Global model layer indices belonging to this group. + layer_ids: Tuple[int, ...] + #: Attention window for this group, or None for full attention. + window_size: Optional[int] + regions: Tuple[KvCacheRegion, ...] + + @property + def bytes_per_page(self) -> int: + """Total bytes this group occupies for a single page slot.""" + return sum(region.size for region in self.regions) + + +@dataclass(frozen=True) +class KvCacheLayout: + """What a connector is handed in place of a single KV cache pool tensor.""" + + tokens_per_block: int + groups: Tuple[KvCacheLayerGroupLayout, ...] + + def group(self, layer_group_id: int) -> KvCacheLayerGroupLayout: + for group in self.groups: + if group.layer_group_id == layer_group_id: + return group + raise KeyError(f"no layer group {layer_group_id} in layout") + + def group_of_layer(self, layer_id: int) -> KvCacheLayerGroupLayout: + """The layer group owning a global model layer index.""" + for group in self.groups: + if layer_id in group.layer_ids: + return group + raise KeyError(f"layer {layer_id} is not covered by this layout") + + +def _global_layer_ids(manager: "KVCacheManagerV2", local_layer_ids) -> List[int]: + """Map V2-internal layer ids to global model layer indices. + + ``pp_layers`` is the local-to-global table the manager already keeps. Models + that map several internal layers onto one model layer (the sparse-attention + virtual-layer path) have no single global index per internal layer, so they + are rejected rather than silently mislabelled. + """ + if hasattr(manager, "_layer_attn_to_layer_id"): + raise NotImplementedError( + "KV connector layout is not supported for managers with virtual " + "attention layers (sparse attention): an internal layer does not map " + "to a single model layer index." + ) + pp_layers = manager.pp_layers + return [int(pp_layers[int(lid)]) for lid in local_layer_ids] + + +def _window_size(init_config, local_layer_id: int) -> Optional[int]: + layers = init_config.layers + if local_layer_id >= len(layers): + raise ValueError(f"no layer config for internal layer {local_layer_id}") + window = getattr(layers[local_layer_id], "window_size", None) + return None if window is None else int(window) + + +def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: + """Describe a ``KVCacheManagerV2``'s GPU pools for a KV connector. + + Built only from V2's public layout API -- ``layer_grouping``, + ``all_buffer_ids``, ``get_aggregated_pages`` and ``pool_group_descs``. No + private storage state is touched, and no assumption is made about dimension + order, kv factor, or the number of pools. + """ + impl = manager.impl + init_config = impl.init_config + + # A pool group's slot count applies to every layer group drawn from it. + # Note LayerGroupId and PoolGroupIndex are distinct index spaces; the + # variants of a pool group name the layer groups it backs. + slots_by_group: Dict[int, int] = {} + for pool_group in impl.pool_group_descs: + for variant in pool_group.slot_desc.variants: + slots_by_group[int(variant.layer_group_id)] = int(pool_group.num_slots) + + buffers_by_layer: Dict[int, List] = {} + for buffer_id in impl.all_buffer_ids: + buffers_by_layer.setdefault(int(buffer_id.layer_id), []).append(buffer_id) + + groups: List[KvCacheLayerGroupLayout] = [] + for layer_group_id, local_layer_ids in enumerate(impl.layer_grouping): + local_layer_ids = [int(lid) for lid in local_layer_ids] + if not local_layer_ids: + continue + + num_slots = slots_by_group[layer_group_id] + global_by_local = dict(zip(local_layer_ids, _global_layer_ids(manager, local_layer_ids))) + + buffer_ids = [b for lid in local_layer_ids for b in buffers_by_layer.get(lid, ())] + + regions: List[KvCacheRegion] = [] + for desc in impl.get_aggregated_pages(buffer_ids): + if int(desc.layer_group_id) != layer_group_id: + continue + regions.append( + KvCacheRegion( + base=int(desc.base), + size=int(desc.size), + stride=int(desc.stride), + num_slots=num_slots, + buffers=tuple( + KvCacheBufferRef( + layer_id=global_by_local[int(b.id.layer_id)], + role=str(b.id.role), + expansion=int(b.expansion), + ) + for b in desc.buffers + ), + ) + ) + + groups.append( + KvCacheLayerGroupLayout( + layer_group_id=layer_group_id, + layer_ids=tuple(global_by_local[lid] for lid in local_layer_ids), + window_size=_window_size(init_config, local_layer_ids[0]), + regions=tuple(regions), + ) + ) + + return KvCacheLayout( + tokens_per_block=int(manager.tokens_per_block), + groups=tuple(groups), + ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index b1a965d20e32..c48cbbdeb781 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -804,10 +804,8 @@ def __init__( self.mapping = mapping self.dtype = dtype self.is_disagg = is_disagg + self.kv_connector_manager = kv_connector_manager - assert kv_connector_manager is None, ( - "kv_connector_manager is not supported for KVCacheManagerV2" - ) assert max_beam_width == 1, "max_beam_width must be 1 for KVCacheManagerV2" self.kv_cache_type = kv_cache_type @@ -1045,7 +1043,21 @@ def append_to_kv_heads_per_layer( logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB") cache_tiers: List[CacheTierConfig] = [GpuCacheTierConfig(quota=int(quota))] - if kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: + if kv_connector_manager is not None and kv_cache_config.host_cache_size is None: + # A KV connector registers device addresses for its pages, and a + # page evicted to another tier has its GPU slot reassigned. The + # automatic host tier below exists only to give the MAX_UTILIZATION + # scheduler's suspend/resume somewhere to spill to, and a connector + # run cannot use that policy (py_executor_creator requires + # GUARANTEED_NO_EVICT), so skip it rather than silently migrating + # pages out from under the connector. An explicitly configured + # host_cache_size is left alone and rejected loudly at bring-up. + host_quota = 0 + logger.info( + "KV cache manager v2 host tier disabled: a KV connector is attached " + "and registers GPU page addresses that tier migration would invalidate." + ) + elif kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: host_quota = kv_cache_config.host_cache_size else: # The V2 MAX_UTILIZATION scheduler relies on suspend/resume to @@ -2501,11 +2513,22 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: kv_cache.num_committed_tokens, self.tokens_per_block ) + # Connector phase 1: ask what the connector can serve past the local + # match and advance the context position over it, before residency + # is decided. The capacity/history bump this implies has to wait for + # the resume below, because `resize` requires an ACTIVE cache and a + # freshly created one is suspended. + position = self._connector_prefix_position(req, kv_cache) + if req.is_disagg_generation_init_state: # Disagg generation receives prompt KV from the context worker; # scratch blocks are only valid for local prefill chunks. kv_cache.enable_swa_scratch_reuse = False - return self._resume_and_restore(req.py_request_id, kv_cache) + if not self._resume_and_restore(req.py_request_id, kv_cache): + return False + if position is not None: + self._reserve_connector_prefix(req, kv_cache, position) + return True else: # Subsequent chunk: cache must exist from first chunk. # It may be suspended (e.g., evicted between chunks), so @@ -2516,6 +2539,171 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: ) return self._resume_and_restore(req.py_request_id, kv_cache) + # ---- KV connector prefix (see docs/source/features/kv-cache-connector.md) ---- + # + # V2 splits "match and take ownership" from "become resident on GPU", and the + # connector gets the same decomposition: + # + # phase 1 _connector_prefix_position ask, and advance the context position + # phase 2 _reserve_connector_prefix cover the offer with pages + # phase 3 _deliver_connector_prefix tell the connector what to transfer + # + # Phases 1 and 2 run in the scheduling pass, which is speculative: a request + # can still be dropped afterwards at the token budget, `resize_context`, + # multimodal alignment or cross attention. Phase 3 runs on the final batch. + + def _connector_may_serve(self, req: LlmRequest) -> bool: + if req.is_dummy: + # Mirrors V1, which skips the query for dummy requests entirely. + return False + if req.is_generation_only_request: + # The connector API rejects these outright; a disagg generation + # server gets its prompt KV from the context worker. + return False + return not req.is_disagg_generation_init_state + + def _connector_prefix_position(self, req: LlmRequest, kv_cache) -> Optional[int]: + """Phase 1. Position through which KV is valid before the forward pass. + + That is ``max(locally committed, connector offer end)``, or None when no + connector prefix applies to this request. The connector is queried at + most once per request: a deferred request re-derives its position here + from the memoised offer end and the *current* local match, without + asking again. + """ + if self.kv_connector_manager is None or self.is_draft: + return None + + if req.py_connector_prefix_end is None: + if not self._connector_may_serve(req): + return None + local_end = kv_cache.num_committed_tokens + num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( + req, local_end + ) + # The last prompt position must be computed locally whatever the + # connector holds, because the first generation step consumes its + # activations. A connector that offers the whole prompt therefore + # loses at most one token, which the forward recomputes. + req.py_connector_prefix_start = local_end + offered_end = local_end + num_tokens + req.py_connector_prefix_end = min(offered_end, req.prompt_len - 1) + # Hand back what that clamp just dropped, because nothing + # downstream can: `_deliver_connector_prefix` and + # `_release_undelivered_connector_prefix` both read the + # *clamped* end, so the tail past it would stay owned by the + # connector for the life of the process -- the exact leak + # those two exist to prevent. + dropped_start = max(local_end, req.py_connector_prefix_end) + if offered_end > dropped_start: + self.kv_connector_manager.cancel_load(req, dropped_start, offered_end) + # Deliberately not conditioned on the clamp above: if the connector + # said it would transfer asynchronously it has already started, and + # the request must be held out of the batch until it reports done. + req.py_connector_load_async = load_async + + # Not `py_connector_prefix_end` on its own: while the request waited, + # another request may have committed blocks that grew the local match + # past the offer. + return max(kv_cache.num_committed_tokens, req.py_connector_prefix_end) + + def _reserve_connector_prefix(self, req: LlmRequest, kv_cache, position: int) -> None: + """Phase 2. Give the offered prefix capacity, history and pages. + + Runs after ``_resume_and_restore`` because ``resize`` asserts the cache + is ACTIVE. Capacity and history move together in one call: after a reuse + match they are both equal to the local match, so raising history alone + would trip "History length cannot be greater than capacity". + + Raising history is what keeps a served prefix from allocating a page per + block in a sliding-window layer group -- ``history_length`` is the sole + input to the stale-range computation. + """ + local_end = kv_cache.num_committed_tokens + if position <= local_end: + # Nothing offered, or the local match has caught up with the offer. + return + + # SWA scratch slots are transient prefill storage; a connector writes + # real cache content into these blocks. Same reason the disagg + # generation path above opts out. + kv_cache.enable_swa_scratch_reuse = False + + if kv_cache.resize(max(kv_cache.capacity, position), position): + req.context_current_position = position + req.set_prepopulated_prompt_len(position, self.tokens_per_block) + return + + # Out of pages. Run local-only rather than leaving the offered range + # unallocated, and hand the offer back instead of silently dropping it. + logger.debug( + "req %s: could not reserve connector prefix up to %d, falling back to " + "the local match at %d", + req.py_request_id, + position, + local_end, + ) + self.kv_connector_manager.cancel_load(req, req.py_connector_prefix_start, position) + req.py_connector_prefix_start = local_end + req.py_connector_prefix_end = local_end + req.py_connector_load_async = False + + def _deliver_connector_prefix(self, req: LlmRequest) -> None: + """Phase 3. Resolve the offer against ownership as it stands now. + + The connector was told it could serve ``[start, end)``. Anything below + the current commit boundary is now locally owned -- shared + ``CommittedPage`` s that this request is not the only writer to -- so it + is handed back rather than transferred, and only ``[committed, end)`` + is reported as an external load. + """ + end = req.py_connector_prefix_end + if end is None: + return + start = req.py_connector_prefix_start + kv_cache = self.kv_cache_map.get(req.py_request_id) + # Never past `end`: the offer bounds what there is to hand back. + committed = min(kv_cache.num_committed_tokens, end) if kv_cache is not None else start + if committed > start: + self.kv_connector_manager.cancel_load(req, start, committed) + recorded = end - committed + # The one place phases 2 and 3 have to agree. The runtime reports + # `context_current_position - recorded` to the connector as the range it + # computed locally, and nothing downstream validates that: the + # subtraction is unguarded (kv_cache_connector.py:480) and connectors + # divide the result into block ordinals rather than checking it (see + # `computed_position // block_size` in + # examples/llm-api/llm_kv_cache_connector.py). So a `recorded` the + # position was never advanced over does not fail here -- it silently + # points the connector at the wrong offset, inside its own code. Fail + # loudly and locally instead. + assert 0 <= recorded <= req.context_current_position, ( + f"req {req.py_request_id}: connector prefix [{start}, {end}) " + f"records {recorded} externally loaded tokens, but the context " + f"position is only {req.context_current_position} -- phase 2 did " + f"not reserve what phase 1 offered" + ) + self.kv_connector_manager.commit_new_matched_tokens( + req, recorded, req.py_connector_load_async + ) + + def _release_undelivered_connector_prefix(self, req: LlmRequest) -> None: + """Hand back an offer for a request that died before it was delivered. + + Phase 1 is speculative, so a request can be asked and then cancelled, + time out, or fail before it ever reaches a batch. The connector took + ownership of remote blocks in the query and would otherwise hold them + for the rest of the process's life. + """ + if self.kv_connector_manager is None or req.py_connector_delivered: + return + end = req.py_connector_prefix_end + start = req.py_connector_prefix_start + if end is None or end <= start: + return + self.kv_connector_manager.cancel_load(req, start, end) + req.py_connector_prefix_end = start + def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: """Resize KV cache to cover context_current_position + num_tokens. @@ -2641,6 +2829,40 @@ def resume_request(self, req: LlmRequest) -> bool: # ---- prepare_resources ---- @nvtx_range("prepare_resources_kv_cache_manager_v2") + def get_page_indices_by_layer_group(self, request: LlmRequest) -> Dict[int, List[int]]: + """Per-layer-group page slot indices for ``request``, by block ordinal. + + ``valid_only=False`` is deliberate: the positionally-aligned form yields + one entry per block ordinal, which is what preserves the ordinal-to-token + -range mapping a connector needs. A block with no page in a given layer + group -- the sliding-window case -- reads back as ``BAD_PAGE_INDEX`` in + place rather than shortening the list, so ordinals stay stable and an + append-delta over successive calls remains valid. + """ + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return {} + return { + layer_group_id: list( + kv_cache.get_aggregated_page_indices(layer_group_id, valid_only=False) + ) + for layer_group_id in range(len(self.impl.layer_grouping)) + } + + def get_connector_page_indices(self, request: LlmRequest) -> List[int]: + """Flat page slot indices for ``request``, for connectors. + + A page index is scoped to a layer group, so there is no correct way to + flatten indices from several groups into one list. With a single group + -- every non-VSWA, non-hybrid model -- that group's indices are the flat + list; with several, connectors must read + ``new_block_ids_by_layer_group`` off the scheduler output instead. + """ + indices_by_group = self.get_page_indices_by_layer_group(request) + if len(indices_by_group) != 1: + return [] + return next(iter(indices_by_group.values())) + def prepare_resources(self, scheduled_batch: ScheduledRequests): if self.is_draft: # Draft V2 manager: mirror the main manager by creating/resizing @@ -2649,6 +2871,39 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self._prepare_draft_resources(scheduled_batch) return + # V2 allocates in KVCacheV2Scheduler (prepare_context / resize_context) + # rather than here, so by this point every scheduled request already has + # its pages. This is the same point in the iteration at which the V1 + # manager drives the connector's scheduler-side hooks, and page indices + # are available, so the connector is driven from here. + if self.kv_connector_manager is not None: + self._run_kv_connector_hooks(scheduled_batch) + + def _run_kv_connector_hooks(self, scheduled_batch: ScheduledRequests) -> None: + """Report freshly allocated pages, then build the connector metadata.""" + for request in scheduled_batch.context_requests: + # Mirror V1, which reports allocation once per sequence, right after + # the sequence is added -- that is the first context chunk. + if not request.is_first_context_chunk: + continue + # A request that was loaded asynchronously re-enters this batch still + # on its first context chunk, with the same pages and nothing left to + # load. Reporting it again would re-record an external load and fire + # a second `update_state_after_alloc` for one allocation. + if request.py_connector_delivered: + continue + self._deliver_connector_prefix(request) + request.py_connector_delivered = True + # Connectors that do not reason about layer groups see the single + # group's indices; with several groups there is no correct flat + # list, so report none and leave them to the layer-group-aware + # metadata carried on RequestData. + self.kv_connector_manager.update_state_after_alloc( + request, self.get_connector_page_indices(request) + ) + + self.kv_connector_manager.build_scheduler_output(scheduled_batch, self) + def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): """Create/resize KV caches in the draft V2 manager for scheduled requests. @@ -3536,6 +3791,7 @@ def release_index_slot(self, request_id: int) -> None: self._early_freed_index_requests.add(request_id) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + self._release_undelivered_connector_prefix(request) if self.conversation_manager is not None: self.conversation_manager.finish_request(request) self._allocated_draft_lens.pop(request.py_request_id, None) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 65f613337fff..8651185d6006 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -13,6 +13,7 @@ CacheTransceiverConfig) from tensorrt_llm.mapping import Mapping +from .kv_cache_manager_v2 import KVCacheManagerV2 from .llm_request import LlmRequest from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, @@ -204,6 +205,26 @@ def create_kv_cache_transceiver( "MambaHybridCacheManagerV2 requires transceiver_runtime='PYTHON' " "with backend='NIXL'; it cannot use the C++ transceiver.") + # The same applies to KVCacheManagerV2 itself, not only to its hybrid + # subclass: its `impl` is the Python V2 core's KVCacheManager, and + # CacheTransceiverCpp is bound to BaseKVCacheManager. Without this check the + # combination reaches BindKvCacheTransceiver and dies on a nanobind + # signature mismatch that names neither the manager nor the way out. + # + # Note `transceiver_runtime` defaults to "auto", which is resolved from the + # *model's* preference (llm_utils._resolve_transceiver_runtime_auto) and + # knows nothing about which cache manager will be built -- so most models + # land here rather than on the Python transceiver. + if isinstance(kv_cache_manager, + KVCacheManagerV2) and not use_python_transceiver: + raise ValueError( + "KVCacheManagerV2 requires transceiver_runtime='PYTHON' with " + "backend='NIXL' for disaggregated serving; it cannot use the C++ " + "transceiver, which is bound to the V1 BaseKVCacheManager. Either " + "set cache_transceiver_config.transceiver_runtime='PYTHON' and " + "backend='NIXL', or select the V1 manager with " + "kv_cache_config.use_kv_cache_manager_v2=False.") + if use_python_transceiver: if isinstance(mamba_cache_manager, CppMambaHybridCacheManager): raise ValueError( diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 041c4933542a..a1b5b1763aea 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1050,6 +1050,18 @@ def __init__( self.py_num_connector_matched_tokens = 0 + # KV connector query state, used by KVCacheManagerV2 (V1 answers the + # query under the block manager's tree mutex and needs no state). + # UNASKED (`py_connector_prefix_end is None`) -> ASKED -> DELIVERED. + # `py_connector_prefix_end` is the absolute prompt position one past + # the last token the connector offered to serve, deliberately not the + # returned delta: the request may be deferred and re-derive a larger + # locally-matched prefix, against which a delta would overshoot. + self.py_connector_prefix_start: Optional[int] = None + self.py_connector_prefix_end: Optional[int] = None + self.py_connector_load_async = False + self.py_connector_delivered = False + self.py_result = PyResult( prompt_len=self.py_prompt_len, max_new_tokens=self.py_max_new_tokens, @@ -1179,7 +1191,12 @@ def reset_for_recompute(self, max_input_len: int) -> None: orig_prompt_len=self.prompt_len, clear_draft_tokens=True) + @property def is_generation_only_request(self): + # Must stay a property: the C++ base exposes this as a read-only + # property (nanobind/batch_manager/bindings.cpp), and a plain method + # here shadows it with something always-truthy for any caller that + # reads it as an attribute. return self.py_llm_request_type == LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY def create_response(self, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ae97f89bb386..8ef2ce1b2f2b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -132,7 +132,7 @@ def _make_single_token_context_graph_batch( or request.py_beam_width != 1 or get_draft_token_length(request) > 0 or request.py_is_first_draft or request.is_context_only_request - or request.is_generation_only_request() + or request.is_generation_only_request or request.py_disaggregated_params is not None or request.py_mm_encoder_event is not None or (request.py_multimodal_data is not None and diff --git a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py index ca0cb4a5c0c7..0b131ea8105a 100644 --- a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py @@ -250,7 +250,7 @@ def append_step_metrics(self, request, iter_counter: int, batch_token_time=None) # - py_decoding_iter == 1 and not yet marked complete: last/only chunk # - Gen-only requests (disagg gen server) are never ctx is_ctx = ( - not request.is_generation_only_request() + not request.is_generation_only_request and not perf.ctx_chunks_complete and request.py_decoding_iter <= 1 ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 19ab0cffc873..3a60536c8d75 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -66,6 +66,7 @@ from ..speculative.speculation_gate import SpeculationGate from .adp_iter_stats import ADPIterStatsBuffer from .connectors.kv_cache_connector import KvCacheConnectorManager +from .connectors.kv_cache_layout import build_kv_cache_layout_v2 from .dwdp import DwdpManager from .error_classification import ErrorBudget from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem @@ -989,8 +990,16 @@ def _maybe_init_kv_connector_manager(self): "connector scheduler / worker hooks and are not " "distinguished from real requests.") + if self.kv_cache_manager is None: + raise ValueError( + "KV Cache Connector requires a KV Cache Manager.") + + is_kv_cache_manager_v2 = isinstance(self.kv_cache_manager, + KVCacheManagerV2) + kv_cache_config = getattr(self.llm_args, 'kv_cache_config', None) - if kv_cache_config is not None and kv_cache_config.host_cache_size: + if not is_kv_cache_manager_v2 and (kv_cache_config is not None and + kv_cache_config.host_cache_size): raise NotImplementedError( "KV Cache Connector is not supported with KV cache host " "offloading (KvCacheConfig.host_cache_size). The connector " @@ -1000,11 +1009,11 @@ def _maybe_init_kv_connector_manager(self): "streams are not synchronized with the internal " "onboard/offload streams.") - if self.kv_cache_manager is None: - raise ValueError( - "KV Cache Connector requires a KV Cache Manager.") - - if getattr(self.kv_cache_manager, 'is_vswa', False): + # VSWA allocates one pool per window size. That is fatal for the V1 + # single-tensor registration, but is the normal case for V2, whose + # layout describes one region set per layer group. + if not is_kv_cache_manager_v2 and getattr(self.kv_cache_manager, + 'is_vswa', False): raise NotImplementedError( "KV Cache Connector is not supported with variable " "sliding-window attention (per-layer max_attention_window " @@ -1020,8 +1029,20 @@ def _maybe_init_kv_connector_manager(self): "per-layer load/save hooks have nothing meaningful to " "transfer for those layers.") - kv_tensor = self.kv_cache_manager.get_unique_primary_pool() - self.kv_connector_manager.worker.register_kv_caches(kv_tensor) + if is_kv_cache_manager_v2: + # A registered region is only a valid address while its page is + # pinned to GPU. V2 migrates pages between cache tiers, and it + # provisions a host tier by default, so reject any non-GPU tier + # until the connector participates in migration. Read the + # resolved tier list rather than KvCacheConfig.host_cache_size: + # the default of None is falsy but still yields a host tier. + self._reject_non_gpu_cache_tiers(self.kv_cache_manager) + layout = build_kv_cache_layout_v2(self.kv_cache_manager) + self.kv_connector_manager.worker.register_kv_cache_layout( + layout) + else: + kv_tensor = self.kv_cache_manager.get_unique_primary_pool() + self.kv_connector_manager.worker.register_kv_caches(kv_tensor) # For each of our layers, we need to register the pre/post hooks. # These are used for methods like `wait_for_layer_load` and `save_kv_layer`. @@ -1034,6 +1055,36 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() + @staticmethod + def _reject_non_gpu_cache_tiers(kv_cache_manager) -> None: + """Reject KVCacheManagerV2 cache tiers below GPU while a connector runs. + + A registered region is only a valid device address while its page is + pinned to GPU; eviction to another tier reassigns the page's slot. Until + the connector participates in migration, allow GPU-only configurations. + + The resolved tier list is read from the manager rather than from + ``KvCacheConfig.host_cache_size`` on purpose: V2 provisions a host tier + automatically when that field is left at its default of ``None``, which + is falsy and would slip past a truthiness check. + """ + from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheTier + + cache_tiers = getattr(kv_cache_manager.impl.init_config, "cache_tiers", + None) + if not cache_tiers: + return + extra = [tier for tier in cache_tiers if tier.tier != CacheTier.GPU_MEM] + if extra: + names = ", ".join(str(tier.tier) for tier in extra) + raise NotImplementedError( + "KV Cache Connector is not supported with KVCacheManagerV2 " + f"cache tiers below GPU (found: {names}). Pages evicted to " + "another tier have their GPU slot reassigned, which would " + "invalidate the addresses registered with the connector. Set " + "KvCacheConfig.host_cache_size=0 and " + "KvCacheConfig.disk_cache_size=0 to run GPU-only.") + def _end_transfer_and_maybe_terminate(self, request: LlmRequest): transfer_failed = request.state == LlmRequestState.DISAGG_TRANS_ERROR if self.kv_cache_transceiver and request in self.active_requests: @@ -6696,7 +6747,7 @@ def _check_gen_cache_transfer_errors_consensus(self) -> None: """Flush generation transfer errors through a TP-uniform path.""" error_requests = [ req for req in self._get_disagg_reqs_in_error_state() - if req.is_generation_only_request() + if req.is_generation_only_request ] local_needs_flush = bool(error_requests) @@ -7419,7 +7470,17 @@ def _send_kv_async(self, scheduled_requests: List[LlmRequest]): def kv_connector_request_finished(req: LlmRequest): try: - cache_block_ids = self.kv_cache_manager.get_cache_indices(req) + # KVCacheManagerV2 has no primary-pool block list; its page + # indices are scoped to a layer group. Without this the lookup + # below raises, the warning path swallows it, and + # `request_finished` is never called at all on V2 -- so a + # connector is never told to save anything. + if isinstance(self.kv_cache_manager, KVCacheManagerV2): + cache_block_ids = self.kv_cache_manager.get_connector_page_indices( + req) + else: + cache_block_ids = self.kv_cache_manager.get_cache_indices( + req) except Exception as e: logger.warning( f"Unable to get cache blocks for request {req.py_request_id}. Skipping asynchronous saving: {e}" @@ -8317,7 +8378,7 @@ def _handle_responses(self, emit_first_iter: bool = True): new_active_requests.append(request) continue - if request.is_generation_only_request() and not request.is_finished: + if request.is_generation_only_request and not request.is_finished: # If request is in transmission, so we don't need to emit a response # Also, for the first iteration with overlap, we should skip since first # token has already been emitted previously diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7f04fb9742f3..3f9c99f84af4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -855,10 +855,21 @@ def drafting_loop_wrapper(model): "KV connector is only supported with guaranteed no evict scheduler policy." ) + # VSWA allocates one pool per window size, which the V1 connector + # registration cannot describe: it hands the worker a single primary + # pool tensor. KVCacheManagerV2 has no such limitation -- its layout + # describes one region set per layer group -- so only reject here when + # V2 is definitively off. `use_kv_cache_manager_v2` is tri-state + # (True / False / "auto"), and under "auto" the manager is not chosen + # yet, so defer: PyExecutor re-checks against the manager it actually + # built (py_executor.py, `_maybe_init_kv_connector_manager`) and rejects + # there if the selection landed on V1. max_attention_window = kv_cache_config.max_attention_window - if uses_vswa_kv_cache_layout(max_attention_window): + if (uses_vswa_kv_cache_layout(max_attention_window) + and kv_cache_config.use_kv_cache_manager_v2 is False): raise NotImplementedError( - "KV connector is not supported with VSWA (Variable Sliding Window Attention)." + "KV connector is not supported with VSWA (Variable Sliding Window Attention) " + "on the V1 KV cache manager. Set kv_cache_config.use_kv_cache_manager_v2=True." ) if mapping.enable_attention_dp: diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index cf12b4b91a7e..e4184495a017 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -305,7 +305,7 @@ def _schedule_loop(self, active_requests, inflight_request_ids): if self._prioritize_first_token_gen: requests_list.sort( key=lambda req: ( - 0 if (req.is_generation_only_request() and req.py_decoding_iter == 0) else 1 + 0 if (req.is_generation_only_request and req.py_decoding_iter == 0) else 1 ) ) @@ -550,6 +550,14 @@ def _try_schedule_context( Returns ``(action, tokens, chunking_flag)``. *tokens* and *chunking_flag* are meaningful only when *action* is ``SCHEDULED``. """ + # No `should_add_sequence` gate here, unlike V1. That predicate stays + # false from the moment an asynchronous load completes until + # `request_finished`, which only runs at the end of generation, so V2 + # would SKIP the request forever and never run the prefill the load was + # for. What keeps a loading request out of the batch instead is its + # DISAGG_GENERATION_TRANS_IN_PROGRESS state, and what stops the + # connector being asked or told twice is the per-request query state + # (see KVCacheManagerV2._connector_prefix_position). if self.chunking_enabled: return self._try_schedule_context_chunked(req, budget) return self._try_schedule_context_full(req, budget) diff --git a/tensorrt_llm/_torch/speculative/suffix_automaton.py b/tensorrt_llm/_torch/speculative/suffix_automaton.py index 4eba0781d46f..b2d6ace23a9c 100644 --- a/tensorrt_llm/_torch/speculative/suffix_automaton.py +++ b/tensorrt_llm/_torch/speculative/suffix_automaton.py @@ -612,7 +612,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # one-shot -- and (b) pin an SA slot for the whole KV-transfer # duration. Defer to the generation-request loop below, which # runs once the request is scheduled with its full history. - if req.is_generation_only_request(): + if req.is_generation_only_request: continue if req.is_first_context_chunk: if req.request_id not in self._initialized_requests: @@ -631,7 +631,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # one-shot init; free_resources clears it on completion. for req in scheduled_batch.generation_requests: if ( - req.is_generation_only_request() + req.is_generation_only_request and not req.is_dummy and req.request_id not in self._initialized_requests ): diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 6a0fbd26c7ac..d55dbb9ff421 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -13,8 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import math import os +import shutil import sys import tempfile import time @@ -23,15 +25,83 @@ import pytest from tensorrt_llm import LLM, DisaggregatedParams, SamplingParams +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ + KvCacheConnectorWorker +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, KvCacheConfig, KvCacheConnectorConfig) from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX from ..conftest import llm_models_root +# Name of the TensorRT-LLM logger. It sets `propagate = False` +# (tensorrt_llm/logger.py:186-187), so pytest's `caplog` only sees its records +# once `caplog.handler` is attached to it directly. +TRTLLM_LOGGER_NAME = "TRT-LLM" + +# Emitted by `_fallback_if_unsupported_kv_cache_manager_v2` +# (tensorrt_llm/_torch/pyexecutor/_util.py:629-631) when a connector run is +# downgraded from V2 to V1. +FALLBACK_WARNING_FRAGMENT = "Falling back to KVCacheManager" + +# V1 `KVCacheManager` methods reached on the KV connector path proper. Each is +# defined in tensorrt_llm/_torch/pyexecutor/resource_manager.py and wraps a +# nanobind method on the C++ manager. `KVCacheManagerV2` implements none of +# them and is not meant to: each assumes a single flat block-id space over one +# primary pool, which cannot describe memory whose page indices are scoped to a +# layer group. V2's connector contract is `register_kv_cache_layout` plus +# `get_page_indices_by_layer_group` / `get_connector_page_indices` instead. +CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS = ( + # Connector bring-up: hands the worker the single primary pool tensor. + # `PyExecutor._maybe_init_kv_connector_manager`. + "get_unique_primary_pool", + # `KvCacheConnectorSchedulerOutputRequest.update_and_build_data` and + # `PyExecutor.kv_connector_request_finished`. + "get_cache_indices", + # `update_and_build_data`, for `RequestData.block_hashes`. + "commit_and_get_block_hashes", + # `update_and_build_data`, for `RequestData.priorities`. + "get_priority_by_block_id", +) + +# Reached only when the connector coexists with disaggregated serving +# (`test_connector_disagg_prefill`), via AsyncTransferManager and the V1 cache +# reuse adapter - not from `KvCacheConnectorManager` itself. None of them is a +# gap for V2, because V2 does not traverse those paths: +# `enable_partial_reuse_for_disagg` excludes V2, so AsyncTransferManager never +# reaches the pin/unpin pair, and the Python transceiver (the only one V2 can +# be driven by) resolves block ids through `_CacheReuseAdapterV2`. +DISAGG_PATH_KV_CACHE_MANAGER_METHODS = ( + "store_blocks_for_reuse", + "unpin_blocks_by_id", + "get_memory_pool_block_indices", + "pin_blocks", # no Python caller today +) + + +@pytest.fixture(scope="function") +def use_kv_cache_manager_v2(request): + """Run each connector test under both KV cache managers. + + Parametrized by an explicit `@pytest.mark.parametrize(..., indirect=True)` + on each test, applied as the innermost decorator so the manager lands first + in the generated test id. It is spelled out per test rather than set as a + fixture `params=` because the test-list validator + (scripts/check_test_list.py) resolves ids from parametrize decorators via + AST and cannot see fixture-level parametrization. + + Selecting V2 must actually reach V2: `_fallback_if_unsupported_kv_cache_manager_v2` + silently substitutes the V1 manager for combinations it cannot serve, and a + connector test that ran on V1 while claiming to test V2 would pass while + exercising nothing. `test_connector_runs_on_kv_cache_manager_v2` guards that. + """ + return request.param + @pytest.fixture(scope="function") -def model_with_connector(): +def model_with_connector(use_kv_cache_manager_v2): with patch("tensorrt_llm._torch.pyexecutor.py_executor_creator.importlib" ) as importlib_mock: mock_scheduler = MagicMock() @@ -56,7 +126,16 @@ def model_fn(*args, **kwargs): "kv_cache_config": KvCacheConfig(free_gpu_memory_fraction=0.1) } - return LLM(*args, **{**default_kwargs, **kwargs}) + merged_kwargs = {**default_kwargs, **kwargs} + + # Tests that supply their own `KvCacheConfig` must still honour the + # manager under test, otherwise the V2 parametrization silently + # degrades into a second V1 run. + kv_cache_config = merged_kwargs.get("kv_cache_config") + if kv_cache_config is not None: + kv_cache_config.use_kv_cache_manager_v2 = use_kv_cache_manager_v2 + + return LLM(*args, **merged_kwargs) yield model_fn, mock_scheduler, mock_worker @@ -68,27 +147,225 @@ def enforce_single_worker(monkeypatch): yield -def generate_and_sleep(model, *args, **kwargs): - # Some KV connector API calls are made after a full response is returned. We want to be able to track these calls. - # However, we don't have any indication of when all the calls are complete. - # To compensate for this, we sleep between the generate call and the return of the outputs. - # TODO(jthomson04): Surely there's a better way to do this? +# Some KV connector API calls are made after a full response is returned +# (`request_finished`, the trailing `get_finished` polls and asynchronous +# saves), and there is no public signal for when they are complete. Instead of +# sleeping a fixed amount, wait until the connector mocks stop recording new +# calls. That returns as soon as the connector goes quiet, and - unlike a fixed +# sleep - stretches automatically when a slower path lengthens the tail. +CONNECTOR_QUIESCE_TIMEOUT_S = 60.0 +CONNECTOR_QUIET_PERIOD_S = 0.5 +CONNECTOR_POLL_INTERVAL_S = 0.01 + +# Fraction of generated tokens a connector-warmed run must reproduce exactly. +# See test_connector_e2e_persistent_cache for why this is not 1.0. +E2E_MIN_TOKEN_AGREEMENT = 0.75 + + +def assert_kv_caches_registered(worker, use_kv_cache_manager_v2): + """The two managers hand the worker its pools through different entry points. + + V1 passes a single pool tensor to `register_kv_caches`; V2 has no such + tensor and passes a `KvCacheLayout` to `register_kv_cache_layout` instead. + Asserting the V1 method unconditionally would silently pass on V2 only if + the connector were never registered at all. + """ + if use_kv_cache_manager_v2: + assert worker.register_kv_cache_layout.call_count == 1 + assert worker.register_kv_caches.call_count == 0 + else: + assert worker.register_kv_caches.call_count == 1 + assert worker.register_kv_cache_layout.call_count == 0 + + +def wait_for_connector_quiescence(scheduler, + worker, + timeout=CONNECTOR_QUIESCE_TIMEOUT_S, + quiet_period=CONNECTOR_QUIET_PERIOD_S): + """Block until no new connector callback lands for `quiet_period` seconds. + + `MagicMock.mock_calls` records every call made on the mock and its children, + so its length is a monotonic progress counter for connector activity. + """ + deadline = time.monotonic() + timeout + + def total_calls(): + return len(scheduler.mock_calls) + len(worker.mock_calls) + + last_seen = total_calls() + quiet_since = time.monotonic() + + while time.monotonic() < deadline: + time.sleep(CONNECTOR_POLL_INTERVAL_S) + + current = total_calls() + if current != last_seen: + last_seen = current + quiet_since = time.monotonic() + elif time.monotonic() - quiet_since >= quiet_period: + return + + raise AssertionError( + f"KV connector callbacks did not go quiet within {timeout}s " + f"({last_seen} calls recorded). The connector is still active or a " + "callback is blocked.") + + +def generate_and_wait(model, scheduler, worker, *args, **kwargs): + """`model.generate`, then block until the connector callbacks settle.""" outputs = model.generate(*args, **kwargs) - time.sleep(1) + wait_for_connector_quiescence(scheduler, worker) return outputs +def test_v2_connector_contract_does_not_reuse_the_v1_methods(): + """The V2 connector path implements none of the V1 accessors, by design. + + Something depends on that, and it does not ask: `update_and_build_data` + reports `block_hashes` and `priorities` empty on V2 by branching on + `isinstance(manager, KVCacheManagerV2)`, not on `hasattr`. Those + short-circuits are only correct while V2 genuinely has no such accessor - + the day one is added (retention priorities are a known gap; see + `test_connector_priorities`) the branch keeps reporting nothing while the + data exists, and this is what says so. + + A static check rather than an end-to-end run: under V2 the connector would + die at the *first* method it reached, so no run can report more than one at + a time. + + Needs no GPU. + """ + stale = [ + name for name in CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS + + DISAGG_PATH_KV_CACHE_MANAGER_METHODS + if not hasattr(KVCacheManager, name) + ] + assert stale == [], ( + f"{stale} are not defined on the V1 KVCacheManager either, so this " + "test is measuring a stale method list rather than a real difference.") + + implemented = [ + name for name in CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS + if hasattr(KVCacheManagerV2, name) + ] + assert implemented == [], ( + f"KVCacheManagerV2 now implements {implemented}. A V1-shaped accessor " + "on V2 is not automatically the right answer - a flat block-id list " + "cannot describe more than one layer group - but if it is, revisit the " + "`is_v2` short-circuits in " + "`KvCacheConnectorSchedulerOutputRequest.update_and_build_data`, which " + "report nothing on the strength of these methods being absent.") + + # The other half of the contract: what V2 offers instead. + for name in ("get_page_indices_by_layer_group", + "get_connector_page_indices"): + assert hasattr(KVCacheManagerV2, name), ( + f"KVCacheManagerV2.{name} is the V2 replacement for the V1 " + "block-id accessors and every connector path on V2 goes through " + "it.") + assert hasattr(KvCacheConnectorWorker, "register_kv_cache_layout"), ( + "The worker ABC must keep a default `register_kv_cache_layout`, or " + "every existing connector becomes abstract and fails to instantiate.") + + +@pytest.mark.threadleak(enabled=False) +def test_connector_runs_on_kv_cache_manager_v2(enforce_single_worker, + monkeypatch, caplog): + """Anti-vacuity guard for the `kv_cache_manager_v2` parametrization. + + Without this, every V2-parametrized test below could pass green while + `_fallback_if_unsupported_kv_cache_manager_v2` silently swapped in the V1 + manager. It asserts positively that V2 is constructed, and that the + downgrade warning is absent. + + Any construction failure is recorded rather than asserted on: what must + hold is that the run reached V2 rather than being papered over by a + fallback, which stays true regardless of how far bring-up gets. + """ + constructed = [] + + def record_construction(cls): + original_init = cls.__init__ + + def recording_init(self, *args, **kwargs): + constructed.append(cls.__name__) + return original_init(self, *args, **kwargs) + + monkeypatch.setattr(cls, "__init__", recording_init) + + record_construction(KVCacheManagerV2) + record_construction(KVCacheManager) + + # The TensorRT-LLM logger sets `propagate = False`, so caplog only sees its + # records once its handler is attached to that logger directly. + trtllm_logger = logging.getLogger(TRTLLM_LOGGER_NAME) + trtllm_logger.addHandler(caplog.handler) + + construction_error = None + llm = None + try: + with patch( + "tensorrt_llm._torch.pyexecutor.py_executor_creator.importlib" + ) as importlib_mock: + connector_module = importlib_mock.import_module.return_value + connector_module.KvConnectorScheduler.return_value = MagicMock() + connector_module.KvConnectorWorker.return_value = MagicMock() + + try: + llm = LLM( + model=f"{llm_models_root()}/Qwen2-0.5B", + backend="pytorch", + kv_connector_config=KvCacheConnectorConfig( + connector_module="", + connector_scheduler_class="KvConnectorScheduler", + connector_worker_class="KvConnectorWorker", + ), + cuda_graph_config=None, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1, + use_kv_cache_manager_v2=True), + ) + # The V2 connector path is knowingly incomplete, so any construction + # failure is an acceptable outcome. What must hold is that the run + # reached V2 instead of being papered over by a V1 fallback. + except Exception as exc: # noqa: BLE001 + construction_error = exc + finally: + trtllm_logger.removeHandler(caplog.handler) + if llm is not None: + llm.shutdown() + + # Report the current frontier so the failure mode is visible in CI output + # instead of being silently swallowed by the except above. + print(f"\n[connector+V2 frontier] construction_error=" + f"{type(construction_error).__name__ if construction_error else None}" + f": {construction_error}") + + assert "KVCacheManagerV2" in constructed, ( + "KVCacheManagerV2 was never constructed, so the connector run silently " + f"fell back to V1 (managers constructed: {constructed}; construction " + f"error: {construction_error!r}). Every kv_cache_manager_v2-" + "parametrized connector test in this file is vacuous until this passes." + ) + + assert FALLBACK_WARNING_FRAGMENT not in caplog.text, ( + f"{FALLBACK_WARNING_FRAGMENT!r} was logged, so the connector was " + "downgraded to the V1 KV cache manager.") + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_simple(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -96,7 +373,8 @@ def test_connector_simple(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, ["Hello, world"], sampling_params) + generate_and_wait(model, scheduler, worker, ["Hello, world"], + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 @@ -154,22 +432,26 @@ def test_connector_simple(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_async_onboard(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 16, True worker.get_finished.side_effect = lambda finished_gen, load_async: ( finished_gen, load_async) - generate_and_sleep(model, [ + generate_and_wait(model, scheduler, worker, [ "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua." ], SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True)) @@ -183,15 +465,18 @@ def test_connector_async_onboard(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_async_save(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -202,7 +487,8 @@ def test_connector_async_save(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, ["Hello, world"], sampling_params) + generate_and_wait(model, scheduler, worker, ["Hello, world"], + sampling_params) assert scheduler.request_finished.call_count == 1 @@ -224,8 +510,12 @@ def test_connector_async_save(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_scheduler_output(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): NUM_INPUT_TOKENS = 48 NUM_TOKENS = 32 BLOCK_SIZE = 32 @@ -234,7 +524,7 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -242,7 +532,8 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) - generate_and_sleep(model, [0] * NUM_INPUT_TOKENS, sampling_params) + generate_and_wait(model, scheduler, worker, [0] * NUM_INPUT_TOKENS, + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 assert len( @@ -292,7 +583,8 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, assert len(scheduler.request_finished.call_args.args[1]) == math.ceil( (NUM_INPUT_TOKENS + NUM_TOKENS) / BLOCK_SIZE) - generate_and_sleep(model, [1] * NUM_INPUT_TOKENS, sampling_params) + generate_and_wait(model, scheduler, worker, [1] * NUM_INPUT_TOKENS, + sampling_params) # The initial computed position should be 0, since we haven't yet onboarded any blocks. assert scheduler.build_connector_meta.call_args_list[0].args[ @@ -301,9 +593,13 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_scheduler_output_chunked_context(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): model_fn, scheduler, worker = model_with_connector CHUNK_SIZE = 128 @@ -313,7 +609,7 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, enable_chunked_prefill=True, max_num_tokens=CHUNK_SIZE) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -321,13 +617,22 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, sampling_params = SamplingParams(max_tokens=BLOCK_SIZE, ignore_eos=True) - generate_and_sleep(model, [0] * (CHUNK_SIZE * 2), sampling_params) + generate_and_wait(model, scheduler, worker, [0] * (CHUNK_SIZE * 2), + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 - assert len( - scheduler.update_state_after_alloc.call_args.args[1]) == math.ceil( - CHUNK_SIZE * 2 / BLOCK_SIZE) + # V1 allocates for the whole prompt when the sequence is added, so every + # block exists on the first chunk. V2 allocates per chunk, which is the + # lower-peak-memory behaviour and the one to keep; the remaining blocks + # arrive as append-deltas in `new_block_ids` on the next chunk, so no + # information is lost. The expectation is split rather than V2 changed. + total_blocks = math.ceil(CHUNK_SIZE * 2 / BLOCK_SIZE) + first_chunk_blocks = (math.ceil(CHUNK_SIZE / BLOCK_SIZE) + if use_kv_cache_manager_v2 else total_blocks) + + assert len(scheduler.update_state_after_alloc.call_args.args[1] + ) == first_chunk_blocks for i, call in enumerate(scheduler.build_connector_meta.call_args_list): sched_output = call.args[0] @@ -342,18 +647,18 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, req = sched_output.cached_requests[0] if i == 0: - # The first prefill chunk. - # All of the prefill tokens and all the blocks should be provided upfront. + # The first prefill chunk. All of the prefill tokens are provided + # upfront on both managers; the blocks are whatever has been + # allocated so far. assert req.computed_position == 0 assert len(req.new_tokens) == CHUNK_SIZE * 2 - assert len(req.new_block_ids) == math.ceil(CHUNK_SIZE * 2 / - BLOCK_SIZE) + assert len(req.new_block_ids) == first_chunk_blocks assert req.num_scheduled_tokens == CHUNK_SIZE elif i == 1: # The second prefill chunk. assert req.computed_position == CHUNK_SIZE assert len(req.new_tokens) == 0 - assert len(req.new_block_ids) == 0 + assert len(req.new_block_ids) == total_blocks - first_chunk_blocks assert req.num_scheduled_tokens == CHUNK_SIZE elif i == 2 and use_overlap_scheduler: assert len(req.new_tokens) == 0 @@ -365,19 +670,289 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, (CHUNK_SIZE * 2 + BLOCK_SIZE) / BLOCK_SIZE) +# The mock scheduler answers every query with the same offer; these helpers +# record what it was asked and when, so a test can assert the connector was +# consulted exactly once per request rather than assuming it. +def record_connector_queries(scheduler, num_matched, load_async=False): + """Answer every query with `num_matched`, recording when each one arrived. + + The third element of each record is `build_connector_meta.call_count` at + query time -- the number of iterations whose connector hooks had already + run. Phase 1 runs during scheduling, before those hooks, so a request asked + in the first scheduling pass records 0. That is how the tests below *prove* + a request was asked before the iteration it eventually ran in, rather than + assuming the scheduler deferred anything. + """ + queries = [] + + def side_effect(request, num_computed_tokens): + queries.append((request.request_id, num_computed_tokens, + scheduler.build_connector_meta.call_count)) + return num_matched, load_async + + scheduler.get_num_new_matched_tokens.side_effect = side_effect + return queries + + +# Sliding-window coverage. `KvCacheConfig.max_attention_window` is repeated +# cyclically across layers (llm_args.py:3761-3765), so a one-element list gives +# every layer the same window -- one V2 layer group with a live window -- while +# a two-element list alternates and produces two. A window equal to +# `max_seq_len` is normalised to "no window" (kv_cache_manager_v2.py:856-858), +# which is how the full-attention half of the VSWA pair is spelled. +SWA_WINDOW = 64 +SWA_MAX_SEQ_LEN = 512 +SWA_NUM_INPUT_TOKENS = 256 + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], + ids=["kv_cache_manager_v2"], + indirect=True) +def test_connector_uniform_sliding_window(enforce_single_worker, + model_with_connector, + use_kv_cache_manager_v2): + """Connector against a KV cache in which every layer slides. + + There was no sliding-window connector coverage before this test. Uniform + SWA is a single layer group, which is the configuration where the + connector's flat `new_block_ids` list is still well defined -- so this pins + the block-reporting contract, and the VSWA test below pins what happens once + that assumption breaks. + + V2 only. The V1 guard rejects *variable* windows + (py_executor_creator.py:845-850), so uniform SWA reaches V1's connector path + and then dies inside it: `commit_and_get_block_hashes` (kv_cache_connector.py:386) + raises "commitAndGetBlockHashesForRequest does not support sliding-window + attention with detached front blocks" (kvCacheManager.cpp:4645) as soon as + the window drops a front block. That is a pre-existing V1 limitation, not + something this work introduces, so it is recorded rather than asserted here. + """ + model_fn, scheduler, worker = model_with_connector + + model = model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW])) + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + scheduler.get_num_new_matched_tokens.return_value = 0, False + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + sched_output = scheduler.build_connector_meta.call_args_list[0].args[0] + assert len(sched_output.new_requests) == 1 + req = sched_output.new_requests[0] + + assert req.computed_position == 0 + assert req.num_scheduled_tokens == SWA_NUM_INPUT_TOKENS + # A single layer group keeps the flat list meaningful on both managers. + assert req.new_block_ids + + if use_kv_cache_manager_v2: + # Anti-vacuity: prove the window really did collapse to one layer + # group, otherwise the assertion above would hold for the wrong reason. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 1 + assert layout.groups[0].window_size == SWA_WINDOW + assert list(req.new_block_ids_by_layer_group) == [0] + assert req.new_block_ids_by_layer_group[0] == req.new_block_ids + else: + assert req.new_block_ids_by_layer_group == {} + + +# Half the prompt, and well past the window, so the pages the offer does *not* +# need are a large enough fraction to assert on. +SWA_OFFER_TOKENS = 128 + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], + ids=["kv_cache_manager_v2"], + indirect=True) +def test_connector_sliding_window_prefix_is_backed_by_history( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """Phase 2 raises `history_length`, not just capacity -- observably. + + Under full attention the two are indistinguishable from outside: capacity is + `prompt_len` whether or not a prefix was served, so the page count says + nothing. Under a sliding window it does. `history_length` is the sole input + to the stale-range computation, so raising it to the offer end tells V2 that + the blocks the window has already left need no pages -- and raising capacity + alone would allocate one for every block of the served prefix. + + So this is the test that would fail if `_reserve_connector_prefix` called + `resize(capacity, None)`; `test_connector_uniform_sliding_window` above + cannot, because it offers nothing. + """ + model_fn, scheduler, worker = model_with_connector + + model = model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW])) + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + record_connector_queries(scheduler, SWA_OFFER_TOKENS) + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + # Anti-vacuity, as in the test above: one layer group with a live window, + # otherwise the page count below is being read off full attention. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 1 + assert layout.groups[0].window_size == SWA_WINDOW + + req = scheduler.build_connector_meta.call_args_list[0].args[0].new_requests[ + 0] + + # The offer was materialized: the position advanced over it, and the + # runtime rolled the reported position back to the local match. + assert req.computed_position == 0 + assert req.num_scheduled_tokens == SWA_NUM_INPUT_TOKENS - SWA_OFFER_TOKENS + + all_blocks = math.ceil(SWA_NUM_INPUT_TOKENS / 32) + page_indices = scheduler.update_state_after_alloc.call_args.args[1] + assert page_indices == req.new_block_ids + # One entry per block ordinal either way -- a block with no page reads back + # as BAD_PAGE_INDEX in place rather than shortening the list, so that block + # ordinals stay aligned to token ranges (kv_cache_manager_v2.py:2628-2636). + # The page count is therefore not the signal; *which* entries are bad is. + assert len(page_indices) == all_blocks + + # With history at the offer end, every block that lies entirely below the + # window got no page. Without the bump history would still be at the local + # match, nothing would be stale, and all 8 blocks would be backed. + stale_blocks = (SWA_OFFER_TOKENS - SWA_WINDOW) // 32 + assert stale_blocks > 0, "test sizes no longer put any block out of window" + assert page_indices[:stale_blocks] == [BAD_PAGE_INDEX] * stale_blocks, ( + f"expected the first {stale_blocks} blocks of a {SWA_OFFER_TOKENS}-token " + f"served prefix to fall outside a {SWA_WINDOW}-token window and hold no " + f"page, got {page_indices}. History was not raised to the offer end.") + + live = page_indices[stale_blocks:] + assert all(index != BAD_PAGE_INDEX for index in live), ( + f"a block inside the window has no page: {page_indices}") + assert len(set(live)) == len(live), ( + f"page slots reported to the connector are not distinct: {live}") + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_vswa_reports_page_indices_per_layer_group( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """VSWA is where the connector's flat block list stops working. + + V1 cannot describe VSWA to a connector at all: it registers a single primary + pool, but VSWA allocates one pool per window size. V2's layout describes one + region set per layer group, so the combination runs there -- but a page index + is scoped to a layer group, and there is no correct way to flatten indices + from several groups into one list. `KVCacheManagerV2` therefore reports an + empty `new_block_ids` and leaves `new_block_ids_by_layer_group` as the only + correct source (kv_cache_manager_v2.py:2519-2526, kv_cache_connector.py:365-377). + + Both halves are asserted here because the rejection is deliberately + conditional on the manager (py_executor_creator.py). Pinning only the V2 half + would let the V1 guard silently disappear. + """ + model_fn, scheduler, worker = model_with_connector + + def build(): + return model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW, SWA_MAX_SEQ_LEN])) + + if not use_kv_cache_manager_v2: + with pytest.raises(NotImplementedError, match="VSWA"): + build() + return + + model = build() + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + # Alternating windows must actually produce two groups, one sliding and one + # full-attention, or the rest of this test is vacuous. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 2 + assert {group.window_size for group in layout.groups} == {SWA_WINDOW, None} + + scheduler.get_num_new_matched_tokens.return_value = 0, False + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + # The multi-group degradation reaches `update_state_after_alloc` too. A + # connector reading only `new_block_ids` sees nothing and saves nothing, + # which is silent unless pinned here. + assert scheduler.update_state_after_alloc.call_args.args[1] == [] + + sched_output = scheduler.build_connector_meta.call_args_list[0].args[0] + assert len(sched_output.new_requests) == 1 + req = sched_output.new_requests[0] + + assert req.new_block_ids == [] + assert sorted(req.new_block_ids_by_layer_group) == [0, 1] + # Ordinals stay positionally aligned across groups: a block with no page in + # the sliding group reads back as BAD_PAGE_INDEX in place rather than + # shortening the list (kv_cache_manager_v2.py:2481-2490). + lengths = { + len(indices) + for indices in req.new_block_ids_by_layer_group.values() + } + assert len(lengths) == 1 + + +def _disagg_transceiver_config(use_kv_cache_manager_v2): + """The transceiver each KV cache manager can actually be driven by. + + `CacheTransceiverCpp` is bound to the V1 `BaseKVCacheManager`, while + `KVCacheManagerV2.impl` is the Python V2 core's manager, so V2 can only use + the Python transceiver -- which in turn only supports NIXL + (kv_cache_transceiver.py, `create_kv_cache_transceiver`). This is spelled + out per manager rather than left at the default because + `transceiver_runtime` defaults to "auto", and "auto" is resolved from the + *model's* preference (llm_utils._resolve_transceiver_runtime_auto), which + knows nothing about which cache manager will be built. Qwen2 declares no + preference, so the default resolves to the C++ transceiver -- which V2 + cannot use, and which is now rejected with an actionable error rather than + a nanobind signature mismatch. + """ + if use_kv_cache_manager_v2: + return CacheTransceiverConfig(backend="NIXL", + transceiver_runtime="PYTHON") + return CacheTransceiverConfig(backend="DEFAULT") + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("save_async", [False, True]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, - save_async): + save_async, use_kv_cache_manager_v2): model_fn, scheduler, worker = model_with_connector - prefill_worker = model_fn( - disable_overlap_scheduler=True, - cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT")) + transceiver_config = _disagg_transceiver_config(use_kv_cache_manager_v2) + + prefill_worker = model_fn(disable_overlap_scheduler=True, + cache_transceiver_config=transceiver_config) - decode_worker = model_fn( - cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), - kv_connector_config=None) + decode_worker = model_fn(cache_transceiver_config=transceiver_config, + kv_connector_config=None) sampling_params = SamplingParams(ignore_eos=True, max_tokens=16) @@ -394,16 +969,20 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, scheduler.request_finished.return_value = False worker.get_finished.return_value = [], [] - result = generate_and_sleep(prefill_worker, [0] * 48, - sampling_params=sampling_params, - disaggregated_params=disaggregated_params) + result = generate_and_wait(prefill_worker, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=disaggregated_params) gen_disagg_params = result.disaggregated_params gen_disagg_params.request_type = "generation_only" - generate_and_sleep(decode_worker, [0] * 48, - sampling_params=sampling_params, - disaggregated_params=gen_disagg_params) + generate_and_wait(decode_worker, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=gen_disagg_params) assert scheduler.build_connector_meta.call_count == 1 @@ -422,6 +1001,9 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_multi_request(enforce_single_worker, model_with_connector): model_fn, scheduler, worker = model_with_connector @@ -446,12 +1028,38 @@ def test_connector_multi_request(enforce_single_worker, model_with_connector): @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [ + pytest.param(False, id="kv_cache_manager_v1"), + pytest.param( + True, + id="kv_cache_manager_v2", + marks=pytest.mark.xfail( + strict=True, + reason= + "KvCacheRetentionConfig does not reach KVCacheManagerV2 at all " + "(per-page priority comes from custom_priority_callback, which " + "V2 never overrides), so the connector reports priorities=None."), + ), +], + indirect=True) def test_connector_priorities(enforce_single_worker, model_with_connector): """Test that retention priorities flow through the connector correctly. This test verifies that when KvCacheRetentionConfig is provided, the RequestData.priorities field is populated with the correct per-block priorities based on the token ranges. + + KNOWN GAP -- `xfail(strict=True)` on `kv_cache_manager_v2`. + `KvCacheRetentionConfig` does not reach KVCacheManagerV2 at all: V2's + per-page priority comes from `custom_priority_callback` + (kv_cache_manager_v2/_core/_kv_cache_manager.py), which KVCacheManagerV2 + never overrides, so every page carries the default priority and the + connector reports `priorities=None`. A user who sets a retention config on + V2 silently gets none of it -- not only through the connector. The + assertions below stay the correct expectation for both managers rather than + being relaxed per manager, so wiring retention into V2 turns this green + instead of needing the test rewritten; `strict=True` is what makes it fail + loudly on that day rather than passing silently. """ BLOCK_SIZE = 32 NUM_INPUT_TOKENS = 64 # 2 blocks @@ -487,9 +1095,11 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, [0] * NUM_INPUT_TOKENS, - sampling_params=sampling_params, - kv_cache_retention_config=retention_config) + generate_and_wait(model, + scheduler, + worker, [0] * NUM_INPUT_TOKENS, + sampling_params=sampling_params, + kv_cache_retention_config=retention_config) # Verify that build_connector_meta was called assert scheduler.build_connector_meta.call_count >= 1 @@ -517,6 +1127,9 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_priorities_default(enforce_single_worker, model_with_connector): """Test that priorities are None when no retention config is provided.""" @@ -530,7 +1143,10 @@ def test_connector_priorities_default(enforce_single_worker, sampling_params = SamplingParams(max_tokens=4, ignore_eos=True) # Generate without retention config - generate_and_sleep(model, [0] * 48, sampling_params=sampling_params) + generate_and_wait(model, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params) first_call = scheduler.build_connector_meta.call_args_list[0] sched_output = first_call.args[0] @@ -544,45 +1160,72 @@ def test_connector_priorities_default(enforce_single_worker, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize( - "llm_kwargs,match", + "llm_kwargs,match_v1,match_v2", [ + # The two managers refuse offloading for the same reason -- a page that + # leaves GPU has its slot reassigned, invalidating what the connector + # registered -- but say so differently, and V2 names the resolved tier + # rather than the config field so it also catches the disk tier. Match + # the manager-specific wording: both messages contain the bare word + # "host", so matching that would still pass if a silent fallback to V1 + # ever crept back in, which is the one thing this parametrization + # exists to rule out. pytest.param( dict(kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1, host_cache_size=1024**3)), - "host", + "host offloading", + "cache tiers below GPU", id="host_offloading", ), pytest.param( dict(max_beam_width=2), "beam", + "beam", id="beam_search", ), pytest.param( dict(enable_attention_dp=True), "attention data parallelism", + "attention data parallelism", id="attention_dp", ), ], ) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_rejects_unsupported_config(enforce_single_worker, - model_with_connector, llm_kwargs, - match): + model_with_connector, + use_kv_cache_manager_v2, + llm_kwargs, match_v1, match_v2): # Configurations the connector cannot handle today must fail loudly at # construction time rather than silently miscompute. This pins the set of # constructor-time exclusions in `_maybe_init_kv_connector_manager`. model_fn, _, _ = model_with_connector + match = match_v2 if use_kv_cache_manager_v2 else match_v1 with pytest.raises(NotImplementedError, match=match): model_fn(**llm_kwargs) @pytest.mark.threadleak(enabled=False) -def test_connector_e2e_persistent_cache(enforce_single_worker): - """Test e2e KV cache connector using PersistentKvCacheConnector from examples. - - Runs generation twice with separate LLM instances sharing a disk-based - connector cache, verifying that outputs are identical (proving cache - save/load works end-to-end). +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_e2e_persistent_cache(enforce_single_worker, + use_kv_cache_manager_v2, monkeypatch): + """End-to-end KV connector test using PersistentKvCacheConnector from examples. + + Runs the same prompt through two separate LLM instances sharing a + disk-backed connector cache and asserts that: + + 1. the first (cold) run matches nothing and writes cache files, + 2. the second (warm) run actually reads blocks back from disk, and + 3. both runs produce identical text and token ids. + + (3) on its own proves nothing - two deterministic runs of the same prompt + agree whether or not the cache is ever consulted - so (2) is what makes + this a real correctness test rather than a tautology. """ examples_dir = os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "examples", "llm-api") @@ -590,9 +1233,29 @@ def test_connector_e2e_persistent_cache(enforce_single_worker): sys.path.insert(0, examples_dir) cache_dir = tempfile.mkdtemp() - os.environ["CONNECTOR_CACHE_FOLDER"] = cache_dir + monkeypatch.setenv("CONNECTOR_CACHE_FOLDER", cache_dir) try: + import llm_kv_cache_connector + + # Record how many tokens the connector served from disk on each run. + # The leader logs this, but the TensorRT-LLM logger does not propagate + # to the root logger, so read it from the return value instead. + matched_tokens = [] + leader_cls = llm_kv_cache_connector.PersistentKvCacheConnectorLeader + original_get_num_new_matched_tokens = ( + leader_cls.get_num_new_matched_tokens) + + def recording_get_num_new_matched_tokens(self, request, + num_computed_tokens): + result = original_get_num_new_matched_tokens( + self, request, num_computed_tokens) + matched_tokens.append(result[0]) + return result + + monkeypatch.setattr(leader_cls, "get_num_new_matched_tokens", + recording_get_num_new_matched_tokens) + kv_connector_config = KvCacheConnectorConfig( connector_module="llm_kv_cache_connector", connector_scheduler_class="PersistentKvCacheConnectorLeader", @@ -605,7 +1268,9 @@ def test_connector_e2e_persistent_cache(enforce_single_worker): kv_connector_config=kv_connector_config, cuda_graph_config=None, disable_overlap_scheduler=True, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1), + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + use_kv_cache_manager_v2=use_kv_cache_manager_v2), ) prompt = ( @@ -620,21 +1285,70 @@ def test_connector_e2e_persistent_cache(enforce_single_worker): sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) llm1 = LLM(**llm_kwargs) - output1 = llm1.generate([prompt], sampling_params) - output1[0].outputs[0].text - del llm1 + try: + output1 = llm1.generate([prompt], sampling_params) + cold_text = output1[0].outputs[0].text + cold_token_ids = list(output1[0].outputs[0].token_ids) + finally: + llm1.shutdown() + + assert matched_tokens and all(count == 0 for count in matched_tokens), ( + "The first run should be a cold miss, but the connector reported " + f"matched token counts {matched_tokens}. The cache directory was " + "not clean, so the comparison below is meaningless.") cache_files = [f for f in os.listdir(cache_dir) if f.endswith(".pt")] assert len(cache_files) > 0, "No cache files written by connector" + matched_tokens.clear() + llm2 = LLM(**llm_kwargs) - llm2.generate([prompt], sampling_params) - del llm2 + try: + output2 = llm2.generate([prompt], sampling_params) + warm_text = output2[0].outputs[0].text + warm_token_ids = list(output2[0].outputs[0].token_ids) + finally: + llm2.shutdown() + + assert matched_tokens and max(matched_tokens) > 0, ( + "The second run read nothing back from the connector cache " + f"(matched token counts {matched_tokens}), so the comparisons " + "below would pass just as well with the connector disabled.") + + assert len(warm_token_ids) == len(cold_token_ids), ( + f"Generation length changed: cold {len(cold_token_ids)} tokens, " + f"warm {len(warm_token_ids)} tokens.") + + # Exact equality is NOT asserted. Reusing cached KV skips prefill for + # the matched blocks, which changes the attention reduction order, so + # the logits differ in the last bits even though the restored K/V are + # bit-identical (the connector round-trips them through torch.save / + # torch.load). Greedy decoding turns a near-tie into a different token. + # Observed on V1: the two runs agreed on 31 of 32 tokens and split on + # the final one ("The company's" vs "The company is"). + # + # A corrupted or misaddressed cache does not look like that - it + # diverges early and degenerates - so requiring a long common prefix + # keeps the test meaningful without making it a coin flip. + common_prefix = 0 + for cold_id, warm_id in zip(cold_token_ids, warm_token_ids): + if cold_id != warm_id: + break + common_prefix += 1 + + min_common_prefix = math.floor( + len(cold_token_ids) * E2E_MIN_TOKEN_AGREEMENT) + assert common_prefix >= min_common_prefix, ( + f"Connector cache reuse diverged at token {common_prefix} of " + f"{len(cold_token_ids)}, below the {min_common_prefix}-token " + "floor. Early divergence indicates the restored KV is wrong, not " + "just numerically different.\n" + f" cold run: {cold_text!r}\n" + f" warm run: {warm_text!r}\n" + f" cold ids: {cold_token_ids}\n" + f" warm ids: {warm_token_ids}") finally: - os.environ.pop("CONNECTOR_CACHE_FOLDER", None) - if examples_dir in sys.path: sys.path.remove(examples_dir) - import shutil shutil.rmtree(cache_dir, ignore_errors=True) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 5b39eea23efa..15e80cf5b417 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -31,6 +31,7 @@ l0_a10: - unittest/_torch/executor/test_disagg_index_mapper_early_release.py - unittest/_torch/executor/test_kv_cache_compression_manager.py - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py + - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/modules/moe/test_communication_factory.py # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no @@ -118,25 +119,59 @@ l0_a10: # usage telemetry - unittest/llmapi/test_llm_telemetry.py::TestTelemetryPyTorchBackend - unittest/llmapi/test_llm_telemetry.py::TestTelemetryArchitectureExtraction - - llmapi/test_llm_api_connector.py::test_connector_simple[True] - - llmapi/test_llm_api_connector.py::test_connector_simple[False] - - llmapi/test_llm_api_connector.py::test_connector_async_onboard[True] - - llmapi/test_llm_api_connector.py::test_connector_async_onboard[False] - - llmapi/test_llm_api_connector.py::test_connector_async_save[True] - - llmapi/test_llm_api_connector.py::test_connector_async_save[False] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[True] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[False] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[True] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[False] - - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[False] - - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[True] - - llmapi/test_llm_api_connector.py::test_connector_multi_request - - llmapi/test_llm_api_connector.py::test_connector_priorities - - llmapi/test_llm_api_connector.py::test_connector_priorities_default - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[host_offloading] - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[beam_search] - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[attention_dp] - - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache + # KV connector. The suite is parametrized over the KV cache manager and both + # halves are gated: the connector is supported on KVCacheManagerV2, so a V2 + # regression has to be a merge-gate failure rather than something found later + # by hand. test_connector_runs_on_kv_cache_manager_v2 is what makes the V2 + # entries meaningful -- the creator silently falls back to the V1 manager for + # combinations it cannot serve, so without it every V2 id below could pass + # while running V1. test_connector_priorities[kv_cache_manager_v2] is + # xfail(strict=True) in the suite, not omitted here, so the retention gap + # stays counted. + - llmapi/test_llm_api_connector.py::test_v2_connector_contract_does_not_reuse_the_v1_methods + - llmapi/test_llm_api_connector.py::test_connector_runs_on_kv_cache_manager_v2 + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_multi_request[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-host_offloading] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-beam_search] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-attention_dp] + - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_vswa_reports_page_indices_per_layer_group[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_multi_request[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-host_offloading] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-beam_search] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-attention_dp] + - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_uniform_sliding_window[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_sliding_window_prefix_is_backed_by_history[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_vswa_reports_page_indices_per_layer_group[kv_cache_manager_v2] # third-party policy checks CPU-only - thirdparty/test_cmake_third_party.py::test_cmake_listfiles - thirdparty/test_git_modules.py::test_gitmodules diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 0294bcd0ecf4..483e555cb9cf 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -21,6 +21,7 @@ from tensorrt_llm._torch.pyexecutor import kv_cache_transceiver as transceiver_module from tensorrt_llm._torch.pyexecutor import py_executor as executor_module +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import BindKvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( @@ -557,6 +558,57 @@ def test_cpp_runtime_rejects_v2_mamba_manager(runtime): ) +@pytest.mark.parametrize("runtime", [None, "CPP", "auto"]) +def test_cpp_runtime_rejects_v2_manager(runtime): + """Plain KVCacheManagerV2 cannot drive the C++ transceiver either. + + `CacheTransceiverCpp` is bound to the V1 `BaseKVCacheManager`, while + `KVCacheManagerV2.impl` is the Python V2 core's manager. Without this guard + the combination reaches `BindKvCacheTransceiver` and dies on a nanobind + signature mismatch that names neither the manager nor the way out. + + `auto` is covered because it is the *default*: it resolves from the model's + preferred runtime, which knows nothing about which cache manager will be + built, so most models land on the C++ transceiver here. + """ + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + manager = object.__new__(KVCacheManagerV2) + + with pytest.raises(ValueError, match="KVCacheManagerV2 requires transceiver_runtime='PYTHON'"): + transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), manager, Mock(), config) + + +def test_python_nixl_transceiver_accepts_v2_manager(monkeypatch): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + expected = object() + constructor = Mock(return_value=expected) + fake_module = SimpleNamespace(KvCacheTransceiverV2=constructor) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.disaggregation.transceiver", fake_module) + manager = object.__new__(KVCacheManagerV2) + + result = transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), manager, Mock(), config) + + assert result is expected + constructor.assert_called_once() + + +def test_v2_mamba_manager_keeps_its_specific_rejection(): + """MambaHybridCacheManagerV2 subclasses KVCacheManagerV2. + + So the general guard must sit after the hybrid one, or the more specific + message -- the one that tells a hybrid user which manager they are on -- + would be shadowed. Both messages share the "requires + transceiver_runtime='PYTHON'" phrase, so this matches the manager name. + """ + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="CPP") + manager = object.__new__(MambaHybridCacheManagerV2) + + with pytest.raises(ValueError, match="MambaHybridCacheManagerV2 requires"): + transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + def test_python_runtime_rejects_cpp_mamba_manager(): config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") manager = object.__new__(CppMambaHybridCacheManager) diff --git a/tests/unittest/_torch/executor/test_kv_cache_layout.py b/tests/unittest/_torch/executor/test_kv_cache_layout.py new file mode 100644 index 000000000000..83a7d8116ea9 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_layout.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Tests for the KV cache layout description handed to a KV connector under +# KVCacheManagerV2 (``connectors/kv_cache_layout.py``). +# +# The layout replaces the single-pool-tensor registration used by the V1 +# manager, whose memory V2 cannot express: V2 has one slot address space per +# pool and one page-index space per layer group. +# +# Tests in TestKvCacheRegionArithmetic need no GPU. The rest construct a real +# KVCacheManagerV2 and therefore allocate device memory pools. + +import gc +import unittest + +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import ( + KvCacheBufferRef, + KvCacheLayerGroupLayout, + KvCacheLayout, + KvCacheRegion, + build_kv_cache_layout_v2, +) +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 +from tensorrt_llm.mapping import Mapping + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + + +def _make_kwargs( + *, + num_layers: int = 4, + num_kv_heads=4, + head_dim=128, + tokens_per_block: int = 8, + max_seq_len: int = 256, + max_batch_size: int = 4, + max_tokens: int = 2048, + dtype=DataType.HALF, + kv_cache_type=CacheType.SELF, + vocab_size: int = 32000, + kv_cache_config=None, +): + return dict( + kv_cache_config=kv_cache_config + or KvCacheConfigV2(max_tokens=max_tokens, enable_block_reuse=False), + kv_cache_type=kv_cache_type, + 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(world_size=1, tp_size=1, rank=0), + dtype=dtype, + vocab_size=vocab_size, + ) + + +class TestKvCacheRegionArithmetic(unittest.TestCase): + """Address arithmetic and lookup helpers. No GPU required.""" + + def _region(self, base=4096, size=256, stride=1024, num_slots=8): + return KvCacheRegion( + base=base, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=0, role="key"),), + ) + + def test_address_of_follows_stride(self): + region = self._region() + self.assertEqual(region.address_of(0), 4096) + self.assertEqual(region.address_of(1), 4096 + 1024) + self.assertEqual(region.address_of(7), 4096 + 7 * 1024) + + def test_address_of_rejects_out_of_range_slot(self): + region = self._region(num_slots=8) + for bad in (-1, 8, 99): + with self.assertRaises(IndexError): + region.address_of(bad) + + def test_as_tensor_rejects_dtype_not_dividing_extent(self): + # size/stride are byte counts; a dtype whose itemsize does not divide + # them cannot produce a correct view, so it must fail loudly rather + # than silently truncate. + region = self._region(size=6, stride=6) + with self.assertRaises(ValueError): + region.as_tensor(dtype=torch.float32) + + def test_bytes_per_page_sums_regions(self): + group = KvCacheLayerGroupLayout( + layer_group_id=0, + layer_ids=(0, 1), + window_size=None, + regions=(self._region(size=256), self._region(base=8192, size=128)), + ) + self.assertEqual(group.bytes_per_page, 384) + + def test_layout_lookup_by_group_and_layer(self): + group_a = KvCacheLayerGroupLayout(0, (0, 2), None, ()) + group_b = KvCacheLayerGroupLayout(1, (1, 3), 128, ()) + layout = KvCacheLayout(tokens_per_block=8, groups=(group_a, group_b)) + + self.assertIs(layout.group(1), group_b) + self.assertIs(layout.group_of_layer(2), group_a) + self.assertIs(layout.group_of_layer(3), group_b) + with self.assertRaises(KeyError): + layout.group(7) + with self.assertRaises(KeyError): + layout.group_of_layer(99) + + +class TestKvCacheRegionAliasing(unittest.TestCase): + """as_tensor must alias the exact bytes address_of names.""" + + def setUp(self): + torch.cuda.init() + + def test_as_tensor_aliases_strided_slots(self): + # Lay out 4 "slots" of 32 bytes each, and describe the middle 8 bytes + # of every slot as a region. Writing through the view must land at + # base + stride * i, and must not disturb neighbouring bytes. + num_slots, stride, offset, size = 4, 32, 8, 8 + backing = torch.zeros(num_slots * stride, dtype=torch.uint8, device="cuda") + + region = KvCacheRegion( + base=backing.data_ptr() + offset, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=0, role="key"),), + ) + view = region.as_tensor() + self.assertEqual(tuple(view.shape), (num_slots, size)) + + for slot in range(num_slots): + view[slot] = slot + 1 + + flat = backing.cpu() + for slot in range(num_slots): + start = slot * stride + self.assertTrue( + bool((flat[start + offset : start + offset + size] == slot + 1).all()), + f"slot {slot} payload not written at the address address_of() names", + ) + # Bytes outside the described range must be untouched. + self.assertTrue(bool((flat[start : start + offset] == 0).all())) + self.assertTrue(bool((flat[start + offset + size : start + stride] == 0).all())) + + +class TestBuildKvCacheLayoutV2(unittest.TestCase): + """The builder against a real KVCacheManagerV2.""" + + def setUp(self): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + + def tearDown(self): + gc.collect() + torch.cuda.empty_cache() + + def test_layout_covers_every_layer_exactly_once(self): + num_layers = 4 + mgr = KVCacheManagerV2(**_make_kwargs(num_layers=num_layers)) + try: + layout = build_kv_cache_layout_v2(mgr) + + self.assertEqual(layout.tokens_per_block, mgr.tokens_per_block) + self.assertTrue(layout.groups, "layout must describe at least one layer group") + + covered = [lid for group in layout.groups for lid in group.layer_ids] + self.assertCountEqual( + covered, + list(mgr.pp_layers), + "every local layer must appear in exactly one layer group", + ) + + # Every layer that owns storage must be reachable through a region. + in_regions = [ + ref.layer_id + for group in layout.groups + for region in group.regions + for ref in region.buffers + ] + self.assertCountEqual(set(in_regions), set(covered)) + finally: + mgr.shutdown() + del mgr + + def test_regions_are_disjoint_and_inside_the_slot(self): + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + pool_groups = list(mgr.impl.pool_group_descs) + self.assertEqual(len(pool_groups), 1, "test config should yield one pool group") + pool_group = pool_groups[0] + self.assertEqual(len(pool_group.pools), 1, "test config should yield one pool") + pool = pool_group.pools[0] + + for group in layout.groups: + spans = [] + for region in group.regions: + self.assertEqual(region.num_slots, int(pool_group.num_slots)) + self.assertEqual(region.stride, int(pool.slot_bytes)) + + offset = region.base - int(pool.base_address) + self.assertGreaterEqual(offset, 0) + self.assertLessEqual( + offset + region.size, + int(pool.slot_bytes), + "a region must lie inside one slot", + ) + spans.append((offset, offset + region.size)) + + spans.sort() + for (_, prev_end), (next_start, _) in zip(spans, spans[1:]): + self.assertLessEqual(prev_end, next_start, "regions must not overlap") + finally: + mgr.shutdown() + del mgr + + def test_region_addresses_agree_with_pool_descriptor(self): + # Cross-check: region base/stride come from get_aggregated_pages, while + # pool base_address/slot_bytes come from pool_group_descs. These are + # independent public APIs and must agree on where slot i lives. + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + pool = list(mgr.impl.pool_group_descs)[0].pools[0] + pool_base, slot_bytes = int(pool.base_address), int(pool.slot_bytes) + + for group in layout.groups: + for region in group.regions: + offset = region.base - pool_base + for slot in (0, 1, region.num_slots - 1): + self.assertEqual( + region.address_of(slot), + pool_base + slot_bytes * slot + offset, + ) + finally: + mgr.shutdown() + del mgr + + def test_uniform_model_yields_one_full_slot_region(self): + # With uniform K/V sizes every buffer in a layer group is adjacent, so + # coalescing should collapse them into a single region spanning the + # whole slot -- the efficient whole-page transfer, derived rather than + # assumed. + mgr = KVCacheManagerV2(**_make_kwargs(num_layers=4)) + try: + layout = build_kv_cache_layout_v2(mgr) + pool = list(mgr.impl.pool_group_descs)[0].pools[0] + + self.assertEqual(len(layout.groups), 1) + group = layout.groups[0] + self.assertEqual(len(group.regions), 1) + region = group.regions[0] + self.assertEqual(region.size, int(pool.slot_bytes)) + self.assertEqual(region.base, int(pool.base_address)) + # 4 layers * (K, V) + self.assertEqual(len(region.buffers), 8) + self.assertEqual( + [ref.role for ref in region.buffers], + ["key", "value"] * 4, + ) + finally: + mgr.shutdown() + del mgr + + def test_full_attention_reports_no_window(self): + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + for group in layout.groups: + self.assertIsNone(group.window_size) + finally: + mgr.shutdown() + del mgr + + def test_mla_layout_has_no_value_buffers(self): + # SELFKONLY carries a single compressed latent per token. Nothing in the + # layout counts K against V, so this must simply come out as a layout + # with no "value" role rather than needing a kv-factor special case. + mgr = KVCacheManagerV2(**_make_kwargs(kv_cache_type=CacheType.SELFKONLY)) + try: + layout = build_kv_cache_layout_v2(mgr) + roles = { + ref.role + for group in layout.groups + for region in group.regions + for ref in region.buffers + } + self.assertIn("key", roles) + self.assertNotIn("value", roles) + finally: + mgr.shutdown() + del mgr + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py new file mode 100644 index 000000000000..2515b70548f6 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py @@ -0,0 +1,640 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the KV connector prefix phases on KVCacheManagerV2. + +KVCacheManagerV2's scheduling pass is speculative: a request can be prepared +and then dropped in the same iteration at the token budget, at +``resize_context``, at multimodal alignment or at cross attention, and retried +later. The connector ABC, meanwhile, promises exactly one +``get_num_new_matched_tokens`` per request and lets connectors take ownership +of remote blocks inside it. + +Reconciling those two is the whole point of the three phases exercised here, +and none of it is observable from the end-to-end connector tests, which never +defer a request. So these drive the phases directly against a stub cache. + +The stub cache models V2's own two-phase split, because that split is what +forces the connector phases apart: ``_create_kv_cache`` returns a **suspended** +cache, ``_KVCache.resize`` asserts the cache is ACTIVE, and +``_resume_and_restore`` is what activates it. So the offer has to be taken -- +and the context position advanced over it -- *before* the resume, while the +capacity and history that back it can only be reserved *after*. ``FakeKvCache`` +therefore starts suspended and refuses ``resize`` until resumed, so a version of +this code that folded the two phases into one would fail here rather than only +on hardware. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + +TOKENS_PER_BLOCK = 32 +PROMPT_LEN = 256 + + +class FakeKvCache: + """The parts of ``_KVCache`` the connector phases touch. + + Starts suspended, as ``_create_kv_cache`` leaves it. ``resize`` reproduces + the real assertions -- ACTIVE status, history never decreasing, history + never above capacity -- since those are exactly what dictate when each + phase may run and how it must call ``resize``. + """ + + def __init__(self, committed=0, resize_ok=True, active=False, trace=None): + self.num_committed_tokens = committed + self.capacity = committed + self.history_length = committed + self.enable_swa_scratch_reuse = True + self.resize_ok = resize_ok + self.is_active = active + self.resize_calls = [] + self.trace = trace if trace is not None else [] + + def resume(self, cuda_stream): + self.is_active = True + self.trace.append("resume") + return True + + def suspend(self): + self.is_active = False + self.trace.append("suspend") + + def resize(self, capacity, history_length=None): + # `_KVCache.resize` asserts ACTIVE. A cache that has just been created, + # or that was suspended when its request was deferred, is not. + assert self.is_active, "resize on a suspended cache" + self.resize_calls.append((capacity, history_length)) + self.trace.append(("resize", capacity, history_length)) + if not self.resize_ok: + return False + if history_length is not None: + if history_length < self.history_length: + raise ValueError("History length cannot be decreased") + if capacity is not None and capacity < history_length: + raise ValueError("History length cannot be greater than capacity") + self.history_length = history_length + if capacity is not None: + self.capacity = capacity + return True + + +class FakeRequest: + def __init__(self, request_id=0, prompt_len=PROMPT_LEN): + self.py_request_id = request_id + self.request_id = request_id + self.prompt_len = prompt_len + self.context_current_position = 0 + self.prepopulated_prompt_len = 0 + self.is_first_context_chunk = True + self.is_dummy = False + self.is_generation_only_request = False + self.is_disagg_generation_init_state = False + self.py_num_connector_matched_tokens = 0 + self.py_connector_prefix_start = None + self.py_connector_prefix_end = None + self.py_connector_load_async = False + self.py_connector_delivered = False + + def set_prepopulated_prompt_len(self, prepopulated_prompt_len, tokens_per_block): + assert prepopulated_prompt_len < self.prompt_len, ( + f"prepopulatedPromptLen ({prepopulated_prompt_len}) >= promptLen ({self.prompt_len})" + ) + self.prepopulated_prompt_len = prepopulated_prompt_len + if prepopulated_prompt_len > 0: + self.context_current_position = prepopulated_prompt_len + + +class FakeConnectorManager: + """Records the calls the phases make, in order.""" + + def __init__(self, num_matched=0, load_async=False, trace=None): + self.num_matched = num_matched + self.load_async = load_async + self.queries = [] + self.commits = [] + self.cancels = [] + self.trace = trace if trace is not None else [] + + def query_num_new_matched_tokens(self, request, num_computed_tokens): + self.queries.append((request.request_id, num_computed_tokens)) + self.trace.append(("query", num_computed_tokens)) + return self.num_matched, self.load_async + + def commit_new_matched_tokens(self, request, num_tokens, load_kv_async): + self.commits.append((request.request_id, num_tokens, load_kv_async)) + self.trace.append(("commit", num_tokens)) + request.py_num_connector_matched_tokens = num_tokens + + def cancel_load(self, request, start, end): + self.cancels.append((request.request_id, start, end)) + self.trace.append(("cancel", start, end)) + + +def make_manager(connector): + """A KVCacheManagerV2 with only the fields the connector phases read. + + Constructing a real one needs a GPU and a pool allocation; the phases under + test are pure request/cache bookkeeping, so bypass __init__ rather than + turning this into an integration test. + """ + manager = object.__new__(KVCacheManagerV2) + manager.kv_connector_manager = connector + manager.is_draft = False + manager.tokens_per_block = TOKENS_PER_BLOCK + manager.kv_cache_map = {} + # Read by `_prepare_context_impl` on the first-chunk path, so that the + # ordering tests below can drive the real thing rather than a re-statement + # of it. + manager.enable_block_reuse = True + manager.conversation_manager = None + manager._stream = SimpleNamespace(cuda_stream=0) + # Page-index buffer plumbing: needs the real IndexMapper and pool tensors, + # and has no bearing on which phase runs when. + manager._restore_page_index_bufs = lambda request_id, kv_cache: None + return manager + + +def prepare(manager, req, kv_cache): + """One scheduling attempt, driving the real ``_prepare_context_impl``. + + Seeding ``kv_cache_map`` skips the ``_create_kv_cache`` branch, which needs + a GPU; everything after it -- the local-match anchor, phase 1, the resume, + phase 2 -- is the production code, so the phase ordering is observed rather + than restated. + """ + manager.kv_cache_map[req.py_request_id] = kv_cache + assert manager._prepare_context_impl(req) + # A memoised re-read of the position the attempt settled on, not a second + # ask -- `py_connector_prefix_end` is set by now, so the connector is not + # consulted again. `TestAskOnce` pins that. + return manager._connector_prefix_position(req, kv_cache) + + +class TestTwoPhaseSplit: + """The connector phases are split because V2's own allocation is. + + V2 separates *match and take ownership*, which needs only the token + sequence, from *become resident on GPU*, which needs slots and can fail. + `_create_kv_cache` returns a suspended cache and `_KVCache.resize` asserts + ACTIVE, so the two connector steps straddle `_resume_and_restore`. Every + test in this class fails if they are folded into one. + """ + + def test_the_offer_is_taken_before_the_cache_is_resident(self): + trace = [] + connector = FakeConnectorManager(num_matched=64, trace=trace) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32, trace=trace) + + prepare(manager, req, kv_cache) + + assert trace == [("query", 32), "resume", ("resize", 96, 96)], ( + "the connector is asked while the cache is still suspended, and the " + "capacity backing its answer is reserved only once it is active" + ) + + def test_phase_one_touches_no_residency(self): + """Asking needs the token sequence and nothing else. + + This is what makes it safe to ask during a speculative scheduling pass: + no slot is claimed, so a request that is then dropped has cost nothing + locally. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + position = manager._connector_prefix_position(req, kv_cache) + + assert position == 96 + assert kv_cache.is_active is False + assert kv_cache.resize_calls == [] + assert kv_cache.capacity == 32 + + def test_phase_two_refuses_a_suspended_cache(self): + """The assertion that forces the split to exist. + + `_KVCache.resize` asserts ACTIVE, so reserving before the resume -- the + shape a single merged step would have -- fails on the first request + that has anything to reserve. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + position = manager._connector_prefix_position(req, kv_cache) + + with pytest.raises(AssertionError, match="suspended"): + manager._reserve_connector_prefix(req, kv_cache, position) + + def test_a_deferred_request_is_resumed_again_before_reserving(self): + """A deferred first chunk is suspended, so the split holds every attempt. + + `resize_context` suspends the cache when it cannot grow it, so the next + attempt starts from the suspended state again -- and must still ask + nothing, resume, and only then reserve. + """ + trace = [] + connector = FakeConnectorManager(num_matched=64, trace=trace) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32, trace=trace) + + prepare(manager, req, kv_cache) + kv_cache.suspend() + trace.clear() + + prepare(manager, req, kv_cache) + + assert trace == ["resume", ("resize", 96, 96)] + assert connector.queries == [(0, 32)] + + +class TestAskOnce: + def test_deferred_request_is_not_asked_again(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + # Iteration 1: prepared, then dropped before it reached the batch. + prepare(manager, req, kv_cache) + # Iteration 2: prepared again. + prepare(manager, req, kv_cache) + + assert connector.queries == [(0, 0)], ( + "the connector takes ownership of remote blocks inside the query, " + "so a second one double-pins them and breaks the ABC's " + "at-most-once promise" + ) + + def test_delivering_records_once_for_the_iteration_that_ran(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + assert connector.commits == [] + + manager._deliver_connector_prefix(req) + assert connector.commits == [(0, 64, False)] + + def test_unasked_request_delivers_nothing(self): + connector = FakeConnectorManager() + manager = make_manager(connector) + req = FakeRequest() + + manager._deliver_connector_prefix(req) + + assert connector.commits == [] + assert connector.cancels == [] + + +class TestOfferEndIsAbsolute: + def test_position_does_not_overshoot_when_the_local_match_grows(self): + """The reason the offer end is memoised rather than the returned delta. + + A deferred request re-derives its local match from the radix tree, and + another request's commit may have grown it in the meantime. Adding the + delta to the new match would set the position past the union of what is + locally computed and what the connector holds, leaving the tokens in + between neither computed nor loaded -- silently garbage KV. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + assert req.context_current_position == 64 + + # Another request committed the first 32 tokens while this one waited. + kv_cache.num_committed_tokens = 32 + position = prepare(manager, req, kv_cache) + + assert position == 64, "0 + 64 offered, so 64 -- not 32 + 64" + assert req.context_current_position == 64 + + def test_position_follows_the_local_match_when_it_overtakes_the_offer(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + + kv_cache.num_committed_tokens = 128 + position = prepare(manager, req, kv_cache) + + assert position == 128 + assert req.context_current_position == 128 + assert connector.queries == [(0, 0)] + + def test_offer_end_is_clamped_below_the_prompt(self): + """A connector holding the whole prompt is the steady state of a repeat. + + The last prompt position must be computed locally regardless, since the + first generation step consumes its activations, so the offer is clamped + rather than rejected -- and the *clamped* delta is what gets recorded, + or the connector is pointed at the wrong offset. + """ + connector = FakeConnectorManager(num_matched=PROMPT_LEN) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert req.py_connector_prefix_end == PROMPT_LEN - 1 + assert req.context_current_position == PROMPT_LEN - 1 + assert connector.commits == [(0, PROMPT_LEN - 1, False)] + # The clamp is the one place an offer shrinks without passing through + # phase 3 or the release path, both of which read the clamped end. So + # it has to hand the remainder back itself or the connector keeps + # ownership of it for the life of the process. + assert connector.cancels == [(0, PROMPT_LEN - 1, PROMPT_LEN)] + + +class TestWriteRangeAtDelivery: + def test_subsumed_head_is_handed_back_and_not_transferred(self): + """``[start, committed)`` is locally owned by the time we deliver. + + Those are committed pages in the radix tree, potentially shared with + other live requests, and V2's rule is that only the owning request + writes into its own padding. So the connector is told to transfer only + what is still privately owned, and the rest is cancelled. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + kv_cache.num_committed_tokens = 32 + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert connector.cancels == [(0, 0, 32)] + assert connector.commits == [(0, 32, False)], "[32, 64), not [0, 64)" + + def test_fully_subsumed_offer_transfers_nothing(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + kv_cache.num_committed_tokens = 128 + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert connector.cancels == [(0, 0, 64)] + assert connector.commits == [(0, 0, False)] + + def test_recorded_delta_restores_the_locally_computed_position(self): + """``computed_position`` is reported as ``end - recorded``. + + That is what a connector uses to find where its blocks start, so the + recorded delta has to be the one anchored at the *current* commit + boundary. + """ + connector = FakeConnectorManager(num_matched=96) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + prepare(manager, req, kv_cache) + kv_cache.num_committed_tokens = 64 + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + recorded = connector.commits[0][1] + assert req.context_current_position - recorded == 64 + + +class TestResidency: + def test_capacity_and_history_move_together(self): + """After a reuse match both equal the local match. + + Raising history alone would trip "History length cannot be greater than + capacity"; raising capacity alone would leave a sliding-window group + allocating a page for every block of the served prefix. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + prepare(manager, req, kv_cache) + + assert kv_cache.resize_calls == [(96, 96)] + assert kv_cache.capacity == 96 + assert kv_cache.history_length == 96 + + def test_existing_capacity_is_not_shrunk(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + kv_cache.capacity = 256 + + prepare(manager, req, kv_cache) + + assert kv_cache.resize_calls == [(256, 64)] + + def test_swa_scratch_reuse_is_disabled_for_a_served_prefix(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + + assert kv_cache.enable_swa_scratch_reuse is False + + def test_empty_offer_touches_nothing(self): + connector = FakeConnectorManager(num_matched=0) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert kv_cache.resize_calls == [] + assert kv_cache.enable_swa_scratch_reuse is True + assert connector.cancels == [] + assert connector.commits == [(0, 0, False)] + + def test_offer_is_handed_back_when_pages_run_out(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32, resize_ok=False) + + prepare(manager, req, kv_cache) + + assert connector.cancels == [(0, 32, 96)] + assert req.context_current_position == 32 + # The request runs local-only from here, and delivering must not then + # record an external load for a range that was never covered. + manager._deliver_connector_prefix(req) + assert connector.commits == [(0, 0, False)] + + +class TestExclusions: + @pytest.mark.parametrize( + "attribute", + ["is_dummy", "is_generation_only_request", "is_disagg_generation_init_state"], + ) + def test_request_kinds_that_are_never_asked(self, attribute): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + setattr(req, attribute, True) + kv_cache = FakeKvCache(committed=0) + + assert prepare(manager, req, kv_cache) is None + assert connector.queries == [] + assert req.context_current_position == 0 + + def test_no_connector_is_a_no_op(self): + manager = make_manager(None) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + assert manager._connector_prefix_position(req, kv_cache) is None + + def test_draft_manager_never_asks(self): + """The draft manager's prepare_resources skips the connector hooks, so + a query there would never be delivered.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + manager.is_draft = True + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + assert manager._connector_prefix_position(req, kv_cache) is None + assert connector.queries == [] + + +class TestUndeliveredOfferIsReleased: + def test_request_that_dies_before_delivery_hands_the_offer_back(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + # Cancelled, timed out, or failed before it ever reached a batch. + manager._release_undelivered_connector_prefix(req) + + assert connector.cancels == [(0, 0, 64)] + + def test_delivered_request_keeps_its_offer_on_free(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + req.py_connector_delivered = True + manager._release_undelivered_connector_prefix(req) + + assert connector.cancels == [] + + def test_unasked_request_releases_nothing(self): + connector = FakeConnectorManager() + manager = make_manager(connector) + + manager._release_undelivered_connector_prefix(FakeRequest()) + + assert connector.cancels == [] + + +class TestAsyncLoad: + def test_async_flag_is_carried_from_query_to_delivery(self): + connector = FakeConnectorManager(num_matched=64, load_async=True) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + assert req.py_connector_load_async is True + + manager._deliver_connector_prefix(req) + assert connector.commits == [(0, 64, True)] + + def test_async_hold_survives_the_clamp(self): + """An async connector starts its transfer inside the query itself. + + So even when the clamp leaves nothing to load, the request still has to + be held out of the batch until the connector reports the transfer done, + or the forward races the writes. + """ + connector = FakeConnectorManager(num_matched=1, load_async=True) + manager = make_manager(connector) + req = FakeRequest(prompt_len=PROMPT_LEN) + kv_cache = FakeKvCache(committed=PROMPT_LEN - 1) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert req.py_connector_prefix_end == PROMPT_LEN - 1 + assert connector.commits == [(0, 0, True)] + # Lossy for an async load -- the transfer began inside the query -- but + # still the only signal the connector gets that the tail is dead. + assert connector.cancels == [(0, PROMPT_LEN - 1, PROMPT_LEN)] + + +class TestPhasesMustAgree: + """Phase 3 refuses to report a load phase 2 never made room for. + + This is the single coupling in the design that nothing downstream checks: + the runtime hands the connector ``context_current_position - recorded`` as + the range it computed locally, the subtraction is unguarded + (``kv_cache_connector.py:480``), and connectors divide the result into block + ordinals rather than validating it. A phase 2 that silently no-ops -- by + regression, or by a future caller reordering the phases -- therefore points + the connector at a negative offset inside its own code rather than failing. + """ + + def test_recording_more_than_the_position_covers_is_an_error(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0) + + prepare(manager, req, kv_cache) + # The shape a silently skipped phase 2 leaves behind: the offer is on + # the request, but no capacity, history or position backs it. + req.context_current_position = 0 + + with pytest.raises(AssertionError, match="phase 2 did not reserve"): + manager._deliver_connector_prefix(req) + + assert connector.commits == [], ( + "the load must not be recorded when the assertion fires, or the " + "connector is left holding an offer the runtime denied" + ) + + def test_a_fully_reserved_offer_is_accepted(self): + """Anti-vacuity: the assertion must not fire on the normal path.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32) + + prepare(manager, req, kv_cache) + manager._deliver_connector_prefix(req) + + assert connector.commits == [(0, 64, False)] 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 new file mode 100644 index 000000000000..a8031da32800 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KV connector prefix phases against a *real* ``KVCacheManagerV2``. + +``test_kv_connector_v2_prefix.py`` drives the same three phases against a stub +cache. That is what makes it fast and exhaustive, and it is the right place for +the ordering and state-machine rules -- but a stub cannot show that a real +``_KVCache`` survives the sequence: that ``resize`` finds real pages for the +offered prefix, that ``history_length`` really moves, that suspend/resume across +a deferral leaves the request askable exactly once, and that the page slots +handed to the connector at delivery are distinct and real. + +The engine-level suite cannot show it either, for a different reason. Whether a +request is deferred there depends on whether it reaches the scheduler in the +same pass as the request that outbids it, and that is a race: the same test was +observed asking both requests in the first pass when run alone, and asking the +second one an iteration later when run after the rest of the suite. Deferral is +therefore driven here directly -- ``prepare_context`` without the +``prepare_resources`` that would have followed it in an iteration where the +request actually ran -- which is exactly what the scheduler does to a request +that loses the token budget after being prepared (scheduler_v2.py:554-558). + +These tests allocate device memory pools. +""" + +import gc +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + +# These build a real manager, which allocates device pools. The directory is +# listed in the GPU-less l0_cpu stage, so the requirement is declared rather +# than left to fail at `torch.cuda.init()`. +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="allocates real KV cache pools" +) + +TOKENS_PER_BLOCK = 32 +PROMPT_LEN = 96 +OFFER_TOKENS = 32 + + +class FakeConnectorManager: + """Records what the phases tell the connector, in order. + + Same shape as the one in ``test_kv_connector_v2_prefix.py``, plus the + ``build_scheduler_output`` that ``prepare_resources`` calls after delivery. + """ + + def __init__(self, num_matched=OFFER_TOKENS, load_async=False): + self.num_matched = num_matched + self.load_async = load_async + self.queries = [] + self.commits = [] + self.cancels = [] + self.allocs = [] + + def query_num_new_matched_tokens(self, request, num_computed_tokens): + self.queries.append((request.py_request_id, num_computed_tokens)) + return self.num_matched, self.load_async + + def commit_new_matched_tokens(self, request, num_tokens, load_kv_async): + self.commits.append((request.py_request_id, num_tokens, load_kv_async)) + request.py_num_connector_matched_tokens = num_tokens + + def cancel_load(self, request, start, end): + self.cancels.append((request.py_request_id, start, end)) + + def update_state_after_alloc(self, request, block_ids): + self.allocs.append((request.py_request_id, list(block_ids))) + + def build_scheduler_output(self, scheduled_batch, kv_cache_manager): + pass + + +def make_manager(connector, **overrides): + kwargs = dict( + kv_cache_config=KvCacheConfig(max_tokens=2048, enable_block_reuse=True), + kv_cache_type=CacheType.SELF, + num_layers=2, + num_kv_heads=4, + head_dim=64, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=256, + max_batch_size=4, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.HALF, + vocab_size=32000, + kv_connector_manager=connector, + ) + kwargs.update(overrides) + return KVCacheManagerV2(**kwargs) + + +def make_request(request_id=1, prompt_len=PROMPT_LEN): + return LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=list(range(prompt_len)), + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + + +def deliver(manager, request): + """One `prepare_resources`, i.e. the request reached the final batch.""" + manager.prepare_resources(SimpleNamespace(context_requests=[request], generation_requests=[])) + + +@pytest.fixture +def connector(): + return FakeConnectorManager() + + +@pytest.fixture +def manager(connector): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + mgr = make_manager(connector) + yield mgr + mgr.shutdown() + del mgr + gc.collect() + torch.cuda.empty_cache() + + +def test_offer_is_backed_by_real_pages(manager, connector): + """Phase 2 against a real `_KVCache`: capacity, history, and slots. + + `history_length` is what the stub suite can only observe as an argument to + a fake `resize`. Here it is read back off the cache that the forward pass + would use, which is the difference between "we called resize correctly" and + "the offered prefix is resident". + """ + request = make_request() + + assert manager.prepare_context(request) + + kv_cache = manager.kv_cache_map[request.py_request_id] + assert request.context_current_position == OFFER_TOKENS + assert kv_cache.capacity >= OFFER_TOKENS + assert kv_cache.history_length == OFFER_TOKENS + assert kv_cache.is_active + + deliver(manager, request) + + assert connector.commits == [(request.py_request_id, OFFER_TOKENS, False)] + assert len(connector.allocs) == 1 + + _, page_indices = connector.allocs[0] + assert len(page_indices) >= OFFER_TOKENS // TOKENS_PER_BLOCK + assert all(index != BAD_PAGE_INDEX for index in page_indices) + assert len(set(page_indices)) == len(page_indices) + + +def test_deferred_request_is_asked_once_and_delivered_when_it_runs(manager, connector): + """Prepared, dropped before the batch, prepared again, then run. + + The cache is suspended in between, which is what `_revert_context_resize` + does to a request that loses its slot after being prepared. Both attempts + must reach the same offer end, and the connector must see exactly one query + and exactly one commit -- it took ownership of remote blocks inside the + first query, and a second would double-pin them. + """ + request = make_request() + + # Iteration 1: prepared, then outbid before it reached the batch. + assert manager.prepare_context(request) + kv_cache = manager.kv_cache_map[request.py_request_id] + kv_cache.suspend() + + assert connector.queries == [(request.py_request_id, 0)] + assert connector.commits == [], "nothing may be recorded until it runs" + + # Iteration 2: prepared again, and this time it runs. + assert manager.prepare_context(request) + assert kv_cache.is_active, "phase 2 must have found an active cache" + assert request.context_current_position == OFFER_TOKENS + assert kv_cache.history_length == OFFER_TOKENS + + deliver(manager, request) + + assert connector.queries == [(request.py_request_id, 0)] + assert connector.commits == [(request.py_request_id, OFFER_TOKENS, False)] + assert len(connector.allocs) == 1 + + +def test_delivering_twice_reports_one_allocation(manager, connector): + """An asynchronously loaded request re-enters on its first context chunk. + + `py_connector_delivered` is what stops that second pass re-recording the + load and firing a second `update_state_after_alloc` for one allocation. + """ + request = make_request() + + assert manager.prepare_context(request) + deliver(manager, request) + deliver(manager, request) + + assert len(connector.commits) == 1 + assert len(connector.allocs) == 1 + + +def test_undelivered_offer_is_released_when_the_request_is_freed(manager, connector): + """Phase 1 is speculative, so an offer can outlive the request. + + Without this the connector holds the remote blocks it took ownership of for + the rest of the process's life. + """ + request = make_request() + + assert manager.prepare_context(request) + manager.free_resources(request) + + assert connector.cancels == [(request.py_request_id, 0, OFFER_TOKENS)] + + +def test_delivered_offer_is_not_released_when_the_request_is_freed(manager, connector): + request = make_request() + + assert manager.prepare_context(request) + deliver(manager, request) + manager.free_resources(request) + + assert connector.cancels == [] diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 1ce6d85df791..37a2ed3b94c9 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -900,7 +900,7 @@ def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank(): transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) transceiver._page_table = SimpleNamespace(layer_groups=[]) request = SimpleNamespace( - is_generation_only_request=lambda: False, + is_generation_only_request=False, prompt_len=0, py_request_id=123, ) @@ -922,7 +922,7 @@ def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) transceiver._page_table = SimpleNamespace(layer_groups=[]) request = SimpleNamespace( - is_generation_only_request=lambda: False, + is_generation_only_request=False, prompt_len=0, py_request_id=123, ) @@ -937,8 +937,13 @@ def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): "max_beam_width, has_connector, expected", [ (2, False, "max_beam_width > 1"), - (1, True, "kv_connector_manager"), - (2, True, "kv_connector_manager, max_beam_width > 1"), + # A KV connector alone no longer forces a V1 fallback: it is supported on + # KVCacheManagerV2 through the pool-layout registration path, so V2 is + # returned unchanged and nothing is raised. + (1, True, None), + # With beam search still incompatible, the connector must not appear in + # the reason list -- it is not what makes this configuration unsupported. + (2, True, "max_beam_width > 1"), ], ) def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( @@ -957,6 +962,15 @@ def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( creator._kv_connector_manager = object() if has_connector else None creator._max_beam_width = max_beam_width + if expected is None: + assert ( + creator._validate_or_fallback_kv_cache_manager_v2( + MambaHybridCacheManagerV2, model_config, KvCacheConfig() + ) + is MambaHybridCacheManagerV2 + ) + return + with pytest.raises(NotImplementedError, match=expected): creator._validate_or_fallback_kv_cache_manager_v2( MambaHybridCacheManagerV2, model_config, KvCacheConfig() diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 7d44477e520e..1ecc71fb25ba 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -179,7 +179,7 @@ def _make_request_stub(req_id: int, prompt_len: int = 4) -> SimpleNamespace: py_draft_tokens=[], py_is_first_draft=False, is_context_only_request=False, - is_generation_only_request=lambda: False, + is_generation_only_request=False, py_disaggregated_params=None, py_multimodal_data=None, py_mm_encoder_event=None, @@ -432,7 +432,7 @@ def test_context_logits_use_final_token_graph_candidate(self) -> None: def test_generation_only_request_in_context_list_falls_back(self) -> None: context = _make_request_stub(1) - context.is_generation_only_request = lambda: True + context.is_generation_only_request = True batch = ScheduledRequests() batch.context_requests_last_chunk = [context] graph_batch, promoted_ids = _make_single_token_context_graph_batch( diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index 0c098fda90a7..bf074942761e 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/tests/unittest/_torch/executor/test_request_utils.py @@ -127,7 +127,7 @@ def test_executor_request_to_llm_request_adopts_context_phase_draft_tokens() -> exclude_last_generation_logits=False, ) - assert llm_request.is_generation_only_request() + assert llm_request.is_generation_only_request assert llm_request.has_draft_tokens() assert llm_request.num_draft_tokens == len(draft_tokens) assert llm_request.draft_tokens == draft_tokens diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py index 99c233a58ce0..dc9e56eaed17 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py @@ -267,6 +267,7 @@ def get_tokens(self, beam: int) -> list: assert beam == 0 return list(self._tokens) + @property def is_generation_only_request(self) -> bool: return self._generation_only diff --git a/tests/unittest/_torch/test_connector.py b/tests/unittest/_torch/test_connector.py index 01675416c5ea..b6fb2f8935bf 100644 --- a/tests/unittest/_torch/test_connector.py +++ b/tests/unittest/_torch/test_connector.py @@ -23,7 +23,7 @@ from tensorrt_llm import mpi_rank from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( - AsyncRequests, KvCacheConnectorManager, + AsyncRequests, KvCacheConnectorManager, KvCacheConnectorScheduler, KvCacheConnectorSchedulerOutputManager) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -177,6 +177,81 @@ def test(): run_across_mpi(mpi_pool_executor, test, 2) +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_connector_manager_query_is_side_effect_free(mpi_pool_executor): + """The query and the commit are separable, and the query records nothing. + + KVCacheManagerV2 asks during a speculative scheduling pass and resolves the + answer in whichever iteration the request actually runs. That only works if + asking is inert: `external_loads` is cleared by every + `build_scheduler_output`, and a request registered as loading is dropped + from the batch. Recording at query time would attribute the load to + whichever iteration happened to ask. + """ + + def test(): + worker = MagicMock() + + if mpi_rank() == 0: + scheduler = MagicMock() + scheduler.get_num_new_matched_tokens.return_value = (16, True) + else: + scheduler = None + + manager = KvCacheConnectorManager(worker, scheduler=scheduler) + + req = MagicMock() + req.request_id = 42 + req.is_generation_only_request = False + req.py_num_connector_matched_tokens = 0 + + assert manager.query_num_new_matched_tokens(req, 32) == (16, True) + + assert manager.new_async_requests.loading_ids == set() + assert manager.scheduler_output_manager.external_loads == {} + assert req.py_num_connector_matched_tokens == 0 + + manager.commit_new_matched_tokens(req, 16, True) + + assert manager.new_async_requests.loading_ids == {42} + assert manager.scheduler_output_manager.external_loads == {42: 16} + assert req.py_num_connector_matched_tokens == 16 + + if mpi_rank() == 0: + assert scheduler.get_num_new_matched_tokens.call_count == 1 + + run_across_mpi(mpi_pool_executor, test, 2) + + +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_connector_manager_cancel_load_reaches_the_leader(mpi_pool_executor): + + def test(): + worker = MagicMock() + scheduler = MagicMock() if mpi_rank() == 0 else None + + manager = KvCacheConnectorManager(worker, scheduler=scheduler) + + req = MagicMock() + req.request_id = 42 + + manager.cancel_load(req, 0, 32) + + if mpi_rank() == 0: + assert scheduler.cancel_load.call_args[0] == (req, 0, 32) + + run_across_mpi(mpi_pool_executor, test, 2) + + +def test_cancel_load_is_additive(): + """Existing connectors predate `cancel_load` and must keep working. + + It is only ever raised by KVCacheManagerV2, so it has to stay optional with + a no-op default rather than becoming another abstract method. + """ + assert "cancel_load" not in KvCacheConnectorScheduler.__abstractmethods__ + + def test_scheduler_output_num_scheduled_tokens_with_mtp(): """Test that num_scheduled_tokens is correctly set for MTP (multi-token prediction).""" NUM_DRAFT_TOKENS = 3 diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index eaca868014cf..2795ef9e1b3d 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -422,7 +422,7 @@ def _build_transceiver_for_kv_slice( prompt_len=prompt_len, py_request_id=0, py_beam_width=beam_width, - is_generation_only_request=lambda: is_generation_only, + is_generation_only_request=is_generation_only, ) return transceiver, req