diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index b6dcc2d963c4..b45d59f440ac 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -12,6 +12,7 @@ from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.kimi_k3_cache_policy import get_kimi_k3_bf16_kv_layer_ids from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import Tensor, replicated from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, prefer_pinned from tensorrt_llm.logger import logger @@ -214,8 +215,13 @@ def install_manager_capabilities(manager) -> None: "get_fp4_mla_hp_pool", "get_disagg_role_mapper_kinds", "get_disagg_ignored_roles", + "_page_buffer_tensors", + "_page_untransferred_roles", + "_page_rebuilds_untransferred", "rebuild_disagg_imported_cache", + "_model_layers_and_tails", "get_disagg_global_layer_ids", + "_page_layer_ids", "_get_runtime_cache_size_layer_components", "_get_generation_request_capacity", ) @@ -945,6 +951,24 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: # worker rebuilds them locally in rebuild_disagg_imported_cache. return frozenset((Role.MLA_V_SCALE, Role.MLA_V_PACKED)) + def _page_buffer_tensors( + self, layer_id: int, role: str, size: int, tokens: int, state: bool + ) -> tuple[Tensor, ...]: + """Latents token-major by head, as the FP4 MLA kernels write them; the high-precision + tail alike on every rank.""" + if role == Role.MLA_HP_TAIL: + return replicated(size) + return self._page_head_split(layer_id, size, tokens, state, token_major_=True) + + def _page_untransferred_roles(self) -> frozenset[str]: + """V scales and packed V derive from K and K scales; the generation worker rebuilds them + locally in rebuild_disagg_imported_cache.""" + return frozenset((str(Role.MLA_V_SCALE), str(Role.MLA_V_PACKED))) + + def _page_rebuilds_untransferred(self) -> bool: + """The untransferred V sidecars are rebuilt after a native import.""" + return True + def rebuild_disagg_imported_cache(self, requests: Iterable["LlmRequest"]) -> None: """Rebuild the derived V sidecars of requests whose prompt KV was imported.""" from .v_cache import rebuild_fp4_mla_disagg_imported_cache @@ -956,7 +980,9 @@ def rebuild_disagg_imported_cache(self, requests: Iterable["LlmRequest"]) -> Non [request.prompt_len for request in requests], ) - def get_disagg_global_layer_ids(self, layer_group_id: int) -> list[int]: + def _model_layers_and_tails(self, layer_group_id: int) -> list[tuple[int, bool]]: + """Per internal layer of the layer group, its model layer and whether it holds that + layer's high-precision tail.""" local_layer_ids = list(self.impl.layer_grouping[layer_group_id]) cache_local = { int(layer_id): self._fp4_mla_compact_to_local[compact_layer] @@ -973,13 +999,30 @@ def get_disagg_global_layer_ids(self, layer_group_id: int) -> list[int]: for layer_id in local_layer_ids: internal_layer = int(layer_id) if internal_layer in cache_local: - result.append(2 * int(self.pp_layers[cache_local[internal_layer]])) + result.append((int(self.pp_layers[cache_local[internal_layer]]), False)) elif internal_layer in hp_local: - result.append(2 * int(self.pp_layers[hp_local[internal_layer]]) + 1) + result.append((int(self.pp_layers[hp_local[internal_layer]]), True)) else: raise ValueError(f"Unknown FP4 MLA V2 internal layer {internal_layer}.") return result + def get_disagg_global_layer_ids(self, layer_group_id: int) -> list[int]: + return [ + 2 * model_layer + int(tail) + for model_layer, tail in self._model_layers_and_tails(layer_group_id) + ] + + def _page_layer_ids(self) -> dict[int, tuple[int, int]]: + """A model layer's cache and its high-precision tail are two internal layers: global + layers ``2 * m`` and ``2 * m + 1``, both of model layer ``m``.""" + layer_ids = {} + for layer_group_id, layers in enumerate(self.impl.layer_grouping): + for layer, (model_layer, tail) in zip( + layers, self._model_layers_and_tails(layer_group_id) + ): + layer_ids[int(layer)] = (2 * model_layer + int(tail), model_layer) + return layer_ids + def _get_runtime_cache_size_layer_components(self) -> tuple[list[int], list[Optional[int]]]: fp4_local_layers = self._fp4_mla_local_layer_indices() bf16_local_layers = self._bf16_mla_local_layer_indices() diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py index 7d87e291667f..58ab590b6519 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py @@ -25,6 +25,8 @@ _estimate_swa_cache_size, _fill_kv_pages, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import Tensor, replicated +from tensorrt_llm._torch.pyexecutor.kv_cache.page_table import virtual_layers from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, prefer_pinned from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType @@ -1103,6 +1105,16 @@ def get_disagg_role_mapper_kinds(self): **{role.role: MapperKind.REPLICATED for role in CSA2CacheRole}, } + def _page_buffer_tensors( + self, layer_id: int, role: str, size: int, tokens: int, state: bool + ) -> tuple[Tensor, ...]: + """Every CSA2 cache holds the same bytes on every rank.""" + return replicated(size) + + def _virtual_layers(self): + """Each model layer keeps a cache per CSA2 attention type.""" + return virtual_layers(self._layer_attn_to_layer_id) + def get_disagg_pool_view_partition(self, layer_id: LayerId, role: DataRole): """Keep remote-tail SWA independently selectable in the transfer layout.""" if not self.bounded_replay_on_generation or role != CSA2CacheRole.SWA.role: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index 1d97a1f2f1cd..aeac2d6fefb4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -29,6 +29,8 @@ KVCacheManagerV2, _fill_kv_pages, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import Tensor +from tensorrt_llm._torch.pyexecutor.kv_cache.page_table import virtual_layers from tensorrt_llm._torch.pyexecutor.resource_manager import scan_kv_buffer_for_invalid_values from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import ( @@ -81,6 +83,7 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata COMPRESS_BLOCK_SCALE_ROLE = DataRole("deepseek_v4_compress_block_scale") +_ATTENTION_TYPE_ROLES = frozenset(str(t.role) for t in DeepseekV4AttentionType) KV_CACHE_COPY_ALIGNMENT = 16 NVFP4_VECTOR_SIZE = 16 @@ -1482,6 +1485,28 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """Exclude storage-only NVFP4 alignment pages from KV transfer.""" return self._disagg_ignored_roles + def _page_buffer_tensors( + self, layer_id: int, role: str, size: int, tokens: int, state: bool + ) -> Tuple[Tensor, ...]: + """The cache of each attention type and the compressed cache's block scales hold MLA + latents of one KV head, which every tensor rank holds: split by head as the base manager + splits K/V. The padding pages never move (``_page_untransferred_roles``).""" + if role == COMPRESS_BLOCK_SCALE_ROLE or role in _ATTENTION_TYPE_ROLES: + return self._page_head_split(layer_id, size, tokens, state) + return super()._page_buffer_tensors(layer_id, role, size, tokens, state) + + def _virtual_layers(self): + """Each model layer keeps a cache per DeepSeek-V4 attention type.""" + return virtual_layers(self._layer_attn_to_layer_id) + + def _page_untransferred_roles(self) -> frozenset[str]: + """Storage-only NVFP4 alignment pages, which hold no state a peer needs.""" + return frozenset(str(role) for role in self._disagg_ignored_roles) + + def _page_quant_mode(self) -> Optional[str]: + """``"fp8_ds_mla"`` for FP8 pages carrying their own scales.""" + return "fp8_ds_mla" if self.use_fp8_ds_mla else super()._page_quant_mode() + def get_indexer_k_cache_buffers(self, layer_idx: int) -> torch.Tensor: """ Get the buffers for the indexer k cache for a specific layer. diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py index d788e5ea7bdc..cdaf2614d3d6 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py @@ -141,6 +141,17 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: roles |= self._kda_replay.disagg_ignored_roles(self._state_layout) return roles + @override + def _page_untransferred_roles(self) -> frozenset[str]: + """The KDA replay buffers, which a generation server seeds from the transferred conv + state instead of receiving.""" + roles = super()._page_untransferred_roles() + if self._kda_replay is not None: + roles |= { + str(role) for role in self._kda_replay.disagg_ignored_roles(self._state_layout) + } + return roles + @override def _setup_model_state(self) -> None: if self._kda_replay is not None: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index b0090bdcee48..d8ec28f27ec1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -38,6 +38,7 @@ Role, _BasePageTableMaterializer, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import Tensor, replicated from tensorrt_llm._utils import ( TensorWrapper, binding_to_torch_dtype, @@ -718,6 +719,17 @@ def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: Role.INDEX_KEY: MapperKind.REPLICATED, } + def _page_buffer_tensors( + self, layer_id: int, role: str, size: int, tokens: int, state: bool + ) -> tuple[Tensor, ...]: + """K/V and their block scales split by head as the backend writes them, head-major (MSA) + or token-major (Triton); index-K alike on every rank.""" + if role == Role.INDEX_KEY: + return replicated(size) + return self._page_head_split( + layer_id, size, tokens, state, token_major_=self._main_kv_layout_name() == "NHD" + ) + def _main_kv_layout_name(self) -> str: """Return the tensor-view layout name for the selected mapper kind.""" return "HND" if self._main_kv_mapper_kind == MapperKind.HND else "NHD" diff --git a/tensorrt_llm/_torch/disaggregation/base/backend.py b/tensorrt_llm/_torch/disaggregation/base/backend.py index c583ab585d95..4ba5b1222f26 100644 --- a/tensorrt_llm/_torch/disaggregation/base/backend.py +++ b/tensorrt_llm/_torch/disaggregation/base/backend.py @@ -98,8 +98,8 @@ class CacheKind(IntEnum): @runtime_checkable class SlotRegion(Protocol): - """A contiguous memory region divided into equal-sized slots; a row given by an address - occupies one slot. + """A contiguous memory region divided into equal-sized slots; a row lying in the region occupies + one of its slots. A backend registers only the regions whose :attr:`memory_kind` it can access, each once as ``(address, nbytes)``; rows in a region it has not registered are rows it does not recognize. @@ -141,10 +141,11 @@ def memory_kind(self) -> MemoryKind: class LayerGroupRows(Protocol): """One layer group's rows of a view, described by aligned arrays. - A row is one transfer unit. When addresses are provided, each row occupies one slot in the - region. These members are all a backend reads of a group, apart from those a capability it - declares names. Row addresses need not be consecutive. A backend given rows it does not - recognize rejects the submission with :class:`InvalidSubmission`; it never skips them. + A row is one transfer unit. When slots are provided, row ``i`` occupies slot ``slots[i]`` in + every region :attr:`region_indices` lists. These members are all a backend reads of a group, + apart from those a capability it declares names. Slots need not be consecutive. A backend given + rows it does not recognize rejects the submission with :class:`InvalidSubmission`; it never + skips them. """ @property @@ -152,6 +153,13 @@ def layer_group_index(self) -> int: """The lender's local layer-group index, not the group's position in :attr:`View.groups`.""" ... + @property + def page_ordinals(self) -> np.ndarray: + """One page ordinal per row, ``int64`` of shape ``[rows]``: the request's page the row + holds, as the lender's published page layout numbers pages. Ordinals need not be + consecutive.""" + ... + @property def object_keys(self) -> np.ndarray | None: """One opaque object key per row, ``uint8`` of shape ``[rows, width]`` with the lender's @@ -161,16 +169,16 @@ def object_keys(self) -> np.ndarray | None: ... @property - def addresses(self) -> np.ndarray | None: - """One address per row, ``int64`` of shape ``[rows]``: row ``i`` is the slot at - ``addresses[i]`` in the :class:`SlotRegion` at index :attr:`region_index`. ``None`` exactly - when :attr:`region_index` is.""" + def slots(self) -> np.ndarray | None: + """One slot per row, ``int64`` of shape ``[rows]``: row ``i`` occupies slot ``slots[i]`` in + every :class:`SlotRegion` :attr:`region_indices` lists, its bytes those slots back to back + in the listed order. ``None`` exactly when :attr:`region_indices` is.""" ... @property - def region_index(self) -> int | None: - """The index, among the lender's regions, of the :class:`SlotRegion` the rows sit in, or - ``None``.""" + def region_indices(self) -> tuple[int, ...] | None: + """The indices, among the lender's regions, of the :class:`SlotRegion` objects the rows lie + in, at least one and distinct, in the order a row's bytes run; or ``None``.""" ... def __len__(self) -> int: ... diff --git a/tensorrt_llm/_torch/disaggregation/blob/backend.py b/tensorrt_llm/_torch/disaggregation/blob/backend.py index b985d207a960..e32d92b47119 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/backend.py +++ b/tensorrt_llm/_torch/disaggregation/blob/backend.py @@ -356,8 +356,8 @@ class _GroupRows(NamedTuple): """One group's arrays, taken from the view at submission.""" object_keys: np.ndarray # uint8 [rows, width] - addresses: np.ndarray # int64 [rows] - region_index: int + addresses: np.ndarray # int64 [rows, regions] + region_indices: tuple[int, ...] @dataclass(eq=False) @@ -534,6 +534,7 @@ def __init__(self, binding: RegionBinding, wake: WakeFn, config: BlobConfig) -> self._name = name = f"blob-backend-{next(_instance_ids)}" driver_capabilities = self._driver.capabilities self._max_rows_per_op = driver_capabilities.max_rows_per_op + self._driver_multi_span = driver_capabilities.multi_span self._driver_reads_hints = driver_capabilities.hints self._driver_cancels = driver_capabilities.cancel # Every call goes to the driver as it is submitted, so the backend can cancel only by @@ -651,7 +652,7 @@ def _submit(self, view: View, hint: Hint, is_fetch: bool) -> BlobAttempt: raise InvalidSubmission(f"the store driver cannot {'get' if is_fetch else 'put'}") groups = [self._checked_group_rows(group) for group in view.groups] if is_fetch: - _refuse_shared_destinations(groups) + _refuse_shared_destinations(groups, self._regions) label = f"attempt:{next(self._attempt_ids)}" attempt = BlobAttempt(self, label, is_fetch, groups, self._call_ranges(groups), hint) batches = [self._row_batch(groups, call) for call in attempt._calls] @@ -671,10 +672,10 @@ def _submit(self, view: View, hint: Hint, is_fetch: bool) -> BlobAttempt: def _checked_group_rows(self, group: LayerGroupRows) -> _GroupRows: where = f"layer group {group.layer_group_index}" - if group.object_keys is None or group.addresses is None or group.region_index is None: - raise InvalidSubmission(f"{where}: rows need object keys, addresses and a region") + if group.object_keys is None or group.slots is None or group.region_indices is None: + raise InvalidSubmission(f"{where}: rows need object keys, slots and regions") keys = np.asarray(group.object_keys) - addresses = np.asarray(group.addresses) + slots = np.asarray(group.slots) rows = len(group) if keys.dtype != np.uint8 or keys.ndim != 2 or keys.shape[0] != rows or keys.shape[1] == 0: raise InvalidSubmission( @@ -682,21 +683,44 @@ def _checked_group_rows(self, group: LayerGroupRows) -> _GroupRows: f"{keys.shape}" ) keys = np.ascontiguousarray(keys) # one key twice is allowed: _call_bounds cuts there - if addresses.dtype != np.int64 or addresses.shape != (rows,): + if slots.dtype != np.int64 or slots.shape != (rows,): raise InvalidSubmission( - f"{where}: addresses must be int64 [{rows}], got {addresses.dtype} " - f"{addresses.shape}" + f"{where}: slots must be int64 [{rows}], got {slots.dtype} {slots.shape}" ) + indices = group.region_indices + if ( + not isinstance(indices, tuple) + or not indices + or any(type(index) is not int for index in indices) + or len(set(indices)) != len(indices) + ): + raise InvalidSubmission( + f"{where}: region indices must be a nonempty tuple of distinct ints" + ) + if len(indices) > 1 and not self._driver_multi_span: + raise InvalidSubmission(f"{where}: the store driver cannot transfer multiple spans") # Refused, never skipped: the driver must not be handed memory it never registered. - if group.region_index not in self._binding.served_region_indices: - raise InvalidSubmission(f"{where}: region {group.region_index} is not registered") - region = self._regions[group.region_index] - # Known limit: numpy lets go of the GIL for arrays of more than about 500 elements, so a - # group that large may wait here behind a busy thread. - slot, misalignment = np.divmod(addresses - region.address, region.slot_bytes) - if np.any((slot < 0) | (slot >= region.slot_count) | (misalignment != 0)): - raise InvalidSubmission(f"{where}: an address is not a slot of its region") - return _GroupRows(keys, addresses, group.region_index) + regions: list[SlotRegion] = [] + for index in indices: + if index not in self._binding.served_region_indices: + raise InvalidSubmission(f"{where}: region {index} is not registered") + region = self._regions[index] + if np.any((slots < 0) | (slots >= region.slot_count)): + raise InvalidSubmission(f"{where}: a slot is outside region {index}") + if region.address + region.slot_bytes * region.slot_count > np.iinfo(np.int64).max: + raise InvalidSubmission(f"{where}: region {index} exceeds int64 addresses") + regions.append(region) + addresses = np.column_stack( + [region.address + slots * region.slot_bytes for region in regions] + ) + for column, region in enumerate(regions): + for previous, other in enumerate(regions[:column]): + if np.any( + (addresses[:, column] < addresses[:, previous] + other.slot_bytes) + & (addresses[:, previous] < addresses[:, column] + region.slot_bytes) + ): + raise InvalidSubmission(f"{where}: a row's spans overlap") + return _GroupRows(keys, addresses, indices) def _call_ranges(self, groups: list[_GroupRows]) -> list[tuple[int, int, int]]: """``(group position, start, stop)`` of each call: one per group, split where the driver @@ -713,13 +737,13 @@ def _call_ranges(self, groups: list[_GroupRows]) -> list[tuple[int, int, int]]: def _row_batch(self, groups: list[_GroupRows], call: _Call) -> RowBatch: group_rows = groups[call.group_position] rows = slice(call.start, call.stop) - region_index = group_rows.region_index - return RowBatch.of_slots( + indices = group_rows.region_indices + return RowBatch.of_spans( self._key_stem, group_rows.object_keys[rows], group_rows.addresses[rows], - self._regions[region_index].slot_bytes, - self._region_extents[region_index], + tuple(self._regions[index].slot_bytes for index in indices), + tuple(self._region_extents[index] for index in indices), ) def _hand_over(self, attempt: BlobAttempt, batches: list[RowBatch]) -> None: @@ -1059,12 +1083,15 @@ def _call_bounds(keys: np.ndarray, step: int) -> list[tuple[int, int]]: return bounds -def _refuse_shared_destinations(groups: list[_GroupRows]) -> None: +def _refuse_shared_destinations(groups: list[_GroupRows], regions: Sequence[SlotRegion]) -> None: """Two fetched rows in one slot would let one report READY over the other's bytes. A publish - only reads its slots, so only a fetch is checked. A Python set, not np.unique, whose sort lets - go of the GIL.""" - seen: set[int] = set() - for rows in groups: - seen.update(rows.addresses.tolist()) - if len(seen) != sum(len(rows.addresses) for rows in groups): + only reads its slots, so only a fetch is checked. Spans are checked across all regions and + groups before any driver call.""" + spans = sorted( + (address, address + regions[index].slot_bytes) + for group in groups + for column, index in enumerate(group.region_indices) + for address in group.addresses[:, column].tolist() + ) + if any(start < previous_end for (_, previous_end), (start, _) in zip(spans, spans[1:])): raise InvalidSubmission("two rows of a fetch write the same slot") diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index 3f5a50091a43..62587798d794 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -229,7 +229,8 @@ class RowBatch(Sequence[RowIO]): Its maker guarantees what a driver would otherwise check row by row: at least one row, no key twice, spans of positive size inside their extent that do not overlap within a row, and each extent inside one registered region. :meth:`of` checks all but the last, which only the driver - can; :meth:`of_slots`, for a maker that checked its rows already, checks nothing. + can; :meth:`of_slots` and :meth:`of_spans`, for a maker that checked its rows already, check + nothing. """ def __init__( @@ -268,6 +269,26 @@ def of_slots( column.flags.writeable = False return cls(_HexKeys(key_stem, object_keys), column, (slot_bytes,), (extent,)) + @classmethod + def of_spans( + cls, + key_stem: str, + object_keys: np.ndarray, + addresses: np.ndarray, + span_bytes: tuple[int, ...], + extents: tuple[Span, ...], + ) -> RowBatch: + """Already checked rows with one address column per span, in payload order. + + ``object_keys`` is ``uint8 [rows, width > 0]`` with no key twice; ``addresses`` is + ``int64 [rows, spans]``. Column ``j`` holds spans of ``span_bytes[j]`` inside + ``extents[j]``. Checks nothing and copies neither array; the driver's address view is + read-only, while the caller's array stays writable. + """ + columns = addresses.view() + columns.flags.writeable = False + return cls(_HexKeys(key_stem, object_keys), columns, span_bytes, extents) + @classmethod def of(cls, rows: Sequence[RowIO]) -> RowBatch: """The batch of ``rows``, built row by row and checked. Span ``j``'s extent is the diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 1254c06001f6..64ccf5699ae8 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -16,12 +16,13 @@ from __future__ import annotations from dataclasses import dataclass, field -from enum import IntEnum -from typing import FrozenSet, List, Optional, Tuple +from typing import FrozenSet, List, Optional import numpy as np from tensorrt_llm._torch.disaggregation.base import CacheKind +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import MapperKind +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import RoleLayout as RoleLayout BUFFER_ENTRY_DTYPE = np.dtype( [ @@ -33,69 +34,6 @@ ) -class MapperKind(IntEnum): - """Transfer semantics of one physical pool's bytes. - - Every PoolView carries ``buffer_entries`` listing - ``(local_layer_id, offset, size)`` per buffer; the view's exact layer - set always comes from those entries (a view may cover a subset of the - LG when V2 splits an LG into multiple pools by buffer-size class, or - when a role class exists only on some layers). The kind selects how - bytes move between heterogeneous topologies: - - INDEXED/HND: Head-major (HND) K/V — the layout written by the TRTLLM - attention kernels and the default for standard V2 managers. - Heterogeneous-head transfer selects one contiguous head-major range - per K/V buffer. - REPLICATED: The pool holds bytes that are identical on every TP rank - (MiniMax M3 index-key, DSA indexer K). Copied without KV-head - remapping using per-layer strides; fan-in routing elects one owning - sender per destination so each peer receives exactly one copy. - NHD: Ordinary K/V whose per-buffer storage is token-major - ``[token, head, dim]``. Heterogeneous-head transfer must select the - corresponding head slice inside every token rather than a single - contiguous head-major range. - - A physical pool may hold roles of different kinds: V2 storage coalesces - buffers purely by ``(life_cycle, buffer size)``, so e.g. MiniMax M3's - replicated index-K shares the K/V pool at TP degrees where their - per-block sizes coincide. The page-table builder therefore emits one - PoolView per ``(physical pool, mapper kind)`` — a view covers exactly - the bytes of one role class, and its per-layer byte ranges come from - ``buffer_entries`` (offsets are per layer because another class may - interleave between layers; only the per-layer size is uniform, recorded - in ``bytes_per_layer``). View count per layer group is bounded by the - number of role classes, never by layer count. - - Mamba state pools use SECTIONED for convolution state, INDEXED for SSM - state, and REPLICATED for auxiliary recurrent state. The policy is chosen - by the layer group's :class:`CacheKind` (PAGED -> ``AttentionPolicy``, - STATE -> ``MambaPolicy``); each policy then picks the mapper for a view - from its mapper kind, copying REPLICATED views whole per layer. - """ - - INDEXED = 0 - HND = INDEXED - REPLICATED = 1 - NHD = 2 - SECTIONED = 3 # Sectioned layout: [Sec0|Sec1|...], each section independently TP-sharded - - -@dataclass(frozen=True) -class RoleLayout: - """Resharding geometry a cache manager declares for one role. - - Carried onto the role's :class:`PoolView` by the page-table builder. - ``section_bytes`` describes a SECTIONED role (one entry per section, in - slot order); ``bytes_per_head`` describes an INDEXED role whose head - count is not derivable from ``RankInfo.attention`` (recurrent SSM state). - Roles with neither need no geometry (REPLICATED, attention K/V). - """ - - section_bytes: Optional[Tuple[int, ...]] = None - bytes_per_head: Optional[int] = None - - @dataclass class PhysicalPool: """Physical storage for a V2 pool's slots. diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index a927aa4bafb5..0001a32702cc 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -977,6 +977,17 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: roles |= self._kda_replay.disagg_ignored_roles(self._state_layout) return roles + @override + def _page_untransferred_roles(self) -> frozenset[str]: + """The KDA replay buffers, which a generation server seeds from the transferred conv + state instead of receiving.""" + roles = super()._page_untransferred_roles() + if self._kda_replay is not None: + roles |= { + str(role) for role in self._kda_replay.disagg_ignored_roles(self._state_layout) + } + return roles + @override def _setup_model_state(self) -> None: if self._kda_replay is not None: @@ -1002,6 +1013,8 @@ def suspend_request(self, request) -> None: replay = self._kda_replay request_id = request.py_request_id cache = self.kv_cache_map.get(request_id) + if cache is not None and cache.is_active: + self._check_suspendable(cache) slot = self._request_id_to_state_index.get(request_id, -1) if replay is not None and slot >= 0 and cache is not None and cache.is_active: if request_id in replay.acceptance_pending: diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py index f304c6d5c16e..1fd09616db7b 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -38,6 +38,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( _get_local_mamba_cache_layout, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import Tensor, replicated from tensorrt_llm._torch.pyexecutor.resource_manager import get_pp_layers from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor from tensorrt_llm.bindings import DataType @@ -191,6 +192,17 @@ def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: QSA_SIDE_PLANES: MapperKind.REPLICATED, } + @override + def _page_buffer_tensors( + self, layer_id: int, role: str, size: int, tokens: int, state: bool + ) -> tuple[Tensor, ...]: + """PLE state is computed from replicated inputs, and QSA index positions are request-wide + coordinates, packed with index-K into the side planes when they would share a pool: every + rank holds the same bytes of each, with no head axis to split.""" + if role in (PLE_CONV_STATE, PLE_NGRAM_CONTEXT, QSA_INDEX_POSITION, QSA_SIDE_PLANES): + return replicated(size) + return super()._page_buffer_tensors(layer_id, role, size, tokens, state) + @override def __init__( self, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/_lender_hooks.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/_lender_hooks.py new file mode 100644 index 000000000000..824704a762fa --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/_lender_hooks.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The only way managers, the scheduler and the executor call a lender (``_ManagerHooks``), and the +writers of their lending records; a lender reaches a manager only via ``sharing/_manager.py``.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping, Protocol, Sequence, runtime_checkable + +import numpy as np + +from tensorrt_llm.logger import logger + +if TYPE_CHECKING: + from tensorrt_llm.runtime.kv_cache_manager_v2 import _KVCache + + from ..llm_request import LlmRequest + from ..scheduler.scheduler_v2 import KVCacheV2Scheduler + from .kv_cache_manager_v2 import KVCacheManagerV2 + + +class _PageLockState(Protocol): + """Live lock state a batch reads per request without another call: ``holding`` counts the + unreleased locks holding a block; a cache whose ``id`` is not in ``locked_caches`` has no lock + holding a block of it or capping its history; ``caps_history`` says whether any lock can cap a + history; ``unreusable`` caches, by ``id``, commit nothing more for reuse.""" + + holding: int + locked_caches: Mapping[int, object] + caps_history: bool + unreusable: Mapping[int, object] + + +# What ``commit_gate`` answers: commit; hold the commit back for a retry at the unlock; or commit +# nothing more for reuse, stopping as usual. +COMMIT, HOLD, REFUSE = 0, 1, 2 + + +@runtime_checkable +class _ManagerHooks(Protocol): + """Lender callbacks on the manager thread for cache lifetime, page locks and history limits. + + Call order. Free, the cache out of the map: ``_on_free``, then ``_defer_close`` instead of a + close. In-place shrink: ``_on_shrink`` right after. Early index-slot release: ``_parking_row``, + the buffers for the cache's beam-0 page indices from then on. Right after a suspended cache + resumed: ``_parked_row`` for those buffers if its slot was released early, then ``_on_resume``. + Reuse reset: ``_before_reset``, the reuse tree's clearing, ``_on_reset``. Shutdown: + ``_on_shutdown``, ``_defer_close`` per mapped cache, ``_keeps_pools``, ``_on_caches_closed`` + once the runtime's shutdown returned. ``_is_locked`` and ``_history_cap`` answer before a + suspension, a commit or a history move, unless ``_page_lock_state``, live state a batch reads + per request without a call, shows no lock on the cache or none capping a history. + + The ``_on_*`` notifications raise nothing: they log their errors, some of which stop the lender + for good. ``_before_reset`` raises ``RuntimeError`` before any reset takes place. The queries + raise into the manager whatever goes wrong.""" + + def _on_free(self, request_id: int, kv_cache: _KVCache) -> None: ... + + def _on_shrink(self, request_id: int, kv_cache: _KVCache) -> None: ... + + def _on_resume(self, request_id: int) -> None: ... + + def _parking_row(self, request_id: int, kv_cache: _KVCache) -> Sequence[memoryview] | None: ... + + def _parked_row(self, request_id: int) -> Sequence[memoryview] | None: ... + + def _before_reset(self) -> None: ... + + def _on_reset(self) -> None: ... + + def _on_shutdown(self) -> None: ... + + def _on_caches_closed(self) -> None: ... + + def _is_locked( + self, kv_cache: _KVCache | None, begin: int = 0, end: int | None = None + ) -> bool: ... + + def _defer_close(self, kv_cache: _KVCache) -> bool: ... + + def _keeps_pools(self) -> bool: ... + + def _history_cap(self, kv_cache: _KVCache) -> int | None: ... + + def _page_lock_state(self) -> _PageLockState: ... + + +def generation_lock_state(manager: KVCacheManagerV2) -> _PageLockState | None: + """The lender's live lock state, or None without a lender; queries propagate errors.""" + return None if manager._sharing is None else manager._sharing._page_lock_state() + + +def history_lock_state(manager: KVCacheManagerV2) -> _PageLockState | None: + """The lender's live lock state where a lock can cap a cache's history, else None, when + ``lent_history`` caps nothing; queries propagate errors.""" + if manager._sharing is None: + return None + state = manager._sharing._page_lock_state() + return state if state.caps_history else None + + +def on_create(manager: KVCacheManagerV2, request_id: int) -> None: + """Discard commits held for an earlier cache allocation of this request.""" + if manager._sharing is None or not manager._held_commits: + return + for key, held in tuple(manager._held_commits.items()): + if held[1].py_request_id == request_id: + del manager._held_commits[key] + + +def on_free(manager: KVCacheManagerV2, request_id: int, kv_cache: _KVCache) -> None: + """Tell the lender, if any, that the request's cache left the map.""" + if manager._sharing is not None: + manager._sharing._on_free(request_id, kv_cache) + + +def defer_close(manager: KVCacheManagerV2, kv_cache: _KVCache) -> bool: + """Whether the lender keeps the cache open, to close it at its last unlock.""" + return manager._sharing is not None and manager._sharing._defer_close(kv_cache) + + +def on_shrink(manager: KVCacheManagerV2, request_id: int, kv_cache: _KVCache) -> None: + """Tell the lender, if any, that the request's cache may have shrunk in place.""" + if manager._sharing is not None: + manager._sharing._on_shrink(request_id, kv_cache) + + +def on_resume(manager: KVCacheManagerV2, request_id: int) -> None: + """Tell the lender, if any, that the request's cache resumed.""" + if manager._sharing is not None: + manager._sharing._on_resume(request_id) + + +def parking_row( + manager: KVCacheManagerV2, request_id: int, kv_cache: _KVCache +) -> Sequence[memoryview] | None: + """Lender-owned buffers, one per pool, for the cache's beam-0 page indices after the request's + index slot's early release; None if no lender gives them, the cache keeping them itself.""" + if manager._sharing is None: + return None + return manager._sharing._parking_row(request_id, kv_cache) + + +def parked_row(manager: KVCacheManagerV2, request_id: int) -> Sequence[memoryview] | None: + """The buffers ``parking_row`` gave for the request, to bind again after its cache resumed; + None without them.""" + return None if manager._sharing is None else manager._sharing._parked_row(request_id) + + +def before_reset(manager: KVCacheManagerV2) -> None: + """Let the lender, if any, refuse a reset of the manager's reuse state: ``RuntimeError``.""" + if manager._sharing is not None: + manager._sharing._before_reset() + + +def on_reset(manager: KVCacheManagerV2) -> None: + """Tell the lender, if any, that the manager cleared its reuse tree.""" + if manager._sharing is not None: + manager._sharing._on_reset() + + +def on_shutdown(manager: KVCacheManagerV2) -> None: + """Tell the lender, if any, that the manager shuts down.""" + if manager._sharing is not None: + manager._sharing._on_shutdown() + + +def keeps_pools(manager: KVCacheManagerV2) -> bool: + """Whether the lender keeps the pools allocated past the manager's shutdown.""" + return manager._sharing is not None and manager._sharing._keeps_pools() + + +def on_caches_closed(manager: KVCacheManagerV2) -> None: + """Notify the lender, if any, after runtime shutdown returns.""" + if manager._sharing is not None: + manager._sharing._on_caches_closed() + + +def check_suspendable(manager: KVCacheManagerV2, kv_cache: _KVCache) -> None: + """Refuse suspension while a lender locks the cache.""" + if manager._sharing is not None and manager._sharing._is_locked(kv_cache): + raise RuntimeError( + f"the KV cache of request {kv_cache.id} cannot be suspended while a lender locks " + "its blocks" + ) + + +def commit_gate( + manager: KVCacheManagerV2, kv_cache: _KVCache, request: LlmRequest, end: int +) -> int: + """``REFUSE`` for a cache the lender keeps out of reuse; ``HOLD``, registering a retry, where + locks prevent committing or advancing history through ``end``; else ``COMMIT``. While neither a + lock nor a refusal names the cache, the lender is asked nothing.""" + if manager._sharing is None: + return COMMIT + state = manager._sharing._page_lock_state() + key = id(kv_cache) + if state.unreusable and key in state.unreusable: + return REFUSE + if not state.holding or key not in state.locked_caches: + return COMMIT + tokens_per_block = manager.tokens_per_block + begin = kv_cache.num_committed_tokens // tokens_per_block + locked = manager._sharing._is_locked(kv_cache, begin, -(-end // tokens_per_block)) + if not locked and manager._lent_history(kv_cache, end) >= end: + return COMMIT + held = manager._held_commits.get(key) + if held is not None and held[1] is request: + end = held[2] + fetch_touched_suffix = held[3] + else: + fetch_touched_suffix = False + manager._held_commits[key] = (kv_cache, request, end, fetch_touched_suffix) + logger.warning( + f"KVCacheManagerV2: request {kv_cache.id}: commit through token {end} skipped while " + "a lender locks blocks it covers or its window would pass" + ) + return HOLD + + +def retry_held_commit(manager: KVCacheManagerV2, kv_cache: _KVCache) -> None: + """Retry after unlock, before deferred close, unless the request now uses another cache.""" + held = manager._held_commits.pop(id(kv_cache), None) + if held is None: + return + request, end, fetch_touched_suffix = held[1:] + current = manager.kv_cache_map.get(request.py_request_id) + if current is None or current is kv_cache: + stopped = manager._retry_held_commit(request, kv_cache, end, fetch_touched_suffix) + retried = manager._held_commits.get(id(kv_cache)) + if fetch_touched_suffix and retried is not None and retried[1] is request: + manager._held_commits[id(kv_cache)] = (*retried[:3], True) + if stopped and current is kv_cache and manager.conversation_manager is not None: + manager.conversation_manager.save_drop_plan(request, kv_cache) + + +def drop_held_commit(manager: KVCacheManagerV2, kv_cache: _KVCache, start: int) -> None: + """Drop the cache's held-back commit when it covers computed tokens from ``start`` on, which a + fetch from ``start`` may overwrite; the request then commits from the bounds it resumes at.""" + held = manager._held_commits.get(id(kv_cache)) + if held is None: + return + end = manager._held_commit_drop_end(held) + if end <= start: + manager._held_commits[id(kv_cache)] = (held[0], held[1], held[2], True) + return + del manager._held_commits[id(kv_cache)] + logger.warning( + f"KVCacheManagerV2: request {kv_cache.id}: commit through token " + f"{end} dropped, as a fetch from token {start} may " + "overwrite what it covers" + ) + + +def lent_history(manager: KVCacheManagerV2, kv_cache: _KVCache, history: int) -> int: + """Cap history at locked pages without rewinding the current history; while no lock holds a + block of the cache, or no lock can cap a history, nothing caps it.""" + if manager._sharing is None: + return history + state = manager._sharing._page_lock_state() + if not state.caps_history or not state.holding or id(kv_cache) not in state.locked_caches: + return history + cap = manager._sharing._history_cap(kv_cache) + return history if cap is None else max(min(history, cap), kv_cache.history_length) + + +def holds_page_locks(manager: KVCacheManagerV2, request_id: int | None) -> bool: + """Whether the request, or any cache when request_id is None, holds lender locks.""" + if manager._sharing is None: + return False + if request_id is None: + return manager._sharing._is_locked(None) + kv_cache = manager.kv_cache_map.get(request_id) + return kv_cache is not None and manager._sharing._is_locked(kv_cache) + + +def scheduler_holds_page_locks(scheduler: KVCacheV2Scheduler, req: LlmRequest | None) -> bool: + """Check both target and draft managers before reclaiming pages.""" + from .kv_cache_manager_v2 import KVCacheManagerV2 + + rid = None if req is None else req.py_request_id + if scheduler.kv_cache_manager._holds_page_locks(rid): + return True + draft = scheduler.draft_kv_cache_manager + return isinstance(draft, KVCacheManagerV2) and draft._holds_page_locks(rid) + + +def fresh_page_record( + manager: KVCacheManagerV2, kv_cache: _KVCache +) -> dict[int, np.ndarray] | None: + """The cache's base page of each block, per pool, as the fresh-page fill records the pages it + gave; ``None`` with the fill off.""" + if manager._fresh_page_fill is None: + return None + count = int(kv_cache.num_blocks) + return { + pool: np.frombuffer( + kv_cache.get_base_page_indices(pool), dtype=np.int32, count=count + ).copy() + for pool in range(int(manager.num_pools)) + } + + +def fill_fresh_pages_past( + manager: KVCacheManagerV2, request_id: int, record: dict[int, np.ndarray] | None +) -> None: + """Run the fresh-page fill on the request's pages that ``record``, from ``fresh_page_record``, + lacks: the fill takes the recorded pages as given and leaves them alone.""" + if record is None: + return + manager._fresh_pages_filled[int(request_id)] = record + manager._fill_fresh_kv_pages(request_id) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a963738f9ff8..e2dcd5d79fde 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -137,6 +137,32 @@ scan_kv_buffer_for_invalid_values, ) from ..scheduler import ScheduledRequests +from . import _lender_hooks +from .page_layout import ( + BufferDesc, + LayerGroupDesc, + PageLayout, + ShardDesc, + Tensor, + grid_place, + head_major, + owned, + replicated, + token_major, +) +from .page_table import ( + DevicePool, + ManagerLayout, + attention_shard, + buffer_expansions, + group_shard, + has_sparse_buffers, + is_non_required_layer, + is_private_layer, + is_recurrent_layer, + layer_configs_by_id, + rank_topology, +) if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata @@ -280,6 +306,12 @@ class Role: ALL = DataRole("all") +# The roles the base manager splits by KV head (_page_buffer_tensors). +_HEAD_SPLIT_ROLES = frozenset( + str(role) for role in (Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) +) + + class BlockReusePolicy(StrEnum): ALL_REUSABLE = "all_reusable" PER_REQUEST = "per_request" @@ -1788,14 +1820,37 @@ class KVCacheManagerV2(BaseResourceManager): # symptom is a small perturbation that flips an argmax a couple of steps into # decode, which reads as ordinary sampling variation rather than as a bug. _supports_partial_reuse = True - # The lender sharing.attach_lender installs, one per manager for its life; None without one. - # On the class so it exists without running __init__. + # Whole-model capabilities include layers outside this manager's local pools. + _model_has_attention_window: bool = False + _model_has_sparse_buffers: bool = False + # One optional lender for the manager's lifetime. _sharing = None # Branch snapshots (block_reuse_config.enable_branch_snapshot); declared on # the class for instances built without running __init__. _commit_branch_snapshots = False _record_branch_snapshots = False + def _check_suspendable(self, kv_cache: _KVCache) -> None: + _lender_hooks.check_suspendable(self, kv_cache) + + def _commit_gate(self, kv_cache: _KVCache, request: LlmRequest, end: int) -> int: + return _lender_hooks.commit_gate(self, kv_cache, request, end) + + def _commit_held_back(self, kv_cache: _KVCache) -> None: + _lender_hooks.retry_held_commit(self, kv_cache) + + def _drop_held_commit(self, kv_cache: _KVCache, start: int) -> None: + _lender_hooks.drop_held_commit(self, kv_cache, start) + + def _lent_history(self, kv_cache: _KVCache, history: int) -> int: + return _lender_hooks.lent_history(self, kv_cache, history) + + def _history_lock_state(self) -> _lender_hooks._PageLockState | None: + return _lender_hooks.history_lock_state(self) + + def _holds_page_locks(self, request_id: int | None) -> bool: + return _lender_hooks.holds_page_locks(self, request_id) + def __init__( self, kv_cache_config: KvCacheConfig, @@ -1862,6 +1917,8 @@ def __init__( ) # Filled on first use; the layer grouping does not change after init. self._layer_group_life_cycles: Optional[List[AttnLifeCycle]] = None + # Global layer id -> the window the model's attention declares it reads (None: none). + self._model_attention_windows = kwargs.get("attention_read_windows", {}) self.kv_cache_type = kv_cache_type if self.layer_specs is None: @@ -2056,6 +2113,7 @@ def __init__( else layer.sliding_window_size for layer in self.layer_specs ] + self._model_attention_window_pattern = tuple(kv_cache_config.max_attention_window or ()) # A D-token match backoff needs D more physical SWA history at the # committed endpoint so commit_min_snapshot still covers the rewind. @@ -2426,6 +2484,8 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._allocation_stream: torch.cuda.Stream | None = None self._on_allocation_stream: set[int] = set() self._allocation_ready: dict[int, torch.cuda.Event] = {} + # By cache id: cache, request, first held end, and whether a fetch touched its suffix. + self._held_commits: dict[int, tuple[_KVCache, LlmRequest, int, bool]] = {} self._request_stats_enabled_ids: set[int] = set() # Lazily built map of layer-group id -> sliding window size, restricted @@ -4250,6 +4310,315 @@ def get_draft_kv_cache_view(self) -> Optional["KVCacheManagerV2"]: return None return self + def get_page_table(self) -> ManagerLayout: + """This manager's page table: its page layout, device pools and the facts of its layer + groups, which whoever lends or moves its pages reads instead of the manager. + + Built from the runtime config and storage and the hooks below: how each buffer's bytes + spread over ranks (``_page_buffer_tensors``), how layers are named (``_page_layer_ids``), + the KV heads (``_page_kv_heads``), which roles never travel + (``_page_untransferred_roles``) and the quantization (``_page_quant_mode``); a subclass + overrides the hook whose answer differs. ``ValueError`` if its layer groups, layer configs + and declarations do not agree.""" + tokens_per_block = int(self.tokens_per_block) + layer_cfg = self._page_layer_configs() + grouping = [[int(layer) for layer in group] for group in self.impl.layer_grouping] + pg_of_lg = [int(g) for g in self.impl.get_life_cycle_pool_group_indices()] + if len(pg_of_lg) != len(grouping): + raise ValueError(f"{len(grouping)} layer groups but {len(pg_of_lg)} pool-group indices") + layer_ids = self._page_layer_ids() + missing = sorted({lid for group in grouping for lid in group} - set(layer_ids)) + if missing: + raise ValueError(f"the manager names no global layer for internal layers {missing}") + global_of = {lid: int(ids[0]) for lid, ids in layer_ids.items()} + model_of = {lid: int(ids[1]) for lid, ids in layer_ids.items()} + kv_heads = {lid: self._page_kv_heads(lid) for group in grouping for lid in group} + + state_groups = set() + for lg, group in enumerate(grouping): + if not group or group[0] not in layer_cfg: + raise ValueError(f"layer group {lg} has no layer config (layers {group})") + if is_recurrent_layer(layer_cfg[group[0]]): + state_groups.add(lg) + + untransferred = self._page_untransferred_roles() + expansions = { + lid: buffer_expansions(cfg, tokens_per_block) for lid, cfg in layer_cfg.items() + } + + # The grid attention heads split over, and the one recurrent state splits over: tensor + # ranks, none under attention data parallelism, times context ranks under Helix. + mapping = self.mapping + attention_dp = bool(getattr(mapping, "enable_attention_dp", False)) + tp_size = 1 if attention_dp else max(int(mapping.tp_size), 1) + tp_rank = 0 if tp_size == 1 else int(mapping.tp_rank) % tp_size + cp_rank, cp_size = int(self._helix_cp_rank), int(self._helix_cp_size) + state_cp_rank, state_cp_size = (0, 1) if attention_dp else (cp_rank, cp_size) + state_shard = ( + ShardDesc() + if tp_size * state_cp_size == 1 + else ShardDesc(tp_size * state_cp_size, tp_rank * state_cp_size + state_cp_rank) + ) + + device_pools: dict[int, tuple[DevicePool, ...]] = {} + buffers: list[BufferDesc] = [] + shards: dict[int, list[ShardDesc]] = {lg: [] for lg in range(len(grouping))} + for pg in self.impl.pool_group_descs: + g = int(pg.pool_group_index) + pools = tuple( + DevicePool( + g, int(p.pool_index), int(p.base_address), int(p.slot_bytes), int(pg.num_slots) + ) + for p in pg.pools + ) + device_pools[g] = tuple(sorted(pools, key=lambda p: p.index)) + # One variant per layer group drawing from this pool group; pool ``i`` of a slot holds + # the ``i``-th coalesced buffer, whose members sit back to back. + for variant in pg.slot_desc.variants: + lg = int(variant.layer_group_id) + state = lg in state_groups + for pool_index, coalesced in enumerate(variant.coalesced_buffers): + size = int(coalesced.single_buffer_size) + for j, buffer_id in enumerate(coalesced.buffer_ids): + lid = int(buffer_id.layer_id) + role = str(buffer_id.role) + transfer = role not in untransferred + expansion = expansions.get(lid, {}).get(role, 1) + try: + tensors = self._page_buffer_tensors( + lid, role, size, tokens_per_block // expansion, state + ) + except ValueError as error: + raise ValueError(f"buffer ({lid}, {role}): {error}") from error + desc = BufferDesc( + layer=global_of.get(lid, lid), + role=role, + pool_group=g, + pool_index=pool_index, + slot_offset=j * size, + size=size, + tensors=tensors, + transferred=transfer, + expansion=expansion, + ) + buffers.append(desc) + if not transfer or not desc.split: + continue + heads, total = kv_heads[lid] + shards[lg].append( + state_shard + if state + else attention_shard(tp_size, tp_rank, heads, total) + ) + + layer_groups: list[LayerGroupDesc] = [] + total_kv_heads: dict[int, int] = {} + for lg, group in enumerate(grouping): + first = layer_cfg[group[0]] + is_state = lg in state_groups + window = read = None + sink_tokens = 0 + if not is_state: + window = getattr(first, "window_size", None) + window = None if window is None else int(window) + read = self._attention_read_window(lg) + read = None if read is None else int(read) + if (read is None) != (window is None) or ( + read is not None and not 0 < read <= window + ): + raise ValueError( + f"layer group {lg} reads a window of {read} tokens, which its kept window " + f"of {window} tokens does not hold" + ) + sink_tokens = int(getattr(first, "num_sink_tokens", None) or 0) + for lid in group: + total = kv_heads[lid][1] + if total: + total_kv_heads[global_of[lid]] = total + # Layer groups with equal slot sizes may share a pool group, each drawing its own + # slots. + g = pg_of_lg[lg] + if g not in device_pools: + raise ValueError(f"layer group {lg} maps to pool group {g}, which has no pools") + named = sorted((global_of[lid], model_of[lid]) for lid in group) + layer_groups.append( + LayerGroupDesc( + kind="state" if is_state else "attention", + kept_window=window, + num_sink_blocks=-(-sink_tokens // tokens_per_block), + pool_group=g, + layers=tuple(layer for layer, _ in named), + shard=group_shard(shards[lg], ShardDesc(tp_size, tp_rank)), + read_window=read, + model_layers=tuple(model_layer for _, model_layer in named), + ) + ) + + page_layout = PageLayout( + tokens_per_page=tokens_per_block, + pool_slot_bytes={ + g: tuple(p.slot_bytes for p in pools) for g, pools in device_pools.items() + }, + layer_groups=tuple(layer_groups), + buffers=tuple(buffers), + topology=rank_topology(mapping), + total_kv_heads_by_layer=total_kv_heads, + dtype=self._page_dtype(), + quant_mode=self._page_quant_mode(), + rebuilds_untransferred=self._page_rebuilds_untransferred(), + ) + return ManagerLayout( + page_layout=page_layout, + tokens_per_block=tokens_per_block * cp_size, + pool_group_of=tuple(pg_of_lg), + windows=tuple(desc.kept_window for desc in layer_groups), + sink_blocks=tuple(desc.num_sink_blocks for desc in layer_groups), + recurrent=tuple(desc.kind == "state" for desc in layer_groups), + sparse=tuple(has_sparse_buffers(layer_cfg[group[0]]) for group in grouping), + layers=tuple(desc.layers for desc in layer_groups), + shards=tuple((desc.shard.count, desc.shard.index) for desc in layer_groups), + page_bytes={g: sum(p.slot_bytes for p in pools) for g, pools in device_pools.items()}, + device_pools=device_pools, + cp=(cp_rank, cp_size), + read_windows=tuple(desc.read_window for desc in layer_groups), + named_groups=tuple( + not any(is_private_layer(layer_cfg[lid]) for lid in group) for group in grouping + ), + prefix_required=tuple( + not any(is_non_required_layer(layer_cfg[lid]) for lid in group) + for group in grouping + ), + ) + + def _page_layer_configs(self) -> dict[int, AttentionLayerConfig | SsmLayerConfig]: + """Internal layer id -> the layer's config, as the manager configured its runtime.""" + return layer_configs_by_id(self.kv_cache_manager_py_config) + + def _page_layer_ids(self) -> dict[int, tuple[int, int]]: + """Internal layer id -> (global layer id, model layer), the same under any pipeline split: + the model layer itself, or ``model_layer * num_types + type`` for a manager keeping several + caches per model layer (``_virtual_layers``). A manager whose internal layers are numbered + otherwise overrides this.""" + internal = sorted({int(lid) for group in self.impl.layer_grouping for lid in group}) + virtual = self._virtual_layers() + if virtual is None: + return {lid: (int(self.pp_layers[lid]),) * 2 for lid in internal} + types_of, num_types = virtual + return { + lid: (types_of[lid][0] * num_types + types_of[lid][1], types_of[lid][0]) + for lid in internal + } + + def _page_kv_heads(self, layer: int) -> tuple[int, int | None]: + """(KV heads on this rank, KV heads over all tensor-parallel ranks or None if unknown) of + internal layer ``layer``. A virtual layer counts the heads of its model layer; an internal + layer past the local layers counts those of the first.""" + virtual = self._virtual_layers() + index = layer + if virtual is not None and layer in virtual[0]: + model_layer = virtual[0][layer][0] + if model_layer in self.pp_layers: + index = list(self.pp_layers).index(model_layer) + per_layer = list(self.num_kv_heads_per_layer or ()) + per_rank = int(per_layer[index if index < len(per_layer) else 0] or 0) if per_layer else 0 + total = self.num_kv_heads + if isinstance(total, int): + return per_rank, int(total) + if total is None or virtual is not None or layer >= len(self.pp_layers): + # An extra internal layer (FP4 MLA's tail) is no model layer. + return per_rank, None + value = total[int(self.pp_layers[layer])] + return per_rank, None if value is None else int(value) + + def _page_buffer_tensors( + self, layer_id: int, role: str, size: int, tokens: int, state: bool + ) -> tuple[Tensor, ...]: + """How the bytes of internal layer ``layer_id``'s ``role`` buffer spread over ranks: the + tensors whose local pieces its ``size`` bytes hold, for ``tokens`` tokens of a page, in a + layer of recurrent ``state`` or of attention. + + The base manager declares the roles it registers itself: K/V and their block scales + head-major (HND) by the layer's KV heads, as the TRTLLM attention kernels write them, and + ``Role.INDEX_KEY`` alike on every rank, since every index-key side cache so far (DSA + indexer-K, MiniMax M3) projects replicated across TP ranks. A buffer that never moves + (``_page_untransferred_roles``) is this rank's own. A manager registering any other role, + or laying these out otherwise, overrides this method, building tensors with + ``page_layout``'s ``head_major``, ``token_major``, ``sectioned`` and ``replicated`` or + declaring its own. + + Raises: + ValueError: For a role that moves and that the manager does not declare: guessing how + its bytes spread would transfer them wrongly between ranks of different parallel + sizes. + """ + if role == Role.INDEX_KEY: + return replicated(size) + if role in _HEAD_SPLIT_ROLES: + return self._page_head_split(layer_id, size, tokens, state) + if role in self._page_untransferred_roles(): + split = ("tp", "cp") if state else ("tp",) + return owned(size, grid_place(rank_topology(self.mapping), split)[0], split) + raise ValueError( + f"{type(self).__name__} registers role {role!r} but does not declare how its bytes " + "spread over ranks; override _page_buffer_tensors (or list the role in " + "_page_untransferred_roles if it never moves)" + ) + + def _page_head_split( + self, + layer_id: int, + size: int, + tokens: int, + state: bool, + *, + head_bytes: Optional[int] = None, + token_major_: bool = False, + ) -> tuple[Tensor, ...]: + """A buffer split by head over the ranks attention splits over (tensor ranks) or recurrent + state does (tensor times context ranks): ``size // head_bytes`` heads, by default the + layer's KV heads on this rank, head-major or, with ``token_major_``, token-major. A buffer + whose heads are unknown is this rank's own share.""" + split = ("tp", "cp") if state else ("tp",) + ranks, _ = grid_place(rank_topology(self.mapping), split) + heads, total = self._page_kv_heads(layer_id) + if head_bytes is None and not state and heads > 0 and size % heads == 0: + head_bytes = size // heads + if head_bytes is None or size % head_bytes: + return owned(size, ranks, split) + heads = size // head_bytes + total = total if total and not state else heads * ranks + if token_major_: + return token_major(size, tokens, heads, total, split) + return head_major(size, heads, total, split) + + def _page_rebuilds_untransferred(self) -> bool: + """Whether the untransferred buffers derive from the others and are rebuilt after the + native transfer imports a request's pages. False by default.""" + return False + + def _page_untransferred_roles(self) -> frozenset[str]: + """Roles whose buffers hold no state a peer needs, such as storage-only buffers that + satisfy a pool's layout: they stay in local copies of whole pages but are never moved. + None by default.""" + return frozenset() + + def _page_dtype(self) -> str | None: + """The element type of the KV data by name, such as ``"BF16"`` or ``"NVFP4"``; None for a + manager declaring none.""" + dtype = getattr(self, "dtype", None) + return None if dtype is None else str(getattr(dtype, "name", dtype)) + + def _page_quant_mode(self) -> str | None: + """How the KV data is quantized: None unquantized, else ``"fp8"`` or ``"nvfp4"``; a manager + with another format (DeepSeek-V4's ``"fp8_ds_mla"``) overrides this.""" + return {"FP8": "fp8", "NVFP4": "nvfp4"}.get(self._page_dtype()) + + def _virtual_layers(self) -> tuple[dict[int, tuple[int, int]], int] | None: + """(internal id -> (model layer, attention type), type count) of a manager keeping several + caches per model layer, which overrides this (``page_table.virtual_layers``); None: one + cache per model layer.""" + return None + @property def draft_prompt_lookahead(self) -> int: """Prompt tokens the configured draft reads ahead, before reuse-backoff opt-outs.""" @@ -4300,6 +4669,15 @@ def _cache_stream(self, request_id: int) -> int: return self._allocation_stream.cuda_stream return self._stream.cuda_stream + def _allocation_fence(self, request_id: int) -> Optional[torch.cuda.Event]: + """An event after the work queued so far on the allocation side stream while the request's + cache is still on it; ``None`` once the execution stream adopted the cache.""" + if request_id not in self._on_allocation_stream: + return None + fence = torch.cuda.Event() + fence.record(self._allocation_stream) + return fence + def _place_on_allocation_stream(self, req: LlmRequest, kv_cache: _KVCache | None) -> None: """Allocate a connector-served context cache on the side stream. @@ -4822,7 +5200,7 @@ def revert_allocate_context(self, req: LlmRequest) -> bool: f"to {pre_cap}" ) self._after_shrink(req.py_request_id, kv_cache) - if pre_cap > 0: + if pre_cap > 0 and not self._holds_page_locks(req.py_request_id): kv_cache.suspend() return True @@ -4870,8 +5248,20 @@ def _restore_page_index_bufs(self, request_id: int, kv_cache: _KVCache) -> None: copy_batch_block_offsets still reads from the host buffer, so we must re-connect the buffers to avoid stale/zero page indices that would cause illegal memory accesses during the forward pass. + + A request whose index slot was released early has no buffer row: + its cache keeps its beam-0 page indices in the lender's parked row + instead, if it has one, which is connected again. """ + if request_id in self._early_freed_index_requests: + parked = _lender_hooks.parked_row(self, request_id) + if parked is not None: + for pool_idx in range(self.num_pools): + kv_cache.set_base_page_index_buf(BeamIndex(0), pool_idx, parked[pool_idx]) + _lender_hooks.on_resume(self, request_id) + return self._set_page_index_bufs(request_id, kv_cache) + _lender_hooks.on_resume(self, request_id) if self.sparse_metadata_batch is not None: index = self.index_mapper.get_index(request_id) @@ -5171,7 +5561,7 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: if reservation is not None: self.free_resources(req) rewind_context_after_cache_drop(req, self.tokens_per_block) - elif req.is_first_context_chunk: + elif req.is_first_context_chunk and not self._holds_page_locks(req.py_request_id): kv_cache.suspend() return False self._fill_fresh_kv_pages(req.py_request_id) @@ -5196,8 +5586,8 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: Allocates capacity for the full prompt (+ draft) and sets ``kv_cache.history_length`` to ``prompt_len``. Returns True on - success, False if preparation or resize failed (cache is suspended - on resize failure). + success, False if preparation or resize failed (an unlocked cache + is suspended on first-chunk resize failure). """ reused = self.prepare_context_cache(req) if reused is None: @@ -5231,7 +5621,7 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: success = kv_cache.resize(capacity, history_len) if not success: - if req.is_first_context_chunk: + if req.is_first_context_chunk and not self._holds_page_locks(req.py_request_id): kv_cache.suspend() return False self._fill_fresh_kv_pages(req.py_request_id) @@ -5303,6 +5693,7 @@ def suspend_request(self, req: LlmRequest) -> None: """Suspend a request's KV cache, allowing pages to migrate to a secondary tier.""" kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is not None and kv_cache.is_active: + self._check_suspendable(kv_cache) kv_cache.suspend() def resume_request(self, req: LlmRequest) -> bool: @@ -5447,6 +5838,23 @@ def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: ] return self._layer_group_life_cycles + def _attention_read_window(self, layer_group_id: int) -> Optional[int]: + """The largest declared read window, ``kept`` for missing, non-positive or virtual layers + and at most ``kept``; overrides return None without retention, else a value in (0, kept].""" + kept = self._life_cycle_by_layer_group()[layer_group_id].window_size + if kept is None: + return None + declared = self._model_attention_windows + if not declared or len(self.kv_cache_manager_py_config.layers) != len(self.pp_layers): + return kept + return max( + kept if window is None or window <= 0 else min(int(window), kept) + for window in ( + declared.get(self.pp_layers[int(lid)]) + for lid in self.impl.layer_grouping[layer_group_id] + ) + ) + def _stale_block_range(self, layer_group_id: int, history_length: int) -> Tuple[int, int]: """Block ordinals ``[start, end)`` that ``layer_group_id`` no longer reads. @@ -5612,12 +6020,17 @@ def _prepare_connector_prefix_reservation(self, req: LlmRequest) -> None: or not self.kv_connector_manager.should_add_sequence(req) ): return - local_end = self.kv_cache_map[req.py_request_id].num_committed_tokens + kv_cache = self.kv_cache_map[req.py_request_id] + # A connector loads past what a lender's fetch served, never over it. + local_end = max(kv_cache.num_committed_tokens, req.py_served_prefix_end) reservation = self.kv_connector_manager.reserve_prefix(req, local_end) if reservation is None: return - # Loads end at full blocks and leave the final prompt token for logits. - end = min(reservation.end, req.prompt_len - 1) + # Loads end at full blocks and leave the final prompt token for logits; the history they + # move to keeps the blocks a lender locks. + end = min( + reservation.end, req.prompt_len - 1, self._lent_history(kv_cache, reservation.end) + ) end = end // self.tokens_per_block * self.tokens_per_block # A connector-served prefix starves a hidden-state drafter exactly like # local reuse does. Cap the reservation at the same claim limit, so @@ -5716,11 +6129,16 @@ def _apply_connector_matched_prefix(self, req: LlmRequest) -> bool: if kv_cache is None: return False - local_end = kv_cache.num_committed_tokens + local_end = max(kv_cache.num_committed_tokens, req.py_served_prefix_end) num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( req, local_end ) - offered_end = min(local_end + num_tokens, req.prompt_len - 1) + # The served end becomes the history, which keeps the blocks a lender locks. + offered_end = min( + local_end + num_tokens, + req.prompt_len - 1, + self._lent_history(kv_cache, local_end + num_tokens), + ) position = local_end if offered_end > local_end: @@ -6903,21 +7321,51 @@ def release_resources( return requests - def try_commit_blocks(self, request: LlmRequest) -> None: + def _retry_held_commit( + self, request: LlmRequest, kv_cache: _KVCache, end: int, fetch_touched_suffix: bool + ) -> bool: + """Keep a fetch from publishing a later, potentially overwritten suffix.""" + return self.try_commit_blocks( + request, kv_cache, held_end=end if fetch_touched_suffix else None + ) + + def _held_commit_drop_end(self, held: tuple[_KVCache, LlmRequest, int, bool]) -> int: + """Use the saved attention prefix when a fetch overlaps a held commit.""" + return held[2] + + def try_commit_blocks( + self, request: LlmRequest, kv_cache: _KVCache | None = None, held_end: int | None = None + ) -> bool: + """Commit computed context blocks; return whether committing stopped. Locks defer the commit + until unlock and a lender's refusal of reuse skips it, as in Mamba.""" should_block_reuse = ( self.enable_block_reuse and self._can_publish_block_reuse and not request.is_dummy_request ) if not should_block_reuse: - return + return False - kv_cache = self.kv_cache_map.get(request.py_request_id) if kv_cache is None: - return + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return False - commit_end = request.context_current_position - if commit_end > kv_cache.num_committed_tokens: + current_position = request.context_current_position + commit_end = current_position if held_end is None else min(held_end, current_position) + commits = commit_end > kv_cache.num_committed_tokens + keep_committing = self._can_commit_generation_blocks(request) + stops = ( + request.context_remaining_length == 0 + and not keep_committing + and commit_end == current_position + ) + gate = _lender_hooks.COMMIT + if commits or stops: + gate = self._commit_gate(kv_cache, request, commit_end) + if gate == _lender_hooks.HOLD: + return False + if commits and gate == _lender_hooks.COMMIT: tokens = self._augment_tokens_for_block_reuse( self._reuse_token_source(request), request, @@ -6927,9 +7375,9 @@ def try_commit_blocks(self, request: LlmRequest) -> None: # TODO: On a disaggregated prefill server, pass is_end=True for # the last context chunk to improve performance. kv_cache.commit(tokens) - keep_committing = self._can_commit_generation_blocks(request) - if request.context_remaining_length == 0 and not keep_committing: + if stops: kv_cache.stop_committing() + return stops def _can_commit_generation_blocks(self, request: LlmRequest) -> bool: """Whether decode KV remains a valid reusable logical prefix.""" @@ -6995,15 +7443,22 @@ def release_index_slot(self, request_id: int) -> None: # mirrored, and the target may release the same request twice. return if kv_cache is not None: - if self.sparse_metadata_batch is not None: - self.sparse_metadata_batch.record_read(self._stream.cuda_stream) - self.sparse_metadata_batch.remove(kv_cache) - for beam_idx in range(int(kv_cache.beam_width)): - for pool_idx in range(self.num_pools): - kv_cache.set_base_page_index_buf(BeamIndex(beam_idx), pool_idx, None) + self._detach_index_slot(kv_cache, request_id) self.index_mapper.remove_sequence(request_id) self._early_freed_index_requests.add(request_id) + def _detach_index_slot(self, kv_cache: _KVCache, released: int | None = None) -> None: + """Detach native index pointers before the row is reused while the cache stays open; at + ``released``'s early slot release, beam 0 binds the lender's parked row, if any.""" + if self.sparse_metadata_batch is not None: + self.sparse_metadata_batch.record_read(self._stream.cuda_stream) + self.sparse_metadata_batch.remove(kv_cache) + parked = None if released is None else _lender_hooks.parking_row(self, released, kv_cache) + for beam_idx in range(int(kv_cache.beam_width)): + for pool_idx in range(self.num_pools): + buffer = parked[pool_idx] if parked is not None and beam_idx == 0 else None + kv_cache.set_base_page_index_buf(BeamIndex(beam_idx), pool_idx, buffer) + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): # A request awaiting preemption can still be cancelled or fail while # its saves drain, and a dead entry here would block every later @@ -7043,9 +7498,11 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self.impl.clear_stats_excluded(request.py_request_id) return kv_cache.discard_pending_stats() - if self._sharing is not None: - self._sharing._on_free(request.py_request_id, kv_cache) - kv_cache.close() + _lender_hooks.on_free(self, request.py_request_id, kv_cache) + if not _lender_hooks.defer_close(self, kv_cache): + kv_cache.close() + elif request.py_request_id not in self._early_freed_index_requests: + self._detach_index_slot(kv_cache) self.impl.clear_stats_excluded(request.py_request_id) if request.py_request_id in self._early_freed_index_requests: self._early_freed_index_requests.discard(request.py_request_id) @@ -7053,11 +7510,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self.index_mapper.remove_sequence(request.py_request_id) def _after_shrink(self, request_id: int, kv_cache: _KVCache) -> None: - """Tell the lender, if one is attached, that ``kv_cache`` may have shrunk in place: blocks - past its capacity lost their pages, and growing again does not bring their contents back. - Every in-place shrink the manager makes calls this right after it.""" - if self._sharing is not None: - self._sharing._on_shrink(request_id, kv_cache) + """Notify the lender after an in-place shrink invalidates pages past capacity.""" + _lender_hooks.on_shrink(self, request_id, kv_cache) def get_layer_page_index_scale(self, layer_idx: int) -> int: """Page-index scale of this layer's KV buffer. Layers in one pool can @@ -7354,7 +7808,14 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool def shutdown(self): if self._sharing is not None: - self._sharing._on_shutdown() + self._held_commits.clear() + _lender_hooks.on_shutdown(self) + # Locked caches close at their last unlock. + for request_id, kv_cache in list(self.kv_cache_map.items()): + if _lender_hooks.defer_close(self, kv_cache): + if request_id not in self._early_freed_index_requests: + self._detach_index_slot(kv_cache) + del self.kv_cache_map[request_id] if self.sparse_metadata_batch is not None: self.sparse_metadata_batch.record_read(self._stream.cuda_stream) self.sparse_metadata_batch.close() @@ -7371,11 +7832,16 @@ def shutdown(self): # its plan, which mutates manager state. if self.conversation_manager is not None: self.conversation_manager.clear() - self.impl.shutdown() - if self._sharing is not None: - # After the runtime's shutdown, which raises while any cache is open, in the map or not; - # a poisoned runtime returns at once, and a close then does nothing. - self._sharing._on_caches_closed() + if _lender_hooks.keeps_pools(self): + logger.warning( + "KVCacheManagerV2: a lender still holds KV caches or the pools at shutdown, so the " + "KV cache pools stay allocated until the process exits" + ) + else: + self.impl.shutdown() + # After the runtime's shutdown, which raises while any cache is open, in the map or + # not; a poisoned runtime returns at once, and a close then does nothing. + _lender_hooks.on_caches_closed(self) # Shut the streaming event manager down last so removals emitted during # cache / impl teardown (via the radix tree's own event-manager # reference) are still flushed before the publisher stops. Do not null @@ -7801,7 +8267,10 @@ def update_context_resources(self, scheduled_batch: ScheduledRequests): layer. In non-overlap scheduler, you should call it together with update_resources(). """ - for req in scheduled_batch.context_requests: + requests = scheduled_batch.context_requests + # The lender is asked only where a lock can cap a history, through a window. + lock_state = self._history_lock_state() if requests else None + for req in requests: if req.py_request_id not in self.kv_cache_map: continue kv_cache = self.kv_cache_map[req.py_request_id] @@ -7834,11 +8303,14 @@ def update_context_resources(self, scheduled_batch: ScheduledRequests): should_commit = is_all_reusable or is_branch_point or req.context_remaining_length == 0 if should_resize: - self._resize_context_history(req, kv_cache, req.context_current_position) - if should_commit: - self.try_commit_blocks(req) + position = req.context_current_position + if lock_state is not None: + position = self._lent_history(kv_cache, position) + self._resize_context_history(req, kv_cache, position) + stopped = should_commit and self.try_commit_blocks(req) if req.context_remaining_length == 0: - if self.conversation_manager is not None: + # Without reuse nothing commits; a held-back commit saves the plan at its retry. + if self.conversation_manager is not None and (stopped or not should_block_reuse): self.conversation_manager.save_drop_plan(req, kv_cache) # Scratch blocks are only for prefill chunks. Disable them at # the context/generation boundary so generation uses normal KV @@ -7872,6 +8344,10 @@ def update_resources( ) # Context request KV cache updates are handled by # update_context_resources, called separately from the executor loop. + lock_state = _lender_hooks.generation_lock_state(self) + has_lender = lock_state is not None + # A lender's lock caps a history only through a window. + caps_history = has_lender and lock_state.caps_history for req in scheduled_batch.generation_requests: if req.py_request_id not in self.kv_cache_map: continue @@ -7915,13 +8391,15 @@ def update_resources( if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) else kv_cache.capacity - rewind_len ) - history_length = ( - None - # Reuse (history's consumer) is disabled under helix, and - # max_beam_num_tokens mixes rank-local and global counts. - if self.kv_compression_manages_history or self._has_cp_helix - else req.max_beam_num_tokens - 1 - ) + # Reuse (history's consumer) is disabled under helix, and + # max_beam_num_tokens mixes rank-local and global counts. + if self.kv_compression_manages_history or self._has_cp_helix: + history_length = None + else: + history_length = req.max_beam_num_tokens - 1 + if caps_history and lock_state.holding: + history_length = self._lent_history(kv_cache, history_length) + num_blocks_before = kv_cache.num_blocks if has_lender else 0 success = kv_cache.resize(new_capacity, history_length) if not success: raise ValueError( @@ -7929,8 +8407,7 @@ def update_resources( f"to capacity {new_capacity} and history length " f"{history_length} tokens at generation update" ) - # Inline check: without a lender this per-request hot path makes no extra call. - if new_capacity is not None and self._sharing is not None: + if has_lender and kv_cache.num_blocks < num_blocks_before: self._after_shrink(req.py_request_id, kv_cache) self._allocated_draft_lens.pop(req.py_request_id, None) self.try_commit_generation_blocks(req) @@ -8086,6 +8563,7 @@ def _create_kv_cache( **priority_kwargs, ) self.kv_cache_map[request_id] = kv_cache + _lender_hooks.on_create(self, request_id) if enable_request_stats and not self.enable_stats: self._request_stats_enabled_ids.add(request_id) if is_dummy: @@ -8188,8 +8666,8 @@ def prefetch_for_context_tokens(self, requests: list) -> bool: return success def reset_reuse_state(self): + _lender_hooks.before_reset(self) self.impl.clear_reusable_blocks() if self.conversation_manager is not None: self.conversation_manager.clear() - if self._sharing is not None: - self._sharing._on_reset() + _lender_hooks.on_reset(self) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 3d4dd7b8b952..cd9fc93c187e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -41,12 +41,15 @@ from typing_extensions import override from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout +from tensorrt_llm._torch.pyexecutor.kv_cache import _lender_hooks from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( _RESERVED_REQUEST_IDS, BlockReusePolicy, KVCacheManagerV2, Role, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import Tensor, grid_place, sectioned +from tensorrt_llm._torch.pyexecutor.kv_cache.page_table import rank_topology from tensorrt_llm._torch.pyexecutor.kv_cache_stats import KVCacheV2IterationStatsReport from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import ( @@ -403,6 +406,35 @@ def get_disagg_role_layouts(self) -> dict[DataRole, RoleLayout]: ), } + @override + def _page_buffer_tensors( + self, layer_id: int, role: str, size: int, tokens: int, state: bool + ) -> tuple[Tensor, ...]: + """Convolution state is a concatenation of sections (Mamba2 ``[x | B | C]``, GDN + ``[Q | K | V]``), each split by byte over the ranks recurrent state splits over; SSM state + is split by head.""" + if role == MambaRole.CONV_STATE: + d_conv_m1 = int(self.conv_state_shape[1]) + ranks, _ = grid_place(rank_topology(self.mapping), ("tp", "cp")) + return sectioned( + [ + int(dim) * d_conv_m1 * self.conv_state_dtype.itemsize + for dim in self.conv_section_dims + ], + ranks, + ("tp", "cp"), + ) + if role == MambaRole.SSM_STATE: + _, head_dim, d_state = self.ssm_state_shape + return self._page_head_split( + layer_id, + size, + tokens, + state, + head_bytes=int(head_dim) * int(d_state) * self.ssm_state_dtype.itemsize, + ) + return super()._page_buffer_tensors(layer_id, role, size, tokens, state) + def _extra_scratch_bytes_per_slot(self) -> int: return 0 @@ -1105,27 +1137,48 @@ def update_mamba_states( ) @override - def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: + def _retry_held_commit( + self, request: LlmRequest, kv_cache: _KVCache, end: int, fetch_touched_suffix: bool + ) -> bool: + """Retry the current recurrent state, since an earlier snapshot may be stale.""" + return self.try_commit_blocks(request, kv_cache) + + @override + def _held_commit_drop_end(self, held: tuple) -> int: + """A recurrent retry uses the request's current endpoint.""" + return held[1].context_current_position + + @override + def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> bool: + """Commit recurrent snapshots when page locks and the lender permit it.""" should_block_reuse = ( self.enable_block_reuse and not self.is_draft and not request.is_dummy_request ) if not should_block_reuse: - return + return False if kv_cache is None: kv_cache = self.kv_cache_map.get(request.py_request_id) if kv_cache is None: - return + return False snapshot_points = request.expect_snapshot_points commit_limit = ( min(max(snapshot_points), request.prompt_len) if snapshot_points else request.prompt_len ) commit_end = min(request.context_current_position, commit_limit) - if ( + commits = ( request.context_current_position in request.expect_snapshot_points and commit_end > kv_cache.num_committed_tokens - ): + ) + stops = request.context_remaining_length == 0 + gate = ( + self._commit_gate(kv_cache, request, commit_end) + if commits or stops + else _lender_hooks.COMMIT + ) + held = gate == _lender_hooks.HOLD + if commits and gate == _lender_hooks.COMMIT: tokens = self._augment_tokens_for_block_reuse( request.get_tokens(DEFAULT_BEAM_INDEX), request, @@ -1140,12 +1193,13 @@ def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: self._branch_snapshots_taken_total += 1 if request.context_current_position >= commit_limit: self._mark_context_position_as_history(request, kv_cache) - if request.context_remaining_length == 0: + if stops and not held: kv_cache.stop_committing() + return stops and not held def _mark_context_position_as_history(self, request: LlmRequest, kv_cache) -> None: """Advance history without making later recurrent state reusable.""" - history_length = request.context_current_position + history_length = self._lent_history(kv_cache, request.context_current_position) if history_length <= kv_cache.history_length: return capacity = max(kv_cache.capacity, history_length) @@ -1157,7 +1211,10 @@ def _mark_context_position_as_history(self, request: LlmRequest, kv_cache) -> No @override def update_context_resources(self, scheduled_batch: ScheduledRequests) -> None: - for request in scheduled_batch.context_requests: + requests = scheduled_batch.context_requests + # A lender can cap history only while a page lock holds a window. + lock_state = self._history_lock_state() if requests else None + for request in requests: kv_cache = self.kv_cache_map.get(request.py_request_id) if kv_cache is None or not kv_cache.is_active: continue @@ -1179,17 +1236,21 @@ def update_context_resources(self, scheduled_batch: ScheduledRequests) -> None: is_all_reusable or is_snapshot_boundary or request.context_remaining_length == 0 ) - if should_resize and not kv_cache.resize(None, request.context_current_position): - raise ValueError( - "Failed to resize history length of V2 Mamba cache for " - f"request {request.py_request_id} to " - f"{request.context_current_position} tokens at context " - "update" - ) - if should_commit: - self.try_commit_blocks(request, kv_cache) + if should_resize: + position = request.context_current_position + if lock_state is not None: + position = self._lent_history(kv_cache, position) + if not kv_cache.resize(None, position): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to " + f"{request.context_current_position} tokens at context " + "update" + ) + stopped = should_commit and self.try_commit_blocks(request, kv_cache) if request.context_remaining_length == 0: - if self.conversation_manager is not None: + # A held commit saves the conversation's drop plan at its retry. + if self.conversation_manager is not None and (stopped or not should_block_reuse): self.conversation_manager.save_drop_plan(request, kv_cache) kv_cache.enable_swa_scratch_reuse = False diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/page_layout.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/page_layout.py new file mode 100644 index 000000000000..1aaf5597931e --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/page_layout.py @@ -0,0 +1,957 @@ +# 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. +"""One rank's page layout as immutable values, and which ranks of an instance hold the same bytes; +standard library only. + +A buffer holds the local pieces of one or more tensors. A tensor names its axes and their extents +over all ranks together; each axis says how its indices spread over the parallel dimensions +(``Axis.split``: whole on every rank, a contiguous run per rank, or every n-th index per rank) and +the byte stride of its local index. Which global indices a rank holds is therefore a progression +per axis (``held``), and how two ranks' pages correspond follows from intersecting them.""" + +from __future__ import annotations + +import collections.abc +import dataclasses +import enum +import math +import numbers +import types +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Literal, Mapping + +__all__ = [ + "GRID_DIMENSIONS", + "Axis", + "BufferDesc", + "LayerGroupDesc", + "MapperKind", + "PageLayout", + "Progression", + "RankTopology", + "RoleLayout", + "ShardDesc", + "Tensor", + "grid_place", + "head_major", + "held", + "owned", + "replica_ranks", + "replicated", + "sectioned", + "token_major", +] + +_QUANT_MODES = ("fp8", "nvfp4", "fp8_ds_mla") + + +def _integer(value: int, name: str, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, numbers.Integral) or value < minimum: + requirement = "positive" if minimum == 1 else "non-negative" + raise ValueError(f"{name} must be a {requirement} integer, got {value!r}") + return int(value) + + +def _frozen(mapping: dict) -> Mapping: + return types.MappingProxyType(mapping) + + +# The default of an ``unchecked`` constructor's mapping argument. +_NO_ENTRIES: Mapping = _frozen({}) + + +class MapperKind(enum.IntEnum): + """How a role's bytes split across ranks, as a manager declares it to the native transfer + (``get_disagg_role_mapper_kinds``). + + Attributes: + INDEXED: Head-major layout; a head range is one contiguous byte range. + HND: Alias of INDEXED for the [head, token, dimension] layout. + REPLICATED: Identical bytes on every tensor-parallel rank. + NHD: Token-major [token, head, dimension] layout; a head range is a slice within every + token. + SECTIONED: Consecutive byte sections, each sharded independently, such as recurrent + convolution state. + + Raises: + ValueError: If constructed from a value that names no mapper kind.""" + + INDEXED = 0 + HND = INDEXED + REPLICATED = 1 + NHD = 2 + SECTIONED = 3 + + +@dataclass(frozen=True) +class RoleLayout: + """Resharding geometry of one role, as a manager declares it to the native transfer + (``get_disagg_role_layouts``); both fields None means no declared head axis. + + Attributes: + section_bytes: Byte sizes of consecutive sections for a SECTIONED role, normalized to a + tuple. When given, their sum must equal the buffer's size. + bytes_per_head: Total bytes for one rank-local head, derived as size/head count for + attention INDEXED/NHD roles. NHD offsets remain token-major, not one contiguous head + slice. For other INDEXED roles, size must be divisible by this value; the quotient is + the per-rank head count. + + Raises: + ValueError: If a section size or bytes_per_head is not a positive integer, including a + boolean.""" + + section_bytes: tuple[int, ...] | None = None + bytes_per_head: int | None = None + + def __post_init__(self) -> None: + if self.section_bytes is not None: + object.__setattr__( + self, + "section_bytes", + tuple(_integer(n, "section size", 1) for n in self.section_bytes), + ) + if self.bytes_per_head is not None: + object.__setattr__( + self, "bytes_per_head", _integer(self.bytes_per_head, "bytes_per_head", 1) + ) + + @classmethod + def unchecked( + cls, section_bytes: tuple[int, ...] | None = None, bytes_per_head: int | None = None + ) -> RoleLayout: + """The constructor's role layout of these values, without its checks and conversions; + see ``PageLayout.unchecked``.""" + layout = object.__new__(cls) + layout.__dict__.update(section_bytes=section_bytes, bytes_per_head=bytes_per_head) + return layout + + +GRID_DIMENSIONS = ("tp", "cp") +"""Parallel dimensions an axis may split over (``Axis.split``): "tp" the tensor-parallel ranks of +one attention data-parallel group, "cp" the context-parallel ranks. Pipeline ranks hold other +layers and data-parallel ranks other requests, so neither splits a page.""" + + +@dataclass(frozen=True) +class Progression: + """The indices ``start, start + step, ..., start + (count - 1) * step``. + + Attributes: + start: First index. + step: Positive distance between consecutive indices. + count: Non-negative number of indices. + """ + + start: int + step: int + count: int + + @property + def last(self) -> int: + """The last index; ``start - step`` when empty.""" + return self.start + (self.count - 1) * self.step + + def __len__(self) -> int: + return self.count + + def __iter__(self): + return iter(range(self.start, self.start + self.count * self.step, self.step)) + + def __contains__(self, index: object) -> bool: + if not isinstance(index, int) or self.count == 0: + return False + offset = index - self.start + return 0 <= offset <= (self.count - 1) * self.step and offset % self.step == 0 + + def intersect(self, other: Progression) -> Progression | None: + """The indices both hold, themselves a progression, or None when they share none.""" + if self.count == 0 or other.count == 0: + return None + lo, hi = max(self.start, other.start), min(self.last, other.last) + if lo > hi: + return None + g = math.gcd(self.step, other.step) + if (other.start - self.start) % g: + return None + lcm = self.step // g * other.step + # x = self.start + self.step * t with x = other.start (mod other.step). + m = other.step // g + t = 0 if m == 1 else (other.start - self.start) // g * pow(self.step // g, -1, m) % m + first = lo + (self.start + self.step * t - lo) % lcm + if first > hi: + return None + return Progression(first, lcm, (hi - first) // lcm + 1) + + +def grid_place(topology: RankTopology, dims: Sequence[str]) -> tuple[int, int]: + """``(size, rank)`` of this rank in the grid of the parallel dimensions ``dims``, outer to + inner, the last minor. Under attention data parallelism each tensor-parallel rank serves its own + requests and holds its buffers whole, so "tp" counts this rank alone.""" + size, rank = 1, 0 + for dim in dims: + if dim == "tp": + n, r = (1, 0) if topology.attention_dp else (topology.tp_size, topology.tp_rank) + elif dim == "cp": + n, r = topology.cp_size, topology.cp_rank + else: + raise ValueError(f"{dim!r} is no grid dimension; one of {GRID_DIMENSIONS}") + size, rank = size * n, rank * n + r + return size, rank + + +@dataclass(frozen=True) +class Axis: + """One dimension of a tensor: its extent over all ranks together, which of its indices a rank + holds, and the byte stride between consecutive indices a rank holds. + + A rank at ``r`` of ``n`` in the grid of the ``split`` dimensions holds every index when + ``split`` is empty; ``held`` indices from ``r * extent // n`` when split contiguously (ranks + hold overlapping or equal runs when ``held * n`` exceeds ``extent``); or indices + ``r, r + n, ...`` below ``extent`` when ``cyclic``. Local index ``k`` (the k-th index held, in + order) lies at byte ``k * stride`` of its tensor. + + Attributes: + name: The axis' name; peers match tensors' axes on it. + extent: Positive number of indices over all ranks together. + stride: Positive bytes between consecutive local indices. + split: Parallel dimensions from ``GRID_DIMENSIONS`` the axis spreads over, outer to inner; + empty when every rank holds every index. Normalized to a tuple. + held: Positive number of indices a rank holds of a contiguous split, at most extent; + required for a contiguous split, None otherwise. + cyclic: Whether a split rank holds every n-th index rather than a contiguous run. + + Raises: + ValueError: If a number is not a positive integer or is a boolean, split names a dimension + twice or one outside ``GRID_DIMENSIONS``, held is missing from a contiguous split or + given for another, or held exceeds extent. + """ + + name: str + extent: int + stride: int + split: tuple[str, ...] = () + held: int | None = None + cyclic: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or not self.name: + raise ValueError(f"an axis needs a name, got {self.name!r}") + object.__setattr__(self, "extent", _integer(self.extent, f"extent of {self.name}", 1)) + object.__setattr__(self, "stride", _integer(self.stride, f"stride of {self.name}", 1)) + split = tuple(str(dim) for dim in self.split) + if len(set(split)) != len(split) or not set(split) <= set(GRID_DIMENSIONS): + raise ValueError(f"axis {self.name} splits over {split}, not over {GRID_DIMENSIONS}") + object.__setattr__(self, "split", split) + object.__setattr__(self, "cyclic", bool(self.cyclic)) + contiguous = bool(split) and not self.cyclic + if self.cyclic and not split: + raise ValueError(f"axis {self.name} is cyclic but not split") + if contiguous != (self.held is not None): + raise ValueError( + f"axis {self.name}: a contiguous split, and only one, says how many indices a " + f"rank holds (held={self.held!r}, split={split}, cyclic={self.cyclic})" + ) + if self.held is not None: + held_ = _integer(self.held, f"indices held of {self.name}", 1) + if held_ > self.extent: + raise ValueError(f"axis {self.name}: {held_} held of {self.extent} indices") + object.__setattr__(self, "held", held_) + + @classmethod + def unchecked( + cls, + name: str, + extent: int, + stride: int, + split: tuple[str, ...] = (), + held: int | None = None, + cyclic: bool = False, + ) -> Axis: + """The constructor's axis of these values, without its checks and conversions; see + ``PageLayout.unchecked``.""" + axis = object.__new__(cls) + axis.__dict__.update( + name=name, extent=extent, stride=stride, split=split, held=held, cyclic=cyclic + ) + return axis + + +@dataclass(frozen=True) +class Tensor: + """A tensor whose local piece a buffer holds: its axes, and where the piece starts. + + Attributes: + axes: The tensor's axes, normalized to a tuple; their order is the order peers list them + in, not the order of their bytes, which their strides give. + offset: Non-negative byte offset of the local piece within its buffer. + name: Tells the tensors of one buffer apart; peers match tensors on it. + + Raises: + ValueError: If there is no axis, two axes share a name, a parallel dimension splits two + axes, or offset is not a non-negative integer. + """ + + axes: tuple[Axis, ...] + offset: int = 0 + name: str = "" + + def __post_init__(self) -> None: + axes = tuple(self.axes) + if not axes or not all(isinstance(a, Axis) for a in axes): + raise ValueError(f"a tensor has one or more Axis values, got {axes!r}") + names = [a.name for a in axes] + if len(set(names)) != len(names): + raise ValueError(f"axis names repeat: {names}") + dims = [dim for a in axes for dim in a.split] + if len(set(dims)) != len(dims): + raise ValueError(f"a parallel dimension splits two axes: {dims}") + object.__setattr__(self, "axes", axes) + object.__setattr__(self, "offset", _integer(self.offset, "tensor offset")) + object.__setattr__(self, "name", str(self.name)) + + @property + def split(self) -> bool: + """Whether some axis spreads over parallel dimensions.""" + return any(a.split for a in self.axes) + + @classmethod + def unchecked(cls, axes: tuple[Axis, ...], offset: int = 0, name: str = "") -> Tensor: + """The constructor's tensor of these values, without its checks and conversions; see + ``PageLayout.unchecked``.""" + tensor = object.__new__(cls) + tensor.__dict__.update(axes=tuple(axes), offset=offset, name=name) + return tensor + + +def held(axis: Axis, topology: RankTopology) -> Progression: + """The global indices of ``axis`` that the rank at ``topology`` holds, in local order.""" + if not axis.split: + return Progression(0, 1, axis.extent) + n, r = grid_place(topology, axis.split) + if axis.cyclic: + return Progression(r, n, len(range(r, axis.extent, n))) + start = r * axis.extent // n + if start + axis.held > axis.extent: + raise ValueError( + f"axis {axis.name}: rank {r} of {n} holds {axis.held} indices from {start}, past " + f"its extent of {axis.extent}" + ) + return Progression(start, 1, axis.held) + + +def _footprint(tensor: Tensor) -> int | None: + """Bytes from its buffer's start to the end of a tensor's local piece, or None when a cyclic + axis leaves the count of its indices to the rank.""" + end = tensor.offset + 1 + for a in tensor.axes: + if a.cyclic: + return None + end += ((a.held if a.split else a.extent) - 1) * a.stride + return end + + +def replicated(size: int) -> tuple[Tensor, ...]: + """A buffer of ``size`` bytes that every rank holds alike.""" + return (Tensor((Axis("byte", size, 1),)),) + + +def owned(size: int, ranks: int, split: Sequence[str] = ("tp",)) -> tuple[Tensor, ...]: + """A buffer of ``size`` bytes that is this rank's own share of ``ranks`` ranks along ``split``, + in a split the manager does not describe further: it corresponds only to the same rank's.""" + return ( + Tensor((Axis("shard", ranks, size, tuple(split), 1), Axis("byte", size, 1)), name="own"), + ) + + +def head_major( + size: int, heads: int, total_heads: int, split: Sequence[str] = ("tp",) +) -> tuple[Tensor, ...]: + """A buffer of ``heads`` heads of ``size // heads`` bytes each ([head, ...], HND), of + ``total_heads`` split contiguously over ``split``.""" + if heads <= 0 or size % heads: + raise ValueError(f"{size} bytes is not a whole number of {heads} heads") + head = size // heads + return ( + Tensor( + (Axis("head", total_heads, head, tuple(split), heads), Axis("byte", head, 1)), + ), + ) + + +def token_major( + size: int, tokens: int, heads: int, total_heads: int, split: Sequence[str] = ("tp",) +) -> tuple[Tensor, ...]: + """A buffer of ``tokens`` tokens, each ``heads`` heads of equal bytes ([token, head, ...], + NHD), of ``total_heads`` split contiguously over ``split``.""" + if tokens <= 0 or heads <= 0 or size % (tokens * heads): + raise ValueError(f"{size} bytes is not a whole number of {tokens} x {heads} token heads") + head = size // (tokens * heads) + return ( + Tensor( + ( + Axis("token", tokens, heads * head), + Axis("head", total_heads, head, tuple(split), heads), + Axis("byte", head, 1), + ) + ), + ) + + +def sectioned( + section_bytes: Sequence[int], ranks: int, split: Sequence[str] = ("tp",) +) -> tuple[Tensor, ...]: + """Consecutive sections of a buffer, each the local ``1 / ranks`` of its own tensor, split + contiguously by byte over the ``ranks`` ranks of ``split`` (Mamba2 conv state ``[x | B | C]``, + GDN ``[Q | K | V]``).""" + out, offset = [], 0 + for i, n in enumerate(section_bytes): + out.append( + Tensor((Axis("byte", n * ranks, 1, tuple(split), n),), offset=offset, name=f"s{i}") + ) + offset += n + return tuple(out) + + +@dataclass(frozen=True) +class ShardDesc: + """Which share of a layer group's content this rank holds. + + The default (count=1, index=0) is the whole content, named alike on every rank. + + Attributes: + count: Positive number of content shards. + index: Zero-based index of this rank's share, less than count. + + Raises: + ValueError: If count is not a positive integer, index is not a non-negative integer, either + is a boolean, or index lies outside count.""" + + count: int = 1 + index: int = 0 + + def __post_init__(self) -> None: + object.__setattr__(self, "count", _integer(self.count, "shard count", 1)) + object.__setattr__(self, "index", _integer(self.index, "shard index")) + if self.index >= self.count: + raise ValueError(f"shard {self.index} of {self.count}") + + @classmethod + def unchecked(cls, count: int = 1, index: int = 0) -> ShardDesc: + """The constructor's shard of these values, without its checks and conversions; see + ``PageLayout.unchecked``.""" + shard = object.__new__(cls) + shard.__dict__.update(count=count, index=index) + return shard + + +@dataclass(frozen=True) +class LayerGroupDesc: + """One local layer group's windows, pools, global layers and content shard. + + A block held only by the kept window may remain unwritten without preventing + resumption when it lies outside the read window. + + A group's transferred buffers split over one grid of ranks, set by its kind: + attention groups over the tensor-parallel ranks of one attention data-parallel + group; state groups over the tensor-parallel ranks, over none under attention + data parallelism, and over tensor times context ranks, context-minor, under + Helix. ``shard`` is this rank's share; each buffer's ``tensors`` say how it + splits. + + Attributes: + kind: "attention" for attention KV data, or "state" for recurrent state. + kept_window: Positive number of tokens whose blocks the manager retains; None means full + attention. + num_sink_blocks: Non-negative number of sink blocks retained by the manager. + pool_group: Non-negative index of the group's device pool group. + layers: Global layer ids in this local group, normalized to a tuple. + shard: The share of the group's content held on this rank. + read_window: Positive number of tokens attention reads, defaulting to kept_window; None + means full attention. A declared read window must lie within the kept window. + model_layers: The model layer each of ``layers`` belongs to, in the same order, normalized + to a tuple; None gives ``layers`` themselves. Layer ids numbering several caches of one + model layer apart, such as one per attention type, all name that model layer. + + Raises: + ValueError: If an index, layer id or model layer is not a non-negative integer, a window is + not a positive integer, any such value is a boolean, read_window exceeds kept_window or + is given when kept_window is None, or model_layers does not name one model layer per + layer.""" + + kind: Literal["attention", "state"] + kept_window: int | None + num_sink_blocks: int + pool_group: int + layers: tuple[int, ...] + shard: ShardDesc = ShardDesc() + read_window: int | None = None + model_layers: tuple[int, ...] | None = None + + def __post_init__(self) -> None: + if self.kept_window is not None: + object.__setattr__(self, "kept_window", _integer(self.kept_window, "kept_window", 1)) + object.__setattr__( + self, "num_sink_blocks", _integer(self.num_sink_blocks, "num_sink_blocks") + ) + object.__setattr__(self, "pool_group", _integer(self.pool_group, "pool_group")) + object.__setattr__(self, "layers", tuple(_integer(n, "layer") for n in self.layers)) + if self.read_window is None: + object.__setattr__(self, "read_window", self.kept_window) + else: + object.__setattr__(self, "read_window", _integer(self.read_window, "read_window", 1)) + if self.read_window is not None and ( + self.kept_window is None or self.read_window > self.kept_window + ): + raise ValueError( + f"a read window of {self.read_window} tokens is not within the kept window of " + f"{self.kept_window} tokens" + ) + if self.model_layers is None: + object.__setattr__(self, "model_layers", self.layers) + else: + model_layers = tuple(_integer(n, "model layer") for n in self.model_layers) + if len(model_layers) != len(self.layers): + raise ValueError( + f"{len(model_layers)} model layers for the {len(self.layers)} layers " + f"{self.layers}" + ) + object.__setattr__(self, "model_layers", model_layers) + + @classmethod + def unchecked( + cls, + kind: Literal["attention", "state"], + kept_window: int | None, + num_sink_blocks: int, + pool_group: int, + layers: tuple[int, ...], + shard: ShardDesc = ShardDesc(), + read_window: int | None = None, + model_layers: tuple[int, ...] | None = None, + ) -> LayerGroupDesc: + """The constructor's layer group of these values, without its checks and conversions; a + read window of None still defaults to the kept window, and model layers of None to the + layers. See ``PageLayout.unchecked``.""" + group = object.__new__(cls) + group.__dict__.update( + kind=kind, + kept_window=kept_window, + num_sink_blocks=num_sink_blocks, + pool_group=pool_group, + layers=layers, + shard=shard, + read_window=kept_window if read_window is None else read_window, + model_layers=layers if model_layers is None else model_layers, + ) + return group + + +@dataclass(frozen=True) +class BufferDesc: + """One role of one global layer within a page's pool slot. + + Attributes: + layer: Non-negative global layer id. + role: Role name, normalized to a string. + pool_group: Non-negative device pool group index. + pool_index: Non-negative pool index within pool_group. + slot_offset: Non-negative byte offset within that pool's slot. + size: Positive byte size of the buffer. + tensors: The tensors whose local pieces the buffer holds, and how they spread over ranks; + normalized to a tuple. Defaults to the whole buffer, alike on every rank. + transferred: Whether peers transfer this role. False marks local-only buffers, which still + accompany copies of whole slots. + expansion: Positive number of sub-pages within one page, each of + PageLayout.tokens_per_page / expansion of the buffer's tokens; consumers accessing bytes + within the buffer must support it. + + Raises: + ValueError: If an index is not a non-negative integer, size or expansion is not a positive + integer, any such value is a boolean, two tensors share a name, a tensor's piece + reaches past size, or the pieces of several tensors do not lie back to back over the + whole buffer.""" + + layer: int + role: str + pool_group: int + pool_index: int + slot_offset: int + size: int + tensors: tuple[Tensor, ...] = () + transferred: bool = True + expansion: int = 1 + + def __post_init__(self) -> None: + for name in ("layer", "pool_group", "pool_index", "slot_offset"): + object.__setattr__(self, name, _integer(getattr(self, name), name)) + object.__setattr__(self, "size", _integer(self.size, "size", 1)) + object.__setattr__(self, "expansion", _integer(self.expansion, "expansion", 1)) + object.__setattr__(self, "role", str(self.role)) + tensors = tuple(self.tensors) or replicated(self.size) + if not all(isinstance(t, Tensor) for t in tensors): + raise ValueError(f"tensors must be Tensor values, got {tensors!r}") + names = [t.name for t in tensors] + if len(set(names)) != len(names): + raise ValueError(f"buffer ({self.layer}, {self.role}): tensor names repeat: {names}") + ends = [_footprint(t) for t in tensors] + for t, end in zip(tensors, ends): + if end is not None and end > self.size: + raise ValueError( + f"buffer ({self.layer}, {self.role}): tensor {t.name!r} reaches byte {end} of " + f"a {self.size}-byte buffer" + ) + if len(tensors) > 1 and None not in ends: + pieces = sorted((t.offset, end) for t, end in zip(tensors, ends)) + bounds = [0, *(b for piece in pieces for b in piece), self.size] + if any(a != b for a, b in zip(bounds[::2], bounds[1::2])): + raise ValueError( + f"buffer ({self.layer}, {self.role}): sections {pieces} do not cover its " + f"{self.size} bytes" + ) + object.__setattr__(self, "tensors", tensors) + + @property + def split(self) -> bool: + """Whether some tensor spreads over parallel dimensions; else every rank holds it alike.""" + return any(t.split for t in self.tensors) + + @classmethod + def unchecked( + cls, + layer: int, + role: str, + pool_group: int, + pool_index: int, + slot_offset: int, + size: int, + tensors: tuple[Tensor, ...] = (), + transferred: bool = True, + expansion: int = 1, + ) -> BufferDesc: + """The constructor's buffer of these values, without its checks and conversions; see + ``PageLayout.unchecked``.""" + buffer = object.__new__(cls) + buffer.__dict__.update( + layer=layer, + role=role, + pool_group=pool_group, + pool_index=pool_index, + slot_offset=slot_offset, + size=size, + tensors=tuple(tensors) or replicated(size), + transferred=transferred, + expansion=expansion, + ) + return buffer + + +@dataclass(frozen=True) +class RankTopology: + """This rank's parallel sizes and ranks, as native RankInfo carries them. + + Attributes: + tp_size: Positive tensor-parallel group size. + tp_rank: Zero-based tensor-parallel rank, less than tp_size. + pp_size: Positive pipeline-parallel group size. + pp_rank: Zero-based pipeline-parallel rank, less than pp_size. + cp_size: Positive context-parallel group size. + cp_rank: Zero-based context-parallel rank, less than cp_size. + dp_size: Positive attention data-parallel group size. + dp_rank: Zero-based attention data-parallel rank, less than dp_size. + attention_dp: Whether attention uses data parallelism; each tensor-parallel rank then + occupies its own attention data-parallel rank. + + Raises: + ValueError: If a size is not a positive integer, a rank is not a non-negative integer, any + such value is a boolean, or a rank lies outside its size.""" + + tp_size: int = 1 + tp_rank: int = 0 + pp_size: int = 1 + pp_rank: int = 0 + cp_size: int = 1 + cp_rank: int = 0 + dp_size: int = 1 + dp_rank: int = 0 + attention_dp: bool = False + + def __post_init__(self) -> None: + for name in ("tp", "pp", "cp", "dp"): + size = _integer(getattr(self, f"{name}_size"), f"{name}_size", 1) + rank = _integer(getattr(self, f"{name}_rank"), f"{name}_rank") + if rank >= size: + raise ValueError(f"{name} rank {rank} of {size}") + object.__setattr__(self, f"{name}_size", size) + object.__setattr__(self, f"{name}_rank", rank) + + @classmethod + def unchecked( + cls, + tp_size: int = 1, + tp_rank: int = 0, + pp_size: int = 1, + pp_rank: int = 0, + cp_size: int = 1, + cp_rank: int = 0, + dp_size: int = 1, + dp_rank: int = 0, + attention_dp: bool = False, + ) -> RankTopology: + """The constructor's topology of these values, without its checks and conversions; see + ``PageLayout.unchecked``.""" + topology = object.__new__(cls) + topology.__dict__.update( + tp_size=tp_size, + tp_rank=tp_rank, + pp_size=pp_size, + pp_rank=pp_rank, + cp_size=cp_size, + cp_rank=cp_rank, + dp_size=dp_size, + dp_rank=dp_rank, + attention_dp=attention_dp, + ) + return topology + + +@dataclass(frozen=True) +class PageLayout: + """How one rank's pages are built and where the rank stands, as immutable values. + + Values contain no addresses or manager ownership. A page uses one slot from each pool of its + device pool group, in pool order. Buffers locate their bytes within those slots. Layer groups + follow the manager's local order, indexed by LentLayerGroupRows.layer_group_index; their layer + ids are global. Mapping inputs are copied into read-only mappings, and layer groups and buffers + into tuples. Equal layouts hash alike and can serve as mapping keys. + + Region maps index Lender.regions, where staging regions precede device regions. A staging slot + concatenates the group's pools' slots in pool order. Regions absent from a lender have no map + entry; both maps are empty outside a lender. + + Attributes: + tokens_per_page: Positive token count of one page. Under Helix, one manager block spans + topology.cp_size pages, one per context-parallel rank. + pool_slot_bytes: Device pool group index -> positive slot byte widths of its pools, in pool + order. + layer_groups: Tuple describing each local layer group's kind, windows, sinks, device pool + group, global layer ids and content shard. + buffers: Tuple of every role buffer of every layer within a page. + topology: This rank's parallel sizes and ranks. + total_kv_heads_by_layer: Global layer id -> positive KV head count over all tensor-parallel + ranks, only for attention layers whose count the manager knows. + dtype: Name of the manager's tensorrt_llm.bindings.DataType member for KV data, such as + "HALF", "BF16", "FP8" or "NVFP4"; None if undeclared. + quant_mode: None for unquantized KV, else "fp8", "nvfp4" or "fp8_ds_mla" (DeepSeek-V4's FP8 + format with its scales within the page). + device_region_index: (pool group, pool index) -> device region index in Lender.regions. + staging_region_index: Pool group -> staging region index in Lender.regions; empty when the + lender stages nothing. + rebuilds_untransferred: Whether the buffers peers do not transfer derive from the others, + which the manager rebuilds after the native transfer imports a request's pages; a + reader that invokes no rebuild, such as the lender, moves them too. + + Raises: + ValueError: If tokens_per_page, a slot width or a KV head count is not a positive integer; a + group, pool, layer or region index is not a non-negative integer; any such value is a + boolean; a buffer's expansion does not divide tokens_per_page; or quant_mode is not + listed above.""" + + tokens_per_page: int + pool_slot_bytes: Mapping[int, tuple[int, ...]] + layer_groups: tuple[LayerGroupDesc, ...] + buffers: tuple[BufferDesc, ...] + topology: RankTopology = RankTopology() + total_kv_heads_by_layer: Mapping[int, int] = field(default_factory=dict) + dtype: str | None = None + quant_mode: Literal["fp8", "nvfp4", "fp8_ds_mla"] | None = None + device_region_index: Mapping[tuple[int, int], int] = field(default_factory=dict) + staging_region_index: Mapping[int, int] = field(default_factory=dict) + rebuilds_untransferred: bool = False + + def __post_init__(self) -> None: + object.__setattr__( + self, "tokens_per_page", _integer(self.tokens_per_page, "tokens_per_page", 1) + ) + pools = { + _integer(g, "pool group"): tuple(_integer(w, "slot width", 1) for w in widths) + for g, widths in self.pool_slot_bytes.items() + } + object.__setattr__(self, "pool_slot_bytes", _frozen(pools)) + heads = { + _integer(layer, "layer"): _integer(count, "KV heads", 1) + for layer, count in self.total_kv_heads_by_layer.items() + } + object.__setattr__(self, "total_kv_heads_by_layer", _frozen(heads)) + device = { + (_integer(g, "pool group"), _integer(p, "pool index")): _integer(i, "region index") + for (g, p), i in self.device_region_index.items() + } + object.__setattr__(self, "device_region_index", _frozen(device)) + staging = { + _integer(g, "pool group"): _integer(i, "region index") + for g, i in self.staging_region_index.items() + } + object.__setattr__(self, "staging_region_index", _frozen(staging)) + object.__setattr__(self, "rebuilds_untransferred", bool(self.rebuilds_untransferred)) + object.__setattr__(self, "layer_groups", tuple(self.layer_groups)) + object.__setattr__(self, "buffers", tuple(self.buffers)) + if self.quant_mode is not None and self.quant_mode not in _QUANT_MODES: + raise ValueError( + f"quant_mode must be None or one of {_QUANT_MODES}, got {self.quant_mode!r}" + ) + for buffer in self.buffers: + if self.tokens_per_page % buffer.expansion: + raise ValueError( + f"buffer ({buffer.layer}, {buffer.role}): expansion {buffer.expansion} does not " + f"divide {self.tokens_per_page} tokens per page" + ) + + @classmethod + def unchecked( + cls, + tokens_per_page: int, + pool_slot_bytes: Mapping[int, tuple[int, ...]], + layer_groups: tuple[LayerGroupDesc, ...], + buffers: tuple[BufferDesc, ...], + topology: RankTopology = RankTopology(), + total_kv_heads_by_layer: Mapping[int, int] = _NO_ENTRIES, + dtype: str | None = None, + quant_mode: Literal["fp8", "nvfp4", "fp8_ds_mla"] | None = None, + device_region_index: Mapping[tuple[int, int], int] = _NO_ENTRIES, + staging_region_index: Mapping[int, int] = _NO_ENTRIES, + rebuilds_untransferred: bool = False, + ) -> PageLayout: + """The constructor's layout of these values, without its checks and conversions. + + Mappings become read-only copies and layer groups and buffers tuples, as in the + constructor; every other value, including those inside layer groups and buffers, is kept + as given. For values the constructor keeps unchanged, such as another layout's fields, the + result equals the constructor's. Values it would refuse or convert are not detected: they + stay as given, the layout then breaks the constraints documented above, and code reading + it may fail or misbehave. Build layouts from values of unknown origin with the constructor. + """ + layout = object.__new__(cls) + layout.__dict__.update( + tokens_per_page=tokens_per_page, + pool_slot_bytes=_frozen(dict(pool_slot_bytes)), + layer_groups=tuple(layer_groups), + buffers=tuple(buffers), + topology=topology, + total_kv_heads_by_layer=_frozen(dict(total_kv_heads_by_layer)), + dtype=dtype, + quant_mode=quant_mode, + device_region_index=_frozen(dict(device_region_index)), + staging_region_index=_frozen(dict(staging_region_index)), + rebuilds_untransferred=rebuilds_untransferred, + ) + return layout + + def __hash__(self) -> int: + values = (getattr(self, f.name) for f in dataclasses.fields(self)) + return hash( + tuple( + tuple(sorted(value.items())) + if isinstance(value, collections.abc.Mapping) + else value + for value in values + ) + ) + + +def _same_group(layout: PageLayout, desc: LayerGroupDesc) -> LayerGroupDesc | None: + """The layer group of ``layout`` holding exactly ``desc``'s layers, of its kind, or None.""" + layers = frozenset(desc.layers) + for other in layout.layer_groups: + if other.kind == desc.kind and frozenset(other.layers) == layers: + return other + return None + + +def _held_indices(layout: PageLayout, desc: LayerGroupDesc) -> tuple | None: + """The indices this rank holds of every split axis of the group's transferred buffers, by + (layer, role, tensor, axis); None when no transferred buffer splits.""" + out = [] + for buffer in layout.buffers: + if ( + buffer.pool_group != desc.pool_group + or buffer.layer not in desc.layers + or not buffer.transferred + or not buffer.split + ): + continue + for tensor in buffer.tensors: + for axis in tensor.axes: + if axis.split: + key = (buffer.layer, buffer.role, tensor.name, axis.name) + out.append((key, held(axis, layout.topology))) + return tuple(sorted(out, key=lambda item: item[0])) if out else None + + +def replica_ranks( + ranks: Sequence[PageLayout], + rank: int, + group: int, + *, + replicated: bool = False, +) -> tuple[int, ...]: + """The flat ranks of one instance holding the same bytes as layer group ``group`` of rank + ``rank``, read from every rank's published layout. + + Ranks hold the same bytes of a layer group when they serve the same requests (the same + attention data-parallel rank); hold the same pages (under Helix an attention group's pages split + by context-parallel rank, while a state group's rows do not); hold the same layers, in a layer + group of the same kind; and hold the same share of the group's content: a buffer every rank + holds alike is the same on every such rank, and split buffers match when the ranks hold the + same indices of every split axis. Without split buffers, the published content shard is the + only basis for matching. + + Args: + ranks: Every rank's layout of one instance, by flat rank. + rank: The flat rank whose layer group is matched. + group: The index of the layer group in that rank's layout. + replicated: Match the group's replicated buffers; otherwise its other buffers. + + Returns: + The matching flat ranks, ascending, ``rank`` among them. + + Raises: + IndexError: ``rank`` or ``group`` names no rank or layer group. + """ + layout = ranks[rank] + desc = layout.layer_groups[group] + topology = layout.topology + paged_by_context = desc.kind == "attention" + held_indices = None if replicated else _held_indices(layout, desc) + out = [] + for other_rank in range(len(ranks)): + other = ranks[other_rank] + other_topology = other.topology + if other_topology.dp_rank != topology.dp_rank: + continue + if paged_by_context and other_topology.cp_rank != topology.cp_rank: + continue + match = _same_group(other, desc) + if match is None: + continue + if not replicated: + other_held = _held_indices(other, match) + if held_indices is None and other_held is None: + if match.shard != desc.shard: + continue + elif other_held != held_indices: + continue + out.append(other_rank) + return tuple(out) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/page_table.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/page_table.py new file mode 100644 index 000000000000..d0bf151a6660 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/page_table.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A ``KVCacheManagerV2``'s page table as plain values: how its pages are laid out (``PageLayout``), +where its device pools lie and what each layer group is; and what a layer config or a rank's place +says about it. The manager builds the table itself (``KVCacheManagerV2.get_page_table``); readers +of the table, such as the lender, never look into the manager.""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Mapping, Sequence + +from .page_layout import PageLayout, RankTopology, ShardDesc + +if TYPE_CHECKING: + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig, SsmLayerConfig + + +@dataclass(frozen=True) +class DevicePool: + """One device pool of pool group ``group``: slot ``s`` lives at ``base + s * slot_bytes``.""" + + group: int + index: int + base: int + slot_bytes: int + num_slots: int + + +@dataclass(frozen=True, eq=False) +class ManagerLayout: + """A manager's page table as plain values, without a manager reference; ``cp`` is (rank, size), + (0, 1) without Helix. ``tokens_per_block`` spans ``size`` physical pages of + ``tokens_per_block // size`` tokens.""" + + page_layout: PageLayout + tokens_per_block: int + pool_group_of: tuple[int, ...] + windows: tuple[int | None, ...] + sink_blocks: tuple[int, ...] + recurrent: tuple[bool, ...] + sparse: tuple[bool, ...] + layers: tuple[tuple[int, ...], ...] + shards: tuple[tuple[int, int], ...] + page_bytes: Mapping[int, int] + device_pools: Mapping[int, tuple[DevicePool, ...]] + cp: tuple[int, int] = (0, 1) + # Per layer group, the window its attention reads, at most the window kept; ``()``: the same. + read_windows: tuple[int | None, ...] = () + # Per layer group, whether it lends by name: none of its layers kept out of prefix reuse. + named_groups: tuple[bool, ...] = field(kw_only=True) + # Per layer group, whether every layer must cover a reused prefix; ``()``: every group's. + prefix_required: tuple[bool, ...] = field(default=(), kw_only=True) + # Whether every windowed layer group's stale range is the runtime's window formula of its + # window and sink blocks at each history, as the manager answered for every history probed. + stale_ranges_by_formula: bool = field(default=False, kw_only=True) + + @property + def num_layer_groups(self) -> int: + """The manager's layer groups.""" + return len(self.pool_group_of) + + @property + def pool_groups(self) -> tuple[int, ...]: + """The device pool groups, sorted; staging regions follow this order.""" + return tuple(sorted(self.device_pools)) + + +def virtual_layers( + attn_to_layer_id: Mapping[tuple[int, enum.Enum], int] | None, +) -> tuple[dict[int, tuple[int, int]], int] | None: + """(internal id -> (model layer, attention type), type count) of a manager keeping several + caches per model layer, from its (model layer, attention type) -> internal id map; None for an + empty map. An internal id shared by several types takes its first: under DeepSeek-V4.1 + long-range layers sharing a compressed cache resolve to the layer that allocated it, which + comes first on every pipeline stage. Every member of the attention-type enum counts, so stages + holding different attention types agree.""" + if not attn_to_layer_id: + return None + types_of: dict[int, tuple[int, int]] = {} + for (model_layer, attn_type), layer in attn_to_layer_id.items(): + types_of.setdefault(int(layer), (int(model_layer), int(attn_type.value))) + kinds = type(next(iter(attn_to_layer_id))[1]) + return types_of, max(int(member.value) for member in kinds) + 1 + + +# -- what a layer config says --------------------------------------------------------------------- + + +def layer_configs_by_id(config: object) -> dict[int, AttentionLayerConfig | SsmLayerConfig]: + """Internal layer id -> layer config of a runtime config; a layer without an explicit id uses + its position in the config.""" + return { + int(getattr(layer, "layer_id", position)): layer + for position, layer in enumerate(config.layers) + } + + +def is_recurrent_layer(layer_config: AttentionLayerConfig | SsmLayerConfig) -> bool: + """Whether a layer config of the runtime describes a recurrent-state layer.""" + from tensorrt_llm.runtime import kv_cache_manager_v2 as runtime + + return isinstance(layer_config, runtime.SsmLayerConfig) + + +def is_private_layer(layer_config: AttentionLayerConfig | SsmLayerConfig) -> bool: + """Whether the native lifecycle excludes this layer from prefix reuse.""" + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy + + return getattr(layer_config, "reuse_policy", None) == AttentionReusePolicy.PRIVATE + + +def is_non_required_layer(layer_config: AttentionLayerConfig | SsmLayerConfig) -> bool: + """Whether prefix coverage is optional or private for this layer.""" + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy + + return ( + getattr(layer_config, "reuse_policy", AttentionReusePolicy.REQUIRED) + != AttentionReusePolicy.REQUIRED + ) + + +def has_sparse_buffers(layer_config: AttentionLayerConfig | SsmLayerConfig) -> bool: + """Whether a layer config of the runtime has sparse buffers (``is_sparse``), whose read-only + pages a cache can lock in host memory.""" + return any(getattr(b, "is_sparse", False) for b in getattr(layer_config, "buffers", ()) or ()) + + +def buffer_expansions( + layer_config: AttentionLayerConfig | SsmLayerConfig, tokens_per_block: int +) -> dict[str, int]: + """Role -> sub-pages per block, from the layer's ``tokens_per_block_override``.""" + out = {} + for buffer in getattr(layer_config, "buffers", ()) or (): + override = getattr(buffer, "tokens_per_block_override", None) + if override is not None: + if tokens_per_block % int(override): + raise ValueError( + f"buffer {buffer.role}: {override} tokens per block does not divide " + f"{tokens_per_block}" + ) + out[str(buffer.role)] = tokens_per_block // int(override) + return out + + +# -- where a rank stands and what it holds -------------------------------------------------------- + + +def rank_topology(mapping: object) -> RankTopology: + """The rank's place as native's ``RankInfo`` reads it from the mapping: under attention data + parallelism each tensor-parallel rank is a data-parallel rank of its own.""" + if mapping is None: + return RankTopology() + adp = bool(getattr(mapping, "enable_attention_dp", False)) + tp_size, tp_rank = int(getattr(mapping, "tp_size", 1)), int(getattr(mapping, "tp_rank", 0)) + return RankTopology( + tp_size=tp_size, + tp_rank=tp_rank, + pp_size=int(getattr(mapping, "pp_size", 1)), + pp_rank=int(getattr(mapping, "pp_rank", 0)), + cp_size=int(getattr(mapping, "cp_size", 1)), + cp_rank=int(getattr(mapping, "cp_rank", 0)), + dp_size=tp_size if adp else int(getattr(mapping, "dp_size", 1)), + dp_rank=tp_rank if adp else 0, + attention_dp=adp, + ) + + +def attention_shard(tp_size: int, tp_rank: int, heads: int, total: int | None) -> ShardDesc: + """The share of a head-sharded attention buffer a rank holds: ``total`` heads over ``tp_size`` + ranks, rank ``r`` holding ``heads`` from ``r * total // tp_size``.""" + if tp_size == 1: + return ShardDesc() + if heads <= 0 or total is None or total <= 0: + return ShardDesc(tp_size, tp_rank) + if any(rank * total // tp_size % heads for rank in range(tp_size)): + # Some rank's heads straddle two shares, so each rank's buffer is a share of its own. + return ShardDesc(tp_size, tp_rank) + # total // heads distinct shares; fewer heads than ranks are repeated, so a single head is one + # share that every rank holds. + count = max(total // heads, 1) + return ShardDesc(count, min((tp_rank * total // tp_size) // heads, count - 1)) + + +# A layer group is whole when every transferred buffer holds the same bytes on every rank +# (replicated, or one KV head shared by all ranks); else it is the share this rank holds. +def group_shard(shards: Sequence[ShardDesc], fallback: ShardDesc) -> ShardDesc: + parts = {s for s in shards if s.count != 1} + if not parts: + return ShardDesc() + return parts.pop() if len(parts) == 1 else fallback diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/__init__.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/__init__.py index 0a10fed3261a..242733e43cb7 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/__init__.py @@ -12,30 +12,27 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Lending a KV cache manager v2's blocks to transfer backends. - -This module is the whole API: callers import only its names, the package's other modules are -private, and ``attach_lender`` selects the implementation. The docstrings of these names state every -caller obligation. - -Roles: the integrator attaches one lender to a manager; the holder lends, polls, marks arrivals and -releases; the backend moves the bytes; the waiter asks ``Lender.resume_bounds``. Lender, -lease and hold methods run only on the manager's thread, one call at a time: the builder, then the -executor loop, then the shutdown thread. Positions are tokens. Ready views describe stable, -accessible lent memory; leases and region holds delimit backend access. Cache ownership remains -with the manager. -""" +"""Public KV-cache lending API; other modules are private. Positions are tokens. +The integrator attaches; orchestration lends; holders poll and return leases; backends move ready +bytes. Calls run serially on the manager thread through shutdown; leases/holds bound access; +managers own caches.""" import typing as _typing from ._types import ( - LayerGroupBlocks, + ALL_ARRIVED, + Arrived, + HeldPages, + HoldReason, Lease, LeaseView, Lender, + LentLayerGroupRows, + MemoryKind, MemoryRegion, RegionHold, ResumeBounds, + RowSelection, StagingOptions, ) @@ -43,149 +40,188 @@ from ..kv_cache_manager_v2 import KVCacheManagerV2 __all__ = [ - "LayerGroupBlocks", + "ALL_ARRIVED", + "Arrived", + "HeldPages", + "HoldReason", + "LentLayerGroupRows", "Lease", "LeaseView", "Lender", + "MemoryKind", "MemoryRegion", "RegionHold", "ResumeBounds", + "RowSelection", "StagingOptions", "attach_lender", ] def attach_lender(manager: "KVCacheManagerV2", *, scope: bytes, staging: StagingOptions) -> Lender: - """Attach the manager's one lender, which relays whole blocks through host staging. - - The lender copies blocks between a request's device pages and host slots on the manager's - execution stream, for backends that can only read and write host memory, and says where the - request may resume after a fetch. ``Lender`` states what its holder and backends keep. A - name covers ``scope``, the committed tokens up to its block's end, their multimodal digests and - the manager's reuse scope, nothing else a request carries, so the decoder self-attention pool of - an encoder-decoder model is not lent by name. - - Caller obligations: - - This implementation reserves the request's device pages before waiting for host regions. - Keep room for running requests in each pool group while fetches are parked: their locked - pages, other unreclaimable pages, the largest generation request's full resume window and - its next step must fit within ``max_util_for_resume``. V2 cannot evict or preempt a parked - fetch. Reserving space for parked fetches alone does not prevent scheduler deadlock. - - Copies use the manager's execution stream, ordered with its forward work. Any writer of - recycled pages on another stream must wait for that stream first. The diagnostic fresh- - page fill (``TRTLLM_KV_FRESH_PAGE_FILL``) already does so and then synchronizes its own - stream; a fetch that grows pages with this diagnostic enabled can wait on the CPU. - Pageable memory under confidential computing can also make a call wait for its stream. - - A failed copy whose completion cannot be established quarantines its slots permanently. - Capacity can erode until a pending lease cannot fit and blocks those behind it. Release - leases that wait too long or free their requests. Quarantined slots, unreleased leases - or region holds, and cleanup errors can keep host memory until process exit. - - From ``lend_write`` until the request resumes or its cache is dropped, keep it among the - executor's active requests and unscheduled until the resume state is applied. To resume at - ``p``, set ``py_served_prefix_end=p``, raise the cache history to ``p`` if necessary and, - back in context state with a chunk spanning the remaining prompt, call - ``set_prepopulated_prompt_len(p, tokens_per_block)``. It moves the prepopulated length - and context position together; moving either alone is not sufficient. A resume at 0 - from a later context position requires dropping the cache and recomputing. - - With routed-expert capture, fetch only before the first context step, not after a - recompute pause; otherwise skipped positions can have no recorded routes. A later - fetch also does not change the response's original ``cached_tokens`` count. - - Under attention data parallelism without a transceiver, do not park every active - request on a rank unless the resulting all-rank forward stall is acceptable. - - SWA scratch reuse must be off for fetches. A window must not abandon an uncomputed block - still held for prefix reuse; with extra speculative KV tokens, keep later consecutive - segments shorter than ``window + tokens_per_block - 1`` or compute the remainder. - Requests requiring context logits, prompt logprobs or additional per-token model - outputs cannot be fetched, because those outputs are not restored by this implementation. - - These device-page reservations, execution-stream ordering, diagnostics, slot quarantine - and executor-specific constraints belong to the current implementation and may change - when its backing mechanism changes; the ``Lender`` contract does not expose them. - - Attach on the thread that builds the executor, to the manager the executor serves. A - manager built only to estimate the KV cache size shuts down first, after which a lender - attached there lends nothing; that shutdown frees the staging memory, or keeps it until - the process exits where a lease or staging hold is still open then, a slot was lost to a - failed copy, an error in a free stopped the lender (``Lease``), or an error arose in the - lender's own work within that shutdown. - - Shut the manager down last, with the device's CUDA context used for the attach current - on the calling thread, after the shutdown steps ``Lender`` lists. A manager - that never shuts down keeps the staging memory and its own page-index host buffer until - the process exits, and a shutdown that raises keeps that buffer until a retried shutdown - returns, or else until the process exits: the runtime's shutdown raises while any cache is - open, also one whose close raised in a free, unless the runtime is poisoned. Its leases, - released ones too, and the lender's records also keep the requests' caches, and with them - the manager's device pools, until the leases and the lender are dropped; drop the last - reference on a thread that has used the device's CUDA context. - - For a request whose one-model draft has a joint-reuse pool of its own, attach a lender to - both managers, the draft pool's under a ``scope`` of its own (laid out alike, both pools - would name blocks alike), fetch through both, since a fetch fills only this manager's - blocks, and resume as ``ResumeBounds`` states, keeping what ``Lender`` states for - such a pair. - - A caller serving requests the input processor may treat differently keeps them apart with - per-item ``multi_modal_uuids``, which names and the multimodal encoder cache both cover, - since a name trusts the multimodal digests the input pipeline computes, as the manager's - own prefix reuse does. A digest covers an item's bytes, not how the input processor treats - them, so blocks with different KV bytes can share a name: a NumPy array and a - ``torch.Tensor`` with the same dtype, shape and bytes share a digest though some input - processors rescale only one, and a request's ``mm_processor_kwargs`` reach the processor - but not the digest. A ``cache_salt`` covers names but not that cache, which some models - keep by default and which keys an item's embedding by its digest and - ``mm_processor_kwargs``: salted requests can still share an embedding, so a ``cache_salt`` - alone keeps them apart only with that cache off - (``multimodal_config.encoder_cache_max_bytes=0``). - - Under ``SaveHiddenStatesDecodingConfig`` a fetch restores only this manager's KV blocks, - so the hidden states saved for a resumed request pair its whole prompt with the rows of - the tokens it computed, then stale rows, misaligned as after the manager's own prefix - reuse: the caller fetches into no such request, or accepts that. - - An attach that raises for a reason its docstring does not list may keep the staging memory - it allocated until the process exits. + """Attach a lifetime lender; encoder-decoder models, including CROSS caches, are unsupported. + + HOST supports NAMED; DEVICE supports NAMED and ALL. HOST/ALL or HOST without staging raises + ``ValueError``. Host regions precede device regions; ``Lender.page_layout`` locates each pool. + Import layout types from ``tensorrt_llm._torch.pyexecutor.kv_cache.page_layout``. + + Names cover ``scope``, committed prefix tokens, multimodal digests and reuse scope, excluding + other request inputs. Naming requires committed blocks, a lossless cold-page codec, no model + recurrent state or context parallelism, and no one-model draft reading past a block's end. + Otherwise only DEVICE/ALL is available. A HOST region establishes named-lending support; + DEVICE naming follows these conditions. After reuse reset, NAMED returns failed leases while + ALL remains available. A loading KV connector disables staging, but does not itself forbid + DEVICE/NAMED. Eagle and MTP-Eagle look ahead one token, vanilla MTP its draft length; PARD, + DFlash and DSpark read none. Checks include shared draft layers, joint-reuse pools and targets + whose draft layers live elsewhere. Unknown upstream lookahead counts as zero, including + DraftTarget: its target names blocks, its unpaired noncommitting draft does not. + + Contract obligations: + These caller duties are part of the public lending contract; a lending method a later + optional protocol adds states the duties of the leases it lends. + + - Follow ``Lender``'s thread, progress, registration and shutdown duties. A lease binds to + the cache its request holds at the lend: do not run the request on a cache with an open + in-place write or recurrent-state lease that ``lend_read`` or ``lend_write`` lent. + Freeing such a cache retires it: its leases keep its pages protected until they are + returned and their rows count nothing, while the request may run on a new cache, the + pages then held twice (``Lender.held_page_counts``). Reads may overlap later steps; + access only returned rows, whose memory remains protected until the lease is returned. + - Attach on every cache-holding rank or none. Each pipeline stage lends only its own + layers; ``lend_read`` and ``lend_write`` lend them after the stage processed the step. + Aggregate attach outcomes across all ranks: any failure fails startup, since an attached + lender cannot be detached. Model-wide exclusions are checked consistently across stages; + ``max_bytes`` sufficiency is stage-local and can differ. For each distributed transfer, + require delivery on every stage, return every stage's leases in the same iteration + using the transfer path's allgather, and intersect bounds. Before ``return_writes``, + settle the transfer's outcome on every participating rank and finish all backend access + to its views. An arrival notification alone is insufficient; a HOST return may enqueue + a copy from staging immediately. + - Coordinate scheduling across cache-holding ranks. Protect every rank's locked requests + from eviction and preemption, and run a batch only after every rank has allocated it. + If any rank cannot allocate, roll back or retry consistently across ranks before forward. + A pipeline root's scheduling result does not establish admission on the other stages; + an integrator attaching lenders to a pipeline executor must supply this coordination. + - Apply settled bounds before scheduling a request fetched through ``lend_write``. An + empty intersection requires dropping its fetched caches and computing from zero. Fetch + both participating joint-reuse target/draft pools before advancing their shared cursor; + P/D receives that lend no draft pool do not intersect its bounds. Separate draft pools + need a lender and distinct scope. + - ALL arrivals need every transferred buffer's bytes this request will read. See + ``Lender.return_writes`` for permitted omissions; return a partially filled receive + with ``reusable=False``, and the lender keeps its cache out of reuse, as it does one + missing window margins. NAMED always transfers the whole object. + - Treat failed leases and empty bounds as request-level outcomes, including permanent local + states such as reuse reset. Validation ``ValueError``/``TypeError`` from lends and + returns and misuse-related ``RuntimeError`` from lease methods and returns are caller + errors. Entry validation in ``return_writes`` changes nothing: correct the entries and + retry; lending, progress and bounds queries remain permitted. Any error once returning + entries begins permits cleanup only, even if other entries returned successfully. + Exception type alone does not establish that validation failed; if unsure, treat it as + an operation failure. ``RuntimeError`` from ``lend_*`` may also report a manager or CUDA + operation failure. After an operation failure + or an unlisted exception, stop lending, progress, lease polling and bounds queries; return + leases and release holds after access ends, treat parked fetch intervals as empty, drop + their caches and compute from zero. + - Do not subtract ``held_page_counts().deduct`` from a V2 manager's admission capacity: a + ``DEFERRED_CLOSE`` page stays with the cache it belongs to and is already outside the + free pages, so subtracting it counts it twice. Fetches can hold device pages while + waiting for host slots. Reserve the entire device allocation of every live cache that + locks prevent from being preempted, plus ``DEFERRED_CLOSE`` pages, the largest generation + resume window and its next step, within each pool group's ``max_util_for_resume``. + ``held_page_counts().unreclaimable`` counts lease rows and deferred-close pages; it omits + unleased pages of live locked caches, so it is insufficient for this reservation. + A parked-fetch quota alone is insufficient. Use + ``held_page_counts`` per pool group only to check that ranks and pipeline stages hold + alike, and warn where they do not. + - Use per-item ``multi_modal_uuids`` to distinguish processor-dependent inputs. Digests trust + input-pipeline bytes, like prefix reuse, and omit differences such as NumPy/Tensor rescaling + or ``mm_processor_kwargs``. ``cache_salt`` separates names but not encoder embeddings: + salt alone suffices only with ``multimodal_config.encoder_cache_max_bytes=0``. + + Implementation constraints: + The V2 recovery recipe and the following reservations, streams, diagnostics, quarantine + and executor constraints depend on this implementation and are not frozen API obligations. + + - Attach on the construction thread. Remote-cache configuration alone does not attach a + lender. A lender cannot lend after its manager shuts down. + - Host staging copies whole named rows on the manager stream; connector writes are unordered + with these copies. Each device pool is one region, fixed for the lender's life: do not + rebalance the pools of a manager with a lender attached. + - Before this worker finishes context, ALL reads end at finished whole blocks and in-place + reads omit groups whose window is shorter than that context, with SWA scratch reuse on or + off; read again after context finishes to obtain them. Recurrent groups lend one state row + at prompt end, between context and generation; their fetch resumes there or nowhere. Under + Helix, each rank lends pages of ``PageLayout.tokens_per_page`` tokens; starts align to + ``cp_size * tokens_per_page`` and ``total_input_len_cp`` is the whole prompt length. + - Locks prevent eviction, preemption, scheduler suspension and deadlock reports; history + cannot abandon locked pages. Allocation-failure rollback in ``revert_allocate_context`` + and first-block ``resize_context`` skips suspension, leaving locked caches active; + explicit suspension and admission paths raise ``RuntimeError`` on locks. Overlapping + commits retry after unlock; freed caches close at their last unlock, and a reset with such + caches open raises ``RuntimeError``. Manager query hooks propagate original exception + types. These V2 errors report calling order, not a failed lease or poisoned domain. + - With a connector, fetch at or beyond the served prefix; connector loads start beyond a + fetched prefix. Reads reject prefix reservations and wait for pending loads, during which + windowed NAMED fetches reject history advancement. Lend outside batch preparation: an + accepted load can be abandoned before forward. + - Keep fetched requests active but unscheduled. To resume at ``p``, set + ``py_served_prefix_end=p``, raise cache history to ``p`` if needed, restore context state + with a chunk covering the remaining prompt, then call + ``set_prepopulated_prompt_len(p, tokens_per_block)`` to move both cursor fields. This + needs ``p`` below prompt length (``total_input_len_cp`` under Helix). A full ALL fetch + into generation-init resumes at prompt end as after a receive. Rewinding to zero requires + dropping and recomputing. + - Routed-expert capture allows fetching only before the first context step, never after + recompute pause; skipped routes are not restored. Later fetches leave ``cached_tokens`` + unchanged. With attention data parallelism and no transceiver, parking every active + request on one rank can stall all ranks' forward. + - Disable SWA scratch reuse for fetches. Retain uncomputed prefix-reuse blocks; with extra + speculative KV tokens, keep later consecutive segments shorter than + ``window + tokens_per_block - 1`` or compute the remainder. Fetches do not restore context + logits, prompt logprobs or additional per-token outputs. With + ``SaveHiddenStatesDecodingConfig``, exclude fetches or accept saved states pairing the + full prompt with computed then stale rows. + - Writers recycling pages off the manager stream must wait for it. Diagnostic + ``TRTLLM_KV_FRESH_PAGE_FILL`` waits and synchronizes its own stream, so growth can block + the CPU, as confidential-computing pageable copies can. Shutdown waits for recorded + copies; the first runtime cache close can synchronize the device. Unknown copy completion + quarantines slots permanently: return leases waiting too long or free their requests. + - Shut down with the attach device's CUDA context current. Without shutdown, staging and + the manager's page-index host buffer survive until exit; failed shutdown retains that + buffer until retry succeeds or exit. Unless poisoned, runtime shutdown raises with any + cache open, including failed closes. Open locks/device holds retain runtime and pools; + open leases/holds, lost slots and cleanup errors can retain host memory until exit. + Even returned leases and lender records retain caches/pools: drop their last references + on a thread that has used the device's CUDA context. Unlisted attach errors may retain + allocated staging until exit. Args: - manager: A ``KVCacheManagerV2`` with a ``mapping`` that commits blocks to its prefix-reuse - tree: block reuse on, and joint reuse for a draft manager; not built for a one-model - draft that reads prompt tokens ahead, and without a KV cache connector. Subclasses are - accepted; one whose ``try_commit_blocks`` keys blocks from other tokens or another reuse - scope than the base manager's changes names silently, at worst putting wrong bytes under - a correct name. - scope: Equal exactly where KV bytes mean the same; at most 65535 bytes. The lender only puts - it, unread, into every name's namespace, so different scopes never share a name. It must - cover every configuration that changes a block's bytes beyond what the manager declares: - the model, its weights and LoRA adapters (the adapter an id maps to), the one-model - draft whose pool it is (its weights, mode and the target layers it reads), element types - and quantization, the attention (windows such as ``max_attention_window``, sinks, sparse - attention), the attention backend's K/V arrangement (head-major or token-major; the base - manager declares head-major whichever backend writes its pages), and a release that - changes what KV bytes mean. An adapter enters a name only through the request's - ``lora_task_id``: instances sharing a ``scope`` map each id to the same adapter, or else - lend no LoRA request by name or take different scopes. The framework's assembly builds - it from its configuration, so no backend guesses which settings it needs. - staging: The staging size, in whole fetches. + manager: A ``KVCacheManagerV2`` with a mapping. Naming requires block reuse, joint reuse + for a draft, and the restrictions above. Subclasses must preserve the committed tokens + and reuse scope used by ``try_commit_blocks`` or names can identify wrong bytes. + scope: At most 65535 bytes, equal exactly where KV bytes mean the same; different scopes + never share names. Cover model/weights, LoRA id-to-adapter mapping, one-model draft + weights/mode/target layers, dtypes/quantization, attention/windows/sinks/sparsity, + backend K/V arrangement and incompatible releases. The base manager declares head-major + regardless of writing backend. Equal scopes must map each ``lora_task_id`` to the same + adapter, or exclude LoRA from names. Framework assembly derives scope; backends do not. + staging: Host capacity in fetch units (``StagingOptions``). Returns: - The lender, installed on ``manager`` for the manager's life. + The lender installed for the manager's life. Raises: - TypeError: ``manager`` is not a ``KVCacheManagerV2`` or has no mapping, ``scope`` is not - ``bytes``, or ``staging`` is not ``StagingOptions``. - ValueError: A lender is attached already; the manager is context-parallel (Helix included, - ``mapping.cp_size > 1``) or pipeline-parallel (``mapping.pp_size > 1``), holds recurrent - state, commits no blocks (block reuse off, or a draft manager without joint reuse), - since a publish lends only committed blocks, has a KV cache connector or a lossy - cold-page codec (``quantization_for_cold_page``), or is built for a - one-model draft that reads prompt tokens ahead; it has a layer group of sparse buffers - (``is_sparse``), whose read-only pages a cache can lock in host memory, though sparse - attention whose buffers are not marked sparse is not refused by this; its layer groups, - layer configs and declarations do not agree; it has no pages to stage; ``scope`` is - longer than 65535 bytes; ``staging.max_bytes`` is below one fetch. One-model Eagle and - MTP-Eagle read 1 prompt token ahead, vanilla MTP its draft length, and the check reads - the manager, not ``is_draft``, so it refuses draft layers sharing the target's manager, - a joint-reuse draft pool, and the target's manager of such a draft whose layers live - elsewhere. Drafts reading no token ahead (PARD, DFlash, DSpark) and managers without - speculative decoding are accepted. A draft whose read-ahead upstream has not established - counts as reading none: for one-model DraftTarget, staging accepts the target's manager - and refuses its unpaired draft pool, which commits no blocks, so a fetch fills only the - target's blocks, as the manager's own prefix reuse does. - MemoryError: Page-locking the staging memory failed. - """ - from ._lender import _attach_staging as _attach + TypeError: Manager is not V2 or lacks a mapping; scope is not bytes or staging is not + ``StagingOptions``. + ValueError: Already attached; CROSS cache; non-Helix context parallelism; model windows or + recurrent state under Helix; model sparse buffers (``is_sparse``), whose pages can live + in host memory; a non-REQUIRED attention group without a window; inconsistent groups, + configs or declarations, or a missing ``Role.ALL`` mapper declaration; no pages to + stage; ``max_bytes`` below one local fetch; a malformed region, or two of one memory + kind that overlap; or scope over 65535 bytes. Unmarked sparse-attention buffers are not + rejected as sparse. + MemoryError: Staging page-lock allocation failed.""" + from ._lender import _attach_lender as _attach return _attach(manager, scope=scope, staging=staging) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_fetch_ledger.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_fetch_ledger.py index ec41ad58d56e..50669190d83d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_fetch_ledger.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_fetch_ledger.py @@ -12,15 +12,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""A staging lender's fetch ledger: per request, the latest fetch, what fetches delivered, the -prefix a fetch kept and a cache whose committed tail a failed copy may have mixed, and the -resume bounds they give.""" +"""Per-request fetch deliveries and safe resume intervals.""" from __future__ import annotations import enum -from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable +from dataclasses import dataclass, field +from heapq import merge +from typing import TYPE_CHECKING, Callable, Sequence import numpy as np @@ -30,122 +29,319 @@ if TYPE_CHECKING: from tensorrt_llm.runtime.kv_cache_manager_v2 import _KVCache - from ...llm_request import LlmRequest from ..kv_cache_manager_v2 import KVCacheManagerV2 from ._layout import ManagerLayout - from ._lender import _LentBlockEntries, _QueuedCopy, _StagingLease + from ._lender import _LeaseBase, _LentBlockEntries, _QueuedCopy -@dataclass(eq=False) +# Read-only masks of whole deliveries' blocks by run, shared by records, which write a mask only +# through ``_writable_mask``; emptied when the masks hold ``_RUN_MASK_BYTES``. +_run_masks: dict[tuple[int, int], np.ndarray] = {} +_RUN_MASK_BYTES = 1 << 22 +_run_mask_bytes = 0 +# The run of a layer group that lent no block. +_NO_RUN = (0, 0) +# Builds a named tuple as the tuple it is, without its Python-level constructor. +_new_tuple = tuple.__new__ + + +def _run_mask(first: int, last: int) -> np.ndarray: + """The shared read-only mask of ``last`` blocks, true from ``first`` on.""" + global _run_mask_bytes + found = _run_masks.get((first, last)) + if found is None: + if _run_mask_bytes + last > _RUN_MASK_BYTES: + _run_masks.clear() + _run_mask_bytes = 0 + found = np.zeros(last, dtype=bool) + found[first:] = True + found.flags.writeable = False + _run_masks[first, last] = found + _run_mask_bytes += last + return found + + +def _writable_mask(masks: list[np.ndarray], lg: int) -> np.ndarray: + """``masks[lg]`` to write, first replaced by a copy of its own where it is read-only.""" + blocks = masks[lg] + if not blocks.flags.writeable: + blocks = masks[lg] = blocks.copy() + return blocks + + +@dataclass(eq=False, slots=True) class _FetchRecord: - """One write lease's fetch into one cache: ``COPY_PENDING`` once the copy of its marked entries - is queued without error, ``COPY_DONE`` once that copy has completed, or at once where its marks - queue no copy and the lease has none. ``_FetchLedger._fetch_state`` derives its ``_FetchState`` - from ``marks_queued``, ``copy`` and whether ``_FetchLedger._latest_fetches`` holds it.""" + """One fetch into a cache, with mark/copy completion and its declared token extent.""" start: int kv: object marks_queued: bool = False copy: _QueuedCopy | None = None + # Where the rows of recurrent state the lease lends stand, its end; None where it lends none. + state_end: int | None = None + # Where the fetch ends; None where unknown. + end: int | None = None + entries: _LentBlockEntries | None = None + # The request's prompt length and multimodal runs at the lend, which its return's bounds read. + prompt: int = 0 + runs: list[tuple[int, int]] | None = None class _FetchState(enum.Enum): - """A fetch's state, which ``_FetchLedger._fetch_state`` derives. - ``_FetchLedger._latest_fetches`` holds each request's latest fetch until it is abandoned, the - request is freed, or resume_bounds, while the lender serves the manager and is not stopped, - finds the request holding another cache. ``Staging._apply_marks`` takes it out - (``_take_out_fetch``) before queuing the copy and puts it back (``_put_back_fetch``) with - ``marks_queued`` set once the copy is queued without error.""" - - # In _latest_fetches, marks not queued: resume_bounds is None and a new fetch into the cache fails - # at the call. Staging._apply_marks sets marks_queued, or abandons it where the copy fails or - # the call raises; _abandon_fetch abandons it too. + """Derived lifecycle of the request's latest fetch.""" + + # In _latest_fetches with marks not queued: poll_resume_bounds is None and a new fetch fails. + # _apply_marks sets marks_queued, or abandons it (failed copy, raise) as _abandon_fetch does. AWAITING_MARKS = "awaiting_marks" # In _latest_fetches, its marked entries' copy queued without error and pending: as - # AWAITING_MARKS for resume_bounds and a new fetch. The copy's completion makes it COPY_DONE. + # AWAITING_MARKS for poll_resume_bounds and a new fetch. The copy's completion makes it COPY_DONE. COPY_PENDING = "copy_pending" - # In _latest_fetches, its copy complete or none queued: resume_bounds answers, and a new fetch into + # In _latest_fetches, its copy complete or none queued: poll_resume_bounds answers, and a new fetch into # the cache may start. The request's next fetch replaces it. COPY_DONE = "copy_done" # Out of _latest_fetches before its marks were queued: none of its entries count. ABANDONED = "abandoned" - # Out of _latest_fetches after its marks were queued: the request's next fetch replaced it, the - # request's free dropped it, or resume_bounds dropped it as a record of a cache the request no - # longer holds. Replacing it drops none of its entries from the cache's deliveries. + # Out of _latest_fetches after its marks were queued: replaced by the next fetch, or dropped at + # the free or by poll_resume_bounds for a replaced cache. Its entries stay delivered. RETIRED = "retired" -@dataclass(eq=False) +@dataclass(eq=False, slots=True) class _Deliveries: - """What the fetches into one cache delivered: per layer group, by block ordinal, the entries - whose last copy from staging was queued without error and that no shrink freed since; - ``lowest_start`` is the lowest fetch start.""" + """Delivered rows and valid token prefixes for one cache; state holds at one endpoint.""" kv: object lowest_start: int - # Index by local layer group, then logical block ordinal (not LeaseView tuple positions). - # Recorded when a copy queues successfully; resume bounds also wait for its completion. + # By local layer group (not view position), then block ordinal; set once a copy queues without + # error, the bounds also awaiting its completion. Recurrent state records nothing here. delivered_block_masks: list[np.ndarray] - # (tokens computed before, max_position, the end of the last run below it), last computed - resume_bounds_memo: tuple[int, int, int] | None = None + # (tokens computed before per layer group, max_position, lowest safe position), last computed + resume_bounds_memo: tuple[tuple[int, ...], int, int] | None = None + # Where the recurrent state the latest fetch delivered stands, or None: none delivered, or a + # later fetch lent its rows and did not deliver them all. + state_end: int | None = None + # Settled fetch extent, including fetches whose margin rows were all absent. + end: int = 0 + # Partial physical pages: (layer group, logical block) -> valid token end. + partial_ends: dict[tuple[int, int], int] = field(default_factory=dict) -# TODO: each per-request record keeps its own reference to the cache it describes, so every -# reader compares it again, and a record of a replaced cache keeps that cache until the record is -# dropped or replaced. +# TODO: each per-request record holds its own cache reference, so every reader compares it again and +# a record of a replaced cache keeps that cache until the record is dropped or replaced. class _FetchLedger: - """A staging lender's per-request fetch records and the resume_bounds they give. ``Staging`` - changes the records through these methods alone and reads them through these methods, apart - from the fields of the ``_Deliveries`` that ``_deliveries_of`` returns.""" + """Per-request fetch state and resume bounds, changed through these methods.""" def __init__( self, layout: ManagerLayout, copy_done: Callable[[_QueuedCopy | None], bool] ) -> None: self._layout = layout + recurrent = getattr(layout, "recurrent", ()) + # The layer groups of recurrent state, which hold one row each, valid at one position. + self._recurrent_groups = frozenset(lg for lg, flag in enumerate(recurrent) if flag) + groups = getattr(layout, "num_layer_groups", len(recurrent)) + self._paged_groups = tuple(lg for lg in range(groups) if lg not in self._recurrent_groups) + # What the closed form of a whole delivery's bounds reads; it applies without Helix and + # where every window's stale range is the sinkless window formula. + self._tokens_per_block = int(getattr(layout, "tokens_per_block", 1)) + # The highest history the window formula reads, as the manager reads it: ``_MAX_TOKENS`` + # less one block. + self._longest = _manager._MAX_TOKENS - self._tokens_per_block + windows = tuple(getattr(layout, "windows", ())) + sinks = tuple(getattr(layout, "sink_blocks", ())) + self._window_of = tuple( + None + if window is None or lg in self._recurrent_groups + else (int(window), _ranges._read_slack(layout, lg)) + for lg, window in enumerate(windows) + ) + self._stale_windows = tuple(int(window) for window in windows if window is not None) + self._closed_form = ( + getattr(layout, "cp", (0, 1))[1] == 1 + and len(windows) == groups + and ( + not self._stale_windows + or ( + bool(getattr(layout, "stale_ranges_by_formula", False)) + and all( + lg < len(sinks) and not sinks[lg] + for lg, window in enumerate(windows) + if window is not None + ) + ) + ) + ) # Whether a copy has completed, as the lender answers: each event asked at most once per # round of its progress. self._copy_done = copy_done # Per-request fetch records, each with the cache it describes: the request's free drops all - # four, and resume_bounds drops one of a cache the request no longer holds. + # four, and poll_resume_bounds drops one of a cache the request no longer holds. self._latest_fetches: dict[int, _FetchRecord] = {} # the latest fetch into each request - self._deliveries: dict[int, _Deliveries] = {} - # Requests whose history was past the committed tokens after a fetch's grow: the cache it - # grew, what the latest such fetch kept of the prefix computed before it, and the history - # its grow left. - self._kept_prefixes: dict[int, tuple[object, int, int]] = {} + # What fetches delivered into each request's cache, as ``_deliveries`` reads it once the + # whole deliveries in ``_pending`` are written into it. + self._delivered: dict[int, _Deliveries] = {} + # Whole deliveries ``_defer_whole_delivery`` settled and no reader has needed yet, by + # request: ``[fetch, spans, resume bounds memo]``; the request's free drops them. + self._pending: dict[int, list] = {} + # Requests a fetch left with history past a group's committed tokens: the cache, per layer + # group the computed prefix the latest such fetch kept, and the history it left. + self._kept_prefixes: dict[int, tuple[object, tuple[int, ...], int]] = {} # Requests whose committed tail block a failed copy may have left with the fetch's bytes in # some pools and the bytes it held before in others: the cache. ResumeBounds is empty for it. self._mixed_tail_caches: dict[int, object] = {} + @property + def _deliveries(self) -> dict[int, _Deliveries]: + """What fetches delivered into each request's cache, by request: ``_delivered``, with every + whole delivery ``_defer_whole_delivery`` holds written into it first.""" + if self._pending: + self._write_pending() + return self._delivered + + def _write_pending(self) -> None: + """Write the records of the whole deliveries ``_defer_whole_delivery`` holds; one that + raises leaves the rest held.""" + pending = self._pending + while pending: + request_id = next(iter(pending)) + fetch, spans, memo = pending.pop(request_id) + self._record_whole_delivery(request_id, fetch, spans, memo) + def _record_fetch( - self, request_id: int, kv: _KVCache, start: int, kept: int, position: int, committed: int + self, + request_id: int, + kv: _KVCache, + start: int, + kept: Sequence[int], + position: int, + committed: Sequence[int], + state_end: int | None = None, + end: int | None = None, + prompt: int = 0, + runs: list[tuple[int, int]] | None = None, ) -> _FetchRecord: - """Record the request's latest fetch, from ``start`` into ``kv``, and, where its grow left - the history at ``position`` past the ``committed`` tokens, ``kept``: what the fetch keeps of - the prefix computed before it. Returns the fetch's record.""" - if position > committed: + """Record and return fetch ``[start, end)`` into ``kv``; growth leaving the history at + ``position`` past committed tokens keeps ``kept`` until computation passes ``position``.""" + if position > min(committed, default=0): + self._kept_prefixes[request_id] = (kv, tuple(kept), position) + fetch = _FetchRecord(start, kv, state_end=state_end, end=end, prompt=prompt, runs=runs) + self._latest_fetches[request_id] = fetch + return fetch + + def _record_next_fetch( + self, + request_id: int, + kv: _KVCache, + start: int, + position: int, + committed: Sequence[int], + state_end: int | None = None, + end: int | None = None, + ) -> _FetchRecord | None: + """``_record_fetch`` of a fetch that grows nothing, reading ``_prefix_kept_by_fetch`` only + where kept; ``None``, recording nothing, while an earlier fetch into ``kv`` is unsettled.""" + previous = self._latest_fetches.get(request_id) + if previous is not None and previous.kv is kv and not self._fetch_settled(previous): + return None + if position > min(committed, default=0): + kept = self._prefix_kept_by_fetch(request_id, kv, start) self._kept_prefixes[request_id] = (kv, kept, position) - fetch = _FetchRecord(start, kv) + # Neither marks queued nor a copy yet. + fetch = _FetchRecord(start, kv, False, None, state_end, end) self._latest_fetches[request_id] = fetch return fetch - def _abandon_fetch(self, lease: _StagingLease) -> None: - """Drop the write lease's fetch record: it delivered nothing, and resume_bounds counts what + def _has_records(self, request_id: int) -> bool: + """Whether a record names the request: a fetch, deliveries, a kept prefix or a mixed tail.""" + return ( + request_id in self._latest_fetches + or request_id in self._delivered + or request_id in self._pending + or request_id in self._kept_prefixes + or request_id in self._mixed_tail_caches + ) + + def _abandon_fetch(self, lease: _LeaseBase) -> None: + """Drop the write lease's fetch record: it delivered nothing, and poll_resume_bounds counts what earlier fetches into the cache delivered.""" if lease._fetch is not None and self._latest_fetches.get(lease._request_id) is lease._fetch: del self._latest_fetches[lease._request_id] - def _unsettled_fetch_refusal(self, request_id: int, kv: _KVCache) -> str | None: - """Why the request's earlier fetch into ``kv`` keeps a new one from starting, or ``None``: - it has not settled on this rank, where the copy its marks queued may still be pending.""" + def _unsettled_fetch_refusal( + self, request_id: int, kv: _KVCache, entries: _LentBlockEntries | None = None + ) -> str | None: + """Refuse another fetch, or a read overlapping an unsettled fetch's rows.""" previous = self._latest_fetches.get(request_id) if previous is not None and previous.kv is kv and not self._fetch_settled(previous): + if entries is not None and previous.entries is not None: + pending = dict( + zip(previous.entries.layer_group_indices, previous.entries.block_ordinals) + ) + if not any( + np.intersect1d(rows, pending.get(lg, ())).size + for lg, rows in zip(entries.layer_group_indices, entries.block_ordinals) + ): + return None return f"request {request_id} already has an unsettled fetch" return None + def _unheld_block( + self, + request_id: int, + kv: _KVCache, + entries: _LentBlockEntries, + top: int, + history: int | None = None, + ) -> tuple[int, int] | None: + """First unreadable ``(layer group, block)`` below ``top``, or None: attention needs + computed or delivered bytes; a cache with deliveries needs a complete STATE snapshot at + ``top``, one without keeps its computed state. ``history``: ``kv``'s, if read.""" + delivered = self._deliveries.get(request_id) + if delivered is not None and delivered.kv is not kv: + delivered = None + tpb = int(self._layout.tokens_per_block) + tainted = self._mixed_tail_caches.get(request_id) is kv + if delivered is not None or tainted: + committed = ( + _manager.committed_by_group(kv, self._layout.num_layer_groups) if tainted else () + ) + for lg, ordinals in zip(entries.layer_group_indices, entries.block_ordinals): + if not len(ordinals): + continue + if lg in self._recurrent_groups: + if delivered is not None and delivered.state_end != top: + return lg, 0 + elif tainted: + tail, partial = divmod(committed[lg], tpb) + if partial and tail in ordinals: + return lg, tail + if history is None: + history = _manager.cache_state(kv).history + advanced = self._kept_prefixes.get(request_id) + if history >= top and (advanced is None or advanced[0] is not kv or history > advanced[2]): + return None + computed = self._computed_tokens(request_id, kv) + if self._known(computed) >= top: + return None + rank, size = getattr(self._layout, "cp", (0, 1)) + page = tpb // size + for lg, ordinals in zip(entries.layer_group_indices, entries.block_ordinals): + if lg in self._recurrent_groups: + continue + for b in ordinals.tolist(): + first = b * tpb + rank * page # the first token of this rank's page + if first >= top: + continue + valid = int(computed[lg]) + if delivered is not None: + blocks = delivered.delivered_block_masks[lg] + if b < len(blocks) and blocks[b]: + valid = max(valid, delivered.partial_ends.get((lg, b), first + page)) + if valid < min(first + page, top): + return lg, b + return None + def _fetch_settled(self, fetch: _FetchRecord) -> bool: - """The fetch's arrived entries are marked and their copy into the request's pages is - done.""" + """The fetch's arrived entries are marked and their copy into the pages is done.""" return fetch.marks_queued and self._copy_done(fetch.copy) def _fetch_state(self, request_id: int, fetch: _FetchRecord) -> _FetchState: @@ -157,7 +353,7 @@ def _fetch_state(self, request_id: int, fetch: _FetchRecord) -> _FetchState: return _FetchState.COPY_DONE return _FetchState.COPY_PENDING if fetch.marks_queued else _FetchState.AWAITING_MARKS - def _take_out_fetch(self, lease: _StagingLease, fetch: _FetchRecord | None) -> bool: + def _take_out_fetch(self, lease: _LeaseBase, fetch: _FetchRecord | None) -> bool: """Take ``fetch``, the write lease's, out of the latest fetches before its marks' copy is queued, if it is still the request's latest; returns whether it was.""" current = fetch is not None and self._latest_fetches.get(lease._request_id) is fetch @@ -165,21 +361,32 @@ def _take_out_fetch(self, lease: _StagingLease, fetch: _FetchRecord | None) -> b del self._latest_fetches[lease._request_id] return current - def _invalidate_delivered_blocks(self, lease: _StagingLease, copied: list[np.ndarray]) -> None: - """Stop counting the delivered entries ``copied`` selects, before their copy is queued: one - that fails partway can leave a row holding each fetch's bytes in different pools. - ``_record_delivered_blocks`` counts them again once the copy is queued without error.""" + def _invalidate_delivered_blocks( + self, lease: _LeaseBase, copied: list[np.ndarray] | None + ) -> None: + """Stop counting the delivered entries ``copied`` selects (``None``: all the lease's) before + their copy is queued, as one failing partway can leave a row with each fetch's bytes in + different pools; ``_record_delivered_blocks`` recounts them once it queues without error.""" delivered = self._deliveries.get(lease._request_id) if delivered is None or delivered.kv is not lease._kv: return - for lg, block_ordinals, mask in zip( - lease._block_entries.layer_group_indices, lease._block_entries.block_ordinals, copied + lent = lease._block_entries + for i, (lg, block_ordinals) in enumerate( + zip(lent.layer_group_indices, lent.block_ordinals) ): + if lg in self._recurrent_groups: + if len(block_ordinals) and (copied is None or copied[i].any()): + delivered.state_end = None + delivered.resume_bounds_memo = None + continue blocks = delivered.delivered_block_masks[lg] - entries = block_ordinals[mask] + entries = block_ordinals if copied is None else block_ordinals[copied[i]] entries = entries[entries < len(blocks)] + for key in tuple(delivered.partial_ends): + if key[0] == lg and np.any(entries == key[1]): + delivered.partial_ends.pop(key) if blocks[entries].any(): - blocks[entries] = False + _writable_mask(delivered.delivered_block_masks, lg)[entries] = False delivered.resume_bounds_memo = None def _record_mixed_tail(self, request_id: int, kv: _KVCache) -> None: @@ -188,11 +395,11 @@ def _record_mixed_tail(self, request_id: int, kv: _KVCache) -> None: self._mixed_tail_caches[request_id] = kv def _put_back_fetch( - self, lease: _StagingLease, fetch: _FetchRecord, copied: list[np.ndarray] + self, lease: _LeaseBase, fetch: _FetchRecord, copied: list[np.ndarray] ) -> None: - """Put ``fetch`` back as the request's latest once its marks' copy is queued without error, - with ``marks_queued`` set and the lease's copy, and add the entries ``copied`` selects to - what fetches delivered into the cache.""" + """Once its marks' copy is queued without error, put ``fetch`` back as the request's latest, + ``marks_queued`` set and with the lease's copy, and add the entries ``copied`` selects to + the cache's deliveries.""" fetch.copy = lease._copy fetch.marks_queued = True self._latest_fetches[lease._request_id] = fetch @@ -205,35 +412,362 @@ def _record_delivered_blocks( entries: _LentBlockEntries, copied: list[np.ndarray], ) -> None: - """Add a fetch's copied entries to what earlier fetches into the same cache delivered, so a - fetch split into consecutive leases counts as one.""" + """Add a fetch's copied entries to the cache's deliveries, so consecutive leases count as + one fetch; recurrent state stands at its end if every state row arrived, else nowhere.""" delivered = self._deliveries.get(request_id) if delivered is None or delivered.kv is not fetch.kv: empty = [np.zeros(0, dtype=bool) for _ in range(self._layout.num_layer_groups)] delivered = _Deliveries(fetch.kv, fetch.start, empty) self._deliveries[request_id] = delivered delivered.lowest_start = min(delivered.lowest_start, fetch.start) + tpb = int(self._layout.tokens_per_block) + rank, size = getattr(self._layout, "cp", (0, 1)) + if fetch.end is not None: + delivered.end = max(delivered.end, fetch.end) + if fetch.state_end is not None: + states = [ + mask + for lg, mask in zip(entries.layer_group_indices, copied) + if lg in self._recurrent_groups + ] + whole = bool(states) and all(len(mask) and mask.all() for mask in states) + delivered.state_end = fetch.state_end if whole else None for lg, block_ordinals, mask in zip( entries.layer_group_indices, entries.block_ordinals, copied ): + if lg in self._recurrent_groups: + continue got = block_ordinals[mask] if not len(got): continue blocks = delivered.delivered_block_masks[lg] if int(got.max()) >= len(blocks): blocks = np.concatenate([blocks, np.zeros(int(got.max()) + 1 - len(blocks), bool)]) + elif not blocks.flags.writeable: + blocks = blocks.copy() blocks[got] = True delivered.delivered_block_masks[lg] = blocks + if fetch.end is not None: + tail = fetch.end // tpb + page_end = tail * tpb + (rank + 1) * (tpb // size) + if fetch.end < page_end and np.any(got == tail): + delivered.partial_ends[lg, tail] = fetch.end delivered.resume_bounds_memo = None + def _defer_whole_delivery( + self, request_id: int, fetch: _FetchRecord, spans: Sequence[tuple[int, int] | None] + ) -> None: + """``_record_whole_delivery`` of ``fetch``, settled now and held in ``_pending`` until a + read of ``_deliveries`` writes it; ``_whole_delivery_bounds`` gives its bounds meanwhile.""" + fetch.copy = None + fetch.marks_queued = True + self._pending[request_id] = [fetch, spans, None] + + def _defer_first_delivery( + self, + request_id: int, + fetch: _FetchRecord, + kv: _KVCache, + spans: Sequence[tuple[int, int] | None], + ) -> bool: + """If ``fetch`` is the request's latest and nothing was delivered into ``kv``, its cache, + yet, settle it as a whole delivery of ``spans`` (held in the closed form, else recorded): + True; else False, changing nothing.""" + if ( + self._latest_fetches.get(request_id) is not fetch + or self._delivered_kv(request_id) is kv + ): + return False + if self._closed_form: + self._defer_whole_delivery(request_id, fetch, spans) + else: + self._record_whole_delivery(request_id, fetch, spans) + return True + + def _delivered_kv(self, request_id: int) -> object | None: + """The cache the request's deliveries describe, a held whole delivery's included, or + ``None`` without deliveries; writes nothing.""" + pending = self._pending.get(request_id) + if pending is not None: + return pending[0].kv + delivered = self._delivered.get(request_id) + return None if delivered is None else delivered.kv + + def _record_whole_delivery( + self, + request_id: int, + fetch: _FetchRecord, + spans: Sequence[tuple[int, int] | None], + memo: tuple[tuple[int, ...], int, int] | None = None, + ) -> None: + """Settle ``fetch``, the request's latest, all lent rows arrived needing no copy, into a + cache without deliveries: ``_put_back_fetch``'s records written directly, ``memo`` their + bounds memo, ``spans`` as ``_LentBlockEntries.spans``. No Helix.""" + fetch.copy = None + fetch.marks_queued = True + # This record replaces any held for the request. + self._pending.pop(request_id, None) + tpb = self._tokens_per_block + end = fetch.end + tail = None if end is None or not end % tpb else end // tpb + masks = [] + partial_ends = {} + for lg, span in enumerate(spans): + run = _NO_RUN if span is None or span[1] <= span[0] else span + mask = _run_masks.get(run) + masks.append(_run_mask(*run) if mask is None else mask) + if tail is not None and run[0] <= tail < run[1]: + partial_ends[lg, tail] = end + # ``spans`` holds every layer group, so it lent a row of each group of recurrent state. + self._delivered[request_id] = _Deliveries( + fetch.kv, + fetch.start, + masks, + memo, + fetch.state_end if self._recurrent_groups else None, + 0 if end is None else max(0, end), + partial_ends, + ) + + def _whole_delivery_bounds( + self, + prompt: int, + runs: Sequence[tuple[int, int]], + request_id: int, + manager: KVCacheManagerV2, + kv: _KVCache, + state: _manager.CacheState, + spans: Sequence[tuple[int, int] | None], + blocks: int | None = None, + ) -> ResumeBounds | None: + """``_resume_bounds`` from the runs while the only deliveries are the latest fetch's whole + delivery of ``spans``, its runs ending at the fetch's end block; ``None`` where other rules + may apply: deliveries past capacity, an endpoint block runs and computed prefix miss, a sink + before a read window, a multimodal run at the endpoint. No Helix; keeps a held one held.""" + if not self._closed_form: + return self._whole_delivery_bounds_by_ranges( + prompt, runs, request_id, manager, kv, state, spans, blocks + ) + held = self._pending.get(request_id) + if held is None: + delivered = self._deliveries[request_id] + delivered_end = delivered.end + lowest_start = delivered.lowest_start + state_end = delivered.state_end + else: + # The values ``_record_whole_delivery`` writes of the fetch. + settled = held[0] + delivered_end = 0 if settled.end is None else max(0, settled.end) + lowest_start = settled.start + state_end = settled.state_end if self._recurrent_groups else None + kept_prefixes = self._kept_prefixes + if kept_prefixes: + advanced = kept_prefixes.get(request_id) + if advanced is not None and advanced[0] is not kv: + del kept_prefixes[request_id] + history = state.history + mixed = self._mixed_tail_caches + if mixed: + tainted = mixed.get(request_id) + if tainted is not None: + if tainted is kv: + return _new_tuple(ResumeBounds, (0, max(history, 1))) + del mixed[request_id] + tpb = self._tokens_per_block + if delivered_end > (_manager.num_blocks(kv) if blocks is None else blocks) * tpb: + # A shrink the manager did not report: the rules drop what lies past the capacity. + return None + computed = self._computed_tokens(request_id, kv, state) + known = self._known(computed) + cap = max(known, min(delivered_end, prompt)) + first, last = known // tpb, -(-cap // tpb) + recurrent = self._recurrent_groups + windows = self._window_of + longest = self._longest + # Per windowed group: its window, slack and the end of the blocks behind its read window at + # ``cap`` (the formula's stale range there starts at block 0). + read_ends = [] + for lg, span in enumerate(spans): + if lg in recurrent: + continue + covered_from = computed[lg] // tpb + if covered_from < first: + covered_from = first + terms = windows[lg] + if terms is not None: + window, slack = terms + behind = min(cap + slack, longest) + 1 - window + read_end = behind // tpb if behind > 0 else 0 + read_ends.append((lg, window, slack, read_end)) + if read_end > covered_from: + covered_from = read_end + if covered_from < last and (span is None or span[0] > covered_from or span[1] < last): + return None + usable = cap + floor = 0 + if runs: + # ``_clip_to_bidirectional_runs``: an endpoint strictly inside a run is the rules' case. + for b, e in runs: + if b < usable < e: + return None + floor = max((e for _, e in runs if e <= usable), default=0) + # ``_min_resume_position``; recurrent state needs none, as it holds at one position. + lowest = None + if not recurrent: + lowest = history + if not _ranges._floor_follows_history(manager): + lowest = min(lowest_start, history) + for window in self._stale_windows: + if min(history, longest) + 1 - window >= tpb: + lowest = history + break + if usable > known and read_ends: + if lowest is None or lowest >= usable: + # A floor a block missing behind a read window sets is at most ``usable``, which + # that window passed, so it cannot raise the bounds; the rules compute it if asked. + if held is None: + delivered.resume_bounds_memo = None + else: + held[2] = None + if lowest is not None: + return _new_tuple(ResumeBounds, (usable, lowest)) + end = state_end + if end is None or usable < end: + return _new_tuple(ResumeBounds, (0, max(history, 1))) + return _new_tuple(ResumeBounds, (end, end)) + end_block = -(-usable // tpb) + # No read window at ``usable`` passed the one partial page, the fetch's end block: each + # group's one candidate is its last block missing behind its window, below its runs. + for lg, window, slack, read_end in read_ends: + valid = computed[lg] + beg, end = valid // tpb, min(end_block, read_end) + if end <= beg: + continue + span = spans[lg] + lent_from, lent_to = span if span is not None and span[1] > span[0] else (0, 0) + if end > lent_to: + missing = end - 1 + elif beg < min(end, lent_from): + missing = min(end, lent_from) - 1 + else: + continue + # The last history whose read window still reads the block, by the formula. + keeping = (missing + 1) * tpb + window - 2 + last_read = (_manager._MAX_TOKENS if keeping >= longest else keeping) - slack + if valid < last_read and last_read + 1 > floor: + floor = last_read + 1 + if held is None: + delivered.resume_bounds_memo = computed, usable, floor + else: + held[2] = computed, usable, floor + if recurrent: + end = state_end + if end is None or usable < end: + return _new_tuple(ResumeBounds, (0, max(history, 1))) + return _new_tuple(ResumeBounds, (end, end)) + return _new_tuple(ResumeBounds, (usable, max(floor, lowest))) + + def _whole_delivery_bounds_by_ranges( + self, + prompt: int, + runs: Sequence[tuple[int, int]], + request_id: int, + manager: KVCacheManagerV2, + kv: _KVCache, + state: _manager.CacheState, + spans: Sequence[tuple[int, int] | None], + blocks: int | None = None, + ) -> ResumeBounds | None: + """``_whole_delivery_bounds`` for a layout whose windows the formula does not give, with + each stale range and last history reading a block asked of ``_ranges``.""" + delivered = self._deliveries[request_id] + advanced = self._kept_prefixes.get(request_id) + if advanced is not None and advanced[0] is not kv: + del self._kept_prefixes[request_id] + tainted = self._mixed_tail_caches.get(request_id) + if tainted is not None and tainted is not kv: + del self._mixed_tail_caches[request_id] + tainted = None + history = state.history + if tainted is not None: + return ResumeBounds(max_position=0, min_position=max(history, 1)) + layout = self._layout + tpb = int(layout.tokens_per_block) + kept = _manager.num_blocks(kv) if blocks is None else blocks + # The record's masks and partial pages end at the block that holds its end. + if delivered.end > kept * tpb: + return None + computed = self._computed_tokens(request_id, kv, state) + known = self._known(computed) + cap = max(known, min(delivered.end, prompt)) + first, last = known // tpb, -(-cap // tpb) + windowed = [] + for lg, span in enumerate(spans): + if lg in self._recurrent_groups: + continue + covered_from = max(first, int(computed[lg]) // tpb) + if layout.windows[lg] is not None: + windowed.append(lg) + read_beg, read_end = _ranges._read_stale_range(manager, layout, lg, cap) + if min(read_beg, last) > covered_from: + return None + covered_from = max(covered_from, read_end) + if covered_from < last and (span is None or span[0] > covered_from or span[1] < last): + return None + usable = cap + clipped, floor = self._clip_to_bidirectional_runs(runs, usable) + if clipped != usable: + return None + if usable > known: + end_block = -(-usable // tpb) + for lg in windowed: + computed_blocks = int(computed[lg]) // tpb + read_beg, read_end = _ranges._read_stale_range(manager, layout, lg, usable) + beg, end = max(computed_blocks, read_beg), min(end_block, read_end) + if end <= beg: + continue + span = spans[lg] + lent_from, lent_to = span if span is not None and span[1] > span[0] else (0, 0) + if end > lent_to: + missing = end - 1 + elif beg < min(end, lent_from): + missing = min(end, lent_from) - 1 + else: + missing = -1 + candidates = {missing} if missing >= 0 else set() + candidates.update( + ordinal + for group, ordinal in delivered.partial_ends + if group == lg and beg <= ordinal < end + ) + for ordinal in sorted(candidates, reverse=True): + last_read = _ranges._last_history_reading(manager, layout, lg, ordinal) + valid = max(int(computed[lg]), delivered.partial_ends.get((lg, ordinal), 0)) + if valid < last_read: + floor = max(floor, last_read + 1) + break + delivered.resume_bounds_memo = computed, usable, floor + if self._recurrent_groups: + end = delivered.state_end + if end is None or usable < end: + return ResumeBounds(max_position=0, min_position=max(history, 1)) + return ResumeBounds(max_position=end, min_position=end) + return ResumeBounds( + max_position=usable, + min_position=max( + floor, int(self._min_resume_position(manager, history, delivered.lowest_start)) + ), + ) + def _deliveries_of(self, request_id: int) -> _Deliveries | None: """What fetches delivered into the request's cache, or ``None``; the record may describe a cache the request no longer holds.""" return self._deliveries.get(request_id) def _drop_deliveries(self, request_id: int) -> None: - """Drop what fetches delivered into the request's cache.""" - self._deliveries.pop(request_id, None) + """Drop what fetches delivered into the request's cache, a whole delivery held included.""" + self._pending.pop(request_id, None) + self._delivered.pop(request_id, None) def _invalidate_delivered_blocks_past_capacity( self, delivered: _Deliveries, kv: _KVCache @@ -241,48 +775,82 @@ def _invalidate_delivered_blocks_past_capacity( """Forget delivered entries at or past the cache's block count: a shrink freed their pages, and a regrow brings pages without their contents.""" kept = _manager.num_blocks(kv) - for blocks in delivered.delivered_block_masks: + capacity_end = kept * int(self._layout.tokens_per_block) + if ( + delivered.end <= capacity_end + and all(len(blocks) <= kept for blocks in delivered.delivered_block_masks) + and all(ordinal < kept for _, ordinal in delivered.partial_ends) + ): + return + end = min(delivered.end, capacity_end) + if end != delivered.end: + delivered.end = end + delivered.resume_bounds_memo = None + delivered.partial_ends = { + key: end for key, end in delivered.partial_ends.items() if key[1] < kept + } + for lg, blocks in enumerate(delivered.delivered_block_masks): if blocks[kept:].any(): - blocks[kept:] = False + _writable_mask(delivered.delivered_block_masks, lg)[kept:] = False delivered.resume_bounds_memo = None def _drop_records(self, request_id: int) -> None: - """Drop the request's four records, at its free.""" + """Drop the request's four records, at its free, a whole delivery held included.""" for records in ( self._latest_fetches, - self._deliveries, + self._pending, + self._delivered, self._kept_prefixes, self._mixed_tail_caches, ): records.pop(int(request_id), None) - def _computed_tokens(self, request_id: int, kv: _KVCache) -> int: - """What resume_bounds counts as computed besides delivered entries: the tokens up to the - history, but only the committed ones and what the latest fetch kept while the history stays - where that fetch's grow left it past them (a windowed fetch moves it to its end).""" - state = _manager.cache_state(kv) + def _computed_tokens( + self, request_id: int, kv: _KVCache, state: _manager.CacheState | None = None + ) -> tuple[int, ...]: + """Computed tokens per layer group, fetched rows excluded: with the committed ones, the kept + prefix while the history stays where the latest fetch grew it, else the history. ``state``: + ``kv``'s, if read.""" + if state is None: + state = _manager.cache_state(kv) + history = state.history advanced = self._kept_prefixes.get(request_id) # Past where that fetch left it, the history moved as the request resumed and computed on. - if advanced is None or advanced[0] is not kv or state.history > advanced[2]: - return max(state.committed, state.history) - return max(state.committed, advanced[1]) + moved = advanced is None or advanced[0] is not kv or history > advanced[2] + if not state.committed: + return (history,) * self._layout.num_layer_groups if moved else tuple(advanced[1]) + committed = _manager.committed_by_group(kv, self._layout.num_layer_groups) + if moved: + return tuple(max(tokens, history) for tokens in committed) + return tuple(max(tokens, kept) for tokens, kept in zip(committed, advanced[1])) + + def _prefix_kept_by_fetch(self, request_id: int, kv: _KVCache, start: int) -> tuple[int, ...]: + """Keep each group's computed prefix up to start; a fetch may overwrite later uncommitted rows. + Delivered rows past that prefix keep separate records.""" + if start == 0: + return (0,) * self._layout.num_layer_groups + return tuple(min(tokens, start) for tokens in self._computed_tokens(request_id, kv)) - def _prefix_kept_by_fetch(self, request_id: int, kv: _KVCache, start: int) -> int: - """What a fetch from ``start`` keeps of what was computed before it, since it may overwrite - every block from ``start`` on. It never passes the history, which no shrink goes below; - delivered entries past it keep their own record.""" - return min(self._computed_tokens(request_id, kv), start) + def _known(self, computed: Sequence[int]) -> int: + """The tokens every layer group but those of recurrent state computed, of ``computed``, one + count per layer group.""" + paged = self._paged_groups + if not paged: + return max(computed, default=0) + return min([computed[lg] for lg in paged]) def _resume_bounds( self, - request: LlmRequest, + prompt: int, + runs: Sequence[tuple[int, int]], request_id: int, manager: KVCacheManagerV2, kv: _KVCache, - history: int, + state: _manager.CacheState, ) -> ResumeBounds | None: - """``Staging.resume_bounds`` of a lender that serves and has not stopped, for the request's - cache ``kv`` at a history of ``history`` tokens.""" + """``poll_resume_bounds`` of a serving, unstopped lender for the request's cache ``kv`` in + ``state``, given its ``prompt`` length and multimodal ``runs``.""" + history = state.history # A record of a replaced cache (a restart) is moot. fetch = self._latest_fetches.get(request_id) if fetch is not None and fetch.kv is not kv: @@ -304,21 +872,35 @@ def _resume_bounds( if tainted is not None: # Empty: the request drops its cache and computes from 0. return ResumeBounds(max_position=0, min_position=max(history, 1)) - known = self._computed_tokens(request_id, kv) + computed = self._computed_tokens(request_id, kv, state) + known = self._known(computed) if delivered is None: # Nothing delivered: what was computed and kept counts, resumed no lower than the # history, which a windowed fetch moved to its end. - usable, floor = self._clip_to_bidirectional_runs(request, known) + usable, floor = self._clip_to_bidirectional_runs(runs, known) return ResumeBounds(max_position=usable, min_position=max(floor, history)) # A shrink the manager did not report still shows as blocks past the capacity. self._invalidate_delivered_blocks_past_capacity(delivered, kv) - if delivered.resume_bounds_memo is None or delivered.resume_bounds_memo[0] != known: - usable = self._max_resume_position(manager, delivered, known) - delivered.resume_bounds_memo = ( - known, - *self._clip_to_bidirectional_runs(request, usable), - ) + if delivered.resume_bounds_memo is None or delivered.resume_bounds_memo[0] != computed: + cap = max(known, min(delivered.end, prompt)) + while True: + usable = self._max_resume_position(manager, delivered, computed, cap) + clipped, floor = self._clip_to_bidirectional_runs(runs, usable) + if clipped == usable or clipped <= known: + usable = clipped + break + # Moving back into a window can require blocks the larger endpoint did not read. + cap = clipped + floor = max(floor, self._delivered_resume_floor(manager, delivered, computed, usable)) + delivered.resume_bounds_memo = computed, usable, floor _, usable, floor = delivered.resume_bounds_memo + if self._recurrent_groups: + # Recurrent state holds at one position alone: where the delivered state stands, every + # block below it there; anything else leaves the request to compute from 0. + end = delivered.state_end + if end is None or usable < end: + return ResumeBounds(max_position=0, min_position=max(history, 1)) + return ResumeBounds(max_position=end, min_position=end) return ResumeBounds( max_position=usable, min_position=max( @@ -326,11 +908,14 @@ def _resume_bounds( ), ) - def _clip_to_bidirectional_runs(self, request: LlmRequest, usable: int) -> tuple[int, int]: - """The end lowered to the start of a run of multimodal tokens it falls strictly inside, and - the lowest floor that then leaves no position strictly inside a run below that end: the end - of the last run at or below it. Runs are those the scheduler keeps within one chunk.""" - runs = _manager.bidirectional_runs(request) + def _clip_to_bidirectional_runs( + self, runs: Sequence[tuple[int, int]], usable: int + ) -> tuple[int, int]: + """The end lowered to the start of a multimodal run it lies strictly inside, and the lowest + floor leaving no position below it inside a run: the end of the last run at or below it. + ``runs``: those the scheduler keeps within one chunk.""" + if not runs: + return int(usable), 0 for b, e in runs: if b < usable < e: usable = b @@ -338,8 +923,8 @@ def _clip_to_bidirectional_runs(self, request: LlmRequest, usable: int) -> tuple def _min_resume_position(self, manager: KVCacheManagerV2, history: int, origin: int) -> int: """The lowest start that needs no restart: the history where a chunk resumed below it may - end below it and raise, or a window has released blocks at it (they have no pages), else - the smaller of the lowest start of a delivered fetch and the history.""" + end below it and raise, or a window has released blocks at it (they have no pages); else the + lower of the history and ``origin``, a delivered fetch's lowest start.""" if _ranges._floor_follows_history(manager): return history for lg in range(self._layout.num_layer_groups): @@ -349,40 +934,83 @@ def _min_resume_position(self, manager: KVCacheManagerV2, history: int, origin: return min(origin, history) def _max_resume_position( - self, manager: KVCacheManagerV2, delivered: _Deliveries, known: int + self, + manager: KVCacheManagerV2, + delivered: _Deliveries, + computed: Sequence[int], + cap: int | None = None, ) -> int: - """The largest start ``P >= known`` where every layer group has what it reads among the - blocks below ``known``, computed before the fetches, and delivered entries: full attention - every block below ``P``, a window its sinks and in-window blocks.""" - # TODO: a fetch none of whose entries a window reads delivers none, so the usable end stays at - # the blocks computed before the fetches, and the interval can come out empty although the - # fetch covered all that a resume at its end reads. + """Largest endpoint whose read rows are computed or delivered, partial pages to their valid + end; a gap a window passed does not count. Helix checks this rank's page alone.""" tpb = int(self._layout.tokens_per_block) - first = known // tpb # the blocks below were computed - last = max((len(blocks) for blocks in delivered.delivered_block_masks), default=0) + known = self._known(computed) + masks = delivered.delivered_block_masks + groups = [lg for lg in range(len(masks)) if lg not in self._recurrent_groups] + if not groups: + return known if cap is None else cap + cap = max((len(masks[lg]) for lg in groups), default=0) * tpb if cap is None else cap + rank, size = getattr(self._layout, "cp", (0, 1)) + if size > 1: + page = tpb // size + usable = cap + for lg in groups: + b = int(computed[lg]) // tpb + while b * tpb + rank * page < usable: + page_end = b * tpb + (rank + 1) * page + valid = int(computed[lg]) + if b < len(masks[lg]) and masks[lg][b]: + valid = max(valid, delivered.partial_ends.get((lg, b), page_end)) + if valid < min(usable, page_end): + usable = max(known, b * tpb + rank * page, valid) + break + b += 1 + return usable + first, last = known // tpb, -(-cap // tpb) if last <= first: return known - # Non-monotonic in P under windows, so each is checked. Blocks below known behind a window's - # history have no pages, but no start at or above the floor reads them. - top = last # no start past a full-attention group's first missing block passes that group - missing = [] # per layer group: missing[j], the undelivered blocks of [first, first + j) - for lg, blocks in enumerate(delivered.delivered_block_masks): + if all(self._layout.windows[lg] is None for lg in groups): + return self._max_resume_without_windows(delivered, computed, groups, known, cap) + top = cap + missing = {} + for lg in groups: have = np.zeros(last - first, dtype=bool) - mine = blocks[first:last] + mine = masks[lg][first:last] have[: len(mine)] = mine - missing.append(np.concatenate([[0], np.cumsum(~have)])) + for (group, ordinal), _ in delivered.partial_ends.items(): + if group == lg and first <= ordinal < last: + have[ordinal - first] = False + have[: max(0, int(computed[lg]) // tpb - first)] = True + missing[lg] = np.concatenate([[0], np.cumsum(~have)]) if self._layout.windows[lg] is None and not have.all(): - top = min(top, first + int(np.argmin(have))) + top = min(top, (first + int(np.argmin(have)) + 1) * tpb) def gap(lg: int, a: int, b: int) -> bool: a, b = max(a, first), min(b, last) - return b > a and missing[lg][b - first] - missing[lg][a - first] > 0 + if b <= a: + return False + prefix = missing[lg] + count = prefix[b - first] - prefix[a - first] + tail = (candidate - 1) // tpb + valid = max(int(computed[lg]), delivered.partial_ends.get((lg, tail), 0)) + if a <= tail < b and candidate <= valid: + count -= prefix[tail + 1 - first] - prefix[tail - first] + return count > 0 - # The largest start first: a fetch whose leases all landed is usable at its end at once. - for end_block in range(top, first, -1): - for lg in range(len(delivered.delivered_block_masks)): - stale_beg, stale_end = _ranges._stale_range( - manager, self._layout, lg, end_block * tpb + # A partial end can be the only valid endpoint, even after a bidirectional rewind. + extra = {top} + extra.update(end for end in delivered.partial_ends.values() if known < end <= top) + extra.update(int(end) for end in computed if known < end <= top) + previous = None + for candidate in merge( + range(top // tpb * tpb, known, -tpb), sorted(extra, reverse=True), reverse=True + ): + if candidate == previous or candidate <= known: + continue + previous = candidate + end_block = -(-candidate // tpb) + for lg in groups: + stale_beg, stale_end = _ranges._read_stale_range( + manager, self._layout, lg, candidate ) if stale_end > stale_beg: bad = gap(lg, first, min(stale_beg, end_block)) or gap(lg, stale_end, end_block) @@ -391,5 +1019,98 @@ def gap(lg: int, a: int, b: int) -> bool: if bad: break else: - return end_block * tpb + return candidate return known + + def _max_resume_without_windows( + self, + delivered: _Deliveries, + computed: Sequence[int], + groups: Sequence[int], + known: int, + cap: int, + ) -> int: + """``_max_resume_position`` in ``(known, cap]`` without windows or Helix: usable where each + group has every block before the endpoint and its last up to it, so each group's first + missing block (partial pages count) bounds it; the largest candidate within all wins.""" + tpb = int(self._layout.tokens_per_block) + masks = delivered.delivered_block_masks + partial_ends = delivered.partial_ends + first, last = known // tpb, -(-cap // tpb) + top, bound = cap, None + for lg in groups: + start = max(first, int(computed[lg]) // tpb) + missing = None + if start < last: + segment = masks[lg][start:last] + if len(segment): + at = int(segment.argmin()) + if not segment[at]: + missing = start + at + if missing is None and len(segment) < last - start: + missing = start + len(segment) + for group, ordinal in partial_ends: + if ( + group == lg + and start <= ordinal < last + and (missing is None or ordinal < missing) + ): + missing = ordinal + if missing is None: + continue + top = min(top, (missing + 1) * tpb) + # An endpoint inside the missing block reads it up to the endpoint alone. + valid = max(int(computed[lg]), partial_ends.get((lg, missing), 0)) + group_bound = max(missing * tpb, valid) + bound = group_bound if bound is None else min(bound, group_bound) + limit = top if bound is None else min(top, bound) + # The scan's candidates: whole blocks up to ``top``, ``top`` itself, partial ends and + # computed ends, each past ``known``. + best = limit // tpb * tpb + if best <= known: + best = known + for end in (top, *partial_ends.values(), *computed): + end = int(end) + if known < end <= limit and end > best: + best = end + return best + + def _delivered_resume_floor( + self, + manager: KVCacheManagerV2, + delivered: _Deliveries, + computed: Sequence[int], + usable: int, + ) -> int: + """The floor of the safe interval ending at ``usable``: a missing block behind a read window + stays behind it at every position the interval permits, rewinds included. Each layer group + computed the blocks below its ``computed[lg]`` tokens.""" + if usable <= self._known(computed): + return 0 + tpb = int(self._layout.tokens_per_block) + floor = 0 + for lg, blocks in enumerate(delivered.delivered_block_masks): + if lg in self._recurrent_groups or self._layout.windows[lg] is None: + continue + first = int(computed[lg]) // tpb + stale_beg, stale_end = _ranges._read_stale_range(manager, self._layout, lg, usable) + beg, end = max(first, stale_beg), min(-(-usable // tpb), stale_end) + if end <= beg: + continue + gaps = np.flatnonzero(~blocks[beg : min(end, len(blocks))]) + missing = beg + int(gaps[-1]) if len(gaps) else -1 + if end > len(blocks): + missing = end - 1 + candidates = {missing} if missing >= 0 else set() + candidates.update( + ordinal + for group, ordinal in delivered.partial_ends + if group == lg and beg <= ordinal < end + ) + for ordinal in sorted(candidates, reverse=True): + last_read = _ranges._last_history_reading(manager, self._layout, lg, ordinal) + valid = max(int(computed[lg]), delivered.partial_ends.get((lg, ordinal), 0)) + if valid < last_read: + floor = max(floor, last_read + 1) + break + return floor diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_identity.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_identity.py index f89dd9571840..ae39671b3173 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_identity.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_identity.py @@ -12,9 +12,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Object keys and staging region labels instances exchange; equal keys mean interchangeable bytes -given the ``scope``, the ``manager`` and the caller obligations ``attach_lender`` states. The -format is a contract between instances: changing it bumps ``KEY_FORMAT_VERSION``.""" +"""Object keys and region labels instances exchange; equal keys mean interchangeable bytes given +scope, manager and ``attach_lender``'s obligations. Format changes bump ``KEY_FORMAT_VERSION``.""" from __future__ import annotations @@ -23,7 +22,8 @@ import numpy as np -from ._layout import ShardDesc, canonical_order +from ..page_layout import ShardDesc +from ._layout import canonical_order KEY_FORMAT_VERSION = 2 REUSE_KEY_BYTES = 32 @@ -38,9 +38,8 @@ def _u16(value: int) -> bytes: def namespace(scope: bytes, layout_id: bytes) -> bytes: - """The 16 bytes every object key starts with, from the caller's ``scope`` and the block layout's - ``layout_id`` (each length-prefixed). ``ValueError`` for an empty ``layout_id`` or a ``scope`` - longer than 65535 bytes.""" + """The 16 bytes every object key starts with, hashed from length-prefixed ``scope`` and + ``layout_id``; ``ValueError`` for an empty ``layout_id`` or a ``scope`` over 65535 bytes.""" if not layout_id: raise ValueError("a namespace without layout_id would mix incompatible bytes") if len(scope) > 0xFFFF: @@ -71,27 +70,29 @@ def __init__( self.layout_id = bytes(layout_id) self.namespace = namespace(bytes(scope), self.layout_id) self.canonical = canonical_order(layers) - self._namespace = np.frombuffer(self.namespace, dtype=np.uint8) self._groups = [] for index, (count, share) in zip(self.canonical, shards): ShardDesc(count, share) # ValueError for a share outside the count suffix = _u16(index) + _u16(count) + _u16(share) - self._groups.append(np.frombuffer(suffix, dtype=np.uint8)) + self._groups.append(suffix) def object_keys(self, layer_group_index: int, reuse_keys: np.ndarray) -> np.ndarray: - """``uint8 (n, OBJECT_KEY_BYTES)``: the namespace, each entry's 32-byte reuse key, then + """``uint8 (n, OBJECT_KEY_BYTES)``: the namespace, each row's 32-byte reuse key, then the group's canonical index, shard count and shard index as big-endian ``u16``.""" if not 0 <= layer_group_index < len(self._groups): raise ValueError(f"unknown layer group {layer_group_index}") reuse_keys = np.ascontiguousarray(reuse_keys, dtype=np.uint8).reshape(-1, REUSE_KEY_BYTES) - out = np.empty((reuse_keys.shape[0], OBJECT_KEY_BYTES), dtype=np.uint8) - out[:, :_NAMESPACE_BYTES] = self._namespace - out[:, _NAMESPACE_BYTES : _NAMESPACE_BYTES + REUSE_KEY_BYTES] = reuse_keys - out[:, _NAMESPACE_BYTES + REUSE_KEY_BYTES :] = self._groups[layer_group_index] - return out + suffix = self._groups[layer_group_index] + payload = b"".join(piece for key in reuse_keys for piece in (self.namespace, key, suffix)) + return np.frombuffer(payload, dtype=np.uint8).reshape(-1, OBJECT_KEY_BYTES) def region_label(self, layer_group_indices: Iterable[int]) -> str: """The layout-derived label for a staging region serving ``layer_group_indices``: equal on instances laid out alike, independent of addresses, capacity and scope.""" groups = sorted({self.canonical[int(g)] for g in layer_group_indices}) return f"{self.layout_id.hex()[:_REGION_LAYOUT_HEX]}:lg{'+'.join(str(g) for g in groups)}" + + def device_region_label(self, layer_group_indices: Iterable[int], pool_index: int) -> str: + """The label of pool ``pool_index``'s device region serving ``layer_group_indices``: equal + on instances laid out alike, distinct from every staging label and from other pools'.""" + return f"{self.region_label(layer_group_indices)}:pool{int(pool_index)}" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py index 26d98f8a6f68..5e86d86871d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py @@ -12,167 +12,29 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The layout of a manager's blocks as the lender needs it, and ``layout_id``, the digest of how a -block's bytes are laid out. Read from the manager's public layout surface and declarations.""" +"""The layout of a manager's blocks as the lender needs it, the page table the manager builds, and +``layout_id``, the digest of how a block's bytes are laid out.""" from __future__ import annotations +import dataclasses import enum import hashlib import json import struct -from dataclasses import dataclass -from typing import TYPE_CHECKING, Literal, Mapping, Sequence +from typing import TYPE_CHECKING, Mapping, Sequence import numpy as np -from . import _manager +from ..page_layout import BufferDesc, PageLayout, replicated +from ..page_table import DevicePool, ManagerLayout # noqa: F401 DevicePool: the lender's pools if TYPE_CHECKING: - from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig, SsmLayerConfig - from ..kv_cache_manager_v2 import KVCacheManagerV2 # A new field left out at its default keeps the version; changing the encoding or the meaning of # an existing field bumps it. -LAYOUT_ID_VERSION = 2 - -# The fallback entry of the manager's per-role declarations (``Role.ALL``). -_ALL_ROLE = "all" - - -class BufferMapper(enum.IntEnum): - """How a buffer's bytes split across tensor-parallel ranks; the values of the mapper kinds - managers declare per role, read by value.""" - - INDEXED = 0 # head-major K/V, the default: a head range is one contiguous range per buffer - REPLICATED = 1 # bytes identical on every TP rank - NHD = 2 # token-major [token, head, dim] K/V: a head range is a slice inside every token - SECTIONED = 3 # [Sec0|Sec1|...], each section sharded independently (recurrent conv state) - - -@dataclass(frozen=True) -class BufferGeometry: - """Resharding geometry of one buffer; all ``None`` means no head axis. ``section_bytes`` sum to - the buffer's size; ``bytes_per_head * num_heads`` equals it.""" - - section_bytes: tuple[int, ...] | None = None - bytes_per_head: int | None = None - num_heads: int | None = None - - def validate(self) -> BufferGeometry: - """A copy with ``section_bytes`` as ints; ``ValueError`` unless every size given is - positive.""" - sections = self.section_bytes - if sections is not None: - sections = tuple(int(s) for s in sections) - if any(s <= 0 for s in sections): - raise ValueError(f"section sizes must be positive, got {sections}") - for name in ("bytes_per_head", "num_heads"): - value = getattr(self, name) - if value is not None and int(value) <= 0: - raise ValueError(f"{name} must be positive, got {value}") - return BufferGeometry(sections, self.bytes_per_head, self.num_heads) - - def check_size(self, size: int) -> None: - """``ValueError`` unless this geometry describes a buffer of ``size`` bytes.""" - if self.section_bytes is not None and sum(self.section_bytes) != size: - raise ValueError(f"sections {self.section_bytes} do not sum to {size} bytes") - if self.bytes_per_head is not None and size % self.bytes_per_head: - raise ValueError( - f"{size} bytes is not a whole number of {self.bytes_per_head}-byte heads" - ) - if ( - self.bytes_per_head is not None - and self.num_heads is not None - and self.bytes_per_head * self.num_heads != size - ): - raise ValueError( - f"{self.num_heads} heads of {self.bytes_per_head} bytes is not {size} bytes" - ) - - -@dataclass(frozen=True) -class ShardDesc: - """Which share (``index`` of ``count``) of a layer group's content this rank holds; ``(1, 0)`` - is the whole content, which every rank names alike. ``ValueError`` for an index outside the - count.""" - - count: int = 1 - index: int = 0 - - def __post_init__(self) -> None: - if self.count < 1 or not 0 <= self.index < self.count: - raise ValueError(f"shard {self.index} of {self.count}") - - -@dataclass(frozen=True) -class LayerGroupDesc: - """One layer group: its kind, window and sink blocks (which blocks exist), its device pool - group, its global layer ids and the share of its content this rank holds.""" - - kind: Literal["attention", "state"] - window: int | None - sink_blocks: int - pool_group: int - layers: tuple[int, ...] - shard: ShardDesc = ShardDesc() - - -@dataclass(frozen=True) -class BufferDesc: - """One buffer of one layer (global id) in a page: ``size`` bytes at ``offset`` in the slot of - ``pool`` = (pool group, pool index), how its bytes split across ranks, whether peers transfer - it, and its sub-pages per block. ``ValueError`` for a non-positive size or expansion.""" - - layer: int - role: str - pool: tuple[int, int] - offset: int - size: int - mapper: BufferMapper = BufferMapper.INDEXED - geometry: BufferGeometry = BufferGeometry() - # False for local-only buffers peers skip; they still ride along wherever whole slots are - # copied. - transfer: bool = True - # The buffer's own tokens per block is tokens_per_block / expansion; a consumer reading inside - # the buffer that does not support the value must reject the layout. - expansion: int = 1 - - def __post_init__(self) -> None: - object.__setattr__(self, "role", str(self.role)) - object.__setattr__(self, "mapper", BufferMapper(self.mapper)) - if not isinstance(self.geometry, BufferGeometry): - raise ValueError(f"geometry must be a buffer geometry, got {self.geometry!r}") - if self.size <= 0 or self.offset < 0: - raise ValueError(f"buffer at offset {self.offset} of {self.size} bytes") - if int(self.expansion) < 1: - raise ValueError(f"expansion must be a positive integer, got {self.expansion}") - self.geometry.check_size(self.size) - - -@dataclass(frozen=True) -class Layout: - """How pages are built, without addresses or slot counts: ``pool_groups[g]`` holds the slot - widths of device pool group ``g`` in pool order. ``ValueError`` for an expansion that does not - divide ``tokens_per_block``.""" - - tokens_per_block: int - pool_groups: Mapping[int, tuple[int, ...]] - layer_groups: tuple[LayerGroupDesc, ...] - buffers: tuple[BufferDesc, ...] - - def __post_init__(self) -> None: - pool_groups = { - int(g): tuple(int(w) for w in widths) for g, widths in dict(self.pool_groups).items() - } - object.__setattr__(self, "pool_groups", pool_groups) - for b in self.buffers: - if self.tokens_per_block % int(b.expansion): - raise ValueError( - f"buffer ({b.layer}, {b.role}): expansion {b.expansion} does not divide " - f"{self.tokens_per_block} tokens per block" - ) +LAYOUT_ID_VERSION = 4 # Enums encode as the member name, floats as "f64:" plus the 16 hex digits of their big-endian @@ -202,8 +64,7 @@ def _canonical(value: object) -> object: def canonical_bytes(document: Mapping) -> bytes: - """``document`` as ASCII JSON, keys sorted, no whitespace; enums by name, floats as ``f64:`` and - their big-endian bits in hex, ``None`` values dropped from objects.""" + """``document`` as ASCII JSON, keys sorted, no whitespace, values encoded by ``_canonical``.""" return json.dumps( _canonical(document), sort_keys=True, @@ -223,47 +84,63 @@ def canonical_order(layers: Sequence[Sequence[int]]) -> list[int]: return canonical -def _geometry_document(geometry: BufferGeometry) -> dict: - return { - "section_bytes": None if geometry.section_bytes is None else list(geometry.section_bytes), - "bytes_per_head": geometry.bytes_per_head, - "num_heads": geometry.num_heads, - } +def _tensors_document(buffer: BufferDesc) -> list[dict] | None: + """The buffer's tensors, or None for the whole buffer alike on every rank.""" + if buffer.tensors == replicated(buffer.size): + return None + out = [] + for tensor in buffer.tensors: + axes = [] + for axis in tensor.axes: + doc = {"name": axis.name, "extent": int(axis.extent), "stride": int(axis.stride)} + if axis.split: + doc["split"] = list(axis.split) + if axis.held is not None: + doc["held"] = int(axis.held) + if axis.cyclic: + doc["cyclic"] = True + axes.append(doc) + doc = {"axes": axes} + if tensor.name: + doc["name"] = tensor.name + if tensor.offset: + doc["offset"] = int(tensor.offset) + out.append(doc) + return out -def _buffer_document(layout: Layout, b: BufferDesc) -> dict: +def _buffer_document(layout: PageLayout, b: BufferDesc) -> dict: # A buffer's position in the page: the slot widths of the pools before it in its group, plus # its offset inside the slot. - widths = layout.pool_groups.get(b.pool[0]) + widths = layout.pool_slot_bytes.get(b.pool_group) if widths is None: - raise ValueError(f"buffer ({b.layer}, {b.role}) lies in pool {b.pool}, not in the layout") + raise ValueError( + f"buffer ({b.layer}, {b.role}) lies in pool {(b.pool_group, b.pool_index)}, not in the layout" + ) doc = { "layer": int(b.layer), "role": b.role, - "offset": sum(widths[: b.pool[1]]) + int(b.offset), + "offset": sum(widths[: b.pool_index]) + int(b.slot_offset), "size": int(b.size), } - # Fields at their defaults are left out, so a new defaulted field keeps existing ids. - if b.mapper is not BufferMapper.INDEXED: - doc["mapper"] = b.mapper.name - geometry = {k: v for k, v in _geometry_document(b.geometry).items() if v is not None} - if geometry: - doc["geometry"] = geometry - if not b.transfer: + tensors = _tensors_document(b) + if tensors is not None: + doc["tensors"] = tensors + if not b.transferred: doc["transfer"] = False if int(b.expansion) != 1: doc["expansion"] = int(b.expansion) return doc -def layout_document(layout: Layout) -> dict: - """The document ``layout_id`` hashes: tokens per block, groups in canonical order (layers, slot - widths, kind), buffers by (layer, role). Addresses, windows, sinks, shards, local numbering and - element types stay out.""" +def layout_document(layout: PageLayout) -> dict: + """Hash canonical global layer groups, slot widths, kinds and each buffer's tensors, but no + addresses, windows, sinks, shards, topology, total head counts, dtypes or local region + numbering.""" canonical = canonical_order([desc.layers for desc in layout.layer_groups]) groups = [None] * len(layout.layer_groups) for local, desc in enumerate(layout.layer_groups): - widths = layout.pool_groups.get(desc.pool_group) + widths = layout.pool_slot_bytes.get(desc.pool_group) if widths is None: raise ValueError( f"layer group {local} maps to pool group {desc.pool_group}, which is not in the " @@ -285,292 +162,24 @@ def layout_document(layout: Layout) -> dict: raise ValueError(f"buffer ({before['layer']}, {before['role']}) appears twice") return { "v": LAYOUT_ID_VERSION, - "tokens_per_block": int(layout.tokens_per_block), + "tokens_per_block": int(layout.tokens_per_page), "groups": groups, "buffers": buffers, } -def layout_id(layout: Layout) -> bytes: +def layout_id(layout: PageLayout) -> bytes: """SHA-256 of ``canonical_bytes(layout_document(layout))``.""" return hashlib.sha256(canonical_bytes(layout_document(layout))).digest() -@dataclass(frozen=True) -class DevicePool: - """One device pool of pool group ``group``: slot ``s`` lives at ``base + s * slot_bytes``.""" - - group: int - index: int - base: int - slot_bytes: int - num_slots: int - - -@dataclass(frozen=True, eq=False) -class ManagerLayout: - """What the lender reads of one manager, as plain values that hold no manager: the address-free - ``layout``, facts per layer group (tuples indexed by the local layer group) and per device pool - group (mappings keyed by the pool group).""" - - layout: Layout - tokens_per_block: int - pool_group_of: tuple[int, ...] - windows: tuple[int | None, ...] - sink_blocks: tuple[int, ...] - recurrent: tuple[bool, ...] - sparse: tuple[bool, ...] - layers: tuple[tuple[int, ...], ...] - shards: tuple[tuple[int, int], ...] - page_bytes: Mapping[int, int] - device_pools: Mapping[int, tuple[DevicePool, ...]] - - @property - def num_layer_groups(self) -> int: - """The manager's layer groups.""" - return len(self.pool_group_of) - - @property - def pool_groups(self) -> tuple[int, ...]: - """The device pool groups, sorted; staging regions follow this order.""" - return tuple(sorted(self.device_pools)) - - -def global_layer_ids(manager: KVCacheManagerV2, internal_ids: Sequence[int]) -> list[int]: - """The global id of each internal layer id, the same under any pipeline split: the declared id; - else for virtual layers ``model_layer * number_of_attention_types + attention_type``; else the - model layer.""" - # One declared id per internal layer of each layer group, in order (FP4 MLA's tail layers). - declared = _manager.declared_global_layer_ids(manager) - if declared is not None: - return [declared[int(lid)] for lid in internal_ids] - virtual = _manager.virtual_layers(manager) - if virtual is None: - return [int(manager.pp_layers[int(lid)]) for lid in internal_ids] - # The number of types counts every member of the enum, so stages holding different attention - # types agree. - inverse, num_types = virtual - return [inverse[int(lid)][0] * num_types + inverse[int(lid)][1] for lid in internal_ids] - - -def _sparse(layer_config: AttentionLayerConfig | SsmLayerConfig) -> bool: - """Whether a layer config of the runtime has sparse buffers (``is_sparse``), whose read-only - pages a cache can lock in host memory.""" - return any(getattr(b, "is_sparse", False) for b in getattr(layer_config, "buffers", ()) or ()) - - -def _expansions( - layer_config: AttentionLayerConfig | SsmLayerConfig, tokens_per_block: int -) -> dict[str, int]: - """Role -> sub-pages per block, from the layer's ``tokens_per_block_override``.""" - out = {} - for buffer in getattr(layer_config, "buffers", ()) or (): - override = getattr(buffer, "tokens_per_block_override", None) - if override is not None: - if tokens_per_block % int(override): - raise ValueError( - f"buffer {buffer.role}: {override} tokens per block does not divide " - f"{tokens_per_block}" - ) - out[str(buffer.role)] = tokens_per_block // int(override) - return out - - -class _Heads: - """Head counts of attention buffers: per rank from the manager, in total from its inputs.""" - - def __init__(self, manager: KVCacheManagerV2) -> None: - mapping = getattr(manager, "mapping", None) - dp = bool(getattr(mapping, "enable_attention_dp", False)) - self.tp_size = 1 if mapping is None or dp else max(int(getattr(mapping, "tp_size", 1)), 1) - self.tp_rank = ( - 0 if self.tp_size == 1 else int(getattr(mapping, "tp_rank", 0)) % self.tp_size - ) - self._per_rank = list(getattr(manager, "num_kv_heads_per_layer", ()) or ()) - self._total = getattr(manager, "num_kv_heads", None) - self._virtual = _manager.virtual_layers(manager) is not None - self._pp_layers = manager.pp_layers - - def per_rank(self, internal: int) -> int: - if not self._per_rank: - return 0 - index = internal if internal < len(self._per_rank) else 0 - return int(self._per_rank[index] or 0) - - def total(self, internal: int) -> int | None: - if isinstance(self._total, int): - return int(self._total) - if self._total is None or self._virtual: - return None - if internal >= len(self._pp_layers): - return None # an extra internal layer (FP4 MLA's tail) is no model layer - value = self._total[int(self._pp_layers[internal])] - return None if value is None else int(value) - - def attention_shard(self, internal: int) -> ShardDesc: - """The share of a head-sharded attention buffer this rank holds: ``T`` heads over ``tp`` - ranks, rank ``r`` holds ``ceil(T / tp)`` from ``r * T // tp``. Where any rank's range starts - inside a share, each rank's buffer is a share of its own.""" - if self.tp_size == 1: - return ShardDesc() - heads, total = self.per_rank(internal), self.total(internal) - if heads <= 0 or total is None or total <= 0: - return ShardDesc(self.tp_size, self.tp_rank) - if any(rank * total // self.tp_size % heads for rank in range(self.tp_size)): - # Some rank's heads straddle two shares, so each rank's buffer is a share of its own. - return ShardDesc(self.tp_size, self.tp_rank) - # total // heads distinct shares; fewer heads than ranks are repeated, so a single head is - # one share that every rank holds. - count = max(total // heads, 1) - index = min((self.tp_rank * total // self.tp_size) // heads, count - 1) - return ShardDesc(count, index) - - def state_shard(self) -> ShardDesc: - return ShardDesc(self.tp_size, self.tp_rank) if self.tp_size > 1 else ShardDesc() - - -def _buffer_geometry( - declared: object, mapper: BufferMapper, size: int, heads: int, state: bool -) -> BufferGeometry: - """A buffer's head geometry: declared sections or head size, completed per buffer. - ``ValueError`` if the declared or completed geometry has a size that is not positive.""" - declared = BufferGeometry( - getattr(declared, "section_bytes", None), getattr(declared, "bytes_per_head", None) - ).validate() - if mapper in (BufferMapper.REPLICATED, BufferMapper.SECTIONED): - return declared - bytes_per_head = declared.bytes_per_head - if bytes_per_head is None and not state and heads > 0 and size % heads == 0: - bytes_per_head = size // heads - if bytes_per_head is None or size % bytes_per_head: - return declared - return BufferGeometry(declared.section_bytes, bytes_per_head, size // bytes_per_head).validate() - - -# A layer group is whole when every transferred buffer holds the same bytes on every rank -# (replicated, or one KV head shared by all ranks); else it is the share this rank holds. -def _group_shard(shards: Sequence[ShardDesc], fallback: ShardDesc) -> ShardDesc: - parts = {s for s in shards if s.count != 1} - if not parts: - return ShardDesc() - return parts.pop() if len(parts) == 1 else fallback - - def derive_layout(manager: KVCacheManagerV2) -> ManagerLayout: - """The layout of a ``KVCacheManagerV2`` and its device pools. ``ValueError`` if its layer - groups, layer configs and declarations do not agree.""" - tokens_per_block = int(manager.tokens_per_block) - layer_cfg = _manager.layer_configs(manager) - grouping = _manager.layer_group_layer_ids(manager) - pg_of_lg = _manager.layer_group_pool_indices(manager) - if len(pg_of_lg) != len(grouping): - raise ValueError(f"{len(grouping)} layer groups but {len(pg_of_lg)} pool-group indices") - internal = sorted({lid for group in grouping for lid in group}) - global_of = dict(zip(internal, global_layer_ids(manager, internal))) - - state_groups = set() - for lg, layer_ids in enumerate(grouping): - if not layer_ids or layer_ids[0] not in layer_cfg: - raise ValueError(f"layer group {lg} has no layer config (layers {layer_ids})") - if _manager.is_recurrent_layer(layer_cfg[layer_ids[0]]): - state_groups.add(lg) - - # TODO: the base manager declares head-major K/V whichever attention backend writes its pages, - # so a manager paired with a token-major backend gets the head-major layout_id unless it - # declares its layout itself. - mappers = { - role: BufferMapper(int(kind)) for role, kind in _manager.role_mapper_kinds(manager).items() - } - default_mapper = mappers.get(_ALL_ROLE, BufferMapper.INDEXED) - role_layouts = _manager.role_layouts(manager) - ignored = _manager.ignored_roles(manager) - heads = _Heads(manager) - expansions = {lid: _expansions(cfg, tokens_per_block) for lid, cfg in layer_cfg.items()} - - device_pools: dict[int, tuple[DevicePool, ...]] = {} - buffers: list[BufferDesc] = [] - shards: dict[int, list[ShardDesc]] = {lg: [] for lg in range(len(grouping))} - for pg in _manager.pool_group_descriptors(manager): - g = int(pg.pool_group_index) - pools = tuple( - DevicePool( - g, int(p.pool_index), int(p.base_address), int(p.slot_bytes), int(pg.num_slots) - ) - for p in pg.pools - ) - device_pools[g] = tuple(sorted(pools, key=lambda p: p.index)) - # One variant per layer group drawing from this pool group; pool ``i`` of a slot holds the - # ``i``-th coalesced buffer, whose members sit back to back. - for variant in pg.slot_desc.variants: - lg = int(variant.layer_group_id) - state = lg in state_groups - for pool_index, coalesced in enumerate(variant.coalesced_buffers): - size = int(coalesced.single_buffer_size) - for j, buffer_id in enumerate(coalesced.buffer_ids): - lid = int(buffer_id.layer_id) - role = str(buffer_id.role) - mapper = mappers.get(role, default_mapper) - transfer = role not in ignored - buffers.append( - BufferDesc( - layer=global_of.get(lid, lid), - role=role, - pool=(g, pool_index), - offset=j * size, - size=size, - mapper=mapper, - geometry=_buffer_geometry( - role_layouts.get(role), mapper, size, heads.per_rank(lid), state - ), - transfer=transfer, - expansion=expansions.get(lid, {}).get(role, 1), - ) - ) - if not transfer or mapper is BufferMapper.REPLICATED: - continue - shards[lg].append(heads.state_shard() if state else heads.attention_shard(lid)) - - layer_groups: list[LayerGroupDesc] = [] - for lg, layer_ids in enumerate(grouping): - first = layer_cfg[layer_ids[0]] - is_state = lg in state_groups - window = None - sink_tokens = 0 - if not is_state: - window = getattr(first, "window_size", None) - window = None if window is None else int(window) - sink_tokens = int(getattr(first, "num_sink_tokens", None) or 0) - # Layer groups with equal slot sizes may share a pool group, each drawing its own slots. - g = pg_of_lg[lg] - if g not in device_pools: - raise ValueError(f"layer group {lg} maps to pool group {g}, which has no pools") - layer_groups.append( - LayerGroupDesc( - kind="state" if is_state else "attention", - window=window, - sink_blocks=-(-sink_tokens // tokens_per_block), - pool_group=g, - layers=tuple(sorted(global_of[lid] for lid in layer_ids)), - shard=_group_shard(shards[lg], ShardDesc(heads.tp_size, heads.tp_rank)), - ) - ) - - layout = Layout( - tokens_per_block=tokens_per_block, - pool_groups={g: tuple(p.slot_bytes for p in pools) for g, pools in device_pools.items()}, - layer_groups=tuple(layer_groups), - buffers=tuple(buffers), - ) - return ManagerLayout( - layout=layout, - tokens_per_block=tokens_per_block, - pool_group_of=tuple(pg_of_lg), - windows=tuple(desc.window for desc in layer_groups), - sink_blocks=tuple(desc.sink_blocks for desc in layer_groups), - recurrent=tuple(desc.kind == "state" for desc in layer_groups), - sparse=tuple(_sparse(layer_cfg[layer_ids[0]]) for layer_ids in grouping), - layers=tuple(desc.layers for desc in layer_groups), - shards=tuple((desc.shard.count, desc.shard.index) for desc in layer_groups), - page_bytes={g: sum(p.slot_bytes for p in pools) for g, pools in device_pools.items()}, - device_pools=device_pools, - ) + """The page table the manager builds (``KVCacheManagerV2.get_page_table``), as a lender + reads it: buffers the manager rebuilds after a native import travel too, since lending invokes + no rebuild.""" + table = manager.get_page_table() + layout = table.page_layout + if not layout.rebuilds_untransferred: + return table + buffers = tuple(dataclasses.replace(b, transferred=True) for b in layout.buffers) + return dataclasses.replace(table, page_layout=dataclasses.replace(layout, buffers=buffers)) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_lender.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_lender.py index afeb89f5b4c4..ccf0ecd64037 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_lender.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_lender.py @@ -12,86 +12,19 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The host-staging implementation behind the mechanism-independent lending API. - -Relays whole blocks between a request's device pages and host staging slots. - -Leases needing slots are granted first come, first served; a lease with no block entries is -ready at its first poll, with no place in line. Lender and lease methods run only on the -manager's thread. Backends never lend or call a lease; their own threads read a view, access the -memory it points to, call no lease or lender method, and tell the holder through their own -channel. - -Copies queue on the manager's execution stream, serial with the forward on the GPU: a step -queuing copies takes about their time longer. With page-locked staging and the fresh-page fill -off, no lend, poll, mark or resume_bounds call waits for them on the CPU; the manager's shutdown -does, and so does a process's first cache close, which waits for the whole device. Under -confidential computing staging is pageable, and a copy can hold the call queuing it until the -stream reaches it. - -The executor's pool rebalance suspends the caches of its active requests, parked ones among -them, and moves their pages, lent ones included: queued copies complete first, a read still -waiting for slots fails at its grant where its pages moved, and a write marked after the -rebalance copies only the block entries whose pages stayed, all that ``resume_bounds`` counts. - -Caller obligations: - - A staging backend, on the manager's thread, takes a hold with ``hold_regions()`` and - registers ``regions`` after the attach and before it touches a lease; it deregisters - before the manager shuts down, and releases the hold once deregistration is confirmed. - - Each iteration of the executor loop the holder polls every open lease and the waiter asks - ``resume_bounds`` for every waiting request, also when no new request arrives. - - At shutdown the executor loop stops; staging backends stop, deregister and release their - holds; the holder marks and releases what they give back; the manager shuts down last. - - A request whose one-model draft has a joint-reuse pool of its own shares its context - cursor with that pool: the caller fetches the same range into both managers, each - through its own lender, resumes within both intervals and publishes both. A fetch into - one alone leaves the other without the prefix while the shared cursor moves past it. - - While requests are parked in a state other than a disaggregated transfer state, the caller - keeps room for the running requests in each pool group: the pages parked requests' caches - lock, the other pages the scheduler cannot reclaim, and, for the generation request that - needs the most, the pages of its cache within its windows, which a resume locks all at - once, and the pages its next step adds stay within ``max_util_for_resume`` of the group's - GPU pages. The pages a fetch grows are outside the V2 scheduler's reach: it neither - evicts, pauses nor preempts a parked request. Its deadlock check counts no pass while any - request is in a disaggregated transfer state, as one parked in the transfer-in-progress - state is, and otherwise raises "V2 scheduler deadlock" after 1000 passes in a row that - schedule and reclaim nothing while a context or generation request waits, so without that - room a request that cannot be admitted, resume or grow can end in that deadlock. A share - of each pool group for the parked fetches alone does not ensure it, with or without a - cache tier below the GPU. The manager's ``get_page_indices_by_layer_group`` lists a - request's pages per layer group, one beam's (under beam search each further beam also - locks a page of its own for every block not wholly inside the prompt), and the runtime's - ``impl.pool_group_descs`` give each pool group's number of GPU pages and its layer groups. - - Under attention data parallelism without a cache transceiver, the caller keeps a rank from - parking all its active requests, or accepts the stall: the executor counts parked requests - as schedulable, so a rank whose active requests are all parked, at its cap of active - requests or without pages for a padding dummy, schedules nothing, and its empty batch - holds every rank's forward until a fetch settles. - - An integrator adding a writer of recycled pages off the manager's stream makes it wait on - that stream first, as the disaggregated receive does: such a writer is not ordered after - staging copies, which precede only later work on that stream, the work a page's new owner - waits for. With a KV cache connector refused, the manager's only such writer is the - fresh-page fill (``TRTLLM_KV_FRESH_PAGE_FILL``, a diagnostic off by default), which - waits for the manager stream before filling and synchronizes the fill stream afterward: - it overwrites no page a queued copy still reads, and a ``lend_write`` that grows the cache - then waits on the CPU for the queued staging copies. - - The caller releases, or frees the request of, a lease that waits too long: a slot lost to - a failed copy is never reused, and a slot is lost when a copy into or out of it may have - been queued without its completion recorded, because recording that completion failed or a - read's grant failed while queuing the copy. Such failures erode staging capacity: a lease - needing, in a pool group, a longer run of slots than any there without a lost slot waits - in line for good without failing, holding up every lease behind it until it is released or - its request is freed. A lost slot also keeps the staging memory until process exit. - -""" +"""HOST/NAMED staging and DEVICE page leases; public contracts live in ``attach_lender`` and +``_types``. Calls run on the manager thread and order copies and events on its stream; waiting reads +recheck pages, writes count only unchanged pages, unsafe copies quarantine slots.""" from __future__ import annotations import collections -import enum +import dataclasses +import numbers +import struct import traceback import weakref -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING, Deque, Mapping, Sequence import numpy as np @@ -102,12 +35,27 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.logger import logger +from ..page_layout import PageLayout from . import _manager, _ranges from ._fetch_ledger import _FetchLedger from ._identity import Identity from ._layout import ManagerLayout, derive_layout, layout_id +from ._page_locks import CheckedBlocks, LockMode, PageLock, PageLocks from ._slots import SlotQueue, SlotRuns, slot_counts -from ._types import LayerGroupBlocks, LeaseView, MemoryRegion, ResumeBounds, StagingOptions +from ._types import ( + ALL_ARRIVED, + HeldPages, + LeaseView, + LentLayerGroupRows, + MemoryKind, + MemoryRegion, + ResumeBounds, + RowSelection, + StagingOptions, + _check_disjoint, + _frozen_int64, + _immutable, +) if TYPE_CHECKING: from tensorrt_llm.runtime.kv_cache_manager_v2 import _KVCache @@ -116,9 +64,8 @@ from ..kv_cache_manager_v2 import KVCacheManagerV2 from ._fetch_ledger import _FetchRecord -# Staging memory and page-index buffers kept until the process exits, by identity, oldest first; -# only a clean shutdown removes an entry. Each change is one dict operation, atomic under the GIL, -# so lenders of managers on different threads need no lock. +# Staging memory and page-index buffers kept until exit, by identity, oldest first; only a clean +# shutdown removes one. GIL-atomic dict changes let lenders on any thread skip a lock. _retained_until_exit: dict[int, object] = {} _REASON_SHUT_DOWN = "the KV cache manager shut down" @@ -132,6 +79,10 @@ "the lender stopped after an error it logged while the KV cache manager shut down" ) _REASON_SUSPENDED = "the request's cache is suspended" +_REASON_RESERVED = ( + "a KV cache connector holds a prefix reservation for the request, a load not yet started into " + "pages its cache already counts" +) _REASON_PLACEHOLDERS = ( "the request carries multimodal data without digests, which its object_keys do not cover" ) @@ -192,54 +143,57 @@ def free(self) -> None: def _allocate_retained(nbytes: int) -> _HostMemory: """One host allocation of ``nbytes`` for the staging regions, put on the keep list.""" - # TODO: staging duplicates the manager's host tier, which backends cannot use in its place: its - # host pools may move when they resize (mremap), the manager lends none of their pages, and its - # cold-page codec may hold a row in a form other than the row's device bytes. memory = _HostMemory(max(int(nbytes), 1)) _retain_until_exit(memory) return memory def _checked_layout(manager: KVCacheManagerV2) -> ManagerLayout: - """The attach's checks of the manager: a v2 manager, no context or pipeline parallelism, no - lender yet, no recurrent state, no sparse buffer.""" + """The attach's checks: a V2 manager without a lender, sparse buffers or non-Helix context + parallelism, and under Helix no recurrent state or window. A stage lends its own layers.""" _manager.require_v2(manager) mapping = _manager.mapping(manager) - # TODO: the layout derives only tensor-parallel shards, so context-parallel ranks, Helix's - # included, would give different pages the same object_keys. + # TODO: under context parallelism other than Helix the lender cannot place a rank's pages. cp_size = int(mapping.cp_size) - if cp_size > 1: - raise ValueError( - f"context parallelism (cp_size={cp_size}) is not supported: its ranks would give " - "different pages the same object_keys" - ) - # TODO: each pipeline stage's staging slots and windows are its own, so one call can raise on - # one stage and lend on another. - pp_size = int(mapping.pp_size) - if pp_size > 1: + if cp_size > 1 and not mapping.has_cp_helix(): raise ValueError( - f"pipeline parallelism (pp_size={pp_size}) is not supported by staging: each stage's " - "slots and windows are its own, so one call could raise on one stage and lend on " - "another" + f"context parallelism (cp_size={cp_size}) is supported only as Helix: other ranks " + "would give different pages the same object_keys" ) if _manager.attached_lender(manager) is not None: raise ValueError("a lender is already attached to this KV cache manager") - if _manager.has_lossy_cold_page_codec(manager): - raise ValueError( - "the KV cache lender cannot publish original KV bytes by name with a lossy " - "cold-page codec (quantization_for_cold_page)" - ) layout = derive_layout(manager) - recurrent = [lg for lg, state in enumerate(layout.recurrent) if state] - if recurrent: - raise ValueError( - f"layer groups {recurrent} hold recurrent state, which the lender does not lend" - ) + unbounded = [ + lg + for lg in range(layout.num_layer_groups) + if not layout.recurrent[lg] + and layout.windows[lg] is None + and not (layout.prefix_required[lg] if layout.prefix_required else True) + ] + if unbounded: + raise ValueError(f"non-REQUIRED layer groups {unbounded} must have a retained window") + if layout.cp[1] > 1: + recurrent = [lg for lg, state in enumerate(layout.recurrent) if state] + windowed = [ + lg + for lg in range(layout.num_layer_groups) + if layout.windows[lg] is not None or layout.sink_blocks[lg] + ] + if ( + recurrent + or windowed + or _manager.model_has_recurrent_state(manager) + or _manager.model_has_attention_window(manager) + ): + raise ValueError( + f"under Helix the lender lends no recurrent state (layer groups {recurrent}) and " + f"no window (layer groups {windowed}): native's transfer under Helix is not shown " + "to carry them end to end either" + ) # TODO: a sparse layer group's page index can name a host slot, and its locked pages move - # between GPU and host while a lease holds their indices: a decoding cache offloads its - # history, and a prefill lock promotes a host page back to the GPU. + # between GPU and host while a lease holds their indices (decode offloads, prefill promotes). sparse = [lg for lg, flag in enumerate(layout.sparse) if flag] - if sparse: + if sparse or _manager.model_has_sparse_buffers(manager): raise ValueError( f"layer groups {sparse} hold sparse buffers, whose read-only pages a cache can lock " "in host memory" @@ -247,100 +201,147 @@ def _checked_layout(manager: KVCacheManagerV2) -> ManagerLayout: return layout -def _check_commits(manager: KVCacheManagerV2) -> None: - """``ValueError`` for a manager that commits no blocks: a publish lends only committed ones.""" +def _no_names(manager: KVCacheManagerV2, layout: ManagerLayout) -> str | None: + """Why nothing is lent by name, or ``None``: a name covers a committed block's tokens alone, so + its bytes must match across ranks, hold no state, read no later token and stay as computed.""" + if layout.cp[1] > 1: + return "under Helix a block's tokens lie on all its ranks" + recurrent = [lg for lg, state in enumerate(layout.recurrent) if state] + # TODO: the manager keeps a request's recurrent state at a commit's end in a snapshot of the + # reuse tree, whose page no request's cache lends, so no state is lent by name. + if recurrent: + return f"layer groups {recurrent} hold recurrent state, which no name covers" + # A pipeline stage without recurrent layers names nothing either, as every stage of the model. + if _manager.model_has_recurrent_state(manager): + return "the model holds recurrent state, which no name covers, on another pipeline stage" if not _manager.commits_blocks(manager): - raise ValueError( - "staging needs a manager that commits blocks: block reuse on, and joint reuse for a " - "draft manager; with block reuse off nothing could ever be published" + return ( + "the manager commits no blocks (block reuse off, or a draft manager without joint " + "reuse)" ) - - -def _check_lookahead(manager: KVCacheManagerV2) -> None: - """``ValueError`` for a manager whose blocks may depend on tokens past their end: a name covers - the tokens up to the block's end, so two requests sharing them could hold different bytes. A - read-ahead upstream has not established, as one-model DraftTarget's, counts as none.""" - # TODO: staging serves no manager built for a draft that reads ahead (Eagle, MTP): a block's - # name covers only the tokens up to its end, so requests sharing them could hold different - # draft bytes under one name. + if _manager.has_lossy_cold_page_codec(manager): + return ( + "the manager's cold-page codec is lossy (quantization_for_cold_page): a committed " + "block brought back from a colder tier holds other bytes than its tokens computed" + ) + # A read-ahead upstream has not established (one-model DraftTarget's) counts as none. A target + # whose draft lives elsewhere names nothing: a named fetch could never fill that draft's pool. lookahead = _manager.prompt_lookahead(manager) if lookahead > 0: - # The target manager of a draft whose layers live elsewhere is refused too: that draft's own - # pool is, so a fetch could fill the target's blocks but never the draft's. - raise ValueError( - "staging needs blocks that depend on no token past their end; this manager is built " - f"for a one-model draft that reads {lookahead} prompt tokens ahead (Eagle or MTP), so " - "requests sharing a block's tokens could hold different bytes under one name" - ) - - -def _check_connector(manager: KVCacheManagerV2) -> None: - """``ValueError`` for a manager with a KV cache connector.""" - # TODO: a KV cache connector serves a prefix at the first context chunk, from the committed - # tokens, so it cannot lower a history a windowed fetch moved, and it loads off the manager's - # stream, unordered with staging copies. - if _manager.kv_connector(manager) is not None: - raise ValueError( - "staging does not serve a manager with a KV cache connector: the connector serves a " - "request's prefix at its first context chunk, measured from the committed tokens, and " - "cannot lower a history a fetch moved" + return ( + f"the manager is built for a one-model draft that reads {lookahead} prompt tokens past " + "a block's end (Eagle or MTP), so requests sharing the block's tokens could hold " + "different bytes under one name" ) + return None -def _attach_staging( +def _attach_lender( manager: KVCacheManagerV2, *, scope: bytes, staging: StagingOptions, cls: type | None = None -) -> Staging: - """``attach_lender`` with the lender class as a parameter (``Staging`` when ``None``), so a - test can attach a subclass that breaks one rule.""" +) -> _ManagerLender: + """``attach_lender`` with a lender class for tests to replace one rule at a time.""" layout = _checked_layout(manager) - _check_commits(manager) - _check_lookahead(manager) - _check_connector(manager) + if any(window is not None for window in layout.windows) and ( + _ranges._stale_ranges_follow_formula(manager, layout) + ): + layout = dataclasses.replace(layout, stale_ranges_by_formula=True) + no_names = _no_names(manager, layout) if not isinstance(scope, bytes): raise TypeError(f"scope must be bytes, got {type(scope).__name__}") if not isinstance(staging, StagingOptions): raise TypeError(f"staging must be StagingOptions, got {type(staging).__name__}") - counts = slot_counts(layout, staging) - identity = Identity(scope, layout_id(layout.layout), layout.layers, layout.shards) - slots = {g: int(counts.get(g, 0)) for g in layout.pool_groups} - sizes = {g: slots[g] * int(layout.page_bytes[g]) for g in layout.pool_groups} - # TODO: an exception after the allocation keeps the staging memory, and the page-index buffer - # once kept, until the process exits. - memory = _allocate_retained(sum(sizes.values())) - base = memory.address - regions = [] + identity = Identity(scope, layout_id(layout.page_layout), layout.layers, layout.shards) + device = _device_regions(manager, layout, identity) + # The staging path relays named blocks alone, so a lender lending nothing by name stages nothing. + no_staging = no_names + if no_staging is None and _manager.connector_loads(manager): + no_staging = ( + "the manager's KV cache connector writes the pages it loads off the manager's stream, " + "unordered with staging copies" + ) + regions: list[MemoryRegion] = [] + slots: dict[int, int] = {} + memory = None offset = 0 - for g in layout.pool_groups: - groups = tuple(lg for lg, pg in enumerate(layout.pool_group_of) if pg == g) - layout_label = identity.region_label(groups) - regions.append( - MemoryRegion( - layout_label, base + offset, sizes[g], int(layout.page_bytes[g]), slots[g], groups + if no_staging is None: + counts = slot_counts(layout, staging) + slots = {g: int(counts.get(g, 0)) for g in layout.pool_groups} + sizes = {g: slots[g] * int(layout.page_bytes[g]) for g in layout.pool_groups} + # TODO: an exception after the allocation keeps the staging memory, and the page-index + # buffer once kept, until the process exits. + memory = _allocate_retained(sum(sizes.values())) + base = memory.address + for g in layout.pool_groups: + groups = tuple(lg for lg, pg in enumerate(layout.pool_group_of) if pg == g) + layout_label = identity.region_label(groups) + regions.append( + MemoryRegion( + layout_label, + base + offset, + sizes[g], + int(layout.page_bytes[g]), + slots[g], + groups, + memory_kind=MemoryKind.HOST, + ) ) - ) - offset += sizes[g] - lender = (cls or Staging)( + offset += sizes[g] + all_regions = tuple(regions + device) + _check_disjoint(all_regions) + lender = (cls or _ManagerLender)( weakref.ref(manager), layout, identity, - tuple(regions), + all_regions, SlotQueue(slots), - weakref.ref(memory), + _no_memory if memory is None else weakref.ref(memory), + no_names=no_names, + no_staging=no_staging, ) lender._retain_index_buffer(manager) logger.info( f"KV cache lender: namespace {identity.namespace.hex()}, staging {offset >> 20} MiB " - f"in {len(regions)} regions" + f"in {len(regions)} regions, {len(device)} device regions" + f"{'' if no_staging is None else '; nothing staged: ' + no_staging}" ) # Installed last, so a failure above leaves nothing attached. _manager.install_lender(manager, lender) return lender +def _no_memory() -> None: + """The staging memory of a lender that stages nothing: none.""" + return None + + +def _device_regions( + manager: KVCacheManagerV2, layout: ManagerLayout, identity: Identity +) -> list[MemoryRegion]: + """One region per device pool, in pool-group and then pool order, which leases in place lend + from. Its pools must not be rebalanced while the lender lives.""" + # TODO: a slot may straddle two physical allocations of its pool. + # TODO: no pool exposes its handle type, so a region cannot say if other processes can import it. + regions = [] + for g in layout.pool_groups: + groups = tuple(lg for lg, pg in enumerate(layout.pool_group_of) if pg == g) + for pool in layout.device_pools[g]: + regions.append( + MemoryRegion( + identity.device_region_label(groups, pool.index), + int(pool.base), + int(pool.num_slots) * int(pool.slot_bytes), + int(pool.slot_bytes), + int(pool.num_slots), + groups, + memory_kind=MemoryKind.DEVICE, + ) + ) + return regions + + class _QueuedCopy: - """Copies queued together, complete once their event says so, and never without one. ``done`` - asks the event at most once per round of the lender's progress, and never again once it - reported completion.""" + """Copies queued together, complete once their event says so and never without one; ``done`` + asks the event at most once per progress round, and not after it reported completion.""" def __init__(self, event: torch.cuda.Event | None) -> None: self._event = event @@ -373,42 +374,206 @@ def _no_cache_reason(request_id: int) -> str: return f"request {request_id} has no KV cache" -@dataclass(eq=False) -class _LentBlockEntries: - """Block entries grouped by local layer group, with aligned ordinals and slot arrays. +_CUDA_EVENT = torch.cuda.Event - Outer position selects ``layer_group_indices``; each inner position selects one block's - ordinal, device page index and, once assigned, staging slot index. - """ + +def _check_ready_event(ready_event: object, staged: bool = False) -> None: + """``TypeError`` unless ``ready_event`` is ``None`` or an event its holder asks with + ``query()``; for a staging lease, whose copies a stream waits for, one with ``wait()`` too.""" + if ready_event is None: + return + if not callable(getattr(ready_event, "query", None)) or ( + staged and not callable(getattr(ready_event, "wait", None)) + ): + needs = "query() and wait()" if staged else "query()" + raise TypeError( + f"ready_event must be an event with {needs} or None, got {type(ready_event).__name__}" + ) + + +class _BothDone: + """Two readiness conditions, done once both are.""" + + __slots__ = ("_ready_event", "_fence") + + def __init__(self, ready_event: torch.cuda.Event | _BothDone, fence: torch.cuda.Event) -> None: + self._ready_event = ready_event + self._fence = fence + + def query(self) -> bool: + return self._fence.query() and self._ready_event.query() + + +def _fenced( + ready_event: torch.cuda.Event | _BothDone | None, fence: torch.cuda.Event +) -> torch.cuda.Event | _BothDone: + """Combine a stream fence with a lease's earlier readiness condition, where there is one.""" + return fence if ready_event is None else _BothDone(ready_event, fence) + + +def _frozen_array(values: np.ndarray) -> np.ndarray: + """``values`` as contiguous ``int64`` on immutable bytes of their own.""" + return np.frombuffer(np.ascontiguousarray(values, dtype=np.int64).tobytes(), dtype=np.int64) + + +def _unshared(values: np.ndarray) -> np.ndarray: + """``values`` on immutable bytes as an array object of its own, for one lease's view: no other + view, record or lock holds it.""" + return _immutable(values).view() + + +# No rows: the ordinals or pages of a layer group that lends none; a view holds a view of it. +_NO_ROWS = _frozen_int64(()) +# The ordinal of the one row of a layer group of recurrent state; a view holds a view of it. +_STATE_ROW = _frozen_int64((0,)) + +# Up to this many blocks, a cache's pages and a fast write's range of its row are read as lists, +# cheaper than NumPy's casts and reductions, whose first call after other work takes microseconds. +_LISTED_BLOCKS = 128 +# Fast lends of blocks below this take their ordinals as views of these bytes, the entries and the +# locks the same view; a lease's view holds a view of its own. +_ORDINAL_BLOCKS = 1 << 14 +_ORDINALS = _frozen_int64(range(_ORDINAL_BLOCKS)) +# Views of ``_ORDINALS`` by range, each made once; emptied when it holds ``_ORDINAL_VIEWS_KEPT``. +_ordinal_views: dict[tuple[int, int], np.ndarray] = {} +_ORDINAL_VIEWS_KEPT = 4096 +# What ``_ManagerLender._return_whole`` gives where it returns nothing, for the rules to return. +_NOT_WHOLE = object() +# The unchecked constructor of views, bound once for fast lends. +_new_view = LeaseView._unchecked +_new_object = object.__new__ +_set_attribute = object.__setattr__ + + +def _lent_rows( + lg: int, ordinals: np.ndarray, slots: np.ndarray, regions: tuple[int, ...] +) -> LentLayerGroupRows: + """``LentLayerGroupRows._unchecked(lg, ordinals, slots=slots, region_indices=regions)`` of + positional arguments, as a fast lend builds each layer group's rows.""" + rows = _new_object(LentLayerGroupRows) + _set_attribute( + rows, + "__dict__", + { + "layer_group_index": lg, + "page_ordinals": ordinals, + "object_keys": None, + "slots": slots, + "region_indices": regions, + }, + ) + return rows + + +# Packers of 0 to ``_LISTED_BLOCKS`` unsigned 64-bit integers, by count. +_PAGE_PACKERS = tuple(struct.Struct(f"={count}Q") for count in range(_LISTED_BLOCKS + 1)) +# The dtypes of fast lends' pages as dtype objects, which NumPy takes without converting them. +_from_buffer = np.frombuffer +_INT64 = np.dtype(np.int64) +_INT32 = np.dtype(np.int32) +_UINT64 = np.dtype(np.uint64) +_UINT32 = np.dtype(np.uint32) + + +def _frozen_pages(pages: Sequence[int]) -> np.ndarray: + """At most ``_LISTED_BLOCKS`` non-negative ``pages`` as ``int64`` on immutable bytes of their + own; ``struct.error`` for a negative one.""" + return np.frombuffer(_PAGE_PACKERS[len(pages)].pack(*pages), _INT64) + + +class _ParkedRow: + """Lender-owned ``int32`` page-index rows, one per layer group, that a cache whose index slot + was released early writes its beam-0 page indices into until it closes.""" + + __slots__ = ("pages", "bytes", "buffers", "source", "kv") + + def __init__(self, groups: int, blocks: int) -> None: + self.pages = np.full((groups, blocks), -1, dtype=np.int32) + self.bytes = memoryview(self.pages).cast("B") + # The buffer the manager binds each layer group's page indices to. + self.buffers = tuple(memoryview(self.pages[lg]) for lg in range(groups)) + # As a lend reads a row: its words, where it starts, the words of a group, its blocks. + self.source = (memoryview(self.pages).cast("B").cast("i"), 0, blocks, blocks) + # The cache bound to the row, which writes into it until it closes. + self.kv: _KVCache | None = None + + +@dataclass(eq=False, slots=True) +class _LentBlockEntries: + """Per-group aligned block ordinals, device pages and staging slots (none until a grant); a fast + lend of every row adds ``spans``, each group's run ``(first, last)`` (empty if + ``last <= first``, ``None`` for state), and ``block_end``, where its block runs end.""" layer_group_indices: list[int] block_ordinals: list[np.ndarray] device_pages: list[np.ndarray] - staging_slots: list[np.ndarray] = field(default_factory=list) + staging_slots: list[np.ndarray] | tuple[()] = () + spans: list[tuple[int, int] | None] | None = None + block_end: int = 0 @property - def num_block_entries(self) -> int: + def num_rows(self) -> int: return sum(len(o) for o in self.block_ordinals) -class Staging: - """``Lender`` over one manager, referenced weakly. What its calls do depends on its mode. - - - Serving the manager: a lend past its range check, a poll and a resume_bounds call return - recyclable leases' slots and grant waiting leases in order before their own work, a mark and a - release after it. - - After the manager's reuse reset: as serving, but a lend that raises no ``ValueError`` fails - at the call. - - Stopped for good by an error its free or shutdown hook logs: a lend that raises no - ``ValueError`` fails at the call, a poll fails a lease whose view no poll returned, a mark - copies nothing, and the shutdown keeps the staging memory; while the manager is served, - progress fails the waiting leases and returns no slot, and resume_bounds gives a request with - a cache an empty interval. - - The manager shut down or gone: a lend that raises no ``ValueError`` fails at the call, a poll - runs no progress and returns the view only once ready, a mark copies nothing and resume_bounds - raises ``ValueError``; a stopped lender's poll still fails a lease whose view no poll - returned. - """ +class _ManagerLender: + """Lender over one weakly held manager: refusals return failed leases, permanent stops raise + ``RuntimeError``; progress recycles and grants staging slots, which DEVICE leases bypass.""" + + # The lender's own state is in slots; the instance dict allows instrumented lenders to + # override methods without changing its storage layout. + __slots__ = ( + "_manager_ref", + "_layout", + "_tokens_per_block", + "_identity", + "_regions", + "_slot_queue", + "_memory_ref", + "_no_names", + "_no_staging", + "_recurrent_groups", + "_region_index_of_pool_group", + "_device_regions_of_pool_group", + "_staging_region_indices", + "_device_slot_limit", + "_any_window", + "_ledger", + "_waiting_for_slots", + "_leases_holding_slots", + "_open_leases", + "_open_holds", + "_quarantined", + "_shut_down", + "_reuse_was_reset", + "_stop_reason", + "_index_buffer", + "_progress_round", + "_locks", + "_named_groups", + "_page_layout", + "_ordinal_ranges", + "_group_regions", + "_fast_groups", + "_fast_layout", + "_state_rows", + "_shrink_epochs", + "_state_group", + "_state_record", + "_state_unverified", + "_index_rows", + "_index_words", + "_index_strides", + "_parked", + "_parked_free", + "_parked_closing", + "_parked_made", + "_bad_row", + "_writes_at_end", + "_margin_groups", + "__dict__", + "__weakref__", + ) def __init__( self, @@ -418,170 +583,428 @@ def __init__( regions: tuple[MemoryRegion, ...], slots: SlotQueue, memory: weakref.ref, + *, + no_names: str | None = None, + no_staging: str | None = None, ) -> None: self._manager_ref = manager self._layout = layout + self._tokens_per_block = int(layout.tokens_per_block) self._identity = identity self._regions = tuple(regions) self._slot_queue = slots - # Only the keep list holds the staging memory strongly; this finds it there at shutdown. + # Only the keep list holds the staging memory strongly; this finds it there at shutdown, + # or ``None`` for a lender that stages nothing. self._memory_ref = memory + self._no_names = no_names + # Why the lender stages nothing, or ``None`` where it stages. + self._no_staging = no_names if no_staging is None else no_staging + recurrent = tuple(getattr(layout, "recurrent", ())) + # The layer groups of recurrent state: each lends one row, its state page, as block 0. + self._recurrent_groups = frozenset(lg for lg, flag in enumerate(recurrent) if flag) self._region_index_of_pool_group = {g: i for i, g in enumerate(layout.pool_groups)} + # Per pool group, the positions of its device regions in ``regions``, in pool order. + self._device_regions_of_pool_group: dict[int, tuple[int, ...]] = {} + for index, region in enumerate(self._regions): + if region.memory_kind is MemoryKind.DEVICE: + g = int(layout.pool_group_of[region.layer_group_indices[0]]) + held = self._device_regions_of_pool_group.get(g, ()) + self._device_regions_of_pool_group[g] = (*held, index) + # Per pool group, the one-region tuple its staged rows name, shared by every lend. + self._staging_region_indices = { + g: (index,) + for g, index in self._region_index_of_pool_group.items() + if index < len(self._regions) and self._regions[index].memory_kind is MemoryKind.HOST + } + # Per layer group, the least ``slot_count`` of the device regions its rows in place span, + # which no slot such a row holds may reach. + self._device_slot_limit: dict[int, int] = {} + for region in self._regions: + if region.memory_kind is MemoryKind.DEVICE: + for lg in region.layer_group_indices: + held = self._device_slot_limit.get(lg, region.slot_count) + self._device_slot_limit[lg] = min(held, region.slot_count) self._any_window = any(w is not None for w in layout.windows) - # The fetch records and the resume_bounds they give. The ledger asks this lender whether a - # copy completed through a weak reference, so no reference cycle keeps the lender, or a - # cache its records hold, past the lender's last reference. + # Fetch records and the bounds they give; the ledger asks this lender about copies through a + # weak reference, so no cycle keeps the lender or a cache its records hold. lender = weakref.ref(self) self._ledger = _FetchLedger(layout, lambda copy: lender()._copy_done(copy)) - # Leases waiting for slots, in order. self._waiting_for_slots: Deque[_StagingLease] = collections.deque() self._leases_holding_slots: list[_StagingLease] = [] # granted, slots not yet returned # Open leases, held strongly: an open lease keeps the staging memory at shutdown even # once its holder has dropped it. self._open_leases: set[_StagingLease] = set() - # Open staging holds, held strongly too: a dropped hold still keeps the memory. - self._open_holds: set[_StagingHold] = set() + # Open region holds, held strongly too: a dropped hold still keeps the memory. + self._open_holds: set[_RegionHold] = set() self._quarantined: list[SlotRuns] = [] # slots a failed copy may still touch; never reused self._shut_down = False - # Whether the manager reset its reuse state; once it has, nothing is lent by name again. self._reuse_was_reset = False self._stop_reason: str | None = None # why lending stopped for good; the first reason stays self._index_buffer: object | None = None # kept until the runtime's shutdown has returned + # The kept page-index buffer fast lends read pages from: as an array, as one flat run of + # words, and the words a layer group, a row and a half row span (a row's blocks). + self._index_rows: np.ndarray | None = None + self._index_words: memoryview | None = None + self._index_strides = (0, 0, 0) + # Parked rows by request; rows of freed requests whose cache may be open; rows ready to + # bind; every row made, kept with the page-index buffer since closing caches write them. + self._parked: dict[int, _ParkedRow] = {} + self._parked_closing: list[_ParkedRow] = [] + self._parked_free: list[_ParkedRow] = [] + self._parked_made: list[_ParkedRow] = [] + # A row's bytes with the bad index in every word, to clear a row handed out again. + self._bad_row = b"" self._progress_round = 0 # rounds of progress, so each asks a copy's event at most once + # The locks on blocks of the manager's caches, and the locked caches it closes later. + self._locks = PageLocks(manager, recurrent, getattr(layout, "pool_group_of", None)) + self._locks.caps_history = self._any_window + # Per layer group, whether it lends by name: one holding a layer the manager keeps out of + # prefix reuse does not. + self._named_groups = tuple(layout.named_groups) + self._page_layout: PageLayout | None = None # built at its first read + self._ordinal_ranges: dict[int, dict[tuple[int, int], tuple[int, np.ndarray]]] = {} + # Per layer group, the positions of the device regions its rows are placed in. + self._group_regions = tuple( + self._device_regions_of_pool_group.get(int(g), ()) + for g in getattr(layout, "pool_group_of", ()) + ) + # Per layer group: its index, the positions of its device regions, its device slot limit, + # whether it holds recurrent state, and its window, as fast lends read them. + windows = tuple(getattr(layout, "windows", ())) + self._fast_groups = tuple( + ( + lg, + regions, + self._device_slot_limit.get(lg, 0), + lg in self._recurrent_groups, + windows[lg] if lg < len(windows) else None, + ) + for lg, regions in enumerate(self._group_regions) + ) + # Whether a lend of every row in place may take a fast lend: no Helix, a device region for + # every layer group, and windows without sinks whose stale ranges follow the window formula. + # TODO: under Helix every in-place lend, mark and resume-bound query takes the rules. + by_formula = getattr(layout, "stale_ranges_by_formula", False) + sinks = tuple(getattr(layout, "sink_blocks", ())) + self._fast_layout = ( + getattr(layout, "cp", (0, 1))[1] == 1 + and bool(self._group_regions) + and all(self._group_regions) + and all( + window is None or (by_formula and lg < len(sinks) and not sinks[lg]) + for lg, window in enumerate(windows) + ) + ) + # One frozen row per state slot a fast lend gave, by slot, for the entries of fast lends; a + # lease's view holds a view of its own. + self._state_rows: dict[int, np.ndarray] = {} + # Per request, in-place shrinks the manager reported since the last free: a fast write's + # marks trust the pages it lent only while this count stays as at the lend. + self._shrink_epochs: dict[int, int] = {} + # The manager's record of each request's recurrent-state slot and its group, read instead of + # the cache; requests resumed since it last matched, which the cache answers for meanwhile. + prepared = _manager.prepared_state_slots(manager()) if manager() is not None else None + self._state_group, self._state_record = prepared if prepared is not None else (None, None) + self._state_unverified: set[int] = set() + # Open fast write leases whose return would resume at their end, by ``id`` of their cache, + # at most one per cache; the hooks that can move a cache's history take them out. + self._writes_at_end: dict[int, _InPlaceLease] = {} + # Layer groups that keep a window wider than the one they read, whose margin rows a fetch + # need not deliver. + self._margin_groups = frozenset( + lg for lg in range(len(windows)) if _ranges._read_slack(layout, lg) > 0 + ) @property def regions(self) -> tuple[MemoryRegion, ...]: """See ``Lender.regions``.""" return self._regions - def hold_regions(self) -> _StagingHold: + @property + def page_layout(self) -> PageLayout: + """See ``Lender.page_layout``.""" + if self._page_layout is None: + # Device regions follow pool-group, then pool order; a staging region serves a group. + device = {} + for g, positions in self._device_regions_of_pool_group.items(): + for pool, position in zip(self._layout.device_pools[g], positions): + device[(g, int(pool.index))] = position + staging = { + int(self._layout.pool_group_of[region.layer_group_indices[0]]): position + for position, region in enumerate(self._regions) + if region.memory_kind is MemoryKind.HOST + } + self._page_layout = dataclasses.replace( + self._layout.page_layout, device_region_index=device, staging_region_index=staging + ) + return self._page_layout + + def held_page_counts(self) -> HeldPages: + """See ``Lender.held_page_counts``.""" + return self._locks.held_page_counts() + + def hold_regions(self) -> _RegionHold: """See ``Lender.hold_regions``.""" if self._serving_manager() is None: # The memory was freed or kept at the shutdown; a later hold changes neither. - return _StagingHold(None) - hold = _StagingHold(self) + return _RegionHold(None) + hold = _RegionHold(self) self._open_holds.add(hold) return hold - def lend_read(self, request: LlmRequest, start: int, end: int) -> _StagingLease: + def lend_read( + self, + request: LlmRequest, + start: int, + end: int, + *, + memory_kind: MemoryKind, + rows: RowSelection, + ready_event: torch.cuda.Event | None, + ) -> _StagingLease | _InPlaceLease: """See ``Lender.lend_read``.""" - start, end = self._whole_blocks(start, end) + # A DEVICE lend of every row passes every check of the kind and rows. + fast = memory_kind is MemoryKind.DEVICE and rows is RowSelection.ALL and self._lends_fast() + if ( + fast + and (ready_event is None or type(ready_event) is _CUDA_EVENT) + and type(start) is int + and type(end) is int + and 0 <= start <= end + and not start % self._tokens_per_block + and not self._shut_down + and self._stop_reason is None + ): + # Every check below passes: the fast lend, or the rules alone where it lends nothing. + manager = self._manager_ref() + if manager is not None: + lease = self._read_every_row_fast( + manager, request, int(request.py_request_id), start, end, ready_event + ) + if lease is not None: + return lease + fast = False + if not fast: + self._check_memory_and_rows(memory_kind, rows) + if ready_event is not None and type(ready_event) is not _CUDA_EVENT: + _check_ready_event(ready_event, memory_kind is MemoryKind.HOST) + named = rows is RowSelection.NAMED + start, end = self._whole_blocks(start, end, named) request_id = int(request.py_request_id) - manager = self._serving_manager() - if manager is None: - return _StagingLease._failed_at_call(self, "read", request_id, _REASON_SHUT_DOWN) - self._progress() - if self._stop_reason is not None: - return _StagingLease._failed_at_call(self, "read", request_id, self._stop_reason) - if self._reuse_was_reset: - return _StagingLease._failed_at_call(self, "read", request_id, _REASON_REUSE_RESET) + manager = self._lending_manager() + self._check_not_stopped() + if fast: + lease = self._read_every_row_fast(manager, request, request_id, start, end, ready_event) + if lease is not None: + return lease + if named and self._reuse_was_reset: + return self._failed_at_call(memory_kind, "read", request_id, _REASON_REUSE_RESET) kv = _manager.kv_of(manager, request_id) if kv is None: - return _StagingLease._failed_at_call( - self, "read", request_id, _no_cache_reason(request_id) + return self._failed_at_call( + memory_kind, "read", request_id, _no_cache_reason(request_id) ) - unnamed = self._unnamed_input_refusal(request) + unnamed = self._unnamed_input_refusal(request) if named else None if unnamed is not None: - return _StagingLease._failed_at_call(self, "read", request_id, unnamed) + return self._failed_at_call(memory_kind, "read", request_id, unnamed) state = _manager.cache_state(kv) - if end > state.committed: - raise ValueError( - f"the range ends at {end}, past the {state.committed} committed tokens" - ) + refusal = self._read_end_refusal(manager, request, state, end, named) + if refusal is not None: + return self._failed_at_call(memory_kind, "read", request_id, refusal) if not state.active: - return _StagingLease._failed_at_call(self, "read", request_id, _REASON_SUSPENDED) - # TODO: a publish leaves out window blocks the request's own window has passed, although - # the prefix tree may still hold their committed pages, so a fetch whose window still keeps - # such a block finds it missing. + return self._failed_at_call(memory_kind, "read", request_id, _REASON_SUSPENDED) + # TODO: a publish leaves out blocks the request's window passed though the prefix tree may + # still hold their pages, so a fetch whose window keeps such a block finds it missing. entries = self._lendable_block_entries( - manager, state.history, self._block_entries_of_range(manager, kv, start, end) + manager, + state.history, + self._block_entries_of_range(manager, kv, start, end, request_id=request_id), ) - counts = self._required_slots_per_pool_group( - [len(block_ordinals) for block_ordinals in entries.block_ordinals] + if named: + entries = self._committed_entries(kv, entries) + if memory_kind is MemoryKind.DEVICE: + entries = self._without_short_windows_mid_context(manager, request, state, entries) + else: + counts = self._required_slots_per_pool_group( + [len(block_ordinals) for block_ordinals in entries.block_ordinals] + ) + refusal = self._fits_refusal(counts) + if refusal is not None: + return self._failed_at_call(memory_kind, "read", request_id, refusal) + if _manager.connector_reserved(manager, request): + return self._failed_at_call(memory_kind, "read", request_id, _REASON_RESERVED) + reuse_keys = ( + self._reuse_keys_for(manager, request, kv, entries.block_ordinals) if named else None ) - self._check_fits(counts) - reuse_keys = self._reuse_keys_for(manager, request, kv, entries.block_ordinals) + if memory_kind is MemoryKind.DEVICE: + refusal = self._ledger._unsettled_fetch_refusal(request_id, kv, entries) + if refusal is None: + refusal = self._unfetched_rows_refusal(request_id, kv, state, entries, end) + if refusal is not None: + return self._failed_at_call(memory_kind, "read", request_id, refusal) + return self._lend_in_place( + manager, + request, + kv, + "read", + entries, + reuse_keys, + start=start, + end=end, + ready_event=ready_event, + ) lease = _StagingLease(self, "read", request_id, kv, entries, reuse_keys) + lease._ready_event = ready_event + lease._range = (start, end) self._admit(lease, counts) return lease - def lend_write(self, request: LlmRequest, start: int, end: int) -> _StagingLease: + def lend_write( + self, + request: LlmRequest, + start: int, + end: int, + *, + memory_kind: MemoryKind, + rows: RowSelection, + ready_event: torch.cuda.Event | None, + ) -> _StagingLease | _InPlaceLease: """See ``Lender.lend_write``.""" - start, end = self._whole_blocks(start, end) + fast = memory_kind is MemoryKind.DEVICE and rows is RowSelection.ALL and self._lends_fast() + if ( + fast + and (ready_event is None or type(ready_event) is _CUDA_EVENT) + and type(start) is int + and type(end) is int + and 0 <= start <= end + and not start % self._tokens_per_block + and not self._shut_down + and self._stop_reason is None + ): + # No Helix on a layout fast lends cover: the prompt is the request's own. + prompt = _manager.prompt_length(request) + manager = self._manager_ref() if end <= prompt else None + if manager is not None: + lease = self._write_every_row_fast( + manager, request, int(request.py_request_id), start, end, prompt, ready_event + ) + if lease is not None: + return lease + fast = False + if not fast: + self._check_memory_and_rows(memory_kind, rows) + if ready_event is not None and type(ready_event) is not _CUDA_EVENT: + _check_ready_event(ready_event, memory_kind is MemoryKind.HOST) + named = rows is RowSelection.NAMED + start, end = self._whole_blocks(start, end, named) request_id = int(request.py_request_id) - manager = self._serving_manager() - if manager is None: - return _StagingLease._failed_at_call(self, "write", request_id, _REASON_SHUT_DOWN) - self._progress() layout = self._layout tpb = int(layout.tokens_per_block) - # The request computes its last prompt token itself, for its logits, so a fetch ends at - # the whole blocks before it. - prompt = _manager.prompt_length(request) - if end > (prompt - 1) // tpb * tpb: + prompt = _manager.prompt_length(request, cp_size=getattr(layout, "cp", (0, 1))[1]) + # The request computes its last prompt token itself, for its logits, so a named fetch ends + # at the whole blocks before it; a fetch of every row may reach the prompt's end. + if named and end > (prompt - 1) // tpb * tpb: raise ValueError( f"the range ends at {end}, past the whole blocks before the request's last prompt " f"token ({prompt} prompt tokens)" ) - # Block entries and their keys come from the layout for a history of ``end``, so every - # ValueError is raised before the cache changes. The entries are counted from their ranges - # and checked to fit before any array is built. - runs = [ - _ranges._needed_ranges(manager, layout, lg, start // tpb, end // tpb, end) - for lg in range(layout.num_layer_groups) - ] - counts = self._required_slots_per_pool_group( - [sum(e - b for b, e in lg_runs) for lg_runs in runs] - ) - self._check_fits(counts) - block_ordinals = [_ranges._ordinals(lg_runs) for lg_runs in runs] - # The checks above read only the call, the request's prompt and the layout, so on a - # manager not shut down a wrong call raises whatever this rank holds for the request. - if self._stop_reason is not None: - return _StagingLease._failed_at_call(self, "write", request_id, self._stop_reason) - if self._reuse_was_reset: - return _StagingLease._failed_at_call(self, "write", request_id, _REASON_REUSE_RESET) + if not named and end > prompt: + raise ValueError(f"the range ends at {end}, past the request's {prompt} prompt tokens") + manager = self._lending_manager() + self._check_not_stopped() + if fast: + lease = self._write_every_row_fast( + manager, request, request_id, start, end, prompt, ready_event + ) + if lease is not None: + return lease + if named and self._reuse_was_reset: + return self._failed_at_call(memory_kind, "write", request_id, _REASON_REUSE_RESET) kv = _manager.kv_of(manager, request_id) if kv is None: - return _StagingLease._failed_at_call( - self, "write", request_id, _no_cache_reason(request_id) + return self._failed_at_call( + memory_kind, "write", request_id, _no_cache_reason(request_id) ) - unnamed = self._unnamed_input_refusal(request) + unnamed = self._unnamed_input_refusal(request) if named else None if unnamed is not None: - return _StagingLease._failed_at_call(self, "write", request_id, unnamed) + return self._failed_at_call(memory_kind, "write", request_id, unnamed) state = _manager.cache_state(kv) - if start < (state.committed // tpb) * tpb: - raise ValueError( - f"the range starts at {start}, inside the committed whole blocks of " - f"{state.committed} tokens" - ) - reuse_keys = self._reuse_keys_for(manager, request, kv, block_ordinals) - if self._any_window and end < state.history: - raise ValueError( - f"the range ends at {end}, below the history of {state.history} tokens its " - "windows keep" + # Count future rows before allocating ordinals or growing the cache. + committed = _manager.committed_by_group(kv, layout.num_layer_groups) + starts = [ + -(-end // tpb) * tpb + if lg not in self._recurrent_groups and tokens > end + else max(start, tokens // tpb * tpb) + for lg, tokens in enumerate(committed) + ] + runs = [ + _ranges._needed_ranges(manager, layout, lg, starts[lg] // tpb, -(-end // tpb), end) + if not named or self._named_groups[lg] + else [] + for lg in range(layout.num_layer_groups) + ] + refusal = None + if memory_kind is MemoryKind.HOST: + counts = self._required_slots_per_pool_group( + [sum(e - b for b, e in lg_runs) for lg_runs in runs] ) + refusal = self._fits_refusal(counts) + if refusal is None: + refusal = self._write_range_refusal(manager, request, state, start, end, prompt) + if refusal is not None: + return self._failed_at_call(memory_kind, "write", request_id, refusal) + block_ordinals = [ + self._range_ordinals(request_id, lg, lg_runs) for lg, lg_runs in enumerate(runs) + ] + reuse_keys = self._reuse_keys_for(manager, request, kv, block_ordinals) if named else None split = self._bidirectional_split_refusal(request, end) if split is not None: - return _StagingLease._failed_at_call(self, "write", request_id, split) + return self._failed_at_call(memory_kind, "write", request_id, split) if not state.active: - return _StagingLease._failed_at_call(self, "write", request_id, _REASON_SUSPENDED) + return self._failed_at_call(memory_kind, "write", request_id, _REASON_SUSPENDED) if self._scratch_reuse_on(kv): - return _StagingLease._failed_at_call(self, "write", request_id, _REASON_SCRATCH) - if self._returns_context_outputs(request): - return _StagingLease._failed_at_call(self, "write", request_id, _REASON_CONTEXT_OUTPUTS) + return self._failed_at_call(memory_kind, "write", request_id, _REASON_SCRATCH) + if named and self._returns_context_outputs(request): + return self._failed_at_call(memory_kind, "write", request_id, _REASON_CONTEXT_OUTPUTS) # After every ValueError: whether an earlier fetch settled is this rank's own timing. unsettled = self._ledger._unsettled_fetch_refusal(request_id, kv) if unsettled is not None: - return _StagingLease._failed_at_call(self, "write", request_id, unsettled) - # With a window the history moves to ``end``, so windows need pages only for the blocks a - # history of that length reads. - position = end if self._any_window else state.history + return self._failed_at_call(memory_kind, "write", request_id, unsettled) + # A named fetch's grow moves a windowed history to ``end``, so windows need pages only for + # what it reads; a fetch of every row writes pages held since admission, history kept. + position = end if named and self._any_window else state.history + if position > state.history and _manager.connector_load_pending(manager, request): + return self._failed_at_call( + memory_kind, + "write", + request_id, + "a window cannot advance history while a connector load is pending", + ) refusal = self._history_refusal(manager, state.history, position, end) if refusal is None: refusal = self._unwritten_pages_refusal(manager, request_id, kv, position) + cap = self._history_cap(kv) + if refusal is None and cap is not None and position > cap: + refusal = f"a window would leave a locked block behind at a history of {position}" + if refusal is not None: + return self._failed_at_call(memory_kind, "write", request_id, refusal) + # HOST copies and DEVICE writers share the same exclusion domain before any grow. + planned = ( + None + if named + else self._block_entries_of_range( + manager, kv, starts, end, request_id=request_id, ordinals=block_ordinals + ) + ) + refusal = None if named else self._pageless(planned) + ordinals = block_ordinals if named else planned.block_ordinals + checked = self._locks.check( + kv, dict(enumerate(ordinals)), LockMode.WRITE, ordered=True, committed=committed + ) + if refusal is None and isinstance(checked, str): + refusal = f"the blocks cannot be locked: {checked}" if refusal is not None: - return _StagingLease._failed_at_call(self, "write", request_id, refusal) + return self._failed_at_call(memory_kind, "write", request_id, refusal) # Read before the grow, which moves a windowed cache's history to ``end``. kept = self._ledger._prefix_kept_by_fetch(request_id, kv, start) # TODO: the V2 scheduler cannot reclaim the pages a parked fetch grew. @@ -589,71 +1012,540 @@ def lend_write(self, request: LlmRequest, start: int, end: int) -> _StagingLease # below it, until the request finishes, so a later resume leaves positions without routes. # TODO: with extra KV tokens the grow locks a page past the history it leaves, which nothing # writes, so a later all-reusable windowed lease that leaves that block behind fails. - # TODO: an exception once the grow resized the cache, in its fill or a step below, can leave - # no record of a windowed cache's moved history, so resume_bounds counts tokens never fetched as - # computed, or a fetch record nothing settles, so resume_bounds stays None. - if not _manager.grow(manager, request, kv, position, end): - return _StagingLease._failed_at_call( - self, "write", request_id, f"no free pages to grow the cache to {end} tokens" + if named: + stream = _manager.cache_stream(manager, request_id) + if ready_event is not None and not ready_event.query(): + wait = getattr(ready_event, "wait", None) + if not callable(wait) or stream is None: + return self._failed_at_call( + memory_kind, + "write", + request_id, + "a named grow needs a completed or stream-waitable ready_event", + ) + wait(stream) + fence = ( + torch.cuda.Event() + if memory_kind is MemoryKind.DEVICE and stream is not None + else None + ) + try: + grown = _manager.grow(manager, request, kv, position, end) + except RuntimeError: + # A diagnostic page fill can fail after resize moved a window's history. Keep + # that unfetched interval in the ledger, so it cannot count as computed tokens. + if _manager.cache_state(kv).history > state.history: + self._ledger._record_fetch( + request_id, + kv, + start, + kept, + position, + committed, + state_end=end if self._recurrent_groups else None, + end=end, + prompt=prompt, + runs=_manager.bidirectional_runs(request), + ) + _manager.drop_held_commit(manager, kv, start) + raise + if not grown: + return self._failed_at_call( + memory_kind, + "write", + request_id, + f"no free pages to grow the cache to {end} tokens", + ) + # From here poll_resume_bounds accounts for the fetch, whatever this lease's outcome. + fetch = self._ledger._record_fetch( + request_id, + kv, + start, + kept, + position, + committed, + state_end=end if self._recurrent_groups else None, + end=end, + prompt=prompt, + runs=_manager.bidirectional_runs(request), + ) + # A commit the cache's locks hold back would commit computed rows this fetch may overwrite. + _manager.drop_held_commit(manager, kv, start) + if named and fence is not None: + # A grow can queue a wait for pages evicted to the host tier. Keep the view behind + # that work, and retain the fetch's prefix accounting if recording its fence fails. + try: + fence.record(stream) + except RuntimeError: + stream.synchronize() + if self._ledger._latest_fetches.get(request_id) is fetch: + del self._ledger._latest_fetches[request_id] + raise + ready_event = _fenced(ready_event, fence) + entries = ( + self._block_entries_of_range( + manager, + kv, + starts, + end, + named=named, + request_id=request_id, + ordinals=block_ordinals, + ) + if planned is None + else planned + ) + fetch.entries = entries + doomed = self._pageless(entries) if planned is None else None + if memory_kind is MemoryKind.DEVICE: + # An ALL write locks the blocks it checked; a NAMED write checks its blocks again after + # its grow. + return self._lend_in_place( + manager, + request, + kv, + "write", + entries, + reuse_keys, + fetch, + doomed, + start=start, + end=end, + checked=None if named else checked, + ready_event=ready_event, ) - # The cache has grown, and resume_bounds accounts for it whatever this lease's outcome. - fetch = self._ledger._record_fetch(request_id, kv, start, kept, position, state.committed) - entries = self._block_entries_of_range(manager, kv, start, end) - doomed = None - for lg, lg_ordinals, lg_slots in zip( - entries.layer_group_indices, entries.block_ordinals, entries.device_pages - ): - if doomed is None and np.any(lg_slots < 0): - missing = lg_ordinals[lg_slots < 0].tolist() - doomed = f"layer group {lg}: blocks {missing[:8]} have no page" lease = _StagingLease(self, "write", request_id, kv, entries, reuse_keys, fetch) + lease._ready_event = ready_event + lease._range = (start, end) if doomed is not None: # It fails at its first poll, which abandons the fetch. - lease._first_poll_failure = doomed + lease._defer_failure(doomed) self._open_leases.add(lease) return lease self._admit(lease, counts) return lease - def resume_bounds(self, request: LlmRequest) -> ResumeBounds | None: - """See ``Lender.resume_bounds``.""" + def poll_resume_bounds(self, request: LlmRequest) -> ResumeBounds | None: + """See ``Lender.poll_resume_bounds``.""" request_id = int(request.py_request_id) manager = self._serving_manager() if manager is None: raise ValueError(f"{_no_cache_reason(request_id)}: {_REASON_SHUT_DOWN}") - self._progress() kv = _manager.kv_of(manager, request_id) if kv is None: raise ValueError(_no_cache_reason(request_id)) - history = _manager.cache_state(kv).history + state = _manager.cache_state(kv) if self._stop_reason is not None: # Empty: no fetch through a stopped lender counts. - return ResumeBounds(max_position=0, min_position=max(history, 1)) - return self._ledger._resume_bounds(request, request_id, manager, kv, history) + return ResumeBounds(max_position=0, min_position=max(state.history, 1)) + return self._ledger._resume_bounds( + _manager.prompt_length(request, manager), + _manager.bidirectional_runs(request), + request_id, + manager, + kv, + state, + ) + + def return_writes( + self, + entries: Sequence[tuple[_StagingLease | _InPlaceLease, object, bool]], + ) -> list[ResumeBounds | None]: + """See ``Lender.return_writes``.""" + if type(entries) is list and len(entries) == 1: + # One entry: checked, then returned, raising what its return raises; an in-place write + # returned with every row passes ``_accepted_return``'s checks here. + entry = entries[0] + if type(entry) is tuple and len(entry) == 3: + lease, arrived, reusable = entry + if ( + arrived is ALL_ARRIVED + and type(reusable) is bool + and type(lease) is _InPlaceLease + and lease._view_returned + and not lease._marked + and not lease._released + and lease._kind == "write" + and lease._lender() is self + ): + return [self._return_write(lease, None, reusable)] + lease, masks, reusable = self._accepted_return(entry, set()) + return [self._return_write(lease, masks, reusable)] + seen: set[int] = set() + accepted = [self._accepted_return(entry, seen) for entry in entries] + results: list[ResumeBounds | None] = [] + error = None + for lease, masks, reusable in accepted: + try: + results.append(self._return_write(lease, masks, reusable)) + except BaseException as raised: # noqa: BLE001 - raised once every lease is returned + error = raised if error is None else error + results.append(None) + if error is not None: + raise error + return results + + def progress(self) -> None: + """See ``Lender.progress``.""" + if self._leases_holding_slots or self._waiting_for_slots: + self._progress() + else: + # Nothing holds or waits for staging slots: the round only counts. + self._progress_round += 1 + + def _accepted_return( + self, entry: object, seen: set[int] + ) -> tuple[_StagingLease | _InPlaceLease, list[np.ndarray] | None, bool]: + """``entry`` checked, with its masks as kept (``None``: every row) and its lease added to + ``seen``; raises what ``return_writes`` raises for it, changing nothing else.""" + try: + lease, arrived, reusable = entry + except (TypeError, ValueError): + raise TypeError( + f"an entry is (lease, arrived, reusable), got {type(entry).__name__}" + ) from None + if type(reusable) is not bool: + raise TypeError(f"reusable must be a bool, got {type(reusable).__name__}") + if type(lease) not in (_InPlaceLease, _StagingLease) or lease._lender() is not self: + raise ValueError(f"the lease was not lent by this lender: {lease!r}") + if lease._kind != "write": + raise RuntimeError("return_writes is for write leases") + if lease._marked or lease._released or id(lease) in seen: + raise RuntimeError(f"the lease was returned already: {lease!r}") + if not lease._view_returned: + raise RuntimeError("return_writes before poll() returned the lease's view") + masks = None if arrived is ALL_ARRIVED else lease._accepted_masks(arrived) + seen.add(id(lease)) + return lease, masks, reusable + + def _return_write( + self, + lease: _StagingLease | _InPlaceLease, + masks: list[np.ndarray] | None, + reusable: bool, + ) -> ResumeBounds | None: + """Count a checked write's arrived rows (``None``: all) into the cache it was lent from, + keeping an unreusable cache out of reuse before the unlock's retried commit; return the + lease even if counting raised; give the bounds, ``None`` while unsettled or where the cache + is no longer the request's.""" + if not reusable or (masks is not None and self._margin_missing(lease, masks)): + self._keep_out_of_reuse(lease) + manager = self._serving_manager() + if manager is not None and lease._fetch is not None: + _manager.drop_held_commit(manager, lease._kv, lease._fetch.start) + if masks is None and type(lease) is _InPlaceLease: + bounds = self._return_whole(lease) + if bounds is not _NOT_WHOLE: + return bounds + lease._marked = True + request_id = lease._request_id + kv = lease._kv + try: + whole = self._count_arrivals(lease, masks) + finally: + if type(lease) is _InPlaceLease and not lease._released: + # ``release`` of a write whose view was returned and that is returned with its + # arrivals: the lease's lock ends. + lease._finish_returned_write() + self._writes_at_end.pop(id(kv), None) + if lease._lock is not None: + self._locks.unlock(lease._lock) + else: + lease.release() + if whole and self._stop_reason is None: + manager = self._serving_manager() + if manager is None or _manager.kv_of(manager, request_id) is not kv: + return None + state, blocks = _manager.cache_extent(kv) + ledger = self._ledger + fetch = lease._fetch + spans = lease._block_entries.spans + bounds = ledger._whole_delivery_bounds( + fetch.prompt, fetch.runs, request_id, manager, kv, state, spans, blocks + ) + if bounds is None: + bounds = ledger._resume_bounds( + fetch.prompt, fetch.runs, request_id, manager, kv, state + ) + return bounds + return self._bounds_after_return(lease) + + def _keep_out_of_reuse(self, lease: _StagingLease | _InPlaceLease) -> None: + """Keep the lease's cache from committing anything more for reuse until it closes, if it is + still open: the request's cache, or one a lock keeps open.""" + kv = lease._kv + manager = self._serving_manager() + if manager is None or kv is None: + return + if kv is _manager.kv_of(manager, lease._request_id) or id(kv) in self._locks.locked_caches: + self._locks.keep_out_of_reuse(kv) + + def _margin_missing( + self, lease: _StagingLease | _InPlaceLease, masks: list[np.ndarray] + ) -> bool: + """Whether a lent window-margin row did not arrive: a row of a group whose kept window is + wider than its read window, which the read window passed at the fetch's end.""" + if not self._margin_groups: + return False + fetch = lease._fetch + manager = self._manager_ref() + if fetch is None or fetch.end is None or manager is None: + return False + entries = lease._block_entries + for lg, ordinals, mask in zip(entries.layer_group_indices, entries.block_ordinals, masks): + if lg not in self._margin_groups or not len(ordinals): + continue + beg, end = _ranges._read_stale_range(manager, self._layout, lg, fetch.end) + if np.any((ordinals >= beg) & (ordinals < end) & ~mask): + return True + return False + + def _return_whole(self, lease: _InPlaceLease) -> ResumeBounds | None | object: + """``_return_write`` of a checked in-place write returned whole as its cache's first + delivery, equal to the rules while the fast lend holds every page it lent + (``_lent_pages_held``), its fetch is the latest and nothing was delivered before; re-reads + the cache only after a retried commit; else ``_NOT_WHOLE``, unchanged.""" + spans = lease._block_entries.spans + if spans is None or self._shut_down or self._stop_reason is not None: + return _NOT_WHOLE + manager = self._manager_ref() + request_id = lease._request_id + kv = lease._kv + if manager is None or not self._lends_fast(): + return _NOT_WHOLE + ledger = self._ledger + try: + if _manager.kv_of(manager, request_id) is not kv: + return _NOT_WHOLE + state, blocks = extent = _manager.cache_extent(kv) + if not self._lent_pages_held(lease, kv, extent) or not ledger._defer_first_delivery( + request_id, lease._fetch, kv, spans + ): + return _NOT_WHOLE + except Exception: # noqa: BLE001 - the rules then return the lease and raise the error + return _NOT_WHOLE + lease._finish_returned_write() + writes_at_end = self._writes_at_end + if writes_at_end: + writes_at_end.pop(id(kv), None) + lock = lease._lock + retried = lock is not None and self._locks.unlock(lock) + if self._stop_reason is not None: + return self._bounds_after_return(lease) + if retried: + manager = self._serving_manager() + if manager is None or _manager.kv_of(manager, request_id) is not kv: + return None + state, blocks = _manager.cache_extent(kv) + fetch = lease._fetch + bounds = ledger._whole_delivery_bounds( + fetch.prompt, fetch.runs, request_id, manager, kv, state, spans, blocks + ) + if bounds is None: + bounds = ledger._resume_bounds(fetch.prompt, fetch.runs, request_id, manager, kv, state) + return bounds + + def _count_arrivals( + self, lease: _StagingLease | _InPlaceLease, masks: list[np.ndarray] | None + ) -> bool: + """Count a returning write lease's arrived rows, every row for ``masks`` ``None``; whether + they are the cache's first delivery, every row arriving into pages the lease still holds.""" + manager = None if self._shut_down else self._manager_ref() + if manager is None or self._stop_reason is not None: + return False + if type(lease) is _StagingLease: + self._apply_marks(lease, self._every_row_mask(lease) if masks is None else masks) + return False + request_id = lease._request_id + kv = _manager.kv_of(manager, request_id) + spans = lease._block_entries.spans + held = spans is not None and self._lends_fast() and self._lent_pages_held(lease, kv) + if ( + held + and (masks is None or all(mask.all() for mask in masks)) + and self._ledger._defer_first_delivery(request_id, lease._fetch, kv, spans) + ): + return True + self._apply_in_place_marks( + lease, self._every_row_mask(lease) if masks is None else masks, lock_held=held + ) + return False + + def _lent_pages_held( + self, + lease: _InPlaceLease, + kv: _KVCache | None, + extent: tuple[_manager.CacheState, int] | None = None, + ) -> bool: + """Whether every row a fast write lent still holds its page, kept by the lock a returned + view holds: same active cache, no shrink reported since the lend, its blocks kept and a + window's history at most the fetch's end. A locked page changes only at a shrink or free, so + ``_locked_block_masks`` is then true for every row. ``extent``: ``kv``'s, if read.""" + if kv is not lease._kv: + return False + epochs = self._shrink_epochs + if (epochs.get(lease._request_id, 0) if epochs else 0) != lease._shrinks: + return False + state, blocks = _manager.cache_extent(kv) if extent is None else extent + if not state.active or blocks < lease._block_end: + return False + return not self._any_window or state.history <= lease._fetch.end + + @staticmethod + def _every_row_mask(lease: _StagingLease | _InPlaceLease) -> list[np.ndarray]: + """A true mask for every row of the lease's view, in group order.""" + return [np.ones(len(o), dtype=bool) for o in lease._block_entries.block_ordinals] + + @staticmethod + def _no_row_mask(lease: _StagingLease | _InPlaceLease) -> list[np.ndarray]: + """A false mask for every row of the lease's view, in group order.""" + return [np.zeros(len(o), dtype=bool) for o in lease._block_entries.block_ordinals] + + def _bounds_after_return(self, lease: _StagingLease | _InPlaceLease) -> ResumeBounds | None: + """``poll_resume_bounds`` of the write lease's cache once it returned, ``None`` where that + cache is no longer the request's, the request has none or the manager shut down.""" + request_id = lease._request_id + manager = self._serving_manager() + kv = None if manager is None else _manager.kv_of(manager, request_id) + if kv is None or kv is not lease._kv: + return None + state = _manager.cache_state(kv) + if self._stop_reason is not None: + return ResumeBounds(max_position=0, min_position=max(state.history, 1)) + fetch = lease._fetch + return self._ledger._resume_bounds(fetch.prompt, fetch.runs, request_id, manager, kv, state) def _on_free(self, request_id: int, kv_cache: _KVCache) -> None: """Manager hook, after the request's cache left the map: its waiting leases fail and its fetch records go. An error in that work, which it logs, stops the lender for good.""" try: - if self._serving_manager() is None: + if self._shut_down: return + request_id = int(request_id) + if self._writes_at_end: + self._lose_end_resume(kv_cache) # Granted leases go on: a read's copy is queued already, and a write's marks copy # nothing into pages other than the ones lent. - self._fail_waiting(int(request_id), "the request was freed") + if self._waiting_for_slots: + self._fail_waiting(request_id, "the request was freed") self._ledger._drop_records(request_id) - self._progress() + self._ordinal_ranges.pop(request_id, None) + self._shrink_epochs.pop(request_id, None) + self._state_unverified.discard(request_id) + parked = self._parked.pop(request_id, None) + if parked is not None: + # The cache writes into its row until it closes, after ``_defer_close`` or at its + # last unlock; rows of closed caches go back to use here, off a slot release's path. + if self._parked_closing: + self._reuse_closed_rows() + self._parked_closing.append(parked) except Exception: - # Logged, not raised into the manager's free, which then closes the cache as usual. The - # error can leave the records disagreeing with the slots where no caller sees it, so the - # lender stops; the reason is set first, by an assignment that cannot raise. + # Logged, not raised into the free, which then closes the cache; records may disagree + # with slots unseen, so the lender stops, setting its reason first, which cannot raise. if self._stop_reason is None: self._stop_reason = _REASON_STOPPED + if self._writes_at_end: + self._lose_end_resumes() logger.error(f"KV cache lender: freeing request {request_id}: {traceback.format_exc()}") + def _on_resume(self, request_id: int) -> None: + """Manager hook after a resume, which may move recurrent state the manager records only at + the next preparation: the cache answers for the slot until the record matches it.""" + record = self._state_record + if record is not None and int(request_id) in record: + self._state_unverified.add(int(request_id)) + + def _parking_row(self, request_id: int, kv_cache: _KVCache) -> tuple[memoryview, ...] | None: + """Manager hook at an early index-slot release: a lender-owned row's buffers, one per layer + group, for the cache's beam-0 page indices from then on and for lends' reads; None, keeping + nothing, without a page-index buffer.""" + if self._index_words is None or self._shut_down: + return None + request_id = int(request_id) + previous = self._parked.get(request_id) + if previous is not None: + # A release of a slot released already: the cache leaves the earlier row at this bind. + self._parked_closing.append(previous) + row = self._unbound_row() + row.kv = kv_cache + self._parked[request_id] = row + return row.buffers + + def _parked_row(self, request_id: int) -> tuple[memoryview, ...] | None: + """Manager hook, right after a cache resumed whose index slot was released early: the + buffers ``_parking_row`` gave for the request, to bind again, or None.""" + row = self._parked.get(int(request_id)) + return None if row is None else row.buffers + + def _unbound_row(self) -> _ParkedRow: + """A row no open cache writes into, every page index in it bad: a row whose cache has + closed, or a new one.""" + if not self._parked_free: + self._reuse_closed_rows() + if self._parked_free: + row = self._parked_free.pop() + # A row a failed bind leaves unbound then gives no page. + row.bytes[:] = self._bad_row + return row + row = _ParkedRow(len(self._index_rows), self._index_strides[2]) + self._bad_row = b"\xff" * row.pages.nbytes + self._parked_made.append(row) + return row + + def _reuse_closed_rows(self) -> None: + """Put the rows of freed requests whose caches have closed back to use.""" + still_open = [] + for row in self._parked_closing: + if _manager.is_closed(row.kv): + row.kv = None + self._parked_free.append(row) + else: + still_open.append(row) + self._parked_closing = still_open + + def _page_row( + self, manager: KVCacheManagerV2, request_id: int, kv: _KVCache + ) -> tuple[memoryview, int, int, int] | None: + """Where active ``kv`` keeps its beam-0 page indices for lends to read: the request's row of + the page-index buffer while it holds one, else the row parked for ``kv`` at its early + release; None without either.""" + words = self._index_words + if words is None: + return None + row = _manager.index_row(manager, request_id) + if row is not None: + # The key half of the slot's row. + group_words, row_words, row_blocks = self._index_strides + return words, row * row_words, group_words, row_blocks + parked = self._parked.get(request_id) + return None if parked is None or parked.kv is not kv else parked.source + + def _state_slot(self, kv: _KVCache, request_id: int | None, lg: int) -> int: + """The request's slot of recurrent group ``lg``: the manager's record, which the cache + confirms once after each resume; the cache's own answer where the manager records none for + the request or the group.""" + record = self._state_record + if record is None or lg != self._state_group or request_id is None: + return _manager.state_slot(kv, lg) + if request_id not in self._state_unverified: + slot = record.get(request_id) + if slot is not None: + return slot + return _manager.state_slot(kv, lg) + slot = _manager.state_slot(kv, lg) + if record.get(request_id) == slot: + self._state_unverified.discard(request_id) + return slot + def _on_shrink(self, request_id: int, kv_cache: _KVCache) -> None: - """Manager hook, right after the request's cache may have shrunk in place: delivered entries - past its blocks stop counting, since the shrink freed their pages. Logs its own errors.""" + """Manager hook right after a possible in-place shrink: delivered entries past its blocks + stop counting, as the shrink freed their pages, and the request's fast writes stop trusting + the pages they lent. Logs its own errors.""" try: + if self._writes_at_end: + self._lose_end_resume(kv_cache) + epochs = self._shrink_epochs + epochs[int(request_id)] = epochs.get(int(request_id), 0) + 1 delivered = self._ledger._deliveries_of(int(request_id)) if ( self._serving_manager() is not None @@ -663,29 +1555,53 @@ def _on_shrink(self, request_id: int, kv_cache: _KVCache) -> None: self._ledger._invalidate_delivered_blocks_past_capacity(delivered, kv_cache) except Exception: # Logged, not raised into the manager's resize. All the request's delivered entries go: - # the cache may grow back before resume_bounds finds the freed ones past its blocks. + # the cache may grow back before poll_resume_bounds finds the freed ones past its blocks. self._ledger._drop_deliveries(int(request_id)) logger.error( f"KV cache lender: shrink of request {request_id}: {traceback.format_exc()}" ) + def _before_reset(self) -> None: + """Manager hook before a reuse reset, which the runtime refuses with a cache open: + ``RuntimeError``, changing nothing, while a lock keeps a freed request's cache open, until + its last unlock closes it.""" + if self._locks.holds_open_caches(): + raise RuntimeError( + "KV cache lender: a lock keeps a freed request's KV cache open until its lease is " + "released, and the reuse state cannot be reset before" + ) + def _on_reset(self) -> None: - """Manager hook after a reuse reset, with every cache closed as an in-place weight update - ends: object_keys stop meaning the bytes computed, so none is lent by name again; waiting - leases - and fetches went at their requests' free; granted publishes finish. Never raises.""" - # TODO: after the manager resets its reuse state the lender lends nothing by name, since - # ``scope`` is fixed at the attach and nothing identifies the weights it computes with from - # then on. + """Stop NAMED lending, as a fixed scope cannot name updated weights; caches are closed and + frees removed waiting leases and fetches; granted publishes finish. Never raises.""" self._reuse_was_reset = True + self._ordinal_ranges.clear() + if self._writes_at_end: + self._lose_end_resumes() + + def _lose_end_resume(self, kv_cache: _KVCache) -> None: + """The open fast write lease of ``kv_cache``, if any, no longer resumes at its end.""" + lease = self._writes_at_end.pop(id(kv_cache), None) + if lease is not None: + lease._at_end = False + + def _lose_end_resumes(self) -> None: + """No open fast write lease resumes at its end any more.""" + for lease in self._writes_at_end.values(): + lease._at_end = False + self._writes_at_end.clear() def _on_shutdown(self) -> None: - """Manager hook, first in its shutdown, acting once: waits for the copies whose completion - it recorded, fails waiting leases, and frees the staging memory unless a lease or a hold is - open, a slot was lost to a failed copy or the lender stopped. An error in that work, which - it logs, stops the lender for good.""" + """Once: wait for recorded copies, fail waiters and stop lending, logging errors; free + staging unless open leases or holds, lost slots or a stop keep it. Device-region holds keep + the runtime and its pools until exit.""" if self._shut_down: return + self._ordinal_ranges.clear() + if self._writes_at_end: + self._lose_end_resumes() + if self._device_regions_of_pool_group and self._open_holds: + _retain_until_exit(_manager.runtime(self._manager_ref())) try: # With page-locked staging and the fresh-page fill off, the lender's only host # wait: no more work comes on the stream. @@ -702,36 +1618,86 @@ def _on_shutdown(self) -> None: ]: self._fail(lease, lease._first_poll_failure) self._shut_down = True - if not self._must_retain_memory(): + memory = self._memory_ref() # None for a lender that stages nothing + if memory is not None and not self._must_retain_memory(): # Every copy on it has completed above, and no lease or hold reaches it any more. - memory = self._memory_ref() _stop_retaining(memory) memory.free() - else: + elif memory is not None: logger.warning( - f"KV cache lender: keeping {sum(p.nbytes for p in self._regions) >> 20} MiB of " - f"staging memory until exit: {len(self._open_leases)} leases open, " - f"{len(self._open_holds)} staging holds open, " + f"KV cache lender: keeping {self._staging_bytes() >> 20} MiB of staging memory " + f"until exit: {len(self._open_leases)} leases open, " + f"{len(self._open_holds)} region holds open, " f"{len(self._quarantined)} slot runs lost to failed copies" f"{'' if self._stop_reason is None else ', and ' + self._stop_reason}" ) except Exception: - # Logged, not raised into the manager's shutdown; the staging memory then stays in - # the keep list until exit. As in the free hook, the records may disagree with the - # slots, so the lender stops before the log line, keeping any earlier stop's reason. + # Logged, not raised into the shutdown; staging stays until exit. As in ``_on_free`` + # records may disagree with slots: the lender stops first, keeping an earlier reason. self._shut_down = True if self._stop_reason is None: self._stop_reason = _REASON_STOPPED_AT_SHUTDOWN + if self._writes_at_end: + self._lose_end_resumes() logger.error(f"KV cache lender: shutting down: {traceback.format_exc()}") def _on_caches_closed(self) -> None: - """Manager hook, once the runtime's shutdown has returned, after which no cache writes into - the page-index buffer, so the lender lets it go. A shutdown that raises skips this call, and - the buffer stays until a retried one returns, or until exit.""" + """Manager hook once the runtime's shutdown returned and no cache writes the page-index + buffer: let it go. A raising shutdown skips this, keeping it until a retry or exit.""" if self._index_buffer is None: return self._stop_retaining_index_buffer() + def _page_lock_state(self) -> PageLocks: + """Live lock registry, whose ``holding`` counts the unreleased locks holding a block.""" + return self._locks + + def _is_locked(self, kv_cache: _KVCache | None, begin: int = 0, end: int | None = None) -> bool: + """Manager hook before a suspension or commit: whether a lock holds a block of ``kv_cache`` + in ``[begin, end)``, or for ``None`` of any cache, freed ones kept open included.""" + if kv_cache is None: + return self._locks.holds_any() + return self._locks.is_locked(kv_cache, begin, end) + + def _defer_close(self, kv_cache: _KVCache) -> bool: + """Manager hook at a free or shutdown, in place of a close: whether locks keep ``kv_cache`` + open until its last unlock. One closed at once frees the row parked for it at its free, as + no release takes a row before the close.""" + deferred = self._locks.defer_close(kv_cache) + closing = self._parked_closing + if not deferred and closing and closing[-1].kv is kv_cache: + row = closing.pop() + row.kv = None + self._parked_free.append(row) + return deferred + + def _keeps_pools(self) -> bool: + """Manager hook, at its shutdown: whether the runtime, whose shutdown frees the pools, must + keep them: a cache ``_defer_close`` kept is still open, or a hold covers device regions.""" + return self._locks.holds_open_caches() or bool( + self._device_regions_of_pool_group and self._open_holds + ) + + def _history_cap(self, kv_cache: _KVCache) -> int | None: + """Manager hook before a history move: the highest history leaving no locked block behind a + window, or ``None``; the manager stops there until the unlock, keeping the lock's page.""" + if not self._any_window: + return None + manager = self._manager_ref() + if manager is None: + return None + cap = None + layout = self._layout + for lg, lowest in self._locks.lowest_locked_blocks(kv_cache, layout.sink_blocks).items(): + if layout.windows[lg] is not None: + last = _ranges._last_history_keeping(manager, layout, lg, lowest) + cap = last if cap is None else min(cap, last) + return cap + + def _staging_bytes(self) -> int: + """The bytes of the staging regions.""" + return sum(p.nbytes for p in self._regions if p.memory_kind is MemoryKind.HOST) + # One rule per method, so a test subclass that breaks exactly one rule overrides one method. def _progress(self) -> None: @@ -744,8 +1710,9 @@ def _progress(self) -> None: self._fail(lease, self._stop_reason) return self._progress_round += 1 - self._recycle() - self._grant_waiting() + if self._leases_holding_slots or self._waiting_for_slots: + self._recycle() + self._grant_waiting() def _copy_done(self, copy: _QueuedCopy | None) -> bool: """``copy`` has completed (or there is none), its event asked at most once this round.""" @@ -759,16 +1726,13 @@ def _recyclable(self, lease: _StagingLease) -> bool: """The lease's slots may return: no backend access possible and no copy on them pending. The copy is asked last, so a lease still lent costs no event query.""" # A failed lease was never ready, so no backend has seen its slots. - if lease._holder_state() is not _HolderState.FAILED: - if not lease._released or lease._marks_due(): - return False - # TODO: every lender call asks each released lease's pending copy again, so N calls while P - # copies pend cost N*P event queries. + if lease._failure is None and (not lease._released or lease._marks_due()): + return False return self._copy_done(lease._copy) - def _locked_block_masks(self, kv: _KVCache | None, lease: _StagingLease) -> list[np.ndarray]: - """Per layer group, entries still in the same active cache's window and still locking the - GPU page lent; a page only held may sit on another tier under the same number.""" + def _locked_block_masks(self, kv: _KVCache | None, lease: _LeaseBase) -> list[np.ndarray]: + """Per layer group, entries still in the same active cache's window locking the GPU page + lent (a held page may sit on another tier), and state rows still its state page.""" entries = lease._block_entries manager = self._manager_ref() state = _manager.cache_state(kv) if kv is not None and kv is lease._kv else None @@ -776,15 +1740,24 @@ def _locked_block_masks(self, kv: _KVCache | None, lease: _StagingLease) -> list for lg, block_ordinals, slots in zip( entries.layer_group_indices, entries.block_ordinals, entries.device_pages ): - same = np.zeros(len(block_ordinals), dtype=bool) - if state is not None and state.active: + if state is None or not state.active: + same = np.zeros(len(block_ordinals), dtype=bool) + elif lg in self._recurrent_groups: + same = slots == self._state_slot(kv, lease._request_id, lg) + else: pages = _manager.locked_pages(kv, lg) - inside = block_ordinals < len(pages) - same[inside] = pages[block_ordinals[inside]] == slots[inside] - stale_beg, stale_end = _ranges._stale_range( - manager, self._layout, lg, state.history - ) - same &= (block_ordinals < stale_beg) | (block_ordinals >= stale_end) + try: + same = pages[block_ordinals] == slots + except IndexError: + # Blocks past the cache's end lock no page. + same = np.zeros(len(block_ordinals), dtype=bool) + inside = block_ordinals < len(pages) + same[inside] = pages[block_ordinals[inside]] == slots[inside] + if self._layout.windows[lg] is not None: + stale_beg, stale_end = _ranges._stale_range( + manager, self._layout, lg, state.history + ) + same &= (block_ordinals < stale_beg) | (block_ordinals >= stale_end) masks.append(same) return masks @@ -811,14 +1784,62 @@ def _bidirectional_split_refusal(self, request: LlmRequest, end: int) -> str | N return None def _scratch_reuse_on(self, kv: _KVCache) -> bool: - """A write target with SWA scratch reuse on, windowed or not: each capacity change keeps its - history within the scratch rewind of the old capacity, which a fetch's grow breaks, and a - window's next chunk overwrites scratch slots.""" + """Whether ``kv`` has SWA scratch reuse on (``_manager.has_scratch_reuse``), which a fetch's + grow breaks.""" return _manager.has_scratch_reuse(kv) - # TODO: the executor gives context outputs only for the positions a request computes, prompt - # logprobs pair the logits with the prompt from its second token wherever they start, and a - # rewind or a recompute pause keeps those held, so a request returning them takes no fetch. + def _committed_entries(self, kv: _KVCache, entries: _LentBlockEntries) -> _LentBlockEntries: + """``entries`` without the blocks past each layer group's committed whole blocks: a group + the prefix claim left out holds none of the claimed blocks a name covers.""" + tpb = int(self._layout.tokens_per_block) + committed = _manager.committed_by_group(kv, self._layout.num_layer_groups) + out = _LentBlockEntries([], [], []) + for lg, ordinals, slots in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + keep = (ordinals < committed[lg] // tpb) & self._named_groups[lg] + out.layer_group_indices.append(lg) + whole = ordinals.ndim == 1 and slots.shape == ordinals.shape and bool(keep.all()) + out.block_ordinals.append(ordinals if whole else _immutable(ordinals[keep])) + out.device_pages.append(slots if whole else _immutable(slots[keep])) + return out + + def _without_short_windows_mid_context( + self, + manager: KVCacheManagerV2, + request: LlmRequest, + state: _manager.CacheState, + entries: _LentBlockEntries, + ) -> _LentBlockEntries: + """Until this worker's context is computed, omit groups whose window is shorter than it. + Their locks would block later chunks from advancing history; lend them after context.""" + context_end = self._mid_context_end(manager, request, state) + if context_end is None: + return entries + out = _LentBlockEntries([], [], []) + for lg, block_ordinals, slots in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + window = self._layout.windows[lg] + keep = window is None or int(window) >= context_end + out.layer_group_indices.append(lg) + out.block_ordinals.append(block_ordinals if keep else block_ordinals[:0]) + out.device_pages.append(slots if keep else slots[:0]) + return out + + def _mid_context_end( + self, manager: KVCacheManagerV2, request: LlmRequest, state: _manager.CacheState + ) -> int | None: + """The pending context end while a windowed cache is still computing it, if any.""" + if not self._any_window: + return None + context_end = _manager.context_end(request, manager) + if state.history >= context_end or _manager.context_computed(request) >= context_end: + return None + return context_end + + # TODO: context outputs cover only computed positions, prompt logprobs pair logits from the + # prompt's second token, and rewinds or recompute pauses keep them: such requests take no fetch. def _returns_context_outputs(self, request: LlmRequest) -> bool: """Whether the request returns context logits or asks for additional model outputs, which a fetch leaves without the fetched positions, or after a context step gapped or repeated.""" @@ -827,9 +1848,8 @@ def _returns_context_outputs(self, request: LlmRequest) -> bool: def _history_refusal( self, manager: KVCacheManagerV2, history: int, position: int, end: int ) -> str | None: - """Why the history the fetch leaves, ``position``, admits no resume at the positions - resume_bounds would count, or ``None``: in a cache without a window whose floor follows the - history, the history stands past ``end``.""" + """Why the history ``position`` the fetch leaves admits no resume the bounds would count, or + ``None``: without a window and with the floor following it, a history past ``end``.""" if self._any_window: return None # the fetch moves the history to ``end`` if not _ranges._floor_follows_history(manager): @@ -841,14 +1861,12 @@ def _history_refusal( def _unwritten_pages_refusal( self, manager: KVCacheManagerV2, request_id: int, kv: _KVCache, position: int ) -> str | None: - """Why moving the history to ``position`` would have the commit store bytes the request - never wrote, or ``None``: a window leaving behind a block whose page holds tokens past those - the request computed keeps that page for the commit (a partial match's copy, a page grown - early, a row an earlier lease missed), unless a fetch into the cache delivered its row.""" + """Refuse a history move that would commit unwritten window pages (partial prefixes, early + growth, incomplete fetches), else ``None``; a delivered row supplies the bytes.""" if not _manager.keeps_stale_pages(manager): return None computed = self._ledger._computed_tokens(request_id, kv) - first = computed // int(self._layout.tokens_per_block) # the first block past them + tpb = int(self._layout.tokens_per_block) delivered = self._ledger._deliveries_of(request_id) fetched = ( delivered.delivered_block_masks @@ -857,12 +1875,11 @@ def _unwritten_pages_refusal( ) for lg in range(self._layout.num_layer_groups): stale_beg, stale_end = _ranges._stale_range(manager, self._layout, lg, position) - beg = max(first, stale_beg) + beg = max(computed[lg] // tpb, stale_beg) # from the first block past them if beg >= stale_end: continue unwritten = _manager.locked_pages(kv, lg)[beg:stale_end] >= 0 if fetched is not None: - # A delivered row's page holds the bytes the fetch copied in. got = fetched[lg][beg : beg + len(unwritten)] unwritten[: len(got)] &= ~got blocks = (beg + np.nonzero(unwritten)[0]).tolist() @@ -871,7 +1888,7 @@ def _unwritten_pages_refusal( # fails where the commit could store no page for those blocks instead. return ( f"layer group {lg} leaves blocks {blocks[:8]} behind its window at {position} " - f"tokens; past the {computed} tokens the request computed, their pages hold " + f"tokens; past the {computed[lg]} tokens the request computed, their pages hold " "bytes it never wrote and no fetch delivered, which the commit would store" ) return None @@ -901,36 +1918,155 @@ def _slots_quarantined(self) -> bool: """Whether slots are out of use for good: runs ``_quarantine`` took, or those of a holding lease whose copy has no record of its completion, which ``_recycle`` never returns.""" return bool(self._quarantined) or any( - lease._copy_state() is _CopyState.WITHOUT_RECORD for lease in self._leases_holding_slots + lease._copy is _COPY_WITHOUT_RECORD for lease in self._leases_holding_slots ) def _retain_index_buffer(self, manager: KVCacheManagerV2) -> None: - """Keep the manager's page-index buffer until the runtime's shutdown has returned: a cache a - lease or a record holds writes its page indices there as it closes, also once the manager is - gone.""" + """Keep the manager's page-index buffer until the runtime's shutdown returned: a cache a + lease or record holds writes it as it closes, even once the manager is gone.""" self._index_buffer = _manager.index_buffer(manager) _retain_until_exit(self._index_buffer) + _retain_until_exit(self._parked_made) + self._read_index_rows(np.asarray(self._index_buffer)) def _stop_retaining_index_buffer(self) -> None: """Once the runtime's shutdown has returned: no cache writes into the page-index buffer any more.""" + self._read_index_rows(None) _stop_retaining(self._index_buffer) self._index_buffer = None + # Nor into a parked row. + _stop_retaining(self._parked_made) + self._parked.clear() + self._parked_closing.clear() + self._parked_free.clear() + self._parked_made = [] + + def _read_index_rows(self, rows: np.ndarray | None) -> None: + """Fast lends read pages from ``rows``, a contiguous ``int32`` ``[group, row, K/V, block]`` + buffer kept as is, and parked rows; from the caches for ``None`` or other arrays.""" + if rows is None or rows.dtype != np.int32 or rows.ndim != 4 or not rows.flags.c_contiguous: + self._index_rows = self._index_words = None + return + _, count, halves, blocks = rows.shape + self._index_rows = rows + self._index_words = memoryview(rows).cast("B").cast("i") + self._index_strides = (count * halves * blocks, halves * blocks, blocks) - def _on_hold_released(self, hold: _StagingHold) -> None: + def _on_hold_released(self, hold: _RegionHold) -> None: """The hold's holder let go; ``KeyError`` for a hold not open, which its guard prevents.""" self._open_holds.remove(hold) - def _check_fits(self, counts: Mapping[int, int]) -> None: - """``ValueError`` if a lease needing ``counts`` entries per pool group can never be - granted.""" + def _fits_refusal(self, counts: Mapping[int, int]) -> str | None: + """Why a lease needing ``counts`` entries per pool group can never be granted here, or + ``None``: this stage's staging holds fewer slots.""" try: self._slot_queue.check_fits_capacity(counts) except ValueError as error: - raise ValueError( + return ( f"the range needs more staging slots than there are ({error}); ranges of at most " - "tokens_per_fetch tokens always fit" - ) from None + "guaranteed_fetch_tokens tokens always fit" + ) + return None + + def _read_end_refusal( + self, + manager: KVCacheManagerV2, + request: LlmRequest, + state: _manager.CacheState, + end: int, + named: bool, + ) -> str | None: + """Why this stage cannot read through ``end`` now, or ``None``.""" + if named: + if end > state.committed: + return f"the range ends at {end}, past the {state.committed} committed tokens" + return None + # A window held for a lock keeps the history below what finished context steps computed. + computed = _manager.context_computed(request) + if end > state.history and computed < end: + return f"the range ends at {end}, past the {state.history} tokens the cache holds" + # A later context step writes the rest of a partial block; generation writes past the + # history alone. + if end % int(self._layout.tokens_per_block): + context_end = _manager.context_end(request, manager) + if context_end > state.history and computed < context_end: + return ( + f"the range ends at {end}: until this worker computes its context to " + f"{context_end} tokens, a read of every row ends at a whole block" + ) + if not self._recurrent_groups: + return None + # The state stands at the cache's history only from the context's last step until the + # request's first generation step, whose speculation leaves it at a checkpoint. + prompt = _manager.prompt_length(request, manager) + if not end == state.history == prompt: + return ( + f"the range ends at {end}: with recurrent state a lease ends where the state " + f"stands, at the prompt's {prompt} tokens once its context is computed and " + f"before a generation step, the cache holding {state.history}" + ) + return None + + def _unfetched_rows_refusal( + self, + request_id: int, + kv: _KVCache, + state: _manager.CacheState, + entries: _LentBlockEntries, + end: int, + ) -> str | None: + """Why a read in place would lend rows below the history nothing wrote, or ``None``: rows a + fetch lent or grew over without delivering, which the request had not computed.""" + found = self._ledger._unheld_block( + request_id, kv, entries, min(end, state.history), state.history + ) + if found is None: + return None + lg, block = found + if lg in self._recurrent_groups: + return f"recurrent state has no complete fetched snapshot at {end} tokens" + return ( + f"layer group {lg}: block {block} holds rows below the history of {state.history} " + "tokens that a fetch lent or grew the cache over and did not deliver" + ) + + def _write_range_refusal( + self, + manager: KVCacheManagerV2, + request: LlmRequest, + state: _manager.CacheState, + start: int, + end: int, + prompt: int, + ) -> str | None: + """Why this stage cannot fetch ``[start, end)`` now, or ``None``; changes nothing.""" + tpb = int(self._layout.tokens_per_block) + if end % tpb and end not in (prompt, state.history): + return ( + f"the range ends at {end}: a fetch of every row ends at a whole block, at the " + f"request's {prompt} prompt tokens or at the {state.history} tokens of history " + "its cache declares" + ) + # A KV cache connector loads past the prefix served to the request, never into it. + served = _manager.served_prefix_end(request) if _manager.connector_loads(manager) else 0 + if start < served: + return ( + f"the range starts at {start}, inside the {served} tokens served to the request " + "from outside the manager, where a KV cache connector loads" + ) + # TODO: lendable state is kept only at the latest position: no fetch resumes mid-prompt. + if self._recurrent_groups and end != prompt: + return ( + f"the range ends at {end}: with recurrent state a fetch ends where the state " + f"stands, at the request's {prompt} prompt tokens" + ) + if self._any_window and end < state.history: + return ( + f"the range ends at {end}, below the history of {state.history} tokens its " + "windows keep" + ) + return None def _free_slots(self, group: int) -> int: """Free slots of a pool group; for tests.""" @@ -948,14 +2084,56 @@ def _serving_manager(self) -> KVCacheManagerV2 | None: return None return self._manager_ref() - def _whole_blocks(self, start: int, end: int) -> tuple[int, int]: + def _lending_manager(self) -> KVCacheManagerV2: + """The manager a lend serves; ``RuntimeError`` once it shut down or is gone, when the + lender lends no more.""" + manager = self._serving_manager() + if manager is None: + raise RuntimeError(f"the KV cache lender lends no more: {_REASON_SHUT_DOWN}") + return manager + + def _check_not_stopped(self) -> None: + """``RuntimeError`` once the lender stopped for good: it lends no more.""" + if self._stop_reason is not None: + raise RuntimeError(f"the KV cache lender lends no more: {self._stop_reason}") + + def _check_memory_and_rows(self, memory_kind: MemoryKind, rows: RowSelection) -> None: + """``TypeError`` for arguments of other types; ``ValueError`` for named rows where nothing + is named, host rows other than named ones, or host memory where nothing is staged.""" + if not isinstance(memory_kind, MemoryKind): + raise TypeError(f"memory_kind must be a MemoryKind, got {type(memory_kind).__name__}") + if not isinstance(rows, RowSelection): + raise TypeError(f"rows must be a RowSelection, got {type(rows).__name__}") + if rows is RowSelection.NAMED and self._no_names is not None: + raise ValueError( + f"{rows} is not lent: this lender lends nothing by name, as {self._no_names}" + ) + # TODO: staging relays named whole blocks only, so host memory lends no other rows. + if memory_kind is MemoryKind.HOST and rows is not RowSelection.NAMED: + raise ValueError( + f"{rows} is not lent in {memory_kind}: staging relays named whole blocks" + ) + # TODO: the manager's own host-tier pages cannot be lent; host lending copies into staging. + if memory_kind is MemoryKind.HOST and self._no_staging is not None: + raise ValueError( + f"{memory_kind} is not lent: this lender stages nothing, as {self._no_staging}" + ) + + def _whole_blocks(self, start: int, end: int, whole_end: bool = True) -> tuple[int, int]: + """Validate integer ``[start, end)``; require an ordered, non-negative range starting, + and with ``whole_end`` also ending, at a whole block.""" + for name, value in (("start", start), ("end", end)): + if isinstance(value, bool) or not isinstance(value, numbers.Integral): + raise TypeError(f"{name} must be an integer, got {type(value).__name__}") start, end = int(start), int(end) - tpb = int(self._layout.tokens_per_block) + tpb = self._tokens_per_block if start < 0 or end < 0 or start > end: raise ValueError(f"bad token range [{start}, {end})") - if start % tpb or end % tpb: + if start % tpb or (whole_end and end % tpb): raise ValueError( - f"a staging lease covers whole blocks of {tpb} tokens, got [{start}, {end})" + f"a lease of named rows covers whole blocks of {tpb} tokens, got [{start}, {end})" + if whole_end + else f"a lease starts at a whole block of {tpb} tokens, got [{start}, {end})" ) return start, end @@ -971,22 +2149,34 @@ def _admit(self, lease: _StagingLease, counts: Mapping[int, int]) -> None: """Record a new lease and grant it now when no lease waits and its slots are free, or at once, out of line, when it has no entries.""" self._open_leases.add(lease) - if lease._block_entries.num_block_entries == 0: - # Nothing to stage: ready at the first poll, without slots or a place in line. - self._grant(lease, None) - return - lease._ticket = self._slot_queue.enqueue(counts) - self._waiting_for_slots.append(lease) - self._grant_waiting(just_admitted=lease) + try: + if lease._block_entries.num_rows == 0: + # Nothing to stage: ready at the first poll, without slots or a place in line. + self._grant(lease, None) + return + lease._ticket = self._slot_queue.enqueue(counts) + self._waiting_for_slots.append(lease) + self._grant_waiting(just_admitted=lease) + except BaseException: + # The caller never received this lease. Copies without completion records still keep + # their slots quarantined, independently of the lease's open state. + if lease._ticket is not None: + self._slot_queue.cancel(lease._ticket) + lease._ticket = None + if lease in self._waiting_for_slots: + self._waiting_for_slots.remove(lease) + self._open_leases.discard(lease) + if lease._kind == "write": + self._ledger._abandon_fetch(lease) + raise def _grant_waiting(self, just_admitted: _StagingLease | None = None) -> None: """Grant the leases in line, strictly in order; ``just_admitted`` was looked up in this call.""" while self._waiting_for_slots: head = self._waiting_for_slots[0] - # TODO: an exception while slots are taken, a MemoryError say, can leave slots taken for - # no lease, or drop the head's ticket with the head still in line, after which every - # round of progress raises "not waiting" until the head fails. + # TODO: an exception while slots are taken can leave them taken for no lease, or the + # head in line ticketless, so each progress round raises "not waiting" until it fails. runs = self._slot_queue.try_grant(head._ticket) if runs is None: return @@ -1006,14 +2196,17 @@ def _grant_waiting(self, just_admitted: _StagingLease | None = None) -> None: if head is just_admitted and head._kind == "write": # The cache grew for it, so it fails at its first poll, which abandons the # fetch, as a lease missing pages does. - head._first_poll_failure = reason + head._defer_failure(reason) else: self._fail(head, reason) logger.warning(f"KV cache lender: {reason}") except BaseException: # An interrupt can leave a copy queued on the slots without its completion. head._slot_runs = runs - self._quarantine(head) + try: + self._quarantine(head) + finally: + self._fail(head, "granting staging slots was interrupted") raise def _grant(self, lease: _StagingLease, runs: SlotRuns | None, recheck: bool = False) -> None: @@ -1028,14 +2221,16 @@ def _grant(self, lease: _StagingLease, runs: SlotRuns | None, recheck: bool = Fa if lease._kind == "read" and runs is not None: # Lost until its copy is recorded: no raise from here on leaves it ready. lease._copy = _COPY_WITHOUT_RECORD - lease._view = self._build_view(lease._block_entries, lease._reuse_keys) + lease._view = self._build_view(lease._block_entries, lease._reuse_keys, *lease._range) if runs is None: return lease._slot_runs = runs self._leases_holding_slots.append(lease) if lease._kind != "read": return - copy, error = self._queue_copy(self._copy_segments(lease._block_entries), to_staging=True) + copy, error = self._queue_copy( + self._copy_segments(lease._block_entries), to_staging=True, after=lease._ready_event + ) if copy is None: self._quarantine(lease) if error is not None: @@ -1081,9 +2276,8 @@ def _quarantine(self, lease: _StagingLease) -> None: # The copy is lost and the runs recorded before the lease leaves the holding list, so a step # that raises leaves its slots out of use and the staging memory kept at the shutdown. lease._copy = _COPY_WITHOUT_RECORD - # TODO: quarantined slots are never reused, so failed copies erode staging capacity: a lease - # needing a longer run than any left between lost slots waits in line for good, holding up - # every lease behind it. + # TODO: quarantined slots are never reused, so a lease needing a longer run than any left + # between lost slots waits in line for good, holding up every lease behind it. if lease._slot_runs is not None: self._quarantined.append(lease._slot_runs) self._leases_holding_slots = [ @@ -1093,8 +2287,6 @@ def _quarantine(self, lease: _StagingLease) -> None: f"KV cache lender: a failed copy took staging slots of request {lease._request_id}" ) - # TODO: every lender call checks every lease holding slots, even with no copy pending, so - # polling H such leases once each costs about H*H checks. def _recycle(self) -> None: """Return the slots of every lease ``_recyclable`` allows.""" # Every lease is judged before any slot returns: an event query that raises returns none @@ -1111,23 +2303,25 @@ def _recycle(self) -> None: self._leases_holding_slots = holding def _on_lease_released(self, lease: _StagingLease) -> None: - """The lease's holder let go. The record changes first and progress runs last, so an - unexpected error leaves the fetch abandoned rather than unsettled.""" + """The holder let go: a write released after its view but before its return delivers none of + its rows; its slots return at a later progress round.""" self._open_leases.discard(lease) if self._serving_manager() is None: return - if lease._slot_state() is _SlotState.WAITING: + if lease._ticket is not None: self._fail(lease, "released while waiting for staging slots") - elif lease._kind == "write" and lease._holder_state() is _HolderState.VIEW_PENDING: + elif lease._kind == "write" and lease._failure is None and not lease._view_returned: # Released before anyone saw it ready: no backend wrote, and no marks are due. lease._first_poll_failure = None self._ledger._abandon_fetch(lease) - self._progress() + elif lease._kind == "write" and lease._view_returned and not lease._marked: + lease._marked = True + if self._stop_reason is None: + self._apply_marks(lease, self._no_row_mask(lease)) def _apply_marks(self, lease: _StagingLease, masks: list[np.ndarray]) -> None: - """Copy the marked entries whose page is still the one lent into the request's pages. The - fetch stays abandoned until that copy is queued without error; then its entries add to the - cache's deliveries.""" + """Copy marked entries whose page is still the one lent into the request's pages; the fetch + stays abandoned until that copy queues without error, then joins the deliveries.""" fetch = lease._fetch current = self._ledger._take_out_fetch(lease, fetch) kv = _manager.kv_of(self._manager_ref(), lease._request_id) @@ -1140,7 +2334,9 @@ def _apply_marks(self, lease: _StagingLease, masks: list[np.ndarray]) -> None: tail = self._covers_committed_tail(kv, lease._block_entries, copy) queued = None try: - queued, error = self._queue_copy(segments, to_staging=False) + queued, error = self._queue_copy( + segments, to_staging=False, after=lease._ready_event + ) lease._copy = queued finally: # A copy may be queued without its completion on the lease, also when the call @@ -1152,41 +2348,121 @@ def _apply_marks(self, lease: _StagingLease, masks: list[np.ndarray]) -> None: self._ledger._record_mixed_tail(lease._request_id, kv) if current and error is None: self._ledger._put_back_fetch(lease, fetch, copy) - self._progress() def _covers_committed_tail( - self, kv: _KVCache, entries: _LentBlockEntries, copied: list[np.ndarray] + self, + kv: _KVCache, + entries: _LentBlockEntries, + copied: list[np.ndarray], + *, + complement: bool = False, ) -> bool: - """The copy writes the block the committed tokens end inside, whose committed tokens - resume_bounds counts without a delivery.""" + """Whether the copied rows (with ``complement``, the others) include a group's partial + committed tail block, counted by the bounds undelivered; state rows are not blocks.""" tpb = int(self._layout.tokens_per_block) - committed = _manager.cache_state(kv).committed - if committed % tpb == 0: - return False - tail = committed // tpb - return any(bool(np.any(o[m] == tail)) for o, m in zip(entries.block_ordinals, copied)) + committed = _manager.committed_by_group(kv, self._layout.num_layer_groups) + return any( + committed[lg] % tpb and bool(np.any(o[~m if complement else m] == committed[lg] // tpb)) + for lg, o, m in zip(entries.layer_group_indices, entries.block_ordinals, copied) + if lg not in self._recurrent_groups + ) + + def _range_ordinals( + self, request_id: int | None, lg: int, runs: Sequence[tuple[int, int]] + ) -> np.ndarray: + """Reuse immutable ordinal ranges for successive chunks; keep at most two spans per group. + Each span reserves one further chunk, and replacement leaves earlier views immutable.""" + if request_id is None: + if not runs: + return _NO_ROWS + return _frozen_array(np.concatenate([np.arange(b, e, dtype=np.int64) for b, e in runs])) + cached = self._ordinal_ranges.setdefault(request_id, {}) + parts = [] + for i, (begin, end) in enumerate(runs): + previous = cached.get((lg, i)) + if previous is None or begin < previous[0] or end > previous[0] + len(previous[1]): + previous = (begin, _frozen_array(np.arange(begin, 2 * end - begin, dtype=np.int64))) + cached[(lg, i)] = previous + base, values = previous + parts.append(values[begin - base : end - base]) + if len(parts) == 1: + return parts[0] + if not parts: + return _NO_ROWS + return _frozen_array(np.concatenate(parts)) def _block_entries_of_range( - self, manager: KVCacheManagerV2, kv: _KVCache, start: int, end: int + self, + manager: KVCacheManagerV2, + kv: _KVCache, + start: int | Sequence[int], + end: int, + *, + named: bool = False, + request_id: int | None = None, + ordinals: Sequence[np.ndarray] | None = None, ) -> _LentBlockEntries: - """The blocks of ``[start, end)`` a history of ``end`` reads, per layer group, with their - device slots (-1 where a block has no page).""" + """Retained rows of ``[start, end)`` at history ``end`` (``start`` one or per group), device + pages -1 if none: state row 0 of an active cache, partial tails, under Helix this rank's + pages that intersect the range.""" layout = self._layout tpb = int(layout.tokens_per_block) + rank, size = getattr(layout, "cp", (0, 1)) + starts = [start] * layout.num_layer_groups if isinstance(start, int) else list(start) entries = _LentBlockEntries([], [], []) for lg in range(layout.num_layer_groups): - block_ordinals = _ranges._needed_ordinals( - manager, layout, lg, start // tpb, end // tpb, end + if lg in self._recurrent_groups: + entries.layer_group_indices.append(lg) + entries.block_ordinals.append(_STATE_ROW) + entries.device_pages.append(_frozen_int64((self._state_slot(kv, request_id, lg),))) + continue + block_ordinals = ( + self._range_ordinals( + request_id, + lg, + _ranges._needed_ranges( + manager, layout, lg, starts[lg] // tpb, -(-end // tpb), end + ), + ) + if ordinals is None + else ordinals[lg] ) - pages = _manager.pages(kv, lg) - slots = np.full(len(block_ordinals), -1, dtype=np.int64) - inside = block_ordinals < len(pages) - slots[inside] = pages[block_ordinals[inside]] + if named and not self._named_groups[lg]: + block_ordinals = block_ordinals[:0] + if size > 1: + # This rank's page of block ``b`` holds tokens from ``b * tpb + rank * page``. + page = tpb // size + block_ordinals = _frozen_array( + block_ordinals[block_ordinals * tpb + rank * page < end] + ) + if not len(block_ordinals): + slots = _NO_ROWS + else: + pages = _manager.pages(kv, lg) + try: + gathered = pages[block_ordinals] + except IndexError: + # Blocks past the cache's end have no page. + gathered = np.full(len(block_ordinals), -1, dtype=np.int64) + inside = block_ordinals < len(pages) + gathered[inside] = pages[block_ordinals[inside]] + slots = _frozen_array(gathered) entries.layer_group_indices.append(lg) entries.block_ordinals.append(block_ordinals) entries.device_pages.append(slots) return entries + @staticmethod + def _pageless(entries: _LentBlockEntries) -> str | None: + """The first layer group's blocks of ``entries`` that have no page, as a failure, or + ``None``.""" + for lg, ordinals, slots in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + if len(slots) and slots.min() < 0: + return f"layer group {lg}: blocks {ordinals[slots < 0].tolist()[:8]} have no page" + return None + def _lendable_block_entries( self, manager: KVCacheManagerV2, history: int, entries: _LentBlockEntries ) -> _LentBlockEntries: @@ -1197,10 +2473,18 @@ def _lendable_block_entries( entries.layer_group_indices, entries.block_ordinals, entries.device_pages ): stale_beg, stale_end = _ranges._stale_range(manager, self._layout, lg, history) - ok = (slots >= 0) & ~((block_ordinals >= stale_beg) & (block_ordinals < stale_end)) + if stale_end <= stale_beg: + # Nothing behind a window: only blocks without a page drop out. + whole = not len(slots) or slots.min() >= 0 + ok = None if whole else slots >= 0 + else: + ok = (slots >= 0) & ~((block_ordinals >= stale_beg) & (block_ordinals < stale_end)) + whole = bool(ok.all()) out.layer_group_indices.append(lg) - out.block_ordinals.append(block_ordinals[ok]) - out.device_pages.append(slots[ok]) + out.block_ordinals.append( + block_ordinals if whole else _frozen_array(block_ordinals[ok]) + ) + out.device_pages.append(slots if whole else _frozen_array(slots[ok])) return out def _reuse_keys_for( @@ -1234,30 +2518,474 @@ def _assign_staging_slots(self, entries: _LentBlockEntries, runs: SlotRuns | Non ) cursor[g] = offset + len(block_ordinals) - def _build_view(self, entries: _LentBlockEntries, reuse_keys: list[np.ndarray]) -> LeaseView: + def _build_view( + self, entries: _LentBlockEntries, reuse_keys: list[np.ndarray], start: int, end: int + ) -> LeaseView: + """Each entry as its staging slot, with its object key; every array a snapshot and an + object of the view's own; the view of range ``[start, end)``.""" groups = [] for lg, block_ordinals, lg_keys, slots in zip( entries.layer_group_indices, entries.block_ordinals, reuse_keys, entries.staging_slots ): - index = self._region_index_of_pool_group[int(self._layout.pool_group_of[lg])] - region = self._regions[index] groups.append( - LayerGroupBlocks( - lg, - block_ordinals, - object_keys=self._object_keys(lg, lg_keys), - addresses=region.address + slots * region.slot_bytes, - region_index=index, + LentLayerGroupRows._unchecked( + int(lg), + _unshared(block_ordinals), + _immutable(self._object_keys(lg, lg_keys)), + slots=_frozen_array(slots), + region_indices=self._staging_region_indices[ + int(self._layout.pool_group_of[lg]) + ], + ) + ) + return LeaseView._unchecked(tuple(groups), start, end) + + def _failed_at_call( + self, memory_kind: MemoryKind, kind: str, request_id: int, reason: str + ) -> _StagingLease | _InPlaceLease: + """A lease failed at the call, of the class leases of ``memory_kind`` have.""" + if memory_kind is MemoryKind.DEVICE: + return _InPlaceLease._failed_at_call(self, kind, request_id, reason) + return _StagingLease._failed_at_call(self, kind, request_id, reason) + + def _slot_limit_refusal(self, entries: _LentBlockEntries) -> str | None: + """The first group's blocks whose page reaches its device slot limit (0 without device + regions), as a failure, or ``None``; compares each group's largest page once.""" + for lg, ordinals, pages in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + limit = self._device_slot_limit.get(lg, 0) + if len(pages) and pages.max() >= limit: + return ( + f"layer group {lg}: blocks {ordinals[pages >= limit].tolist()[:8]} have pages " + f"past the {limit} slots of its device regions" + ) + return None + + def _lend_in_place( + self, + manager: KVCacheManagerV2, + request: LlmRequest, + kv: _KVCache, + kind: str, + entries: _LentBlockEntries, + reuse_keys: list[np.ndarray] | None, + fetch: _FetchRecord | None = None, + doomed: str | None = None, + *, + start: int, + end: int, + checked: CheckedBlocks | None = None, + ready_event: torch.cuda.Event | None = None, + ) -> _InPlaceLease: + """Lock current pages; the lease is ready once ``ready_event`` and, while the cache is on + the allocation side stream, the manager's fence completed, a read once pending connector + loads finished. Refused locks, missing pages or pages past the slot limit fail the lease and + abandon a write's fetch. ``checked``: this lend's check of exactly these blocks.""" + request_id = int(request.py_request_id) + lease = _InPlaceLease(self, kind, request_id, kv, entries, fetch) + if doomed is None: + doomed = self._slot_limit_refusal(entries) + if doomed is not None and fetch is not None: + lease._defer_failure(doomed) + return lease + if doomed is None: + lease._view = self._in_place_view(entries, reuse_keys, start, end) + if kind == "read": + mode = LockMode.READ + if _manager.connector_load_pending(manager, request): + lease._loading = request + else: + mode = LockMode.WRITE + blocks = ( + dict(zip(entries.layer_group_indices, entries.block_ordinals)) + if checked is None + else checked.blocks + ) + fence = _manager.allocation_fence(manager, request_id) + ready = ready_event if fence is None else _fenced(ready_event, fence) + try: + lease._lock = self._locks.lock( + kv, blocks, mode, ready=ready, ordered=True, checked=checked + ) + except ValueError as error: + doomed = f"the blocks cannot be locked: {error}" + if doomed is not None: + if fetch is not None: + lease._defer_failure(doomed) + else: + self._fail_in_place(lease, doomed) + return lease + + def _in_place_view( + self, entries: _LentBlockEntries, reuse_keys: list[np.ndarray] | None, start: int, end: int + ) -> LeaseView: + """The view of device-page entries as snapshots of its own, with optional object keys; under + Helix, ordinal ``b * size + rank`` names this rank's page.""" + rank, size = getattr(self._layout, "cp", (0, 1)) + groups = [] + for i, (lg, block_ordinals, pages) in enumerate( + zip(entries.layer_group_indices, entries.block_ordinals, entries.device_pages) + ): + group = int(self._layout.pool_group_of[lg]) + groups.append( + LentLayerGroupRows._unchecked( + int(lg), + _frozen_array(block_ordinals * size + rank) + if size > 1 + else _unshared(block_ordinals), + None + if reuse_keys is None + else _immutable(self._object_keys(lg, reuse_keys[i])), + slots=_unshared(pages), + region_indices=self._device_regions_of_pool_group[group], ) ) - return LeaseView(tuple(groups)) + return LeaseView._unchecked(tuple(groups), start, end) + + # Fast lends of every row in place: what the rules give where they lend every block and refuse + # nothing, or ``None``, the rules then lending as they would have. + + def _lends_fast(self) -> bool: + """Whether this lender and ledger have a layout covered by the fast path.""" + return ( + self._fast_layout + and type(self) is _ManagerLender + and type(self._ledger) is _FetchLedger + ) + + def _read_every_row_fast( + self, + manager: KVCacheManagerV2, + request: LlmRequest, + request_id: int, + start: int, + end: int, + ready_event: torch.cuda.Event | None, + ) -> _InPlaceLease | None: + """``lend_read``'s lease of every row in place for an active cache that no fetch record + names, holds the range and a page for each of its blocks, and with a window has its context + computed; ``None`` elsewhere, having changed nothing the rules read.""" + kv = _manager.kv_of(manager, request_id) + if kv is None or self._ledger._has_records(request_id): + return None + state = _manager.cache_state(kv) + history = state.history + if ( + not state.active + or self._read_end_refusal(manager, request, state, end, False) is not None + ): + return None + # A computed range past history can be valid under the general rules, but has no pages in + # this cache that the fast path can lend. + if end > history or self._mid_context_end(manager, request, state) is not None: + return None + reserved, loading = _manager.connector_state(manager, request) + if reserved: + return None + # A window keeps at the cache's history, at or past the range's end, what a read lends. + rows = self._every_row( + kv, request_id, start, end, kept_at=history, row=self._page_row(manager, request_id, kv) + ) + if rows is None: + return None + entries, view, blocks = rows + lease = _InPlaceLease(self, "read", request_id, kv, entries) + lease._view = view + if loading: + lease._loading = request + locks = self._locks + fence = _manager.allocation_fence(manager, request_id) + ready = ready_event if fence is None else _fenced(ready_event, fence) + lock = locks.lock_fresh(kv, blocks, LockMode.READ, ready) + if lock is None: + try: + lock = locks.lock(kv, blocks, LockMode.READ, ready=ready, ordered=True) + except ValueError as error: + self._fail_in_place(lease, f"the blocks cannot be locked: {error}") + return lease + lease._lock = lock + return lease + + def _write_every_row_fast( + self, + manager: KVCacheManagerV2, + request: LlmRequest, + request_id: int, + start: int, + end: int, + prompt: int, + ready_event: torch.cuda.Event | None, + ) -> _InPlaceLease | None: + """``lend_write``'s lease of every row in place for an active cache without a fetch to + settle or SWA scratch reuse, a range its prompt and history allow, and a page for each block + it lends; ``None`` elsewhere, having changed nothing the rules read.""" + kv = _manager.kv_of(manager, request_id) + if kv is None: + return None + state = _manager.cache_state(kv) + if ( + not state.active + or self._write_range_refusal(manager, request, state, start, end, prompt) is not None + ): + return None + if self._history_refusal(manager, state.history, state.history, end) is not None: + return None + if self._any_window and ( + self._unwritten_pages_refusal(manager, request_id, kv, state.history) is not None + ): + return None + if self._bidirectional_split_refusal(request, end) is not None: + return None + runs = _manager.bidirectional_runs(request) + if self._scratch_reuse_on(kv): + return None + count = len(self._group_regions) + committed = (0,) * count if not state.committed else _manager.committed_by_group(kv, count) + # The pages from the cache's row, which it keeps current itself. + row = self._page_row(manager, request_id, kv) + rows = self._every_row(kv, request_id, start, end, committed, kept_at=end, row=row) + if rows is None: + return None + entries, view, blocks = rows + # The rows start past each group's committed blocks, so only another lock, a hand-over at a + # free, or with a window the other locks' history cap can refuse them. + checked = None + if self._locks.is_locked(kv): + checked = self._locks.check( + kv, blocks, LockMode.WRITE, ordered=True, committed=committed + ) + if isinstance(checked, str): + return None + if self._any_window: + cap = self._history_cap(kv) + if cap is not None and state.history > cap: + return None + # Its whole return resumes at ``end`` where no record of the request's could lower the + # bounds and the rows below ``start`` are computed. + at_end = start <= state.history and not self._ledger._has_records(request_id) + # Last, as it records the fetch where no earlier fetch into the cache is unsettled. + fetch = self._ledger._record_next_fetch( + request_id, + kv, + start, + state.history, + committed, + end if self._recurrent_groups else None, + end, + ) + if fetch is None: + return None + if checked is not None: + # Only a locked cache has a commit held back, which would commit computed rows this + # fetch may overwrite. + _manager.drop_held_commit(manager, kv, start) + fetch.entries = entries + fetch.prompt = prompt + fetch.runs = runs + lease = _InPlaceLease(self, "write", request_id, kv, entries, fetch) + lease._view = view + lease._shrinks = self._shrink_epochs.get(request_id, 0) + lease._block_end = entries.block_end + locks = self._locks + fence = _manager.allocation_fence(manager, request_id) + ready = ready_event if fence is None else _fenced(ready_event, fence) + lock = locks.lock_fresh(kv, blocks, LockMode.WRITE, ready) if checked is None else None + if lock is None: + try: + lock = locks.lock( + kv, + blocks, + LockMode.WRITE, + ready=ready, + ordered=True, + checked=checked, + committed=committed, + ) + except ValueError as error: + lease._defer_failure(f"the blocks cannot be locked: {error}") + return lease + lease._lock = lock + if at_end: + lease._at_end = True + self._writes_at_end[id(kv)] = lease + return lease + + def _every_row( + self, + kv: _KVCache, + request_id: int, + start: int, + end: int, + committed: Sequence[int] | None = None, + *, + kept_at: int | None = None, + row: tuple[memoryview, int, int, int] | None = None, + ) -> tuple[_LentBlockEntries, LeaseView, dict[int, np.ndarray]] | None: + """``(entries, view, blocks to lock)`` of ``[start, end)``, given no Helix and sinkless + formula windows: windows from what ``kept_at`` keeps, writes from ``committed`` whole blocks + (none past ``end``), a state row per recurrent group, pages from ``row`` or the cache, view + arrays viewing the others'; ``None`` if a page is absent or not below its limit.""" + tpb = self._tokens_per_block + last = -(-end // tpb) + listed = None + if row is not None: + words, row_at, group_words, row_blocks = row + if last > row_blocks: + return None + layer_groups, block_ordinals, device_pages, groups, spans = [], [], [], [], [] + blocks = {} + block_end = 0 + for lg, regions, limit, recurrent, window in self._fast_groups: + if recurrent: + slot = self._state_slot(kv, request_id, lg) + if not 0 <= slot < limit: + return None + ordinals, slots = _STATE_ROW, self._state_rows.get(slot) + if slots is None: + slots = self._state_rows[slot] = _frozen_int64((slot,)) + spans.append(None) + else: + first = start // tpb + if committed is not None: + tokens = committed[lg] + first = last if tokens > end else max(first, tokens // tpb) + if window is not None: + # The window formula without sinks: blocks before (history + 1 - window) // tpb + # are behind the window. + behind = min(kept_at, _manager._MAX_TOKENS - tpb) + 1 - window + if behind > 0: + first = max(first, behind // tpb) + spans.append((first, last)) + if last <= first: + ordinals = slots = _NO_ROWS + else: + block_end = last + if last > _ORDINAL_BLOCKS: + ordinals = self._fresh_ordinals(request_id, lg, first, last) + else: + ordinals = _ordinal_views.get((first, last)) + if ordinals is None: + if len(_ordinal_views) >= _ORDINAL_VIEWS_KEPT: + _ordinal_views.clear() + ordinals = _ordinal_views[first, last] = _ORDINALS[first:last] + if row is not None: + at = row_at + lg * group_words + first + if last - first <= _LISTED_BLOCKS: + held = words[at : at + last - first].tolist() + if max(held) >= limit: + return None + try: + # ``_frozen_pages``: the pages as int64 on immutable bytes. + slots = _from_buffer(_PAGE_PACKERS[len(held)].pack(*held), _INT64) + except struct.error: + return None + else: + held = np.frombuffer(words, _INT32, last - first, 4 * at) + # As unsigned, a page of -1 is past every limit too. + if held.view(_UINT32).max() >= limit: + return None + slots = _frozen_array(held) + else: + if listed is None: + listed = _manager.num_blocks(kv) <= _LISTED_BLOCKS + if listed: + pages = _manager.page_list(kv, lg) + if last > len(pages): + return None + held = pages[first:last] + if max(held) >= limit: + return None + try: + slots = _frozen_pages(held) + except struct.error: + return None + else: + pages = _manager.pages(kv, lg) + if last > len(pages): + return None + slots = _frozen_array(pages[first:last]) + if slots.view(_UINT64).max() >= limit: + return None + layer_groups.append(lg) + block_ordinals.append(ordinals) + device_pages.append(slots) + blocks[lg] = ordinals + groups.append(_lent_rows(lg, ordinals.view(), slots.view(), regions)) + entries = _LentBlockEntries( + layer_groups, block_ordinals, device_pages, spans=spans, block_end=block_end + ) + return entries, _new_view(tuple(groups), start, end), blocks + + def _fresh_ordinals(self, request_id: int, lg: int, first: int, last: int) -> np.ndarray: + """Blocks ``[first, last)`` past ``_ORDINALS`` as frozen ordinals: from the request's + templates if any, else on immutable bytes of the range's size.""" + if request_id in self._ordinal_ranges: + return self._range_ordinals(request_id, lg, [(first, last)]) + return _frozen_array(np.arange(first, last, dtype=np.int64)) + + def _fail_in_place(self, lease: _InPlaceLease, reason: str) -> None: + """Fail an in-place lease never seen ready: its lock ends, and a write abandons its + fetch.""" + if lease.failure is not None: + return + lease._set_failure(reason) + if lease._lock is not None: + self._locks.unlock(lease._lock) + if lease._kind == "write": + self._ledger._abandon_fetch(lease) + + def _apply_in_place_marks( + self, lease: _InPlaceLease, masks: list[np.ndarray], lock_held: bool = False + ) -> None: + """Invalidate all lent entries (writes may touch any row), then count arrived ones whose + pages still match (all with ``lock_held``); a lent, undelivered committed tail taints the + cache.""" + fetch = lease._fetch + current = self._ledger._take_out_fetch(lease, fetch) + kv = _manager.kv_of(self._manager_ref(), lease._request_id) + entries = lease._block_entries + if lock_held: + delivered = masks + else: + lent = self._locked_block_masks(kv, lease) + delivered = [mask & held for mask, held in zip(masks, lent)] + self._ledger._invalidate_delivered_blocks(lease, None) + if kv is lease._kv and self._covers_committed_tail(kv, entries, delivered, complement=True): + self._ledger._record_mixed_tail(lease._request_id, kv) + if current: + self._ledger._put_back_fetch(lease, fetch, delivered) + + def _on_in_place_released(self, lease: _InPlaceLease) -> None: + """The holder let go: the lock ends; a write released after its view but before its return + delivers no rows, and one released before its view, which no backend wrote through, abandons + its fetch.""" + if lease._at_end: + lease._at_end = False + self._writes_at_end.pop(id(lease._kv), None) + try: + if lease._kind == "write" and lease._view_returned and not lease._marked: + lease._marked = True + manager = self._serving_manager() + if manager is not None and lease._fetch is not None: + _manager.drop_held_commit(manager, lease._kv, lease._fetch.start) + if manager is not None and self._stop_reason is None: + self._apply_in_place_marks(lease, self._no_row_mask(lease)) + finally: + if lease._lock is not None: + self._locks.unlock(lease._lock) + if ( + lease._kind == "write" + and not lease._view_returned + and self._serving_manager() is not None + ): + self._ledger._abandon_fetch(lease) def _copy_segments( self, entries: _LentBlockEntries, mask: list[np.ndarray] | None = None ) -> list[list[int]]: - """``[staging address, device address, bytes]`` for every pool of every row ``mask`` keeps - (one boolean array per layer group; ``None``: every row), by device pool group. Segments - that continue each other on both sides are merged.""" + """``[staging address, device address, bytes]`` per pool of each row ``mask`` keeps + (``None``: all), by pool group, merging segments continuous on both sides.""" by_group: dict[int, list[int]] = {} for i, lg in enumerate(entries.layer_group_indices): by_group.setdefault(int(self._layout.pool_group_of[lg]), []).append(i) @@ -1284,19 +3012,22 @@ def _copy_segments( offset += width return out - # TODO: copies run serially with the forward passes on the execution stream, and a copy call per - # row and pool, which entries share only where they continue each other in a pool group of one - # pool, keeps small entries and entries spanning several pools below the host-to-device bandwidth. + # TODO: copies run serially with forward passes on the execution stream, and one call per row + # and pool keeps small or multi-pool entries below the host-to-device bandwidth. def _queue_copy( - self, segments: Sequence[Sequence[int]], to_staging: bool + self, + segments: Sequence[Sequence[int]], + to_staging: bool, + after: torch.cuda.Event | None = None, ) -> tuple[_QueuedCopy | None, str | None]: - """Queue async copies on the manager's execution stream, then record an event covering - every copy queued, even after one failed; ``(None, error)`` if it could record none. With - page-locked staging the CPU does not wait for them; pageable staging may hold the call.""" - # The execution stream orders a copy after the forward passes that wrote its pages and - # before later work on it, which a page's new owner waits for; the lender does not rely on a - # release waiting for the copies, and a writer off that stream is not ordered after them. + """Queue async copies on the manager's execution stream after ``after``, the lease's ready + event, then record one event, the lender's record of them, covering every copy queued even + after a failure; ``(None, error)`` without one. Pageable staging may hold the call.""" + # The lender does not rely on a release waiting for the copies; a writer off this stream is + # not ordered after them. stream = _manager.stream(self._manager_ref()) + if after is not None: + stream.wait_event(after) handle = drv.CUstream(stream.cuda_stream) error = None for host, device, nbytes in segments: @@ -1320,115 +3051,86 @@ def _queue_copy( class _LeaseBase: - """A lease's failure, and the checks of a write's one mark.""" + """A lease's failure, whether its return resumes at its end, and the checks of the masks its + return counts.""" + + def __init__( + self, + lender: _ManagerLender, + kind: str, + request_id: int, + kv: _KVCache | None, + entries: _LentBlockEntries | None, + fetch: _FetchRecord | None, + ) -> None: + self._lender = weakref.ref(lender) + self._kind = kind + self._request_id = request_id + self._kv = kv + self._block_entries = entries + self._fetch = fetch + self._view: LeaseView | None = None + self._view_returned = False + self._released = False + self._marked = False + self._failure: str | None = None + self._first_poll_failure: str | None = None + # Set only for a fast write in place; the lender's hooks clear it. + self._at_end = False + + def _defer_failure(self, reason: str) -> None: + """Fail on the first poll so a fetched write is abandoned when its caller sees failure.""" + self._first_poll_failure = reason - # Set by each lease class's ``__init__``; the methods below use them. - _failure: str | None - _kind: str - _view_returned: bool - _marked: bool - _view: LeaseView | None + def _set_failure(self, reason: str) -> None: + """Record a terminal failure and remove any reason deferred to the first poll.""" + self._failure = reason + self._first_poll_failure = None @property def failure(self) -> str | None: """See ``Lease.failure``.""" return self._failure - def _accept_marks(self, masks: Sequence[np.ndarray]) -> list[np.ndarray]: - """Copies of ``masks`` for the one mark of a write whose view ``poll()`` returned; - ``ValueError`` unless one bool mask of shape ``(len(blocks),)`` per group collection.""" - if self._kind != "write": - raise RuntimeError("mark_arrived is for write leases") - if not self._view_returned: - raise RuntimeError("mark_arrived before poll() returned the view") - if self._marked: - raise RuntimeError("mark_arrived called twice") - masks, groups = list(masks), self._view.groups + @property + def writable(self) -> bool: + """See ``Lease.writable``.""" + return self._kind == "write" + + @property + def return_resumes_at_end(self) -> bool: + """See ``Lease.return_resumes_at_end``.""" + return self._at_end and self._view_returned + + def _accepted_masks(self, masks: Sequence[np.ndarray]) -> list[np.ndarray]: + """``masks`` of a write's arrived rows as its return keeps them, a read-only exact + ``numpy.ndarray`` as is, others copied; ``ValueError`` unless one bool mask of shape + ``(len(blocks),)`` per group. Changes nothing.""" + masks, groups = list(masks), self._block_entries.block_ordinals if len(masks) != len(groups): raise ValueError(f"{len(masks)} masks for {len(groups)} groups") out = [] - for blocks, mask in zip(groups, masks): - mask = np.asarray(mask) + for lg, blocks, mask in zip(self._block_entries.layer_group_indices, groups, masks): + if type(mask) is not np.ndarray: + mask = np.array(mask) + elif mask.flags.writeable: + mask = mask.copy() if mask.dtype != np.bool_ or mask.shape != (len(blocks),): raise ValueError( - f"layer group {blocks.layer_group_index}: the mask must be bool of shape " + f"layer group {lg}: the mask must be bool of shape " f"({len(blocks)},), got {mask.dtype} {mask.shape}" ) - out.append(mask.copy()) - self._marked = True + out.append(mask) return out -class _LeaseState(enum.Enum): - """A staging lease's state, which ``_StagingLease._state`` derives from its fields. ``release`` - moves any state to ``RELEASED``. Its slots and its copy vary apart from the state. The slots: - none, waiting in order in ``Staging._waiting_for_slots``, or held in - ``Staging._leases_holding_slots`` until ``_recycle`` returns them or ``_quarantine`` keeps them - from reuse. The copy (``_copy``): none, ``_COPY_WITHOUT_RECORD``, or a ``_QueuedCopy``, pending - or done. Three axes each derive from their own fields alone, released or not: ``_HolderState``, - whether the lease failed and, if not, whether a poll returned its view; ``_SlotState``, whether - it waits in line for slots; ``_CopyState``, whether its copy is ``_COPY_WITHOUT_RECORD``. - ``Staging._recyclable`` reads the first, ``Staging._on_lease_released`` the first two and - ``Staging._slots_quarantined`` the third.""" - - # In line for slots: _grant_waiting grants it in turn, or dooms a fresh write whose grant - # raised; _fail fails it. - WAITING = "waiting" - # Without slots or a view: its first poll fails it, as does the shutdown. - FAILS_AT_FIRST_POLL = "fails_at_first_poll" - # Its view and any slots: poll makes it ready, a read once its copy into them has landed, and - # fails a read whose copy is lost. - GRANTED = "granted" - # Poll returned its view: a backend may use its slots until release; mark_arrived marks a write. - READY = "ready" - # A write mark_arrived marked: its marked entries still lent go to the request's pages. - MARKED = "marked" - # Failed at the call, in line, at its grant, at a poll or at the shutdown. - FAILED = "failed" - # Let go: _recycle returns its slots once neither a backend nor a copy can reach them. - RELEASED = "released" - - -class _HolderState(enum.Enum): - """Whether a staging lease failed and, if not, whether a poll returned its view, which - ``_StagingLease._holder_state`` derives from ``_failure`` and ``_view_returned`` alone, released - or not.""" - - # Its failure is set. - FAILED = "failed" - # No failure, and no poll has returned its view. - VIEW_PENDING = "view_pending" - # No failure, and a poll returned its view. - VIEW_RETURNED = "view_returned" - - -class _SlotState(enum.Enum): - """Whether a staging lease waits in line for slots, which ``_StagingLease._slot_state`` derives - from ``_ticket`` alone.""" - - # A ticket: in line for slots. - WAITING = "waiting" - # No ticket. - NOT_WAITING = "not_waiting" - - -class _CopyState(enum.Enum): - """Whether a staging lease's copy is ``_COPY_WITHOUT_RECORD``, which - ``_StagingLease._copy_state`` derives from ``_copy`` alone.""" - - # _COPY_WITHOUT_RECORD: a copy that may be queued with no record of its completion. - WITHOUT_RECORD = "without_record" - # No copy, or one queued with the record of its completion. - NONE_OR_RECORDED = "none_or_recorded" - - class _StagingLease(_LeaseBase): - """A staging lease; holds its lender weakly and has no finalizer. Backends may read its view's - arrays on their own threads until release, and never call it.""" + """Host staging lease; see ``Lease`` for thread and access lifetimes. + Holds its lender weakly and has no finalizer.""" def __init__( self, - lender: Staging, + lender: _ManagerLender, kind: str, request_id: int, kv: _KVCache | None = None, @@ -1436,27 +3138,20 @@ def __init__( reuse_keys: list[np.ndarray] | None = None, fetch: _FetchRecord | None = None, ) -> None: - self._lender = weakref.ref(lender) - self._kind = kind - self._request_id = request_id - self._kv = kv - self._block_entries = entries + super().__init__(lender, kind, request_id, kv, entries, fetch) self._reuse_keys = reuse_keys - self._fetch = fetch - self._first_poll_failure: str | None = None # why it fails at its first poll + # The caller's event after which the pages are free of work in flight; copies wait for it. + self._ready_event: torch.cuda.Event | None = None + # The lend's token range, which its view gives. + self._range = (0, 0) self._ticket: int | None = None # its place in line while it waits for slots # The granted slots; _recycle returns them and leaves this field set. self._slot_runs: SlotRuns | None = None - self._view: LeaseView | None = None self._copy: _QueuedCopy | None = None - self._view_returned = False - self._released = False - self._marked = False - self._failure: str | None = None @classmethod def _failed_at_call( - cls, lender: Staging, kind: str, request_id: int, reason: str + cls, lender: _ManagerLender, kind: str, request_id: int, reason: str ) -> _StagingLease: """A lease failed at the call; open until released, like any other.""" lease = cls(lender, kind, request_id) @@ -1465,47 +3160,25 @@ def _failed_at_call( lender._open_leases.add(lease) return lease - def _set_failure(self, reason: str) -> None: - self._failure = reason - self._first_poll_failure = None - - def _state(self) -> _LeaseState: - """Its state, from the first of these fields that is set; see ``_LeaseState``.""" - if self._released: - return _LeaseState.RELEASED - if self._failure is not None: - return _LeaseState.FAILED - if self._first_poll_failure is not None: - return _LeaseState.FAILS_AT_FIRST_POLL - if self._ticket is not None: - return _LeaseState.WAITING - if self._marked: - return _LeaseState.MARKED - return _LeaseState.READY if self._view_returned else _LeaseState.GRANTED - - def _holder_state(self) -> _HolderState: - """See ``_HolderState``.""" - if self._failure is not None: - return _HolderState.FAILED - return _HolderState.VIEW_RETURNED if self._view_returned else _HolderState.VIEW_PENDING - - def _slot_state(self) -> _SlotState: - """See ``_SlotState``.""" - return _SlotState.WAITING if self._ticket is not None else _SlotState.NOT_WAITING - - def _copy_state(self) -> _CopyState: - """See ``_CopyState``.""" - if self._copy is _COPY_WITHOUT_RECORD: - return _CopyState.WITHOUT_RECORD - return _CopyState.NONE_OR_RECORDED - def _marks_due(self) -> bool: - """A write whose view a poll returned and that has no mark: ``mark_arrived`` may still come - after release, and the copy it queues reads the lease's slots.""" + """A write whose view a poll returned and that was not returned yet: its return may still + queue a copy reading the lease's slots.""" return self._kind == "write" and self._view_returned and not self._marked def __repr__(self) -> str: - return f"" + if self._released: + state = "released" + elif self._failure is not None: + state = "failed" + elif self._first_poll_failure is not None: + state = "fails_at_first_poll" + elif self._ticket is not None: + state = "waiting" + elif self._marked: + state = "marked" + else: + state = "ready" if self._view_returned else "granted" + return f"" def poll(self) -> LeaseView | None: """See ``Lease.poll``.""" @@ -1521,7 +3194,6 @@ def poll(self) -> LeaseView | None: return None if lender is None or lender._serving_manager() is None: return self._poll_not_serving() - lender._progress() if self._failure is not None: return None if self._first_poll_failure is not None: @@ -1534,17 +3206,6 @@ def poll(self) -> LeaseView | None: self._view_returned = True return self._view - def mark_arrived(self, masks: Sequence[np.ndarray]) -> None: - """See ``Lease.mark_arrived``.""" - masks = self._accept_marks(masks) - lender = self._lender() - if ( - lender is not None - and lender._serving_manager() is not None - and lender._stop_reason is None - ): - lender._apply_marks(self, masks) - def release(self) -> None: """See ``Lease.release``.""" if self._released: @@ -1564,11 +3225,89 @@ def _poll_not_serving(self) -> LeaseView | None: return self._view -class _StagingHold: - """A hold on the staging memory; holds its lender weakly and has no finalizer. Without a lender - it is inert.""" +class _InPlaceLease(_LeaseBase): + """Locked request pages; poll waits for the lock event and pending connector loads. + See ``Lease`` for thread and access lifetimes. Holds the lender weakly, without a finalizer.""" + + def __init__( + self, + lender: _ManagerLender, + kind: str, + request_id: int, + kv: _KVCache | None = None, + entries: _LentBlockEntries | None = None, + fetch: _FetchRecord | None = None, + ) -> None: + super().__init__(lender, kind, request_id, kv, entries, fetch) + self._lock: PageLock | None = None + # The fetch ledger reads a lease's copy of its marks; marks in place copy nothing. + self._copy = None + # The request a KV cache connector was loading into at the lend, until it has finished. + self._loading: LlmRequest | None = None + # For a fast write: the request's shrink count at the lend, and the end of its lent blocks. + self._shrinks = 0 + self._block_end = 0 + + @classmethod + def _failed_at_call( + cls, lender: _ManagerLender, kind: str, request_id: int, reason: str + ) -> _InPlaceLease: + """A lease failed at the call, with no lock and no fetch.""" + lease = cls(lender, kind, request_id) + lease._set_failure(reason) + return lease + + def _finish_returned_write(self) -> None: + """Close a returned write's view while its lock remains held until the lender unlocks.""" + self._marked = True + self._released = True + self._at_end = False + + def __repr__(self) -> str: + state = "released" if self._released else "held" if self._failure is None else "failed" + return f"" + + def poll(self) -> LeaseView | None: + """See ``Lease.poll``.""" + if self._released: + raise RuntimeError("poll after release") + if self._failure is not None: + return None + lender = self._lender() + if lender is not None and lender._stop_reason is not None and not self._view_returned: + lender._fail_in_place(self, lender._stop_reason) + return None + if self._first_poll_failure is not None: + if lender is not None: + lender._fail_in_place(self, self._first_poll_failure) + else: + self._set_failure(self._first_poll_failure) + return None + if not self._view_returned and self._loading is not None: + manager = lender._serving_manager() if lender is not None else None + if manager is not None and _manager.connector_load_pending(manager, self._loading): + return None + self._loading = None + if not self._view_returned and not self._lock.is_ready(): + return None + self._view_returned = True + return self._view + + def release(self) -> None: + """See ``Lease.release``.""" + if self._released: + return + self._released = True + lender = self._lender() + if lender is not None: + lender._on_in_place_released(self) + + +class _RegionHold: + """A hold on the lender's regions: its staging memory and the runtime with its device pools. + Holds its lender weakly and has no finalizer; without a lender it is inert.""" - def __init__(self, lender: Staging | None) -> None: + def __init__(self, lender: _ManagerLender | None) -> None: self._lender = weakref.ref(lender) if lender is not None else None self._released = False diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py index aa1f46cd2813..a2ea841f7e3c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py @@ -12,18 +12,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The boundary for sharing's V2 operations, runtime metadata and private lifecycle hooks. - -Manager and runtime operations, including layout declarations, prefix resize and fresh-page fill, -go through this facade, as do reads of runtime layer configs, layer grouping and pool descriptors. -The layout builder reads the manager's head counts and other structural fields directly. Manager -and runtime imports stay inside the functions; the private hook protocol is checked before a -lender is installed. -""" +"""Facade for V2 operations, runtime metadata and lifecycle hooks; heavy imports stay local. +Layout also reads structural manager fields directly.""" from __future__ import annotations -from typing import TYPE_CHECKING, NamedTuple, Protocol, Sequence, runtime_checkable +from typing import TYPE_CHECKING, NamedTuple, Sequence import numpy as np @@ -31,44 +25,26 @@ import torch from tensorrt_llm.mapping import Mapping - from tensorrt_llm.runtime.kv_cache_manager_v2 import ( - AttentionLayerConfig, - PoolGroupDesc, - SsmLayerConfig, - _KVCache, - ) + from tensorrt_llm.runtime.kv_cache_manager_v2 import _KVCache from ...llm_request import LlmRequest + from .._lender_hooks import _ManagerHooks from ..kv_cache_manager_v2 import KVCacheManagerV2 # The largest 32-bit count: the manager's window ranges compute in 32-bit token counts. _MAX_TOKENS = 2**31 - 1 +# ``BlockReusePolicy.ALL_REUSABLE``, imported at the first call that needs it. +_all_reusable = None class CacheState(NamedTuple): - """A request's runtime cache as the lender reads it: tokens committed, the history length its - sliding windows keep, and whether it is active (not suspended).""" + """Committed tokens, retained history and activity (not suspended) of a runtime cache.""" committed: int history: int active: bool -@runtime_checkable -class _ManagerHooks(Protocol): - """Callbacks a lender supplies for its manager's cache lifecycle.""" - - def _on_free(self, request_id: int, kv_cache: _KVCache) -> None: ... - - def _on_shrink(self, request_id: int, kv_cache: _KVCache) -> None: ... - - def _on_reset(self) -> None: ... - - def _on_shutdown(self) -> None: ... - - def _on_caches_closed(self) -> None: ... - - def mapping(manager: KVCacheManagerV2) -> Mapping: """The manager's mapping; ``TypeError`` for a manager without one.""" value = getattr(manager, "mapping", None) @@ -84,12 +60,39 @@ def has_lossy_cold_page_codec(manager: KVCacheManagerV2) -> bool: def require_v2(manager: object) -> None: """Raise ``TypeError`` unless ``manager`` is a ``KVCacheManagerV2`` (subclasses included).""" + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from ..kv_cache_manager_v2 import KVCacheManagerV2 if not isinstance(manager, KVCacheManagerV2): raise TypeError( f"the KV cache lender needs a KVCacheManagerV2, got {type(manager).__name__}" ) + if getattr(manager, "kv_cache_type", None) == CacheType.CROSS: + raise ValueError( + "the KV cache lender does not support CROSS caches or encoder-decoder models" + ) + + +def model_has_recurrent_state(manager: KVCacheManagerV2) -> bool: + """Whether the manager serves a model with recurrent state: a hybrid manager, on every pipeline + stage, whichever layers this stage holds.""" + from ..mamba_cache_manager import MambaHybridCacheManagerV2 + + return isinstance(manager, MambaHybridCacheManagerV2) + + +def model_has_attention_window(manager: KVCacheManagerV2) -> bool: + """Whether any pipeline stage retains a configured or backend-defined window.""" + pattern = manager._model_attention_window_pattern + return manager._model_has_attention_window or any( + 0 < window < manager.max_seq_len for window in pattern[: manager.num_layers] + ) + + +def model_has_sparse_buffers(manager: KVCacheManagerV2) -> bool: + """Whether any pipeline stage holds sparse buffers.""" + return manager._model_has_sparse_buffers def commits_blocks(manager: KVCacheManagerV2) -> bool: @@ -98,40 +101,42 @@ def commits_blocks(manager: KVCacheManagerV2) -> bool: return bool(manager.enable_block_reuse) and bool(manager._can_publish_block_reuse) -def kv_connector(manager: KVCacheManagerV2) -> object | None: - """The manager's KV cache connector, or ``None`` without one.""" - return manager.kv_connector_manager +def connector_loads(manager: KVCacheManagerV2) -> bool: + """Whether a connector loads KV; capacity-only connectors transfer nothing.""" + connector = manager.kv_connector_manager + return connector is not None and not connector.capacity_only def context_moves_history(manager: KVCacheManagerV2) -> bool: - """Whether the manager's context update moves a request's history to the end of each chunk, - raising where that would move it back: under a block reuse policy other than all-reusable, or - when the manager commits no blocks; otherwise the update only commits.""" - from ..kv_cache_manager_v2 import BlockReusePolicy - + """Whether the context update moves the history to each chunk's end, raising where that moves it + back: when no blocks commit or the reuse policy is not all-reusable; else it only commits.""" if not commits_blocks(manager): return True - return manager.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE + global _all_reusable + if _all_reusable is None: + from ..kv_cache_manager_v2 import BlockReusePolicy + + _all_reusable = BlockReusePolicy.ALL_REUSABLE + return manager.block_reuse_policy != _all_reusable def resize_ends_at_chunk(manager: KVCacheManagerV2) -> bool: - """Whether the manager's context resize sets a cache's capacity from the chunk it runs, as a - draft manager's does, so a chunk resumed below the history may leave the capacity below it and - raise there.""" + """Whether the context resize sets capacity from the chunk it runs, as a draft manager's does, + so a chunk resumed below the history may leave the capacity below it and raise.""" return bool(manager.is_draft) def keeps_stale_pages(manager: KVCacheManagerV2) -> bool: - """Whether a block a window leaves behind before the manager commits it keeps its page, which - the commit then stores whole under the block's key: the manager commits blocks without the - minimal snapshot, as under the all-reusable policy.""" + """Whether a block a window leaves behind before its commit keeps its page, which the commit + then stores whole under the block's key: the manager commits blocks without the minimal + snapshot, as under the all-reusable policy.""" return commits_blocks(manager) and not bool(manager.impl.commit_min_snapshot) def prompt_lookahead(manager: KVCacheManagerV2) -> int: - """Prompt tokens past a position the one-model draft the manager was built for reads (Eagle 1, - vanilla MTP its draft length), wherever its layers live, also where reuse matches are not backed - off by them; 0 for a draft reading none or whose read-ahead upstream has not established.""" + """Prompt tokens past a position the manager's one-model draft reads (Eagle 1, vanilla MTP its + draft length), wherever its layers live, also where reuse matches are not backed off by them; 0 + for a draft reading none or whose read-ahead upstream has not established.""" return max(int(manager.draft_prompt_lookahead), int(manager.reuse_match_backoff)) @@ -141,161 +146,260 @@ def attached_lender(manager: KVCacheManagerV2) -> object | None: def install_lender(manager: KVCacheManagerV2, lender: _ManagerHooks) -> None: - """Attach ``lender`` for the manager's life: its free, in-place shrinks, reuse reset and - shutdown then call the lender's ``_on_free``, ``_on_shrink``, ``_on_reset``, ``_on_shutdown`` - and, once the runtime's shutdown has returned, ``_on_caches_closed``.""" - hooks = ("_on_free", "_on_shrink", "_on_reset", "_on_shutdown", "_on_caches_closed") + """Install ``lender``'s hooks (``_lender_hooks._ManagerHooks``) for the manager's lifetime; + ``TypeError`` unless it implements them all.""" + from .._lender_hooks import _ManagerHooks + + hooks = ( + "_on_free", + "_on_shrink", + "_on_resume", + "_parking_row", + "_parked_row", + "_before_reset", + "_on_reset", + "_on_shutdown", + "_on_caches_closed", + "_is_locked", + "_defer_close", + "_keeps_pools", + "_history_cap", + "_page_lock_state", + ) if not isinstance(lender, _ManagerHooks) or any( not callable(getattr(lender, name, None)) for name in hooks ): - raise TypeError("the KV cache lender must implement all five manager lifecycle hooks") + raise TypeError("the KV cache lender must implement all manager hooks") manager._sharing = lender -def layer_configs(manager: KVCacheManagerV2) -> dict[int, object]: - """Internal layer id -> layer config, preferring the manager's config to the runtime's. - A layer without an explicit id uses its position in the config.""" - config = getattr(manager, "kv_cache_manager_py_config", None) - if config is None: - config = manager.impl.init_config - out: dict[int, object] = {} - for position, layer in enumerate(config.layers): - out[int(getattr(layer, "layer_id", position))] = layer - return out +def layer_group_pool_indices(manager: KVCacheManagerV2) -> list[int]: + """The pool-group index of each runtime layer group, in layer-group order.""" + return [int(value) for value in manager.impl.get_life_cycle_pool_group_indices()] -def layer_group_layer_ids(manager: KVCacheManagerV2) -> list[list[int]]: - """Internal layer ids per runtime layer group, in layer-group order.""" - return [[int(layer_id) for layer_id in group] for group in manager.impl.layer_grouping] +def index_buffer(manager: KVCacheManagerV2) -> torch.Tensor: + """The host buffer the manager's caches write their page indices into, through a raw pointer + each keeps until it is detached or closed, whether or not the manager still exists.""" + return manager.host_kv_cache_block_offsets -def pool_group_descriptors(manager: KVCacheManagerV2) -> Sequence[PoolGroupDesc]: - """Runtime pool-group descriptors, in the runtime's order.""" - return manager.impl.pool_group_descs +def index_row(manager: KVCacheManagerV2, request_id: int) -> int | None: + """The row of ``index_buffer``, ``[layer group, row, 0, block]``, into which the request's + active cache writes its beam-0 page indices itself at each page change, on the thread making it; + ``None`` where the manager released the row before the free, to the next request it admits.""" + if request_id in manager._early_freed_index_requests: + return None + return int(manager.index_mapper.get_index(request_id)) * int(manager.max_beam_width) -def declared_global_layer_ids(manager: KVCacheManagerV2) -> dict[int, int] | None: - """Internal-to-global layer ids, or ``None`` without the manager's declaration.""" - getter = getattr(manager, "get_disagg_global_layer_ids", None) - if getter is None: - return None - table: dict[int, int] = {} - for layer_group, layer_ids in enumerate(manager.impl.layer_grouping): - ids = [int(value) for value in getter(layer_group)] - if len(ids) != len(layer_ids): - raise ValueError( - f"layer group {layer_group}: {len(ids)} declared global ids for {len(layer_ids)} layers" - ) - table.update(zip((int(layer_id) for layer_id in layer_ids), ids)) - return table +def runtime(manager: KVCacheManagerV2) -> object: + """Runtime owning the device pools and freeing them on destruction.""" + return manager.impl -def disagg_declaration(manager: KVCacheManagerV2, name: str) -> dict[str, object]: - """Role-keyed layout declarations from the named manager getter.""" - getter = getattr(manager, name, None) - if getter is None: - return {} - return {str(role): value for role, value in dict(getter()).items()} +def context_computed(request: LlmRequest) -> int: + """The prompt tokens the request's context steps computed: its context position, which the + executor moves past each context step once it ran.""" + return int(request.context_current_position) -def role_mapper_kinds(manager: KVCacheManagerV2) -> dict[str, object]: - """The manager's per-role disaggregation mapper declarations.""" - return disagg_declaration(manager, "get_disagg_role_mapper_kinds") +def connector_load_pending(manager: KVCacheManagerV2, request: LlmRequest) -> bool: + """Whether the connector has unfinished off-stream prefix writes into ``request``.""" + connector = manager.kv_connector_manager + return connector is not None and bool(connector.has_pending_load(request)) -def role_layouts(manager: KVCacheManagerV2) -> dict[str, object]: - """The manager's per-role disaggregation layout declarations.""" - return disagg_declaration(manager, "get_disagg_role_layouts") +def connector_reserved(manager: KVCacheManagerV2, request: LlmRequest) -> bool: + """Whether the connector reserved a prefix already in history but not yet loading.""" + connector = manager.kv_connector_manager + return connector is not None and connector.get_prefix_reservation(request) is not None -def ignored_roles(manager: KVCacheManagerV2) -> frozenset[str]: - """Roles omitted by the manager's disaggregation declaration.""" - getter = getattr(manager, "get_disagg_ignored_roles", None) - return frozenset(str(role) for role in (getter() if getter else ())) +def connector_state(manager: KVCacheManagerV2, request: LlmRequest) -> tuple[bool, bool]: + """``(connector_reserved, connector_load_pending)`` of the request in one call.""" + connector = manager.kv_connector_manager + if connector is None: + return False, False + return ( + connector.get_prefix_reservation(request) is not None, + bool(connector.has_pending_load(request)), + ) -def layer_group_pool_indices(manager: KVCacheManagerV2) -> list[int]: - """The pool-group index of each runtime layer group, in layer-group order.""" - return [int(value) for value in manager.impl.get_life_cycle_pool_group_indices()] +def served_prefix_end(request: LlmRequest) -> int: + """End of the prefix served by a connector or a fetch the request resumed from.""" + return int(request.py_served_prefix_end) -def index_buffer(manager: KVCacheManagerV2) -> torch.Tensor: - """The host buffer the manager's caches write their page indices into, through a raw pointer - each keeps until it is detached or closed, whether or not the manager still exists.""" - return manager.host_kv_cache_block_offsets +def stream(manager: KVCacheManagerV2) -> torch.cuda.Stream: + """Execution stream ordering copies after forward writes and before later work, which new page + owners wait for; writers on other streams need explicit ordering. A cache still on the + allocation side stream is ordered behind ``allocation_fence`` instead, until adopted.""" + return manager._stream -def virtual_layers(manager: KVCacheManagerV2) -> tuple[dict[int, tuple[int, int]], int] | None: - """Virtual layers (DeepSeek-V4): internal layer id -> (model layer, attention type value), and - the number of attention types the enum defines. ``None`` for a manager without them.""" - virtual = getattr(manager, "_layer_attn_to_layer_id", None) - if not virtual: - return None - inverse: dict[int, tuple[int, int]] = {} - for (model_layer, attn_type), layer_id in virtual.items(): - inverse[int(layer_id)] = (int(model_layer), int(attn_type.value)) - # Every member of the enum counts, so stages holding different attention types agree. - attn_type_class = type(next(iter(virtual))[1]) - return inverse, max(int(member.value) for member in attn_type_class) + 1 +def cache_stream(manager: KVCacheManagerV2, request_id: int) -> torch.cuda.Stream | None: + """The stream that orders operations on a request's cache during a named grow.""" + import torch + cache_stream_for = getattr(manager, "_cache_stream", None) + if cache_stream_for is not None: + return torch.cuda.ExternalStream( + cache_stream_for(request_id), device=manager._stream.device + ) + return getattr(manager, "_stream", None) -def stream(manager: KVCacheManagerV2) -> torch.cuda.Stream: - """The manager's execution stream: a copy queued on it runs after the forward passes that wrote - the pages and before later work on the stream, which a page's new owner waits for; a writer - off the stream is not ordered after it.""" - return manager._stream + +def allocation_fence(manager: KVCacheManagerV2, request_id: int) -> torch.cuda.Event | None: + """The manager's event after the allocation work queued for the request's cache while that + cache is still on the allocation side stream; ``None`` once the execution stream adopted it.""" + # A cache off the side stream, the common case, costs one set lookup. + if request_id not in manager._on_allocation_stream: + return None + return manager._allocation_fence(request_id) def kv_of(manager: KVCacheManagerV2, request_id: int) -> _KVCache | None: - """The request's runtime cache, or ``None`` when it has none (never had one, freed, shut down). - Caches are compared by identity to see a free or a replaced cache.""" + """Request cache, or ``None`` when absent/freed/shut down; compare identity to detect replacement.""" return manager.kv_cache_map.get(request_id) +_new_tuple = tuple.__new__ + + def cache_state(kv: _KVCache) -> CacheState: """Committed tokens, history length and activity of runtime cache ``kv``, as plain values.""" - return CacheState(int(kv.num_committed_tokens), int(kv.history_length), bool(kv.is_active)) + # Built as the tuple it is, without the NamedTuple's Python-level constructor. + return _new_tuple( + CacheState, (int(kv.num_committed_tokens), int(kv.history_length), bool(kv.is_active)) + ) + + +def cache_extent(kv: _KVCache) -> tuple[CacheState, int]: + """``(cache_state(kv), num_blocks(kv))`` in one call.""" + return _new_tuple( + CacheState, (int(kv.num_committed_tokens), int(kv.history_length), bool(kv.is_active)) + ), int(kv.num_blocks) def has_scratch_reuse(kv: _KVCache) -> bool: - """Whether runtime cache ``kv`` has SWA scratch reuse on: each capacity change then keeps its - history within the scratch rewind of the old capacity, and its window blocks may sit in scratch - slots, which the next chunk overwrites.""" + """Whether ``kv`` has SWA scratch reuse on: capacity changes keep its history within the old + capacity's scratch rewind, and its window blocks may sit in slots the next chunk overwrites.""" return bool(kv.enable_swa_scratch_reuse) def pages(kv: _KVCache, layer_group: int) -> np.ndarray: - """``int64`` page (pool-group slot) of each block ordinal of ``layer_group``, -1 where the block - has no page.""" - return np.fromiter( - kv.get_aggregated_page_indices(layer_group, valid_only=False), dtype=np.int64 - ) + """Read-only view of the page (pool-group slot) each block of an active cache locks, -1 for + none, read on the manager's thread, valid until the cache next changes: copy what must outlive + it. A window block behind the history, its held page maybe on another tier, counts as none.""" + return np.asarray(kv.get_base_page_indices(layer_group))[: kv.num_blocks] + + +def page_list(kv: _KVCache, layer_group: int) -> list[int]: + """The page each block of ``layer_group`` holds, -1 where none, as a list: for an active cache + ``pages``'s, except at window blocks behind the history, whose held pages it gives though it + does not lock them. For few blocks cheaper than that view.""" + return kv.get_aggregated_page_indices(layer_group) + + +def held_pages(kv: _KVCache, layer_group: int, *, recurrent: bool = False) -> int: + """Distinct ordinary, active STATE or scratch slots held by this group, across beams.""" + slots = set() + for beam in range(int(kv.beam_width)): + if recurrent: + slots.add(int(kv.get_ssm_block_base_index(layer_group, beam_id=beam))) + else: + slots.update(kv.get_aggregated_page_indices(layer_group, beam_id=beam, valid_only=True)) + scratch = kv.get_scratch_desc(layer_group) + if scratch is not None: + slots.update(scratch.slot_ids) + slots.discard(-1) + return len(slots) + + +def state_slot(kv: _KVCache, layer_group: int) -> int: + """Active cache's in-place recurrent-state slot for this group, or -1.""" + return int(kv.get_ssm_block_base_index(layer_group)) + + +def prepared_state_slots(manager: KVCacheManagerV2) -> tuple[int, dict[int, int]] | None: + """The layer group of the manager's recurrent state and its record of each request's slot, + written from the cache at each preparation of its resources and dropped at its free, left stale + by a resume that moves the state until the next; ``None`` without one.""" + layer_group = getattr(manager, "ssm_layer_group_id", None) + record = getattr(manager, "_request_id_to_state_index", None) + if layer_group is None or not isinstance(record, dict): + return None + return int(layer_group), record def locked_pages(kv: _KVCache, layer_group: int) -> np.ndarray: - """``int64`` page of each block of an active cache from its locked base page indices, -1 where - the block has none: a window block behind the history keeps at most a held page, which may be - on another tier.""" - return np.array(kv.get_base_page_indices(layer_group)[: kv.num_blocks], dtype=np.int64) + """The same read-only view as ``pages``.""" + return np.asarray(kv.get_base_page_indices(layer_group))[: kv.num_blocks] + + +def committed_by_group(kv: _KVCache, num_layer_groups: int) -> tuple[int, ...]: + """Count a claimed prefix only in groups whose reuse status is complete; a group also counts + tokens committed past its claim endpoint, OPTIONAL/PRIVATE groups included.""" + committed = int(kv.num_committed_tokens) + if committed == 0: + return (0,) * num_layer_groups + statuses = kv.reuse_status # one per layer group, indexed by its id + return tuple( + committed if statuses[lg].complete or committed > int(statuses[lg].endpoint) else 0 + for lg in range(num_layer_groups) + ) -def num_blocks(kv: _KVCache) -> int: - """Block ordinals of runtime cache ``kv``: a shrink drops those past its capacity, with their - pages.""" - return int(kv.num_blocks) +def committed_blocks( + kv: _KVCache, num_layer_groups: int, committed: Sequence[int] | None = None +) -> tuple[int, ...]: + """Per layer group, its whole committed blocks (``committed_by_group``, rounded down); + ``committed``: ``committed_by_group``'s tokens, if read.""" + tpb = int(kv.tokens_per_block) + if committed is None: + committed = committed_by_group(kv, num_layer_groups) + return tuple(tokens // tpb for tokens in committed) -def is_recurrent_layer(layer_config: AttentionLayerConfig | SsmLayerConfig) -> bool: - """Whether a layer config of the runtime describes a recurrent-state layer.""" - from tensorrt_llm.runtime import kv_cache_manager_v2 as runtime +def close_cache(kv: _KVCache) -> None: + """Close the cache, returning its pages to the manager's pools.""" + kv.close() - return isinstance(layer_config, runtime.SsmLayerConfig) + +def commit_held_back(manager: KVCacheManagerV2, kv: _KVCache) -> None: + """Retry a deferred commit; remaining overlapping locks defer it again.""" + manager._commit_held_back(kv) + + +def has_held_commit(manager: KVCacheManagerV2, kv: _KVCache) -> bool: + """Whether the manager holds back a commit of ``kv`` for ``commit_held_back`` to retry; true + for a manager that keeps no record of the commits it holds back.""" + held = getattr(manager, "_held_commits", None) + return held is None or id(kv) in held + + +def drop_held_commit(manager: KVCacheManagerV2, kv: _KVCache, start: int) -> None: + """Drop a commit held back for ``kv`` that would commit computed tokens from ``start`` on, which + a fetch from ``start`` may overwrite; nothing where none is held back.""" + if getattr(manager, "_held_commits", None): + manager._drop_held_commit(kv, start) + + +def is_closed(kv: _KVCache) -> bool: + """Whether runtime cache ``kv`` has closed, after which it writes no page index anywhere.""" + return kv.status == kv.Status.CLOSED + + +def num_blocks(kv: _KVCache) -> int: + """Block ordinals of ``kv``: a shrink drops those past its capacity, with their pages.""" + return int(kv.num_blocks) def stale_range(manager: KVCacheManagerV2, layer_group: int, history: int) -> tuple[int, int]: - """Block ordinals ``[beg, end)`` behind a windowed group's window at ``history``, sinks - excepted. The range adds a block's tokens to the history in 32 bits, so a longer history is read - as ``_MAX_TOKENS`` less one block, which lowers only ``end``.""" + """Block ordinals ``[beg, end)`` behind a windowed group's window at ``history``, sinks aside; + in 32 bits a longer history reads as ``_MAX_TOKENS`` less a block, lowering only ``end``.""" # TODO: with a window near ``_MAX_TOKENS``, a history past that bound lowers ``end`` below # blocks the window has passed, which then count as inside it. longest = _MAX_TOKENS - int(manager.tokens_per_block) @@ -319,9 +423,8 @@ def has_encoder_input(request: LlmRequest) -> bool: def bidirectional_runs(request: LlmRequest) -> list[tuple[int, int]]: - """The runs ``[b, e)`` of multimodal tokens the scheduler keeps within one context chunk: none - unless the request's multimodal data sets ``mm_bidirectional_blocks`` and holds a - ``multimodal_embed_mask_cumsum``, the gate the scheduler uses.""" + """Runs ``[b, e)`` of multimodal tokens the scheduler keeps in one context chunk, by its gate: + none without ``mm_bidirectional_blocks`` set and a ``multimodal_embed_mask_cumsum``.""" data = request.py_multimodal_data if not isinstance(data, dict) or not data.get("mm_bidirectional_blocks", False): return [] @@ -339,11 +442,33 @@ def returns_context_outputs(request: LlmRequest) -> bool: return bool(request.py_return_context_logits) or bool(request.py_additional_outputs) -def prompt_length(request: LlmRequest) -> int: - """The request's prompt length.""" +def prompt_length( + request: LlmRequest, manager: KVCacheManagerV2 | None = None, *, cp_size: int = 1 +) -> int: + """Prompt tokens, using whole-rank ``total_input_len_cp`` under Helix.""" + if cp_size > 1 or (manager is not None and helix(manager)[1] > 1): + return int(request.total_input_len_cp) return int(request.prompt_len) +def context_end(request: LlmRequest, manager: KVCacheManagerV2 | None = None) -> int: + """Where this worker's context ends, at most the prompt, or all of it under ``manager``'s Helix; + a request leaving the prompt's tail to another worker reports the tail's start.""" + prompt = prompt_length(request, manager) + if manager is not None and helix(manager)[1] > 1: + return prompt + end = int(request.context_current_position) + int(request.context_remaining_length) + return min(prompt, end) + + +def helix(manager: KVCacheManagerV2) -> tuple[int, int]: + """Helix ``(rank, size)``, or ``(0, 1)``: block ``b`` spans ``size`` pages of the manager's + ``tokens_per_block`` tokens, rank ``j``'s from ``(b * size + j) * tokens_per_block``.""" + if not getattr(manager, "_has_cp_helix", False): + return 0, 1 + return int(manager._helix_cp_rank), int(manager._helix_cp_size) + + def block_keys( manager: KVCacheManagerV2, request: LlmRequest, kv: _KVCache, count: int ) -> list[bytes]: @@ -389,28 +514,14 @@ def grow( manager: KVCacheManagerV2, request: LlmRequest, kv: _KVCache, position: int, end: int ) -> bool: """Cover ``[0, end)`` and move the history to ``position`` with the manager's own resize, then - run its fresh-page fill (a diagnostic, off unless set) on the pages the resize added alone. + run its fresh-page fill (a diagnostic, off unless set) on the pages the resize added alone; ``False``, the cache unchanged, when pages run out.""" - before = _pages_by_pool(manager, kv) + from .._lender_hooks import fill_fresh_pages_past, fresh_page_record + + before = fresh_page_record(manager, kv) if not manager._resize_for_served_prefix(request, kv, position, end): return False - # TODO: grow writes the fresh-page fill's private record (``_fresh_pages_filled``) in the - # fill's own keying: per request, one int32 page array per pool. - if before is not None: - # The fill takes as fresh the pages its record lacks, and a joint draft pool's own resizes - # keep no record: the pages before the grow become it, so the fill leaves them alone. - manager._fresh_pages_filled[int(request.py_request_id)] = before - manager._fill_fresh_kv_pages(request.py_request_id) + # A joint draft pool's own resizes keep no fill record: the pages before the grow stand for + # it, so the fill leaves them alone. + fill_fresh_pages_past(manager, request.py_request_id, before) return True - - -def _pages_by_pool(manager: KVCacheManagerV2, kv: _KVCache) -> dict[int, np.ndarray] | None: - """Runtime cache ``kv``'s base page of each block, per pool, as the fresh-page fill records - them; ``None`` with the fill off.""" - if manager._fresh_page_fill is None: - return None - count = int(kv.num_blocks) - return { - pool: np.frombuffer(kv.get_base_page_indices(pool), dtype=np.int32, count=count).copy() - for pool in range(int(manager.num_pools)) - } diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_page_locks.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_page_locks.py new file mode 100644 index 000000000000..280a13a8fe98 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_page_locks.py @@ -0,0 +1,559 @@ +# 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. +"""Per-cache block locks on the manager thread. A STATE lock holds block 0 and defers every commit, +which snapshots its private page; writes may lock STATE past committed history.""" + +from __future__ import annotations + +import enum +import types +from typing import TYPE_CHECKING, Mapping, NamedTuple, Optional, Sequence + +import numpy as np + +from . import _manager +from ._types import HeldPages, HoldReason, _integer + +if TYPE_CHECKING: + import weakref + + import torch + + from tensorrt_llm.runtime.kv_cache_manager_v2 import _KVCache + +_NO_BLOCKS = np.zeros(0, dtype=np.int64) +_NO_BLOCKS.flags.writeable = False + + +class LockMode(enum.Enum): + """How a lock holds its blocks: per block, one writer or any number of readers.""" + + READ = "read" + WRITE = "write" + + +class PageLock: + """A lock on blocks of one runtime cache, and the event, if any, its holder waits for before + any access to their pages. Holds at most that one event, dropped once it reports completion.""" + + __slots__ = ("_kv", "_mode", "_blocks", "_ready", "_released", "_holds") + + def __init__( + self, + kv: _KVCache, + mode: LockMode, + blocks: dict[int, np.ndarray], + ready: Optional[torch.cuda.Event], + ) -> None: + self._kv: _KVCache | None = kv + self._mode = mode + self._blocks = blocks + self._ready: torch.cuda.Event | None = ready + self._released = False + # Whether it holds a block, which ``PageLocks.holding`` counts. + self._holds = any(map(len, blocks.values())) + + @property + def mode(self) -> LockMode: + """``LockMode.READ`` or ``LockMode.WRITE``.""" + return self._mode + + @property + def blocks(self) -> Mapping[int, np.ndarray]: + """Layer-group index -> sorted, distinct, read-only ``int64`` block ordinals.""" + return types.MappingProxyType(self._blocks) + + @property + def released(self) -> bool: + """Whether ``PageLocks.unlock`` ended the lock.""" + return self._released + + def is_ready(self) -> bool: + """Whether the lock is held and its event, if any, has completed; asks the event without + waiting, and not again once it reported completion.""" + if self._released: + return False + if self._ready is not None and self._ready.query(): + self._ready = None + return self._ready is None + + +class CheckedBlocks(NamedTuple): + """The blocks ``PageLocks.lock`` would take of a cache in a mode, as it takes them, and the + registry's version when ``PageLocks.check`` found so.""" + + kv: object + mode: LockMode + blocks: dict[int, np.ndarray] + version: int + + +class _CacheLocks: + """One runtime cache's locks, whether it is closed at its last unlock, and the pages per pool + group it added to its hold reason's counts when they were last counted.""" + + __slots__ = ("kv", "locks", "close_at_last_unlock", "held", "counted_in") + + def __init__(self, kv: _KVCache) -> None: + self.kv = kv + self.locks: list[PageLock] = [] + self.close_at_last_unlock = False + self.held: dict[int, int] = {} + # The registry's counts ``held`` is part of, or ``None``. + self.counted_in: dict[int, int] | None = None + + +def _ordinals(layer_group: int, values: Sequence[int]) -> np.ndarray: + """``values`` as sorted, distinct, read-only ``int64`` block ordinals.""" + array = np.asarray(values) + if array.size == 0: + return _NO_BLOCKS + if array.ndim != 1 or not np.issubdtype(array.dtype, np.integer): + raise ValueError( + f"layer group {layer_group}: block ordinals must be integers in one dimension, got " + f"{array.dtype} {array.shape}" + ) + if array.min() < 0: + raise ValueError(f"layer group {layer_group}: block ordinals must be non-negative") + unique = np.unique(array.astype(np.int64)) + unique.flags.writeable = False + return unique + + +class PageLocks: + """The locks on the runtime caches of one manager, which is referenced weakly, and the locked + caches handed over to be closed at their last unlock.""" + + def __init__( + self, + manager: weakref.ref, + recurrent: Sequence[bool] = (), + pool_group_of: Optional[Sequence[int]] = None, + ) -> None: + self._manager_ref = manager + # By ``id`` of the runtime cache, which the entry holds, so the id stays its own. + self._caches: dict[int, _CacheLocks] = {} + # The same entries, for the manager's lender hooks (``_PageLockState.locked_caches``). + self.locked_caches: Mapping[int, object] = self._caches + # Whether a lock can cap a cache's history: only a lock on a windowed layer group's block, + # which the lender's layout may not have. + self.caps_history = True + self.version = 0 + # Unreleased locks that hold a block, of any cache. + self.holding = 0 + # The pool group of each layer group, as given or as the manager answers at the first need. + self._pool_group_of: list[int] | None = ( + None if pool_group_of is None else list(pool_group_of) + ) + # The layer groups of recurrent state, by their flags in layer-group order. + self._recurrent = frozenset(lg for lg, flag in enumerate(recurrent) if flag) + # Pages per pool group of open caches' locks and of caches closing at their last unlock, as + # last counted; held caches locked, unlocked or freed since then wait in ``_dirty``. + self._locked_pages: dict[int, int] = {} + self._deferred_pages: dict[int, int] = {} + self._dirty: dict[int, _CacheLocks] = {} + self._held_version = 0 + self._held: HeldPages | None = None + # Caches kept out of reuse, by ``id``, each entry holding its cache until the cache closes. + self.unreusable: dict[int, object] = {} + + def lock( + self, + kv: _KVCache, + blocks: Mapping[int, Sequence[int]], + mode: LockMode, + level: Optional[int] = None, + *, + ready: Optional[torch.cuda.Event] = None, + ordered: bool = False, + checked: Optional[CheckedBlocks] = None, + committed: Optional[Sequence[int]] = None, + ) -> PageLock: + """Lock an open cache's blocks (``level=None``): per block one writer or many readers, + writes only private uncommitted blocks (STATE row 0 aside); refusals raise, changing + nothing. ``ordered``: sorted, distinct, non-negative, read-only ``int64`` arrays in no + lease's view, taken as is; ``checked``: as ``check`` states.""" + if level is not None: + raise ValueError(f"only level=None, in place, is supported, got level={level!r}") + if not isinstance(mode, LockMode): + raise TypeError(f"mode must be a LockMode, got {type(mode).__name__}") + manager = self._manager_ref() + if manager is None: + raise RuntimeError("the KV cache manager is gone") + if _manager.is_closed(kv): + raise ValueError("the cache is closed") + cache = self._caches.get(id(kv)) + if ( + checked is not None + and checked.kv is kv + and checked.mode is mode + and checked.version == self.version + ): + if cache is not None and cache.close_at_last_unlock: + raise ValueError("the cache was freed and stays open only for the locks it holds") + wanted = checked.blocks + if mode is LockMode.WRITE: + self._check_write_past_committed(kv, wanted, self._pool_groups()) + elif ( + ordered + and checked is None + and cache is None + and (mode is LockMode.READ or (committed is not None and not any(committed))) + ): + # No lock holds the cache and a write has no committed block to start past, so only + # the layer-group count can refuse ordered blocks. + if self._pool_group_of is None: + self._pool_group_of = _manager.layer_group_pool_indices(manager) + wanted = dict(blocks) + if wanted and max(wanted) >= len(self._pool_group_of): + raise ValueError( + f"layer group {max(wanted)}: the manager has {len(self._pool_group_of)} " + "layer groups" + ) + else: + # A checked write may have been invalidated by a lock change and a commit. Its + # caller's committed snapshot no longer establishes which blocks remain private. + live_committed = ( + None + if checked is not None and checked.kv is kv and mode is LockMode.WRITE + else committed + ) + wanted = self._wanted(manager, kv, blocks, mode, ordered, live_committed) + if cache is None: + cache = self._caches[id(kv)] = _CacheLocks(kv) + lock = PageLock(kv, mode, wanted, ready) + cache.locks.append(lock) + self.version += 1 + self.holding += lock._holds + if not cache.close_at_last_unlock: + self._dirty[id(cache)] = cache + return lock + + def lock_fresh( + self, + kv: _KVCache, + blocks: dict[int, np.ndarray], + mode: LockMode, + ready: Optional[torch.cuda.Event], + ) -> PageLock | None: + """``lock`` of a fast lend's ``blocks``, the dict taken as is: ``ordered`` blocks of the + manager's layer groups, a write's past its committed ones, which nothing refuses while no + lock holds ``kv``; else ``None``, locking nothing, for ``lock`` to decide.""" + caches = self._caches + if _manager.is_closed(kv): + return None + if id(kv) in caches: + return None + lock = PageLock(kv, mode, blocks, ready) + cache = caches[id(kv)] = _CacheLocks(kv) + cache.locks.append(lock) + self.version += 1 + self.holding += lock._holds + if lock._holds: + self._dirty[id(cache)] = cache + return lock + + def refusal( + self, kv: _KVCache, blocks: Mapping[int, Sequence[int]], mode: LockMode + ) -> str | None: + """Why ``lock`` would refuse ``blocks`` of ``kv`` in ``mode`` now, or ``None``; changes + nothing.""" + found = self.check(kv, blocks, mode) + return found if isinstance(found, str) else None + + def check( + self, + kv: _KVCache, + blocks: Mapping[int, Sequence[int]], + mode: LockMode, + *, + ordered: bool = False, + committed: Optional[Sequence[int]] = None, + ) -> str | CheckedBlocks: + """Why ``lock`` would refuse ``blocks`` of ``kv`` in ``mode`` now, or, for its ``checked``, + what it would take of them, valid while no lock, unlock or commit comes between; changes + nothing. ``committed``: ``kv``'s ``_manager.committed_by_group``, if read.""" + manager = self._manager_ref() + if manager is None: + return "the KV cache manager is gone" + if _manager.is_closed(kv): + return "the cache is closed" + try: + wanted = self._wanted(manager, kv, blocks, mode, ordered, committed) + except ValueError as error: + return str(error) + return CheckedBlocks(kv, mode, wanted, self.version) + + def _wanted( + self, + manager: object, + kv: _KVCache, + blocks: Mapping[int, Sequence[int]], + mode: LockMode, + ordered: bool = False, + committed: Optional[Sequence[int]] = None, + ) -> dict[int, np.ndarray]: + """``blocks`` as checked ordinals per layer group, ``ordered`` ones taken as they are; + ``ValueError`` for what ``lock`` refuses. ``committed`` as for ``check``.""" + if self._pool_group_of is None: + self._pool_group_of = _manager.layer_group_pool_indices(manager) + cache = self._caches.get(id(kv)) + if cache is not None and cache.close_at_last_unlock: + raise ValueError("the cache was freed and stays open only for the locks it holds") + if ordered: + # Keyed by ``int`` layer-group indices, as the lender keys what it locks. + wanted = dict(blocks) + if wanted and max(wanted) >= len(self._pool_group_of): + raise ValueError( + f"layer group {max(wanted)}: the manager has {len(self._pool_group_of)} " + "layer groups" + ) + else: + wanted = {} + for layer_group, values in blocks.items(): + lg = _integer("a layer-group index", layer_group, 0) + if lg >= len(self._pool_group_of): + raise ValueError( + f"layer group {lg}: the manager has {len(self._pool_group_of)} layer groups" + ) + wanted[lg] = _ordinals(lg, values) + # With no committed token, every block is past the committed whole blocks. + if mode is LockMode.WRITE and (committed is None or any(committed)): + self._check_write_past_committed(kv, wanted, self._pool_group_of, committed) + for held in cache.locks if cache is not None else (): + if LockMode.WRITE not in (mode, held.mode): + continue + for lg, ordinals in wanted.items(): + both = np.intersect1d(ordinals, held._blocks.get(lg, _NO_BLOCKS)) + if len(both): + raise ValueError( + f"layer group {lg}: blocks {both[:8]} are locked for " + f"{held.mode.value}, which a {mode.value} lock conflicts with" + ) + return wanted + + def _check_write_past_committed( + self, + kv: _KVCache, + wanted: Mapping[int, np.ndarray], + pool_group_of: Sequence[int], + committed: Optional[Sequence[int]] = None, + ) -> None: + """Refuse writes to whole committed blocks, reading current counts unless given a snapshot.""" + count = len(pool_group_of) + whole = ( + _manager.committed_blocks(kv, count) + if committed is None + else _manager.committed_blocks(kv, count, committed) + ) + for lg, ordinals in wanted.items(): + held = whole[lg] + earlier = [int(ordinal) for ordinal in ordinals if ordinal < held] + if lg not in self._recurrent and earlier: + raise ValueError( + f"layer group {lg}: a write lock covers only blocks past the {held} " + f"committed whole blocks, got blocks {earlier[:8]}" + ) + + def unlock(self, lock: PageLock) -> bool: + """Release idempotently, retrying a held-back commit before a last unlock closes the cache; + whether it retried one. ``ValueError`` for an unreleased lock not held here.""" + if lock._released: + return False + key = id(lock._kv) + cache = self._caches.get(key) + # A lock equals only itself. + if cache is None or lock not in cache.locks: + raise ValueError("the lock is not held in this registry") + self.version += 1 + self.holding -= lock._holds + lock._released = True + lock._ready = None + lock._kv = None + locks = cache.locks + cache.locks = [] if len(locks) == 1 else [held for held in locks if held is not lock] + if not cache.close_at_last_unlock: + self._dirty[id(cache)] = cache + manager = self._manager_ref() + kv = cache.kv + retried = False + try: + if manager is not None and _manager.has_held_commit(manager, kv): + retried = True + _manager.commit_held_back(manager, kv) + finally: + if not cache.locks: + if cache.close_at_last_unlock: + _manager.close_cache(kv) + self.unreusable.pop(key, None) + del self._caches[key] + self._forget(cache) + return retried + + def is_locked(self, kv: _KVCache, begin: int = 0, end: Optional[int] = None) -> bool: + """Whether a lock holds a ``kv`` block in ``[begin, end)`` (``None``: no end) in any group, + or a row of recurrent state whatever the range.""" + cache = self._caches.get(id(kv)) + if cache is None: + return False + for lock in cache.locks: + for lg, ordinals in lock._blocks.items(): + if lg in self._recurrent and len(ordinals): + return True + first = np.searchsorted(ordinals, begin) + if first < len(ordinals) and (end is None or ordinals[first] < end): + return True + return False + + def holds_any(self) -> bool: + """Whether a lock holds a block of any cache, open or handed over by ``defer_close``.""" + return self.holding > 0 + + def lowest_locked_blocks(self, kv: _KVCache, floors: Sequence[int]) -> dict[int, int]: + """Layer-group index -> the lowest block ordinal of ``kv`` a lock holds there at or past + ``floors[lg]``, for each layer group where a lock holds one.""" + lowest: dict[int, int] = {} + cache = self._caches.get(id(kv)) + for lock in cache.locks if cache is not None else (): + for lg, ordinals in lock._blocks.items(): + at = int(np.searchsorted(ordinals, floors[lg])) + if at < len(ordinals): + lowest[lg] = min(lowest.get(lg, int(ordinals[at])), int(ordinals[at])) + return lowest + + def defer_close(self, kv: _KVCache) -> bool: + """Whether a lock keeps ``kv`` open, held here and closed at its last unlock, its pages + counted as ``DEFERRED_CLOSE`` until then; asked at a free or shutdown instead of a close.""" + if not self.is_locked(kv): + self.unreusable.pop(id(kv), None) + return False + cache = self._caches[id(kv)] + if cache.close_at_last_unlock: + return True + cache.close_at_last_unlock = True + self.version += 1 + self._dirty[id(cache)] = cache + return True + + def holds_open_caches(self) -> bool: + """Whether a cache handed over by ``defer_close`` is still open.""" + return any(cache.close_at_last_unlock for cache in self._caches.values()) + + def keep_out_of_reuse(self, kv: _KVCache) -> None: + """Keep ``kv`` from committing anything more for reuse until it closes.""" + self.unreusable[id(kv)] = kv + + def held_page_counts(self) -> HeldPages: + """The pages the caches here hold by reason, the same value while no count changes; only + caches locked, unlocked or freed since the last call are counted again.""" + if self._dirty: + self._count_dirty() + held = self._held + if held is not None and held.version == self._held_version: + return held + locked = dict(self._locked_pages) + deferred = dict(self._deferred_pages) + unreclaimable = dict(locked) + for group, pages in deferred.items(): + unreclaimable[group] = unreclaimable.get(group, 0) + pages + by_reason = {} + if locked: + by_reason[HoldReason.LOCKED] = types.MappingProxyType(locked) + if deferred: + by_reason[HoldReason.DEFERRED_CLOSE] = types.MappingProxyType(deferred) + # TODO: pages shared with other caches cannot be told apart, so deduct stays an upper bound. + held = self._held = HeldPages( + types.MappingProxyType(deferred), + types.MappingProxyType(unreclaimable), + types.MappingProxyType(by_reason), + self._held_version, + ) + return held + + def _pool_groups(self) -> list[int]: + """The pool group of each layer group, asked of the manager at the first need.""" + if self._pool_group_of is None: + self._pool_group_of = _manager.layer_group_pool_indices(self._manager_ref()) + return self._pool_group_of + + def _locked_pages_of(self, cache: _CacheLocks) -> dict[int, int]: + """Per pool group, the distinct blocks ``cache``'s locks hold; one lock's are distinct.""" + pool_group_of = self._pool_groups() + counts: dict[int, int] = {} + locks = cache.locks + if len(locks) == 1: + for lg, ordinals in locks[0]._blocks.items(): + if len(ordinals): + group = pool_group_of[lg] + counts[group] = counts.get(group, 0) + len(ordinals) + return counts + by_layer_group: dict[int, list[np.ndarray]] = {} + for lock in locks: + for lg, ordinals in lock._blocks.items(): + if len(ordinals): + by_layer_group.setdefault(lg, []).append(ordinals) + for lg, arrays in by_layer_group.items(): + pages = len(arrays[0]) if len(arrays) == 1 else len(np.unique(np.concatenate(arrays))) + group = pool_group_of[lg] + counts[group] = counts.get(group, 0) + pages + return counts + + def _count_dirty(self) -> None: + """Count again the pages of each held cache locked, unlocked or freed since the last count, + a cache closing at its last unlock's as every page it holds.""" + dirty, self._dirty = self._dirty, {} + for cache in dirty.values(): + if cache.close_at_last_unlock: + self._recount(cache, self._deferred_pages, self._deferred_pages_of(cache)) + else: + self._recount(cache, self._locked_pages, self._locked_pages_of(cache)) + + def _forget(self, cache: _CacheLocks) -> None: + """Drop a cache whose last lock ended from the counts, at once where it was counted, so no + closed cache stays referenced here until a query.""" + self._dirty.pop(id(cache), None) + if cache.counted_in is not None: + self._recount(cache, None, {}) + + def _deferred_pages_of(self, cache: _CacheLocks) -> dict[int, int]: + """Per pool group, every page of the runtime cache ``cache`` closes at its last unlock.""" + held: dict[int, int] = {} + for lg, group in enumerate(self._pool_groups()): + pages = _manager.held_pages(cache.kv, lg, recurrent=lg in self._recurrent) + if pages: + held[group] = held.get(group, 0) + pages + return held + + def _recount( + self, cache: _CacheLocks, counts: dict[int, int] | None, held: dict[int, int] + ) -> None: + """Count ``held`` as ``cache``'s pages in ``counts`` (``None``: nowhere) in place of what it + counted before; the version moves if that changes anything.""" + before, was = cache.held, cache.counted_in + if before == held and (was is counts or not held): + return + if was is not None: + for group, pages in before.items(): + left = was[group] - pages + if left: + was[group] = left + else: + del was[group] + if counts is not None: + for group, pages in held.items(): + counts[group] = counts.get(group, 0) + pages + cache.held = held if counts is not None else {} + cache.counted_in = counts if held and counts is not None else None + self._held_version += 1 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_ranges.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_ranges.py index 82cf26e03168..7962620376b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_ranges.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_ranges.py @@ -12,15 +12,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Which block ordinals of a layer group a history still reads under its window, and whether a -resume must stay at or past the history. Other modules call these through the module object -(``_ranges._stale_range(...)``), so each has one binding, which every caller reads.""" +"""Retained and read-window block ranges, and the history required for resumption.""" from __future__ import annotations -from typing import TYPE_CHECKING, Sequence - -import numpy as np +from typing import TYPE_CHECKING from . import _manager @@ -29,15 +25,65 @@ from ._layout import ManagerLayout +def _formula_stale_range( + window: int, sink_blocks: int, tokens_per_block: int, history: int +) -> tuple[int, int]: + """The runtime's stale range of a ``window`` with ``sink_blocks`` at ``history``, read at most + at ``_MAX_TOKENS`` less one block as the manager reads it, in 32-bit integer division: covered + sinks, then the blocks before the window ending at the history's next token.""" + history = min(int(history), _manager._MAX_TOKENS - tokens_per_block) + start = min(-(-history // tokens_per_block), sink_blocks) + behind = history + 1 - window + window_start = behind // tokens_per_block if behind >= 0 else -(-behind // tokens_per_block) + return start, max(start, window_start) + + def _stale_range( manager: KVCacheManagerV2, layout: ManagerLayout, lg: int, history: int ) -> tuple[int, int]: - """Block ordinals ``[beg, end)`` behind layer group ``lg``'s window at ``history``.""" - if layout.windows[lg] is None: + """Block ordinals ``[beg, end)`` behind layer group ``lg``'s window at ``history``: by the + window formula where the layout found the manager follows it, else as the manager answers.""" + window = layout.windows[lg] + if window is None: return 0, 0 + if getattr(layout, "stale_ranges_by_formula", False): + return _formula_stale_range( + window, layout.sink_blocks[lg], layout.tokens_per_block, history + ) return _manager.stale_range(manager, lg, history) +def _probed_histories(window: int, sink_blocks: int, tokens_per_block: int) -> list[int]: + """Histories at which the window formula's terms change, or meet the manager's bounds.""" + tpb = tokens_per_block + probes = set(range(3 * tpb + 2)) + for k in range(8): + for d in range(-2, 3): + probes.add(window + k * tpb + d) + probes.add((k + 1) * tpb + window - 2 + d) + probes.add((sink_blocks + k) * tpb + d) + longest = _manager._MAX_TOKENS - tpb + probes.update((10007, 1 << 20, longest - 1, longest, longest + 1, _manager._MAX_TOKENS)) + return sorted(h for h in probes if 0 <= h <= _manager._MAX_TOKENS) + + +def _stale_ranges_follow_formula(manager: KVCacheManagerV2, layout: ManagerLayout) -> bool: + """Whether the manager's stale range of every windowed layer group equals the window formula at + every probed history, without Helix.""" + if getattr(layout, "cp", (0, 1))[1] != 1: + return False + tpb = int(layout.tokens_per_block) + for lg, window in enumerate(layout.windows): + if window is None: + continue + sinks = int(layout.sink_blocks[lg]) + for history in _probed_histories(int(window), sinks, tpb): + formula = _formula_stale_range(int(window), sinks, tpb, history) + if tuple(_manager.stale_range(manager, lg, history)) != formula: + return False + return True + + def _needed_ranges( manager: KVCacheManagerV2, layout: ManagerLayout, @@ -46,41 +92,64 @@ def _needed_ranges( end_block: int, history: int, ) -> list[tuple[int, int]]: - """Ordinal ranges ``[beg, end)`` of ``lg`` in ``[start_block, end_block)`` that a history of - ``history`` tokens still reads, in order and none empty: the whole range for full attention; - the sinks and the window otherwise. Arithmetic only, so a range of any length costs nothing.""" - # TODO: DeepSeek-V4 widens every window by the draft length ``D`` under any speculative - # decoding, so at 128 tokens per block and ``D >= 2`` a fetch asks for margin rows that a - # publish from a history at least ``D - 1`` tokens past the fetch's end leaves out. + """Nonempty retained ordinal ranges within ``[start_block, end_block)`` at ``history``. + Full attention keeps the whole range; windows keep sinks and their retained suffix.""" stale_beg, stale_end = _stale_range(manager, layout, lg, history) runs = ((start_block, min(end_block, stale_beg)), (max(start_block, stale_end), end_block)) return [(beg, end) for beg, end in runs if end > beg] -def _ordinals(runs: Sequence[tuple[int, int]]) -> np.ndarray: - """``int64`` ordinals of ``runs``, in order.""" - if not runs: - return np.zeros(0, dtype=np.int64) - return np.concatenate([np.arange(beg, end, dtype=np.int64) for beg, end in runs]) +def _read_slack(layout: ManagerLayout, lg: int) -> int: + """Tokens by which layer group ``lg`` keeps a window wider than the one its attention reads.""" + read_windows = getattr(layout, "read_windows", ()) + if not read_windows or layout.windows[lg] is None: + return 0 + return int(layout.windows[lg]) - int(read_windows[lg]) -def _needed_ordinals( - manager: KVCacheManagerV2, - layout: ManagerLayout, - lg: int, - start_block: int, - end_block: int, - history: int, -) -> np.ndarray: - """Ordinals of ``lg`` in ``[start_block, end_block)`` that a history of ``history`` tokens - still reads: all of them for full attention; the sinks and the window otherwise.""" - return _ordinals(_needed_ranges(manager, layout, lg, start_block, end_block, history)) +def _read_stale_range( + manager: KVCacheManagerV2, layout: ManagerLayout, lg: int, history: int +) -> tuple[int, int]: + """Unread ordinals at ``history``, excluding sinks, using the read-window slack.""" + return _stale_range(manager, layout, lg, history + _read_slack(layout, lg)) + + +def _last_history_keeping( + manager: KVCacheManagerV2, layout: ManagerLayout, lg: int, block: int +) -> int: + """The largest history at which layer group ``lg``'s window still keeps ``block``, a block past + its sinks: the stale range's end, which grows with the history, stays at or below it.""" + tpb = int(layout.tokens_per_block) + if getattr(layout, "stale_ranges_by_formula", False) and block >= layout.sink_blocks[lg]: + # Past the sinks the stale end is the window's start, (history + 1 - window) // tpb. + last = (block + 1) * tpb + int(layout.windows[lg]) - 2 + return _manager._MAX_TOKENS if last >= _manager._MAX_TOKENS - tpb else last + low = block * tpb # inside the window: no history this short leaves the block behind + high = low + tpb + int(layout.windows[lg]) + while _stale_range(manager, layout, lg, high)[1] <= block: + if high >= _manager._MAX_TOKENS: + return _manager._MAX_TOKENS + high = min(2 * high, _manager._MAX_TOKENS) + while high - low > 1: + middle = (low + high) // 2 + if _stale_range(manager, layout, lg, middle)[1] <= block: + low = middle + else: + high = middle + return low + + +def _last_history_reading( + manager: KVCacheManagerV2, layout: ManagerLayout, lg: int, block: int +) -> int: + """The largest history at which layer group ``lg``'s attention still reads ``block``: the last + the kept window keeps it, less the slack between the windows.""" + return _last_history_keeping(manager, layout, lg, block) - _read_slack(layout, lg) # TODO: a draft pool also gives up a resume below its history whose chunk would reach it, as # every unchunked prefill's does. def _floor_follows_history(manager: KVCacheManagerV2) -> bool: """Whether a request may resume only at or past its history: a chunk resumed below it may end - below it, where the manager's context update raises, or leave a draft pool's capacity below it, - where the pool's context resize raises.""" + there (the context update raises) or leave a draft pool's capacity there (its resize raises).""" return _manager.context_moves_history(manager) or _manager.resize_ends_at_chunk(manager) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_slots.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_slots.py index 02bac5be3341..d7d0145713f3 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_slots.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_slots.py @@ -30,12 +30,11 @@ from ._layout import ManagerLayout -def slots_per_fetch(layout: ManagerLayout, tokens_per_fetch: int) -> dict[int, int]: - """Slots per device pool group that any whole-block range of at most ``tokens_per_fetch`` tokens can - need: ``ceil(tokens_per_fetch / tpb)`` per layer group, with a window at most its sink blocks plus - ``ceil(window / tpb)``. ``ValueError`` if ``tokens_per_fetch < 1``, ``TypeError`` if not an int.""" +def slots_per_fetch(layout: ManagerLayout, guaranteed_fetch_tokens: int) -> dict[int, int]: + """Slots per device pool group any whole-block range of at most ``guaranteed_fetch_tokens`` + tokens can need; ``ValueError`` if that is below 1, ``TypeError`` if it is not an int.""" tpb = int(layout.tokens_per_block) - target = -(-_integer("tokens_per_fetch", tokens_per_fetch) // tpb) + target = -(-_integer("guaranteed_fetch_tokens", guaranteed_fetch_tokens) // tpb) slots: dict[int, int] = {} for lg, group in enumerate(layout.pool_group_of): blocks = target @@ -49,24 +48,23 @@ def slots_per_fetch(layout: ManagerLayout, tokens_per_fetch: int) -> dict[int, i def slot_counts(layout: ManagerLayout, options: StagingOptions) -> dict[int, int]: - """Slots per device pool group: ``planned_fetches`` fetches, or ``max_bytes`` split by each group's - share of one fetch and rounded down to whole slots. ``ValueError`` if ``max_bytes`` is below - one fetch, so every group holds at least one fetch's slots.""" - slots = slots_per_fetch(layout, options.tokens_per_fetch) + """Slots per device pool group: ``capacity_in_fetches`` fetches, or ``max_bytes`` split by each + group's share of one fetch in whole slots; ``ValueError`` for ``max_bytes`` below a fetch.""" + slots = slots_per_fetch(layout, options.guaranteed_fetch_tokens) fetch_bytes_by_pool_group = {g: slots[g] * int(layout.page_bytes[g]) for g in slots} bytes_per_fetch = sum(fetch_bytes_by_pool_group.values()) if bytes_per_fetch < 1: raise ValueError("the KV cache manager has no pages to stage") - nbytes = options.planned_fetches * bytes_per_fetch + nbytes = options.capacity_in_fetches * bytes_per_fetch if options.max_bytes is not None: if options.max_bytes < bytes_per_fetch: raise ValueError( - f"max_bytes of {options.max_bytes} is below one fetch of {options.tokens_per_fetch} " + f"max_bytes of {options.max_bytes} is below one fetch of {options.guaranteed_fetch_tokens} " f"tokens, which needs {bytes_per_fetch} bytes" ) nbytes = min(nbytes, options.max_bytes) # Each group's share of the bytes is its share of one fetch, so a cap keeps every group's - # slots of one fetch; uncapped, a group gets exactly planned_fetches times its slots. + # slots of one fetch; uncapped, a group gets exactly capacity_in_fetches times its slots. return { g: (nbytes * fetch_bytes_by_pool_group[g] // bytes_per_fetch) // int(layout.page_bytes[g]) for g in slots @@ -148,10 +146,8 @@ def num_slots(self, group: int) -> int: return self._num_slots[group] def check_fits_capacity(self, counts: Mapping[int, int]) -> None: - """Validate total capacity, not current free runs or queue position; queue nothing. - - ``ValueError`` for an unknown group or more slots than a group holds. - """ + """Validate total capacity without queueing; free runs and queue position are ignored. + Raise ValueError for an unknown group or a count exceeding its capacity.""" self._checked_counts(counts) def enqueue(self, counts: Mapping[int, int]) -> int: @@ -189,7 +185,6 @@ def cancel(self, ticket: int) -> None: def give_back(self, runs: SlotRuns) -> None: """Return one lease's slots, all or none. ``ValueError`` for a run not held (freed twice or never granted).""" - # Every run is checked before any is returned, so a bad one returns none. for g, (start, count) in runs.by_pool_group.items(): alloc = self._allocators.get(g) if alloc is None: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_types.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_types.py index 2526e37d0b26..30a174f12e21 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_types.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_types.py @@ -12,30 +12,80 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The lender's public types and protocols. Numpy and the standard library only, so any side can -import them without a live KV cache manager.""" +"""Public lending types and protocols, defined with NumPy and the standard library on the +contract's memory kinds and the shared page layout; their values hold no manager.""" from __future__ import annotations +import enum import numbers +import struct from dataclasses import dataclass -from typing import TYPE_CHECKING, NamedTuple, Optional, Protocol, Sequence, Tuple, runtime_checkable +from typing import ( + TYPE_CHECKING, + List, + Mapping, + NamedTuple, + Optional, + Protocol, + Sequence, + Tuple, + Union, + runtime_checkable, +) import numpy as np +from tensorrt_llm._torch.disaggregation.base.memory import MemoryKind + +from ..page_layout import PageLayout + if TYPE_CHECKING: + import torch + from ...llm_request import LlmRequest -# A block entry's fixed-length opaque object key. +# A named row's fixed-length opaque object key. _OBJECT_KEY_BYTES = 54 +_INT64_MAX = (1 << 63) - 1 + + +class _AllArrived: + """The type of ``ALL_ARRIVED``, which has that one value; it pickles to the same value.""" + + __slots__ = () + _value: Optional["_AllArrived"] = None + + def __new__(cls) -> "_AllArrived": + if cls._value is None: + cls._value = super().__new__(cls) + return cls._value + + def __repr__(self) -> str: + return "ALL_ARRIVED" + + def __reduce__(self) -> str: + return "ALL_ARRIVED" + + +# Every row of a write lease's view arrived: ``Lender.return_writes`` takes it in place of masks. +ALL_ARRIVED = _AllArrived() + +def _immutable(array: np.ndarray) -> np.ndarray: + """``array`` where its memory is immutable bytes, which no one can write or make writable; else + its values in such bytes of their own.""" + root = array + while isinstance(root, np.ndarray) and root.base is not None: + root = root.base + if isinstance(root, bytes): + return array + return np.frombuffer(array.tobytes(), dtype=array.dtype).reshape(array.shape) -def _freeze(array: np.ndarray) -> np.ndarray: - """``array`` marked read-only (a view when it is not already).""" - if array.flags.writeable: - array = array.view() - array.flags.writeable = False - return array + +def _frozen_int64(values: Sequence[int]) -> np.ndarray: + """Build contiguous int64 metadata directly on immutable bytes.""" + return np.frombuffer(struct.pack(f"={len(values)}q", *values), dtype=np.int64) def _integer(name: str, value: object, least: int = 1) -> int: @@ -46,61 +96,98 @@ def _integer(name: str, value: object, least: int = 1) -> int: return int(value) +def _entry_integers(name: str, values: object, shape: Tuple[int, ...]) -> np.ndarray: + """``values`` as an immutable ``int64`` array of ``shape``; ``ValueError`` for another dtype or + shape.""" + array = np.asarray(values) + if not np.issubdtype(array.dtype, np.integer) or array.shape != shape: + raise ValueError( + f"{name} must be integers of shape {shape}, got {array.dtype} {array.shape}" + ) + return _immutable(np.ascontiguousarray(array, dtype=np.int64)) + + +def _region_indices(values: object) -> Tuple[int, ...]: + """``values`` as a tuple of at least one distinct non-negative integer.""" + try: + items = tuple(values) + except TypeError: + raise TypeError( + f"region_indices must be a sequence of integers, got {type(values).__name__}" + ) from None + indices = tuple(_integer(f"region_indices[{i}]", item, 0) for i, item in enumerate(items)) + if not indices: + raise ValueError("region_indices must name at least one region") + if len(set(indices)) != len(indices): + raise ValueError(f"region_indices must be distinct, got {indices}") + return indices + + @dataclass(frozen=True) class StagingOptions: - """Host staging capacity, planned in fetch-sized units. + """Host capacity in fetch units, sized by whole blocks, windows and sinks. - The lender computes the bytes of one fetch from the layer groups and their windows, and a range - of at most ``tokens_per_fetch`` tokens always fits. A capacity budget, not concurrency: a lease - takes one contiguous run of slots per pool group and leases that need slots are granted strictly - in order, so holes that released leases leave can make a lease wait although enough slots are - free in total; a lease with no block entries is ready at its first poll, with no place in line. + Ranges up to ``guaranteed_fetch_tokens`` always fit; larger ones may fit their slot budget. Each + lease needs a contiguous run per pool group, granted FIFO: fragmentation can delay it despite + sufficient total free slots. Empty leases are ready at first poll without joining the queue. Attributes: - tokens_per_fetch: The token count used to size one fetch unit, rounded up to whole blocks - and adjusted for each layer group's window and sinks. Not a lease's range limit: - larger ranges may fit when their required slots fit. - planned_fetches: The capacity multiplier in fetch units, before the byte cap. Not a - limit on concurrent leases or a guarantee that this many fetches run together. - max_bytes: A cap on the staging bytes, or ``None``; ``attach_lender`` raises - ``ValueError`` if it is below one fetch. + guaranteed_fetch_tokens: Tokens per sizing unit, rounded to whole blocks per group. + capacity_in_fetches: Multiplier before the byte cap; no concurrency limit or guarantee. + max_bytes: Optional local byte cap; attach raises ``ValueError`` below a stage-local fetch. Raises: - TypeError: A field is not an integer (``bool`` included). - ValueError: A field is not positive. - """ + TypeError: A field is not an integer, including bool. + ValueError: A field is not positive.""" - tokens_per_fetch: int - planned_fetches: int = 1 + guaranteed_fetch_tokens: int + capacity_in_fetches: int = 1 max_bytes: Optional[int] = None def __post_init__(self) -> None: object.__setattr__( - self, "tokens_per_fetch", _integer("tokens_per_fetch", self.tokens_per_fetch) + self, + "guaranteed_fetch_tokens", + _integer("guaranteed_fetch_tokens", self.guaranteed_fetch_tokens), ) object.__setattr__( - self, "planned_fetches", _integer("planned_fetches", self.planned_fetches) + self, "capacity_in_fetches", _integer("capacity_in_fetches", self.capacity_in_fetches) ) if self.max_bytes is not None: object.__setattr__(self, "max_bytes", _integer("max_bytes", self.max_bytes)) +class RowSelection(enum.Enum): + """Which rows to lend, independently of ``MemoryKind``. + + NAMED selects whole blocks with object keys: reads require committed blocks; writes leave the + final prompt token for computation. ALL includes uncommitted and partial blocks, can reach + prompt end, and has ``object_keys=None``.""" + + NAMED = "named" + ALL = "all" + + @dataclass(frozen=True) class MemoryRegion: - """One memory region containing equal-sized block entries. + """A region of equal-sized slots; one row may span several regions. - A region describes where a backend may access lent memory, not the manager's cache ownership - or its allocation strategy. Its address and extent remain stable while a region hold is open. + Describes backend access, not cache ownership or allocation. Register only accessible memory + kinds; ``RegionHold`` defines address and backing-memory lifetime. Attributes: - layout_label: An opaque layout-derived label, equal for regions laid out alike; not an - allocation identity. - address: The region's base virtual address. - nbytes: The region's extent in bytes. - slot_bytes: The byte extent of one block entry. - slot_count: The number of entries, back to back from ``address``. - layer_group_indices: Manager-local layer-group indices whose entries use this region. - """ + layout_label: Opaque label equal for equal layouts, independent of allocation identity. + address: Base virtual address, positive; ``address + nbytes`` fits in ``int64``. + nbytes: Byte extent, positive. + slot_bytes: Bytes per slot, positive. + slot_count: Back-to-back slots starting at ``address``, positive, within ``nbytes``. + layer_group_indices: Manager-local groups using this region, a tuple of non-negative ints. + memory_kind: Memory kind. + + Raises: + TypeError: ``memory_kind`` is not a ``MemoryKind``, an integer field holds another type or + a bool, or ``layer_group_indices`` is not a tuple. + ValueError: An integer field is out of its range, or the slots exceed ``nbytes``.""" layout_label: str address: int @@ -108,227 +195,330 @@ class MemoryRegion: slot_bytes: int slot_count: int layer_group_indices: Tuple[int, ...] + memory_kind: MemoryKind + + def __post_init__(self) -> None: + if not isinstance(self.memory_kind, MemoryKind): + raise TypeError( + f"memory_kind must be a MemoryKind, got {type(self.memory_kind).__name__}" + ) + for name in ("address", "nbytes", "slot_bytes", "slot_count"): + object.__setattr__(self, name, _integer(name, getattr(self, name))) + if self.slot_bytes * self.slot_count > self.nbytes: + raise ValueError( + f"{self.slot_count} slots of {self.slot_bytes} bytes exceed nbytes {self.nbytes}" + ) + if self.address + self.nbytes > _INT64_MAX: + raise ValueError(f"address {self.address} plus nbytes {self.nbytes} exceeds int64") + indices = self.layer_group_indices + if type(indices) is not tuple: + raise TypeError(f"layer_group_indices must be a tuple, got {type(indices).__name__}") + indices = tuple(_integer("a layer group index", index, 0) for index in indices) + object.__setattr__(self, "layer_group_indices", indices) + + +def _check_disjoint(regions: Sequence[MemoryRegion]) -> None: + """``ValueError`` where two regions of one memory kind overlap; each kind sorted by address.""" + spans_by_kind: dict = {} + for position, region in enumerate(regions): + spans = spans_by_kind.setdefault(region.memory_kind, []) + spans.append((region.address, region.address + region.nbytes, position)) + for kind, spans in spans_by_kind.items(): + spans.sort() + for (_, end, first), (start, _, second) in zip(spans, spans[1:]): + if start < end: + raise ValueError(f"regions {first} and {second} overlap in {kind.name} memory") @dataclass(frozen=True, eq=False) -class LayerGroupBlocks: - """One layer group's lent KV blocks, described by aligned read-only arrays. +class LentLayerGroupRows: + """One layer group's rows in aligned immutable arrays, copied only when not already immutable. - Entry ``i`` describes this rank's KV for logical token block ``block_ordinals[i]`` across the - layers in this group, not a token tensor row or a physical page. Entries need not have - consecutive ordinals or addresses. Placement fields are either all given or all ``None``. + Each row describes this rank's logical page across the group's layers. Recurrent state has + at most one row, ordinal 0, at the range's end. Ordinals and slots need not be consecutive. + ``slots`` and ``region_indices`` are given together or not at all, and object keys are + independent of them. In rows a lender lends, every slot is below the ``slot_count`` of each + ``region_indices`` region: the lender fails a lease rather than lend a slot that is not. - Equal object keys mean interchangeable bytes under the scope and caller obligations of - ``attach_lender``. Keys are opaque identities, not payload checksums. Their format is not API, - but their width, 54 bytes, is. Consumers store them and compare them only for equality. - Changing the format or the layout makes previously stored objects miss, never hit wrongly. + Under ``attach_lender``'s scope and caller duties, equal keys mean interchangeable bytes. + Keys are opaque equality-only identities, not checksums: their width (54 bytes) is API, their + format is not. Format/layout changes make stored objects miss rather than hit incorrectly. Attributes: - layer_group_index: The manager's local index of the layer group. - block_ordinals: ``int64 (n,)`` logical block ordinals; block ``b`` covers token positions - ``[b * tokens_per_block, (b + 1) * tokens_per_block)``. - object_keys: ``uint8 (n, 54)`` complete opaque object keys, distinct from the manager's - 32-byte reuse keys. - addresses: ``int64 (n,)`` virtual addresses; entry ``i`` occupies ``slot_bytes`` bytes at - ``addresses[i]`` in the region identified by ``region_index``. - region_index: The position of this group's region in ``Lender.regions``, not a pool-group - or layer-group identifier. - """ + layer_group_index: Manager-local group index. + page_ordinals: Logical page ordinals, converted to ``int64 (n,)``. Without Helix, page + ``b`` covers ``[b * t, (b + 1) * t)`` for ``t=PageLayout.tokens_per_page``; under Helix, + this rank's in-place page of block ``b`` has ordinal ``b * cp_size + cp_rank``. + object_keys: ``uint8 (n, 54)`` opaque keys, distinct from 32-byte reuse keys, or ``None``. + slots: ``int64 (n,)`` non-negative slots. Row ``i`` concatenates ``slot_bytes`` bytes + from ``address + slots[i] * slot_bytes`` in each ``region_indices`` region. + region_indices: Nonempty, distinct positions in ``Lender.regions``, in row-byte order; not + pool or layer-group IDs. + + Raises: + TypeError: A noninteger index (including bool), or nonsequence ``region_indices``. + ValueError: Wrong array shape, invalid key or slot dtype, negative index or slot, + ``slots`` without ``region_indices`` or the reverse, or empty or duplicate regions. + + Performance (non-normative): + Construction fixes each array's dtype and shape, makes it immutable, and checks that + ``layer_group_index``, ``region_indices`` and ``slots`` hold non-negative integers and that + ``region_indices`` is non-empty and distinct, so holders need not check these again; + ``page_ordinals`` gets no value check. A lender gives all rows of a pool group that lie in + the same regions one ``region_indices`` tuple object. Read the arrays with whole-array + operations and keep references instead of looping over rows or copying defensively; the + lender neither changes nor reuses them.""" layer_group_index: int - block_ordinals: np.ndarray + page_ordinals: np.ndarray object_keys: Optional[np.ndarray] = None - addresses: Optional[np.ndarray] = None - region_index: Optional[int] = None + slots: Optional[np.ndarray] = None + region_indices: Optional[Tuple[int, ...]] = None def __post_init__(self) -> None: object.__setattr__( self, "layer_group_index", _integer("layer_group_index", self.layer_group_index, 0) ) - block_ordinals = _freeze(np.ascontiguousarray(self.block_ordinals, dtype=np.int64)) - if block_ordinals.ndim != 1: - raise ValueError( - f"block_ordinals must be one-dimensional, got shape {block_ordinals.shape}" - ) - object.__setattr__(self, "block_ordinals", block_ordinals) - placed = ( - self.object_keys is not None, - self.addresses is not None, - self.region_index is not None, - ) - if any(placed) and not all(placed): - raise ValueError("object_keys, addresses and region_index are all given or all None") - if not all(placed): - return - object_keys = np.asarray(self.object_keys) - if object_keys.dtype != np.uint8 or object_keys.shape != ( - len(block_ordinals), - _OBJECT_KEY_BYTES, - ): + page_ordinals = _immutable(np.ascontiguousarray(self.page_ordinals, dtype=np.int64)) + if page_ordinals.ndim != 1: raise ValueError( - f"object_keys must be uint8 ({len(block_ordinals)}, {_OBJECT_KEY_BYTES}), " - f"got {object_keys.dtype} {object_keys.shape}" + f"page_ordinals must be one-dimensional, got shape {page_ordinals.shape}" ) - addresses = np.asarray(self.addresses) - if ( - not np.issubdtype(addresses.dtype, np.integer) - or addresses.shape != block_ordinals.shape - ): - raise ValueError( - f"addresses must be integers of shape {block_ordinals.shape}, " - f"got {addresses.dtype} {addresses.shape}" - ) - object.__setattr__(self, "object_keys", _freeze(np.ascontiguousarray(object_keys))) - object.__setattr__( - self, - "addresses", - _freeze(np.ascontiguousarray(addresses, dtype=np.int64)), - ) - object.__setattr__(self, "region_index", _integer("region_index", self.region_index, 0)) + object.__setattr__(self, "page_ordinals", page_ordinals) + if (self.slots is None) != (self.region_indices is None): + raise ValueError("slots and region_indices are given together or not at all") + if self.object_keys is not None: + object_keys = np.asarray(self.object_keys) + if object_keys.dtype != np.uint8 or object_keys.shape != ( + len(page_ordinals), + _OBJECT_KEY_BYTES, + ): + raise ValueError( + f"object_keys must be uint8 ({len(page_ordinals)}, {_OBJECT_KEY_BYTES}), " + f"got {object_keys.dtype} {object_keys.shape}" + ) + object.__setattr__(self, "object_keys", _immutable(np.ascontiguousarray(object_keys))) + if self.slots is not None: + slots = _entry_integers("slots", self.slots, page_ordinals.shape) + if np.any(slots < 0): + raise ValueError(f"slots must be non-negative, got {slots[slots < 0][:8]}") + object.__setattr__(self, "slots", slots) + object.__setattr__(self, "region_indices", _region_indices(self.region_indices)) def __len__(self) -> int: - return int(self.block_ordinals.shape[0]) - - def select(self, mask: np.ndarray) -> LayerGroupBlocks: - """The block entries where ``mask`` is true, keeping the arrays aligned. + return int(self.page_ordinals.shape[0]) - Args: - mask: One boolean per entry, ``True`` for each entry to keep. + def select(self, mask: np.ndarray) -> LentLayerGroupRows: + """Select rows with one boolean per row, preserving order and array alignment. - Returns: - The same layer group and region with the kept entries, in their original order. - - Raises: - ValueError: ``mask`` is not boolean with one value per block entry. - """ + Returns the same group and regions in immutable arrays of its own; raises ``ValueError`` + for a nonboolean or wrong-shape mask.""" mask = np.asarray(mask) - if mask.dtype != np.bool_ or mask.shape != self.block_ordinals.shape: + if mask.dtype != np.bool_ or mask.shape != self.page_ordinals.shape: raise ValueError( - f"mask must be bool of shape {self.block_ordinals.shape}, got {mask.dtype} {mask.shape}" + f"mask must be bool of shape {self.page_ordinals.shape}, got {mask.dtype} {mask.shape}" ) - return LayerGroupBlocks( + return LentLayerGroupRows( self.layer_group_index, - self.block_ordinals[mask], + self.page_ordinals[mask], None if self.object_keys is None else self.object_keys[mask], - None if self.addresses is None else self.addresses[mask], - self.region_index, + None if self.slots is None else self.slots[mask], + self.region_indices, ) + @classmethod + def _unchecked( + cls, + layer_group_index: int, + page_ordinals: np.ndarray, + object_keys: Optional[np.ndarray] = None, + slots: Optional[np.ndarray] = None, + region_indices: Optional[Tuple[int, ...]] = None, + ) -> LentLayerGroupRows: + """The rows the constructor makes of these values, taken as they are, for values that pass + its checks unchanged: an ``int`` index, contiguous ``int64 (n,)`` ordinals and slots and + ``uint8 (n, 54)`` keys on immutable bytes, slots and region indices given together or not + at all, non-negative slots, and a tuple of distinct non-negative ``int`` region indices.""" + rows = object.__new__(cls) + object.__setattr__( + rows, + "__dict__", + { + "layer_group_index": layer_group_index, + "page_ordinals": page_ordinals, + "object_keys": object_keys, + "slots": slots, + "region_indices": region_indices, + }, + ) + return rows + @dataclass(frozen=True, eq=False) class LeaseView: - """What a ready lease lends: at most one block collection per layer group. - - When returned by ``Lease.poll``, its memory is stable and ready for the requested access. - All addresses belong to the lender's regions. Backends may read its immutable arrays on their - own threads and access only the described memory, until they relinquish access and release - the lease. + """Ready rows of a lease of token range ``[start, end)``, at most one collection per layer + group. + + ``Lease.poll`` grants stable access only to the described memory, until the lease is returned. + Of each row's page only the range's tokens are lent: a read keeps them unchanged, a write + delivers them, and the page's other bytes carry no promise. A recurrent-state row holds the + state at ``end``; a lend whose state stands short of its end narrows ``end`` to it. A write + view's rows never overlap. The row collections and immutable arrays remain valid afterward: the + lender never changes or reuses them. Holders must not reassign array headers (shape, dtype or + strides), which would change every holder's view. All rows use the requested memory kind in + ``Lender.regions``. Attributes: - groups: Block collections, at most one per layer group. Tuple positions need not equal - ``layer_group_index``. + groups: Collections whose tuple positions need not equal ``layer_group_index``. + start: The lend's ``start``. + end: The lend's ``end``. Raises: - TypeError: A collection is not a ``LayerGroupBlocks``. - ValueError: A layer group appears twice. - """ + TypeError: A collection is not ``LentLayerGroupRows``, or a bound is not an integer. + ValueError: A layer group repeats, or not ``0 <= start <= end``.""" - groups: Tuple[LayerGroupBlocks, ...] + groups: Tuple[LentLayerGroupRows, ...] + start: int + end: int def __post_init__(self) -> None: groups = tuple(self.groups) seen = set() - for blocks in groups: - if not isinstance(blocks, LayerGroupBlocks): - raise TypeError(f"groups hold LayerGroupBlocks, got {type(blocks).__name__}") - if blocks.layer_group_index in seen: - raise ValueError(f"layer group {blocks.layer_group_index} appears twice") - seen.add(blocks.layer_group_index) + for rows in groups: + if not isinstance(rows, LentLayerGroupRows): + raise TypeError(f"groups hold LentLayerGroupRows, got {type(rows).__name__}") + if rows.layer_group_index in seen: + raise ValueError(f"layer group {rows.layer_group_index} appears twice") + seen.add(rows.layer_group_index) object.__setattr__(self, "groups", groups) + start = _integer("start", self.start, 0) + object.__setattr__(self, "start", start) + object.__setattr__(self, "end", _integer("end", self.end, start)) @property - def num_block_entries(self) -> int: - """Total entries across groups; an ordinal in two groups contributes two entries.""" - return sum(len(blocks) for blocks in self.groups) + def num_rows(self) -> int: + """Total rows across groups; the same ordinal in two groups counts twice.""" + return sum(len(rows) for rows in self.groups) def row_masks(self, fill_value: bool = False) -> Tuple[np.ndarray, ...]: - """Make one writable boolean mask per group, filled with ``fill_value``. - - Mask ``i`` selects entries of ``groups[i]``; its element ``j`` refers to - ``groups[i].block_ordinals[j]``, not to ordinal ``j``. Empty groups keep an empty mask. + """Return writable boolean masks filled with ``fill_value``, in ``groups`` order. - Args: - fill_value: What every block entry's mask value starts as. + Mask ``i[j]`` selects row ``groups[i].page_ordinals[j]``, not ordinal ``j``. Empty + groups retain empty masks. These are the masks ``Lender.return_writes`` accepts.""" + return tuple(np.full(len(rows), bool(fill_value), dtype=bool) for rows in self.groups) - Returns: - One mask per collection, in ``groups`` order: the shape ``mark_arrived`` takes. - """ - return tuple(np.full(len(blocks), bool(fill_value), dtype=bool) for blocks in self.groups) + @classmethod + def _unchecked(cls, groups: Tuple[LentLayerGroupRows, ...], start: int, end: int) -> LeaseView: + """The view the constructor makes of these values, taken as they are, for a tuple of + ``LentLayerGroupRows`` of distinct layer groups and ``int`` bounds ``0 <= start <= end``.""" + view = object.__new__(cls) + object.__setattr__(view, "__dict__", {"groups": groups, "start": start, "end": end}) + return view -# TODO: the caller repeats by hand the manager's own skip past a served prefix (raising the -# history, setting py_served_prefix_end and the prepopulated length), and no public manager method -# takes a position to raise a request's history to. +# TODO: V2 has no public method to resume a request at a specified position. class ResumeBounds(NamedTuple): - """Inclusive local bounds on token positions for resuming after a fetch. - - Resume at a position ``p`` with ``min_position <= p <= max_position``, no lower than the - request's current context position. The bounds include only previously valid cache contents - and fetched block entries that have been incorporated into the request's cache. + """Inclusive local resume interval, counting valid original and incorporated fetched data. - The interval belongs to this rank and this manager. Across ranks, and across a target and - its joint-reuse draft pool sharing the request's cursor, use the largest minimum and smallest - maximum. Fetch into both managers before advancing their shared cursor. An empty effective - interval requires dropping the request's cache in every manager fetched into and computing - from 0, without another fetch into those caches. + Choose ``min_position <= p <= max_position``, no lower than the current context cursor. + Bounds exclude unavailable history, unsupported rewinds and positions strictly inside + bidirectional multimodal runs. The maximum need not be block-aligned; local computation or + an ALL fetch can reach prompt end, while a NAMED fetch leaves its last token for computation. - A minimum excludes positions whose required history is no longer available or which the - manager cannot rewind to. Bounds never allow resuming strictly inside a bidirectional - multimodal run. The maximum need not be a whole-block boundary; it can reach the prompt length - after the remaining prompt tokens have been computed locally. A fetch leaves the last prompt - token for local computation. + Intersect bounds across the request's cache-holding ranks, including pipeline stages and its + own attention data-parallel rank: take the largest minimum and smallest maximum. Intersect + participating joint-reuse target/draft pools and fetch both before advancing their shared + cursor; P/D receives that lend no draft pool exclude its bounds. An empty intersection requires + dropping every fetched cache and computing from zero, without refetching into those caches. Attributes: - max_position: The last permitted resume position. - min_position: The first permitted resume position, unless the context position is higher. - """ + max_position: Last permitted position. + min_position: First permitted position, unless the current cursor is higher.""" max_position: int min_position: int +class HoldReason(enum.Enum): + """Why lending holds a page, as ``HeldPages.by_reason`` counts it; a page has one reason. + + Attributes: + LOCKED: A lease ``lend_read`` or ``lend_write`` lent locks the page in a cache its request + still holds. + DEFERRED_CLOSE: The page belongs to a cache freed while a lease locked some of it, a cache + the request left for a new one included; it stays allocated until that cache's last + lease is returned.""" + + LOCKED = "locked" + DEFERRED_CLOSE = "deferred_close" + + +@dataclass(frozen=True) +class HeldPages: + """Pages lending holds on one rank and pipeline stage, per local pool group (keys of + ``PageLayout.pool_slot_bytes``; zero groups omitted). + + A page counts once per cache however many of that cache's leases lock it; a page several caches + share counts once per cache, so the counts are upper bounds. + + Attributes: + deduct: Pages ``DEFERRED_CLOSE`` holds, including pages the manager may also count against + live requests. They stay with the cache they belong to, already outside its free pages, + so admission does not subtract them again. + unreclaimable: Lease rows ``LOCKED`` holds and pages ``DEFERRED_CLOSE`` holds together. + This omits unleased pages of live locked caches: a scheduler that cannot preempt a + locked cache must reserve its entire allocation, not just these rows. + by_reason: Per ``HoldReason``, the pages it holds; reasons holding none are omitted. + version: Increases whenever a count changes.""" + + deduct: Mapping[int, int] + unreclaimable: Mapping[int, int] + by_reason: Mapping[HoldReason, Mapping[int, int]] + version: int + + +# What ``Lender.return_writes`` takes of a write lease's rows: masks, or ``ALL_ARRIVED``. +Arrived = Union[Sequence[np.ndarray], _AllArrived] + + @runtime_checkable class Lease(Protocol): - """One lent range: poll until the view or a final failure, mark a write once, then release. + """One lent range: poll until it gives its view or fails, then return it. - Methods run on the manager's thread, one call at a time. Backend threads may inspect the - ready view and access its memory, but call no lender, lease or hold method. They report - completion to the holder through their own channel. + Follow ``Lender``'s thread, progress and shutdown duties. Backend threads only inspect ready + views and access the memory they describe, reporting completion through their own channel; + they call no API methods. Caller obligations: - - Poll every open lease each executor iteration, including iterations without new requests. - Do not block waiting for a new request while a lease or backend needs progress. - - Release every lease, failed ones too, only after every backend access has ended. - Dropping an unreleased lease does not release its memory. - - Combine every rank's outcome before deciding for the distributed request. - - A failure is final. A lease failed at the lending call changed no request cache state. - - An exception not listed by the called method is not recoverable through normal lending. - Stop lending, polling, marking and querying bounds through that lender; release leases - and holds once access ends, and treat parked fetches as empty resume intervals. - A failure during manager cleanup can stop a lender permanently. A view already returned - remains held until release; otherwise the lease reports failure. - """ + - Poll a lease until it returns its view or fails, and no further: the view, and access to + the memory it describes, stay valid without another call until the lease is returned. + Combine ranks' outcomes for distributed requests; do not wait for arrivals while + backends need progress. + - Return every lease once backend access ends, failed ones too: a write lease whose view + was returned through ``Lender.return_writes``, which counts its arrived rows into the + cache it was lent from, any other through ``release``. A lease of a cache its request + no longer holds is returned the same way and counts nothing. Dropping a lease does not + return it. + - Failures are final and request-level: compute locally or refuse. Local state, cache, + layers, progress, capacity and reuse reset can fail a lease at the call, without mutation. + - ``poll`` after the lease was returned raises ``RuntimeError`` for caller misuse. + ``RuntimeError`` from ``lend_*`` means the lender stopped or shut down. + - After that error or an unlisted exception, follow ``attach_lender``'s failure duties; + unlisted errors may stop lending permanently, and a stopped lender keeps returned views + held and fails unseen views. + + Performance (non-normative): + Call on the manager thread. ``poll`` does not wait: a DEVICE lease asks its ready event and + any pending connector load, a HOST lease reads what ``Lender.progress`` found and asks its + copy's event at most once per round of it. A returned view costs nothing to keep, and + ``return_resumes_at_end`` is an attribute read.""" def poll(self) -> Optional[LeaseView]: - """Advance pending work and return a stable, ready view. + """Return the view once ready, the same view at every later call, or ``None`` while + pending or failed. - Returns: - The same view on every call once ready; ``None`` while pending, or permanently on failure. - Read views contain readable bytes, and write views are ready for backend writes. - - Raises: - RuntimeError: The lease was released. - """ + Read views hold readable bytes; write views permit backend writes. Raises ``RuntimeError`` + once the lease was returned.""" ... @property @@ -336,187 +526,333 @@ def failure(self) -> Optional[str]: """Why the lease will never be ready, once that is known; for logs only.""" ... - def mark_arrived(self, masks: Sequence[np.ndarray]) -> None: - """Incorporate the indicated arrived block entries into the request's cache. - - Write leases only, once, after ``poll()`` returned the view, before or after release. Each - selected entry must have arrived whole and have no further backend writes in flight. - Returning does not guarantee that incorporation has completed: ``resume_bounds`` reports - when the request may use the result. Entries no longer belonging to the request are ignored; - the holder still marks and releases after the request exits. A write released before its view - was returned is abandoned and needs no mark. - - Args: - masks: One boolean array per entry of ``LeaseView.groups``, in tuple order. Positions - select entries, not logical block ordinals. ``LeaseView.row_masks()`` creates - all-false masks. - - Raises: - RuntimeError: A read lease, a second call, or a call before a view was returned. - ValueError: A wrong number, shape or dtype of masks; nothing is recorded. - """ + @property + def writable(self) -> bool: + """Whether ``lend_write`` lent the lease, so a backend may write its view; fixed for the + lease's life.""" ... def release(self) -> None: - """Confirm that the backend has stopped accessing all lent memory. + """Return the lease once all backend access ended; idempotent and legal in every state. + + Required even on failure; it does not cancel in-flight access. It ends the protection of the + view's memory, not the view: its arrays stay as they were. A write lease released after its + view was returned delivers none of its rows, as if returned through ``Lender.return_writes`` + with no row arrived; released before, it abandons its write. Its unlock retries a commit its + locks held back, and a freed request's cache closes with its last lock; an error in that + work raises from ``release``.""" + ... - Required for every lease, failed ones too; legal in every state and idempotent. Release is - not cancellation of an access still in flight. A ready write still needs ``mark_arrived``, - which may follow release. - """ + @property + def return_resumes_at_end(self) -> bool: + """Whether returning this lease now through ``Lender.return_writes``, every row arrived, + gives bounds whose interval holds its view's ``end``. Only a write lease lent in place whose + view ``poll`` returned can be ``True``; it turns ``False`` for good once its cache shrinks or + is freed, the manager resets reuse or shuts down, or the lender stops.""" ... @runtime_checkable class RegionHold(Protocol): - """A backend's lifetime hold on all memory regions. + """Keep all regions' addresses and backing memory valid, including across shutdown. - Take a hold on the manager's thread before registering ``Lender.regions``. Their addresses - and backing memory remain valid while the hold is open, including across manager shutdown. - Deregister before shutdown, and release the hold only after deregistration is confirmed. - A backend unable to confirm deregistration keeps its hold. Dropping a hold does not release - it. - """ + Take on the manager thread, at construction the construction thread, before registering + accessible regions. Deregister before shutdown; release only after confirmation, retaining + unconfirmed holds. Dropping a hold does not release it.""" def release(self) -> None: - """Confirm deregistration and that the backend cannot access the regions. + """Confirm deregistration and no possible backend access, on the manager thread. - Runs on the manager's thread, is legal in every state, and is idempotent. - """ + Idempotent and legal in every state.""" ... -# TODO: object_keys exist only in ready views, so a flow cannot look up remote hits before it reserves -# pages for them, and a publish lends a windowed group's block entries only for the window at its end: a fork -# at an earlier block or a longer next turn can miss block entries and compute from 0. +# TODO: keys require reserved pages, and publication includes only the end's retained window; +# earlier forks or longer turns can miss rows and require full recomputation. @runtime_checkable class Lender(Protocol): - """Lend a request's KV blocks without transferring cache ownership. + """Lend KV blocks while the manager retains cache ownership. - Lender, lease and hold methods run on the manager's thread, one call at a time. Backend - threads access only a ready lease's described memory. A returned view's addresses are stable - and ready, belong to ``regions``, and remain held until all backend access has ended. + Lender, lease and hold calls run serially on the manager thread. Backend threads access only + ready views and their described memory, held until access ends. See ``Lease`` for failure + handling, ``RegionHold`` for registration lifetime and ``ResumeBounds`` for paired managers. Caller obligations: - - Take a region hold before registration or the first lease's backend access. Deregister - and release the hold only after the backend can no longer access the regions. - - Poll every open lease and query bounds for every waiting request each executor iteration, - even without new requests. - - Release a lease that waits too long, or free its request; progress is not guaranteed for - a pending lease. - - At shutdown stop the executor loop and backends, confirm deregistration, mark arrived - entries and release leases and holds, then shut down the manager. - - Fetch into both a target and its joint-reuse draft pool before advancing their shared - request cursor; combine their bounds as ``ResumeBounds`` specifies. - """ + - Hold regions before registration or first backend access; deregister and release only + after backend access is impossible. + - Call ``progress`` once each executor iteration while a lease is open or a fetch is + unsettled, even without new requests: only it advances leases waiting for their view and + fetches waiting to settle. Pending leases have no progress guarantee: return those + waiting too long or free their requests. + - Pass each lend the event after which no work in flight touches the lent pages, or + ``None`` where none does; the lender records no event of its own. + - Return each write lease whose view was returned through ``return_writes``, saying + whether its arrived rows may enter reuse. Where an entry for the request's current + cache gives no bounds, query ``poll_resume_bounds`` once each iteration, after + ``progress``, until it gives them, and no more then. + - Shutdown: stop executor/backends, confirm deregistration, return leases and release holds, + then shut down the manager. + - Fetch both participating joint-reuse pools before moving their shared cursor; a P/D receive + that lends no draft pool does not intersect that pool's bounds. + + Performance (non-normative): + Call on the manager thread; DEVICE lends pages without copying KV bytes, and HOST queues + its copies on the execution stream without waiting when staging is page-locked. Only + ``progress`` runs a round of lender progress, which has work only while HOST leases or + their fetches are open. A lend's cost grows with the layer groups, and its view's arrays + with the rows lent; ``return_writes`` of a DEVICE lease every row of which arrived, lent + without a rule-specific path, grows little with the rows. ``held_page_counts`` counts again + only the caches whose locks changed, or that were freed while locked, since its previous + call, and costs O(1) when there are none. Constructing a ``torch.cuda.Event`` lets go of + the GIL: give lends events made once and reused, not one made each iteration, and record a + reused event again only after its previous record completed or on the same stream, as a + waiting lease reads only the latest record.""" @property def regions(self) -> Tuple[MemoryRegion, ...]: - """The memory regions to which every ready view's addresses belong. + """Regions containing every ready view, each with one memory kind, fixed for the lender's + life: a lending method a later optional protocol adds gives its own. + + Register only accessible kinds; leases refer only to their requested kind. Registering a + region exposes all of it: any party holding what a backend registered can access the whole + region, as leases and holds bind this process alone. Holds preserve region positions, + addresses and extents. Unsupported kinds are not listed and requests for them raise + ``ValueError``. Equal layout labels mean equal layouts. Regions are well formed and regions + of one memory kind do not overlap. Address ordering, adjacency, a slot's lying within one + allocation and import into another process are not guaranteed: a backend needing one + validates it. See ``RegionHold`` for deregistration before shutdown; open leases/holds can + retain memory beyond it.""" + ... - Each address and extent remains stable while a region hold is open. A layout label is equal - for regions laid out alike. No ordering, adjacency or single-allocation guarantee is made. - A backend requiring one contiguous region must validate that requirement before it starts. + @property + def page_layout(self) -> PageLayout: + """This rank's page layout, fixed for the lender's life: each local layer group's kind, + windows and global layer ids, each buffer's place in a page, the rank's topology, element + type and quantization, and the position in ``regions`` of each pool's region.""" + ... - Deregister before manager shutdown and release the hold after confirmation. Unreleased - leases or holds can retain memory beyond shutdown. - """ + def held_page_counts(self) -> HeldPages: + """Pages lending holds on this rank and pipeline stage, per local pool group (``HeldPages``). + Compare them across ranks and stages to find holds that differ. These diagnostics alone do + not size recovery room: live locked caches may retain unleased pages too. The value stays + the same, ``version`` included, while no count changes, also after the manager's shutdown.""" ... def hold_regions(self) -> RegionHold: - """Return a new lifetime hold on every region. - - The hold does not register regions or grant a lease. Take it on the manager's thread before - registration. A hold taken after manager shutdown is inert. + """Return a new hold on all regions, taken before registration on the manager thread. - Returns: - The hold, open until its ``release``. - """ + It neither registers regions nor grants a lease; it remains open until release and is + inert if taken after manager shutdown.""" ... - def lend_read(self, request: LlmRequest, start: int, end: int) -> Lease: - """Lend committed whole blocks covering token positions ``[start, end)`` for reads. - - Only blocks committed to the manager's prefix-reuse tree are published by name; generated - tokens are not committed. A windowed group lends only retained blocks required at ``end``; - a block without backing memory is omitted. Access only the returned view's entries. - - The lease may fail immediately or while waiting if the request's cache is missing, inactive - or changed, its inputs cannot be named, or the manager stopped or reset reuse. After a reuse - reset the lender cannot publish again: new weights require a new manager and lender under a - scope distinguishing those weights. A previously granted publish may finish. + def lend_read( + self, + request: LlmRequest, + start: int, + end: int, + *, + memory_kind: MemoryKind, + rows: RowSelection, + ready_event: Optional[torch.cuda.Event], + ) -> Lease: + """Lend token range ``[start, end)`` for reads in the requested ``memory_kind``. + + ``rows=NAMED`` publishes the whole blocks the manager has committed, generated or not. It + lends a layer group holding a PRIVATE layer (CSA2's) nothing, its collection empty, and an + OPTIONAL group the prefix claim did not select (CSA2's) none of the claimed blocks. ALL + lends retained blocks, including the uncommitted partial tail, without names. Several beams + lend beam 0's pages, which the manager expands at its first generation step. Windowed groups + include only backed blocks required at ``end``. Access only returned rows. Until this + worker's context is computed, an in-place read of the cache's own pages leaves out each + layer group whose window is shorter than that context (no rows, no lock), so later context + steps move its window on; other groups are lent, and a read once the context is computed + takes them all. This worker's context ends at the request's + ``context_current_position + context_remaining_length``, at most the prompt; under Helix, at + the prompt. + + The lease fails at the call where this rank's cache does not hold the range: NAMED past its + committed tokens; ALL past what it holds or a finished context step computed, or, until this + worker's context is computed, inside a block; with recurrent state, anywhere but the + prompt's end between the context and the first generation step; or needing more than this + rank's capacity. A read in place also fails at the call over rows of an unsettled fetch into + the request, and over rows below the history that a fetch lent or grew the cache over + without delivering them. Missing or inactive caches, unnameable inputs and, by name, a reuse + reset of this rank's manager fail it at the call too; a freed, suspended or replaced cache, + the manager's shutdown and the lender's stop can fail it while it waits. After reset no new + named publication is possible; new weights need a new manager, lender and distinguishing + scope. Granted publishes may finish. Args: - request: The request whose blocks are read. - start: Inclusive first token, a multiple of ``tokens_per_block``. - end: Exclusive end token, a multiple of ``tokens_per_block``, at most the committed tokens. + request: Request to read; the lease binds to the cache it holds now. + start: Nonnegative whole-block boundary. + end: Exclusive end, at least ``start``; a whole-block boundary for NAMED. + memory_kind: Required ``MemoryKind``. + rows: Required ``RowSelection``. + ready_event: Required: an event with ``query()``, and ``wait()`` where rows are not the + cache's own pages, such as a ``torch.cuda.Event``, recorded after every piece of + work in flight that writes the range's pages, or ``None`` where none is in flight. + ``poll`` gives no view before it completes. Such work includes the context steps + that computed the range and any later move or write of the pages, such as a commit, + a resume or a fetch into them; where these run in order, an event recorded after the + last covers them all. Until the manager adopts the request's cache into its + execution order, a lease in place also waits for the manager's allocation fence, so + the event need only cover work in that execution order. Returns: - The lease, possibly already failed. + A lease, possibly already failed. Raises: - ValueError: Negative, reversed or non-whole-block bounds, an end past committed tokens, - or a range exceeding the attached implementation's capacity. - """ + TypeError: Invalid ``memory_kind``, ``rows`` or ``ready_event`` type. + ValueError: Unsupported kind/selection or invalid bounds above, alike on every rank. + RuntimeError: The manager shut down, the lender stopped after an error, or a manager + or CUDA operation failed. Stop lending through it after this error.""" ... - def lend_write(self, request: LlmRequest, start: int, end: int) -> Lease: - """Lend writable memory for fetched blocks covering token positions ``[start, end)``. + def lend_write( + self, + request: LlmRequest, + start: int, + end: int, + *, + memory_kind: MemoryKind, + rows: RowSelection, + ready_event: Optional[torch.cuda.Event], + ) -> Lease: + """Lend writable ``memory_kind`` for fetched token range ``[start, end)``. + + ALL includes the unnamed partial tail; windows lend retained rows required at ``end``. Each + group skips its committed whole blocks. Once committed tokens pass the claimed reuse prefix, + every group counts all of them, previously incomplete groups included; a group committed + past ``end`` lends no rows. Before that, unselected OPTIONAL and PRIVATE groups hold none of + the claimed prefix, so fetching ALL from 0 fills their gaps. NAMED always omits groups + holding PRIVATE layers. Non-REQUIRED attention groups need windows. Multiple beams lend + beam 0; the manager expands the rest at the first generation step. + Window margins beyond ``LayerGroupDesc.read_window`` may stay unwritten without lowering + ``poll_resume_bounds``; the lender then keeps the cache out of reuse. Follow + ``return_writes`` for other permitted omissions. Extend a prefix only after every rank's + prior fetch settles. + + Do not schedule an unsettled fetch. Original history must be valid; computed tokens below + ``start`` can contribute, but retained windows constrain resumption and bytes from ``start`` + may be overwritten. Resume from settled bounds, not the requested extent. Abandoned rows + contribute nothing; failure after this call need not undo cache changes. + + Failures include ``lend_read``'s reasons, exhaustion, unsettled prior fetches, incompatible + history/reuse, unnamed inputs, split bidirectional runs and unrestorable per-token outputs. + Rank-local failures also include an ALL end inside a block unless at prompt end or declared + history (a CSA2 remote tail's start), a start inside an externally served prefix, a + recurrent end short of the prompt, or a windowed end below history. These failures precede + mutation. - A windowed group lends the retained blocks required at ``end``. The holder marks entries - that arrived whole and releases after backend access ends. Consecutive settled fetches can - contribute to one prefix; lend the next segment only after every rank reports settled bounds. - Abandoning a lease contributes none of its entries. + Args: + request: Request fetched into; the lease binds to the cache it holds now. + start: Nonnegative whole-block boundary. + end: Exclusive end, at least ``start``. NAMED ends at a whole block before the final + prompt token; ALL ends no later than the prompt's end. + memory_kind: Required ``MemoryKind``. + rows: Required ``RowSelection``. + ready_event: Required: an event with ``query()``, and ``wait()`` where rows are not the + cache's own pages, such as a ``torch.cuda.Event``, recorded after every piece of + work in flight that still uses or resets the range's pages, such as the request's + admission and a reset of its recurrent state, or ``None`` where none is in + flight. A NAMED grow also requires ``wait()`` if the event is not yet complete, so + the grow can follow that work in the cache's execution order. ``poll`` gives no + view of the range's pages before it completes; a view of other memory may come + first, and the rows counted from it reach the pages after it. + Until the manager adopts the request's cache into its execution order, a lease in + place also waits for the manager's allocation fence, so the event need only cover + work in that execution order. - The request must not execute while a fetch is unsettled. Previously computed tokens below - ``start`` may contribute to the prefix; retained windows may limit where it can resume. - Contents from ``start`` onward may be overwritten. Fetch only into a cache whose history is - backed by valid data. Use ``resume_bounds``, not the fetch extent, to decide where to continue. - Failure after the lending call need not undo cache changes. + Returns: + A lease, possibly already failed. - The lease may fail for the reasons in ``lend_read``, resource exhaustion, an unsettled prior - fetch, incompatible cache history or reuse, unnamed inputs, a bidirectional multimodal run - split by ``end``, or requested per-token outputs that a fetch cannot restore. + Raises: + TypeError: Invalid ``memory_kind``, ``rows`` or ``ready_event`` type. + ValueError: Unsupported kind/selection or invalid bounds above, alike on every rank. + RuntimeError: As ``lend_read``; stop lending through it after this error.""" + ... + + def return_writes( + self, entries: Sequence[Tuple[Lease, Arrived, bool]] + ) -> List[Optional[ResumeBounds]]: + """Return write leases whose views were returned, counting the rows that arrived, and give + each lease's resume bounds. + + Each entry is ``(lease, arrived, reusable)``: a write lease this lender lent whose view + ``poll`` returned; ``ALL_ARRIVED``, or one boolean mask per ``LeaseView.groups`` + collection, in tuple order, selecting rows rather than ordinals (``LeaseView.row_masks`` + gives masks to fill); and whether the arrived rows may enter reuse. A read-only mask of + dtype ``bool`` whose type is exactly ``numpy.ndarray`` is kept as it is, and its giver + writes neither it nor memory it shares afterwards; any other mask is copied. + + Call only after the transfer's outcome is settled on every participating rank and all + backend access to the returned views has ended. An arrival notification alone establishes + neither condition. Returning a lease may let the cache consume its returned memory + immediately; ``reusable`` does not replace either completion requirement. + + An entry counts its arrived rows into the cache the lease was lent from, then returns the + lease as ``Lease.release`` does. Arrived rows must have no writes in flight. For ALL, every + ``transferred=True`` buffer's bytes the request reads must have arrived; buffers may stay + unwritten if not transferred, read only by one-model draft layers absent at the peer, or + overwritten before the request reads them, such as CSA2 remote-tail decoder rows, and such + an entry passes ``reusable=False``. NAMED requires the whole object. ``reusable=False`` + keeps the cache from committing anything more for reuse until it closes; so does a lent + window-margin row that did not arrive, and a later ``True`` undoes neither. Rows of a cache + the request no longer holds count nothing; still return the lease, also after the request + exits. Only counted rows make data usable. + + Every entry is checked before any is returned. The validation errors listed below change + nothing: correct the entries and retry; lending, ``progress`` and ``poll_resume_bounds`` + remain permitted. Any error once returning entries begins raises after all entries' + returns have been attempted. It permits cleanup only: stop lending, progress, lease + polling and bounds queries, even for other entries returned successfully. Follow + ``attach_lender``'s failure duties: finish access, release remaining leases and holds, + drop affected fetched caches and compute from zero. Exception + type alone does not distinguish validation from operation failure; if unsure, use cleanup + only. Returning entries need not undo cache changes made before an error. Args: - request: The request fetched into. - start: Inclusive first token, a whole-block boundary no earlier than the committed tokens - rounded down to whole blocks. - end: Exclusive whole-block boundary, no later than the whole blocks before the last prompt - token and, for a windowed cache, no earlier than its retained history. + entries: The ``(lease, arrived, reusable)`` entries, each lease at most once. Returns: - The lease, possibly already failed. + Per entry, in order, the bounds of the lease's cache as ``poll_resume_bounds`` gives + them once the entries are returned; ``None`` where they are not settled yet, as when + another fetch into the request is pending or the entry's rows are not counted yet (then + query ``poll_resume_bounds``), or where that cache is no longer the request's or the + request has none (no bounds apply). Raises: - ValueError: Negative, reversed or non-whole-block bounds, a start within committed whole - blocks, an end past the last eligible prompt block, an end below a windowed cache's - retained history, or a range exceeding the attached implementation's capacity. - Validation precedes changes. - """ + TypeError: Entry validation: an entry is not a ``(lease, arrived, reusable)`` triple, + or ``reusable`` is not a ``bool``. + ValueError: Entry validation: a lease this lender did not lend, or masks of the wrong + count, shape or dtype. + RuntimeError: Entry validation: a read lease, a lease returned before, in this call or + earlier, or one whose view ``poll`` did not return (caller misuse).""" ... - def resume_bounds(self, request: LlmRequest) -> Optional[ResumeBounds]: - """Advance work and return local bounds once the request's fetches have settled. - - Query only while retained history not advanced by this lender is backed by valid cache - contents. The bounds cannot detect a history that runs ahead of its computed data. + def progress(self) -> None: + """Advance the lender's pending work: leases waiting for their view, fetches waiting to + settle, and what returned leases leave to free. - Only entries incorporated into the request's cache count as fetched. Marking arrival alone - does not make data usable. Abandoned or failed entries contribute nothing; an overwritten - partially committed block can make the interval empty. Shrinking a cache invalidates delivered - entries beyond its new capacity permanently: growing it again does not restore their contents. + Call it once each executor iteration while a lease is open or a fetch is unsettled; + ``Lease.poll`` and ``poll_resume_bounds`` report what it found.""" + ... - Combine rank and target/draft intervals as ``ResumeBounds`` specifies. If the effective - interval is empty, drop the request's cache in every manager fetched into and compute from 0. + def poll_resume_bounds(self, request: LlmRequest) -> Optional[ResumeBounds]: + """Return local bounds of the request's current cache once all its fetches settle, as + ``progress`` last found. - Args: - request: The request fetched into. + History not advanced by this lender must already contain valid data: bounds cannot detect + fabricated history. Only counted rows count; a fetch whose rows were not counted yet is + unsettled. Failed or abandoned rows, and rows of the request's earlier caches, count + nothing, and overwriting a partially committed block may empty the interval. Shrinking + permanently invalidates deliveries beyond new capacity; regrowing does not restore them. + Apply ``ResumeBounds``'s intersection/drop-and-recompute rule. Returns: - ``None`` while a fetch is unsettled, otherwise the inclusive permitted resume interval. + ``None`` while unsettled, otherwise inclusive bounds for ``request``; an empty + interval, its maximum below its minimum, once the lender stopped for good. Raises: - ValueError: The request has no cache, including after manager shutdown. - """ + ValueError: No request cache, including after manager shutdown.""" ... diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e0502b4cd181..fa5703029851 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -97,6 +97,7 @@ hard_kill_on_rank_crash, propagate_hard_kill, start_rank_crash_kill_watchdog) from .hang_diagnostics import create_executor_hang_diagnostics +from .kv_cache import _lender_hooks from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from .kv_cache.mamba_cache_manager import (BaseMambaCacheManager, MixedMambaHybridCacheManager) @@ -5689,9 +5690,9 @@ def _sync_and_process_resource_governor_queue(self): def _can_pause_for_rebalance(self) -> bool: """Gate KV pool rebalance to the cases the hook supports. - Scope: no in-flight disagg transfer, no KV connector, no drafter, not - during warmup or shutdown. Honors the ``enable_kv_pool_rebalance`` - opt-in flag (default off). + Scope: no in-flight disagg transfer, no KV connector, no attached KV + cache lender, no drafter, not during warmup or shutdown. Honors the + ``enable_kv_pool_rebalance`` opt-in flag (default off). Beam search *is* supported. It needs no special handling on this side: ``KvCache::suspend`` / ``resume`` walk every beam already @@ -5726,6 +5727,12 @@ def _can_pause_for_rebalance(self) -> bool: return False if self.kv_cache_transceiver is not None: return False + if self.kv_cache_manager._sharing is not None: + # A lender exposes fixed-address device regions for its lifetime; + # adjust() may relocate those pools. The attach contract requires + # every cache-holding rank to attach or none, keeping this gate + # identical across ranks before the rebalance collectives. + return False if self.kv_connector_manager is not None: # Rebalance suspends every active request and runs a defragmenting # migration, which reassigns slot_id -- the same field a tier @@ -10100,9 +10107,15 @@ def _remove_inflight_ids(self, scheduled_requests: ScheduledRequests): scheduled_requests.added_inflight_req_ids = [] def reset_prefix_cache(self): - self.kv_cache_manager.reset_reuse_state() + managers = [self.kv_cache_manager] if self.enable_joint_kv_cache_reuse: - self.draft_kv_cache_manager.reset_reuse_state() + managers.append(self.draft_kv_cache_manager) + # Every lender must allow the reset before either reuse tree changes. + for manager in managers: + if isinstance(manager, KVCacheManagerV2): + _lender_hooks.before_reset(manager) + for manager in managers: + manager.reset_reuse_state() # R3: invalidate cached routes together with the KV reuse state. Guard the # engine lookup too -- minimal executors (unit tests) may have no engine. route_capture = getattr(getattr(self, "model_engine", None), diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 9c17e72d3fb3..ec1d50dccb0c 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -21,6 +21,7 @@ from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy from tensorrt_llm.logger import logger +from ..kv_cache import _lender_hooks from ..llm_request import ( LlmRequest, LlmRequestState, @@ -1688,9 +1689,11 @@ def _try_schedule_generation( # GPU pages so other requests can resume(). # Skip if already suspended — suspending again is a no-op # that frees no pages. - if self.kv_cache_manager.is_request_active( - req.py_request_id - ) and not self._has_pending_connector_load(req): + if ( + self.kv_cache_manager.is_request_active(req.py_request_id) + and not self._has_pending_connector_load(req) + and not self._holds_page_locks(req) + ): logger.debug( f"[V2Scheduler] Self-evicting request {req.py_request_id} " f"(state={req.state.name}) to free GPU pages" @@ -1746,6 +1749,8 @@ def _suspend_request(self, req: LlmRequest) -> None: remains "active" on device, which could cause ensure_batch to fail if it needs to load a different adapter into a full cache. """ + if self._holds_page_locks(req): + return self._clear_request_runtime_state(req) self.kv_cache_manager.suspend_request(req) if self.draft_kv_cache_manager is not None: @@ -1879,6 +1884,7 @@ def _detect_deadlock( if ( any(self._has_pending_connector_load(r) for r in active_requests) or self.kv_cache_manager.has_pending_preemption() + or self._holds_page_locks(None) ): self._stalled_schedules = 0 return @@ -2004,6 +2010,9 @@ def _has_pending_connector_load(self, req: LlmRequest) -> bool: connector = self.kv_cache_manager.kv_connector_manager return connector is not None and connector.has_pending_load(req) + def _holds_page_locks(self, req: LlmRequest | None) -> bool: + return _lender_hooks.scheduler_holds_page_locks(self, req) + def _is_evictable(self, req: LlmRequest, inflight_request_ids: set[int]) -> bool: """A started request whose KV cache is still active on GPU. @@ -2014,7 +2023,7 @@ def _is_evictable(self, req: LlmRequest, inflight_request_ids: set[int]) -> bool return False if not self._is_started_request(req): return False - if self._has_pending_connector_load(req): + if self._has_pending_connector_load(req) or self._holds_page_locks(req): return False return self.kv_cache_manager.is_request_active(req.py_request_id) @@ -2023,7 +2032,7 @@ def _is_recompute_pause_candidate( ) -> bool: if req.request_id in inflight_request_ids: return False - if self._has_pending_connector_load(req): + if self._has_pending_connector_load(req) or self._holds_page_locks(req): return False # is_generation_in_progress_state also includes GENERATION_TO_COMPLETE, # which is outside the schedulable range and may still be finalizing. diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index b597efc96bc0..30d4ab2a338c 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -33,6 +33,7 @@ l0_b200: backend: pytorch tests: # ------------- PyTorch tests --------------- + - unittest/_torch/executor/kv_cache/sharing - llmapi/test_llm_api_pytorch_t5.py::test_t5_v2_decode_beyond_encoder_length[False] - llmapi/test_llm_api_pytorch_t5.py::test_t5_v2_decode_beyond_encoder_length[True] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_block_reuse_is_disabled_for_v2 diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 8ec539b91268..417220051c5f 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -34,7 +34,9 @@ l0_b300: # rebind mpi4py's serializer at import time, which applies process-wide on # every stage that collects them, including where the tests themselves skip. # The sparse/dsa/test_cute_dsl_* files ignored below keep their own entries (pre-existing). - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py --ignore=unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_msa_proxy_split.py TIMEOUT (90) + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py --ignore=unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_msa_proxy_split.py --ignore=unittest/_torch/attention/test_prims_ts_attention_backend.py TIMEOUT (90) + # Give the cold-JIT PrimTS suite its own 60-minute test budget. + - unittest/_torch/attention/test_prims_ts_attention_backend.py # Run cold-JIT MSA coverage separately from the attention aggregate. # Cold-JIT runs take about 25 minutes; allow a 45-minute budget. - unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py TIMEOUT (45) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index e03bb4e02f5b..795bd621404d 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -1942,6 +1942,14 @@ def test_disagg_page_table_accepts_deepseek_v4_roles(self): for role_name in pool_view.pool_role: assert role_name.startswith("deepseek_v4_"), role_name assert saw_pool_view + # The manager's own page table moves the same roles, each split by head. + transferred = [ + b for b in cache_manager.get_page_table().page_layout.buffers if b.transferred + ] + assert transferred + for buffer in transferred: + assert buffer.role.startswith("deepseek_v4_"), buffer.role + assert buffer.split, buffer.role finally: cache_manager.shutdown() diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py index c2b78b4996e0..18c102726b1c 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py @@ -288,12 +288,12 @@ def _logical_page(view: torch.Tensor) -> tuple: @pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) def test_transfer_views_do_not_depend_on_the_kv_head_split(dtype: DataType) -> None: """Disaggregated peers compare the replicated side roles each layer registers and copy - their pages byte for byte, so roles, mapper kinds and every side view's page (shape, - dtype, bytes) must not depend on how many KV heads a rank holds; where pages sit in - their pools may. At TP2 and TEP4 an FP8 K/V page (1 head x 256 x 32 tokens) is the - size of an index page; under NVFP4 the same holds at TP1 and attention DP (2 heads x - 128 bytes x 32 tokens). Both pack on every layout; bf16, which can never collide, on - none.""" + their pages byte for byte, so roles, mapper kinds, the tensors the page table declares and + every side view's page (shape, dtype, bytes) must not depend on how many KV heads a rank + holds; where pages sit in their pools may. At TP2 and TEP4 an FP8 K/V page (1 head x 256 x + 32 tokens) is the size of an index page; under NVFP4 the same holds at TP1 and attention DP + (2 heads x 128 bytes x 32 tokens). Both pack on every layout; bf16, which can never collide, + on none.""" seen = {} for name in _LAYOUTS: cm = _manager(dtype, draft=False, layout=name) @@ -308,7 +308,15 @@ def test_transfer_views_do_not_depend_on_the_kv_head_split(dtype: DataType) -> N views = {layer: cm.get_index_k_buffer(layer) for layer in _QSA_LAYERS} views["position"] = cm.get_qsa_position_buffer() logical = {key: _logical_page(view) for key, view in views.items()} - seen[name] = (roles, logical, cm.get_disagg_role_mapper_kinds()) + # How each side buffer spreads over ranks, as the page table declares it. + tensors = { + (layer, role): cm._page_buffer_tensors( + cm.layer_offsets[layer], role, size, _TPB, False + ) + for layer, layer_roles in roles.items() + for role, size in layer_roles + } + seen[name] = (roles, logical, cm.get_disagg_role_mapper_kinds(), tensors) finally: cm.shutdown() first, *rest = seen.values() diff --git a/tests/unittest/_torch/disaggregation/blob/cpu_fixtures.py b/tests/unittest/_torch/disaggregation/blob/cpu_fixtures.py index f130e5dd8396..b99343b68c3f 100644 --- a/tests/unittest/_torch/disaggregation/blob/cpu_fixtures.py +++ b/tests/unittest/_torch/disaggregation/blob/cpu_fixtures.py @@ -136,8 +136,8 @@ def slot(self, index: int) -> int: return self.address + index * self.slot_bytes def slots(self, *indexes: int) -> np.ndarray: - """The addresses of slots ``indexes``, as a view's ``addresses``.""" - return np.array([self.slot(index) for index in indexes], dtype=np.int64) + """Slot indices, as a view's ``slots``.""" + return np.array(indexes, dtype=np.int64) def read(self, index: int) -> bytes: return self.buffer.read(index * self.slot_bytes, self.slot_bytes) @@ -155,20 +155,21 @@ def __init__( self, layer_group_index: int, object_keys: np.ndarray | None, - addresses: np.ndarray | None, - region_index: int | None, + slots: np.ndarray | None, + region_indices: tuple[int, ...] | None, rows: int | None = None, ) -> None: self.layer_group_index = layer_group_index self.object_keys = object_keys - self.addresses = addresses - self.region_index = region_index + self.slots = slots + self.region_indices = region_indices + self.page_ordinals = np.arange(0 if slots is None else len(slots), dtype=np.int64) self.rows = rows def __len__(self) -> int: if self.rows is not None: return self.rows - return 0 if self.addresses is None else len(self.addresses) + return 0 if self.slots is None else len(self.slots) class CpuView(_Sealable): @@ -483,7 +484,7 @@ def __init__( def rows(self, names: Sequence[str], layer_group_index: int = 0) -> Rows: keys = np.stack([object_key(name) for name in names]) if names else _no_keys() slots = self.region.slots(*(self.names.index(name) for name in names)) - return Rows(layer_group_index, keys, slots, 0) + return Rows(layer_group_index, keys, slots, (0,)) def view(self, *names: str) -> CpuView: return CpuView([self.rows(names)]) diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_backend.py b/tests/unittest/_torch/disaggregation/blob/test_blob_backend.py index 5359a43c0f78..3da14d5f3546 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_blob_backend.py +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_backend.py @@ -18,6 +18,7 @@ import threading import time from collections.abc import Callable +from contextlib import closing from dataclasses import replace import numpy as np @@ -45,6 +46,7 @@ from tensorrt_llm._torch.disaggregation.base.backend import ( Capabilities, InvalidSubmission, + LayerGroupRows, MemoryKind, SubmissionRejected, ) @@ -1055,8 +1057,8 @@ def _view_with(endpoint: Endpoint, **changes) -> CpuView: fields = { "layer_group_index": 0, "object_keys": rows.object_keys, - "addresses": rows.addresses, - "region_index": 0, + "slots": rows.slots, + "region_indices": (0,), } fields.update(changes) return CpuView([Rows(**fields)]) @@ -1064,18 +1066,27 @@ def _view_with(endpoint: Endpoint, **changes) -> CpuView: INVALID_SUBMISSIONS: dict[str, Callable[[Endpoint], tuple[CpuView, object]]] = { "hint": lambda e: (e.view("A"), "peer-0"), - "no_keys": lambda e: (_view_with(e, object_keys=None, addresses=None, region_index=None), None), + "no_keys": lambda e: (_view_with(e, object_keys=None, slots=None, region_indices=None), None), + "no_slots": lambda e: (_view_with(e, slots=None), None), + "no_regions": lambda e: (_view_with(e, region_indices=None), None), "key_dtype": lambda e: (_view_with(e, object_keys=np.zeros((1, 4), dtype=np.int32)), None), "key_width": lambda e: (_view_with(e, object_keys=np.zeros((1, 0), dtype=np.uint8)), None), - "address_dtype": lambda e: (_view_with(e, addresses=e.region.slots(0).astype(np.uint64)), None), + "slot_dtype": lambda e: (_view_with(e, slots=e.region.slots(0).astype(np.uint64)), None), + "slot_shape": lambda e: (_view_with(e, slots=e.region.slots(0)[:, None]), None), "len_disagrees_with_arrays": lambda e: (_view_with(e, rows=2), None), - "past_the_region": lambda e: (_view_with(e, addresses=e.region.slots(2)), None), - "before_the_region": lambda e: (_view_with(e, addresses=e.region.slots(-1)), None), - "misaligned": lambda e: (_view_with(e, addresses=e.region.slots(0) + 1), None), - "unknown_region_index": lambda e: (_view_with(e, region_index=1), None), - "negative_region_index": lambda e: (_view_with(e, region_index=-1), None), + "past_the_region": lambda e: (_view_with(e, slots=e.region.slots(2)), None), + "before_the_region": lambda e: (_view_with(e, slots=e.region.slots(-1)), None), + "unknown_region_index": lambda e: (_view_with(e, region_indices=(1,)), None), + "negative_region_index": lambda e: (_view_with(e, region_indices=(-1,)), None), + "empty_regions": lambda e: (_view_with(e, region_indices=()), None), + "repeated_region": lambda e: (_view_with(e, region_indices=(0, 0)), None), + "region_list": lambda e: (_view_with(e, region_indices=[0]), None), + "region_bool": lambda e: (_view_with(e, region_indices=(True,)), None), + "region_string": lambda e: (_view_with(e, region_indices=("0",)), None), "same_slot_in_a_group": lambda e: ( - CpuView([Rows(0, np.stack([object_key("A"), object_key("B")]), e.region.slots(0, 0), 0)]), + CpuView( + [Rows(0, np.stack([object_key("A"), object_key("B")]), e.region.slots(0, 0), (0,))] + ), None, ), "same_slot_across_groups": lambda e: (CpuView([e.rows(["A"], 0), e.rows(["A"], 1)]), None), @@ -1086,13 +1097,80 @@ def test_object_keys_in_any_memory_order_are_taken(store: MemoryStore) -> None: commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) reader = Endpoint(MemoryDriver(store), ["A", "B"]) rows = reader.rows(["A", "B"]) - view = CpuView([Rows(0, np.asfortranarray(rows.object_keys), rows.addresses, 0)]) + view = CpuView([Rows(0, np.asfortranarray(rows.object_keys), rows.slots, (0,))]) attempt = reader.backend.fetch(view, None) assert masks(reader.wait_outcome(attempt)) == [[True, True]] reader.wait_quiesced(attempt) reader.close() +@pytest.mark.parametrize("indices", [(0, 1), (1, 0)]) +@pytest.mark.parametrize("row_limit", [None, 1]) +def test_rows_transfer_the_same_slot_from_every_listed_region( + store: MemoryStore, indices: tuple[int, ...], row_limit: int | None +) -> None: + other = CpuRegion(3, 3) + capabilities = replace(DEFAULT_CAPABILITIES, max_rows_per_op=row_limit) + with closing( + Endpoint(MemoryDriver(store, capabilities=capabilities), ["A", "B", "C"], others=[other]) + ) as endpoint: + first = {name: name.encode() * SLOT_BYTES for name in endpoint.names} + second = {name: name.lower().encode() * 3 for name in endpoint.names} + for slot, name in enumerate(endpoint.names): + endpoint.write(name, first[name]) + other.write(slot, second[name]) + rows = Rows( + 7, + np.stack([object_key("C"), object_key("A")]), + np.array([2, 0], dtype=np.int64), + indices, + ) + assert isinstance(rows, LayerGroupRows) + view = CpuView([rows]) + published = endpoint.backend.publish(view, None) + assert masks(endpoint.wait_outcome(published)) == [[True, True]] + endpoint.wait_quiesced(published) + for name in ("C", "A"): + parts = (first[name], second[name]) + assert store.committed_payload(store_key(name)) == b"".join( + parts[index] for index in indices + ) + for call in endpoint.driver.calls_to("put"): + for row in call.rows: + assert tuple(span.nbytes for span in row.spans) == tuple( + (SLOT_BYTES, 3)[index] for index in indices + ) + + for slot, name in enumerate(endpoint.names): + endpoint.write(name, b"\0" * SLOT_BYTES) + other.write(slot, b"\0" * 3) + fetched = endpoint.backend.fetch(view, None) + assert masks(endpoint.wait_outcome(fetched)) == [[True, True]] + endpoint.wait_quiesced(fetched) + for slot, name in ((2, "C"), (0, "A")): + assert endpoint.read(name) == first[name] + assert other.read(slot) == second[name] + assert endpoint.read("B") == b"\0" * SLOT_BYTES and other.read(1) == b"\0" * 3 + + +@pytest.mark.parametrize("failure", ["driver", "slot_bound", "shared_destination"]) +def test_multi_region_rows_are_rejected_before_any_driver_access( + store: MemoryStore, failure: str +) -> None: + other = CpuRegion(3, 1 if failure == "slot_bound" else 2) + capabilities = replace(DEFAULT_CAPABILITIES, multi_span=failure != "driver") + endpoint = Endpoint(MemoryDriver(store, capabilities=capabilities), ["A", "B"], others=[other]) + rows = Rows(0, object_key("A")[None, :], np.array([1], dtype=np.int64), (0, 1)) + groups = [rows] + if failure == "shared_destination": + groups.append(Rows(1, object_key("B")[None, :], rows.slots, (1,))) + with pytest.raises(InvalidSubmission): + endpoint.backend.fetch(CpuView(groups), None) + assert store_calls(endpoint) == [] + assert endpoint.backend.stats()["attempts_active"] == 0 + endpoint.close() + + @pytest.mark.parametrize("case", INVALID_SUBMISSIONS) def test_invalid_submissions_are_rejected_before_the_driver(store: MemoryStore, case: str) -> None: reader = Endpoint(MemoryDriver(store), ["A", "B"]) @@ -1128,7 +1206,7 @@ def test_a_publish_may_read_one_slot_twice(store: MemoryStore) -> None: writer = Endpoint(MemoryDriver(store), ["A", "B"]) writer.write("A", b"aaaaaaaa") keys = np.stack([object_key("A"), object_key("B")]) - view = CpuView([Rows(0, keys, writer.region.slots(0, 0), 0)]) + view = CpuView([Rows(0, keys, writer.region.slots(0, 0), (0,))]) attempt = writer.backend.publish(view, None) assert masks(writer.wait_outcome(attempt)) == [[True, True]] assert store.committed_payload(store_key("B")) == b"aaaaaaaa" @@ -1577,7 +1655,7 @@ def test_rows_in_a_region_the_backend_did_not_register_are_rejected(store: Memor device = CpuRegion(SLOT_BYTES, 1, memory_kind=MemoryKind.DEVICE) reader = Endpoint(MemoryDriver(store, capabilities=HOST_ONLY), ["A"], others=[device]) assert reader.binding.live_labels() == frozenset({"region:0"}) - view = CpuView([Rows(0, object_key("A")[None, :], device.slots(0), 1)]) + view = CpuView([Rows(0, object_key("A")[None, :], device.slots(0), (1,))]) with pytest.raises(InvalidSubmission, match="region 1 is not registered"): reader.backend.fetch(view, None) with pytest.raises(ValueError, match="region 1 is not registered"): @@ -1618,7 +1696,7 @@ def test_rows_naming_one_key_at_two_addresses_never_share_a_call( 0, np.stack([object_key(key) for key in keys]), endpoint.region.slots(*range(len(keys))), - 0, + (0,), ) ] ) diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py index 622d6f7c5001..5bc3c799141c 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py @@ -132,6 +132,19 @@ def test_a_batch_of_slots_builds_nothing_until_read() -> None: assert list(batch) == [batch[0], batch[1], batch[2]] # read row by row, as a Sequence[RowIO] +def test_a_batch_of_spans_keeps_the_listed_payload_order_without_copying() -> None: + keys = np.array([[1], [2]], dtype=np.uint8) + addresses = np.array([[80, 8], [96, 12]], dtype=np.int64) + extents = (Span(80, 24), Span(8, 8)) + batch = RowBatch.of_spans("s/", keys, addresses, (8, 4), extents) + assert np.shares_memory(batch.addresses, addresses) + assert addresses.flags.writeable and not batch.addresses.flags.writeable + assert batch.spans_per_row == 2 and batch.row_bytes == 12 + assert batch.extents == extents + assert batch[0] == RowIO("s/01", (Span(80, 8), Span(8, 4))) + assert list(batch[1:]) == [RowIO("s/02", (Span(96, 8), Span(12, 4)))] + + def test_a_slice_of_a_batch_is_the_batch_of_those_rows() -> None: names = np.arange(8, dtype=np.uint8).reshape(4, 2) batch = RowBatch.of_slots("s/", names, np.arange(4, dtype=np.int64) * 8, 8, Span(0, 32)) diff --git a/tests/unittest/_torch/disaggregation/test_backend_contract.py b/tests/unittest/_torch/disaggregation/test_backend_contract.py index 69b7651a0fe0..4e0e6b0e1d7e 100644 --- a/tests/unittest/_torch/disaggregation/test_backend_contract.py +++ b/tests/unittest/_torch/disaggregation/test_backend_contract.py @@ -67,13 +67,13 @@ def nbytes(self) -> int: @dataclasses.dataclass(frozen=True, eq=False) class _LayerGroupRows: layer_group_index: int - row_count: int + page_ordinals: np.ndarray object_keys: typing.Optional[np.ndarray] = None - addresses: typing.Optional[np.ndarray] = None - region_index: typing.Optional[int] = None + slots: typing.Optional[np.ndarray] = None + region_indices: typing.Optional[typing.Tuple[int, ...]] = None def __len__(self) -> int: - return self.row_count + return len(self.page_ordinals) @dataclasses.dataclass(frozen=True, eq=False) @@ -86,22 +86,17 @@ def row_masks(self, fill_value: bool = False) -> typing.Tuple[np.ndarray, ...]: def _keyed_view(*row_counts: int) -> _View: """One group per row count, in order: ``_keyed_view(2, 3)`` has two rows in layer group 0 - and three in layer group 1. Rows occupy consecutive slots in region 0; their keys encode - their slot indices. The arrays are read-only, as a view's metadata is immutable.""" - region = _SlotRegion( - "pool", address=4096, slot_bytes=64, slot_count=sum(row_counts), memory_kind=MemoryKind.HOST - ) - groups, first_slot_index = [], 0 + and three in layer group 1. Rows hold pages 0, 1, ... and occupy consecutive slots of region 0; + their keys encode their slots. The arrays are read-only, as a view's metadata is immutable.""" + groups, first_slot = [], 0 for layer_group_index, row_count in enumerate(row_counts): - slot_indices = np.arange(first_slot_index, first_slot_index + row_count, dtype=np.int64) - object_keys = slot_indices.astype(" int", + "page_ordinals": "-> np.ndarray", "object_keys": "-> np.ndarray | None", - "addresses": "-> np.ndarray | None", - "region_index": "-> int | None", + "slots": "-> np.ndarray | None", + "region_indices": "-> tuple[int, ...] | None", "__len__": "(self) -> 'int'", }, "SlotRegion": { @@ -246,7 +242,7 @@ def test_minimal_structural_implementations_pass_runtime_protocol_checks(): region = _SlotRegion( "pool", address=4096, slot_bytes=64, slot_count=8, memory_kind=MemoryKind.HOST ) - view = _View((_keyed_view(2).groups[0], _LayerGroupRows(1, 3))) + view = _View((_keyed_view(2).groups[0], _LayerGroupRows(1, np.arange(3)))) assert isinstance(region, SlotRegion) and region.nbytes == 512 assert isinstance(view, View) assert all(isinstance(group, LayerGroupRows) for group in view.groups) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/conftest.py b/tests/unittest/_torch/executor/kv_cache/sharing/conftest.py index d58a98dbfc27..b9dd7506b6b4 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/conftest.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/conftest.py @@ -19,6 +19,7 @@ from __future__ import annotations import ctypes +import functools import gc import hashlib import inspect @@ -26,13 +27,27 @@ import threading from contextlib import contextmanager from types import SimpleNamespace -from typing import Callable, Dict, Generator, Iterator, List, Optional, Sequence, Tuple +from typing import ( + TYPE_CHECKING, + Callable, + Dict, + Generator, + Iterator, + List, + Optional, + Sequence, + Tuple, +) from unittest.mock import patch import numpy as np import pytest import torch +if TYPE_CHECKING: + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import Lease, LeaseView + TPB = 32 MAX_SEQ_LEN = 256 SENTINEL = 0xEE @@ -59,10 +74,7 @@ def allreduce(value, op=None): @contextmanager def _collectives_for(mapping): - """Stand-in collectives while a multi-rank manager is built in a single process.""" - if mapping is None or mapping.world_size == 1: - yield - return + """Stand-in collectives while a manager is built outside a process group.""" from tensorrt_llm._torch.distributed.communicator import Distributed with patch.object(Distributed, "get", return_value=_RanksAgree()): @@ -89,11 +101,13 @@ def make_manager( joint_kv_cache_reuse: bool = False, kv_connector_manager=None, cold_page_codec_provider=None, + max_beam_width: int = 1, + attention_read_windows=None, **kv_cache_config, ): """A small real ``KVCacheManagerV2`` on its own execution stream, FP16 unless ``dtype``, built for ``spec_config`` as a target or (``is_draft``) a draft pool, with ``kv_connector_manager`` if - given; further keywords go to its ``KvCacheConfig``.""" + given and up to ``max_beam_width`` beams; further keywords go to its ``KvCacheConfig``.""" import tensorrt_llm.bindings from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -126,13 +140,14 @@ def make_manager( joint_kv_cache_reuse=joint_kv_cache_reuse, kv_connector_manager=kv_connector_manager, cold_page_codec_provider=cold_page_codec_provider, + max_beam_width=max_beam_width, + attention_read_windows=attention_read_windows or {}, ) -def _make_deepseek_v4_manager(spec_config=None, fp8_ds_mla=False): - """A four-layer DeepSeek-V4 manager, FP8 with one KV head: one layer of each compression kind - (1, 4, 128) and a last ratio-4 layer the draft layers copy; with ``spec_config`` its one-model - draft layers follow. ``fp8_ds_mla``: the footer-scale cache, in SM100+'s 256-token blocks.""" +def _make_deepseek_v4_manager(spec_config=None, fp8_ds_mla=False, compress_ratios=None): + """A small FP8 DeepSeek-V4 manager; default ratios cover 1, 4 and 128. + ``fp8_ds_mla`` selects footer scales in SM100+'s 256-token blocks.""" from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager from tensorrt_llm._torch.speculative.utils import get_num_spec_layers from tensorrt_llm.bindings import DataType @@ -140,7 +155,7 @@ def _make_deepseek_v4_manager(spec_config=None, fp8_ds_mla=False): from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig from tensorrt_llm.mapping import Mapping - compress_ratios = [1, 4, 128, 4] + compress_ratios = [1, 4, 128, 4] if compress_ratios is None else list(compress_ratios) draft_layers = get_num_spec_layers(spec_config) if spec_config is not None else 0 draft_ratios = [compress_ratios[-1]] * draft_layers return DeepseekV4CacheManager( @@ -171,8 +186,9 @@ def _make_deepseek_v4_manager(spec_config=None, fp8_ds_mla=False): ) -def _make_hybrid_manager(): - """Nemotron-H shaped, few layers: Mamba2 state at layers 0, 2 and 4, attention at layer 3.""" +def _make_hybrid_manager(mapping=None): + """Nemotron-H shaped, few layers: Mamba2 state at layers 0, 2 and 4, attention at layer 3; on + one rank of ``mapping`` when given.""" from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( MambaHybridCacheManagerV2, ) @@ -184,33 +200,81 @@ def _make_hybrid_manager(): pattern = "M-M*M-" mamba_mask = [c == "M" for c in pattern] attn_mask = [c == "*" for c in pattern] - return MambaHybridCacheManagerV2( - mamba_d_state=128, - mamba_d_conv=4, - mamba_num_heads=128, - mamba_n_groups=8, - mamba_head_dim=80, - mamba_num_layers=sum(mamba_mask), - mamba_layer_mask=mamba_mask, - mamba_cache_dtype=torch.bfloat16, - mamba_ssm_cache_dtype=torch.float32, - kv_cache_config=KvCacheConfig( - max_tokens=2048, - enable_block_reuse=True, - mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), - ), - kv_cache_type=CacheType.SELF, - num_layers=sum(attn_mask), - num_kv_heads=8, - head_dim=128, - tokens_per_block=32, - max_seq_len=1024, - max_batch_size=2, - mapping=Mapping(world_size=1, rank=0, tp_size=1), - layer_mask=attn_mask, - vocab_size=1024, - dtype=DataType.BF16, + mapping = mapping or Mapping(world_size=1, rank=0, tp_size=1) + with _collectives_for(mapping): + return MambaHybridCacheManagerV2( + mamba_d_state=128, + mamba_d_conv=4, + mamba_num_heads=128, + mamba_n_groups=8, + mamba_head_dim=80, + mamba_num_layers=sum(mamba_mask), + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig( + max_tokens=2048, + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + ), + kv_cache_type=CacheType.SELF, + num_layers=sum(attn_mask), + num_kv_heads=8, + head_dim=128, + tokens_per_block=32, + max_seq_len=1024, + max_batch_size=2, + mapping=mapping, + layer_mask=attn_mask, + vocab_size=1024, + dtype=DataType.BF16, + ) + + +def _make_gdn_manager(mapping=None, pattern="MMM*"): + """Qwen3-Next shaped, smaller: gated delta-net state, its convolution state in [Q | K | V] + sections of 4 key heads, 4 key heads and 8 value heads of 64, at each ``M`` of ``pattern``, and + full attention of 2 KV heads at each ``*``; on one rank of ``mapping`` when given.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, ) + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MambaStateConfig + from tensorrt_llm.mapping import Mapping + + mamba_mask = [c == "M" for c in pattern] + attn_mask = [c == "*" for c in pattern] + mapping = mapping or Mapping(world_size=1, rank=0, tp_size=1) + with _collectives_for(mapping): + return MambaHybridCacheManagerV2( + mamba_d_state=64, + mamba_d_conv=4, + mamba_num_heads=8, + mamba_n_groups=4, + mamba_head_dim=64, + mamba_num_layers=sum(mamba_mask), + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig( + max_tokens=2048, + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + ), + kv_cache_type=CacheType.SELF, + num_layers=sum(attn_mask), + num_kv_heads=2, + head_dim=128, + tokens_per_block=32, + max_seq_len=1024, + max_batch_size=2, + mapping=mapping, + layer_mask=attn_mask, + vocab_size=1024, + dtype=DataType.BF16, + conv_state_layout="q_k_v", + ) @contextmanager @@ -273,21 +337,34 @@ def factory(**kwargs): @pytest.fixture def deepseek_v4_manager(): - """``with deepseek_v4_manager(spec_config=None, fp8_ds_mla=False) as mgr:`` a small - ``DeepseekV4CacheManager``; callers skip before Blackwell.""" + """``with deepseek_v4_manager(spec_config=None, fp8_ds_mla=False, compress_ratios=None) as + mgr:`` a small ``DeepseekV4CacheManager`` with a layer per entry of ``compress_ratios``, each + its layer's compression ratio, ``[1, 4, 128, 4]`` by default. Callers skip before Blackwell.""" - def factory(spec_config=None, fp8_ds_mla=False): - return _managed(lambda: _make_deepseek_v4_manager(spec_config, fp8_ds_mla)) + def factory(spec_config=None, fp8_ds_mla=False, compress_ratios=None): + return _managed(lambda: _make_deepseek_v4_manager(spec_config, fp8_ds_mla, compress_ratios)) return factory @pytest.fixture def hybrid_manager(): - """``with hybrid_manager() as mgr:`` a manager whose layer groups include recurrent state.""" + """``with hybrid_manager(mapping=None) as mgr:`` a manager whose layer groups include recurrent + state.""" - def factory(): - return _managed(_make_hybrid_manager) + def factory(mapping=None): + return _managed(lambda: _make_hybrid_manager(mapping)) + + return factory + + +@pytest.fixture +def gdn_manager(): + """``with gdn_manager(mapping=None, pattern="MMM*") as mgr:`` a Qwen3-Next shaped manager of + gated delta-net state and attention layers.""" + + def factory(mapping=None, pattern="MMM*"): + return _managed(lambda: _make_gdn_manager(mapping, pattern)) return factory @@ -351,7 +428,7 @@ def pool_group_of(mgr) -> List[int]: def pool_group_ids(mgr) -> List[int]: - """Device pool groups in index order: the order of a staging lender's regions.""" + """Device pool groups in the order of the lender's host staging regions.""" return sorted(int(pg.pool_group_index) for pg in mgr.impl.pool_group_descs) @@ -486,6 +563,46 @@ def admitted(mgr, request_id: int, tokens: Sequence[int]): return request +def ready_view(lease: Lease, mgr: KVCacheManagerV2, tries: int = 3) -> LeaseView: + """The lease's view once the copies queued so far have run, a round of its lender's progress + before each poll, as an executor iteration gives; fails if it never comes.""" + for _ in range(tries): + mgr._stream.synchronize() + lender = lease._lender() + if lender is not None: + lender.progress() + view = lease.poll() + if view is not None: + return view + raise AssertionError(f"the lease never became ready: {lease.failure}") + + +def return_write(lender, lease: Lease, arrived, reusable: bool = True): + """``lender.return_writes`` of the one entry ``(lease, arrived, reusable)``: its bounds.""" + (bounds,) = lender.return_writes([(lease, arrived, reusable)]) + return bounds + + +def settled_bounds(lender, mgr: KVCacheManagerV2, request, tries: int = 8): + """The request's bounds once its fetches settle, a round of the lender's progress before each + query; fails if they never do.""" + for _ in range(tries): + mgr._stream.synchronize() + lender.progress() + bounds = lender.poll_resume_bounds(request) + if bounds is not None: + return bounds + raise AssertionError("the request's fetches never settled") + + +def stream_event(mgr: KVCacheManagerV2) -> torch.cuda.Event: + """An event recorded after the work queued so far on the manager's execution stream: a lend's + ``ready_event`` where that work writes or uses the lent pages.""" + event = torch.cuda.Event() + event.record(mgr._stream) + return event + + def host_bytes(address: int, length: int) -> bytes: return ctypes.string_at(address, length) @@ -513,13 +630,25 @@ def staged_row(byte: int, layer_group: int, ordinal: int, nbytes: int) -> bytes: return np.random.default_rng(seed).integers(0, 256, nbytes, dtype=np.uint8).tobytes() +def staged_slots(lender, rows) -> Tuple[object, List[int]]: + """The staging region holding ``rows`` and the address of each row's slot in it, after checking + that the rows are staged pages of their layer group.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import MemoryKind + + (index,) = rows.region_indices + region = lender.regions[index] + assert region.memory_kind is MemoryKind.HOST + assert rows.layer_group_index in region.layer_group_indices + return region, [region.address + slot * region.slot_bytes for slot in rows.slots.tolist()] + + def stage(lender, view, byte: int) -> None: """Fill every slot of a staging view as a backend's transfer would, each block entry with its own bytes (``staged_row``).""" for blocks in view.groups: - assert blocks.layer_group_index in lender.regions[blocks.region_index].layer_group_indices - length = lender.regions[blocks.region_index].slot_bytes - for ordinal, address in zip(blocks.block_ordinals.tolist(), blocks.addresses.tolist()): + region, addresses = staged_slots(lender, blocks) + length = region.slot_bytes + for ordinal, address in zip(blocks.page_ordinals.tolist(), addresses): ctypes.memmove( address, staged_row(byte, blocks.layer_group_index, ordinal, length), length ) @@ -531,25 +660,18 @@ def relay(src_lender, src_view, dst_lender, dst_view, deliver=None) -> Tuple[np. entries (default all); returns the masks.""" by_object_key = {} for blocks in src_view.groups: - assert ( - blocks.layer_group_index in src_lender.regions[blocks.region_index].layer_group_indices - ) - length = src_lender.regions[blocks.region_index].slot_bytes - for object_key, address in zip(blocks.object_keys, blocks.addresses.tolist()): - by_object_key[object_key.tobytes()] = (address, length) + region, addresses = staged_slots(src_lender, blocks) + for object_key, address in zip(blocks.object_keys, addresses): + by_object_key[object_key.tobytes()] = (address, region.slot_bytes) masks = [] for i, blocks in enumerate(dst_view.groups): - assert ( - blocks.layer_group_index in dst_lender.regions[blocks.region_index].layer_group_indices - ) + region, addresses = staged_slots(dst_lender, blocks) mask = np.zeros(len(blocks), dtype=bool) - for entry_index, (object_key, address) in enumerate( - zip(blocks.object_keys, blocks.addresses.tolist()) - ): + for entry_index, (object_key, address) in enumerate(zip(blocks.object_keys, addresses)): if deliver is not None and not deliver(i, entry_index): continue src, length = by_object_key[object_key.tobytes()] - assert length == dst_lender.regions[blocks.region_index].slot_bytes + assert length == region.slot_bytes ctypes.memmove(address, src, length) mask[entry_index] = True masks.append(mask) @@ -743,6 +865,13 @@ def retained() -> Tuple[object, ...]: return _lender._retained() +def staging_regions(lender) -> Tuple[object, ...]: + """The lender's host regions, its staging: what a staging backend registers.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import MemoryKind + + return tuple(region for region in lender.regions if region.memory_kind is MemoryKind.HOST) + + def staging_memory(regions) -> List[object]: """The kept host allocations holding any of ``regions``, compared by address.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender @@ -908,23 +1037,42 @@ def run(): pool_pages=pool_pages, prefill=prefill, published=published, + ready_view=ready_view, reclaim=reclaim, relay=relay, retained=retained, + return_write=return_write, + settled_bounds=settled_bounds, stage=stage, staged_row=staged_row, + staged_slots=staged_slots, stale_blocks=stale_blocks, staging_kept=staging_kept, staging_memory=staging_memory, + staging_regions=staging_regions, + stream_event=stream_event, tier_used=tier_used, warm_tier_moves=warm_tier_moves, windows=windows, ) +def _lend_keywords() -> None: + """The kit's keyword sets of a lend, each with ``ready_event=None``, made at the first use.""" + if hasattr(_KIT, "DEVICE_ALL"): + return + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import MemoryKind, RowSelection + + for kind in MemoryKind: + for rows in RowSelection: + keywords = dict(memory_kind=kind, rows=rows, ready_event=None) + setattr(_KIT, f"{kind.name}_{rows.name}", keywords) + + @pytest.fixture def kit(): """The helpers and oracles above, for test modules that cannot import this file by name.""" + _lend_keywords() return _KIT @@ -967,6 +1115,107 @@ def _let_go_of_what_the_test_kept(request: pytest.FixtureRequest) -> Iterator[No gc.collect() +@pytest.fixture(autouse=True) +def ordered_locks(monkeypatch: pytest.MonkeyPatch) -> List[Dict[int, List[int]]]: + """Every lock and lock check asked for with ``ordered`` blocks of a real runtime cache, in every + test, each checked for what ``ordered`` promises: ``int`` layer groups, sorted, distinct, + non-negative ``int64`` on immutable bytes. The blocks of each, as lists, in order. Stand-in + caches, whose tests may hand the lender entries of their own, are not checked.""" + seen: List[Dict[int, List[int]]] = [] + locks = sys.modules.get(f"{_SHARING}._page_locks") + runtime = sys.modules.get("tensorrt_llm.runtime.kv_cache_manager_v2") + if locks is None or runtime is None: + return seen + lock, check = locks.PageLocks.lock, locks.PageLocks.check + + def promised(kv, blocks): + if not isinstance(kv, runtime._KVCache): + return + for lg, ordinals in blocks.items(): + assert type(lg) is int and isinstance(ordinals, np.ndarray), (lg, type(ordinals)) + assert ordinals.dtype == np.int64 and ordinals.ndim == 1, (lg, ordinals) + root = ordinals + while isinstance(root, np.ndarray) and root.base is not None: + root = root.base + assert not ordinals.flags.writeable and isinstance(root, bytes), ( + f"layer group {lg}: ordinals on {type(root).__name__}, " + f"writable: {ordinals.flags.writeable}" + ) + values = ordinals.tolist() + assert values == sorted(set(values)) and min(values, default=0) >= 0, (lg, values) + seen.append({lg: ordinals.tolist() for lg, ordinals in blocks.items()}) + + def lock_ordered(self, kv, blocks, mode, level=None, *, ordered=False, **kwargs): + if ordered: + promised(kv, blocks) + return lock(self, kv, blocks, mode, level, ordered=ordered, **kwargs) + + def check_ordered(self, kv, blocks, mode, *, ordered=False, **kwargs): + if ordered: + promised(kv, blocks) + return check(self, kv, blocks, mode, ordered=ordered, **kwargs) + + fresh = locks.PageLocks.lock_fresh + + def lock_fresh_ordered(self, kv, blocks, mode, ready): + # A fast lend's blocks, which a fresh cache's lock takes as ``ordered`` promises. + taken = fresh(self, kv, blocks, mode, ready) + if taken is not None: + promised(kv, blocks) + return taken + + monkeypatch.setattr(locks.PageLocks, "lock", lock_ordered) + monkeypatch.setattr(locks.PageLocks, "check", check_ordered) + monkeypatch.setattr(locks.PageLocks, "lock_fresh", lock_fresh_ordered) + return seen + + +@pytest.fixture(autouse=True) +def _lent_slots_are_the_cache_s_pages(monkeypatch: pytest.MonkeyPatch) -> None: + """An oracle beside the lender's sources of pages: every in-place lease a test lends over a real + runtime cache lends, for each block of a layer group, the page the cache's aggregated page + indices give it on beam 0 and the page its base page indices lock, and for a layer group of + recurrent state the cache's state slot. Checked as each lease is lent; stand-in caches are not + checked.""" + lender = sys.modules.get(f"{_SHARING}._lender") + runtime = sys.modules.get("tensorrt_llm.runtime.kv_cache_manager_v2") + if lender is None or runtime is None: + return + + def check(self, lease): + kv = getattr(lease, "_kv", None) + if not isinstance(lease, lender._InPlaceLease) or lease._view is None: + return + if not isinstance(kv, runtime._KVCache) or not kv.is_active: + return + for rows, ordinals in zip(lease._view.groups, lease._block_entries.block_ordinals): + lg = rows.layer_group_index + if lg in self._recurrent_groups: + expected = [[int(kv.get_ssm_block_base_index(lg))] * len(ordinals)] * 2 + else: + aggregated = [int(p) for p in kv.get_aggregated_page_indices(lg, valid_only=False)] + base = [int(p) for p in kv.get_base_page_indices(lg)[: kv.num_blocks]] + expected = [ + [pages[b] if b < len(pages) else -1 for b in ordinals.tolist()] + for pages in (aggregated, base) + ] + assert [rows.slots.tolist()] * 2 == expected, ( + f"layer group {lg}: lent slots {rows.slots.tolist()} for blocks " + f"{ordinals.tolist()}, the cache's aggregated and base pages there {expected}" + ) + + for name in ("lend_read", "lend_write"): + lend = getattr(lender._ManagerLender, name) + + @functools.wraps(lend) + def checked(self, *args, _lend=lend, **kwargs): + lease = _lend(self, *args, **kwargs) + check(self, lease) + return lease + + monkeypatch.setattr(lender._ManagerLender, name, checked) + + @pytest.fixture def retained_cleanup() -> Callable[[bool], Iterator[None]]: """A fresh teardown generator for tests of the autouse cleanup itself.""" diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_host_tier.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_host_tier.py index 429c7e5c9571..2af0dd599a86 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_host_tier.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_host_tier.py @@ -12,9 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The staging lender beside the KV cache's host tier. Each check first proves its move: pages to -host and back by their request's suspension or others' pressure, bystander blocks to host while lent -pages stay, lent pages to other slots by a pool rebalance, held and late-marked pages to host.""" +"""The staging lender beside the KV cache's host tier; each check first proves the move it tests.""" import ctypes from contextlib import contextmanager @@ -39,27 +37,24 @@ SCOPE = b"host-tier-suite" -def attach(mgr, *, tokens_per_fetch=END, planned_fetches=1): +def attach(mgr, *, guaranteed_fetch_tokens=END, capacity_in_fetches=1): return attach_lender( - mgr, scope=SCOPE, staging=StagingOptions(tokens_per_fetch, planned_fetches) + mgr, scope=SCOPE, staging=StagingOptions(guaranteed_fetch_tokens, capacity_in_fetches) ) -def ready_view(lease, mgr, tries=3): - for _ in range(tries): - mgr._stream.synchronize() - view = lease.poll() - if view is not None: - return view - raise AssertionError(f"the lease never became ready: {lease.failure}") - - def staged(lender, view): """Per run, the bytes of every row's staging slot.""" out = [] for run in view.groups: - length = lender.regions[run.region_index].slot_bytes - out.append([ctypes.string_at(a, length) for a in run.addresses.tolist()]) + (index,) = run.region_indices + region = lender.regions[index] + out.append( + [ + ctypes.string_at(region.address + slot * region.slot_bytes, region.slot_bytes) + for slot in run.slots.tolist() + ] + ) return out @@ -99,11 +94,11 @@ def test_a_fetch_stays_usable_across_its_request_s_trip_to_host( with host_tier_manager() as mgr_a, host_tier_manager() as mgr_b: source = kit.published(mgr_a, SOURCE, PROMPT) lender_a = attach(mgr_a) - publish = lender_a.lend_read(source, 0, END) - publish_view = ready_view(publish, mgr_a) + publish = lender_a.lend_read(source, 0, END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) target = kit.admitted(mgr_b, TARGET, PROMPT) lender_b = attach(mgr_b) - lease = lender_b.lend_write(target, 0, END) + lease = lender_b.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() assert view is not None kit.fill_sentinel(mgr_b, target) @@ -113,13 +108,12 @@ def test_a_fetch_stays_usable_across_its_request_s_trip_to_host( with kit.held_stream(mgr_b._stream, strict=False) as gate: # The copy into the lent pages waits behind the gate; the suspension and the move to # host come after it on the stream, so the move carries the fetched bytes. - lease.mark_arrived(masks) - lease.release() + assert kit.return_write(lender_b, lease, masks) is None, "settled before its copy ran" with to_host_and_back(kit, mgr_b, target, dev) as (others, old): - assert lender_b.resume_bounds(target) is None, "settled before its copy ran" + lender_b.progress() + assert lender_b.poll_resume_bounds(target) is None, "settled before its copy ran" gate.open() - mgr_b._stream.synchronize() - assert lender_b.resume_bounds(target) == (END, 0), "not usable once its copy ran" + assert kit.settled_bounds(lender_b, mgr_b, target) == (END, 0), "not usable" for ordinal in range(BLOCKS): got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) want = kit.digest(kit.page(mgr_a, source, 0, ordinal)) @@ -136,12 +130,12 @@ def test_a_publish_after_a_trip_to_host_reads_the_pages_where_they_are( source = kit.published(mgr, SOURCE, PROMPT) original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] lender = attach(mgr) - first = lender.lend_read(source, 0, END) - assert kit.digest(staged(lender, ready_view(first, mgr))) == kit.digest([original]) + first = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + assert kit.digest(staged(lender, kit.ready_view(first, mgr))) == kit.digest([original]) first.release() with to_host_and_back(kit, mgr, source): - again = lender.lend_read(source, 0, END) - got = kit.digest(staged(lender, ready_view(again, mgr))) + again = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + got = kit.digest(staged(lender, kit.ready_view(again, mgr))) assert got == kit.digest([original]), "read a stale page" again.release() @@ -149,18 +143,16 @@ def test_a_publish_after_a_trip_to_host_reads_the_pages_where_they_are( # -- (b) other requests' pressure while a fetch is lent ----------------------------------------- -@pytest.mark.parametrize("moment", ["copy_queued", "before_marks"]) +@pytest.mark.parametrize("moment", ["copy_queued", "before_the_return"]) def test_other_requests_pushing_blocks_to_host_leave_a_fetch_where_its_resume_bounds_says( kit, host_tier_manager, moment, attach=attach ): - """``moment``: ``copy_queued``, the target stays active and its marks' copy waits on the stream - while other requests push committed blocks to host; ``before_marks``, the target itself goes to - host and back between the grant and the marks.""" + """Others push blocks to host behind the return's copy, or the target goes before its return.""" with host_tier_manager() as mgr_a, host_tier_manager() as mgr_b: source = kit.published(mgr_a, SOURCE, PROMPT) lender_a = attach(mgr_a) - publish = lender_a.lend_read(source, 0, END) - publish_view = ready_view(publish, mgr_a) + publish = lender_a.lend_read(source, 0, END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) original = [kit.page(mgr_a, source, 0, o) for o in range(BLOCKS)] got = kit.digest(staged(lender_a, publish_view)) assert got == kit.digest([original]), "the relay carries rows off their blocks" @@ -169,7 +161,7 @@ def test_other_requests_pushing_blocks_to_host_leave_a_fetch_where_its_resume_bo mgr_b.free_resources(kit.published(mgr_b, 10 + i, tokens)) target = kit.admitted(mgr_b, TARGET, PROMPT) lender_b = attach(mgr_b) - lease = lender_b.lend_write(target, 0, END) + lease = lender_b.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() assert view is not None kit.fill_sentinel(mgr_b, target) @@ -183,29 +175,30 @@ def test_other_requests_pushing_blocks_to_host_leave_a_fetch_where_its_resume_bo host_before = kit.tier_used(mgr_b, 1) try: with kit.held_stream(mgr_b._stream, strict=False) as gate: - lease.mark_arrived(masks) - lease.release() + bounds = kit.return_write(lender_b, lease, masks) + assert bounds is None, f"settled before its copy ran: {bounds}" assert others.allocate(kit.pool_pages(mgr_b) - BLOCKS, fill=False, chunk=1) moved = kit.tier_used(mgr_b, 1) - host_before assert moved >= len(BYSTANDERS) * BLOCKS, "no bystander block reached host" assert kit.pages(kv, 0)[:BLOCKS] == lent, "the lent pages moved" - assert lender_b.resume_bounds(target) is None, "settled before its copy ran" + lender_b.progress() + assert lender_b.poll_resume_bounds(target) is None, "settled before its copy" gate.open() - assert lender_b.resume_bounds(target) == (END, 0) + assert kit.settled_bounds(lender_b, mgr_b, target) == (END, 0) assert all(dev.read(0, slot) == sentinel for slot in others.pages()) finally: others.free() else: with to_host_and_back(kit, mgr_b, target) as (others, old): assert old == lent - lease.mark_arrived(masks) - lease.release() + bounds = kit.return_write(lender_b, lease, masks) + assert bounds == (0, 0), f"the return counts rows that never landed: {bounds}" mgr_b._stream.synchronize() assert all(dev.read(0, slot) == sentinel for slot in old), ( - "the marks' copy reached another request's page" + "the return's copy reached another request's page" ) - assert lender_b.resume_bounds(target) == (0, 0), "counts rows that never landed" + assert kit.settled_bounds(lender_b, mgr_b, target) == (0, 0), "rows never landed" for ordinal in range(BLOCKS): got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) expected = kit.page(mgr_a, source, 0, ordinal) if moment == "copy_queued" else sentinel @@ -223,16 +216,17 @@ def use_as_a_backend(kit, mgr, lender, lease, reading, original): """What a backend does with a lease that became ready: a read's rows are compared with the blocks' bytes, a write's rows are filled.""" mgr._stream.synchronize() + lender.progress() view = lease.poll() if view is None: return if reading: (run,) = view.groups - want = [[original[o] for o in run.block_ordinals.tolist()]] + want = [[original[o] for o in run.page_ordinals.tolist()]] assert kit.digest(staged(lender, view)) == kit.digest(want), SERVED else: kit.stage(lender, view, 0x5A) - lease.mark_arrived(view.row_masks(True)) + kit.return_write(lender, lease, view.row_masks(True)) def test_a_cache_on_host_fails_every_lend_at_the_call_and_lends_again_once_resumed( @@ -246,7 +240,7 @@ def test_a_cache_on_host_fails_every_lend_at_the_call_and_lends_again_once_resum original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] target = kit.admitted(mgr, TARGET, BYSTANDERS[0]) assert mgr.resize_context(target, target.context_remaining_length) - lender = attach(mgr, planned_fetches=2) + lender = attach(mgr, capacity_in_fetches=2) writer = target kv, kv_writer = kit.kv(mgr, source), kit.kv(mgr, writer) old = [p for p in kit.pages(kv, 0) if p >= 0] @@ -263,7 +257,10 @@ def test_a_cache_on_host_fails_every_lend_at_the_call_and_lends_again_once_resum assert kit.tier_used(mgr, 1) >= host_before + moved, "its pages did not reach host" dev = kit.DevicePages(mgr) theirs = {slot: dev.read(0, slot) for slot in others.pages()} - failed = [lender.lend_read(source, 0, END), lender.lend_write(writer, 0, END)] + failed = [ + lender.lend_read(source, 0, END, **kit.HOST_NAMED), + lender.lend_write(writer, 0, END, **kit.HOST_NAMED), + ] for lease, reading in zip(failed, (True, False)): use_as_a_backend(kit, mgr, lender, lease, reading, original) mgr._stream.synchronize() @@ -273,20 +270,21 @@ def test_a_cache_on_host_fails_every_lend_at_the_call_and_lends_again_once_resum assert lease.failure is not None, "not failed at the call" assert lease.poll() is None groups = kit.pool_group_ids(mgr) - assert [lender._free_slots(g) for g in groups] == [p.slot_count for p in lender.regions] + assert [lender._free_slots(g) for g in groups] == [ + p.slot_count for p in kit.staging_regions(lender) + ] for lease in failed: lease.release() others.free_all_but(set(old)) assert mgr.resume_request(source) and mgr.resume_request(target) now = [p for p in kit.pages(kv, 0) if p >= 0] assert not set(now) & set(old), "it came back to its old pages" - read = lender.lend_read(source, 0, END) - view = ready_view(read, mgr) + read = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + view = kit.ready_view(read, mgr) assert kit.digest(staged(lender, view)) == kit.digest([original]) - write = lender.lend_write(writer, 0, END) - write_view = ready_view(write, mgr) - write.mark_arrived(write_view.row_masks()) - write.release() + write = lender.lend_write(writer, 0, END, **kit.HOST_NAMED) + write_view = kit.ready_view(write, mgr) + kit.return_write(lender, write, write_view.row_masks()) read.release() finally: others.free() @@ -301,8 +299,8 @@ def test_a_committed_block_has_one_name_on_the_gpu_and_after_a_trip_to_host( with host_tier_manager() as mgr: first = kit.published(mgr, SOURCE, PROMPT) lender = attach(mgr) - lease = lender.lend_read(first, 0, END) - view = ready_view(lease, mgr) + lease = lender.lend_read(first, 0, END, **kit.HOST_NAMED) + view = kit.ready_view(lease, mgr) first_names, first_bytes = names(view), staged(lender, view) lease.release() old = kit.pages(kit.kv(mgr, first), 0)[:BLOCKS] @@ -320,8 +318,8 @@ def test_a_committed_block_has_one_name_on_the_gpu_and_after_a_trip_to_host( assert kv.num_committed_tokens == END, "the blocks on host were not reused" now = kit.pages(kv, 0)[:BLOCKS] assert not set(now) & set(old), "the blocks came back to their old pages" - again = lender.lend_read(second, 0, END) - view = ready_view(again, mgr) + again = lender.lend_read(second, 0, END, **kit.HOST_NAMED) + view = kit.ready_view(again, mgr) assert names(view) == first_names, "a block's name changed with its tier" assert kit.digest(staged(lender, view)) == kit.digest(first_bytes) again.release() @@ -395,9 +393,9 @@ def test_a_rebalance_under_a_fetch_moves_no_copy_off_the_pages_still_lent( shape = dict(windows=[WINDOW, 256], num_layers=3, max_tokens=4 * kit.POOL_TOKENS) with host_tier_manager(**shape) as mgr_a, host_tier_manager(**shape) as mgr_b: source = kit.published(mgr_a, SOURCE, WINDOWED_PROMPT) - lender_a = attach(mgr_a, tokens_per_fetch=WINDOWED_END) - publish = lender_a.lend_read(source, 0, WINDOWED_END) - publish_view = ready_view(publish, mgr_a) + lender_a = attach(mgr_a, guaranteed_fetch_tokens=WINDOWED_END) + publish = lender_a.lend_read(source, 0, WINDOWED_END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) full = kit.windows(mgr_b).index(None) pg_full = kit.pool_group_of(mgr_b)[full] device = [] @@ -411,8 +409,8 @@ def test_a_rebalance_under_a_fetch_moves_no_copy_off_the_pages_still_lent( while stats(0)[pg_full].total - stats(0)[pg_full].free < 72 and fillers.allocate(7): pass target = kit.admitted(mgr_b, TARGET, WINDOWED_PROMPT) - lender_b = attach(mgr_b, tokens_per_fetch=WINDOWED_END) - lease = lender_b.lend_write(target, 0, WINDOWED_END) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=WINDOWED_END) + lease = lender_b.lend_write(target, 0, WINDOWED_END, **kit.HOST_NAMED) view = lease.poll() assert view is not None kit.fill_sentinel(mgr_b, target) @@ -434,15 +432,14 @@ def test_a_rebalance_under_a_fetch_moves_no_copy_off_the_pages_still_lent( assert now[full] != lent[full] and gone, "the rebalance left the lent pages where they were" guard = DeviceCopyGuard(_lender.drv, device, lambda: page_ranges(kit, mgr_b, kv, range(2))) monkeypatch.setattr(_lender, "drv", guard) - lease.mark_arrived(masks) - lease.release() + kit.return_write(lender_b, lease, masks) mgr_b._stream.synchronize() - assert guard.stray == [], "the marks' copy wrote outside the pages the target locks" + assert guard.stray == [], "the return's copy wrote outside the pages the target locks" assert guard.into_device, "nothing was copied: the check saw no copy" stayed = [ (run.layer_group_index, o) for run in view.groups - for o in run.block_ordinals.tolist() + for o in run.page_ordinals.tolist() if now[run.layer_group_index][o] == lent[run.layer_group_index][o] >= 0 ] wrong = [ @@ -453,7 +450,7 @@ def test_a_rebalance_under_a_fetch_moves_no_copy_off_the_pages_still_lent( assert stayed and not wrong, ( f"rows {wrong} of {stayed}, whose pages stayed, hold other bytes" ) - assert lender_b.resume_bounds(target) == (0, WINDOWED_END), "counts rows that never landed" + assert kit.settled_bounds(lender_b, mgr_b, target) == (0, WINDOWED_END), "rows never landed" publish.release() @@ -463,8 +460,8 @@ def test_a_rebalance_under_a_fetch_moves_no_copy_off_the_pages_still_lent( REB_PROMPTS = { "queued_read": list(range(11_000, 11_097)), # a read whose copy is queued behind the stream "waiting_read": list(range(12_000, 12_097)), # a read waiting for slots - "unmarked_write": list(range(13_000, 13_097)), # a write granted, marked after the rebalance - "marked_write": list(range(14_000, 14_097)), # a write marked before, its copy queued + "held_write": list(range(13_000, 13_097)), # a write granted, returned after the rebalance + "queued_write": list(range(14_000, 14_097)), # a write returned before, its copy queued } REB_GATE_SECONDS = 4.0 # One pool group, its quota halved with resize(level 0): every page above the new capacity moves @@ -526,8 +523,9 @@ def staged_right(kit, mgr, lender, view, request): """Per row of a read's view, whether its slot holds the block's published bytes.""" out = {} for run in view.groups: - nbytes = lender.regions[run.region_index].slot_bytes - for ordinal, address in zip(run.block_ordinals.tolist(), run.addresses.tolist()): + region, addresses = kit.staged_slots(lender, run) + nbytes = region.slot_bytes + for ordinal, address in zip(run.page_ordinals.tolist(), addresses): want = published_bytes(kit, mgr, request, run.layer_group_index, ordinal, nbytes) out[f"{run.layer_group_index}:{ordinal}"] = kit.host_bytes(address, nbytes) == want return out @@ -579,8 +577,8 @@ def rebalance_with_leases_in_flight( ): """The executor's rebalance (suspend every cache, ``resize`` or ``adjust`` the pools, resume) with four leases in flight: a read whose copy waits behind a held stream, a read waiting for - slots, a granted write marked after the rebalance and a write marked before, its copy queued. It - moves their pages, or with ``stay`` not the waiting read's and unmarked write's. A driver guard + slots, a granted write returned after the rebalance and one returned before, its copy queued. It + moves their pages, or with ``stay`` not the waiting read's and held write's. A driver guard sees every copy to or from a device pool; bytes are checked against what ``published`` wrote.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection @@ -593,11 +591,13 @@ def rebalance_with_leases_in_flight( # The publisher of what the writes fetch. pubs = { k: kit.published(mgr_a, 30 + i, REB_PROMPTS[k]) - for i, k in enumerate(("unmarked_write", "marked_write")) + for i, k in enumerate(("held_write", "queued_write")) + } + lender_a = attach(mgr_a, guaranteed_fetch_tokens=REB_END, capacity_in_fetches=2) + pub_leases = { + k: lender_a.lend_read(p, 0, REB_END, **kit.HOST_NAMED) for k, p in pubs.items() } - lender_a = attach(mgr_a, tokens_per_fetch=REB_END, planned_fetches=2) - pub_leases = {k: lender_a.lend_read(p, 0, REB_END) for k, p in pubs.items()} - pub_views = {k: ready_view(lease, mgr_a) for k, lease in pub_leases.items()} + pub_views = {k: kit.ready_view(lease, mgr_a) for k, lease in pub_leases.items()} stats = mgr_b.impl.get_storage_statistics def used(group): @@ -607,11 +607,11 @@ def used(group): totals = [int(stats(0)[g].total) for g in groups] total = totals[moving] goal = int(total * 0.6) - lender_b = attach(mgr_b, tokens_per_fetch=REB_END, planned_fetches=3) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=REB_END, capacity_in_fetches=3) leases, views, masks, early = {}, {}, {}, {} def lend_write(k, request): - leases[k] = lender_b.lend_write(request, 0, REB_END) + leases[k] = lender_b.lend_write(request, 0, REB_END, **kit.HOST_NAMED) views[k] = leases[k].poll() assert views[k] is not None, leases[k].failure kit.fill_sentinel(mgr_b, request) # a row that never lands shows @@ -620,8 +620,8 @@ def lend_write(k, request): if stay: # Admitted and grown before the fillers: their pages sit low, below what moves. early["waiting_read"] = kit.published(mgr_b, 22, REB_PROMPTS["waiting_read"]) - early["unmarked_write"] = kit.admitted(mgr_b, 23, REB_PROMPTS["unmarked_write"]) - lend_write("unmarked_write", early["unmarked_write"]) + early["held_write"] = kit.admitted(mgr_b, 23, REB_PROMPTS["held_write"]) + lend_write("held_write", early["held_write"]) # Fillers take the low slots (a fresh pool hands out slot ids in order), so the leased # requests' pages sit above them; another pool group may fill first, so stop at 70% of it. fillers = kit.Requests(mgr_b) @@ -635,17 +635,17 @@ def lend_write(k, request): "queued_read": kit.published(mgr_b, 21, REB_PROMPTS["queued_read"]), "waiting_read": early.get("waiting_read") or kit.published(mgr_b, 22, REB_PROMPTS["waiting_read"]), - "unmarked_write": early.get("unmarked_write") - or kit.admitted(mgr_b, 23, REB_PROMPTS["unmarked_write"]), - "marked_write": kit.admitted(mgr_b, 24, REB_PROMPTS["marked_write"]), + "held_write": early.get("held_write") + or kit.admitted(mgr_b, 23, REB_PROMPTS["held_write"]), + "queued_write": kit.admitted(mgr_b, 24, REB_PROMPTS["queued_write"]), } bystanders = [kit.published(mgr_b, 25 + i, t) for i, t in enumerate(BYSTANDERS)] - for k in ("unmarked_write", "marked_write"): + for k in ("held_write", "queued_write"): if k not in leases: lend_write(k, reqs[k]) rows = { - k: [(run.layer_group_index, run.block_ordinals.tolist()) for run in views[k].groups] - for k in ("unmarked_write", "marked_write") + k: [(run.layer_group_index, run.page_ordinals.tolist()) for run in views[k].groups] + for k in ("held_write", "queued_write") } fillers.free() dev = kit.DevicePages(mgr_b) # built before the gate: building it waits for the device @@ -667,10 +667,14 @@ def lend_write(k, request): guard = PageCopyGuard(_lender.drv, device, lambda: leased_ranges(kit, mgr_b, caches), block) monkeypatch.setattr(_lender, "drv", guard) with kit.gated_stream(mgr_b._stream, open_after=REB_GATE_SECONDS) as gate: - leases["marked_write"].mark_arrived(masks["marked_write"]) # its copy queues - leases["marked_write"].release() - leases["queued_read"] = lender_b.lend_read(reqs["queued_read"], 0, REB_END) - leases["waiting_read"] = lender_b.lend_read(reqs["waiting_read"], 0, REB_END) + # Its return queues the copy of its arrived rows. + kit.return_write(lender_b, leases["queued_write"], masks["queued_write"]) + leases["queued_read"] = lender_b.lend_read( + reqs["queued_read"], 0, REB_END, **kit.HOST_NAMED + ) + leases["waiting_read"] = lender_b.lend_read( + reqs["waiting_read"], 0, REB_END, **kit.HOST_NAMED + ) waiting = ( leases["waiting_read"].poll() is None and leases["waiting_read"].failure is None ) @@ -686,7 +690,7 @@ def lend_write(k, request): lowest = min( s for k in reqs - if not (stay and k in ("waiting_read", "unmarked_write")) + if not (stay and k in ("waiting_read", "held_write")) for s in before[k][full] if s >= 0 ) @@ -720,19 +724,19 @@ def lend_write(k, request): dev = kit.DevicePages(mgr_b) # the pools changed size dev_a = kit.DevicePages(mgr_a) # The queued read's copy was queued before the rebalance. - view_q = ready_view(leases["queued_read"], mgr_b) if granted else None + view_q = kit.ready_view(leases["queued_read"], mgr_b) if granted else None staged_q = ( staged_right(kit, mgr_b, lender_b, view_q, reqs["queued_read"]) if view_q else None ) - # The unmarked write is marked after the rebalance. - leases["unmarked_write"].mark_arrived(masks["unmarked_write"]) - leases["unmarked_write"].release() + # The held write is returned after the rebalance. + kit.return_write(lender_b, leases["held_write"], masks["held_write"]) mgr_b._stream.synchronize() - # The waiting read is granted once the queued read's slots return. + # The waiting read is granted at a round of progress once slots return. leases["queued_read"].release() view_w = None for _ in range(3): mgr_b._stream.synchronize() + lender_b.progress() view_w = leases["waiting_read"].poll() if view_w is not None or leases["waiting_read"].failure is not None: break @@ -740,13 +744,14 @@ def lend_write(k, request): staged_right(kit, mgr_b, lender_b, view_w, reqs["waiting_read"]) if view_w else None ) mgr_b._stream.synchronize() + lender_b.progress() resume_bounds = {} - for k in ("unmarked_write", "marked_write"): - got = lender_b.resume_bounds(reqs[k]) + for k in ("held_write", "queued_write"): + got = lender_b.poll_resume_bounds(reqs[k]) resume_bounds[k] = None if got is None else tuple(got) arrived = { k: landed(kit, mgr_b, reqs[k], rows[k], mgr_a, pubs[k], dev, dev_a) - for k in ("unmarked_write", "marked_write") + for k in ("held_write", "queued_write") } # Bystanders keep their own bytes wherever their pages went. bystanders_kept = [] @@ -760,7 +765,7 @@ def lend_write(k, request): bystanders_kept.append(dev.read(lg, slot) == want) # ResumeBounds counts only what landed: every block a start at max_position reads. unbacked = {} - for k in ("unmarked_write", "marked_write"): + for k in ("held_write", "queued_write"): if resume_bounds[k] is None: unbacked[k] = ["unsettled"] continue @@ -777,8 +782,8 @@ def lend_write(k, request): assert all(resumed), f"a cache did not resume: {resumed}" assert waiting and granted, "the leases were not in the states the scenario needs" if stay: - assert any(moved["queued_read"]) and any(moved["marked_write"]), moved - assert not any(moved["waiting_read"]) and not any(moved["unmarked_write"]), moved + assert any(moved["queued_read"]) and any(moved["queued_write"]), moved + assert not any(moved["waiting_read"]) and not any(moved["held_write"]), moved else: assert all(any(m) for m in moved.values()), f"leased pages stayed: {moved}" # What the docstrings promise. @@ -790,7 +795,7 @@ def lend_write(k, request): ) assert all(staged_q.values()), f"the queued read staged wrong bytes: {staged_q}" # Both writes cover [0, REB_END) at a history of REB_END, so they lend the same rows. - for k in ("unmarked_write", "marked_write"): + for k in ("held_write", "queued_write"): assert set(arrived[k]) == read_rows, ( f"the {k} lent rows {sorted(arrived[k])}, not {sorted(read_rows)}" ) @@ -806,17 +811,17 @@ def lend_write(k, request): assert view_w is not None, ( f"the waiting read failed with its pages in place: {leases['waiting_read'].failure}" ) - assert all(arrived["unmarked_write"].values()), ( - f"the unmarked write lost rows that stayed: {arrived['unmarked_write']}" + assert all(arrived["held_write"].values()), ( + f"the held write lost rows that stayed: {arrived['held_write']}" ) - assert resume_bounds["unmarked_write"][0] == REB_END, ( - f"resume_bounds misses rows that landed: {resume_bounds['unmarked_write']}" + assert resume_bounds["held_write"][0] == REB_END, ( + f"resume_bounds misses rows that landed: {resume_bounds['held_write']}" ) assert all(not v for v in unbacked.values()), ( f"resume_bounds counts rows that never landed: {unbacked}" ) - assert all(arrived["marked_write"].values()), ( - f"the queued write's rows were lost: {arrived['marked_write']}" + assert all(arrived["queued_write"].values()), ( + f"the queued write's rows were lost: {arrived['queued_write']}" ) assert all(bystanders_kept), "a bystander's pages changed" @@ -829,7 +834,7 @@ def test_a_rebalance_with_every_kind_of_lease_in_flight( rebalance_with_leases_in_flight(kit, host_tier_manager, monkeypatch, mechanism, stay=stay) -# -- (g) a mark arriving after the window passed the lent blocks -------------------------------- +# -- (g) a return arriving after the window passed the lent blocks ------------------------------ LATE_PROMPT = list(range(2000, 2161)) # five whole blocks and one token LATE_END = 2 * TPB # the lease covers blocks 0 and 1 @@ -860,9 +865,9 @@ def locked_by(kit, mgr, requests): return out -def late_mark(kit, host_tier_manager, attach, windows, fillers, keep): +def late_return(kit, host_tier_manager, attach, windows, fillers, keep): """One run, ``fillers`` one-block requests allocated first and kept or freed by ``keep``, which - moves the slots the target gets. What the marks' copy wrote, or ``None`` when no lent block went + moves the slots the target gets. What the return copied, or ``None`` when no lent block went to host under its lent GPU slot's number.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender @@ -876,15 +881,15 @@ def late_mark(kit, host_tier_manager, attach, windows, fillers, keep): if not keep: first.free() target = kit.admitted(mgr, TARGET, LATE_PROMPT) - lender = attach(mgr, tokens_per_fetch=LATE_END) - lease = lender.lend_write(target, 0, LATE_END) + lender = attach(mgr, guaranteed_fetch_tokens=LATE_END) + lease = lender.lend_write(target, 0, LATE_END, **kit.HOST_NAMED) view = lease.poll() assert view is not None, lease.failure kv = kit.kv(mgr, target) lent = kit.pages(kv, sliding)[:2] assert [int(s) for s in kv.get_base_page_indices(sliding)[:2]] == lent kit.stage(lender, view, FETCHED) - # Before the marks, a served prefix moves the history past the lent blocks of the + # Before the return, a served prefix moves the history past the lent blocks of the # sliding group, as a KV connector's reservation does: they keep their pages only held. assert mgr._resize_for_served_prefix(target, kv, LATE_HISTORY, LATE_HISTORY + 1) stale = kit.stale_blocks(mgr, sliding, kv.history_length) @@ -903,8 +908,7 @@ def late_mark(kit, host_tier_manager, attach, windows, fillers, keep): assert kit.tier_used(mgr, 1) >= host_before + len(moved), "nothing reached host" aliased = [o for o in moved if now[o] == lent[o]] if not aliased: - lease.mark_arrived(view.row_masks(False)) - lease.release() + kit.return_write(lender, lease, view.row_masks(False)) return None dev = kit.DevicePages(mgr) watched = {lent[o]: owners[(group, lent[o])] for o in aliased} @@ -912,8 +916,7 @@ def late_mark(kit, host_tier_manager, attach, windows, fillers, keep): log = DeviceCopyGuard(_lender.drv) with pytest.MonkeyPatch.context() as patch: patch.setattr(_lender, "drv", log) - lease.mark_arrived(view.row_masks(True)) - lease.release() + kit.return_write(lender, lease, view.row_masks(True)) mgr._stream.synchronize() written = [device_slot(mgr, dst) for dst in log.copies] hit = sorted({w for w in written if w is not None and w in owners}) @@ -924,15 +927,15 @@ def late_mark(kit, host_tier_manager, attach, windows, fillers, keep): first.free() -def test_a_late_mark_writes_no_slot_its_window_left(kit, host_tier_manager, attach=attach): +def test_a_late_return_writes_no_slot_its_window_left(kit, host_tier_manager, attach=attach): # Slots go out lowest first, so a short search finds a lent block on host under its lent GPU # slot's number while another request locks that GPU slot. runs = [(w, n, k) for w in ([WINDOW], [WINDOW, 256]) for k in (True, False) for n in range(4)] for windows, fillers, keep in runs: - result = late_mark(kit, host_tier_manager, attach, windows, fillers, keep) + result = late_return(kit, host_tier_manager, attach, windows, fillers, keep) if result is not None: break else: raise RuntimeError("no run put a lent block on host under its lent GPU slot's number") - assert not result["hit"], f"the marks' copy wrote slots other requests lock: {result['hit']}" + assert not result["hit"], f"the return's copy wrote slots other requests lock: {result['hit']}" assert not result["changed"], f"requests {result['changed']} now hold the fetch" diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_in_place_lending.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_in_place_lending.py new file mode 100644 index 000000000000..ac4ca01beeaa --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_in_place_lending.py @@ -0,0 +1,3360 @@ +# 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. +"""Leases in place over real KV cache managers (``gpu``) and over stand-ins (``cpu_only``).""" + +import random +import weakref +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + ALL_ARRIVED, + LeaseView, + LentLayerGroupRows, + MemoryKind, + MemoryRegion, + RowSelection, + StagingOptions, + _lender, + _manager, + _ranges, + _types, + attach_lender, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import Identity +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import DevicePool, derive_layout +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._page_locks import LockMode +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._slots import SlotQueue +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttnLifeCycle + +gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") + +TPB = 32 +WINDOW = 64 +PROMPT = list(range(1000, 1097)) # three whole blocks and one token +PROMPT_LEN = len(PROMPT) +OTHER_PROMPT = list(range(5000, 5097)) +LONG_PROMPT = list(range(7000, 7000 + 6 * TPB)) # six whole blocks +LONG = 10 * TPB +END = 96 +BLOCKS = END // TPB +SOURCE, TARGET, OTHER = 1, 2, 3 +RID = 7 # the stand-ins' request +DEVICE, HOST = MemoryKind.DEVICE, MemoryKind.HOST +NAMED, ALL = RowSelection.NAMED, RowSelection.ALL +STALE, FRESH = 0x3C, 0xA7 + + +# -- stand-ins ---------------------------------------------------------------------------------- + + +class Event: + """A stand-in CUDA event: complete once ``done`` is set; waiting on it fails the test.""" + + def __init__(self, done=False): + self.done = done + self.queries = 0 + + def query(self): + self.queries += 1 + return self.done + + def synchronize(self): + raise AssertionError("a lease waited on its event") + + +class Connector: + """A stand-in KV cache connector manager, loading into the requests in ``loading`` and holding + a prefix reservation for those in ``reserved``.""" + + def __init__(self, capacity_only=False): + self.capacity_only = capacity_only + self.loading = set() + self.reserved = set() + + def has_pending_load(self, request): + return request.py_request_id in self.loading + + def get_prefix_reservation(self, request): + return SimpleNamespace() if request.py_request_id in self.reserved else None + + +class Manager: + """A stand-in manager: a layer group per window, ``None`` for full attention, on pool groups 0 + and 1 in turn, with the runtime's stale ranges, and the requests' caches by request.""" + + # No cache of this stand-in is on an allocation side stream. + _on_allocation_stream = frozenset() + + def _commit_held_back(self, kv): + """No commit is held back here.""" + + def __init__(self, windows=(None,) * 3, sinks=None, connector=None): + sinks = sinks or (0,) * len(windows) + cycles = [AttnLifeCycle.make(w, s * TPB, TPB) for w, s in zip(windows, sinks)] + self.tokens_per_block = TPB + self.impl = SimpleNamespace( + get_life_cycle_pool_group_indices=lambda: [lg % 2 for lg in range(len(windows))] + ) + self.kv_cache_map = {} + self.kv_connector_manager = connector + self.enable_block_reuse = False + self._can_publish_block_reuse = False + + def stale_block_range(lg, history): + stale = cycles[lg].get_stale_range(history, TPB) + return stale.beg, stale.end + + self._stale_block_range = stale_block_range + + +def layout(windows=(None,) * 3, sinks=None): + """The layout of ``Manager(windows, sinks)``: pool group 0 of two pools, 1 of one, ten slots + each.""" + count = len(windows) + pool_group_of = tuple(lg % 2 for lg in range(count)) + return SimpleNamespace( + tokens_per_block=TPB, + windows=tuple(windows), + sink_blocks=tuple(sinks or (0,) * count), + recurrent=(False,) * count, + pool_groups=tuple(sorted(set(pool_group_of))), + pool_group_of=pool_group_of, + num_layer_groups=count, + named_groups=(True,) * count, + device_pools={ + 0: (DevicePool(0, 0, 1 << 30, 256, 10), DevicePool(0, 1, 2 << 30, 128, 10)), + 1: (DevicePool(1, 0, 3 << 30, 512, 10),), + }, + ) + + +def identity(scope=b"scope", count=3): + return Identity(scope, bytes(range(32)), [(lg,) for lg in range(count)], [(1, 0)] * count) + + +def lender_over(manager, lay=None, no_staging=None): + """A lender over ``manager`` and ``lay``, ``layout()`` if not given: a staging region per pool + group, none where ``no_staging`` says why, then the device regions an attach lists.""" + lay = layout() if lay is None else lay + ident = identity(count=lay.num_layer_groups) + regions = [] + for g in lay.pool_groups if no_staging is None else (): + groups = tuple(lg for lg, pg in enumerate(lay.pool_group_of) if pg == g) + regions.append( + MemoryRegion(f"staging{g}", 4096 * (g + 1), 4096, 1024, 4, groups, MemoryKind.HOST) + ) + regions += _lender._device_regions(manager, lay, ident) + ref = weakref.ref(manager) + slots = SlotQueue({g: 4 for g in lay.pool_groups}) + return _lender._ManagerLender( + ref, lay, ident, tuple(regions), slots, ref, no_staging=no_staging + ) + + +def cache(blocks=4, history=None, committed=0, pages=None): + """A stand-in active cache of ``history`` tokens, ``committed`` by default, whose blocks lock + ``pages``, ``0, 1, ...`` by default, in every layer group, its committed tokens counting in + each.""" + pages = list(range(blocks)) if pages is None else list(pages) + return SimpleNamespace( + Status=SimpleNamespace(OPEN=0, CLOSED=1), + status=0, + num_committed_tokens=committed, + reuse_status=[SimpleNamespace(complete=True, endpoint=committed)] * 4, + history_length=committed if history is None else history, + is_active=True, + num_blocks=blocks, + tokens_per_block=TPB, + enable_swa_scratch_reuse=False, + get_base_page_indices=lambda lg: pages, + get_aggregated_page_indices=lambda lg, valid_only=False: pages, + ) + + +def entries(ordinals=(0, 1, 2), pages=(5, 6, 7)): + """The same blocks on the same pages in each of the three layer groups.""" + o, p = np.array(ordinals, dtype=np.int64), np.array(pages, dtype=np.int64) + return _lender._LentBlockEntries( + [0, 1, 2], [o.copy() for _ in range(3)], [p.copy() for _ in range(3)] + ) + + +def request(computed=0, served=0, prompt=PROMPT_LEN, end=None): + """Stand-in request ``RID`` of a ``prompt``-token prompt whose context on this worker runs to + ``end``, the prompt's by default, ``computed`` tokens of it computed, ``served`` served.""" + return SimpleNamespace( + py_request_id=RID, + prompt_len=prompt, + py_multimodal_data=None, + py_served_prefix_end=served, + context_current_position=computed, + context_remaining_length=(prompt if end is None else end) - computed, + ) + + +def in_place(lender, manager, kv, kind, lent=None, fetch=None, end=END, **kwargs): + """``_lend_in_place`` of ``request()`` over ``[0, end)``, ``entries()`` by default.""" + lent = entries() if lent is None else lent + return lender._lend_in_place( + manager, request(), kv, kind, lent, None, fetch, start=0, end=end, **kwargs + ) + + +def return_write(lender, req, lease, arrived): + """``return_writes`` of the one entry ``(lease, arrived, True)``: its bounds, which + ``poll_resume_bounds`` then gives ``req`` too.""" + (bounds,) = lender.return_writes([(lease, arrived, True)]) + assert lender.poll_resume_bounds(req) == bounds, bounds + return bounds + + +def refused(lease, why): + """``lease`` failed at the call for ``why`` and holds no lock; released.""" + assert lease.failure is not None and why in lease.failure, lease.failure + assert lease._lock is None, "a lease failed at the call holds a lock" + lease.release() + + +def inputs_refuse_nothing(monkeypatch, lender): + """``lender`` with no refusal from a request's inputs: reuse keys, unnamed input, context + outputs or a bidirectional split.""" + monkeypatch.setattr(lender, "_reuse_keys_for", lambda *args: None) + monkeypatch.setattr(lender, "_unnamed_input_refusal", lambda request: None) + monkeypatch.setattr(lender, "_returns_context_outputs", lambda request: False) + monkeypatch.setattr(lender, "_bidirectional_split_refusal", lambda request, end: None) + + +@pytest.fixture +def recorded(monkeypatch): + """The CUDA events made while the test runs: a lend makes none, as its readiness is its + caller's event.""" + events = [] + + class Made: + def __init__(self, *args, **kwargs): + events.append(self) + + def record(self, stream=None): + pass + + def query(self): + return True + + monkeypatch.setattr(torch.cuda, "Event", Made) + return events + + +# -- real managers ------------------------------------------------------------------------------ + + +def attach(mgr, fetches=1): + return attach_lender( + mgr, scope=b"in-place", staging=StagingOptions(END, capacity_in_fetches=fetches) + ) + + +def kept_page_ordinals(window, end, tokens_per_block, cached=0): + """Pages intersecting the next query's retained token interval after ``end`` tokens.""" + first_token = 0 if window is None else max(0, end + 1 - window) + return list( + range( + max(cached // tokens_per_block, first_token // tokens_per_block), + (end + tokens_per_block - 1) // tokens_per_block, + ) + ) + + +def device_tensor(address: int, nbytes: int) -> torch.Tensor: + from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor + + return convert_to_torch_tensor(TensorWrapper(int(address), torch.uint8, shape=(int(nbytes),))) + + +def slot_parts(lender, run, slot): + """``(address, nbytes)`` of ``slot`` in each region of ``run``, in region order.""" + regions = lender.regions + return [ + (regions[r].address + int(slot) * regions[r].slot_bytes, regions[r].slot_bytes) + for r in run.region_indices + ] + + +def entry_bytes(lender, run, i) -> bytes: + return b"".join( + device_tensor(address, nbytes).cpu().numpy().tobytes() + for address, nbytes in slot_parts(lender, run, run.slots[i]) + ) + + +def write_through(lender, view, byte, masks=None) -> None: + """Fill each entry of ``view`` that ``masks`` marks, every entry by default, with ``byte`` + through its addresses, on a stream of its own.""" + side = torch.cuda.Stream() + with torch.cuda.stream(side): + for run, mask in zip(view.groups, masks or [None] * len(view.groups)): + for i in range(len(run)): + if mask is None or mask[i]: + for address, nbytes in slot_parts(lender, run, run.slots[i]): + device_tensor(address, nbytes).fill_(byte) + side.synchronize() + + +def device_regions(lender): + return tuple(r for r in lender.regions if r.memory_kind is MemoryKind.DEVICE) + + +def rows_giving(patched, lender, mgr, requests, block, page): + """Make the rows of the page-index buffer that the lender reads a write's pages from give + ``page`` at ``block`` of each of ``requests``, in every layer group: a copy of the buffer, the + manager's own rows unchanged.""" + rows = lender._index_rows.copy() + for request in requests: + rows[:, _manager.index_row(mgr, request.py_request_id), 0, block] = page + # Set to themselves first, so the context's end puts back the manager's rows. + for name in ("_index_rows", "_index_words"): + patched.setattr(lender, name, getattr(lender, name)) + lender._read_index_rows(rows) + + +def pages_giving(patched, lender, mgr, requests, block, page, readers=("pages", "page_list")): + """Make every way the lender reads the pages it lends give ``page`` at ``block`` of each of + ``requests``: the cache's ``readers`` and its page-index row.""" + caches = [mgr.kv_cache_map[r.py_request_id] for r in requests] + + def giving(read): + def pages(kv_cache, lg): + found = np.array(read(kv_cache, lg), dtype=np.int64) + if any(kv_cache is lent for lent in caches) and len(found) > block: + found[block] = page + return found + + return pages + + for name in readers: + patched.setattr(_manager, name, giving(getattr(_manager, name))) + rows_giving(patched, lender, mgr, requests, block, page) + + +def step(mgr, request, tokens): + """One context step of ``tokens``, as the executor ends it: the position moves, then the + manager updates the cache.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + request.context_chunk_size = tokens + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + mgr.update_context_resources(batch) + + +def computed_to(mgr, request, end): + """``request``, admitted first where the manager holds no cache of it, with its context + computed to ``end`` in one step, as the executor runs it.""" + if request.py_request_id not in mgr.kv_cache_map: + assert mgr.prepare_context(request) + tokens = end - request.context_current_position + assert mgr.resize_context(request, tokens) + step(mgr, request, tokens) + assert mgr.kv_cache_map[request.py_request_id].history_length == end + return request + + +def generate(mgr, request, token=7): + """One generation step as the executor runs it: the scheduler grows the cache for a token, the + sampler adds the token, then the manager updates the cache.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + assert mgr.try_allocate_generation(request) + request.add_new_token(token, 0) + batch = ScheduledRequests() + batch.append_generation_request(request) + mgr.update_resources(batch) + + +def receiving(kit, mgr, request_id, tokens, prepare=False, **attributes): + """A generation-init request of ``tokens`` with ``attributes``, admitted to receive them, and + the event a receive into it waits for: the admission's, or with ``prepare`` one recorded after + the preparation a hybrid manager's executor makes, which resets its state slots.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestType + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + request = kit.make_request( + request_id, tokens, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + ) + for name, value in attributes.items(): + setattr(request, name, value) + assert mgr.prepare_disagg_gen_init(request) + if not prepare: + return request, mgr.take_disagg_receive_ready(request) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + mgr.prepare_resources(batch) + # The preparation waits for the admission's event and resets the state slots on this stream. + reset = torch.cuda.Event() + reset.record(torch.cuda.current_stream()) + torch.cuda.current_stream().synchronize() + return request, reset + + +def two_beams(request_id, **kwargs): + """A request of ``PROMPT`` searching two beams.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig + + request = LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=list(PROMPT), + sampling_config=SamplingConfig(2), + is_streaming=False, + **kwargs, + ) + assert request.py_beam_width == 2, request.py_beam_width + return request + + +def slots_by_block(view): + """``(layer group, block ordinal) -> slot`` of each entry of ``view``.""" + return { + (run.layer_group_index, int(ordinal)): int(slot) + for run in view.groups + for ordinal, slot in zip(run.page_ordinals.tolist(), run.slots.tolist()) + } + + +def holds(dev, lent, byte_of) -> bool: + """Whether each lent page holds the bytes its block was given.""" + return all( + dev.read(lg, slot) == bytes([byte_of(ordinal)]) * dev.page_bytes(lg) + for (lg, ordinal), slot in lent.items() + ) + + +def state_group(mgr) -> int: + (lg,) = [lg for lg, flag in enumerate(derive_layout(mgr).recurrent) if flag] + return lg + + +def state_tensors(mgr, slot): + """The state each Mamba layer's kernels read for ``slot``: SSM and convolution states.""" + return [t[slot] for t in list(mgr.all_ssm_states) + list(mgr.all_conv_states)] + + +def helix_rank(rank): + from tensorrt_llm.mapping import CpType, Mapping + + return Mapping(world_size=2, rank=rank, cp_size=2, cp_config={"cp_type": CpType.HELIX}) + + +def buffers(lender, view): + """``(buffer, block ordinal, bytes)`` of each public buffer descriptor of each lent row, + padding included.""" + layout = lender.page_layout + for rows in view.groups: + group = layout.layer_groups[rows.layer_group_index] + for buffer in layout.buffers: + if buffer.pool_group != group.pool_group or buffer.layer not in group.layers: + continue + index = layout.device_region_index[buffer.pool_group, buffer.pool_index] + assert index in rows.region_indices, (index, rows.region_indices, buffer) + region = lender.regions[index] + for ordinal, slot in zip(rows.page_ordinals.tolist(), rows.slots.tolist()): + address = region.address + slot * region.slot_bytes + buffer.slot_offset + yield buffer, ordinal, device_tensor(address, buffer.size) + + +def fill_buffers(lender, view, stream, byte): + with torch.cuda.stream(stream): + for _, _, tensor in buffers(lender, view): + tensor.fill_(byte) + stream.synchronize() + + +def remote_tail_layout(): + from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout + + return CSA2Layout((0, 2, 2, 1, 1, 1), (1, 3), (1, 3, 4), candidate_source_layer_id=3) + + +def csa2_manager(kit, model, reuse=False, **kwargs): + """``with csa2_manager(...) as mgr:`` a small FP8 CSA2 manager of ``model`` for disaggregated + serving, of 128-token blocks, ``kwargs`` over its defaults.""" + from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.mapping import Mapping + + config = dict( + num_layers=len(model.compress_ratios), + tokens_per_block=128, + max_seq_len=512, + max_batch_size=2, + max_num_tokens=1024, + vocab_size=129280, + dtype=DataType.BF16, + mapping=Mapping(), + layout=model, + is_disagg=True, + ) + config.update(kwargs) + return kit.managed( + lambda: CSA2CacheManager( + KvCacheConfig(max_gpu_total_bytes=128 << 20, enable_block_reuse=reuse, dtype="fp8"), + CacheType.SELFKONLY, + **config, + ) + ) + + +# -- regions and the attach --------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_an_attach_lists_one_device_region_per_pool_after_the_staging_regions(): + manager = Manager() + regions = _lender._device_regions(manager, layout(), identity()) + assert [(r.address, r.slot_bytes, r.slot_count, r.nbytes) for r in regions] == [ + (1 << 30, 256, 10, 2560), + (2 << 30, 128, 10, 1280), + (3 << 30, 512, 10, 5120), + ] + assert [r.layer_group_indices for r in regions] == [(0, 2), (0, 2), (1,)] + assert all(r.memory_kind is MemoryKind.DEVICE for r in regions) + labels = [r.layout_label for r in regions] + staging = identity().region_label((0, 2)) + assert labels[0] == f"{staging}:pool0" and labels[1] == f"{staging}:pool1" + assert len(set(labels)) == 3 and staging not in labels + other = _lender._device_regions(manager, layout(), identity(b"another scope")) + assert [r.layout_label for r in other] == labels, "a label follows the layout, not the scope" + # A manager naming no rebalance policy gets them too: a lender's pools are not rebalanced. + assert _lender._device_regions(SimpleNamespace(), layout(), identity()) == regions + lender = lender_over(manager) + assert [r.memory_kind for r in lender.regions] == [HOST, HOST, DEVICE, DEVICE, DEVICE] + assert lender._device_regions_of_pool_group == {0: (2, 3), 1: (4,)} + + +@gpu +def test_the_device_regions_are_the_manager_s_pools_after_the_staging_regions(kit, real_manager): + with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: # two pool groups + lender = attach(mgr) + assert lender.regions == kit.staging_regions(lender) + device_regions(lender) + group_of = kit.pool_group_of(mgr) + expected = [] + for pg in sorted(mgr.impl.pool_group_descs, key=lambda d: int(d.pool_group_index)): + g = int(pg.pool_group_index) + groups = tuple(lg for lg, pool_group in enumerate(group_of) if pool_group == g) + for pool in sorted(pg.pools, key=lambda p: int(p.pool_index)): + slots = int(pg.num_slots) + expected.append( + (int(pool.base_address), slots * int(pool.slot_bytes), slots, groups) + ) + got = [ + (r.address, r.nbytes, r.slot_count, r.layer_group_indices) + for r in device_regions(lender) + ] + assert got == expected + labels = [r.layout_label for r in lender.regions] + assert len(set(labels)) == len(labels), "two regions share a label" + + +@gpu +def test_device_memory_is_lent_with_pool_rebalancing_configured(kit, real_manager): + """The executor does not rebalance a lender's manager's pools, so the lender lends in place.""" + with real_manager(enable_kv_pool_rebalance=True) as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + regions = [(r.memory_kind, r.layout_label) for r in lender.regions] + assert device_regions(lender), f"no device region among {regions}" + lease = lender.lend_read(source, 0, END, **kit.DEVICE_NAMED) + assert lease.failure is None, lease.failure + assert kit.ready_view(lease, mgr).num_rows > 0 + lease.release() + + +def attach_checks(monkeypatch, plain, **attributes): + """A stand-in manager with ``attributes`` whose layout is ``plain``, for the attach's checks.""" + count = plain.num_layer_groups + manager = SimpleNamespace( + mapping=SimpleNamespace(cp_size=1, pp_size=1, has_cp_helix=lambda: False), + _model_has_attention_window=False, + _model_has_sparse_buffers=False, + _model_attention_window_pattern=(), + num_layers=count, + kv_cache_manager_py_config=SimpleNamespace(layers=[]), + impl=SimpleNamespace(layer_grouping=[[] for _ in range(count)]), + **attributes, + ) + monkeypatch.setattr(_manager, "require_v2", lambda manager: None) + monkeypatch.setattr(_lender, "derive_layout", lambda manager: plain) + return manager + + +def checked_layout(sparse): + """A layout of full-attention layer groups, those ``sparse`` marks holding sparse buffers.""" + count = len(sparse) + return SimpleNamespace( + cp=(0, 1), + num_layer_groups=count, + recurrent=(False,) * count, + windows=(None,) * count, + sink_blocks=(0,) * count, + sparse=tuple(sparse), + ) + + +@pytest.mark.cpu_only +def test_a_lossy_cold_page_codec_names_nothing_and_passes_the_attach(monkeypatch): + """A lossy cold-page codec keeps the lender from naming blocks, not from attaching.""" + plain = checked_layout((False,)) + manager = attach_checks(monkeypatch, plain, _has_lossy_cold_page_codec=True) + monkeypatch.setattr(_manager, "commits_blocks", lambda manager: True) + monkeypatch.setattr(_manager, "prompt_lookahead", lambda manager: 0) + monkeypatch.setattr(_manager, "model_has_recurrent_state", lambda manager: False) + assert _lender._checked_layout(manager) is plain + reason = _lender._no_names(manager, plain) + assert reason is not None and "lossy" in reason, reason + manager._has_lossy_cold_page_codec = False + assert _lender._no_names(manager, plain) is None + + +@pytest.mark.cpu_only +def test_an_attach_refuses_sparse_layer_groups_and_names_them(monkeypatch): + manager = attach_checks(monkeypatch, checked_layout((False, True, True))) + with pytest.raises(ValueError, match=r"layer groups \[1, 2\] hold sparse buffers"): + _lender._checked_layout(manager) + dense = checked_layout((False,) * 3) + monkeypatch.setattr(_lender, "derive_layout", lambda manager: dense) + assert _lender._checked_layout(manager) is dense + + +@gpu +def test_a_lossy_cold_page_codec_still_lends_a_computed_request_s_pages_in_place(kit, real_manager): + from tensorrt_llm.llmapi.llm_args import ColdPageQuantizationCompressionConfig + from tensorrt_llm.runtime.kv_cache_manager_v2 import create_default_kv_cache_cold_page_codec + + provider = SimpleNamespace( + config=ColdPageQuantizationCompressionConfig(), + create_cold_page_codec=lambda *args, **kwargs: create_default_kv_cache_cold_page_codec(), + ) + with real_manager(cold_page_codec_provider=provider) as mgr: + assert mgr._has_lossy_cold_page_codec, "the manager took the codec for a lossless one" + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + lease = lender.lend_read(source, 0, len(PROMPT), **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + view = kit.ready_view(lease, mgr) + kv = kit.kv(mgr, source) + for run in view.groups: + lg = run.layer_group_index + assert run.object_keys is None + assert run.page_ordinals.tolist() == list(range(BLOCKS + 1)) + assert run.slots.tolist() == kit.pages(kv, lg)[: BLOCKS + 1] + for i, ordinal in enumerate(run.page_ordinals.tolist()): + assert entry_bytes(lender, run, i) == kit.page(mgr, source, lg, ordinal) + lease.release() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "no_staging", [None, "its connector loads"], ids=["stages", "stages_nothing"] +) +def test_device_memory_is_lent_by_name_or_not_and_host_memory_only_by_name(no_staging): + lender = lender_over(Manager(), no_staging=no_staging) + lender._check_memory_and_rows(DEVICE, NAMED) + lender._check_memory_and_rows(DEVICE, ALL) + with pytest.raises(ValueError, match="RowSelection.ALL is not lent.*staging relays named"): + lender._check_memory_and_rows(HOST, ALL) + if no_staging is None: + lender._check_memory_and_rows(HOST, NAMED) + else: + with pytest.raises(ValueError, match="stages nothing, as its connector loads"): + lender._check_memory_and_rows(HOST, NAMED) + + +# -- views -------------------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_an_in_place_view_places_each_row_by_its_page_in_its_pool_group_regions(): + lender = lender_over(Manager()) + lent = entries(ordinals=(0, 3), pages=(9, 4)) + view = lender._in_place_view(lent, None, 0, 0) + assert [g.layer_group_index for g in view.groups] == [0, 1, 2] + assert [g.region_indices for g in view.groups] == [(2, 3), (4,), (2, 3)] + assert view.groups[0].region_indices is view.groups[2].region_indices, "one per pool group" + for group in view.groups: + assert isinstance(group, LentLayerGroupRows) + assert group.slots.tolist() == [9, 4] and group.page_ordinals.tolist() == [0, 3] + assert group.object_keys is None + keys = [np.full((2, 32), lg, dtype=np.uint8) for lg in range(3)] + named = lender._in_place_view(lent, keys, 0, 0) + for lg, group in enumerate(named.groups): + np.testing.assert_array_equal(group.object_keys, lender._identity.object_keys(lg, keys[lg])) + # Entry i's bytes: slot ``slots[i]`` of each region, in region order. + regions = lender.regions + addresses = [ + regions[r].address + 9 * regions[r].slot_bytes for r in view.groups[0].region_indices + ] + assert addresses == [(1 << 30) + 9 * 256, (2 << 30) + 9 * 128] + + +@pytest.mark.cpu_only +def test_every_lend_of_a_pool_group_names_its_regions_by_one_tuple(): + """Rows of one pool group lent in place name one tuple of its device regions, and staged rows + one tuple of its staging region, lend after lend.""" + lender = lender_over(Manager()) + lent = entries(ordinals=(0, 3), pages=(9, 4)) + first, second = (lender._in_place_view(lent, None, 0, 0) for _ in range(2)) + assert all(a.region_indices is b.region_indices for a, b in zip(first.groups, second.groups)) + lender._assign_staging_slots(lent, None) + keys = [np.full((2, 32), lg, dtype=np.uint8) for lg in range(3)] + staged = [lender._build_view(lent, keys, 0, 0) for _ in range(2)] + assert all(a.region_indices is b.region_indices for a, b in zip(*(v.groups for v in staged))) + assert [group.region_indices for group in staged[0].groups] == [(0,), (1,), (0,)] + assert {lender.regions[r].memory_kind for r in (0, 1)} == {MemoryKind.HOST} + + +@gpu +def test_an_in_place_read_lends_the_request_s_own_pages_by_name(kit, real_manager): + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + staged = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + staged_view = kit.ready_view(staged, mgr) + lease = lender.lend_read(source, 0, END, **kit.DEVICE_NAMED) + view = kit.ready_view(lease, mgr) + kv = kit.kv(mgr, source) + descs = {int(pg.pool_group_index): pg for pg in mgr.impl.pool_group_descs} + expected_groups = list(range(len(kit.pool_group_of(mgr)))) + assert [group.layer_group_index for group in view.groups] == expected_groups + assert [group.layer_group_index for group in staged_view.groups] == expected_groups + for run, staged_run in zip(view.groups, staged_view.groups): + assert run.page_ordinals.tolist() == [0, 1, 2] + np.testing.assert_array_equal(run.page_ordinals, staged_run.page_ordinals) + kinds = {lender.regions[r].memory_kind for r in run.region_indices} + staged_kinds = {lender.regions[r].memory_kind for r in staged_run.region_indices} + assert (kinds, staged_kinds) == ({MemoryKind.DEVICE}, {MemoryKind.HOST}) + np.testing.assert_array_equal(run.object_keys, staged_run.object_keys) + lg = run.layer_group_index + pages = kit.pages(kv, lg) + assert run.slots.tolist() == [pages[o] for o in run.page_ordinals.tolist()] + pools = sorted(descs[kit.pool_group_of(mgr)[lg]].pools, key=lambda p: int(p.pool_index)) + for i, ordinal in enumerate(run.page_ordinals.tolist()): + manager_parts = [ + (int(p.base_address) + pages[ordinal] * int(p.slot_bytes), int(p.slot_bytes)) + for p in pools + ] + assert slot_parts(lender, run, run.slots[i]) == manager_parts + assert entry_bytes(lender, run, i) == kit.page(mgr, source, lg, ordinal) + lease.release() + staged.release() + + +def checked_by_the_constructors(view): + """``view`` rebuilt by the public constructors, which check what the lender's own construction + of it does not: equal, else ``AssertionError``; values out of the rows' contract raise.""" + names = "layer_group_index page_ordinals object_keys slots region_indices".split() + groups = tuple(LentLayerGroupRows(*(getattr(rows, n) for n in names)) for rows in view.groups) + rebuilt = LeaseView(groups, view.start, view.end) + assert repr(rebuilt) == repr(view) + return rebuilt + + +@gpu +def test_a_read_in_place_leaves_out_blocks_behind_a_window_and_blocks_without_a_page( + kit, real_manager, monkeypatch +): + """A read of every row lends no block behind a window, none without a page, no negative slot.""" + with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: + source = kit.published(mgr, SOURCE, LONG_PROMPT) + lender = attach(mgr) + kv = kit.kv(mgr, source) + history, blocks = kv.history_length, len(LONG_PROMPT) // TPB + stale = {lg: kit.stale_blocks(mgr, lg, history) for lg in range(kit.num_layer_groups(mgr))} + assert any(begin < end for begin, end in stale.values()), f"inconclusive: {stale}" + pageless = blocks - 2 # inside every window + pages_giving(monkeypatch, lender, mgr, [source], pageless, -1) + lease = lender.lend_read(source, 0, len(LONG_PROMPT), **kit.DEVICE_ALL) + view = kit.ready_view(lease, mgr) + for run in view.groups: + lg = run.layer_group_index + begin, end = stale[lg] + expected = [b for b in range(blocks) if not begin <= b < end and b != pageless] + assert run.page_ordinals.tolist() == expected, (lg, stale[lg]) + assert min(run.slots.tolist(), default=0) >= 0, run.slots.tolist() + pages = kit.pages(kv, lg) + assert run.slots.tolist() == [pages[b] for b in expected] + checked_by_the_constructors(view) + lease.release() + + +def view_arrays(view): + """Every array of ``view``, by (layer group, field).""" + return { + (group.layer_group_index, name): getattr(group, name) + for group in view.groups + for name in ("page_ordinals", "object_keys", "slots") + if getattr(group, name) is not None + } + + +def check_snapshot(view, before, records): + """``view``'s arrays read as ``before`` holds them, share no memory with ``records`` and can be + neither written nor made writable.""" + arrays = view_arrays(view) + assert arrays.keys() == before.keys(), (sorted(arrays), sorted(before)) + for key, array in arrays.items(): + np.testing.assert_array_equal(array, before[key], err_msg=f"{key} changed") + shared = [i for i, record in enumerate(records) if np.shares_memory(array, record)] + assert not shared, f"{key} shares memory with records {shared}" + for write in (lambda: setattr(array.flags, "writeable", True), lambda: array.fill(0)): + try: + write() + except ValueError: + continue + raise AssertionError(f"{key} took a write") + + +@pytest.mark.cpu_only +def test_an_in_place_view_is_a_snapshot_that_outlives_its_lease(): + """A release and a re-lend leave a view's arrays, copied from mutable records, as they were.""" + manager = Manager() + lender = lender_over(manager) + kv = cache() + keys = [np.full((3, 32), lg, dtype=np.uint8) for lg in range(3)] + lease = lender._lend_in_place(manager, request(), kv, "read", entries(), keys, start=0, end=END) + view = lease.poll() + assert view is not None, lease.failure + assert {name for _, name in view_arrays(view)} == {"page_ordinals", "object_keys", "slots"} + before = {key: array.copy() for key, array in view_arrays(view).items()} + records = [*lease._block_entries.block_ordinals, *lease._block_entries.device_pages] + lease.release() + again = lender._lend_in_place(manager, request(), kv, "read", entries(), keys, start=0, end=END) + assert again.poll() is not None + again.release() + check_snapshot(view, before, records) + + +@pytest.mark.cpu_only +def test_a_staging_view_is_a_snapshot_of_its_own(): + lender = lender_over(Manager()) + lent = entries() + lender._assign_staging_slots(lent, None) + keys = [np.full((3, 32), lg, dtype=np.uint8) for lg in range(3)] + view = lender._build_view(lent, keys, 0, 0) + assert {name for _, name in view_arrays(view)} == {"page_ordinals", "object_keys", "slots"} + assert all( + lender.regions[g.region_indices[0]].memory_kind is MemoryKind.HOST for g in view.groups + ) + before = {key: array.copy() for key, array in view_arrays(view).items()} + lender._build_view(lent, keys, 0, 0) + check_snapshot(view, before, [*lent.block_ordinals, *lent.staging_slots]) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("build", ["staging", "in_place"]) +def test_views_of_the_same_frozen_entries_share_no_array_object(build): + """Two views of the same frozen entries hold arrays of their own, not each other's or theirs.""" + lender = lender_over(Manager()) + # Entries as the lender keeps them, one group's rows the module's shared empty rows. + kept = _lender._LentBlockEntries( + [0, 1, 2], + [_types._frozen_int64((0, 1, 2)), _lender._NO_ROWS, _types._frozen_int64((4,))], + [_types._frozen_int64((5, 6, 7)), _lender._NO_ROWS, _types._frozen_int64((8,))], + ) + keys = [np.full((len(o), 32), lg, dtype=np.uint8) for lg, o in enumerate(kept.block_ordinals)] + if build == "staging": + lender._assign_staging_slots(kept, None) + views = [lender._build_view(kept, keys, 0, 0) for _ in range(2)] + else: + views = [lender._in_place_view(kept, keys, 0, 0) for _ in range(2)] + held = {id(a) for a in (*kept.block_ordinals, *kept.device_pages, *kept.staging_slots)} + lent = [id(array) for view in views for array in view_arrays(view).values()] + assert len(set(lent)) == len(lent), "two views share an array" + assert not held & set(lent), "a view holds an array of the entries" + + +@gpu +@pytest.mark.parametrize( + "lend", + ["DEVICE_ALL", "DEVICE_NAMED", "HOST_NAMED"], + ids=["in_place_every_row", "in_place_named", "staged_named"], +) +def test_no_two_leases_share_an_array_and_a_reshaped_one_reaches_no_other(kit, real_manager, lend): + """Leases share no array with each other or a lock; a reshaped one changes no other lease.""" + import warnings + + keywords = getattr(kit, lend) + end = END if keywords["rows"] is NAMED else len(PROMPT) + with real_manager() as one, real_manager() as two: + lenders = tuple(attach(mgr, fetches=3) for mgr in (one, two)) + requests = [ + (one, lenders[0], kit.published(one, SOURCE, PROMPT)), + (one, lenders[0], kit.published(one, OTHER, PROMPT)), + (two, lenders[1], kit.published(two, SOURCE, PROMPT)), + ] + leases = [lender.lend_read(request, 0, end, **keywords) for _, lender, request in requests] + views = [kit.ready_view(lease, mgr) for lease, (mgr, _, _) in zip(leases, requests)] + lent = [id(array) for view in views for array in view_arrays(view).values()] + assert len(set(lent)) == len(lent), "two leases share an array" + locks = [getattr(lease, "_lock", None) for lease in leases] + locked = {id(a) for lock in locks if lock is not None for a in lock.blocks.values()} + assert not locked & set(lent), "a lock holds an array a view holds" + lengths = [[len(group) for group in view.groups] for view in views] + mine = views[0].groups[0].page_ordinals + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + mine.shape = (1, len(mine)) + assert [[len(group) for group in view.groups] for view in views[1:]] == lengths[1:] + if keywords["memory_kind"] is DEVICE: + assert lenders[0]._is_locked(kit.kv(one, requests[0][2]), BLOCKS - 1, BLOCKS) + later = lenders[0].lend_read(requests[1][2], 0, end, **keywords) + assert [len(group) for group in kit.ready_view(later, one).groups] == lengths[1] + for lease in (*leases, later): + lease.release() + + +# -- readiness ---------------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_an_in_place_read_is_ready_once_its_event_completes_and_never_waits(): + manager = Manager() + lender = lender_over(manager) + kv = cache() + ready = Event() + assert not lender._is_locked(None) + lease = in_place(lender, manager, kv, "read", ready_event=ready) + assert isinstance(lease, _lender._InPlaceLease) and lease.failure is None + assert lender._locks.is_locked(kv) and lender._is_locked(kv, 2, 3) + assert lender._is_locked(None), "a lock went unseen by the query over every cache" + assert lease.poll() is None and lease.poll() is None and ready.queries == 2 + ready.done = True + view = lease.poll() + assert view is not None and lease.poll() is view and ready.queries == 3 + with pytest.raises(RuntimeError, match="return_writes is for write leases"): + return_write(lender, request(), lease, view.row_masks()) + assert lender._locks.is_locked(kv), "a refused return let the read go" + lease.release() + lease.release() + assert not lender._locks.is_locked(kv) and not lender._is_locked(None) + with pytest.raises(RuntimeError, match="after release"): + lease.poll() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("kind", ["write", "read"]) +@pytest.mark.parametrize("given", [False, True], ids=["no_event", "an_event"]) +def test_a_lease_waits_for_the_event_its_caller_gives_and_for_none_without_one( + recorded, kind, given +): + """A lease is ready once its caller's event completed, at once without one.""" + manager = Manager() + lender = lender_over(manager) + ready = Event() if given else None + lease = in_place(lender, manager, cache(), kind, ready_event=ready) + if given: + assert lease.poll() is None, f"a {kind} was ready before its caller's event completed" + ready.done = True + assert lease.poll() is not None, f"a {kind} was not ready" + assert recorded == [], "the lender made an event of its own" + lease.release() + + +@pytest.mark.cpu_only +def test_a_named_grow_waits_for_the_callers_pending_event_before_recycling_pages(kit, monkeypatch): + """The grow's cache stream follows work still using pages it may recycle.""" + manager = Manager((None,)) + manager._stream = SimpleNamespace(device="cuda:0") + cache_stream = object() + manager._cache_stream = lambda request_id: cache_stream + lender = lender_over(manager, layout((None,))) + manager.kv_cache_map[RID] = cache(history=2 * TPB) + order = [] + + class Pending: + def query(self): + return False + + def wait(self, stream): + assert stream is cache_stream + order.append("wait") + + def grow(*args): + order.append("grow") + return False + + monkeypatch.setattr(torch.cuda, "ExternalStream", lambda handle, device: handle) + monkeypatch.setattr(_manager, "grow", grow) + monkeypatch.setattr(_manager, "keeps_stale_pages", lambda manager: False) + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: True) + inputs_refuse_nothing(monkeypatch, lender) + served = request(served=2 * TPB) + lease = lender.lend_write( + served, 2 * TPB, 3 * TPB, **{**kit.DEVICE_NAMED, "ready_event": Pending()} + ) + assert lease.failure is not None and order == ["wait", "grow"] + lease.release() + + order.clear() + lease = lender.lend_write( + served, 2 * TPB, 3 * TPB, **{**kit.DEVICE_NAMED, "ready_event": Event()} + ) + assert "stream-waitable" in lease.failure and order == [] + lease.release() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("failure", ["fence", "fill"]) +def test_a_failure_after_named_grow_keeps_unfetched_history_out_of_resume_bounds( + kit, monkeypatch, failure +): + """A failed fill or fence cannot make a grown, unfetched window look computed.""" + manager = Manager((2 * TPB,)) + manager._stream = SimpleNamespace(device="cuda:0") + order = [] + + class CacheStream: + def synchronize(self): + order.append("synchronize") + + cache_stream = CacheStream() + manager._cache_stream = lambda request_id: cache_stream + lender = lender_over(manager, layout((2 * TPB,))) + manager.kv_cache_map[RID] = cache(history=2 * TPB) + + class FailingFence: + def record(self, stream): + assert stream is cache_stream + order.append("record") + if failure == "fence": + raise RuntimeError("record failed") + + def grow(manager, request, kv, position, end): + kv.history_length = position + order.append("grow") + if failure == "fill": + raise RuntimeError("fill failed") + return True + + monkeypatch.setattr(torch.cuda, "ExternalStream", lambda handle, device: handle) + monkeypatch.setattr(torch.cuda, "Event", FailingFence) + monkeypatch.setattr(_manager, "grow", grow) + monkeypatch.setattr(_manager, "keeps_stale_pages", lambda manager: False) + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: True) + inputs_refuse_nothing(monkeypatch, lender) + with pytest.raises(RuntimeError, match="record failed|fill failed"): + lender.lend_write( + request(served=2 * TPB), + 2 * TPB, + 3 * TPB, + **{**kit.DEVICE_NAMED, "ready_event": Event(done=True)}, + ) + assert order == (["grow", "record", "synchronize"] if failure == "fence" else ["grow"]) + assert RID in lender._ledger._kept_prefixes + assert (RID in lender._ledger._latest_fetches) is (failure == "fill") + assert lender._open_lease_count() == 0 + + +@pytest.mark.cpu_only +def test_a_named_device_write_waits_for_the_cache_stream_after_grow(kit, monkeypatch): + """The caller's pre-grow event alone cannot release a view whose grow queued more work.""" + manager = Manager((None,)) + manager._stream = SimpleNamespace(device="cuda:0") + cache_stream = object() + manager._cache_stream = lambda request_id: cache_stream + lender = lender_over(manager, layout((None,))) + manager.kv_cache_map[RID] = cache(history=2 * TPB) + grown = [] + fences = [] + + class Fence: + def __init__(self): + self.done = False + fences.append(self) + + def record(self, stream): + assert stream is cache_stream + assert grown, "the fence was recorded before the grow" + + def query(self): + return self.done + + def grow(manager, request, kv, position, end): + grown.append((position, end)) + return True + + monkeypatch.setattr(torch.cuda, "Event", Fence) + monkeypatch.setattr(torch.cuda, "ExternalStream", lambda handle, device: handle) + monkeypatch.setattr(_manager, "grow", grow) + monkeypatch.setattr(_manager, "keeps_stale_pages", lambda manager: False) + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: True) + inputs_refuse_nothing(monkeypatch, lender) + served = request(served=2 * TPB) + lease = lender.lend_write( + served, 2 * TPB, 3 * TPB, **{**kit.DEVICE_NAMED, "ready_event": Event(done=True)} + ) + try: + assert lease.failure is None, lease.failure + assert grown == [(2 * TPB, 3 * TPB)] and len(fences) == 1 + assert lease.poll() is None, "the pre-grow event released a view before the grow completed" + fences[0].done = True + view = lease.poll() + assert view is not None, lease.failure + return_write(lender, served, lease, view.row_masks(True)) + finally: + lease.release() + + +@gpu +def test_a_read_after_its_pages_went_to_host_and_back_waits_for_the_resume(kit, host_tier_manager): + """A read given an event recorded after a resume waits for it and lends the pages held now.""" + with host_tier_manager() as mgr: + kit.warm_tier_moves(mgr) + lender = attach(mgr) + source = kit.published(mgr, SOURCE, PROMPT) + kv = kit.kv(mgr, source) + old = [p for p in kit.pages(kv, 0) if p >= 0] + others = kit.Requests(mgr) + try: + mgr.suspend_request(source) + assert others.allocate(kit.pool_pages(mgr), chunk=1) + assert set(old) <= others.pages(), "inconclusive: its pages stayed on the GPU" + others.free_all_but(set(old)) + assert mgr.resume_request(source) + with kit.held_stream(mgr._stream, strict=False): + ready = kit.stream_event(mgr) + lease = lender.lend_read(source, 0, END, **{**kit.DEVICE_ALL, "ready_event": ready}) + assert lease.failure is None, lease.failure + early = lease.poll() + assert early is None, "a read went out before the work queued after the resume" + view = kit.ready_view(lease, mgr) + for run in view.groups: + now = kit.pages(kv, run.layer_group_index) + assert run.slots.tolist() == [now[b] for b in run.page_ordinals.tolist()] + assert not set(run.slots.tolist()) & set(old), "lent the pages held before" + lease.release() + finally: + others.free() + + +@gpu +def test_a_named_device_fetch_on_a_host_tier_waits_for_grow_after_its_caller_event( + kit, host_tier_manager +): + """A DEVICE view waits for the cache stream even when its caller's event finished first.""" + with host_tier_manager() as mgr: + kit.warm_tier_moves(mgr) + lender = attach(mgr) + target = kit.admitted(mgr, TARGET, PROMPT) + ready = kit.stream_event(mgr) + ready.synchronize() + with kit.held_stream(mgr._stream, strict=False): + lease = lender.lend_write(target, 0, END, **{**kit.DEVICE_NAMED, "ready_event": ready}) + assert lease.failure is None, lease.failure + assert lease.poll() is None, "the view went out before the grow's stream work" + view = kit.ready_view(lease, mgr) + assert view.num_rows == BLOCKS + return_write(lender, target, lease, view.row_masks(True)) + + +@gpu +def test_a_recycled_page_is_lent_for_writing_only_after_the_work_queued_on_it(kit, real_manager): + """A write onto freed pages goes out after the fill queued on them, and its bytes survive it.""" + # No host tier: a freed page goes straight back to the pool. A process's first cache close + # waits for the whole device, the held stream included: close one before the hold. + with real_manager(max_tokens=kit.POOL_TOKENS, host_cache_size=0, max_batch_size=16) as mgr: + warm = kit.make_request(900, list(range(40_000, 40_000 + TPB))) + assert mgr.prepare_context(warm) and mgr.resize_context(warm, TPB) + mgr.free_resources(warm) + torch.cuda.synchronize() + lender = attach(mgr) + dev = kit.DevicePages(mgr) # built before the hold: building it waits for the device + old = kit.make_request(OTHER, OTHER_PROMPT) + assert mgr.prepare_context(old) and mgr.resize_context(old, 4 * TPB) + old_pages = [p for p in kit.pages(kit.kv(mgr, old), 0) if p >= 0] + others = kit.Requests(mgr, dev) + # Every other page, in requests of a few blocks, then of one: only the old request's go free. + others.allocate(kit.pool_pages(mgr)) + while others.allocate(1, chunk=1): + pass + with kit.held_stream(mgr._stream, strict=False): + for slot in old_pages: + dev.fill_async(0, slot, STALE) + mgr.free_resources(old) + target = kit.admitted(mgr, TARGET, PROMPT) # all it takes from here, it takes freed + ready = kit.stream_event(mgr) + lease = lender.lend_write(target, 0, END, **{**kit.DEVICE_NAMED, "ready_event": ready}) + assert lease.failure is None, lease.failure + lent = kit.pages(kit.kv(mgr, target), 0)[:BLOCKS] + assert set(lent) <= set(old_pages), "inconclusive: the write took other pages" + assert lease.poll() is None, "the view went out with work still queued on its pages" + view = kit.ready_view(lease, mgr) + write_through(lender, view, FRESH) + mgr._stream.synchronize() + page_bytes = dev.page_bytes(0) + for ordinal in range(BLOCKS): + assert kit.page(mgr, target, 0, ordinal) == bytes([FRESH]) * page_bytes, ( + "the work queued on a recycled page overwrote what was written through the view" + ) + return_write(lender, target, lease, view.row_masks()) + others.free() + + +@gpu +def test_a_chunked_context_lends_each_range_once_the_step_that_computed_it_ran(kit, real_manager): + """Each step's range, lent with the event after the step, goes out once the step ran.""" + with real_manager(enable_block_reuse=False) as mgr: + lender = attach(mgr) + dev = kit.DevicePages(mgr) # built before any hold: building it waits for the device + request = kit.make_request(SOURCE, PROMPT) + assert mgr.prepare_context(request) + assert mgr.resize_context(request, request.context_remaining_length) + kv = kit.kv(mgr, request) + leases = [] + for block in range(BLOCKS): + with kit.held_stream(mgr._stream, strict=False): + for lg in range(kit.num_layer_groups(mgr)): + dev.fill_async(lg, kit.pages(kv, lg)[block], 0x10 + block) + step(mgr, request, TPB) + ready = kit.stream_event(mgr) + lease = lender.lend_read( + request, + block * TPB, + (block + 1) * TPB, + **{**kit.DEVICE_ALL, "ready_event": ready}, + ) + assert lease.failure is None, lease.failure + assert lease.poll() is None, f"block {block}'s view went out with its step queued" + leases.append(lease) + # Before the last step a read may neither end past what the steps computed nor inside a + # block a later step writes: either fails at the call. + for end, why in ((len(PROMPT), "tokens the cache holds"), (2 * TPB + 5, "whole block")): + refused(lender.lend_read(request, 0, end, **kit.DEVICE_ALL), why) + step(mgr, request, len(PROMPT) - BLOCKS * TPB) + ready = kit.stream_event(mgr) + last = lender.lend_read( + request, 2 * TPB, len(PROMPT), **{**kit.DEVICE_ALL, "ready_event": ready} + ) + assert sorted(slots_by_block(kit.ready_view(last, mgr))) == [ + (lg, ordinal) + for lg in range(kit.num_layer_groups(mgr)) + for ordinal in (BLOCKS - 1, BLOCKS) + ] + last.release() + for block, lease in enumerate(leases): + lent = slots_by_block(kit.ready_view(lease, mgr)) + assert {ordinal for _, ordinal in lent} == {block} + assert holds(dev, lent, lambda o: 0x10 + o), f"block {block} lost its step's bytes" + lease.release() + + +@gpu +def test_a_read_after_a_commit_rebases_its_pages_waits_for_the_new_pages(kit, real_manager): + """A read after a commit rebased its page waits for the writes to the new page.""" + with real_manager(block_reuse_config={"policy": "per_request"}, host_cache_size=0) as mgr: + lender = attach(mgr) + dev = kit.DevicePages(mgr) + groups = range(kit.num_layer_groups(mgr)) + + def allocate_pair(first_id, tokens): + requests = [kit.make_request(first_id + i, tokens) for i in range(2)] + for request in requests: + assert mgr.prepare_context(request) + kv = kit.kv(mgr, request) + kv.enable_swa_scratch_reuse = False + assert mgr.resize_context(request, len(tokens)) + assert kv.num_committed_tokens == 0 + return requests + + # Warm the native rebase and first close before holding the stream. + warm = allocate_pair(900, list(range(20_000, 20_000 + 2 * TPB))) + for request in warm: + for lg in groups: + for slot in kit.pages(kit.kv(mgr, request), lg): + if slot >= 0: + dev.fill_async(lg, slot, 0x11) + for request in reversed(warm): + step(mgr, request, 2 * TPB) + for request in warm: + mgr.free_resources(request) + mgr._stream.synchronize() + + # Both requests are admitted before either publishes their shared prefix. + first, canonical = allocate_pair(SOURCE, PROMPT[: 2 * TPB]) + first_kv, canonical_kv = (kit.kv(mgr, request) for request in (first, canonical)) + old_pages = {lg: kit.pages(first_kv, lg)[:2] for lg in groups} + canonical_pages = {lg: kit.pages(canonical_kv, lg)[:2] for lg in groups} + assert all( + len(old_pages[lg]) == len(canonical_pages[lg]) == 2 + and min(old_pages[lg] + canonical_pages[lg]) >= 0 + and set(old_pages[lg]).isdisjoint(canonical_pages[lg]) + for lg in groups + ), f"the requests have no distinct private pages: A={old_pages}, B={canonical_pages}" + for lg in groups: + dev.fill_async(lg, old_pages[lg][0], 0x40) + for slot in canonical_pages[lg]: + dev.fill_async(lg, slot, 0xEE) + step(mgr, first, TPB) + blocked_write = torch.cuda.Event() + blocked_write.record(mgr._stream) + mgr._stream.synchronize() + assert first.context_current_position == TPB, first.context_current_position + assert first_kv.num_committed_tokens == 0 + assert all( + dev.read(lg, canonical_pages[lg][0]) == bytes([0xEE]) * dev.page_bytes(lg) + for lg in groups + ), "B's pending canonical pages were not initialized to stale bytes" + + lease = None + try: + with kit.held_stream(mgr._stream, strict=False, watchdog_seconds=20.0) as gate: + for lg in groups: + for block, slot in enumerate(canonical_pages[lg]): + dev.fill_async(lg, slot, 0x40 + block) + blocked_write.record(mgr._stream) + step(mgr, canonical, 2 * TPB) + for lg in groups: + dev.fill_async(lg, old_pages[lg][1], 0x41) + step(mgr, first, TPB) + rebased = {lg: kit.pages(first_kv, lg)[0] for lg in groups} + assert all( + rebased[lg] == canonical_pages[lg][0] != old_pages[lg][0] for lg in groups + ), f"no native rebase: old A={old_pages}, B={canonical_pages}, new A={rebased}" + assert first_kv.num_committed_tokens == canonical_kv.num_committed_tokens == 2 * TPB + assert not gate.opened and not blocked_write.query(), "the writes ran early" + ready = kit.stream_event(mgr) + lease = lender.lend_read(first, 0, TPB, **{**kit.DEVICE_ALL, "ready_event": ready}) + assert lease.failure is None, lease.failure + early = lease.poll() is not None + assert not gate.opened and not blocked_write.query(), "a lend waited for them" + assert not early, f"A's read was ready before B wrote its rebased pages: {rebased}" + lent = slots_by_block(kit.ready_view(lease, mgr)) + assert lent == {(lg, 0): canonical_pages[lg][0] for lg in groups}, lent + assert holds(dev, lent, lambda ordinal: 0x40), "the canonical page has stale bytes" + finally: + if lease is not None: + lease.release() + + +@pytest.mark.cpu_only +def test_a_stopped_lender_fails_an_in_place_lease_no_poll_returned(): + manager = Manager() + lender = lender_over(manager) + kv = cache() + seen = in_place(lender, manager, kv, "read", entries((0,), (1,)), end=TPB) + view = seen.poll() + assert view is not None, seen.failure + pending = in_place(lender, manager, kv, "read", entries((2,), (3,)), ready_event=Event()) + lender._stop_reason = "stopped" + assert pending.poll() is None and pending.failure == "stopped" + assert not lender._is_locked(kv, 2, 3), "a failed lease kept its lock" + assert seen.poll() is view, "a view a poll returned before the stop was taken back" + seen.release() + pending.release() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("name", ["lend_read", "lend_write"]) +def test_a_device_lease_failed_at_the_call_is_an_in_place_lease_holding_nothing(kit, name): + """A request this rank holds no cache for: its lease fails at the call, holding nothing.""" + manager = Manager() + lender = lender_over(manager) + lend = getattr(lender, name) + lease = lend(request(), 0, TPB, **kit.DEVICE_ALL) + assert isinstance(lease, _lender._InPlaceLease) and lease._lock is None + assert lease.failure == _lender._no_cache_reason(RID) and lease.poll() is None + assert lender._open_lease_count() == 0, "a lease in place kept the staging memory" + lease.release() + lease.release() + staged = lend(request(), 0, TPB, **kit.HOST_NAMED) + assert isinstance(staged, _lender._StagingLease) + assert staged.failure == _lender._no_cache_reason(RID) + staged.release() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("cause", ["shut_down", "gone", "stopped"]) +@pytest.mark.parametrize("name", ["lend_read", "lend_write"]) +def test_a_lender_that_lends_no_more_raises_past_the_argument_checks(kit, name, cause): + """A lender that lends no more raises RuntimeError at a lend, after the argument checks.""" + manager = Manager() + lender = lender_over(manager) + if cause == "shut_down": + lender._shut_down = True + elif cause == "gone": + del manager + else: + lender._stop_reason = _lender._REASON_STOPPED + lend = getattr(lender, name) + for keywords in (kit.DEVICE_ALL, kit.HOST_NAMED): + with pytest.raises(RuntimeError, match="lends no more"): + lend(request(), 0, TPB, **keywords) + with pytest.raises(ValueError, match="bad token range"): + lend(request(), TPB, 0, **kit.DEVICE_ALL) + assert lender._open_lease_count() == 0, "a lend that raised left a lease open" + + +# -- reads -------------------------------------------------------------------------------------- + + +@gpu +@pytest.mark.parametrize("delivery", ["local", "complete", "missing"]) +def test_an_undelivered_committed_tail_cannot_be_lent_for_reading(kit, real_manager, delivery): + with real_manager() as mgr: + prompt = PROMPT[: TPB + 4] + request = kit.published(mgr, SOURCE, prompt) + lender = attach(mgr) + kv = kit.kv(mgr, request) + assert kv.num_committed_tokens == kv.history_length == len(prompt) + if delivery != "local": + write = lender.lend_write(request, TPB, len(prompt), **kit.DEVICE_ALL) + view = kit.ready_view(write, mgr) + assert all(group.page_ordinals.tolist() == [1] for group in view.groups) + return_write(lender, request, write, view.row_masks(delivery == "complete")) + read = lender.lend_read(request, TPB, len(prompt), **kit.DEVICE_ALL) + try: + if delivery == "missing": + assert read.failure is not None and read.poll() is None + assert not lender._is_locked(kv) + else: + assert kit.ready_view(read, mgr) is not None, read.failure + finally: + read.release() + + +@gpu +def test_a_read_of_a_two_beam_request_lends_the_pages_native_sends(kit, real_manager): + """A read of a two-beam prompt lends the first beam's pages, not the second's tail copy.""" + with real_manager(max_beam_width=2, enable_block_reuse=False) as mgr: + lender = attach(mgr) + request = computed_to(mgr, two_beams(SOURCE), len(PROMPT)) + kv = kit.kv(mgr, request) + assert mgr.try_allocate_generation(request), "no pages for the second beam" + assert kv.beam_width == 2, f"the cache holds {kv.beam_width} beams" + lease = lender.lend_read(request, 0, len(PROMPT), **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + view = kit.ready_view(lease, mgr) + for run in view.groups: + lg = run.layer_group_index + first = kit.pages(kv, lg) + second = [ + int(p) for p in kv.get_aggregated_page_indices(lg, beam_id=1, valid_only=False) + ] + assert second[BLOCKS] >= 0 and second[BLOCKS] != first[BLOCKS], (lg, first, second) + ordinals = run.page_ordinals.tolist() + assert ordinals == list(range(BLOCKS + 1)), ordinals + assert run.slots.tolist() == first[: BLOCKS + 1], (run.slots.tolist(), first) + lease.release() + + +# -- locks and refusals ------------------------------------------------------------------------- + + +@gpu +def test_a_write_whose_blocks_are_locked_fails_at_the_call_and_leaves_no_fetch(kit, real_manager): + """A named write over a locked block fails at the call, before its grow and its fetch record.""" + with real_manager() as mgr: + lender = attach(mgr) + target = kit.admitted(mgr, TARGET, PROMPT) + kv = kit.kv(mgr, target) + held = lender._locks.lock(kv, {0: [1]}, LockMode.READ) + before = (kv.capacity, lender.poll_resume_bounds(target)) + lease = lender.lend_write(target, 0, END, **kit.DEVICE_NAMED) + assert lease.failure is not None and "cannot be locked" in lease.failure + assert lease.poll() is None + assert (kv.capacity, lender.poll_resume_bounds(target)) == before + lease.release() + lender._locks.unlock(held) + + +@gpu +def test_a_write_of_every_row_over_a_block_a_read_holds_fails_at_the_call(kit, real_manager): + """A write of every row over a read-locked block fails at the call, recording no fetch.""" + with real_manager() as mgr: + lender = attach(mgr) + target = kit.admitted(mgr, TARGET, PROMPT) + assert mgr.resize_context(target, target.context_remaining_length) + kv = kit.kv(mgr, target) + before = lender.poll_resume_bounds(target) + held = lender._locks.lock(kv, {0: [1]}, LockMode.READ) + lease = lender.lend_write(target, 0, END, **kit.DEVICE_ALL) + assert lease.failure is not None and "locked for read" in lease.failure, lease.failure + assert lease.poll() is None + lease.release() + lender._locks.unlock(held) + assert not lender._is_locked(kv) and lender.poll_resume_bounds(target) == before + lease = lender.lend_write(target, 0, END, **kit.DEVICE_ALL) + view = kit.ready_view(lease, mgr) + assert {lg: o.tolist() for lg, o in lease._lock.blocks.items()} == { + run.layer_group_index: run.page_ordinals.tolist() for run in view.groups + } + return_write(lender, target, lease, view.row_masks(False)) + + +@gpu +def test_reads_in_place_share_their_blocks_and_none_with_a_write(kit, real_manager): + """Two reads of the same rows hold them together; a read and a write refuse each other.""" + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + kv = kit.kv(mgr, source) + writer = lender._locks.lock(kv, {0: [BLOCKS]}, LockMode.WRITE) + refused(lender.lend_read(source, 0, len(PROMPT), **kit.DEVICE_ALL), "locked for write") + lender._locks.unlock(writer) + first = lender.lend_read(source, 0, len(PROMPT), **kit.DEVICE_ALL) + second = lender.lend_read(source, TPB, len(PROMPT), **kit.DEVICE_ALL) + for lease in (first, second): + assert lease.failure is None, lease.failure + kit.ready_view(lease, mgr) + refused(lender.lend_write(source, 0, len(PROMPT), **kit.DEVICE_ALL), "locked for read") + first.release() + assert lender._is_locked(kv, BLOCKS, BLOCKS + 1), "the second read let go with the first" + second.release() + assert not lender._is_locked(kv) + + +@gpu +def test_a_write_in_place_over_a_block_without_a_page_fails_at_the_call( + kit, real_manager, monkeypatch +): + """A write of every row over a block without a page fails at the call, recording no fetch.""" + with real_manager() as mgr: + lender = attach(mgr) + target = kit.admitted(mgr, TARGET, PROMPT) + assert mgr.resize_context(target, target.context_remaining_length) + kv = kit.kv(mgr, target) + before = lender.poll_resume_bounds(target) + with monkeypatch.context() as patched: + readers = ("pages", "page_list", "locked_pages") + pages_giving(patched, lender, mgr, [target], 1, -1, readers) + lease = lender.lend_write(target, 0, END, **kit.DEVICE_ALL) + assert lease.failure is not None and "have no page" in lease.failure, lease.failure + assert not lender._is_locked(kv) + assert lender.poll_resume_bounds(target) == before, "the failed write left a fetch record" + lease.release() + + +@gpu +def test_a_page_at_its_pool_s_slot_count_is_never_lent_in_place(kit, real_manager, monkeypatch): + """A page at its pool's slot count fails a read at the call and a write at its first poll.""" + with real_manager() as mgr: + lender = attach(mgr) + source = kit.published(mgr, SOURCE, PROMPT) + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) # matches none of the source's blocks + assert mgr.resize_context(target, target.context_remaining_length) + caches = (kit.kv(mgr, source), kit.kv(mgr, target)) + limit = min(region.slot_count for region in device_regions(lender)) + before = lender.poll_resume_bounds(target) + with monkeypatch.context() as patched: + pages_giving(patched, lender, mgr, [source, target], 1, limit) + read = lender.lend_read(source, 0, len(PROMPT), **kit.DEVICE_ALL) + write = lender.lend_write(target, 0, END, **kit.DEVICE_ALL) + reason = f"blocks [1] have pages past the {limit} slots" + assert read.failure is not None and reason in read.failure, read.failure + assert write.failure is None and write.poll() is None + assert write.failure is not None and reason in write.failure, write.failure + assert not any(lender._is_locked(lent) for lent in caches) + assert lender.poll_resume_bounds(target) == before, "the failed write left a fetch record" + for lease in (read, write): + lease.release() + # Without that page the same lends lend the request's own pages. + for request, lease in ( + (source, lender.lend_read(source, 0, len(PROMPT), **kit.DEVICE_ALL)), + (target, lender.lend_write(target, 0, END, **kit.DEVICE_ALL)), + ): + (rows,) = kit.ready_view(lease, mgr).groups + pages = _manager.pages(kit.kv(mgr, request), rows.layer_group_index) + assert rows.slots.tolist() == pages[rows.page_ordinals].tolist() + if request is target: + return_write(lender, target, lease, (np.ones(len(rows), dtype=bool),)) + else: + lease.release() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("early_release", [False, True]) +def test_post_grow_refusals_wait_for_poll_and_read_lock_refusals_fail_at_call(early_release): + """A fetch's lock refusal fails at its first poll, a read's at the call.""" + manager = Manager() + lender = lender_over(manager) + kv = manager.kv_cache_map[RID] = cache() + reader = in_place(lender, manager, kv, "read") + fetch = lender._ledger._record_fetch(RID, kv, 0, (0, 0, 0), 0, (0, 0, 0)) + writer = in_place(lender, manager, kv, "write", fetch=fetch) + assert writer.failure is None and writer._lock is None + assert lender._ledger._latest_fetches.get(RID) is fetch + if early_release: + writer.release() + assert writer.failure is None + else: + assert writer.poll() is None + assert writer.failure is not None and "cannot be locked" in writer.failure + assert not lender._ledger._latest_fetches, "the refused fetch stayed unsettled" + reader.release() + assert not lender._is_locked(kv) + missing = lender._ledger._record_fetch(RID, kv, 0, (0, 0, 0), 0, (0, 0, 0)) + doomed = in_place(lender, manager, kv, "write", fetch=missing, doomed="blocks [1] have no page") + assert doomed.failure is None and doomed._lock is None + assert lender._ledger._latest_fetches.get(RID) is missing + if early_release: + doomed.release() + assert doomed.failure is None + else: + assert doomed.poll() is None and doomed.failure == "blocks [1] have no page" + assert not lender._ledger._latest_fetches + for lease in (writer, doomed, doomed, reader): + lease.release() + blocker = in_place(lender, manager, kv, "write") + refused_read = in_place(lender, manager, kv, "read") + assert refused_read.failure is not None and "cannot be locked" in refused_read.failure + assert refused_read.poll() is None + refused_read.release() + blocker.release() + assert not lender._locks.is_locked(kv) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("why", ["locked", "pageless"]) +def test_a_write_in_place_refused_at_the_call_leaves_the_records_as_they_were( + kit, monkeypatch, why +): + """A write over a read lock, or every row over a block without a page, fails at the call.""" + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: True) + monkeypatch.setattr(_manager, "keeps_stale_pages", lambda manager: False) + monkeypatch.setattr( + _lender._ManagerLender, "_bidirectional_split_refusal", lambda self, request, end: None + ) + manager = Manager() + lender = lender_over(manager) + pages = (0, 1, 2, -1) if why == "pageless" else None + kv = manager.kv_cache_map[RID] = cache(history=2 * TPB, pages=pages) + held = lender._locks.lock(kv, {0: [3]}, LockMode.READ) if why == "locked" else None + before = lender.poll_resume_bounds(request()) + lease = lender.lend_write(request(), 2 * TPB, PROMPT_LEN, **kit.DEVICE_ALL) + expected = "cannot be locked" if why == "locked" else "have no page" + assert lease.failure is not None and expected in lease.failure, lease.failure + assert lender.poll_resume_bounds(request()) == before, "the write left a record" + lease.release() + if held is not None: + lender._locks.unlock(held) + + +@pytest.mark.cpu_only +def test_a_lease_in_place_lends_no_slot_past_its_regions_and_fails_instead(): + """A page at its regions' least slot count fails a read at the call and a write at its poll.""" + lay = layout() + lay.device_pools = { + 0: (DevicePool(0, 0, 1 << 30, 256, 12), DevicePool(0, 1, 2 << 30, 128, 10)), + 1: (DevicePool(1, 0, 3 << 30, 512, 6),), + } + manager = Manager() + lender = lender_over(manager, lay) + limits = (10, 6, 10) # pool group 0's smaller pool bounds layer groups 0 and 2 + kv = manager.kv_cache_map[RID] = cache() + highest = [limit - 1 for limit in limits] + + def rows(pages): + return _lender._LentBlockEntries( + [0, 1, 2], [np.array([0, 1], dtype=np.int64)] * 3, [np.array([0, p]) for p in pages] + ) + + lease = in_place(lender, manager, kv, "read", rows(highest)) + view = lease.poll() + assert view is not None, lease.failure + assert [group.slots.tolist() for group in view.groups] == [[0, p] for p in highest] + lease.release() + for lg, limit in enumerate(limits): + pages = list(highest) + pages[lg] = limit + reason = f"layer group {lg}: blocks [1] have pages past the {limit} slots" + read = in_place(lender, manager, kv, "read", rows(pages)) + assert read.failure is not None and read.failure.startswith(reason), read.failure + assert read._lock is None and read.poll() is None and not lender._is_locked(kv) + read.release() + fetch = lender._ledger._record_fetch(RID, kv, 0, (0, 0, 0), 0, (0, 0, 0)) + write = in_place(lender, manager, kv, "write", rows(pages), fetch) + assert write.failure is None and write._lock is None and not lender._is_locked(kv) + assert write.poll() is None and write.failure.startswith(reason), write.failure + assert RID not in lender._ledger._latest_fetches, "the refused write's fetch stayed open" + write.release() + + +@pytest.mark.cpu_only +def test_the_lock_mode_follows_the_lease_kind(): + manager = Manager() + lender = lender_over(manager) + kv = cache() + read = in_place(lender, manager, kv, "read") + other = in_place(lender, manager, kv, "read") + assert read._lock.mode is other._lock.mode is LockMode.READ and other.failure is None + for lease in (read, other): + lease.release() + write = in_place(lender, manager, kv, "write") + assert write._lock.mode is LockMode.WRITE + write.release() + + +# -- writes and returns ------------------------------------------------------------------------- + + +@gpu +@pytest.mark.parametrize("read_only", [False, True], ids=["writable_mask", "read_only_mask"]) +@pytest.mark.parametrize( + "rows,usable", [((0, 1), 64), ((0, 2), 32), ((), 0)], ids=["prefix", "gap", "nothing"] +) +def test_an_in_place_write_lands_in_the_pages_and_counts_only_marked_rows( + kit, real_manager, rows, usable, read_only +): + with real_manager() as mgr: + lender = attach(mgr) + target = kit.admitted(mgr, TARGET, PROMPT) + lease = lender.lend_write(target, 0, END, **kit.DEVICE_NAMED) + view = kit.ready_view(lease, mgr) + assert view.num_rows == BLOCKS and all(r.object_keys is not None for r in view.groups) + write_through(lender, view, FRESH) + mask = np.isin(np.arange(BLOCKS), rows) + mask.flags.writeable = not read_only + assert return_write(lender, target, lease, (mask,)) == (usable, 0) + mgr._stream.synchronize() + page_bytes = kit.DevicePages(mgr).page_bytes(0) + for ordinal in range(BLOCKS): + assert kit.page(mgr, target, 0, ordinal) == bytes([FRESH]) * page_bytes, ordinal + + +@gpu +def test_a_receive_of_every_row_writes_the_admitted_pages_up_to_the_prompt_end(kit, real_manager): + with real_manager() as mgr: + lender = attach(mgr) + request, admitted = receiving(kit, mgr, TARGET, PROMPT) + assert admitted is not None, "the admission recorded no event" + keywords = {**kit.DEVICE_ALL, "ready_event": admitted} + with pytest.raises(ValueError, match="past the request's"): + lender.lend_write(request, 0, len(PROMPT) + 1, **keywords) + refused( + lender.lend_write(request, 0, len(PROMPT) - 2, **keywords), f"ends at {len(PROMPT) - 2}" + ) + lease = lender.lend_write(request, 0, len(PROMPT), **keywords) + assert lease.failure is None, lease.failure + view = kit.ready_view(lease, mgr) + assert [run.page_ordinals.tolist() for run in view.groups] == [list(range(BLOCKS + 1))] + assert all(run.object_keys is None for run in view.groups) + write_through(lender, view, FRESH) + assert return_write(lender, request, lease, view.row_masks(True)) == (len(PROMPT), 0) + mgr._stream.synchronize() + page_bytes = kit.DevicePages(mgr).page_bytes(0) + for ordinal in range(BLOCKS + 1): + assert kit.page(mgr, request, 0, ordinal) == bytes([FRESH]) * page_bytes + + +@gpu +def test_a_receive_may_end_at_a_remote_tail_s_start_and_resumes_no_further(kit, real_manager): + """A receive may end at a CSA2 remote tail's start inside a block, and resumes there only.""" + tail = 2 * TPB + 6 + with real_manager(enable_block_reuse=False) as mgr: + lender = attach(mgr) + request, admitted = receiving(kit, mgr, TARGET, PROMPT, py_csa2_remote_tail_start=tail) + kv = kit.kv(mgr, request) + assert kv.history_length == tail, f"the cache declares {kv.history_length} tokens" + keywords = {**kit.DEVICE_ALL, "ready_event": admitted} + refused(lender.lend_write(request, 0, tail - 1, **keywords), f"ends at {tail - 1}") + lease = lender.lend_write(request, 0, tail, **keywords) + assert lease.failure is None, lease.failure + view = kit.ready_view(lease, mgr) + assert [run.page_ordinals.tolist() for run in view.groups] == [[0, 1, 2]] + bounds = return_write(lender, request, lease, view.row_masks(True)) + assert tuple(bounds) == (tail, tail) + + +@gpu +def test_a_fetch_in_place_settles_in_the_lease_call_that_ends_it(kit, real_manager, monkeypatch): + """The lease call that ends a fetch in place settles it, as for a twin queried every step.""" + with real_manager() as mgr: + lender = attach(mgr) + quiet = kit.admitted(mgr, TARGET, PROMPT) + queried = kit.admitted(mgr, OTHER, OTHER_PROMPT) + leases = {r: lender.lend_write(r, 0, END, **kit.DEVICE_NAMED) for r in (quiet, queried)} + views = {} + for _ in range(3): # iterations: poll every open lease, query only ``queried``'s bounds + mgr._stream.synchronize() + for request, lease in leases.items(): + view = lease.poll() + if view is not None: + views[request] = view + assert lender.poll_resume_bounds(queried) is None, "settled before its return" + assert len(views) == 2, [lease.failure for lease in leases.values()] + returned = lender.return_writes( + [(lease, views[request].row_masks(True), True) for request, lease in leases.items()] + ) + assert returned == [(END, 0), (END, 0)], returned + assert lender.poll_resume_bounds(quiet) == lender.poll_resume_bounds(queried) == (END, 0) + # A release before the view, and a first poll that fails, leave the bounds as they were. + released = kit.admitted(mgr, 4, list(range(9000, 9097))) + failed = kit.admitted(mgr, 5, list(range(11000, 11097))) + assert mgr.resize_context(failed, failed.context_remaining_length) + before = {r: lender.poll_resume_bounds(r) for r in (released, failed)} + lease = lender.lend_write(released, 0, END, **kit.DEVICE_NAMED) + assert lease.failure is None, lease.failure + lease.release() + assert lender.poll_resume_bounds(released) == before[released] + limit = min(region.slot_count for region in device_regions(lender)) + with monkeypatch.context() as patched: + pages_giving(patched, lender, mgr, [failed], 1, limit) + lease = lender.lend_write(failed, 0, END, **kit.DEVICE_ALL) + assert lease.failure is None and lease.poll() is None and lease.failure is not None + assert lender.poll_resume_bounds(failed) == before[failed] + lease.release() + + +@gpu +def test_a_write_released_after_its_view_settles_as_one_returned_with_no_row_arrived( + kit, real_manager +): + """A write released after its view settles as a twin returned with no row arrived.""" + with real_manager() as mgr: + lender = attach(mgr) + released = kit.admitted(mgr, TARGET, PROMPT) + returned = kit.admitted(mgr, OTHER, OTHER_PROMPT) + bounds = {} + for request in (released, returned): + lease = lender.lend_write(request, 0, END, **kit.DEVICE_NAMED) + view = kit.ready_view(lease, mgr) + assert lender.poll_resume_bounds(request) is None, "settled before the lease ended" + if request is released: + lease.release() + bounds[request] = lender.poll_resume_bounds(request) + else: + bounds[request] = return_write(lender, request, lease, view.row_masks(False)) + assert bounds[released] is not None, "a release after the view left the fetch unsettled" + assert bounds[released] == bounds[returned], bounds + assert bounds[released].max_position == 0, bounds + + +@gpu +@pytest.mark.parametrize("reuse", [False, True], ids=["reuse_off", "reuse_on"]) +def test_a_receive_into_beam_zero_reaches_the_second_beam_at_the_first_generation( + kit, real_manager, reuse +): + """A two-beam receive into beam 0's pages reaches beam 1 at its first generation allocation.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestType + + with real_manager(max_beam_width=2, enable_block_reuse=reuse) as mgr: + lender = attach(mgr) + request = two_beams(TARGET, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY) + assert mgr.prepare_disagg_gen_init(request) + ready = mgr.take_disagg_receive_ready(request) + lease = lender.lend_write( + request, 0, len(PROMPT), **{**kit.DEVICE_ALL, "ready_event": ready} + ) + assert lease.failure is None, lease.failure + view = kit.ready_view(lease, mgr) + assert [run.page_ordinals.tolist() for run in view.groups] == [list(range(BLOCKS + 1))] + write_through(lender, view, FRESH) + bounds = return_write(lender, request, lease, view.row_masks(True)) + assert bounds.max_position == len(PROMPT) + kv = kit.kv(mgr, request) + mine = set(kit.pages(kv, 0)) + dev = kit.DevicePages(mgr) + # Every other page holds SENTINEL before the allocation, so a copy into one shows. + for slot in set(range(kit.pool_pages(mgr))) - mine: + dev.fill_async(0, slot, kit.SENTINEL) + mgr._stream.synchronize() + assert mgr.try_allocate_generation(request), "no pages for the second beam" + assert kv.beam_width == 2, f"the cache holds {kv.beam_width} beams" + first = kit.pages(kv, 0) + second = [int(p) for p in kv.get_aggregated_page_indices(0, beam_id=1, valid_only=False)] + assert len(second) == BLOCKS + 1 and second[BLOCKS] >= 0, (first, second) + assert second[BLOCKS] not in mine, f"beam 1's tail page {second[BLOCKS]} is beam 0's" + received = kit.digest(bytes([FRESH]) * dev.page_bytes(0)) + for ordinal in range(BLOCKS): + # -1: the second beam reads beam 0's page for the block. + page = first[ordinal] if second[ordinal] == -1 else second[ordinal] + assert kit.digest(dev.read(0, page)) == received, f"beam 1's block {ordinal}" + assert kit.digest(dev.read(0, second[BLOCKS])) == received, "beam 1's own tail page" + assert kit.digest(dev.read(0, first[BLOCKS])) == received, "beam 0's tail block" + + +# Rows that arrived in each layer group of ``write_with_its_view``'s lease, of blocks 1 to 3. +ARRIVED_ROWS = ((False, True, True), (True, True, True), (True, False, True)) + + +class BoolMask(np.ndarray): + """A subclass of ``numpy.ndarray``.""" + + +def frozen(values, dtype=bool): + """``values`` as a read-only array of ``dtype``.""" + array = np.array(values, dtype=dtype) + array.flags.writeable = False + return array + + +def write_with_its_view(): + """A write in place of blocks 1 to 3 of a cache that commits one block and five tokens of the + next, after a poll returned its view; with its manager, lender and fetch.""" + manager = Manager() + lender = lender_over(manager) + kv = manager.kv_cache_map[RID] = cache(blocks=4, committed=TPB + 5) + fetch = lender._ledger._record_fetch( + RID, kv, TPB, (TPB,) * 3, kv.history_length, (kv.num_committed_tokens,) * 3 + ) + lease = in_place(lender, manager, kv, "write", entries((1, 2, 3), (1, 2, 3)), fetch, 4 * TPB) + assert lease.poll() is not None + return manager, lender, fetch, lease + + +@pytest.mark.cpu_only +def test_marks_in_place_count_only_marked_rows_still_lent_and_taint_an_unmarked_tail(): + manager, lender, _, lease = write_with_its_view() + masks = lease.poll().row_masks() + masks[0][1:] = True # layer group 0: blocks 2 and 3, not the committed tail block 1 + masks[1][:] = True + masks[2][:] = True + bounds = return_write(lender, request(), lease, masks) + delivered = lender._ledger._deliveries_of(RID) + assert [m.nonzero()[0].tolist() for m in delivered.delivered_block_masks] == [ + [2, 3], + [1, 2, 3], + [1, 2, 3], + ] + assert bounds.max_position == 0, f"a request with a mixed tail block resumes: {bounds}" + assert not lender._locks.is_locked(manager.kv_cache_map[RID]) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("form", ["read_only", "writable", "read_only_subclass", "list"]) +def test_a_return_keeps_a_read_only_bool_mask_as_given_and_copies_any_other(monkeypatch, form): + """A read-only bool ``numpy.ndarray`` mask is kept as given, any other as an unshared copy.""" + manager, lender, _, lease = write_with_its_view() + given = { + "read_only": lambda: [frozen(rows) for rows in ARRIVED_ROWS], + "writable": lambda: [np.array(rows) for rows in ARRIVED_ROWS], + "read_only_subclass": lambda: [frozen(rows).view(BoolMask) for rows in ARRIVED_ROWS], + "list": lambda: [list(rows) for rows in ARRIVED_ROWS], + }[form]() + taken = [] + count = lender._return_write + + def recording_count(lease, masks, reusable): + taken.extend(masks) + return count(lease, masks, reusable) + + monkeypatch.setattr(lender, "_return_write", recording_count) + assert return_write(lender, request(), lease, given) is not None, "the fetch did not settle" + assert len(taken) == len(given) == 3 + for mask, kept in zip(given, taken): + if form == "read_only": + assert kept is mask, "a read-only bool mask was copied" + else: + assert type(kept) is np.ndarray and kept.dtype == np.bool_, type(kept) + assert not np.shares_memory(kept, np.asarray(mask)), f"a {form} mask was kept" + assert kept.tolist() == list(mask) + delivered = lender._ledger._deliveries_of(RID) + assert [m.nonzero()[0].tolist() for m in delivered.delivered_block_masks] == [ + [2, 3], + [1, 2, 3], + [1, 3], + ] + assert not lender._locks.is_locked(manager.kv_cache_map[RID]), "the return left the lease lent" + + +@pytest.mark.cpu_only +def test_a_return_refuses_a_read_only_mask_of_another_dtype_or_length_and_changes_nothing(): + """A read-only mask of integers or of another length is refused, the fetch and lease kept.""" + manager, lender, _, lease = write_with_its_view() + kv = manager.kv_cache_map[RID] + good = [frozen(rows) for rows in ARRIVED_ROWS] + for bad in ( + [frozen([1, 1, 1], np.int64)] + good[1:], + good[:1] + [frozen([True, True])] + good[2:], + ): + with pytest.raises(ValueError, match="the mask must be bool"): + return_write(lender, request(), lease, bad) + assert lender.poll_resume_bounds(request()) is None, "a refused return settled the fetch" + assert lender._locks.is_locked(kv), "a refused return let the lease go" + assert return_write(lender, request(), lease, good) is not None + assert not lender._locks.is_locked(kv) + + +@pytest.mark.cpu_only +def test_a_write_released_before_its_view_abandons_its_fetch_and_one_after_delivers_nothing(): + """A write released before its view abandons its fetch; one released after it counts no row.""" + manager = Manager() + lender = lender_over(manager) + # The pages ``entries()`` lends: an arrived row would count. + kv = manager.kv_cache_map[RID] = cache(pages=(5, 6, 7, 8)) + for seen in (False, True): + fetch = lender._ledger._record_fetch(RID, kv, 0, (0, 0, 0), 0, (0, 0, 0)) + lease = in_place(lender, manager, kv, "write", fetch=fetch) + if seen: + assert lease.poll() is not None + lease.release() + latest = lender._ledger._latest_fetches.get(RID) + assert (latest is fetch) == seen + if seen: + assert lender._ledger._fetch_settled(fetch), "a release after the view left it open" + delivered = lender._ledger._deliveries_of(RID).delivered_block_masks + assert not any(blocks.any() for blocks in delivered), "a release counted rows" + with pytest.raises(RuntimeError, match="returned already"): + return_write(lender, request(), lease, lease._view.row_masks(True)) + delivered = lender._ledger._deliveries_of(RID).delivered_block_masks + assert not any(blocks.any() for blocks in delivered), "a refused return counted rows" + assert not lender._locks.is_locked(kv) + + +# -- what a lent cache keeps -------------------------------------------------------------------- + + +@gpu +def test_a_freed_request_keeps_its_lent_pages_until_the_lease_is_released(kit, real_manager): + """A freed request's lent cache stays open, each of its pages counted, until the release.""" + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + lease = lender.lend_read(source, 0, END, **kit.DEVICE_NAMED) + view = kit.ready_view(lease, mgr) + kv = kit.kv(mgr, source) + held = {} + for lg, group in enumerate(kit.pool_group_of(mgr)): + held[group] = held.get(group, 0) + sum(page >= 0 for page in kit.pages(kv, lg)) + lent = sum(len(run) for run in view.groups) + assert sum(held.values()) > lent, f"inconclusive: {held} pages, {lent} of them lent" + assert dict(lender.held_page_counts().deduct) == {}, "a live request's pages counted" + before = [entry_bytes(lender, run, i) for run in view.groups for i in range(len(run))] + mgr.free_resources(source) + assert not kit.closed(kv), "the free closed a cache whose pages are lent" + assert mgr._holds_page_locks(None), "the kept cache's lock went unseen" + retained = dict(lender.held_page_counts().deduct) + assert retained == {g: n for g, n in held.items() if n}, (retained, held) + kit.overwrite_free_pages(mgr) + after = [entry_bytes(lender, run, i) for run in view.groups for i in range(len(run))] + assert after == before, "a freed request's lent page went to another request" + lease.release() + assert kit.closed(kv), "the release did not close the freed request's cache" + assert not mgr._holds_page_locks(None) + assert dict(lender.held_page_counts().deduct) == {}, "a closed cache's pages still count" + + +def live_slots(kv, lg, recurrent): + """The slots layer group ``lg`` of the live cache ``kv`` holds over its beams: its pages, its + recurrent state's slot, its scratch slots.""" + slots = set() + for beam in range(int(kv.beam_width)): + if recurrent: + slots.add(int(kv.get_ssm_block_base_index(lg, beam_id=beam))) + else: + slots.update( + int(p) for p in kv.get_aggregated_page_indices(lg, beam_id=beam, valid_only=True) + ) + scratch = kv.get_scratch_desc(lg) + slots.update(int(s) for s in (scratch.slot_ids if scratch is not None else ())) + slots.discard(-1) + return slots + + +@gpu +@pytest.mark.parametrize("kind", ["recurrent_state", "swa_scratch", "two_beams"]) +def test_a_freed_cache_counts_its_state_scratch_and_beams_as_retained( + kit, real_manager, hybrid_manager, kind +): + """A freed lent cache counts as retained its state slot, scratch slots and both beams' pages.""" + if kind == "recurrent_state": + manager = hybrid_manager() + elif kind == "swa_scratch": + manager = real_manager(windows=[WINDOW, 256], swa_scratch_reuse=True) + else: + manager = real_manager(max_beam_width=2, enable_block_reuse=False) + with manager as mgr: + lender = attach(mgr) + recurrent = derive_layout(mgr).recurrent + if kind == "two_beams": + request = two_beams(SOURCE) + else: + request = kit.make_request(SOURCE, LONG_PROMPT if kind == "swa_scratch" else PROMPT) + end = 2 * TPB if kind == "swa_scratch" else len(PROMPT) + computed_to(mgr, request, end) + kv = kit.kv(mgr, request) + if kind == "two_beams": + assert mgr.try_allocate_generation(request), "no pages for the second beam" + assert kv.beam_width == 2, f"the cache holds {kv.beam_width} beams" + lease = lender.lend_read(request, 0, end, **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + kit.ready_view(lease, mgr) + if kind == "swa_scratch": + # The next chunk, in flight at the free, takes the 64-token window past its first two + # blocks, which the runtime backs with scratch slots, not pages. + assert mgr.resize_context(request, 4 * TPB), "no capacity for the next chunk" + held = {} + for lg, group in enumerate(kit.pool_group_of(mgr)): + slots = live_slots(kv, lg, recurrent[lg]) + if slots: + held[group] = held.get(group, 0) + len(slots) + if kind == "recurrent_state": + assert any(recurrent), "inconclusive: the manager holds no recurrent state" + elif kind == "swa_scratch": + (windowed,) = [lg for lg, window in enumerate(kit.windows(mgr)) if window == WINDOW] + scratch = kv.get_scratch_desc(windowed) + assert scratch is not None and len(scratch.slot_ids), (kv.history_length, kv.capacity) + # The count includes slots that only the scratch descriptor names. + pages = set(kit.pages(kv, windowed)) + assert not pages & {int(s) for s in scratch.slot_ids}, (pages, scratch.slot_ids) + else: + (lg,) = range(kit.num_layer_groups(mgr)) + assert held == {kit.pool_group_of(mgr)[lg]: BLOCKS + 2}, f"inconclusive: {held}" + mgr.free_resources(request) + assert not kit.closed(kv), "the free closed a cache whose pages are lent" + assert dict(lender.held_page_counts().deduct) == held + lease.release() + assert kit.closed(kv), "the release did not close the freed request's cache" + assert dict(lender.held_page_counts().deduct) == {} + + +@gpu +def test_a_hold_with_device_regions_keeps_the_runtime_past_the_shutdown(kit, real_manager): + with real_manager() as mgr: + lender = attach(mgr) + hold = lender.hold_regions() + runtime = mgr.impl + mgr.shutdown() + assert lender._keeps_pools() + assert any(kept is runtime for kept in kit.retained()), "the runtime is not kept" + hold.release() + assert not lender._keeps_pools() + # The block's exit shuts the manager down again, and with nothing held the runtime's own runs. + + +@pytest.mark.cpu_only +def test_a_lender_reports_the_pages_its_leases_keep_past_their_requests_free(): + """The lender counts a lent cache's pages from its request's free to its last release.""" + manager = Manager() + lender = lender_over(manager) + pages = [5, 6, 7, -1] + kv = SimpleNamespace( + Status=SimpleNamespace(OPEN=0, CLOSED=1), + status=0, + num_committed_tokens=0, + tokens_per_block=TPB, + beam_width=1, + closed=False, + get_scratch_desc=lambda lg: None, + get_aggregated_page_indices=lambda lg, beam_id=0, valid_only=False: ( + [p for p in pages if p >= 0] if valid_only else pages + ), + ) + + def close(): + kv.closed = True + kv.status = kv.Status.CLOSED + + kv.close = close + lease = in_place(lender, manager, kv, "read") + assert dict(lender.held_page_counts().deduct) == {} + assert lender._defer_close(kv), "a locked cache was not kept past its free" + # Pool group 0 holds layer groups 0 and 2, three pages each; pool group 1 layer group 1. + assert dict(lender.held_page_counts().deduct) == {0: 6, 1: 3} + lease.release() + assert kv.closed and dict(lender.held_page_counts().deduct) == {} + + +# -- host fetches beside reads in place --------------------------------------------------------- + + +@pytest.fixture +def mixed_host_case(monkeypatch): + """Request ``RID`` after its first context step: a history of a block, its position past it.""" + manager = Manager() + manager.enable_block_reuse = False + manager._can_publish_block_reuse = False + manager.is_draft = False + manager._fresh_page_fill = None + manager._resize_for_served_prefix = lambda req, kv, position, end: True + manager._stale_block_range = lambda lg, history: (0, 0) + manager.kv_cache_map[RID] = cache(history=TPB) + lender = lender_over(manager) + req = request(computed=TPB) + req.py_return_context_logits = False + req.py_additional_outputs = None + req.try_get_encoder_output_len = lambda: None + monkeypatch.setattr( + _manager, "block_keys", lambda m, r, k, n: [bytes([i]) * 32 for i in range(n)] + ) + copied = Event() + monkeypatch.setattr( + lender, + "_queue_copy", + lambda segments, to_staging, after=None: (_lender._QueuedCopy(copied), None), + ) + return manager, lender, req, copied + + +@pytest.mark.cpu_only +def test_host_fetch_refuses_pages_a_device_reader_still_holds(kit, mixed_host_case): + _, lender, req, _ = mixed_host_case + read = lender.lend_read(req, 0, TPB, **kit.DEVICE_ALL) + assert read.poll() is not None + fetch = lender.lend_write(req, 0, 2 * TPB, **kit.HOST_NAMED) + assert fetch.failure is not None and "locked" in fetch.failure + assert lender._ledger._latest_fetches == {} + fetch.release() + read.release() + + +@pytest.mark.cpu_only +def test_host_fetch_excludes_later_device_reads_until_its_marked_copy_finishes( + kit, mixed_host_case +): + """A read of rows a host fetch lent fails at the call until progress saw its copy complete.""" + _, lender, req, copied = mixed_host_case + fetch = lender.lend_write(req, 0, 2 * TPB, **kit.HOST_NAMED) + view = fetch.poll() + assert view is not None, fetch.failure + for returned in (False, True): + if returned: + bounds = return_write(lender, req, fetch, view.row_masks(True)) + assert bounds is None, f"settled before its copy ran: {bounds}" + refused(lender.lend_read(req, 0, TPB, **kit.DEVICE_ALL), "unsettled") + copied.done = True + lender.progress() + ready = Event() + read = lender.lend_read(req, 0, TPB, **{**kit.DEVICE_ALL, "ready_event": ready}) + assert read.failure is None, read.failure + assert read.poll() is None + ready.done = True + assert read.poll() is not None + read.release() + + +@pytest.mark.cpu_only +def test_a_host_fetch_leaves_nonoverlapping_device_rows_readable(kit, mixed_host_case): + _, lender, req, copied = mixed_host_case + fetch = lender.lend_write(req, TPB, 2 * TPB, **kit.HOST_NAMED) + view = fetch.poll() + assert view is not None + read = lender.lend_read(req, 0, TPB, **kit.DEVICE_ALL) + assert read.poll() is not None, read.failure + return_write(lender, req, fetch, view.row_masks(True)) + copied.done = True + read.release() + + +@pytest.mark.cpu_only +def test_a_device_read_of_a_row_another_group_still_fetches_fails_at_the_call(kit, mixed_host_case): + """A read in place overlapping rows another layer group still fetches fails at the call.""" + manager, lender, req, copied = mixed_host_case + # Block 0 committed, a claim complete in layer group 0 alone. + kv = manager.kv_cache_map[RID] + kv.num_committed_tokens = TPB + kv.reuse_status = [SimpleNamespace(complete=c, endpoint=TPB) for c in (True, False, False)] + fetch = lender.lend_write(req, 0, 2 * TPB, **kit.HOST_NAMED) + view = fetch.poll() + assert view is not None, fetch.failure + pending = {g.layer_group_index: g.page_ordinals.tolist() for g in view.groups} + assert pending == {0: [1], 1: [0, 1], 2: [0, 1]}, pending + refused(lender.lend_read(req, 0, TPB, **kit.DEVICE_ALL), "unsettled") + return_write(lender, req, fetch, view.row_masks(True)) + copied.done = True + lender.progress() + read = lender.lend_read(req, 0, TPB, **kit.DEVICE_ALL) + assert read.failure is None, read.failure + read.release() + + +# -- a request that runs on --------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_a_read_past_the_history_needs_a_finished_step_that_computed_its_end(kit, recorded): + """A read past the history ends where the request's steps computed, after its caller's event.""" + manager = Manager() + lender = lender_over(manager) + manager.kv_cache_map[RID] = cache(history=TPB) + why = "past the 32 tokens the cache holds" + refused(lender.lend_read(request(computed=TPB), 0, 2 * TPB, **kit.DEVICE_ALL), why) + ran = Event() + keywords = {**kit.DEVICE_ALL, "ready_event": ran} + lease = lender.lend_read(request(computed=2 * TPB), 0, 2 * TPB, **keywords) + assert lease.failure is None, lease.failure + assert recorded == [] + assert lease.poll() is None, "the view went out before the step that computed it ran" + ran.done = True + assert lease.poll() is not None + lease.release() + refused(lender.lend_read(request(computed=2 * TPB), 0, 3 * TPB, **kit.DEVICE_ALL), why) + + +@pytest.mark.cpu_only +def test_until_the_prompt_is_computed_a_read_of_every_row_ends_at_a_whole_block(kit, recorded): + """A read of every row ends at a whole block until the prompt is computed.""" + manager = Manager() + lender = lender_over(manager) + manager.kv_cache_map[RID] = cache(history=3 * TPB) + refused( + lender.lend_read(request(), 0, 2 * TPB + 5, **kit.DEVICE_ALL), + "until this worker computes its context to 97 tokens", + ) + lender.lend_read(request(), 0, 2 * TPB, **kit.DEVICE_ALL).release() + # The prompt's last step ran, while a window held the history below the prompt's end. + computed = request(computed=PROMPT_LEN) + lender.lend_read(computed, 0, 2 * TPB + 5, **kit.DEVICE_ALL).release() + # Generating, past the prompt: the history holds it. + manager.kv_cache_map[RID] = cache(blocks=5, history=PROMPT_LEN + 3) + lease = lender.lend_read(computed, 0, PROMPT_LEN + 2, **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + lease.release() + refused( + lender.lend_read(computed, 0, PROMPT_LEN + 4, **kit.DEVICE_ALL), + "past the 100 tokens the cache holds", + ) + assert recorded == [] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("computed", ["history", "position"]) +@pytest.mark.parametrize("scratch", [False, True], ids=["scratch_off", "scratch_on"]) +def test_a_read_before_the_prompt_is_computed_leaves_out_windows_shorter_than_the_prompt( + kit, recorded, scratch, computed +): + """Until the prompt is computed a read lends and holds no window shorter than the prompt.""" + windows = (None, 64, PROMPT_LEN) + manager = Manager(windows) + lender = lender_over(manager, layout(windows)) + kv = manager.kv_cache_map[RID] = cache(history=2 * TPB) + kv.enable_swa_scratch_reuse = scratch + lease = lender.lend_read(request(computed=2 * TPB), 0, 2 * TPB, **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + lent = {run.layer_group_index: run.page_ordinals.tolist() for run in lease.poll().groups} + assert lent == {0: [0, 1], 1: [], 2: [0, 1]}, lent + assert {lg: blocks.tolist() for lg, blocks in lease._lock.blocks.items()} == lent + cap = lender._history_cap(kv) + assert cap is None or cap >= PROMPT_LEN, f"the read holds the context's history at {cap}" + lease.release() + # The context's last step turned scratch reuse off; its history reached the prompt's end, or + # stayed at the last whole block while the request's context position reached it. + kv.enable_swa_scratch_reuse = False + kv.history_length = PROMPT_LEN if computed == "history" else 3 * TPB + ran = request(computed=3 * TPB if computed == "history" else PROMPT_LEN) + after = lender.lend_read(ran, 0, PROMPT_LEN, **kit.DEVICE_ALL) + assert after.failure is None, after.failure + lent = {run.layer_group_index: run.page_ordinals.tolist() for run in after.poll().groups} + assert lent == {0: [0, 1, 2, 3], 1: [1, 2, 3], 2: [0, 1, 2, 3]}, lent + # The 64-token window leaves block 1 behind at a history of 127. + assert lender._history_cap(kv) == 126, lender._history_cap(kv) + after.release() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "position,remaining,helix,end", + [ + (TPB, PROMPT_LEN - TPB, False, PROMPT_LEN), + (2 * TPB + 6, 0, False, 2 * TPB + 6), + (2 * TPB, PROMPT_LEN, False, PROMPT_LEN), + (2 * TPB + 6, 0, True, LONG), + ], + ids=["to_the_prompt", "short_of_the_prompt", "past_the_prompt", "helix"], +) +def test_this_worker_s_context_ends_where_its_remaining_rows_end_within_the_prompt( + position, remaining, helix, end +): + """The context ends at position plus remaining rows, within the prompt; under Helix, at it.""" + manager = SimpleNamespace(_has_cp_helix=helix, _helix_cp_rank=1, _helix_cp_size=2) + req = SimpleNamespace( + prompt_len=PROMPT_LEN, + total_input_len_cp=LONG, + context_current_position=position, + context_remaining_length=remaining, + ) + assert _manager.context_end(req, manager) == end + assert _manager.context_computed(req) == position + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "tail,lent", + [ + (2 * TPB + 6, {0: [0, 1, 2], 1: [0, 1, 2], 2: [2]}), + (4 * TPB, {0: [0, 1, 2, 3], 1: [2, 3], 2: [3]}), + ], + ids=["inside_a_block", "at_a_block"], +) +def test_a_context_ending_short_of_the_prompt_is_computed_once_this_worker_computed_it( + kit, recorded, tail, lent +): + """A context ending short of the prompt, computed to its end, lends what each window keeps.""" + windows = (None, 64, 2) + manager = Manager(windows) + lender = lender_over(manager, layout(windows)) + manager.kv_cache_map[RID] = cache(blocks=-(-tail // TPB), history=tail) + ran = request(computed=tail, prompt=LONG, end=tail) + lease = lender.lend_read(ran, 0, tail, **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + got = {run.layer_group_index: run.page_ordinals.tolist() for run in lease.poll().groups} + assert got == lent, got + lease.release() + + +@pytest.mark.cpu_only +def test_while_a_context_ending_short_of_the_prompt_runs_both_rules_hold_against_its_end( + kit, recorded +): + """Mid-context, a read ends at a whole block and leaves out windows shorter than the end.""" + tail = 4 * TPB + windows = (None, 64, 2, tail) + manager = Manager(windows) + lender = lender_over(manager, layout(windows)) + manager.kv_cache_map[RID] = cache(blocks=3, history=3 * TPB) + running = request(computed=3 * TPB, prompt=LONG, end=tail) + refused( + lender.lend_read(running, 0, 2 * TPB + 5, **kit.DEVICE_ALL), + "until this worker computes its context to 128 tokens", + ) + # A request at the same step whose context runs to the prompt leaves out the tail's window too. + for req, lent in ( + (running, {0: [0, 1, 2], 1: [], 2: [], 3: [0, 1, 2]}), + (request(computed=3 * TPB, prompt=LONG), {0: [0, 1, 2], 1: [], 2: [], 3: []}), + ): + lease = lender.lend_read(req, 0, 3 * TPB, **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + got = {run.layer_group_index: run.page_ordinals.tolist() for run in lease.poll().groups} + assert got == lent, (req.context_remaining_length, got) + lease.release() + + +@gpu +@pytest.mark.parametrize("scratch", [False, True], ids=["scratch_off", "scratch_on"]) +@pytest.mark.parametrize("reuse", [False, True], ids=["resized", "committed"]) +def test_a_read_before_the_context_ends_leaves_out_a_window_shorter_than_the_prompt( + kit, real_manager, scratch, reuse +): + """A mid-context read leaves out the 64-token window, so each step moves its history on.""" + with real_manager( + windows=[WINDOW, 256], enable_block_reuse=reuse, swa_scratch_reuse=scratch + ) as mgr: + lender = attach(mgr) + groups = range(kit.num_layer_groups(mgr)) + (windowed,) = [lg for lg, window in enumerate(kit.windows(mgr)) if window == WINDOW] + request = kit.make_request(SOURCE, LONG_PROMPT) + assert mgr.prepare_context(request) + kv = kit.kv(mgr, request) + assert kv.enable_swa_scratch_reuse == scratch + + def chunk(tokens): + assert mgr.resize_context(request, tokens), f"no capacity for {tokens} tokens" + step(mgr, request, tokens) + + chunk(2 * TPB) + first = {lg: kit.pages(kv, lg)[:2] for lg in groups} + assert all(min(first[lg]) >= 0 for lg in groups), f"inconclusive: {first}" + lease = lender.lend_read(request, 0, 2 * TPB, **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + lent = slots_by_block(kit.ready_view(lease, mgr)) + longer = [(lg, block) for lg in groups if lg != windowed for block in range(2)] + assert sorted(lent) == longer, f"lent {sorted(lent)}, not the longer window's {longer}" + for _ in range(3): + chunk(TPB) + position = request.context_current_position + stale = kit.stale_blocks(mgr, windowed, position) + assert stale[1] > 0, f"inconclusive: at {position} the window passes no block, {stale}" + assert kv.history_length == position, "the read held the history back" + assert kit.pages(kv, windowed)[0] < 0, "the window kept block 0" + chunk(len(LONG_PROMPT) - position) + lease.release() + assert kv.history_length == len(LONG_PROMPT), kv.history_length + stale = kit.stale_blocks(mgr, windowed, kv.history_length) + pages = kit.pages(kv, windowed) + kept = {(windowed, block) for block in range(stale[1], len(pages)) if pages[block] >= 0} + assert kept, f"inconclusive: the window keeps no backed block, stale {stale}, {pages}" + after = lender.lend_read(request, 0, len(LONG_PROMPT), **kit.DEVICE_ALL) + assert after.failure is None, after.failure + lent = slots_by_block(kit.ready_view(after, mgr)) + assert {key for key in lent if key[0] == windowed} == kept, sorted(lent) + after.release() + + +@gpu +def test_generation_holds_a_window_where_it_keeps_a_lent_block_until_the_release(kit, real_manager): + """Generation holds a window's history where it keeps a lent block until the release.""" + from tensorrt_llm.bindings import LlmRequestState + + with real_manager(windows=[WINDOW, 256], enable_block_reuse=False) as mgr: + lender = attach(mgr) + dev = kit.DevicePages(mgr) + (windowed,) = [lg for lg, window in enumerate(kit.windows(mgr)) if window == WINDOW] + # The largest history at which the window keeps block 0. + keeping = max(h for h in range(4 * WINDOW) if kit.stale_blocks(mgr, windowed, h)[1] <= 0) + request = kit.make_request(SOURCE, LONG_PROMPT[:keeping]) + assert mgr.prepare_context(request) + kv = kit.kv(mgr, request) + kv.enable_swa_scratch_reuse = False + assert mgr.resize_context(request, request.context_remaining_length) + for lg in range(kit.num_layer_groups(mgr)): + for block in range(2): + dev.write(lg, kit.pages(kv, lg)[block], bytes([0x30 + block]) * dev.page_bytes(lg)) + step(mgr, request, keeping) + request.state = LlmRequestState.GENERATION_IN_PROGRESS + lease = lender.lend_read(request, 0, 2 * TPB, **kit.DEVICE_ALL) + lent = slots_by_block(kit.ready_view(lease, mgr)) + assert (windowed, 0) in lent, f"inconclusive: lent {sorted(lent)}" + for _ in range(3): + generate(mgr, request) + last = request.max_beam_num_tokens - 1 + stale = kit.stale_blocks(mgr, windowed, last) + assert stale[1] > 0, f"inconclusive: at {last} the window passes no block, {stale}" + assert kv.history_length == keeping, (last, keeping) + assert kit.pages(kv, windowed)[0] == lent[(windowed, 0)], "the window let a lent page go" + kit.overwrite_free_pages(mgr) + assert holds(dev, lent, lambda ordinal: 0x30 + ordinal), "another request took a lent page" + lease.release() + generate(mgr, request) + assert kv.history_length == request.max_beam_num_tokens - 1 + assert kit.pages(kv, windowed)[0] < 0, "the release left the window holding the block" + + +# -- beside a KV cache connector ---------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_a_read_waits_for_a_load_the_connector_has_pending(kit, recorded): + connector = Connector() + manager = Manager(connector=connector) + lender = lender_over(manager) + manager.kv_cache_map[RID] = cache(history=2 * TPB) + connector.loading.add(RID) + lease = lender.lend_read(request(), 0, 2 * TPB, **kit.DEVICE_ALL) + assert lease.failure is None and recorded == [] + for _ in range(2): + assert lease.poll() is None, "a view went out while the connector was loading" + connector.loading.clear() + view = lease.poll() + assert view is not None and lease.poll() is view + lease.release() + # A lease taken with no load pending waits for its event alone. + quiet = lender.lend_read(request(), 0, 2 * TPB, **kit.DEVICE_ALL) + connector.loading.add(RID) # a load the connector starts later goes past what the cache holds + assert quiet.poll() is not None + quiet.release() + + +@pytest.mark.cpu_only +def test_a_read_fails_at_the_call_while_the_connector_holds_a_prefix_reservation(kit, recorded): + """A read fails at the call while a prefix reservation counts unfilled tokens in the history.""" + connector = Connector() + manager = Manager(connector=connector) + lender = lender_over(manager) + kv = manager.kv_cache_map[RID] = cache(history=2 * TPB) + connector.reserved.add(RID) + lease = lender.lend_read(request(), 0, 2 * TPB, **kit.DEVICE_ALL) + assert lease.failure is not None and "prefix reservation" in lease.failure, lease.failure + assert not lender._is_locked(kv) and lease.poll() is None + lease.release() + connector.reserved.clear() + lease = lender.lend_read(request(), 0, 2 * TPB, **kit.DEVICE_ALL) + assert lease.failure is None and lender._is_locked(kv) + lease.release() + + +class PastTheRule(Exception): + """Raised by the first check of ``lend_write`` after the connector's rule.""" + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("kind", [None, "capacity_only", "loads"]) +def test_a_fetch_starts_at_or_past_the_prefix_served_beside_a_loading_connector( + kit, recorded, monkeypatch, kind +): + """Beside a loading connector a fetch starts at or past the prefix served, else fails.""" + + def past(self, request, end): + raise PastTheRule() + + monkeypatch.setattr(_lender._ManagerLender, "_bidirectional_split_refusal", past) + connector = None if kind is None else Connector(capacity_only=kind == "capacity_only") + manager = Manager(connector=connector) + lender = lender_over(manager) + manager.kv_cache_map[RID] = cache(history=2 * TPB) + served = request(served=2 * TPB) + if kind == "loads": + refused( + lender.lend_write(served, TPB, PROMPT_LEN, **kit.DEVICE_ALL), + "inside the 64 tokens served to the request", + ) + else: + with pytest.raises(PastTheRule): + lender.lend_write(served, TPB, PROMPT_LEN, **kit.DEVICE_ALL) + with pytest.raises(PastTheRule): + lender.lend_write(served, 2 * TPB, PROMPT_LEN, **kit.DEVICE_ALL) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("window", [None, 64], ids=["full", "window"]) +@pytest.mark.parametrize("rows", [NAMED, ALL], ids=["named", "all"]) +@pytest.mark.parametrize("pending", [False, True], ids=["load_done", "load_pending"]) +def test_a_fetch_cannot_advance_a_window_past_a_connectors_pending_destinations( + kit, recorded, monkeypatch, window, rows, pending +): + """A named fetch under a window fails at the call while a connector load is pending.""" + connector = Connector() + if pending: + connector.loading.add(RID) + manager = Manager((window,), connector=connector) + lender = lender_over(manager, layout((window,))) + kv = manager.kv_cache_map[RID] = cache(history=2 * TPB) + grown = [] + + def grow(manager, request, cache, position, end): + grown.append((position, end)) + cache.history_length = position + return True + + monkeypatch.setattr(_manager, "grow", grow) + monkeypatch.setattr(_manager, "keeps_stale_pages", lambda manager: False) + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: True) + inputs_refuse_nothing(monkeypatch, lender) + served = request(served=2 * TPB) + before = lender.poll_resume_bounds(served) + lease = lender.lend_write(served, 2 * TPB, 3 * TPB, **getattr(kit, f"DEVICE_{rows.name}")) + try: + if pending and rows is NAMED and window is not None: + assert lease.failure is not None and "connector load" in lease.failure + assert lease.poll() is None and not lender._is_locked(kv) + assert not grown and kv.history_length == 2 * TPB + assert lender.poll_resume_bounds(served) == before, "the fetch left a record" + else: + view = lease.poll() + assert view is not None, lease.failure + lender.return_writes([(lease, ALL_ARRIVED, True)]) + assert len(grown) == (rows is NAMED) + assert kv.history_length == (3 * TPB if rows is NAMED and window else 2 * TPB) + finally: + lease.release() + + +# -- a window under a lock ---------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_the_last_history_keeping_a_block_is_the_last_whose_window_does(): + """Over random windows, sinks and blocks, the window keeps the block there, not one further.""" + rng = random.Random(0) + for _ in range(2000): + window = rng.randint(1, 300) + sinks = rng.choice([0, 0, 1, 2]) + block = rng.randint(sinks, 40) + manager = Manager((window,), (sinks,)) + lay = layout((window,), (sinks,)) + last = _ranges._last_history_keeping(manager, lay, 0, block) + keeps = _ranges._stale_range(manager, lay, 0, last)[1] <= block + passes = _ranges._stale_range(manager, lay, 0, last + 1)[1] > block + assert keeps and passes, (window, sinks, block, last) + + +@pytest.mark.cpu_only +def test_the_history_cap_keeps_the_lowest_locked_block_past_the_sinks_in_every_window(recorded): + """The cap keeps each window's lowest locked block past its sinks; other locks set none.""" + windows, sinks = (None, 64, 96), (0, 1, 0) + manager = Manager(windows, sinks) + lay = layout(windows, sinks) + lender = lender_over(manager, lay) + kv = cache(blocks=8) + assert lender._history_cap(kv) is None + lender._locks.lock(kv, {0: [0, 1], 1: [0]}, LockMode.READ) # full attention, and a sink + assert lender._history_cap(kv) is None + lock = lender._locks.lock(kv, {1: [0, 2, 3], 2: [4, 5]}, LockMode.READ) + + def last(lg, block): + return _ranges._last_history_keeping(manager, lay, lg, block) + + assert lender._history_cap(kv) == min(last(1, 2), last(2, 4)) + lender._locks.unlock(lock) + lender._locks.lock(kv, {2: [1]}, LockMode.READ) + assert lender._history_cap(kv) == last(2, 1) + assert lender._history_cap(cache()) is None, "another cache's locks set the cap" + + +@pytest.mark.cpu_only +def test_a_named_fetch_whose_grow_would_pass_a_locked_block_fails_at_the_call( + kit, recorded, monkeypatch +): + """A named fetch whose grow would leave a locked window block behind fails before the grow.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _fetch_ledger + + windows = (64,) + manager = Manager(windows) + lay = layout(windows) + lender = lender_over(manager, lay) + kv = manager.kv_cache_map[RID] = cache(blocks=4, history=2 * TPB, committed=2 * TPB) + grown = [] + monkeypatch.setattr(_manager, "grow", lambda *args: grown.append(args) or True) + monkeypatch.setattr(_manager, "keeps_stale_pages", lambda manager: False) + monkeypatch.setattr( + _fetch_ledger._FetchLedger, "_unsettled_fetch_refusal", lambda self, rid, kv: None + ) + inputs_refuse_nothing(monkeypatch, lender) + held = lender._locks.lock(kv, {0: [0]}, LockMode.READ) + end = 3 * TPB + assert _ranges._stale_range(manager, lay, 0, end)[1] > 0, "inconclusive: block 0 stays kept" + lease = lender.lend_write(request(), 2 * TPB, end, **kit.DEVICE_NAMED) + assert lease.failure is not None and "leave a locked block behind" in lease.failure + assert grown == [], "the cache grew for a fetch that failed at the call" + lease.release() + lender._locks.unlock(held) + + +# -- recurrent state ---------------------------------------------------------------------------- + + +@gpu +def test_a_fetch_from_a_held_commit_endpoint_caps_attention_retry(kit, real_manager): + """Unlock cannot publish a suffix while its replacement fetch remains unsettled.""" + with real_manager() as mgr: + lender = attach(mgr) + source = kit.make_request(SOURCE, PROMPT) + assert mgr.prepare_context(source) + kv = kit.kv(mgr, source) + read = lender._locks.lock(kv, {0: [0]}, LockMode.READ) + assert mgr.resize_context(source, 2 * TPB) + step(mgr, source, 2 * TPB) + assert source.context_current_position == 2 * TPB + assert mgr._held_commits[id(kv)][2] == 2 * TPB + assert mgr.resize_context(source, TPB) + step(mgr, source, TPB) + assert source.context_current_position == 3 * TPB + assert mgr._held_commits[id(kv)][2] == 2 * TPB + + fetch = lender.lend_write(source, 2 * TPB, 3 * TPB, **kit.HOST_NAMED) + try: + assert fetch.failure is None, fetch.failure + assert id(kv) in mgr._held_commits + lender._locks.unlock(read) + assert kv.num_committed_tokens == 2 * TPB + finally: + if lender._locks.is_locked(kv): + lender._locks.unlock(read) + fetch.release() + + +@gpu +def test_a_held_attention_commit_catches_up_without_a_fetch(kit, real_manager): + """Unlock can publish the later computed prefix when nothing replaced it.""" + prompt = PROMPT[: 3 * TPB] + with real_manager() as mgr: + lender = attach(mgr) + source = kit.make_request(SOURCE, prompt) + assert mgr.prepare_context(source) + kv = kit.kv(mgr, source) + read = lender._locks.lock(kv, {0: [0]}, LockMode.READ) + assert mgr.resize_context(source, 2 * TPB) + step(mgr, source, 2 * TPB) + assert mgr.resize_context(source, TPB) + step(mgr, source, TPB) + assert source.context_current_position == 3 * TPB + assert kv.num_committed_tokens == 0 + lender._locks.unlock(read) + assert kv.num_committed_tokens == 3 * TPB + + +@gpu +def test_a_state_group_s_regions_are_the_pools_native_names_and_its_row_the_state_slot( + kit, hybrid_manager +): + from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( + build_page_table_from_manager, + ) + from tensorrt_llm._torch.disaggregation.resource.page import CacheKind + + with hybrid_manager() as mgr: + lender = attach(mgr) + assert kit.staging_regions(lender) == () + lg = state_group(mgr) + page_table = build_page_table_from_manager(mgr) + native = page_table.layer_groups[lg] + assert native.kind == CacheKind.STATE + pools = sorted( + page_table.pool_groups[native.pool_group_idx].pools, key=lambda p: int(p.base_address) + ) + source = computed_to(mgr, kit.make_request(SOURCE, PROMPT), len(PROMPT)) + lease = lender.lend_read(source, 0, len(PROMPT), **kit.DEVICE_ALL) + view = kit.ready_view(lease, mgr) + run = next(g for g in view.groups if g.layer_group_index == lg) + regions = sorted((lender.regions[r] for r in run.region_indices), key=lambda r: r.address) + assert [(r.address, r.slot_bytes, r.slot_count) for r in regions] == [ + (int(p.base_address), int(p.slot_bytes), int(p.num_slots)) for p in pools + ] + slot = kit.kv(mgr, source).get_ssm_block_base_index(lg) + assert run.page_ordinals.tolist() == [0] and run.slots.tolist() == [slot] + assert run.object_keys is None + # Each state tensor of the slot lies in the slot's bytes of one region: the kernels read + # and write exactly what the row lends. + for number, tensor in enumerate(state_tensors(mgr, slot)): + tensor.view(torch.uint8).fill_(number + 1) + torch.cuda.synchronize() + parts = slot_parts(lender, run, slot) + for number, tensor in enumerate(state_tensors(mgr, slot)): + pointer = tensor.data_ptr() + base, nbytes = next((b, n) for b, n in parts if b <= pointer < b + n) + offset, size = pointer - base, tensor.numel() * tensor.element_size() + assert offset + size <= nbytes + got = device_tensor(base, nbytes)[offset : offset + size].cpu() + assert bool((got == number + 1).all()), f"state tensor {number} is not where lent" + lease.release() + + +@gpu +def test_a_lease_with_recurrent_state_ends_where_the_state_stands(kit, hybrid_manager): + """A fetch with a state row that ends short of the prompt fails at the call.""" + with hybrid_manager() as mgr: + lender = attach(mgr) + target, received = receiving(kit, mgr, TARGET, PROMPT, prepare=True) + keywords = {**kit.DEVICE_ALL, "ready_event": received} + refused(lender.lend_write(target, 0, 3 * TPB, **keywords), "where the state stands") + + +@gpu +def test_a_locked_state_row_holds_the_commit_that_would_snapshot_it_and_its_suspension( + kit, hybrid_manager +): + """A lock on the state row holds the snapshot commit and the suspension; the unlock commits.""" + with hybrid_manager() as mgr: + lender = attach(mgr) + source = computed_to(mgr, kit.make_request(SOURCE, PROMPT), 2 * TPB) + kv = kit.kv(mgr, source) + source.expect_snapshot_points = [2 * TPB] # the commit here snapshots the state + lg = state_group(mgr) + slot = kv.get_ssm_block_base_index(lg) + lock = lender._locks.lock(kv, {lg: [0]}, LockMode.READ) + with pytest.raises(RuntimeError, match="cannot be suspended"): + mgr.suspend_request(source) + mgr.try_commit_blocks(source) + assert kv.num_committed_tokens == 0, "a commit snapshot the state a lock holds" + assert kv.get_ssm_block_base_index(lg) == slot, "a held snapshot moved the locked row" + assert source.context_current_position == kv.history_length == 2 * TPB + lender._locks.unlock(lock) + assert kv.num_committed_tokens == 2 * TPB, "the unlock left the snapshot commit undone" + assert source.context_current_position == kv.history_length == 2 * TPB + + +@gpu +def test_a_stale_held_state_snapshot_is_not_committed_at_a_later_chunk(kit, hybrid_manager): + """A recurrent snapshot cannot be committed after its live state has moved past it.""" + with hybrid_manager() as mgr: + lender = attach(mgr) + source = computed_to(mgr, kit.make_request(SOURCE, PROMPT), 2 * TPB) + kv = kit.kv(mgr, source) + source.expect_snapshot_points = [2 * TPB] + lg = state_group(mgr) + first = lender._locks.lock(kv, {lg: [0]}, LockMode.READ) + second = lender._locks.lock(kv, {lg: [0]}, LockMode.READ) + + mgr.try_commit_blocks(source) + assert mgr._held_commits[id(kv)][2] == 2 * TPB + computed_to(mgr, source, 3 * TPB) + assert source.context_current_position == 3 * TPB + assert kv.num_committed_tokens == 0 + + lender._locks.unlock(first) + assert lender._locks.is_locked(kv) + assert id(kv) not in mgr._held_commits + assert kv.num_committed_tokens == 0 + lender._locks.unlock(second) + assert kv.num_committed_tokens == 0 + assert id(kv) not in mgr._held_commits + assert source.context_current_position == kv.history_length == 3 * TPB + + +@gpu +def test_a_held_state_snapshot_uses_the_latest_live_snapshot_before_stopping(kit, hybrid_manager): + """A final snapshot at the live history replaces an earlier locked snapshot.""" + with hybrid_manager() as mgr: + lender = attach(mgr) + source = computed_to(mgr, kit.make_request(SOURCE, PROMPT), 2 * TPB) + kv = kit.kv(mgr, source) + source.expect_snapshot_points = [2 * TPB, len(PROMPT)] + lg = state_group(mgr) + first = lender._locks.lock(kv, {lg: [0]}, LockMode.READ) + second = lender._locks.lock(kv, {lg: [0]}, LockMode.READ) + + mgr.try_commit_blocks(source) + computed_to(mgr, source, len(PROMPT)) + assert kv.num_committed_tokens == 0 + assert kv.history_length == len(PROMPT) + lender._locks.unlock(first) + assert id(kv) in mgr._held_commits + assert kv.num_committed_tokens == 0 + lender._locks.unlock(second) + assert kv.num_committed_tokens == len(PROMPT) + assert id(kv) not in mgr._held_commits + + +@pytest.mark.cpu_only +def test_a_fetch_inside_a_held_snapshot_uses_the_recurrent_retry_position(): + """A recurrent fetch drops a held commit that could publish the current state.""" + from tensorrt_llm._torch.pyexecutor.kv_cache import _lender_hooks + + kv = SimpleNamespace(id=RID) + source = request(computed=3 * TPB) + mgr = SimpleNamespace( + _held_commits={id(kv): (kv, source, 2 * TPB, False)}, + _held_commit_drop_end=lambda held: held[1].context_current_position, + ) + + _lender_hooks.drop_held_commit(mgr, kv, 3 * TPB) + assert id(kv) in mgr._held_commits + _lender_hooks.drop_held_commit(mgr, kv, 2 * TPB) + assert id(kv) not in mgr._held_commits + + +# -- Helix -------------------------------------------------------------------------------------- + + +@gpu +@pytest.mark.parametrize("rank", [0, 1]) +def test_each_helix_rank_lends_its_pages_of_the_prompt_and_resumes_at_its_end( + kit, real_manager, rank +): + """Rank ``r`` lends its page of each 64-token block ``b`` as block ``2 * b + r``.""" + from tensorrt_llm._torch.pyexecutor.request_utils import partition_context_for_helix + + with real_manager(mapping=helix_rank(rank), enable_block_reuse=False) as mgr: + lender = attach(mgr) + assert kit.staging_regions(lender) == () and lender.regions + # The rank's slice of the prompt, as the executor builds it, and the whole prompt's length. + tokens, _, total, _ = partition_context_for_helix(PROMPT, rank, 2, TPB) + request, admitted = receiving(kit, mgr, TARGET, tokens, total_input_len_cp=total) + assert request.prompt_len < total == len(PROMPT), "inconclusive: no rank-local prompt" + kv = kit.kv(mgr, request) + assert kv.history_length == len(PROMPT) + with pytest.raises(ValueError, match="lends nothing by name"): + lender.lend_write(request, 0, 2 * TPB, **{**kit.DEVICE_NAMED, "ready_event": admitted}) + keywords = {**kit.DEVICE_ALL, "ready_event": admitted} + with pytest.raises(ValueError, match="starts at a whole block of 64 tokens"): + lender.lend_write(request, TPB, len(PROMPT), **keywords) + lease = lender.lend_write(request, 0, len(PROMPT), **keywords) + assert lease.failure is None, lease.failure + view = kit.ready_view(lease, mgr) + (run,) = view.groups + expected = [p for p in range(4) if p % 2 == rank] + assert run.page_ordinals.tolist() == expected + pages = kit.pages(kv, 0) + assert run.slots.tolist() == [pages[p // 2] for p in expected] + bounds = return_write(lender, request, lease, view.row_masks(True)) + assert tuple(bounds) == (len(PROMPT), len(PROMPT)) + + +@gpu +@pytest.mark.parametrize("pp_rank", [0, 1]) +def test_helix_rejects_a_window_on_another_real_pipeline_stage(real_manager, pp_rank): + from tensorrt_llm.mapping import CpType, Mapping + + mapping = Mapping( + world_size=4, rank=2 * pp_rank, cp_size=2, pp_size=2, cp_config={"cp_type": CpType.HELIX} + ) + with real_manager( + mapping=mapping, windows=[256, 64], num_layers=2, enable_block_reuse=False + ) as mgr: + local = mgr.max_attention_window_vec + assert local == ([None] if pp_rank == 0 else [64]), (pp_rank, mgr.pp_layers, local) + with pytest.raises(ValueError, match="Helix"): + attach(mgr) + assert mgr._sharing is None + + +@gpu +def test_helix_ignores_a_window_pattern_entry_past_the_model(real_manager): + with real_manager( + mapping=helix_rank(0), windows=[256, 64], num_layers=1, enable_block_reuse=False + ) as manager: + assert manager.pp_layers == [0] and manager.max_attention_window_vec == [None] + assert manager._model_attention_window_pattern == (256, 64) + assert attach(manager).regions + + +@gpu +@pytest.mark.parametrize("rank", [0, 1]) +def test_helix_attaches_with_a_declared_read_window_and_full_kv_retention(kit, real_manager, rank): + declared = {0: 64, 1: 64} + with real_manager( + mapping=helix_rank(rank), + windows=[256], + attention_read_windows=declared, + enable_block_reuse=False, + ) as manager: + assert manager._model_attention_windows == declared + assert manager.max_attention_window_vec == [None, None] + layout = derive_layout(manager) + assert layout.windows and all(window is None for window in layout.windows) + lender = attach(manager) + assert manager._sharing is lender + assert lender.regions and kit.staging_regions(lender) == () + + +# -- bytes a receive need not write ------------------------------------------------------------- + + +@gpu +@pytest.mark.parametrize("mark_rows", [True, False], ids=["arrived", "not_arrived"]) +def test_a_generation_only_draft_layer_can_remain_unwritten(kit, real_manager, mark_rows): + """A receive leaving the draft layer's buffers unwritten resumes as its rows are marked.""" + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + spec = MTPDecodingConfig(max_draft_len=1) + spec.num_nextn_predict_layers = 1 + tokens = list(range(PROMPT_LEN)) + with ( + real_manager(num_layers=2) as source, + real_manager(num_layers=2, spec_config=spec, enable_block_reuse=False) as destination, + ): + sender, receiver = attach(source), attach(destination) + source_request = kit.published(source, 1, tokens) + destination_request, ready = receiving(kit, destination, 2, tokens) + assert kit.kv(destination, destination_request).history_length == PROMPT_LEN + source_layers = {b.layer for b in sender.page_layout.buffers} + destination_layers = {b.layer for b in receiver.page_layout.buffers} + assert source_layers == {0, 1}, source_layers + assert destination_layers == {0, 1, 2}, (destination_layers, destination.pp_layers) + assert destination.max_draft_len == 1 and not destination.enable_block_reuse + read = sender.lend_read(source_request, 0, PROMPT_LEN, **kit.DEVICE_ALL) + keywords = {**kit.DEVICE_ALL, "ready_event": ready} + write = receiver.lend_write(destination_request, 0, PROMPT_LEN, **keywords) + bounds = None + try: + source_view = kit.ready_view(read, source) + destination_view = kit.ready_view(write, destination) + fill_buffers(receiver, destination_view, destination._stream, STALE) + sent = { + (b.layer, b.role, ordinal): tensor + for b, ordinal, tensor in buffers(sender, source_view) + } + copied, omitted = [], [] + with torch.cuda.stream(destination._stream): + for buffer, ordinal, tensor in buffers(receiver, destination_view): + peer = sent.get((buffer.layer, buffer.role, ordinal)) + if peer is None: + assert buffer.layer == 2, (buffer, ordinal) + omitted.append(tensor) + else: + assert peer.numel() == tensor.numel(), (buffer, peer.shape, tensor.shape) + tensor.copy_(peer) + copied.append((tensor, peer)) + destination._stream.synchronize() + assert copied and omitted, (len(copied), len(omitted)) + assert all(torch.equal(tensor, peer) for tensor, peer in copied) + assert all(bool(torch.all(tensor == STALE)) for tensor in omitted) + bounds = return_write( + receiver, destination_request, write, destination_view.row_masks(mark_rows) + ) + finally: + write.release() + read.release() + assert bounds is not None, write.failure + if mark_rows: + assert bounds.max_position == PROMPT_LEN, (bounds, write.failure) + else: + assert bounds.max_position < PROMPT_LEN, (bounds, write.failure) + + +@gpu +@skip_pre_blackwell +def test_nvfp4_storage_padding_can_remain_unwritten(kit, monkeypatch): + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig + from tensorrt_llm.mapping import Mapping + + monkeypatch.setenv("TRTLLM_NVFP4_MLA_RESIDUAL_QUANTIZATION", "1") + prompt = 257 + with kit.managed( + lambda: DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig(max_tokens=2048, enable_block_reuse=False), + kv_cache_type=CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + max_seq_len=512, + max_batch_size=2, + max_input_len=512, + max_num_tokens=1024, + vocab_size=129280, + mapping=Mapping(), + dtype=DataType.NVFP4, + compressor_dtype=DataType.FLOAT, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[128] + ), + execution_stream=torch.cuda.Stream(), + ) + ) as manager: + lender = attach(manager) + request, ready = receiving(kit, manager, 3, list(range(prompt))) + assert kit.kv(manager, request).history_length == prompt + padding = [b for b in lender.page_layout.buffers if not b.transferred] + assert padding and all("padding" in b.role for b in padding), padding + lease = lender.lend_write(request, 0, prompt, **{**kit.DEVICE_ALL, "ready_event": ready}) + bounds = None + try: + view = kit.ready_view(lease, manager) + fill_buffers(lender, view, manager._stream, STALE) + written, omitted = [], [] + with torch.cuda.stream(manager._stream): + for buffer, _, tensor in buffers(lender, view): + if buffer.transferred: + tensor.fill_(FRESH) + written.append(tensor) + else: + omitted.append(tensor) + manager._stream.synchronize() + assert written and omitted, (len(written), len(omitted), padding) + assert all(bool(torch.all(tensor == FRESH)) for tensor in written) + assert all(bool(torch.all(tensor == STALE)) for tensor in omitted) + bounds = return_write(lender, request, lease, view.row_masks(True)) + finally: + lease.release() + assert bounds is not None and bounds.max_position == prompt, (bounds, lease.failure) + + +@gpu +@skip_pre_blackwell +@pytest.mark.parametrize("tail", [256, 262], ids=["at_a_page", "inside_a_page"]) +def test_csa2_remote_tail_resumes_with_decoder_buffers_unwritten(kit, monkeypatch, tail): + monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "1") + monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0") + model = remote_tail_layout() + prompt, split = tail + 128, 3 + with csa2_manager(kit, model, bounded_replay_on_generation=True) as manager: + lender = attach(manager) + request, ready = receiving( + kit, + manager, + 4, + list(range(prompt)), + py_csa2_remote_tail_mode="destination", + py_csa2_remote_tail_start=tail, + py_csa2_remote_tail_split=split, + ) + cache = kit.kv(manager, request) + assert cache.history_length == tail, (cache.history_length, tail, cache.capacity) + assert manager.bounded_replay_on_generation and manager._decoder_replay_window > 0 + swa_layers = sorted({b.layer for b in lender.page_layout.buffers if b.role == "csa2_swa"}) + assert len(swa_layers) == len(model.compress_ratios), swa_layers + decoder = set(swa_layers[split:]) + lease = lender.lend_write(request, 0, tail, **{**kit.DEVICE_ALL, "ready_event": ready}) + bounds = None + try: + view = kit.ready_view(lease, manager) + fill_buffers(lender, view, manager._stream, STALE) + written, omitted = [], [] + with torch.cuda.stream(manager._stream): + for buffer, _, tensor in buffers(lender, view): + if buffer.role == "csa2_swa" and buffer.layer in decoder: + omitted.append(tensor) + else: + tensor.fill_(FRESH) + written.append(tensor) + manager._stream.synchronize() + assert written and omitted, (len(written), len(omitted), decoder) + assert all(bool(torch.all(tensor == FRESH)) for tensor in written) + assert all(bool(torch.all(tensor == STALE)) for tensor in omitted) + bounds = return_write(lender, request, lease, view.row_masks(True)) + finally: + lease.release() + assert bounds == (tail, tail), (bounds, lease.failure, cache.history_length) + + +# -- CSA2 --------------------------------------------------------------------------------------- + + +@gpu +@skip_pre_blackwell +@pytest.mark.parametrize("policy", ["required", "hit", "partial-eviction", "evicted"]) +def test_a_csa2_receive_of_every_row_lends_each_group_what_native_receives( + kit, monkeypatch, policy +): + """Each CSA2 group receives its retained pages beyond its independently restored prefix.""" + from tensorrt_llm._torch.disaggregation.resource.cache_reuse import _CacheReuseAdapterV2 + from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( + build_page_table_from_manager, + ) + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType + from tensorrt_llm.bindings import SamplingConfig + from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection + + monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "1") + monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0" if policy == "required" else "1") + tpb, prompt = 128, 257 + + def request(request_id, kind): + return LlmRequest( + request_id=request_id, + max_new_tokens=2, + input_tokens=list(range(prompt)), + sampling_config=SamplingConfig(), + is_streaming=False, + llm_request_type=kind, + ) + + with csa2_manager( + kit, remote_tail_layout(), reuse=True, max_batch_size=16, max_num_tokens=512 * 16 + ) as mgr: + lender = attach(mgr) + warm = request(SOURCE, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + assert mgr.prepare_context(warm) + assert mgr.resize_context(warm, 2 * tpb) + cache = kit.kv(mgr, warm) + cache.commit(list(range(2 * tpb))) + optional = list(mgr._encoder_optional_groups) + missing = {"partial-eviction": optional[:1], "evicted": optional}.get(policy, []) + if missing: + _introspection.drop_optional_pages(cache, 2 * tpb, missing) + mgr.free_resources(warm) + target = request(TARGET, LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY) + assert mgr.prepare_disagg_gen_init(target) + kv = kit.kv(mgr, target) + assert kv.num_committed_tokens == 2 * tpb, "inconclusive: no prefix on the generation side" + assert [kv.reuse_status[g].complete for g in optional] == [ + g not in missing for g in optional + ] + ready = mgr.take_disagg_receive_ready(target) + lease = lender.lend_write(target, 0, prompt, **{**kit.DEVICE_ALL, "ready_event": ready}) + assert lease.failure is None, lease.failure + view = kit.ready_view(lease, mgr) + lent = {run.layer_group_index: run.page_ordinals.tolist() for run in view.groups} + page_table = build_page_table_from_manager(mgr) + cached = _CacheReuseAdapterV2(mgr).get_cached_token_count_per_layer_group( + target, page_table.layer_groups + ) + native = { + lg: kept_page_ordinals(group.sliding_window_size, prompt, tpb, cached[lg]) + for lg, group in enumerate(page_table.layer_groups) + } + assert lent == native, (policy, lent, native, cached) + bounds = return_write(lender, target, lease, view.row_masks(True)) + assert bounds.max_position == prompt + + +@gpu +@skip_pre_blackwell +@pytest.mark.parametrize("tail", [256, 262], ids=["at_a_block", "inside_a_block"]) +def test_a_csa2_remote_tail_source_lends_what_native_sends_once_its_context_is_computed( + kit, monkeypatch, tail +): + """A completed remote-tail source lends each group's retained pages below the tail.""" + from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( + build_page_table_from_manager, + ) + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType + from tensorrt_llm.bindings import SamplingConfig + + monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "1") + monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0") + tpb = 128 + prompt = tail + tpb # the generation side computes the rest + with csa2_manager(kit, remote_tail_layout(), bounded_replay_on_generation=True) as mgr: + lender = attach(mgr) + request = LlmRequest( + request_id=SOURCE, + max_new_tokens=1, + input_tokens=list(range(prompt)), + sampling_config=SamplingConfig(), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + ) + request.py_csa2_remote_tail_mode = "source" + request.py_csa2_remote_tail_start = tail + assert mgr.prepare_context(request) + assert request.context_remaining_length == tail, "inconclusive: no tail on this worker" + computed_to(mgr, request, tail) + assert request.context_remaining_length == 0 and request.prompt_len == prompt + lease = lender.lend_read(request, 0, tail, **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + view = kit.ready_view(lease, mgr) + lent = {run.layer_group_index: run.page_ordinals.tolist() for run in view.groups} + page_table = build_page_table_from_manager(mgr) + native = { + lg: kept_page_ordinals(group.sliding_window_size, tail, tpb) + for lg, group in enumerate(page_table.layer_groups) + } + assert lent == native, (lent, native) + windowed = [ + lg for lg, d in enumerate(lender.page_layout.layer_groups) if d.kept_window is not None + ] + assert windowed and all(lent[lg] for lg in windowed), (windowed, lent) + lease.release() + + +@gpu +@skip_pre_blackwell +@pytest.mark.parametrize( + "hole", + [None, "margin", "csa2_swa", "csa2_compressor_kv"], + ids=[ + "all_kept", + "missing_margin_page", + "missing_swa_read_page", + "missing_compressor_read_page", + ], +) +def test_a_real_csa2_receive_requires_its_kept_pages_to_resume(kit, monkeypatch, hole): + """A receive resumes with all retained pages; an absent margin or read page bounds it.""" + from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "0") + monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0") + tpb, prompt = 128, 514 + model = CSA2Layout((0, 2, 2, 1), (1, 3), (1, 3), window_size=256, index_topk=4) + spec = DSparkDecodingConfig(max_draft_len=8) + with csa2_manager(kit, model, max_seq_len=1024, vocab_size=8192, spec_config=spec) as mgr: + lender = attach(mgr) + request, ready = receiving(kit, mgr, TARGET, list(range(prompt))) + lease = lender.lend_write(request, 0, prompt, **{**kit.DEVICE_ALL, "ready_event": ready}) + view = kit.ready_view(lease, mgr) + masks, margin_roles, missing = [], set(), 0 + for run in view.groups: + desc = lender.page_layout.layer_groups[run.layer_group_index] + roles = {b.role for b in lender.page_layout.buffers if b.layer in desc.layers} + if "csa2_swa" in roles: + read_window = model.window_size + elif roles & {"csa2_compressor_kv", "csa2_compressor_score"}: + read_window = desc.kept_window - mgr.max_draft_len - mgr.reuse_match_backoff + else: + read_window = None + read_pages = set(kept_page_ordinals(read_window, prompt, tpb)) + kept_pages = set(kept_page_ordinals(desc.kept_window, prompt, tpb)) + lent = set(run.page_ordinals.tolist()) + assert read_pages <= kept_pages == lent, (roles, read_pages, kept_pages, lent) + margin = kept_pages - read_pages + if margin: + margin_roles.update(roles) + if hole == "margin" and margin and not missing: + kept_pages.remove(min(margin)) + missing += 1 + elif hole in roles: + kept_pages.remove(min(read_pages)) + missing += 1 + masks.append(np.isin(run.page_ordinals, list(kept_pages))) + assert margin_roles == {"csa2_swa", "csa2_compressor_kv", "csa2_compressor_score"} + assert (missing > 0) == (hole is not None) + write_through(lender, view, FRESH, masks) + bounds = return_write(lender, request, lease, masks) + if hole is None: + assert tuple(bounds) == (prompt, prompt) + else: + assert bounds.max_position < prompt or bounds.max_position < bounds.min_position diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_in_place_rows.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_in_place_rows.py new file mode 100644 index 000000000000..e67815c580ed --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_in_place_rows.py @@ -0,0 +1,1807 @@ +# 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. +"""Rows lent in place: recurrent state, Helix pages, windowed fetches, index and parked rows.""" + +import itertools +import random +import weakref +from functools import partial +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.pyexecutor.kv_cache import _lender_hooks +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + ALL_ARRIVED, + MemoryKind, + RowSelection, + StagingOptions, + _fetch_ledger, + _lender, + _manager, + _ranges, + attach_lender, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import Identity +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import DevicePool, derive_layout +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._slots import SlotQueue + +gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") + +TPB = 32 +PROMPT = 97 # three whole blocks and one token +TOKENS = list(range(1000, 1000 + PROMPT)) +STATE_SLOT = 5 +DEVICE, HOST = MemoryKind.DEVICE, MemoryKind.HOST +NAMED, ALL = RowSelection.NAMED, RowSelection.ALL +NO_NAMES = "layer groups [1] hold recurrent state, which no name covers" + + +class Manager: + """A stand-in manager: its layer groups' pool groups, its caches by request.""" + + _on_allocation_stream = frozenset() + + def _commit_held_back(self, kv): + """No commit is held back here.""" + + def __init__(self, pool_group_of): + self.impl = SimpleNamespace(get_life_cycle_pool_group_indices=lambda: list(pool_group_of)) + self.kv_cache_map = {} + self.kv_connector_manager = None + + +def layout(count=1, tpb=TPB, pages=16, **fields): + """A stand-in layout of ``count`` full-attention layer groups, each on a pool of its own.""" + lay = SimpleNamespace( + tokens_per_block=tpb, + windows=(None,) * count, + sink_blocks=(0,) * count, + recurrent=(False,) * count, + pool_groups=tuple(range(count)), + pool_group_of=tuple(range(count)), + num_layer_groups=count, + named_groups=(True,) * count, + cp=(0, 1), + device_pools={g: (DevicePool(g, 0, (g + 1) << 30, 256, pages),) for g in range(count)}, + ) + vars(lay).update(fields) + return lay + + +def hybrid_layout(): + """Layer group 0 of attention; layer group 1 of recurrent state, whose slots span two pools.""" + lay = layout(2, recurrent=(False, True)) + lay.device_pools[1] = (DevicePool(1, 0, 2 << 30, 4096, 8), DevicePool(1, 1, 3 << 30, 512, 8)) + return lay + + +def lender_over(manager, lay, no_names=NO_NAMES): + count = lay.num_layer_groups + ident = Identity(b"scope", bytes(range(32)), [(lg,) for lg in range(count)], [(1, 0)] * count) + ref = weakref.ref(manager) + regions = tuple(_lender._device_regions(manager, lay, ident)) + return _lender._ManagerLender(ref, lay, ident, regions, SlotQueue({}), ref, no_names=no_names) + + +def helix(rank): + """A stand-in manager under Helix over two ranks, as rank ``rank`` sees it, and its lender.""" + manager = Manager((0,)) + manager.enable_block_reuse = manager.is_draft = False + manager._has_cp_helix, manager._helix_cp_rank, manager._helix_cp_size = True, rank, 2 + lay = layout(tpb=2 * TPB, cp=(rank, 2)) + return manager, lender_over(manager, lay, no_names="under Helix") + + +def cache(blocks, pages=None, history=PROMPT, committed=0, tpb=TPB, groups=1, paged=None): + """A stand-in active cache of ``blocks`` blocks on ``pages`` in the groups ``paged`` flags, its + claim complete in each, its state page ``kv.state_slot``.""" + pages = list(range(blocks) if pages is None else pages) + + def pages_of(lg, valid_only=False): + return pages if paged is None or paged[lg] else [-1] * blocks + + kv = SimpleNamespace( + Status=SimpleNamespace(CLOSED="CLOSED"), + status="OPEN", + num_committed_tokens=committed, + reuse_status=[SimpleNamespace(complete=True, endpoint=committed)] * groups, + history_length=history, + is_active=True, + num_blocks=blocks, + tokens_per_block=tpb, + enable_swa_scratch_reuse=False, + state_slot=STATE_SLOT, + get_base_page_indices=pages_of, + get_aggregated_page_indices=pages_of, + ) + kv.get_ssm_block_base_index = lambda lg: kv.state_slot + return kv + + +def hybrid_cache(blocks=4, committed=0, history=PROMPT): + return cache(blocks, range(10, 10 + blocks), history, committed, groups=2, paged=(True, False)) + + +class Event: + """A stand-in CUDA event, complete at once; waiting on it fails the test.""" + + def query(self): + return True + + def synchronize(self): + raise AssertionError("a lease waited on its event") + + +READY = Event() + + +@pytest.fixture +def floor_follows_history(monkeypatch): + """The stand-in managers' resume floor follows the history.""" + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: True) + + +def return_write(lender, req, lease, arrived): + """The return's bounds, which ``poll_resume_bounds`` gives too, memoized and computed afresh.""" + (bounds,) = lender.return_writes([(lease, arrived, True)]) + assert lender.poll_resume_bounds(req) == bounds, bounds + delivered = lender._ledger._deliveries_of(int(req.py_request_id)) + if delivered is not None: + delivered.resume_bounds_memo = None + assert lender.poll_resume_bounds(req) == bounds, ("computed afresh", bounds) + return bounds + + +def request(request_id=7, prompt_len=PROMPT): + return SimpleNamespace( + py_request_id=request_id, + prompt_len=prompt_len, + py_multimodal_data=None, + context_current_position=0, + context_remaining_length=prompt_len, + ) + + +# -- recurrent state ----------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_a_state_group_lends_one_row_its_state_page_as_block_zero(): + manager = Manager((0, 1)) + lender = lender_over(manager, hybrid_layout()) + kv = hybrid_cache() + entries = lender._block_entries_of_range(manager, kv, 0, PROMPT) + assert [o.tolist() for o in entries.block_ordinals] == [[0, 1, 2, 3], [0]] + assert [p.tolist() for p in entries.device_pages] == [[10, 11, 12, 13], [STATE_SLOT]] + view = lender._in_place_view(entries, None, 0, 0) + attention, state = view.groups + assert (attention.region_indices, state.region_indices) == ((0,), (1, 2)) + assert state.page_ordinals.tolist() == [0] and state.slots.tolist() == [STATE_SLOT] + assert state.object_keys is None + kv.state_slot = -1 + assert lender._block_entries_of_range(manager, kv, 0, PROMPT).device_pages[1].tolist() == [-1] + + +@pytest.mark.cpu_only +def test_a_lender_naming_nothing_refuses_named_rows_and_lends_every_row_in_place(): + lender = lender_over(Manager((0, 1)), hybrid_layout()) + for memory in (DEVICE, HOST): + with pytest.raises(ValueError, match="lends nothing by name, as layer groups"): + lender._check_memory_and_rows(memory, NAMED) + lender._check_memory_and_rows(DEVICE, ALL) + with pytest.raises(ValueError, match="staging relays named whole blocks"): + lender._check_memory_and_rows(HOST, ALL) + named = lender_over(Manager((0, 1)), hybrid_layout(), no_names=None) + named._check_memory_and_rows(DEVICE, NAMED) # the check reads the lender's reason alone + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "history,end,lent", + [ + (2 * TPB, 2 * TPB, False), + (PROMPT, PROMPT, True), + (PROMPT + 3, PROMPT + 3, False), + (PROMPT, 2 * TPB, False), + ], + ids=["mid_context", "context_computed", "after_generation", "short_of_the_state"], +) +def test_a_read_lends_a_state_row_only_where_the_state_stands_at_the_prompt_end( + floor_follows_history, history, end, lent +): + """A read with a state row ends at the prompt's end; elsewhere it fails at the call.""" + manager = Manager((0, 1)) + lender = lender_over(manager, hybrid_layout()) + kv = manager.kv_cache_map[7] = hybrid_cache(blocks=5, history=history) + lease = lender.lend_read(request(), 0, end, memory_kind=DEVICE, rows=ALL, ready_event=READY) + if lent: + assert lease.failure is None, lease.failure + else: + assert lease.failure is not None and "where the state stands" in lease.failure + assert not lender._is_locked(kv), "a lease failed at the call locks blocks" + lease.release() + + +@pytest.mark.cpu_only +def test_a_state_row_lock_holds_every_commit_and_takes_a_write_past_any_commit( + floor_follows_history, +): + manager = Manager((0, 1)) + lender = lender_over(manager, hybrid_layout()) + kv = hybrid_cache(committed=2 * TPB) # two whole blocks committed + entries = lender._block_entries_of_range(manager, kv, 2 * TPB, PROMPT) + lease = lender._lend_in_place( + manager, request(), kv, "write", entries, None, start=0, end=PROMPT, ready_event=READY + ) + assert lease.failure is None, lease.failure + assert lease._lock.blocks[1].tolist() == [0] + # A commit of any blocks snapshots the state page, so the lock holds them all. + assert lender._is_locked(kv, 100, 200) and lender._is_locked(kv, 0, 1) + lease.release() + assert not lender._is_locked(kv, 100, 200) + + +def deliver(lender, manager, kv, masks_of): + """A fetch of every row of ``[0, PROMPT)`` returned with the rows ``masks_of`` gives.""" + committed = (kv.num_committed_tokens,) * 2 + fetch = lender._ledger._record_fetch( + 7, kv, 0, (0, 0), kv.history_length, committed, state_end=PROMPT, end=PROMPT, prompt=PROMPT + ) + entries = lender._block_entries_of_range(manager, kv, 0, PROMPT) + req = request() + lease = lender._lend_in_place( + manager, req, kv, "write", entries, None, fetch, start=0, end=PROMPT, ready_event=READY + ) + view = lease.poll() + assert view is not None, lease.failure + return_write(lender, req, lease, masks_of(view)) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("case", ["everything", "state_missing", "block_missing", "state_moved"]) +def test_a_fetch_with_recurrent_state_resumes_where_the_state_stands_or_nowhere( + floor_follows_history, case +): + """Resume at the prompt's end with the state and every block below, else nowhere.""" + manager = Manager((0, 1)) + lender = lender_over(manager, hybrid_layout()) + kv = manager.kv_cache_map[7] = hybrid_cache() + + def masks_of(view): + masks = view.row_masks(True) + if case == "state_missing": + masks[1][:] = False + if case == "block_missing": + masks[0][1] = False + if case == "state_moved": + kv.state_slot = STATE_SLOT + 1 # the state now lives on another page + return masks + + deliver(lender, manager, kv, masks_of) + got = lender.poll_resume_bounds(request()) + if case == "everything": + assert tuple(got) == (PROMPT, PROMPT) + else: + assert got.max_position < got.min_position, got + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("state_only", [True, False]) +def test_replacing_a_delivered_state_revokes_resume_until_the_new_state_arrives( + floor_follows_history, state_only +): + manager = Manager((0, 1)) + lender = lender_over(manager, hybrid_layout()) + kv = manager.kv_cache_map[7] = hybrid_cache() + req = request() + deliver(lender, manager, kv, lambda view: view.row_masks(True)) + assert tuple(lender.poll_resume_bounds(req)) == (PROMPT, PROMPT) + + entries = lender._block_entries_of_range(manager, kv, 0, PROMPT) + lease = SimpleNamespace(_request_id=7, _kv=kv, _block_entries=entries) + arrived = [np.ones(len(ordinals), dtype=bool) for ordinals in entries.block_ordinals] + if state_only: + arrived[0][:] = False + lender._ledger._invalidate_delivered_blocks(lease, arrived if state_only else None) + bounds = lender.poll_resume_bounds(req) + assert bounds.max_position < bounds.min_position + + fetch = _fetch_ledger._FetchRecord(0, kv, state_end=PROMPT, end=PROMPT) + lender._ledger._record_delivered_blocks(7, fetch, entries, arrived) + assert tuple(lender.poll_resume_bounds(req)) == (PROMPT, PROMPT) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("committed", [0, PROMPT]) +@pytest.mark.parametrize("delivery", ["local", "complete", "missing", "overwritten"]) +def test_a_read_requires_delivered_state_even_with_a_computed_attention_prefix( + floor_follows_history, committed, delivery +): + manager = Manager((0, 1)) + manager.enable_block_reuse = manager._can_publish_block_reuse = manager.is_draft = False + lender = lender_over(manager, hybrid_layout()) + manager.kv_cache_map[7] = hybrid_cache(committed=committed) + req = request() + + def fetch(complete): + write = lender.lend_write(req, 0, PROMPT, memory_kind=DEVICE, rows=ALL, ready_event=READY) + view = write.poll() + assert view is not None, write.failure + masks = view.row_masks(True) + masks[1][:] = complete + return_write(lender, req, write, masks) + + if delivery == "overwritten": + fetch(True) + if delivery != "local": + fetch(delivery == "complete") + read = lender.lend_read(req, 0, PROMPT, memory_kind=DEVICE, rows=ALL, ready_event=READY) + try: + if delivery in ("missing", "overwritten"): + assert read.failure is not None and "recurrent state" in read.failure + assert read.poll() is None + bounds = lender.poll_resume_bounds(req) + assert bounds.max_position < bounds.min_position + else: + assert read.poll() is not None, read.failure + finally: + read.release() + + +@pytest.mark.cpu_only +def test_a_committed_tail_check_reads_no_state_row(): + lender = lender_over(Manager((0, 1)), hybrid_layout()) + kv = hybrid_cache(committed=5) # five tokens of block 0 committed: its tail block is 0 + lent = _lender._LentBlockEntries( + [0, 1], [np.array([1], np.int64), np.array([0], np.int64)], [np.array([11]), np.array([5])] + ) + assert not lender._covers_committed_tail(kv, lent, [np.array([True]), np.array([True])]) + + +# -- Helix --------------------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "end,ordinals", + [(PROMPT, ([0, 2], [1, 3])), (90, ([0, 2], [1])), (64, ([0], [1]))], + ids=["prompt", "inside_rank_0_page", "whole_block"], +) +def test_each_helix_rank_lends_its_pages_of_the_range_under_their_own_blocks(end, ordinals): + """Page ``j`` of a block ``b`` of 64 tokens lies on rank ``j``, lent as block ``b * 2 + j``.""" + for rank in (0, 1): + manager, lender = helix(rank) + kv = cache(2, (3 + rank, 9 + rank), tpb=2 * TPB) + entries = lender._block_entries_of_range(manager, kv, 0, end) + (run,) = lender._in_place_view(entries, None, 0, 0).groups + assert run.page_ordinals.tolist() == ordinals[rank], (rank, end) + assert run.slots.tolist() == [(3, 9)[b // 2] + rank for b in ordinals[rank]] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("kind", ["recurrent", "window", "sinks"]) +def test_helix_refuses_recurrent_state_and_windows_at_the_attach(monkeypatch, kind): + """Under Helix the attach refuses a state row, a window or sinks, and passes a plain layout.""" + lay = SimpleNamespace( + cp=(0, 2), + num_layer_groups=2, + recurrent=(False, kind == "recurrent"), + windows=(None, 64 if kind == "window" else None), + sink_blocks=(0, 1 if kind == "sinks" else 0), + sparse=(False, False), + prefix_required=(), + ) + plain = SimpleNamespace(**vars(lay)) + plain.recurrent, plain.windows, plain.sink_blocks = (False, False), (None, None), (0, 0) + manager = SimpleNamespace( + mapping=SimpleNamespace(cp_size=2, pp_size=1, has_cp_helix=lambda: True), + _model_has_attention_window=False, + _model_has_sparse_buffers=False, + _model_attention_window_pattern=(), + num_layers=2, + kv_cache_manager_py_config=SimpleNamespace(layers=[]), + impl=SimpleNamespace(layer_grouping=[[], []]), + ) + monkeypatch.setattr(_manager, "require_v2", lambda manager: None) + monkeypatch.setattr(_lender, "derive_layout", lambda manager: lay) + monkeypatch.setattr(_manager, "model_has_recurrent_state", lambda manager: False) + with pytest.raises( + ValueError, match="under Helix the lender lends no recurrent state" + ) as raised: + _lender._checked_layout(manager) + assert "not shown to carry them end to end" in str(raised.value), "the refusal gives no reason" + monkeypatch.setattr(_lender, "derive_layout", lambda manager: plain) + assert _lender._checked_layout(manager) is plain + + +@pytest.mark.cpu_only +def test_a_helix_range_starts_at_a_whole_block_of_every_rank(): + _owner, lender = helix(0) + with pytest.raises(ValueError, match="starts at a whole block of 64 tokens"): + lender._whole_blocks(TPB, PROMPT, whole_end=False) + assert lender._whole_blocks(2 * TPB, PROMPT, whole_end=False) == (2 * TPB, PROMPT) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "rank,delivered,cap,usable", + [ + (0, [True, True], PROMPT, PROMPT), + (1, [True, True], PROMPT, PROMPT), + (1, [True, False], PROMPT, 96), # its page of block 1 starts at token 96 + (0, [True, False], PROMPT, 64), + (1, [True, False], 90, 90), # nothing of its page of block 1 lies below 90 + (1, [False, True], PROMPT, TPB), + ], +) +def test_a_helix_rank_resumes_up_to_its_first_missing_page_below_the_prompt( + rank, delivered, cap, usable +): + _owner, lender = helix(rank) + deliveries = _fetch_ledger._Deliveries(None, 0, [np.array(delivered)]) + assert lender._ledger._max_resume_position(None, deliveries, (0,), cap) == usable + + +@pytest.mark.cpu_only +def test_a_helix_lease_marks_and_resumes_per_rank(floor_follows_history): + for rank, expected in ((0, PROMPT), (1, 96)): + manager, lender = helix(rank) + kv = manager.kv_cache_map[7] = cache(2, (3 + rank, 9 + rank), tpb=2 * TPB) + fetch = lender._ledger._record_fetch( + 7, kv, 0, (0,), kv.history_length, (0,), end=PROMPT, prompt=PROMPT + ) + entries = lender._block_entries_of_range(manager, kv, 0, PROMPT) + req = SimpleNamespace( + py_request_id=7, prompt_len=49, total_input_len_cp=PROMPT, py_multimodal_data=None + ) + lease = lender._lend_in_place( + manager, req, kv, "write", entries, None, fetch, start=0, end=PROMPT, ready_event=READY + ) + masks = lease.poll().row_masks(True) + masks[0][-1] = rank == 0 # rank 1's page of block 1, token 96, does not arrive + got = return_write(lender, req, lease, masks) + assert got.max_position == expected, (rank, got) + assert got.min_position == PROMPT # the history, where the floor follows it + + +# -- windowed fetches --------------------------------------------------------------------------- + + +def _native_stale(history, tpb, window, sinks): + """AttnLifeCycle::getStaleRange, including C++ division's truncation toward zero.""" + begin = min(-(-history // tpb), sinks) + if window is None: + return begin, begin + numerator = history + 1 - window + end = numerator // tpb if numerator >= 0 else -((-numerator) // tpb) + return begin, max(begin, end) + + +def _native_needed(position, tpb, window, sinks): + """Required ordinals at the actual token position, independently of the ledger's search.""" + begin, end = _native_stale(position, tpb, window, sinks) + return [b for b in range(-(-position // tpb)) if not begin <= b < end] + + +def _window_case(windows, prompt, *, tpb=TPB, sinks=None, known=0, history=None, state=False): + sinks = tuple(sinks or (0,) * len(windows)) + windows = tuple(windows) + count = len(windows) + int(state) + manager = Manager(tuple(range(count))) + manager.tokens_per_block = tpb + manager.enable_block_reuse = False + manager._stale_block_range = lambda lg, h: _native_stale(h, tpb, windows[lg], sinks[lg]) + blocks = -(-prompt // tpb) + lay = layout( + count, + tpb, + max(blocks + 1, STATE_SLOT + 1), + windows=windows + (None,) * int(state), + sink_blocks=sinks + (0,) * int(state), + recurrent=(False,) * len(windows) + (True,) * int(state), + ) + lender = lender_over(manager, lay, no_names="this manager commits no blocks") + kv = cache(blocks, None, known if history is None else history, known, tpb, count) + manager.kv_cache_map[7] = kv + return manager, lender, kv, request(prompt_len=prompt) + + +def _fetch_all(lender, req, start, end, masks_of=lambda view: view.row_masks(True)): + lease = lender.lend_write(req, start, end, memory_kind=DEVICE, rows=ALL, ready_event=READY) + assert lease.failure is None, lease.failure + view = lease.poll() + assert view is not None + masks = masks_of(view) + return_write(lender, req, lease, masks) + return view, masks + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("delivery", ["local", "complete", "missing"]) +def test_a_read_refuses_a_cache_whose_committed_tail_was_not_delivered( + delivery, floor_follows_history +): + prompt = TPB + 4 + _owner, lender, kv, req = _window_case((None,), prompt, known=prompt, history=prompt) + if delivery != "local": + _fetch_all(lender, req, TPB, prompt, lambda view: view.row_masks(delivery == "complete")) + lease = lender.lend_read(req, TPB, prompt, memory_kind=DEVICE, rows=ALL, ready_event=READY) + try: + if delivery == "missing": + assert lease.failure is not None + assert lease.poll() is None and not lender._is_locked(kv) + else: + assert lease.poll() is not None, lease.failure + finally: + lease.release() + + +def _assert_interval(bounds, tpb, windows, sinks, known, delivered, runs=()): + def safe(position): + if any(begin < position < end for begin, end in runs): + return False + for lg, (window, sink) in enumerate(zip(windows, sinks)): + missing = [ + b + for b in _native_needed(position, tpb, window, sink) + if min((b + 1) * tpb, position) > known and b not in delivered[lg] + ] + if missing: + return False + return True + + if bounds.max_position < bounds.min_position: + return + assert bounds.min_position >= known, (bounds, known) + for position in range(bounds.min_position, bounds.max_position + 1): + assert safe(position), (position, bounds, known, delivered, runs) + # Completed bidirectional runs set the scheduler's floor even for a one-token run. + floor_limit = max(known, max((e for _, e in runs if e <= bounds.max_position), default=0)) + floor = bounds.max_position + while floor > floor_limit and safe(floor - 1): + floor -= 1 + assert bounds.min_position == floor, (bounds, floor, known, delivered, runs) + + +@pytest.mark.cpu_only +def test_a_fetch_of_every_row_may_end_at_the_history_its_cache_declares(floor_follows_history): + """A fetch may end at the declared history inside a block, and resumes there at most.""" + manager, lender, kv, req = _window_case((None,), 5 * TPB, history=100) + lease = lender.lend_write(req, 0, 90, memory_kind=DEVICE, rows=ALL, ready_event=READY) + assert lease.failure is not None and "ends at 90" in lease.failure, lease.failure + assert lease._lock is None, "a lease failed at the call holds a lock" + lease.release() + with pytest.raises(ValueError, match="past the request's"): + lender.lend_write(req, 0, 5 * TPB + 1, memory_kind=DEVICE, rows=ALL, ready_event=READY) + view, _ = _fetch_all(lender, req, 0, 100) + assert view.groups[0].page_ordinals.tolist() == [0, 1, 2, 3] + assert lender.poll_resume_bounds(req).max_position == 100 + assert kv.history_length == 100 + + +@pytest.mark.cpu_only +def test_a_fetch_skips_a_group_committed_past_its_end_and_fills_the_other_group( + floor_follows_history, +): + """A 34-token fetch must not overwrite another group's 36-token committed prefix.""" + _owner, lender, kv, req = _window_case((None, None), 64, tpb=8, known=36, history=0) + kv.reuse_status = [ + SimpleNamespace(complete=True, endpoint=36), + SimpleNamespace(complete=False, endpoint=36), + ] + prefix, _ = _fetch_all(lender, req, 0, 32) + assert [g.page_ordinals.tolist() for g in prefix.groups] == [[], [0, 1, 2, 3]] + kv.history_length = 34 + view, _ = _fetch_all(lender, req, 32, 34) + assert [g.page_ordinals.tolist() for g in view.groups] == [[], [4]] + assert lender.poll_resume_bounds(req) == (34, 34) + + view, _ = _fetch_all(lender, req, 32, 40) + assert [g.layer_group_index for g in view.groups] == [0, 1] + assert all(g.page_ordinals.tolist() == [4] for g in view.groups) + + +def _names_in_place(lender, monkeypatch): + """``lender`` lending by name, with a grow that moves the cache's history where asked.""" + lender._no_names = None + monkeypatch.setattr( + _manager, "block_keys", lambda manager, req, kv, n: [bytes([i]) * 32 for i in range(n)] + ) + + def grow(manager, req, kv, position, end): + kv.history_length = position + return True + + monkeypatch.setattr(_manager, "grow", grow) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("fetch", ["named_window", "every_row"]) +def test_a_read_of_every_row_fails_over_rows_a_fetch_moving_the_history_left_unwritten( + floor_follows_history, monkeypatch, fetch +): + """A read through the history fails over a block no fetch delivered, until a refetch does.""" + if fetch == "named_window": + _owner, lender, kv, req = _window_case((80,), 80, tpb=8) + _names_in_place(lender, monkeypatch) + req.try_get_encoder_output_len = lambda: None + req.py_return_context_logits, req.py_additional_outputs = False, None + end, missing, rows = 64, 6, NAMED + else: + _owner, lender, kv, req = _window_case((None,), 64, tpb=8, history=48) + end, missing, rows = 48, 3, ALL + + def deliver(missing_blocks): + lease = lender.lend_write(req, 0, end, memory_kind=DEVICE, rows=rows, ready_event=READY) + view = lease.poll() + assert view is not None, lease.failure + return_write(lender, req, lease, [~np.isin(view.groups[0].page_ordinals, missing_blocks)]) + + deliver([missing]) + assert kv.history_length == end + assert lender.poll_resume_bounds(req).max_position < end + read = lender.lend_read(req, 0, end, memory_kind=DEVICE, rows=ALL, ready_event=READY) + assert read.failure is not None and f"block {missing} holds rows" in read.failure, read.failure + assert read._lock is None, "a lease failed at the call holds a lock" + read.release() + # A request that computed on past the fetch's history wrote every row below it. + kv.history_length = end + 8 + read = lender.lend_read(req, 0, end, memory_kind=DEVICE, rows=ALL, ready_event=READY) + assert read.failure is None, read.failure + read.release() + kv.history_length = end + deliver([]) + assert lender.poll_resume_bounds(req).max_position == end + read = lender.lend_read(req, 0, end, memory_kind=DEVICE, rows=ALL, ready_event=READY) + assert read.failure is None and read.poll() is not None + assert missing in read.poll().groups[0].page_ordinals.tolist() + read.release() + + +@pytest.mark.cpu_only +def test_a_read_of_every_row_fails_over_committed_rows_a_fetch_wrote_and_did_not_deliver( + floor_follows_history, +): + """An undelivered partial committed tail cannot be read; a disjoint delivered block can.""" + _owner, lender, kv, req = _window_case((None,), 64, tpb=8, known=28, history=40) + _fetch_all(lender, req, 24, 48, lambda view: [view.groups[0].page_ordinals != 3]) + read = lender.lend_read(req, 0, 40, memory_kind=DEVICE, rows=ALL, ready_event=READY) + assert read.failure is not None and "block 3 holds rows" in read.failure, read.failure + read.release() + read = lender.lend_read(req, 32, 40, memory_kind=DEVICE, rows=ALL, ready_event=READY) + assert read.failure is None and read.poll().groups[0].page_ordinals.tolist() == [4] + read.release() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("boundary_arrives", [False, True], ids=["boundary_missing", "complete"]) +def test_a_partial_prompt_fetch_checks_the_window_at_its_actual_end( + floor_follows_history, boundary_arrives +): + manager, lender, kv, req = _window_case((64,), PROMPT) + + def masks_of(view): + assert view.groups[0].page_ordinals.tolist() == [1, 2, 3] + return [np.array([boundary_arrives, True, True])] + + _fetch_all(lender, req, 0, PROMPT, masks_of) + bounds = lender.poll_resume_bounds(req) + assert bounds.max_position == (PROMPT if boundary_arrives else 0), bounds + delivered = [{2, 3} | ({1} if boundary_arrives else set())] + _assert_interval(bounds, TPB, (64,), (0,), 0, delivered) + assert kv.history_length == 0 and manager.kv_cache_map[7] is kv + + +@pytest.mark.cpu_only +def test_every_position_in_a_windowed_fetch_s_resume_interval_has_its_rows(floor_follows_history): + manager, lender, kv, req = _window_case((64,), 4 * TPB) + view, _ = _fetch_all(lender, req, 0, req.prompt_len) + assert view.groups[0].page_ordinals.tolist() == [2, 3] + bounds = lender.poll_resume_bounds(req) + assert tuple(bounds) == (128, 127), bounds + _assert_interval(bounds, TPB, (64,), (0,), 0, [{2, 3}]) + # A second manager may lower the combined maximum to any point in this interval. + assert bounds.min_position <= 127 < bounds.max_position + assert kv.history_length == 0 and manager.kv_cache_map[7] is kv + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "kept,reads,sinks,prompt,runs,segments,expected", + [ + pytest.param( + (64,), + (64,), + (0,), + 160, + ((96, 150),), + ((0, 96, ({2},)), (96, 160, ({3},))), + (0, 0), + id="one_run", + ), + pytest.param( + (64, 160), + (64, 96), + (0, 1), + 384, + ((128, 184), (224, 280), (320, 360)), + ( + (0, 64, ({0, 1}, None)), + (64, 128, ({3}, None)), + (128, 192, ({4}, None)), + (192, 224, ({6}, None)), + (224, 288, ({7}, None)), + (288, 320, ({9}, None)), + (320, 384, ({10}, None)), + ), + (64, 0), + id="three_runs", + ), + ], +) +def test_bidirectional_clips_recheck_rows_from_public_fetches( + floor_follows_history, kept, reads, sinks, prompt, runs, segments, expected +): + manager, lender, kv, req = _window_case(kept, prompt, sinks=sinks) + lender._layout.read_windows = reads + mask = np.zeros(prompt, dtype=np.int64) + for begin, end in runs: + mask[begin:end] = 1 + req.py_multimodal_data = { + "mm_bidirectional_blocks": True, + "multimodal_embed_mask_cumsum": np.cumsum(mask), + } + valid = [{} for _ in kept] + for start, end, arrived in segments: + view, marks = _fetch_all( + lender, + req, + start, + end, + lambda view: [ + np.ones(len(g.page_ordinals), dtype=bool) + if arrived[g.layer_group_index] is None + else np.isin(g.page_ordinals, list(arrived[g.layer_group_index])) + for g in view.groups + ], + ) + for group, delivered in zip(view.groups, marks): + pages = valid[group.layer_group_index] + for ordinal, got in zip(group.page_ordinals.tolist(), delivered.tolist()): + pages.pop(ordinal, None) + if got: + pages[ordinal] = min(end, (ordinal + 1) * TPB) + + def safe(position): + if any(begin < position < end for begin, end in runs): + return False + for lg, window in enumerate(reads): + # Intersect actual token intervals; include native's in-flight token. + needed = set(range(sinks[lg])) | { + token // TPB for token in range(max(0, position + 1 - window), position + 1) + } + for ordinal in needed: + if ordinal * TPB < position and valid[lg].get(ordinal, 0) < min( + position, (ordinal + 1) * TPB + ): + return False + return True + + permitted = [position for position in range(prompt + 1) if safe(position)] + high = max(permitted) + low = high + floor = max((end for _, end in runs if end <= high), default=0) + while low > floor and safe(low - 1): + low -= 1 + assert (high, low) == expected + bounds = lender.poll_resume_bounds(req) + assert tuple(bounds) == expected, bounds + assert all(safe(position) for position in range(bounds.min_position, bounds.max_position + 1)) + assert kv.history_length == 0 and manager.kv_cache_map[7] is kv + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "attention_arrives,state_arrives", + [(False, True), (True, False), (True, True)], + ids=["attention_boundary_missing", "state_missing", "complete"], +) +def test_a_state_fetch_needs_the_attention_window_at_the_state_end( + floor_follows_history, attention_arrives, state_arrives +): + manager, lender, kv, req = _window_case((64,), PROMPT, history=PROMPT, state=True) + + def masks_of(view): + assert [g.page_ordinals.tolist() for g in view.groups] == [[1, 2, 3], [0]] + return [np.array([attention_arrives, True, True]), np.array([state_arrives])] + + _fetch_all(lender, req, 0, PROMPT, masks_of) + assert kv.history_length == PROMPT + bounds = lender.poll_resume_bounds(req) + if attention_arrives and state_arrives: + assert tuple(bounds) == (PROMPT, PROMPT), bounds + else: + assert bounds.max_position < bounds.min_position, bounds + assert manager.kv_cache_map[7] is kv + + +@pytest.mark.cpu_only +def test_public_windowed_fetch_bounds_match_a_native_formula_oracle(floor_follows_history): + """Small exhaustive masks and seeded layouts check every permitted token, not only the end.""" + cases = [] + tpb = 4 + for window, sink, known in itertools.product((None, 1, 3, 4, 5, 7, 8), (0, 1), (0, 3, 5)): + prompt = 3 * tpb + 1 + ordinals = [b for b in _native_needed(prompt, tpb, window, sink) if b >= known // tpb] + for bits in itertools.product((False, True), repeat=len(ordinals)): + cases.append((tpb, (window,), (sink,), known, prompt, (bits,), ())) + rng = random.Random(290053) + for _ in range(200): + tpb = rng.choice((4, 8, 16, 32)) + windows = tuple( + rng.choice((None, 1, tpb - 1, tpb + 1, 2 * tpb - 1, 3 * tpb + 3)) + for _ in range(rng.randint(1, 3)) + ) + sinks = tuple(rng.randrange(3) for _ in windows) + prompt = rng.randint(2 * tpb, 7 * tpb + 1) + known = rng.choice((0, tpb - 1, tpb + 1, prompt // 2)) + bits = tuple( + tuple( + rng.random() < 0.7 for b in _native_needed(prompt, tpb, w, s) if b >= known // tpb + ) + for w, s in zip(windows, sinks) + ) + runs = () + if rng.random() < 0.4 and prompt > known + 2: + begin = rng.randint(known, prompt - 2) + runs = ((begin, rng.randint(begin + 1, prompt)),) + cases.append((tpb, windows, sinks, known, prompt, bits, runs)) + + for case in cases: + tpb, windows, sinks, known, prompt, bits, runs = case + manager, lender, kv, req = _window_case(windows, prompt, tpb=tpb, sinks=sinks, known=known) + if runs: + mask = np.zeros(prompt, dtype=np.int64) + for begin, end in runs: + mask[begin:end] = 1 + req.py_multimodal_data = { + "mm_bidirectional_blocks": True, + "multimodal_embed_mask_cumsum": np.cumsum(mask), + } + view, masks = _fetch_all( + lender, req, known // tpb * tpb, prompt, lambda view: [np.array(b, bool) for b in bits] + ) + delivered = [ + {int(b) for b, arrived in zip(g.page_ordinals, m) if arrived} + for g, m in zip(view.groups, masks) + ] + bounds = lender.poll_resume_bounds(req) + tainted_tail = known % tpb and any( + known // tpb in g.page_ordinals and known // tpb not in got + for g, got in zip(view.groups, delivered) + ) + if tainted_tail: + assert bounds.max_position < bounds.min_position, (case, bounds) + continue + _assert_interval(bounds, tpb, windows, sinks, known, delivered, runs) + assert bounds.max_position <= prompt, (case, bounds) + if any(len(g) for g in view.groups) and all(m.all() for m in masks): + assert bounds.max_position == prompt and bounds.min_position <= prompt, (case, bounds) + assert kv.history_length == known and manager.kv_cache_map[7] is kv + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("window", [None, 72]) +def test_a_complete_refetch_replaces_an_earlier_partial_page(floor_follows_history, window): + _owner, lender, _kv, req = _window_case((window,), 160, history=65) + _fetch_all(lender, req, 0, 65) + assert lender.poll_resume_bounds(req) == (65, 65) + _fetch_all(lender, req, 0, 160) + assert lender.poll_resume_bounds(req) == (160, 65) + + +@pytest.mark.cpu_only +def test_a_partial_page_behind_the_window_still_bounds_every_rewind(floor_follows_history): + _owner, lender, _kv, req = _window_case((72,), 192, history=65) + lender._layout.read_windows = (64,) + _fetch_all(lender, req, 0, 65) + _fetch_all(lender, req, 96, 192, lambda view: [view.groups[0].page_ordinals < 5]) + assert lender.poll_resume_bounds(req) == (160, 159) + mask = np.zeros(193, dtype=int) + mask[152:168] = 1 + req.py_multimodal_data = { + "mm_bidirectional_blocks": True, + "multimodal_embed_mask_cumsum": np.cumsum(mask), + } + lender._ledger._deliveries_of(7).resume_bounds_memo = None + assert lender.poll_resume_bounds(req) == (65, 65) + + +@pytest.mark.cpu_only +def test_a_fetch_with_no_read_history_uses_its_declared_extent(floor_follows_history): + _owner, lender, _kv, req = _window_case((9,), 160) + lender._layout.read_windows = (1,) + _fetch_all(lender, req, 0, 32, lambda view: view.row_masks(False)) + assert lender.poll_resume_bounds(req) == (32, 32) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("arrived", [True, False]) +def test_a_state_only_stage_constrains_the_pipeline_at_its_delivered_state( + floor_follows_history, arrived +): + """A stage of state alone pins the pipeline's resume to the prompt's end, or empties it.""" + _owner, lender, _kv, req = _window_case((), PROMPT, history=PROMPT, state=True) + _fetch_all(lender, req, 0, PROMPT, lambda view: view.row_masks(arrived)) + state = lender.poll_resume_bounds(req) + _other, attention_lender, _other_kv, other_req = _window_case((None,), PROMPT) + _fetch_all(attention_lender, other_req, 0, PROMPT) + attention = attention_lender.poll_resume_bounds(other_req) + assert attention.max_position == PROMPT and attention.min_position < PROMPT, attention + pipeline = ( + min(state.max_position, attention.max_position), + max(state.min_position, attention.min_position), + ) + if arrived: + assert state == (PROMPT, PROMPT) and pipeline == (PROMPT, PROMPT) + else: + assert state.max_position < state.min_position and pipeline[0] < pipeline[1] + + +@pytest.mark.cpu_only +def test_a_holder_changing_its_view_cannot_change_the_arrival_mask_shape(floor_follows_history): + """Arrivals are checked against the lease's own entries; a refused return leaves it lent.""" + _owner, lender, _kv, req = _window_case((None,), PROMPT) + lease = lender.lend_write(req, 0, 96, memory_kind=DEVICE, rows=ALL, ready_event=READY) + view = lease.poll() + assert view is not None + object.__setattr__(view.groups[0], "page_ordinals", np.zeros(1, dtype=np.int64)) + with pytest.raises(ValueError, match="shape"): + lender.return_writes([(lease, [np.array([True])], True)]) + assert return_write(lender, req, lease, [np.zeros(3, dtype=bool)]).max_position == 0 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("stop", ["error", "shutdown", "gone"]) +def test_invalid_write_ends_are_checked_before_the_lender_stops(floor_follows_history, stop): + manager, lender, _, req = _window_case((None,), PROMPT) + if stop == "error": + lender._stop_reason = "injected" + elif stop == "shutdown": + lender._shut_down = True + else: + del manager + with pytest.raises(ValueError, match="past the request"): + lender.lend_write(req, 0, PROMPT + 1, memory_kind=DEVICE, rows=ALL, ready_event=READY) + with pytest.raises(RuntimeError, match="lends no more"): + lender.lend_write(req, 0, 96, memory_kind=DEVICE, rows=ALL, ready_event=READY) + + +@pytest.mark.cpu_only +def test_partial_refetch_intervals_match_the_native_window_rule_at_every_position( + floor_follows_history, +): + rng = random.Random(31) + for _ in range(300): + read = rng.randint(1, 160) + kept = read + rng.choice((0, 1, 8, 32)) + partial = rng.randint(1, 95) + prompt = 192 + sink = rng.choice((0, 0, 1)) + _owner, lender, _kv, req = _window_case((kept,), prompt, history=partial, sinks=(sink,)) + lender._layout.read_windows = (read,) + valid = {} + for start, end in ((0, partial), (rng.choice((0, 96)), prompt)): + lease = lender.lend_write( + req, start, end, memory_kind=DEVICE, rows=ALL, ready_event=READY + ) + view = lease.poll() + assert view is not None, lease.failure + (group,) = view.groups + marks = np.array([rng.random() < 0.8 for _ in group.page_ordinals], dtype=bool) + for ordinal, arrived in zip(group.page_ordinals.tolist(), marks): + valid.pop(ordinal, None) + if arrived: + valid[ordinal] = min(end, (ordinal + 1) * TPB) + return_write(lender, req, lease, [marks]) + safe = [ + p + for p in range(partial, prompt + 1) + if all( + valid.get(b, 0) >= min(p, (b + 1) * TPB) for b in _native_needed(p, TPB, read, sink) + ) + ] + bounds = lender.poll_resume_bounds(req) + if not safe: + assert bounds.max_position < bounds.min_position, (read, kept, partial, valid, bounds) + continue + assert bounds.max_position == max(safe), (read, kept, partial, valid, bounds, safe) + assert all(p in safe for p in range(bounds.min_position, bounds.max_position + 1)) + + +@pytest.mark.cpu_only +def test_refetch_progress_tracks_helix_physical_pages(floor_follows_history): + for rank in (0, 1): + manager, lender = helix(rank) + manager.kv_cache_map[7] = cache(3, (3, 9, 12), history=96, tpb=2 * TPB) + req = request(prompt_len=80) + req.total_input_len_cp = 160 + _fetch_all(lender, req, 0, 96) + assert lender.poll_resume_bounds(req) == (96, 96) + _fetch_all(lender, req, 64, 160) + assert lender.poll_resume_bounds(req) == (160, 96) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "kept,reads,floor", + [((None, 64), (), 191), ((128, 64), (64, 32), 159)], + ids=["full_attention_and_window", "two_windows"], +) +def test_a_partial_page_a_refetch_leaves_raises_the_floor_of_its_own_layer_group_alone( + floor_follows_history, kept, reads, floor +): + """A refetch from 96 re-lends block 3 in the wider group alone; the other's stops at 100.""" + _owner, lender, _kv, req = _window_case(kept, 192, history=100) + lender._layout.read_windows = reads + first, _ = _fetch_all(lender, req, 0, 100) + assert first.groups[1].page_ordinals.tolist() == [1, 2, 3] + second, _ = _fetch_all(lender, req, 96, 192) + assert [g.page_ordinals.tolist() for g in second.groups] == [[3, 4, 5], [4, 5]] + assert lender.poll_resume_bounds(req) == (192, floor) + + +@pytest.mark.cpu_only +def test_a_group_whose_block_a_refetch_left_undelivered_keeps_the_resume_below_it( + floor_follows_history, +): + _owner, lender, _kv, req = _window_case((None, None), 192, history=100) + _fetch_all(lender, req, 0, 100) + _fetch_all(lender, req, 96, 192, lambda view: [view.row_masks(True)[0], view.row_masks()[1]]) + assert lender.poll_resume_bounds(req).max_position == 96 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("history", [80, 95]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_a_helix_rank_s_partly_fetched_page_bounds_a_later_fetch_s_resume( + floor_follows_history, rank, history +): + """A rank resumes where its own pages hold every token: rank 0 at the history, rank 1 at 96.""" + manager, lender = helix(rank) + manager.kv_cache_map[7] = cache(3, (3, 9, 12), history=history, tpb=2 * TPB) + req = request(prompt_len=80) + req.total_input_len_cp = 160 + first, _ = _fetch_all(lender, req, 0, history) + assert first.groups[0].page_ordinals.tolist() == ([0, 2] if rank == 0 else [1]) + second, _ = _fetch_all(lender, req, 128, 160) + assert second.groups[0].page_ordinals.tolist() == ([4] if rank == 0 else []) + assert lender.poll_resume_bounds(req).max_position == (history if rank == 0 else 96) + + +@pytest.mark.cpu_only +def test_a_shrink_past_a_partial_page_bounds_the_resume_at_the_blocks_it_keeps( + floor_follows_history, +): + _owner, lender, kv, req = _window_case((None,), 160, history=100) + _fetch_all(lender, req, 0, 100) + assert lender.poll_resume_bounds(req).max_position == 100 + kv.num_blocks = 3 + lender._on_shrink(7, kv) + assert lender.poll_resume_bounds(req).max_position == 96 + + +@pytest.mark.cpu_only +def test_a_shorter_refetch_keeps_what_the_earlier_fetch_delivered_past_it(floor_follows_history): + _owner, lender, _kv, req = _window_case((None,), 160) + _fetch_all(lender, req, 0, 160, lambda view: [view.groups[0].page_ordinals > 0]) + assert lender.poll_resume_bounds(req).max_position == 0 + _fetch_all(lender, req, 0, 32) + assert lender.poll_resume_bounds(req) == (160, 0) + + +@pytest.mark.cpu_only +def test_refetch_intervals_match_the_native_rule_over_layer_groups_and_shorter_refetches( + floor_follows_history, +): + """Random refetches into one to three groups resume where the native window rule permits.""" + rng = random.Random(83) + prompt = 192 + checked = 0 + for _ in range(600): + count = rng.randint(1, 3) + reads, kept, sinks = [], [], [] + for _ in range(count): + if rng.random() < 0.4: + reads.append(None) + kept.append(None) + sinks.append(0) + else: + read = rng.randint(1, 160) + reads.append(read) + kept.append(read + rng.choice((0, 1, 8, 32))) + sinks.append(rng.choice((0, 0, 1))) + known = rng.choice((0, 0, rng.randint(1, 95), rng.randint(1, 160))) + complete = [rng.random() < 0.6 for _ in range(count)] + history = rng.choice((0, known, rng.randint(0, known))) + _owner, lender, kv, req = _window_case( + tuple(kept), prompt, history=history, sinks=tuple(sinks), known=known + ) + lender._layout.read_windows = tuple(reads) + kv.reuse_status = [SimpleNamespace(complete=c, endpoint=known) for c in complete] + committed = [known if c else 0 for c in complete] + computed = [max(c, history) for c in committed] + valid = [{} for _ in range(count)] + tainted = False + extent = 0 # the furthest end the fetches declare + ends = sorted({e for e in (history, prompt, *range(TPB, prompt, TPB)) if e >= history}) + for _ in range(rng.randint(1, 4)): + end = rng.choice(ends) + start = rng.choice([s for s in range(0, prompt, TPB) if s <= end]) + lease = lender.lend_write( + req, start, end, memory_kind=DEVICE, rows=ALL, ready_event=READY + ) + if lease.failure is not None and "committed tokens" in lease.failure: + # A block the fetch lends would hold committed tokens past its end. + assert any(end < c and c // TPB == end // TPB for c in committed), lease.failure + lease.release() + continue + view = lease.poll() + assert view is not None, (lease.failure, start, end, history, known, complete) + extent = max(extent, end) + if history > min(committed): + computed = [max(c, min(t, start)) for c, t in zip(committed, computed)] + masks = [] + for group in view.groups: + lg = group.layer_group_index + arrived = np.array([rng.random() < 0.8 for _ in group.page_ordinals], dtype=bool) + for ordinal, got in zip(group.page_ordinals.tolist(), arrived): + if got: + valid[lg][ordinal] = min(end, (ordinal + 1) * TPB) + else: + valid[lg].pop(ordinal, None) + tainted |= ordinal * TPB < committed[lg] + masks.append(arrived) + return_write(lender, req, lease, masks) + + def safe(position): + for lg in range(count): + for b in _native_needed(position, TPB, reads[lg], sinks[lg]): + need = min(position, (b + 1) * TPB) + if need > computed[lg] and valid[lg].get(b, 0) < need: + return False + return True + + bounds = lender.poll_resume_bounds(req) + case = (reads, kept, sinks, known, complete, history, valid, computed, bounds) + # No resume passes the furthest end the fetches declare, or what every group computed. + top = min(prompt, max(min(computed), extent)) + permitted = [p for p in range(history, top + 1) if safe(p)] + if tainted or not permitted: + assert bounds.max_position < bounds.min_position, case + continue + assert bounds.max_position == max(permitted), case + low = bounds.max_position + while low > history and safe(low - 1): + low -= 1 + assert bounds.min_position == low, case + checked += 1 + assert checked > 150, f"only {checked} cases resumed anywhere" + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("prefix", ["fetched", "computed"]) +@pytest.mark.parametrize("valid_tokens", [TPB - 2, TPB - 1], ids=["read_hole", "last_read_covered"]) +def test_a_partial_prefix_keeps_the_pipeline_resume_only_when_it_covers_the_last_read( + prefix, valid_tokens, floor_follows_history +): + _owner, lender, _kv, req = _window_case((2,), 64, history=valid_tokens) + lender._layout.read_windows = (1,) + if prefix == "fetched": + _fetch_all(lender, req, 0, valid_tokens) + _fetch_all(lender, req, 32, 64) + bounds = lender.poll_resume_bounds(req) + covered = valid_tokens == TPB - 1 + assert bounds == (64, valid_tokens if covered else 32), (prefix, valid_tokens, bounds) + for position in range(bounds.min_position, bounds.max_position + 1): + required = _native_needed(position, TPB, 1, 0) + assert all( + min(position, (b + 1) * TPB) <= (valid_tokens if b == 0 else 64) for b in required + ) + + _peer_owner, peer, _, peer_req = _window_case((None,), 64, history=valid_tokens) + if prefix == "fetched": + _fetch_all(peer, peer_req, 0, valid_tokens) + _fetch_all(peer, peer_req, 32, 64, lambda view: view.row_masks(False)) + peer_bounds = peer.poll_resume_bounds(peer_req) + assert peer_bounds == (valid_tokens, valid_tokens), peer_bounds + intersection = ( + min(bounds.max_position, peer_bounds.max_position), + max(bounds.min_position, peer_bounds.min_position), + ) + assert (intersection[0] >= intersection[1]) == covered, intersection + + +# -- page-index rows, parked rows and state slots of real managers ------------------------------- + + +def attach(mgr, staging_tokens=3 * TPB): + return attach_lender(mgr, scope=b"rows", staging=StagingOptions(staging_tokens)) + + +@pytest.fixture +def cache_page_reads(monkeypatch): + """How many times the lender asked a cache for its pages.""" + reads = [0] + for name in ("pages", "page_list"): + ask = getattr(_manager, name) + + def counted(kv, layer_group, _ask=ask): + reads[0] += 1 + return _ask(kv, layer_group) + + monkeypatch.setattr(_manager, name, counted) + return reads + + +def row_matches_cache(mgr, lender, request_id): + """Whether the request's page-index row holds its cache's pages; ``None`` without one.""" + row = _manager.index_row(mgr, request_id) + if row is None: + return None + kv = mgr.kv_cache_map[request_id] + blocks = int(kv.num_blocks) + for lg in range(len(mgr.impl.layer_grouping)): + cache = np.asarray(kv.get_base_page_indices(lg))[:blocks] + if not np.array_equal(cache, lender._index_rows[lg, row, 0, :blocks]): + return False + return True + + +def parked_matches_cache(mgr, lender, request_id): + """Whether the request's parked row holds its cache's pages, then -1; ``None`` without a row.""" + row = lender._parked.get(request_id) + if row is None: + return None + kv = mgr.kv_cache_map[request_id] + blocks = int(kv.num_blocks) + for lg in range(len(mgr.impl.layer_grouping)): + cache = np.asarray(kv.get_base_page_indices(lg))[:blocks] + if ( + not np.array_equal(cache, row.pages[lg, :blocks]) + or (row.pages[lg, blocks:] != -1).any() + ): + return False + return True + + +def assert_lent_pages_are_the_cache_s(kit, mgr, lender, request, view): + """Each row of ``view``, recurrent state's aside, lends the page the cache gives its block.""" + kv = kit.kv(mgr, request) + for run in view.groups: + if run.layer_group_index not in lender._recurrent_groups: + cache = np.asarray(kv.get_base_page_indices(run.layer_group_index)) + assert run.slots.tolist() == cache[run.page_ordinals].tolist() + + +def receive_target(kit, mgr, request_id, tokens): + """A generation-init request admitted to receive ``tokens``, or ``None`` without room.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestType + + request = kit.make_request( + request_id, tokens, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + ) + if not mgr.prepare_disagg_gen_init(request): + return None + mgr.take_disagg_receive_ready(request) + return request + + +def read_every_row(kit, mgr, lender, request, tokens, cache_page_reads): + """A fast read of every row of ``request``'s first ``tokens``, asking the cache for no page.""" + before = cache_page_reads[0] + lease = lender.lend_read(request, 0, tokens, **kit.DEVICE_ALL) + assert lease.failure is None and lease._block_entries.spans is not None, lease.failure + assert cache_page_reads[0] == before, "the read asked the cache for pages" + assert_lent_pages_are_the_cache_s(kit, mgr, lender, request, kit.ready_view(lease, mgr)) + return lease + + +def prepare(mgr, *requests): + """The requests' resources prepared as for a batch's context requests before a forward.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + batch = ScheduledRequests() + for request in requests: + batch.append_context_request(request) + mgr.prepare_resources(batch) + + +def chunk(mgr, request, tokens): + """One context chunk of ``tokens``: allocated, prepared, computed and updated.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + assert mgr.resize_context(request, tokens) + prepare(mgr, request) + request.context_chunk_size = tokens + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + mgr.update_context_resources(batch) + + +def computed_request(kit, mgr, request_id, tokens): + """A request whose prompt the manager prepared and computed in one chunk.""" + request = kit.admitted(mgr, request_id, tokens) + chunk(mgr, request, request.context_remaining_length) + assert request.context_remaining_length == 0 + return request + + +@gpu +@pytest.mark.parametrize("groups", [1, 2], ids=["one_group", "two_groups"]) +@pytest.mark.parametrize("blocks", [3, _lender._LISTED_BLOCKS + 6], ids=["listed", "viewed"]) +def test_a_write_of_every_row_lends_its_row_s_pages_without_asking_the_cache( + kit, real_manager, cache_page_reads, blocks, groups +): + """A receive's write of every row lends its row's pages, asking the cache for none.""" + tokens = list(range(2000, 2000 + blocks * TPB + 1)) + windows = None if groups == 1 else [(blocks + 3) * TPB, (blocks + 4) * TPB] + with real_manager( + max_tokens=4 * (blocks + 2) * TPB, max_seq_len=(blocks + 4) * TPB, windows=windows + ) as mgr: + lender = attach(mgr) + assert len(mgr.impl.layer_grouping) == groups, "inconclusive: another number of groups" + target = receive_target(kit, mgr, 7, tokens) + assert target is not None + mgr._stream.synchronize() + assert row_matches_cache(mgr, lender, 7) + lease = lender.lend_write(target, 0, len(tokens), **kit.DEVICE_ALL) + assert lease.failure is None and lease._block_entries.spans is not None, lease.failure + assert cache_page_reads[0] == 0, "the write asked the cache for pages" + assert_lent_pages_are_the_cache_s(kit, mgr, lender, target, lease.poll()) + bounds = kit.return_write(lender, lease, ALL_ARRIVED) + assert bounds.max_position == len(tokens), bounds + mgr._stream.synchronize() + mgr.free_resources(target) + + +@gpu +def test_the_rows_hold_the_cache_s_pages_through_the_executor_s_operations(kit, host_tier_manager): + """A random walk of the executor's operations: every request's row holds its cache's pages.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + rng = random.Random(11) + done = dict.fromkeys( + ("compared", "released", "reused", "resumed", "generated", "parked_compared"), 0 + ) + done.update(dict.fromkeys(("parked_resumed", "deferred", "parked_reused"), 0)) + with host_tier_manager() as mgr: + lender = attach(mgr) + live, rows_of, seen_rows, receives, next_id = {}, {}, set(), set(), 100 + held, parked_seen, closing = {}, set(), [] + for _ in range(1000): + action = rng.random() + if action < 0.2 and len(live) < 6: + next_id += 1 + tokens = list(range(next_id * 1000, next_id * 1000 + rng.choice((17, 33, 64)))) + if rng.random() < 0.5: + request = receive_target(kit, mgr, next_id, tokens) + if request is None: + continue + receives.add(next_id) + else: + request = kit.make_request(next_id, tokens) + if not mgr.prepare_context(request): + continue + kit.kv(mgr, request).enable_swa_scratch_reuse = False + live[next_id] = request + row = _manager.index_row(mgr, next_id) + done["reused"] += row in seen_rows + seen_rows.add(row) + rows_of[next_id] = row + elif action < 0.42 and live: + rid = rng.choice(list(live)) + request = live[rid] + kv = kit.kv(mgr, request) + if rid not in receives and request.context_remaining_length and kv.is_active: + step = min(TPB, request.context_remaining_length) + if mgr.resize_context(request, step): + request.context_chunk_size = step + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + mgr.update_context_resources(batch) + elif action < 0.6 and live: + rid = rng.choice(list(live)) + request = live[rid] + computed = rid not in receives and not request.context_remaining_length + if computed and mgr.try_allocate_generation(request): + request.add_new_token(7, 0) + batch = ScheduledRequests() + batch.append_generation_request(request) + mgr.update_resources(batch) + done["generated"] += 1 + elif action < 0.68 and live: + rid = rng.choice(list(live)) + request = live[rid] + kv = kit.kv(mgr, request) + # A lender refuses the suspension of a cache whose blocks a read holds. + if kv.is_active and rid not in held: + mgr._stream.synchronize() + mgr.suspend_request(request) + # Others push the suspended pages to host, so the resume can move them. + for orid, other in live.items(): + computed = orid not in receives and not other.context_remaining_length + if other is not request and computed and kit.kv(mgr, other).is_active: + mgr.try_allocate_generation(other) + if mgr.resume_request(request): + done["resumed"] += 1 + done["parked_resumed"] += rows_of[rid] is None + elif action < 0.76 and live: + rid = rng.choice(list(live)) + request = live[rid] + if not request.context_remaining_length and rows_of[rid] is not None: + mgr.release_index_slot(rid) + rows_of[rid] = None + done["released"] += 1 + parked = lender._parked.get(rid) + assert parked is not None, f"request {rid}: no row parked at its release" + done["parked_reused"] += id(parked) in parked_seen + parked_seen.add(id(parked)) + assert all(parked is not row for row, _ in closing), "an open cache's row" + elif action < 0.82 and held: + lease = held.pop(rng.choice(list(held))) + lease.release() + elif live: + rid = rng.choice(list(live)) + request = live[rid] + computed = rid not in receives and not request.context_remaining_length + kv = kit.kv(mgr, request) + if computed and kv.is_active and rid not in held and rng.random() < 0.4: + lease = lender.lend_read(request, 0, request.prompt_len, **kit.DEVICE_ALL) + if lease.failure is None: + kit.ready_view(lease, mgr) + held[rid] = lease + else: + lease.release() + parked = lender._parked.get(rid) + mgr._stream.synchronize() + mgr.free_resources(live.pop(rid)) + if rid in held: + done["deferred"] += 1 + assert not kit.closed(kv), "a free closed a cache whose blocks a read holds" + if parked is not None: + closing.append((parked, kv)) + rows_of.pop(rid) + receives.discard(rid) + closing = [(row, kv) for row, kv in closing if not kit.closed(kv)] + mgr._stream.synchronize() + for rid in live: + if not kit.kv(mgr, live[rid]).is_active: + continue + if rows_of[rid] is None: + assert row_matches_cache(mgr, lender, rid) is None, f"request {rid}: a row" + assert parked_matches_cache(mgr, lender, rid), f"request {rid}: parked row" + done["parked_compared"] += 1 + else: + assert row_matches_cache(mgr, lender, rid), f"request {rid}: its row differs" + done["compared"] += 1 + mgr._stream.synchronize() + for lease in held.values(): + lease.release() + for request in live.values(): + mgr.free_resources(request) + counts = str(done) + assert done["compared"] >= 500 and done["released"] >= 5 and done["reused"] >= 5, counts + assert done["resumed"] >= 5 and done["generated"] >= 20, counts + assert done["parked_compared"] >= 50 and done["parked_resumed"] >= 2, counts + assert done["deferred"] >= 3 and done["parked_reused"] >= 2, counts + + +def manager_of(kind, real_manager, hybrid_manager, deepseek_v4_manager): + """A manager of ``kind`` and the prompt length its tests read.""" + if kind == "one_group": + return real_manager(), 3 * TPB + 1 + if kind == "two_groups": + # Windows longer than any prompt here give two layer groups that keep every block. + return real_manager(windows=[6 * TPB, 7 * TPB], max_seq_len=8 * TPB), 3 * TPB + 1 + if kind == "hybrid": + return hybrid_manager(), 3 * TPB + 1 + return deepseek_v4_manager(), 3 * 128 + 5 + + +@gpu +@pytest.mark.parametrize( + "kind", + ["one_group", "two_groups", "hybrid", pytest.param("deepseek_v4", marks=skip_pre_blackwell)], +) +def test_a_read_after_the_slot_s_release_lends_the_parked_row_s_pages( + kit, real_manager, hybrid_manager, deepseek_v4_manager, cache_page_reads, kind +): + """A released slot's parked row keeps its pages past the slot's next request, and lends them.""" + context, length = manager_of(kind, real_manager, hybrid_manager, deepseek_v4_manager) + with context as mgr: + lender = attach(mgr, staging_tokens=128 if kind == "deepseek_v4" else 3 * TPB) + if kind != "one_group": + assert len(mgr.impl.layer_grouping) > 1, "inconclusive: one layer group" + publish = kit.published if kind != "hybrid" else partial(computed_request, kit) + source = publish(mgr, 7, list(range(1000, 1000 + length))) + slot_row = _manager.index_row(mgr, 7) + mgr._stream.synchronize() + mgr.release_index_slot(7) + assert parked_matches_cache(mgr, lender, 7), "the parked row differs from the cache" + nxt = publish(mgr, 8, list(range(5000, 5000 + length))) + assert _manager.index_row(mgr, 8) == slot_row, "inconclusive: the slot's row went unused" + assert parked_matches_cache(mgr, lender, 7), "the slot's next request changed the row" + read_every_row(kit, mgr, lender, source, length, cache_page_reads).release() + mgr._stream.synchronize() + for request in (source, nxt): + mgr.free_resources(request) + + +@gpu +def test_a_freed_request_s_row_goes_back_to_use_only_once_its_cache_closes( + kit, real_manager, cache_page_reads +): + """A row parked for a cache whose close a read defers goes back to use only once it closes.""" + prompts = {rid: list(range(rid * 1000, rid * 1000 + 2 * TPB + 1)) for rid in (1, 2, 3)} + with real_manager() as mgr: + lender = attach(mgr) + first = kit.published(mgr, 1, prompts[1]) + mgr.release_index_slot(1) + row = lender._parked[1] + held = read_every_row(kit, mgr, lender, first, len(prompts[1]), cache_page_reads) + closing = kit.kv(mgr, first) + mgr._stream.synchronize() + mgr.free_resources(first) + assert not kit.closed(closing), "inconclusive: the free closed the cache" + second = kit.published(mgr, 2, prompts[2]) + mgr.release_index_slot(2) + assert lender._parked[2] is not row, "a row went back to use before its cache closed" + held.release() + assert kit.closed(closing), "inconclusive: the last unlock left the cache open" + third = kit.published(mgr, 3, prompts[3]) + mgr.release_index_slot(3) + assert lender._parked[3] is row, "inconclusive: the closed cache's row was not reused" + for request in (second, third): + rid = request.py_request_id + assert parked_matches_cache(mgr, lender, rid), rid + read_every_row(kit, mgr, lender, request, len(prompts[rid]), cache_page_reads).release() + mgr._stream.synchronize() + for request in (second, third): + mgr.free_resources(request) + + +@gpu +def test_a_row_handed_out_again_gives_no_page_until_its_cache_binds_it( + kit, real_manager, monkeypatch +): + """A row handed out again holds no page; unbound, a read lends the cache's own pages.""" + with real_manager() as mgr: + lender = attach(mgr) + first = kit.published(mgr, 1, list(range(1000, 1000 + 2 * TPB + 1))) + mgr.release_index_slot(1) + row = lender._parked[1] + mgr._stream.synchronize() + mgr.suspend_request(first) + assert (row.pages[0, :2] >= 0).all(), "inconclusive: the suspension left no page" + mgr.free_resources(first) + second = kit.published(mgr, 2, list(range(5000, 5000 + 2 * TPB + 1))) + give = _lender_hooks.parking_row + + def unbound(manager, request_id, kv_cache): + give(manager, request_id, kv_cache) + return None + + monkeypatch.setattr(_lender_hooks, "parking_row", unbound) + mgr.release_index_slot(2) + assert lender._parked[2] is row, "inconclusive: the closed cache's row was not reused" + assert (row.pages == -1).all(), "a row went back to use holding pages" + lease = lender.lend_read(second, 0, 2 * TPB + 1, **kit.DEVICE_ALL) + assert_lent_pages_are_the_cache_s(kit, mgr, lender, second, kit.ready_view(lease, mgr)) + lease.release() + mgr._stream.synchronize() + mgr.free_resources(second) + + +@gpu +def test_a_parked_row_bound_to_another_cache_lends_nothing(kit, real_manager, cache_page_reads): + """A row parked for another cache than the request's lends no page: the read asks the cache.""" + with real_manager() as mgr: + lender = attach(mgr) + prompt = list(range(1000, 1000 + 2 * TPB + 1)) + source = kit.published(mgr, 1, prompt) + mgr.release_index_slot(1) + lender._parked[1].kv = object() + before = cache_page_reads[0] + lease = lender.lend_read(source, 0, len(prompt), **kit.DEVICE_ALL) + view = kit.ready_view(lease, mgr) + assert cache_page_reads[0] > before, "the read lent a row bound to another cache" + assert_lent_pages_are_the_cache_s(kit, mgr, lender, source, view) + lease.release() + lender._parked[1].kv = kit.kv(mgr, source) + mgr._stream.synchronize() + mgr.free_resources(source) + + +def state_group(mgr) -> int: + (lg,) = [lg for lg, flag in enumerate(derive_layout(mgr).recurrent) if flag] + return lg + + +@pytest.fixture +def cache_reads(monkeypatch): + """How many times the lender asked a cache for its state slot.""" + reads = [0] + ask = _manager.state_slot + + def counted(kv, layer_group): + reads[0] += 1 + return ask(kv, layer_group) + + monkeypatch.setattr(_manager, "state_slot", counted) + return reads + + +def slot_of(kit, mgr, lender, request, cache_reads): + """The state slot a lend of ``request`` reads, checked, and which of record or cache gave it.""" + lg = state_group(mgr) + kv = kit.kv(mgr, request) + before = cache_reads[0] + slot = lender._state_slot(kv, request.py_request_id, lg) + assert slot == kv.get_ssm_block_base_index(lg), (request.py_request_id, slot) + return slot, ("cache" if cache_reads[0] > before else "record") + + +def lent_slot(kit, lender, mgr, request, kind): + """The state row's slot of an in-place lease of every row of ``request``'s prompt.""" + lend = lender.lend_read if kind == "read" else lender.lend_write + lease = lend(request, 0, PROMPT, **kit.DEVICE_ALL) + assert lease.failure is None and lease._first_poll_failure is None, lease.failure + view = lease.poll() + assert view is not None, lease.failure + (row,) = [run for run in view.groups if run.layer_group_index == state_group(mgr)] + return lease, int(row.slots[0]) + + +@gpu +def test_the_record_answers_after_each_preparation_of_a_chunked_prefill( + kit, hybrid_manager, cache_reads +): + """The cache answers before a request's first preparation, its record at every chunk after.""" + with hybrid_manager() as mgr: + lender = attach(mgr) + first, second = kit.admitted(mgr, 7, TOKENS), kit.admitted(mgr, 8, TOKENS) + assert slot_of(kit, mgr, lender, first, cache_reads)[1] == "cache" + sources = [] + for tokens in (TPB, TPB, PROMPT - 2 * TPB): + chunk(mgr, first, tokens) + sources.append(slot_of(kit, mgr, lender, first, cache_reads)[1]) + if second.context_remaining_length: + chunk(mgr, second, second.context_remaining_length) + sources.append(slot_of(kit, mgr, lender, second, cache_reads)[1]) + assert set(sources) == {"record"}, sources + mine = slot_of(kit, mgr, lender, first, cache_reads)[0] + theirs = slot_of(kit, mgr, lender, second, cache_reads)[0] + assert mine != theirs, "inconclusive: both requests hold one slot" + reads = cache_reads[0] + lease, slot = lent_slot(kit, lender, mgr, first, "read") + assert slot == mine and cache_reads[0] == reads, "the read asked the cache" + lease.release() + mgr._stream.synchronize() + for request in (first, second): + mgr.free_resources(request) + + +@gpu +def test_a_resume_leaves_the_record_unanswered_until_the_cache_confirms_it( + kit, hybrid_manager, cache_reads +): + """After a resume the cache answers until it confirms the record; a stale one is never lent.""" + with hybrid_manager() as mgr: + lender = attach(mgr) + request = kit.admitted(mgr, 7, TOKENS) + chunk(mgr, request, PROMPT) + slot, source = slot_of(kit, mgr, lender, request, cache_reads) + assert source == "record" + mgr.suspend_request(request) + assert mgr.resume_request(request) + assert slot_of(kit, mgr, lender, request, cache_reads)[1] == "cache" + assert slot_of(kit, mgr, lender, request, cache_reads)[1] == "record", "never confirmed" + mgr.suspend_request(request) + assert mgr.resume_request(request) + record = mgr._request_id_to_state_index + record[request.py_request_id] = slot + 1 # the slot before a move the record missed + for _ in range(2): + assert slot_of(kit, mgr, lender, request, cache_reads)[1] == "cache" + lease, lent = lent_slot(kit, lender, mgr, request, "read") + assert lent == slot, "a lease lent the record's stale slot" + lease.release() + prepare(mgr, request) + assert record[request.py_request_id] == slot + assert slot_of(kit, mgr, lender, request, cache_reads)[1] == "cache" + assert slot_of(kit, mgr, lender, request, cache_reads)[1] == "record" + mgr.suspend_request(request) + assert mgr.resume_request(request) + assert request.py_request_id in lender._state_unverified + mgr._stream.synchronize() + mgr.free_resources(request) + assert request.py_request_id not in lender._state_unverified, "the free kept its mark" + + +@gpu +def test_a_freed_slot_serves_its_next_request_from_that_request_s_record( + kit, hybrid_manager, cache_reads +): + """The sole freed slot has no new record until its next request's preparation.""" + with hybrid_manager() as mgr: + lender = attach(mgr) + first = kit.admitted(mgr, 7, TOKENS) + chunk(mgr, first, PROMPT) + slot = slot_of(kit, mgr, lender, first, cache_reads)[0] + second = kit.admitted(mgr, 8, list(range(5000, 5000 + PROMPT))) + chunk(mgr, second, PROMPT) + assert slot_of(kit, mgr, lender, second, cache_reads)[0] != slot + mgr._stream.synchronize() + mgr.free_resources(first) + assert first.py_request_id not in mgr._request_id_to_state_index + nxt = kit.admitted(mgr, 9, list(range(9000, 9000 + PROMPT))) + assert slot_of(kit, mgr, lender, nxt, cache_reads)[1] == "cache" + chunk(mgr, nxt, PROMPT) + again, source = slot_of(kit, mgr, lender, nxt, cache_reads) + assert again == slot, "the sole freed state slot was not reused" + assert source == "record" + lease, lent = lent_slot(kit, lender, mgr, nxt, "read") + assert lent == again + lease.release() + mgr._stream.synchronize() + for request in (second, nxt): + mgr.free_resources(request) + + +@gpu +def test_a_receive_target_lends_the_cache_s_slot_until_its_preparation_then_the_record_s( + kit, hybrid_manager, cache_reads +): + """A receive target's write lends the cache's slot before its preparation, its record after.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestType + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + with hybrid_manager() as mgr: + lender = attach(mgr) + target = kit.make_request( + 7, TOKENS, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + ) + assert mgr.prepare_disagg_gen_init(target) + mgr._stream.synchronize() + assert slot_of(kit, mgr, lender, target, cache_reads)[1] == "cache" + batch = ScheduledRequests() + batch.context_requests_last_chunk = [target] + mgr.prepare_resources(batch) + mgr._stream.synchronize() + slot, source = slot_of(kit, mgr, lender, target, cache_reads) + assert source == "record" + reads = cache_reads[0] + lease, lent = lent_slot(kit, lender, mgr, target, "write") + assert lent == slot and cache_reads[0] == reads, "the write asked the cache" + bounds = kit.return_write(lender, lease, ALL_ARRIVED) + assert bounds is not None and bounds.max_position == PROMPT, bounds + mgr._stream.synchronize() + mgr.free_resources(target) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py index 972fb1ee126f..1c551a579867 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py @@ -19,7 +19,12 @@ import dataclasses import enum import hashlib +import inspect import json +import subprocess +import sys +import types +import weakref from types import SimpleNamespace from typing import Dict, List, Optional, Sequence @@ -28,21 +33,37 @@ import torch from utils.util import skip_pre_blackwell, skip_pre_hopper -from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout -from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import Identity -from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import ( - LAYOUT_ID_VERSION, +from tensorrt_llm._torch.pyexecutor.kv_cache import page_layout, page_table, sharing +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import ( + Axis, BufferDesc, - BufferGeometry, - BufferMapper, LayerGroupDesc, - Layout, + PageLayout, + RankTopology, ShardDesc, - _Heads, + Tensor, + head_major, + replicated, + sectioned, + token_major, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.page_table import attention_shard +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + MemoryKind, + MemoryRegion, + StagingOptions, + _layout, + _lender, + _manager, + attach_lender, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import Identity +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import ( + LAYOUT_ID_VERSION, canonical_bytes, canonical_order, derive_layout, - global_layer_ids, layout_document, layout_id, ) @@ -68,22 +89,23 @@ def make_layout( shard=ShardDesc(), group_ids=(0, 1), slot_bytes=(100, 40), -) -> Layout: + **change, +) -> PageLayout: g0, g1 = group_ids pool_groups = {g0: tuple(slot_bytes), g1: (64,)} layer_groups = ( LayerGroupDesc( kind="attention", - window=windows[0], - sink_blocks=0, + kept_window=windows[0], + num_sink_blocks=0, pool_group=g0, layers=layers[0], shard=shard, ), LayerGroupDesc( kind="attention", - window=windows[1], - sink_blocks=0, + kept_window=windows[1], + num_sink_blocks=0, pool_group=g1, layers=layers[1], shard=shard, @@ -92,12 +114,13 @@ def make_layout( buffers = [] for layer in layers[0]: buffers += [ - BufferDesc(layer, "key", (g0, 0), 50 * (layer // 2), 50), - BufferDesc(layer, "scale", (g0, 1), 20 * (layer // 2), 20), + BufferDesc(layer, "key", g0, 0, 50 * (layer // 2), 50), + BufferDesc(layer, "scale", g0, 1, 20 * (layer // 2), 20), ] for layer in layers[1]: - buffers.append(BufferDesc(layer, "key", (g1, 0), 0, 64)) - return Layout(tokens_per_block, pool_groups, layer_groups, tuple(buffers)) + buffers.append(BufferDesc(layer, "key", g1, 0, 0, 64)) + layout = PageLayout(tokens_per_block, pool_groups, layer_groups, tuple(buffers)) + return dataclasses.replace(layout, **change) def _replace_buffers(layout, **change): @@ -125,10 +148,11 @@ def _first_group(layout, **change): ), id="layer_group_order", ), - pytest.param(lambda: _first_group(make_layout(), sink_blocks=3), id="sinks"), + pytest.param(lambda: _first_group(make_layout(), num_sink_blocks=3), id="sinks"), pytest.param(lambda: make_layout(windows=(None, 128)), id="window"), pytest.param(lambda: make_layout(windows=(64, 64)), id="full_to_window"), pytest.param(lambda: make_layout(shard=ShardDesc(2, 1)), id="shard"), + pytest.param(lambda: _first_group(make_layout(), model_layers=(0, 0)), id="model_layers"), pytest.param( lambda: dataclasses.replace(make_layout(), buffers=make_layout().buffers[::-1]), id="buffer_order", @@ -137,8 +161,9 @@ def _first_group(layout, **change): ) def test_layout_id_leaves_out_what_does_not_decide_how_bytes_read(variant): """Local numbering and order, sinks and windows (which blocks exist, not how their bytes read), - and the shard position, which rides in the name so every rank computes the same ``layout_id``. - Buffers are described by global layer, so their listing order does not matter either.""" + the shard position, which rides in the name so every rank computes the same ``layout_id``, and + the model layer each layer belongs to. Buffers are described by global layer, so their listing + order does not matter either.""" assert layout_id(variant()) == layout_id(make_layout()) @@ -149,17 +174,18 @@ def test_layout_id_leaves_out_what_does_not_decide_how_bytes_read(variant): lambda: make_layout(tokens_per_block=64), lambda: make_layout(layers=((0, 1), (2,))), lambda: make_layout(slot_bytes=(104, 40)), - lambda: _replace_buffers(make_layout(), mapper=BufferMapper.REPLICATED), - lambda: _replace_buffers(make_layout(), geometry=BufferGeometry(bytes_per_head=10)), - lambda: _replace_buffers(make_layout(), geometry=BufferGeometry(section_bytes=(15, 5))), - lambda: _replace_buffers(make_layout(), transfer=False), + lambda: _replace_buffers(make_layout(), tensors=head_major(20, 2, 4)), + lambda: _replace_buffers(make_layout(), tensors=head_major(20, 4, 8)), + lambda: _replace_buffers(make_layout(), tensors=token_major(20, 2, 2, 4)), + lambda: _replace_buffers(make_layout(), tensors=sectioned((15, 5), 2)), + lambda: _replace_buffers(make_layout(), transferred=False), lambda: _replace_buffers(make_layout(), expansion=2), lambda: _replace_buffers(make_layout(), role="scale2"), lambda: _first_group(make_layout(), kind="state"), lambda: dataclasses.replace( make_layout(), buffers=tuple( - dataclasses.replace(b, offset=b.offset + 10) if b.role == "scale" else b + dataclasses.replace(b, slot_offset=b.slot_offset + 10) if b.role == "scale" else b for b in make_layout().buffers ), ), @@ -168,8 +194,9 @@ def test_layout_id_leaves_out_what_does_not_decide_how_bytes_read(variant): "tokens_per_block", "layer_membership", "slot_width", - "mapper", - "head_geometry", + "split_by_head", + "head_size", + "token_major", "sections", "transfer", "expansion", @@ -193,7 +220,7 @@ def test_canonical_order_sorts_groups_by_smallest_layer(): @pytest.mark.cpu_only -def test_layout_id_is_the_hash_of_the_documented_version_2_encoding(): +def test_layout_id_is_the_hash_of_the_documented_version_4_encoding(): """The canonical encoding, spelled out here byte for byte so a change to it cannot pass unnoticed.""" layout = make_layout(layers=((0,), (1,))) @@ -203,9 +230,9 @@ def test_layout_id_is_the_hash_of_the_documented_version_2_encoding(): '{"layer":0,"offset":100,"role":"scale","size":20},' '{"layer":1,"offset":0,"role":"key","size":64}],' '"groups":[{"layers":[0],"slot_bytes":[100,40]},{"layers":[1],"slot_bytes":[64]}],' - '"tokens_per_block":32,"v":2}' + '"tokens_per_block":32,"v":4}' ).encode() - assert LAYOUT_ID_VERSION == 2 + assert LAYOUT_ID_VERSION == 4 assert canonical_bytes(layout_document(layout)) == expected assert layout_id(layout) == hashlib.sha256(expected).digest() @@ -216,16 +243,19 @@ def test_named_keys_at_their_defaults_are_left_out_and_others_are_written(): plain = layout_document(layout)["buffers"][1] assert set(plain) == {"layer", "role", "offset", "size"} changed = layout_document( - _replace_buffers( - layout, - mapper=BufferMapper.SECTIONED, - geometry=BufferGeometry(section_bytes=(15, 5)), - transfer=False, - expansion=4, - ) + _replace_buffers(layout, tensors=sectioned((15, 5), 2), transferred=False, expansion=4) )["buffers"][1] - assert changed["mapper"] == "SECTIONED" - assert changed["geometry"] == {"section_bytes": [15, 5]} + assert changed["tensors"] == [ + { + "axes": [{"name": "byte", "extent": 30, "stride": 1, "split": ["tp"], "held": 15}], + "name": "s0", + }, + { + "axes": [{"name": "byte", "extent": 10, "stride": 1, "split": ["tp"], "held": 5}], + "name": "s1", + "offset": 15, + }, + ] assert changed["transfer"] is False assert changed["expansion"] == 4 assert layout_document(_first_group(layout, kind="state"))["groups"][0]["kind"] == "state" @@ -233,10 +263,9 @@ def test_named_keys_at_their_defaults_are_left_out_and_others_are_written(): @pytest.mark.cpu_only -def test_enums_encode_by_name_and_the_hnd_alias_is_indexed(): - hnd = _replace_buffers(make_layout(), mapper=MapperKind.HND) - indexed = _replace_buffers(make_layout(), mapper=BufferMapper.INDEXED) - assert layout_id(hnd) == layout_id(indexed) == layout_id(make_layout()) +def test_enums_encode_by_name_and_a_whole_buffer_is_its_default_tensor(): + whole = _replace_buffers(make_layout(), tensors=replicated(20)) + assert whole == make_layout() and layout_id(whole) == layout_id(make_layout()) class Color(enum.IntEnum): RED = 3 @@ -258,7 +287,7 @@ def test_a_document_refuses_what_the_layout_does_not_hold(): layout = make_layout() with pytest.raises(ValueError, match="maps to pool group 9"): layout_document(_first_group(layout, pool_group=9)) - outside = layout.buffers + (BufferDesc(7, "key", (9, 0), 0, 8),) + outside = layout.buffers + (BufferDesc(7, "key", 9, 0, 0, 8),) with pytest.raises(ValueError, match="not in the layout"): layout_document(dataclasses.replace(layout, buffers=outside)) twice = layout.buffers + (layout.buffers[0],) @@ -270,19 +299,35 @@ def test_a_document_refuses_what_the_layout_does_not_hold(): @pytest.mark.cpu_only -def test_a_geometry_that_does_not_fit_its_buffer_is_refused(): - with pytest.raises(ValueError, match="sum"): - BufferDesc(0, "conv", (0, 0), 0, 30, geometry=BufferGeometry(section_bytes=(10, 10))) - with pytest.raises(ValueError, match="heads"): - BufferDesc(0, "k", (0, 0), 0, 30, geometry=BufferGeometry(bytes_per_head=8)) - with pytest.raises(ValueError, match="heads"): - BufferDesc(0, "k", (0, 0), 0, 32, geometry=BufferGeometry(bytes_per_head=8, num_heads=2)) +def test_tensors_that_do_not_fit_their_buffer_are_refused(): + with pytest.raises(ValueError, match="reaches byte 40"): + BufferDesc(0, "conv", 0, 0, 0, 30, tensors=sectioned((20, 20), 2)) + with pytest.raises(ValueError, match="reaches byte 32"): + BufferDesc(0, "k", 0, 0, 0, 30, tensors=head_major(32, 4, 8)) + with pytest.raises(ValueError, match="whole number"): + head_major(30, 8, 8) + with pytest.raises(ValueError, match="repeat"): + BufferDesc(0, "k", 0, 0, 0, 30, tensors=replicated(30) + replicated(30)) + + +@pytest.mark.cpu_only +def test_an_axis_says_how_its_split_holds(): + with pytest.raises(ValueError, match="says how many indices"): + Axis("head", 8, 4, ("tp",)) + with pytest.raises(ValueError, match="says how many indices"): + Axis("head", 8, 4, held=2) + with pytest.raises(ValueError, match="not split"): + Axis("row", 8, 4, cyclic=True) + with pytest.raises(ValueError, match="splits over"): + Axis("head", 8, 4, ("pp",), 2) + with pytest.raises(ValueError, match="splits two axes"): + Tensor((Axis("a", 4, 4, ("tp",), 2), Axis("b", 4, 1, ("tp",), 2))) @pytest.mark.cpu_only def test_an_expansion_must_divide_the_block(): with pytest.raises(ValueError, match="positive"): - BufferDesc(0, "k", (0, 0), 0, 32, expansion=0) + BufferDesc(0, "k", 0, 0, 0, 32, expansion=0) with pytest.raises(ValueError, match="does not divide"): _replace_buffers(make_layout(), expansion=3) @@ -294,26 +339,18 @@ def test_a_shard_is_a_share_of_its_count(shard): ShardDesc(*shard) -@pytest.mark.cpu_only -@pytest.mark.parametrize( - "sections", - [(np.int64(4), np.int64(4)), np.array([4, 4]), [4, 4], (4, 4)], - ids=["numpy-ints", "numpy-array", "list", "tuple"], -) -def test_a_validated_geometry_holds_its_sections_as_a_tuple_of_ints(sections): - checked = BufferGeometry(section_bytes=sections, bytes_per_head=2).validate() - assert checked == BufferGeometry(section_bytes=(4, 4), bytes_per_head=2) - assert type(checked.section_bytes) is tuple - assert all(type(s) is int for s in checked.section_bytes) - - -@pytest.mark.cpu_only -@pytest.mark.parametrize( - "fields", [{"section_bytes": (4, 0)}, {"bytes_per_head": 0}, {"num_heads": -1}] -) -def test_validation_refuses_sizes_that_are_not_positive(fields): - with pytest.raises(ValueError, match="positive"): - BufferGeometry(**fields).validate() +def described(b): + """(kind, heads on this rank, bytes per head) a buffer's tensors amount to: "REPLICATED" alike + on every rank, "SECTIONED" split by byte, else "NHD" or "INDEXED" by where its head axis lies; + heads and bytes are None without a head axis.""" + if not b.split: + return "REPLICATED", None, None + (tensor, *rest) = b.tensors + axes = {a.name: a for a in tensor.axes} + if rest or "head" not in axes: + return "SECTIONED", None, None + kind = "NHD" if tensor.axes[0].name == "token" else "INDEXED" + return kind, axes["head"].held, b.size // axes["head"].held # -- a stand-in manager -------------------------------------------------------------------------- @@ -332,7 +369,29 @@ def coalesced(size, members): class Manager: """One layer group per entry of ``groups``: ``(layer config class, [internal layer ids], [(role, size, override)])``; group ``i`` draws from pool group ``pool_group_of[i]`` (its own by - default), one pool per role, and has window ``windows[i]`` and ``sinks[i]`` sink tokens.""" + default), one pool per role, and has window ``windows[i]`` and ``sinks[i]`` sink tokens. It + builds its page table with ``KVCacheManagerV2``'s methods, from the state that manager keeps.""" + + get_page_table = KVCacheManagerV2.get_page_table + _page_head_split = KVCacheManagerV2._page_head_split + _page_rebuilds_untransferred = KVCacheManagerV2._page_rebuilds_untransferred + _page_layer_configs = KVCacheManagerV2._page_layer_configs + _page_layer_ids = KVCacheManagerV2._page_layer_ids + _page_kv_heads = KVCacheManagerV2._page_kv_heads + _page_dtype = KVCacheManagerV2._page_dtype + _page_quant_mode = KVCacheManagerV2._page_quant_mode + + def _virtual_layers(self): + """The model layers and attention types a test names in ``_layer_attn_to_layer_id``.""" + return page_table.virtual_layers(getattr(self, "_layer_attn_to_layer_id", None)) + + _model_has_attention_window = False + _model_has_sparse_buffers = False + _model_attention_window_pattern = () + _has_cp_helix = False + _helix_cp_rank = 0 + _helix_cp_size = 1 + dtype = None def __init__( self, @@ -343,8 +402,8 @@ def __init__( attention_dp=False, num_kv_heads=4, pp_layers: Optional[Sequence[int]] = None, - mappers: Optional[Dict[str, MapperKind]] = None, - layouts: Optional[Dict[str, RoleLayout]] = None, + mappers: Optional[Dict[str, str]] = None, + layouts: Optional[Dict[str, Dict[str, object]]] = None, ignored=(), pool_group_of: Optional[Sequence[int]] = None, windows: Optional[Sequence[Optional[int]]] = None, @@ -359,10 +418,11 @@ def __init__( self.mapping = SimpleNamespace( tp_size=tp_size, tp_rank=tp_rank, enable_attention_dp=attention_dp ) + self.num_layers = len(internal) self.num_kv_heads = num_kv_heads self.num_kv_heads_per_layer = [heads_per_rank] * len(internal) self.pp_layers = list(pp_layers if pp_layers is not None else internal) - self.mappers = {"all": MapperKind.INDEXED, **(mappers or {})} + self.mappers = {"all": "head_major", **(mappers or {})} self.layouts = layouts or {} self.ignored = frozenset(ignored) layers: List[object] = [] @@ -403,21 +463,43 @@ def __init__( ], slot_desc=SimpleNamespace(variants=[variant]), ) + self.windows = windows + self.kv_cache_manager_py_config = SimpleNamespace(layers=layers) self.impl = SimpleNamespace( layer_grouping=[list(lids) for _, lids, _ in groups], pool_group_descs=list(descs.values()), get_life_cycle_pool_group_indices=lambda: list(pool_group_of), - init_config=SimpleNamespace(layers=layers), + init_config=self.kv_cache_manager_py_config, ) - def get_disagg_role_mapper_kinds(self): - return self.mappers - - def get_disagg_role_layouts(self): - return self.layouts + def _attention_read_window(self, layer_group): + """Attention reads the whole window it keeps.""" + return self.windows[layer_group] + + def _page_buffer_tensors(self, layer_id, role, size, tokens, state): + """Each role spread as ``mappers`` names it ("all" for the roles it does not list): + "replicated", "sectioned" by ``layouts``' ``section_bytes``, or split by head, + "head_major" or "token_major", ``layouts``' ``bytes_per_head`` or the layer's KV heads + apart.""" + kind = self.mappers.get(role, self.mappers["all"]) + layout = self.layouts.get(role, {}) + if kind == "replicated": + return replicated(size) + if kind == "sectioned": + split = ("tp", "cp") if state else ("tp",) + ranks, _ = page_layout.grid_place(page_table.rank_topology(self.mapping), split) + return sectioned(layout["section_bytes"], ranks, split) + return self._page_head_split( + layer_id, + size, + tokens, + state, + head_bytes=layout.get("bytes_per_head"), + token_major_=kind == "token_major", + ) - def get_disagg_ignored_roles(self): - return self.ignored + def _page_untransferred_roles(self): + return frozenset(self.ignored) def kv_group(lids=(0, 1), heads_per_rank=1, extra=()): @@ -441,10 +523,10 @@ def recurrent_state(): [(SsmLayerConfig, [0, 1], [("conv_state", 96, None), ("ssm_state", 128, None)])], tp_size=2, tp_rank=1, - mappers={"conv_state": MapperKind.SECTIONED}, + mappers={"conv_state": "sectioned"}, layouts={ - "conv_state": RoleLayout(section_bytes=(32, 32, 32)), - "ssm_state": RoleLayout(bytes_per_head=32), + "conv_state": dict(section_bytes=(32, 32, 32)), + "ssm_state": dict(bytes_per_head=32), }, ) @@ -455,12 +537,12 @@ def recurrent_state(): def check_token_major_kv_has_its_own_layout_id(): """A manager declaring token-major K/V (NHD) and one declaring the default head-major K/V (HND) lay a block's bytes out differently, so they never share a ``layout_id``.""" - head_major = derive_layout(attention(1, 0, 4)).layout - token_major = derive_layout(attention(1, 0, 4, mappers={"all": MapperKind.NHD})).layout + head_major = derive_layout(attention(1, 0, 4)).page_layout + token_major = derive_layout(attention(1, 0, 4, mappers={"all": "token_major"})).page_layout assert layout_id(head_major) != layout_id(token_major), TOKEN_MAJOR_SHARES - assert {b.mapper for b in token_major.buffers} == {BufferMapper.NHD} - assert {b.mapper for b in head_major.buffers} == {BufferMapper.INDEXED} - hnd = derive_layout(attention(1, 0, 4, mappers={"all": MapperKind.HND})).layout + assert {described(b)[0] for b in token_major.buffers} == {"NHD"} + assert {described(b)[0] for b in head_major.buffers} == {"INDEXED"} + hnd = derive_layout(attention(1, 0, 4, mappers={"all": "head_major"})).page_layout assert layout_id(hnd) == layout_id(head_major), "HND is the default head-major layout" @@ -492,12 +574,12 @@ def test_kv_heads_name_the_share_each_rank_holds(case): derive_layout(attention(tp_size, r, num_kv_heads, attention_dp=attention_dp)) for r in range(tp_size) ] - assert [geo.layout.layer_groups[0].shard for geo in geometries] == shares + assert [geo.page_layout.layer_groups[0].shard for geo in geometries] == shares assert [geo.shards for geo in geometries] == [((s.count, s.index),) for s in shares] - assert len({layout_id(geo.layout) for geo in geometries}) == 1 + assert len({layout_id(geo.page_layout) for geo in geometries}) == 1 keys = np.full((1, 32), ord("k"), dtype=np.uint8) names = [ - Identity(b"head-shares", layout_id(geo.layout), geo.layers, geo.shards) + Identity(b"head-shares", layout_id(geo.page_layout), geo.layers, geo.shards) .object_keys(0, keys)[0] .tobytes() for geo in geometries @@ -507,15 +589,6 @@ def test_kv_heads_name_the_share_each_rank_holds(case): assert (names[rank] == names[other]) == (share == other_share) -def heads_on(rank, tp_size, total, per_rank): - """``_Heads`` of a rank that holds ``per_rank`` of ``total`` KV heads over ``tp_size`` ranks.""" - mapping = SimpleNamespace(tp_size=tp_size, tp_rank=rank, enable_attention_dp=False) - manager = SimpleNamespace( - mapping=mapping, num_kv_heads_per_layer=[per_rank], num_kv_heads=total, pp_layers=[0] - ) - return _Heads(manager) - - # (KV heads T, TP size, heads per rank) -> each rank's share. Rank r holds its heads from # r * T // tp; where any rank's range starts inside a share, every rank is a share of its own. UNALIGNED_HEADS = { @@ -530,7 +603,7 @@ def heads_on(rank, tp_size, total, per_rank): @pytest.mark.cpu_only def test_unaligned_head_ranges_give_every_rank_its_own_share(): for (total, tp_size, per_rank), shares in UNALIGNED_HEADS.items(): - got = [heads_on(r, tp_size, total, per_rank).attention_shard(0) for r in range(tp_size)] + got = [attention_shard(tp_size, r, per_rank, total) for r in range(tp_size)] assert got == shares, f"{total} heads over {tp_size} ranks: {got}" @@ -542,17 +615,17 @@ def test_a_group_of_replicated_buffers_is_whole_and_one_head_sharded_buffer_make tp_size=2, tp_rank=1, num_kv_heads=8, - mappers={"index_key": MapperKind.REPLICATED}, + mappers={"index_key": "replicated"}, ) - assert derive_layout(replicated_only).layout.layer_groups[0].shard == ShardDesc() + assert derive_layout(replicated_only).page_layout.layer_groups[0].shard == ShardDesc() mixed = Manager( [kv_group(heads_per_rank=4, extra=[index])], tp_size=2, tp_rank=1, num_kv_heads=8, - mappers={"index_key": MapperKind.REPLICATED}, + mappers={"index_key": "replicated"}, ) - assert derive_layout(mixed).layout.layer_groups[0].shard == ShardDesc(2, 1) + assert derive_layout(mixed).page_layout.layer_groups[0].shard == ShardDesc(2, 1) @pytest.mark.cpu_only @@ -562,12 +635,12 @@ def test_local_only_buffers_do_not_decide_the_identity(): [(AttentionLayerConfig, [0], [("index_key", 64, None), scratch])], tp_size=2, tp_rank=1, - mappers={"index_key": MapperKind.REPLICATED}, + mappers={"index_key": "replicated"}, ignored=("scratch",), ) - layout = derive_layout(manager).layout + layout = derive_layout(manager).page_layout assert layout.layer_groups[0].shard == ShardDesc() - assert {b.role: b.transfer for b in layout.buffers} == {"index_key": True, "scratch": False} + assert {b.role: b.transferred for b in layout.buffers} == {"index_key": True, "scratch": False} @pytest.mark.cpu_only @@ -580,9 +653,9 @@ def test_groups_of_one_manager_are_judged_one_by_one(): tp_size=2, tp_rank=0, num_kv_heads=8, - mappers={"index_key": MapperKind.REPLICATED}, + mappers={"index_key": "replicated"}, ) - groups = derive_layout(manager).layout.layer_groups + groups = derive_layout(manager).page_layout.layer_groups assert [g.shard for g in groups] == [ShardDesc(2, 0), ShardDesc()] @@ -617,10 +690,12 @@ def test_a_group_whose_buffers_hold_different_shares_is_each_rank_alone(case): names = set() for rank in range(tp_size): geometry = derive_layout(mixed_shares(tp_size, rank, heads)) - assert geometry.layout.layer_groups[0].shard == ShardDesc(tp_size, rank), ( + assert geometry.page_layout.layer_groups[0].shard == ShardDesc(tp_size, rank), ( f"rank {rank}: the group's share is not the rank's own" ) - identity = Identity(b"scope", layout_id(geometry.layout), geometry.layers, geometry.shards) + identity = Identity( + b"scope", layout_id(geometry.page_layout), geometry.layers, geometry.shards + ) names |= {row.tobytes() for row in identity.object_keys(0, keys)} assert len(names) == tp_size * len(keys), "two ranks name a block alike" @@ -628,22 +703,259 @@ def test_a_group_whose_buffers_hold_different_shares_is_each_rank_alone(case): @pytest.mark.cpu_only def test_recurrent_state_is_described_and_sharded_by_rank(): geometry = derive_layout(recurrent_state()) - layout = geometry.layout + layout = geometry.page_layout (group,) = layout.layer_groups assert group.kind == "state" and group.shard == ShardDesc(2, 1) assert geometry.recurrent == (True,) and geometry.windows == (None,) by_role = {b.role: b for b in layout.buffers} - assert by_role["conv_state"].mapper is BufferMapper.SECTIONED - assert by_role["conv_state"].geometry == BufferGeometry(section_bytes=(32, 32, 32)) - assert by_role["ssm_state"].geometry == BufferGeometry(bytes_per_head=32, num_heads=4) + assert by_role["conv_state"].tensors == sectioned((32, 32, 32), 2, ("tp", "cp")) + assert by_role["ssm_state"].tensors == head_major(128, 4, 8, ("tp", "cp")) + + +def tensor_ranks_per_group(topology): + """Tensor ranks in the attention data-parallel group of a rank, by its topology.""" + return topology.tp_size // topology.dp_size if topology.attention_dp else topology.tp_size + + +def state_grid(topology): + """``(ranks, place)`` of the grid recurrent state splits over, by a rank's topology: tensor + ranks, none under attention data parallelism, tensor times context ranks (context-minor) under + Helix.""" + if topology.attention_dp: + return 1, 0 + if topology.cp_size > 1: + return ( + topology.tp_size * topology.cp_size, + topology.tp_rank * topology.cp_size + topology.cp_rank, + ) + return topology.tp_size, topology.tp_rank + + +def attention_and_state(tp_size, tp_rank, attention_dp, num_kv_heads, cp_size=1, cp_rank=0): + """A rank of ``num_kv_heads`` KV heads in layers 0 and 1 and recurrent state in layers 2 and + 3; under Helix, context rank ``cp_rank`` of ``cp_size``.""" + heads = -(-num_kv_heads // (1 if attention_dp else tp_size)) + manager = Manager( + [ + kv_group(lids=(0, 1), heads_per_rank=heads), + (SsmLayerConfig, [2, 3], [("conv_state", 96, None), ("ssm_state", 128, None)]), + ], + tp_size=tp_size, + tp_rank=tp_rank, + attention_dp=attention_dp, + num_kv_heads=num_kv_heads, + mappers={"conv_state": "sectioned"}, + layouts={ + "conv_state": dict(section_bytes=(32, 32, 32)), + "ssm_state": dict(bytes_per_head=32), + }, + ) + if cp_size > 1: + manager._has_cp_helix = True + manager._helix_cp_rank, manager._helix_cp_size = cp_rank, cp_size + manager.mapping.cp_size, manager.mapping.cp_rank = cp_size, cp_rank + return manager + + +# (TP size, attention DP, Helix context ranks, KV heads) of every instance the sweep builds. +SHARD_SWEEP = [ + (tp_size, attention_dp, cp_size, num_kv_heads) + for tp_size in (1, 2, 4, 8) + for attention_dp in ((False, True) if tp_size > 1 else (False,)) + for cp_size in ((1,) if attention_dp else (1, 2, 4)) + for num_kv_heads in (1, 2, 4, 6, 8, 12, 16) +] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "tp_size, attention_dp, cp_size, num_kv_heads", + SHARD_SWEEP, + ids=[f"tp{t}{'_adp' if a else ''}_cp{c}_heads{h}" for t, a, c, h in SHARD_SWEEP], +) +def test_published_shards_are_the_shares_the_topology_gives( + tp_size, attention_dp, cp_size, num_kv_heads +): + """Every rank's published shards name what its topology gives it. + + Attention: a rank at place ``p`` of an attention data-parallel group of ``n`` tensor ranks + holds the heads from ``p * T // n``; its share's first head, ``index * T // count``, is that + head, and two ranks have equal shares exactly when they hold the same heads. State: the share + is the rank's place in the grid recurrent state splits over. + """ + layouts = [ + derive_layout( + attention_and_state(tp_size, tp_rank, attention_dp, num_kv_heads, cp_size, cp_rank) + ).page_layout + for tp_rank in range(tp_size) + for cp_rank in range(cp_size) + ] + first_heads = [] + for layout in layouts: + attention, state = layout.layer_groups + topology = layout.topology + ranks = tensor_ranks_per_group(topology) + first = (topology.tp_rank % ranks) * num_kv_heads // ranks + first_heads.append(first) + share = attention.shard + assert share.index * num_kv_heads // share.count == first, (topology, share) + count, index = state_grid(topology) + assert state.shard == (ShardDesc(count, index) if count > 1 else ShardDesc()), ( + topology, + state.shard, + ) + for one, first_one in zip(layouts, first_heads): + for other, first_other in zip(layouts, first_heads): + same_share = one.layer_groups[0].shard == other.layer_groups[0].shard + assert same_share == (first_one == first_other), (one.topology, other.topology) + + +# -- which ranks of an instance hold the same bytes --------------------------------------------- + + +def replica_instance(tp=1, pp=1, cp=1, attention_dp=False, groups_of=None): + """Every rank's layout by flat rank, ``pp * (tp * cp) + tp * cp + cp``. + + ``groups_of(pp_rank, tp_rank, cp_rank)`` gives each rank's ``(kind, layers, shard)`` groups. + """ + layouts = [] + for pp_rank in range(pp): + for tp_rank in range(tp): + for cp_rank in range(cp): + groups = tuple( + LayerGroupDesc(kind, None, 0, g, tuple(layers), shard) + for g, (kind, layers, shard) in enumerate(groups_of(pp_rank, tp_rank, cp_rank)) + ) + layouts.append( + PageLayout( + tokens_per_page=4, + pool_slot_bytes={g: (64,) for g in range(len(groups))}, + layer_groups=groups, + buffers=(), + topology=RankTopology( + tp_size=tp, + tp_rank=tp_rank, + pp_size=pp, + pp_rank=pp_rank, + cp_size=cp, + cp_rank=cp_rank, + dp_size=tp if attention_dp else 1, + dp_rank=tp_rank if attention_dp else 0, + attention_dp=attention_dp, + ), + ) + ) + return layouts + + +def two_heads_over_four_ranks(): + """Ranks 0 and 1 hold KV head 0, ranks 2 and 3 KV head 1.""" + return replica_instance( + tp=4, groups_of=lambda p, t, c: [("attention", (0, 1), ShardDesc(2, t // 2))] + ) + + +@pytest.mark.cpu_only +def test_ranks_of_equal_shards_are_replicas_and_replicated_buffers_are_on_every_rank(): + ranks = two_heads_over_four_ranks() + assert [page_layout.replica_ranks(ranks, r, 0) for r in range(4)] == [ + (0, 1), + (0, 1), + (2, 3), + (2, 3), + ] + assert page_layout.replica_ranks(ranks, 3, 0, replicated=True) == (0, 1, 2, 3) + + +@pytest.mark.cpu_only +def test_replicas_are_listed_by_ascending_flat_rank(): + """The lowest rank holding a share, which publishes it to a store, comes first.""" + ranks = two_heads_over_four_ranks() + assert [page_layout.replica_ranks(ranks, r, 0)[0] for r in range(4)] == [0, 0, 2, 2] + assert page_layout.replica_ranks(ranks, 2, 0, replicated=True)[0] == 0 + assert page_layout.replica_ranks(ranks, 3, 0) == (2, 3) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "tp_size,heads,expected", + [ + (8, (2, 4), [(0, 1), (0, 1), (2, 3), (2, 3), (4, 5), (4, 5), (6, 7), (6, 7)]), + (6, (2, 3), [(0, 1), (0, 1), (2,), (3,), (4, 5), (4, 5)]), + (8, (4, 8), [(rank,) for rank in range(8)]), + ], +) +def test_mixed_head_replication_matches_every_buffer(tp_size, heads, expected): + ranks = [ + derive_layout(mixed_shares(tp_size, rank, heads)).page_layout for rank in range(tp_size) + ] + assert [page_layout.replica_ranks(ranks, rank, 0) for rank in range(tp_size)] == expected + + +@pytest.mark.cpu_only +def test_each_attention_data_parallel_rank_holds_its_own_requests(): + ranks = replica_instance( + tp=4, attention_dp=True, groups_of=lambda p, t, c: [("attention", (0, 1), ShardDesc())] + ) + for r in range(4): + assert page_layout.replica_ranks(ranks, r, 0) == (r,) + assert page_layout.replica_ranks(ranks, r, 0, replicated=True) == (r,) + + +@pytest.mark.cpu_only +def test_pipeline_stages_hold_other_layers_and_groups_of_another_kind_other_bytes(): + def groups_of(pp_rank, tp_rank, cp_rank): + if pp_rank == 0: + return [("attention", (0, 1), ShardDesc()), ("state", (2,), ShardDesc())] + return [("attention", (3, 4), ShardDesc()), ("attention", (2,), ShardDesc())] + + ranks = replica_instance(tp=2, pp=2, groups_of=groups_of) + assert page_layout.replica_ranks(ranks, 1, 0) == (0, 1) + assert page_layout.replica_ranks(ranks, 3, 0) == (2, 3) + # Layer 2 is state on stage 0 and attention on stage 1: no rank of the other stage matches. + assert page_layout.replica_ranks(ranks, 0, 1) == (0, 1) + assert page_layout.replica_ranks(ranks, 2, 1) == (2, 3) + + +@pytest.mark.cpu_only +def test_helix_ranks_hold_their_own_attention_pages_and_split_state_by_context_rank(): + def groups_of(pp_rank, tp_rank, cp_rank): + return [("attention", (0,), ShardDesc()), ("state", (1,), ShardDesc(2, cp_rank))] + + ranks = replica_instance(cp=2, groups_of=groups_of) + assert [page_layout.replica_ranks(ranks, r, 0) for r in range(2)] == [(0,), (1,)] + assert [page_layout.replica_ranks(ranks, r, 1) for r in range(2)] == [(0,), (1,)] + # Replicated state is the same on both context ranks. + assert page_layout.replica_ranks(ranks, 1, 1, replicated=True) == (0, 1) + + +@pytest.mark.cpu_only +def test_replicas_are_every_rank_of_the_instance_holding_the_bytes(): + """Which of them a caller may choose from is the caller's to filter.""" + parameters = inspect.signature(page_layout.replica_ranks).parameters + assert list(parameters) == ["ranks", "rank", "group", "replicated"] + assert parameters["replicated"].kind is inspect.Parameter.KEYWORD_ONLY + assert parameters["replicated"].default is False @pytest.mark.cpu_only def test_head_geometry_is_described_per_buffer(): - layout = derive_layout(attention(2, 0, num_kv_heads=8)).layout + layout = derive_layout(attention(2, 0, num_kv_heads=8)).page_layout for b in layout.buffers: - assert b.mapper is BufferMapper.INDEXED - assert b.geometry == BufferGeometry(bytes_per_head=HEAD_BYTES, num_heads=4) + assert b.tensors == head_major(b.size, 4, 8) + + +@pytest.mark.cpu_only +def test_layer_configs_are_keyed_by_their_layer_id_or_position(): + layers = [ + SimpleNamespace(layer_id="7"), + SimpleNamespace(layer_id=4), + SimpleNamespace(), + SimpleNamespace(layer_id=0), + ] + actual = page_table.layer_configs_by_id(SimpleNamespace(layers=layers)) + assert list(actual) == [7, 4, 2, 0] + assert all(actual[layer_id] is layer for layer_id, layer in zip([7, 4, 2, 0], layers)) @pytest.mark.cpu_only @@ -651,10 +963,10 @@ def test_layers_are_global_ids_stable_across_pipeline_stages(): """Stage 2 of a pipeline holds model layers 10 and 11 as internal layers 0 and 1.""" manager = Manager([kv_group()], pp_layers=[10, 11]) geometry = derive_layout(manager) - assert geometry.layout.layer_groups[0].layers == (10, 11) + assert geometry.page_layout.layer_groups[0].layers == (10, 11) assert geometry.layers == ((10, 11),) - assert {b.layer for b in geometry.layout.buffers} == {10, 11} - assert global_layer_ids(manager, [1, 0]) == [11, 10] + assert {b.layer for b in geometry.page_layout.buffers} == {10, 11} + assert manager._page_layer_ids() == {0: (10, 10), 1: (11, 11)} # A manager without the map is refused, not named by its stage-local ids, which another # stage reuses for other layers. del manager.pp_layers @@ -678,46 +990,185 @@ class AttentionType(enum.Enum): (4, AttentionType.COMPRESSED): 1, (5, AttentionType.SLIDING): 2, } - assert global_layer_ids(manager, [0, 1, 2]) == [12, 13, 15] + assert manager._page_layer_ids() == {0: (12, 4), 1: (13, 4), 2: (15, 5)} + assert derive_layout(manager).layers == ((12, 13, 15),) + + +@pytest.mark.cpu_only +def test_an_internal_layer_several_attention_types_share_takes_the_first_s_global_id(): + """An internal layer several model layers share is named by its first entry, as native does.""" + + class AttentionType(enum.Enum): + SLIDING = 0 + COMPRESSED = 1 + INDEXER = 2 + + manager = Manager([kv_group(lids=(0, 1, 2))]) + manager._layer_attn_to_layer_id = { + (4, AttentionType.SLIDING): 0, + (4, AttentionType.COMPRESSED): 1, + (5, AttentionType.SLIDING): 2, + (5, AttentionType.COMPRESSED): 1, # shares model layer 4's compressed cache + } + assert manager._page_layer_ids() == {0: (12, 4), 1: (13, 4), 2: (15, 5)} assert derive_layout(manager).layers == ((12, 13, 15),) @pytest.mark.cpu_only -def test_declared_global_layer_ids_name_extra_internal_layers(): +def test_the_layout_carries_the_rank_heads_element_type_and_read_windows(): + """The layout carries native's rank info, KV heads, element type and read windows.""" + manager = Manager( + [kv_group(lids=(0, 1), heads_per_rank=2), kv_group(lids=(2,), heads_per_rank=2)], + tp_size=4, + tp_rank=3, + attention_dp=False, + num_kv_heads=8, + pp_layers=[10, 11, 12], + windows=[None, 64], + ) + manager.mapping.pp_size, manager.mapping.pp_rank = 2, 1 + manager.dtype = SimpleNamespace(name="FP8") + layout = derive_layout(manager).page_layout + assert layout.topology == RankTopology(tp_size=4, tp_rank=3, pp_size=2, pp_rank=1) + assert dict(layout.total_kv_heads_by_layer) == {10: 8, 11: 8, 12: 8} + assert (layout.dtype, layout.quant_mode) == ("FP8", "fp8") + assert [(g.kept_window, g.read_window) for g in layout.layer_groups] == [(None, None), (64, 64)] + adp = Manager([kv_group(heads_per_rank=8)], tp_size=4, tp_rank=3, attention_dp=True) + topology = derive_layout(adp).page_layout.topology + assert (topology.dp_size, topology.dp_rank, topology.attention_dp) == (4, 3, True) + state = derive_layout(recurrent_state()).page_layout + assert (dict(state.total_kv_heads_by_layer), state.dtype, state.quant_mode) == ({}, None, None) + + +@pytest.mark.cpu_only +def test_a_declared_read_window_narrows_the_kept_one(): + """A declared read window narrows the kept one; one the kept window does not hold is refused.""" + manager = Manager([kv_group(lids=(0,)), kv_group(lids=(1,))], windows=[None, 72]) + manager._attention_read_window = lambda lg: None if lg == 0 else 64 + geometry = derive_layout(manager) + groups = geometry.page_layout.layer_groups + assert [(g.kept_window, g.read_window) for g in groups] == [(None, None), (72, 64)] + assert geometry.windows == (None, 72) and geometry.read_windows == (None, 64) + for liar in (lambda lg: None if lg == 0 else 80, lambda lg: 64): + manager._attention_read_window = liar + with pytest.raises(ValueError, match="does not hold"): + derive_layout(manager) + + +@pytest.mark.cpu_only +def test_a_manager_naming_its_own_layers_names_extra_internal_layers(): """A manager whose internal layers outnumber its model layers (FP4 MLA keeps its tail on - extra ones) declares the ids per layer group; they are used as declared, in group order.""" + extra ones) names them itself (``_page_layer_ids``); a layer it leaves unnamed is refused.""" manager = Manager( [kv_group(lids=(1, 0)), (AttentionLayerConfig, [2, 3], [("tail", 64, None)])], pp_layers=[0, 1], - mappers={"tail": MapperKind.REPLICATED}, + mappers={"tail": "replicated"}, ) - declared = {0: [2, 0], 1: [1, 3]} - manager.get_disagg_global_layer_ids = lambda lg: declared[lg] - assert global_layer_ids(manager, [0, 1, 2, 3]) == [0, 2, 1, 3] - layout = derive_layout(manager).layout + named = {1: (2, 1), 0: (0, 0), 2: (1, 0), 3: (3, 1)} + manager._page_layer_ids = lambda: named + layout = derive_layout(manager).page_layout assert [g.layers for g in layout.layer_groups] == [(0, 2), (1, 3)] assert {(b.layer, b.role) for b in layout.buffers if b.role == "tail"} == { (1, "tail"), (3, "tail"), } - manager.get_disagg_global_layer_ids = lambda lg: declared[lg][:1] - with pytest.raises(ValueError, match="declared global ids"): - global_layer_ids(manager, [0]) + manager._page_layer_ids = lambda: {k: v for k, v in named.items() if k != 3} + with pytest.raises(ValueError, match=r"names no global layer for internal layers \[3\]"): + derive_layout(manager) @pytest.mark.cpu_only -def test_a_buffer_with_its_own_tokens_per_block_is_expanded(): +def test_layer_groups_publish_each_layer_s_model_layer(): + """A manager numbering its layers by model layer publishes them as their own model layers: + a pipeline stage's layers, attention and recurrent state alike.""" manager = Manager( - [kv_group(extra=[("index_key", 64, 2)])], mappers={"index_key": MapperKind.REPLICATED} + [ + kv_group(lids=(1, 0)), + (SsmLayerConfig, [2, 3], [("conv_state", 96, None), ("ssm_state", 128, None)]), + ], + pp_layers=[7, 6, 8, 9], + mappers={"conv_state": "sectioned"}, + layouts={"conv_state": dict(section_bytes=(32, 32, 32))}, ) - layout = derive_layout(manager).layout + layout = derive_layout(manager).page_layout + assert [(g.kind, g.layers, g.model_layers) for g in layout.layer_groups] == [ + ("attention", (6, 7), (6, 7)), + ("state", (8, 9), (8, 9)), + ] + + +@pytest.mark.cpu_only +def test_virtual_layers_publish_the_model_layer_they_belong_to(): + """Layer ``3 * m + t`` of model layer ``m`` and attention type ``t`` is model layer ``m``, + listed beside its layer id; an internal layer two model layers share is its first's.""" + + class AttentionType(enum.Enum): + SLIDING = 0 + COMPRESSED = 1 + INDEXER = 2 + + manager = Manager([kv_group(lids=(2, 0)), kv_group(lids=(1, 3))]) + manager._layer_attn_to_layer_id = { + (4, AttentionType.SLIDING): 0, + (4, AttentionType.COMPRESSED): 1, + (5, AttentionType.SLIDING): 2, + (6, AttentionType.COMPRESSED): 3, + (6, AttentionType.INDEXER): 1, # shares model layer 4's compressed cache + } + layout = derive_layout(manager).page_layout + assert [(g.layers, g.model_layers) for g in layout.layer_groups] == [ + ((12, 15), (4, 5)), + ((13, 19), (4, 6)), + ] + assert {lid: ids[1] for lid, ids in manager._page_layer_ids().items()} == { + 0: 4, + 1: 4, + 2: 5, + 3: 6, + } + + +@pytest.mark.cpu_only +def test_layers_a_manager_names_publish_the_model_layers_named_beside_them(): + """A manager naming its layers (FP4 MLA: ``2 * m`` for model layer ``m``'s cache and + ``2 * m + 1`` for its tail) names each one's model layer beside it.""" + manager = Manager( + [kv_group(lids=(1, 0)), (AttentionLayerConfig, [2, 3], [("tail", 64, None)])], + pp_layers=[0, 1], + mappers={"tail": "replicated"}, + ) + manager._page_layer_ids = lambda: {1: (2, 1), 0: (0, 0), 2: (1, 0), 3: (3, 1)} + layout = derive_layout(manager).page_layout + assert [(g.layers, g.model_layers) for g in layout.layer_groups] == [ + ((0, 2), (0, 1)), + ((1, 3), (0, 1)), + ] + + +@pytest.mark.cpu_only +def test_a_layer_group_names_one_model_layer_per_layer(): + group = LayerGroupDesc("attention", None, 0, 0, [12, 13, 15], model_layers=[4, 4, 5]) + assert group.model_layers == (4, 4, 5) and type(group.model_layers) is tuple + plain = LayerGroupDesc("attention", None, 0, 0, [3, 1]) + assert plain.model_layers == (3, 1) == plain.layers + assert group != LayerGroupDesc("attention", None, 0, 0, (12, 13, 15)) + assert hash(group) != hash(LayerGroupDesc("attention", None, 0, 0, (12, 13, 15))) + for bad in ([4, 4], [4, 4, 5, 6], [4, -1, 5], [4, True, 5], [4, 4.0, 5]): + with pytest.raises(ValueError, match="model layer"): + LayerGroupDesc("attention", None, 0, 0, (12, 13, 15), model_layers=bad) + + +@pytest.mark.cpu_only +def test_a_buffer_with_its_own_tokens_per_block_is_expanded(): + manager = Manager([kv_group(extra=[("index_key", 64, 2)])], mappers={"index_key": "replicated"}) + layout = derive_layout(manager).page_layout assert {b.role: b.expansion for b in layout.buffers} == {"key": 1, "value": 1, "index_key": 2} plain = derive_layout( Manager( [kv_group(extra=[("index_key", 64, None)])], - mappers={"index_key": MapperKind.REPLICATED}, + mappers={"index_key": "replicated"}, ) - ).layout + ).page_layout assert layout_id(layout) != layout_id(plain) @@ -739,7 +1190,7 @@ def test_the_manager_layout_holds_per_group_and_per_pool_group_facts(): assert geometry.num_layer_groups == 3 assert geometry.tokens_per_block == TPB assert geometry.pool_group_of == (3, 3, 1) - assert [d.pool_group for d in geometry.layout.layer_groups] == [3, 3, 1] + assert [d.pool_group for d in geometry.page_layout.layer_groups] == [3, 3, 1] assert geometry.pool_groups == (1, 3) assert geometry.windows == (None, 8, None) assert geometry.sink_blocks == (0, 2, 0) # 5 sink tokens take two 4-token blocks @@ -747,7 +1198,7 @@ def test_the_manager_layout_holds_per_group_and_per_pool_group_facts(): assert geometry.layers == ((0,), (1,), (2,)) assert geometry.shards == ((1, 0),) * 3 assert dict(geometry.page_bytes) == {3: 2 * HEAD_BYTES, 1: 96} - assert dict(geometry.layout.pool_groups) == {3: (HEAD_BYTES, HEAD_BYTES), 1: (96,)} + assert dict(geometry.page_layout.pool_slot_bytes) == {3: (HEAD_BYTES, HEAD_BYTES), 1: (96,)} pools = { g: [(p.group, p.index, p.base, p.slot_bytes, p.num_slots) for p in ps] for g, ps in geometry.device_pools.items() @@ -758,6 +1209,20 @@ def test_the_manager_layout_holds_per_group_and_per_pool_group_facts(): } +@pytest.mark.cpu_only +def test_a_helix_rank_s_layout_counts_the_tokens_of_its_own_page(): + """Under Helix the layout counts one page's tokens; the lender's block is ``cp_size`` pages.""" + manager = attention(1, 0, 4) + manager._has_cp_helix, manager._helix_cp_rank, manager._helix_cp_size = True, 1, 2 + manager.mapping.cp_size, manager.mapping.cp_rank = 2, 1 + manager.impl.init_config.tokens_per_block = 2 * TPB # the runtime's block: both ranks' pages + geometry = derive_layout(manager) + layout = geometry.page_layout + assert (layout.topology.cp_size, layout.topology.cp_rank) == (2, 1), layout.topology + assert layout.tokens_per_page == manager.tokens_per_block == TPB, layout.tokens_per_page + assert geometry.cp == (1, 2) and geometry.tokens_per_block == 2 * TPB, geometry + + @pytest.mark.cpu_only def test_a_state_group_is_told_by_its_config_class_not_its_name(): """Only the runtime's state layer config makes a group a state group, whatever another layer @@ -801,20 +1266,20 @@ def test_a_manager_whose_parts_disagree_is_refused(): tp_size=2, tp_rank=1, num_kv_heads=8, - mappers={"index_key": MapperKind.REPLICATED}, + mappers={"index_key": "replicated"}, ), - "069d747522257b726c6b28d917cb2c260350e26c55550bb687bea7350a76f51c", + "fd26ca16e4400324b2590f51fdff89ce93369776249ffdeb08e10ecaf07c3652", ), ( recurrent_state, - "3eba9af0efca5fe1e8701398e86b43ada86f07ba249f483be435925ffeaab6fa", + "ddbb463ba7a5ec34551df4ac7de306d27cd28d4bbf984e9645004ec26f005747", ), ], ids=["head_sharded_with_expanded_replicated_role", "recurrent_state"], ) def test_the_derived_layout_id_is_pinned(build, expected): """The derivation and the encoding together: a change here moves every name built on it.""" - assert layout_id(derive_layout(build()).layout).hex() == expected + assert layout_id(derive_layout(build()).page_layout).hex() == expected # -- real managers ------------------------------------------------------------------------------- @@ -837,13 +1302,13 @@ def native_pool_groups(mgr): def check_pool_groups(mgr, geometry): - layout = geometry.layout + layout = geometry.page_layout pools, counts = native_pool_groups(mgr) - assert set(layout.pool_groups) == set(pools) + assert set(layout.pool_slot_bytes) == set(pools) assert geometry.pool_groups == tuple(sorted(pools)) for g, native in pools.items(): widths = tuple(w for _, w in native) - assert layout.pool_groups[g] == widths + assert layout.pool_slot_bytes[g] == widths assert geometry.page_bytes[g] == sum(widths) # Slot i at base + i * the layout's slot width: a pool is its slots back to back. assert [(p.base, p.slot_bytes * p.num_slots) for p in geometry.device_pools[g]] == [ @@ -857,10 +1322,9 @@ def check_buffers(mgr, layout): for b in layout.buffers: assert (b.layer, b.role) not in seen seen[(b.layer, b.role)] = b - width = layout.pool_groups[b.pool[0]][b.pool[1]] - assert 0 <= b.offset and b.offset + b.size <= width - internal = sorted({int(bid.layer_id) for bid in mgr.impl.all_buffer_ids}) - global_of = dict(zip(internal, global_layer_ids(mgr, internal))) + width = layout.pool_slot_bytes[b.pool_group][b.pool_index] + assert 0 <= b.slot_offset and b.slot_offset + b.size <= width + global_of = {lid: ids[0] for lid, ids in mgr._page_layer_ids().items()} native = {(global_of[int(bid.layer_id)], str(bid.role)) for bid in mgr.impl.all_buffer_ids} assert set(seen) == native @@ -869,12 +1333,12 @@ def check_buffers(mgr, layout): def test_a_full_attention_cache_is_one_layer_group_over_one_pool_group(real_manager): with real_manager() as mgr: geometry = derive_layout(mgr) - layout = geometry.layout + layout = geometry.page_layout - assert layout.tokens_per_block == geometry.tokens_per_block == REAL_TPB + assert layout.tokens_per_page == geometry.tokens_per_block == REAL_TPB (group,) = layout.layer_groups assert group.kind == "attention" - assert group.window is None + assert group.kept_window is None assert group.layers == (0, 1) g = int(mgr.impl.get_life_cycle_pool_group_indices()[0]) assert group.pool_group == g @@ -888,9 +1352,8 @@ def test_a_full_attention_cache_is_one_layer_group_over_one_pool_group(real_mana # K and V are head-sharded (INDEXED): four heads of ``tokens_per_block * head_dim`` fp16 # values on this rank. for b in layout.buffers: - assert b.mapper is BufferMapper.INDEXED - assert b.transfer and b.expansion == 1 - assert b.geometry == BufferGeometry(bytes_per_head=REAL_TPB * 64 * 2, num_heads=4) + assert described(b) == ("INDEXED", 4, REAL_TPB * 64 * 2) + assert b.transferred and b.expansion == 1 assert group.shard == ShardDesc() @@ -898,7 +1361,7 @@ def test_a_full_attention_cache_is_one_layer_group_over_one_pool_group(real_mana def test_two_windows_are_two_layer_groups_with_their_own_windows(real_manager): with real_manager(windows=[WINDOW, 256]) as mgr: geometry = derive_layout(mgr) - layout = geometry.layout + layout = geometry.page_layout assert len(layout.layer_groups) == 2 native = [int(g) for g in mgr.impl.get_life_cycle_pool_group_indices()] @@ -906,9 +1369,9 @@ def test_two_windows_are_two_layer_groups_with_their_own_windows(real_manager): by_layer = {} for lg, group in enumerate(layout.layer_groups): assert group.pool_group == native[lg] - assert geometry.windows[lg] == group.window + assert geometry.windows[lg] == group.kept_window for layer in group.layers: - by_layer[layer] = group.window + by_layer[layer] = group.kept_window # Layer 0 slides; layer 1's window reaches max_seq_len and normalizes to full attention. assert mgr.max_seq_len <= 256 assert by_layer == {0: WINDOW, 1: None} @@ -919,13 +1382,13 @@ def test_two_windows_are_two_layer_groups_with_their_own_windows(real_manager): @gpu def test_layout_id_ignores_slot_counts_and_addresses(real_manager): with real_manager(max_tokens=2048) as small: - small_id = layout_id(derive_layout(small).layout) + small_id = layout_id(derive_layout(small).page_layout) small_pools = native_pool_groups(small) with real_manager(max_tokens=8192) as large: large_geometry = derive_layout(large) large_pools = native_pool_groups(large) assert large_pools[1] != small_pools[1], "the two caches must differ in slot counts" - assert layout_id(large_geometry.layout) == small_id + assert layout_id(large_geometry.page_layout) == small_id @gpu @@ -933,19 +1396,21 @@ def test_layout_id_leaves_the_window_out(real_manager): """One window for every layer: the same single layer group, only its window differs. The window says which blocks exist, not how a block's bytes read, so the id stays.""" with real_manager(windows=[WINDOW]) as mgr: - before = derive_layout(mgr).layout + before = derive_layout(mgr).page_layout with real_manager(windows=[2 * WINDOW]) as mgr: - after = derive_layout(mgr).layout - assert [g.window for g in before.layer_groups] != [g.window for g in after.layer_groups] + after = derive_layout(mgr).page_layout + assert [g.kept_window for g in before.layer_groups] != [ + g.kept_window for g in after.layer_groups + ] assert layout_id(after) == layout_id(before) @gpu def test_layout_id_follows_the_tokens_per_block(real_manager): with real_manager() as mgr: - before = derive_layout(mgr).layout + before = derive_layout(mgr).page_layout with real_manager(tokens_per_block=64) as mgr: - after = derive_layout(mgr).layout + after = derive_layout(mgr).page_layout assert len(after.layer_groups) == len(before.layer_groups) assert layout_id(after) != layout_id(before) @@ -957,9 +1422,9 @@ def test_layout_id_leaves_the_element_type_out(real_manager): from tensorrt_llm.bindings import DataType with real_manager(dtype=DataType.HALF) as mgr: - fp16, fp16_type = derive_layout(mgr).layout, mgr.dtype + fp16, fp16_type = derive_layout(mgr).page_layout, mgr.dtype with real_manager(dtype=DataType.BF16) as mgr: - bf16, bf16_type = derive_layout(mgr).layout, mgr.dtype + bf16, bf16_type = derive_layout(mgr).page_layout, mgr.dtype assert fp16_type != bf16_type, "the two caches hold the same element type: proves nothing" assert layout_id(bf16) == layout_id(fp16) @@ -1006,9 +1471,104 @@ def test_fp4_mla_buffers_carry_the_declared_global_layer_ids(kit, monkeypatch): """The tail lives on internal layers past the model's; in the layout its buffers take the ids the manager declares (``2 * layer + 1``), next to the cache's ``2 * layer``.""" with fp4_mla(kit, monkeypatch, "triton") as mgr: - buffers = {(b.layer, b.role) for b in derive_layout(mgr).layout.buffers} + # A cache with more layers than the model reads the windows it keeps. + mgr._model_attention_windows = dict.fromkeys(mgr.pp_layers, 1) + derived = derive_layout(mgr) + buffers = {(b.layer, b.role) for b in derived.page_layout.buffers} assert {layer for layer, role in buffers if role == "key"} == {0, 2, 4} assert {layer for layer, role in buffers if role == "mla_hp_tail"} == {1, 3, 5} + assert any(window is not None for window in derived.windows), derived.windows + assert derived.read_windows == derived.windows, (derived.windows, derived.read_windows) + + +@gpu +@skip_pre_hopper +def test_fp4_mla_publishes_each_layer_s_model_layer(kit, monkeypatch): + """The cache and the tail of model layer ``m``, layers ``2 * m`` and ``2 * m + 1``, are both + model layer ``m``.""" + with fp4_mla(kit, monkeypatch, "triton") as mgr: + mgr._model_attention_windows = dict.fromkeys(mgr.pp_layers, 1) + groups = derive_layout(mgr).page_layout.layer_groups + published = {layer: m for g in groups for layer, m in zip(g.layers, g.model_layers)} + assert published == {layer: layer // 2 for layer in range(6)}, published + + +@gpu +@skip_pre_hopper +def test_fp4_mla_lending_requires_the_sidecars_decode_reads(kit, monkeypatch): + """A lender fetch must deliver V sidecars because it has no native receive rebuild.""" + with fp4_mla(kit, monkeypatch, "triton") as mgr: + mgr._model_attention_windows = dict.fromkeys(mgr.pp_layers, 1) + layout = derive_layout(mgr).page_layout + required = [b for b in layout.buffers if b.role in ("mla_v_scale", "mla_v_packed")] + assert "mla_v_scale" in {b.role for b in required} + assert all(b.transferred for b in required) + + +@gpu +@skip_pre_blackwell +def test_retained_precision_twins_describe_their_original_layer_s_head_stride(kit): + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.mapping import Mapping + + with kit.managed( + lambda: KVCacheManagerV2( + kv_cache_config=KvCacheConfig( + dtype="mixed_nvfp4_bfloat16", + mixed_precision_first_n=1, + mixed_precision_last_n=65, + enable_block_reuse=True, + pool_ratio=None, + max_gpu_total_bytes=128 << 20, + ), + kv_cache_type=CacheType.SELF, + num_layers=2, + num_kv_heads=[2, 4], + head_dim=64, + tokens_per_block=128, + max_seq_len=1024, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1), + dtype=DataType.BF16, + ) + ) as mgr: + layout = derive_layout(mgr).page_layout + for original, internal in mgr.high_layers.items(): + global_layer = mgr._page_layer_ids()[internal][0] + actual = mgr.get_buffers(original, "HND", mixed_precision_pool="high") + head_bytes = actual.stride(-3) * actual.element_size() + buffers = [b for b in layout.buffers if b.layer == global_layer] + assert buffers and all(described(b)[2] == head_bytes for b in buffers) + + +@pytest.mark.cpu_only +def test_fp4_mla_installs_its_model_layers_beside_its_layer_ids(): + """A manager the FP4 MLA policy installs its methods on names each layer and its model layer, + and lends the roles native's transfer leaves to the receiver's rebuild, as an FP4 MLA manager + does.""" + from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import ( + Fp4MlaV2CacheLayoutPolicy, + ) + + # Model layers 4 and 7: their caches are internal layers 0 and 1, their tails 2 and 3. + manager = SimpleNamespace( + _fp4_mla_hp_pool_size=0, + impl=SimpleNamespace(layer_grouping=[[0, 1], [2, 3]]), + _cache_manager_layer_ids=[0, 1], + _hp_manager_layer_ids=[2, 3], + _bf16_mla_manager_layer_ids=[], + _fp4_mla_compact_to_local=[0, 1], + pp_layers=[4, 7], + ) + Fp4MlaV2CacheLayoutPolicy.install_manager_capabilities(manager) + assert manager._page_layer_ids() == {0: (8, 4), 1: (14, 7), 2: (9, 4), 3: (15, 7)} + # Native's transfer skips the V sidecars, which the receiver rebuilds; the page table says so, + # and a lender, which invokes no rebuild, lends them too. + assert manager._page_untransferred_roles() + assert manager._page_rebuilds_untransferred() M3_TPB = 128 @@ -1057,24 +1617,17 @@ def check_minimax_m3_nvfp4(mgr, layout): """Every buffer once, with the size, mapper and heads its storage has: FP8 K/V on dense layers, half-width K/V plus one scale byte per 16 elements on the sparse layer, replicated index-K.""" check_buffers(mgr, layout) - fp8 = (M3_TPB * M3_KV_HEADS * M3_HEAD_DIM, BufferMapper.INDEXED, M3_KV_HEADS) - packed = (fp8[0] // 2, BufferMapper.INDEXED, M3_KV_HEADS) - scale = (fp8[0] // 16, BufferMapper.INDEXED, M3_KV_HEADS) + fp8 = (M3_TPB * M3_KV_HEADS * M3_HEAD_DIM, "INDEXED", M3_KV_HEADS) + packed = (fp8[0] // 2, "INDEXED", M3_KV_HEADS) + scale = (fp8[0] // 16, "INDEXED", M3_KV_HEADS) expected = {(layer, role): fp8 for layer in range(4) for role in ("key", "value")} expected.update({(M3_SPARSE_LAYER, role): packed for role in ("key", "value")}) expected.update( {(M3_SPARSE_LAYER, role): scale for role in ("key_block_scale", "value_block_scale")} ) - expected[(M3_SPARSE_LAYER, "index_key")] = ( - M3_TPB * M3_INDEX_DIM * 2, - BufferMapper.REPLICATED, - None, - ) - described = { - (b.layer, b.role): (b.size, b.mapper, b.geometry.num_heads) for b in layout.buffers - } - assert described == expected - assert all(b.transfer and b.expansion == 1 for b in layout.buffers) + expected[(M3_SPARSE_LAYER, "index_key")] = (M3_TPB * M3_INDEX_DIM * 2, "REPLICATED", None) + assert {(b.layer, b.role): (b.size, *described(b)[:2]) for b in layout.buffers} == expected + assert all(b.transferred and b.expansion == 1 for b in layout.buffers) assert sorted(layer for group in layout.layer_groups for layer in group.layers) == [0, 1, 2, 3] assert all(group.shard == ShardDesc() for group in layout.layer_groups) @@ -1095,9 +1648,14 @@ def swap(old, new): "scale dropped": buffers(b for b in layout.buffers if b is not scale), "scale on a dense layer": buffers((*layout.buffers, dataclasses.replace(scale, layer=0))), "scale as wide as its data": swap( - scale, dataclasses.replace(scale, size=scale.size * 8, geometry=BufferGeometry()) + scale, + dataclasses.replace( + scale, size=scale.size * 8, tensors=head_major(scale.size * 8, 1, M3_KV_HEADS) + ), + ), + "index-K head-sharded": swap( + index, dataclasses.replace(index, tensors=token_major(index.size, M3_TPB, 1, 1)) ), - "index-K head-sharded": swap(index, dataclasses.replace(index, mapper=BufferMapper.NHD)), } @@ -1107,8 +1665,8 @@ def test_minimax_m3_nvfp4_describes_its_packed_kv_and_block_scales(kit): with minimax_m3(kit, "nvfp4") as mgr: geometry = derive_layout(mgr) check_pool_groups(mgr, geometry) - check_minimax_m3_nvfp4(mgr, geometry.layout) - for name, lie in minimax_m3_nvfp4_lies(geometry.layout).items(): + check_minimax_m3_nvfp4(mgr, geometry.page_layout) + for name, lie in minimax_m3_nvfp4_lies(geometry.page_layout).items(): with pytest.raises(AssertionError): check_minimax_m3_nvfp4(mgr, lie) pytest.fail(f"the {name} lie went unnoticed") @@ -1131,7 +1689,7 @@ def test_minimax_m3_nvfp4_and_fp8_caches_have_their_own_layout_ids(kit): builds = [] for kv_dtype, max_tokens in (("nvfp4", 1024), ("fp8", 1024), ("nvfp4", 16384)): with minimax_m3(kit, kv_dtype, max_tokens) as mgr: - builds.append((kv_dtype, derive_layout(mgr).layout, native_pool_groups(mgr)[1])) + builds.append((kv_dtype, derive_layout(mgr).page_layout, native_pool_groups(mgr)[1])) assert builds[0][2] != builds[2][2], "the two NVFP4 caches must differ in slot counts" check_minimax_m3_ids(builds, lambda layout, counts: layout_id(layout)) lies = { @@ -1150,6 +1708,29 @@ def test_minimax_m3_nvfp4_and_fp8_caches_have_their_own_layout_ids(kit): pytest.fail(f"{name} went unnoticed") +@gpu +@skip_pre_blackwell +@pytest.mark.parametrize("draft_len", [1, 3]) +@pytest.mark.parametrize("swa_only", [False, True], ids=["compressed", "swa_only"]) +def test_a_drafting_deepseek_v4_manager_uses_its_full_kept_windows( + deepseek_v4_manager, draft_len, swa_only +): + """Without a backend read declaration, lending requires the draft reserves too.""" + from tensorrt_llm.llmapi.llm_args import NGramDecodingConfig + + spec_config = NGramDecodingConfig(max_draft_len=draft_len, max_matching_ngram_size=2) + with deepseek_v4_manager( + spec_config=spec_config, compress_ratios=(1,) * 4 if swa_only else None + ) as mgr: + assert mgr.max_draft_len == draft_len, mgr.max_draft_len + groups = derive_layout(mgr).page_layout.layer_groups + windows = [(g.kept_window, g.read_window) for g in groups] + assert any(w is not None for w, _ in windows) + assert any(w is None for w, _ in windows) == (not swa_only) + for window, read in windows: + assert read == window, windows + + @gpu @skip_pre_blackwell @pytest.mark.parametrize("fp8_ds_mla", [False, True], ids=["fp8", "fp8_ds_mla"]) @@ -1159,12 +1740,13 @@ def test_the_deepseek_v4_layout_covers_every_virtual_layer_and_pool( with deepseek_v4_manager(fp8_ds_mla=fp8_ds_mla) as mgr: assert mgr.use_fp8_ds_mla == fp8_ds_mla geometry = derive_layout(mgr) - layout = geometry.layout + layout = geometry.page_layout life_cycles = mgr._life_cycle_by_layer_group() assert len(layout.layer_groups) == len(life_cycles) >= 2 assert list(geometry.windows) == [lc.window_size for lc in life_cycles] assert None in geometry.windows and any(w is not None for w in geometry.windows) + assert geometry.read_windows == geometry.windows, "no draft keeps a margin" # Several pools per group: a page is their slots back to back. native = {int(pg.pool_group_index): pg for pg in mgr.impl.pool_group_descs} @@ -1173,14 +1755,14 @@ def test_the_deepseek_v4_layout_covers_every_virtual_layer_and_pool( widths = tuple( int(p.slot_bytes) for p in sorted(pg.pools, key=lambda p: int(p.pool_index)) ) - assert layout.pool_groups[g] == widths + assert layout.pool_slot_bytes[g] == widths assert geometry.page_bytes[g] == sum(widths) # Buffers and layer groups name layers by global layer id: for virtual layers, model layer # times the number of attention types plus a type the internal layer holds. internal = sorted(int(i) for group in mgr.impl.layer_grouping for i in group) - global_of = dict(zip(internal, global_layer_ids(mgr, internal))) - assert len(set(global_of.values())) == len(internal) + global_of = {lid: ids[0] for lid, ids in mgr._page_layer_ids().items()} + assert sorted(global_of) == internal and len(set(global_of.values())) == len(internal) virtual = mgr._layer_attn_to_layer_id num_types = max(t.value for t in type(next(iter(virtual))[1])) + 1 for (model_layer, attn_type), layer_id in virtual.items(): @@ -1195,3 +1777,879 @@ def test_the_deepseek_v4_layout_covers_every_virtual_layer_and_pool( assert layers == sorted(global_of.values()) # One KV head: every layer group is the same content on every rank. assert all(group.shard == ShardDesc() for group in layout.layer_groups) + + +@gpu +@skip_pre_blackwell +def test_deepseek_v4_publishes_each_virtual_layer_s_model_layer(deepseek_v4_manager): + """Each cache of model layer ``m``, one per attention type, is model layer ``m``; so is the + last model layer's, as a draft layer after the target model's would be.""" + with deepseek_v4_manager(compress_ratios=(1, 4, 128, 4, 4)) as mgr: + groups = derive_layout(mgr).page_layout.layer_groups + virtual = mgr._layer_attn_to_layer_id + num_types = max(t.value for t in type(next(iter(virtual))[1])) + 1 + published = {layer: m for g in groups for layer, m in zip(g.layers, g.model_layers)} + assert published == {layer: layer // num_types for layer in published}, published + assert set(published.values()) == set(range(5)) == {m for m, _ in virtual} + assert any(layer != m for layer, m in published.items()) + + +@gpu +@skip_pre_blackwell +def test_csa2_publishes_each_virtual_layer_s_model_layer(kit, monkeypatch): + """CSA2 numbers each model layer's SWA, global and compressor caches apart; each is that + model layer.""" + from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager + from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.mapping import Mapping + + monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "0") + monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0") + config = KvCacheConfig( + max_gpu_total_bytes=128 << 20, host_cache_size=0, enable_block_reuse=False, dtype="fp8" + ) + with kit.managed( + lambda: CSA2CacheManager( + config, + CacheType.SELFKONLY, + num_layers=4, + tokens_per_block=128, + max_seq_len=2048, + max_batch_size=2, + max_num_tokens=2048, + mapping=Mapping(), + dtype=DataType.BF16, + vocab_size=8192, + layout=CSA2Layout((0, 2, 2, 1), (1, 3), (1, 3), window_size=256, index_topk=4), + ) + ) as mgr: + groups = derive_layout(mgr).page_layout.layer_groups + virtual = mgr._layer_attn_to_layer_id + num_types = max(t.value for t in type(next(iter(virtual))[1])) + 1 + published = {layer: m for g in groups for layer, m in zip(g.layers, g.model_layers)} + assert published == {layer: layer // num_types for layer in published}, published + assert set(published.values()) == {m for m, _ in virtual} + assert any(layer != m for layer, m in published.items()) + + +@gpu +def test_a_hybrid_stage_publishes_its_model_layers_for_state_and_attention(gdn_manager): + """On pipeline stage 1 of 2, a gated delta-net and attention manager holds model layers 4 to + 7 as its internal layers 0 to 3; both kinds of group publish the model layers.""" + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=2, rank=1, pp_size=2) + with gdn_manager(mapping=mapping, pattern="MMM*MMM*") as mgr: + assert list(mgr.pp_layers) == [4, 5, 6, 7], mgr.pp_layers + groups = derive_layout(mgr).page_layout.layer_groups + assert {g.kind for g in groups} == {"state", "attention"} + for g in groups: + assert g.model_layers == g.layers, g + assert sorted(layer for g in groups for layer in g.layers) == [4, 5, 6, 7] + + +@pytest.mark.cpu_only +def test_private_reuse_policy_stays_internal_and_disables_named_rows(): + """A PRIVATE layer group lends no block by name, and the public layout does not show it.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._slots import SlotQueue + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy + + manager = Manager([kv_group(lids=(0,)), kv_group(lids=(1,))]) + manager.impl.init_config.layers[1].reuse_policy = AttentionReusePolicy.PRIVATE + internal = derive_layout(manager) + assert internal.named_groups == (True, False) + assert internal.prefix_required == (True, False) + assert all(not hasattr(group, "reuse_policy") for group in internal.page_layout.layer_groups) + ref = weakref.ref(manager) + lender = _lender._ManagerLender(ref, internal, None, (), SlotQueue({}), ref) + pages = list(range(4)) + kv = SimpleNamespace( + num_committed_tokens=4 * TPB, + reuse_status=[SimpleNamespace(complete=True, endpoint=4 * TPB)] * 2, + num_blocks=len(pages), + get_base_page_indices=lambda lg: pages, + get_aggregated_page_indices=lambda lg, valid_only=False: pages, + ) + + def blocks(entries): + return dict(zip(entries.layer_group_indices, map(list, entries.block_ordinals))) + + every = lender._block_entries_of_range(manager, kv, 0, 4 * TPB) + assert blocks(every) == {0: [0, 1, 2, 3], 1: [0, 1, 2, 3]} + assert blocks(lender._committed_entries(kv, every)) == {0: [0, 1, 2, 3], 1: []} + named = lender._block_entries_of_range(manager, kv, 0, 4 * TPB, named=True) + assert blocks(named) == {0: [0, 1, 2, 3], 1: []} + + +@pytest.mark.cpu_only +def test_a_layout_names_the_groups_lent_by_name_or_builds_no_lender(): + """Which groups lend by name has no default: no layout and no lender builds without it.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._slots import SlotQueue + + manager = Manager([kv_group(lids=(0,)), kv_group(lids=(1,))]) + internal = derive_layout(manager) + fields = {f.name: getattr(internal, f.name) for f in dataclasses.fields(internal)} + del fields["named_groups"] + with pytest.raises(TypeError, match="named_groups"): + _layout.ManagerLayout(**fields) + unnamed = SimpleNamespace(**fields, num_layer_groups=2, pool_groups=internal.pool_groups) + ref = weakref.ref(manager) + with pytest.raises(AttributeError, match="named_groups"): + _lender._ManagerLender(ref, unnamed, None, (), SlotQueue({}), ref) + + +@gpu +@pytest.mark.parametrize("explicit", [True, False], ids=["model_window", "backend_fallback"]) +def test_mtp_read_windows_follow_the_model_or_the_actual_backend_fallback( + kit, real_manager, explicit +): + from tensorrt_llm._torch.models.modeling_exaone_moe import get_exaone_attention_window + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestType + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + config = SimpleNamespace( + num_hidden_layers=2, + layer_types=["sliding_attention", "full_attention"], + sliding_windows=[256, 0], + sliding_window=256, + mtp_layer_types=None, + ) + windows = { + i: get_exaone_attention_window(config, i, i >= 2) if explicit else None for i in range(3) + } + assert windows == {0: 256 if explicit else None, 1: None, 2: None} + with real_manager( + num_layers=2, + max_seq_len=2048, + max_tokens=4096, + tokens_per_block=128, + windows=[256, 2048, 2048], + spec_config=MTPDecodingConfig(max_draft_len=1), + attention_read_windows=windows, + block_reuse_config={"policy": "per_request"}, + ) as mgr: + lender = attach_lender(mgr, scope=b"mtp-window", staging=StagingOptions(511)) + groups = lender.page_layout.layer_groups + swa = next(i for i, g in enumerate(groups) if g.kept_window is not None) + assert (groups[swa].kept_window, groups[swa].read_window) == (257, 256 if explicit else 257) + req = kit.make_request( + 91, list(range(511)), llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + ) + assert mgr.prepare_disagg_gen_init(req) + event = mgr.take_disagg_receive_ready(req) + lease = lender.lend_write(req, 0, 511, **{**kit.DEVICE_ALL, "ready_event": event}) + mgr._stream.synchronize() + view = lease.poll() + assert view is not None, lease.failure + kit.fill_sentinel(mgr, req) + mgr._stream.synchronize() + masks = [ + np.isin(g.page_ordinals, (2, 3)) | (g.layer_group_index != swa) for g in view.groups + ] + (bounds,) = lender.return_writes([(lease, masks, True)]) + assert bounds == lender.poll_resume_bounds(req), bounds + if explicit: + assert bounds == (511, 511) + else: + assert bounds.max_position < bounds.min_position + + +@gpu +def test_a_configured_window_below_the_model_s_is_the_window_attention_reads(real_manager): + """A configured window below the model's is kept, and is then the window attention reads.""" + with real_manager( + num_layers=2, + max_seq_len=2048, + max_tokens=4096, + tokens_per_block=128, + windows=[256, 2048], + attention_read_windows={0: 4096, 1: None}, + block_reuse_config={"policy": "per_request"}, + ) as mgr: + lender = attach_lender(mgr, scope=b"configured-window", staging=StagingOptions(511)) + (group,) = [g for g in lender.page_layout.layer_groups if g.kept_window is not None] + assert group.kept_window < 4096 and group.read_window == group.kept_window, group + + +@gpu +@skip_pre_blackwell +@pytest.mark.parametrize("draft_len", [0, 8]) +def test_csa2_lender_uses_kept_swa_and_compressor_windows(kit, monkeypatch, draft_len): + from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager + from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, KvCacheConfig + from tensorrt_llm.mapping import Mapping + + monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "0") + monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0") + spec = DSparkDecodingConfig(max_draft_len=draft_len) if draft_len else None + config = KvCacheConfig( + max_gpu_total_bytes=128 << 20, host_cache_size=0, enable_block_reuse=False, dtype="fp8" + ) + with kit.managed( + lambda: CSA2CacheManager( + config, + CacheType.SELFKONLY, + num_layers=4, + tokens_per_block=128, + max_seq_len=2048, + max_batch_size=2, + max_num_tokens=2048, + mapping=Mapping(), + dtype=DataType.BF16, + vocab_size=8192, + layout=CSA2Layout((0, 2, 2, 1), (1, 3), (1, 3), window_size=256, index_topk=4), + spec_config=spec, + ) + ) as mgr: + public = derive_layout(mgr).page_layout + margin = mgr.max_draft_len + mgr.reuse_match_backoff + seen = set() + for desc in public.layer_groups: + roles = {buffer.role for buffer in public.buffers if buffer.layer in desc.layers} + seen.update(roles) + if roles & {"csa2_global", "csa2_index"}: + assert desc.read_window is None + else: + assert desc.read_window == desc.kept_window + if "csa2_swa" in roles: + assert desc.read_window == mgr.layout.window_size + margin + else: + assert desc.read_window == 2 + margin + kinds = ("swa", "global", "index", "compressor_kv", "compressor_score") + assert seen == {f"csa2_{kind}" for kind in kinds} + + +@gpu +@skip_pre_blackwell +def test_csa2_lender_requires_the_full_bounded_replay_window(kit, monkeypatch): + from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager + from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, KvCacheConfig + from tensorrt_llm.mapping import Mapping + + monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "1") + monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0") + read_window, replay_window = 256, 512 + spec = DSparkDecodingConfig(max_draft_len=8) + spec.__dict__["draft_is_embedded_in_target"] = True + with kit.managed( + lambda: CSA2CacheManager( + KvCacheConfig( + max_gpu_total_bytes=128 << 20, + host_cache_size=0, + enable_block_reuse=False, + dtype="fp8", + ), + CacheType.SELFKONLY, + num_layers=4, + tokens_per_block=128, + max_seq_len=2048, + max_batch_size=2, + max_num_tokens=2048, + mapping=Mapping(), + dtype=DataType.BF16, + vocab_size=8192, + layout=CSA2Layout((0, 2, 2, 1), (1, 3), (1, 3), window_size=read_window, index_topk=4), + spec_config=spec, + pretrained_config=SimpleNamespace(sliding_window=replay_window), + ) + ) as mgr: + assert mgr._decoder_replay_window == replay_window + layout = derive_layout(mgr).page_layout + swa = [ + group + for group in layout.layer_groups + if any( + buffer.role == "csa2_swa" and buffer.layer in group.layers + for buffer in layout.buffers + ) + ] + margin = mgr.max_draft_len + mgr.reuse_match_backoff + assert all(group.read_window == group.kept_window for group in swa) + assert {group.kept_window for group in swa} == { + read_window + margin, + replay_window + margin, + } + + +@pytest.mark.cpu_only +def test_csa2_decoder_replay_reads_rebuilt_rows_inside_its_kept_margin(): + """A replay query can read rebuilt SWA rows outside the final query's read window.""" + from tensorrt_llm._torch.attention.backends.sparse.csa2.decoder_replay import ( + plan_decoder_replay, + ) + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import ( + gather_rows, + row_bytes, + store_rows, + ) + + prompt, read_window, replay_window, tokens_per_block = 768, 256, 512, 128 + seq_lens = torch.tensor([prompt], dtype=torch.int32) + metadata = SimpleNamespace( + is_cuda_graph=False, + mapping=None, + num_tokens=prompt, + num_contexts=1, + num_seqs=1, + seq_lens=seq_lens, + seq_lens_cuda=seq_lens, + prompt_lens=[prompt], + kv_cache_params=SimpleNamespace(num_cached_tokens_per_seq=[0]), + ) + request = SimpleNamespace(prompt_len=prompt, py_ced_replay=None) + plan = plan_decoder_replay(metadata, replay_window, [request], private_decoder=True) + floor = prompt - replay_window + assert plan is not None and plan.swa_floors == [floor] + assert plan.rows.tolist() == list(range(floor, prompt)) + + positions = plan.rows + requests = torch.zeros(len(positions), dtype=torch.int64) + floors = torch.tensor(plan.swa_floors, dtype=torch.int64) + pages = torch.arange(prompt // tokens_per_block, dtype=torch.int32).unsqueeze(0) + reads = torch.empty((len(positions), read_window), dtype=torch.int64) + writes = torch.empty(len(positions), dtype=torch.int64) + visible = torch.empty(len(positions), dtype=torch.int64) + CSA2TrtllmMetadata._refresh_swa_tensor_outputs( + positions, + requests, + floors, + tokens_per_block, + read_window, + [(pages, reads, writes, visible, 0)], + read_floors=floors, + ) + assert writes.tolist() == list(range(floor, prompt)) + assert reads[0, :-1].eq(-1).all() and reads[0, -1] == floor + assert reads[read_window, 0] == floor + 1 + + pool = torch.empty((prompt, row_bytes(512, "swa")), dtype=torch.uint8) + old_rows = torch.full((prompt, 512), -1, dtype=torch.bfloat16) + rebuilt_rows = torch.ones((len(positions), 512), dtype=torch.bfloat16) + store_rows(pool, torch.arange(prompt), old_rows, "swa") + store_rows(pool, writes, rebuilt_rows, "swa") + assert gather_rows(pool, torch.tensor([floor - 1]), 512, "swa").eq(-1).all() + assert gather_rows(pool, reads[0], 512, "swa")[:-1].eq(0).all() + assert gather_rows(pool, reads[read_window], 512, "swa").eq(1).all() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("policy", ["OPTIONAL", "PRIVATE"]) +@pytest.mark.parametrize("window", [None, 64], ids=["full_attention", "windowed"]) +def test_attach_requires_a_window_for_every_non_required_group(monkeypatch, policy, window): + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy + + manager = Manager([kv_group(lids=(0,)), kv_group(lids=(1,))], windows=[None, window]) + manager.mapping.cp_size = 1 + manager.impl.init_config.layers[1].reuse_policy = getattr(AttentionReusePolicy, policy) + monkeypatch.setattr(_manager, "require_v2", lambda owner: None) + if window is None: + with pytest.raises(ValueError, match="non-REQUIRED.*window"): + _lender._checked_layout(manager) + else: + assert _lender._checked_layout(manager).windows == (None, 64) + + +class BaseDeclarations(Manager): + """The stand-in declaring its buffers as the base manager does.""" + + _page_buffer_tensors = KVCacheManagerV2._page_buffer_tensors + + +@pytest.mark.cpu_only +def test_a_role_the_manager_does_not_declare_is_refused(): + """The base manager declares only the roles it registers; a role a manager adds without + declaring how its bytes spread over ranks is refused rather than guessed.""" + manager = BaseDeclarations([kv_group(extra=[("side", 64, None)])]) + with pytest.raises(ValueError, match="registers role 'side' but does not declare"): + derive_layout(manager) + + +@pytest.mark.cpu_only +def test_a_role_that_never_moves_is_the_rank_s_own(): + manager = BaseDeclarations([kv_group(extra=[("side", 64, None)])], ignored={"side"}) + layout = derive_layout(manager).page_layout + sides = [b for b in layout.buffers if b.role == "side"] + assert sides and all(not b.transferred for b in sides) + assert all(b.tensors == page_layout.owned(64, 1) for b in sides) + keys = [b for b in layout.buffers if b.role == "key"] + assert all(described(b)[0] == "INDEXED" for b in keys) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("attention_dp", [False, True], ids=["tp", "attention_dp"]) +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_a_state_group_splits_over_the_grid_its_kind_sets(tp_size, attention_dp): + """State groups split over the TP ranks, and over none under attention data parallelism.""" + grid = 1 if attention_dp else tp_size + for rank in range(tp_size): + manager = Manager( + [(SsmLayerConfig, [0, 1], [("ssm_state", 128, None)])], + tp_size=tp_size, + tp_rank=rank, + attention_dp=attention_dp, + layouts={"ssm_state": dict(bytes_per_head=32)}, + ) + (group,) = derive_layout(manager).page_layout.layer_groups + assert group.shard == (ShardDesc(grid, rank) if grid > 1 else ShardDesc()), (rank, group) + + +@gpu +def test_a_v2_manager_whose_tensors_overrun_a_buffer_attaches_no_lender(real_manager, monkeypatch): + """A V2 manager declaring a tensor that reaches past its buffer fails to attach and keeps no + lender.""" + with real_manager() as mgr: + monkeypatch.setattr( + mgr, + "_page_buffer_tensors", + lambda layer, role, size, tokens, state: replicated(size + 1), + ) + with pytest.raises(ValueError, match="reaches byte"): + attach_lender(mgr, scope=b"overrun", staging=StagingOptions(64)) + assert mgr._sharing is None + + +@gpu +def test_a_manager_whose_regions_overlap_attaches_no_lender(real_manager, monkeypatch): + """Two device regions that overlap fail the attach, which keeps no lender.""" + regions_of = _lender._device_regions + + def overlapping(*args): + regions = regions_of(*args) + first = regions[0] + return [*regions, dataclasses.replace(first, address=first.address + first.nbytes - 1)] + + with real_manager() as mgr: + monkeypatch.setattr(_lender, "_device_regions", overlapping) + with pytest.raises(ValueError, match="overlap in DEVICE memory"): + attach_lender(mgr, scope=b"overlap", staging=StagingOptions(64)) + assert mgr._sharing is None + + +# -- the public layout types --------------------------------------------------------------------- + +LAYOUT_TYPES = ( + PageLayout, + LayerGroupDesc, + BufferDesc, + RankTopology, + ShardDesc, + Tensor, + Axis, +) +# Each layout type's public names beyond its fields and ``unchecked``. +LAYOUT_TYPE_EXTRAS = {BufferDesc: {"split"}, Tensor: {"split"}} + + +def pinned_fields(cls): + """``name`` or ``name=default`` per field, in order; a factory default shows as ``{}``.""" + out = [] + for f in dataclasses.fields(cls): + default = f.default_factory() if f.default_factory is not dataclasses.MISSING else f.default + out.append(f.name if default is dataclasses.MISSING else f"{f.name}={default!r}") + return " ".join(out) + + +@pytest.mark.cpu_only +def test_the_layout_types_keep_their_module_fields_and_members(): + assert sorted(page_layout.__all__) == sorted( + [ + *(t.__name__ for t in LAYOUT_TYPES), + "GRID_DIMENSIONS", + "MapperKind", + "Progression", + "RoleLayout", + "grid_place", + "head_major", + "held", + "owned", + "replica_ranks", + "replicated", + "sectioned", + "token_major", + ] + ) + assert not hasattr(sharing, "PageLayout") and "PageLayout" not in sharing.__all__ + assert _layout.ManagerLayout is page_table.ManagerLayout + assert _layout.DevicePool is page_table.DevicePool + assert pinned_fields(PageLayout) == ( + "tokens_per_page pool_slot_bytes layer_groups buffers topology=RankTopology(tp_size=1, " + "tp_rank=0, pp_size=1, pp_rank=0, cp_size=1, cp_rank=0, dp_size=1, dp_rank=0, " + "attention_dp=False) total_kv_heads_by_layer={} dtype=None quant_mode=None " + "device_region_index={} staging_region_index={} rebuilds_untransferred=False" + ) + assert pinned_fields(LayerGroupDesc) == ( + "kind kept_window num_sink_blocks pool_group layers shard=ShardDesc(count=1, index=0) " + "read_window=None model_layers=None" + ) + assert pinned_fields(BufferDesc) == ( + "layer role pool_group pool_index slot_offset size tensors=() transferred=True expansion=1" + ) + assert pinned_fields(Tensor) == "axes offset=0 name=''" + assert pinned_fields(Axis) == "name extent stride split=() held=None cyclic=False" + assert pinned_fields(RankTopology) == ( + "tp_size=1 tp_rank=0 pp_size=1 pp_rank=0 cp_size=1 cp_rank=0 dp_size=1 dp_rank=0 " + "attention_dp=False" + ) + assert pinned_fields(ShardDesc) == "count=1 index=0" + layout = every_field_set() + group, buffer = layout.layer_groups[0], layout.buffers[0] + tensor, axis = buffer.tensors[0], buffer.tensors[0].axes[0] + for value in (layout, group, buffer, layout.topology, group.shard, tensor, axis): + cls = type(value) + assert cls.__dataclass_params__.frozen, cls + named = {f.name for f in dataclasses.fields(cls)} + beyond = {name for obj in (cls, value) for name in dir(obj) if not name.startswith("_")} + assert beyond - named == {"unchecked", *LAYOUT_TYPE_EXTRAS.get(cls, ())}, cls + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("unchecked", [False, True], ids=["checked", "unchecked"]) +def test_a_layout_holds_its_own_read_only_copies_and_compares_by_value(unchecked): + pools, heads, regions = {0: (100, 40), 1: (64,)}, {0: 8}, {(0, 0): 2} + groups, buffers = list(make_layout().layer_groups), list(make_layout().buffers) + build = PageLayout.unchecked if unchecked else PageLayout + layout = build( + 32, pools, groups, buffers, total_kv_heads_by_layer=heads, device_region_index=regions + ) + pools[5], heads[9], regions[(9, 9)] = (1,), 1, 9 + groups.pop() + buffers.clear() + same = make_layout(total_kv_heads_by_layer={0: 8}, device_region_index={(0, 0): 2}) + assert layout == same and hash(layout) == hash(same), "the caller's values changed the layout" + assert len({layout, same, dataclasses.replace(same, staging_region_index={0: 1})}) == 2 + assert type(layout.layer_groups) is tuple and type(layout.buffers) is tuple + for mapping in ( + layout.pool_slot_bytes, + layout.total_kv_heads_by_layer, + layout.device_region_index, + layout.staging_region_index, + ): + assert isinstance(mapping, types.MappingProxyType) + with pytest.raises(TypeError): + mapping[7] = 0 + + +@pytest.mark.cpu_only +def test_a_read_window_defaults_to_the_kept_window_and_lies_within_it(): + assert LayerGroupDesc("attention", None, 0, 0, (0,)).read_window is None + assert LayerGroupDesc("attention", 133, 0, 0, (0,)).read_window == 133 + assert LayerGroupDesc("attention", 133, 0, 0, (0,), read_window=128).read_window == 128 + shorter = dataclasses.replace(LayerGroupDesc("attention", None, 0, 0, (0,)), kept_window=64) + assert shorter.read_window == 64, "a replaced window left the read window behind" + for window, read in ((64, 65), (64, 0), (64, -1), (None, 8)): + with pytest.raises(ValueError, match=r"read.window"): + LayerGroupDesc("attention", window, 0, 0, (0,), read_window=read) + kept = LayerGroupDesc("attention", 133, 0, 0, (0,)) + with pytest.raises(ValueError, match="read window of 133 tokens"): + dataclasses.replace(kept, kept_window=64) + assert dataclasses.replace(kept, kept_window=64, read_window=None).read_window == 64 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("axis", ["tp", "pp", "cp", "dp"]) +def test_a_topology_rank_lies_within_its_size(axis): + assert getattr(RankTopology(**{f"{axis}_size": 4, f"{axis}_rank": 3}), f"{axis}_rank") == 3 + for size, rank in ((4, 4), (4, -1), (0, 0)): + with pytest.raises(ValueError, match=axis): + RankTopology(**{f"{axis}_size": size, f"{axis}_rank": rank}) + + +@pytest.mark.cpu_only +def test_the_quantization_is_one_the_layout_names(): + import typing + + literal, none = typing.get_args(typing.get_type_hints(PageLayout)["quant_mode"]) + assert none is type(None) and typing.get_args(literal) == ("fp8", "nvfp4", "fp8_ds_mla") + for mode in (None, *typing.get_args(literal)): + assert make_layout(quant_mode=mode).quant_mode == mode + for mode in ("FP8", "int8", ""): + with pytest.raises(ValueError, match="quant_mode"): + make_layout(quant_mode=mode) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "mapping,expected", + [ + ( + dict(tp_size=4, tp_rank=3, pp_size=2, pp_rank=1, cp_size=1, cp_rank=0, dp_size=1), + RankTopology(tp_size=4, tp_rank=3, pp_size=2, pp_rank=1), + ), + ( + dict(tp_size=4, tp_rank=3, enable_attention_dp=True, dp_size=4), + RankTopology(tp_size=4, tp_rank=3, dp_size=4, dp_rank=3, attention_dp=True), + ), + (dict(tp_size=1, tp_rank=0, cp_size=2, cp_rank=1), RankTopology(cp_size=2, cp_rank=1)), + ], + ids=["tensor_and_pipeline", "attention_data_parallel", "helix"], +) +def test_the_topology_is_what_native_rank_info_reads_from_the_mapping(mapping, expected): + """As native's ``RankInfo``: under attention DP, ``dp_size`` and ``dp_rank`` are the TP ones.""" + assert page_table.rank_topology(SimpleNamespace(**mapping)) == expected + assert page_table.rank_topology(None) == RankTopology() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "dtype,quant", + [("BF16", None), ("HALF", None), ("FP8", "fp8"), ("NVFP4", "nvfp4"), (None, None)], +) +def test_the_quantization_follows_the_element_type_the_manager_declares(dtype, quant): + manager = attention(1, 0, 4) + if dtype is not None: + manager.dtype = SimpleNamespace(name=dtype) + assert (manager._page_dtype(), manager._page_quant_mode()) == (dtype, quant) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("ds_mla,quant", [(True, "fp8_ds_mla"), (False, "fp8")]) +def test_deepseek_v4_names_its_fp8_pages_with_their_scales(ds_mla, quant): + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager import ( + DeepseekV4CacheManager, + ) + + manager = DeepseekV4CacheManager.__new__(DeepseekV4CacheManager) + manager.dtype, manager.use_fp8_ds_mla = SimpleNamespace(name="FP8"), ds_mla + assert manager._page_quant_mode() == quant + + +@pytest.mark.cpu_only +def test_the_layout_id_leaves_out_the_topology_heads_element_type_and_regions(): + layout = make_layout() + shorter = dataclasses.replace(layout.layer_groups[1], read_window=59) + rich = make_layout( + topology=RankTopology(tp_size=2, tp_rank=1, pp_size=2, pp_rank=1), + total_kv_heads_by_layer={0: 8}, + dtype="FP8", + quant_mode="fp8", + device_region_index={(0, 0): 4}, + staging_region_index={0: 0}, + layer_groups=(layout.layer_groups[0], shorter), + ) + assert layout_id(rich) == layout_id(layout) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("staged", [True, False], ids=["staging", "in_place_alone"]) +def test_a_lender_s_layout_gives_each_pool_its_region_positions(staged): + """Staging regions come first, one per pool group, then one device region per pool.""" + manager = Manager([kv_group(lids=(0,)), kv_group(lids=(1,))]) + geometry = derive_layout(manager) + ident = Identity(b"scope", layout_id(geometry.page_layout), geometry.layers, geometry.shards) + staging = [ + MemoryRegion(f"s{g}", 4096 * (g + 1), 4096, 1024, 4, (g,), MemoryKind.HOST) for g in (0, 1) + ] + regions = tuple((staging if staged else []) + _lender._device_regions(manager, geometry, ident)) + ref = weakref.ref(manager) + lender = _lender._ManagerLender(ref, geometry, ident, regions, None, ref) + public = lender.page_layout + first = 2 if staged else 0 + order = [(0, 0), (0, 1), (1, 0), (1, 1)] + assert dict(public.device_region_index) == {pool: first + i for i, pool in enumerate(order)} + bases = {(p.group, p.index): p.base for pools in geometry.device_pools.values() for p in pools} + for pool, position in public.device_region_index.items(): + region = regions[position] + assert region.memory_kind is MemoryKind.DEVICE and region.address == bases[pool], pool + assert dict(public.staging_region_index) == ({0: 0, 1: 1} if staged else {}) + assert dataclasses.replace(public, device_region_index={}, staging_region_index={}) == ( + geometry.page_layout + ) + assert lender.page_layout is public, "the layout is built once" + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("bad", [True, False, 1.5, "2", -1, 0]) +def test_layout_values_are_refused_rather_than_truncated(bad): + """Indices refuse booleans, non-integers and negatives; sizes refuse zero as well.""" + buffer = dict(layer=0, role="kv", pool_group=0, pool_index=0, slot_offset=0, size=16) + names = ["size", "expansion"] + if isinstance(bad, bool) or bad != 0: + names += ["layer", "pool_group", "pool_index", "slot_offset"] + with pytest.raises(ValueError, match="model layer"): + LayerGroupDesc("attention", None, 0, 0, (0,), model_layers=(bad,)) + for name in names: + with pytest.raises(ValueError, match=name): + BufferDesc(**{**buffer, name: bad}) + for construct in ( + lambda: make_layout(tokens_per_page=bad), + lambda: make_layout(pool_slot_bytes={0: (bad,)}), + lambda: make_layout(total_kv_heads_by_layer={0: bad}), + lambda: ShardDesc(count=bad), + lambda: RankTopology(tp_size=bad), + lambda: LayerGroupDesc("attention", bad, 0, 0, (0,)), + ): + with pytest.raises(ValueError): + construct() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "sections", [(16, 16), (32, 16, 32), (16, 8, 32)], ids=["short", "overlapping", "gap"] +) +def test_a_buffer_of_several_tensors_must_be_covered_by_them_back_to_back(sections): + """Sections that leave bytes out, or overlap, would move part of a buffer or a byte twice.""" + assert BufferDesc(0, "conv", 0, 0, 0, 64, sectioned((16, 16, 32), 1)).size == 64 + tensors = tuple( + dataclasses.replace(t, offset=offset) + for t, offset in zip(sectioned(sections, 1), (0, 16, 32)) + ) + with pytest.raises(ValueError, match="do not cover its 64 bytes"): + BufferDesc(0, "conv", 0, 0, 0, 64, tensors) + + +def every_field_set(): + """A layout with each of its own fields away from the default, and parts of every kind.""" + return make_layout( + layer_groups=( + LayerGroupDesc( + "attention", 128, 1, 0, (0, 2), ShardDesc(2, 1), read_window=96, model_layers=(0, 1) + ), + LayerGroupDesc("state", None, 0, 1, (1,)), + ), + buffers=( + BufferDesc(0, "key", 0, 0, 0, 64, token_major(64, 4, 2, 8)), + BufferDesc(2, "scale", 0, 1, 8, 8, transferred=False), + BufferDesc(1, "conv", 1, 0, 0, 48, sectioned((16, 32), 2), expansion=2), + BufferDesc( + 1, + "rows", + 1, + 0, + 0, + 48, + ( + Tensor( + (Axis("row", 12, 8, ("tp",), cyclic=True), Axis("byte", 8, 1)), + offset=8, + name="r", + ), + ), + ), + ), + topology=RankTopology(4, 3, 2, 1, 2, 1, 4, 3, True), + total_kv_heads_by_layer={0: 8, 2: 8}, + dtype="FP8", + quant_mode="fp8", + device_region_index={(0, 0): 2, (0, 1): 3, (1, 0): 4}, + staging_region_index={0: 0, 1: 1}, + ) + + +def rebuilt(value): + """``value`` rebuilt through the unchecked constructors from plain values.""" + if dataclasses.is_dataclass(value): + parts = {f.name: rebuilt(getattr(value, f.name)) for f in dataclasses.fields(value)} + return type(value).unchecked(**parts) + if isinstance(value, types.MappingProxyType): + return {key: rebuilt(item) for key, item in value.items()} + return tuple(map(rebuilt, value)) if isinstance(value, tuple) else value + + +def typed(value): + """``value`` as nested ``(type name, parts)``, equal only for equal values of equal types.""" + name = type(value).__name__ + if dataclasses.is_dataclass(value): + return name, tuple(typed(getattr(value, f.name)) for f in dataclasses.fields(value)) + if isinstance(value, types.MappingProxyType): + return name, tuple(sorted((typed(key), typed(item)) for key, item in value.items())) + return name, tuple(map(typed, value)) if isinstance(value, tuple) else value + + +@pytest.mark.cpu_only +def test_each_unchecked_constructor_takes_the_constructor_s_fields_and_defaults(): + for cls in LAYOUT_TYPES: + parameters = list(inspect.signature(cls.unchecked).parameters.values()) + assert [p.name for p in parameters] == [f.name for f in dataclasses.fields(cls)], cls + for p, f in zip(parameters, dataclasses.fields(cls)): + if f.default is not dataclasses.MISSING: + assert p.default == f.default, (cls, p.name, p.default) + elif f.default_factory is not dataclasses.MISSING: + assert p.default == f.default_factory(), (cls, p.name, p.default) + else: + assert p.default is inspect.Parameter.empty, (cls, p.name, p.default) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("build", [make_layout, every_field_set], ids=["defaults", "every_field"]) +def test_an_unchecked_rebuild_equals_the_checked_layout(build): + """Rebuilt from plain values without checks, a layout and each of its parts come back alike.""" + layout = build() + again = rebuilt(layout) + assert again == layout and hash(again) == hash(layout) and typed(again) == typed(layout) + tensors = [t for b in layout.buffers for t in b.tensors] + parts = ( + *layout.layer_groups, + *layout.buffers, + layout.topology, + *tensors, + *(a for t in tensors for a in t.axes), + ) + for value in (layout, ShardDesc(2, 1), *parts): + fields = {f.name: getattr(value, f.name) for f in dataclasses.fields(value)} + unchecked = type(value).unchecked(**fields) + assert typed(unchecked) == typed(value) == typed(type(value)(**fields)), value + assert LayerGroupDesc.unchecked("attention", 64, 0, 0, (0,)).read_window == 64 + assert LayerGroupDesc.unchecked("attention", 64, 0, 0, (0, 2)).model_layers == (0, 2) + + +@pytest.mark.cpu_only +def test_unchecked_constructors_keep_values_the_constructor_refuses_or_converts(): + """Values the constructor refuses or converts stay as given; a list makes them unhashable.""" + group = dict(kind="attention", kept_window=64, num_sink_blocks=0, pool_group=0, layers=(0,)) + odd = BufferDesc(0, "kv", 0, 0, 0, 8, expansion=3) + refused = [ + (BufferDesc, dict(layer=-1, role="kv", pool_group=0, pool_index=0, slot_offset=0, size=8)), + (LayerGroupDesc, dict(group, read_window=128)), + (LayerGroupDesc, dict(group, model_layers=(0, 1))), + (RankTopology, dict(tp_size=2, tp_rank=2)), + (ShardDesc, dict(count=1, index=1)), + (PageLayout, dict(tokens_per_page=32, pool_slot_bytes={}, layer_groups=(), buffers=(odd,))), + ] + for cls, values in refused: + with pytest.raises(ValueError): + cls(**values) + kept = cls.unchecked(**values) + assert all(getattr(kept, name) == value for name, value in values.items()), cls + with pytest.raises(ValueError): + dataclasses.replace(kept) + listed = LayerGroupDesc.unchecked("attention", None, 0, 0, [0, 2]) + assert listed.layers == [0, 2] + assert LayerGroupDesc("attention", None, 0, 0, [0, 2]).layers == (0, 2) + with pytest.raises(TypeError, match="unhashable"): + hash(listed) + + +@pytest.mark.cpu_only +def test_page_layout_imports_with_only_the_standard_library_available(): + script = """ +import importlib.util +import sys +before = set(sys.modules) +spec = importlib.util.spec_from_file_location("_page_layout_leaf", sys.argv[1]) +module = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = module +spec.loader.exec_module(module) +assert module.PageLayout(32, {}, (), ()).tokens_per_page == 32 +unexpected = { + name.split(".")[0] for name in set(sys.modules) - before + if name != spec.name and name.split(".")[0] not in sys.stdlib_module_names +} +assert not unexpected, unexpected +""" + result = subprocess.run( + [sys.executable, "-I", "-S", "-c", script, page_layout.__file__], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py index e5f47987b30d..7714ef884fde 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py @@ -12,39 +12,45 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""``derive_layout`` against the page table native builds from the same live manager, for three -model shapes with small caches and no checkpoint: each transferred buffer's layer, pool, offset, -size and mapper, and each layer group's kind, window, heads, pools and view roles; and each -buffer's offset and size against the runtime's address and page size of its layer and role.""" +"""Compare live lender layouts with native page tables across model and parallel layouts.""" from __future__ import annotations import dataclasses from collections import defaultdict +from types import SimpleNamespace +from typing import Callable import pytest import torch +from utils.util import skip_pre_blackwell, skip_pre_hopper from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, MapperKind -from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import ( - BufferMapper, +from tensorrt_llm._torch.disaggregation.resource.page import CacheKind +from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import ( + MapperKind, ShardDesc, - derive_layout, + head_major, + replicated, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + MemoryKind, + StagingOptions, + attach_lender, ) from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType -from tensorrt_llm.mapping import Mapping +from tensorrt_llm.mapping import CpType, Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") -# Both dense shapes: bf16, a 4096-token pool, sequences of up to 2048 tokens. +# The dense shapes: bf16, a 4096-token pool, sequences of up to 2048 tokens. DENSE = dict(dtype=DataType.BF16, max_tokens=4096, max_seq_len=2048) -# TinyLlama-1.1B: 22 layers, 4 KV heads of 64. Rank 1 of TP2 holds heads 2 and 3. -TINYLLAMA_TP2_RANK1 = dict( - num_layers=22, num_kv_heads=4, head_dim=64, mapping=Mapping(world_size=2, rank=1, tp_size=2) -) +# TinyLlama-1.1B: 22 layers, 4 KV heads of 64. +TINYLLAMA = dict(num_layers=22, num_kv_heads=4, head_dim=64) +# Its heads on five layers, which two pipeline stages split unevenly. +FIVE_LAYERS = dict(num_layers=5, num_kv_heads=4, head_dim=64) # DeepSeek-V3-Lite: 30 MLA layers, one 576-wide latent per token (512 + 64 rope), K only. DSV3_LITE_MLA = dict( kv_cache_type=CacheType.SELFKONLY, @@ -53,15 +59,318 @@ head_dim=512 + 64, tokens_per_block=64, ) -# shape -> (keywords for conftest's ``real_manager``, or ``None`` for its ``hybrid_manager``: the -# Nemotron-Nano-9B-v2 shape with fewer layers; the shard every layer group names) -SHAPES = { - "tinyllama_tp2_rank1": ({**DENSE, **TINYLLAMA_TP2_RANK1}, ShardDesc(2, 1)), - "dsv3_lite_mla": ({**DENSE, **DSV3_LITE_MLA}, ShardDesc()), - "nemotron_h_hybrid": (None, ShardDesc()), + + +def tp(rank, size=2, **kwargs): + return Mapping(world_size=size, rank=rank, tp_size=size, **kwargs) + + +def pp2(rank): + return Mapping(world_size=2, rank=rank, pp_size=2) + + +def helix2(rank): + return Mapping(world_size=2, rank=rank, cp_size=2, cp_config={"cp_type": CpType.HELIX}) + + +def minimax_m3(kit, kv_dtype, implementation="msa"): + """A four-layer MiniMax-M3 manager: layers 0 to 2 dense, layer 3 sparse with a bf16 index-K.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3KVCacheManagerV2 + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + + return kit.managed( + lambda: MiniMaxM3KVCacheManagerV2( + kv_cache_config=KvCacheConfig( + max_tokens=1024, dtype=kv_dtype, enable_block_reuse=True, host_cache_size=0 + ), + kv_cache_type=CacheType.SELF, + num_layers=4, + num_kv_heads=2, + head_dim=128, + tokens_per_block=128, + max_seq_len=512, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1), + dtype={"nvfp4": DataType.NVFP4, "fp8": DataType.FP8}[kv_dtype], + vocab_size=32000, + max_num_tokens=512, + sparse_layer_ids=[3], + disable_index_value_layer_ids=[3], + sparse_index_dim=128, + sparse_attention_config=SimpleNamespace( + implementation=implementation, indexer_kv_dtype="bf16" + ), + ) + ) + + +def csa2(f): + """The six-layer CSA2 manager native's CSA2 transfer test builds, bounded replay off.""" + from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager + from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + + f.monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "0") + f.monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0") + layout = CSA2Layout((0, 2, 2, 1, 1, 1), (1, 3), (1, 3, 4), candidate_source_layer_id=3) + return f.kit.managed( + lambda: CSA2CacheManager( + KvCacheConfig(max_gpu_total_bytes=128 << 20, enable_block_reuse=False, dtype="fp8"), + CacheType.SELFKONLY, + num_layers=len(layout.compress_ratios), + tokens_per_block=128, + max_seq_len=512, + max_batch_size=2, + max_num_tokens=1024, + vocab_size=129280, + dtype=DataType.BF16, + mapping=Mapping(), + layout=layout, + is_disagg=True, + ) + ) + + +def deepseek_v4_mtp(f): + """The DeepSeek-V4 manager with one MTP layer, which it keeps as a virtual layer of its own.""" + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + spec_config = MTPDecodingConfig(max_draft_len=1) + spec_config.num_nextn_predict_layers = 1 # the checkpoint sets it at model load + return f.dsv4(spec_config=spec_config) + + +def kinds(layout): + return [desc.kind for desc in layout.layer_groups] + + +def mapper_of(b): + """The native mapper kind a buffer's tensors amount to.""" + if not b.split: + return "REPLICATED" + (tensor, *rest) = b.tensors + names = [a.name for a in tensor.axes] + if "shard" in names: + return "INDEXED" # a rank's own bytes, its heads undeclared + if rest or "head" not in names: + return "SECTIONED" + return "NHD" if names[0] == "token" else "INDEXED" + + +def sections_of(b): + """A sectioned buffer's section sizes on this rank, as native's ``section_bytes``; None + otherwise.""" + if mapper_of(b) != "SECTIONED": + return None + return tuple(t.axes[0].held for t in b.tensors) + + +def head_bytes_of(b): + """Bytes of one head on this rank of a buffer split by head, as native's ``bytes_per_head``; + None otherwise.""" + head = next((a for a in b.tensors[0].axes if a.name == "head"), None) + if mapper_of(b) not in ("INDEXED", "NHD") or head is None: + return None + return b.size // head.held + + +def roles(layout): + return {(b.role, mapper_of(b)) for b in layout.buffers} + + +def check_tp2(mgr, layout): + heads = set(mgr.num_kv_heads_per_layer) + assert (mgr.mapping.tp_size, mgr.mapping.tp_rank, heads) == (2, 1, {2}), (mgr.mapping, heads) + + +def check_tp4_replicated(mgr, layout): + assert (mgr.mapping.tp_size, mgr.mapping.tp_rank) == (4, 3) + assert set(mgr.num_kv_heads_per_layer) == {1} + + +def check_attention_dp(mgr, layout): + heads = set(mgr.num_kv_heads_per_layer) + assert mgr.mapping.enable_attention_dp and heads == {4}, (mgr.mapping, heads) + + +def check_stage(rank): + def check(mgr, layout): + held = len(mgr.pp_layers) + assert mgr.mapping.pp_rank == rank and 0 < held < 5, (mgr.mapping, mgr.pp_layers) + + return check + + +def check_vswa(mgr, layout): + windows = sorted(desc.kept_window or 0 for desc in layout.layer_groups) + assert windows == [0, 64], f"layer group windows {windows}, not one full and one of 64" + + +def check_mla(tp_size): + def check(mgr, layout): + heads = set(mgr.num_kv_heads_per_layer) + assert (mgr.kv_factor, mgr.mapping.tp_size, heads) == (1, tp_size, {1}), mgr.mapping + + return check + + +def check_helix(mgr, layout): + assert mgr._has_cp_helix and (mgr._helix_cp_rank, mgr._helix_cp_size) == (1, 2), mgr.mapping + # A page holds tokens_per_block tokens of each head's K and V, two bytes each (fp16). + per_head = {head_bytes_of(b) for b in layout.buffers} + assert per_head == {layout.tokens_per_page * mgr.head_dim * 2}, (per_head, mgr.head_dim) + + +def check_beams(mgr, layout): + assert mgr.max_beam_width == 2, mgr.max_beam_width + + +def check_hybrid(tp_size): + def check(mgr, layout): + assert mgr.mapping.tp_size == tp_size, mgr.mapping + assert set(kinds(layout)) == {"attention", "state"}, kinds(layout) + + return check + + +def check_deepseek_v4(mgr, layout): + widths = {g: len(w) for g, w in layout.pool_slot_bytes.items()} + assert mgr._layer_attn_to_layer_id and max(widths.values()) >= 2, widths + + +def check_deepseek_v4_mtp(mgr, layout): + check_deepseek_v4(mgr, layout) + model_layers = {model_layer for model_layer, _ in mgr._layer_attn_to_layer_id} + assert mgr.max_draft_len == 1 and max(model_layers) == 4, (mgr.max_draft_len, model_layers) + + +def check_csa2(mgr, layout): + assert mgr._physical_roles and mgr._decoder_replay_window == 0, mgr._decoder_replay_window + assert any(desc.kept_window is not None for desc in layout.layer_groups), kinds(layout) + + +def check_gdn(tp_size, stage_kinds): + def check(mgr, layout): + assert mgr.mapping.tp_size == tp_size, mgr.mapping + assert set(kinds(layout)) == stage_kinds, kinds(layout) + sections = {sections_of(b) for b in layout.buffers if mapper_of(b) == "SECTIONED"} + assert sections and all(len(s) == 3 and s[0] == s[1] < s[2] for s in sections), sections + + return check + + +def check_minimax_m3(scales, *, nhd=False): + def check(mgr, layout): + held = roles(layout) + if nhd: + assert any(mapper == "NHD" for _, mapper in held), sorted(held) + assert ("index_key", "REPLICATED") in held, sorted(held) + assert (("key_block_scale", "INDEXED") in held) == scales, sorted(held) + + return check + + +@dataclasses.dataclass(frozen=True) +class Case: + """One rank's manager, what it must be, and the shard, heads and quantization it names.""" + + build: Callable + check: Callable + shard: ShardDesc = ShardDesc() + total_heads: int | None = None + quant: str | None = None + + +CASES = { + "tinyllama_tp2_rank1": Case( + lambda f: f.real(**DENSE, **TINYLLAMA, mapping=tp(1)), check_tp2, ShardDesc(2, 1), 4 + ), + "tp4_two_heads_rank3": Case( + lambda f: f.real(**{**DENSE, **TINYLLAMA, "num_kv_heads": 2}, mapping=tp(3, size=4)), + check_tp4_replicated, + ShardDesc(2, 1), + 2, + ), + "minimax_m3_triton_fp8": Case( + lambda f: minimax_m3(f.kit, "fp8", implementation="triton"), + check_minimax_m3(scales=False, nhd=True), + total_heads=2, + quant="fp8", + ), + "attention_dp_rank1": Case( + lambda f: f.real(**DENSE, **TINYLLAMA, mapping=tp(1, enable_attention_dp=True)), + check_attention_dp, + total_heads=4, + ), + "pp2_stage0": Case( + lambda f: f.real(**DENSE, **FIVE_LAYERS, mapping=pp2(0)), check_stage(0), total_heads=4 + ), + "pp2_stage1": Case( + lambda f: f.real(**DENSE, **FIVE_LAYERS, mapping=pp2(1)), check_stage(1), total_heads=4 + ), + "vswa": Case( + lambda f: f.real(**DENSE, **FIVE_LAYERS, windows=[64, DENSE["max_seq_len"]]), + check_vswa, + total_heads=4, + ), + "dsv3_lite_mla": Case(lambda f: f.real(**DENSE, **DSV3_LITE_MLA), check_mla(1), total_heads=1), + # One latent head, which every tensor-parallel rank holds whole. + "dsv3_lite_mla_tp2_rank1": Case( + lambda f: f.real(**DENSE, **DSV3_LITE_MLA, mapping=tp(1)), check_mla(2), total_heads=1 + ), + "helix_rank1": Case( + lambda f: f.real(mapping=helix2(1), enable_block_reuse=False), check_helix, total_heads=4 + ), + "beam2": Case( + lambda f: f.real(**DENSE, **TINYLLAMA, max_beam_width=2), check_beams, total_heads=4 + ), + "nemotron_h_hybrid": Case(lambda f: f.hybrid(), check_hybrid(1), total_heads=8), + "nemotron_h_hybrid_tp2_rank1": Case( + lambda f: f.hybrid(mapping=tp(1)), check_hybrid(2), ShardDesc(2, 1), 8 + ), + "deepseek_v4": Case(lambda f: f.dsv4(), check_deepseek_v4, quant="fp8"), + "deepseek_v4_mtp": Case(deepseek_v4_mtp, check_deepseek_v4_mtp, quant="fp8"), + "csa2_full_prefill": Case(csa2, check_csa2), + "qwen3_next_gdn": Case(lambda f: f.gdn(), check_gdn(1, {"attention", "state"}), total_heads=2), + "qwen3_next_gdn_tp2_rank1": Case( + lambda f: f.gdn(mapping=tp(1)), check_gdn(2, {"attention", "state"}), ShardDesc(2, 1), 2 + ), + # Layers 0 and 1 of "MMM*": a stage holding recurrent state alone. + "qwen3_next_gdn_pp2_stage0": Case( + lambda f: f.gdn(mapping=pp2(0)), check_gdn(1, {"state"}), total_heads=2 + ), + "qwen3_next_gdn_pp2_stage1": Case( + lambda f: f.gdn(mapping=pp2(1)), check_gdn(1, {"attention", "state"}), total_heads=2 + ), + "minimax_m3_fp8": Case( + lambda f: minimax_m3(f.kit, "fp8"), + check_minimax_m3(scales=False), + total_heads=2, + quant="fp8", + ), + "minimax_m3_nvfp4": Case( + lambda f: minimax_m3(f.kit, "nvfp4"), + check_minimax_m3(scales=True), + total_heads=2, + quant="nvfp4", + ), +} +MARKS = { + "deepseek_v4": [skip_pre_blackwell], + "deepseek_v4_mtp": [skip_pre_blackwell], + "csa2_full_prefill": [skip_pre_blackwell], + "minimax_m3_fp8": [skip_pre_hopper], + "minimax_m3_triton_fp8": [skip_pre_hopper], + "minimax_m3_nvfp4": [skip_pre_hopper], } +def attach(mgr): + return attach_lender( + mgr, scope=b"page-table", staging=StagingOptions(int(mgr.tokens_per_block)) + ) + + def _as_list(values): return None if values is None else [int(v) for v in values] @@ -102,40 +411,38 @@ def native_groups(page_table): return out -def layout_groups(geometry): - """The same per layer group from ``derive_layout``'s result, without local-only buffers (native - skips ignored roles); transferred buffers no group claims are listed last. Attention heads are - the head-sharded buffers' ``num_heads``: native's attention views have no per-head size.""" - layout = geometry.layout +def public_groups(layout, regions): + """Project public transferred buffers into native's group, pool and role representation.""" out, claimed = [], set() for desc in layout.layer_groups: g = desc.pool_group - pools = [(p.base, p.slot_bytes, p.num_slots) for p in geometry.device_pools[g]] - assert layout.pool_groups[g] == tuple(w for _, w, _ in pools) - base_of = {p.index: p.base for p in geometry.device_pools[g]} + pools = [] + for p in range(len(layout.pool_slot_bytes.get(g, ()))): + region = regions[layout.device_region_index[(g, p)]] + pools.append((region.address, region.slot_bytes, region.slot_count)) state = desc.kind == "state" rows, roles, geometries, heads = [], defaultdict(list), defaultdict(set), set() + whole = set() for i, b in enumerate(layout.buffers): - if b.layer not in desc.layers or b.pool[0] != g or not b.transfer: + if b.layer not in desc.layers or b.pool_group != g or not b.transferred: continue - assert b.expansion == 1, "native reads no expansion" claimed.add(i) - key = (base_of[b.pool[1]], b.mapper.name) - rows.append((b.layer, key[0], b.offset, b.size, key[1])) + key = (pools[b.pool_index][0], mapper_of(b)) + rows.append((b.layer, key[0], b.slot_offset, b.size, key[1])) roles[key].append(b.role) if state: - section = b.geometry.section_bytes - section = None if section is None else tuple(section) - geometries[key].add((section, b.geometry.bytes_per_head)) - elif b.mapper in (BufferMapper.INDEXED, BufferMapper.NHD): - heads.add(b.geometry.num_heads) + geometries[key].add((sections_of(b), head_bytes_of(b))) + elif key[1] in ("INDEXED", "NHD"): + heads.add(b.size // head_bytes_of(b)) + elif key[1] == "REPLICATED": + whole.add(layout.total_kv_heads_by_layer.get(b.layer)) + heads = heads or whole assert all(len(v) == 1 for v in geometries.values()), f"roles disagree: {dict(geometries)}" - assert len(heads) <= 1, heads out.append( { "kind": desc.kind, - "window": desc.window, - "heads": None if state else (heads.pop() if heads else 0), + "window": desc.kept_window, + "heads": None if state else (heads.pop() if len(heads) == 1 else sorted(heads)), "layers": sorted(desc.layers), "pools": pools, "rows": sorted(rows), @@ -144,17 +451,33 @@ def layout_groups(geometry): } ) stray = [ - (b.layer, b.role) for i, b in enumerate(layout.buffers) if b.transfer and i not in claimed + (b.layer, b.role) + for i, b in enumerate(layout.buffers) + if b.transferred and i not in claimed ] if stray: out.append({"unclaimed buffers": stray}) return out -def lies(geometry): - """Layout fields that each misstate one thing native reads.""" - layout = geometry.layout - first = next(b for b in layout.buffers if b.transfer) +# RankTopology and Mapping name the same rank coordinates. +PLACE = ("tp_size", "tp_rank", "pp_size", "pp_rank", "cp_size", "cp_rank", "dp_size", "dp_rank") + + +def topology(layout): + return tuple(getattr(layout.topology, name) for name in PLACE) + (layout.topology.attention_dp,) + + +def manager_topology(mgr): + mapping = mgr.mapping + coordinates = tuple(getattr(mapping, name) for name in PLACE[:6]) + dp_rank = mapping.tp_rank if mapping.enable_attention_dp else 0 + return coordinates + (mapping.dp_size, dp_rank, mapping.enable_attention_dp) + + +def lies(layout): + """Public layouts that each misstate one thing native reads.""" + first = next(b for b in layout.buffers if b.transferred) def buffers(**change): return tuple(dataclasses.replace(b, **change) if b is first else b for b in layout.buffers) @@ -166,51 +489,82 @@ def first_group(**change): ) other = ( - BufferMapper.INDEXED if first.mapper is BufferMapper.REPLICATED else BufferMapper.REPLICATED + head_major(first.size, 1, 1) if mapper_of(first) == "REPLICATED" else replicated(first.size) ) + window = 128 if layout.layer_groups[0].kept_window == 64 else 64 out = { - "offset": {"buffers": buffers(offset=first.offset + first.size)}, - "layer": {"buffers": buffers(layer=first.layer + 1000)}, - "mapper": {"buffers": buffers(mapper=other)}, - "extra buffer": { - "buffers": layout.buffers + (dataclasses.replace(first, layer=first.layer + 1000),) - }, - "window": {"layer_groups": first_group(window=64)}, + "offset": dict(buffers=buffers(slot_offset=first.slot_offset + first.size)), + "layer": dict(buffers=buffers(layer=first.layer + 1000)), + "mapper": dict(buffers=buffers(tensors=other)), + "extra buffer": dict( + buffers=layout.buffers + (dataclasses.replace(first, layer=first.layer + 1000),) + ), + # The read window follows the misstated kept window, which it may not exceed. + "window": dict(layer_groups=first_group(kept_window=window, read_window=None)), } # The first layer group pointed at another device pool group, when the layout has one. mine = layout.layer_groups[0].pool_group - others = sorted(g for g in layout.pool_groups if g != mine) + others = sorted(g for g in layout.pool_slot_bytes if g != mine) if others: - out["pool group"] = {"layer_groups": first_group(pool_group=others[0])} - return out + out["pool group"] = dict(layer_groups=first_group(pool_group=others[0])) + # Two pools' regions traded, when there are two. + pools = sorted(layout.device_region_index) + if len(pools) >= 2: + traded = dict(layout.device_region_index) + traded[pools[0]], traded[pools[1]] = traded[pools[1]], traded[pools[0]] + out["regions"] = dict(device_region_index=traded) + return {name: dataclasses.replace(layout, **change) for name, change in out.items()} -@pytest.mark.parametrize("shape", list(SHAPES)) -def test_layout_matches_the_page_table_native_builds(real_manager, hybrid_manager, shape): - kwargs, shard = SHAPES[shape] - with real_manager(**kwargs) if kwargs else hybrid_manager() as mgr: - geometry = derive_layout(mgr) +@pytest.mark.parametrize("case", [pytest.param(name, marks=MARKS.get(name, [])) for name in CASES]) +def test_the_public_layout_matches_what_native_builds( + case, kit, real_manager, hybrid_manager, gdn_manager, deepseek_v4_manager, monkeypatch +): + spec = CASES[case] + fixtures = SimpleNamespace( + real=real_manager, + hybrid=hybrid_manager, + gdn=gdn_manager, + dsv4=deepseek_v4_manager, + kit=kit, + monkeypatch=monkeypatch, + ) + with spec.build(fixtures) as mgr: + lender = attach(mgr) + layout, regions = lender.page_layout, lender.regions + spec.check(mgr, layout) page_table = build_page_table_from_manager(mgr) native = native_groups(page_table) - derived = layout_groups(geometry) + derived = public_groups(layout, regions) - layout = geometry.layout - assert layout.tokens_per_block == page_table.tokens_per_block + # Each pool's region is a device region of its slot width. + for (g, p), position in layout.device_region_index.items(): + region = regions[position] + assert region.memory_kind is MemoryKind.DEVICE, (g, p, region) + assert region.slot_bytes == layout.pool_slot_bytes[g][p], (g, p, region) + assert layout.tokens_per_page * layout.topology.cp_size == page_table.tokens_per_block assert len(derived) == len(native), derived[len(native) :] assert all(group["rows"] for group in native) for lg, (ours, theirs) in enumerate(zip(derived, native)): for key in theirs: assert ours[key] == theirs[key], f"layer group {lg}: {key}" - assert [d.shard for d in layout.layer_groups] == [shard] * len(native) - assert geometry.shards == ((shard.count, shard.index),) * len(native) - # The lender finds a group's pools through either; they must agree. - assert tuple(d.pool_group for d in layout.layer_groups) == geometry.pool_group_of - assert geometry.recurrent == tuple(g["kind"] == "state" for g in native) + assert [d.shard for d in layout.layer_groups] == [spec.shard] * len(native) + assert all(b.expansion == 1 for b in layout.buffers if b.transferred), "native reads none" + assert topology(layout) == manager_topology(mgr), (layout.topology, mgr.mapping) + assert layout.dtype == mgr.dtype.name + assert layout.quant_mode == spec.quant, layout.quant_mode + if spec.total_heads is not None: + expected = { + int(layer.global_layer_id): spec.total_heads + for group in page_table.layer_groups + if group.kind != CacheKind.STATE + for layer in group.local_layers + } + assert dict(layout.total_kv_heads_by_layer) == expected # The comparison is not vacuous: each misstated field shows. - for name, change in lies(geometry).items(): - lie = dataclasses.replace(geometry, layout=dataclasses.replace(layout, **change)) - assert layout_groups(lie) != native, f"the {name} lie went unnoticed" + for name, lie in lies(layout).items(): + assert public_groups(lie, regions) != native, f"the {name} lie went unnoticed" # Every buffer sits where the runtime keeps its layer and role. Native's views list roles # per pool and mapper, so a layer's K and V trading places shows only here. @@ -219,9 +573,9 @@ def test_layout_matches_the_page_table_native_builds(real_manager, hybrid_manage for lg in page_table.layer_groups for ll in lg.local_layers } - base = {(p.group, p.index): p.base for ps in geometry.device_pools.values() for p in ps} for b in layout.buffers: lid = internal[b.layer] address = int(mgr.impl.get_mem_pool_base_address(lid, b.role, PageIndexMode.SHARED)) + base = regions[layout.device_region_index[(b.pool_group, b.pool_index)]].address page = int(mgr.impl.get_page_stride(lid, b.role)) * b.expansion - assert (address - base[b.pool], page) == (b.offset, b.size), (b.layer, b.role) + assert (address - base, page) == (b.slot_offset, b.size), (b.layer, b.role) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_lender_fast_paths.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_lender_fast_paths.py new file mode 100644 index 000000000000..48bdca5a0a85 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_lender_fast_paths.py @@ -0,0 +1,2293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The lender's fast paths against eager references and against its rules alone, on stand-ins.""" + +import ast +import copy +import numbers +import random +import struct +import textwrap +import traceback +import weakref +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + ALL_ARRIVED, + MemoryKind, + RowSelection, + _fetch_ledger, + _lender, + _manager, + _ranges, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import Identity +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import DevicePool +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._page_locks import LockMode +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._slots import SlotQueue +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._types import LentLayerGroupRows + +pytestmark = pytest.mark.cpu_only +# The sources whose methods the eager-reference differentials run, read by path. +_SHARING = Path(_lender.__file__).resolve().parent +TPB = 8 +RID = 7 +# Slots per device pool: past every page a stand-in cache draws. +SLOTS = 64 +DEVICE, ALL = MemoryKind.DEVICE, RowSelection.ALL + + +def _load(file, qualified, ns): + source = (_SHARING / file).read_text() + node = ast.parse(source) + for part in qualified.split("."): + node = next(child for child in node.body if getattr(child, "name", None) == part) + first = min([node.lineno] + [d.lineno for d in node.decorator_list]) + code = textwrap.dedent("\n".join(source.splitlines()[first - 1 : node.end_lineno])) + exec(compile("from __future__ import annotations\n" + code, str(_SHARING / file), "exec"), ns) + return ns[node.name] + + +def _sources(immutable_mutant=None): + manager = SimpleNamespace( + # A cache's committed tokens: none where no layer group committed any. + cache_state=lambda kv: SimpleNamespace( + history=kv.history, + active=kv.active, + committed=max(getattr(kv, "committed", ()), default=0), + ), + committed_by_group=lambda kv, count: kv.committed, + pages=lambda kv, lg: kv.pages[lg], + locked_pages=lambda kv, lg: kv.pages[lg], + state_slot=lambda kv, lg: kv.state_slots[lg], + num_blocks=lambda kv: kv.num_blocks, + stale_range=lambda mgr, lg, history: ( + mgr.stale_at(lg, history) if hasattr(mgr, "stale_at") else mgr.stale[lg] + ), + context_end=lambda request, mgr: request.prompt, + context_computed=lambda request: request.context_current_position, + bidirectional_runs=lambda request: [], + _MAX_TOKENS=2**31 - 1, + ) + ns = dict(__name__=__name__, np=np, struct=struct, numbers=numbers, traceback=traceback) + ns.update(dataclass=dataclass, field=field, _manager=manager, REUSE_KEY_BYTES=32) + ns.update(_OBJECT_KEY_BYTES=54, OBJECT_KEY_BYTES=54, logger=SimpleNamespace(error=pytest.fail)) + for name in ("_immutable", "_frozen_int64", "_integer", "_entry_integers", "_region_indices"): + _load("_types.py", name, ns) + if immutable_mutant == "identity": + ns["_immutable"] = lambda array: array + elif immutable_mutant == "readonly_flag": + good = ns["_immutable"] + ns["_immutable"] = lambda array: array if not array.flags.writeable else good(array) + for name in ("LentLayerGroupRows", "LeaseView"): + _load("_types.py", name, ns) + _load("_lender.py", "_LentBlockEntries", ns) + _load("_lender.py", "_frozen_array", ns) + _load("_lender.py", "_unshared", ns) + ns["_NO_ROWS"] = ns["_frozen_int64"](()) + ns["_STATE_ROW"] = ns["_frozen_int64"]((0,)) + ranges = SimpleNamespace() + for name in ("_stale_range", "_needed_ranges", "_last_history_keeping"): + setattr(ranges, name, _load("_ranges.py", name, ns)) + ns["_ranges"] = ranges + _load("_identity.py", "Identity.object_keys", ns) + result = SimpleNamespace( + object_keys=ns["object_keys"], + ns=ns, + manager=manager, + ranges=ranges, + entries=ns["_LentBlockEntries"], + rows=ns["LentLayerGroupRows"], + view=ns["LeaseView"], + ) + _load("_fetch_ledger.py", "_writable_mask", ns) + ledger_methods = """_unheld_block _computed_tokens _known _prefix_kept_by_fetch + _invalidate_delivered_blocks_past_capacity""" + for name in ledger_methods.split(): + setattr(result, name, _load("_fetch_ledger.py", f"_FetchLedger.{name}", ns)) + lender_methods = """_locked_block_masks _history_cap _without_short_windows_mid_context + _mid_context_end + _range_ordinals _block_entries_of_range _lendable_block_entries _committed_entries + _build_view _in_place_view _on_free _on_reset _on_shutdown""" + for name in lender_methods.split(): + setattr(result, name, _load("_lender.py", f"_ManagerLender.{name}", ns)) + return result + + +@pytest.fixture(scope="module") +def sources(): + return _sources() + + +def _bind(obj, sources, *names): + for name in names: + setattr(obj, name, getattr(sources, name).__get__(obj)) + + +def _outcome(call): + try: + return "ok", call() + except Exception as error: # noqa: BLE001 - compare complete exception outcomes + return "raise", type(error), str(error) + + +def _root(array): + """The object at the base of ``array``'s chain of views.""" + while isinstance(array, np.ndarray) and array.base is not None: + array = array.base + return array + + +def _slow_unheld(ledger, rid, kv, entries, top): + """The eager path: STATE and mixed tails, then computed prefixes, then delivered rows.""" + delivered = ledger._deliveries.get(rid) + if delivered is not None and delivered.kv is not kv: + delivered = None + tpb = ledger._layout.tokens_per_block + for lg, ordinals in zip(entries.layer_group_indices, entries.block_ordinals): + if not len(ordinals): + continue + if lg in ledger._recurrent_groups: + if delivered is not None and delivered.state_end != top: + return lg, 0 + elif ledger._mixed_tail_caches.get(rid) is kv: + tail, partial = divmod(kv.committed[lg], tpb) + if partial and tail in ordinals: + return lg, tail + computed = ledger._computed_tokens(rid, kv) + if ledger._known(computed) >= top: + return None + rank, size = ledger._layout.cp + page = tpb // size + for lg, ordinals in zip(entries.layer_group_indices, entries.block_ordinals): + if lg in ledger._recurrent_groups: + continue + for block in ordinals.tolist(): + first = block * tpb + rank * page + if first >= top: + continue + valid = computed[lg] + if delivered is not None: + mask = delivered.delivered_block_masks[lg] + if block < len(mask) and mask[block]: + valid = max(valid, delivered.partial_ends.get((lg, block), first + page)) + if valid < min(first + page, top): + return lg, block + return None + + +def _ledger(sources, layout, kv, *, advanced=None, delivered=None, tainted=False): + ledger = SimpleNamespace( + _layout=layout, + _recurrent_groups=frozenset( + lg for lg, recurrent in enumerate(layout.recurrent) if recurrent + ), + _paged_groups=tuple(lg for lg, recurrent in enumerate(layout.recurrent) if not recurrent), + _kept_prefixes={} if advanced is None else {7: advanced}, + _deliveries={} if delivered is None else {7: delivered}, + _mixed_tail_caches={7: kv} if tainted else {}, + ) + _bind(ledger, sources, "_known", "_prefix_kept_by_fetch", "_computed_tokens") + return ledger + + +def _delivery(kv, masks, top=0): + return SimpleNamespace( + kv=kv, + delivered_block_masks=masks, + state_end=top, + end=top, + partial_ends={}, + resume_bounds_memo=("memo",), + ) + + +def _delivery_key(delivered): + return ( + tuple(mask.tobytes() for mask in delivered.delivered_block_masks), + delivered.end, + tuple(sorted(delivered.partial_ends.items())), + delivered.resume_bounds_memo, + ) + + +def test_unheld_block_and_ledger_shortcuts_match_eager_paths(sources): + rng = random.Random(20261016) + for _ in range(600): + count, tpb = rng.randint(1, 3), rng.choice((4, 8)) + size = rng.choice((1, 2)) + layout = SimpleNamespace( + num_layer_groups=count, + tokens_per_block=tpb, + recurrent=tuple(rng.random() < 0.25 for _ in range(count)), + cp=(rng.randrange(size), size), + ) + kv = SimpleNamespace( + history=rng.randint(0, 64), + active=True, + committed=tuple(rng.randint(0, 48) for _ in range(count)), + num_blocks=rng.randint(0, 10), + ) + other = copy.copy(kv) + advanced = rng.choice( + ( + None, + (kv, (0,) * count, kv.history), + (kv, (0,) * count, rng.randint(0, 64)), + (other, (0,) * count, kv.history), + ) + ) + delivered = _delivery( + rng.choice((kv, kv, other)), + [ + np.array([rng.random() < 0.7 for _ in range(rng.randint(0, 14))], dtype=bool) + for _ in range(count) + ], + rng.randint(0, 72), + ) + delivered.partial_ends = {(rng.randrange(count), rng.randrange(14)): rng.randrange(72)} + entries = sources.entries( + list(range(count)), + [ + np.array(rng.sample(range(14), rng.randint(0, 7)), dtype=np.int64) + for _ in range(count) + ], + [], + ) + top, tainted = rng.randint(0, 72), rng.random() < 0.25 + ledger = _ledger( + sources, layout, kv, advanced=advanced, delivered=delivered, tainted=tainted + ) + before = _delivery_key(delivered) + kept_before = dict(ledger._kept_prefixes) + tainted_before = dict(ledger._mixed_tail_caches) + entries_before = tuple(array.tobytes() for array in entries.block_ordinals) + old_result = _outcome(lambda: _slow_unheld(ledger, 7, kv, entries, top)) + assert old_result == _outcome(lambda: sources._unheld_block(ledger, 7, kv, entries, top)) + assert old_result == _outcome( + lambda: sources._unheld_block(ledger, 7, kv, entries, top, kv.history) + ) + assert ledger._kept_prefixes == kept_before and ledger._mixed_tail_caches == tainted_before + assert tuple(array.tobytes() for array in entries.block_ordinals) == entries_before + assert _delivery_key(delivered) == before + start = rng.choice((0, 0, rng.randint(0, 64))) + assert ledger._prefix_kept_by_fetch(7, kv, start) == tuple( + min(tokens, start) for tokens in ledger._computed_tokens(7, kv) + ) + old, new = copy.deepcopy(delivered), copy.deepcopy(delivered) + kept = kv.num_blocks + old.end = min(old.end, kept * tpb) + old.partial_ends = {key: end for key, end in old.partial_ends.items() if key[1] < kept} + if old.end != delivered.end: + old.resume_bounds_memo = None + for mask in old.delivered_block_masks: + if mask[kept:].any(): + mask[kept:] = False + old.resume_bounds_memo = None + sources._invalidate_delivered_blocks_past_capacity(ledger, new, kv) + assert _delivery_key(new) == _delivery_key(old) + sources._invalidate_delivered_blocks_past_capacity(ledger, new, kv) + assert _delivery_key(new) == _delivery_key(old) + + +@pytest.mark.parametrize( + "case", ["before_fetch_end", "at_fetch_end", "after_fetch_end", "missing_state", "mixed_tail"] +) +def test_unheld_block_guards_history_state_and_mixed_tail(sources, case): + """A read refuses the first unreadable block across history and fetch boundaries.""" + kv = SimpleNamespace(history=16, active=True, committed=(5, 0)) + delivered = _delivery(kv, [np.array([True, True, False, True]), np.ones(1, bool)], 8) + lg, blocks = (1, [0]) if case == "missing_state" else (0, [1] if case == "mixed_tail" else [2]) + entries = sources.entries([lg], [np.array(blocks, dtype=np.int64)], []) + ledger = _ledger( + sources, + layout([False, True], tpb=4), + kv, + delivered=delivered, + tainted=case == "mixed_tail", + advanced=None if case in ("missing_state", "mixed_tail") else (kv, (0, 0), 16), + ) + kv.history = {"before_fetch_end": 15, "after_fetch_end": 17}.get(case, 16) + expected = { + "before_fetch_end": (0, 2), + "at_fetch_end": (0, 2), + "after_fetch_end": None, + "missing_state": (1, 0), + "mixed_tail": (0, 1), + }[case] + assert _slow_unheld(ledger, 7, kv, entries, 16) == expected + assert sources._unheld_block(ledger, 7, kv, entries, 16) == expected + + +def test_taint_from_a_replaced_cache_does_not_block_the_current_cache(sources): + kv = SimpleNamespace(history=8, active=True, committed=(5,)) + ledger = _ledger(sources, layout([False], tpb=4), kv) + ledger._mixed_tail_caches[7] = object() + entries = sources.entries([0], [np.array([1], dtype=np.int64)], []) + assert _slow_unheld(ledger, 7, kv, entries, 8) is None + assert sources._unheld_block(ledger, 7, kv, entries, 8) is None + + +def _slow_masks(lender, kv, lease): + masks = [] + for lg, ordinals, slots in zip( + lease._block_entries.layer_group_indices, + lease._block_entries.block_ordinals, + lease._block_entries.device_pages, + ): + same = np.zeros(len(ordinals), dtype=bool) + if kv is not None and kv is lease._kv and kv.active: + if lg in lender._recurrent_groups: + same[:] = slots == kv.state_slots[lg] + else: + pages = kv.pages[lg] + inside = ordinals < len(pages) + same[inside] = pages[ordinals[inside]] == slots[inside] + begin, end = ( + lender._manager_ref().stale[lg] + if lender._layout.windows[lg] is not None + else (0, 0) + ) + same &= (ordinals < begin) | (ordinals >= end) + masks.append(same) + return masks + + +def test_locked_masks_match_eager_paths(sources): + rng = random.Random(18) + for _ in range(600): + groups = rng.randint(1, 3) + windows = [rng.choice((None, 8, 16)) for _ in range(groups)] + manager = SimpleNamespace( + stale=[tuple(sorted((rng.randrange(4), rng.randrange(8)))) for _ in range(groups)] + ) + layout = SimpleNamespace(windows=windows) + lender = SimpleNamespace( + _layout=layout, + _manager_ref=lambda: manager, + _recurrent_groups=frozenset(lg for lg in range(groups) if rng.random() < 0.2), + _state_slot=lambda kv, request_id, lg: kv.state_slots[lg], + ) + kv = SimpleNamespace( + active=rng.random() < 0.8, + history=rng.randrange(64), + pages=[ + np.array([rng.randrange(20) for _ in range(rng.randrange(12))], dtype=np.int64) + for _ in range(groups) + ], + state_slots=[rng.randrange(20) for _ in range(groups)], + ) + ordinals = [ + np.array(rng.sample(range(15), rng.randrange(8)), dtype=np.int64) for _ in range(groups) + ] + entries = sources.entries( + list(range(groups)), + ordinals, + [np.array([rng.randrange(20) for _ in blocks], dtype=np.int64) for blocks in ordinals], + ) + lease = SimpleNamespace( + _kv=kv if rng.random() < 0.85 else object(), _block_entries=entries, _request_id=7 + ) + current = kv if rng.random() < 0.9 else None + before = [a.tobytes() for a in entries.block_ordinals + entries.device_pages] + expected = _slow_masks(lender, current, lease) + actual = sources._locked_block_masks(lender, current, lease) + assert [a.tobytes() for a in actual] == [a.tobytes() for a in expected] + assert before == [a.tobytes() for a in entries.block_ordinals + entries.device_pages] + + +@pytest.mark.parametrize( + "blocks", [[6, 1, 0], [0, 2, 6], [5, 0, 9, 3, 1]], ids=["down", "up", "mixed"] +) +def test_ordinals_past_the_cache_s_end_in_any_order_get_and_lock_no_page(sources, blocks): + manager = SimpleNamespace(stale=[(1, 3)]) + lender = _metadata_lender(sources, layout([False], windows=(8,)), manager) + kv = SimpleNamespace(active=True, history=16, pages=[np.array([10, 11, 12, 13])]) + ordinals = [np.array(blocks, dtype=np.int64)] + supplied = sources._block_entries_of_range(lender, manager, kv, 0, 16, ordinals=ordinals) + assert supplied.device_pages[0].tolist() == [10 + b if b < 4 else -1 for b in blocks] + slots = [np.array([10 + b if b < 4 else 10 for b in blocks])] + lease = SimpleNamespace(_kv=kv, _block_entries=sources.entries([0], ordinals, slots)) + expected = [[b < 4 and not 1 <= b < 3 for b in blocks]] + assert [a.tolist() for a in _slow_masks(lender, kv, lease)] == expected + assert [a.tolist() for a in sources._locked_block_masks(lender, kv, lease)] == expected + + +def _bytes_key(array): + return None if array is None else (array.dtype.str, array.shape, array.tobytes()) + + +def _entry_key(entries): + return ( + tuple(entries.layer_group_indices), + tuple(map(_bytes_key, entries.block_ordinals)), + tuple(map(_bytes_key, entries.device_pages)), + ) + + +def _assert_built_frozen(out, incoming): + """Every array of ``out`` other than ``incoming``'s own lies on immutable bytes, as the locks a + lend asks for take its ordinals.""" + given = {id(array) for array in (*incoming.block_ordinals, *incoming.device_pages)} + for array in (*out.block_ordinals, *out.device_pages): + if id(array) not in given: + root = _root(array) + assert not array.flags.writeable and isinstance(root, bytes), type(root).__name__ + + +def _view_key(view): + return tuple( + ( + g.layer_group_index, + _bytes_key(g.page_ordinals), + _bytes_key(g.object_keys), + _bytes_key(g.slots), + g.region_indices, + ) + for g in view.groups + ) + + +def _assert_frozen_view(view): + for group in view.groups: + for array in (group.page_ordinals, group.object_keys, group.slots): + if array is not None: + assert not array.flags.writeable + with pytest.raises(ValueError): + array.flags.writeable = True + + +def _metadata_lender(sources, layout, manager): + lender = SimpleNamespace( + _layout=layout, + _ordinal_ranges={}, + _recurrent_groups=frozenset(lg for lg, flag in enumerate(layout.recurrent) if flag), + _named_groups=layout.named_groups, + _manager_ref=lambda: manager, + _any_window=any(window is not None for window in layout.windows), + _regions=[ + SimpleNamespace(address=0x10000 + i * 0x10000, slot_bytes=128 << i) for i in range(2) + ], + _region_index_of_pool_group={0: 0, 1: 1}, + _device_regions_of_pool_group={0: (2, 3), 1: (4, 5)}, + _staging_region_indices={0: (0,), 1: (1,)}, + # A manager without a record of state slots: the cache answers. + _state_slot=lambda kv, request_id, lg: kv.state_slots[lg], + ) + lender._object_keys = lambda lg, keys: np.pad(keys, ((0, 0), (0, 22)), constant_values=lg) + _bind(lender, sources, "_range_ordinals", "_mid_context_end") + return lender + + +def _slow_entries(sources, lender, manager, kv, starts, end, named): + """The eager path: ordinal arrays made anew, device pages filled through a bounds mask.""" + layout, tpb = lender._layout, lender._layout.tokens_per_block + out = sources.entries([], [], []) + rank, size = layout.cp + for lg, start in enumerate(starts): + if lg in lender._recurrent_groups: + blocks, slots = np.zeros(1, dtype=np.int64), np.array([kv.state_slots[lg]], np.int64) + else: + runs = sources.ranges._needed_ranges( + manager, layout, lg, start // tpb, -(-end // tpb), end + ) + blocks = np.array( + [block for begin, stop in runs for block in range(begin, stop)], np.int64 + ) + if named and not layout.named_groups[lg]: + blocks = blocks[:0] + if size > 1: + blocks = blocks[blocks * tpb + rank * (tpb // size) < end] + slots = np.full(len(blocks), -1, dtype=np.int64) + inside = blocks < len(kv.pages[lg]) + slots[inside] = kv.pages[lg][blocks[inside]] + out.layer_group_indices.append(lg) + out.block_ordinals.append(blocks) + out.device_pages.append(slots) + return out + + +def _slow_view(sources, lender, entries, keys, device): + """The eager view; the public row constructor freezes its inputs.""" + groups = [] + rank, size = lender._layout.cp + for index, (lg, blocks) in enumerate(zip(entries.layer_group_indices, entries.block_ordinals)): + group = lender._layout.pool_group_of[lg] + if device: + groups.append( + sources.rows( + lg, + blocks * size + rank if size > 1 else blocks, + object_keys=None if keys is None else lender._object_keys(lg, keys[index]), + slots=entries.device_pages[index], + region_indices=lender._device_regions_of_pool_group[group], + ) + ) + else: + groups.append( + sources.rows( + lg, + blocks, + object_keys=lender._object_keys(lg, keys[index]), + slots=entries.staging_slots[index], + region_indices=(lender._region_index_of_pool_group[group],), + ) + ) + return sources.view(tuple(groups), 0, 0) + + +def test_cached_ranges_and_views_match_eager_paths(sources): + rng = random.Random(23) + for _ in range(160): + count, tpb, size = rng.randint(1, 3), rng.choice((4, 8)), rng.choice((1, 2)) + layout = SimpleNamespace( + num_layer_groups=count, + tokens_per_block=tpb, + windows=[rng.choice((None, 8, 16)) for _ in range(count)], + recurrent=[rng.random() < 0.2 for _ in range(count)], + named_groups=[rng.random() < 0.8 for _ in range(count)], + pool_group_of=[lg % 2 for lg in range(count)], + cp=(rng.randrange(size), size), + sink_blocks=[rng.randrange(3) for _ in range(count)], + ) + manager = SimpleNamespace() + manager.stale_at = lambda lg, h: ( + min(-(-h // tpb), layout.sink_blocks[lg]), + max( + min(-(-h // tpb), layout.sink_blocks[lg]), + int((h + 1 - (layout.windows[lg] or h + 1)) / tpb), + ), + ) + lender = _metadata_lender(sources, layout, manager) + kv = SimpleNamespace( + state_slots=[rng.randrange(32) for _ in range(count)], + history=rng.randrange(64), + active=True, + ) + snapshots = [] + for step in range(3): + end = rng.randint(1, 48) + starts = [rng.randrange(end + 1) for _ in range(count)] + kv.pages = [ + np.array([rng.randrange(32) for _ in range(rng.randrange(15))], np.int64) + for _ in range(count) + ] + named = rng.random() < 0.5 + entries = sources._block_entries_of_range( + lender, manager, kv, starts, end, named=named, request_id=7 + ) + expected = _slow_entries(sources, lender, manager, kv, starts, end, named) + assert _entry_key(entries) == _entry_key(expected) + assert len(lender._ordinal_ranges.get(7, {})) <= 2 * count + # A later call reuses ordinal templates but must reread this cache's current slots. + for pages in kv.pages: + pages += 100 + refreshed = sources._block_entries_of_range( + lender, manager, kv, starts, end, named=named, request_id=7 + ) + assert _entry_key(refreshed) == _entry_key( + _slow_entries(sources, lender, manager, kv, starts, end, named) + ) + entries.staging_slots = [ + np.arange(len(blocks), dtype=np.int64) + step for blocks in entries.block_ordinals + ] + keys = [ + np.full((len(blocks), 32), step + lg, np.uint8) + for lg, blocks in zip(entries.layer_group_indices, entries.block_ordinals) + ] + view = sources._build_view(lender, entries, keys, 0, 0) + assert _view_key(view) == _view_key(_slow_view(sources, lender, entries, keys, False)) + _assert_frozen_view(view) + snapshots.append((view, _view_key(view))) + device = sources.entries(entries.layer_group_indices, [], []) + for blocks, pages in zip(entries.block_ordinals, entries.device_pages): + device.block_ordinals.append(blocks[pages >= 0]) + device.device_pages.append(pages[pages >= 0]) + device_view = sources._in_place_view(lender, device, None, 0, 0) + assert _view_key(device_view) == _view_key( + _slow_view(sources, lender, device, None, True) + ) + _assert_frozen_view(device_view) + snapshots.append((device_view, _view_key(device_view))) + for values in ( + entries.staging_slots + keys + device.block_ordinals + device.device_pages + ): + values[:] = 0 + assert all(_view_key(old) == value for old, value in snapshots) + # The window filter keeps each group's rows and the history limit of the eager path. + prompt = rng.randrange(64) + computed = prompt if rng.random() < 0.3 else rng.randrange(prompt + 1) + request = SimpleNamespace( + py_request_id=7, prompt=prompt, context_current_position=computed + ) + filtered = sources._without_short_windows_mid_context( + lender, manager, request, SimpleNamespace(history=kv.history), entries + ) + keep = [ + kv.history >= request.prompt + or computed >= request.prompt + or window is None + or window >= request.prompt + for window in layout.windows + ] + wanted = sources.entries( + entries.layer_group_indices, + [ + blocks if flag else blocks[:0] + for blocks, flag in zip(entries.block_ordinals, keep) + ], + [pages if flag else pages[:0] for pages, flag in zip(entries.device_pages, keep)], + ) + assert _entry_key(filtered) == _entry_key(wanted) + lowest = {lg: rng.randrange(5) for lg in range(count)} + lowest = {lg: block for lg, block in lowest.items() if block >= layout.sink_blocks[lg]} + lender._locks = SimpleNamespace(lowest_locked_blocks=lambda cache, floor: lowest) + limits = [ + sources.ranges._last_history_keeping(manager, layout, lg, block) + for lg, block in lowest.items() + if layout.windows[lg] is not None + ] + assert sources._history_cap(lender, kv) == min(limits, default=None) + + +def _storage(array): + return np.frombuffer(_root(array), dtype=np.uint8) + + +def test_ordinal_templates_share_capacity_and_leave_old_arrays_unchanged(sources): + lender = _metadata_lender(sources, layout([False]), object()) + first = lender._range_ordinals(7, 0, [(0, 4)]) + second = lender._range_ordinals(7, 0, [(4, 8)]) + assert np.shares_memory(_storage(first), _storage(second)) + before = first.tobytes(), second.tobytes() + third = lender._range_ordinals(7, 0, [(8, 12)]) + assert third.tolist() == [8, 9, 10, 11] + assert not np.shares_memory(_storage(first), _storage(third)) + assert (first.tobytes(), second.tobytes()) == before + for values in (first, second, third): + with pytest.raises(ValueError): + values.flags.writeable = True + assert lender._range_ordinals(None, 0, [(1, 3), (8, 10)]).tolist() == [1, 2, 8, 9] + assert set(lender._ordinal_ranges) == {7} + large = lender._range_ordinals(8, 0, [(10_000, 10_004)]) + assert large.tolist() == [10_000, 10_001, 10_002, 10_003] + assert _storage(large).nbytes <= 2 * len(large) * large.itemsize, ( + "range storage must be bounded by row count, not absolute ordinal" + ) + + +@pytest.mark.parametrize("mutant", ["identity", "readonly_flag"]) +def test_metadata_mutants_expose_writable_or_changing_old_views(mutant): + correct, broken = _sources(), _sources(mutant) + for methods, should_fail in ((correct, False), (broken, True)): + lender = _metadata_lender(methods, layout([False]), object()) + raw_blocks, raw_slots = np.array([1, 2]), np.array([3, 4]) + blocks, slots = raw_blocks.view(), raw_slots.view() + blocks.flags.writeable = slots.flags.writeable = False + entries = methods.entries([0], [blocks], [slots]) + view = methods._in_place_view(lender, entries, None, 0, 0) + before = _view_key(view) + raw_blocks[:] = 9 + raw_slots[:] = 8 + if should_fail: + assert _view_key(view) != before, "the mutant failed to expose a template alias" + view.groups[0].page_ordinals.flags.writeable = True + assert view.groups[0].page_ordinals.flags.writeable + else: + assert _view_key(view) == before + _assert_frozen_view(view) + + +def test_ordinal_cache_is_cleared_on_free_reset_and_shutdown(sources): + lender = SimpleNamespace( + _ordinal_ranges={7: {}, 8: {}}, + _shrink_epochs={7: 2, 8: 1}, + _writes_at_end={}, + _state_unverified={7, 8}, + _parked={}, + _parked_closing=[], + _stop_reason=None, + _shut_down=False, + _waiting_for_slots=(), + _serving_manager=lambda: object(), + _fail_waiting=lambda rid, reason: pytest.fail("no lease waits for slots"), + _ledger=SimpleNamespace(_drop_records=lambda rid: None), + _progress=lambda: pytest.fail("the free hook runs no round of progress"), + ) + sources._on_free(lender, 7, object()) + assert lender._ordinal_ranges == {8: {}} and lender._shrink_epochs == {8: 1} + assert lender._state_unverified == {8} + sources._on_reset(lender) + assert lender._ordinal_ranges == {} and lender._reuse_was_reset + lender._ordinal_ranges[7] = {} + lender._shut_down = False + lender._device_regions_of_pool_group = {} + lender._leases_holding_slots = lender._waiting_for_slots = lender._open_leases = () + lender._progress_round = 0 + lender._recycle = lambda: None + lender._memory_ref = lambda: None + sources._on_shutdown(lender) + assert lender._ordinal_ranges == {} and lender._shut_down + + +def _slow_object_keys(identity, group, keys): + if not 0 <= group < len(identity._groups): + raise ValueError(f"unknown layer group {group}") + keys = np.ascontiguousarray(keys, dtype=np.uint8).reshape(-1, 32) + out = np.empty((len(keys), 54), dtype=np.uint8) + out[:, :16] = np.frombuffer(identity.namespace, dtype=np.uint8) + out[:, 16:48] = keys + out[:, 48:] = identity._groups[group] + return out + + +def test_object_key_bytes_and_snapshots_match_eager_layout(sources): + rng = random.Random(29) + for _ in range(200): + count = rng.choice((0, 0, 1, 4, 15)) + raw = np.frombuffer(rng.randbytes(count * 64), dtype=np.uint8).copy().reshape(count, 64) + keys = raw[:, ::2] + suffix = struct.pack(">HHH", rng.randrange(65536), 4, 2) + identity = SimpleNamespace( + namespace=bytes(range(16)), _groups=[np.frombuffer(suffix, dtype=np.uint8).copy()] + ) + expected = _slow_object_keys(identity, 0, keys) + result = sources.object_keys(identity, 0, keys) + assert result.shape == (count, 54) and result.dtype == np.uint8 + assert result.tobytes() == expected.tobytes() + assert not result.flags.writeable + with pytest.raises(ValueError): + result.flags.writeable = True + assert all(row[:16].tobytes() == identity.namespace for row in result) + assert np.array_equal(result[:, 16:48], keys) + assert np.array_equal(result[:, 48:], np.broadcast_to(identity._groups[0], (count, 6))) + snapshot = result.tobytes() + raw[:] = 0 + identity._groups[0][:] = 0 + identity.namespace = b"!" * 16 + assert result.tobytes() == snapshot + for invalid in (-1, 1): + assert _outcome(lambda: sources.object_keys(identity, invalid, np.zeros((0, 32)))) == ( + _outcome(lambda: _slow_object_keys(identity, invalid, np.zeros((0, 32)))) + ) + malformed = np.zeros(31, np.uint8) + assert _outcome(lambda: sources.object_keys(identity, 0, malformed)) == ( + _outcome(lambda: _slow_object_keys(identity, 0, malformed)) + ) + + +def test_an_all_write_preserves_records_and_does_not_grow(sources, kit): + manager_api = copy.copy(sources.manager) + ns = dict( + sources.ns, + _manager=manager_api, + MemoryKind=MemoryKind, + RowSelection=RowSelection, + LockMode=LockMode, + _CUDA_EVENT=Event, + _check_ready_event=lambda event: None, + ) + lend_write = _load("_lender.py", "_ManagerLender.lend_write", ns.copy()) + manager_api.prompt_length = lambda request, **kwargs: request.prompt + manager_api.kv_of = lambda manager, rid: manager.kv + manager_api.connector_load_pending = lambda *args: False + manager_api.grow = lambda *args: pytest.fail("ALL writes must not grow their cache") + # The stand-in manager holds back no commit for a fetch to drop. + manager_api.drop_held_commit = lambda *args: None + + def run(start, end, pages): + kv = SimpleNamespace(pages=[pages.copy()], committed=(0,), active=True, history=0) + manager = SimpleNamespace(kv=kv, stale=[(0, 0)]) + lender = _metadata_lender(sources, layout([False], tpb=4), manager) + lender._check_memory_and_rows = lambda *args: None + lender._whole_blocks = lambda start, end, whole_end=True: (start, end) + lender._lending_manager = lambda: manager + lender._progress = lender._check_not_stopped = lambda: None + lender._lends_fast = lambda: False # the rules alone, which this test compares + lender._required_slots_per_pool_group = lambda counts: {0: sum(counts)} + lender._write_range_refusal = lender._bidirectional_split_refusal = lambda *args: None + lender._scratch_reuse_on = lambda cache: False + lender._history_refusal = lender._unwritten_pages_refusal = lambda *args: None + lender._history_cap = sources._history_cap.__get__(lender) + lender._locks = SimpleNamespace( + check=lambda cache, blocks, mode, **kwargs: SimpleNamespace(blocks=blocks) + ) + lender._pageless = lambda entries: ( + "missing" if any((slots < 0).any() for slots in entries.device_pages) else None + ) + lender._ledger = SimpleNamespace( + _unsettled_fetch_refusal=lambda *args: None, + _prefix_kept_by_fetch=lambda *args: (0,), + _record_fetch=lambda *args, **kwargs: SimpleNamespace( + record=(args[0], args[2:], kwargs) + ), + ) + lender._lend_in_place = ( + lambda mgr, req, cache, kind, entries, keys, fetch, doomed, **kwargs: SimpleNamespace( + entries=entries, fetch=fetch, doomed=doomed, checked=kwargs.get("checked") + ) + ) + + lender._block_entries_of_range = sources._block_entries_of_range.__get__(lender) + request = SimpleNamespace(py_request_id=7, prompt=48) + lease = lend_write(lender, request, start, end, **kit.DEVICE_ALL) + assert np.array_equal(kv.pages[0], pages) and kv.history == 0 + assert lease.doomed is None and lease.fetch.entries is lease.entries + expected_ordinals = list(range(start // 4, -(-end // 4))) + assert lease.entries.layer_group_indices == [0] + assert lease.entries.block_ordinals[0].tolist() == expected_ordinals + assert lease.entries.device_pages[0].tolist() == pages[expected_ordinals].tolist() + # The write locks the same ordinal array it lends, without another cache lookup. + assert list(lease.checked.blocks) == [0] + assert lease.checked.blocks[0] is lease.entries.block_ordinals[0] + assert lease.fetch.record == ( + 7, + (start, (0,), 0, (0,)), + {"state_end": None, "end": end, "prompt": 48, "runs": []}, + ) + + rng = random.Random(31) + for _ in range(100): + start = 4 * rng.randrange(12) + end = rng.randint(start, 48) + pages = np.array(rng.sample(range(1000), 12), dtype=np.int64) + run(start, end, pages) + + +def _slow_lendable(sources, lender, manager, entries): + out = sources.entries([], [], []) + for lg, blocks, pages in zip( + entries.layer_group_indices, entries.block_ordinals, entries.device_pages + ): + begin, end = (0, 0) if lender._layout.windows[lg] is None else manager.stale[lg] + kept = [ + (block, page) + for block, page in zip(blocks.tolist(), pages.tolist()) + if page >= 0 and not begin <= block < end + ] + out.layer_group_indices.append(lg) + out.block_ordinals.append(np.array([block for block, _ in kept], dtype=np.int64)) + out.device_pages.append(np.array([page for _, page in kept], dtype=np.int64)) + return out + + +def test_lendable_and_committed_rows_match_scalar_filters(sources): + rng = random.Random(5820) + covered = set() + for _ in range(300): + count, tpb = rng.randrange(5), rng.choice((4, 8)) + # A window with nothing behind it yet has an empty stale range. + manager = SimpleNamespace( + stale=[ + tuple(sorted(rng.sample(range(16), 2))) if rng.random() < 0.8 else (3, 3) + for _ in range(count) + ] + ) + named = [rng.random() < 0.75 for _ in range(count)] + lender = SimpleNamespace( + _layout=SimpleNamespace( + windows=[rng.choice((None, 8, 16)) for _ in range(count)], + num_layer_groups=count, + tokens_per_block=tpb, + ), + _named_groups=named, + ) + kv = SimpleNamespace(committed=tuple(rng.randrange(24) * tpb for _ in range(count))) + blocks = [ + np.array([rng.randrange(20) for _ in range(rng.randrange(15))], dtype=np.int64) + for _ in range(count) + ] + entries = sources.entries( + list(range(count)), + blocks, + [np.array([rng.randrange(-2, 20) for _ in group], dtype=np.int64) for group in blocks], + ) + before = _entry_key(entries) + expected = _slow_lendable(sources, lender, manager, entries) + history = rng.randrange(100) + actual = sources._lendable_block_entries(lender, manager, history, entries) + assert _entry_key(actual) == _entry_key(expected), ( + "lendable rows include stale/missing pages" + ) + _assert_built_frozen(actual, entries) + for incoming in (entries, actual): + expected = sources.entries([], [], []) + for lg, ordinals, slots in zip( + incoming.layer_group_indices, incoming.block_ordinals, incoming.device_pages + ): + rows = list(zip(ordinals.tolist(), slots.tolist())) + kept = [(b, s) for b, s in rows if named[lg] and b * tpb + tpb <= kv.committed[lg]] + expected.layer_group_indices.append(lg) + expected.block_ordinals.append(np.array([b for b, _ in kept], dtype=np.int64)) + expected.device_pages.append(np.array([s for _, s in kept], dtype=np.int64)) + covered.add( + "empty" + if not rows + else "whole" + if len(kept) == len(rows) + else "none" + if not kept + else "partial" + ) + if not named[lg]: + covered.add("unnamed") + if any(a > b for a, b in zip(ordinals.tolist(), ordinals.tolist()[1:])): + covered.add("unsorted") + committed = sources._committed_entries(lender, kv, incoming) + assert _entry_key(committed) == _entry_key(expected), ( + "committed rows differ from prefix" + ) + _assert_built_frozen(committed, incoming) + assert _entry_key(entries) == before + assert covered == {"empty", "whole", "none", "partial", "unnamed", "unsorted"} + + +def test_committed_filter_keeps_the_eager_shapes_and_error_outcomes(sources): + lender = SimpleNamespace( + _layout=SimpleNamespace(tokens_per_block=4, num_layer_groups=1), _named_groups=[True] + ) + + def eager(kv, entries): + ordinals, slots = entries.block_ordinals[0], entries.device_pages[0] + keep = (ordinals < kv.committed[0] // 4) & lender._named_groups[0] + return _entry_key(sources.entries([0], [ordinals[keep]], [slots[keep]])) + + shapes = ((), (0,), (1,), (3,), (1, 3), (3, 1), (2, 2), (0, 2)) + for shape in shapes: + for slot_shape in shapes: + for committed in (0, 4, 64): + for named in (False, True): + lender._named_groups[0] = named + ordinals = np.arange(int(np.prod(shape)), dtype=np.int64).reshape(shape) + slots = np.arange(int(np.prod(slot_shape)), dtype=np.int64).reshape(slot_shape) + entries = sources.entries([0], [ordinals], [slots]) + kv = SimpleNamespace(committed=(committed,)) + expected = _outcome(lambda: eager(kv, entries)) + actual = _outcome( + lambda: _entry_key(sources._committed_entries(lender, kv, entries)) + ) + assert actual == expected, (shape, slot_shape, committed, named) + + +@pytest.mark.parametrize("device", [False, True], ids=["HOST", "DEVICE"]) +def test_named_pipeline_reuses_frozen_metadata_and_preserves_old_views( + sources, monkeypatch, device +): + manager = SimpleNamespace(stale=[(0, 0), (0, 0)]) + lender = _metadata_lender(sources, layout([False, False], tpb=4), manager) + identity = SimpleNamespace( + namespace=bytes(range(16)), + _groups=[np.frombuffer(struct.pack(">HHH", lg, 4, 2), dtype=np.uint8) for lg in range(2)], + ) + emitted = {} + + def object_keys(lg, keys): + result = sources.object_keys(identity, lg, keys) + emitted[lg] = result + return result + + lender._object_keys = object_keys + observations = [] + immutable = sources.ns["_immutable"] + + def observe(array): + result = immutable(array) + observations.append((np.shares_memory(array, result), array.flags.writeable)) + return result + + monkeypatch.setitem(sources.ns, "_immutable", observe) + kv = SimpleNamespace( + pages=[np.arange(8, dtype=np.int64) + 20 + lg * 100 for lg in range(2)], + committed=(32, 32), + ) + snapshots = [] + for start, end in ((0, 16), (16, 32)): + raw = sources._block_entries_of_range( + lender, manager, kv, start, end, named=True, request_id=7 + ) + lendable = sources._lendable_block_entries(lender, manager, 32, raw) + entries = sources._committed_entries(lender, kv, lendable) + for blocks, slots, original_blocks, original_slots in zip( + entries.block_ordinals, entries.device_pages, raw.block_ordinals, raw.device_pages + ): + assert np.shares_memory(blocks, original_blocks), "NAMED all-keep ordinals were copied" + assert np.shares_memory(slots, original_slots), "NAMED all-keep slots were copied" + entries.staging_slots = [np.arange(4, dtype=np.int64) for _ in range(2)] + keys = [np.arange(128, dtype=np.uint8).reshape(4, 32) + lg for lg in range(2)] + view = ( + sources._in_place_view(lender, entries, keys, 0, 0) + if device + else sources._build_view(lender, entries, keys, 0, 0) + ) + for lg, rows in enumerate(view.groups): + assert rows.page_ordinals.tolist() == list(range(start // 4, end // 4)) + assert np.shares_memory(rows.page_ordinals, entries.block_ordinals[lg]) + assert np.shares_memory(rows.object_keys, emitted[lg]) + expected_keys = b"".join( + identity.namespace + row.tobytes() + identity._groups[lg].tobytes() + for row in keys[lg] + ) + assert rows.object_keys.tobytes() == expected_keys + if device: + assert np.shares_memory(rows.slots, entries.device_pages[lg]) + assert rows.slots.tolist() == kv.pages[lg][start // 4 : end // 4].tolist() + else: + assert rows.slots.tolist() == list(range(4)) + assert rows.region_indices is lender._staging_region_indices[lg] + _assert_frozen_view(view) + assert observations and all(shared and not writable for shared, writable in observations), ( + "production metadata must enter DTOs already immutable without another copy" + ) + snapshots.append((view, _view_key(view))) + for values in entries.staging_slots + keys + kv.pages: + values[:] = 0 + assert all(_view_key(old) == before for old, before in snapshots) + assert all(_view_key(old) == before for old, before in snapshots) + + +def test_two_groups_and_two_spans_reuse_bounded_storage(sources): + lender = _metadata_lender(sources, layout([False, False]), object()) + first_runs = [[(0, 2), (100, 104)], [(20, 23), (200, 205)]] + next_runs = [[(0, 2), (104, 108)], [(23, 26), (205, 210)]] + old_entries = sources.entries([0, 1], [], []) + for lg, runs in enumerate(first_runs): + values = lender._range_ordinals(7, lg, runs) + assert values.tolist() == [b for begin, end in runs for b in range(begin, end)] + old_entries.block_ordinals.append(values) + old_entries.device_pages.append(sources.ns["_frozen_int64"](range(len(values)))) + old_view = sources._in_place_view(lender, old_entries, None, 0, 0) + before = _view_key(old_view) + storage = [_storage(values) for _, values in lender._ordinal_ranges[7].values()] + for lg, runs in enumerate(next_runs): + values = lender._range_ordinals(7, lg, runs) + assert values.tolist() == [b for begin, end in runs for b in range(begin, end)] + later = [_storage(values) for _, values in lender._ordinal_ranges[7].values()] + assert all(any(np.shares_memory(a, b) for b in later) for a in storage), ( + "retained span storage must survive consecutive chunks across layer groups" + ) + rows = sum(end - begin for runs in next_runs for begin, end in runs) + assert sum(buf.nbytes for buf in later) <= 2 * rows * np.dtype(np.int64).itemsize + for lg, runs in enumerate([[(0, 2), (108, 112)], [(26, 29), (210, 215)]]): + values = lender._range_ordinals(7, lg, runs) + assert values.tolist() == [b for begin, end in runs for b in range(begin, end)] + assert _view_key(old_view) == before + _assert_frozen_view(old_view) + + +def _slow_resume_without_windows( + tpb: int, + groups: tuple[int, ...], + masks: list[np.ndarray], + partial_ends: dict[tuple[int, int], int], + computed: tuple[int, ...], + cap: int | None, +) -> int: + known = min((computed[lg] for lg in groups), default=max(computed, default=0)) + if not groups: + return known if cap is None else cap + if cap is None: + cap = max(len(masks[lg]) for lg in groups) * tpb + for end in range(cap, known, -1): + if all( + all( + max( + computed[lg], + partial_ends[(lg, ordinal)] + if (lg, ordinal) in partial_ends + else (ordinal + 1) * tpb + if ordinal < len(masks[lg]) and masks[lg][ordinal] + else 0, + ) + >= min(end, (ordinal + 1) * tpb) + for ordinal in range(known // tpb, -(-end // tpb)) + ) + for lg in groups + ): + return end + return known + + +def test_resume_bound_without_windows_matches_scalar_validity(sources): + """The largest resume endpoint has all rows computed or delivered in every paged group.""" + current = _load("_fetch_ledger.py", "_FetchLedger._max_resume_position", sources.ns.copy()) + closed = _load( + "_fetch_ledger.py", "_FetchLedger._max_resume_without_windows", sources.ns.copy() + ) + known_of = _load("_fetch_ledger.py", "_FetchLedger._known", sources.ns.copy()) + rng = random.Random(126) + seen = set() + for _ in range(4000): + tpb, count = rng.choice((4, 8, 32)), rng.randint(1, 3) + recurrent = frozenset(lg for lg in range(count) if rng.random() < 0.25) + layout = SimpleNamespace(tokens_per_block=tpb, windows=[None] * count, cp=(0, 1)) + paged = tuple(lg for lg in range(count) if lg not in recurrent) + ledger = SimpleNamespace(_layout=layout, _recurrent_groups=recurrent, _paged_groups=paged) + ledger._known = known_of.__get__(ledger) + ledger._max_resume_without_windows = closed.__get__(ledger) + blocks = rng.randint(0, 24) + masks = [] + for _ in range(count): + arrived = rng.choice((1.0, 1.0, 0.95, 0.7, 0.3)) + length = max(0, blocks + rng.randint(-3, 3)) + masks.append(np.array([rng.random() < arrived for _ in range(length)], dtype=bool)) + partial_ends = {} + for _ in range(rng.choice((0, 0, 1, 2, 3))): + lg, ordinal = rng.randrange(count), rng.randrange(blocks + 2) + partial_ends[(lg, ordinal)] = ordinal * tpb + rng.randint(1, tpb - 1) + delivered = SimpleNamespace(delivered_block_masks=masks, partial_ends=partial_ends) + computed = tuple( + rng.choice((0, 0, rng.randint(0, (blocks + 2) * tpb))) for _ in range(count) + ) + cap = rng.choice((None, rng.randint(0, (blocks + 3) * tpb))) + want = _slow_resume_without_windows(tpb, paged, masks, partial_ends, computed, cap) + got = current(ledger, None, delivered, computed, cap) + assert got == want, (tpb, sorted(recurrent), masks, partial_ends, computed, cap) + known = ledger._known(computed) + seen.add( + "known" + if got == known + else "partial end" + if got in partial_ends.values() + else "computed" + if got in computed + else "block" + if got % tpb == 0 + else "cap" + ) + assert seen == {"known", "partial end", "computed", "block", "cap"}, seen + + +class Event: + """A stand-in CUDA event, complete at once, named by who made it.""" + + def __init__(self, name): + self.name = name + + def query(self): + return True + + def synchronize(self): + raise AssertionError("a lease waited on its event") + + +class Connector: + """A stand-in KV cache connector manager.""" + + def __init__(self, capacity_only, loading, reserved): + self.capacity_only = capacity_only + self.loading = loading + self.reserved = reserved + + def has_pending_load(self, request): + return self.loading + + def get_prefix_reservation(self, request): + return SimpleNamespace() if self.reserved else None + + +def runtime_stale_range(window, sinks, history): + """The runtime's stale range of a ``window``-token window with ``sinks`` sink blocks at + ``history``, in its 32-bit arithmetic, the division truncating toward zero.""" + covered = (history + TPB - 1) // TPB + start = min(covered, sinks) + if window is None: + return start, start + behind = history + 1 - window + window_start = int(behind / TPB) + return start, max(start, window_start) + + +class Manager: + """A stand-in manager: its layer groups' pools, windows and sinks, its caches and two rules of + its configuration; a commit it holds back for a cache runs at the unlock.""" + + # No cache of this stand-in is on an allocation side stream. + _on_allocation_stream = frozenset() + + def __init__(self, groups, connector, floor_follows_history, keeps_stale_pages, windows, sinks): + self.impl = SimpleNamespace(get_life_cycle_pool_group_indices=lambda: list(range(groups))) + self.kv_cache_map = {} + self._held_commits = {} + self.kv_connector_manager = connector + self.floor_follows_history = floor_follows_history + self.keeps_stale_pages = keeps_stale_pages + self.tokens_per_block = TPB + self.windows = windows + self.sinks = sinks + + def _stale_block_range(self, layer_group, history): + return runtime_stale_range(self.windows[layer_group], self.sinks[layer_group], history) + + def _commit_held_back(self, kv): + """Commit what the locks held back: more tokens, the committed blocks on other pages.""" + commit = getattr(kv, "commit_at_unlock", None) + if self._held_commits.pop(id(kv), None) is None or commit is None: + return + kv.commit_at_unlock = None + kv.num_committed_tokens = commit + for pages in kv.pages: + whole = min(commit // TPB, len(pages)) + pages[:whole] = np.where(pages[:whole] >= 0, (pages[:whole] + 17) % SLOTS, -1) + + +class Cache: + """A stand-in runtime cache: per layer group its blocks' pages, -1 where none and past its + blocks, as the page-index buffer holds them; a state slot per group of recurrent state.""" + + def __init__(self, rng, groups, history, committed, statuses, plain): + self.num_blocks = max(rng.randint(0, 7), -(-history // TPB) if plain else 0) + self.pages = [] + for _ in range(groups): + missing = not plain and rng.random() < 0.25 + pages = [ + -1 if missing and rng.random() < 0.3 else rng.randrange(SLOTS) + for _ in range(self.num_blocks) + ] + self.pages.append(np.array(pages + [-1] * 4, dtype=np.int32)) + self.state_slots = [rng.choice((3, 9) if plain else (-1, 3, 3, 9)) for _ in range(groups)] + self.history_length = history + self.Status = SimpleNamespace(OPEN=0, CLOSED=1) + self.status = self.Status.OPEN + self.num_committed_tokens = committed + self.reuse_status = statuses + self.is_active = plain or rng.random() < 0.9 + self.tokens_per_block = TPB + self.enable_swa_scratch_reuse = not plain and rng.random() < 0.1 + self.commit_at_unlock = None + + def get_base_page_indices(self, layer_group, beam_id=0): + view = self.pages[layer_group][:] + view.flags.writeable = False + return view + + def get_aggregated_page_indices(self, layer_group, beam_id=0, valid_only=False): + pages = [int(p) for p in self.pages[layer_group][: self.num_blocks]] + return [p for p in pages if p >= 0] if valid_only else pages + + def get_ssm_block_base_index(self, layer_group, beam_id=0): + return self.state_slots[layer_group] + + +class Rules(_lender._ManagerLender): + """A lender of a class of its own, which lends and returns through the rules alone.""" + + +def layout( + recurrent, num_slots=SLOTS, windows=None, sinks=None, read_windows=(), formula=False, tpb=TPB +): + """One layer group per pool group, each with a device pool of ``num_slots`` slots or its pool + group's count of them; ``formula``: stale ranges by the window formula, as an attach finds.""" + count = len(recurrent) + if isinstance(num_slots, int): + num_slots = (num_slots,) * count + return SimpleNamespace( + tokens_per_block=tpb, + windows=tuple(windows or (None,) * count), + sink_blocks=tuple(sinks or (0,) * count), + read_windows=tuple(read_windows), + stale_ranges_by_formula=formula, + recurrent=tuple(recurrent), + pool_groups=tuple(range(count)), + pool_group_of=tuple(range(count)), + num_layer_groups=count, + named_groups=(True,) * count, + cp=(0, 1), + device_pools={ + g: (DevicePool(g, 0, (g + 1) << 30, 256, num_slots[g]),) for g in range(count) + }, + ) + + +def lender_over(cls, manager, lay, no_names=None): + """A lender of class ``cls`` over ``manager`` with layout ``lay``, its identity and regions.""" + count = lay.num_layer_groups + ident = Identity(b"scope", bytes(range(32)), [(lg,) for lg in range(count)], [(1, 0)] * count) + ref = weakref.ref(manager) + regions = tuple(_lender._device_regions(manager, lay, ident)) + return cls(ref, lay, ident, regions, SlotQueue({}), ref, no_names=no_names) + + +def plain_world(cls, lay, history, rng, floor=False, no_names=None): + """A lender of class ``cls`` with layout ``lay`` over a manager of its windows and sinks, and + an active cache of ``history`` tokens with every page, drawn from ``rng``.""" + count = lay.num_layer_groups + manager = Manager(count, None, floor, False, lay.windows, lay.sink_blocks) + lender = lender_over(cls, manager, lay, no_names) + status = SimpleNamespace(complete=True, endpoint=0) + kv = Cache(rng, count, history, 0, [status] * count, True) + manager.kv_cache_map[RID] = kv + return lender, manager, kv + + +def request_of(prompt, position=0, remaining=None, runs=None, served=0): + """A request of ``prompt`` tokens whose context computed ``position``, ``remaining`` to go.""" + return SimpleNamespace( + py_request_id=RID, + prompt_len=prompt, + total_input_len_cp=prompt, + py_multimodal_data=runs, + py_served_prefix_end=served, + context_current_position=position, + context_remaining_length=prompt - position if remaining is None else remaining, + ) + + +def world_of(lender, manager, kv, request, start, end, kind="read", ready_event=None): + """A ``kind`` lend of ``[start, end)`` by ``lender`` over ``manager``, its cache and request.""" + return SimpleNamespace(lease=None, **locals()) + + +def world(seed, cls): + """A lend by a lender of class ``cls``, the same draw for the same seed: locks, fetch records, + windows, sinks and pools short of slots now and then; half the draws plain, with every page.""" + rng = random.Random(seed) + plain = rng.random() < 0.5 + count = rng.randint(1, 3) + recurrent = [rng.random() < 0.3 for _ in range(count)] + windowed = rng.random() < 0.4 + windows = [ + rng.choice((None, rng.randint(1, 6 * TPB))) if windowed and not state else None + for state in recurrent + ] + sinks = [int(w is not None and rng.random() < 0.15) for w in windows] + read_windows = () + if any(w is not None for w in windows) and rng.random() < 0.4: + read_windows = tuple(w if w is None else rng.randint(1, w) for w in windows) + connector = ( + None + if rng.random() < 0.6 + else Connector(rng.random() < 0.3, rng.random() < 0.3, rng.random() < 0.2) + ) + manager = Manager( + count, connector, rng.random() < 0.5, rng.random() < 0.5, tuple(windows), tuple(sinks) + ) + # A few pools hold fewer slots than the pages and state slots the cache draws. + num_slots = [SLOTS if rng.random() < 0.9 else rng.randint(1, SLOTS - 1) for _ in range(count)] + formula = cls is _lender._ManagerLender + lay = layout(recurrent, num_slots, windows, sinks, read_windows, formula=formula) + no_names = "this lender names nothing" if any(recurrent) else None + lender = lender_over(cls, manager, lay, no_names) + prompt = rng.randint(1, 6 * TPB) + history = prompt if plain else rng.choice((prompt, prompt, rng.randint(0, 7 * TPB))) + committed = rng.choice((0, 0, rng.randint(0, history))) + statuses = [ + SimpleNamespace(complete=rng.random() < 0.8, endpoint=rng.randint(0, committed)) + for _ in range(count) + ] + kv = Cache(rng, count, history, committed, statuses, plain) + manager.kv_cache_map[RID] = kv + ready_event = rng.choice((None, None, Event("caller"))) + kind = rng.choice(("read", "write")) + if plain: + end = rng.choice((prompt, TPB * rng.randint(0, prompt // TPB))) + start = TPB * rng.randint(0, end // TPB) + else: + start = TPB * rng.randint(0, 6) + end = rng.choice((prompt, prompt, history, rng.randint(0, 7 * TPB))) + runs = None + if rng.random() < 0.3: + # A run of multimodal tokens, often one the range's end falls inside. + begin = rng.randint(max(0, end - TPB), max(0, end - 1)) if rng.random() < 0.7 else 0 + length = rng.randint(2, 2 * TPB) + runs = { + "mm_bidirectional_blocks": True, + "multimodal_embed_mask_cumsum": np.cumsum( + [1 if begin <= i < begin + length else 0 for i in range(prompt)] + ), + } + position = rng.randint(0, prompt) + served = rng.choice((0, 0, rng.randint(0, prompt))) + remaining = rng.choice((prompt - position, rng.randint(0, prompt - position))) + request = request_of(prompt, position, remaining, runs, served) + for _ in range(0 if plain else rng.choice((0, 0, 1, 2))): + lg = rng.randrange(count) + blocks = sorted(rng.sample(range(8), rng.randint(1, 3))) + mode = rng.choice((LockMode.READ, LockMode.WRITE)) + try: + lender._locks.lock(kv, {lg: blocks}, mode, ready=Event("held")) + except ValueError: + pass + ledger = lender._ledger + if not plain and rng.random() < 0.2: + fetch = _fetch_ledger._FetchRecord( + rng.choice((0, TPB)), kv, marks_queued=rng.random() < 0.5, end=rng.randint(0, prompt) + ) + ledger._latest_fetches[RID] = fetch + if not plain and rng.random() < 0.15: + masks = [np.array([rng.random() < 0.7 for _ in range(6)]) for _ in range(count)] + ledger._deliveries[RID] = _fetch_ledger._Deliveries(kv, 0, masks) + if not plain and rng.random() < 0.1: + ledger._kept_prefixes[RID] = (kv, (rng.randint(0, history),) * count, history) + if not plain and rng.random() < 0.05: + ledger._mixed_tail_caches[RID] = kv + return world_of(lender, manager, kv, request, start, end, kind, ready_event) + + +def mid_context_world(cls, window, history, position, end, prompt): + """A read of ``[0, end)`` over a full layer group and one of a ``window``-token window, its + cache of ``history`` tokens of ``prompt``, whose context computed ``position`` of them.""" + lay = layout([False, False], windows=(None, window), formula=cls is _lender._ManagerLender) + lender, manager, kv = plain_world(cls, lay, history, random.Random(history)) + return world_of(lender, manager, kv, request_of(prompt, position), 0, end) + + +def windowed_receive_world(seed, cls): + """A receive of every row over windows of the formula without sinks, now and then a full group + beside them and narrower read windows, into a cache whose history lies up to the prompt.""" + rng = random.Random(seed) + count = rng.randint(1, 3) + windows = [rng.randint(1, 6 * TPB) for _ in range(count)] + if count > 1 and rng.random() < 0.3: + windows[rng.randrange(count)] = None + read_windows = () + if rng.random() < 0.6: + read_windows = tuple(w if w is None else rng.randint(1, w) for w in windows) + floor = rng.random() < 0.5 + fast = cls is _lender._ManagerLender + lay = layout([False] * count, SLOTS, windows, None, read_windows, formula=fast) + prompt = rng.randint(1, 6 * TPB) + history = rng.randint(0, prompt) + lender, manager, kv = plain_world(cls, lay, history, rng, floor) + kv.num_blocks = -(-prompt // TPB) + kv.pages = [ + np.array([rng.randrange(SLOTS) for _ in range(kv.num_blocks)] + [-1] * 4, dtype=np.int32) + for _ in range(count) + ] + end = rng.choice((prompt, prompt, TPB * rng.randint(0, prompt // TPB))) + start = TPB * rng.randint(0, end // TPB) + return world_of(lender, manager, kv, request_of(prompt), start, end, "write") + + +def hybrid_world(cls, kind, num_slots, page, state_slot): + """A ``kind`` lend of a 4-block prompt over a group of recurrent state in slot ``state_slot`` + and one of attention on pages 0, 1, ``page`` and 3, their pools of ``num_slots`` slots.""" + prompt, no_names = 4 * TPB, "this lender names nothing" + lay = layout([True, False], num_slots) + lender, manager, kv = plain_world(cls, lay, prompt, random.Random(0), no_names=no_names) + kv.num_blocks = 4 + kv.pages = [np.array([-1] * 8, dtype=np.int32), np.array([0, 1, page, 3] + [-1] * 4, np.int32)] + kv.state_slots = [state_slot, -1] + return world_of(lender, manager, kv, request_of(prompt, prompt), 0, prompt, kind) + + +def row_lender(blocks): + """A lender of one layer group over a cache of ``blocks + 1`` blocks: lender, cache, pages.""" + history, rng = (blocks + 1) * TPB, random.Random(blocks) + lender, _, kv = plain_world(_lender._ManagerLender, layout([False]), history, rng) + return lender, kv, kv.pages[0] + + +def buffer_row(lender, row): + """Row ``row`` of the lender's page-index buffer, its key half, as a lend reads a row.""" + group_words, row_words, row_blocks = lender._index_strides + return lender._index_words, row * row_words, group_words, row_blocks + + +def array_key(array): + """``array``'s dtype, shape, values and writability, and the type of the memory under it.""" + if array is None: + return None + root = type(_root(array)).__name__ + return array.dtype.str, array.shape, array.tolist(), bool(array.flags.writeable), root + + +def lease_key(lease): + view, lock, fetch, entries = lease._view, lease._lock, lease._fetch, lease._block_entries + return dict( + cls=type(lease).__name__, + kind=lease._kind, + failure=lease.failure, + first_poll_failure=lease._first_poll_failure, + loading=lease._loading is not None, + view=None + if view is None + else [ + ( + rows.layer_group_index, + array_key(rows.page_ordinals), + array_key(rows.slots), + rows.region_indices, + array_key(rows.object_keys), + ) + for rows in view.groups + ], + lock=None + if lock is None + else ( + lock.mode.value, + {lg: array_key(blocks) for lg, blocks in lock.blocks.items()}, + getattr(lock._ready, "name", lock._ready), + ), + entries=None + if entries is None + else ( + entries.layer_group_indices, + [array_key(o) for o in entries.block_ordinals], + [array_key(p) for p in entries.device_pages], + ), + fetch=None + if fetch is None + else ( + fetch.start, + fetch.end, + fetch.state_end, + fetch.marks_queued, + fetch.entries is entries, + ), + ) + + +def state_key(w): + ledger, locks = w.lender._ledger, w.lender._locks + fetch = ledger._latest_fetches.get(RID) + kept = ledger._kept_prefixes.get(RID) + return dict( + progress_round=w.lender._progress_round, + latest_fetch=None + if fetch is None + else (fetch.start, fetch.end, fetch.state_end, fetch.marks_queued), + kept_prefix=None if kept is None else (kept[1], kept[2]), + deliveries=RID in ledger._deliveries, + mixed_tail=RID in ledger._mixed_tail_caches, + locks_version=locks.version, + locks_holding=locks.holding, + locks=[ + (lock.mode.value, {lg: blocks.tolist() for lg, blocks in lock.blocks.items()}) + for cache in locks._caches.values() + for lock in cache.locks + ], + ) + + +def outcome(w): + lend = w.lender.lend_read if w.kind == "read" else w.lender.lend_write + try: + lease = lend( + w.request, w.start, w.end, memory_kind=DEVICE, rows=ALL, ready_event=w.ready_event + ) + except Exception as error: # noqa: BLE001 - compare complete exception outcomes + return ("raise", type(error).__name__, str(error)), state_key(w) + w.lease = lease + return ("lease", lease_key(lease)), state_key(w) + + +@pytest.fixture(autouse=True) +def stand_in_rules(monkeypatch): + """The stand-in manager's own answers for the facade's configuration rules.""" + monkeypatch.setattr(_manager, "keeps_stale_pages", lambda manager: manager.keeps_stale_pages) + monkeypatch.setattr( + _ranges, "_floor_follows_history", lambda manager: manager.floor_follows_history + ) + + +@pytest.fixture +def taken(monkeypatch): + """Fast lends that gave a lease, by kind, and the whole-delivery bounds fast returns gave.""" + counts = {"read": 0, "write": 0, "bounds": 0} + for owner, name, key in ( + (_lender._ManagerLender, "_read_every_row_fast", "read"), + (_lender._ManagerLender, "_write_every_row_fast", "write"), + (_fetch_ledger._FetchLedger, "_whole_delivery_bounds", "bounds"), + ): + + def counted(self, *args, _taken=getattr(owner, name), _key=key, **kwargs): + found = _taken(self, *args, **kwargs) + counts[_key] += found is not None + return found + + monkeypatch.setattr(owner, name, counted) + return counts + + +def test_a_windowed_read_mid_context_lends_short_windows_as_the_rules_do(taken): + """Mid-context reads over short windows are never lent fast and lend what the rules lend.""" + kept = 0 + prompt = 6 * TPB + edges = sorted( + { + v + for b in range(1, prompt // TPB + 1) + for v in (b * TPB - 1, b * TPB, b * TPB + 1) + if TPB <= v < prompt + } + ) + boundary_windows = { + 1, + TPB - 1, + TPB, + TPB + 1, + 2 * TPB - 1, + 2 * TPB, + 2 * TPB + 1, + prompt - TPB - 1, + prompt - TPB, + prompt - TPB + 1, + prompt - 1, + } + windows = sorted(boundary_windows | set(random.Random(0).sample(range(1, prompt), 16))) + for window in windows: + for history in edges: + for position in edges: + for end in range(TPB, min(history, position) + 1, TPB): + fast = mid_context_world( + _lender._ManagerLender, window, history, position, end, prompt + ) + rules = mid_context_world(Rules, window, history, position, end, prompt) + assert fast.lender._lends_fast() + got, want = outcome(fast), outcome(rules) + assert got == want, (window, history, position, end) + behind = history + 1 - window + kept += behind < 0 or behind // TPB < end // TPB + assert kept >= 100, f"inconclusive: {kept} cases where the window keeps rows of the range" + assert taken["read"] == 0, "a read mid-context with a short window lent fast" + + +@pytest.mark.parametrize("listed_blocks", [64, 3], ids=["pages_listed", "pages_viewed"]) +def test_fast_lends_give_what_the_rules_give(monkeypatch, taken, listed_blocks): + """On drawn worlds, a lend that may take a fast path gives what the rules alone give.""" + monkeypatch.setattr(_lender, "_LISTED_BLOCKS", listed_blocks) + fell_back = past_the_slots = windowed = 0 + for seed in range(3000): + fast, rules = world(seed, _lender._ManagerLender), world(seed, Rules) + assert fast.lender._lends_fast() == fast.lender._fast_layout + assert not rules.lender._lends_fast() + before = taken["read"] + taken["write"] + got, want = outcome(fast), outcome(rules) + assert got == want, f"seed {seed}" + took = taken["read"] + taken["write"] > before + fell_back += not took + windowed += took and any(w is not None for w in fast.manager.windows) + past_the_slots += "slots of its device regions" in repr(want) + assert min(taken["read"], taken["write"]) >= 400 and fell_back >= 800, (taken, fell_back) + assert windowed >= 100, windowed + assert past_the_slots >= 50, past_the_slots + + +# A lend's bounds the call refuses, or takes as another integer type, by case: (start, end) of a +# prompt of ``prompt`` tokens. +ODD_BOUNDS = { + "unaligned_start": lambda prompt: (1, prompt), + "start_past_end": lambda prompt: (TPB, 0), + "negative_start": lambda prompt: (-TPB, prompt), + "end_past_prompt": lambda prompt: (0, prompt + 1), + "numpy_bounds": lambda prompt: (np.int64(0), np.int64(prompt)), +} + + +@pytest.mark.parametrize("case", sorted(ODD_BOUNDS)) +@pytest.mark.parametrize("kind", ["read", "write"]) +def test_a_lend_the_fast_entry_could_take_checks_its_bounds_as_the_rules_do(kind, case): + """Odd bounds of a lend the fast entry could take give what the rules alone give.""" + compared = 0 + for seed in range(400): + fast, rules = world(seed, _lender._ManagerLender), world(seed, Rules) + if not fast.lender._lends_fast(): + continue + for w in (fast, rules): + w.kind = kind + w.start, w.end = ODD_BOUNDS[case](w.request.prompt_len) + assert outcome(fast) == outcome(rules), f"seed {seed}" + compared += 1 + assert compared >= 100, compared + + +# How the returns' draws change the world between the lend and the return, with their weights. +PERTURBATIONS = ( + ("none", 40), + ("shrink_reported", 6), + ("shrink_reported_and_regrown", 6), + ("shrink_unreported", 5), + ("cache_replaced", 5), + ("cache_freed", 5), + ("cache_suspended", 4), + ("history_moved", 6), + ("commit_at_unlock", 6), + ("lender_stopped", 3), + ("earlier_deliveries", 5), + ("fetch_replaced", 4), +) + + +def perturb(w, rng): + """Change ``w``'s world as the draw says, the same way in both worlds of a seed; the name.""" + name = rng.choices([p for p, _ in PERTURBATIONS], [n for _, n in PERTURBATIONS])[0] + kv, lender, ledger = w.kv, w.lender, w.lender._ledger + if name == "shrink_reported": + lender._on_shrink(RID, kv) + elif name == "shrink_reported_and_regrown": + if kv.num_blocks: + kv.num_blocks -= 1 + lender._on_shrink(RID, kv) + kv.num_blocks += 1 + # A block regrown takes a page of its own. + for pages in kv.pages: + if pages[kv.num_blocks - 1] >= 0: + pages[kv.num_blocks - 1] = (pages[kv.num_blocks - 1] + 29) % SLOTS + elif name == "shrink_unreported": + kv.num_blocks = max(0, kv.num_blocks - rng.randint(1, 3)) + elif name == "cache_replaced": + status = SimpleNamespace(complete=True, endpoint=0) + other = Cache(rng, len(kv.pages), kv.history_length, 0, [status] * len(kv.pages), True) + w.manager.kv_cache_map[RID] = other + elif name == "cache_freed": + del w.manager.kv_cache_map[RID] + lender._on_free(RID, kv) + elif name == "cache_suspended": + kv.is_active = False + elif name == "history_moved": + kv.history_length = rng.randint(0, 7 * TPB) + elif name == "commit_at_unlock": + kv.commit_at_unlock = rng.randint(kv.num_committed_tokens, 7 * TPB) + w.manager._held_commits[id(kv)] = (kv, w.request) + elif name == "lender_stopped": + lender._stop_reason = "the lender stopped" + elif name == "earlier_deliveries": + masks = [np.array([rng.random() < 0.7 for _ in range(8)]) for _ in kv.pages] + ledger._deliveries[RID] = _fetch_ledger._Deliveries(kv, 0, masks) + elif name == "fetch_replaced": + ledger._latest_fetches[RID] = _fetch_ledger._FetchRecord( + 0, kv, marks_queued=True, end=rng.randint(0, w.request.prompt_len) + ) + return name + + +# What reads or changes the request's records between its return and the lender's next round, +# with weights: nothing, the round, or a reader of the deliveries a fast return holds until then. +AFTER_RETURN = ( + ("none", 40), + ("progress", 10), + ("records_read", 10), + ("shrink_reported", 10), + ("cache_freed", 10), +) + + +def after_return(w, rng): + """Read or change ``w``'s records after the return as the draw says, the same way in both + worlds of a seed; the name.""" + name = rng.choices([p for p, _ in AFTER_RETURN], [n for _, n in AFTER_RETURN])[0] + kv, lender = w.kv, w.lender + if name == "progress": + lender.progress() + elif name == "records_read": + lender._ledger._deliveries_of(RID) + elif name == "shrink_reported": + if kv.num_blocks: + kv.num_blocks -= 1 + lender._on_shrink(RID, kv) + elif name == "cache_freed": + if w.manager.kv_cache_map.get(RID) is kv: + del w.manager.kv_cache_map[RID] + lender._on_free(RID, kv) + return name + + +def arrivals(view, rng): + """What the draw says arrived of ``view``: every row, as the constant or as true masks, writable + or read-only, some rows, or none.""" + how = rng.choice(("constant", "true", "true_read_only", "some", "some", "none")) + if how == "constant": + return ALL_ARRIVED + masks = [np.full(len(rows), how != "none", dtype=bool) for rows in view.groups] + if how == "some": + masks = [np.array([rng.random() < 0.7 for _ in range(len(m))], dtype=bool) for m in masks] + if how == "true_read_only": + for mask in masks: + mask.flags.writeable = False + return masks + + +def label(w, obj): + """``obj`` named by what it is in ``w``'s world, so the two worlds compare.""" + if obj is w.kv: + return "cache" + if obj is None: + return None + return type(obj).__name__ + + +def records_key(w): + """The ledger's records of the request, the request's shrink count and the locks.""" + ledger = w.lender._ledger + fetch = ledger._latest_fetches.get(RID) + delivered = ledger._deliveries.get(RID) + kept = ledger._kept_prefixes.get(RID) + return dict( + latest=None + if fetch is None + else ( + label(w, fetch.kv), + fetch.start, + fetch.end, + fetch.state_end, + fetch.marks_queued, + fetch.copy, + ), + deliveries=None + if delivered is None + else ( + label(w, delivered.kv), + delivered.lowest_start, + [mask.tolist() for mask in delivered.delivered_block_masks], + delivered.state_end, + delivered.end, + sorted(delivered.partial_ends.items()), + delivered.resume_bounds_memo, + ), + kept=None if kept is None else (label(w, kept[0]), kept[1], kept[2]), + mixed=label(w, ledger._mixed_tail_caches.get(RID)), + shrinks=w.lender._shrink_epochs.get(RID), + state=state_key(w), + ) + + +def return_outcome(w, rng): + """Poll ``w``'s lease, change the world as drawn, return it with what the draw says arrived, + then query the bounds again: what each gave and the records left.""" + lease = w.lease + view = lease.poll() + if view is None: + return ("no view", lease.failure) + perturbation = perturb(w, rng) + arrived = arrivals(view, rng) + got = _outcome(lambda: w.lender.return_writes([(lease, arrived, True)])) + after = after_return(w, rng) + try: + again = ("bounds", w.lender.poll_resume_bounds(w.request)) + except ValueError as error: + again = ("raise", str(error)) + return perturbation, got, after, again, (lease._released, lease._marked), records_key(w) + + +def generic_bounds(w): + """``w`` with its ledger giving a whole delivery's bounds through ``_ranges``.""" + w.lender._ledger._closed_form = False + return w + + +# How a fast return's ledger gives a whole delivery's bounds: in closed form, as the layouts fast +# lends cover allow, or through ``_ranges``, as for any other layout. +BOUNDS_FORMS = {"closed_form": lambda w: w, "by_ranges": generic_bounds} + + +@pytest.mark.parametrize("form", sorted(BOUNDS_FORMS)) +def test_fast_returns_give_what_the_rules_give(taken, form): + """Returns after drawn changes to the world give the bounds and records the rules give.""" + compared = fast_taken = fell_back = windowed = committed = agreed = 0 + perturbed = {name: 0 for name, _ in PERTURBATIONS} + afterwards = {name: 0 for name, _ in AFTER_RETURN} + for seed in range(7000): + fast, rules = BOUNDS_FORMS[form](world(seed, _lender._ManagerLender)), world(seed, Rules) + got, want = outcome(fast), outcome(rules) + assert got == want, f"seed {seed}" + if fast.lease is None or fast.kind != "write": + continue + before = taken["bounds"] + got = return_outcome(fast, random.Random(seed * 7919 + 1)) + want = return_outcome(rules, random.Random(seed * 7919 + 1)) + assert got == want, f"seed {seed}" + if got[0] == "no view": + continue + compared += 1 + perturbed[got[0]] += 1 + afterwards[got[2]] += 1 + if got[2] == "none" and got[1][0] == "ok" and got[1][1][0] is not None: + # The bounds a return gives are those the next query gives. + assert got[3] == ("bounds", got[1][1][0]), f"seed {seed}" + agreed += 1 + took = taken["bounds"] > before + fast_taken += took + fell_back += not took + windowed += took and any(w is not None for w in fast.manager.windows) + committed += got[0] == "commit_at_unlock" + assert compared >= 1000 and fast_taken >= 300 and fell_back >= 300, ( + compared, + fast_taken, + fell_back, + ) + assert windowed >= 30 and committed >= 5 and agreed >= 300, (windowed, committed, agreed) + assert min(perturbed.values()) >= 10, perturbed + assert min(afterwards.values()) >= 50, afterwards + + +def test_a_write_view_lends_no_region_slot_twice(): + """No write view over caches of distinct pages, fast or by the rules, lends a slot twice.""" + checked = 0 + for seed in range(2000): + for cls in (_lender._ManagerLender, Rules): + w = world(seed, cls) + pages = [p[: w.kv.num_blocks] for p in w.kv.pages] + if w.kind != "write" or any( + len(set(p[p >= 0].tolist())) < (p >= 0).sum() for p in pages + ): + continue + outcome(w) + view = None if w.lease is None else w.lease._view + if view is None or not view.num_rows: + continue + lent = [ + (region, slot) + for rows in view.groups + for slot in rows.slots.tolist() + for region in rows.region_indices + ] + assert len(lent) == len(set(lent)), f"seed {seed}" + checked += 1 + assert checked >= 300, checked + + +@pytest.mark.parametrize("form", sorted(BOUNDS_FORMS)) +def test_a_fast_return_over_windows_gives_the_floor_the_rules_give(kit, taken, form): + """Fast returns over windows with read slack give the floor and bounds the rules give.""" + compared = floored = left = 0 + for seed in range(3000): + outcomes = [] + for cls in (_lender._ManagerLender, Rules): + w = windowed_receive_world(seed, cls) + if cls is _lender._ManagerLender: + BOUNDS_FORMS[form](w) + lease = w.lender.lend_write(w.request, w.start, w.end, **kit.DEVICE_ALL) + view = lease.poll() + if view is None: + outcomes.append(("no view", lease.failure)) + continue + before = taken["bounds"] + got = w.lender.return_writes([(lease, ALL_ARRIVED, True)]) + delivered = w.lender._ledger._deliveries.get(RID) + took = taken["bounds"] > before + uncomputed = took and delivered is not None and delivered.resume_bounds_memo is None + outcomes.append((got, w.lender.poll_resume_bounds(w.request), took, uncomputed)) + fast, rules = outcomes + assert fast[:2] == rules[:2], f"seed {seed}" + if fast[0] == "no view": + continue + compared += 1 + bounds = fast[0][0] + floored += fast[2] and bounds is not None and bounds.min_position > w.kv.history_length + left += fast[3] + assert compared >= 1500, compared + assert floored >= 50, floored + if form == "closed_form": + assert left >= 50, left + + +@pytest.mark.parametrize("replace", ["class", "instance", "ledger"]) +@pytest.mark.parametrize("refusal", ["", "the replaced rule refuses"]) +def test_fast_lends_use_canonical_range_rules(monkeypatch, kit, replace, refusal): + """Fast lends use the range rules directly; a ledger subclass takes the general path.""" + lender, manager, kv = plain_world( + _lender._ManagerLender, layout([False]), TPB, random.Random(4) + ) + w = world_of(lender, manager, kv, request_of(TPB, position=TPB), 0, TPB) + assert w.lender._lends_fast() + if replace == "class": + monkeypatch.setattr(_lender._ManagerLender, "_read_end_refusal", lambda *args: refusal) + monkeypatch.setattr(_lender._ManagerLender, "_write_range_refusal", lambda *args: refusal) + elif replace == "instance": + w.lender._read_end_refusal = lambda *args: refusal + w.lender._write_range_refusal = lambda *args: refusal + else: + w.lender._ledger = type("Ledger", (_fetch_ledger._FetchLedger,), {})( + w.lender._layout, w.lender._ledger._copy_done + ) + assert w.lender._lends_fast() == (replace != "ledger") + if replace != "ledger": + for lend in (w.lender.lend_read, w.lender.lend_write): + lease = lend(w.request, 0, w.request.prompt_len, **kit.DEVICE_ALL) + assert lease.failure == refusal, lease.failure + lease.release() + + +def test_a_replaced_mark_rule_takes_effect_on_a_return(): + """A fast write returned through the general path counts only still held rows.""" + seen = [] + for seed in range(400): + w = world(seed, _lender._ManagerLender) + if not w.lender._lends_fast() or w.kind != "write": + continue + outcome(w) + if w.lease is None or w.lease.poll() is None: + continue + w.lender._locked_block_masks = lambda kv, lease: [ + np.zeros(len(o), dtype=bool) for o in lease._block_entries.block_ordinals + ] + w.lender._fast_layout = False + assert not w.lender._lends_fast() + w.lender.return_writes([(w.lease, ALL_ARRIVED, True)]) + delivered = w.lender._ledger._deliveries.get(RID) + seen.append(delivered is None or not any(m.any() for m in delivered.delivered_block_masks)) + assert len(seen) >= 50 and all(seen), (len(seen), seen.count(False)) + + +def test_a_lender_s_own_state_leaves_its_attributes_free(): + """Lender state leaves the instance dict available, and subclasses take the general path.""" + w = world(4, _lender._ManagerLender) + assert not vars(w.lender) + w.lender.note = "set on the lender" + assert w.lender._lends_fast() + + assert not world(4, Rules).lender._lends_fast() + assert w.lender._lends_fast() + + +@pytest.mark.parametrize("source", ["cache", "row"]) +@pytest.mark.parametrize("first", [0, 1]) +@pytest.mark.parametrize("offset", [-1, 0, 1, 2]) +@pytest.mark.parametrize( + "boundary", [_lender._LISTED_BLOCKS, _lender._ORDINAL_BLOCKS], ids=["listed", "ordinals"] +) +def test_fast_rows_hold_each_block_around_the_shared_ordinals(first, offset, boundary, source): + """Rows near the shared-ordinal and listed-page bounds hold each block's page, in own arrays.""" + last = boundary + offset + lender, kv, pages = row_lender(last) + row = None + if source == "row": + # The cache's row of a page-index buffer between others: its pages, then the bad index. + rows = np.full((1, 3, 2, len(pages) + 4), -1, dtype=np.int32) + rows[0, 1, 0, : len(pages)] = pages + lender._read_index_rows(rows) + row = buffer_row(lender, 1) + entries, view, blocks = lender._every_row(kv, RID, first * TPB, last * TPB, row=row) + assert entries.block_ordinals[0].tolist() == list(range(first, last)) + assert entries.device_pages[0].tolist() == kv.pages[0][first:last].tolist() + assert entries.spans == [(first, last)] + assert blocks == {0: entries.block_ordinals[0]} + (rows,) = view.groups + assert rows.page_ordinals.tolist() == list(range(first, last)) + assert rows.slots.tolist() == entries.device_pages[0].tolist() + assert rows.page_ordinals is not entries.block_ordinals[0] + assert rows.slots is not entries.device_pages[0] + + +@pytest.mark.parametrize("blocks", [3, _lender._LISTED_BLOCKS + 3], ids=["listed", "viewed"]) +def test_a_range_past_the_end_of_its_row_is_not_read_from_the_buffer(blocks): + """A buffer row a block shorter than the range lends none of it; the shorter range, all.""" + lender, kv, pages = row_lender(blocks) + # Every row and half of a row holds the cache's pages of its first ``blocks - 1`` blocks. + rows = np.empty((1, 3, 2, blocks - 1), dtype=np.int32) + rows[...] = pages[: blocks - 1] + lender._read_index_rows(rows) + shorter = lender._every_row(kv, RID, 0, (blocks - 1) * TPB, row=buffer_row(lender, 1)) + assert shorter is not None, "inconclusive: the row lends nothing" + assert shorter[0].device_pages[0].tolist() == pages[: blocks - 1].tolist() + assert lender._every_row(kv, RID, 0, blocks * TPB, row=buffer_row(lender, 1)) is None + + +@pytest.mark.parametrize("past", ["no_page", "at_the_slot_count"]) +@pytest.mark.parametrize("blocks", [3, _lender._LISTED_BLOCKS + 3], ids=["listed", "viewed"]) +def test_a_row_with_a_page_outside_the_pool_lends_nothing_from_the_buffer(blocks, past): + """A buffer row with no page, or a page at the slot count, for one block lends nothing.""" + lender, kv, pages = row_lender(blocks) + rows = np.full((1, 3, 2, len(pages)), -1, dtype=np.int32) + rows[0, 1, 0] = pages + lender._read_index_rows(rows) + entries = lender._every_row(kv, RID, 0, blocks * TPB, row=buffer_row(lender, 1))[0] + assert entries.device_pages[0].tolist() == pages[:blocks].tolist() + ((_, _, limit, _, _),) = lender._fast_groups + rows[0, 1, 0, blocks // 2] = -1 if past == "no_page" else limit + assert lender._every_row(kv, RID, 0, blocks * TPB, row=buffer_row(lender, 1)) is None + + +def test_the_views_of_the_shared_ordinals_stay_bounded(monkeypatch): + """At most ``_ORDINAL_VIEWS_KEPT`` views of the shared ordinals stay, each of its own range.""" + monkeypatch.setattr(_lender, "_ordinal_views", {}) + monkeypatch.setattr(_lender, "_ORDINAL_VIEWS_KEPT", 4) + lender, kv, _ = row_lender(12) + emptied = 0 + for first in range(3): + for last in range(first + 1, 12): + held = len(_lender._ordinal_views) + entries = lender._every_row(kv, RID, first * TPB, last * TPB)[0] + assert entries.block_ordinals[0].tolist() == list(range(first, last)) + emptied += len(_lender._ordinal_views) <= held + assert len(_lender._ordinal_views) <= 4 + assert emptied >= 5, emptied + + +def view_arrays(view): + """The arrays of ``view``'s rows.""" + return [ + array + for rows in view.groups + for array in (rows.page_ordinals, rows.slots, rows.object_keys) + if array is not None + ] + + +def lender_arrays(lease): + """The arrays the lender keeps for the lease: its entries' and its lock's.""" + entries = lease._block_entries + held = [*entries.block_ordinals, *entries.device_pages, *entries.staging_slots] + return held + ([] if lease._lock is None else list(lease._lock.blocks.values())) + + +@pytest.mark.parametrize("listed_blocks", [64, 3], ids=["pages_listed", "pages_viewed"]) +@pytest.mark.parametrize("cls", [_lender._ManagerLender, Rules], ids=["fast", "rules"]) +def test_each_view_holds_array_objects_of_its_own(monkeypatch, kit, taken, listed_blocks, cls): + """Views of two reads of the same rows share no array object with each other or the lender.""" + monkeypatch.setattr(_lender, "_LISTED_BLOCKS", listed_blocks) + compared = 0 + for seed in range(600): + w = world(seed, cls) + if w.kind != "read" or not 0 <= w.start <= w.end: + continue + leases = [w.lender.lend_read(w.request, w.start, w.end, **kit.DEVICE_ALL) for _ in range(2)] + if all(lease._view is not None for lease in leases): + first, second = (view_arrays(lease._view) for lease in leases) + assert len({id(a) for a in first + second}) == len(first) + len(second), f"seed {seed}" + kept = {id(a) for lease in leases for a in lender_arrays(lease)} + assert not kept & {id(a) for a in first + second}, f"seed {seed}" + compared += 1 + for lease in leases: + lease.release() + assert compared >= 100, compared + assert taken["read"] >= (100 if cls is _lender._ManagerLender else 0), taken + + +@pytest.mark.parametrize("listed_blocks", [64, 3], ids=["pages_listed", "pages_viewed"]) +@pytest.mark.parametrize("kind", ["read", "write"]) +@pytest.mark.parametrize("past", ["page", "state_slot"]) +def test_a_slot_at_the_limit_leaves_the_lend_to_the_rules( + monkeypatch, taken, listed_blocks, kind, past +): + """A page or state slot at its group's device slot limit leaves the lend to the rules.""" + monkeypatch.setattr(_lender, "_LISTED_BLOCKS", listed_blocks) + slots_per_group = (4, 6) # the state group's pool, then the attention group's + limit = slots_per_group[1 if past == "page" else 0] + for value in (limit - 1, limit, limit + 1): + page, slot = (value, 2) if past == "page" else (2, value) + fast, rules = ( + hybrid_world(cls, kind, slots_per_group, page, slot) + for cls in (_lender._ManagerLender, Rules) + ) + before = taken[kind] + got, want = outcome(fast), outcome(rules) + assert got == want + assert taken[kind] == before + (value < limit) + (shape, lease), state = got + assert shape == "lease" + if value < limit: + assert lease["failure"] is None and lease["first_poll_failure"] is None + assert [view[2] for view in lease["view"]] == [ + (" 0 and ledger._kept_prefixes.get(RID) is not before + ) + outcomes.append( + ( + fetch is None, + None if latest is None else (latest.start, latest.end, latest.state_end), + latest is fetch, + ledger._kept_prefixes.get(RID, (None,))[1:], + ) + ) + assert outcomes[0] == outcomes[1], f"seed {seed}" + refused += outcomes[0][0] + assert refused >= 20 and kept_past_start >= 20, (refused, kept_past_start) + + +@pytest.mark.parametrize("sinks", [0, 1, 3]) +def test_the_window_formula_answers_as_the_stand_in_runtime_and_the_last_keeping_history(sinks): + """The window formula gives the runtime's stale ranges and the searched last keeping history.""" + for window in (1, 2, TPB - 1, TPB, TPB + 1, 3 * TPB + 5, 64 * TPB): + manager = Manager(1, None, False, False, (window,), (sinks,)) + searched = layout([False], windows=(window,), sinks=(sinks,), formula=False) + formula = layout([False], windows=(window,), sinks=(sinks,), formula=True) + for history in range(0, 6 * window + 8 * TPB): + assert _ranges._formula_stale_range(window, sinks, TPB, history) == ( + runtime_stale_range(window, sinks, history) + ), (window, sinks, history) + for block in range(sinks, sinks + 12): + assert _ranges._last_history_keeping(manager, formula, 0, block) == ( + _ranges._last_history_keeping(manager, searched, 0, block) + ), (window, sinks, block) + huge = _manager._MAX_TOKENS - 3 * TPB + manager = Manager(1, None, False, False, (huge,), (0,)) + for block in (0, 1, 5): + assert _ranges._last_history_keeping( + manager, layout([False], windows=(huge,), formula=True), 0, block + ) == _ranges._last_history_keeping( + manager, layout([False], windows=(huge,), formula=False), 0, block + ) + + +def test_an_attach_probe_finds_a_manager_off_the_formula(): + """The attach probe accepts a manager on the window formula and refuses one a block off.""" + lay = layout([False, False], windows=(3 * TPB + 2, None), formula=False) + manager = Manager(2, None, False, False, lay.windows, (0, 0)) + assert _ranges._stale_ranges_follow_formula(manager, lay) + honest = manager._stale_block_range + + def off_by_one(lg, history): + beg, end = honest(lg, history) + return (beg, end + 1) if history == 3 * TPB + 2 + TPB - 1 else (beg, end) + + manager._stale_block_range = off_by_one + assert not _ranges._stale_ranges_follow_formula(manager, lay) + + +def test_a_fast_lend_builds_its_rows_as_the_unchecked_constructor_does(): + """A fast lend's rows are what ``LentLayerGroupRows._unchecked`` makes of the same values.""" + ordinals, slots = np.arange(3, dtype=np.int64), np.arange(5, 8, dtype=np.int64) + fast = _lender._lent_rows(2, ordinals, slots, (4, 5)) + unchecked = LentLayerGroupRows._unchecked(2, ordinals, slots=slots, region_indices=(4, 5)) + assert type(fast) is type(unchecked) + assert vars(fast).keys() == vars(unchecked).keys() + assert all(vars(fast)[name] is vars(unchecked)[name] for name in vars(unchecked)) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_page_locks.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_page_locks.py new file mode 100644 index 000000000000..bf22ba95f1db --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_page_locks.py @@ -0,0 +1,2371 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The page-lock registry, the manager's lock hooks and the generation update under locks.""" + +from __future__ import annotations + +import ast +import copy +import enum +import gc +import numbers +import random +import traceback +import types +import weakref +from collections import Counter +from pathlib import Path +from types import SimpleNamespace +from typing import Callable, List + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + ALL_ARRIVED, + HoldReason, + StagingOptions, + _lender, + _manager, + _page_locks, + attach_lender, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._page_locks import LockMode, PageLocks + +gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") + +READ, WRITE = LockMode.READ, LockMode.WRITE +TPB = 32 + + +# -- the registry over stand-ins ------------------------------------------------------------------ + + +class Manager: + """A stand-in manager: three layer groups, the first and last on pool group 0. It notes each + cache it retries a held-back commit of, with that cache's closes then; with ``held``, a record + of held commits by cache ``id`` as the KV cache manager keeps, it retries only those.""" + + def __init__(self, pool_group_of=(0, 1, 0), held=None): + self.impl = SimpleNamespace(get_life_cycle_pool_group_indices=lambda: list(pool_group_of)) + self.retried = [] + if held is not None: + self._held_commits = held + + def _commit_held_back(self, kv): + held = getattr(self, "_held_commits", None) + if held is None or held.pop(id(kv), None) is not None: + self.retried.append((kv, kv.closes)) + + +class Cache: + """A stand-in runtime cache that counts its closes: its committed tokens in each of three layer + groups and, per beam, a page per block of each layer group, -1 for none.""" + + class Status(enum.Enum): + OPEN = 0 + CLOSED = 1 + + def __init__(self, committed_tokens=0, beams=None): + self.status = self.Status.OPEN + self.num_committed_tokens = committed_tokens + self.reuse_status = [SimpleNamespace(complete=True, endpoint=committed_tokens)] * 3 + self.tokens_per_block = TPB + self._beams = beams or [{}] + self.beam_width = len(self._beams) + self.closes = 0 + + def close(self): + self.closes += 1 + self.status = self.Status.CLOSED + + def get_scratch_desc(self, lg): + return None + + def get_aggregated_page_indices(self, lg, beam_id=0, valid_only=False): + pages = self._beams[beam_id].get(lg, []) + return [p for p in pages if p >= 0] if valid_only else pages + + +class Event: + """A stand-in CUDA event: complete once ``done`` is set; waiting on it fails the test.""" + + def __init__(self, done=False): + self.done = done + self.queries = 0 + + def query(self): + self.queries += 1 + return self.done + + def synchronize(self): + raise AssertionError("a lock waited on its event") + + +@pytest.fixture +def recorded(monkeypatch): + """The CUDA events made while the test runs.""" + events, make = [], torch.cuda.Event + + def counted(*args, **kwargs): + events.append(make(*args, **kwargs)) + return events[-1] + + monkeypatch.setattr(torch.cuda, "Event", counted) + return events + + +MANAGER = Manager() # the registry holds its manager weakly + + +def registry(manager=MANAGER, recurrent=()): + return PageLocks(weakref.ref(manager), recurrent) + + +def locked(owner): + """The locked pages per pool group ``owner.held_page_counts()`` gives.""" + return dict(owner.held_page_counts().by_reason.get(HoldReason.LOCKED, {})) + + +def deducted(owner): + return dict(owner.held_page_counts().deduct) + + +@pytest.mark.cpu_only +def test_readers_share_blocks_and_a_writer_shares_none(recorded): + locks, kv, other = registry(), Cache(), Cache() + first = locks.lock(kv, {0: [1, 2, 3]}, READ) + second = locks.lock(kv, {0: [3, 4], 1: [0]}, READ) + assert {first.mode, second.mode} == {READ} + for blocks in ({0: [3]}, {0: [0, 4]}, {1: [0]}): + with pytest.raises(ValueError, match="locked for read"): + locks.lock(kv, blocks, WRITE) + writer = locks.lock(kv, {0: [5, 6], 2: [3]}, WRITE) # other blocks, another layer group + for blocks, mode in (({0: [6]}, READ), ({0: [6]}, WRITE), ({2: [0, 3]}, READ)): + with pytest.raises(ValueError, match="locked for write"): + locks.lock(kv, blocks, mode) + # Another cache's blocks and another layer group's ordinals are other blocks. + locks.lock(other, {0: [1, 2, 3, 5, 6]}, WRITE) + locks.lock(kv, {1: [5, 6]}, WRITE) + locks.unlock(writer) + again = locks.lock(kv, {0: [5, 6], 2: [3]}, WRITE) + assert first.is_ready() and second.is_ready() and again.is_ready(), "a lock without an event" + assert recorded == [], "a lock made an event" + + +@pytest.mark.cpu_only +def test_a_refused_lock_changes_nothing(recorded): + locks, kv = registry(), Cache(committed_tokens=2 * TPB) + reader = locks.lock(kv, {0: [2, 3]}, READ) + before = (locked(locks), locks.is_locked(kv, 4), len(recorded)) + refused = [ + lambda: locks.lock(kv, {0: [3, 9]}, WRITE), # a reader holds block 3 + lambda: locks.lock(kv, {0: [1, 9]}, WRITE), # block 1 is committed + lambda: locks.lock(kv, {1: [9], 0: [-1]}, READ), + lambda: locks.lock(kv, {3: [9]}, READ), # three layer groups + lambda: locks.lock(kv, {0: [9]}, READ, level=1), + ] + for call in refused: + with pytest.raises(ValueError): + call() + assert (locked(locks), locks.is_locked(kv, 4), len(recorded)) == before + locks.unlock(reader) + assert not locks.is_locked(kv) and locked(locks) == {} + + +@pytest.mark.cpu_only +def test_a_lock_another_registry_holds_is_not_released_here(): + """Unlocking another registry's lock raises and changes neither registry.""" + mine, theirs, kv = registry(), registry(), Cache() + held = mine.lock(kv, {0: [1, 2]}, READ) + foreign = theirs.lock(kv, {0: [5]}, READ) + before = (mine.holding, mine.version, mine.is_locked(kv, 1, 3)) + with pytest.raises(ValueError, match="not held in this registry"): + mine.unlock(foreign) + assert (mine.holding, mine.version, mine.is_locked(kv, 1, 3)) == before + assert not foreign.released and theirs.is_locked(kv, 5, 6) + for locks, lock in ((mine, held), (theirs, foreign)): + locks.unlock(lock) + assert not mine.is_locked(kv) and not theirs.is_locked(kv) + + +@pytest.mark.cpu_only +def test_a_writer_locks_only_blocks_past_the_committed_whole_blocks(): + locks = registry() + kv = Cache(committed_tokens=2 * TPB + 5) # 2 whole blocks, and 5 tokens of the third + with pytest.raises(ValueError, match="past the 2 committed whole blocks"): + locks.lock(kv, {0: [2, 3], 1: [1, 2]}, WRITE) + writer = locks.lock(kv, {0: [2, 3], 1: [2]}, WRITE) + reader = locks.lock(kv, {0: [0, 1], 1: [0, 1]}, READ) # committed blocks are read + assert locks.is_locked(kv, 0, 1) and locks.is_locked(kv, 3, 4) + for lock in (writer, reader): + locks.unlock(lock) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "blocks, mode, level, error", + [ + ({0: [1]}, "read", None, TypeError), + ({0: [1]}, READ, 0, ValueError), + ({True: [1]}, READ, None, TypeError), + ({1.0: [1]}, READ, None, TypeError), + ({-1: [1]}, READ, None, ValueError), + ({3: [1]}, READ, None, ValueError), + ({0: [1.5]}, READ, None, ValueError), + ({0: [[1]]}, READ, None, ValueError), + ({0: [True]}, READ, None, ValueError), + ({0: [0, -2]}, READ, None, ValueError), + ], +) +def test_a_lock_takes_integer_blocks_of_known_layer_groups_in_place( + recorded, blocks, mode, level, error +): + locks, kv = registry(), Cache() + with pytest.raises(error): + locks.lock(kv, blocks, mode, level) + assert not locks.is_locked(kv) and recorded == [] + + +@pytest.mark.cpu_only +def test_a_lock_holds_sorted_distinct_read_only_ordinals(): + locks, kv = registry(), Cache() + given = np.array([7, 3, 3, 5], dtype=np.int32) + lock = locks.lock(kv, {2: given, 0: range(2), 1: []}, READ) + assert {lg: o.tolist() for lg, o in lock.blocks.items()} == {2: [3, 5, 7], 0: [0, 1], 1: []} + for ordinals in lock.blocks.values(): + assert ordinals.dtype == np.int64 and not ordinals.flags.writeable + with pytest.raises(TypeError): + lock.blocks[3] = np.zeros(1) # a read-only mapping + given[0] = 0 + assert lock.blocks[2].tolist() == [3, 5, 7], "the lock shares the caller's array" + + +def frozen(values): + """Ordinals as the lender makes them: sorted, distinct, ``int64`` on immutable bytes.""" + return np.frombuffer(np.asarray(values, dtype=np.int64).tobytes(), dtype=np.int64) + + +@pytest.mark.cpu_only +def test_a_lock_of_ordered_ordinals_takes_them_as_they_are_and_still_refuses(recorded): + """A lock takes the lender's ordered arrays as they are and refuses them as any blocks.""" + locks, kv = registry(), Cache(committed_tokens=2 * TPB) + given = {0: frozen([2, 3, 5]), 1: frozen([]), 2: frozen([0])} + reader = locks.lock(kv, given, READ, ordered=True) + assert all(reader.blocks[lg] is given[lg] for lg in given) + assert locks.is_locked(kv, 5, 6) and not locks.is_locked(kv, 4, 5) + refused = ( + ({0: frozen([3, 9])}, WRITE, "locked for read"), + ({0: frozen([1, 9])}, WRITE, "past the 2 committed whole blocks"), + ({3: frozen([9])}, READ, "the manager has 3 layer groups"), + ) + for blocks, mode, match in refused: + with pytest.raises(ValueError, match=match): + locks.lock(kv, blocks, mode, ordered=True) + found = locks.check(kv, blocks, mode, ordered=True) + assert isinstance(found, str) and found == locks.refusal(kv, blocks, mode) + # Committed tokens the caller read stand for the cache's own count. + past = {0: frozen([1, 9])} + for committed in (None, (2 * TPB,) * 3): + assert "past the 2" in locks.check(kv, past, WRITE, ordered=True, committed=committed) + assert isinstance( + locks.check(kv, past, WRITE, ordered=True, committed=(0,) * 3), _page_locks.CheckedBlocks + ) + assert reader.is_ready() + locks.unlock(reader) + assert recorded == [] + + +@pytest.mark.cpu_only +def test_a_checked_lock_takes_what_its_check_found_until_a_lock_or_unlock_comes_between(): + """A lock takes its check's result while no lock or unlock came between, else checks again.""" + locks, kv, other = registry(), Cache(), Cache() + blocks = {0: frozen([4, 5]), 2: frozen([1])} + checked = locks.check(kv, blocks, WRITE, ordered=True) + assert isinstance(checked, _page_locks.CheckedBlocks) + assert (checked.kv, checked.mode, checked.version) == (kv, WRITE, locks.version) + assert all(checked.blocks[lg] is blocks[lg] for lg in blocks) + writer = locks.lock(kv, blocks, WRITE, ordered=True, checked=checked) + assert all(writer.blocks[lg] is blocks[lg] for lg in blocks) + locks.unlock(writer) + checked = locks.check(kv, blocks, WRITE, ordered=True) + reader = locks.lock(kv, {0: [5]}, READ) # comes between the check and the lock + with pytest.raises(ValueError, match="locked for read"): + locks.lock(kv, blocks, WRITE, ordered=True, checked=checked) + elsewhere = locks.check(other, blocks, WRITE, ordered=True) + as_reader = locks.check(kv, {0: frozen([5])}, READ, ordered=True) + for found in (elsewhere, as_reader): + assert isinstance(found, _page_locks.CheckedBlocks) + with pytest.raises(ValueError, match="locked for read"): + locks.lock(kv, {0: frozen([5])}, WRITE, ordered=True, checked=found) + locks.unlock(reader) + assert not locks.is_locked(kv) and not locks.is_locked(other) + + +@pytest.mark.cpu_only +def test_a_checked_lock_cannot_reopen_a_cache_deferred_at_free(): + locks, kv = registry(), Cache() + first = locks.lock(kv, {0: [0]}, READ) + checked = locks.check(kv, {0: [1]}, WRITE) + assert isinstance(checked, _page_locks.CheckedBlocks) + + assert locks.defer_close(kv) + with pytest.raises(ValueError, match="freed"): + locks.lock(kv, {0: [1]}, WRITE, checked=checked) + assert locks.holding == 1 + locks.unlock(first) + assert kv.closes == 1 + + +@pytest.mark.cpu_only +def test_a_closed_cache_refuses_checked_and_new_locks(): + locks, kv = registry(), Cache() + checked = locks.check(kv, {0: [0]}, WRITE) + assert isinstance(checked, _page_locks.CheckedBlocks) + assert not locks.defer_close(kv) + kv.close() + + with pytest.raises(ValueError, match="closed"): + locks.lock(kv, {0: [0]}, WRITE, checked=checked) + with pytest.raises(ValueError, match="closed"): + locks.lock(kv, {0: [0]}, READ) + assert "closed" in locks.check(kv, {0: [0]}, READ) + assert locks.lock_fresh(kv, {0: frozen([0])}, READ, None) is None + assert not locks.is_locked(kv) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("intervening", ["none", "held", "released"]) +def test_a_checked_write_rechecks_blocks_committed_since_check(intervening): + locks, kv = registry(), Cache() + committed = (0,) * 3 + checked = locks.check(kv, {0: [1]}, WRITE, ordered=True, committed=committed) + assert isinstance(checked, _page_locks.CheckedBlocks) + if intervening != "none": + other = locks.lock(kv, {0: [3]}, READ) + if intervening == "released": + locks.unlock(other) + kv.num_committed_tokens = 2 * TPB + + with pytest.raises(ValueError, match="past the 2 committed whole blocks"): + locks.lock(kv, {0: [1]}, WRITE, checked=checked, ordered=True, committed=committed) + assert not locks.is_locked(kv, 1, 2) + + +@pytest.mark.cpu_only +def test_an_ordered_lock_on_an_unlocked_cache_takes_and_refuses_what_its_check_does(): + """Random ordered blocks: a lock on an unlocked cache takes and refuses what its check does.""" + rng = random.Random(0) + locks = registry() + for _ in range(600): + tokens = rng.choice((0, 0, rng.randint(1, 3 * TPB))) + kv = Cache(committed_tokens=tokens) + groups = rng.sample(range(4), rng.randint(0, 3)) + blocks = {lg: frozen(sorted(rng.sample(range(6), rng.randint(0, 3)))) for lg in groups} + mode = rng.choice((READ, WRITE)) + committed = rng.choice((None, (0,) * 3, (tokens,) * 3)) + found = locks.check(kv, blocks, mode, ordered=True, committed=committed) + try: + lock = locks.lock(kv, blocks, mode, ordered=True, committed=committed) + except ValueError as error: + assert found == str(error) + else: + assert isinstance(found, _page_locks.CheckedBlocks) + assert lock.blocks.keys() == blocks.keys() + assert all(lock.blocks[lg] is blocks[lg] for lg in blocks) + locks.unlock(lock) + assert not locks.is_locked(kv) + + +@pytest.mark.cpu_only +def test_readiness_only_queries_the_event_and_drops_it_once_complete(recorded): + locks, kv = registry(), Cache() + given = Event() + lock = locks.lock(kv, {0: [0]}, WRITE, ready=given) + assert recorded == [], "a lock made an event besides the one it was given" + assert not lock.is_ready() and not lock.is_ready() + assert given.queries == 2 + given.done = True + assert lock.is_ready() and lock.is_ready() + assert given.queries == 3, "a completed event was queried again" + held = weakref.ref(given) + del given + assert held() is None, "the lock keeps its event after it completed" + # Without an event, nothing in flight touches the pages: the lock is ready at once. + mine = locks.lock(kv, {0: [1]}, WRITE) + assert recorded == [] and mine.is_ready() + + +@pytest.mark.cpu_only +def test_unlock_is_idempotent_and_ends_access(): + locks, kv = registry(), Cache() + event = Event(done=True) + lock = locks.lock(kv, {0: [0, 1]}, WRITE, ready=event) + held = weakref.ref(event) + del event + locks.unlock(lock) + locks.unlock(lock) + assert lock.released and not lock.is_ready() + assert held() is None, "an unlocked lock keeps its event" + assert not locks.is_locked(kv) + again = locks.lock(kv, {0: [0, 1]}, WRITE) # the blocks are free again + with pytest.raises(ValueError, match="not held"): + registry(Manager()).unlock(again) + assert locks.is_locked(kv) + + +@pytest.mark.cpu_only +def test_the_holding_count_follows_every_lock_and_unlock(recorded): + """Through random locks and unlocks, ``holding`` counts the locks that hold a block.""" + rng = random.Random(5) + locks = registry() + caches = [Cache() for _ in range(3)] + held = [] + for _ in range(400): + if held and rng.random() < 0.4: + locks.unlock(held.pop(rng.randrange(len(held)))) + else: + blocks = { + lg: sorted(rng.sample(range(8), rng.randint(0, 2))) + for lg in range(3) + if rng.random() < 0.6 + } + try: + held.append(locks.lock(rng.choice(caches), blocks, rng.choice((READ, WRITE)))) + except ValueError: + pass + assert locks.holding == sum(any(map(len, lock.blocks.values())) for lock in held) + assert locks.holds_any() == any(locks.is_locked(kv) for kv in caches) + assert recorded == [] + + +@pytest.mark.cpu_only +def test_is_locked_reads_the_ordinal_range_across_layer_groups(): + locks, kv = registry(), Cache() + assert not locks.is_locked(kv) and not locks.holds_any() + lock = locks.lock(kv, {0: [4, 9], 2: [6]}, READ) + cases = {(0, None): True, (0, 4): False, (4, 5): True, (5, 6): False, (5, 7): True} + cases.update({(7, 9): False, (10, None): False, (9, 10): True, (6, 6): False}) + assert {r: locks.is_locked(kv, *r) for r in cases} == cases + assert not locks.is_locked(Cache()) and locks.holds_any() + empty = locks.lock(kv, {1: []}, READ) + locks.unlock(lock) + assert not locks.is_locked(kv), "a lock of no block locks the cache" + assert not locks.holds_any(), "a lock of no block counts as holding one" + assert not locks.defer_close(kv) + locks.unlock(empty) + + +@pytest.mark.cpu_only +def test_a_freed_locked_cache_closes_once_at_its_last_unlock(): + locks, kv, free = registry(), Cache(), Cache() + first = locks.lock(kv, {0: [0, 1]}, READ) + second = locks.lock(kv, {1: [0]}, READ) + assert not locks.defer_close(free), "an unlocked cache was kept for the manager to close" + assert locks.defer_close(kv) and locks.defer_close(kv) + assert locks.holds_open_caches() and kv.closes == 0 + with pytest.raises(ValueError, match="freed"): + locks.lock(kv, {2: [0]}, READ) + locks.unlock(first) + assert kv.closes == 0 and locks.holds_open_caches(), "closed while a lock still holds it" + assert HoldReason.LOCKED not in locks.held_page_counts().by_reason, "counted as locked" + assert locks.holds_any(), "a freed cache's lock went unseen" + locks.unlock(second) + locks.unlock(second) + locks.unlock(first) + assert kv.closes == 1 and not locks.holds_open_caches() and not locks.holds_any() + assert free.closes == 0, "the registry closed a cache it did not keep" + + +@pytest.mark.cpu_only +def test_shutdown_keeps_each_locked_cache_until_its_release(): + """At the manager's shutdown every mapped cache is offered; the locked ones stay open.""" + manager = Manager() + locks = registry(manager) + caches = [Cache(), Cache(), Cache()] + lock = locks.lock(caches[1], {0: [3]}, WRITE) + kept = [locks.defer_close(kv) for kv in caches] + assert kept == [False, True, False] and locks.holds_open_caches() + del manager # the manager is gone after its shutdown; the close does not need it + locks.unlock(lock) + assert [kv.closes for kv in caches] == [0, 1, 0] and not locks.holds_open_caches() + + +@pytest.mark.cpu_only +def test_the_lender_refuses_a_reuse_reset_while_a_lock_keeps_a_freed_cache_open(): + """The lender's reset hook refuses only while a lock keeps a freed cache open.""" + locks, kv = registry(), Cache() + lender = SimpleNamespace(_locks=locks) # all the hook reads + lock = locks.lock(kv, {0: [0]}, READ) + _lender._ManagerLender._before_reset(lender) # the request still holds its cache + assert locks.defer_close(kv) + with pytest.raises(RuntimeError, match="until its lease is released"): + _lender._ManagerLender._before_reset(lender) + locks.unlock(lock) + assert kv.closes == 1 + _lender._ManagerLender._before_reset(lender) + + +@pytest.mark.cpu_only +def test_a_state_row_is_locked_whatever_the_range_and_written_past_any_commit(): + """A recurrent layer group's row 0 is locked for any range and written past any commit.""" + locks = registry(recurrent=(False, True, False)) + kv = Cache(committed_tokens=2 * TPB) + writer = locks.lock(kv, {1: [0]}, WRITE) + with pytest.raises(ValueError, match="committed whole blocks"): + locks.lock(kv, {0: [0]}, WRITE) + assert locks.is_locked(kv, 7, 9) and locks.is_locked(kv) + with pytest.raises(ValueError, match="locked for write"): + locks.lock(kv, {1: [0]}, READ) + assert locked(locks) == {1: 1} + locks.unlock(writer) + assert not locks.is_locked(kv, 7, 9) + with pytest.raises(ValueError, match="committed whole blocks"): + registry().lock(Cache(committed_tokens=2 * TPB), {1: [0]}, WRITE) + + +@pytest.mark.cpu_only +def test_deferred_close_counts_include_every_page_of_each_cache_until_close(): + """A kept cache counts every page of each beam until it closes, a live locked one none.""" + locks = registry() + kept = Cache( + beams=[{0: [3, 4, 5, -1], 1: [7, 8], 2: [9]}, {0: [3, 4, 10, -1], 1: [7, 8], 2: [9]}] + ) + other = Cache(beams=[{0: [3, 11], 1: [12]}]) + live = Cache(beams=[{0: [20, 21], 1: [22], 2: [23]}]) + assert deducted(locks) == {} + first = locks.lock(kept, {0: [0]}, READ) + second = locks.lock(kept, {1: [1]}, READ) + sharing = locks.lock(other, {0: [0]}, READ) + locks.lock(live, {0: [0, 1], 1: [0]}, READ) + assert deducted(locks) == {}, "a live cache's pages counted as retained" + assert locks.defer_close(kept) and locks.defer_close(other) + # kept: layer group 0 holds 3, 4, 5 and beam 1's 10, layer group 2 holds 9; other: 3 and 11. + assert deducted(locks) == {0: 4 + 1 + 2, 1: 2 + 1} + locks.unlock(first) + assert deducted(locks) == {0: 7, 1: 3}, "an unlocked page of a kept cache freed" + locks.unlock(sharing) + assert other.closes == 1 and deducted(locks) == {0: 5, 1: 2} + locks.unlock(second) + assert kept.closes == 1 and deducted(locks) == {} + + +@pytest.mark.cpu_only +def test_the_lowest_locked_block_past_each_floor_spans_every_lock(): + locks, kv = registry(), Cache() + assert locks.lowest_locked_blocks(kv, (0, 0, 0)) == {} + first = locks.lock(kv, {0: [0, 5, 9], 2: [1]}, READ) + locks.lock(kv, {0: [3, 4], 1: []}, READ) + locks.lock(Cache(), {1: [0]}, READ) # another cache's lock counts for that cache alone + assert locks.lowest_locked_blocks(kv, (0, 0, 0)) == {0: 0, 2: 1} + # Blocks below a floor, a window's sinks, are passed over; a group with none past it is left out. + assert locks.lowest_locked_blocks(kv, (1, 0, 2)) == {0: 3} + assert locks.lowest_locked_blocks(kv, (6, 0, 0)) == {0: 9, 2: 1} + locks.unlock(first) + assert locks.lowest_locked_blocks(kv, (0, 0, 0)) == {0: 3} + + +@pytest.mark.cpu_only +def test_each_unlock_has_the_manager_retry_a_held_commit_before_the_close(): + """Every unlock retries a held commit; the last closes the cache after it, even if it raises.""" + manager = Manager() + locks, kv = registry(manager), Cache() + first, last = locks.lock(kv, {0: [0]}, READ), locks.lock(kv, {2: [1]}, READ) + assert locks.defer_close(kv) + locks.unlock(first) + assert manager.retried == [(kv, 0)] and kv.closes == 0 + locks.unlock(last) + assert manager.retried == [(kv, 0), (kv, 0)], "the cache closed before the retry" + assert kv.closes == 1 + locks.unlock(last) + assert len(manager.retried) == 2, "a repeated unlock retried again" + + def failing(kv): + raise RuntimeError("retry failed") + + manager._commit_held_back = failing + other = Cache() + lock = locks.lock(other, {0: [0]}, READ) + assert locks.defer_close(other) + with pytest.raises(RuntimeError, match="retry failed"): + locks.unlock(lock) + assert other.closes == 1 and not locks.is_locked(other), "a failed retry kept the cache" + + +@pytest.mark.cpu_only +def test_an_unlock_retries_a_commit_only_where_the_manager_holds_one_back(): + """An unlock retries only a commit the manager's record holds back, the last unlock too.""" + manager = Manager(held={}) + locks, kv, other = registry(manager), Cache(), Cache() + first, last = locks.lock(kv, {0: [0]}, READ), locks.lock(kv, {2: [1]}, READ) + unheld = locks.lock(other, {0: [0]}, READ) + locks.unlock(unheld) + assert manager.retried == [], "an unlock retried a commit nobody held back" + manager._held_commits[id(kv)] = (kv, None) + assert locks.defer_close(kv) + locks.unlock(first) + assert manager.retried == [(kv, 0)] and kv.closes == 0, "a held commit was not retried" + manager._held_commits[id(kv)] = (kv, None) + locks.unlock(last) + assert manager.retried == [(kv, 0), (kv, 0)], "the last unlock closed before the retry" + assert kv.closes == 1 and manager._held_commits == {} + + +@pytest.mark.cpu_only +def test_a_registry_needs_its_manager_to_lock(): + manager = Manager() + locks = registry(manager) + del manager + with pytest.raises(RuntimeError, match="gone"): + locks.lock(Cache(), {0: [0]}, READ, ready=Event()) + + +@pytest.mark.cpu_only +def test_the_module_reads_the_manager_and_caches_through_the_facade_alone(): + """The registry reaches a manager or a cache only through ``_manager``.""" + names = {"layer_group_pool_indices", "committed_blocks", "close_cache", "commit_held_back"} + source = Path(_page_locks.__file__).read_text() + assert {n for n in names if f"_manager.{n}(" in source} == names + assert all(callable(getattr(_manager, name)) for name in names) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("state_slots", [(17, 18), (17, -1)], ids=["every_beam", "one_beam"]) +def test_retained_counts_include_active_state_and_scratch_slots_across_beams(state_slots): + """A kept cache counts its scratch slots and each beam's state slot, -1 for none.""" + locks = registry(recurrent=(False, True, False)) + kv = Cache(beams=[{0: [3, 4]}, {0: [3, 10]}]) + kv.get_ssm_block_base_index = lambda lg, beam_id=0: state_slots[beam_id] + kv.get_scratch_desc = lambda lg: SimpleNamespace(slot_ids=[12, 13]) if lg == 0 else None + held = locks.lock(kv, {0: [0], 1: [0]}, READ) + assert locks.defer_close(kv) + assert deducted(locks) == {0: 5, 1: sum(slot >= 0 for slot in state_slots)} + locks.unlock(held) + assert kv.closes == 1 and deducted(locks) == {} + + +@pytest.mark.cpu_only +def test_a_failed_deferred_close_keeps_its_cache_and_capacity_registered(): + locks, kv = registry(), Cache(beams=[{0: [3]}]) + held = locks.lock(kv, {0: [0]}, READ) + assert locks.defer_close(kv) + + def fail_close(): + raise RuntimeError("injected close failure") + + kv.close = fail_close + with pytest.raises(RuntimeError, match="injected close failure"): + locks.unlock(held) + assert held.released and not locks.is_locked(kv) + assert locks.holds_open_caches() and deducted(locks) == {0: 1} + with pytest.raises(ValueError, match="freed"): + locks.lock(kv, {0: [0]}, READ) + + +def scanned(locks): + """What a full scan of the registry finds: per pool group, the distinct locked blocks of each + open cache, and every page each cache closing at its last unlock holds.""" + locked, deferred = {}, {} + for cache in locks._caches.values(): + if cache.close_at_last_unlock: + for lg, group in enumerate(locks._pool_group_of): + pages = _manager.held_pages(cache.kv, lg, recurrent=lg in locks._recurrent) + if pages: + deferred[group] = deferred.get(group, 0) + pages + continue + by_layer_group = {} + for lock in cache.locks: + for lg, ordinals in lock._blocks.items(): + by_layer_group.setdefault(lg, []).append(ordinals) + for lg, arrays in by_layer_group.items(): + pages = len(np.unique(np.concatenate(arrays))) + if pages: + group = locks._pool_group_of[lg] + locked[group] = locked.get(group, 0) + pages + return locked, deferred + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("counted", [False, True], ids=["never_counted", "counted"]) +def test_a_cache_whose_last_lock_ended_is_not_kept_for_a_later_count(counted): + """A cache whose last lock ended is no longer referenced by the registry, counted or not.""" + manager = Manager(held={}) # no commit held back, so none retried that would keep the cache + locks = registry(manager) + kv = Cache() + gone = weakref.ref(kv) + lock = locks.lock(kv, {0: [0]}, READ) + if counted: + assert locked(locks) == {0: 1} + locks.unlock(lock) + del kv, lock + gc.collect() + assert gone() is None, "the registry kept a cache whose last lock ended" + assert locked(locks) == {} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("seed", range(6)) +def test_held_page_counts_follow_each_change_as_a_full_scan_finds_them(seed): + """Random locks, unlocks and frees: the counts equal a full scan and change only with it.""" + rng = random.Random(seed) + locks = registry(recurrent=(False, True, False)) + caches = [] + for i in range(4): + cache = Cache(beams=[{lg: [rng.randrange(-1, 30) for _ in range(6)] for lg in range(3)}]) + cache.get_ssm_block_base_index = lambda lg, beam_id=0, slot=40 + i: slot + caches.append(cache) + held, changes = [], 0 + value = locks.held_page_counts() + for _ in range(120): + roll, cache = rng.random(), rng.choice(caches) + if roll < 0.5: + blocks = {lg: rng.sample(range(8), rng.randint(0, 3)) for lg in (0, 2)} + if rng.random() < 0.3: + blocks[1] = [0] + try: + held.append(locks.lock(cache, blocks, rng.choice([READ, READ, WRITE]))) + except ValueError: + pass + elif roll < 0.85 and held: + locks.unlock(held.pop(rng.randrange(len(held)))) + else: + locks.defer_close(cache) + new = locks.held_page_counts() + locked, deferred = scanned(locks) + reasons = {HoldReason.LOCKED: locked, HoldReason.DEFERRED_CLOSE: deferred} + assert {r: dict(c) for r, c in new.by_reason.items()} == { + r: c for r, c in reasons.items() if c + } + assert dict(new.deduct) == deferred + unreclaimable = dict(locked) + for group, pages in deferred.items(): + unreclaimable[group] = unreclaimable.get(group, 0) + pages + assert dict(new.unreclaimable) == unreclaimable + if (new.by_reason, new.deduct) == (value.by_reason, value.deduct): + assert new is value, "an unchanged count gave a new value" + else: + assert new.version > value.version, "a changed count kept its version" + changes += 1 + value = new + assert changes >= 20, changes + + +# -- the generation update over cache models, its methods compiled from the package's sources --- + +_CACHE = Path(_page_locks.__file__).resolve().parents[1] +_TPB = 4 +_STATES = SimpleNamespace(GENERATION_COMPLETE="complete", CONTEXT_INIT="context_init") +_TYPES = SimpleNamespace(CROSS="cross", SELF="self") +_LIFECYCLE_HOOKS = ( + "on_free defer_close before_reset on_reset on_shutdown keeps_pools on_caches_closed" +) + + +def _source(path: Path, qualified: str) -> ast.AST: + """The source node of a module-level or class-level definition.""" + node = ast.parse(path.read_text()) + for name in qualified.split("."): + node = next(child for child in node.body if getattr(child, "name", None) == name) + return node + + +def _reads_holding(test: ast.AST) -> bool: + """Whether an ``if`` test reads the lock state's ``holding``.""" + return any(isinstance(n, ast.Attribute) and n.attr == "holding" for n in ast.walk(test)) + + +def _method(path: Path, qualified: str, namespace: dict, eager: bool = False) -> Callable: + """A source definition compiled into ``namespace``. ``eager`` opens both lender gates of + ``update_resources``: every active cache asks the lent history and every resize notifies.""" + node = copy.deepcopy(_source(path, qualified)) + if eager: + changed = Counter() + for child in ast.walk(node): + if isinstance(child, ast.If) and _reads_holding(child.test): + child.test = ast.Name("has_lender", ast.Load()) + changed["holding gate"] += 1 + elif isinstance(child, ast.If) and any( + isinstance(n, ast.Name) and n.id == "num_blocks_before" + for n in ast.walk(child.test) + ): + child.test = ast.parse("has_lender and new_capacity is not None", mode="eval").body + changed["shrink gate"] += 1 + assert changed == {"holding gate": 1, "shrink gate": 1}, (qualified, dict(changed)) + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + module = ast.Module(body=[future, node], type_ignores=[]) + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return namespace[node.name] + + +@pytest.fixture(scope="module") +def sources() -> SimpleNamespace: + namespace = { + "np": np, + "SimpleNamespace": SimpleNamespace, + "traceback": traceback, + "logger": SimpleNamespace( + error=lambda *args: pytest.fail(f"unexpected lender error: {args}"), + warning=lambda *args: None, + ), + "CacheTypeCpp": _TYPES, + "LlmRequestState": _STATES, + "StreamingKVCacheEventManager": type("StreamingEvents", (), {}), + "_update_kv_cache_draft_token_location": lambda *args, **kwargs: None, + "_manager": SimpleNamespace(num_blocks=lambda kv: kv.num_blocks), + "_ranges": SimpleNamespace( + _last_history_keeping=lambda manager, layout, lg, block: ( + (block + 1) * _TPB + layout.windows[lg] - 2 + ) + ), + } + hooks = SimpleNamespace() + for name in ( + _LIFECYCLE_HOOKS + " on_shrink lent_history holds_page_locks generation_lock_state" + ).split(): + setattr(hooks, name, _method(_CACHE / "_lender_hooks.py", name, namespace)) + namespace["_lender_hooks"] = hooks + methods = SimpleNamespace(hooks=hooks, namespace=namespace) + for name in ("is_locked", "holds_any", "lowest_locked_blocks"): + locks = _CACHE / "sharing/_page_locks.py" + setattr(methods, name, _method(locks, f"PageLocks.{name}", namespace)) + for name in ("_history_cap", "_on_shrink", "_page_lock_state"): + lender = _CACHE / "sharing/_lender.py" + setattr(methods, name, _method(lender, f"_ManagerLender.{name}", namespace)) + ledger = _CACHE / "sharing/_fetch_ledger.py" + namespace["_writable_mask"] = _method(ledger, "_writable_mask", namespace) + methods.invalidate = _method( + ledger, "_FetchLedger._invalidate_delivered_blocks_past_capacity", namespace + ) + manager = _CACHE / "kv_cache_manager_v2.py" + for name in ("update_resources", "free_resources", "reset_reuse_state", "shutdown"): + setattr(methods, name, _method(manager, f"KVCacheManagerV2.{name}", namespace)) + update = "KVCacheManagerV2.update_resources" + methods.eager_update = _method(manager, update, dict(namespace), eager=True) + return methods + + +class _Locks: + """The real lock queries over admissions and unlocks modeled between generation updates.""" + + def __init__(self, sources: SimpleNamespace) -> None: + self._caches = {} + self.locked_caches = self._caches + self.caps_history = True + self.unreusable = {} + self.version = self.holding = 0 + self._recurrent = frozenset() + for name in ("is_locked", "holds_any", "lowest_locked_blocks"): + setattr(self, name, getattr(sources, name).__get__(self)) + + def add(self, kv: _Cache, block: int) -> None: + entry = self._caches.setdefault(id(kv), SimpleNamespace(kv=kv, locks=[])) + entry.locks.append(SimpleNamespace(_blocks={0: np.array([block], dtype=np.int64)})) + self.version += 1 + self.holding += 1 + + def clear(self) -> None: + if self._caches: + self._caches.clear() + self.version += 1 + self.holding = 0 + + def locked(self, kv: _Cache, block: int) -> bool: + return self.is_locked(kv, block, block + 1) + + +class _Cache: + """Native resize rules relevant here: block rounding, fresh regrowth and stale windows.""" + + def __init__(self, rid: int, window: int | None, locks: _Locks) -> None: + self.id = rid + self.capacity = self.history_length = self.serial = 0 + self.pages = [] + self.window, self.locks = window, locks + self.is_active = True + self.fail_next = False + self.trace, self.violations = [], [] + self.drops = 0 + + @property + def num_blocks(self) -> int: + return len(self.pages) + + def stale_end(self, history: int) -> int: + return 0 if self.window is None else max(0, (history + 1 - self.window) // _TPB) + + def resize(self, capacity: int | None, history: int | None) -> bool: + if self.fail_next: + self.fail_next = False + return False + end = self.capacity if capacity is None else capacity + past = self.history_length if history is None else history + assert self.history_length <= past <= end + before, after = self.num_blocks, -(-end // _TPB) + if after < before: + self.drops += 1 + for block in range(after, before): + if self.locks.locked(self, block): + self.violations.append(("tail", block)) + del self.pages[after:] + while len(self.pages) < after: + self.pages.append(("fresh", self.serial)) + self.serial += 1 + for block in range(min(self.stale_end(past), after)): + if self.pages[block] is not None and self.locks.locked(self, block): + self.violations.append(("window", block)) + self.pages[block] = None + self.capacity, self.history_length = end, past + self.trace.append((capacity, history, tuple(self.pages))) + return True + + +class _Ledger: + def __init__(self, layout: SimpleNamespace, sources: SimpleNamespace) -> None: + self._layout = layout + self._deliveries = {} + self._invalidate_delivered_blocks_past_capacity = sources.invalidate.__get__(self) + + def _deliveries_of(self, rid: int) -> SimpleNamespace | None: + return self._deliveries.get(rid) + + def _drop_deliveries(self, rid: int) -> None: + self._deliveries.pop(rid, None) + + def deliver(self, kv: _Cache, prompt: int) -> None: + top = min(kv.num_blocks, -(-prompt // _TPB)) + delivered = self._deliveries.setdefault( + kv.id, + SimpleNamespace( + kv=kv, + delivered_block_masks=[np.zeros(0, dtype=bool)], + end=0, + partial_ends={}, + resume_bounds_memo=("memo",), + ), + ) + mask = delivered.delivered_block_masks[0] + if len(mask) < top: + mask = np.pad(mask, (0, top - len(mask))) + for block in range(kv.stale_end(kv.history_length), top): + mask[block] = True + kv.pages[block] = ("fetched", kv.id, block) + delivered.delivered_block_masks[0] = mask + delivered.end = max(delivered.end, min(prompt, top * _TPB)) + if delivered.end % _TPB: + delivered.partial_ends[(0, delivered.end // _TPB)] = delivered.end + + +class _Manager: + def __init__( + self, + sources: SimpleNamespace, + window: int | None, + *, + attached: bool = True, + draft: bool = False, + managed: bool = False, + helix: bool = False, + ) -> None: + self.sources, self.counts = sources, Counter() + self.kv_cache_map, self._allocated_draft_lens = {}, {} + self._inflight_draft_lens = {} + self.tokens_per_block = _TPB + self.is_draft = draft + self.kv_compression_manages_history, self._has_cp_helix = managed, helix + self.kv_cache_type, self._kv_reserve_draft_tokens = _TYPES.SELF, 3 + self.locks = _Locks(sources) + self.locks.caps_history = window is not None + layout = SimpleNamespace(windows=[window], sink_blocks=[0], tokens_per_block=_TPB) + lender = SimpleNamespace( + _layout=layout, + _locks=self.locks, + _any_window=window is not None, + _shrink_epochs={}, + _writes_at_end={}, + _lose_end_resume=lambda kv_cache: None, + ) + lender._manager_ref = lambda: self + lender._serving_manager = lambda: self + lender._page_lock_state = sources._page_lock_state.__get__(lender) + lender._history_cap = sources._history_cap.__get__(lender) + lender._on_shrink = sources._on_shrink.__get__(lender) + lender._is_locked = lambda kv, begin=0, end=None: ( + self.locks.holds_any() if kv is None else self.locks.is_locked(kv, begin, end) + ) + self.ledger = lender._ledger = _Ledger(layout, sources) + self._sharing = lender if attached else None + + def _holds_page_locks(self, rid: int | None) -> bool: + return self.sources.hooks.holds_page_locks(self, rid) + + def _lent_history(self, kv: _Cache, history: int) -> int: + return self.sources.hooks.lent_history(self, kv, history) + + def _after_shrink(self, rid: int, kv: _Cache) -> None: + self.counts["shrink"] += 1 + self.sources.hooks.on_shrink(self, rid, kv) + + def try_commit_generation_blocks(self, request: SimpleNamespace) -> None: + """Commits nothing, as a manager with a lender attached does.""" + + +def _request(rid: int, tokens: int) -> SimpleNamespace: + return SimpleNamespace( + py_request_id=rid, + max_beam_num_tokens=tokens, + py_rewind_len=0, + py_num_accepted_draft_tokens=0, + state="generation", + ) + + +def _observe(manager: _Manager) -> tuple: + caches = tuple( + (rid, kv.capacity, kv.history_length, *map(tuple, (kv.pages, kv.trace, kv.violations))) + for rid, kv in manager.kv_cache_map.items() + ) + records = tuple( + ( + rid, + tuple(tuple(mask) for mask in d.delivered_block_masks), + d.end, + tuple(sorted(d.partial_ends.items())), + d.resume_bounds_memo, + ) + for rid, d in manager.ledger._deliveries.items() + ) + return caches, records, tuple(sorted(manager._allocated_draft_lens.items())) + + +def _run( + sources: SimpleNamespace, + update: Callable, + seed: int, + window: int | None, + *, + attached: bool = True, + draft: bool = False, + managed: bool = False, +) -> tuple: + rng = random.Random(seed) + manager = _Manager(sources, window, attached=attached, draft=draft, managed=managed) + requests, prompts, observations = [], {}, [] + for rid in range(4): + prompt = prompts[rid] = rng.randint(5, 23) + kv = manager.kv_cache_map[rid] = _Cache(rid, window, manager.locks) + kv.resize(prompt, 0 if managed else prompt) + requests.append(_request(rid, prompt)) + if attached: + manager.ledger.deliver(kv, prompt) + batch = SimpleNamespace(generation_requests=requests) + for _ in range(12): + # No admission or unlock occurs within update_resources. + manager.locks.clear() + for req in requests: + kv = manager.kv_cache_map[req.py_request_id] + runtime = rng.choice([0, 0, 1, 3, 5]) + accepted = rng.randint(0, runtime) + allocated = max(runtime, manager._kv_reserve_draft_tokens) if draft else runtime + kv.resize(kv.capacity + 1 + allocated, None) + req.py_rewind_len, req.py_num_accepted_draft_tokens = runtime - accepted, accepted + req.max_beam_num_tokens += 1 + accepted + if rng.random() < 0.5: + manager._allocated_draft_lens[req.py_request_id] = allocated + if attached and rng.random() < 0.3: + manager.ledger.deliver(kv, prompts[kv.id]) + low = kv.stale_end(kv.history_length) + top = min(kv.num_blocks, prompts[kv.id] // _TPB) + if attached and low < top and rng.random() < 0.3: + manager.locks.add(kv, rng.randrange(low, top)) + update(manager, batch) + observations.append(_observe(manager)) + assert not any(kv.violations for kv in manager.kv_cache_map.values()) + return observations, manager.counts, sum(kv.drops for kv in manager.kv_cache_map.values()) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "window,draft,managed,attached", + [ + (None, False, False, True), + (8, False, False, True), + (8, True, False, True), + (None, False, True, True), + (None, False, False, False), + ], +) +def test_generation_matches_eager_reference(sources, window, draft, managed, attached): + """Schedules of reserves, fallback widths, fetch masks, windows, locks and regrowth.""" + totals = Counter() + for seed in range(40): + options = dict(draft=draft, managed=managed, attached=attached) + reference, old_counts, _ = _run(sources, sources.eager_update, seed, window, **options) + observed, new_counts, drops = _run( + sources, sources.update_resources, seed, window, **options + ) + assert observed == reference, (seed, options) + assert new_counts["shrink"] == (drops if attached else 0) + totals.update(new_counts) + totals["old_shrink"] += old_counts["shrink"] + if attached: + assert 0 < totals["shrink"] < totals["old_shrink"] + else: + assert totals["shrink"] == 0 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("batch_size", [64, 256]) +def test_unlocked_batch_preserves_generation_results(sources, batch_size): + managers = [_Manager(sources, None) for _ in range(2)] + for manager, update in zip(managers, (sources.eager_update, sources.update_resources)): + requests = [] + for rid in range(batch_size): + kv = manager.kv_cache_map[rid] = _Cache(rid, None, manager.locks) + kv.resize(8, 0) + requests.append(_request(rid, 8)) + update(manager, SimpleNamespace(generation_requests=requests)) + assert _observe(managers[0]) == _observe(managers[1]) + assert managers[0].counts["shrink"] == batch_size + assert managers[1].counts["shrink"] == 0 + + +@pytest.mark.cpu_only +def test_in_flight_draft_width_sets_the_rewind(sources): + """The oldest in-flight width (2), not the next step's (4), sets the rewind and is consumed.""" + for update in (sources.eager_update, sources.update_resources): + manager = _Manager(sources, None) + kv = manager.kv_cache_map[1] = _Cache(1, None, manager.locks) + kv.resize(20, 12) + manager._allocated_draft_lens[1] = 4 + manager._inflight_draft_lens[1] = [2] + req = _request(1, 14) + req.py_num_accepted_draft_tokens = 1 + update(manager, SimpleNamespace(generation_requests=[req])) + assert (kv.capacity, kv.history_length) == (19, 13), update + assert manager._inflight_draft_lens == {} and manager._allocated_draft_lens == {}, update + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "capacity,rewind,state,notifies", + [ + (10, 1, "generation", 0), + (9, 1, "generation", 1), + (17, 4, "generation", 1), + (17, 0, "generation", 0), + (17, 4, _STATES.GENERATION_COMPLETE, 0), + (17, 4, _STATES.CONTEXT_INIT, 0), + ], +) +def test_shrink_notification_tracks_actual_blocks(sources, capacity, rewind, state, notifies): + manager = _Manager(sources, None, managed=True) + kv = manager.kv_cache_map[1] = _Cache(1, None, manager.locks) + kv.resize(capacity, 0) + manager.ledger.deliver(kv, capacity) + req = _request(1, 16) + req.py_rewind_len, req.state = rewind, state + manager._allocated_draft_lens[1] = rewind + sources.update_resources(manager, SimpleNamespace(generation_requests=[req])) + assert manager.counts["shrink"] == notifies == kv.drops + assert not manager._allocated_draft_lens + + +@pytest.mark.cpu_only +def test_failed_resize_keeps_allocation_and_skips_notification(sources): + manager = _Manager(sources, None) + kv = manager.kv_cache_map[1] = _Cache(1, None, manager.locks) + kv.resize(17, 0) + kv.fail_next = True + manager._allocated_draft_lens[1] = 4 + with pytest.raises(ValueError, match="generation update"): + sources.update_resources(manager, SimpleNamespace(generation_requests=[_request(1, 12)])) + assert kv.capacity == 17 and kv.history_length == 0 + assert manager._allocated_draft_lens == {1: 4} + assert manager.counts["shrink"] == 0 + + +def _changing_batch(sources, update, *, unlock=False): + manager = _Manager(sources, 4) + requests = [] + for rid in (1, 2): + kv = manager.kv_cache_map[rid] = _Cache(rid, 4, manager.locks) + kv.resize(20, 0) + requests.append(_request(rid, 16)) + if unlock: + manager.locks.add(manager.kv_cache_map[2], 0) + + def changing_requests(): + yield requests[0] + if unlock: + manager.locks.clear() + else: + manager.locks.add(manager.kv_cache_map[2], 0) + yield requests[1] + + update(manager, SimpleNamespace(generation_requests=changing_requests())) + return manager + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("unlock", [False, True], ids=["lock", "unlock"]) +def test_a_lock_or_unlock_during_the_batch_counts_for_the_requests_after_it(sources, unlock): + """A lock or an unlock between two requests of a batch counts for the second.""" + reference = _changing_batch(sources, sources.eager_update, unlock=unlock) + current = _changing_batch(sources, sources.update_resources, unlock=unlock) + assert _observe(current) == _observe(reference) + if not unlock: + assert current.kv_cache_map[2].history_length == 6 + + +@pytest.mark.cpu_only +def test_missing_and_suspended_caches_keep_their_allocations(sources): + manager = _Manager(sources, None) + kv = manager.kv_cache_map[1] = _Cache(1, None, manager.locks) + kv.resize(17, 0) + kv.is_active = False + manager._allocated_draft_lens.update({1: 4, 2: 4}) + sources.update_resources( + manager, SimpleNamespace(generation_requests=[_request(1, 12), _request(2, 12)]) + ) + assert manager._allocated_draft_lens == {1: 4, 2: 4} + assert kv.trace == [(17, 0, tuple(kv.pages))] + + +def _lifecycle(*, deferred=False, refuse=False, runtime_error=False): + """Minimal lifecycle dependencies; assertions in callbacks observe the actual call order.""" + events = [] + manager = SimpleNamespace( + _pending_preemption={}, + kv_connector_manager=None, + is_draft=False, + conversation_manager=None, + _allocated_draft_lens={}, + _inflight_draft_lens={}, + _request_stats_enabled_ids=set(), + _fresh_pages_filled={}, + _disagg_receive_ready={}, + _on_allocation_stream=set(), + _allocation_ready={}, + _early_freed_index_requests=set(), + _held_commits={}, + sparse_metadata_batch=None, + _stream=None, + event_manager=None, + index_mapper=SimpleNamespace(remove_sequence=lambda rid: None), + ) + kv = SimpleNamespace(close=lambda: events.append("close"), discard_pending_stats=lambda: None) + manager.kv_cache_map = {1: kv} + manager._detach_index_slot = lambda cache: events.append("detach") + + def on_free(rid, cache): + assert rid not in manager.kv_cache_map + assert "close" not in events + events.append("free") + + def before_reset(): + events.append("before_reset") + if refuse: + raise RuntimeError("locked cache") + + def shutdown(): + events.append("runtime_shutdown") + if runtime_error: + raise RuntimeError("runtime shutdown failed") + + manager._sharing = SimpleNamespace( + _on_free=on_free, + _defer_close=lambda cache: (events.append("defer_close"), deferred)[1], + _before_reset=before_reset, + _on_reset=lambda: events.append("reset"), + _on_shutdown=lambda: events.append("shutdown"), + _keeps_pools=lambda: (events.append("keeps_pools"), deferred)[1], + _on_caches_closed=lambda: events.append("caches_closed"), + ) + manager.impl = SimpleNamespace( + clear_stats_excluded=lambda rid: None, + clear_reusable_blocks=lambda: events.append("clear_reuse"), + shutdown=shutdown, + ) + return manager, events + + +# Per double, the source functions that run with it as their first argument. +_DOUBLED = { + "manager": ( + ("kv_cache_manager_v2.py", "KVCacheManagerV2.update_resources"), + ("_lender_hooks.py", "generation_lock_state"), + ("_lender_hooks.py", "holds_page_locks"), + ("_lender_hooks.py", "lent_history"), + ("_lender_hooks.py", "on_shrink"), + ), + "lifecycle": tuple( + ("kv_cache_manager_v2.py", f"KVCacheManagerV2.{name}") + for name in ("free_resources", "reset_reuse_state", "shutdown") + ) + + tuple(("_lender_hooks.py", name) for name in _LIFECYCLE_HOOKS.split()), + "lender": tuple( + ("sharing/_lender.py", f"_ManagerLender.{name}") + for name in ("_history_cap", "_on_shrink", "_page_lock_state") + ), + "locks": tuple( + ("sharing/_page_locks.py", f"PageLocks.{name}") + for name in ("is_locked", "holds_any", "lowest_locked_blocks") + ), + "ledger": ( + ("sharing/_fetch_ledger.py", "_FetchLedger._invalidate_delivered_blocks_past_capacity"), + ), +} + + +def _first_argument_reads(node: ast.FunctionDef) -> set[str]: + """Attributes a function loads from its first argument: ``self``, or a hook's manager.""" + first = node.args.args[0].arg + return { + n.attr + for n in ast.walk(node) + if isinstance(n, ast.Attribute) + and isinstance(n.ctx, ast.Load) + and isinstance(n.value, ast.Name) + and n.value.id == first + } + + +@pytest.mark.cpu_only +def test_doubles_carry_every_attribute_their_source_functions_read(sources): + """A source function that reads a new attribute fails here, naming the double and attribute.""" + manager = _Manager(sources, 8) + doubles = { + "manager": manager, + "lifecycle": _lifecycle()[0], + "lender": manager._sharing, + "locks": manager.locks, + "ledger": manager.ledger, + } + missing = [ + f"{kind} double lacks {attr}, which {qualified} reads" + for kind, functions in _DOUBLED.items() + for path, qualified in functions + for attr in sorted(_first_argument_reads(_source(_CACHE / path, qualified))) + if not hasattr(doubles[kind], attr) + ] + assert not missing, missing + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("deferred", [False, True]) +def test_free_notifies_before_close_or_deferred_detach(sources, deferred): + manager, events = _lifecycle(deferred=deferred) + sources.free_resources(manager, _request(1, 8)) + assert events == ["free", "defer_close", "detach" if deferred else "close"] + assert manager.kv_cache_map == {} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("refuse", [False, True]) +def test_reset_preflight_precedes_any_reset(sources, refuse): + manager, events = _lifecycle(refuse=refuse) + if refuse: + with pytest.raises(RuntimeError, match="locked cache"): + sources.reset_reuse_state(manager) + assert events == ["before_reset"] + else: + sources.reset_reuse_state(manager) + assert events == ["before_reset", "clear_reuse", "reset"] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("deferred,runtime_error", [(False, False), (True, False), (False, True)]) +def test_shutdown_notifies_caches_closed_only_after_runtime_returns( + sources, deferred, runtime_error +): + manager, events = _lifecycle(deferred=deferred, runtime_error=runtime_error) + if runtime_error: + with pytest.raises(RuntimeError, match="runtime shutdown failed"): + sources.shutdown(manager) + else: + sources.shutdown(manager) + if deferred: + assert events == ["shutdown", "defer_close", "detach", "keeps_pools"] + else: + expected = ["shutdown", "defer_close", "close", "keeps_pools", "runtime_shutdown"] + assert events == expected + ([] if runtime_error else ["caches_closed"]) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "hook,adapter,args", + [ + ("_defer_close", "defer_close", (object(),)), + ("_keeps_pools", "keeps_pools", ()), + ("_history_cap", "lent_history", (SimpleNamespace(history_length=0), 8)), + ("_is_locked", "holds_page_locks", (None,)), + ("_page_lock_state", "generation_lock_state", ()), + ], +) +def test_query_errors_propagate_to_the_manager(sources, hook, adapter, args): + def failed_query(*_query_args): + raise ValueError("query failed") + + queries = {hook: failed_query} + if hook != "_page_lock_state": + # A lock holds a block of every cache and can cap a history, so the adapters ask their + # query. + every = type("EveryCache", (dict,), {"__contains__": lambda self, key: True})() + queries.setdefault( + "_page_lock_state", + lambda: SimpleNamespace( + holding=1, locked_caches=every, caps_history=True, unreusable={} + ), + ) + manager = SimpleNamespace(_sharing=SimpleNamespace(**queries)) + with pytest.raises(ValueError, match="query failed"): + getattr(sources.hooks, adapter)(manager, *args) + + +def _real_registry(sources, manager): + """The registry compiled from its source; only manager primitives and events are modeled.""" + primitives = SimpleNamespace( + layer_group_pool_indices=lambda owner: [0], + committed_blocks=lambda kv, count: [0] * count, + commit_held_back=lambda owner, kv: owner._commit_held_back(kv), + # As the facade answers: a manager keeping no record of its held commits retries each. + has_held_commit=lambda owner, kv: ( + getattr(owner, "_held_commits", None) is None or id(kv) in owner._held_commits + ), + is_closed=lambda kv: False, + close_cache=lambda kv: None, + ) + namespace = dict( + sources.namespace, + enum=enum, + numbers=numbers, + types=types, + _manager=primitives, + _NO_BLOCKS=np.zeros(0, dtype=np.int64), + ) + _method(_CACHE / "sharing/_types.py", "_integer", namespace) + for name in ("LockMode", "PageLock", "_CacheLocks", "_ordinals", "PageLocks"): + _method(_CACHE / "sharing/_page_locks.py", name, namespace) + registry = namespace["PageLocks"](weakref.ref(manager)) + registry.locked = lambda kv, block: registry.is_locked(kv, block, block + 1) + manager._commit_held_back = lambda kv: None + manager.locks = manager._sharing._locks = registry + return registry, namespace["LockMode"] + + +def _real_generation_change(sources, update, change): + manager = _Manager(sources, 4) + registry, mode = _real_registry(sources, manager) + first = manager.kv_cache_map[1] = _Cache(1, 4, registry) + second = manager.kv_cache_map[2] = _Cache(2, 4, registry) + first.resize(17, 0) + second.resize(20, 0) + req_one, req_two = _request(1, 12), _request(2, 16) + req_one.py_rewind_len = 4 + held = registry.lock(first, {0: [0]}, mode.READ) if change in ("retry", "unlock") else None + admitted = [] + + def admit(): + if not admitted: + admitted.append(registry.lock(second, {0: [0]}, mode.READ)) + + if change == "retry": + manager._commit_held_back = lambda kv: admit() if kv is first else None + if change in ("resize", "retry", "unlock"): + original_resize = first.resize + + def resize(capacity, history_length): + result = original_resize(capacity, history_length) + if change == "resize": + admit() + else: + registry.unlock(held) + return result + + first.resize = resize + else: + original_shrink = manager._after_shrink + + def after_shrink(rid, kv): + original_shrink(rid, kv) + if kv is first: + admit() + + manager._after_shrink = after_shrink + update(manager, SimpleNamespace(generation_requests=[req_one, req_two])) + return manager + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("change", ["resize", "after_shrink", "retry", "unlock"]) +def test_production_registry_reentry_preserves_generation_results(sources, change): + reference = _real_generation_change(sources, sources.eager_update, change) + current = _real_generation_change(sources, sources.update_resources, change) + assert _observe(current) == _observe(reference) + if change != "unlock": + assert current.kv_cache_map[2].history_length == 6 + assert not any(kv.violations for kv in current.kv_cache_map.values()) + + +@pytest.mark.cpu_only +def test_production_registry_refusal_idempotence_and_retry_exception(sources): + manager = _Manager(sources, None) + registry, mode = _real_registry(sources, manager) + kv = _Cache(1, None, registry) + state = sources.hooks.generation_lock_state(manager) + before = state.version + assert state.holding == 0 + held = registry.lock(kv, {0: [0]}, mode.READ) + assert state.version != before and state.holding == 1 + stable = state.version + assert registry.refusal(kv, {0: [0]}, mode.WRITE) is not None + with pytest.raises(ValueError, match="conflicts"): + registry.lock(kv, {0: [0]}, mode.WRITE) + with pytest.raises(ValueError, match="only level=None"): + registry.lock(kv, {0: [0]}, mode.READ, level=0) + assert state.version == stable + callback_observed = [] + + def failed_retry(cache): + callback_observed.append((state.version != stable, registry.is_locked(cache))) + raise RuntimeError("retry failed") + + manager._commit_held_back = failed_retry + with pytest.raises(RuntimeError, match="retry failed"): + registry.unlock(held) + assert callback_observed == [(True, False)] + assert held.released and not registry.holds_any() and state.holding == 0 + stable = state.version + registry.unlock(held) + assert state.version == stable and callback_observed == [(True, False)] + + +@pytest.mark.cpu_only +def test_window_model_and_history_cap_follow_native_token_windows(sources): + for window in (1, 3, 4, 5, 8, 13): + manager = _Manager(sources, window) + cache = _Cache(1, window, manager.locks) + layout = manager._sharing._layout + + def retained(history): + return {token // _TPB for token in list(range(history + 1))[-window:]} + + for history in range(64): + assert cache.stale_end(history) == min(retained(history)) + namespace = dict( + sources.namespace, + _manager=SimpleNamespace( + stale_range=lambda owner, lg, history: (0, min(retained(history))), + _MAX_TOKENS=2**31 - 1, + ), + ) + _method(_CACHE / "sharing/_ranges.py", "_stale_range", namespace) + actual_cap = _method(_CACHE / "sharing/_ranges.py", "_last_history_keeping", namespace) + model_cap = sources.namespace["_ranges"]._last_history_keeping + for block in range(5): + last = max( + history + for history in range((block + 2) * _TPB + window) + if block in retained(history) + ) + assert actual_cap(manager, layout, 0, block) == last + assert model_cap(manager, layout, 0, block) == last + + +# -- the manager's lock hooks over real managers ------------------------------------------------ + +PROMPT = list(range(1000, 1097)) # three whole blocks and one token +OTHER_PROMPT = list(range(5000, 5097)) +ALIGNED = list(range(2000, 2000 + 4 * TPB)) # four whole blocks +SOURCE, OTHER, THIRD = 1, 2, 3 + + +def attach(mgr): + return attach_lender(mgr, scope=b"page-lock-hooks", staging=StagingOptions(3 * TPB)) + + +def warnings_of(monkeypatch) -> List[str]: + """The warnings the manager module logs from now on.""" + seen: List[str] = [] + warn = kv_cache_manager_v2.logger.warning + + def spy(message, *args, **kwargs): + seen.append(str(message)) + return warn(message, *args, **kwargs) + + monkeypatch.setattr(kv_cache_manager_v2.logger, "warning", spy) + return seen + + +def compute_context(mgr, request, tokens=None): + """The request's next context chunk, ``tokens`` or the rest, computed and its update run.""" + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + request.context_chunk_size = tokens or request.context_remaining_length + request.move_to_next_context_chunk() + batch = ScheduledRequests() + batch.append_context_request(request) + mgr.update_context_resources(batch) + + +def computed_under_a_lock(kit, mgr, lender, request_id): + """A request computed whole while a lock holds its first block, its last commit held back: + the request, its cache and the lock.""" + request = kit.admitted(mgr, request_id, PROMPT) + kv = kit.kv(mgr, request) + assert mgr.resize_context(request, request.context_remaining_length) + lock = lender._locks.lock(kv, {0: [0]}, READ) + compute_context(mgr, request) + assert kv.num_committed_tokens == 0, "inconclusive: the lock held no commit back" + return request, kv, lock + + +def matched(kit, mgr, request_id, tokens=PROMPT) -> int: + """The prompt tokens a new request with ``tokens`` finds in the reuse tree; freed again.""" + probe = kit.admitted(mgr, request_id, tokens) + tokens = kit.kv(mgr, probe).num_committed_tokens + mgr.free_resources(probe) + return tokens + + +@gpu +def test_a_locked_cache_refuses_suspension_until_its_unlock(kit, real_manager): + with real_manager() as mgr: + request = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + kv = kit.kv(mgr, request) + lock = lender._locks.lock(kv, {0: [1]}, READ) + with pytest.raises(RuntimeError, match="cannot be suspended"): + mgr.suspend_request(request) + assert kv.is_active, "the refused suspension changed the cache" + lender._locks.unlock(lock) + mgr.suspend_request(request) + assert not kv.is_active + assert mgr.resume_request(request) + + +@gpu +def test_a_freed_locked_cache_stays_open_until_its_last_unlock(kit, real_manager): + """A freed locked cache keeps its pages until its last unlock, whose close spares its row.""" + with real_manager() as mgr: + request = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + kv = kit.kv(mgr, request) + index = mgr.index_mapper.get_index(SOURCE) + pages = [kit.pages(kv, lg) for lg in range(kit.num_layer_groups(mgr))] + locks = [lender._locks.lock(kv, {0: [0, 1]}, READ), lender._locks.lock(kv, {0: [2]}, READ)] + assert locked(lender) == {kit.pool_group_of(mgr)[0]: 3} + mgr.free_resources(request) + assert kit.kv(mgr, request) is None and not kit.closed(kv), "the free closed a locked cache" + assert [kit.pages(kv, lg) for lg in range(kit.num_layer_groups(mgr))] == pages + assert lender._keeps_pools() + taker = kit.admitted(mgr, OTHER, OTHER_PROMPT) + assert mgr.index_mapper.get_index(OTHER) == index, "inconclusive: the row went elsewhere" + assert mgr.resize_context(taker, taker.context_remaining_length) + row = kit.index_row(mgr, taker).values + assert -1 not in row[: len(pages[0])], "the taker's row shows no pages of its own" + lender._locks.unlock(locks[0]) + assert not kit.closed(kv), "closed while a lock still holds it" + lender._locks.unlock(locks[1]) + lender._locks.unlock(locks[1]) + assert kit.closed(kv) and not lender._keeps_pools() + assert kit.index_row(mgr, taker).values == row, "the close wrote into the taker's row" + + +@gpu +def test_a_shutdown_keeps_a_locked_cache_and_the_pools_until_its_unlock( + kit, real_manager, monkeypatch +): + seen = warnings_of(monkeypatch) + with real_manager() as mgr: + request = kit.published(mgr, SOURCE, PROMPT) + other = kit.published(mgr, OTHER, OTHER_PROMPT) + lender = attach(mgr) + kv, other_kv = kit.kv(mgr, request), kit.kv(mgr, other) + lock = lender._locks.lock(kv, {0: [0]}, READ) + mgr.shutdown() # the runtime's own shutdown raises while a cache is open + assert kit.closed(other_kv) and not kit.closed(kv) + assert any("pools stay allocated" in message for message in seen), seen + assert lender._keeps_pools() and mgr.kv_cache_map == {} + lender._locks.unlock(lock) + assert kit.closed(kv) and not lender._keeps_pools() + # The block's exit shuts the manager down again; with no cache open the runtime's shutdown runs. + + +@gpu +def test_a_commit_that_reaches_a_locked_block_waits_for_its_unlock(kit, real_manager, monkeypatch): + """A commit reaching a locked block is skipped, with a warning, until the unlock freeing it.""" + seen = warnings_of(monkeypatch) + with real_manager() as mgr: + lender = attach(mgr) + request = kit.admitted(mgr, SOURCE, PROMPT) + kv = kit.kv(mgr, request) + assert mgr.resize_context(request, request.context_remaining_length) + compute_context(mgr, request, 2 * TPB) + assert kv.num_committed_tokens == 2 * TPB + # Locks on committed blocks leave a commit of later blocks alone. + on_committed = lender._locks.lock(kv, {0: [0, 1]}, READ) + on_tail = lender._locks.lock(kv, {0: [3]}, READ) + compute_context(mgr, request) + assert request.context_remaining_length == 0 + assert kv.num_committed_tokens == 2 * TPB, "a commit reached a locked block" + assert any("skipped while a lender locks" in message for message in seen), seen + lender._locks.unlock(on_tail) + assert kv.num_committed_tokens == request.prompt_len, "the unlock left the commit undone" + lender._locks.unlock(on_committed) + + +@gpu +def test_a_commit_held_back_reaches_the_reuse_tree_at_the_last_unlock(kit, real_manager): + with real_manager() as mgr: + lender = attach(mgr) + request, kv, lock = computed_under_a_lock(kit, mgr, lender, SOURCE) + assert matched(kit, mgr, OTHER) == 0, "the held-back blocks reached the reuse tree" + assert not mgr.try_commit_blocks(request) + assert kv.num_committed_tokens == 0, "a commit reached a locked block" + lender._locks.unlock(lock) + assert kv.num_committed_tokens == request.prompt_len + assert matched(kit, mgr, THIRD) == 3 * TPB, "the prompt missed the reuse tree" + + +@gpu +@pytest.mark.parametrize("freed", [False, True], ids=["active", "freed_before_unlock"]) +def test_a_held_last_commit_saves_a_conversation_plan_only_for_an_active_request( + kit, real_manager, freed +): + """The unlock saves an active conversation's plan; a freed request commits before its close.""" + with real_manager(block_reuse_config={"policy": "per_conversation"}) as mgr: + lender = attach(mgr) + request = kit.make_request(SOURCE, PROMPT) + request.py_conversation_params = SimpleNamespace(conversation_id="talk") + assert mgr.prepare_context(request) + kv = kit.kv(mgr, request) + assert mgr.resize_context(request, request.context_remaining_length) + lock = lender._locks.lock(kv, {0: [0]}, READ) + compute_context(mgr, request) # saving the plan here would raise + state = mgr.conversation_manager._conversation_states["talk"] + assert kv.num_committed_tokens == 0 and not state.planned_drop_handles + if freed: + mgr.free_resources(request) + assert "talk" not in mgr.conversation_manager._conversation_states + assert not kit.closed(kv), "the free closed a locked cache" + lender._locks.unlock(lock) + if freed: + assert kit.closed(kv), "the last unlock left the freed cache open" + assert "talk" not in mgr.conversation_manager._conversation_states + assert matched(kit, mgr, OTHER) == 3 * TPB, "the freed request's prompt missed the tree" + else: + assert kv.num_committed_tokens == request.prompt_len + assert len(state.planned_drop_handles) == 1, "the unlock saved no drop plan" + assert state.current_request_id is None + + +@gpu +@pytest.mark.parametrize("path", ["fast", "rules"]) +@pytest.mark.parametrize("order", ["returned_then_unlocked", "unlocked_then_returned", "abandoned"]) +def test_a_fetch_over_rows_a_held_commit_covers_keeps_them_out_of_the_reuse_tree( + kit, real_manager, monkeypatch, path, order +): + """A fetch over held-back rows, however it ends and is lent, leaves none in the reuse tree.""" + fast_lends = [] + fast = _lender._ManagerLender._write_every_row_fast + + def counted(self, *args, **kwargs): + lease = fast(self, *args, **kwargs) + fast_lends.append(lease is not None) + return lease + + monkeypatch.setattr(_lender._ManagerLender, "_write_every_row_fast", counted) + with real_manager() as mgr: + lender = attach(mgr) + if path == "rules": + lender._lends_fast = lambda: False + request, kv, lock = computed_under_a_lock(kit, mgr, lender, SOURCE) + lease = lender.lend_write(request, TPB, len(PROMPT), **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + assert any(fast_lends) == (path == "fast"), f"inconclusive: fast lends {fast_lends}" + view = kit.ready_view(lease, mgr) + arrived = [np.isin(rows.page_ordinals, (1,)) for rows in view.groups] + if order == "unlocked_then_returned": + lender._locks.unlock(lock) + if order == "abandoned": + lease.release() + else: + lender.return_writes([(lease, arrived, True)]) + if order != "unlocked_then_returned": + lender._locks.unlock(lock) + assert kv.num_committed_tokens <= TPB, "the retried commit took rows the fetch overwrote" + assert matched(kit, mgr, OTHER) <= TPB, "rows the fetch overwrote reached the reuse tree" + + +@gpu +@pytest.mark.parametrize("path", ["fast", "rules"]) +@pytest.mark.parametrize("finish", ["mark_before_release", "release_before_mark"]) +def test_a_commit_held_after_lending_is_dropped_before_a_write_unlock( + kit, real_manager, monkeypatch, path, finish +): + """A commit held after the fetch was lent must not enter reuse when that fetch unlocks.""" + fast_lends = [] + fast = _lender._ManagerLender._write_every_row_fast + + def counted(self, *args, **kwargs): + lease = fast(self, *args, **kwargs) + fast_lends.append(lease is not None) + return lease + + monkeypatch.setattr(_lender._ManagerLender, "_write_every_row_fast", counted) + with real_manager() as mgr: + lender = attach(mgr) + if path == "rules": + lender._lends_fast = lambda: False + request, kv, lock = computed_under_a_lock(kit, mgr, lender, SOURCE) + lease = lender.lend_write(request, TPB, len(PROMPT), **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + assert any(fast_lends) == (path == "fast"), f"inconclusive: fast lends {fast_lends}" + assert kit.ready_view(lease, mgr) is not None + assert id(kv) not in mgr._held_commits, "the lend left the earlier commit held" + assert not mgr.try_commit_blocks(request) + assert id(kv) in mgr._held_commits, "inconclusive: no commit was held after the lend" + + if finish == "mark_before_release": + lender.return_writes([(lease, ALL_ARRIVED, True)]) + lease.release() + else: + lease.release() + lender._locks.unlock(lock) + + assert id(kv) not in mgr._held_commits + assert kv.num_committed_tokens <= TPB, "the unlock committed rows the fetch could overwrite" + assert matched(kit, mgr, OTHER) <= TPB, "rows the fetch could overwrite reached reuse" + + +@gpu +@pytest.mark.parametrize("lock_count", [1, 2]) +def test_a_fetch_from_where_the_computed_rows_end_keeps_the_held_commit( + kit, real_manager, lock_count +): + """A fetch from where the computed tokens end leaves the held commit to the last unlock.""" + with real_manager() as mgr: + lender = attach(mgr) + request = kit.admitted(mgr, SOURCE, PROMPT) + kv = kit.kv(mgr, request) + assert mgr.resize_context(request, request.context_remaining_length) + locks = [lender._locks.lock(kv, {0: [0]}, READ) for _ in range(lock_count)] + compute_context(mgr, request, 2 * TPB) + assert kv.num_committed_tokens == 0, "inconclusive: the lock held no commit back" + lease = lender.lend_write(request, 2 * TPB, len(PROMPT), **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + kit.ready_view(lease, mgr) + lender.return_writes([(lease, ALL_ARRIVED, True)]) + request.context_current_position = len(PROMPT) + for lock in locks[:-1]: + lender._locks.unlock(lock) + assert kv.num_committed_tokens == 0 + assert mgr._held_commits[id(kv)][2:] == (2 * TPB, True) + lender._locks.unlock(locks[-1]) + assert kv.num_committed_tokens == 2 * TPB, "a fetch past the held rows dropped their commit" + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("manager", ["paged", "hybrid"]) +@pytest.mark.parametrize("case", ["committed", "held_back", "no_reuse"]) +def test_a_context_end_saves_its_drop_plan_unless_a_lock_holds_its_commit_back(manager, case): + """A context's end saves its drop plan unless a lock holds its commit back.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + BlockReusePolicy, + CacheTypeCpp, + KVCacheManagerV2, + ) + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + + no_reuse = case == "no_reuse" + plans = SimpleNamespace(saved=[]) + plans.save_drop_plan = lambda request, kv_cache: plans.saved.append(request.py_request_id) + request = SimpleNamespace( + py_request_id=SOURCE, + is_dummy_request=no_reuse and manager == "paged", + context_current_position=len(PROMPT), + context_remaining_length=0, + expect_snapshot_points=[], + ) + stand_in = SimpleNamespace( + kv_cache_map={SOURCE: SimpleNamespace(is_active=True, resize=lambda *args: True)}, + kv_cache_type=CacheTypeCpp.SELF, + enable_block_reuse=True, + _can_publish_block_reuse=True, + _commit_branch_snapshots=False, + is_draft=no_reuse and manager == "hybrid", + block_reuse_policy=BlockReusePolicy.PER_CONVERSATION, + conversation_manager=plans, + _lent_history=lambda kv_cache, history: history, + _history_lock_state=lambda: None, + _resize_context_history=lambda req, kv_cache, history: None, + try_commit_blocks=lambda req, kv_cache=None: case == "committed", + ) + cls = KVCacheManagerV2 if manager == "paged" else MambaHybridCacheManagerV2 + cls.update_context_resources(stand_in, SimpleNamespace(context_requests=[request])) + assert plans.saved == ([] if case == "held_back" else [SOURCE]), plans.saved + + +@gpu +def test_a_first_chunk_out_of_pages_keeps_its_locked_cache_active(kit, real_manager): + """A first chunk out of pages is suspended unless a lease lends its matched prefix in place.""" + # Index slots for a one-block request per page, so other requests can take every page. + with real_manager(max_tokens=kit.POOL_TOKENS, host_cache_size=0, max_batch_size=64) as mgr: + kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr) + # The source's three whole blocks, then two blocks of its own: its first chunk needs pages. + target = kit.make_request(OTHER, PROMPT[: 3 * TPB] + list(range(7000, 7000 + 2 * TPB))) + assert mgr.prepare_context(target) and target.is_first_context_chunk + kv = kit.kv(mgr, target) + assert kv.num_committed_tokens == 3 * TPB, "inconclusive: no prefix matched" + held = [page for page in kit.pages(kv, 0) if page >= 0] + assert len(held) < target.prompt_len // TPB, "inconclusive: the chunk needs no new page" + lease = lender.lend_read(target, 0, 3 * TPB, **kit.DEVICE_NAMED) + assert lease.failure is None, lease.failure + others = kit.Requests(mgr) + while others.allocate(1, chunk=1): + pass + # A pool group with fewer free or evictable pages than the chunk's new blocks refuses it. + needed = target.prompt_len // TPB - len(held) + tiers = [str(tier) for tier in mgr.impl.cache_tier_list] + gpu_pages = [ + (int(s.total), int(s.free), int(s.evictable)) + for s in mgr.impl.get_storage_statistics(0) + ] + assert len(tiers) == 1 and any(free + evict < needed for _, free, evict in gpu_pages), ( + f"inconclusive: the chunk could still get its {needed} new pages; tiers {tiers}, GPU " + f"(total, free, evictable) per pool group {gpu_pages}, {len(others.held)} others" + ) + assert not mgr.resize_context(target, target.context_remaining_length) + assert kv.is_active, "the failed resize suspended a cache a lease lends" + lease.release() + others.free() + + +@gpu +def test_a_reverted_context_chunk_keeps_its_locked_cache_active(kit, real_manager): + """A reverted context allocation keeps its cache active while a lease lends its blocks.""" + with real_manager(enable_block_reuse=False) as mgr: + lender = attach(mgr) + request = kit.admitted(mgr, OTHER, OTHER_PROMPT) + assert mgr.resize_context(request, 2 * TPB) + compute_context(mgr, request, 2 * TPB) + kv = kit.kv(mgr, request) + lease = lender.lend_read(request, 0, 2 * TPB, **kit.DEVICE_ALL) + assert lease.failure is None, lease.failure + assert mgr.resize_context(request, TPB) and kv.capacity > 2 * TPB + assert mgr.revert_allocate_context(request) + assert kv.is_active and kv.capacity == 2 * TPB + lease.release() + + +@gpu +@pytest.mark.parametrize("hybrid", [False, True], ids=["paged", "hybrid"]) +def test_no_context_step_commit_or_stop_makes_an_event_for_a_lender( + kit, real_manager, hybrid_manager, recorded, hybrid +): + """A chunked context's steps, commits and stop make as many events with a lender as without.""" + made = [] + for lender in (False, True): + with hybrid_manager() if hybrid else real_manager() as mgr: + if lender: + attach(mgr) + before = len(recorded) + request = kit.admitted(mgr, SOURCE, PROMPT) + assert mgr.resize_context(request, request.context_remaining_length) + compute_context(mgr, request, 2 * TPB) + compute_context(mgr, request) + if not hybrid: + assert kit.kv(mgr, request).num_committed_tokens == request.prompt_len + made.append(len(recorded) - before) + mgr.free_resources(request) + assert made[0] == made[1], f"the lender made events: {made}" + + +@gpu +def test_a_read_in_place_lends_what_the_context_steps_computed(kit, real_manager): + """A read past a held-back history lends up to the computed tokens and waits for its event.""" + with real_manager(windows=[2 * TPB, _manager._MAX_TOKENS], enable_block_reuse=False) as mgr: + lender = attach(mgr) + short = lender._layout.windows.index(2 * TPB) + request = kit.admitted(mgr, OTHER, OTHER_PROMPT) + kv = kit.kv(mgr, request) + assert mgr.resize_context(request, request.context_remaining_length) + held = lender._locks.lock(kv, {short: [0]}, READ) + compute_context(mgr, request, 3 * TPB) + computed = request.context_current_position + assert computed == 3 * TPB and kv.history_length < computed, ( + f"inconclusive: the history {kv.history_length} reached the context position" + ) + + def lend(end): + ready = kit.stream_event(mgr) + return lender.lend_read(request, 0, end, **{**kit.DEVICE_ALL, "ready_event": ready}) + + lease = lend(computed) + assert lease.failure is None, lease.failure + lent = {rows.layer_group_index: len(rows) for rows in kit.ready_view(lease, mgr).groups} + assert lent[short] == 0 and sum(lent.values()) == 3, lent + lease.release() + past = lend(computed + 1) + assert "past the" in (past.failure or ""), past.failure + past.release() + inside = lend(kv.history_length - 1) + assert "ends at a whole block" in (inside.failure or ""), inside.failure + inside.release() + lender._locks.unlock(held) + compute_context(mgr, request) + whole = lend(len(OTHER_PROMPT)) + assert whole.failure is None, whole.failure + lent = {rows.layer_group_index: len(rows) for rows in kit.ready_view(whole, mgr).groups} + assert lent[short] > 0 and sum(lent.values()) > 4, lent + whole.release() + mgr._stream.synchronize() + mgr.free_resources(request) + + +@gpu +@pytest.mark.parametrize("lender", [True, False], ids=["lender", "no_lender"]) +def test_a_hybrid_manager_records_no_receive_event_after_its_slot_reset( + kit, hybrid_manager, recorded, lender +): + """A generation-init admission's resource preparation makes no event and leaves none.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestType + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + with hybrid_manager() as mgr: + if lender: + attach(mgr) + request = kit.make_request( + SOURCE, PROMPT, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + ) + assert request.is_disagg_generation_init_state + assert mgr.prepare_disagg_gen_init(request) + assert SOURCE in mgr._disagg_receive_ready + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + before = len(recorded) + mgr.prepare_resources(batch) + assert len(recorded) == before, "the slot reset made an event" + assert mgr.take_disagg_receive_ready(request) is None, "an event is left to take" + + +@pytest.mark.cpu_only +def test_kda_suspend_checks_page_locks_before_changing_replay_state(): + """A borrowed KDA cache cannot lose replay bookkeeping during a rejected suspension.""" + from tensorrt_llm._torch.modules.kimi_kda.cache_manager import KimiK3HybridCacheManager + + manager = object.__new__(KimiK3HybridCacheManager) + replay = SimpleNamespace( + acceptance_pending=set(), unseeded_requests={SOURCE}, suspended_requests=set() + ) + cache = SimpleNamespace(is_active=True) + manager._kda_replay = replay + manager.kv_cache_map = {SOURCE: cache} + manager._request_id_to_state_index = {SOURCE: 0} + + def reject_suspend(candidate): + assert candidate is cache + raise RuntimeError("borrowed cache") + + manager._check_suspendable = reject_suspend + with pytest.raises(RuntimeError, match="borrowed cache"): + manager.suspend_request(SimpleNamespace(py_request_id=SOURCE)) + + assert replay.unseeded_requests == {SOURCE} + assert replay.suspended_requests == set() + assert manager._request_id_to_state_index == {SOURCE: 0} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("locked", [False, True]) +def test_disagg_resize_pressure_returns_false_and_preserves_a_locked_cache(kit, locked): + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestType + + request = kit.make_request( + SOURCE, PROMPT, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + ) + suspended = [] + kv = SimpleNamespace( + capacity=TPB, resize=lambda *args: False, suspend=lambda: suspended.append(True) + ) + + def check_suspendable(cache): + assert not locked, "ordinary allocation pressure must not try suspending a locked cache" + + manager = SimpleNamespace( + prepare_context_cache=lambda req: 0, + enable_block_reuse=False, + kv_cache_map={SOURCE: kv}, + _has_cp_helix=False, + num_extra_kv_tokens=0, + _holds_page_locks=lambda request_id: locked, + _check_suspendable=check_suspendable, + ) + assert request.is_first_context_chunk + assert not kv_cache_manager_v2.KVCacheManagerV2.prepare_disagg_gen_init(manager, request) + assert manager.kv_cache_map[SOURCE] is kv + assert suspended == ([] if locked else [True]) + + +@gpu +@pytest.mark.parametrize("windows", [None, [2 * TPB, 4096]], ids=["full", "windowed"]) +def test_a_context_step_asks_no_lock_query_while_no_lock_holds_a_block_of_its_cache( + kit, real_manager, monkeypatch, windows +): + """Only a locked cache's step asks about locks, and about its history only through a window.""" + asked = {"_is_locked": 0, "_history_cap": 0} + for name in asked: + query = getattr(_lender._ManagerLender, name) + + def counted(self, *args, _query=query, _name=name, **kwargs): + asked[_name] += 1 + return _query(self, *args, **kwargs) + + monkeypatch.setattr(_lender._ManagerLender, name, counted) + with real_manager(windows=windows) as mgr: + lender = attach(mgr) + source = kit.published(mgr, SOURCE, PROMPT) + kv_cache = kit.kv(mgr, source) + # Asked directly too: a manager with block reuse moves the history through its commits. + assert mgr._lent_history(kv_cache, len(PROMPT)) == len(PROMPT) + assert not mgr._commit_gate(kv_cache, source, len(PROMPT)) + assert asked == {"_is_locked": 0, "_history_cap": 0}, asked + read = lender.lend_read(source, 0, 3 * TPB, **kit.DEVICE_ALL) + assert read.failure is None and lender._locks.holding == 1 + other = kit.published(mgr, OTHER, OTHER_PROMPT) + other_cache = kit.kv(mgr, other) + assert mgr._lent_history(other_cache, len(OTHER_PROMPT)) == len(OTHER_PROMPT) + assert not mgr._commit_gate(other_cache, other, len(OTHER_PROMPT)) + assert asked == {"_is_locked": 0, "_history_cap": 0}, ("asked about another cache", asked) + mgr._commit_gate(kv_cache, source, len(PROMPT)) + mgr._lent_history(kv_cache, len(PROMPT)) + assert asked["_is_locked"] > 0, asked + assert (asked["_history_cap"] > 0) == (windows is not None), asked + read.release() + before = dict(asked) + third = kit.published(mgr, THIRD, PROMPT[:40]) + assert asked == before, (before, asked) + mgr._stream.synchronize() + for request in (source, other, third): + mgr.free_resources(request) + + +@gpu +@pytest.mark.parametrize("windows", [None, [2 * TPB, 4096]], ids=["full", "windowed"]) +def test_a_context_step_keeps_a_block_a_lock_holds_inside_its_window(kit, real_manager, windows): + """A context step moves a history only as far as its window keeps a locked block.""" + with real_manager(windows=windows, enable_block_reuse=False) as mgr: + lender = attach(mgr) + state = mgr._history_lock_state() + assert (state is lender._locks) == (windows is not None), state + windowed = [lg for lg, window in enumerate(kit.windows(mgr)) if window is not None] + assert bool(windowed) == (windows is not None), "inconclusive: the windows were not kept" + lg = windowed[0] if windowed else 0 + request = kit.admitted(mgr, SOURCE, PROMPT) + kv = kit.kv(mgr, request) + assert mgr.resize_context(request, request.context_remaining_length) + lock = lender._locks.lock(kv, {lg: [0]}, READ) + compute_context(mgr, request) + cap = lender._history_cap(kv) + if windows is None: + assert cap is None and kv.history_length == len(PROMPT), (cap, kv.history_length) + else: + assert cap is not None and cap < len(PROMPT), "inconclusive: the window keeps block 0" + assert kv.history_length == cap, (kv.history_length, cap) + lender._locks.unlock(lock) + mgr._stream.synchronize() + mgr.free_resources(request) + + +@gpu +@pytest.mark.parametrize("caps", [True, False], ids=["a_lock_can_cap", "no_lock_can_cap"]) +def test_a_hybrid_context_step_moves_a_history_only_as_far_as_the_lender_lets_it( + kit, hybrid_manager, monkeypatch, caps +): + """A hybrid context step asks the lender's cap only where its lock state says a lock can cap.""" + asked = [] + + def lent_history(kv_cache, history): + asked.append(history) + return min(history, TPB) + + with hybrid_manager() as mgr: + attach(mgr) + state = SimpleNamespace(holding=1, locked_caches={}, caps_history=True, unreusable={}) + monkeypatch.setattr(mgr, "_history_lock_state", lambda: state if caps else None) + monkeypatch.setattr(mgr, "_lent_history", lent_history) + request = kit.admitted(mgr, SOURCE, PROMPT) + assert mgr.resize_context(request, request.context_remaining_length) + compute_context(mgr, request, 2 * TPB) + history = kit.kv(mgr, request).history_length + if caps: + assert asked == [2 * TPB] and history == TPB, (asked, history) + else: + assert asked == [] and history == 2 * TPB, (asked, history) + mgr._stream.synchronize() + mgr.free_resources(request) + + +def refuse_reuse(kit, lender, request, end, reusable=False): + """A write lease of ``[0, end)`` in place, returned with every row and ``reusable``.""" + lease = lender.lend_write(request, 0, end, **kit.DEVICE_ALL) + assert lease.poll() is not None, lease.failure + lender.return_writes([(lease, ALL_ARRIVED, reusable)]) + + +@gpu +@pytest.mark.parametrize( + "path", ["context", "native_entry", "later_true", "retry", "retry_after_exit", "retired"] +) +def test_a_cache_kept_out_of_reuse_commits_nothing_on_any_path(kit, real_manager, path): + """A ``reusable=False`` return keeps its cache out of the reuse tree on every commit path.""" + with real_manager(block_reuse_config={"policy": "per_conversation"}) as mgr: + lender = attach(mgr) + request = kit.make_request(SOURCE, ALIGNED) + request.py_conversation_params = SimpleNamespace(conversation_id="talk") + assert mgr.prepare_context(request) + kv = kit.kv(mgr, request) + kv.enable_swa_scratch_reuse = False + state = mgr.conversation_manager._conversation_states["talk"] + assert mgr.resize_context(request, request.context_remaining_length) + again = None + if path == "retired": + lease = lender.lend_write(request, 0, len(ALIGNED), **kit.DEVICE_ALL) + assert lease.poll() is not None, lease.failure + mgr.free_resources(request) + again = kit.make_request(SOURCE, ALIGNED) + assert mgr.prepare_context(again) + lender.return_writes([(lease, ALL_ARRIVED, False)]) + assert kit.closed(kv) + elif path.startswith("retry"): + lock = lender._locks.lock(kv, {0: [0]}, READ) + compute_context(mgr, request) + assert id(kv) in mgr._held_commits, "inconclusive: no commit held back" + lease = lender.lend_write(request, 4 * TPB, 4 * TPB, **kit.DEVICE_ALL) + assert lease.poll() is not None, lease.failure + if path == "retry_after_exit": + mgr.free_resources(request) + lender.return_writes([(lease, ALL_ARRIVED, False)]) + lender._locks.unlock(lock) + else: + refuse_reuse(kit, lender, request, len(ALIGNED)) + if path == "later_true": + refuse_reuse(kit, lender, request, len(ALIGNED), reusable=True) + if path == "native_entry": + request.context_chunk_size = request.context_remaining_length + request.move_to_next_context_chunk() + assert mgr.try_commit_blocks(request), "the commit did not stop" + else: + compute_context(mgr, request) + if path != "native_entry": + assert state.current_request_id is None, "a refused commit saved no plan" + assert kv.num_committed_tokens == 0, "a cache kept out of reuse committed" + assert matched(kit, mgr, OTHER, ALIGNED) == 0, "a cache kept out of reuse was matched" + if again is None: + if path != "retry_after_exit": + mgr.free_resources(request) + again = kit.make_request(SOURCE, ALIGNED) + assert mgr.prepare_context(again) + assert mgr.resize_context(again, again.context_remaining_length) + compute_context(mgr, again) + assert kit.kv(mgr, again).num_committed_tokens == len(ALIGNED), "a new cache stayed out" + mgr.free_resources(again) + + +@gpu +def test_a_staged_refusal_keeps_a_held_commit_out_of_reuse_after_the_free(kit, real_manager): + """A staged ``reusable=False`` return, which holds no lock, keeps a held commit out after a free.""" + with real_manager() as mgr: + lender = attach(mgr) + request = kit.admitted(mgr, SOURCE, ALIGNED) + kv = kit.kv(mgr, request) + assert mgr.resize_context(request, request.context_remaining_length) + lock = lender._locks.lock(kv, {0: [0]}, READ) + compute_context(mgr, request, 2 * TPB) + assert id(kv) in mgr._held_commits, "inconclusive: no commit held back" + staged = lender.lend_write(request, 2 * TPB, 3 * TPB, **kit.HOST_NAMED) + view = None + for _ in range(4): + lender.progress() + view = staged.poll() + if view is not None: + break + mgr._stream.synchronize() + assert view is not None, staged.failure + lender.return_writes([(staged, ALL_ARRIVED, False)]) + mgr._stream.synchronize() + lender.progress() + assert id(kv) in mgr._held_commits, "inconclusive: the return retried the commit" + mgr.free_resources(request) + lender._locks.unlock(lock) + assert kit.closed(kv) + assert matched(kit, mgr, OTHER, ALIGNED) == 0, "a cache kept out of reuse was matched" + + +@gpu +@pytest.mark.parametrize("replacement_freed", [False, True]) +def test_a_held_commit_of_a_freed_cache_commits_nothing_into_the_request_s_next_cache( + kit, real_manager, replacement_freed +): + """Replacing a cache invalidates its held commit, also after the replacement is freed.""" + with real_manager() as mgr: + lender = attach(mgr) + request = kit.admitted(mgr, SOURCE, ALIGNED) + kv = kit.kv(mgr, request) + assert mgr.resize_context(request, request.context_remaining_length) + lock = lender._locks.lock(kv, {0: [0]}, READ) + compute_context(mgr, request, 2 * TPB) + assert id(kv) in mgr._held_commits, "inconclusive: no commit held back" + mgr.free_resources(request) + again = kit.admitted(mgr, SOURCE, ALIGNED) + assert kit.kv(mgr, again).num_committed_tokens == 0, "inconclusive: the new cache reused" + request.context_chunk_size = request.context_remaining_length + request.move_to_next_context_chunk() + if replacement_freed: + mgr.free_resources(again) + lender._locks.unlock(lock) + assert kit.closed(kv) + if not replacement_freed: + assert kit.kv(mgr, again).num_committed_tokens == 0 + mgr.free_resources(again) + assert matched(kit, mgr, OTHER, ALIGNED) == 0, "the retired cache entered the reuse tree" + + +@gpu +@pytest.mark.parametrize("path", ["context", "free"]) +def test_a_hybrid_cache_kept_out_of_reuse_commits_no_snapshot(kit, hybrid_manager, path): + """A hybrid cache a ``reusable=False`` return keeps out of reuse commits no snapshot.""" + with hybrid_manager() as mgr: + lender = attach(mgr) + request = kit.admitted(mgr, SOURCE, ALIGNED) + assert mgr.resize_context(request, request.context_remaining_length) + refuse_reuse(kit, lender, request, len(ALIGNED)) + request.expect_snapshot_points = [2 * TPB, len(ALIGNED)] + kv = kit.kv(mgr, request) + compute_context(mgr, request, 2 * TPB) + if path == "context": + compute_context(mgr, request) + else: + mgr.free_resources(request) + assert kv.num_committed_tokens == 0, "a snapshot of a cache kept out of reuse committed" + assert matched(kit, mgr, OTHER, ALIGNED) == 0 + if path == "context": + mgr.free_resources(request) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py index d08fa1ea6cd5..8632494d9fa7 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py @@ -12,11 +12,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The lender's public surface and boundaries: the nine names, their members, what importing the +"""The lender's public surface and boundaries: the twelve names, their members, what importing the package loads, and who may import its private modules. No GPU.""" import ast import dataclasses +import enum import inspect import json import os @@ -33,13 +34,18 @@ import tensorrt_llm from tensorrt_llm._torch.pyexecutor.kv_cache import sharing from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( - LayerGroupBlocks, + ALL_ARRIVED, + HeldPages, + HoldReason, Lease, LeaseView, Lender, + LentLayerGroupRows, + MemoryKind, MemoryRegion, RegionHold, ResumeBounds, + RowSelection, StagingOptions, attach_lender, ) @@ -47,8 +53,8 @@ SH = "tensorrt_llm._torch.pyexecutor.kv_cache.sharing" MGR = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" PUBLIC = sorted( - "LayerGroupBlocks Lease MemoryRegion RegionHold ResumeBounds LeaseView Lender StagingOptions " - "attach_lender".split() + "ALL_ARRIVED Arrived HeldPages HoldReason LentLayerGroupRows Lease MemoryKind MemoryRegion " + "RegionHold ResumeBounds RowSelection LeaseView Lender StagingOptions attach_lender".split() ) MODULES = [ "__init__.py", @@ -57,6 +63,7 @@ "_layout.py", "_lender.py", "_manager.py", + "_page_locks.py", "_ranges.py", "_slots.py", ] @@ -67,14 +74,18 @@ # Manager and runtime members only the manager facade may touch. INTERNALS = { "impl", + "kv_cache_type", "init_config", "layer_grouping", "pool_group_descs", "_has_lossy_cold_page_codec", - "get_disagg_global_layer_ids", - "get_disagg_role_mapper_kinds", - "get_disagg_role_layouts", - "get_disagg_ignored_roles", + # The hooks a manager builds its page table with; the lender reads the table it returns. + "_page_buffer_tensors", + "_page_head_split", + "_page_untransferred_roles", + "_page_rebuilds_untransferred", + "_page_layer_ids", + "_page_kv_heads", "get_life_cycle_pool_group_indices", "_reuse_token_source", "_augment_tokens_for_block_reuse", @@ -87,11 +98,25 @@ "block_reuse_policy", "_draft_prompt_lookahead", "kv_connector_manager", + "capacity_only", "commit_min_snapshot", "reuse_match_backoff", "_stream", "_layer_attn_to_layer_id", + "use_fp8_ds_mla", + "_attention_read_window", "_sharing", + "_commit_held_back", + "take_disagg_receive_ready", + "reuse_status", + "_has_cp_helix", + "_helix_cp_rank", + "_helix_cp_size", + "total_input_len_cp", + "get_ssm_block_base_index", + "py_served_prefix_end", + "has_pending_load", + "get_prefix_reservation", "host_kv_cache_block_offsets", "kv_cache_map", "py_multimodal_data", @@ -215,9 +240,18 @@ def private_importers(root: Path, pkg_dir: Path) -> dict: return offenders +# The one module of the transfer contract the package imports: the kinds of memory it lends. +CONTRACT_MEMORY = "tensorrt_llm._torch.disaggregation.base.memory" + + def disaggregation_imports(source: str, package: str) -> List[str]: + """Disaggregation modules a source imports, apart from the contract's memory module.""" prefix = "tensorrt_llm._torch.disaggregation" - return sorted(m for m in imported_modules(source, package) if m.startswith(prefix)) + return sorted( + m + for m in imported_modules(source, package) + if m.startswith(prefix) and CONTRACT_MEMORY not in (m, m.rpartition(".")[0]) + ) def internal_reads(source: str) -> List[str]: @@ -339,7 +373,7 @@ def subpackages(pkg_dir: Path) -> List[str]: @pytest.mark.cpu_only -def test_the_package_is_its_init_and_eight_private_modules(tmp_path): +def test_the_package_is_its_init_and_nine_private_modules(tmp_path): assert [p.name for p in package_files()] == MODULES assert subpackages(PKG_DIR) == [] (tmp_path / "__pycache__").mkdir() @@ -354,13 +388,17 @@ def test_the_protocols_keep_their_members_and_stay_unrelated(): def members(cls): return {n for n in vars(cls) if not n.startswith("_")} - assert members(Lease) == {"poll", "failure", "mark_arrived", "release"} + assert members(Lease) == {"poll", "failure", "writable", "release", "return_resumes_at_end"} assert members(Lender) == { "regions", + "page_layout", + "held_page_counts", "hold_regions", "lend_read", "lend_write", - "resume_bounds", + "return_writes", + "progress", + "poll_resume_bounds", } assert members(RegionHold) == {"release"} assert Lease not in RegionHold.__mro__ and RegionHold not in Lease.__mro__ @@ -368,6 +406,9 @@ def members(cls): assert Protocol in protocol.__mro__ assert not isinstance(object(), protocol), "checkable at run time" assert isinstance(Lease.failure, property) and isinstance(Lender.regions, property) + assert isinstance(Lease.return_resumes_at_end, property) + assert isinstance(Lease.writable, property) + assert isinstance(Lender.page_layout, property) # The checks are structural: a lease passes as a hold. lease = object.__new__(_lender._StagingLease) assert isinstance(lease, RegionHold) @@ -376,21 +417,63 @@ def params(func): return list(inspect.signature(func).parameters) assert params(Lease.poll) == ["self"] - assert params(Lease.mark_arrived) == ["self", "masks"] assert params(Lease.release) == ["self"] assert params(RegionHold.release) == ["self"] assert params(Lender.hold_regions) == ["self"] - assert params(Lender.lend_read) == ["self", "request", "start", "end"] - assert params(Lender.lend_write) == ["self", "request", "start", "end"] - assert params(Lender.resume_bounds) == ["self", "request"] + assert params(Lender.held_page_counts) == ["self"] + + def shape(func): + """The signature, annotated on its keyword-only parameters alone.""" + signature = inspect.signature(func) + kept = [ + p if p.kind is p.KEYWORD_ONLY else p.replace(annotation=p.empty) + for p in signature.parameters.values() + ] + return str(signature.replace(parameters=kept, return_annotation=signature.empty)) + + keywords = "*, memory_kind: 'MemoryKind', rows: 'RowSelection', ready_event" + for cls, event in ( + (Lender, "'Optional[torch.cuda.Event]'"), + (_lender._ManagerLender, "'torch.cuda.Event | None'"), + ): + lend = f"(self, request, start, end, {keywords}: {event})" + assert shape(cls.lend_read) == shape(cls.lend_write) == lend, cls + assert [shape(cls.return_writes), shape(cls.progress)] == ["(self, entries)", "(self)"], cls + assert params(cls.poll_resume_bounds) == ["self", "request"], cls + + +@pytest.mark.cpu_only +def test_all_arrived_is_one_value_that_pickles_and_copies_to_itself(): + import copy + import pickle + + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _types + + assert sharing.ALL_ARRIVED is ALL_ARRIVED is _types.ALL_ARRIVED + assert type(ALL_ARRIVED)() is ALL_ARRIVED, "the type makes a second value" + assert pickle.loads(pickle.dumps(ALL_ARRIVED)) is ALL_ARRIVED + assert copy.copy(ALL_ARRIVED) is ALL_ARRIVED and copy.deepcopy(ALL_ARRIVED) is ALL_ARRIVED + assert repr(ALL_ARRIVED) == "ALL_ARRIVED" + assert not hasattr(ALL_ARRIVED, "__dict__"), "it carries state" + assert not isinstance(ALL_ARRIVED, (tuple, list, np.ndarray)), "it reads as masks" def public_names(obj) -> Set[str]: return {n for n in dir(obj) if not n.startswith("_")} -LEASE_MEMBERS = {"poll", "failure", "mark_arrived", "release"} -STAGING_MEMBERS = {"regions", "hold_regions", "lend_read", "lend_write", "resume_bounds"} +LEASE_MEMBERS = {"poll", "failure", "writable", "release", "return_resumes_at_end"} +LENDER_MEMBERS = { + "regions", + "page_layout", + "held_page_counts", + "hold_regions", + "lend_read", + "lend_write", + "return_writes", + "progress", + "poll_resume_bounds", +} @pytest.mark.cpu_only @@ -404,11 +487,13 @@ class Owner: # anything a weak reference can point to owner = Owner() ref = weakref.ref(owner) - layout = SimpleNamespace(pool_groups=(), windows=()) # all the constructors read + # What constructors read. + layout = SimpleNamespace(pool_groups=(), windows=(), named_groups=(), tokens_per_block=1) shown = { - "staging lender": (_lender.Staging(ref, layout, None, (), None, ref), STAGING_MEMBERS), + "lender": (_lender._ManagerLender(ref, layout, None, (), None, ref), LENDER_MEMBERS), "staging lease": (_lender._StagingLease(owner, "read", 1), LEASE_MEMBERS), - "staging hold": (_lender._StagingHold(owner), {"release"}), + "region hold": (_lender._RegionHold(owner), {"release"}), + "in-place lease": (_lender._InPlaceLease(owner, "read", 1), LEASE_MEMBERS), } assert {what: public_names(obj) for what, (obj, _) in shown.items()} == { what: members for what, (_, members) in shown.items() @@ -417,8 +502,8 @@ class Owner: # anything a weak reference can point to leaky = _lender._StagingLease(owner, "read", 1) leaky.view = None assert public_names(leaky) - LEASE_MEMBERS == {"view"} - hooked = type("Hooked", (_lender.Staging,), {"on_free": _lender.Staging._on_free}) - assert public_names(hooked(ref, layout, None, (), None, ref)) - STAGING_MEMBERS == {"on_free"} + hooked = type("Hooked", (_lender._ManagerLender,), {"on_free": _lender._ManagerLender._on_free}) + assert public_names(hooked(ref, layout, None, (), None, ref)) - LENDER_MEMBERS == {"on_free"} @pytest.mark.cpu_only @@ -431,20 +516,21 @@ def fields(cls): for f in dataclasses.fields(cls) ) - assert fields(StagingOptions) == "tokens_per_fetch planned_fetches=1 max_bytes=None" + assert fields(StagingOptions) == "guaranteed_fetch_tokens capacity_in_fetches=1 max_bytes=None" assert ( fields(MemoryRegion) - == "layout_label address nbytes slot_bytes slot_count layer_group_indices" + == "layout_label address nbytes slot_bytes slot_count layer_group_indices memory_kind" ) - assert ( - fields(LayerGroupBlocks) - == "layer_group_index block_ordinals object_keys=None addresses=None region_index=None" + assert fields(LentLayerGroupRows) == ( + "layer_group_index page_ordinals object_keys=None slots=None region_indices=None" ) - assert fields(LeaseView) == "groups" - for cls in (StagingOptions, MemoryRegion, LayerGroupBlocks, LeaseView): + assert fields(LeaseView) == "groups start end" + assert fields(HeldPages) == "deduct unreclaimable by_reason version" + assert [reason.name for reason in HoldReason] == ["LOCKED", "DEFERRED_CLOSE"] + for cls in (StagingOptions, MemoryRegion, LentLayerGroupRows, LeaseView, HeldPages): assert cls.__dataclass_params__.frozen, cls - assert MemoryRegion("p", 4096, 8192, 4096, 2, (0,)) == MemoryRegion( - "p", 4096, 8192, 4096, 2, (0,) + assert MemoryRegion("p", 4096, 8192, 4096, 2, (0,), MemoryKind.HOST) == MemoryRegion( + "p", 4096, 8192, 4096, 2, (0,), MemoryKind.HOST ) assert ResumeBounds._fields == ("max_position", "min_position") assert issubclass(ResumeBounds, tuple) @@ -453,8 +539,8 @@ def beyond_fields(cls): named = {f.name for f in dataclasses.fields(cls)} return {n for n in dir(cls) if not n.startswith("_")} - named - assert beyond_fields(LayerGroupBlocks) == {"select"} - assert beyond_fields(LeaseView) == {"num_block_entries", "row_masks"} + assert beyond_fields(LentLayerGroupRows) == {"select"} + assert beyond_fields(LeaseView) == {"num_rows", "row_masks"} assert beyond_fields(MemoryRegion) == beyond_fields(StagingOptions) == set() assert {n for n in dir(ResumeBounds) if not n.startswith("_")} == { "max_position", @@ -471,30 +557,28 @@ def test_public_names_are_independent_of_the_lending_mechanism(): assert not hasattr(sharing, name), name assert not hasattr(Lender, "hold_staging") assert not hasattr(LeaseView, "block_masks") - assert not hasattr(LayerGroupBlocks, "staging_addresses") + assert not hasattr(LentLayerGroupRows, "staging_addresses") @pytest.mark.cpu_only -def test_views_and_regions_match_the_backend_contract_member_names(): - """Exercise the structural surface without importing the unmerged backend contract.""" - region = MemoryRegion("layout", 4096, 128, 64, 2, (3,)) - group = LayerGroupBlocks( +def test_rows_views_and_regions_carry_the_backend_contract_s_members(): + """Rows, views and regions are what the backend contract describes, and rows carry no member + the contract leaves out.""" + from tensorrt_llm._torch.disaggregation.base import backend as contract + + region = MemoryRegion("layout", 4096, 128, 64, 2, (3,), MemoryKind.HOST) + group = LentLayerGroupRows( layer_group_index=3, - block_ordinals=np.array([0, 7]), + page_ordinals=np.array([0, 7]), object_keys=np.zeros((2, 54), dtype=np.uint8), - addresses=np.array([4096, 4160]), - region_index=0, - ) - view = LeaseView((group,)) - assert all( - hasattr(region, name) - for name in ("layout_label", "address", "nbytes", "slot_bytes", "slot_count") + slots=np.array([0, 1]), + region_indices=(0,), ) - assert all( - hasattr(group, name) - for name in ("layer_group_index", "object_keys", "addresses", "region_index", "__len__") - ) - assert all(hasattr(view, name) for name in ("groups", "row_masks")) + view = LeaseView((group,), 0, 0) + assert isinstance(region, contract.SlotRegion) and isinstance(view, contract.View) + assert isinstance(group, contract.LayerGroupRows) + members = {name for name in vars(contract.LayerGroupRows) if not name.startswith("_")} + assert {f.name for f in dataclasses.fields(LentLayerGroupRows)} == members assert list(inspect.signature(LeaseView.row_masks).parameters) == ["self", "fill_value"] false_masks = view.row_masks() true_masks = view.row_masks(fill_value=True) @@ -503,112 +587,85 @@ def test_views_and_regions_match_the_backend_contract_member_names(): assert masks[0].shape == (2,) and masks[0].dtype == np.bool_ assert masks[0].flags.writeable and np.all(masks[0] == fill) assert region.nbytes == region.slot_bytes * region.slot_count - assert np.all(group.addresses >= region.address) - assert np.all(group.addresses + region.slot_bytes <= region.address + region.nbytes) + assert np.all(group.slots < region.slot_count) @pytest.mark.cpu_only def test_interface_docstrings_state_guarantees_without_staging_mechanisms(): forbidden = re.compile(r"\b(staging|stream|copy|copies|host|device|gpu)\b", re.IGNORECASE) - for cls in (MemoryRegion, LayerGroupBlocks, LeaseView, ResumeBounds, Lease, RegionHold, Lender): + for cls in ( + MemoryRegion, + LentLayerGroupRows, + LeaseView, + ResumeBounds, + Lease, + RegionHold, + Lender, + RowSelection, + HeldPages, + HoldReason, + ): entries = [cls] + [member for name, member in vars(cls).items() if not name.startswith("_")] for member in entries: - doc = inspect.getdoc(member) or "" + # A trailing non-normative performance note may name mechanisms; guarantees precede it. + doc = (inspect.getdoc(member) or "").split("Performance (non-normative):")[0] assert not forbidden.search(doc), (cls.__name__, doc) -@pytest.mark.cpu_only -@pytest.mark.parametrize( - "invalid_hook", ["_on_free", "_on_shrink", "_on_reset", "_on_shutdown", "_on_caches_closed"] +# The private hooks the manager calls on its lender. +MANAGER_HOOKS = ( + "_on_free", + "_on_shrink", + "_on_resume", + "_parking_row", + "_parked_row", + "_before_reset", + "_on_reset", + "_on_shutdown", + "_on_caches_closed", + "_is_locked", + "_defer_close", + "_keeps_pools", + "_history_cap", + "_page_lock_state", ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("invalid_hook", MANAGER_HOOKS) @pytest.mark.parametrize("replacement", [None, 1, "absent"]) -def test_install_checks_every_private_lifecycle_hook_before_attaching(invalid_hook, replacement): +def test_install_checks_every_private_manager_hook_before_attaching(invalid_hook, replacement): from types import SimpleNamespace from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager - hooks = { - name: lambda *args: None - for name in ("_on_free", "_on_shrink", "_on_reset", "_on_shutdown", "_on_caches_closed") - } + hooks = {name: lambda *args: None for name in MANAGER_HOOKS} if replacement == "absent": hooks.pop(invalid_hook) else: hooks[invalid_hook] = replacement manager = SimpleNamespace() - with pytest.raises(TypeError, match="all five manager lifecycle hooks"): + with pytest.raises(TypeError, match="all manager hooks"): _manager.install_lender(manager, SimpleNamespace(**hooks)) assert "_sharing" not in vars(manager) @pytest.mark.cpu_only -def test_install_accepts_the_complete_private_lifecycle_protocol(): +def test_install_accepts_the_complete_private_hook_protocol(): from types import SimpleNamespace - from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + from tensorrt_llm._torch.pyexecutor.kv_cache import _lender_hooks + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender, _manager - lender = SimpleNamespace( - **{ - name: lambda *args: None - for name in ("_on_free", "_on_shrink", "_on_reset", "_on_shutdown", "_on_caches_closed") - } - ) + lender = SimpleNamespace(**{name: lambda *args: None for name in MANAGER_HOOKS}) manager = SimpleNamespace() - assert isinstance(lender, _manager._ManagerHooks) + assert isinstance(lender, _lender_hooks._ManagerHooks) + assert set(MANAGER_HOOKS) <= set(vars(_lender_hooks._ManagerHooks)), "outside the protocol" + assert all(callable(getattr(_lender._ManagerLender, name)) for name in MANAGER_HOOKS) _manager.install_lender(manager, lender) assert manager._sharing is lender -@pytest.mark.cpu_only -@pytest.mark.parametrize("config_source", ["manager", "runtime", "runtime_none"]) -def test_manager_facade_reads_layer_configs_by_id_with_the_existing_fallback(config_source): - from types import SimpleNamespace - - from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager - - layers = [ - SimpleNamespace(layer_id="7"), - SimpleNamespace(layer_id=4), - SimpleNamespace(), - SimpleNamespace(layer_id=0), - ] - config = SimpleNamespace(layers=layers) - if config_source == "manager": - manager = SimpleNamespace(kv_cache_manager_py_config=config) - else: - manager = SimpleNamespace(impl=SimpleNamespace(init_config=config)) - if config_source == "runtime_none": - manager.kv_cache_manager_py_config = None - actual = _manager.layer_configs(manager) - assert list(actual) == [7, 4, 2, 0] - assert all(actual[layer_id] is layer for layer_id, layer in zip([7, 4, 2, 0], layers)) - - -@pytest.mark.cpu_only -def test_manager_facade_preserves_layer_group_order_and_normalizes_ids(): - from types import SimpleNamespace - - from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager - - grouping = [[np.int64(5), "7"], [2]] - manager = SimpleNamespace(impl=SimpleNamespace(layer_grouping=grouping)) - actual = _manager.layer_group_layer_ids(manager) - assert actual == [[5, 7], [2]] - assert all(type(layer_id) is int for group in actual for layer_id in group) - assert actual is not grouping and all(a is not b for a, b in zip(actual, grouping)) - - -@pytest.mark.cpu_only -def test_manager_facade_returns_the_runtime_pool_group_descriptors_unchanged(): - from types import SimpleNamespace - - from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager - - descriptors = [object(), object()] - manager = SimpleNamespace(impl=SimpleNamespace(pool_group_descs=descriptors)) - assert _manager.pool_group_descriptors(manager) is descriptors - - @pytest.mark.cpu_only def test_the_attach_functions_keep_their_signatures(): kind = inspect.Parameter @@ -684,6 +741,44 @@ def test_the_native_path_loads_nothing_of_the_package(): assert result["loaded"] == [] +def class_definitions(source: str, name: str) -> int: + """How many classes called ``name`` a source defines, nested ones included.""" + tree = ast.parse(source) + return sum(isinstance(node, ast.ClassDef) and node.name == name for node in ast.walk(tree)) + + +@pytest.mark.cpu_only +def test_the_package_and_the_contract_share_one_memory_kind(): + from tensorrt_llm._torch.disaggregation.base import memory + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender, _types + + for holder in (sharing, _types, _lender): + assert holder.MemoryKind is memory.MemoryKind, holder.__name__ + assert MemoryKind.__module__ == memory.__name__ == CONTRACT_MEMORY + assert memory.__all__ == ["MemoryKind"] + assert [p.name for p in package_files() if class_definitions(p.read_text(), "MemoryKind")] == [] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("first", ["lender", "base_memory"]) +def test_the_lender_and_the_contract_s_memory_module_import_in_either_order(first): + order = [f"{SH}._lender", CONTRACT_MEMORY] + if first == "base_memory": + order.reverse() + result = run_python( + "import json, sys\n" + f"import {order[0]}\n" + f"loaded = sorted(m for m in sys.modules if m.startswith('{SH}'))\n" + f"import {order[1]}\n" + f"s, memory = sys.modules['{SH}'], sys.modules['{CONTRACT_MEMORY}']\n" + f"same = [m.MemoryKind is memory.MemoryKind for m in (s, sys.modules['{SH}._lender'])]\n" + "print(json.dumps({'file': s.__file__, 'loaded': loaded, 'same': same}))\n" + ) + assert result["same"] == [True, True] + if first == "base_memory": + assert result["loaded"] == [], "importing the contract's memory module loads the lender" + + @pytest.mark.cpu_only def test_only_the_manager_facade_touches_manager_internals(): offenders = {} @@ -696,6 +791,99 @@ def test_only_the_manager_facade_touches_manager_internals(): assert offenders == {} +# The manager side's one module that calls a lender's hooks, and the manager's own module. +HOOKS_MODULE = ROOT / "_torch/pyexecutor/kv_cache/_lender_hooks.py" +MANAGER_MODULE = ROOT / "_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py" +# A manager's records of lending; the first belongs to the manager and the hooks module alone. +LENDING_RECORDS = ("_held_commits", "_fresh_pages_filled") +# The methods that change a dict. +DICT_WRITES = {"pop", "popitem", "setdefault", "clear", "update", "__setitem__", "__delitem__"} + + +def hook_uses(source: str) -> List[str]: + """The lender's manager hooks a source names as attributes.""" + return sorted( + node.attr + for node in ast.walk(ast.parse(source)) + if isinstance(node, ast.Attribute) and node.attr in MANAGER_HOOKS + ) + + +def record_writes(source: str, records) -> List[str]: + """``record:line`` of each assignment, ``del`` or dict update a source makes to ``records``.""" + hits = [] + for node in ast.walk(ast.parse(source)): + written = None + if isinstance(node, (ast.Attribute, ast.Subscript)) and isinstance( + node.ctx, (ast.Store, ast.Del) + ): + written = node if isinstance(node, ast.Attribute) else node.value + elif ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in DICT_WRITES + ): + written = node.func.value + if isinstance(written, ast.Attribute) and written.attr in records: + hits.append(f"{written.attr}:{node.lineno}") + return sorted(hits) + + +def scan_outside_the_package(names, scan) -> dict: + """Non-empty ``scan(source)`` of each other module outside the package naming ``names``.""" + found = {} + for path in ROOT.rglob("*.py"): + resolved = path.resolve() + if PKG_DIR in resolved.parents or resolved in (HOOKS_MODULE, MANAGER_MODULE): + continue + source = path.read_text(encoding="utf-8", errors="replace") + if not any(name in source for name in names): + continue + try: + hits = scan(source) + except SyntaxError: + continue # not importable, so it runs nothing + if hits: + found[path.relative_to(ROOT).as_posix()] = hits + return found + + +@pytest.mark.cpu_only +def test_the_manager_side_calls_a_lender_s_hooks_through_the_hooks_module_alone(): + """Outside the package only the hooks module names a lender's manager hooks.""" + assert hook_uses(MANAGER_MODULE.read_text()) == [] + assert scan_outside_the_package(MANAGER_HOOKS, hook_uses) == {} + assert set(hook_uses(HOOKS_MODULE.read_text())) == set(MANAGER_HOOKS) + + +@pytest.mark.cpu_only +def test_a_manager_s_lending_records_are_written_on_the_manager_side_alone(): + """Only the manager's module and the hooks module write a manager's lending records.""" + package = { + path.name: record_writes(path.read_text(), LENDING_RECORDS) for path in package_files() + } + assert {name: hits for name, hits in package.items() if hits} == {} + lender_state = LENDING_RECORDS[:1] + found = scan_outside_the_package(lender_state, lambda s: record_writes(s, lender_state)) + assert found == {} + + +@pytest.mark.cpu_only +def test_the_manager_records_no_readiness_event_for_a_lender(): + """No module records readiness for a lender: it takes its readiness from the caller.""" + hooks = HOOKS_MODULE.read_text() + for gone in ("record_context_ready", "discard_context_ready", "record_receive_after_reset"): + assert gone not in hooks, gone + managers = ROOT / "_torch/pyexecutor/kv_cache/mamba_cache_manager" + for path in (MANAGER_MODULE, *sorted(managers.glob("*.py"))): + source = path.read_text() + assert "_context_ready" not in source and "record_receive_after_reset" not in source + package = {path.name: path.read_text() for path in package_files()} + for name, source in package.items(): + for gone in ("take_disagg_receive_ready", "_context_ready", "record_event"): + assert gone not in source, (name, gone) + + @pytest.mark.cpu_only def test_the_context_output_check_reads_the_request_alone(): """The fetch's context-output check reads two request attributes and never the executor's @@ -721,7 +909,8 @@ def __init__(self, logits, outputs): def written_files() -> List[Path]: - return package_files() + sorted(TESTS_DIR.glob("*.py")) + contract_memory = ROOT.parent.joinpath(*CONTRACT_MEMORY.split(".")).with_suffix(".py") + return package_files() + [contract_memory] + sorted(TESTS_DIR.glob("*.py")) @pytest.mark.cpu_only @@ -744,122 +933,393 @@ def header(years: str) -> str: # -- the public types' own rules -------------------------------------------------------------- -def staging_blocks(n=3, layer_group_index=0, region_index=0): - return LayerGroupBlocks( +def staged_rows(n=3, layer_group_index=0, region=0): + """``n`` named rows in slots ``0 .. n - 1`` of region ``region``.""" + return LentLayerGroupRows( layer_group_index, np.arange(n), np.repeat(np.arange(n, dtype=np.uint8)[:, None], 54, axis=1), - np.arange(n, dtype=np.int64) * 4096, - region_index, + np.arange(n, dtype=np.int64), + (region,), ) @pytest.mark.cpu_only -def test_layer_group_blocks_carry_all_three_placements_or_none_in_their_shapes(): - block_ordinals = np.arange(3) +def test_memory_kinds_and_row_selections_are_plain_enums(): + assert [(m.name, m.value) for m in MemoryKind] == [("HOST", "host"), ("DEVICE", "device")] + assert [(m.name, m.value) for m in RowSelection] == [("NAMED", "named"), ("ALL", "all")] + assert set(MemoryKind.__members__) == {"HOST", "DEVICE"} + assert set(RowSelection.__members__) == {"NAMED", "ALL"} + for cls in (MemoryKind, RowSelection): + assert cls.__bases__ == (enum.Enum,), "a member would equal its value" + assert MemoryKind.HOST != "host" and RowSelection.NAMED != "named" + assert MemoryKind("device") is MemoryKind.DEVICE and RowSelection("all") is RowSelection.ALL + + +@pytest.mark.cpu_only +def test_a_region_carries_one_memory_kind(): + for memory in MemoryKind: + assert MemoryRegion("p", 4096, 8192, 4096, 2, (0,), memory).memory_kind is memory + host, device = (MemoryRegion("p", 4096, 8192, 4096, 2, (0,), m) for m in MemoryKind) + assert host != device + for wrong in ("host", None, 0, RowSelection.NAMED): + with pytest.raises(TypeError, match="memory_kind must be a MemoryKind"): + MemoryRegion("p", 4096, 8192, 4096, 2, (0,), wrong) + with pytest.raises(TypeError): + MemoryRegion("p", 4096, 8192, 4096, 2, (0,)) # no default + + +_REGION = dict( + layout_label="p", + address=4096, + nbytes=8192, + slot_bytes=4096, + slot_count=2, + layer_group_indices=(0,), + memory_kind=MemoryKind.HOST, +) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "field, value, error", + [ + ("address", True, TypeError), + ("nbytes", 8192.0, TypeError), + ("address", 0, ValueError), + ("nbytes", 0, ValueError), + ("slot_bytes", -1, ValueError), + ("slot_count", 0, ValueError), + ("slot_count", 3, ValueError), + ("address", (1 << 63) - 8192, ValueError), + ("layer_group_indices", [0], TypeError), + ("layer_group_indices", (0, -1), ValueError), + ("layer_group_indices", (False,), TypeError), + ], +) +def test_a_region_with_a_field_out_of_its_range_is_refused(field, value, error): + with pytest.raises(error): + MemoryRegion(**{**_REGION, field: value}) + + +@pytest.mark.cpu_only +def test_a_region_keeps_its_integers_as_ints_up_to_the_int64_end(): + end = (1 << 63) - 1 + region = MemoryRegion(**{**_REGION, "address": np.int64(end - 8192)}) + assert type(region.address) is int and region.address + region.nbytes == end + indices = MemoryRegion( + **{**_REGION, "layer_group_indices": (np.int32(2), 0)} + ).layer_group_indices + assert indices == (2, 0) and all(type(index) is int for index in indices) + + +@pytest.mark.cpu_only +def test_regions_of_one_memory_kind_that_overlap_by_a_byte_are_refused(): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _types + + def region(address, kind=MemoryKind.DEVICE): + return MemoryRegion("r", address, 4096, 4096, 1, (0,), kind) + + _types._check_disjoint([region(8192), region(4096), region(4096, MemoryKind.HOST)]) + overlapping = [region(1 << 20), region(4096, MemoryKind.HOST), region(4096), region(8191)] + with pytest.raises(ValueError, match="regions 2 and 3 overlap in DEVICE memory"): + _types._check_disjoint(overlapping) + + +@pytest.mark.cpu_only +def test_layer_group_rows_take_slots_with_their_regions_or_neither_and_names_apart(): + page_ordinals = np.arange(3) object_keys = np.zeros((3, 54), np.uint8) - addresses = np.arange(3, dtype=np.int64) + slots = np.array([5, 0, 2]) refused = [ - (ValueError, (0, block_ordinals, object_keys, None, None)), - (ValueError, (0, block_ordinals, None, addresses, 0)), - (ValueError, (0, block_ordinals, object_keys, addresses, None)), - (ValueError, (0, np.zeros((3, 1)))), - (ValueError, (0, block_ordinals, np.zeros((3, 53), np.uint8), addresses, 0)), - (ValueError, (0, block_ordinals, np.zeros((3, 54), np.int8), addresses, 0)), - (ValueError, (0, block_ordinals, object_keys, addresses[:2], 0)), - (ValueError, (-1, block_ordinals)), - (ValueError, (0, block_ordinals, object_keys, addresses, -1)), - (TypeError, (True, block_ordinals)), - (TypeError, (0, block_ordinals, object_keys, addresses, 1.0)), + # Slots without their regions, or regions without slots. + (ValueError, dict(slots=slots)), + (ValueError, dict(object_keys=object_keys, region_indices=(0,))), + # Shapes, dtypes and values, object_keys checked with slots or without. + (ValueError, dict(page_ordinals=np.zeros((3, 1)))), + (ValueError, dict(object_keys=np.zeros((3, 53), np.uint8))), + (ValueError, dict(object_keys=np.zeros((3, 54), np.int8))), + (ValueError, dict(object_keys=object_keys[:2], slots=slots, region_indices=(0,))), + (ValueError, dict(layer_group_index=-1)), + (ValueError, dict(slots=slots[:2], region_indices=(0,))), + (ValueError, dict(slots=slots.astype(np.float64), region_indices=(0,))), + (ValueError, dict(slots=np.array([1, -1, 2]), region_indices=(0,))), + (ValueError, dict(slots=slots, region_indices=())), + (ValueError, dict(slots=slots, region_indices=(1, 1))), + (ValueError, dict(slots=slots, region_indices=(0, -1))), + (TypeError, dict(layer_group_index=True)), + (TypeError, dict(slots=slots, region_indices=(1.0,))), + (TypeError, dict(slots=slots, region_indices=(True,))), + (TypeError, dict(slots=slots, region_indices=1)), ] - for error, args in refused: + for error, fields in refused: with pytest.raises(error): - LayerGroupBlocks(*args) - unplaced = LayerGroupBlocks(1, block_ordinals) - assert (unplaced.object_keys, unplaced.addresses, unplaced.region_index) == ( - None, + LentLayerGroupRows(**{"layer_group_index": 0, "page_ordinals": page_ordinals, **fields}) + unplaced = LentLayerGroupRows(1, page_ordinals) + assert len(unplaced) == 3 and unplaced.page_ordinals.dtype == np.int64 + placements = ("object_keys", "slots", "region_indices") + assert [getattr(unplaced, name) for name in placements] == [None] * 3 + named = LentLayerGroupRows(1, page_ordinals, object_keys) + assert named.object_keys.shape == (3, 54) and (named.slots, named.region_indices) == ( None, None, ) - assert len(unplaced) == 3 and unplaced.block_ordinals.dtype == np.int64 - - -@pytest.mark.cpu_only -def test_layer_group_blocks_arrays_are_read_only_views(): - block_ordinals = np.arange(3, dtype=np.int64) - blocks = LayerGroupBlocks( - 0, block_ordinals, np.zeros((3, 54), np.uint8), np.arange(3, dtype=np.int64), 0 + by_slot = LentLayerGroupRows( + 1, page_ordinals, slots=slots.astype(np.int32), region_indices=[np.int64(4), 1] ) - for array in (blocks.block_ordinals, blocks.object_keys, blocks.addresses): - assert not array.flags.writeable + assert by_slot.object_keys is None + assert by_slot.slots.dtype == np.int64 and by_slot.slots.tolist() == [5, 0, 2] + assert by_slot.region_indices == (4, 1) + assert all(type(index) is int for index in by_slot.region_indices) + keyed = LentLayerGroupRows(1, page_ordinals, object_keys, slots=slots, region_indices=(0,)) + assert keyed.object_keys.shape == (3, 54) and keyed.region_indices == (0,) + + +def check_immutable(rows, *writable): + """No array of ``rows`` can be written or made writable, or shares memory with ``writable``.""" + arrays = [ + getattr(rows, name) + for name in ("page_ordinals", "object_keys", "slots") + if getattr(rows, name) is not None + ] + for array in arrays: + with pytest.raises(ValueError): + array.flags.writeable = True with pytest.raises(ValueError): array[0] = 1 - assert block_ordinals.flags.writeable, "the caller's own array is left as it was" + assert not any(np.shares_memory(array, other) for other in writable) + + +@pytest.mark.cpu_only +def test_layer_group_rows_arrays_are_immutable_copies_of_writable_ones(): + page_ordinals = np.arange(3, dtype=np.int64) + keys = np.zeros((3, 54), np.uint8) + staging_slots = np.arange(3, dtype=np.int64) + slots = np.array([4, 0, 9], dtype=np.int64) + staged = LentLayerGroupRows(0, page_ordinals, keys, staging_slots, (0,)) + by_slot = LentLayerGroupRows(0, page_ordinals, keys, slots=slots, region_indices=(1, 0)) + check_immutable(staged, page_ordinals, keys, staging_slots) + check_immutable(by_slot, page_ordinals, keys, slots) + assert all(a.flags.writeable for a in (page_ordinals, keys, staging_slots, slots)) + # A caller's read-only view of its writable array is copied too. + frozen = page_ordinals.view() + frozen.flags.writeable = False + check_immutable(LentLayerGroupRows(0, frozen), page_ordinals) + # Arrays already immutable are kept, not copied again. + kept = LentLayerGroupRows(0, by_slot.page_ordinals, slots=by_slot.slots, region_indices=(0,)) + assert kept.page_ordinals is by_slot.page_ordinals and kept.slots is by_slot.slots + for picked in ( + staged.select(np.array([True, False, True])), + by_slot.select(np.array([False, True, True])), + ): + check_immutable(picked, page_ordinals, keys, staging_slots, slots) @pytest.mark.cpu_only -def test_select_keeps_the_marked_block_entries_in_order(): - blocks = staging_blocks(4, layer_group_index=2, region_index=1) - picked = blocks.select(np.array([True, False, True, True])) - assert picked.block_ordinals.tolist() == [0, 2, 3] - assert picked.addresses.tolist() == [0, 8192, 12288] - np.testing.assert_array_equal(picked.object_keys, blocks.object_keys[[0, 2, 3]]) - assert picked.object_keys.shape == (3, 54) and ( - picked.layer_group_index, - picked.region_index, - ) == (2, 1) - assert LayerGroupBlocks(0, np.arange(2)).select(np.array([False, True])).object_keys is None +def test_select_keeps_the_marked_rows_in_order(): + rows = staged_rows(4, layer_group_index=2, region=1) + picked = rows.select(np.array([True, False, True, True])) + assert picked.page_ordinals.tolist() == [0, 2, 3] + assert picked.slots.tolist() == [0, 2, 3] + np.testing.assert_array_equal(picked.object_keys, rows.object_keys[[0, 2, 3]]) + assert picked.object_keys.shape == (3, 54) and picked.layer_group_index == 2 + assert picked.region_indices == (1,) + assert LentLayerGroupRows(0, np.arange(2)).select(np.array([False, True])).object_keys is None with pytest.raises(ValueError): - blocks.select(np.array([True, False])) + rows.select(np.array([True, False])) with pytest.raises(ValueError): - blocks.select(np.array([1, 0, 1, 1])) + rows.select(np.array([1, 0, 1, 1])) + + +@pytest.mark.cpu_only +def test_select_keeps_the_slots_aligned_and_their_regions(): + keys = np.arange(4 * 54, dtype=np.uint8).reshape(4, 54) + by_slot = LentLayerGroupRows( + 3, np.array([0, 1, 4, 5]), keys, slots=np.array([7, 3, 0, 9]), region_indices=(2, 0) + ) + picked = by_slot.select(np.array([False, True, True, False])) + assert picked.page_ordinals.tolist() == [1, 4] and picked.slots.tolist() == [3, 0] + np.testing.assert_array_equal(picked.object_keys, keys[[1, 2]]) + assert (picked.layer_group_index, picked.region_indices) == (3, (2, 0)) + assert not picked.slots.flags.writeable + unnamed = LentLayerGroupRows(3, np.arange(2), slots=np.array([1, 0]), region_indices=(0,)) + none = unnamed.select(np.array([False, False])) + assert none.object_keys is None and none.slots.tolist() == [] and none.region_indices == (0,) + + +def on_immutable_bytes(values, dtype=np.int64) -> np.ndarray: + """``values`` as an array on immutable bytes, as the lender builds its own.""" + array = np.ascontiguousarray(values, dtype=dtype) + return np.frombuffer(array.tobytes(), dtype=dtype).reshape(array.shape) + + +@pytest.mark.cpu_only +def test_rows_and_views_built_without_checks_are_what_the_constructors_make(): + """On values that pass the constructors' checks, the unchecked builds are what they make.""" + rng = np.random.default_rng(126) + seen = set() + for _ in range(300): + n = int(rng.integers(0, 6)) + layer_group_index = int(rng.integers(0, 4)) + ordinals = on_immutable_bytes(np.sort(rng.choice(64, n, replace=False))) + keys = ( + None + if rng.random() < 0.4 + else on_immutable_bytes(rng.integers(0, 256, (n, 54)), np.uint8) + ) + placed = bool(rng.random() < 0.6) + seen.add(placed) + fields = dict(object_keys=keys) + if placed: + regions = tuple(int(r) for r in rng.permutation(5)[: int(rng.integers(1, 4))]) + fields.update( + slots=on_immutable_bytes(rng.integers(0, 1000, n)), region_indices=regions + ) + public = LentLayerGroupRows(layer_group_index, ordinals, **fields) + lent = LentLayerGroupRows._unchecked(layer_group_index, ordinals, **fields) + assert type(lent) is LentLayerGroupRows + for field in dataclasses.fields(LentLayerGroupRows): + mine, theirs = getattr(lent, field.name), getattr(public, field.name) + if isinstance(theirs, np.ndarray): + assert mine is theirs, field.name + else: + assert (mine, type(mine)) == (theirs, type(theirs)), field.name + assert len(lent) == len(public) and repr(lent) == repr(public) + mask = rng.random(n) < 0.5 + assert repr(lent.select(mask)) == repr(public.select(mask)) + with pytest.raises(dataclasses.FrozenInstanceError): + lent.slots = None + other = LentLayerGroupRows(layer_group_index + 4, on_immutable_bytes(np.arange(2))) + groups = (lent, other) if rng.random() < 0.5 else (other, lent) + view, built = LeaseView(groups, 0, 0), LeaseView._unchecked(groups, 0, 0) + assert type(built) is LeaseView and built.groups is view.groups + assert built.num_rows == view.num_rows + assert [m.tolist() for m in built.row_masks(True)] == [ + m.tolist() for m in view.row_masks(True) + ] + assert seen == {True, False} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("name", ["lend_read", "lend_write"]) +def test_a_lend_checks_memory_and_rows_before_anything_else(name): + """Wrong lend arguments raise before the range check; right ones reach the shut-down error.""" + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + class Owner: # anything a weak reference can point to + pass + + gone = weakref.ref(Owner()) + assert gone() is None + layout = SimpleNamespace(pool_groups=(), windows=(), named_groups=(), tokens_per_block=4) + lender = _lender._ManagerLender(gone, layout, None, (), None, gone) + lend = getattr(lender, name) + request = SimpleNamespace(py_request_id=1, prompt_len=97) + host, named = MemoryKind.HOST, RowSelection.NAMED + for call in ( + lambda: lend(request, 0, 4), + lambda: lend(request, 0, 4, memory_kind=host), + lambda: lend(request, 0, 4, rows=named), + lambda: lend(request, 0, 4, host, named), + ): + with pytest.raises(TypeError): + call() + for memory, rows, error, why in ( + ("host", named, TypeError, "memory_kind must be a MemoryKind"), + (None, named, TypeError, "memory_kind must be a MemoryKind"), + (named, named, TypeError, "memory_kind must be a MemoryKind"), + (host, "named", TypeError, "rows must be a RowSelection"), + (host, host, TypeError, "rows must be a RowSelection"), + (host, RowSelection.ALL, ValueError, "RowSelection.ALL is not lent"), + ): + with pytest.raises(error, match=why): + # A reversed range as well. + lend(request, 4, 0, memory_kind=memory, rows=rows, ready_event=None) + with pytest.raises(TypeError, match="ready_event"): + lend(request, 4, 0, memory_kind=host, rows=named) + for wrong in (1, "event", object()): + with pytest.raises(TypeError, match="ready_event must be an event"): + lend(request, 4, 0, memory_kind=host, rows=named, ready_event=wrong) + + class Asked: # an event its holder asks, which no stream can wait for + def query(self): + return True + + class Waited(Asked): + def wait(self, stream=None): + pass + + # A staging lease's copies wait for its event on a stream; an in-place lease only asks it. + with pytest.raises(TypeError, match=r"an event with query\(\) and wait\(\) or None"): + lend(request, 4, 0, memory_kind=host, rows=named, ready_event=Asked()) + for kind, event in ((host, Waited()), (MemoryKind.DEVICE, Asked())): + with pytest.raises(ValueError, match="bad token range"): + lend(request, 4, 0, memory_kind=kind, rows=named, ready_event=event) + for boundary, value in (("start", 0.9), ("start", True), ("end", 4.9), ("end", "4")): + start, end = (value, 4) if boundary == "start" else (0, value) + with pytest.raises(TypeError, match=f"{boundary} must be an integer"): + lend( + request, + start, + end, + memory_kind=MemoryKind.DEVICE, + rows=RowSelection.ALL, + ready_event=None, + ) + assert lender._open_lease_count() == 0, "a refused lend opened a lease" + with pytest.raises(ValueError, match="bad token range"): + lend(request, 4, 0, memory_kind=host, rows=named, ready_event=None) + with pytest.raises(RuntimeError, match=f"lends no more: {_lender._REASON_SHUT_DOWN}"): + lend(request, 0, 4, memory_kind=host, rows=named, ready_event=None) + assert lender._open_lease_count() == 0, "a lend that raised opened a lease" @pytest.mark.cpu_only -def test_a_lease_view_holds_one_block_collection_per_layer_group(): - view = LeaseView([staging_blocks(3, 0), staging_blocks(2, 1, region_index=1)]) - assert isinstance(view.groups, tuple) and view.num_block_entries == 5 +def test_a_lease_view_holds_one_row_collection_per_layer_group(): + view = LeaseView([staged_rows(3, 0), staged_rows(2, 1, region=1)], 0, 0) + assert isinstance(view.groups, tuple) and view.num_rows == 5 masks = view.row_masks() assert [m.shape for m in masks] == [(3,), (2,)] assert all(m.dtype == np.bool_ and not m.any() and m.flags.writeable for m in masks) assert all(m.all() for m in view.row_masks(True)) - assert LeaseView(()).num_block_entries == 0 and LeaseView(()).row_masks() == () + assert LeaseView((), 0, 0).num_rows == 0 and LeaseView((), 0, 0).row_masks() == () with pytest.raises(ValueError): - LeaseView((staging_blocks(3, 0), staging_blocks(1, 0))) + LeaseView((staged_rows(3, 0), staged_rows(1, 0)), 0, 0) with pytest.raises(TypeError): - LeaseView((object(),)) + LeaseView((object(),), 0, 0) @pytest.mark.cpu_only -def test_sparse_block_entries_keep_group_position_masks_and_placement_alignment(): +def test_sparse_rows_keep_group_position_masks_and_placement_alignment(): keys = np.arange(3 * 54, dtype=np.uint8).reshape(3, 54) - blocks = LayerGroupBlocks( + rows = LentLayerGroupRows( layer_group_index=7, - block_ordinals=np.array([0, 3, 9]), + page_ordinals=np.array([0, 3, 9]), object_keys=keys, - addresses=np.array([12288, 4096, 8192]), - region_index=2, + slots=np.array([3, 1, 2]), + region_indices=(2,), ) - empty = LayerGroupBlocks(layer_group_index=2, block_ordinals=np.array([], dtype=np.int64)) - repeated = LayerGroupBlocks(layer_group_index=4, block_ordinals=np.array([3, 9])) - view = LeaseView(groups=(blocks, empty, repeated)) - # Five group/block entries, but only three distinct token block ordinals. - assert view.num_block_entries == 5 + empty = LentLayerGroupRows(layer_group_index=2, page_ordinals=np.array([], dtype=np.int64)) + repeated = LentLayerGroupRows(layer_group_index=4, page_ordinals=np.array([3, 9])) + view = LeaseView(groups=(rows, empty, repeated), start=0, end=0) + # Five layer-group/page rows, but only three distinct page ordinals. + assert view.num_rows == 5 masks = view.row_masks() assert [mask.shape for mask in masks] == [(3,), (0,), (2,)] masks[0][1] = True picked = view.groups[0].select(masks[0]) - assert picked.block_ordinals.tolist() == [3] - assert picked.addresses.tolist() == [4096] + assert picked.page_ordinals.tolist() == [3] + assert picked.slots.tolist() == [1] np.testing.assert_array_equal(picked.object_keys, keys[[1]]) - assert (picked.layer_group_index, picked.region_index) == (7, 2) + assert (picked.layer_group_index, picked.region_indices) == (7, (2,)) assert not masks[2].any() @pytest.mark.cpu_only -def test_renamed_bounds_keep_maximum_first_and_region_fields_accept_keywords(): +def test_resume_bounds_keep_maximum_first_and_region_fields_accept_keywords(): bounds = ResumeBounds(max_position=96, min_position=32) - assert bounds._fields == ("max_position", "min_position") assert tuple(bounds) == (96, 32) assert ResumeBounds(0, 1).min_position > ResumeBounds(0, 1).max_position region = MemoryRegion( @@ -869,27 +1329,59 @@ def test_renamed_bounds_keep_maximum_first_and_region_fields_accept_keywords(): slot_bytes=4096, slot_count=2, layer_group_indices=(7, 2), + memory_kind=MemoryKind.HOST, ) assert region.layout_label == "layout:lg0+1" assert region.layer_group_indices == (7, 2) + assert region.memory_kind is MemoryKind.HOST @pytest.mark.cpu_only def test_staging_options_take_positive_integers(): options = StagingOptions(256) - assert (options.tokens_per_fetch, options.planned_fetches, options.max_bytes) == (256, 1, None) - assert StagingOptions(256, planned_fetches=4, max_bytes=1 << 30).max_bytes == 1 << 30 - for bad in (dict(tokens_per_fetch=0), dict(tokens_per_fetch=8, planned_fetches=0)): + assert dataclasses.astuple(options) == (256, 1, None) + assert StagingOptions(256, capacity_in_fetches=4, max_bytes=1 << 30).max_bytes == 1 << 30 + for bad in ( + dict(guaranteed_fetch_tokens=0), + dict(guaranteed_fetch_tokens=8, capacity_in_fetches=0), + ): with pytest.raises(ValueError): StagingOptions(**bad) with pytest.raises(ValueError): StagingOptions(8, max_bytes=0) for bad in ( - dict(tokens_per_fetch=True), - dict(tokens_per_fetch=8.0), - dict(tokens_per_fetch=8, max_bytes=1.5), + dict(guaranteed_fetch_tokens=True), + dict(guaranteed_fetch_tokens=8.0), + dict(guaranteed_fetch_tokens=8, max_bytes=1.5), ): with pytest.raises(TypeError): StagingOptions(**bad) with pytest.raises(dataclasses.FrozenInstanceError): - options.planned_fetches = 2 + options.capacity_in_fetches = 2 + + +@pytest.mark.cpu_only +def test_attach_refuses_a_cross_attention_manager_before_allocating(): + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.bindings.internal.batch_manager import CacheType + + manager = object.__new__(KVCacheManagerV2) + manager.kv_cache_type = CacheType.CROSS + with pytest.raises(ValueError, match="CROSS"): + attach_lender(manager, scope=b"cross", staging=StagingOptions(32)) + assert manager._sharing is None + + +@pytest.mark.cpu_only +def test_a_view_takes_its_lend_s_range_and_refuses_another(): + """A view holds integer bounds ``0 <= start <= end`` and refuses any other.""" + rows = staged_rows(2, 0) + view = LeaseView((rows,), 32, 64) + assert (view.start, view.end) == (32, 64) + assert (LeaseView((rows,), np.int64(0), 0).start, view.num_rows) == (0, 2) + for start, end in ((1.0, 2), (0, "2"), (True, 2)): + with pytest.raises(TypeError): + LeaseView((rows,), start, end) + for start, end in ((-1, 2), (3, 2)): + with pytest.raises(ValueError): + LeaseView((rows,), start, end) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_return_writes.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_return_writes.py new file mode 100644 index 000000000000..a46f433c8c02 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_return_writes.py @@ -0,0 +1,1575 @@ +# 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. +"""Returned write leases and the resume bounds they leave: windows, claimed prefixes, stages.""" + +import random +import weakref +from contextlib import ExitStack +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.pyexecutor.kv_cache import _lender_hooks +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + ALL_ARRIVED, + HoldReason, + MemoryKind, + MemoryRegion, + ResumeBounds, + RowSelection, + StagingOptions, + _fetch_ledger, + _lender, + _manager, + _ranges, + attach_lender, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._identity import Identity +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._layout import DevicePool +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._page_locks import LockMode +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._slots import SlotQueue +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttnLifeCycle + +gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools") + +TPB = 32 +PROMPT = 97 # three whole blocks and one token +TOKENS = list(range(1000, 1000 + PROMPT)) +OTHER_TOKENS = list(range(5000, 5000 + PROMPT)) # no block shared with TOKENS +DEVICE, HOST = MemoryKind.DEVICE, MemoryKind.HOST +NAMED, ALL = RowSelection.NAMED, RowSelection.ALL + + +class Rules(_lender._ManagerLender): + """A lender of a class of its own: every lend and every return goes through the rules.""" + + +def attach(mgr, rules=False, staging_tokens=4 * TPB): + staging = StagingOptions(staging_tokens) + if rules: + return _lender._attach_lender(mgr, scope=b"returns", staging=staging, cls=Rules) + return attach_lender(mgr, scope=b"returns", staging=staging) + + +def receiving(kit, mgr, request_id, tokens=TOKENS, hybrid=False): + """A generation-init request admitted to receive ``tokens``, with recurrent state prepared.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestType + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + request = kit.make_request( + request_id, tokens, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + ) + assert mgr.prepare_disagg_gen_init(request) + mgr.take_disagg_receive_ready(request) + if hybrid: + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + mgr.prepare_resources(batch) + return request + + +def ready_write(lender, mgr, request, end=None, start=0, memory_kind=DEVICE, rows=ALL): + """A write lease of ``[start, end)`` and its view once ready.""" + end = request.prompt_len if end is None else end + event = torch.cuda.Event() + event.record(mgr._stream) + lease = lender.lend_write( + request, start, end, memory_kind=memory_kind, rows=rows, ready_event=event + ) + assert lease.failure is None and lease._first_poll_failure is None, lease.failure + mgr._stream.synchronize() + for _ in range(4): + lender.progress() + view = lease.poll() + if view is not None: + assert (view.start, view.end) == (start, end), (view.start, view.end) + return lease, view + mgr._stream.synchronize() + raise AssertionError(f"the write never became ready: {lease.failure}") + + +def ready_in_place(kit, lender, request, start=0, end=None): + """A fast write lease of ``[start, end)`` and its view, ready without a round of progress.""" + end = request.prompt_len if end is None else end + lease = lender.lend_write(request, start, end, **kit.DEVICE_ALL) + view = lease.poll() + assert view is not None, lease.failure + return lease, view + + +def records(lender, request_id): + """The ledger's records of the request, as plain values.""" + ledger = lender._ledger + fetch = ledger._latest_fetches.get(request_id) + delivered = ledger._deliveries.get(request_id) + return dict( + latest=None + if fetch is None + else (fetch.start, fetch.end, fetch.state_end, fetch.marks_queued), + deliveries=None + if delivered is None + else ( + delivered.lowest_start, + [mask.tolist() for mask in delivered.delivered_block_masks], + delivered.state_end, + delivered.end, + sorted(delivered.partial_ends.items()), + ), + mixed=request_id in ledger._mixed_tail_caches, + ) + + +@pytest.fixture +def fast_returns(monkeypatch): + """The returns that took the fast path: whole deliveries recorded, and bounds in closed form.""" + taken = {"recorded": 0, "bounds": 0} + ledger = _fetch_ledger._FetchLedger + record, bounds = ledger._record_whole_delivery, ledger._whole_delivery_bounds + + def counted_record(self, *args, **kwargs): + taken["recorded"] += 1 + return record(self, *args, **kwargs) + + def counted_bounds(self, *args, **kwargs): + found = bounds(self, *args, **kwargs) + taken["bounds"] += found is not None + return found + + monkeypatch.setattr(ledger, "_record_whole_delivery", counted_record) + monkeypatch.setattr(ledger, "_whole_delivery_bounds", counted_bounds) + return taken + + +def arrivals_of(view, kind): + """``view``'s arrivals: the constant, masks writable or read-only, or all but each last row.""" + if kind == "constant": + return ALL_ARRIVED + masks = [np.ones(len(rows), dtype=bool) for rows in view.groups] + if kind == "all_but_last": + for mask in masks: + if len(mask): + mask[-1] = False + if kind == "read_only": + for mask in masks: + mask.flags.writeable = False + return masks + + +@gpu +@pytest.mark.parametrize("arrived", ["constant", "masks", "read_only", "all_but_last"]) +@pytest.mark.parametrize("hybrid", [False, True], ids=["paged", "hybrid"]) +def test_a_fast_return_gives_what_the_rules_give( + kit, real_manager, hybrid_manager, fast_returns, hybrid, arrived +): + """A receive returned fast and through the rules gives the same bounds and records.""" + outcomes = [] + for rules in (False, True): + with hybrid_manager() if hybrid else real_manager() as mgr: + lender = attach(mgr, rules=rules) + request = receiving(kit, mgr, 7, hybrid=hybrid) + lease, view = ready_write(lender, mgr, request) + bounds = lender.return_writes([(lease, arrivals_of(view, arrived), True)]) + outcomes.append( + (bounds, lender.poll_resume_bounds(request), records(lender, 7), lease._released) + ) + mgr._stream.synchronize() + mgr.free_resources(request) + assert outcomes[0] == outcomes[1] + (bounds,), again, _, released = outcomes[0] + assert bounds == again and released + if arrived != "all_but_last": + assert bounds.max_position == PROMPT + assert fast_returns == {"recorded": 1, "bounds": 1}, fast_returns + else: + assert fast_returns["recorded"] == 0 + + +def invalid_entries(write, view, request, read): + """Entries of a batch that is invalid in one place each, with what it raises.""" + good = (write, ALL_ARRIVED, True) + return [ + ("read lease", [good, (read, ALL_ARRIVED, True)], RuntimeError, "for write leases"), + ("duplicate", [good, good], RuntimeError, "returned already"), + ("mask count", [(write, [], True)], ValueError, "masks for"), + ( + "mask dtype", + [(write, [np.ones(len(r), dtype=np.int8) for r in view.groups], True)], + ValueError, + "must be bool", + ), + ( + "mask shape", + [(write, [np.ones(len(r) + 1, dtype=bool) for r in view.groups], True)], + ValueError, + "must be bool", + ), + ("not a triple", [good, (write, ALL_ARRIVED)], TypeError, "an entry is"), + ("reusable not a bool", [(write, ALL_ARRIVED, 1)], TypeError, "reusable must be a bool"), + ("request first", [(request, write, ALL_ARRIVED)], TypeError, "reusable must be a bool"), + ("not a lease", [(object(), ALL_ARRIVED, True)], ValueError, "not lent by this"), + ] + + +@gpu +@pytest.mark.parametrize("hybrid", [False, True], ids=["paged", "hybrid"]) +def test_a_batch_with_one_invalid_entry_raises_and_changes_nothing( + kit, real_manager, hybrid_manager, hybrid +): + """A batch or a call with one invalid entry raises before any entry is returned.""" + with hybrid_manager() if hybrid else real_manager() as mgr: + lender = attach(mgr) + request = receiving(kit, mgr, 7, hybrid=hybrid) + other = receiving(kit, mgr, 8, hybrid=hybrid) + write, view = ready_write(lender, mgr, request) + second, _ = ready_write(lender, mgr, other) + source = receiving(kit, mgr, 9, hybrid=hybrid) + read = lender.lend_read(source, 0, PROMPT, **kit.DEVICE_ALL) + assert read.poll() is not None, read.failure + before = {7: records(lender, 7), 8: records(lender, 8)} + version = lender._locks.version + for name, batch, error, match in invalid_entries(write, view, request, read): + alone = [batch] if len(batch) == 1 else [] + for entries in [[(second, ALL_ARRIVED, True), *batch], *alone]: + with pytest.raises(error, match=match): + lender.return_writes(entries) + assert not write._released and not write._marked, name + assert not second._released, name + assert {7: records(lender, 7), 8: records(lender, 8)} == before, name + assert lender._locks.version == version, name + lender.progress() + assert lender.poll_resume_bounds(request) is None, name + bounds = lender.return_writes([(write, ALL_ARRIVED, True), (second, ALL_ARRIVED, True)]) + assert all(isinstance(b, ResumeBounds) for b in bounds), bounds + for lease in (write, read): + with pytest.raises(RuntimeError): + lender.return_writes([(lease, ALL_ARRIVED, True)]) + read.release() + unpolled = receiving(kit, mgr, 10, hybrid=hybrid) + pending = lender.lend_write(unpolled, 0, PROMPT, **kit.DEVICE_ALL) + with pytest.raises(RuntimeError, match="before poll"): + lender.return_writes([(pending, ALL_ARRIVED, True)]) + pending.release() + with pytest.raises(RuntimeError, match="returned already"): + lender.return_writes([(pending, ALL_ARRIVED, True)]) + mgr._stream.synchronize() + for r in (request, other, source, unpolled): + mgr.free_resources(r) + + +@gpu +def test_a_partial_arrival_resumes_where_the_delivered_prefix_ends(kit, real_manager): + """Rows that did not arrive count nothing: the bounds stop at the first block missing.""" + with real_manager() as mgr: + lender = attach(mgr) + request = receiving(kit, mgr, 7) + lease, view = ready_write(lender, mgr, request) + masks = view.row_masks(True) + masks[0][2:] = False # blocks 2 and 3 of the only layer group missing + bounds = kit.return_write(lender, lease, masks) + assert bounds.max_position == 2 * TPB, bounds + assert lease._released and not lender._locks.is_locked(kit.kv(mgr, request)) + mgr._stream.synchronize() + mgr.free_resources(request) + + +@gpu +@pytest.mark.parametrize("hybrid", [False, True], ids=["paged", "hybrid"]) +def test_a_lease_of_a_retired_cache_counts_nothing_and_gives_no_bounds( + kit, real_manager, hybrid_manager, hybrid +): + """A lease of a freed request's old cache keeps its pages, counts no row and gives no bounds.""" + with hybrid_manager() if hybrid else real_manager() as mgr: + lender = attach(mgr) + moved = receiving(kit, mgr, 7, hybrid=hybrid) + lease, _ = ready_write(lender, mgr, moved) + old = kit.kv(mgr, moved) + mgr._stream.synchronize() + mgr.free_resources(moved) + assert not kit.closed(old), "the lock kept no cache open" + held = lender.held_page_counts() + assert HoldReason.LOCKED not in held.by_reason and held.deduct == held.unreclaimable + assert sum(held.deduct.values()) > 0, held + again = receiving(kit, mgr, 7, hybrid=hybrid) + assert kit.kv(mgr, again) is not old + fresh = lender.poll_resume_bounds(again) + assert not lease.return_resumes_at_end, "a retired cache's lease resumes at its end" + bounds = kit.return_write(lender, lease, ALL_ARRIVED) + assert bounds is None and kit.closed(old) and lease._released + assert lender.held_page_counts().unreclaimable == {} + assert lender.poll_resume_bounds(again) == fresh, "a retired cache's rows counted" + lease, _ = ready_write(lender, mgr, again) + bounds = kit.return_write(lender, lease, ALL_ARRIVED) + assert bounds is not None and bounds.max_position == PROMPT, bounds + mgr._stream.synchronize() + mgr.free_resources(again) + gone = receiving(kit, mgr, 8, hybrid=hybrid) + lease, _ = ready_write(lender, mgr, gone) + mgr._stream.synchronize() + mgr.free_resources(gone) + assert lender.return_writes([(lease, ALL_ARRIVED, True)]) == [None] + with pytest.raises(ValueError, match="has no KV cache"): + lender.poll_resume_bounds(gone) + + +@gpu +def test_staged_and_in_place_leases_return_in_one_batch(kit, real_manager): + """The in-place fetch of a batch settles in the call, the staged one after its copy's round.""" + with real_manager() as mgr: + lender = attach(mgr) + in_place = receiving(kit, mgr, 7) + staged = kit.admitted(mgr, 8, OTHER_TOKENS) + place, _ = ready_write(lender, mgr, in_place) + stage, view = ready_write(lender, mgr, staged, end=3 * TPB, memory_kind=HOST, rows=NAMED) + assert place.return_resumes_at_end and not stage.return_resumes_at_end + bounds = lender.return_writes([(place, ALL_ARRIVED, True), (stage, ALL_ARRIVED, True)]) + assert bounds[0] is not None and bounds[0].max_position == PROMPT, bounds + settled = bounds[1] + for _ in range(8): + if settled is not None: + break + mgr._stream.synchronize() + lender.progress() + settled = lender.poll_resume_bounds(staged) + assert settled is not None and settled.max_position >= 3 * TPB, settled + mgr._stream.synchronize() + for r in (in_place, staged): + mgr.free_resources(r) + + +@gpu +@pytest.mark.parametrize("request_count", [1, 3], ids=["single", "batch"]) +def test_an_error_in_one_return_raises_once_every_entry_is_returned( + kit, real_manager, monkeypatch, request_count +): + """An unlock error still returns all entries; the caller then drops fetched caches.""" + with real_manager() as mgr: + lender = attach(mgr) + requests = [receiving(kit, mgr, rid) for rid in range(7, 7 + request_count)] + leases = [ready_write(lender, mgr, r)[0] for r in requests] + broken = kit.kv(mgr, requests[0]) + retry = mgr._commit_held_back + + def failing(kv): + if kv is broken: + raise RuntimeError("the retried commit failed") + return retry(kv) + + monkeypatch.setattr(mgr, "_commit_held_back", failing) + # The unlock retries a commit only where the manager holds one back for the cache. + monkeypatch.setitem(mgr._held_commits, id(broken), (broken, requests[0], 0, False)) + with pytest.raises(RuntimeError, match="the retried commit failed"): + lender.return_writes([(lease, ALL_ARRIVED, True) for lease in leases]) + assert all(lease._released for lease in leases) + assert not lender._locks.holds_any() + mgr._stream.synchronize() + for r in requests: + mgr.free_resources(r) + + +class Interrupted(BaseException): + """An interrupt, as ``KeyboardInterrupt`` is, raised where a call runs.""" + + +@gpu +@pytest.mark.parametrize("error", [MemoryError, Interrupted]) +def test_a_return_whose_count_raises_still_returns_its_lease(kit, real_manager, monkeypatch, error): + """A copy error still returns every lease; the caller drops caches without querying bounds.""" + with real_manager() as mgr: + lender = attach(mgr) + in_place = receiving(kit, mgr, 7) + staged = kit.admitted(mgr, 8, OTHER_TOKENS) + stage, _ = ready_write(lender, mgr, staged, end=3 * TPB, memory_kind=HOST, rows=NAMED) + place, _ = ready_write(lender, mgr, in_place) + queue, tried = lender._queue_copy, [] + + def failing(segments, to_staging, after=None): + if not to_staging: + tried.append(len(segments)) + raise error("the copy into the pages failed") + return queue(segments, to_staging, after=after) + + monkeypatch.setattr(lender, "_queue_copy", failing) + with pytest.raises(error, match="into the pages"): + lender.return_writes([(stage, ALL_ARRIVED, True), (place, ALL_ARRIVED, True)]) + assert tried and tried[0] > 0, "inconclusive: no copy into the pages was queued" + assert stage._released and place._released, "a return left its lease open" + assert not lender._locks.holds_any() + mgr._stream.synchronize() + for r in (in_place, staged): + mgr.free_resources(r) + + +@gpu +def test_a_release_whose_count_raises_still_ends_its_lock(kit, real_manager, monkeypatch): + """A release whose count of arrivals raises still ends the lease's lock before raising.""" + with real_manager() as mgr: + lender = attach(mgr) + request = receiving(kit, mgr, 7) + lease, _ = ready_write(lender, mgr, request) + + def failing(*args, **kwargs): + raise MemoryError("counting the arrivals failed") + + monkeypatch.setattr(lender, "_apply_in_place_marks", failing) + assert lender._locks.holds_any() + with pytest.raises(MemoryError, match="counting the arrivals"): + lease.release() + assert not lender._locks.holds_any(), "the release kept the lease's lock" + mgr._stream.synchronize() + mgr.free_resources(request) + + +PAGE_READS = ("pages", "page_list", "locked_pages", "state_slot") + + +@pytest.fixture +def page_reads(monkeypatch): + """Calls of the manager facade's page and state-slot reads, by name, while counting is on.""" + counts = {name: 0 for name in PAGE_READS} + counting = [False] + for name in PAGE_READS: + read = getattr(_manager, name) + + def counted(*args, _read=read, _name=name, **kwargs): + if counting[0]: + counts[_name] += 1 + return _read(*args, **kwargs) + + monkeypatch.setattr(_manager, name, counted) + return counts, counting + + +@gpu +@pytest.mark.parametrize("hybrid", [False, True], ids=["paged", "hybrid"]) +def test_a_fast_write_lend_and_its_return_ask_the_cache_for_no_page( + kit, real_manager, hybrid_manager, page_reads, fast_returns, hybrid +): + """A fast write lend and its return of every row read no page and no state slot of the cache.""" + counts, counting = page_reads + with hybrid_manager() if hybrid else real_manager() as mgr: + lender = attach(mgr) + recurrent = sum(lender._layout.recurrent) + request = receiving(kit, mgr, 7, hybrid=hybrid) + counting[0] = True + lease = lender.lend_write(request, 0, PROMPT, **kit.DEVICE_ALL) + assert lease._block_entries.spans is not None, "inconclusive: the lend was not fast" + assert sum(counts.values()) == 0, counts + assert recurrent == hybrid, "inconclusive: no recurrent state where expected" + for name in counts: + counts[name] = 0 + assert lease.poll() is not None + lender.return_writes([(lease, ALL_ARRIVED, True)]) + lender.progress() + counting[0] = False + assert sum(counts.values()) == 0, counts + assert fast_returns == {"recorded": 0, "bounds": 1}, fast_returns + mgr._stream.synchronize() + mgr.free_resources(request) + + +@gpu +@skip_pre_blackwell +@pytest.mark.parametrize("tokens", [128 + 5, 3 * 128, 7 * 128 + 1]) +def test_deepseek_v4_windowed_groups_take_the_fast_paths_as_the_rules_lend_and_return( + kit, deepseek_v4_manager, fast_returns, monkeypatch, tokens +): + """DeepSeek-V4's windowed layer groups lend and return fast as the rules do.""" + taken = {"read": 0, "write": 0} + for name in ("_read_every_row_fast", "_write_every_row_fast"): + fast = getattr(_lender._ManagerLender, name) + + def counted(self, *args, _fast=fast, _kind=name.split("_")[1], **kwargs): + lease = _fast(self, *args, **kwargs) + taken[_kind] += lease is not None + return lease + + monkeypatch.setattr(_lender._ManagerLender, name, counted) + prompt = list(range(20000, 20000 + tokens)) + outcomes = [] + for rules in (False, True): + with deepseek_v4_manager() as mgr: + lender = attach(mgr, rules=rules, staging_tokens=128) + assert any(w is not None for w in lender._layout.windows), "inconclusive: no window" + assert lender._layout.stale_ranges_by_formula + source = kit.published(mgr, 9, prompt) + read = lender.lend_read( + source, 0, tokens, **{**kit.DEVICE_ALL, "ready_event": kit.stream_event(mgr)} + ) + read_view = kit.ready_view(read, mgr) + target = receiving(kit, mgr, 7, prompt) + write, view = ready_write(lender, mgr, target) + key = [ + [(rows.layer_group_index, rows.page_ordinals.tolist()) for rows in v.groups] + for v in (read_view, view) + ] + locks = [ + {lg: b.tolist() for lg, b in lease._lock.blocks.items()} for lease in (read, write) + ] + bounds = lender.return_writes([(write, ALL_ARRIVED, True)]) + outcomes.append((key, locks, bounds, records(lender, 7))) + read.release() + mgr._stream.synchronize() + for r in (source, target): + mgr.free_resources(r) + assert outcomes[0] == outcomes[1] + assert taken == {"read": 1, "write": 1}, taken + assert fast_returns == {"recorded": 1, "bounds": 1}, fast_returns + assert outcomes[0][2][0].max_position == tokens + + +@gpu +def test_a_staged_read_becomes_ready_only_after_a_round_of_progress(kit, real_manager): + """A staged read gives its view once a round of progress saw its copy done, never by polling.""" + with real_manager() as mgr: + lender = attach(mgr) + source = kit.published(mgr, 9, TOKENS) + with ExitStack() as held: + held.enter_context(kit.held_stream(mgr._stream, strict=False)) + lease = lender.lend_read(source, 0, 3 * TPB, **kit.HOST_NAMED) + assert lease.failure is None, lease.failure + assert lease.poll() is None + mgr._stream.synchronize() + for _ in range(3): + assert lease.poll() is None, "a poll advanced the lender" + lender.progress() + assert lease.poll() is not None + lease.release() + mgr._stream.synchronize() + mgr.free_resources(source) + + +@gpu +def test_a_new_lender_counts_the_pages_its_fast_lend_keeps_and_locks(kit, real_manager): + """A fast write lend's pages count as locked, and all its cache's once a free keeps it open.""" + with real_manager() as mgr: + lender = attach(mgr) + request = receiving(kit, mgr, 7) + lease, view = ready_write(lender, mgr, request) + assert lender._locks.holds_any(), "inconclusive: the lend took no lock" + kv = kit.kv(mgr, request) + held = {} + for lg, group in enumerate(kit.pool_group_of(mgr)): + held[group] = held.get(group, 0) + sum(page >= 0 for page in kit.pages(kv, lg)) + locked = dict(lender.held_page_counts().by_reason.get(HoldReason.LOCKED, {})) + assert sum(locked.values()) == sum(len(rows) for rows in view.groups) > 0, locked + mgr._stream.synchronize() + mgr.free_resources(request) + assert not kit.closed(kv), "the free closed a cache whose pages are lent" + assert dict(lender.held_page_counts().deduct) == {g: n for g, n in held.items() if n} + lender.return_writes([(lease, ALL_ARRIVED, True)]) + assert kit.closed(kv) and dict(lender.held_page_counts().deduct) == {} + + +@gpu +@pytest.mark.parametrize("hybrid", [False, True], ids=["paged", "hybrid"]) +def test_a_whole_return_holds_its_record_until_a_reader_needs_it( + kit, real_manager, hybrid_manager, fast_returns, monkeypatch, hybrid +): + """A fast whole return holds its record: rounds keep it, a reader writes it, a free drops it.""" + with hybrid_manager() if hybrid else real_manager() as mgr: + lender = attach(mgr) + ledger = lender._ledger + request = receiving(kit, mgr, 7, hybrid=hybrid) + lease, _ = ready_in_place(kit, lender, request) + bounds = kit.return_write(lender, lease, ALL_ARRIVED) + assert fast_returns == {"recorded": 0, "bounds": 1}, fast_returns + assert 7 in ledger._pending and 7 not in ledger._delivered, "the return wrote its record" + assert ledger._fetch_settled(ledger._latest_fetches[7]), "the fetch settled later" + assert ledger._has_records(7) + rounds = lender._progress_round + with monkeypatch.context() as patched: + patched.setattr(ledger, "_write_pending", lambda: pytest.fail("a round wrote a record")) + lender.progress() + lender.progress() + assert lender._progress_round == rounds + 2 and 7 in ledger._pending + assert lender.poll_resume_bounds(request) == bounds + assert not ledger._pending and 7 in ledger._delivered, "the reader left the record held" + assert fast_returns == {"recorded": 1, "bounds": 1}, fast_returns + other = receiving(kit, mgr, 8, hybrid=hybrid) + lease, _ = ready_in_place(kit, lender, other) + lender.return_writes([(lease, ALL_ARRIVED, True)]) + assert 8 in ledger._pending, "inconclusive: the return held no record" + mgr._stream.synchronize() + mgr.free_resources(other) + assert not ledger._pending and not ledger._has_records(8), "the free kept a held record" + assert fast_returns == {"recorded": 1, "bounds": 2}, fast_returns + mgr.free_resources(request) + assert not ledger._has_records(7) + + +@gpu +def test_a_held_record_names_its_request_without_its_fetch(kit, real_manager): + """A held whole delivery names its request and cache, and goes with the request's deliveries.""" + with real_manager() as mgr: + lender = attach(mgr) + ledger = lender._ledger + request = receiving(kit, mgr, 7) + lease, _ = ready_in_place(kit, lender, request) + lender.return_writes([(lease, ALL_ARRIVED, True)]) + # The held record alone left to name the request. + fetch = ledger._latest_fetches.pop(7) + ledger._kept_prefixes.pop(7, None) + assert 7 in ledger._pending and 7 not in ledger._delivered + assert 7 not in ledger._mixed_tail_caches + assert ledger._has_records(7), "a held record does not name its request" + assert ledger._delivered_kv(7) is fetch.kv + ledger._drop_deliveries(7) + assert 7 not in ledger._deliveries and ledger._delivered_kv(7) is None + mgr._stream.synchronize() + mgr.free_resources(request) + assert not ledger._has_records(7) and not ledger._pending + + +@gpu +@pytest.mark.parametrize( + "reader", ["second_fetch", "masks", "partial", "staged_none", "free", "shrink", "read_in_place"] +) +def test_what_reads_a_held_record_sees_it_written(kit, real_manager, reader): + """Every reader of a held record sees what the rules, which write it at once, would give.""" + outcomes = [] + for rules in (False, True): + with real_manager() as mgr: + lender = attach(mgr, rules=rules) + request = receiving(kit, mgr, 7) + first, _ = ready_in_place(kit, lender, request) + outcome = [lender.return_writes([(first, ALL_ARRIVED, True)])] + if reader in ("second_fetch", "masks", "partial"): + second, view = ready_in_place(kit, lender, request, start=TPB) + arrived = ALL_ARRIVED if reader == "second_fetch" else view.row_masks(True) + if reader == "partial": + for mask in arrived: + mask[-1:] = False + outcome.append(lender.return_writes([(second, arrived, True)])) + elif reader == "staged_none": + staged, view = ready_write( + lender, mgr, request, end=3 * TPB, memory_kind=HOST, rows=NAMED + ) + outcome.append(lender.return_writes([(staged, view.row_masks(False), True)])) + elif reader == "shrink": + # The manager's report that the cache may have shrunk in place. + lender._on_shrink(7, kit.kv(mgr, request)) + elif reader == "read_in_place": + read = lender.lend_read(request, 0, 3 * TPB, **kit.DEVICE_ALL) + outcome.append(read.failure) + read.release() + mgr._stream.synchronize() + if reader == "free": + mgr.free_resources(request) + outcome.append(lender._ledger._has_records(7)) + else: + outcome.append(lender.poll_resume_bounds(request)) + outcome.append(records(lender, 7)) + outcomes.append(outcome) + if reader != "free": + mgr.free_resources(request) + assert outcomes[0] == outcomes[1] + if reader == "free": + assert outcomes[0][1] is False, outcomes[0] + + +@gpu +def test_a_shrink_hook_that_raises_drops_a_held_record(kit, real_manager, monkeypatch): + """A shrink report whose own work raises drops the request's deliveries, a held record too.""" + with real_manager() as mgr: + lender = attach(mgr) + request = receiving(kit, mgr, 7) + lease, _ = ready_in_place(kit, lender, request) + lender.return_writes([(lease, ALL_ARRIVED, True)]) + assert 7 in lender._ledger._pending, "inconclusive: the return held no record" + raised = [] + + def num_blocks(cache): + raised.append(cache) + raise MemoryError("the block count failed") + + monkeypatch.setattr(_lender._manager, "num_blocks", num_blocks) + lender._on_shrink(7, kit.kv(mgr, request)) + monkeypatch.undo() + assert raised, "inconclusive: the hook's work did not raise" + assert 7 not in lender._ledger._deliveries, "the rows the shrink freed stay delivered" + mgr._stream.synchronize() + mgr.free_resources(request) + + +@gpu +@pytest.mark.parametrize("hybrid", [False, True], ids=["paged", "hybrid"]) +def test_a_write_that_resumes_at_its_end_returns_bounds_holding_it( + kit, real_manager, hybrid_manager, monkeypatch, hybrid +): + """A whole return of a write that says it resumes at its end gives bounds holding that end.""" + calls = [] + for name, value in list(vars(_manager).items()): + if not isinstance(value, type) and getattr(value, "__module__", None) == _manager.__name__: + + def counted(*args, _value=value, _name=name, **kwargs): + calls.append(_name) + return _value(*args, **kwargs) + + monkeypatch.setattr(_manager, name, counted) + ends = [PROMPT] if hybrid else [PROMPT, 2 * TPB, 3 * TPB] + said = [] + with hybrid_manager() if hybrid else real_manager() as mgr: + lender = attach(mgr) + rid = 20 + for admit in ("received", "admitted"): + for start in (0, TPB): + for end in ends: + rid += 1 + if admit == "received": + request = receiving(kit, mgr, rid, hybrid=hybrid) + else: + request = kit.admitted(mgr, rid, TOKENS) + assert mgr.resize_context(request, request.context_remaining_length) + lease = lender.lend_write(request, start, end, **kit.DEVICE_ALL) + assert not lease.return_resumes_at_end, "a lease before its view says so" + if lease.poll() is not None: + before = len(calls) + at_end = lease.return_resumes_at_end + assert len(calls) == before, calls[before:] + bounds = kit.return_write(lender, lease, ALL_ARRIVED) + if at_end: + assert bounds is not None, (admit, start, end) + assert bounds.min_position <= end <= bounds.max_position, ( + admit, + start, + end, + bounds, + ) + assert not lease.return_resumes_at_end, "a returned lease says so" + said.append(at_end) + else: + lease.release() + mgr._stream.synchronize() + mgr.free_resources(request) + assert said.count(True) >= 2 and said.count(False) >= (0 if hybrid else 2), said + + +@gpu +@pytest.mark.parametrize("event", ["shrink", "free", "reset", "stop", "shutdown"]) +def test_a_lease_stops_resuming_at_its_end_once_its_cache_can_move( + kit, real_manager, monkeypatch, event +): + """A shrink, a free, a reuse reset, the lender's stop or the shutdown ends it for good.""" + with real_manager() as mgr: + lender = attach(mgr) + request = receiving(kit, mgr, 7) + lease, _ = ready_in_place(kit, lender, request) + assert lease.return_resumes_at_end, "inconclusive: the lease does not resume at its end" + mgr._stream.synchronize() + if event == "shrink": + lender._on_shrink(7, kit.kv(mgr, request)) + elif event == "free": + mgr.free_resources(request) + elif event == "reset": + # The manager's report; its runtime resets nothing while a cache is open. + lender._on_reset() + elif event == "stop": + + def failing(request_id): + raise MemoryError("dropping the records failed") + + monkeypatch.setattr(lender._ledger, "_drop_records", failing) + other = receiving(kit, mgr, 8) + mgr.free_resources(other) + assert lender._stop_reason is not None, "inconclusive: the lender did not stop" + else: + mgr.shutdown() + assert not lease.return_resumes_at_end, event + lease.release() + if event not in ("free", "shutdown"): + mgr.free_resources(request) + + +# -- the ledger's helpers ------------------------------------------------------------------------ + +LEDGER_TPB = 4 +RUN = [False] * 2 + [True] * 6 # blocks 2 to 7 of 8 + + +def ledger(recurrent=(False, False)): + layout = SimpleNamespace( + tokens_per_block=LEDGER_TPB, num_layer_groups=len(recurrent), recurrent=recurrent, cp=(0, 1) + ) + return _fetch_ledger._FetchLedger(layout, lambda copy: True) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("recurrent", "computed", "known"), + [ + ((False,), (7,), 7), + ((False, False, False), (5, 9, 2), 2), + ((False, True, False), (5, 1, 9), 5), + ((True, False, False), (0, 9, 4), 4), + ((True, True), (3, 8), 8), + ], +) +def test_the_known_prefix_is_the_least_a_group_without_recurrent_state_computed( + recurrent, computed, known +): + """The least any group but those of recurrent state computed; with state alone, the most.""" + assert ledger(recurrent)._known(computed) == known + + +def request_with_runs(runs, length): + """A request whose multimodal tokens form ``runs``; ``None`` for one without multimodal data.""" + if runs is None: + return SimpleNamespace(py_multimodal_data=None) + mask = np.zeros(length, dtype=np.int64) + for begin, end in runs: + mask[begin:end] = 1 + data = {"mm_bidirectional_blocks": True, "multimodal_embed_mask_cumsum": np.cumsum(mask)} + return SimpleNamespace(py_multimodal_data=data) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("runs", "usable", "clipped", "floor"), + [ + (None, 40, 40, 0), + ([], 40, 40, 0), + ([(10, 20)], 40, 40, 20), + ([(10, 20)], 15, 10, 0), + ([(2, 5), (10, 20)], 15, 10, 5), + ([(2, 5), (10, 20)], 20, 20, 20), + ([(2, 5), (30, 50)], 40, 30, 5), + ], +) +def test_an_end_is_clipped_to_the_multimodal_runs_with_its_floor(runs, usable, clipped, floor): + """An end inside a run drops to its start; the floor is the end of the last run below it.""" + runs = _manager.bidirectional_runs(request_with_runs(runs, 60)) + assert ledger((False,))._clip_to_bidirectional_runs(runs, usable) == (clipped, floor) + + +def whole(led, request_id, kv): + """The record of a whole delivery of blocks 2 to 7 of both layer groups into ``kv``.""" + fetch = _fetch_ledger._FetchRecord(2 * LEDGER_TPB, kv, end=8 * LEDGER_TPB) + led._record_whole_delivery(request_id, fetch, [(2, 8), (2, 8)]) + return led._deliveries[request_id] + + +def masks_of(delivered): + return [mask.tolist() for mask in delivered.delivered_block_masks] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("change", ["shrink", "invalidate", "deliver"]) +def test_a_change_to_one_record_of_a_run_leaves_the_others_as_they_were(change): + """Records of one run share a read-only mask; a change to one leaves the other and the mask.""" + led = ledger() + kv1, kv2 = SimpleNamespace(num_blocks=8), SimpleNamespace(num_blocks=8) + one, two = whole(led, 1, kv1), whole(led, 2, kv2) + shared = one.delivered_block_masks[0] + assert two.delivered_block_masks[0] is shared, "inconclusive: the records share no mask" + assert not shared.flags.writeable + if change == "shrink": + kv1.num_blocks = 5 + led._invalidate_delivered_blocks_past_capacity(one, kv1) + expected = [RUN[:5] + [False] * 3] * 2 + elif change == "invalidate": + ordinals = np.arange(2, 8) + lease = SimpleNamespace( + _request_id=1, + _kv=kv1, + _block_entries=SimpleNamespace( + layer_group_indices=[0, 1], block_ordinals=[ordinals, ordinals] + ), + ) + copied = [ordinals == 2, np.zeros(6, dtype=bool)] + led._invalidate_delivered_blocks(lease, copied) + expected = [[False] * 3 + [True] * 5, RUN] + else: + # A later fetch into the same cache delivers blocks 0 and 1 of the first layer group. + fetch = _fetch_ledger._FetchRecord(0, kv1, end=2 * LEDGER_TPB) + entries = SimpleNamespace( + layer_group_indices=[0, 1], block_ordinals=[np.array([0, 1]), np.array([0, 1])] + ) + copied = [np.array([True, True]), np.array([False, False])] + led._record_delivered_blocks(1, fetch, entries, copied) + expected = [[True] * 8, RUN] + assert masks_of(one) == expected + assert masks_of(two) == [RUN, RUN], "the change reached another record of the run" + assert shared.tolist() == RUN and not shared.flags.writeable + + +@pytest.mark.cpu_only +def test_the_shared_masks_hold_a_bounded_number_of_bytes(monkeypatch): + """Masks of ever longer runs are the runs' blocks, and never hold more bytes than the bound.""" + monkeypatch.setattr(_fetch_ledger, "_run_masks", {}) + monkeypatch.setattr(_fetch_ledger, "_run_mask_bytes", 0) + monkeypatch.setattr(_fetch_ledger, "_RUN_MASK_BYTES", 64) + emptied = 0 + for last in range(1, 40): + held = len(_fetch_ledger._run_masks) + assert _fetch_ledger._run_mask(1, last).tolist() == [False] + [True] * (last - 1) + emptied += len(_fetch_ledger._run_masks) <= held + kept = sum(len(mask) for mask in _fetch_ledger._run_masks.values()) + assert _fetch_ledger._run_mask_bytes == kept <= 64 + assert emptied >= 5, emptied + + +# -- stand-in managers: kept and read windows, claimed prefixes, pipeline stages ------------------ + +LONG_PROMPT = 5 * TPB +RID = 7 +KEPT, READ = 72, 64 # a window kept eight tokens wider than the one attention reads +GROUPS = 3 +# CSA2 with block reuse on: the global layer group REQUIRED, the encoder's OPTIONAL and not +# selected, the decoder's sliding window PRIVATE. +GLOBAL_ONLY = (True, False, False) + + +class Event: + """A stand-in CUDA event, complete once ``done`` is set; waiting on it fails the test.""" + + def __init__(self, done=True): + self.done = done + + def query(self): + return self.done + + def synchronize(self): + raise AssertionError("a lease waited on its event") + + +class Manager: + """A stand-in manager: a layer group per window, each on a pool group of its own, the runtime's + stale ranges for the windows, no block reuse; its caches by request.""" + + tokens_per_block = TPB + enable_block_reuse = False + _can_publish_block_reuse = False + kv_connector_manager = None + _on_allocation_stream = frozenset() + + def __init__(self, windows=(None,) * GROUPS, sinks=None): + sinks = sinks or (0,) * len(windows) + cycles = [AttnLifeCycle.make(w, s * TPB, TPB) for w, s in zip(windows, sinks)] + self.impl = SimpleNamespace( + get_life_cycle_pool_group_indices=lambda: list(range(len(windows))) + ) + self.kv_cache_map = {} + + def stale_block_range(lg, history): + stale = cycles[lg].get_stale_range(history, TPB) + return stale.beg, stale.end + + self._stale_block_range = stale_block_range + + def _commit_held_back(self, kv): + """No commit is held back here.""" + + +def layout(windows=(None,) * GROUPS, read_windows=(), sinks=None, named=None): + count = len(windows) + return SimpleNamespace( + tokens_per_block=TPB, + windows=tuple(windows), + read_windows=tuple(read_windows), + sink_blocks=tuple(sinks or (0,) * count), + recurrent=(False,) * count, + pool_groups=tuple(range(count)), + pool_group_of=tuple(range(count)), + num_layer_groups=count, + named_groups=tuple(named or (True,) * count), + cp=(0, 1), + device_pools={g: (DevicePool(g, 0, (g + 1) << 30, 256, 16),) for g in range(count)}, + ) + + +def lender_over(manager, lay, layers=None, staged=False): + """A lender over ``manager`` and ``lay``: the device regions an attach lists, after a staging + region of four slots per pool group with ``staged``.""" + count = lay.num_layer_groups + layers = layers or [(lg,) for lg in range(count)] + ident = Identity(b"scope", bytes(range(32)), layers, [(1, 0)] * count) + ref = weakref.ref(manager) + regions = [] + for g in lay.pool_groups if staged else (): + groups = tuple(lg for lg, pg in enumerate(lay.pool_group_of) if pg == g) + regions.append( + MemoryRegion(f"staging{g}", 4096 * (g + 1), 4096, 1024, 4, groups, MemoryKind.HOST) + ) + regions += _lender._device_regions(manager, lay, ident) + slots = SlotQueue({g: 4 for g in lay.pool_groups} if staged else {}) + return _lender._ManagerLender(ref, lay, ident, tuple(regions), slots, ref) + + +def cache(committed=0, complete=(), endpoint=None, history=LONG_PROMPT, blocks=5): + """A stand-in active cache on pages ``10, 11, ...``, ``committed`` tokens committed, its claim + ending at ``endpoint`` and complete in the layer groups ``complete`` flags.""" + pages = list(range(10, 10 + blocks)) + claim = committed if endpoint is None else endpoint + return SimpleNamespace( + Status=SimpleNamespace(OPEN=0, CLOSED=1), + status=0, + num_committed_tokens=committed, + reuse_status=[SimpleNamespace(complete=c, endpoint=claim) for c in complete], + history_length=history, + is_active=True, + num_blocks=blocks, + tokens_per_block=TPB, + enable_swa_scratch_reuse=False, + get_base_page_indices=lambda lg: pages, + get_aggregated_page_indices=lambda lg, valid_only=False: pages, + ) + + +def request(prompt=LONG_PROMPT): + return SimpleNamespace( + py_request_id=RID, + prompt_len=prompt, + py_multimodal_data=None, + py_served_prefix_end=0, + context_current_position=0, + context_remaining_length=prompt, + ) + + +def return_write(lender, req, lease, arrived): + """The bounds ``return_writes`` gives for ``lease``, which ``req``'s query gives too.""" + (bounds,) = lender.return_writes([(lease, arrived, True)]) + assert bounds == lender.poll_resume_bounds(req), bounds + return bounds + + +def lent(lease): + """The view of ``lease`` and its block ordinals per layer group.""" + view = lease.poll() + assert view is not None, lease.failure + return view, {run.layer_group_index: run.page_ordinals.tolist() for run in view.groups} + + +def unnamed_keys(lender, monkeypatch): + """``lender`` lending by name keys of zeros, its request's input needing no name.""" + monkeypatch.setattr(lender, "_unnamed_input_refusal", lambda request: None) + monkeypatch.setattr( + lender, "_reuse_keys_for", lambda m, r, k, o: [np.zeros((len(v), 32), np.uint8) for v in o] + ) + + +@pytest.mark.cpu_only +def test_the_read_window_leaves_out_what_the_runtime_s_stale_range_for_it_does(): + """The read window's stale range is the runtime's own for a window of its length.""" + rng = random.Random(0) + for _ in range(3000): + read = rng.randint(1, 300) + kept = read + rng.choice([0, 1, 2, 3, 8, 40]) + sinks = rng.choice([0, 0, 1, 2]) + history = rng.randint(0, 1200) + manager = Manager(windows=(kept,), sinks=(sinks,)) + lay = layout((kept,), (read,), (sinks,)) + cycle = AttnLifeCycle.make(read, sinks * TPB, TPB) + stale = cycle.get_stale_range(history, TPB) + got = _ranges._read_stale_range(manager, lay, 0, history) + if stale.end > stale.beg or got[1] > got[0]: + assert got == (stale.beg, stale.end), (kept, read, sinks, history) + block = rng.randint(sinks, 40) + last = _ranges._last_history_reading(manager, lay, 0, block) + if last >= block * TPB: + assert cycle.get_stale_range(last, TPB).end <= block, (kept, read, sinks, block, last) + assert cycle.get_stale_range(last + 1, TPB).end > block, (kept, read, sinks, block, last) + + +def fetch_window(kit, delivered, history=LONG_PROMPT, prompt=LONG_PROMPT, world=None): + """Blocks a fetch of five lends under the window pair, and its bounds with ``delivered`` in.""" + manager = Manager(windows=(KEPT,)) + lender = lender_over(manager, layout((KEPT,), (READ,))) + manager.kv_cache_map[RID] = cache(history=history) + req = request(prompt) + lease = lender.lend_write(req, 0, 5 * TPB, **kit.DEVICE_ALL) + view = lease.poll() + assert view is not None, lease.failure + (run,) = view.groups + bounds = return_write(lender, req, lease, [np.isin(run.page_ordinals, list(delivered))]) + if world is not None: + world.update(manager=manager, lender=lender, kv=manager.kv_cache_map[RID]) + return run.page_ordinals.tolist(), tuple(bounds) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "delivered", + [(3, 4), (2, 3, 4), (2, 4)], + ids=["margin_missing", "every_row", "read_row_missing"], +) +def test_a_margin_row_that_did_not_arrive_lowers_no_bound_and_keeps_the_cache_out_of_reuse( + kit, delivered +): + """Block 2 of the kept window is a margin row the read window leaves out at the prompt's end.""" + world = {} + lent_blocks, bounds = fetch_window(kit, delivered, world=world) + assert lent_blocks == [2, 3, 4], f"the fetch lends what the kept window holds: {lent_blocks}" + if delivered == (2, 4): + assert bounds[0] < bounds[1] and bounds[0] < LONG_PROMPT, bounds + return + assert bounds == (LONG_PROMPT, LONG_PROMPT), bounds + manager = world["manager"] + manager._sharing = world["lender"] + gate = _lender_hooks.commit_gate(manager, world["kv"], None, LONG_PROMPT) + assert gate == (_lender_hooks.REFUSE if delivered == (3, 4) else _lender_hooks.COMMIT), gate + + +@pytest.mark.cpu_only +def test_a_resume_below_the_end_starts_where_the_read_window_leaves_the_margin_behind( + kit, monkeypatch +): + """A resume below the end starts past 158, the last history reading the undelivered block 2.""" + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: False) + lent_blocks, bounds = fetch_window(kit, (3, 4), history=0, prompt=LONG_PROMPT + 1) + assert lent_blocks == [2, 3, 4], lent_blocks + assert bounds == (LONG_PROMPT, LONG_PROMPT - 1), bounds + + +@pytest.mark.cpu_only +def test_a_group_counts_a_claimed_prefix_as_native_s_transfer_does_and_its_own_commits_always(): + """Each group's committed blocks of a claim are what native's transfer counts as cached.""" + from tensorrt_llm._torch.disaggregation.resource.cache_reuse import _CacheReuseAdapterV2 + from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup + + manager = Manager() + manager.enable_block_reuse = True + adapter = _CacheReuseAdapterV2(manager) + groups = [AttentionLayerGroup(pool_group_idx=g, kv_head_num_per_rank=1) for g in range(GROUPS)] + req = request() + rng = random.Random(0) + for _ in range(2000): + endpoint = rng.choice([0, rng.randint(1, 8) * TPB, rng.randint(1, 8 * TPB)]) + committed = rng.choice([0, endpoint, endpoint // TPB * TPB, rng.randint(0, 9 * TPB)]) + complete = [endpoint > 0 and rng.random() < 0.6 for _ in range(GROUPS)] + kv = cache(committed, complete, endpoint) + got = _manager.committed_by_group(kv, GROUPS) + if committed <= endpoint: + manager.kv_cache_map[RID] = kv + want = adapter.get_cached_token_count_per_layer_group(req, groups) + assert [t // TPB * TPB for t in got] == want, (committed, endpoint, complete, got) + else: + assert got == (committed,) * GROUPS, (committed, endpoint, complete, got) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("history", [LONG_PROMPT, 2 * TPB], ids=["generation_init", "context"]) +@pytest.mark.parametrize("missing", [None, 1, 0], ids=["all_delivered", "private_hole", "global"]) +def test_each_layer_group_is_fetched_from_its_own_committed_blocks(kit, missing, history): + """A claim complete in the global group alone: it is fetched from block 2, the others from 0.""" + manager = Manager() + lender = lender_over(manager, layout()) + kv = manager.kv_cache_map[RID] = cache(2 * TPB, GLOBAL_ONLY, history=history) + lease = lender.lend_write(request(), 0, LONG_PROMPT, **kit.DEVICE_ALL) + view, blocks = lent(lease) + assert blocks == {0: [2, 3, 4], 1: [0, 1, 2, 3, 4], 2: [0, 1, 2, 3, 4]}, blocks + assert {lg: b.tolist() for lg, b in lease._lock.blocks.items()} == blocks + masks = view.row_masks(True) + if missing is not None: + masks[missing][0] = False + bounds = return_write(lender, request(), lease, masks) + if missing == 1: + assert bounds.max_position == 0, f"a group without block 0 resumes: {bounds}" + elif missing == 0: + assert bounds.max_position == 2 * TPB, f"block 2 of the global group: {bounds}" + else: + assert tuple(bounds) == (LONG_PROMPT, history), bounds + assert kv.history_length == history + + +@pytest.mark.cpu_only +def test_blocks_the_cache_committed_past_its_claim_are_fetched_in_no_layer_group(kit): + """Blocks committed past a claim the global group alone took whole are held in every group.""" + manager = Manager() + lender = lender_over(manager, layout()) + manager.kv_cache_map[RID] = cache(3 * TPB, GLOBAL_ONLY, endpoint=2 * TPB) + lease = lender.lend_write(request(), 0, LONG_PROMPT, **kit.DEVICE_ALL) + _, blocks = lent(lease) + assert blocks == {lg: [3, 4] for lg in range(GROUPS)}, blocks + lease.release() + + +@pytest.mark.cpu_only +def test_a_write_lock_keeps_off_each_layer_group_s_own_committed_blocks(): + manager = Manager() + lender = lender_over(manager, layout()) + kv = cache(2 * TPB, GLOBAL_ONLY) + refusal = lender._locks.refusal(kv, {0: [1], 1: [0]}, LockMode.WRITE) + assert refusal is not None and "layer group 0" in refusal, refusal + assert lender._locks.refusal(kv, {0: [2], 1: [0], 2: [1]}, LockMode.WRITE) is None + + +@pytest.mark.cpu_only +def test_a_named_read_lends_no_block_a_layer_group_does_not_hold(kit, monkeypatch): + """A read by name of the two claimed blocks lends them in the global group alone.""" + manager = Manager() + lender = lender_over(manager, layout()) + manager.kv_cache_map[RID] = cache(2 * TPB, GLOBAL_ONLY) + unnamed_keys(lender, monkeypatch) + lease = lender.lend_read(request(), 0, 2 * TPB, **kit.DEVICE_NAMED) + _, blocks = lent(lease) + assert blocks == {0: [0, 1], 1: [], 2: []}, blocks + lease.release() + + +@pytest.mark.cpu_only +def test_an_undelivered_block_taints_only_where_a_group_s_committed_tokens_end_inside_it(kit): + """A block missing where a group committed none of its tokens only keeps the resume below it.""" + manager = Manager() + lender = lender_over(manager, layout()) + manager.kv_cache_map[RID] = cache(2 * TPB + 5, GLOBAL_ONLY) + lease = lender.lend_write(request(), 0, LONG_PROMPT, **kit.DEVICE_ALL) + view, blocks = lent(lease) + assert blocks == {0: [2, 3, 4], 1: [0, 1, 2, 3, 4], 2: [0, 1, 2, 3, 4]}, blocks + masks = view.row_masks(True) + masks[1][2] = False + assert return_write(lender, request(), lease, masks).max_position == 2 * TPB + + +@pytest.mark.cpu_only +def test_the_floor_reads_each_window_from_the_tokens_its_layer_group_computed(kit, monkeypatch): + """A window whose group committed the blocks behind it lets a resume start at the history.""" + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: False) + manager = Manager(windows=(64, None, None)) + lender = lender_over(manager, layout((64, None, None))) + manager.kv_cache_map[RID] = cache(3 * TPB, (True, False, False), history=3 * TPB) + lease = lender.lend_write(request(), 0, LONG_PROMPT, **kit.DEVICE_ALL) + view, blocks = lent(lease) + assert blocks == {0: [3, 4], 1: [0, 1, 2, 3, 4], 2: [0, 1, 2, 3, 4]}, blocks + bounds = return_write(lender, request(), lease, view.row_masks(True)) + assert tuple(bounds) == (LONG_PROMPT, 3 * TPB), bounds + + +@pytest.mark.cpu_only +def test_a_window_s_unwritten_pages_are_found_from_the_tokens_its_layer_group_computed( + monkeypatch, +): + """A grow past what a window's group computed refuses, naming that group.""" + monkeypatch.setattr(_manager, "keeps_stale_pages", lambda manager: True) + manager = Manager(windows=(64, 64, None)) + lender = lender_over(manager, layout((64, 64, None))) + kv = manager.kv_cache_map[RID] = cache(2 * TPB, (True, False, False), history=2 * TPB) + lender._ledger._kept_prefixes[RID] = (kv, (0, 0, 0), 2 * TPB) + assert lender._ledger._computed_tokens(RID, kv) == (2 * TPB, 0, 0) + refusal = lender._unwritten_pages_refusal(manager, RID, kv, 4 * TPB) + assert refusal is not None and refusal.startswith("layer group 1 "), refusal + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("memory", [HOST, DEVICE]) +def test_named_reads_exclude_private_rows_even_after_the_global_commit_advanced( + kit, monkeypatch, memory +): + manager = Manager() + lender = lender_over(manager, layout(named=(True, True, False))) + manager.kv_cache_map[RID] = cache(2 * TPB, (False,) * GROUPS, endpoint=0) + unnamed_keys(lender, monkeypatch) + admitted = [] + if memory is HOST: + monkeypatch.setattr(lender, "_fits_refusal", lambda counts: None) + monkeypatch.setattr(lender, "_admit", lambda lease, counts: admitted.append(lease)) + named = kit.HOST_NAMED if memory is HOST else kit.DEVICE_NAMED + lease = lender.lend_read(request(), 0, 2 * TPB, **named) + if memory is HOST: + assert admitted == [lease] and lease.failure is None + entries = lease._block_entries + blocks = { + lg: o.tolist() for lg, o in zip(entries.layer_group_indices, entries.block_ordinals) + } + else: + _, blocks = lent(lease) + assert blocks == {0: [0, 1], 1: [0, 1], 2: []} + lease.release() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("memory", [HOST, DEVICE]) +def test_named_writes_lend_no_rows_of_a_private_layer_group(kit, monkeypatch, memory): + """A named fetch lends a PRIVATE layer group nothing in either memory; the others arrive.""" + manager = Manager() + lender = lender_over(manager, layout(named=(True, True, False)), staged=True) + manager.kv_cache_map[RID] = cache(2 * TPB, (False,) * GROUPS, endpoint=0, history=2 * TPB) + unnamed_keys(lender, monkeypatch) + monkeypatch.setattr(lender, "_returns_context_outputs", lambda request: False) + monkeypatch.setattr(_manager, "grow", lambda manager, request, kv, position, end: True) + monkeypatch.setattr( + lender, + "_queue_copy", + lambda segments, to_staging, after=None: (_lender._QueuedCopy(Event()), None), + ) + named = kit.HOST_NAMED if memory is HOST else kit.DEVICE_NAMED + lease = lender.lend_write(request(), 0, 4 * TPB, **named) + view, blocks = lent(lease) + assert blocks == {0: [2, 3], 1: [2, 3], 2: []}, blocks + assert [len(run.object_keys) for run in view.groups] == [2, 2, 0] + if memory is HOST: + assert [lender._free_slots(g) for g in range(GROUPS)] == [2, 2, 4] + lender.return_writes([(lease, view.row_masks(True), True)]) + delivered = lender._ledger._deliveries_of(RID).delivered_block_masks + assert [m.nonzero()[0].tolist() for m in delivered] == [[2, 3], [2, 3], []] + + +BLOCKS = -(-PROMPT // TPB) +STEPS = (TPB, 2 * TPB, PROMPT) # the context's steps: the position each computes through + + +class Stage: + """One pipeline stage, ``lag`` iterations behind the first in processing each context step: its + manager, lender and request, the event recorded after its last step, and its lease.""" + + def __init__(self, index, lag=0): + self.index = index + self.lag = lag + self.manager = Manager(windows=(None,)) + self.lender = lender_over(self.manager, layout((None,)), [(index,)]) + self.request = request(PROMPT) + self.ready = None + self.lease = None + self.seen = False # a poll returned the view; the backend's access ends one iteration later + self.released_in = None + + def process(self, iteration): + """The context step this stage reaches in ``iteration``, and the event recorded after it.""" + step = iteration - self.lag + if 0 <= step < len(STEPS): + self.manager.kv_cache_map[RID].history_length = STEPS[step] + self.request.context_current_position = STEPS[step] + self.request.context_remaining_length = PROMPT - STEPS[step] + self.ready = Event() + + def processed_the_prompt(self): + return self.manager.kv_cache_map[RID].history_length == PROMPT + + def status(self): + """This stage's part of the per-iteration allgather: delivered once its backend's access + to a ready view ended, failed once its lease failed, else pending.""" + if self.lease is None: + return "pending" + view = self.lease.poll() # every open lease, every iteration + if self.lease.failure is not None: + return "failed" + if view is None: + return "pending" + if self.seen: + return "delivered" + self.seen = True + return "pending" + + +def agree(stages, iteration): + """Once every stage's lease settled, the outcome and every stage's release in this iteration.""" + statuses = [stage.status() for stage in stages] + if "pending" in statuses: + return None + for stage in stages: + stage.lease.release() + stage.released_in = iteration + return "delivered" if set(statuses) == {"delivered"} else "failed" + + +def send_the_prompt(stages, early_release=None, lend_together=False): + """Each stage lends a read of the prompt once it processed it (``lend_together``: once the first + did), polled until the stages agree; ``early_release``'s stage releases its lease early.""" + for stage in stages: + stage.manager.kv_cache_map[RID] = cache(history=0, blocks=BLOCKS) + for iteration in range(len(STEPS) + max(stage.lag for stage in stages) + 4): + for stage in stages: + stage.process(iteration) + ready = [stage.processed_the_prompt() for stage in stages] + for stage in stages: + due = any(ready) if lend_together else stage.processed_the_prompt() + if stage.lease is None and due: + stage.lease = stage.lender.lend_read( + stage.request, 0, PROMPT, memory_kind=DEVICE, rows=ALL, ready_event=stage.ready + ) + outcome = agree(stages, iteration) + if outcome is not None: + return outcome + early = stages[early_release] if early_release is not None else None + if early is not None and early.seen and early.released_in is None: + early.lease.release() # its own backend is done, the other stage's is not + early.released_in = iteration + raise AssertionError("the stages never agreed") + + +@pytest.mark.cpu_only +def test_each_stage_lends_once_it_processed_the_step_and_all_release_in_one_iteration(): + stages = [Stage(0), Stage(1, lag=1)] + assert send_the_prompt(stages) == "delivered" + assert stages[0].released_in == stages[1].released_in is not None + assert not any(stage.lender._is_locked(None) for stage in stages), "a lock outlived its lease" + + +@pytest.mark.cpu_only +def test_a_stage_that_has_not_processed_the_step_fails_its_lease_at_the_call(): + """A read lent past what a stage holds fails at the call; the stages agree it failed.""" + stages = [Stage(0), Stage(1, lag=1)] + assert send_the_prompt(stages, lend_together=True) == "failed" + assert "past the" in stages[1].lease.failure and stages[0].lease.failure is None + assert stages[0].released_in == stages[1].released_in is not None + + +@pytest.mark.cpu_only +def test_a_lease_released_before_the_stages_agree_fails_the_driver_loudly(): + """The driver's next poll of a lease released before the stages agree raises.""" + stages = [Stage(0), Stage(1, lag=2)] + with pytest.raises(RuntimeError, match="poll after release"): + send_the_prompt(stages, early_release=0) + assert stages[0].released_in is not None and stages[1].released_in is None + if stages[1].lease is not None: + stages[1].lease.release() + + +def fetch_the_prompt(kit, stages, delivered_blocks, histories=(PROMPT, PROMPT), end=PROMPT): + """Each stage's write of ``[0, end)`` returned with ``delivered_blocks``; the bounds.""" + bounds = [] + for stage, blocks, history in zip(stages, delivered_blocks, histories): + stage.manager.kv_cache_map[RID] = cache(history=history, blocks=BLOCKS) + req = request(PROMPT) + lease = stage.lender.lend_write(req, 0, end, **kit.DEVICE_ALL) + view = lease.poll() + assert view is not None, lease.failure + arrived = [np.isin(run.page_ordinals, list(blocks)) for run in view.groups] + bounds.append(return_write(stage.lender, req, lease, arrived)) + return bounds + + +def intersect(bounds): + """The rule ``ResumeBounds`` states across ranks: the largest minimum, the smallest maximum.""" + return ResumeBounds( + max_position=min(b.max_position for b in bounds), + min_position=max(b.min_position for b in bounds), + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("partial", [False, True], ids=["every_block", "one_stage_short"]) +def test_the_stages_resume_at_the_intersection_of_their_bounds(kit, partial): + """A stage short of a block keeps both from resuming at the prompt: the bounds are disjoint.""" + stages = [Stage(0), Stage(1)] + second = range(2) if partial else range(BLOCKS) + bounds = fetch_the_prompt(kit, stages, [range(BLOCKS), second]) + agreed = intersect(bounds) + if not partial: + assert bounds == [(PROMPT, PROMPT)] * 2 and agreed == (PROMPT, PROMPT) + return + assert bounds[0] == (PROMPT, PROMPT), bounds[0] + assert bounds[1].max_position < PROMPT, f"a short fetch resumes at the prompt: {bounds[1]}" + assert agreed.max_position < agreed.min_position, f"the stages could diverge: {agreed}" + + +@pytest.mark.cpu_only +def test_a_stage_s_lender_refuses_a_lease_another_stage_lent(kit): + """A lease another lender lent is refused and left to return through its own.""" + first, second = Stage(0), Stage(1) + second.manager.kv_cache_map[RID] = cache(history=PROMPT, blocks=BLOCKS) + lease = second.lender.lend_write(request(PROMPT), 0, PROMPT, **kit.DEVICE_ALL) + assert lease.poll() is not None, lease.failure + with pytest.raises(ValueError, match="not lent by this lender"): + first.lender.return_writes([(lease, ALL_ARRIVED, True)]) + assert return_write(second.lender, request(PROMPT), lease, ALL_ARRIVED) == (PROMPT, PROMPT) + + +@pytest.mark.cpu_only +def test_pipeline_intersection_uses_the_larger_of_real_stage_lower_bounds(kit): + bounds = fetch_the_prompt(kit, [Stage(0), Stage(1)], [range(BLOCKS)] * 2, (0, 64), end=96) + assert bounds == [(96, 0), (96, 64)] + assert intersect(bounds) == (96, 64) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("hybrid", [False, True], ids=["attention_model", "hybrid_model"]) +def test_a_stage_of_a_model_with_recurrent_state_names_nothing_whatever_layers_it_holds( + monkeypatch, hybrid +): + """Every stage of a model with recurrent state names nothing; a model without names blocks.""" + stage = SimpleNamespace(cp=(0, 1), recurrent=(False, False)) + monkeypatch.setattr(_manager, "model_has_recurrent_state", lambda manager: hybrid) + monkeypatch.setattr(_manager, "commits_blocks", lambda manager: True) + monkeypatch.setattr(_manager, "prompt_lookahead", lambda manager: 0) + monkeypatch.setattr(_manager, "has_lossy_cold_page_codec", lambda manager: False) + reason = _lender._no_names(SimpleNamespace(), stage) + if hybrid: + assert reason is not None and "recurrent state" in reason, reason + else: + assert reason is None, reason + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("feature", ["window", "state", "sparse"]) +@pytest.mark.parametrize("pp_rank", [0, 1]) +def test_attach_rejects_model_wide_features_on_every_pipeline_stage(monkeypatch, feature, pp_rank): + """A stage without the unsupported layer rejects the same model as the stage holding it.""" + stage = SimpleNamespace( + cp=(0, 2), + num_layer_groups=1, + recurrent=(False,), + windows=(None,), + sink_blocks=(0,), + sparse=(False,), + ) + manager = SimpleNamespace( + mapping=SimpleNamespace(cp_size=2, pp_rank=pp_rank, has_cp_helix=lambda: True), + _model_has_attention_window=feature == "window", + _model_has_sparse_buffers=feature == "sparse", + _model_attention_window_pattern=(), + num_layers=2, + ) + monkeypatch.setattr(_manager, "require_v2", lambda manager: None) + monkeypatch.setattr(_manager, "layer_configs", lambda manager: {}) + monkeypatch.setattr(_manager, "layer_group_layer_ids", lambda manager: []) + monkeypatch.setattr(_lender, "derive_layout", lambda manager: stage) + monkeypatch.setattr(_manager, "model_has_recurrent_state", lambda manager: feature == "state") + with pytest.raises(ValueError, match="sparse buffers" if feature == "sparse" else "Helix"): + _lender._checked_layout(manager) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "pattern,num_layers,expected", + [((256, 64), 2, True), ((256, 64), 1, False), ((256,), 2, False), ((), 2, False)], +) +def test_model_window_detection_keeps_windows_on_other_pipeline_stages( + pattern, num_layers, expected +): + manager = SimpleNamespace( + max_seq_len=256, + num_layers=num_layers, + _model_has_attention_window=False, + _model_has_sparse_buffers=False, + _model_attention_window_pattern=pattern, + ) + assert _manager.model_has_attention_window(manager) is expected + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "missing, detector", + [ + ("_model_has_attention_window", _manager.model_has_attention_window), + ("_model_attention_window_pattern", _manager.model_has_attention_window), + ("num_layers", _manager.model_has_attention_window), + ("_model_has_sparse_buffers", _manager.model_has_sparse_buffers), + ], +) +def test_model_features_require_explicit_metadata(missing, detector): + manager = SimpleNamespace( + max_seq_len=256, + num_layers=2, + _model_has_attention_window=False, + _model_has_sparse_buffers=False, + _model_attention_window_pattern=(256, 64), + ) + delattr(manager, missing) + with pytest.raises(AttributeError, match=missing): + detector(manager) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_slots.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_slots.py index b13004823378..2c9176d87af1 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_slots.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_slots.py @@ -12,21 +12,17 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Staging sizing and the slot queue alone, over plain layouts, the usable end resume_bounds finds and +"""Staging sizing and the slot queue alone, over plain layouts, the usable end poll_resume_bounds finds and how it and fetch ends keep out of runs of multimodal tokens, with the runtime's own life cycles behind a stand-in manager, the delivered rows a shrink hook drops when its own work raises, the -slots a read's grant or a write's marks keep when a copy's completion cannot be recorded, also when +slots a read's grant or a write's return keep when a copy's completion cannot be recorded, also when the quarantine itself raises, a write whose grant and failure both raise, a lender stopped by an -error its free or shutdown hook logs, the records of a replaced cache resume_bounds drops, the states +error its free or shutdown hook logs, the records of a replaced cache poll_resume_bounds drops, the states a fetch walks through, and the bytes the test kit stages: no real manager, no memory, no GPU.""" -import ast import gc -import inspect -import itertools import random import weakref -from pathlib import Path from types import SimpleNamespace import numpy as np @@ -36,6 +32,7 @@ StagingOptions, _fetch_ledger, _lender, + _manager, _ranges, ) from tensorrt_llm._torch.pyexecutor.kv_cache.sharing._slots import ( @@ -88,24 +85,25 @@ def one_fetch(lay, tokens_per_fetch): def test_tinyllama_stages_eight_whole_prompts(): assert slots_per_fetch(TINYLLAMA, 2048) == {0: 64} assert one_fetch(TINYLLAMA, 2048) == 44 * MiB - assert slot_counts(TINYLLAMA, StagingOptions(2048, planned_fetches=8)) == {0: 8 * 64} + assert slot_counts(TINYLLAMA, StagingOptions(2048, capacity_in_fetches=8)) == {0: 8 * 64} def test_deepseek_v4_windows_take_their_window_and_the_full_group_the_range(): assert slots_per_fetch(DSV4_PRO, 4096) == {1: 1, 0: 32, 2: 1} assert 74 * MiB < one_fetch(DSV4_PRO, 4096) < 75 * MiB - assert slot_counts(DSV4_PRO, StagingOptions(4096, planned_fetches=8)) == {1: 8, 0: 256, 2: 8} + options = StagingOptions(4096, capacity_in_fetches=8) + assert slot_counts(DSV4_PRO, options) == {1: 8, 0: 256, 2: 8} def test_deepseek_v4_long_ranges_capped_at_one_fetch_keep_one_fetch(): one = one_fetch(DSV4_PRO, 1 << 20) assert slots_per_fetch(DSV4_PRO, 1 << 20) == {1: 1, 0: 8192, 2: 1} assert 4 * GiB < one < 5 * GiB - options = StagingOptions(1 << 20, planned_fetches=8, max_bytes=one) + options = StagingOptions(1 << 20, capacity_in_fetches=8, max_bytes=one) assert slot_counts(DSV4_PRO, options) == {1: 1, 0: 8192, 2: 1} # A cap between one and eight fetches: every group keeps its rows of one fetch and more. cap = 960 * GiB // 8 // 4 - counts = slot_counts(DSV4_PRO, StagingOptions(1 << 20, planned_fetches=8, max_bytes=cap)) + counts = slot_counts(DSV4_PRO, StagingOptions(1 << 20, capacity_in_fetches=8, max_bytes=cap)) assert counts[0] >= 8192 and counts[1] >= 1 and counts[2] >= 1 assert sum(counts[g] * DSV4_PAGES[g] for g in counts) <= cap @@ -118,14 +116,14 @@ def test_short_ranges_take_no_more_than_their_blocks_even_in_a_window(): def test_a_window_wider_than_the_range_takes_the_range_blocks(): lay = layout(32, [(1024, 0, 0, 100)]) # a window of 32 blocks over a range of 8 assert slots_per_fetch(lay, 256) == {0: 8} - assert slot_counts(lay, StagingOptions(256, planned_fetches=2)) == {0: 16} + assert slot_counts(lay, StagingOptions(256, capacity_in_fetches=2)) == {0: 16} def test_sink_blocks_count_with_the_window(): lay = layout(32, [(64, 1, 0, 100), (None, 0, 1, 10)]) assert slots_per_fetch(lay, 320) == {0: 1 + 2, 1: 10} assert one_fetch(lay, 320) == 3 * 100 + 10 * 10 - assert slot_counts(lay, StagingOptions(320, planned_fetches=3)) == {0: 9, 1: 30} + assert slot_counts(lay, StagingOptions(320, capacity_in_fetches=3)) == {0: 9, 1: 30} def test_layer_groups_sharing_a_pool_group_add_up(): @@ -138,10 +136,10 @@ def test_a_partial_block_counts_as_a_whole_one(): assert slots_per_fetch(TINYLLAMA, 1) == {0: 1} -@pytest.mark.parametrize("tokens_per_fetch", [0, -32]) -def test_tokens_per_fetch_must_be_positive(tokens_per_fetch): +@pytest.mark.parametrize("guaranteed_fetch_tokens", [0, -32]) +def test_guaranteed_fetch_tokens_must_be_positive(guaranteed_fetch_tokens): with pytest.raises(ValueError): - slots_per_fetch(TINYLLAMA, tokens_per_fetch) + slots_per_fetch(TINYLLAMA, guaranteed_fetch_tokens) def test_max_bytes_below_one_fetch_is_rejected(): @@ -153,8 +151,9 @@ def test_max_bytes_below_one_fetch_is_rejected(): ) -def test_max_bytes_above_planned_fetches_changes_nothing(): - options = StagingOptions(4096, planned_fetches=2, max_bytes=100 * one_fetch(DSV4_PRO, 4096)) +def test_max_bytes_above_capacity_in_fetches_changes_nothing(): + one = one_fetch(DSV4_PRO, 4096) + options = StagingOptions(4096, capacity_in_fetches=2, max_bytes=100 * one) assert slot_counts(DSV4_PRO, options) == {1: 2, 0: 64, 2: 2} @@ -174,7 +173,7 @@ def random_layout(rng): @pytest.mark.parametrize("seed", range(40)) -def test_any_whole_block_range_of_at_most_tokens_per_fetch_fits_one_fetch(seed): +def test_any_whole_block_range_of_at_most_guaranteed_fetch_tokens_fits_one_fetch(seed): # Rows a range [start, end) needs: its blocks that a history of end still reads, as the # runtime's own life cycle says; the longest range ending at end needs the most. rng = random.Random(seed) @@ -194,28 +193,28 @@ def test_any_whole_block_range_of_at_most_tokens_per_fetch_fits_one_fetch(seed): @pytest.mark.parametrize("seed", range(40)) -def test_counts_are_planned_fetches_fetches_or_the_cap_rounded_down(seed): +def test_counts_are_capacity_in_fetches_fetches_or_the_cap_rounded_down(seed): rng = random.Random(seed) lay, _ = random_layout(rng) tokens_per_fetch = rng.randint(1, 12 * lay.tokens_per_block) rows = slots_per_fetch(lay, tokens_per_fetch) one = one_fetch(lay, tokens_per_fetch) - planned_fetches = rng.randint(1, 9) - assert slot_counts(lay, StagingOptions(tokens_per_fetch, planned_fetches)) == { - g: planned_fetches * rows[g] for g in rows + capacity_in_fetches = rng.randint(1, 9) + assert slot_counts(lay, StagingOptions(tokens_per_fetch, capacity_in_fetches)) == { + g: capacity_in_fetches * rows[g] for g in rows } - cap = rng.randint(one, planned_fetches * one) - counts = slot_counts(lay, StagingOptions(tokens_per_fetch, planned_fetches, max_bytes=cap)) + cap = rng.randint(one, capacity_in_fetches * one) + counts = slot_counts(lay, StagingOptions(tokens_per_fetch, capacity_in_fetches, max_bytes=cap)) assert sum(counts[g] * lay.page_bytes[g] for g in counts) <= cap # As many whole fetches as the cap holds hold slots at once. whole = cap // one - assert all(whole * rows[g] <= counts[g] <= planned_fetches * rows[g] for g in rows) + assert all(whole * rows[g] <= counts[g] <= capacity_in_fetches * rows[g] for g in rows) assert all(counts[g] * one <= cap * rows[g] < (counts[g] + 1) * one for g in rows) def test_a_fractional_fetch_budget_keeps_slots_for_a_shorter_range() -> None: lay = layout(32, [(None, 0, 0, 8)]) - slots = SlotQueue(slot_counts(lay, StagingOptions(320, planned_fetches=3, max_bytes=200))) + slots = SlotQueue(slot_counts(lay, StagingOptions(320, capacity_in_fetches=3, max_bytes=200))) assert slots.num_slots(0) == 25 for rows in (10, 10, 5): assert slots.try_grant(slots.enqueue({0: rows})) is not None @@ -223,7 +222,7 @@ def test_a_fractional_fetch_budget_keeps_slots_for_a_shorter_range() -> None: @pytest.mark.parametrize("max_bytes_in_fetches", [None, 1, 2.5]) -def test_planned_fetches_ranges_hold_slots_at_once_and_one_more_waits(max_bytes_in_fetches): +def test_capacity_in_fetches_ranges_hold_slots_at_once_and_one_more_waits(max_bytes_in_fetches): rows = slots_per_fetch(DSV4_PRO, 4096) max_bytes = None at_once = 3 @@ -231,7 +230,7 @@ def test_planned_fetches_ranges_hold_slots_at_once_and_one_more_waits(max_bytes_ max_bytes = int(max_bytes_in_fetches * one_fetch(DSV4_PRO, 4096)) at_once = int(max_bytes_in_fetches) slots = SlotQueue( - slot_counts(DSV4_PRO, StagingOptions(4096, planned_fetches=3, max_bytes=max_bytes)) + slot_counts(DSV4_PRO, StagingOptions(4096, capacity_in_fetches=3, max_bytes=max_bytes)) ) for _ in range(at_once): assert slots.try_grant(slots.enqueue(rows)) is not None @@ -401,13 +400,16 @@ def test_zero_counts_need_no_slots(): class _Owner: """What a stand-in lender's weak references point at.""" + enable_block_reuse = False + _can_publish_block_reuse = False + is_draft = False + _OWNER = _Owner() -def staging_over(windows, sinks, tpb): - """A staging lender and a manager stand-in answering what resume_bounds asks: the tokens per block - and each layer group's stale range, from the runtime's own life cycle.""" +def lender_over(windows, sinks, tpb): + """Lender and manager stand-ins using the runtime's stale-range rule.""" cycles = [AttnLifeCycle.make(w, s * tpb, tpb) for w, s in zip(windows, sinks)] def stale_block_range(lg, history): @@ -418,8 +420,9 @@ def stale_block_range(lg, history): lay = SimpleNamespace( tokens_per_block=tpb, windows=tuple(windows), pool_groups=(), num_layer_groups=len(windows) ) + lay.named_groups = (True,) * len(windows) ref = weakref.ref(_OWNER) - return _lender.Staging(ref, lay, None, (), None, ref), manager + return _lender._ManagerLender(ref, lay, None, (), None, ref), manager def reads_all(manager, windows, blocks, first, end_block, tpb): @@ -457,14 +460,14 @@ def test_max_resume_position_is_the_largest_start_every_layer_group_reads(): row[rng.randrange(len(row))] = False row[: origin // tpb] = False # fetches start at origin or later blocks.append(row) - lender, manager = staging_over(windows, sinks, tpb) + lender, manager = lender_over(windows, sinks, tpb) first = known // tpb last = max((len(b) for b in blocks), default=0) ends = range(first + 1, last + 1) starts = [e for e in ends if reads_all(manager, windows, blocks, first, e, tpb)] want = starts[-1] * tpb if starts else known delivered = _fetch_ledger._Deliveries(None, origin, [b.copy() for b in blocks]) - got = lender._ledger._max_resume_position(manager, delivered, known) + got = lender._ledger._max_resume_position(manager, delivered, [known] * len(windows)) assert got == want, ( f"windows {windows}, sinks {sinks}, tpb {tpb}, origin {origin}, known {known}, " f"delivered {[np.nonzero(b)[0].tolist() for b in blocks]}: usable until {got}, " @@ -472,53 +475,32 @@ def test_max_resume_position_is_the_largest_start_every_layer_group_reads(): ) -def test_max_resume_position_after_a_whole_fetch_checks_its_end_alone(monkeypatch): - """Two windowed layer groups, nothing computed and 4096 blocks delivered in each: the end of the - delivered rows is usable, found with one stale-range lookup per layer group, not one per block - below it, so resume_bounds after each lease of a split fetch costs no more as the fetch grows.""" - lender, manager = staging_over([4096, 8192], [0, 0], 32) +def test_max_resume_position_after_a_whole_fetch_is_the_delivered_end(): + """With both windows delivered, the last delivered block bounds recovery.""" + lender, manager = lender_over([4096, 8192], [0, 0], 32) blocks = [np.ones(4096, dtype=bool), np.ones(4096, dtype=bool)] - real_stale, calls = _ranges._stale_range, [] - - def counted(*args): - calls.append(args) - return real_stale(*args) - - monkeypatch.setattr(_ranges, "_stale_range", counted) got = lender._ledger._max_resume_position( - manager, _fetch_ledger._Deliveries(None, 0, blocks), 0 + manager, _fetch_ledger._Deliveries(None, 0, blocks), (0, 0) ) assert got == 4096 * 32, f"usable until {got}" - assert len(calls) == 2, f"{len(calls)} stale-range lookups for one usable end" -def test_max_resume_position_checks_no_start_past_a_missing_full_attention_block(monkeypatch): - """One full-attention layer group, nothing computed and 8192 blocks delivered but one: no start - past the missing block has every block it reads, so none is checked, and resume_bounds after each - later lease of a split fetch that lost an early block checks no more starts as it grows.""" - lender, manager = staging_over([None], [0], 32) - real_stale, calls = _ranges._stale_range, [] - - def counted(*args): - calls.append(args) - return real_stale(*args) - - monkeypatch.setattr(_ranges, "_stale_range", counted) - for hole, usable, checked in ((0, 0, 0), (4000, 4000 * 32, 1)): +def test_a_missing_full_attention_block_bounds_the_resume_position(): + """A missing full-attention page clips recovery, including a hole at the start.""" + lender, manager = lender_over([None], [0], 32) + for hole, usable in ((0, 0), (4000, 4000 * 32)): blocks = np.ones(8192, dtype=bool) blocks[hole] = False - calls.clear() got = lender._ledger._max_resume_position( - manager, _fetch_ledger._Deliveries(None, 0, [blocks]), 0 + manager, _fetch_ledger._Deliveries(None, 0, [blocks]), (0,) ) assert got == usable, f"usable until {got} with block {hole} missing" - assert len(calls) == checked, f"{len(calls)} starts checked with block {hole} missing" def test_a_shrink_hook_that_raises_forgets_what_the_request_was_delivered(monkeypatch): """A shrink hook whose own work raises, which it logs: the request's delivered rows all go, - since the cache may grow back before resume_bounds sees the freed ones past its blocks.""" - lender, _ = staging_over([None], [0], 32) + since the cache may grow back before poll_resume_bounds sees the freed ones past its blocks.""" + lender, _ = lender_over([None], [0], 32) kv = object() lender._ledger._deliveries[7] = _fetch_ledger._Deliveries(kv, 0, [np.ones(3, dtype=bool)]) lender._ledger._mixed_tail_caches[7] = kv @@ -553,11 +535,11 @@ def strictly_inside(runs, position): def test_resume_bounds_and_fetch_ends_keep_out_of_bidirectional_runs(): """Random bidirectional runs of multimodal tokens: a fetch end strictly inside a run is refused, - others not; resume_bounds lowers its end to the largest position at or below it not strictly inside + others not; poll_resume_bounds lowers its end to the largest position at or below it not strictly inside one, with a floor at the end of the last run at or below that, so none from floor to end is.""" rng = random.Random(0) # With a window too, which the refusal must not read: the model's window can be narrower. - lender, windowed = (staging_over(w, [0] * len(w), 32)[0] for w in ([None], [64, None])) + lender, windowed = (lender_over(w, [0] * len(w), 32)[0] for w in ([None], [64, None])) for _ in range(3000): length = rng.randint(1, 400) runs, pos = [], rng.randint(0, 20) @@ -567,7 +549,9 @@ def test_resume_bounds_and_fetch_ends_keep_out_of_bidirectional_runs(): pos = stop + rng.randint(1, 40) request = multimodal(runs, length) usable = rng.randint(0, length) - got, floor = lender._ledger._clip_to_bidirectional_runs(request, usable) + got, floor = lender._ledger._clip_to_bidirectional_runs( + _manager.bidirectional_runs(request), usable + ) want = max(p for p in range(usable + 1) if not strictly_inside(runs, p)) case = f"runs {runs}, usable {usable}" assert got == want, f"{case}: the end lowered to {got}, not {want}" @@ -587,13 +571,15 @@ def test_resume_bounds_and_fetch_ends_keep_out_of_bidirectional_runs(): def test_a_request_without_bidirectional_runs_keeps_its_interval_and_ends(): """No flag, no multimodal data, or no mask: nothing lowers, rises or fails.""" - lender, _ = staging_over([64, None], [0, 0], 32) + lender, _ = lender_over([64, None], [0, 0], 32) for request in ( multimodal([(1, 300)], 400, flagged=False), SimpleNamespace(py_multimodal_data=None), SimpleNamespace(py_multimodal_data={"mm_bidirectional_blocks": True}), ): - assert lender._ledger._clip_to_bidirectional_runs(request, 200) == (200, 0) + assert lender._ledger._clip_to_bidirectional_runs( + _manager.bidirectional_runs(request), 200 + ) == (200, 0) assert lender._bidirectional_split_refusal(request, 200) is None @@ -626,7 +612,7 @@ def append(self, item): super().append(item) -def _stand_in_view(rows, keys): +def _stand_in_view(rows, keys, *bounds): return "view" @@ -640,10 +626,11 @@ def test_a_read_grant_keeps_the_slots_its_queued_copy_may_write_out_of_use(monke pool_group_of=(0,), windows=(None,), num_layer_groups=1, + named_groups=(True,), tokens_per_block=1, ) ref = weakref.ref(_OWNER) - lender = _lender.Staging(ref, lay, None, (), SlotQueue({0: 4}), ref) + lender = _lender._ManagerLender(ref, lay, None, (), SlotQueue({0: 4}), ref) fired = [] view = _stand_in_view if fault == "before": @@ -651,7 +638,7 @@ def test_a_read_grant_keeps_the_slots_its_queued_copy_may_write_out_of_use(monke monkeypatch.setattr(lender, "_build_view", view) monkeypatch.setattr(lender, "_copy_segments", lambda rows: []) monkeypatch.setattr( - lender, "_queue_copy", lambda segments, to_staging: (None, "planted: no event") + lender, "_queue_copy", lambda segments, to_staging, after=None: (None, "planted: no event") ) if fault == "record": lender._quarantined = _AppendRaisingOnce(fired) @@ -685,7 +672,7 @@ def __iter__(self): return super().__iter__() -def _unrecorded_copy(segments, to_staging): +def _unrecorded_copy(segments, to_staging, after=None): raise MemoryError("planted: the copy is queued, its record is not") @@ -695,6 +682,7 @@ def _one_pool_group(): pool_group_of=(0,), windows=(None,), num_layer_groups=1, + named_groups=(True,), tokens_per_block=1, ) @@ -713,6 +701,63 @@ def _write_target(): return kv, manager +@pytest.mark.parametrize("kind", ["read", "write"]) +def test_an_interrupted_immediate_grant_leaves_no_unreachable_lease(monkeypatch, kind): + """An interrupted lend cannot return its lease, but must keep possibly touched slots lost.""" + kv, manager = _write_target() + ref = weakref.ref(manager) + lender = _lender._ManagerLender(ref, _one_pool_group(), None, (), SlotQueue({0: 4}), ref) + rows = _lender._LentBlockEntries([0], [np.arange(2)], [np.arange(2)]) + fetch = _fetch_ledger._FetchRecord(0, kv) if kind == "write" else None + if fetch is not None: + lender._ledger._latest_fetches[7] = fetch + lease = _lender._StagingLease(lender, kind, 7, kv, rows, None, fetch) + unreachable = weakref.ref(lease) + + def interrupted(*args): + raise KeyboardInterrupt("planted") + + monkeypatch.setattr(lender, "_build_view", interrupted) + with pytest.raises(KeyboardInterrupt, match="planted"): + lender._admit(lease, {0: 2}) + assert lender._open_lease_count() == 0, "the caller never received this lease" + assert 7 not in lender._ledger._latest_fetches, "an unreachable fetch stayed unsettled" + assert lender._free_slots(0) == 2 and len(lender._quarantined) == 1 + assert lender._must_retain_memory(), "an interrupted grant may have queued a copy" + del lease + gc.collect() + assert unreachable() is None, "the lender retained a lease no caller can release" + + +def test_an_interrupted_older_grant_cancels_the_new_unreturned_lease(monkeypatch): + """An older waiter keeps its failure; the new lend cannot leave a second ticket behind.""" + kv, manager = _write_target() + ref = weakref.ref(manager) + lender = _lender._ManagerLender(ref, _one_pool_group(), None, (), SlotQueue({0: 4}), ref) + rows = _lender._LentBlockEntries([0], [np.arange(2)], [np.arange(2)]) + older = _lender._StagingLease(lender, "write", 8, kv, rows) + older._ticket = lender._slot_queue.enqueue({0: 2}) + lender._waiting_for_slots.append(older) + lender._open_leases.add(older) + fetch = _fetch_ledger._FetchRecord(0, kv) + lender._ledger._latest_fetches[7] = fetch + current = _lender._StagingLease(lender, "write", 7, kv, rows, None, fetch) + + def interrupted(*args): + raise KeyboardInterrupt("planted") + + monkeypatch.setattr(lender, "_build_view", interrupted) + with pytest.raises(KeyboardInterrupt, match="planted"): + lender._admit(current, {0: 2}) + assert older.failure is not None and older.poll() is None + assert lender._open_lease_count() == 1, "the new caller never received its lease" + assert lender._slot_queue.num_waiting == 0 and not lender._waiting_for_slots + assert 7 not in lender._ledger._latest_fetches + assert lender._free_slots(0) == 2 and len(lender._quarantined) == 1 + older.release() + assert lender._open_lease_count() == 0 + + @pytest.mark.parametrize("fault", ["none", "call", "record", "holding", "log", "append", "partial"]) def test_a_read_whose_grant_and_quarantine_both_raise_is_never_ready(monkeypatch, fault): """A read granted 2 of 4 slots: its grant raises once its view is set (recording its copy, @@ -722,17 +767,17 @@ def test_a_read_whose_grant_and_quarantine_both_raise_is_never_ready(monkeypatch release, and where a copy may be queued the shutdown would keep the staging memory. With no fault it is ready once its copy landed.""" ref = weakref.ref(_OWNER) - lender = _lender.Staging(ref, _one_pool_group(), None, (), SlotQueue({0: 4}), ref) + lender = _lender._ManagerLender(ref, _one_pool_group(), None, (), SlotQueue({0: 4}), ref) planted = [[], []] # the grant's fault, then the quarantine's landed = SimpleNamespace(query=lambda: True) monkeypatch.setattr(lender, "_build_view", _stand_in_view) monkeypatch.setattr(lender, "_copy_segments", lambda rows: []) if fault == "none": copied = (_lender._QueuedCopy(landed), None) - monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging: copied) + monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging, after=None: copied) elif fault == "partial": copied = (_lender._QueuedCopy(landed), "planted: a copy call failed") - monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging: copied) + monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging, after=None: copied) failing = raising_once(lender._fail, MemoryError("planted"), planted[0]) monkeypatch.setattr(lender, "_fail", failing) else: @@ -761,6 +806,7 @@ def test_a_read_whose_grant_and_quarantine_both_raise_is_never_ready(monkeypatch assert all(planted), "a planted fault did not strike: proves nothing" view = read.poll() read.release() + lender.progress() # the round that returns a released lease's slots if fault == "none": assert view == "view" and read.failure is None, "a read whose copy landed was not ready" assert lender._free_slots(0) == 4, "the slots of a read whose copy landed stayed out of use" @@ -780,21 +826,21 @@ def test_a_read_whose_grant_and_quarantine_both_raise_is_never_ready(monkeypatch def test_marks_whose_copy_and_quarantine_both_fail_keep_the_slots_out_of_use( monkeypatch, copy, fault ): - """A write granted 2 of 4 slots is marked; its marks' copy raises once queued or returns with no - record of its completion, then its quarantine raises too, at its call or the slots' record: - after the release the slots stay out of use and the shutdown would keep the staging memory.""" + """A write returned after its copy and then its quarantine fail keeps its slots out of use.""" kv, manager = _write_target() ref = weakref.ref(manager) - lender = _lender.Staging(ref, _one_pool_group(), None, (), SlotQueue({0: 4}), ref) + lender = _lender._ManagerLender(ref, _one_pool_group(), None, (), SlotQueue({0: 4}), ref) fired = [] monkeypatch.setattr( - lender, "_build_view", lambda rows, keys: SimpleNamespace(groups=rows.block_ordinals) + lender, "_build_view", lambda rows, *keys: SimpleNamespace(groups=rows.block_ordinals) ) monkeypatch.setattr(lender, "_copy_segments", lambda rows, mask: []) if copy == "raised": monkeypatch.setattr(lender, "_queue_copy", _unrecorded_copy) else: - monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging: (None, "no event")) + monkeypatch.setattr( + lender, "_queue_copy", lambda segments, to_staging, after=None: (None, "no event") + ) if fault == "call": failing = raising_once(lender._quarantine, MemoryError("planted"), fired) monkeypatch.setattr(lender, "_quarantine", failing) @@ -804,10 +850,11 @@ def test_marks_whose_copy_and_quarantine_both_fail_keep_the_slots_out_of_use( write = _lender._StagingLease(lender, "write", 7, kv, rows) lender._admit(write, {0: 2}) view = write.poll() + arrived = [np.ones(len(run), dtype=bool) for run in view.groups] with pytest.raises(MemoryError, match="planted"): - write.mark_arrived([np.ones(len(run), dtype=bool) for run in view.groups]) + lender.return_writes([(write, arrived, True)]) assert fired, "the quarantine did not raise: proves nothing" - write.release() + lender.progress() assert lender._free_slots(0) == 2, "slots a queued copy may still read came back" assert lender._must_retain_memory(), ( "the shutdown would free staging a queued copy may still read" @@ -830,19 +877,17 @@ def append(self, item): def test_a_write_whose_grant_and_failure_both_raise_is_never_ready(monkeypatch): - """A write waits for the slots a read holds, and the read's release grants it: entering the - holding list raises, then failing the write raises too. As after a manager hook logged that - error, the write never returns a view over the slots it gave back for the next grant.""" + """A write whose grant and then its failure raise never returns a view over its slots.""" kv, manager = _write_target() ref = weakref.ref(manager) - lender = _lender.Staging(ref, _one_pool_group(), None, (), SlotQueue({0: 2}), ref) + lender = _lender._ManagerLender(ref, _one_pool_group(), None, (), SlotQueue({0: 2}), ref) landed = SimpleNamespace(query=lambda: True) monkeypatch.setattr( - lender, "_build_view", lambda rows, keys: SimpleNamespace(groups=rows.block_ordinals) + lender, "_build_view", lambda rows, *keys: SimpleNamespace(groups=rows.block_ordinals) ) monkeypatch.setattr(lender, "_copy_segments", lambda rows, mask=None: []) copied = (_lender._QueuedCopy(landed), None) - monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging: copied) + monkeypatch.setattr(lender, "_queue_copy", lambda segments, to_staging, after=None: copied) read = _lender._StagingLease( lender, "read", @@ -869,8 +914,9 @@ def recycling(): # the holding list, once rebuilt, raises at the write's grant monkeypatch.setattr(lender, "_recycle", recycling) monkeypatch.setattr(lender, "_fail", raising_once(lender._fail, MemoryError("planted"), failed)) - with pytest.raises(MemoryError, match="planted"): # what a manager hook logs and goes on - read.release() + read.release() + with pytest.raises(MemoryError, match="planted"): + lender.progress() assert appended and failed, "a planted fault did not strike: proves nothing" assert write.poll() is None, "a write was ready over slots it gave back" @@ -894,6 +940,7 @@ def _cache(blocks=2, committed=0): pages = list(range(blocks)) return SimpleNamespace( num_committed_tokens=committed, + reuse_status=[SimpleNamespace(complete=True, endpoint=committed)] * 2, history_length=committed, is_active=True, num_blocks=blocks, @@ -921,20 +968,22 @@ def _a_write_waits(monkeypatch, fault, error=MemoryError): """Two layer groups on pool groups 0 and 1 of four slots each: a read holds slots 0 and 1 of both with its copy pending, a write into request 8 whose view a poll returned holds slot 2, a write whose view no poll returned holds slot 3, and a write into request 7 waits for 2 + 2; - request 10's cache has 8 committed blocks. The read is released and its copy lands. ``error`` - is planted in pool group 1, changed second, for its next take or return of slots as ``fault`` - says: the take of the waiting write's slots or the return of the read's.""" + request 10's cache has 8 committed blocks, request 11's 8 blocks none. The read is released and + its copy lands, no round of progress since. ``error`` is planted where ``fault`` says: pool + group 1's next ``"take"`` or ``"give"`` of slots, or the ledger's next ``"drop"`` of records.""" manager = _Owner() manager.kv_cache_map = {7: _cache(), 8: _cache(), 10: _cache(blocks=8, committed=8)} + manager.kv_cache_map[11] = _cache(blocks=8) layout = SimpleNamespace( pool_groups=(0, 1), pool_group_of=(0, 1), windows=(None, None), num_layer_groups=2, + named_groups=(True, True), tokens_per_block=1, ) memory = _Memory() - lender = _lender.Staging( + lender = _lender._ManagerLender( weakref.ref(manager), layout, None, (), SlotQueue({0: 4, 1: 4}), weakref.ref(memory) ) s = SimpleNamespace( @@ -943,12 +992,12 @@ def _a_write_waits(monkeypatch, fault, error=MemoryError): landed = [] event = SimpleNamespace(query=lambda: bool(landed), synchronize=lambda: None) - def memcpy(segments, to_staging): + def memcpy(segments, to_staging, after=None): s.copies.append(to_staging) return _lender._QueuedCopy(event), None monkeypatch.setattr( - lender, "_build_view", lambda rows, keys: SimpleNamespace(groups=rows.block_ordinals) + lender, "_build_view", lambda rows, *keys: SimpleNamespace(groups=rows.block_ordinals) ) monkeypatch.setattr(lender, "_copy_segments", lambda rows, mask=None: []) monkeypatch.setattr(lender, "_queue_copy", memcpy) @@ -972,7 +1021,11 @@ def lent(kind, request_id, count): assert s.waiting.poll() is None and s.waiting.failure is None, "the write did not wait" s.read.release() # its copy pends, so its slots stay out of use landed.append(True) - if fault != "none": + if fault == "drop": + ledger = lender._ledger + planted = raising_once(ledger._drop_records, error("planted"), s.fired) + monkeypatch.setattr(ledger, "_drop_records", planted) + elif fault != "none": alloc = lender._slot_queue._allocators[1] planted = raising_once(getattr(alloc, fault), error("planted"), s.fired) monkeypatch.setattr(alloc, fault, planted) @@ -989,50 +1042,40 @@ def _freed_while_a_write_waits(monkeypatch, fault, error=MemoryError): @pytest.mark.parametrize( "fault, error", - [ - ("none", None), - ("take", MemoryError), - ("give", MemoryError), - ("take", KeyError), - ("give", KeyError), - ], - ids=["none", "take", "give", "take_key_error", "give_key_error"], + [("none", None), ("drop", MemoryError), ("drop", KeyError)], + ids=["none", "drop", "drop_key_error"], ) -def test_an_error_the_free_hook_logs_stops_the_lender(monkeypatch, fault, error): - """Another request's free returns the slots of a released read whose copy landed and grants - them to the write waiting for them, and a planted MemoryError or KeyError strikes between the - changes of the two pool groups: at the take of the write's slots or at the return of the read's. - The hook logs it, raises nothing into the free, and stops the lender for good: the next lender - call fails the waiting write, a poll fails a lease no poll returned the view of and returns a - view one did, lends fail at the call after their range checks, resume_bounds is an empty interval, - a mark copies nothing, no slot returns or is granted, and the shutdown keeps the staging memory - with nothing open. ResumeBounds of a request without a cache still raises. A range needing more - slots than a part has still raises in ``lend_write``, which checks the stop after its slot - check, and fails at the call in ``lend_read``, which checks the stop before any check of the - cache, each where it checks a reuse reset. Without the fault the free grants the waiting write, - and that range raises in both lends.""" +def test_an_error_the_free_hook_logs_stops_the_lender(kit, monkeypatch, fault, error): + """An error the free hook logs stops the lender for good; without one, progress grants.""" s = _freed_while_a_write_waits(monkeypatch, fault, error) lender, waiting = s.lender, s.waiting big = _request(10, prompt_len=9) # [0, 8) needs 8 slots of each pool group, which has 4 + # A fetch lends past the committed blocks alone, so it goes into a cache without any. + fresh = _request(11, prompt_len=9) if fault == "none": assert not s.errors, s.errors - assert waiting.poll() is not None and waiting.failure is None, "the free granted nothing" + lender.progress() # the next executor iteration + assert waiting.poll() is not None and waiting.failure is None, "the round granted nothing" assert s.granted.poll() is not None and s.seen.poll() is s.seen_view - for lend in (lender.lend_read, lender.lend_write): - with pytest.raises(ValueError, match="more staging slots"): - lend(big, 0, 8) - for lease in (s.seen, s.granted, waiting): - lease.mark_arrived([np.full(len(run), lease is s.seen) for run in lease._view.groups]) + for lend, request in ((lender.lend_read, big), (lender.lend_write, fresh)): + lease = lend(request, 0, 8, **kit.HOST_NAMED) + assert lease.failure is not None and "more staging slots" in lease.failure, ( + f"{lend.__name__}: {lease.failure}" + ) lease.release() - assert s.copies == [True, False], "no copy of the marked rows into the request's pages" + for lease in (s.seen, s.granted, waiting): + arrived = [np.full(len(run), lease is s.seen) for run in lease._view.groups] + lender.return_writes([(lease, arrived, True)]) + assert s.copies == [True, False], "no copy of the arrived rows into the request's pages" lender._on_shutdown() assert s.memory.freed == [True], "the shutdown kept staging that nothing holds" return assert len(s.errors) == 1 and "planted" in s.errors[0], s.errors - free = (0, 2) if fault == "take" else (2, 0) # pool group 0 changed, pool group 1 not - assert (lender._free_slots(0), lender._free_slots(1)) == free, "no fault between the groups" + free = (0, 0) # the released read still holds its slots: no round of progress ran requests = [_request(7), _request(8)] - intervals = [lender.resume_bounds(request) for request in requests] # progress comes first + intervals = [lender.poll_resume_bounds(request) for request in requests] + assert waiting.failure is None, "poll_resume_bounds ran a round of progress" + lender.progress() # the next executor iteration assert waiting.failure is not None, "the waiting write neither became ready nor failed" stop = waiting.failure assert stop == _lender._REASON_STOPPED, stop @@ -1043,30 +1086,30 @@ def test_an_error_the_free_hook_logs_stops_the_lender(monkeypatch, fault, error) ) no_cache = _request(9) with pytest.raises(ValueError, match="has no KV cache"): - lender.resume_bounds(no_cache) + lender.poll_resume_bounds(no_cache) assert s.granted.poll() is None and s.granted.failure == stop, "a view went out after the stop" assert s.seen.poll() is s.seen_view, "a view a poll returned before the stop was taken back" for lend in (lender.lend_read, lender.lend_write): - lease = lend(no_cache, 0, 2) - assert lease.failure == stop, f"{lend.__name__} did not fail at the call: {lease.failure}" - lease.release() - with pytest.raises(ValueError): - lender.lend_read(no_cache, 2, 0) + for request, end in ((no_cache, 2), (big, 8)): + with pytest.raises(RuntimeError, match="lends no more") as raised: + lend(request, 0, end, **kit.HOST_NAMED) + assert stop in str(raised.value), f"{lend.__name__}: {raised.value}" with pytest.raises(ValueError): - lender.lend_write(no_cache, 0, 4) # past the whole blocks before the last prompt token - with pytest.raises(ValueError, match="more staging slots"): - lender.lend_write(big, 0, 8) - read = lender.lend_read(big, 0, 8) - assert read.failure == stop, f"lend_read of the range did not fail at the call: {read.failure}" - read.release() - marks = [np.ones(len(run), dtype=bool) for run in s.seen_view.groups] - s.seen.mark_arrived(marks) - assert s.copies == [True], "a mark queued a copy into the request's pages" - with pytest.raises(RuntimeError, match="twice"): - s.seen.mark_arrived(marks) + lender.lend_read(no_cache, 2, 0, **kit.HOST_NAMED) + with pytest.raises(ValueError): # past the whole blocks before the last prompt token + lender.lend_write(no_cache, 0, 4, **kit.HOST_NAMED) + arrived = [np.ones(len(run), dtype=bool) for run in s.seen_view.groups] + (bounds,) = lender.return_writes([(s.seen, arrived, True)]) + assert s.copies == [True], "a return queued a copy into the request's pages" + assert bounds is not None and bounds.max_position < bounds.min_position, ( + f"request 8: the return gave {bounds}, not an empty interval" + ) + with pytest.raises(RuntimeError, match="returned already"): + lender.return_writes([(s.seen, arrived, True)]) for lease in (s.seen, s.granted, waiting): lease.release() assert lender._open_lease_count() == 0 + lender.progress() assert (lender._free_slots(0), lender._free_slots(1)) == free, "slots returned or granted" assert lender._must_retain_memory(), "the shutdown would free the staging memory" warnings = [] @@ -1078,13 +1121,13 @@ def test_an_error_the_free_hook_logs_stops_the_lender(monkeypatch, fault, error) assert len(s.errors) == 1, f"the shutdown logged an error: {s.errors[1:]}" -@pytest.mark.parametrize("fault", ["take", "give"]) -def test_a_shutdown_right_after_the_free_hook_stopped_the_lender(monkeypatch, fault): +@pytest.mark.parametrize("error", [MemoryError, KeyError], ids=["memory_error", "key_error"]) +def test_a_shutdown_right_after_the_free_hook_stopped_the_lender(monkeypatch, error): """As above, with the manager's shutdown the first call after the free: the shutdown itself fails the waiting write, before any poll, a later poll fails a lease no poll returned the view of and returns a view one did, the staging memory stays, and the shutdown logs no error of its own.""" - s = _freed_while_a_write_waits(monkeypatch, fault) + s = _freed_while_a_write_waits(monkeypatch, "drop", error) s.lender._on_shutdown() assert s.waiting.failure == _lender._REASON_SHUT_DOWN, ( "the shutdown left the waiting write in line" @@ -1111,7 +1154,7 @@ def test_an_error_the_shutdown_logs_stops_the_lender(monkeypatch, before): s = _a_write_waits(monkeypatch, "give") struck = s.fired else: - s = _freed_while_a_write_waits(monkeypatch, "take") + s = _freed_while_a_write_waits(monkeypatch, "drop") struck = [] warned = raising_once(_lender.logger.warning, MemoryError("planted"), struck) monkeypatch.setattr(_lender.logger, "warning", warned) @@ -1123,7 +1166,7 @@ def test_an_error_the_shutdown_logs_stops_the_lender(monkeypatch, before): assert s.seen.poll() is s.seen_view, "a view a poll returned was taken back" assert not s.memory.freed, "the shutdown freed staging whose slot records may be wrong" with pytest.raises(ValueError, match="shut down"): - s.lender.resume_bounds(_request(7)) + s.lender.poll_resume_bounds(_request(7)) first = _lender._REASON_STOPPED_AT_SHUTDOWN if before == "live" else _lender._REASON_STOPPED assert s.granted.failure == first, f"not the first reason: {s.granted.failure}" @@ -1145,12 +1188,8 @@ def test_a_clean_shutdown_does_not_stop_the_lender(monkeypatch): @pytest.mark.parametrize("hook", ["free", "shutdown"]) def test_a_hook_whose_error_log_raises_still_stops_the_lender(monkeypatch, hook): - """A planted MemoryError strikes in a hook's own work, at the take in the free or at the return - in the shutdown as above, and the hook's error log line raises a KeyError, which goes out into - the manager's free or shutdown: the lender has stopped all the same, so the waiting write and a - lease no poll returned the view of fail at their next poll, and a view a poll returned keeps - returning.""" - s = _a_write_waits(monkeypatch, "take" if hook == "free" else "give") + """A free or shutdown hook whose error log line raises has still stopped the lender.""" + s = _a_write_waits(monkeypatch, "drop" if hook == "free" else "give") logged = [] log = raising_once(s.errors.append, KeyError("log line"), logged) monkeypatch.setattr(_lender.logger, "error", log) @@ -1191,10 +1230,11 @@ def _records_of(monkeypatch, kind, owner): pool_group_of=(0, 1), windows=(None, None), num_layer_groups=2, + named_groups=(True, True), tokens_per_block=1, ) ref = weakref.ref(manager) - lender = _lender.Staging(ref, layout, None, (), SlotQueue({0: 4, 1: 4}), ref) + lender = _lender._ManagerLender(ref, layout, None, (), SlotQueue({0: 4, 1: 4}), ref) cache = kv if owner == "current" else old rows = np.zeros(6, dtype=bool) rows[2:] = True @@ -1202,9 +1242,9 @@ def _records_of(monkeypatch, kind, owner): "fetch": (lender._ledger._latest_fetches, _fetch_ledger._FetchRecord(2, cache)), "delivered": ( lender._ledger._deliveries, - _fetch_ledger._Deliveries(cache, 2, [rows, rows.copy()]), + _fetch_ledger._Deliveries(cache, 2, [rows, rows.copy()], end=6), ), - "advanced": (lender._ledger._kept_prefixes, (cache, 2, 4)), + "advanced": (lender._ledger._kept_prefixes, (cache, (2, 2), 4)), "tainted": (lender._ledger._mixed_tail_caches, cache), }[kind] records[7] = record @@ -1223,7 +1263,7 @@ def test_resume_bounds_drops_the_records_of_a_cache_the_request_no_longer_holds( request, answering as without records; the same record of the current cache changes resume_bounds and stays.""" s = _records_of(monkeypatch, kind, owner) - resume_bounds = s.lender.resume_bounds(_request(7)) + resume_bounds = s.lender.poll_resume_bounds(_request(7, prompt_len=8)) # all 8 blocks if owner == "current": assert resume_bounds == CURRENT_RECORDS[kind] != NO_RECORDS, ( f"resume_bounds {resume_bounds}" @@ -1236,20 +1276,42 @@ def test_resume_bounds_drops_the_records_of_a_cache_the_request_no_longer_holds( assert 7 not in s.records, "resume_bounds kept a record of a replaced cache" +@pytest.mark.parametrize("prompt_len,usable", [(10, 10), (12, 12), (14, 12)]) +def test_resume_bounds_stop_at_the_prompt_inside_a_delivered_last_block( + monkeypatch, prompt_len, usable +): + """Three delivered 4-token blocks: a request resumes where its prompt ends inside the last.""" + monkeypatch.setattr(_ranges, "_floor_follows_history", lambda manager: True) + manager = _Owner() + kv = _cache(blocks=3) + manager.kv_cache_map = {7: kv} + layout = SimpleNamespace( + pool_groups=(0, 1), + pool_group_of=(0, 1), + windows=(None, None), + num_layer_groups=2, + named_groups=(True, True), + tokens_per_block=4, + ) + ref = weakref.ref(manager) + lender = _lender._ManagerLender(ref, layout, None, (), SlotQueue({0: 4, 1: 4}), ref) + rows = np.ones(3, dtype=bool) + lender._ledger._deliveries[7] = _fetch_ledger._Deliveries(kv, 0, [rows, rows.copy()], end=12) + assert lender.poll_resume_bounds(_request(7, prompt_len=prompt_len)) == (usable, 0) + + # -- a fetch's states ------------------------------------------------------------------------- def test_a_fetch_walks_through_the_states_its_enum_names(monkeypatch): - """Write leases into a stand-in cache, ``_fetch_state`` read after each call: the first fetch - open, delivered with its copy pending, settled once the copy completed, then replaced by the - next, which is abandoned when its lease is released before a poll returned the view; a third is - abandoned when the copy of its marks fails. A new fetch is refused while one is unsettled.""" + """Write leases into a stand-in cache walk a fetch through every state of ``_fetch_state``.""" states = _fetch_ledger._FetchState lay = SimpleNamespace( pool_groups=(0,), pool_group_of=(0,), windows=(None,), num_layer_groups=1, + named_groups=(True,), tokens_per_block=1, ) kv = SimpleNamespace( @@ -1262,17 +1324,17 @@ def test_a_fetch_walks_through_the_states_its_enum_names(monkeypatch): manager = _Owner() manager.kv_cache_map = {7: kv} ref = weakref.ref(manager) - lender = _lender.Staging(ref, lay, None, (), SlotQueue({0: 4}), ref) + lender = _lender._ManagerLender(ref, lay, None, (), SlotQueue({0: 4}), ref) landed, failure = [], {} event = SimpleNamespace(query=lambda: bool(landed)) monkeypatch.setattr( - lender, "_build_view", lambda rows, keys: SimpleNamespace(groups=rows.block_ordinals) + lender, "_build_view", lambda rows, *keys: SimpleNamespace(groups=rows.block_ordinals) ) monkeypatch.setattr(lender, "_copy_segments", lambda rows, mask: []) monkeypatch.setattr( lender, "_queue_copy", - lambda segments, to_staging: (_lender._QueuedCopy(event), failure.get("error")), + lambda segments, to_staging, after: (_lender._QueuedCopy(event), failure.get("error")), ) seen = set() @@ -1290,20 +1352,20 @@ def lent(): lender._admit(lease, {0: 2}) return lease - def marked(lease): - lease.mark_arrived([np.ones(len(run), dtype=bool) for run in lease.poll().groups]) + def returned(lease): + arrived = [np.ones(len(run), dtype=bool) for run in lease.poll().groups] + lender.return_writes([(lease, arrived, True)]) first = lent() state(first, states.AWAITING_MARKS) assert lender._ledger._unsettled_fetch_refusal(7, kv), "a fetch started while another was open" - marked(first) + returned(first) state(first, states.COPY_PENDING) assert lender._ledger._unsettled_fetch_refusal(7, kv), ( "a fetch started while another's copy was pending" ) - first.release() landed.append(True) - lender._progress() # as a lend, a poll or a resume_bounds call does first; it asks the copy again + lender.progress() # the next executor iteration's round asks the copy again state(first, states.COPY_DONE) assert lender._ledger._unsettled_fetch_refusal(7, kv) is None, ( "a settled fetch kept a new one from starting" @@ -1315,69 +1377,47 @@ def marked(lease): state(second, states.ABANDONED) third = lent() failure["error"] = "planted" - marked(third) + returned(third) state(third, states.ABANDONED) - third.release() assert seen == set(states), f"states never seen: {set(states) - seen}" -# -- the hooks the manager calls -------------------------------------------------------------- - -# The manager's module, beside the sharing package. -MANAGER_SOURCE = Path(_lender.__file__).resolve().parents[1] / "kv_cache_manager_v2.py" -# Each lender method the manager calls on its ``_sharing``: positional arguments, keywords. -HOOKS = { - "_on_free": (2, ()), - "_on_shrink": (2, ()), - "_on_shutdown": (0, ()), - "_on_caches_closed": (0, ()), - "_on_reset": (0, ()), -} - - -def hook_calls(source): - """Each method ``source`` calls on an attribute named ``_sharing``, with the positional - arguments and the keywords of each of its calls.""" - calls = {} - for node in ast.walk(ast.parse(source)): - if ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and isinstance(node.func.value, ast.Attribute) - and node.func.value.attr == "_sharing" - ): - shape = (len(node.args), tuple(k.arg for k in node.keywords)) - calls.setdefault(node.func.attr, set()).add(shape) - return calls - - -def unbound_hooks(cls, calls): - """The methods of ``calls`` that ``cls`` lacks, or whose signature cannot bind one of their - calls.""" - unbound = [] - for name, shapes in sorted(calls.items()): - try: - signature = inspect.signature(getattr(cls, name)) - for args, keywords in shapes: - signature.bind(None, *[None] * args, **dict.fromkeys(keywords)) - except (AttributeError, TypeError, ValueError): - unbound.append(name) - return unbound - - -def test_the_lender_has_every_hook_the_manager_calls(monkeypatch): - """The manager calls five lender methods on its ``_sharing``, and ``Staging`` has each, taking - the arguments the manager passes. The check finds a ``Staging`` planted without one of them, - and one planted with a hook that takes fewer arguments.""" - calls = hook_calls(MANAGER_SOURCE.read_text()) - assert calls == {name: {shape} for name, shape in HOOKS.items()}, calls - assert unbound_hooks(_lender.Staging, calls) == [] - with monkeypatch.context() as patched: - patched.delattr(_lender.Staging, "_on_reset") - assert unbound_hooks(_lender.Staging, calls) == ["_on_reset"], "a missing hook passed" - with monkeypatch.context() as patched: - patched.setattr(_lender.Staging, "_on_free", lambda self, request_id: None) - assert unbound_hooks(_lender.Staging, calls) == ["_on_free"], "a short hook passed" +# -- the manager's lender calls --------------------------------------------------------------- + + +@pytest.mark.parametrize("locked", [False, True], ids=["unlocked", "locked"]) +def test_suspending_a_request_observes_its_lender_lock(locked): + """The manager leaves a locked cache active and suspends it after the lock clears.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + + suspended = [] + cache = SimpleNamespace(id=7, is_active=True, suspend=lambda: suspended.append(True)) + manager = object.__new__(KVCacheManagerV2) + manager.kv_cache_map = {7: cache} + manager._sharing = SimpleNamespace(_is_locked=lambda candidate: locked) + request = SimpleNamespace(py_request_id=7) + + if locked: + with pytest.raises(RuntimeError, match="cannot be suspended"): + manager.suspend_request(request) + assert suspended == [] + else: + manager.suspend_request(request) + assert suspended == [True] + + +def test_a_manager_shrink_notifies_its_lender(): + """The manager forwards an in-place shrink with the same cache identity.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + + calls = [] + cache = object() + manager = object.__new__(KVCacheManagerV2) + manager._sharing = SimpleNamespace( + _on_shrink=lambda request_id, kv_cache: calls.append((request_id, kv_cache)) + ) + manager._after_shrink(7, cache) + assert calls == [(7, cache)] # -- the bytes the test kit stages ------------------------------------------------------------ @@ -1391,45 +1431,6 @@ def test_the_pieces_of_a_staged_row_differ(kit): assert len(set(pieces)) == len(pieces), "a staged row repeats a piece, so a swap would not show" -# -- a lease's axes ------------------------------------------------------------------------------ - - -def test_each_axis_of_a_lease_follows_its_own_fields_alone(): - """Each axis of a staging lease follows its own fields alone: the holder axis the failure and - whether a poll returned the view, the slot axis the ticket, the copy axis the copy. Checked - over every combination of those four fields, ticket 0 (the first a slot queue issues) and an - empty failure among them, and of four lease fields no axis reads.""" - lender, _ = staging_over([None], [0], 32) - holder, slots, copies = _lender._HolderState, _lender._SlotState, _lender._CopyState - queued = _lender._QueuedCopy(SimpleNamespace(query=lambda: False)) - recorded = {None: True, _lender._COPY_WITHOUT_RECORD: False, queued: True} - for failure, seen, ticket, copy, released, marked, runs, doomed in itertools.product( - [None, "failed", ""], - [False, True], - [None, 0, 3], - recorded, - [False, True], - [False, True], - [None, SlotRuns({0: (0, 1)})], - [None, "doomed"], - ): - lease = _lender._StagingLease(lender, "write", 7) - lease._failure, lease._view_returned, lease._ticket = failure, seen, ticket - lease._copy, lease._released, lease._marked = copy, released, marked - lease._slot_runs, lease._first_poll_failure = runs, doomed - failed = failure is not None - want = ( - holder.FAILED if failed else holder.VIEW_RETURNED if seen else holder.VIEW_PENDING, - slots.WAITING if ticket is not None else slots.NOT_WAITING, - copies.NONE_OR_RECORDED if recorded[copy] else copies.WITHOUT_RECORD, - ) - got = (lease._holder_state(), lease._slot_state(), lease._copy_state()) - assert got == want, ( - f"failure {failure}, seen {seen}, ticket {ticket}, copy {copy}, released {released}, " - f"marked {marked}, runs {runs}, doomed {doomed}: {got}" - ) - - class _Cache: """A stand-in cache a weak reference can follow.""" @@ -1438,7 +1439,7 @@ def test_a_lender_and_its_records_go_with_its_last_reference(): """With the cyclic collector off, a lender and a cache its fetch records hold go as soon as the last reference to the lender does: its ledger reaches it only through a weak reference, so no reference cycle keeps either until a collection.""" - lender, _ = staging_over([None], [0], 32) + lender, _ = lender_over([None], [0], 32) cache = _Cache() lender._ledger._kept_prefixes[7] = (cache, 0, 1) lender_ref, cache_ref = weakref.ref(lender), weakref.ref(cache) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py index 46df1d1206fe..24de9d3b1e4e 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py @@ -12,15 +12,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The staging lender's contract over real KV cache managers, driven through the public API; a few -oracles read lender internals the API does not show. Allocates device pools.""" +"""The lender's contract over real KV cache managers; a few oracles read what the API hides.""" import gc import mmap import threading +import time import weakref from types import SimpleNamespace -from typing import List, Set +from typing import List import numpy as np import pytest @@ -28,9 +28,11 @@ from utils.util import skip_pre_blackwell from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + ALL_ARRIVED, Lease, LeaseView, Lender, + MemoryKind, RegionHold, ResumeBounds, StagingOptions, @@ -53,9 +55,9 @@ SCOPE = b"lender-contract-suite" -def attach(mgr, *, tokens_per_fetch, planned_fetches=1, max_bytes=None, scope=SCOPE): +def attach(mgr, *, guaranteed_fetch_tokens, capacity_in_fetches=1, max_bytes=None, scope=SCOPE): """The public attach, as an integrator makes it.""" - options = StagingOptions(tokens_per_fetch, planned_fetches, max_bytes) + options = StagingOptions(guaranteed_fetch_tokens, capacity_in_fetches, max_bytes) return attach_lender(mgr, scope=scope, staging=options) @@ -65,16 +67,16 @@ def attach_breaking(rules): defines the method: a lender breaking exactly those rules.""" def attach_rule_breaker( - mgr, *, tokens_per_fetch, planned_fetches=1, max_bytes=None, scope=SCOPE + mgr, *, guaranteed_fetch_tokens, capacity_in_fetches=1, max_bytes=None, scope=SCOPE ): from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _fetch_ledger, _lender chosen = dict(rules(mgr) if callable(rules) else rules) on_ledger = {n: f for n, f in chosen.items() if _on_ledger(n)} lender_rules = {n: f for n, f in chosen.items() if n not in on_ledger} - cls = type("RuleBreaker", (_lender.Staging,), lender_rules) - options = StagingOptions(tokens_per_fetch, planned_fetches, max_bytes) - lender = _lender._attach_staging(mgr, scope=scope, staging=options, cls=cls) + cls = type("RuleBreaker", (_lender._ManagerLender,), lender_rules) + options = StagingOptions(guaranteed_fetch_tokens, capacity_in_fetches, max_bytes) + lender = _lender._attach_lender(mgr, scope=scope, staging=options, cls=cls) if on_ledger: # A ledger of the rule-breaking class in place of the real one, before it records # anything; its rules reach the lender through _lender_ref. @@ -93,27 +95,16 @@ def _on_ledger(name): return not name.startswith("__") and name in vars(_fetch_ledger._FetchLedger) -def ready_view(lease, mgr, tries=3): - """The lease's view once the copies queued so far have run; fails if it never comes.""" - for _ in range(tries): - mgr._stream.synchronize() - view = lease.poll() - if view is not None: - return view - raise AssertionError(f"the lease never became ready: {lease.failure}") - - def by_group(view): - return {run.layer_group_index: run.block_ordinals.tolist() for run in view.groups} + return {run.layer_group_index: run.page_ordinals.tolist() for run in view.groups} -def one_fetch_bytes(kit, mgr, tokens_per_fetch): - """Bytes one fetch of ``tokens_per_fetch`` tokens needs, from the device pages: every block for full - attention, a window's blocks for a window (no sinks here).""" +def one_fetch_bytes(kit, mgr, guaranteed_fetch_tokens): + """Bytes one fetch of ``guaranteed_fetch_tokens`` tokens takes: a window's blocks for a window.""" dev = kit.DevicePages(mgr) total = 0 for lg, window in enumerate(kit.windows(mgr)): - blocks = -(-tokens_per_fetch // TPB) + blocks = -(-guaranteed_fetch_tokens // TPB) if window is not None: blocks = min(blocks, -(-window // TPB)) total += blocks * dev.page_bytes(lg) @@ -126,13 +117,11 @@ def check_staged(kit, mgr, lender, request, view): dev = kit.DevicePages(mgr) rows = 0 for run in view.groups: - region = lender.regions[run.region_index] - assert run.layer_group_index in lender.regions[run.region_index].layer_group_indices + region, addresses = kit.staged_slots(lender, run) assert region.slot_bytes == dev.page_bytes(run.layer_group_index) slots = kit.pages(kv, run.layer_group_index) - for address, ordinal in zip(run.addresses.tolist(), run.block_ordinals.tolist()): + for address, ordinal in zip(addresses, run.page_ordinals.tolist()): assert region.address <= address < region.address + region.nbytes - assert (address - region.address) % region.slot_bytes == 0 device = kit.digest(dev.read(run.layer_group_index, slots[ordinal])) assert kit.digest(kit.host_bytes(address, region.slot_bytes)) == device rows += 1 @@ -143,7 +132,7 @@ def state_of(kit, mgr, lender, request): """Capacity, history, committed tokens and resume_bounds of the request, or what stands for none.""" kv = kit.kv(mgr, request) try: - resume_bounds = lender.resume_bounds(request) + resume_bounds = lender.poll_resume_bounds(request) except ValueError: resume_bounds = "no cache" if kv is None: @@ -151,6 +140,35 @@ def state_of(kit, mgr, lender, request): return (kv.capacity, kv.history_length, kv.num_committed_tokens), resume_bounds +def return_and_settle(lender, mgr, request, lease, arrived): + """Return ``lease`` with ``arrived``, run its copy and a round of progress: the bounds then.""" + (given,) = lender.return_writes([(lease, arrived, True)]) + mgr._stream.synchronize() + lender.progress() + bounds = lender.poll_resume_bounds(request) + assert given is None or given == bounds, f"the return gave {given}, then {bounds}" + return bounds + + +def check_names_nothing(kit, lender, words, refused=None, in_place=None, mgr=None, source=None): + """Nothing staged; named lends refuse, while a cache-backed source lends device rows.""" + assert lender.regions and kit.staging_regions(lender) == () + assert kit.staging_memory(lender.regions) == [] + missing_id = SOURCE if source is None else SOURCE + 10_000 + request = kit.make_request(missing_id, PROMPT) # never admitted: no cache + for lend in (lender.lend_read, lender.lend_write): + for keywords in refused or (kit.HOST_NAMED, kit.DEVICE_NAMED): + message = refusal(lambda: lend(request, 0, TPB, **keywords)) + assert all(word in message for word in words), message + lease = lender.lend_read(request, 0, TPB, **(in_place or kit.DEVICE_ALL)) + assert lease.failure == f"request {missing_id} has no KV cache", lease.failure + lease.release() + if source is not None: + lease = lender.lend_read(source, 0, TPB, **(in_place or kit.DEVICE_ALL)) + assert kit.ready_view(lease, mgr).num_rows > 0, lease.failure + lease.release() + + # -- attach ----------------------------------------------------------------------------------- @@ -161,18 +179,19 @@ def test_a_manager_takes_one_lender_for_its_life(real_manager): attach_lender(mgr, scope=scope, staging=StagingOptions(END)) with pytest.raises(TypeError, match="staging must be StagingOptions"): attach_lender(mgr, scope=SCOPE, staging=END) - lender = attach(mgr, tokens_per_fetch=END) # the refused calls attached nothing + lender = attach(mgr, guaranteed_fetch_tokens=END) # the refused calls attached nothing assert isinstance(lender, Lender) with pytest.raises(ValueError, match="already attached"): - attach(mgr, tokens_per_fetch=END) + attach(mgr, guaranteed_fetch_tokens=END) -def test_recurrent_state_is_refused(kit, hybrid_manager): - before = len(kit.retained()) +def test_a_manager_holding_recurrent_state_gets_a_lender_naming_nothing(kit, hybrid_manager): + """No name covers recurrent state, which the facade reads from the manager's kind.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + with hybrid_manager() as mgr: - with pytest.raises(ValueError, match="recurrent"): - attach(mgr, tokens_per_fetch=256) - assert len(kit.retained()) == before, "a refused attach allocated staging" + assert _manager.model_has_recurrent_state(mgr) + check_names_nothing(kit, attach(mgr, guaranteed_fetch_tokens=256), ["recurrent state"]) def sparse_layer(monkeypatch, index): @@ -225,7 +244,7 @@ def test_a_layer_group_of_sparse_buffers_is_refused(kit, host_tier_manager, monk assert sparse and (0 in sparse) == (case == "first"), ( f"sparse buffers in layer groups {sparse}: the check proves nothing" ) - message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) + message = refusal(lambda: attach(mgr, guaranteed_fetch_tokens=END)) assert f"layer groups {sparse} hold sparse buffers" in message assert "_sharing" not in vars(mgr) assert len(kit.retained()) == before, "a refused attach kept something" @@ -273,10 +292,7 @@ def _defaulting_a_missing_mapping(manager): # reading of the manager's mapping that would let one of them through. REFUSED = { "context_parallel": ( - [ - (lambda: stand_in("HELIX"), ValueError, "context parallelism"), - (lambda: stand_in("ULYSSES"), ValueError, "context parallelism"), - ], + [(lambda: stand_in("ULYSSES"), ValueError, "context parallelism")], _without_parallelism, ), "not_v2_or_unmapped": ( @@ -299,7 +315,7 @@ def tripwire(manager): before = len(_lender._retained()) for make, kind, words in REFUSED[case][0]: staging = make() - assert words in refusal(lambda: attach(staging, tokens_per_fetch=END), kind) + assert words in refusal(lambda: attach(staging, guaranteed_fetch_tokens=END), kind) assert "_sharing" not in vars(staging) assert len(_lender._retained()) == before, "a refused attach allocated staging" @@ -316,9 +332,9 @@ def raising(*args, **kwargs): with monkeypatch.context() as patched: patched.setattr(_lender, "logger", SimpleNamespace(info=raising)) with pytest.raises(RuntimeError, match="planted"): - attach(mgr, tokens_per_fetch=END) + attach(mgr, guaranteed_fetch_tokens=END) assert "_sharing" not in vars(mgr), "a failed attach left its lender installed" - assert isinstance(attach(mgr, tokens_per_fetch=END), Lender) + assert isinstance(attach(mgr, guaranteed_fetch_tokens=END), Lender) def pipeline_stage(rank): @@ -329,15 +345,21 @@ def pipeline_stage(rank): @pytest.mark.parametrize("rank", [0, 1], ids=["first_stage", "last_stage"]) -def test_staging_refuses_a_pipeline_stage(kit, real_manager, rank, attach=attach): - """Each stage's staging slots and windows are its own, so one call could raise on one stage and - lend on another: staging refuses every stage at the attach, which then changed nothing.""" +def test_each_pipeline_stage_attaches_a_lender_of_its_own(kit, real_manager, rank, attach=attach): + """Each stage's lender lends that stage's layers under their global ids.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager + with real_manager(mapping=pipeline_stage(rank), num_layers=4, windows=[256, WINDOW]) as mgr: - assert mgr.mapping.pp_size == 2 and len(mgr.pp_layers) == 2 - before = len(kit.retained()) - message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) - assert "pipeline parallelism" in message, message - assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + assert mgr.mapping.pp_size == 2 and len(mgr.pp_layers) == 2, mgr.pp_layers + lender = attach(mgr, guaranteed_fetch_tokens=END) + assert vars(mgr)["_sharing"] is lender + assert not _manager.model_has_recurrent_state(mgr) + layers = sorted(layer for group in lender._layout.layers for layer in group) + assert layers == [int(layer) for layer in mgr.pp_layers] == [2 * rank, 2 * rank + 1] + source = kit.published(mgr, SOURCE, PROMPT) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + kit.ready_view(lease, mgr) + lease.release() def unpaired_draft(mgr) -> None: @@ -352,40 +374,117 @@ def unpaired_draft(mgr) -> None: @pytest.mark.parametrize("case", list(COMMITS_NOTHING)) -def test_staging_refuses_a_manager_that_commits_no_blocks(kit, real_manager, case, attach=attach): - """Staging publishes committed blocks only, so a manager that commits none is refused at the - attach, which then changed nothing.""" +def test_a_manager_that_commits_no_blocks_gets_a_lender_naming_nothing( + kit, real_manager, case, attach=attach +): + """A name covers committed blocks only.""" manager_kwargs, adjust = COMMITS_NOTHING[case] with real_manager(**manager_kwargs) as mgr: if adjust is not None: adjust(mgr) - before = len(kit.retained()) - message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) - assert "block reuse" in message, message - assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + check_names_nothing(kit, attach(mgr, guaranteed_fetch_tokens=END), ["commits no blocks"]) + + +@pytest.mark.parametrize("case", ["block_reuse_off", "draft"]) +def test_mamba_final_context_saves_drop_plan_without_reusable_blocks(case): + """A final Mamba context chunk finishes its conversation when no blocks can commit.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import BlockReusePolicy + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + saved = [] + resizes = [] + request = SimpleNamespace( + py_request_id=7, + is_dummy_request=False, + context_current_position=END, + context_remaining_length=0, + is_last_context_chunk=True, + expect_snapshot_points=[], + ) + cache = SimpleNamespace( + is_active=True, + resize=lambda capacity, history: resizes.append((capacity, history)) or True, + enable_swa_scratch_reuse=True, + ) + manager = object.__new__(MambaHybridCacheManagerV2) + manager.enable_block_reuse = case != "block_reuse_off" + manager.is_draft = case == "draft" + manager.block_reuse_policy = BlockReusePolicy.ALL_REUSABLE + manager._sharing = None + manager.kv_cache_map = {request.py_request_id: cache} + manager.conversation_manager = SimpleNamespace( + save_drop_plan=lambda req, kv: saved.append((req, kv)) + ) + batch = ScheduledRequests() + batch.append_context_request(request) + + manager.update_context_resources(batch) + + assert resizes == [(None, END)] + assert saved == [(request, cache)] + assert cache.enable_swa_scratch_reuse is False + + +def test_helix_gets_past_the_parallelism_check(monkeypatch): + """A Helix manager reaches the layout, where a tripwire stops the attach.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + def tripwire(manager): + raise AssertionError(PAST_THE_CHECKS) + + monkeypatch.setattr(_lender, "derive_layout", tripwire) + manager = stand_in("HELIX") + with pytest.raises(AssertionError, match=PAST_THE_CHECKS): + attach(manager, guaranteed_fetch_tokens=END) + assert "_sharing" not in vars(manager) -@pytest.mark.parametrize("capacity_only", [False, True]) -def test_staging_refuses_a_manager_with_a_kv_cache_connector( +@pytest.mark.parametrize("capacity_only", [False, True], ids=["loads", "capacity_only"]) +def test_a_lender_beside_a_kv_cache_connector_that_loads_stages_nothing( kit, real_manager, capacity_only, attach=attach ): - """A KV cache connector serves a request's prefix at its first context chunk, measured from the - committed tokens, so it cannot lower a history a windowed fetch moved, and its loads run off the - manager's stream. Staging refuses either transfer or capacity-only connectors at the attach, - which then changed nothing. A stand-in exposing the connector mode suffices here.""" + """A loading connector writes pages off the manager's stream; a capacity-only one does not.""" connector = SimpleNamespace(capacity_only=capacity_only) with real_manager(windows=[WINDOW, 256], kv_connector_manager=connector) as mgr: assert mgr.kv_connector_manager is connector - before = len(kit.retained()) - message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) - assert "KV cache connector" in message, message - assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + lender = attach(mgr, guaranteed_fetch_tokens=END) + device = tuple(r for r in lender.regions if r.memory_kind is MemoryKind.DEVICE) + assert device and lender.regions == kit.staging_regions(lender) + device + if capacity_only: + assert kit.staging_regions(lender), "a connector transferring nothing took staging away" + else: + words = ["stages nothing", "KV cache connector"] + check_names_nothing(kit, lender, words, [kit.HOST_NAMED], kit.DEVICE_NAMED) -def test_lender_refuses_configured_lossy_cold_page_quantization_before_allocating( - kit, real_manager, monkeypatch, attach=attach +def test_a_lender_beside_a_kv_cache_connector_lends_in_place_with_pool_rebalancing_configured( + kit, real_manager, attach=attach ): - from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + """The rebalance flag leaves the lender's device regions to it.""" + connector = SimpleNamespace( + capacity_only=False, + prefix_reservations_enabled=False, + get_prefix_reservation=lambda request: None, + has_pending_load=lambda request: False, + release_unstarted_prefix_loads=lambda request: None, + release_prefix_reservation=lambda request: None, + reset_request_state=lambda request: None, + ) + with real_manager(enable_kv_pool_rebalance=True, kv_connector_manager=connector) as mgr: + lender = attach(mgr, guaranteed_fetch_tokens=END) + assert any(region.memory_kind is MemoryKind.DEVICE for region in lender.regions) + source = kit.published(mgr, SOURCE, PROMPT) + check_names_nothing( + kit, lender, ["stages nothing"], [kit.HOST_NAMED], mgr=mgr, source=source + ) + mgr.free_resources(source) + + +def test_a_lossy_cold_page_codec_names_nothing_and_lends_in_place(kit, real_manager, attach=attach): + """A lossy codec brings a block back from a colder tier with other bytes.""" from tensorrt_llm.llmapi.llm_args import ColdPageQuantizationCompressionConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import create_default_kv_cache_cold_page_codec @@ -395,17 +494,11 @@ def test_lender_refuses_configured_lossy_cold_page_quantization_before_allocatin create_cold_page_codec=lambda *args, **kwargs: create_default_kv_cache_cold_page_codec(), ) with real_manager(cold_page_codec_provider=provider) as mgr: - assert mgr._has_lossy_cold_page_codec - before = len(kit.retained()) - - def tripwire(*args, **kwargs): - raise AssertionError("a refused codec must not derive layouts or allocate memory") - - monkeypatch.setattr(_lender, "derive_layout", tripwire) - monkeypatch.setattr(_lender, "_allocate_retained", tripwire) - with pytest.raises(ValueError, match="lossy.*cold-page codec"): - attach(mgr, tokens_per_fetch=END) - assert "_sharing" not in vars(mgr) and len(kit.retained()) == before + assert mgr._has_lossy_cold_page_codec, "the manager took the codec for a lossless one" + lender = attach(mgr, guaranteed_fetch_tokens=END) + source = kit.published(mgr, SOURCE, PROMPT) + check_names_nothing(kit, lender, ["lossy"], mgr=mgr, source=source) + mgr.free_resources(source) @pytest.mark.parametrize("custom_lossless", [False, True]) @@ -421,7 +514,7 @@ def test_lender_accepts_default_and_custom_lossless_cold_page_codecs( ) with real_manager(cold_page_codec_provider=provider) as mgr: assert not mgr._has_lossy_cold_page_codec - assert isinstance(attach(mgr, tokens_per_fetch=END), Lender) + assert isinstance(attach(mgr, guaranteed_fetch_tokens=END), Lender) def eagle3(): @@ -484,7 +577,7 @@ def ngram(): spec_config=ngram(), is_draft=True, joint_kv_cache_reuse=True ), } -READ_AHEAD_WORDS = "past their end" +READ_AHEAD_WORDS = "prompt tokens past a block's end" OVER_REFUSED = "the attach refused it" @@ -496,29 +589,21 @@ def accepted(call): raise AssertionError(f"{OVER_REFUSED}: {error}") from error -def check_a_draft_reading_ahead_is_refused(kit, mgr, attach, ahead): - """Staging refuses ``mgr``, built for a draft reading ``ahead`` prompt tokens past a position, - and the refused attach changed nothing.""" - before = len(kit.retained()) - message = refusal(lambda: attach(mgr, tokens_per_fetch=END)) - assert READ_AHEAD_WORDS in message and str(ahead) in message, message - assert "_sharing" not in vars(mgr) and len(kit.retained()) == before - - @pytest.mark.parametrize("case", list(READS_AHEAD)) -def test_staging_refuses_a_manager_built_for_a_draft_reading_past_a_block( +def test_a_manager_built_for_a_draft_reading_past_a_block_gets_a_lender_naming_nothing( kit, real_manager, case, attach=attach ): make, ahead = READS_AHEAD[case] with real_manager(**make()) as mgr: assert mgr.draft_prompt_lookahead == ahead assert mgr.reuse_match_backoff == ahead - check_a_draft_reading_ahead_is_refused(kit, mgr, attach, ahead) + lender = attach(mgr, guaranteed_fetch_tokens=END) + check_names_nothing(kit, lender, [READ_AHEAD_WORDS, str(ahead)]) @skip_pre_blackwell @pytest.mark.parametrize("fp8_ds_mla", [False, True], ids=["fp8", "fp8_ds_mla"]) -def test_staging_refuses_a_deepseek_v4_manager_whose_mtp_layers_read_ahead( +def test_a_deepseek_v4_manager_whose_mtp_layers_read_ahead_gets_a_lender_naming_nothing( kit, deepseek_v4_manager, fp8_ds_mla ): from tensorrt_llm._torch.speculative import draft_prompt_lookahead @@ -529,7 +614,8 @@ def test_staging_refuses_a_deepseek_v4_manager_whose_mtp_layers_read_ahead( with deepseek_v4_manager(spec_config=spec_config, fp8_ds_mla=fp8_ds_mla) as mgr: assert mgr.draft_prompt_lookahead == 1 assert mgr.reuse_match_backoff == 0 - check_a_draft_reading_ahead_is_refused(kit, mgr, attach, 1) + lender = attach(mgr, guaranteed_fetch_tokens=END) + check_names_nothing(kit, lender, [READ_AHEAD_WORDS, "1"]) @pytest.mark.parametrize("case", list(READS_NONE)) @@ -539,83 +625,86 @@ def test_staging_accepts_a_manager_whose_blocks_read_no_token_past_their_end( with real_manager(**READS_NONE[case]()) as mgr: assert mgr.draft_prompt_lookahead == 0 assert mgr.reuse_match_backoff == 0 - assert isinstance(accepted(lambda: attach(mgr, tokens_per_fetch=END)), Lender) + assert isinstance(accepted(lambda: attach(mgr, guaranteed_fetch_tokens=END)), Lender) # -- sizing ----------------------------------------------------------------------------------- -# planned_fetches, max_bytes in fetches (None for no cap), and the fetches staging then holds at once +# capacity_in_fetches, max_bytes in fetches (None: no cap), and the fetches staging holds at once SIZING = { - "planned_fetches": (2, None, 2), + "capacity_in_fetches": (2, None, 2), "max_bytes_of_one_fetch": (4, 1, 1), "max_bytes_of_two_and_a_half_fetches": (4, 2.5, 2), } @pytest.mark.parametrize("case", list(SIZING)) -def test_staging_holds_planned_fetches_fetches_at_once_capped_by_max_bytes(kit, real_manager, case): - planned_fetches, cap, holding = SIZING[case] +def test_staging_holds_capacity_in_fetches_fetches_at_once_capped_by_max_bytes( + kit, real_manager, case +): + capacity, cap, holding = SIZING[case] with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) one = one_fetch_bytes(kit, mgr, END) with pytest.raises(ValueError, match="one fetch"): - attach(mgr, tokens_per_fetch=END, planned_fetches=4, max_bytes=one - 1) + attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=4, max_bytes=one - 1) max_bytes = None if cap is None else int(cap * one) lender = attach( - mgr, tokens_per_fetch=END, planned_fetches=planned_fetches, max_bytes=max_bytes + mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=capacity, max_bytes=max_bytes ) - (region,) = lender.regions + (region,) = kit.staging_regions(lender) assert region.slot_bytes == kit.DevicePages(mgr).page_bytes(0) assert region.nbytes == region.slot_count * region.slot_bytes assert holding * one <= region.nbytes <= (max_bytes or holding * one) - leases = [lender.lend_read(source, 0, END) for _ in range(holding + 1)] + leases = [lender.lend_read(source, 0, END, **kit.HOST_NAMED) for _ in range(holding + 1)] for lease in leases[:holding]: - ready_view(lease, mgr) + kit.ready_view(lease, mgr) assert leases[-1].poll() is None and leases[-1].failure is None, "one fetch too many" leases[0].release() - check_staged(kit, mgr, lender, source, ready_view(leases[-1], mgr)) + check_staged(kit, mgr, lender, source, kit.ready_view(leases[-1], mgr)) for lease in leases[1:]: lease.release() @pytest.mark.parametrize("num_layers", [2, 3], ids=["one_pool_group", "two_pool_groups"]) -def test_a_windowed_range_of_tokens_per_fetch_always_fits(kit, real_manager, num_layers): +def test_a_windowed_range_of_guaranteed_fetch_tokens_always_fits(kit, real_manager, num_layers): with real_manager(windows=[WINDOW, 256], num_layers=num_layers) as mgr: source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) target = kit.admitted(mgr, TARGET, OTHER_WINDOWED_PROMPT) - lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + lender = attach(mgr, guaranteed_fetch_tokens=WINDOWED_END) dev = kit.DevicePages(mgr) rows = {} for lg, (window, g) in enumerate(zip(kit.windows(mgr), kit.pool_group_of(mgr))): rows[g] = rows.get(g, 0) + (WINDOW // TPB if window else WINDOWED_END // TPB) groups = kit.pool_group_ids(mgr) - assert len(lender.regions) == len(groups) == num_layers - 1 + regions = kit.staging_regions(lender) + assert len(regions) == len(groups) == num_layers - 1 group_of = kit.pool_group_of(mgr) - for g, region in zip(groups, lender.regions): + for g, region in zip(groups, regions): assert (region.slot_count, region.slot_bytes) == (rows[g], dev.group_page_bytes(g)) assert region.nbytes == region.slot_count * region.slot_bytes assert region.layer_group_indices == tuple( lg for lg, pg in enumerate(group_of) if pg == g ) # The regions' layer groups are disjoint and cover every layer group. - held = sorted(lg for region in lender.regions for lg in region.layer_group_indices) + held = sorted(lg for region in regions for lg in region.layer_group_indices) assert held == list(range(len(group_of))) - read = lender.lend_read(source, 0, WINDOWED_END) # takes every slot - assert check_staged(kit, mgr, lender, source, ready_view(read, mgr)) == sum(rows.values()) + read = lender.lend_read(source, 0, WINDOWED_END, **kit.HOST_NAMED) # takes every slot + view = kit.ready_view(read, mgr) + assert check_staged(kit, mgr, lender, source, view) == sum(rows.values()) read.release() - write = lender.lend_write(target, 0, WINDOWED_END) - view = ready_view(write, mgr) - assert view.num_block_entries == sum(rows.values()) - write.mark_arrived(view.row_masks()) - write.release() + write = lender.lend_write(target, 0, WINDOWED_END, **kit.HOST_NAMED) + view = kit.ready_view(write, mgr) + assert view.num_rows == sum(rows.values()) + kit.return_write(lender, write, view.row_masks()) def test_the_staging_memory_is_the_parts_and_nothing_more(kit, real_manager, attach=attach): # Three layers: two pool groups, so two regions. with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: - lender = attach(mgr, tokens_per_fetch=WINDOWED_END, planned_fetches=2) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=WINDOWED_END, capacity_in_fetches=2) + regions = kit.staging_regions(lender) assert len(regions) == len(kit.pool_group_ids(mgr)) == 2 (memory,) = kit.staging_memory(regions) assert memory.address <= regions[0].address < memory.address + memory.nbytes @@ -651,7 +740,8 @@ def test_the_staging_memory_is_pinned_at_its_size_and_freed_once(kit, real_manag spy = FreeSpy(_lender.cudart) monkeypatch.setattr(_lender, "cudart", spy) with real_manager() as mgr: - regions = attach(mgr, tokens_per_fetch=END, planned_fetches=3).regions + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=3) + regions = kit.staging_regions(lender) base, total = regions[0].address, sum(p.nbytes for p in regions) whole_pages = -(-total // mmap.PAGESIZE) * mmap.PAGESIZE assert 1 << (total - 1).bit_length() > whole_pages, "a power of two would pass unseen" @@ -722,10 +812,10 @@ def test_pageable_staging_relays_bytes_and_frees_nothing_page_locked( source = kit.published(mgr_a, SOURCE, PROMPT) target = kit.admitted(mgr_b, TARGET, PROMPT) lender_a, lender_b = ( - attach(mgr_a, tokens_per_fetch=END), - attach(mgr_b, tokens_per_fetch=END), + attach(mgr_a, guaranteed_fetch_tokens=END), + attach(mgr_b, guaranteed_fetch_tokens=END), ) - regions = lender_a.regions + lender_b.regions + regions = kit.staging_regions(lender_a) + kit.staging_regions(lender_b) locked = [kit.pinned_range(region.address) for region in regions] assert locked == [None] * len(regions), f"pageable staging was page-locked: {locked}" # A pageable buffer freed early can come back to the next lender at the same address. @@ -740,20 +830,18 @@ def test_pageable_staging_relays_bytes_and_frees_nothing_page_locked( assert kept is not None and kept.ctypes.data == memory.address, ( "the pageable staging is not an array the lender keeps" ) - publish = lender_a.lend_read(source, 0, END) - publish_view = ready_view(publish, mgr_a) - fetch = lender_b.lend_write(target, 0, END) + publish = lender_a.lend_read(source, 0, END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) + fetch = lender_b.lend_write(target, 0, END, **kit.HOST_NAMED) view = fetch.poll() assert view is not None, f"the write lease was not ready: {fetch.failure}" kit.fill_sentinel(mgr_b, target) - fetch.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) - mgr_b._stream.synchronize() - assert lender_b.resume_bounds(target) == (END, 0) + masks = kit.relay(lender_a, publish_view, lender_b, view) + assert return_and_settle(lender_b, mgr_b, target, fetch, masks) == (END, 0) for ordinal in range(BLOCKS): got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) assert got == kit.digest(kit.page(mgr_a, source, 0, ordinal)), "a fetched page differs" - for lease in (fetch, publish): - lease.release() + publish.release() for mgr in (mgr_a, mgr_b): mgr.shutdown() assert spy.freed == [], "pageable staging was freed as page-locked memory" @@ -769,13 +857,13 @@ def test_pageable_staging_relays_bytes_and_frees_nothing_page_locked( def test_a_publish_lends_the_committed_pages_under_their_reuse_keys(kit, real_manager, tokens): with real_manager() as mgr: source = kit.published(mgr, SOURCE, tokens) - lender = attach(mgr, tokens_per_fetch=END) - lease = lender.lend_read(source, 0, END) + lender = attach(mgr, guaranteed_fetch_tokens=END) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) assert isinstance(lease, Lease) - view = ready_view(lease, mgr) + view = kit.ready_view(lease, mgr) assert isinstance(view, LeaseView) and lease.poll() is view, "one view, every poll" (run,) = view.groups - assert run.block_ordinals.tolist() == list(range(BLOCKS)) and run.region_index == 0 + assert run.page_ordinals.tolist() == list(range(BLOCKS)) and run.region_indices == (0,) assert check_staged(kit, mgr, lender, source, view) == BLOCKS kv = kit.kv(mgr, source) assert [bytes(name[16:48]) for name in run.object_keys] == kit.chain_keys(kv, tokens)[ @@ -801,9 +889,9 @@ def test_a_multimodal_publish_is_named_by_the_keys_the_manager_commits(kit, real second = kit.published( mgr, OTHER, PROMPT, multimodal_hashes=[list(range(8, 0, -1))], **image ) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) - leases = [lender.lend_read(r, 0, END) for r in (first, second)] - names = [ready_view(lease, mgr).groups[0].object_keys for lease in leases] + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) + leases = [lender.lend_read(r, 0, END, **kit.HOST_NAMED) for r in (first, second)] + names = [kit.ready_view(lease, mgr).groups[0].object_keys for lease in leases] keys = [[bytes(n[16:48]) for n in run_names] for run_names in names] kv = kit.kv(mgr, first) augmented = list(mgr._augment_tokens_for_block_reuse(PROMPT, first)) @@ -825,10 +913,10 @@ def test_a_publish_is_named_under_the_request_s_reuse_scope(kit, real_manager, s with real_manager() as mgr: plain = kit.published(mgr, SOURCE, PROMPT) scoped = kit.published(mgr, OTHER, PROMPT, **scope) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) - leases = [lender.lend_read(r, 0, END) for r in (plain, scoped)] + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) + leases = [lender.lend_read(r, 0, END, **kit.HOST_NAMED) for r in (plain, scoped)] keys = [ - [bytes(n[16:48]) for n in ready_view(lease, mgr).groups[0].object_keys] + [bytes(n[16:48]) for n in kit.ready_view(lease, mgr).groups[0].object_keys] for lease in leases ] kv = kit.kv(mgr, scoped) @@ -845,24 +933,24 @@ def test_a_publish_is_named_under_the_request_s_reuse_scope(kit, real_manager, s def test_a_windowed_publish_lends_only_what_the_window_at_its_end_reads(kit, real_manager): with real_manager(windows=[WINDOW, 256]) as mgr: source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) - lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + lender = attach(mgr, guaranteed_fetch_tokens=WINDOWED_END) windows = kit.windows(mgr) sliding, full = windows.index(WINDOW), windows.index(None) - lease = lender.lend_read(source, 0, WINDOWED_END) - view = ready_view(lease, mgr) + lease = lender.lend_read(source, 0, WINDOWED_END, **kit.HOST_NAMED) + view = kit.ready_view(lease, mgr) beg, end = kit.stale_blocks(mgr, sliding, WINDOWED_END) assert end > beg, "the window must leave early blocks behind" blocks = range(WINDOWED_END // TPB) in_window = [o for o in blocks if not beg <= o < end] assert by_group(view) == {full: list(blocks), sliding: in_window} - addresses = np.concatenate([run.addresses for run in view.groups]) - assert len(set(addresses.tolist())) == len(addresses), "every row has its own slot" + addresses = [a for run in view.groups for a in kit.staged_slots(lender, run)[1]] + assert len(set(addresses)) == len(addresses), "every row has its own slot" assert check_staged(kit, mgr, lender, source, view) == len(addresses) lease.release() # The request's own window has since passed what a window at 96 reads: none is lent. - early = lender.lend_read(source, 0, END) - rows = by_group(ready_view(early, mgr)) + early = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + rows = by_group(kit.ready_view(early, mgr)) assert rows[full] == list(range(BLOCKS)) assert rows.get(sliding, []) == [] early.release() @@ -871,12 +959,12 @@ def test_a_windowed_publish_lends_only_what_the_window_at_its_end_reads(kit, rea def test_a_range_without_rows_is_ready_at_once_and_waits_for_no_one(kit, real_manager): with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - holder = lender.lend_read(source, 0, END) # every slot - ready_view(holder, mgr) - empty = lender.lend_read(source, END, END) + lender = attach(mgr, guaranteed_fetch_tokens=END) + holder = lender.lend_read(source, 0, END, **kit.HOST_NAMED) # every slot + kit.ready_view(holder, mgr) + empty = lender.lend_read(source, END, END, **kit.HOST_NAMED) view = empty.poll() - assert view is not None and view.num_block_entries == 0 + assert view is not None and view.num_rows == 0 assert len(view.groups) == kit.num_layer_groups(mgr) empty.release() holder.release() @@ -888,16 +976,73 @@ def test_a_range_without_rows_is_ready_at_once_and_waits_for_no_one(kit, real_ma def test_a_publish_is_ready_only_once_its_copy_ran(kit, real_manager, attach=attach): with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) + lender = attach(mgr, guaranteed_fetch_tokens=END) with kit.held_stream(mgr._stream) as gate: - lease = lender.lend_read(source, 0, END) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) assert lease.poll() is None, "ready before its copy into staging ran" + lender.progress() assert lease.poll() is None and lease.failure is None gate.open() - check_staged(kit, mgr, lender, source, ready_view(lease, mgr)) + check_staged(kit, mgr, lender, source, kit.ready_view(lease, mgr)) + lease.release() + + +def test_a_publish_is_ready_only_once_its_ready_event_completed(kit, real_manager): + """A publish's copy into staging waits for its ``ready_event``.""" + with real_manager() as mgr: + source = kit.published(mgr, SOURCE, PROMPT) + lender = attach(mgr, guaranteed_fetch_tokens=END) + elsewhere = torch.cuda.Stream() + with kit.held_stream(elsewhere, strict=False): + ready = torch.cuda.Event() + ready.record(elsewhere) + lease = lender.lend_read(source, 0, END, **{**kit.HOST_NAMED, "ready_event": ready}) + # Long enough for a copy queued without waiting for the event to have run. + deadline = time.monotonic() + 0.5 + while not mgr._stream.query() and time.monotonic() < deadline: + time.sleep(0.01) + for _ in range(3): + lender.progress() + assert lease.poll() is None, "ready before its ready event completed" + assert lease.failure is None, lease.failure + check_staged(kit, mgr, lender, source, kit.ready_view(lease, mgr)) lease.release() +def test_a_fetch_s_copy_into_the_pages_waits_for_its_ready_event(kit, real_manager): + """A fetch's view is ready at once; its copy into the pages waits for its ``ready_event``.""" + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, OTHER_PROMPT) + lender = attach(mgr, guaranteed_fetch_tokens=END) + elsewhere = torch.cuda.Stream() + with kit.held_stream(elsewhere, strict=False): + ready = torch.cuda.Event() + ready.record(elsewhere) + lease = lender.lend_write(target, 0, END, **{**kit.HOST_NAMED, "ready_event": ready}) + view = lease.poll() + assert view is not None, lease.failure + kit.stage(lender, view, 0x5A) + settled = kit.return_write(lender, lease, ALL_ARRIVED) + assert settled is None, "settled before its ready event" + # Long enough for a copy queued without waiting for the event to have run. + deadline = time.monotonic() + 0.5 + while not mgr._stream.query() and time.monotonic() < deadline: + time.sleep(0.01) + for _ in range(3): + lender.progress() + assert lender.poll_resume_bounds(target) is None, "settled before its ready event" + assert kit.settled_bounds(lender, mgr, target).max_position == END + kv = kit.kv(mgr, target) + dev = kit.DevicePages(mgr) + for run in view.groups: + lg = run.layer_group_index + slots = kit.pages(kv, lg) + for ordinal in run.page_ordinals.tolist(): + got = kit.digest(dev.read(lg, slots[ordinal])) + staged = kit.digest(kit.staged_row(0x5A, lg, ordinal, dev.page_bytes(lg))) + assert got == staged, f"layer group {lg} block {ordinal} holds other bytes" + + # -- no slot reused while a copy or a backend may touch it ----------------------------------- @@ -907,56 +1052,60 @@ def test_a_slot_returns_only_after_the_copy_into_it_ran(kit, real_manager, attac target = kit.admitted(mgr, TARGET, OTHER_PROMPT) # Grown ahead, so the held section only lends. assert mgr._resize_for_served_prefix(target, kit.kv(mgr, target), 0, END) - lender = attach(mgr, tokens_per_fetch=END) # one fetch: three slots + lender = attach(mgr, guaranteed_fetch_tokens=END) # one fetch: three slots with kit.held_stream(mgr._stream) as gate: - read = lender.lend_read(source, 0, END) # every slot, its copy held + read = lender.lend_read(source, 0, END, **kit.HOST_NAMED) # every slot, its copy held read.release() - write = lender.lend_write(target, 0, END) + write = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + lender.progress() assert write.poll() is None, "slots handed out while a copy into them was queued" gate.open() - view = ready_view(write, mgr) - write.mark_arrived(view.row_masks()) - write.release() + view = kit.ready_view(write, mgr) + kit.return_write(lender, write, view.row_masks()) def test_leases_wait_for_slots_strictly_in_order(kit, real_manager, attach=attach): """Waiting leases are granted in order, by the lender's own calls, without new lends.""" with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) # three slots - first = lender.lend_read(source, 0, 2 * TPB) - ready_view(first, mgr) - second = lender.lend_read(source, 0, END) - third = lender.lend_read(source, 0, TPB) + lender = attach(mgr, guaranteed_fetch_tokens=END) # three slots + first = lender.lend_read(source, 0, 2 * TPB, **kit.HOST_NAMED) + kit.ready_view(first, mgr) + second = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + third = lender.lend_read(source, 0, TPB, **kit.HOST_NAMED) + lender.progress() assert second.poll() is None assert third.poll() is None, "a later lease that fits jumped the one waiting ahead" first.release() - ready_view(second, mgr) + kit.ready_view(second, mgr) + lender.progress() assert third.poll() is None second.release() - assert by_group(ready_view(third, mgr)) == {0: [0]} + assert by_group(kit.ready_view(third, mgr)) == {0: [0]} third.release() -def test_planned_fetches_is_a_budget_and_holes_can_make_a_lease_wait(kit, real_manager): +def test_capacity_in_fetches_is_a_budget_and_holes_can_make_a_lease_wait(kit, real_manager): with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) (group,) = kit.pool_group_ids(mgr) - (region,) = lender.regions + (region,) = kit.staging_regions(lender) assert region.slot_count == 2 * BLOCKS - first = lender.lend_read(source, 0, 2 * TPB) # slots 0 and 1 - second = lender.lend_read(source, 0, TPB) # slot 2 - third = lender.lend_read(source, 0, TPB) # slot 3 + first = lender.lend_read(source, 0, 2 * TPB, **kit.HOST_NAMED) # slots 0 and 1 + second = lender.lend_read(source, 0, TPB, **kit.HOST_NAMED) # slot 2 + third = lender.lend_read(source, 0, TPB, **kit.HOST_NAMED) # slot 3 for lease in (first, second, third): - ready_view(lease, mgr) + kit.ready_view(lease, mgr) first.release() + lender.progress() assert lender._free_slots(group) == 4 # slots 0, 1, 4 and 5 - fetch = lender.lend_read(source, 0, END) + fetch = lender.lend_read(source, 0, END, **kit.HOST_NAMED) assert fetch.poll() is None and fetch.failure is None, "no three free slots in a row" second.release() # slots 0 to 2 are free in a row - (run,) = ready_view(fetch, mgr).groups - assert ((run.addresses - region.address) // region.slot_bytes).tolist() == [0, 1, 2] + (run,) = kit.ready_view(fetch, mgr).groups + assert run.slots.tolist() == [0, 1, 2] + assert kit.staged_slots(lender, run)[0] == region for lease in (third, fetch): lease.release() @@ -968,20 +1117,20 @@ def test_planned_fetches_is_a_budget_and_holes_can_make_a_lease_wait(kit, real_m def test_a_fetch_is_usable_once_its_copy_lands(kit, real_manager, local_blocks, attach=attach): with real_manager() as mgr_a, real_manager() as mgr_b: source = kit.published(mgr_a, SOURCE, PROMPT) - lender_a = attach(mgr_a, tokens_per_fetch=END) - publish = lender_a.lend_read(source, 0, END) - publish_view = ready_view(publish, mgr_a) + lender_a = attach(mgr_a, guaranteed_fetch_tokens=END) + publish = lender_a.lend_read(source, 0, END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) if local_blocks: kit.published(mgr_b, 7, PROMPT[: local_blocks * TPB] + list(range(9000, 9033))) target = kit.admitted(mgr_b, TARGET, PROMPT) kv = kit.kv(mgr_b, target) local, history = kv.num_committed_tokens, kv.history_length assert local == local_blocks * TPB - lender_b = attach(mgr_b, tokens_per_fetch=END) - assert lender_b.resume_bounds(target) == (max(local, history), history) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=END) + assert lender_b.poll_resume_bounds(target) == (max(local, history), history) - lease = lender_b.lend_write(target, local, END) - assert lender_b.resume_bounds(target) is None, "a fetch counts from its call" + lease = lender_b.lend_write(target, local, END, **kit.HOST_NAMED) + assert lender_b.poll_resume_bounds(target) is None, "a fetch counts from its call" assert kv.capacity >= END and kv.history_length == history view = lease.poll() assert view is not None, "a write is ready at its first poll after the grant" @@ -989,10 +1138,12 @@ def test_a_fetch_is_usable_once_its_copy_lands(kit, real_manager, local_blocks, kit.fill_sentinel(mgr_b, target, first_block=local_blocks) masks = kit.relay(lender_a, publish_view, lender_b, view) with kit.held_stream(mgr_b._stream) as gate: - lease.mark_arrived(masks) - assert lender_b.resume_bounds(target) is None, "settled before the copy into pages ran" + assert kit.return_write(lender_b, lease, masks) is None, "settled before its copy ran" + lender_b.progress() + assert lender_b.poll_resume_bounds(target) is None, "settled before its copy ran" gate.open() - resume_bounds = lender_b.resume_bounds(target) + lender_b.progress() + resume_bounds = lender_b.poll_resume_bounds(target) assert isinstance(resume_bounds, ResumeBounds) and all( type(v) is int for v in resume_bounds ) @@ -1000,7 +1151,6 @@ def test_a_fetch_is_usable_once_its_copy_lands(kit, real_manager, local_blocks, for ordinal in range(BLOCKS): got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) assert got == kit.digest(kit.page(mgr_a, source, 0, ordinal)), "a fetched page differs" - lease.release() publish.release() @@ -1012,24 +1162,22 @@ def test_a_fetch_is_usable_once_its_copy_lands(kit, real_manager, local_blocks, def test_a_partial_arrival_is_usable_up_to_its_first_gap(kit, real_manager, rows, usable): with real_manager() as mgr_a, real_manager() as mgr_b: source = kit.published(mgr_a, SOURCE, PROMPT) - lender_a = attach(mgr_a, tokens_per_fetch=END) - publish = lender_a.lend_read(source, 0, END) - publish_view = ready_view(publish, mgr_a) + lender_a = attach(mgr_a, guaranteed_fetch_tokens=END) + publish = lender_a.lend_read(source, 0, END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) target = kit.admitted(mgr_b, TARGET, PROMPT) - lender_b = attach(mgr_b, tokens_per_fetch=END) - lease = lender_b.lend_write(target, 0, END) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=END) + lease = lender_b.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() kit.fill_sentinel(mgr_b, target) kit.relay(lender_a, publish_view, lender_b, view, deliver=lambda run, row: row in rows) - lease.mark_arrived((np.isin(np.arange(BLOCKS), rows),)) - mgr_b._stream.synchronize() - assert lender_b.resume_bounds(target) == (usable, 0) + arrived = (np.isin(np.arange(BLOCKS), rows),) + assert return_and_settle(lender_b, mgr_b, target, lease, arrived) == (usable, 0) sentinel = bytes([kit.SENTINEL]) * kit.DevicePages(mgr_b).page_bytes(0) for ordinal in range(BLOCKS): expected = kit.page(mgr_a, source, 0, ordinal) if ordinal in rows else sentinel got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) - assert got == kit.digest(expected), "an unmarked row was copied" - lease.release() + assert got == kit.digest(expected), "a row that did not arrive was copied" publish.release() @@ -1038,101 +1186,139 @@ def test_a_windowed_fetch_is_usable_only_from_its_floor(kit, real_manager, parti windows = [WINDOW, 256] with real_manager(windows=windows) as mgr_a, real_manager(windows=windows) as mgr_b: source = kit.published(mgr_a, SOURCE, WINDOWED_PROMPT) - lender_a = attach(mgr_a, tokens_per_fetch=WINDOWED_END) - publish = lender_a.lend_read(source, 0, WINDOWED_END) - publish_view = ready_view(publish, mgr_a) + lender_a = attach(mgr_a, guaranteed_fetch_tokens=WINDOWED_END) + publish = lender_a.lend_read(source, 0, WINDOWED_END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) target = kit.admitted(mgr_b, TARGET, WINDOWED_PROMPT) - lender_b = attach(mgr_b, tokens_per_fetch=WINDOWED_END) - lease = lender_b.lend_write(target, 0, WINDOWED_END) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=WINDOWED_END) + lease = lender_b.lend_write(target, 0, WINDOWED_END, **kit.HOST_NAMED) assert kit.kv(mgr_b, target).history_length == WINDOWED_END, "the history moves to the end" view = lease.poll() assert by_group(view) == by_group(publish_view) kit.fill_sentinel(mgr_b, target) deliver = (lambda run, row: row < 2) if partial else None masks = kit.relay(lender_a, publish_view, lender_b, view, deliver) - lease.mark_arrived(masks) - mgr_b._stream.synchronize() - resume_bounds = lender_b.resume_bounds(target) + resume_bounds = return_and_settle(lender_b, mgr_b, target, lease, masks) assert resume_bounds.min_position == WINDOWED_END, "the window released the early blocks" if partial: assert resume_bounds.max_position < resume_bounds.min_position, "empty: compute from 0" for run in view.groups: # the two rows of each run that arrived hold the source's bytes - for ordinal in run.block_ordinals.tolist()[:2]: + for ordinal in run.page_ordinals.tolist()[:2]: lg = run.layer_group_index got = kit.digest(kit.page(mgr_b, target, lg, ordinal)) assert got == kit.digest(kit.page(mgr_a, source, lg, ordinal)) else: assert resume_bounds.max_position == WINDOWED_END for run in view.groups: - for ordinal in run.block_ordinals.tolist(): + for ordinal in run.page_ordinals.tolist(): lg = run.layer_group_index got = kit.digest(kit.page(mgr_b, target, lg, ordinal)) assert got == kit.digest(kit.page(mgr_a, source, lg, ordinal)) - lease.release() publish.release() -def test_mark_arrived_takes_one_mask_per_run_once_on_a_write_seen_ready(kit, real_manager): - with real_manager() as mgr: +def read_only(array): + """``array`` made read-only.""" + array.flags.writeable = False + return array + + +def test_return_writes_takes_one_mask_per_run_of_each_write_seen_ready_once(kit, real_manager): + """A wrong entry raises before any lease returns; a write seen ready returns once.""" + with real_manager() as mgr, real_manager() as other_mgr: source = kit.published(mgr, SOURCE, PROMPT) target = kit.admitted(mgr, TARGET, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) - read = lender.lend_read(source, 0, END) - read_view = ready_view(read, mgr) - with pytest.raises(RuntimeError): - read.mark_arrived(read_view.row_masks()) - write = lender.lend_write(target, 0, END) - with pytest.raises(RuntimeError): - write.mark_arrived(()) # before poll() gave the view + whole = kit.admitted(mgr, OTHER, list(range(8000, 8097))) + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=3) + read = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + read_view = kit.ready_view(read, mgr) + with pytest.raises(RuntimeError, match="for write leases"): + kit.return_write(lender, read, read_view.row_masks()) + write = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + with pytest.raises(RuntimeError, match="before poll"): + kit.return_write(lender, write, ()) # before poll() gave the view view = write.poll() + foreign_target = kit.admitted(other_mgr, TARGET, OTHER_PROMPT) + foreign_lender = attach(other_mgr, guaranteed_fetch_tokens=END) + foreign = foreign_lender.lend_write(foreign_target, 0, END, **kit.HOST_NAMED) + foreign_view = foreign.poll() bad = ( - (), - (np.ones(BLOCKS - 1, bool),), - (np.ones(BLOCKS, bool),) * 2, - (np.ones(BLOCKS, np.int64),), # one per row, but not bool + ((write, (), True), ValueError), + ((write, (np.ones(BLOCKS - 1, bool),), True), ValueError), + ((write, (np.ones(BLOCKS, bool),) * 2, True), ValueError), + ((write, (np.ones(BLOCKS, np.int64),), True), ValueError), # one per row, not bool + # read-only, but one row short + ((write, (read_only(np.ones(BLOCKS - 1, bool)),), True), ValueError), + # read-only and one per row, but not bool + ((write, (read_only(np.ones(BLOCKS, np.int64)),), True), ValueError), + ((foreign, foreign_view.row_masks(), True), ValueError), # another lender's + (write, TypeError), ) - for masks in bad: - with pytest.raises(ValueError): - write.mark_arrived(masks) - assert lender.resume_bounds(target) is None, "a refused mask recorded something" - write.mark_arrived(view.row_masks()) - with pytest.raises(RuntimeError): - write.mark_arrived(view.row_masks()) - mgr._stream.synchronize() + for entry, error in bad: + with pytest.raises(error): + lender.return_writes([entry]) + assert lender.poll_resume_bounds(target) is None, "a refused entry recorded something" + assert write.poll() is view, "a refused entry returned the lease" kv = kit.kv(mgr, target) - assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) - write.release() - read.release() + own = (kv.num_committed_tokens, kv.history_length) + assert kit.return_write(lender, write, view.row_masks()) == own, "no row arrived" + with pytest.raises(RuntimeError, match="returned already"): + kit.return_write(lender, write, view.row_masks()) + with pytest.raises(RuntimeError): + write.poll() + every = lender.lend_write(whole, 0, END, **kit.HOST_NAMED) + every_view = every.poll() + kit.fill_sentinel(mgr, whole) + kit.stage(lender, every_view, 0x5A) + assert return_and_settle(lender, mgr, whole, every, ALL_ARRIVED) == (END, 0) + nbytes = kit.DevicePages(mgr).page_bytes(0) + got = kit.digest([kit.page(mgr, whole, 0, o) for o in range(BLOCKS)]) + assert got == kit.digest([kit.staged_row(0x5A, 0, o, nbytes) for o in range(BLOCKS)]), ( + "ALL_ARRIVED left a row out" + ) + for lease in (read, foreign): + lease.release() -@pytest.mark.parametrize( - "mark_first", [True, False], ids=["mark_then_release", "release_then_mark"] -) -def test_either_order_of_marks_and_release_copies_the_same(kit, real_manager, mark_first): +@pytest.mark.parametrize("ending", ["returned_writable", "returned_read_only", "released"]) +def test_a_write_seen_ready_gives_back_its_slots_once_no_copy_reads_them(kit, real_manager, ending): + """Its slots come back once its copy ran; released without a return, it delivers no row.""" with real_manager() as mgr: first = kit.admitted(mgr, TARGET, PROMPT) second = kit.admitted(mgr, OTHER, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END) # the second waits for the first's slots - lease = lender.lend_write(first, 0, END) + lender = attach(mgr, guaranteed_fetch_tokens=END) # the second waits for the first's slots + lease = lender.lend_write(first, 0, END, **kit.HOST_NAMED) view = lease.poll() kit.fill_sentinel(mgr, first) kit.stage(lender, view, 0x5A) - waiting = lender.lend_write(second, 0, END) - if not mark_first: - lease.release() - assert waiting.poll() is None, "slots returned before the marks" - with kit.held_stream(mgr._stream) as gate: - lease.mark_arrived(view.row_masks(True)) - lease.release() - assert waiting.poll() is None, "slots returned before the copy out of them ran" - gate.open() - waiting_view = ready_view(waiting, mgr) + waiting = lender.lend_write(second, 0, END, **kit.HOST_NAMED) nbytes = kit.DevicePages(mgr).page_bytes(0) + if ending == "released": + kv = kit.kv(mgr, first) + lease.release() + bounds = lender.poll_resume_bounds(first) + assert bounds == (kv.num_committed_tokens, kv.history_length), ( + f"resume_bounds {bounds} of a write released without a return" + ) + with pytest.raises(RuntimeError, match="returned already"): + kit.return_write(lender, lease, view.row_masks(True)) + expected = [bytes([kit.SENTINEL]) * nbytes] * BLOCKS + else: + masks = view.row_masks(True) + if ending == "returned_read_only": + masks = tuple(read_only(mask) for mask in masks) + with kit.held_stream(mgr._stream) as gate: + assert kit.return_write(lender, lease, masks) is None, "settled before its copy ran" + lender.progress() + assert waiting.poll() is None, "slots returned before the copy out of them ran" + gate.open() + expected = [kit.staged_row(0x5A, 0, o, nbytes) for o in range(BLOCKS)] + waiting_view = kit.ready_view(waiting, mgr) got = kit.digest([kit.page(mgr, first, 0, o) for o in range(BLOCKS)]) - assert got == kit.digest([kit.staged_row(0x5A, 0, o, nbytes) for o in range(BLOCKS)]) - assert lender.resume_bounds(first) == (END, 0) - waiting.mark_arrived(waiting_view.row_masks()) - waiting.release() + assert got == kit.digest(expected), "the pages hold other bytes than the rows delivered" + if ending != "released": + assert lender.poll_resume_bounds(first) == (END, 0) + kit.return_write(lender, waiting, waiting_view.row_masks()) # -- the request exits during a lease --------------------------------------------------------- @@ -1141,21 +1327,21 @@ def test_either_order_of_marks_and_release_copies_the_same(kit, real_manager, ma def test_arrivals_for_a_freed_request_copy_nothing(kit, real_manager, attach=attach): with real_manager(max_tokens=kit.POOL_TOKENS) as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - lease = lender.lend_write(target, 0, END) + lender = attach(mgr, guaranteed_fetch_tokens=END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() assert view is not None kv = kit.kv(mgr, target) - lent = {kit.pages(kv, 0)[o] for o in view.groups[0].block_ordinals.tolist()} + lent = {kit.pages(kv, 0)[o] for o in view.groups[0].page_ordinals.tolist()} kit.stage(lender, view, 0x5A) mgr.free_resources(target) with pytest.raises(ValueError): - lender.resume_bounds(target) + lender.poll_resume_bounds(target) others = kit.Requests(mgr) try: assert others.allocate(kit.pool_pages(mgr)) assert lent <= others.pages(), "no lent page was reused: the check proves nothing" - lease.mark_arrived(view.row_masks(True)) + assert kit.return_write(lender, lease, view.row_masks(True)) is None mgr._stream.synchronize() dev = kit.DevicePages(mgr) sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) @@ -1181,12 +1367,12 @@ def test_a_commit_during_a_fetch_copies_nothing_into_pages_it_returned( returns the lent pages to the pool: arrivals reach neither those pages nor the shared ones.""" with real_manager(max_tokens=kit.POOL_TOKENS, max_batch_size=16) as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - lease = lender.lend_write(target, 0, END) + lender = attach(mgr, guaranteed_fetch_tokens=END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() assert view is not None kv = kit.kv(mgr, target) - lent = {kit.pages(kv, 0)[o] for o in view.groups[0].block_ordinals.tolist()} + lent = {kit.pages(kv, 0)[o] for o in view.groups[0].page_ordinals.tolist()} kit.stage(lender, view, 0x5A) source = kit.published(mgr, SOURCE, PROMPT) shared = kit.pages(kit.kv(mgr, source), 0)[:BLOCKS] @@ -1199,7 +1385,7 @@ def test_a_commit_during_a_fetch_copies_nothing_into_pages_it_returned( while others.allocate(1, chunk=1): pass assert lent <= others.pages(), "no returned page was reused: the check proves nothing" - lease.mark_arrived(view.row_masks(True)) + kit.return_write(lender, lease, view.row_masks(True)) mgr._stream.synchronize() dev = kit.DevicePages(mgr) sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) @@ -1220,10 +1406,10 @@ def test_a_waiting_read_fails_when_its_cache_changes(kit, real_manager, change, with real_manager() as mgr: holder = kit.published(mgr, SOURCE, PROMPT) other = kit.published(mgr, OTHER, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - head = lender.lend_read(holder, 0, END) - ready_view(head, mgr) - waiting = lender.lend_read(other, 0, 2 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=END) + head = lender.lend_read(holder, 0, END, **kit.HOST_NAMED) + kit.ready_view(head, mgr) + waiting = lender.lend_read(other, 0, 2 * TPB, **kit.HOST_NAMED) assert waiting.poll() is None and waiting.failure is None if change == "freed": mgr.free_resources(other) @@ -1231,10 +1417,11 @@ def test_a_waiting_read_fails_when_its_cache_changes(kit, real_manager, change, else: mgr.suspend_request(other) head.release() + lender.progress() assert waiting.poll() is None and waiting.failure is not None, "copied a moved cache" waiting.release() - again = lender.lend_read(holder, 0, END) # its slots came back - ready_view(again, mgr) + again = lender.lend_read(holder, 0, END, **kit.HOST_NAMED) # its slots came back + kit.ready_view(again, mgr) again.release() @@ -1250,11 +1437,11 @@ def test_a_publish_granted_before_its_request_is_freed_copies_what_it_lent(kit, source = kit.published(mgr, SOURCE, PROMPT) original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] lent = set(kit.pages(kit.kv(mgr, source), 0)[:BLOCKS]) - lender = attach(mgr, tokens_per_fetch=END) + lender = attach(mgr, guaranteed_fetch_tokens=END) dev = kit.DevicePages(mgr) # built ahead: building it waits for the device others = kit.Requests(mgr, dev) with kit.held_stream(mgr._stream, strict=False) as gate: - lease = lender.lend_read(source, 0, END) # granted, its copy queued behind the gate + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) # granted, its copy held mgr.free_resources(source) assert lease.failure is None and kit.kv(mgr, source) is None # Other requests take every page and overwrite it, on the stream behind the copy. @@ -1263,9 +1450,9 @@ def test_a_publish_granted_before_its_request_is_freed_copies_what_it_lent(kit, gate.open() sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) assert all(dev.read(0, slot) == sentinel for slot in lent), "the pages were not overwritten" - (run,) = ready_view(lease, mgr).groups - length = lender.regions[run.region_index].slot_bytes - got = kit.digest([kit.host_bytes(a, length) for a in run.addresses.tolist()]) + (run,) = kit.ready_view(lease, mgr).groups + region, addresses = kit.staged_slots(lender, run) + got = kit.digest([kit.host_bytes(a, region.slot_bytes) for a in addresses]) assert got == kit.digest(original), "the copy read a later owner's writes" lease.release() others.free() @@ -1287,33 +1474,34 @@ def test_an_abandoned_fetch_leaves_only_its_growth( with real_manager(windows=[WINDOW, 256] if windowed else None) as mgr: target = kit.admitted(mgr, TARGET, tokens) second = kit.admitted(mgr, OTHER, other) - lender = attach(mgr, tokens_per_fetch=end) + lender = attach(mgr, guaranteed_fetch_tokens=end) head = None if waiting: # a publish takes every slot first, so the write waits in line holder = kit.published(mgr, SOURCE, list(range(8000, 8000 + len(tokens)))) - head = lender.lend_read(holder, 0, end) - ready_view(head, mgr) - lease = lender.lend_write(target, 0, end) - assert lender.resume_bounds(target) is None + head = lender.lend_read(holder, 0, end, **kit.HOST_NAMED) + kit.ready_view(head, mgr) + lease = lender.lend_write(target, 0, end, **kit.HOST_NAMED) + assert lender.poll_resume_bounds(target) is None if waiting: assert lease.poll() is None and lease.failure is None, "the write did not wait" lease.release() # before any grant or view: no backend wrote, the fetch is abandoned if head is not None: - head.release() # the slots come back; the released write takes none of them + head.release() + lender.progress() # the slots come back; the released write takes none of them kv = kit.kv(mgr, target) own = (kv.num_committed_tokens, kv.history_length) - assert lender.resume_bounds(target) == own, "the abandoned fetch still counts" + assert lender.poll_resume_bounds(target) == own, "the abandoned fetch still counts" if windowed: assert kv.history_length == end, "the growth moved the window's history" before = state_of(kit, mgr, lender, target) - with pytest.raises(ValueError, match="below the history"): - lender.lend_write(target, 0, END) + below = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + assert below.failure is not None and "below the history" in below.failure + below.release() assert state_of(kit, mgr, lender, target) == before - again = lender.lend_write(second, 0, end) + again = lender.lend_write(second, 0, end, **kit.HOST_NAMED) view = again.poll() assert view is not None, "the abandoned fetch kept its slots" - again.mark_arrived(view.row_masks()) - again.release() + kit.return_write(lender, again, view.row_masks()) # -- a fetch split into leases, and a shrink of the cache it went into ------------------------- @@ -1333,20 +1521,16 @@ def stale_anywhere(kit, mgr, history): def fetch_segment(kit, mgr, lender, target, start, end, byte, arrived=None): - """One write lease over ``[start, end)``: every slot staged from ``byte``, each row with its - own bytes (``kit.staged_row``), the rows of block ordinals ``arrived`` (default all) marked, - released, the copy into pages run.""" - lease = lender.lend_write(target, start, end) + """A write over ``[start, end)`` staged from ``byte``, returned with ``arrived``: its bounds.""" + lease = lender.lend_write(target, start, end, **kit.HOST_NAMED) view = lease.poll() assert view is not None, f"the write lease was not ready: {lease.failure}" kit.stage(lender, view, byte) - masks = tuple( - np.ones(len(run), dtype=bool) if arrived is None else np.isin(run.block_ordinals, arrived) - for run in view.groups - ) - lease.mark_arrived(masks) - lease.release() - mgr._stream.synchronize() + if arrived is None: + masks = ALL_ARRIVED + else: + masks = tuple(np.isin(run.page_ordinals, arrived) for run in view.groups) + return return_and_settle(lender, mgr, target, lease, masks) SEGMENTS = { @@ -1369,7 +1553,7 @@ def test_consecutive_leases_add_up_to_one_fetch(kit, real_manager, case, attach= target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) kv = kit.kv(mgr, target) assert kv.num_committed_tokens == 0, "no local prefix: every usable token was fetched" - lender = attach(mgr, tokens_per_fetch=64) + lender = attach(mgr, guaranteed_fetch_tokens=64) bytes_of = {} for i, ((start, end, arrived), usable) in enumerate(zip(segments, expected)): byte = 0x11 * (i + 1) @@ -1378,7 +1562,7 @@ def test_consecutive_leases_add_up_to_one_fetch(kit, real_manager, case, attach= if arrived is None or ordinal in arrived: bytes_of[ordinal] = byte assert kv.num_committed_tokens == 0, "a fetch committed" - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) floor = kv.history_length if stale_anywhere(kit, mgr, kv.history_length) else 0 assert resume_bounds is not None and resume_bounds.min_position == floor assert resume_bounds.max_position >= usable, ( @@ -1419,7 +1603,7 @@ def test_a_later_lease_past_the_window_and_a_block_fails_with_extra_kv_tokens( with real_manager(spec_config=pard(), windows=[WINDOW, 256]) as mgr: assert mgr.num_extra_kv_tokens == 1, "no extra KV token: nothing checked" target = kit.admitted(mgr, TARGET, WINDOWED_PROMPT) - lender = attach(mgr, tokens_per_fetch=3 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=3 * TPB) fetch_segment(kit, mgr, lender, target, 0, 2 * TPB, 0x5A) before = state_of(kit, mgr, lender, target) assert before[1].max_position == 2 * TPB, f"resume_bounds after the first lease {before[1]}" @@ -1427,7 +1611,7 @@ def test_a_later_lease_past_the_window_and_a_block_fails_with_extra_kv_tokens( assert kit.pages(kit.kv(mgr, target), sliding)[2] >= 0, ( "block 2 has no page: nothing checked" ) - lease = lender.lend_write(target, 2 * TPB, 2 * TPB + span) + lease = lender.lend_write(target, 2 * TPB, 2 * TPB + span, **kit.HOST_NAMED) if span >= WINDOW + TPB - 1: assert lease.failure is not None and "never wrote" in lease.failure, ( f"a later lease of {span} tokens was lent past an unwritten page: {lease.failure}" @@ -1435,11 +1619,11 @@ def test_a_later_lease_past_the_window_and_a_block_fails_with_extra_kv_tokens( assert state_of(kit, mgr, lender, target) == before, ( "the failed lease changed the cache" ) + lease.release() else: - view = ready_view(lease, mgr) + view = kit.ready_view(lease, mgr) kit.stage(lender, view, 0x6B) - lease.mark_arrived(view.row_masks(True)) - mgr._stream.synchronize() + resume_bounds = return_and_settle(lender, mgr, target, lease, view.row_masks(True)) for lg in range( kit.num_layer_groups(mgr) ): # what resume_bounds counts, the window's alone @@ -1449,23 +1633,21 @@ def test_a_later_lease_past_the_window_and_a_block_fails_with_extra_kv_tokens( f"layer group {lg}: fetched blocks {ordinals} lack their staged bytes" ) end = 2 * TPB + span - resume_bounds = lender.resume_bounds(target) assert resume_bounds == (end, end), ( f"resume_bounds {tuple(resume_bounds)} after the later lease" ) - lease.release() def test_an_abandoned_segment_keeps_what_earlier_ones_delivered(kit, real_manager, attach=attach): with real_manager() as mgr: target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) - lender = attach(mgr, tokens_per_fetch=64) + lender = attach(mgr, guaranteed_fetch_tokens=64) fetch_segment(kit, mgr, lender, target, 0, 64, 0x11) - assert lender.resume_bounds(target) == (64, 0) - second = lender.lend_write(target, 64, 128) - assert lender.resume_bounds(target) is None, "a fetch counts from its call" + assert lender.poll_resume_bounds(target) == (64, 0) + second = lender.lend_write(target, 64, 128, **kit.HOST_NAMED) + assert lender.poll_resume_bounds(target) is None, "a fetch counts from its call" second.release() # before anyone saw it ready: abandoned, nothing written - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds == (64, 0), ( f"resume_bounds {tuple(resume_bounds)} after an abandoned segment; the first segment's " "delivered prefix was dropped" @@ -1495,19 +1677,19 @@ def test_a_lost_later_segment_empties_the_interval_only_where_a_window_released_ with real_manager(windows=[WINDOW, 256]) as mgr: target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) kv = kit.kv(mgr, target) - lender = attach(mgr, tokens_per_fetch=2 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) fetch_segment(kit, mgr, lender, target, 0, first_end, 0x11) for lg in range(kit.num_layer_groups(mgr)): assert holds(kit, mgr, target, lg, range(first_end // TPB), staged_bytes(kit, 0x11)) - later = lender.lend_write(target, first_end, later_end) + later = lender.lend_write(target, first_end, later_end, **kit.HOST_NAMED) assert kv.history_length == later_end, "the later lease did not move the history" if lost == "abandoned": later.release() # before its first poll + resume_bounds = lender.poll_resume_bounds(target) else: - later.mark_arrived(later.poll().row_masks(False)) - later.release() - mgr._stream.synchronize() - resume_bounds = lender.resume_bounds(target) + # No row arrived: the return copies nothing, so it settles the fetch. + resume_bounds = kit.return_write(lender, later, later.poll().row_masks(False)) + assert resume_bounds == lender.poll_resume_bounds(target), resume_bounds assert stale_anywhere(kit, mgr, later_end) == (expected is None) if expected is None: assert resume_bounds.min_position == later_end > resume_bounds.max_position, ( @@ -1533,34 +1715,37 @@ def tp_rank(rank): def test_whether_an_earlier_fetch_settled_is_each_rank_s_own(kit, real_manager, attach=attach): - """Two managers stand for the two tensor-parallel ranks of one request, which a caller fetches - in two leases. Both ranks mark the first; rank 1's copy into its pages still waits behind - earlier work on its stream when both lend the second. Rank 0 gets the lease and rank 1 a lease - failed at the call, which changed nothing: an outcome the caller combines, where a ValueError - on one rank alone would leave that rank's loop while the other waits for it. Lent only once - every rank's resume_bounds is not None, the second lease is granted on both.""" + """Two TP ranks: the one whose first fetch is unsettled fails the second lease at the call.""" with ( real_manager(num_kv_heads=4, mapping=tp_rank(0)) as rank0, real_manager(num_kv_heads=4, mapping=tp_rank(1)) as rank1, ): ranks = (rank0, rank1) targets = [kit.admitted(mgr, TARGET, SPLIT_PROMPT) for mgr in ranks] - lenders = [attach(mgr, tokens_per_fetch=2 * TPB, planned_fetches=2) for mgr in ranks] - firsts = [lender.lend_write(t, 0, 2 * TPB) for lender, t in zip(lenders, targets)] - views = [ready_view(lease, mgr) for lease, mgr in zip(firsts, ranks)] + lenders = [ + attach(mgr, guaranteed_fetch_tokens=2 * TPB, capacity_in_fetches=2) for mgr in ranks + ] + firsts = [ + lender.lend_write(t, 0, 2 * TPB, **kit.HOST_NAMED) + for lender, t in zip(lenders, targets) + ] + views = [kit.ready_view(lease, mgr) for lease, mgr in zip(firsts, ranks)] for lender, view in zip(lenders, views): kit.stage(lender, view, 0x5A) # the backend filled the first segment on both ranks - firsts[0].mark_arrived(views[0].row_masks(True)) + kit.return_write(lenders[0], firsts[0], views[0].row_masks(True)) rank0._stream.synchronize() # rank 0's copy into its pages has run seconds, outcomes = [], [] with kit.held_stream(rank1._stream, strict=False): - firsts[1].mark_arrived(views[1].row_masks(True)) # rank 1's copy stays queued + # Rank 1's copy stays queued. + kit.return_write(lenders[1], firsts[1], views[1].row_masks(True)) + for lender in lenders: + lender.progress() before = [ state_of(kit, mgr, lender, t) for mgr, lender, t in zip(ranks, lenders, targets) ] for lender, target in zip(lenders, targets): try: - lease = lender.lend_write(target, 2 * TPB, 4 * TPB) + lease = lender.lend_write(target, 2 * TPB, 4 * TPB, **kit.HOST_NAMED) except ValueError as error: outcomes.append(f"ValueError: {error}") continue @@ -1579,19 +1764,23 @@ def test_whether_an_earlier_fetch_settled_is_each_rank_s_own(kit, real_manager, lease.release() seconds = [] # The caller lends the next segment once every rank's resume_bounds is not None. - for mgr in ranks: + for mgr, lender in zip(ranks, lenders): mgr._stream.synchronize() - assert all(lender.resume_bounds(t) is not None for lender, t in zip(lenders, targets)) + lender.progress() + assert all( + lender.poll_resume_bounds(t) is not None for lender, t in zip(lenders, targets) + ) for rank, (mgr, target) in enumerate(zip(ranks, targets)): assert holds(kit, mgr, target, 0, [0, 1], staged_bytes(kit, 0x5A)), ( f"rank {rank}: the first segment's rows landed off their blocks" ) seconds = [ - lender.lend_write(t, 2 * TPB, 4 * TPB) for lender, t in zip(lenders, targets) + lender.lend_write(t, 2 * TPB, 4 * TPB, **kit.HOST_NAMED) + for lender, t in zip(lenders, targets) ] assert [lease.failure for lease in seconds] == [None, None] - for lease, mgr in zip(seconds, ranks): - lease.mark_arrived(ready_view(lease, mgr).row_masks()) + for lease, mgr, lender in zip(seconds, ranks, lenders): + kit.return_write(lender, lease, kit.ready_view(lease, mgr).row_masks()) finally: for mgr in ranks: mgr._stream.synchronize() @@ -1610,9 +1799,9 @@ def test_a_resume_with_the_fresh_page_fill_on_keeps_the_fetched_blocks( monkeypatch.setenv("TRTLLM_KV_FRESH_PAGE_FILL", str(FRESH_FILL)) with real_manager() as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) + lender = attach(mgr, guaranteed_fetch_tokens=END) fetch_segment(kit, mgr, lender, target, 0, 64, 0x5A) - assert lender.resume_bounds(target).max_position == 64 + assert lender.poll_resume_bounds(target).max_position == 64 # The executor resumes the request where the fetch left it usable. target.context_chunk_size = target.prompt_len - target.context_current_position target.set_prepopulated_prompt_len(64, TPB) @@ -1657,10 +1846,10 @@ def test_a_draft_pool_fetch_fills_only_the_pages_its_grow_added(kit, real_manage assert request.py_request_id not in draft._fresh_pages_filled, ( "the draft pool's own resizes ran the fill: the check proves nothing" ) - lender = attach(draft, tokens_per_fetch=2 * TPB, scope=SCOPE + b"/draft-pool") + lender = attach(draft, guaranteed_fetch_tokens=2 * TPB, scope=SCOPE + b"/draft-pool") blocks_before = kv.num_blocks - lease = lender.lend_write(request, LOCAL, 4 * TPB) - view = ready_view(lease, draft) + lease = lender.lend_write(request, LOCAL, 4 * TPB, **kit.HOST_NAMED) + view = kit.ready_view(lease, draft) torch.cuda.synchronize() # the fill runs on the current stream filled = kit.digest(np.full(dev.page_bytes(0) // 2, FRESH_FILL, np.float16).tobytes()) added = range(blocks_before, 4) @@ -1671,11 +1860,9 @@ def test_a_draft_pool_fetch_fills_only_the_pages_its_grow_added(kit, real_manage "the fetch's fill overwrote the blocks the draft computed" ) kit.stage(lender, view, 0x5A) - lease.mark_arrived(view.row_masks(True)) - lease.release() - draft._stream.synchronize() + bounds = return_and_settle(lender, draft, request, lease, view.row_masks(True)) assert holds(kit, draft, request, 0, [2, 3], staged_bytes(kit, 0x5A)) - assert lender.resume_bounds(request) == (4 * TPB, LOCAL) + assert bounds == (4 * TPB, LOCAL) def test_a_second_draft_pool_fetch_fills_only_the_pages_its_own_grow_added( @@ -1695,9 +1882,9 @@ def test_a_second_draft_pool_fetch_fills_only_the_pages_its_own_grow_added( request = kit.admitted(draft, TARGET, LONG_PROMPT) kv = kit.kv(draft, request) dev = kit.DevicePages(draft) - lender = attach(draft, tokens_per_fetch=2 * TPB, scope=SCOPE + b"/draft-pool") + lender = attach(draft, guaranteed_fetch_tokens=2 * TPB, scope=SCOPE + b"/draft-pool") fetch_segment(kit, draft, lender, request, 0, LOCAL, 0x5A) - assert lender.resume_bounds(request).max_position == LOCAL + assert lender.poll_resume_bounds(request).max_position == LOCAL # The draft pool computes blocks 2 to 4 as the executor runs it after the resume. chunk = range(LOCAL // TPB, 5) resumed(kit, draft, request, LOCAL) @@ -1719,8 +1906,8 @@ def test_a_second_draft_pool_fetch_fills_only_the_pages_its_own_grow_added( f"the fill's record covers {known} blocks, the draft's whole chunk: nothing to check" ) blocks_before = kv.num_blocks - lease = lender.lend_write(request, 5 * TPB, 7 * TPB) - view = ready_view(lease, draft) + lease = lender.lend_write(request, 5 * TPB, 7 * TPB, **kit.HOST_NAMED) + view = kit.ready_view(lease, draft) torch.cuda.synchronize() # the fill runs on the current stream assert holds(kit, draft, request, 0, chunk, chunk_bytes), ( "the second fetch's fill overwrote the chunk the draft computed" @@ -1731,12 +1918,10 @@ def test_a_second_draft_pool_fetch_fills_only_the_pages_its_own_grow_added( kit.digest(kit.page(draft, request, 0, o)) == filled for o in grown ), "the fill did not run on the pages the second fetch grew" kit.stage(lender, view, 0x6B) - lease.mark_arrived(view.row_masks(True)) - lease.release() - draft._stream.synchronize() + bounds = return_and_settle(lender, draft, request, lease, view.row_masks(True)) assert holds(kit, draft, request, 0, [0, 1], staged_bytes(kit, 0x5A)) assert holds(kit, draft, request, 0, [5, 6], staged_bytes(kit, 0x6B)) - assert lender.resume_bounds(request) == (7 * TPB, 5 * TPB) + assert bounds == (7 * TPB, 5 * TPB) @pytest.mark.parametrize("taker", ["context", "fetch"]) @@ -1756,7 +1941,7 @@ def test_the_fresh_page_fill_lets_a_freed_publisher_s_copy_run_first( source = kit.published(mgr, SOURCE, PROMPT) original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] lent = set(kit.pages(kit.kv(mgr, source), 0)[:BLOCKS]) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) fetching = kit.admitted(mgr, TARGET, OTHER_PROMPT) if taker == "fetch" else None others = kit.Requests(mgr) # Other requests take every free page. Their fills load the fill's kernels before the @@ -1784,7 +1969,7 @@ def recording_fill(request_id): write = None try: with kit.gated_stream(mgr._stream, open_after=10.0) as gate: - read = lender.lend_read(source, 0, END) # its copy waits behind the gate + read = lender.lend_read(source, 0, END, **kit.HOST_NAMED) # its copy held mgr.free_resources(source) if taker == "context": taken = kit.make_request(TARGET, OTHER_PROMPT) @@ -1792,15 +1977,15 @@ def recording_fill(request_id): assert mgr.resize_context(taken, taken.context_remaining_length) else: taken = fetching - write = lender.lend_write(taken, 0, END) + write = lender.lend_write(taken, 0, END, **kit.HOST_NAMED) assert write.failure is None, write.failure finally: del mgr._fill_fresh_kv_pages assert began and not began[0], "the copy ran before the fill: the check proves nothing" assert lent & set(kit.pages(kit.kv(mgr, taken), 0)), "no freed page was taken" - (run,) = ready_view(read, mgr).groups - length = lender.regions[run.region_index].slot_bytes - got = kit.digest([kit.host_bytes(a, length) for a in run.addresses.tolist()]) + (run,) = kit.ready_view(read, mgr).groups + region, addresses = kit.staged_slots(lender, run) + got = kit.digest([kit.host_bytes(a, region.slot_bytes) for a in addresses]) assert got == kit.digest(original), "the fill overwrote a page before its copy ran" read.release() if write is not None: @@ -1816,29 +2001,27 @@ def test_a_rollback_voids_what_the_fetch_delivered(kit, real_manager, ask_betwee ``ask_between``: resume_bounds is asked between the rollback and the regrow.""" with real_manager() as mgr_a, real_manager(**SMALL_POOL) as mgr_b: source = kit.published(mgr_a, SOURCE, PROMPT) - lender_a = attach(mgr_a, tokens_per_fetch=END) - publish = lender_a.lend_read(source, 0, END) - publish_view = ready_view(publish, mgr_a) + lender_a = attach(mgr_a, guaranteed_fetch_tokens=END) + publish = lender_a.lend_read(source, 0, END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) source_pages = [kit.digest(kit.page(mgr_a, source, 0, b)) for b in range(BLOCKS)] target = kit.admitted(mgr_b, TARGET, PROMPT) - lender_b = attach(mgr_b, tokens_per_fetch=END) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=END) kv = kit.kv(mgr_b, target) - lease = lender_b.lend_write(target, 0, END) + lease = lender_b.lend_write(target, 0, END, **kit.HOST_NAMED) assert target.py_ctx_pre_resize_cap == 0, "the rollback would not undo the growth" view = lease.poll() kit.fill_sentinel(mgr_b, target) - lease.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) - lease.release() + masks = kit.relay(lender_a, publish_view, lender_b, view) + assert return_and_settle(lender_b, mgr_b, target, lease, masks) == (END, 0) publish.release() - mgr_b._stream.synchronize() assert [kit.digest(kit.page(mgr_b, target, 0, b)) for b in range(BLOCKS)] == source_pages - assert lender_b.resume_bounds(target) == (END, 0) assert mgr_b.revert_allocate_context(target) is True assert kit.kv(mgr_b, target) is kv, "the rollback replaced the cache" assert kv.capacity == 0 and kit.pages(kv, 0) == [] and kv.num_committed_tokens == 0 if ask_between: - resume_bounds = lender_b.resume_bounds(target) + resume_bounds = lender_b.poll_resume_bounds(target) assert resume_bounds == (0, 0), ( f"resume_bounds {tuple(resume_bounds)} counts freed blocks" ) @@ -1846,7 +2029,7 @@ def test_a_rollback_voids_what_the_fetch_delivered(kit, real_manager, ask_betwee assert mgr_b.resize_context(target, END) regrown = [kit.digest(kit.page(mgr_b, target, 0, b)) for b in range(BLOCKS)] assert regrown != source_pages, "the regrown pages still hold the fetched bytes" - resume_bounds = lender_b.resume_bounds(target) + resume_bounds = lender_b.poll_resume_bounds(target) assert resume_bounds == (0, 0), ( f"resume_bounds {tuple(resume_bounds)} counts blocks whose pages now hold {regrown}" ) @@ -1866,22 +2049,20 @@ def test_a_rollback_to_a_local_prefix_keeps_only_the_prefix(kit, real_manager, a local = 2 * TPB with real_manager() as mgr_a, real_manager(**SMALL_POOL) as mgr_b: source = kit.published(mgr_a, SOURCE, PROMPT) - lender_a = attach(mgr_a, tokens_per_fetch=END) - publish = lender_a.lend_read(source, 0, END) - publish_view = ready_view(publish, mgr_a) + lender_a = attach(mgr_a, guaranteed_fetch_tokens=END) + publish = lender_a.lend_read(source, 0, END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) kit.published(mgr_b, 7, PROMPT[:local] + list(range(9000, 9033))) target = kit.admitted(mgr_b, TARGET, PROMPT) - lender_b = attach(mgr_b, tokens_per_fetch=END) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=END) kv = kit.kv(mgr_b, target) assert kv.num_committed_tokens == local - lease = lender_b.lend_write(target, local, END) + lease = lender_b.lend_write(target, local, END, **kit.HOST_NAMED) view = lease.poll() kit.fill_sentinel(mgr_b, target, first_block=2) - lease.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) - lease.release() + masks = kit.relay(lender_a, publish_view, lender_b, view) + assert return_and_settle(lender_b, mgr_b, target, lease, masks) == (END, local) publish.release() - mgr_b._stream.synchronize() - assert lender_b.resume_bounds(target) == (END, local) got = kit.digest(kit.page(mgr_b, target, 0, 2)) assert got == kit.digest(kit.page(mgr_a, source, 0, 2)), ( "the fetched block 2 holds bytes other than the publisher's" @@ -1889,7 +2070,7 @@ def test_a_rollback_to_a_local_prefix_keeps_only_the_prefix(kit, real_manager, a assert mgr_b.revert_allocate_context(target) is True assert kit.kv(mgr_b, target) is kv assert (kv.capacity, kv.is_active, kv.num_committed_tokens) == (local, False, local) - resume_bounds = lender_b.resume_bounds(target) + resume_bounds = lender_b.poll_resume_bounds(target) assert resume_bounds == (local, local), ( f"resume_bounds {tuple(resume_bounds)} counts block 2, which the rollback freed" ) @@ -1913,25 +2094,23 @@ def test_a_generation_shrink_voids_the_fetched_rows_past_the_capacity( with real_manager() as mgr_a, real_manager(**SMALL_POOL) as mgr_b: source = kit.published(mgr_a, SOURCE, PROMPT) - lender_a = attach(mgr_a, tokens_per_fetch=END) - publish = lender_a.lend_read(source, 0, END) - publish_view = ready_view(publish, mgr_a) + lender_a = attach(mgr_a, guaranteed_fetch_tokens=END) + publish = lender_a.lend_read(source, 0, END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) source_pages = [kit.digest(kit.page(mgr_a, source, 0, b)) for b in range(BLOCKS)] target = kit.admitted(mgr_b, TARGET, PROMPT) target.py_draft_tokens = [0] * GENERATION_DRAFT assert mgr_b.try_allocate_generation(target) kv = kit.kv(mgr_b, target) assert kv.capacity == 1 + GENERATION_DRAFT - lender_b = attach(mgr_b, tokens_per_fetch=END) - lease = lender_b.lend_write(target, 0, END) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=END) + lease = lender_b.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() kit.fill_sentinel(mgr_b, target) - lease.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) - lease.release() + masks = kit.relay(lender_a, publish_view, lender_b, view) + assert return_and_settle(lender_b, mgr_b, target, lease, masks) == (END, 0) publish.release() - mgr_b._stream.synchronize() assert [kit.digest(kit.page(mgr_b, target, 0, b)) for b in range(BLOCKS)] == source_pages - assert lender_b.resume_bounds(target) == (END, 0) if path == "revert_allocate_generation": mgr_b.revert_allocate_generation(target) @@ -1954,7 +2133,7 @@ def test_a_generation_shrink_voids_the_fetched_rows_past_the_capacity( assert all(r != s for r, s in zip(regrown[kept:], source_pages[kept:])), ( "a regrown page still holds the fetched bytes" ) - resume_bounds = lender_b.resume_bounds(target) + resume_bounds = lender_b.poll_resume_bounds(target) assert resume_bounds == (kept * TPB, 0), ( f"resume_bounds {tuple(resume_bounds)} counts blocks whose pages now hold {regrown[kept:]}" ) @@ -2041,11 +2220,11 @@ def test_a_fetch_keeps_the_tokens_computed_before_it( conversation = "kept-chunk" if policy == "per_conversation" else None with real_manager(block_reuse_config={"policy": policy}) as mgr: target = computed(kit, mgr, TARGET, SPLIT_PROMPT, conversation=conversation) - lender = attach(mgr, tokens_per_fetch=2 * TPB) - assert lender.resume_bounds(target) == (LOCAL, LOCAL) + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) + assert lender.poll_resume_bounds(target) == (LOCAL, LOCAL) if outcome == "abandoned": - lease = lender.lend_write(target, LOCAL, 4 * TPB) - assert lender.resume_bounds(target) is None, "a fetch counts from its call" + lease = lender.lend_write(target, LOCAL, 4 * TPB, **kit.HOST_NAMED) + assert lender.poll_resume_bounds(target) is None, "a fetch counts from its call" lease.release() # before its first poll: no backend wrote expected = (LOCAL, LOCAL) else: @@ -2053,7 +2232,7 @@ def test_a_fetch_keeps_the_tokens_computed_before_it( assert holds(kit, mgr, target, 0, [2, 3], staged_bytes(kit, 0x5A)) expected = (4 * TPB, LOCAL) assert holds(kit, mgr, target, 0, [0, 1], chunk_bytes), "the fetch changed the chunk" - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds == expected, ( f"resume_bounds {tuple(resume_bounds)}: the tokens computed before the fetch were dropped" ) @@ -2067,7 +2246,7 @@ def test_a_windowed_fetch_counts_what_its_window_still_holds(kit, real_manager, target = computed(kit, mgr, TARGET, SPLIT_PROMPT) kv = kit.kv(mgr, target) sliding = kit.windows(mgr).index(96) - lender = attach(mgr, tokens_per_fetch=2 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) fetch_segment(kit, mgr, lender, target, LOCAL, 4 * TPB, 0x5A) assert kv.history_length == 4 * TPB, "the fetch did not move the history to its end" assert kit.stale_blocks(mgr, sliding, 4 * TPB) == (0, 1) @@ -2078,7 +2257,7 @@ def test_a_windowed_fetch_counts_what_its_window_still_holds(kit, real_manager, assert holds(kit, mgr, target, 1 - sliding, [2, 3], staged_bytes(kit, 0x5A)), ( "the full-attention group's fetched blocks do not hold the bytes staged for them" ) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds == (4 * TPB, 4 * TPB), ( f"resume_bounds {tuple(resume_bounds)}: the tokens computed before the fetch were dropped" ) @@ -2102,13 +2281,13 @@ def test_a_fetch_counts_the_rows_each_window_reads_whatever_its_start( windows, start, expected = WINDOW_STARTS[case] with real_manager(windows=windows) as mgr: target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) - lender = attach(mgr, tokens_per_fetch=4 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=4 * TPB) fetch_segment(kit, mgr, lender, target, start, 4 * TPB, 0x5A) sliding = kit.windows(mgr).index(WINDOW) assert holds(kit, mgr, target, sliding, [2, 3], staged_bytes(kit, 0x5A)), ( "the window's fetched blocks lack their staged bytes" ) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.min_position == 4 * TPB, f"resume_bounds {tuple(resume_bounds)}" if expected is None: assert resume_bounds.max_position < resume_bounds.min_position, ( @@ -2130,11 +2309,11 @@ def test_an_abandoned_windowed_fetch_claims_nothing(kit, real_manager, attach=at target = computed(kit, mgr, TARGET, SPLIT_PROMPT) kv = kit.kv(mgr, target) sliding = kit.windows(mgr).index(WINDOW) - lender = attach(mgr, tokens_per_fetch=2 * TPB) - lender.lend_write(target, LOCAL, 4 * TPB).release() # before its first poll + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) + lender.lend_write(target, LOCAL, 4 * TPB, **kit.HOST_NAMED).release() # before a poll assert kv.history_length == 4 * TPB, "the fetch did not move the history to its end" assert all(p < 0 for p in kit.pages(kv, sliding)[:2]), "the window kept the chunk" - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.min_position == 4 * TPB assert resume_bounds.max_position < resume_bounds.min_position, ( f"resume_bounds {tuple(resume_bounds)} claims a start whose blocks were released or never " @@ -2151,23 +2330,24 @@ def test_a_fetch_from_below_the_history_keeps_only_what_lies_below_it( with real_manager(**PER_REQUEST) as mgr: target = computed(kit, mgr, TARGET, SPLIT_PROMPT) - lender = attach(mgr, tokens_per_fetch=3 * TPB) - lease = lender.lend_write(target, TPB, 4 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=3 * TPB) + lease = lender.lend_write(target, TPB, 4 * TPB, **kit.HOST_NAMED) view = lease.poll() assert by_group(view) == {0: [1, 2, 3]} kit.stage(lender, view, 0x5A) driver = FailingDriver(_lender.drv, fail_on=2) + arrived = tuple(np.isin(run.page_ordinals, [1, 3]) for run in view.groups) with monkeypatch.context() as patched: patched.setattr(_lender, "drv", driver) # Blocks 1 and 3 arrived: their staging slots are apart, so each takes its own copy. - lease.mark_arrived(tuple(np.isin(run.block_ordinals, [1, 3]) for run in view.groups)) - lease.release() + given = kit.return_write(lender, lease, arrived) assert driver.calls >= 2, "a single copy: nothing failed partway" assert holds(kit, mgr, target, 0, [0], chunk_bytes) assert holds(kit, mgr, target, 0, [1], staged_bytes(kit, 0x5A)), ( "the copy did not bring block 1 its own row" ) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) + assert given == resume_bounds, f"the return of an abandoned fetch gave {given}" assert resume_bounds.max_position <= TPB, ( f"resume_bounds {tuple(resume_bounds)} counts block 1, which the abandoned fetch overwrote" ) @@ -2180,31 +2360,30 @@ def test_a_fetch_from_below_the_history_keeps_only_what_lies_below_it( def test_a_refetch_whose_copy_fails_partway_counts_none_of_the_rows_it_was_to_overwrite( kit, real_manager, monkeypatch, attach=attach ): - """``[0, 128)`` fetched and delivered, nothing committed; then ``[32, 128)`` fetched again with - blocks 1 and 3 marked, its copy failing after block 1's. Neither marked block counts any more, - since a copy failing partway can leave a row holding each fetch's bytes; block 0 still does.""" + """A refetch of blocks 1 to 3 whose copy fails after block 1's leaves only block 0 counted.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender with real_manager() as mgr: target = kit.admitted(mgr, TARGET, SPLIT_PROMPT) - lender = attach(mgr, tokens_per_fetch=4 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=4 * TPB) fetch_segment(kit, mgr, lender, target, 0, 4 * TPB, 0x11) - assert lender.resume_bounds(target) == (4 * TPB, 0) - lease = lender.lend_write(target, TPB, 4 * TPB) + assert lender.poll_resume_bounds(target) == (4 * TPB, 0) + lease = lender.lend_write(target, TPB, 4 * TPB, **kit.HOST_NAMED) view = lease.poll() assert by_group(view) == {0: [1, 2, 3]} kit.stage(lender, view, 0x22) driver = FailingDriver(_lender.drv, fail_on=2) + arrived = tuple(np.isin(run.page_ordinals, [1, 3]) for run in view.groups) with monkeypatch.context() as patched: patched.setattr(_lender, "drv", driver) # Blocks 1 and 3: their staging slots are apart, so each takes its own copy. - lease.mark_arrived(tuple(np.isin(run.block_ordinals, [1, 3]) for run in view.groups)) - lease.release() + given = kit.return_write(lender, lease, arrived) assert driver.calls >= 2, "a single copy: nothing failed partway" assert holds(kit, mgr, target, 0, [1], staged_bytes(kit, 0x22)), ( "the copy did not bring block 1 the second fetch's row" ) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) + assert given == resume_bounds, f"the return of an abandoned fetch gave {given}" assert resume_bounds.max_position <= TPB, ( f"resume_bounds {tuple(resume_bounds)} counts rows the failed copy was to overwrite" ) @@ -2235,15 +2414,7 @@ def discard_pending_stats(self): def test_a_copy_failing_over_the_block_the_committed_tokens_end_inside_empties_the_interval( kit, real_manager, monkeypatch, failure, attach=attach ): - """All-reusable: the target's local match ends inside block 1, and a fetch from block 1 marks - blocks 1 and 3. Block 1's copy is queued, and lands, before block 3's fails (the second memcpy - returns an error, or raises an ``Exception`` or an interruption that is none): the stricter - order for the rule, which reads the marks, not how far the copy got. Block 1 may then hold the - fetch's bytes in some pools and its old bytes in others, so the interval is empty, then and - later, though not for a cache that replaces it under the same id after a free that raised before - the lender's free hook, which keeps the records. The control counts block 1, which holds its - staged bytes. A lender tainting only the failing copy call's rows, or only a block the copy left - incomplete, keeps counting block 1.""" + """A copy failing past the block the committed tokens end inside empties the interval.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender with real_manager() as mgr: @@ -2253,8 +2424,8 @@ def test_a_copy_failing_over_the_block_the_committed_tokens_end_inside_empties_t assert kv.num_committed_tokens == len(SHARED) and len(SHARED) % TPB, ( "no partial match: the check proves nothing" ) - lender = attach(mgr, tokens_per_fetch=3 * TPB) - lease = lender.lend_write(target, TPB, 4 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=3 * TPB) + lease = lender.lend_write(target, TPB, 4 * TPB, **kit.HOST_NAMED) view = lease.poll() assert view is not None, lease.failure assert by_group(view) == {0: [1, 2, 3]} @@ -2262,27 +2433,28 @@ def test_a_copy_failing_over_the_block_the_committed_tokens_end_inside_empties_t planted = {"raised": MemoryError, "interrupted": Interrupted}.get(failure) driver = FailingDriver(_lender.drv, fail_on=2 if failure else None, raises=planted) # Blocks 1 and 3: their staging slots are apart, so each takes its own copy. - masks = tuple(np.isin(run.block_ordinals, [1, 3]) for run in view.groups) + masks = tuple(np.isin(run.page_ordinals, [1, 3]) for run in view.groups) with monkeypatch.context() as patched: patched.setattr(_lender, "drv", driver) if planted is None: - lease.mark_arrived(masks) + kit.return_write(lender, lease, masks) else: with pytest.raises(planted, match="planted"): - lease.mark_arrived(masks) + kit.return_write(lender, lease, masks) lease.release() mgr._stream.synchronize() + lender.progress() assert driver.calls >= 2, "a single copy: nothing failed partway" for run in view.groups: # block 1, in every layer group the view has assert holds(kit, mgr, target, run.layer_group_index, [1], staged_bytes(kit, 0x5A)), ( f"group {run.layer_group_index}: the copy did not reach block 1: the check proves nothing" ) if failure is None: - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds == (2 * TPB, TPB), f"resume_bounds {tuple(resume_bounds)}" return for _ in range(2): # and on a later call - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.max_position < resume_bounds.min_position, ( f"resume_bounds {tuple(resume_bounds)} {TAIL_MIXED}" ) @@ -2296,7 +2468,7 @@ def test_a_copy_failing_over_the_block_the_committed_tokens_end_inside_empties_t kv.close() mgr.index_mapper.remove_sequence(TARGET) target = kit.admitted(mgr, TARGET, MATCHED_PROMPT) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.max_position >= resume_bounds.min_position, ( f"resume_bounds {tuple(resume_bounds)}: the cache that replaced the failed one stays empty" ) @@ -2305,10 +2477,7 @@ def test_a_copy_failing_over_the_block_the_committed_tokens_end_inside_empties_t def test_a_copy_failing_over_that_block_in_a_later_layer_group_empties_the_interval( kit, real_manager, monkeypatch, attach=attach ): - """All-reusable with two layer groups: the target's local match ends inside block 1, which a - window keeps. ``[32, 128)`` is delivered whole, then fetched again with blocks 1 and 3 marked in - the second layer group alone, its copy failing after block 1's. The rule reads the marks in - every layer group, so the interval is empty.""" + """Arrivals in a later layer group alone, whose copy fails past block 1, empty the interval.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender with real_manager(windows=[4 * TPB, 256]) as mgr: @@ -2317,13 +2486,13 @@ def test_a_copy_failing_over_that_block_in_a_later_layer_group_empties_the_inter assert kit.kv(mgr, target).num_committed_tokens == len(SHARED), ( "no partial match: the check proves nothing" ) - lender = attach(mgr, tokens_per_fetch=3 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=3 * TPB) fetch_segment(kit, mgr, lender, target, TPB, 4 * TPB, 0x11) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.max_position >= resume_bounds.min_position, ( f"resume_bounds {tuple(resume_bounds)} is empty before the refetch: the check proves nothing" ) - lease = lender.lend_write(target, TPB, 4 * TPB) + lease = lender.lend_write(target, TPB, 4 * TPB, **kit.HOST_NAMED) view = lease.poll() assert view is not None, lease.failure assert by_group(view) == {0: [1, 2, 3], 1: [1, 2, 3]} @@ -2333,12 +2502,11 @@ def test_a_copy_failing_over_that_block_in_a_later_layer_group_empties_the_inter first, *later = view.groups masks = ( np.zeros(len(first), dtype=bool), - *(np.isin(r.block_ordinals, [1, 3]) for r in later), + *(np.isin(r.page_ordinals, [1, 3]) for r in later), ) with monkeypatch.context() as patched: patched.setattr(_lender, "drv", driver) - lease.mark_arrived(masks) - lease.release() + kit.return_write(lender, lease, masks) mgr._stream.synchronize() assert driver.calls >= 2, "a single copy: nothing failed partway" assert holds(kit, mgr, target, first.layer_group_index, [1], staged_bytes(kit, 0x11)), ( @@ -2347,7 +2515,7 @@ def test_a_copy_failing_over_that_block_in_a_later_layer_group_empties_the_inter assert holds(kit, mgr, target, later[0].layer_group_index, [1], staged_bytes(kit, 0x5A)), ( "the copy did not reach block 1 of the second layer group: the check proves nothing" ) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.max_position < resume_bounds.min_position, ( f"resume_bounds {tuple(resume_bounds)} {TAIL_MIXED}" ) @@ -2360,13 +2528,14 @@ def test_a_second_fetch_keeps_what_the_first_kept(kit, real_manager, attach=atta with real_manager(windows=[WINDOW, 256], **PER_REQUEST) as mgr: target = computed(kit, mgr, TARGET, WINDOWED_PROMPT) kv = kit.kv(mgr, target) - lender = attach(mgr, tokens_per_fetch=2 * TPB) - lender.lend_write(target, LOCAL, 4 * TPB).release() # before its first poll + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) + lender.lend_write(target, LOCAL, 4 * TPB, **kit.HOST_NAMED).release() # before a poll assert kv.history_length == 4 * TPB + lender.progress() # the abandoned fetch's slots come back fetch_segment(kit, mgr, lender, target, 4 * TPB, WINDOWED_END, 0x5A) for lg in range(kit.num_layer_groups(mgr)): assert holds(kit, mgr, target, lg, [4], staged_bytes(kit, 0x5A)), f"layer group {lg}" - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.min_position == WINDOWED_END assert resume_bounds.max_position < resume_bounds.min_position, ( f"resume_bounds {tuple(resume_bounds)} counts block 3, which nothing computed or fetched" @@ -2382,11 +2551,11 @@ def test_a_rollback_keeps_the_computed_chunk_and_voids_the_fetches( target = computed(kit, mgr, TARGET, SPLIT_PROMPT, upto=TPB) target.py_ctx_pre_resize_cap = None # a later resize grew nothing kv = kit.kv(mgr, target) - lender = attach(mgr, tokens_per_fetch=TPB) + lender = attach(mgr, guaranteed_fetch_tokens=TPB) fetch_segment(kit, mgr, lender, target, TPB, 2 * TPB, 0x11) - assert lender.resume_bounds(target) == (2 * TPB, TPB) + assert lender.poll_resume_bounds(target) == (2 * TPB, TPB) fetch_segment(kit, mgr, lender, target, 2 * TPB, 3 * TPB, 0x22) - assert lender.resume_bounds(target) == (3 * TPB, TPB) + assert lender.poll_resume_bounds(target) == (3 * TPB, TPB) for lg in range(kit.num_layer_groups(mgr)): # the rows that resume_bounds counts for ordinal, byte in ((1, 0x11), (2, 0x22)): assert holds(kit, mgr, target, lg, [ordinal], staged_bytes(kit, byte)), ( @@ -2396,7 +2565,7 @@ def test_a_rollback_keeps_the_computed_chunk_and_voids_the_fetches( assert mgr.revert_allocate_context(target) is True assert kit.kv(mgr, target) is kv, "the rollback replaced the cache" assert (kv.capacity, kv.history_length, kv.num_blocks) == (TPB, TPB, 1) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds == (TPB, TPB), ( f"resume_bounds {tuple(resume_bounds)} counts block 1, which the rollback freed" ) @@ -2510,14 +2679,14 @@ def test_resume_bounds_claims_only_starts_the_context_update_takes( assert mgr.prepare_context(target) kit.kv(mgr, target).enable_swa_scratch_reuse = False # a fetch target needs it off kv = kit.kv(mgr, target) - lender = attach(mgr, tokens_per_fetch=end - start) + lender = attach(mgr, guaranteed_fetch_tokens=end - start) fetch_segment(kit, mgr, lender, target, start, end, 0x5A, arrived) delivered = range(start // TPB, end // TPB) if arrived is None else arrived for lg in range(kit.num_layer_groups(mgr)): assert holds(kit, mgr, target, lg, delivered, staged_bytes(kit, 0x5A)), f"group {lg}" history = kv.history_length assert not stale_anywhere(kit, mgr, history), "a window released blocks" - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) floor = resume_bounds.min_position if resume_bounds.max_position >= floor: # The lowest start claimed: from a higher one the chunk ends no lower. @@ -2545,9 +2714,9 @@ def test_a_later_fetch_keeps_what_the_request_computed_after_resuming( with real_manager(**PER_REQUEST) as mgr: target = computed(kit, mgr, TARGET, LONG_PROMPT) kv = kit.kv(mgr, target) - lender = attach(mgr, tokens_per_fetch=2 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) fetch_segment(kit, mgr, lender, target, LOCAL, 4 * TPB, 0x5A) - first = lender.resume_bounds(target) + first = lender.poll_resume_bounds(target) assert first == (4 * TPB, LOCAL), f"resume_bounds {tuple(first)} after the first fetch" resumed(kit, mgr, target, first.max_position) assert next_chunk(kit, mgr, target, 2 * TPB) is None @@ -2556,7 +2725,7 @@ def test_a_later_fetch_keeps_what_the_request_computed_after_resuming( assert holds(kit, mgr, target, 0, [0, 1, 4, 5], chunk_bytes) assert holds(kit, mgr, target, 0, [2, 3], staged_bytes(kit, 0x5A)) assert holds(kit, mgr, target, 0, [6], staged_bytes(kit, 0x6B)) - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds == (7 * TPB, 6 * TPB), ( f"resume_bounds {tuple(resume_bounds)}: the tokens computed after resuming were dropped" ) @@ -2605,8 +2774,8 @@ def test_a_joint_draft_pool_resumes_where_both_lenders_allow( ) # The draft pool under a scope of its own, as attach_lender asks of such a pair. lenders = [ - attach(target, tokens_per_fetch=PAIRED_END), - attach(draft, tokens_per_fetch=PAIRED_END, scope=SCOPE + b"/draft-pool"), + attach(target, guaranteed_fetch_tokens=PAIRED_END), + attach(draft, guaranteed_fetch_tokens=PAIRED_END, scope=SCOPE + b"/draft-pool"), ] request = kit.make_request(TARGET, SPLIT_PROMPT) assert schedule(scheduler, request) == [TARGET] @@ -2615,13 +2784,12 @@ def test_a_joint_draft_pool_resumes_where_both_lenders_allow( request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS # parked to fetch intervals = [] for mgr, lender, delivered in zip(pools, lenders, (True, draft_delivers)): - lease = lender.lend_write(request, 0, PAIRED_END) + lease = lender.lend_write(request, 0, PAIRED_END, **kit.HOST_NAMED) view = lease.poll() kit.stage(lender, view, byte) - lease.mark_arrived(view.row_masks(delivered)) - lease.release() - mgr._stream.synchronize() - intervals.append(lender.resume_bounds(request)) + intervals.append( + return_and_settle(lender, mgr, request, lease, view.row_masks(delivered)) + ) # Each interval covers only its own pool's blocks: the request resumes where both allow. usable = min(interval.max_position for interval in intervals) assert max(interval.min_position for interval in intervals) <= usable @@ -2703,21 +2871,19 @@ def test_a_windowed_joint_draft_pool_resumes_only_where_its_context_resize_holds draft_kv_cache_manager=draft, ) lenders = [ - attach(target, tokens_per_fetch=PAIRED_END), - attach(draft, tokens_per_fetch=PAIRED_END, scope=SCOPE + b"/draft-pool"), + attach(target, guaranteed_fetch_tokens=PAIRED_END), + attach(draft, guaranteed_fetch_tokens=PAIRED_END, scope=SCOPE + b"/draft-pool"), ] request = kit.make_request(TARGET, SPLIT_PROMPT) assert schedule(scheduler, request) == [TARGET] request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS # parked to fetch intervals = [] for mgr, lender, ordinals in zip(pools, lenders, JOINT_DELIVERIES[case]): - lease = lender.lend_write(request, 0, PAIRED_END) - view = ready_view(lease, mgr) + lease = lender.lend_write(request, 0, PAIRED_END, **kit.HOST_NAMED) + view = kit.ready_view(lease, mgr) kit.stage(lender, view, byte) - lease.mark_arrived(tuple(np.isin(run.block_ordinals, ordinals) for run in view.groups)) - lease.release() - mgr._stream.synchronize() - intervals.append(lender.resume_bounds(request)) + arrived = tuple(np.isin(run.page_ordinals, ordinals) for run in view.groups) + intervals.append(return_and_settle(lender, mgr, request, lease, arrived)) floor = max(interval.min_position for interval in intervals) usable = min(interval.max_position for interval in intervals) if floor <= usable: @@ -2851,7 +3017,8 @@ def test_a_request_past_its_first_chunk_resumes_where_the_documented_step_puts_i ) scopes = (SCOPE, SCOPE + b"/draft-pool") # the draft pool under a scope of its own lenders = [ - attach(m, tokens_per_fetch=RESUME_END - LOCAL, scope=s) for m, s in zip(pools, scopes) + attach(m, guaranteed_fetch_tokens=RESUME_END - LOCAL, scope=s) + for m, s in zip(pools, scopes) ] request = kit.make_request(TARGET, LONG_PROMPT) assert schedule(scheduler, request) == [TARGET] @@ -2863,7 +3030,7 @@ def test_a_request_past_its_first_chunk_resumes_where_the_documented_step_puts_i intervals = [] for mgr, lender in zip(pools, lenders): fetch_segment(kit, mgr, lender, request, LOCAL, RESUME_END, byte) - intervals.append(tuple(lender.resume_bounds(request))) + intervals.append(tuple(lender.poll_resume_bounds(request))) assert intervals == [(RESUME_END, LOCAL)] * len(pools), f"resume_bounds {intervals}" at = RESUME_END request.state = LlmRequestState.CONTEXT_INIT @@ -2911,7 +3078,7 @@ def test_a_request_resumes_at_an_unaligned_point_past_its_context_position( ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, TPB), ) sibling = kit.published(mgr, SOURCE, SIBLING_PROMPT) - lender = attach(mgr, tokens_per_fetch=TPB) + lender = attach(mgr, guaranteed_fetch_tokens=TPB) request = kit.make_request(TARGET, prompt) assert schedule(scheduler, request) == [TARGET] kv = kit.kv(mgr, request) @@ -2922,7 +3089,7 @@ def test_a_request_resumes_at_an_unaligned_point_past_its_context_position( assert request.context_current_position == position, "the first chunk was not LOCAL tokens" request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS # parked to fetch fetch_segment(kit, mgr, lender, request, 3 * TPB, 4 * TPB, byte) - resume_bounds = lender.resume_bounds(request) + resume_bounds = lender.poll_resume_bounds(request) assert tuple(resume_bounds) == (4 * TPB, 3 * TPB), f"resume_bounds {tuple(resume_bounds)}" at = resume_bounds.min_position if below else SHIFTED_AT request.state = LlmRequestState.CONTEXT_INIT @@ -2966,7 +3133,7 @@ def test_a_lease_failed_at_the_call_changed_nothing(kit, real_manager, state): with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) target = kit.admitted(mgr, TARGET, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) if state == "no_cache": source, target = kit.make_request(8, PROMPT), kit.make_request(9, OTHER_PROMPT) elif state == "suspended": @@ -2976,7 +3143,13 @@ def test_a_lease_failed_at_the_call_changed_nothing(kit, real_manager, state): else: mgr.shutdown() before = [state_of(kit, mgr, lender, r) for r in (source, target)] - leases = [lender.lend_read(source, 0, END), lender.lend_write(target, 0, END)] + leases = [] + for lend, request in ((lender.lend_read, source), (lender.lend_write, target)): + if state == "shut_down": + with pytest.raises(RuntimeError, match="lends no more"): + lend(request, 0, END, **kit.HOST_NAMED) + else: + leases.append(lend(request, 0, END, **kit.HOST_NAMED)) for lease in leases: assert lease.failure is not None, "not failed at the call" assert lease.poll() is None @@ -2989,20 +3162,19 @@ def test_a_lease_failed_at_the_call_changed_nothing(kit, real_manager, state): def test_a_write_the_pool_cannot_grow_fails_at_the_call_and_changed_nothing(kit, real_manager): with real_manager(max_tokens=kit.POOL_TOKENS) as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) + lender = attach(mgr, guaranteed_fetch_tokens=END) others = kit.Requests(mgr) assert others.allocate(kit.pool_pages(mgr)), "others could not take the whole pool" before = state_of(kit, mgr, lender, target) - lease = lender.lend_write(target, 0, END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) assert "no free pages" in str(lease.failure) and lease.poll() is None assert state_of(kit, mgr, lender, target) == before lease.release() others.free() - lease = lender.lend_write(target, 0, END) # with pages free again it goes through + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) # pages free: it goes through view = lease.poll() assert view is not None - lease.mark_arrived(view.row_masks()) - lease.release() + kit.return_write(lender, lease, view.row_masks()) def test_a_write_the_pool_cannot_grow_leaves_a_computed_prefix_as_it_was( @@ -3013,13 +3185,13 @@ def test_a_write_the_pool_cannot_grow_leaves_a_computed_prefix_as_it_was( fails at the call, and the cache and its resume_bounds stay as they were: the chunk still counts.""" with real_manager(max_tokens=kit.POOL_TOKENS, **PER_REQUEST) as mgr: target = computed(kit, mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=2 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) others = kit.Requests(mgr) free = kit.pool_pages(mgr) - int(kit.kv(mgr, target).num_blocks) # past the chunk's pages assert others.allocate(free), "others could not take the rest of the pool" before = state_of(kit, mgr, lender, target) assert before[1] == (LOCAL, LOCAL), f"resume_bounds before the fetch {before[1]}" - lease = lender.lend_write(target, TPB, 3 * TPB) + lease = lender.lend_write(target, TPB, 3 * TPB, **kit.HOST_NAMED) assert "no free pages" in str(lease.failure) and lease.poll() is None, lease.failure after = state_of(kit, mgr, lender, target) assert after == before, f"the failed grow changed the cache or its resume_bounds: {after}" @@ -3027,14 +3199,16 @@ def test_a_write_the_pool_cannot_grow_leaves_a_computed_prefix_as_it_was( others.free() +@pytest.mark.parametrize("keywords", ["HOST_NAMED", "DEVICE_NAMED"], ids=["host", "device"]) +@pytest.mark.parametrize("poll", [True, False], ids=["poll_failure", "release_before_poll"]) def test_a_write_left_with_a_block_without_a_page_fails_at_its_first_poll( - kit, real_manager, monkeypatch + kit, real_manager, monkeypatch, keywords, poll ): from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager with real_manager() as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) + lender = attach(mgr, guaranteed_fetch_tokens=END) kv = kit.kv(mgr, target) def without_block_one(read): @@ -3049,11 +3223,15 @@ def pages(kv_cache, lg): with monkeypatch.context() as patched: for name in ("pages", "locked_pages"): patched.setattr(_manager, name, without_block_one(getattr(_manager, name))) - lease = lender.lend_write(target, 0, END) + lease = lender.lend_write(target, 0, END, **getattr(kit, keywords)) assert lease.failure is None, "the cache grew: nothing fails at the call after that" - assert lender.resume_bounds(target) is None - assert lease.poll() is None and lease.failure is not None - assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + assert lender.poll_resume_bounds(target) is None + if poll: + assert lease.poll() is None and lease.failure is not None + else: + lease.release() + assert lease.failure is None + assert lender.poll_resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) lease.release() @@ -3068,32 +3246,32 @@ def test_a_write_whose_grant_raises_within_the_call_fails_at_its_first_poll( left without a page. A read changes no cache, so one whose grant raises fails at the call.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender - view = _lender.Staging._build_view + view = _lender._ManagerLender._build_view planted = [] - def raising_once(self, rows, keys): + def raising_once(self, rows, keys, *bounds): if not planted: planted.append(True) raise MemoryError("planted") - return view(self, rows, keys) + return view(self, rows, keys, *bounds) with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) target = kit.admitted(mgr, TARGET, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END) + lender = attach(mgr, guaranteed_fetch_tokens=END) kv = kit.kv(mgr, target) with monkeypatch.context() as patched: - patched.setattr(_lender.Staging, "_build_view", raising_once) - read = lender.lend_read(source, 0, END) + patched.setattr(_lender._ManagerLender, "_build_view", raising_once) + read = lender.lend_read(source, 0, END, **kit.HOST_NAMED) assert planted and "planted" in str(read.failure), "the read's grant did not raise" planted.clear() - lease = lender.lend_write(target, 0, END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) assert planted, "the write's grant ran after the call: the check proves nothing" assert kv.capacity >= END, "the write did not grow the cache" assert lease.failure is None, f"{GROWN_THEN_FAILED}: {lease.failure}" - assert lender.resume_bounds(target) is None + assert lender.poll_resume_bounds(target) is None assert lease.poll() is None and "planted" in str(lease.failure), lease.failure - assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + assert lender.poll_resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) lease.release() read.release() @@ -3128,22 +3306,20 @@ def test_a_fetch_from_a_partially_matched_block(kit, real_manager, case, attach= stale_beg, stale_end = kit.stale_blocks(mgr, window_lg, end) leaves = stale_beg <= 1 < stale_end assert leaves == (case == "window_leaves_it") - lender = attach(mgr, tokens_per_fetch=end - start) + lender = attach(mgr, guaranteed_fetch_tokens=end - start) before = (kv.capacity, kv.history_length, kv.num_committed_tokens) - lease = lender.lend_write(target, start, end) + lease = lender.lend_write(target, start, end, **kit.HOST_NAMED) if lease.failure is None: # Granted: the fetch lands whole, and the target resumes where resume_bounds allows. - view = ready_view(lease, mgr) + view = kit.ready_view(lease, mgr) kit.stage(lender, view, 0x5A) - lease.mark_arrived(view.row_masks(True)) - lease.release() - mgr._stream.synchronize() + return_and_settle(lender, mgr, target, lease, view.row_masks(True)) for run in view.groups: # every row the resume reads - ordinals = run.block_ordinals.tolist() + ordinals = run.page_ordinals.tolist() assert holds( kit, mgr, target, run.layer_group_index, ordinals, staged_bytes(kit, 0x5A) ), f"layer group {run.layer_group_index}: a fetched row lacks its staged bytes" - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.max_position >= resume_bounds.min_position, tuple(resume_bounds) resumed(kit, mgr, target, resume_bounds.min_position) else: @@ -3191,22 +3367,20 @@ def test_a_fetch_past_pages_grown_before_it_fails_at_the_call(kit, real_manager, assert stale_beg <= 1 < stale_end and kit.pages(kv, window_lg)[1] >= 0, ( "block 1 has no page or the window keeps it: the check proves nothing" ) - lender = attach(mgr, tokens_per_fetch=end) + lender = attach(mgr, guaranteed_fetch_tokens=end) before = (kv.capacity, kv.history_length, kv.num_committed_tokens) - lease = lender.lend_write(target, 0, end) + lease = lender.lend_write(target, 0, end, **kit.HOST_NAMED) if lease.failure is None: # Granted: the fetch lands whole, and the target resumes where resume_bounds allows. - view = ready_view(lease, mgr) + view = kit.ready_view(lease, mgr) kit.stage(lender, view, 0x5A) - lease.mark_arrived(view.row_masks(True)) - lease.release() - mgr._stream.synchronize() + return_and_settle(lender, mgr, target, lease, view.row_masks(True)) for run in view.groups: # every row the resume reads - ordinals = run.block_ordinals.tolist() + ordinals = run.page_ordinals.tolist() assert holds( kit, mgr, target, run.layer_group_index, ordinals, staged_bytes(kit, 0x5A) ), f"layer group {run.layer_group_index}: a fetched row lacks its staged bytes" - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.max_position >= resume_bounds.min_position, tuple(resume_bounds) resumed(kit, mgr, target, resume_bounds.min_position) else: @@ -3253,19 +3427,17 @@ def test_a_fetch_past_pages_grown_before_it_goes_through_where_the_commit_keeps_ assert stale_beg <= 1 < stale_end and kit.pages(kv, window_lg)[1] >= 0, ( "block 1 has no page or the window keeps it: the check proves nothing" ) - lender = attach(mgr, tokens_per_fetch=end) - lease = lender.lend_write(target, 0, end) + lender = attach(mgr, guaranteed_fetch_tokens=end) + lease = lender.lend_write(target, 0, end, **kit.HOST_NAMED) assert lease.failure is None, ( f"the fetch failed although the commit keeps no page a window passed: {lease.failure}" ) - view = ready_view(lease, mgr) + view = kit.ready_view(lease, mgr) kit.stage(lender, view, 0x5A) - lease.mark_arrived(view.row_masks(True)) - lease.release() - mgr._stream.synchronize() + return_and_settle(lender, mgr, target, lease, view.row_masks(True)) passed = kit.pages(kv, window_lg)[stale_beg:stale_end] assert all(slot < 0 for slot in passed), f"the window's passed blocks kept pages: {passed}" - resume_bounds = lender.resume_bounds(target) + resume_bounds = lender.poll_resume_bounds(target) assert resume_bounds.max_position == end, tuple(resume_bounds) resumed(kit, mgr, target, resume_bounds.max_position) kept = [o for o in range(end // TPB) if not stale_beg <= o < stale_end] @@ -3297,21 +3469,16 @@ def test_a_later_lease_past_window_rows_an_earlier_one_missed_fails_at_the_call( # Block 0 first, so the check starts past it and reads block 1's delivered row there. assert next_chunk(kit, mgr, target, tokens=TPB) is None assert kv.num_committed_tokens == TPB, "block 0 was not committed" - lender = attach(mgr, tokens_per_fetch=2 * TPB) - first = lender.lend_write(target, TPB, 2 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) + first = lender.lend_write(target, TPB, 2 * TPB, **kit.HOST_NAMED) kit.fill_sentinel(mgr, target, 1) # what a page nothing wrote holds - view = ready_view(first, mgr) + view = kit.ready_view(first, mgr) kit.stage(lender, view, 0x11) delivered = case == "delivered" - first.mark_arrived( - tuple( - np.full(len(run), delivered or run.layer_group_index != sliding) - for run in view.groups - ) + arrived = tuple( + np.full(len(run), delivered or run.layer_group_index != sliding) for run in view.groups ) - first.release() - mgr._stream.synchronize() - resume_bounds = lender.resume_bounds(target) + resume_bounds = return_and_settle(lender, mgr, target, first, arrived) assert resume_bounds == ((2 if delivered else 1) * TPB, TPB), tuple(resume_bounds) if case == "computed": # from where resume_bounds lets it resume, through the lease's end resumed(kit, mgr, target, resume_bounds.max_position) @@ -3320,7 +3487,7 @@ def test_a_later_lease_past_window_rows_an_earlier_one_missed_fails_at_the_call( assert kit.stale_blocks(mgr, sliding, 4 * TPB) == (0, 2), "nothing left behind: no check" assert min(kit.pages(kv, sliding)[:2]) >= 0, "blocks 0 and 1 have no page: no check" before = state_of(kit, mgr, lender, target) - later = lender.lend_write(target, 2 * TPB, 4 * TPB) + later = lender.lend_write(target, 2 * TPB, 4 * TPB, **kit.HOST_NAMED) if case == "missed": assert later.failure is not None and "never wrote" in later.failure, ( f"a later lease was lent past window rows an earlier one missed: {later.failure}" @@ -3331,12 +3498,9 @@ def test_a_later_lease_past_window_rows_an_earlier_one_missed_fails_at_the_call( ) else: assert later.failure is None, f"the later lease failed: {later.failure}" - view = ready_view(later, mgr) + view = kit.ready_view(later, mgr) kit.stage(lender, view, 0x22) - later.mark_arrived(view.row_masks(True)) - later.release() - mgr._stream.synchronize() - resume_bounds = lender.resume_bounds(target) + resume_bounds = return_and_settle(lender, mgr, target, later, view.row_masks(True)) assert resume_bounds == (4 * TPB, 4 * TPB), tuple(resume_bounds) # The target computes the rest of its prompt from where resume_bounds lets it resume, and the # context update commits it. @@ -3368,20 +3532,27 @@ def test_a_read_refuses_a_bad_range_at_the_call(kit, real_manager): # Sized for the rest of its prompt, as before its next chunk: block 2 has a page, no KV. assert mgr.resize_context(sized, sized.context_remaining_length) assert kit.kv(mgr, sized).num_committed_tokens == 2 * TPB - lender = attach(mgr, tokens_per_fetch=4 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=4 * TPB) for request, start, end, why in ( (source, -TPB, TPB, "bad token range"), (source, TPB, 0, "bad token range"), (source, 0, 50, "whole blocks"), (source, 16, 48, "whole blocks"), + ): + with pytest.raises(ValueError, match=why): + lender.lend_read(request, start, end, **kit.HOST_NAMED) + # Past what this rank committed: another pipeline stage may hold it, so the lease fails. + for request, start, end, why in ( (source, 0, 4 * TPB, "past the 97 committed tokens"), (source, 0, 1 << 50, "past the 97 committed tokens"), # of any length (sized, 0, END, "past the 64 committed tokens"), # inside the prompt's whole blocks ): - with pytest.raises(ValueError, match=why): - lender.lend_read(request, start, end) - lease = lender.lend_read(source, 0, END) - ready_view(lease, mgr) + lease = lender.lend_read(request, start, end, **kit.HOST_NAMED) + assert lease.failure is not None and why in lease.failure, (why, lease.failure) + assert lease.poll() is None + lease.release() + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + kit.ready_view(lease, mgr) lease.release() @@ -3391,30 +3562,35 @@ def test_a_write_refuses_a_bad_range_before_changing_anything(kit, real_manager) matched = kit.admitted(mgr, TARGET, PROMPT) # its first two blocks match locally fresh = kit.admitted(mgr, OTHER, OTHER_PROMPT) assert kit.kv(mgr, matched).num_committed_tokens == 2 * TPB - lender = attach(mgr, tokens_per_fetch=4 * TPB, planned_fetches=2) + lender = attach(mgr, guaranteed_fetch_tokens=4 * TPB, capacity_in_fetches=2) before = [state_of(kit, mgr, lender, r) for r in (matched, fresh)] bad = [ (fresh, -TPB, TPB, "bad token range"), (fresh, TPB, 0, "bad token range"), (fresh, 0, 50, "covers whole blocks"), (fresh, 16, 48, "covers whole blocks"), - (matched, TPB, END, "inside the committed whole blocks"), (fresh, 0, 4 * TPB, "past the whole blocks before the request's last prompt token"), ] for request, start, end, why in bad: with pytest.raises(ValueError, match=why): - lender.lend_write(request, start, end) + lender.lend_write(request, start, end, **kit.HOST_NAMED) assert [state_of(kit, mgr, lender, r) for r in (matched, fresh)] == before - leases = [lender.lend_write(matched, 2 * TPB, END), lender.lend_write(fresh, 0, END)] + # A fetch from inside what this rank committed lends from the committed blocks' end on. + leases = [ + lender.lend_write(matched, TPB, END, **kit.HOST_NAMED), + lender.lend_write(fresh, 0, END, **kit.HOST_NAMED), + ] unsettled = state_of(kit, mgr, lender, fresh) - again = lender.lend_write(fresh, 0, END) # one unsettled fetch per cache, on this rank + again = lender.lend_write(fresh, 0, END, **kit.HOST_NAMED) # one unsettled fetch per cache assert again.failure is not None and "unsettled" in again.failure, again.failure assert again.poll() is None and state_of(kit, mgr, lender, fresh) == unsettled again.release() for lease in leases: view = lease.poll() - lease.mark_arrived(view.row_masks()) - lease.release() + if lease is leases[0]: + lent = sorted({o for g in view.groups for o in g.page_ordinals.tolist()}) + assert lent and min(lent) == 2, f"a fetch lent committed blocks: {lent}" + kit.return_write(lender, lease, view.row_masks()) ALIGNED_PROMPT = list(range(11000, 11128)) # exactly four whole blocks @@ -3428,34 +3604,58 @@ def test_a_fetch_ends_at_the_whole_blocks_before_the_last_prompt_token(kit, real with real_manager(windows=[WINDOW, 256] if windowed else None) as mgr: target = kit.admitted(mgr, TARGET, ALIGNED_PROMPT) end = len(ALIGNED_PROMPT) - lender = attach(mgr, tokens_per_fetch=end) + lender = attach(mgr, guaranteed_fetch_tokens=end) before = state_of(kit, mgr, lender, target) for past in (end, 1 << 50): # 1 << 50: whole blocks whose ordinals would take 256 TiB with pytest.raises(ValueError, match="past the whole blocks before"): - lender.lend_write(target, 0, past) + lender.lend_write(target, 0, past, **kit.HOST_NAMED) assert state_of(kit, mgr, lender, target) == before - lease = lender.lend_write(target, 0, end - TPB) # the block of the last token left out + lease = lender.lend_write(target, 0, end - TPB, **kit.HOST_NAMED) # the last block left out view = lease.poll() assert view is not None - lease.mark_arrived(view.row_masks()) - lease.release() + kit.return_write(lender, lease, view.row_masks()) def test_a_replaced_cache_takes_a_new_fetch(kit, real_manager): - with real_manager() as mgr: - target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) - first = lender.lend_write(target, 0, END) - first_view = first.poll() # unsettled: seen ready, not marked - mgr.free_resources(target) + with real_manager( + max_tokens=kit.POOL_TOKENS, host_cache_size=0, max_batch_size=64, max_util_for_resume=1.0 + ) as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - second = lender.lend_write(target, 0, END) - view = second.poll() - assert view is not None and lender.resume_bounds(target) is None - second.mark_arrived(view.row_masks()) - first.mark_arrived(first_view.row_masks()) - for lease in (first, second): - lease.release() + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) + first = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + dev = kit.DevicePages(mgr) + others = kit.Requests(mgr, dev) + second = None + try: + first_view = first.poll() + assert first_view is not None + old_pages = set(kit.pages(kit.kv(mgr, target), 0)) + assert len(old_pages) == BLOCKS and min(old_pages) >= 0 + kit.stage(lender, first_view, 0x5A) + remaining = kit.pool_pages(mgr) - len(old_pages) + assert remaining < mgr.max_batch_size, "fixture has too few request slots" + assert others.allocate(remaining, fill=False, chunk=1) + occupied = others.pages() + assert old_pages.isdisjoint(occupied) + assert len(old_pages | occupied) == kit.pool_pages(mgr) + mgr.free_resources(target) + torch.cuda.synchronize() + target = kit.admitted(mgr, TARGET, PROMPT) + second = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + view = second.poll() + assert view is not None and lender.poll_resume_bounds(target) is None + assert old_pages == set(kit.pages(kit.kv(mgr, target), 0)), "old pages were not reused" + kit.fill_sentinel(mgr, target) + assert kit.return_write(lender, first, first_view.row_masks(True)) is None + sentinel = bytes([kit.SENTINEL]) * dev.page_bytes(0) + assert all(dev.read(0, s) == sentinel for s in old_pages), "old arrivals landed" + assert lender.poll_resume_bounds(target) is None + kit.return_write(lender, second, view.row_masks()) + finally: + first.release() + if second is not None: + second.release() + others.free() def test_a_windowed_target_with_scratch_reuse_fails_at_the_call(kit, real_manager, attach=attach): @@ -3464,10 +3664,10 @@ def test_a_windowed_target_with_scratch_reuse_fails_at_the_call(kit, real_manage assert mgr.prepare_context(target) kv = kit.kv(mgr, target) assert kv.enable_swa_scratch_reuse, "the target must start with scratch reuse on" - lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + lender = attach(mgr, guaranteed_fetch_tokens=WINDOWED_END) before = state_of(kit, mgr, lender, target) try: - lease = lender.lend_write(target, 0, WINDOWED_END) + lease = lender.lend_write(target, 0, WINDOWED_END, **kit.HOST_NAMED) except RuntimeError as error: # the manager refusing the grow a lender let through raise AssertionError(f"not failed at the call: {error}") from error assert lease.failure is not None and "scratch" in lease.failure, "not failed at the call" @@ -3475,11 +3675,10 @@ def test_a_windowed_target_with_scratch_reuse_fails_at_the_call(kit, real_manage assert state_of(kit, mgr, lender, target) == before lease.release() kv.enable_swa_scratch_reuse = False - lease = lender.lend_write(target, 0, WINDOWED_END) # with it off the write goes through + lease = lender.lend_write(target, 0, WINDOWED_END, **kit.HOST_NAMED) # off: it goes through view = lease.poll() assert view is not None - lease.mark_arrived(view.row_masks()) - lease.release() + kit.return_write(lender, lease, view.row_masks()) def test_a_windowless_target_with_scratch_reuse_fails_at_the_call(kit, real_manager, attach=attach): @@ -3492,10 +3691,10 @@ def test_a_windowless_target_with_scratch_reuse_fails_at_the_call(kit, real_mana assert mgr.prepare_context(target) kv = kit.kv(mgr, target) assert kv.enable_swa_scratch_reuse, "the target must start with scratch reuse on" - lender = attach(mgr, tokens_per_fetch=PAIRED_END) + lender = attach(mgr, guaranteed_fetch_tokens=PAIRED_END) before = state_of(kit, mgr, lender, target) try: - lease = lender.lend_write(target, 0, PAIRED_END) + lease = lender.lend_write(target, 0, PAIRED_END, **kit.HOST_NAMED) except RuntimeError as error: # the manager refusing the grow a lender let through raise AssertionError(f"not failed at the call: {error}") from error try: @@ -3507,11 +3706,10 @@ def test_a_windowless_target_with_scratch_reuse_fails_at_the_call(kit, real_mana finally: lease.release() kv.enable_swa_scratch_reuse = False - lease = lender.lend_write(target, 0, PAIRED_END) # with it off the write goes through + lease = lender.lend_write(target, 0, PAIRED_END, **kit.HOST_NAMED) # off: it goes through view = lease.poll() assert view is not None - lease.mark_arrived(view.row_masks()) - lease.release() + kit.return_write(lender, lease, view.row_masks()) def test_a_fetch_into_a_request_returning_context_outputs_fails_at_the_call( @@ -3528,19 +3726,19 @@ def test_a_fetch_into_a_request_returning_context_outputs_fails_at_the_call( for request in (held, fresh, outputs, plain): assert mgr.prepare_context(request) held.py_result.append_context_logits(torch.zeros(TPB, 1, 8, device="cuda")) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) for request, returns in ((held, True), (fresh, True), (outputs, True), (plain, False)): before = state_of(kit, mgr, lender, request) - lease = lender.lend_write(request, 0, END) + lease = lender.lend_write(request, 0, END, **kit.HOST_NAMED) if returns: assert lease.failure is not None, "not failed at the call" assert "context" in lease.failure, lease.failure assert lease.poll() is None assert state_of(kit, mgr, lender, request) == before, "the lease changed the cache" + lease.release() else: assert lease.failure is None, "refused a request returning no context outputs" - lease.mark_arrived(ready_view(lease, mgr).row_masks()) - lease.release() + kit.return_write(lender, lease, kit.ready_view(lease, mgr).row_masks()) @pytest.mark.parametrize("policy", ["per_request", "per_conversation"]) @@ -3554,9 +3752,9 @@ def test_a_fetch_ending_below_a_history_the_update_moved_fails_at_the_call( conversation = "history-past" if policy == "per_conversation" else None with real_manager(block_reuse_config={"policy": policy}) as mgr: target = computed(kit, mgr, TARGET, SPLIT_PROMPT, upto=4 * TPB, conversation=conversation) - lender = attach(mgr, tokens_per_fetch=TPB) + lender = attach(mgr, guaranteed_fetch_tokens=TPB) before = state_of(kit, mgr, lender, target) - lease = lender.lend_write(target, 2 * TPB, 3 * TPB) + lease = lender.lend_write(target, 2 * TPB, 3 * TPB, **kit.HOST_NAMED) try: assert lease.failure is not None, f"not failed at the call, from {before}" assert "stands past" in lease.failure, lease.failure @@ -3564,7 +3762,7 @@ def test_a_fetch_ending_below_a_history_the_update_moved_fails_at_the_call( assert state_of(kit, mgr, lender, target) == before, "the fetch changed the cache" finally: lease.release() - assert lender.resume_bounds(target) == (4 * TPB, 4 * TPB) + assert lender.poll_resume_bounds(target) == (4 * TPB, 4 * TPB) def test_a_draft_pool_fetch_ending_below_its_history_fails_at_the_call( @@ -3591,9 +3789,9 @@ def test_a_draft_pool_fetch_ending_below_its_history_fails_at_the_call( kv = kit.kv(mgr, request) kv.resize(None, at) assert (kv.history_length, kv.num_committed_tokens) == (at, 0) - lender = attach(mgr, tokens_per_fetch=TPB, scope=scope) + lender = attach(mgr, guaranteed_fetch_tokens=TPB, scope=scope) before = state_of(kit, mgr, lender, request) - lease = lender.lend_write(request, 2 * TPB, 3 * TPB) + lease = lender.lend_write(request, 2 * TPB, 3 * TPB, **kit.HOST_NAMED) try: failures.append(lease.failure) if lease.failure is not None: @@ -3606,25 +3804,29 @@ def test_a_draft_pool_fetch_ending_below_its_history_fails_at_the_call( ) -def test_a_range_longer_than_staging_holds_is_refused(kit, real_manager, attach=attach): +def test_a_range_longer_than_staging_holds_fails_at_the_call(kit, real_manager, attach=attach): + """A range longer than this rank's staging holds fails the lease at the call.""" with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) target = kit.admitted(mgr, TARGET, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=2 * TPB) # two slots + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) # two slots before = state_of(kit, mgr, lender, target) - assert "more staging slots" in refusal(lambda: lender.lend_read(source, 0, END)) - assert lender._open_lease_count() == 0, "a refused publish left a lease open" - assert "more staging slots" in refusal(lambda: lender.lend_write(target, 0, END)) + for lend, request in ((lender.lend_read, source), (lender.lend_write, target)): + lease = lend(request, 0, END, **kit.HOST_NAMED) + assert "more staging slots" in str(lease.failure), lease.failure + assert lease.poll() is None + lease.release() + assert lender._open_lease_count() == 0, "a refused lease stayed open after its release" assert state_of(kit, mgr, lender, target) == before - # A range of tokens_per_fetch tokens fits exactly. - read = lender.lend_read(source, 0, 2 * TPB) - ready_view(read, mgr) + # A range of guaranteed_fetch_tokens tokens fits exactly. + read = lender.lend_read(source, 0, 2 * TPB, **kit.HOST_NAMED) + kit.ready_view(read, mgr) read.release() - write = lender.lend_write(target, 0, 2 * TPB) + lender.progress() # the read's slots come back + write = lender.lend_write(target, 0, 2 * TPB, **kit.HOST_NAMED) view = write.poll() assert view is not None - write.mark_arrived(view.row_masks()) - write.release() + kit.return_write(lender, write, view.row_masks()) # What this rank holds for a request that fails a call whose range is right: no KV cache, a reset @@ -3634,12 +3836,9 @@ def test_a_range_longer_than_staging_holds_is_refused(kit, real_manager, attach= @pytest.mark.parametrize("state", REQUEST_STATES) def test_a_wrong_call_raises_whatever_the_request_s_state(kit, real_manager, state, attach=attach): - """A range that is not whole blocks or, for a fetch, ends past the whole blocks before the - request's last prompt token or needs more slots than a region has is the caller's error, the same - on every rank: on a manager not shut down it raises ``ValueError`` whatever this rank holds for - the request, also where a right range fails at the call, and changes nothing.""" + """A wrong range raises whatever the request's state, which fails a right or too long range.""" with real_manager() as mgr: - lender = attach(mgr, tokens_per_fetch=2 * TPB) # two slots + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) # two slots if state == "no_cache": target = kit.make_request(TARGET, OTHER_PROMPT) elif state == "reuse_reset": @@ -3653,17 +3852,18 @@ def test_a_wrong_call_raises_whatever_the_request_s_state(kit, real_manager, sta (lender.lend_read, 50, "whole blocks"), (lender.lend_write, 50, "whole blocks"), (lender.lend_write, 4 * TPB, "past the whole blocks before the request's last prompt"), - (lender.lend_write, END, "more staging slots"), ): - assert why in refusal(lambda: lend(target, 0, end)) + assert why in refusal(lambda: lend(target, 0, end, **kit.HOST_NAMED)) assert state_of(kit, mgr, lender, target) == before, "a wrong call changed the cache" - lease = lender.lend_write(target, 0, 2 * TPB) # a right range: the request's state decides - try: - assert lease.failure is not None and lease.poll() is None, ( - f"a right range went through for a request in state {state}: this proves nothing" - ) - finally: - lease.release() + # A right range, and one past the two slots of staging: the request's state decides. + why = {"no_cache": "no KV cache", "reuse_reset": RESET, "unnamed": "without digests"}[state] + for end in (2 * TPB, END): + lease = lender.lend_write(target, 0, end, **kit.HOST_NAMED) + try: + assert why in str(lease.failure) and lease.poll() is None, (end, lease.failure) + finally: + lease.release() + assert state_of(kit, mgr, lender, target) == before, "a failed lease changed the cache" # -- resume_bounds -------------------------------------------------------------------------------- @@ -3672,28 +3872,30 @@ def test_a_wrong_call_raises_whatever_the_request_s_state(kit, real_manager, sta def test_resume_bounds_without_a_fetch_is_the_cache_s_own_interval(kit, real_manager): with real_manager(windows=[WINDOW, 256]) as mgr: source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) - lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + lender = attach(mgr, guaranteed_fetch_tokens=WINDOWED_END) kv = kit.kv(mgr, source) - resume_bounds = lender.resume_bounds(source) + resume_bounds = lender.poll_resume_bounds(source) assert resume_bounds == (max(kv.num_committed_tokens, kv.history_length), kv.history_length) assert all(type(v) is int for v in resume_bounds) with pytest.raises(ValueError): - lender.resume_bounds(kit.make_request(9, PROMPT)) + lender.poll_resume_bounds(kit.make_request(9, PROMPT)) -def test_resume_bounds_waits_on_nothing_and_settles_by_being_asked(kit, real_manager): +def test_resume_bounds_waits_on_nothing_and_settles_at_a_round_of_progress(kit, real_manager): + """A returned write's fetch settles at the first round of progress after its copy ran.""" with real_manager() as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - lease = lender.lend_write(target, 0, END) + lender = attach(mgr, guaranteed_fetch_tokens=END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() with kit.held_stream(mgr._stream) as gate: - lease.mark_arrived(view.row_masks(True)) - assert lender.resume_bounds(target) is None - assert lender.resume_bounds(target) is None + assert kit.return_write(lender, lease, view.row_masks(True)) is None + for _ in range(2): + lender.progress() + assert lender.poll_resume_bounds(target) is None gate.open() - assert lender.resume_bounds(target) == (END, 0) - lease.release() + lender.progress() + assert lender.poll_resume_bounds(target) == (END, 0) # -- a reset of the manager's reuse state ---------------------------------------------------- @@ -3708,23 +3910,27 @@ def test_a_reuse_reset_stops_lending_by_name(kit, real_manager, attach=attach): source = kit.published(mgr, SOURCE, PROMPT) original = [kit.page(mgr, source, 0, o) for o in range(BLOCKS)] target = kit.admitted(mgr, TARGET, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) - publish = lender.lend_read(source, 0, END) # granted: its copy is queued - fetch = lender.lend_write(target, 0, END) - fetch_view = fetch.poll() # seen ready, not marked + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) + publish = lender.lend_read(source, 0, END, **kit.HOST_NAMED) # granted: its copy is queued + fetch = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + fetch_view = fetch.poll() # seen ready, not returned for request in (source, target): # the manager resets only with every cache closed mgr.free_resources(request) mgr.reset_reuse_state() # what reset_prefix_cache does when a weight update ends - (run,) = ready_view(publish, mgr).groups # granted before the reset: it finishes - length = lender.regions[run.region_index].slot_bytes - staged = [kit.host_bytes(a, length) for a in run.addresses.tolist()] + (run,) = kit.ready_view(publish, mgr).groups # granted before the reset: it finishes + region, addresses = kit.staged_slots(lender, run) + staged = [kit.host_bytes(a, region.slot_bytes) for a in addresses] assert kit.digest(staged) == kit.digest(original), "the granted publish lost its bytes" - fetch.mark_arrived(fetch_view.row_masks(True)) # its request is gone: nothing is copied + # Its request is gone: nothing is copied, and the request has no bounds. + assert kit.return_write(lender, fetch, fetch_view.row_masks(True)) is None # The same tokens, computed after the reset, would go out under the names of the old bytes. source = kit.published(mgr, SOURCE, PROMPT) target = kit.admitted(mgr, TARGET, OTHER_PROMPT) before = [state_of(kit, mgr, lender, r) for r in (source, target)] - later = [lender.lend_read(source, 0, END), lender.lend_write(target, 0, END)] + later = [ + lender.lend_read(source, 0, END, **kit.HOST_NAMED), + lender.lend_write(target, 0, END, **kit.HOST_NAMED), + ] for lease in later: assert lease.poll() is None and RESET in str(lease.failure), ( "lent by name after the reset" @@ -3754,11 +3960,11 @@ def recording_free(memory): with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=END) + regions = kit.staging_regions(lender) monkeypatch.setattr(_lender._HostMemory, "free", recording_free) with kit.held_stream(mgr._stream, strict=False) as gate: - lease = lender.lend_read(source, 0, END) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) event = lease._copy._event assert event is not None synchronize = event.synchronize @@ -3778,18 +3984,18 @@ def test_shutdown_fails_waiting_leases_and_later_calls_only_end_records(kit, rea with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) target = kit.admitted(mgr, TARGET, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - held = lender.lend_read(source, 0, END) - ready_view(held, mgr) - waiting = lender.lend_read(source, 0, END) + lender = attach(mgr, guaranteed_fetch_tokens=END) + held = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + kit.ready_view(held, mgr) + waiting = lender.lend_read(source, 0, END, **kit.HOST_NAMED) assert waiting.poll() is None mgr.shutdown() assert waiting.failure is not None and waiting.poll() is None - for lease in (lender.lend_read(source, 0, END), lender.lend_write(target, 0, END)): - assert lease.failure is not None and lease.poll() is None - lease.release() + for lend, request in ((lender.lend_read, source), (lender.lend_write, target)): + with pytest.raises(RuntimeError, match="lends no more"): + lend(request, 0, END, **kit.HOST_NAMED) with pytest.raises(ValueError): - lender.resume_bounds(source) + lender.poll_resume_bounds(source) for lease in (held, held, waiting): lease.release() mgr.shutdown() # a second shutdown frees nothing and returns @@ -3801,16 +4007,16 @@ def test_an_unreleased_lease_keeps_the_staging_memory_until_exit( ): with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=END) + regions = kit.staging_regions(lender) assert kit.staging_memory(regions) if failed: - lease = lender.lend_read(kit.make_request(9, PROMPT), 0, END) + lease = lender.lend_read(kit.make_request(9, PROMPT), 0, END, **kit.HOST_NAMED) assert lease.failure is not None address, content = None, None else: - lease = lender.lend_read(source, 0, END) - address = int(ready_view(lease, mgr).groups[0].addresses[0]) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + address = kit.staged_slots(lender, kit.ready_view(lease, mgr).groups[0])[1][0] content = kit.host_bytes(address, regions[0].slot_bytes) warnings = kit.lender_warnings(monkeypatch) mgr.shutdown() @@ -3832,8 +4038,8 @@ def planted(self): raise RuntimeError("planted") with real_manager() as mgr: - regions = attach(mgr, tokens_per_fetch=END).regions - monkeypatch.setattr(_lender.Staging, "_must_retain_memory", planted) + regions = kit.staging_regions(attach(mgr, guaranteed_fetch_tokens=END)) + monkeypatch.setattr(_lender._ManagerLender, "_must_retain_memory", planted) mgr.shutdown() assert kit.staging_kept(regions) @@ -3851,8 +4057,8 @@ def released(hold) -> None: def test_a_hold_dropped_unreleased_still_keeps_the_staging_memory(kit, real_manager, attach=attach): with real_manager() as mgr: - lender = attach(mgr, tokens_per_fetch=END) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=END) + regions = kit.staging_regions(lender) assert kit.staging_memory(regions) lender.hold_regions() # dropped unreleased gc.collect() @@ -3862,8 +4068,8 @@ def test_a_hold_dropped_unreleased_still_keeps_the_staging_memory(kit, real_mana def test_a_hold_releases_once_in_every_state(kit, real_manager, attach=attach): with real_manager() as mgr: - lender = attach(mgr, tokens_per_fetch=END) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=END) + regions = kit.staging_regions(lender) assert kit.staging_memory(regions) first, second = lender.hold_regions(), lender.hold_regions() released(first) @@ -3900,15 +4106,15 @@ def test_a_lease_and_a_hold_keep_the_memory_until_both_go( and a failed lease released at once: only with both gone does the shutdown free the memory.""" with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=END) + regions = kit.staging_regions(lender) refs = [weakref.ref(m) for m in kit.staging_memory(regions)] assert refs - lender.lend_read(kit.make_request(9, PROMPT), 0, END).release() # failed, released + lender.lend_read(kit.make_request(9, PROMPT), 0, END, **kit.HOST_NAMED).release() # failed hold = lender.hold_regions() assert isinstance(hold, RegionHold) - lease = lender.lend_read(source, 0, END) - ready_view(lease, mgr) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + kit.ready_view(lease, mgr) ends = {"lease": lease.release, "hold": lambda: released(hold)} for what in order: ends[what]() @@ -3939,20 +4145,20 @@ def no_event(self, stream=None): with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=END) + regions = kit.staging_regions(lender) (group,) = kit.pool_group_ids(mgr) assert kit.staging_memory(regions) hold = lender.hold_regions() with monkeypatch.context() as patched: patched.setattr(torch.cuda.Event, "record", no_event) - lease = lender.lend_read(source, 0, END) # its copy has no event: slots lost + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) # no event: slots lost assert lease.failure is not None mgr._stream.synchronize() ends = {"lease": lease.release, "hold": lambda: released(hold)} for what in order: ends[what]() - lender.resume_bounds(source) # any call makes progress + lender.progress() assert lender._free_slots(group) == 0, "slots a copy may still touch were reused" warnings = kit.lender_warnings(monkeypatch) mgr.shutdown() @@ -3963,85 +4169,6 @@ def no_event(self, stream=None): assert kit.staging_kept(regions), "a release after the shutdown freed it" -# -- progress --------------------------------------------------------------------------------- - - -class EventQueries: - """Counts ``torch.cuda.Event.query`` per event within each lender call made through ``call``.""" - - def __init__(self, monkeypatch): - self.asked: List[int] = [] - self.last: Set[int] = set() # the events the latest call asked - query = torch.cuda.Event.query - - def counted(event): - self.asked.append(id(event)) - return query(event) - - monkeypatch.setattr(torch.cuda.Event, "query", counted) - - def call(self, method, *args): - first = len(self.asked) - result = method(*args) - asked = self.asked[first:] - assert len(asked) == len(set(asked)), "a copy's event was asked twice in one lender call" - self.last = set(asked) - return result - - -HELD_READS = 4 - - -def test_each_call_asks_each_copy_s_event_at_most_once( - kit, real_manager, monkeypatch, attach=attach -): - """Each lender call asks each copy's event at most once, a lease still lent costs no query - beyond its own poll's, and a copy reported complete is never asked again: a call asks only its - own copy and the pending copies of released leases, not one per held lease.""" - with real_manager() as mgr: - sources = [ - kit.published(mgr, 10 + i, [10000 * (i + 1) + t for t in range(END + 1)]) - for i in range(HELD_READS) - ] - target = kit.admitted(mgr, TARGET, PROMPT) - assert mgr._resize_for_served_prefix(target, kit.kv(mgr, target), 0, END) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=HELD_READS + 1) - events = EventQueries(monkeypatch) - with kit.held_stream(mgr._stream) as gate: - reads = [events.call(lender.lend_read, source, 0, END) for source in sources] - write = events.call(lender.lend_write, target, 0, END) - view = events.call(write.poll) - events.call(write.mark_arrived, view.row_masks(True)) - events.call(write.release) - released = id(write._copy._event) - for _ in range(2): - for lease in reads: - assert events.call(lease.poll) is None - own = id(lease._copy._event) - assert events.last == {own, released}, ( - f"a read's poll asked {len(events.last)} copies, not its own and the " - "released one" - ) - assert events.call(lender.resume_bounds, target) is None - assert events.last == {released}, ( - f"resume_bounds asked {len(events.last)} copies, not the released one" - ) - gate.open() - mgr._stream.synchronize() - assert events.call(lender.resume_bounds, target) is not None - for lease in reads: - assert events.call(lease.poll) is not None - first = len(events.asked) - for _ in range(3): - for lease in reads: - assert events.call(lease.poll) is not None - assert events.call(lender.resume_bounds, target) is not None - assert len(events.asked) == first, "a copy reported complete was asked again" - for lease in reads: - events.call(lease.release) - assert len(events.asked) == first, "a copy reported complete was asked again" - - # -- a copy that fails partway ---------------------------------------------------------------- @@ -4079,48 +4206,56 @@ def test_a_publish_whose_copy_fails_partway_fails_and_frees_its_slots_after_the_ # Three layers: two pool groups, so the copy is several memcpys. with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: source = kit.published(mgr, SOURCE, WINDOWED_PROMPT) - lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + lender = attach(mgr, guaranteed_fetch_tokens=WINDOWED_END) groups = kit.pool_group_ids(mgr) driver = FailingDriver(_lender.drv, fail_on=2) monkeypatch.setattr(_lender, "drv", driver) with kit.held_stream(mgr._stream) as gate: - lease = lender.lend_read(source, 0, WINDOWED_END) # every slot + lease = lender.lend_read(source, 0, WINDOWED_END, **kit.HOST_NAMED) # every slot assert driver.calls >= 2, "a single memcpy: nothing failed partway" assert lease.failure is not None and lease.poll() is None lease.release() + lender.progress() assert [lender._free_slots(g) for g in groups] == [0] * len(groups), ( "slots came back while a copy into them was queued" ) gate.open() - lender.resume_bounds(source) # any call makes progress - assert [lender._free_slots(g) for g in groups] == [p.slot_count for p in lender.regions] + lender.progress() + assert [lender._free_slots(g) for g in groups] == [ + p.slot_count for p in kit.staging_regions(lender) + ] -def test_marks_whose_copy_fails_partway_abandon_the_fetch_and_keep_the_lease_ready( +def test_arrivals_whose_copy_fails_partway_abandon_the_fetch_and_return_the_lease( kit, real_manager, monkeypatch ): + """A return whose copy fails partway abandons the fetch and still returns the lease.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: target = kit.admitted(mgr, TARGET, WINDOWED_PROMPT) - lender = attach(mgr, tokens_per_fetch=WINDOWED_END) + lender = attach(mgr, guaranteed_fetch_tokens=WINDOWED_END) groups = kit.pool_group_ids(mgr) - lease = lender.lend_write(target, 0, WINDOWED_END) + lease = lender.lend_write(target, 0, WINDOWED_END, **kit.HOST_NAMED) view = lease.poll() kit.stage(lender, view, 0x5A) driver = FailingDriver(_lender.drv, fail_on=2) monkeypatch.setattr(_lender, "drv", driver) kv = kit.kv(mgr, target) with kit.held_stream(mgr._stream) as gate: - lease.mark_arrived(view.row_masks(True)) # raises nothing + bounds = kit.return_write(lender, lease, view.row_masks(True)) # raises nothing assert driver.calls >= 2, "a single memcpy: nothing failed partway" - assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) - assert lease.poll() is view, "the lease stays ready" - lease.release() + assert bounds == (kv.num_committed_tokens, kv.history_length), bounds + assert lender.poll_resume_bounds(target) == bounds + with pytest.raises(RuntimeError): + lease.poll() # returned + lender.progress() assert [lender._free_slots(g) for g in groups] == [0] * len(groups) gate.open() - lender.resume_bounds(target) - assert [lender._free_slots(g) for g in groups] == [p.slot_count for p in lender.regions] + lender.progress() + assert [lender._free_slots(g) for g in groups] == [ + p.slot_count for p in kit.staging_regions(lender) + ] # -- errors inside the lender ----------------------------------------------------------------- @@ -4130,18 +4265,19 @@ def test_an_event_query_that_raises_returns_no_slot_twice( kit, real_manager, monkeypatch, attach=attach ): """Two released publishes' copies have run, and the slot return's query of the later one's - event raises: that call raises having returned no slot, and the next one returns every slot - once.""" + event raises: that call returns no slot, and shutdown returns every slot once.""" with real_manager() as mgr: first = kit.published(mgr, SOURCE, PROMPT) second = kit.published(mgr, OTHER, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) (group,) = kit.pool_group_ids(mgr) with kit.held_stream(mgr._stream): - earlier = lender.lend_read(first, 0, END) # granted and released first: asked first - later = lender.lend_read(second, 0, END) + # Granted and released first: asked first. + earlier = lender.lend_read(first, 0, END, **kit.HOST_NAMED) + later = lender.lend_read(second, 0, END, **kit.HOST_NAMED) earlier.release() later.release() + lender.progress() assert lender._free_slots(group) == 0 query = torch.cuda.Event.query planted = later._copy._event @@ -4154,20 +4290,16 @@ def raises_once(event): monkeypatch.setattr(torch.cuda.Event, "query", raises_once) with pytest.raises(RuntimeError, match="planted"): - lender.resume_bounds(first) # any call makes progress - try: - lender.resume_bounds(first) - except ValueError as error: - raise AssertionError(f"a slot came back twice: {error}") from None + lender.progress() + assert lender._free_slots(group) == 0 + mgr.shutdown() assert lender._free_slots(group) == lender.regions[0].slot_count - again = lender.lend_read(first, 0, END) - ready_view(again, mgr) - again.release() -def test_marks_that_raise_leave_the_fetch_abandoned_and_the_slots_returnable( +def test_arrivals_that_raise_leave_the_fetch_abandoned_and_the_slots_returnable( kit, real_manager, monkeypatch ): + """Arrivals whose count raises abandon the fetch; a release that follows counts.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender def planted(*args): @@ -4175,77 +4307,75 @@ def planted(*args): with real_manager() as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - other = kit.admitted(mgr, OTHER, OTHER_PROMPT) - lender = attach(mgr, tokens_per_fetch=END) + lender = attach(mgr, guaranteed_fetch_tokens=END) (group,) = kit.pool_group_ids(mgr) - lease = lender.lend_write(target, 0, END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() - kv = kit.kv(mgr, target) with monkeypatch.context() as patched: - patched.setattr(_lender.Staging, "_locked_block_masks", planted) + patched.setattr(_lender._ManagerLender, "_locked_block_masks", planted) with pytest.raises(RuntimeError, match="planted"): - lease.mark_arrived(view.row_masks(True)) - assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + kit.return_write(lender, lease, view.row_masks(True)) lease.release() - waiting = lender.lend_write(other, 0, END) - waiting_view = waiting.poll() - assert waiting_view is not None, "the lease's slots never came back" + assert lender._open_lease_count() == 0 + mgr.shutdown() + assert lender._free_slots(group) == lender.regions[0].slot_count + + +def test_a_release_returns_its_lease_without_running_progress(kit, real_manager, monkeypatch): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender + + def planted(*args): + raise RuntimeError("planted") + + with real_manager() as mgr: + target = kit.admitted(mgr, TARGET, PROMPT) + lender = attach(mgr, guaranteed_fetch_tokens=END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + assert lease.poll() is not None opened = lender._open_lease_count() with monkeypatch.context() as patched: - patched.setattr(_lender.Staging, "_progress", planted) + patched.setattr(_lender._ManagerLender, "_progress", planted) + lease.release() + assert lender._open_lease_count() == opened - 1 with pytest.raises(RuntimeError, match="planted"): - waiting.release() # counted before the progress that raises - assert lender._open_lease_count() == opened - 1 - waiting.mark_arrived(waiting_view.row_masks()) - lender.resume_bounds(other) - assert lender._free_slots(group) == lender.regions[0].slot_count + lender.progress() + mgr.shutdown() -def test_marks_whose_copy_raises_once_queued_lose_the_slots_it_reads( - kit, real_manager, monkeypatch, attach=attach, planted=MemoryError +@pytest.mark.parametrize("interrupted", [False, True], ids=["memory_error", "interrupted"]) +def test_arrivals_whose_copy_raises_once_queued_lose_the_slots_it_reads( + kit, real_manager, monkeypatch, interrupted, attach=attach ): - """The marked rows' copy is queued, then creating its event raises: the mark raises, the fetch - stays abandoned, and the slots the copy reads are lost, not returned at the release.""" + """A return whose queued copy's event raises, or is interrupted, loses the slots it reads.""" + planted = Interrupted if interrupted else MemoryError def no_event(*args, **kwargs): raise planted("planted") with real_manager() as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) + lender = attach(mgr, guaranteed_fetch_tokens=END) (group,) = kit.pool_group_ids(mgr) - lease = lender.lend_write(target, 0, END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() - kv = kit.kv(mgr, target) with kit.held_stream(mgr._stream): with monkeypatch.context() as patched: patched.setattr(torch.cuda, "Event", no_event) with pytest.raises(planted, match="planted"): - lease.mark_arrived(view.row_masks(True)) # the copy is queued, then raises - assert lender.resume_bounds(target) == (kv.num_committed_tokens, kv.history_length) + # The copy is queued, then raises. + kit.return_write(lender, lease, view.row_masks(True)) lease.release() assert lender._free_slots(group) == 0, ( "slots came back while a copy out of them was queued" ) - lender.resume_bounds(target) # any call makes progress assert lender._free_slots(group) == 0, "slots a copy may still read were reused" + mgr.shutdown() class Interrupted(BaseException): """An interruption that is no ``Exception``, as ``KeyboardInterrupt`` is.""" -def test_marks_whose_copy_is_interrupted_once_queued_lose_the_slots_it_reads( - kit, real_manager, monkeypatch, attach=attach -): - """``test_marks_whose_copy_raises_once_queued_lose_the_slots_it_reads`` with an interruption - that is no ``Exception`` in place of its ``MemoryError``: the slots a queued copy reads are lost - whatever the mark raises.""" - test_marks_whose_copy_raises_once_queued_lose_the_slots_it_reads( - kit, real_manager, monkeypatch, attach, planted=Interrupted - ) - - SLOTS_BACK = "slots came back while a copy on them was queued" @@ -4256,75 +4386,67 @@ def interrupted_event(*args, **kwargs): def test_an_interrupt_in_a_read_s_grant_keeps_its_slots_out_of_use( kit, real_manager, monkeypatch, attach=attach ): - """A read waits for slots another read holds, whose release grants it: the copy into them is - queued, then creating its event is interrupted. The release raises; after the read's release the - slots stay out of use, no grant hands them out, and the shutdown keeps the staging memory.""" + """An interrupt as a waiting read's grant queues its copy keeps the slots out of use.""" with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=END) + regions = kit.staging_regions(lender) (group,) = kit.pool_group_ids(mgr) - head = lender.lend_read(source, 0, END) - ready_view(head, mgr) - waiting = lender.lend_read(source, 0, END) + head = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + kit.ready_view(head, mgr) + waiting = lender.lend_read(source, 0, END, **kit.HOST_NAMED) assert waiting.poll() is None, "the read was granted at once: the check proves nothing" + head.release() with kit.held_stream(mgr._stream): with monkeypatch.context() as patched: patched.setattr(torch.cuda, "Event", interrupted_event) with pytest.raises(Interrupted, match="planted"): - head.release() # grants the waiting read: its copy is queued, then interrupted + lender.progress() # grants the waiting read; its copy's event is interrupted runs = waiting._slot_runs waiting.release() assert lender._free_slots(group) == 0, SLOTS_BACK - lender.resume_bounds(source) # any call makes progress assert any(r is runs for r in lender._quarantined), SLOTS_BACK - later = lender.lend_read(source, 0, END) - assert later.poll() is None and later.failure is None, "a grant handed the slots out" - later.release() warnings = kit.lender_warnings(monkeypatch) mgr.shutdown() assert kit.staging_kept(regions), "staging freed at shutdown with slots out of use" assert warnings, "keeping the memory until exit is logged" -def interrupting_as_it_takes_a_copy(lease): - """Turn ``lease`` into one whose first setting of ``_copy`` to a copy raises ``Interrupted``.""" +def interrupting_as_it_takes_a_copy(monkeypatch, lease): + """Make the first setting of ``lease``'s ``_copy`` to a copy raise ``Interrupted``.""" + cls = type(lease) + assign = cls.__setattr__ + interrupted = [] - class Interrupting(type(lease)): - def __setattr__(self, name, value): - if name == "_copy" and value is not None and not self.__dict__.get("_interrupted"): - self.__dict__["_interrupted"] = True - raise Interrupted("planted") - super().__setattr__(name, value) + def interrupting(self, name, value): + if self is lease and name == "_copy" and value is not None and not interrupted: + interrupted.append(True) + raise Interrupted("planted") + assign(self, name, value) - lease.__class__ = Interrupting + monkeypatch.setattr(cls, "__setattr__", interrupting) def test_an_interrupt_before_a_write_takes_its_copy_keeps_the_slots_out_of_use( kit, real_manager, monkeypatch, attach=attach ): - """A write's marked rows are copied into its pages, and the lease taking that copy is - interrupted: the mark raises, the slots the copy reads stay out of use after the release, no - grant hands them out, and the shutdown keeps the staging memory.""" + """An interrupt as a returned write takes its copy keeps the slots out of use.""" with real_manager() as mgr: target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=END) + regions = kit.staging_regions(lender) (group,) = kit.pool_group_ids(mgr) - lease = lender.lend_write(target, 0, END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) view = lease.poll() - interrupting_as_it_takes_a_copy(lease) + interrupting_as_it_takes_a_copy(monkeypatch, lease) with kit.held_stream(mgr._stream): with pytest.raises(Interrupted, match="planted"): - lease.mark_arrived(view.row_masks(True)) # the copy is queued, then interrupted + # The copy is queued, then interrupted. + kit.return_write(lender, lease, view.row_masks(True)) runs = lease._slot_runs lease.release() assert lender._free_slots(group) == 0, SLOTS_BACK - lender.resume_bounds(target) # any call makes progress assert any(r is runs for r in lender._quarantined), SLOTS_BACK - later = lender.lend_write(target, 0, END) - assert later.poll() is None and later.failure is None, "a grant handed the slots out" - later.release() warnings = kit.lender_warnings(monkeypatch) mgr.shutdown() assert kit.staging_kept(regions), "staging freed at shutdown with slots out of use" @@ -4334,24 +4456,21 @@ def test_an_interrupt_before_a_write_takes_its_copy_keeps_the_slots_out_of_use( def test_a_grant_that_raises_fails_only_its_own_lease( kit, real_manager, monkeypatch, attach=attach ): - """Two reads wait in line for the slots a third holds, and the first one's grant raises when - the third's release frees them: that read fails, the healthy one behind it is granted in the - same release and stages its source's bytes, resume_bounds answers and a later read goes through.""" + """A grant that raises fails its own lease alone; the read behind it is granted that round.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager with real_manager() as mgr: holder = kit.published(mgr, SOURCE, PROMPT) broken = kit.published(mgr, OTHER, OTHER_PROMPT) fetched = kit.admitted(mgr, TARGET, list(range(8000, 8097))) - lender = attach(mgr, tokens_per_fetch=END) - write = lender.lend_write(fetched, 0, END) - write.mark_arrived(write.poll().row_masks()) - write.release() - head = lender.lend_read(holder, 0, END) - ready_view(head, mgr) - waiting = lender.lend_read(broken, 0, 2 * TPB) + lender = attach(mgr, guaranteed_fetch_tokens=END) + write = lender.lend_write(fetched, 0, END, **kit.HOST_NAMED) + kit.return_write(lender, write, write.poll().row_masks()) + head = lender.lend_read(holder, 0, END, **kit.HOST_NAMED) + kit.ready_view(head, mgr) + waiting = lender.lend_read(broken, 0, 2 * TPB, **kit.HOST_NAMED) assert waiting.poll() is None - behind = lender.lend_read(holder, 0, TPB) # healthy, in line behind the broken one + behind = lender.lend_read(holder, 0, TPB, **kit.HOST_NAMED) # healthy, in line behind it assert behind.poll() is None broken_kv = kit.kv(mgr, broken) @@ -4374,33 +4493,26 @@ def call(kv_cache, *rest): patched.setattr(_manager, "kv_of", kv_of) for name in ("cache_state", "pages", "locked_pages"): patched.setattr(_manager, name, raising_for_broken(getattr(_manager, name))) - head.release() # grants both waiting reads, and the first one's grant raises - # The lender's own record, read before any other call: a poll would grant it itself. + head.release() + lender.progress() # grants both waiting reads, and the first one's grant raises granted = behind._view is not None assert waiting.poll() is None and waiting.failure is not None assert granted and behind.failure is None, ( - f"the read behind the broken one was not granted in that release: {behind.failure}" + f"the read behind the broken one was not granted in that round: {behind.failure}" ) - assert isinstance(lender.resume_bounds(fetched), ResumeBounds) - view = ready_view(behind, mgr) - assert 0 < view.num_block_entries == check_staged(kit, mgr, lender, holder, view) + assert isinstance(lender.poll_resume_bounds(fetched), ResumeBounds) + view = kit.ready_view(behind, mgr) + assert 0 < view.num_rows == check_staged(kit, mgr, lender, holder, view) behind.release() waiting.release() - again = lender.lend_read(holder, 0, END) - ready_view(again, mgr) + again = lender.lend_read(holder, 0, END, **kit.HOST_NAMED) + kit.ready_view(again, mgr) again.release() -@pytest.mark.parametrize("fault", ["none", "take"], ids=["control", "take_raises"]) +@pytest.mark.parametrize("fault", ["none", "cancel"], ids=["control", "cancel_raises"]) def test_an_error_in_the_free_hook_stops_the_lender(kit, real_manager, monkeypatch, fault): - """Two pool groups. A released publish's copy into its slots is pending while a fetch waits for - them; another request's free through the manager then returns the publish's slots and grants - them to the fetch, and the take of the fetch's second pool group raises a planted MemoryError - once its first pool group changed. The free returns, and the lender has stopped: the waiting - fetch and a lease no poll returned the view of fail, a view a poll returned stays, lends fail at - the call after their range checks, a mark copies nothing, resume_bounds is an empty interval, no - slot returns, and the shutdown keeps the staging memory with nothing open. Without the fault - the free grants the waiting fetch and a mark copies its rows.""" + """An error the free hook meets stops the lender for good; without it the next round grants.""" from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender with real_manager(windows=[WINDOW, 256], num_layers=3) as mgr: @@ -4409,94 +4521,100 @@ def test_an_error_in_the_free_hook_stops_the_lender(kit, real_manager, monkeypat seen_target = kit.admitted(mgr, 4, PROMPT) granted_target = kit.admitted(mgr, 5, list(range(8000, 8097))) other = kit.admitted(mgr, OTHER, list(range(9000, 9097))) - lender = attach(mgr, tokens_per_fetch=WINDOWED_END, planned_fetches=3) - regions = lender.regions + lender = attach(mgr, guaranteed_fetch_tokens=WINDOWED_END, capacity_in_fetches=3) + regions = kit.staging_regions(lender) groups = kit.pool_group_ids(mgr) - assert len(regions) == len(groups) == 2, "a grant changes one pool group, then the other" - seen = lender.lend_write(seen_target, 0, 2 * TPB) + seen = lender.lend_write(seen_target, 0, 2 * TPB, **kit.HOST_NAMED) seen_view = seen.poll() - granted = lender.lend_write(granted_target, 0, TPB) # granted, its view never polled + granted = lender.lend_write(granted_target, 0, TPB, **kit.HOST_NAMED) # never polled assert seen_view is not None and granted.failure is None kit.fill_sentinel(mgr, seen_target) kit.stage(lender, seen_view, 0x5A) # Grown ahead, so the held section only lends. kv = kit.kv(mgr, target) assert mgr._resize_for_served_prefix(target, kv, WINDOWED_END, WINDOWED_END) + assert mgr._resize_for_served_prefix(other, kit.kv(mgr, other), TPB, TPB) with kit.held_stream(mgr._stream) as gate: - publish = lender.lend_read(source, 0, WINDOWED_END) # its copy held - fetch = lender.lend_write(target, 0, WINDOWED_END) + publish = lender.lend_read(source, 0, WINDOWED_END, **kit.HOST_NAMED) # its copy held + fetch = lender.lend_write(target, 0, WINDOWED_END, **kit.HOST_NAMED) assert fetch.poll() is None and fetch.failure is None, "the fetch did not wait" + behind = lender.lend_write(other, 0, TPB, **kit.HOST_NAMED) + assert behind.poll() is None and behind.failure is None, "the write did not wait" publish.release() # its copy pends, so its slots stay out of use gate.open() - taken = [] + cancelled = [] - def planted(take): - def call(start, count): - taken.append(count) - if len(taken) == 2: # the fetch's second pool group, its first one changed + def planted(cancel): + def call(ticket): + cancelled.append(ticket) + if len(cancelled) == 1: # the freed request's write, failed before its cancel raise MemoryError("planted") - return take(start, count) + return cancel(ticket) return call - if fault == "take": - for alloc in lender._slot_queue._allocators.values(): - monkeypatch.setattr(alloc, "take", planted(alloc.take)) + if fault == "cancel": + monkeypatch.setattr(lender._slot_queue, "cancel", planted(lender._slot_queue.cancel)) mgr.free_resources(other) # the manager's free runs the lender's hook assert kit.kv(mgr, other) is None + assert behind.failure is not None, "a write waiting for slots outlived its request" if fault == "none": - assert fetch.poll() is not None and fetch.failure is None, "the free granted nothing" + lender.progress() + assert fetch.poll() is not None and fetch.failure is None, "the round granted nothing" assert granted.poll() is not None - seen.mark_arrived(seen_view.row_masks(True)) + kit.return_write(lender, seen, seen_view.row_masks(True)) mgr._stream.synchronize() for run in seen_view.groups: - nbytes = lender.regions[run.region_index].slot_bytes - for ordinal in run.block_ordinals.tolist(): + nbytes = kit.staged_slots(lender, run)[0].slot_bytes + for ordinal in run.page_ordinals.tolist(): got = kit.digest(kit.page(mgr, seen_target, run.layer_group_index, ordinal)) want = kit.digest(kit.staged_row(0x5A, run.layer_group_index, ordinal, nbytes)) - assert got == want, "a mark copied nothing" + assert got == want, "a return copied nothing" for lease in (fetch, granted): - lease.mark_arrived(lease.poll().row_masks()) - for lease in (seen, granted, fetch): - lease.release() + kit.return_write(lender, lease, lease.poll().row_masks()) + behind.release() mgr.shutdown() assert not kit.staging_kept(regions), "nothing open, yet the staging memory stays" return - assert len(taken) == 2, "the planted fault did not strike: proves nothing" + assert len(cancelled) == 1, "the planted fault did not strike: proves nothing" free = {g: lender._free_slots(g) for g in groups} - intervals = [lender.resume_bounds(r) for r in (target, seen_target)] # progress comes first + lender.progress() assert fetch.failure is not None, "the waiting fetch neither became ready nor failed" stop = fetch.failure assert stop == _lender._REASON_STOPPED, stop assert fetch.poll() is None - for interval in intervals: + for request in (target, seen_target): + interval = lender.poll_resume_bounds(request) assert interval is not None and interval.max_position < interval.min_position, ( f"resume_bounds {interval} after the stop, not an empty interval" ) assert granted.poll() is None and granted.failure == stop, "a view went out after the stop" assert seen.poll() is seen_view, "a view a poll returned before the stop was taken back" for lend, request in ((lender.lend_read, source), (lender.lend_write, target)): - lease = lend(request, 0, WINDOWED_END) - assert lease.failure == stop, ( - f"{lend.__name__} did not fail at the call: {lease.failure}" - ) - lease.release() + with pytest.raises(RuntimeError, match="lends no more") as raised: + lend(request, 0, WINDOWED_END, **kit.HOST_NAMED) + assert stop in str(raised.value), f"{lend.__name__}: {raised.value}" with pytest.raises(ValueError): - lender.lend_read(source, 0, TPB + 1) # not whole blocks + lender.lend_read(source, 0, TPB + 1, **kit.HOST_NAMED) # not whole blocks with pytest.raises(ValueError): - lender.lend_write(target, 0, WINDOWED_END + TPB) # past the last prompt token's block - seen.mark_arrived(seen_view.row_masks(True)) + # Past the last prompt token's block. + lender.lend_write(target, 0, WINDOWED_END + TPB, **kit.HOST_NAMED) + bounds = kit.return_write(lender, seen, seen_view.row_masks(True)) + assert bounds is not None and bounds.max_position < bounds.min_position, ( + f"a return after the stop gave {bounds}, not an empty interval" + ) with pytest.raises(RuntimeError): - seen.mark_arrived(seen_view.row_masks(True)) + kit.return_write(lender, seen, seen_view.row_masks(True)) mgr._stream.synchronize() for run in seen_view.groups: - sentinel = bytes([kit.SENTINEL]) * lender.regions[run.region_index].slot_bytes - for ordinal in run.block_ordinals.tolist(): + sentinel = bytes([kit.SENTINEL]) * kit.staged_slots(lender, run)[0].slot_bytes + for ordinal in run.page_ordinals.tolist(): got = kit.digest(kit.page(mgr, seen_target, run.layer_group_index, ordinal)) - assert got == kit.digest(sentinel), "a mark copied rows after the stop" - for lease in (seen, granted, fetch): + assert got == kit.digest(sentinel), "a return copied rows after the stop" + for lease in (granted, fetch, behind): lease.release() assert lender._open_lease_count() == 0 + lender.progress() assert {g: lender._free_slots(g) for g in groups} == free, "slots returned or granted" warnings = kit.lender_warnings(monkeypatch) mgr.shutdown() @@ -4509,21 +4627,15 @@ def call(start, count): # -- names ------------------------------------------------------------------------------------ -def _tp2(rank): - from tensorrt_llm.mapping import Mapping - - return Mapping(world_size=2, tp_size=2, rank=rank) - - def names_by_rank(kit, real_manager, attach, num_kv_heads, scopes=(SCOPE, SCOPE)): """Each rank of a TP2 pair publishes ``PROMPT``: its names, and its region names.""" names, region_labels = [], [] for rank, scope in enumerate(scopes): - with real_manager(num_kv_heads=num_kv_heads, mapping=_tp2(rank)) as mgr: + with real_manager(num_kv_heads=num_kv_heads, mapping=tp_rank(rank)) as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END, scope=scope) - lease = lender.lend_read(source, 0, END) - (run,) = ready_view(lease, mgr).groups + lender = attach(mgr, guaranteed_fetch_tokens=END, scope=scope) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + (run,) = kit.ready_view(lease, mgr).groups names.append(np.array(run.object_keys)) region_labels.append([p.layout_label for p in lender.regions]) lease.release() @@ -4551,11 +4663,10 @@ def test_block_names_follow_the_scope_and_region_labels_only_the_layout(kit, rea for scope in (b"model-a", b"model-b"): with real_manager() as mgr: source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END, scope=scope) - lease = lender.lend_read(source, 0, END) - names.append( - ([bytes(n) for n in ready_view(lease, mgr).groups[0].object_keys], lender.regions) - ) + lender = attach(mgr, guaranteed_fetch_tokens=END, scope=scope) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + view = kit.ready_view(lease, mgr) + names.append(([bytes(n) for n in view.groups[0].object_keys], lender.regions)) lease.release() (first, first_regions), (second, second_regions) = names assert [n[16:] for n in first] == [n[16:] for n in second], "same blocks, same keys" @@ -4564,7 +4675,7 @@ def test_block_names_follow_the_scope_and_region_labels_only_the_layout(kit, rea "laid out alike" ) with real_manager(tokens_per_block=2 * TPB) as mgr: - other_layout = attach(mgr, tokens_per_fetch=END, scope=b"model-a").regions + other_layout = attach(mgr, guaranteed_fetch_tokens=END, scope=b"model-a").regions assert [p.layout_label for p in other_layout] != [p.layout_label for p in first_regions], ( "another layout" ) @@ -4605,16 +4716,19 @@ def test_a_request_whose_kv_its_names_miss_is_not_lent(kit, real_manager, case, target = carrying(kit, TARGET, OTHER_PROMPT, inputs) assert mgr.prepare_context(target) control = kit.published(mgr, OTHER, CONTROL_PROMPT, **covered) - lender = attach(mgr, tokens_per_fetch=END, planned_fetches=2) + lender = attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2) before = [state_of(kit, mgr, lender, r) for r in (source, target)] - refused = [lender.lend_read(source, 0, END), lender.lend_write(target, 0, END)] + refused = [ + lender.lend_read(source, 0, END, **kit.HOST_NAMED), + lender.lend_write(target, 0, END, **kit.HOST_NAMED), + ] for lease in refused: assert lease.poll() is None and why in str(lease.failure), ( f"lent a request carrying {case.replace('_', ' ')}" ) assert [state_of(kit, mgr, lender, r) for r in (source, target)] == before - lent = lender.lend_read(control, 0, END) - assert check_staged(kit, mgr, lender, control, ready_view(lent, mgr)) == BLOCKS + lent = lender.lend_read(control, 0, END, **kit.HOST_NAMED) + assert check_staged(kit, mgr, lender, control, kit.ready_view(lent, mgr)) == BLOCKS for lease in (*refused, lent): lease.release() @@ -4655,9 +4769,9 @@ def test_a_fetch_ending_inside_a_bidirectional_run_fails_at_the_call( target = carrying(kit, TARGET, WINDOWED_PROMPT, inputs) assert mgr.prepare_context(target) kit.kv(mgr, target).enable_swa_scratch_reuse = False # a fetch target needs it off - lender = attach(mgr, tokens_per_fetch=end) + lender = attach(mgr, guaranteed_fetch_tokens=end) before = state_of(kit, mgr, lender, target) - lease = lender.lend_write(target, 0, end) + lease = lender.lend_write(target, 0, end, **kit.HOST_NAMED) if refused: assert lease.poll() is None and "multimodal" in str(lease.failure), ( f"lent a fetch ending inside the run {run}: {case}" @@ -4666,7 +4780,7 @@ def test_a_fetch_ending_inside_a_bidirectional_run_fails_at_the_call( "the failed lease changed the cache" ) else: - ready_view(lease, mgr) + kit.ready_view(lease, mgr) lease.release() @@ -4693,10 +4807,10 @@ def test_resume_bounds_keeps_resumes_out_of_bidirectional_runs( target = carrying(kit, TARGET, WINDOWED_PROMPT, inputs) assert mgr.prepare_context(target) kit.kv(mgr, target).enable_swa_scratch_reuse = False # a fetch target needs it off - lender = attach(mgr, tokens_per_fetch=end) + lender = attach(mgr, guaranteed_fetch_tokens=end) fetch_segment(kit, mgr, lender, target, 0, end, 0x5A, arrived) for _ in range(2): # the second answer comes from the record resume_bounds keeps - resume_bounds = tuple(lender.resume_bounds(target)) + resume_bounds = tuple(lender.poll_resume_bounds(target)) assert resume_bounds == expected, ( f"resume_bounds {resume_bounds} lets the request resume inside a run of {runs}: {case}" ) @@ -4711,9 +4825,9 @@ def test_resume_bounds_keeps_a_resume_out_of_the_run_the_context_position_is_in( with real_manager(**PER_REQUEST) as mgr: inputs = bidirectional([(40, 100)], True, len(SPLIT_PROMPT)) target = computed(kit, mgr, TARGET, SPLIT_PROMPT, **inputs) - lender = attach(mgr, tokens_per_fetch=2 * TPB) - lender.lend_write(target, LOCAL, 4 * TPB).release() # before its first poll: abandoned - resume_bounds = lender.resume_bounds(target) + lender = attach(mgr, guaranteed_fetch_tokens=2 * TPB) + lender.lend_write(target, LOCAL, 4 * TPB, **kit.HOST_NAMED).release() # abandoned unpolled + resume_bounds = lender.poll_resume_bounds(target) assert tuple(resume_bounds) == (40, LOCAL), ( f"resume_bounds {tuple(resume_bounds)} lets the request resume inside the run [40, 100)" ) @@ -4732,24 +4846,25 @@ def test_threads_take_turns_and_the_lender_starts_none(kit, real_manager, monkey target = kit.admitted(mgr, TARGET, OTHER_PROMPT) before = set(threading.enumerate()) built = kit.on_thread( - lambda: attach(mgr, tokens_per_fetch=END, planned_fetches=2), "builder", cuda=True + lambda: attach(mgr, guaranteed_fetch_tokens=END, capacity_in_fetches=2), + "builder", + cuda=True, ) lender = built["value"] - regions = lender.regions + regions = kit.staging_regions(lender) (memory,) = kit.staging_memory(regions) base = regions[0].address if pinned: assert kit.pinned_range(base) is not None def executor_loop(): - read = lender.lend_read(source, 0, END) - view = ready_view(read, mgr) + read = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + view = kit.ready_view(read, mgr) # A backend reads the arrays on its own thread until the release. seen = kit.on_thread(lambda: [bytes(n) for n in view.groups[0].object_keys], "backend") - write = lender.lend_write(target, 0, END) - write.mark_arrived(write.poll().row_masks(True)) + write = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + kit.return_write(lender, write, write.poll().row_masks(True)) read.release() - write.release() return seen["value"] looped = kit.on_thread(executor_loop, "executor-loop", cuda=True) @@ -4761,93 +4876,7 @@ def executor_loop(): if pinned: assert kit.pinned_range(base) is None, "the shutdown thread did not free pinned memory" with pytest.raises(ValueError): - lender.resume_bounds(source) - - -# -- the staging lease's states -------------------------------------------------------------- - -UNMARKED_SLOTS_BACK = "a write released unmarked gave its slots back while a backend may write" - - -def test_a_staging_lease_walks_through_the_states_its_enum_names( - kit, real_manager, monkeypatch, attach=attach -): - """Reads and writes through every ``_LeaseState``, ``_state()`` read after each call: a read - granted, ready and released, slots back; a waiting read granted; a write marked and released, - slots back; a doomed write failed at its first poll; a lease failed at the call; a write - released ready and unmarked, slots still held; a waiting read failed at its request's free.""" - from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _lender - - states = _lender._LeaseState - seen = set() - - def state(lease, expected): - got = lease._state() - seen.add(got) - assert got is expected, f"{lease!r}: {got.name}, not {expected.name}" - - def holding(lease): - return any(held is lease for held in lender._leases_holding_slots) - - view = _lender.Staging._build_view - planted = [] - - def raising_once(self, rows, keys): - if not planted: - planted.append(True) - raise MemoryError("planted") - return view(self, rows, keys) - - with real_manager() as mgr: - source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - read = lender.lend_read(source, 0, END) - state(read, states.GRANTED) - ready_view(read, mgr) - state(read, states.READY) - waiting = lender.lend_read(source, 0, END) - state(waiting, states.WAITING) - read.release() - state(read, states.RELEASED) - assert not holding(read), "a released read kept its slots" - state(waiting, states.GRANTED) - ready_view(waiting, mgr) - waiting.release() - write = lender.lend_write(kit.admitted(mgr, TARGET, OTHER_PROMPT), 0, END) - state(write, states.GRANTED) - write_view = write.poll() - state(write, states.READY) - write.mark_arrived(write_view.row_masks()) - state(write, states.MARKED) - write.release() - state(write, states.RELEASED) - mgr._stream.synchronize() - lender.resume_bounds(source) # any call makes progress - assert not holding(write), "a marked write kept its slots once its copy had landed" - with monkeypatch.context() as patched: - patched.setattr(_lender.Staging, "_build_view", raising_once) - doomed = lender.lend_write(kit.admitted(mgr, 4, list(range(8000, 8097))), 0, END) - state(doomed, states.FAILS_AT_FIRST_POLL) - assert doomed.poll() is None - state(doomed, states.FAILED) - doomed.release() - failed = lender.lend_read(kit.make_request(9, PROMPT), 0, END) - state(failed, states.FAILED) - failed.release() - unmarked = lender.lend_write(kit.admitted(mgr, 5, list(range(9000, 9097))), 0, END) - assert unmarked.poll() is not None - state(unmarked, states.READY) - unmarked.release() - state(unmarked, states.RELEASED) - lender.resume_bounds(source) - assert holding(unmarked), UNMARKED_SLOTS_BACK - freed = kit.published(mgr, 6, list(range(10000, 10097))) - in_line = lender.lend_read(freed, 0, END) - state(in_line, states.WAITING) - mgr.free_resources(freed) - state(in_line, states.FAILED) - in_line.release() - assert seen == set(states), f"states never seen: {set(states) - seen}" + lender.poll_resume_bounds(source) # -- the manager's page-index buffer, kept from the attach until the runtime's shutdown returns @@ -4859,7 +4888,7 @@ def test_the_index_buffer_is_kept_from_the_attach_until_the_shutdown( with real_manager() as mgr: before = len(kit.retained()) buffer = mgr.host_kv_cache_block_offsets - attach(mgr, tokens_per_fetch=END) + attach(mgr, guaranteed_fetch_tokens=END) assert any(o is buffer for o in kit.retained()), "not kept from the attach" mgr.shutdown() assert not any(o is buffer for o in kit.retained()), "kept past the shutdown" @@ -4894,10 +4923,9 @@ def test_the_index_buffer_stays_kept_through_a_close_that_raises(kit, real_manag with real_manager() as mgr: buffer = mgr.host_kv_cache_block_offsets target = kit.admitted(mgr, TARGET, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - lease = lender.lend_write(target, 0, END) - lease.mark_arrived(ready_view(lease, mgr).row_masks()) - lease.release() + lender = attach(mgr, guaranteed_fetch_tokens=END) + lease = lender.lend_write(target, 0, END, **kit.HOST_NAMED) + kit.return_write(lender, lease, kit.ready_view(lease, mgr).row_masks()) kv = kit.kv(mgr, target) request_id = next(rid for rid, cache in mgr.kv_cache_map.items() if cache is kv) mgr.kv_cache_map[request_id] = ClosingSecondTime(kv) # the released lease's own cache @@ -4919,9 +4947,9 @@ def test_the_index_buffer_stays_kept_past_a_free_whose_close_raised( with real_manager() as mgr: buffer = mgr.host_kv_cache_block_offsets source = kit.published(mgr, SOURCE, PROMPT) - lender = attach(mgr, tokens_per_fetch=END) - lease = lender.lend_read(source, 0, END) - ready_view(lease, mgr) + lender = attach(mgr, guaranteed_fetch_tokens=END) + lease = lender.lend_read(source, 0, END, **kit.HOST_NAMED) + kit.ready_view(lease, mgr) lease.release() # a released read holds its cache until it is dropped kv = kit.kv(mgr, source) mgr.kv_cache_map[source.py_request_id] = ClosingSecondTime(kv) @@ -4962,13 +4990,15 @@ def check_a_lease_outliving_its_manager(kit): kv = kit.kv(mgr, request) blocks = int(kv.num_blocks) lender = attach_lender(mgr, scope=b"index-buffer", staging=StagingOptions(END)) - lease = lender.lend_read(request, 0, END) + lease = lender.lend_read(request, 0, END, **kit.HOST_NAMED) mgr._stream.synchronize() # a staging read is ready once its copy into staging ran + lender.progress() view = lease.poll() assert view is not None (run,) = view.groups - address, slot_bytes = int(run.addresses[0]), lender.regions[run.region_index].slot_bytes - regions, content = lender.regions, kit.host_bytes(address, slot_bytes) + region, addresses = kit.staged_slots(lender, run) + address, slot_bytes = addresses[0], region.slot_bytes + regions, content = kit.staging_regions(lender), kit.host_bytes(address, slot_bytes) row = kit.index_row(mgr, request) own = list(kv.get_base_page_indices(0)[:blocks]) assert row.values[:blocks] == own and -1 not in own, "the cache writes elsewhere" @@ -5018,32 +5048,27 @@ def test_a_deepseek_v4_prompt_is_fetched_through_staging_byte_for_byte( end = 3 * tpb prompt = list(range(3000, 3000 + end + 1)) source = kit.published(mgr_a, SOURCE, prompt) - lender_a = attach(mgr_a, tokens_per_fetch=end) - publish = lender_a.lend_read(source, 0, end) - publish_view = ready_view(publish, mgr_a) + lender_a = attach(mgr_a, guaranteed_fetch_tokens=end) + publish = lender_a.lend_read(source, 0, end, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, mgr_a) assert len(publish_view.groups) == kit.num_layer_groups(mgr_a) - assert ( - check_staged(kit, mgr_a, lender_a, source, publish_view) - == publish_view.num_block_entries - ) + assert check_staged(kit, mgr_a, lender_a, source, publish_view) == publish_view.num_rows target = kit.admitted(mgr_b, TARGET, prompt) - lender_b = attach(mgr_b, tokens_per_fetch=end) + lender_b = attach(mgr_b, guaranteed_fetch_tokens=end) assert [p.layout_label for p in lender_a.regions] == [ p.layout_label for p in lender_b.regions ] - lease = lender_b.lend_write(target, 0, end) + lease = lender_b.lend_write(target, 0, end, **kit.HOST_NAMED) view = lease.poll() assert by_group(view) == by_group(publish_view) kit.fill_sentinel(mgr_b, target) - lease.mark_arrived(kit.relay(lender_a, publish_view, lender_b, view)) - mgr_b._stream.synchronize() - assert lender_b.resume_bounds(target) == (end, end) + masks = kit.relay(lender_a, publish_view, lender_b, view) + assert return_and_settle(lender_b, mgr_b, target, lease, masks) == (end, end) for run in view.groups: - for ordinal in run.block_ordinals.tolist(): + for ordinal in run.page_ordinals.tolist(): lg = run.layer_group_index got = kit.digest(kit.page(mgr_b, target, lg, ordinal)) assert got == kit.digest(kit.page(mgr_a, source, lg, ordinal)) - lease.release() publish.release() @@ -5054,42 +5079,35 @@ def ngram_drafting(draft_len): return NGramDecodingConfig(max_draft_len=draft_len, max_matching_ngram_size=2) -# Drafting without read-ahead widens each DeepSeek-V4 window by the draft length, so a fetch of -# [0, 512) also asks for that margin; layer group 0 has windows of 128. case -> (spec config, -# tokens the publisher computes past 512, rows the publish lacks per layer group, resume_bounds) +# The lender conservatively reads DeepSeek-V4's full kept window, including the draft margin. +# case -> (spec config, tokens the publisher computes past 512, rows the publish lacks per group) DSV4_TPB = 128 DSV4_END = 4 * DSV4_TPB DSV4_PROMPT = list(range(60_000, 60_000 + DSV4_END + 2)) # four whole blocks and two tokens DSV4_TARGET = DSV4_PROMPT[:DSV4_END] + [69_999, 69_998] # the same four blocks, then others DSV4_DRAFT_MARGIN = { # A window of 128 at 512 keeps block 3 alone, which a publisher two tokens past still keeps. - "no_speculation": (lambda: None, 2, {}, (DSV4_END, DSV4_END)), - # A draft length of 2 keeps block 2 too, which only a publisher whose history stands at the - # fetch's end still keeps; past it, the interval is empty. - "draft2_publisher_at_the_end": (ngram, 0, {}, (DSV4_END, DSV4_END)), - "draft2_publisher_one_past": (ngram, 1, {0: [2]}, (0, DSV4_END)), + "no_speculation": (lambda: None, 2, {}), + # A draft length of 2 keeps block 2 too. A publisher past the fetch's end lacks it, so + # the conservative read contract cannot resume there until that margin row arrives. + "draft2_publisher_at_the_end": (ngram, 0, {}), + "draft2_publisher_one_past": (ngram, 1, {0: [2]}), # A draft length of 3 keeps block 2 for a publisher one token past, not two. - "draft3_publisher_one_past": (lambda: ngram_drafting(3), 1, {}, (DSV4_END, DSV4_END)), - "draft3_publisher_two_past": (lambda: ngram_drafting(3), 2, {0: [2]}, (0, DSV4_END)), + "draft3_publisher_one_past": (lambda: ngram_drafting(3), 1, {}), + "draft3_publisher_two_past": (lambda: ngram_drafting(3), 2, {0: [2]}), # PARD also grows the cache by the extra tokens embedded DSpark does. - "pard2_publisher_two_past": (pard, 2, {0: [2]}, (0, DSV4_END)), + "pard2_publisher_two_past": (pard, 2, {0: [2]}), } NOT_COMPUTED = "holds bytes other than the ones the publisher computed" @skip_pre_blackwell @pytest.mark.parametrize("case", list(DSV4_DRAFT_MARGIN)) -def test_a_deepseek_v4_fetch_under_speculative_decoding_asks_for_the_window_margin( +def test_a_deepseek_v4_fetch_requires_every_page_in_its_conservative_kept_window( kit, deepseek_v4_manager, case ): - """DeepSeek-V4 keeps every window the draft length wider under any speculative decoding, also - one reading no token ahead, which staging accepts. A fetch asks for every row the cache's - windows keep at its end, the margin's included; the rows the publish holds arrive with the bytes - the publisher computed, and resume_bounds counts a position only where a resume there finds those - bytes in every page the cache keeps below it. At 128 tokens per block and a draft length of 2 or - more, a publisher whose history stands at least the draft length less one token past the fetch's - end has passed a margin row, so the interval is empty and the request computes from 0.""" - make, tail, lacking, expected = DSV4_DRAFT_MARGIN[case] + """A missing kept-margin page prevents resume until a later fetch delivers that page.""" + make, tail, lacking = DSV4_DRAFT_MARGIN[case] with ( deepseek_v4_manager(spec_config=make()) as pub, deepseek_v4_manager(spec_config=make()) as mgr, @@ -5097,28 +5115,25 @@ def test_a_deepseek_v4_fetch_under_speculative_decoding_asks_for_the_window_marg assert kit.windows(mgr)[0] == DSV4_TPB + mgr.max_draft_len assert mgr.reuse_match_backoff == 0 source = kit.published(pub, SOURCE, DSV4_PROMPT[: DSV4_END + tail]) - publisher = attach(pub, tokens_per_fetch=DSV4_END) - publish = publisher.lend_read(source, 0, DSV4_END) - publish_view = ready_view(publish, pub) - assert ( - check_staged(kit, pub, publisher, source, publish_view) - == publish_view.num_block_entries - ) + publisher = attach(pub, guaranteed_fetch_tokens=DSV4_END) + publish = publisher.lend_read(source, 0, DSV4_END, **kit.HOST_NAMED) + publish_view = kit.ready_view(publish, pub) + assert check_staged(kit, pub, publisher, source, publish_view) == publish_view.num_rows held = {name.tobytes() for run in publish_view.groups for name in run.object_keys} target = kit.admitted(mgr, TARGET, DSV4_TARGET) - lender = attach(mgr, tokens_per_fetch=DSV4_END) - lease = lender.lend_write(target, 0, DSV4_END) - view = ready_view(lease, mgr) + lender = attach(mgr, guaranteed_fetch_tokens=DSV4_END) + lease = lender.lend_write(target, 0, DSV4_END, **kit.HOST_NAMED) + view = kit.ready_view(lease, mgr) kit.fill_sentinel(mgr, target) def published_row(i, row): return view.groups[i].object_keys[row].tobytes() in held - lease.mark_arrived(kit.relay(publisher, publish_view, lender, view, published_row)) - lease.release() + masks = kit.relay(publisher, publish_view, lender, view, published_row) + usable, floor = return_and_settle(lender, mgr, target, lease, masks) publish.release() - mgr._stream.synchronize() - usable, floor = lender.resume_bounds(target) + for group in lender.page_layout.layer_groups: + assert group.read_window == group.kept_window, group kv = kit.kv(mgr, target) keys = kit.chain_keys(kv, DSV4_TARGET) dev = kit.DevicePages(mgr) @@ -5131,26 +5146,53 @@ def wrong(lg, ordinal): for i, run in enumerate(view.groups): lg = run.layer_group_index - for row, ordinal in enumerate(run.block_ordinals.tolist()): + for row, ordinal in enumerate(run.page_ordinals.tolist()): got = wrong(lg, ordinal) if published_row(i, row) else None assert got is None, f"fetched block {ordinal} of layer group {lg} {NOT_COMPUTED}" if floor <= usable: - for lg in range(kit.num_layer_groups(mgr)): + for lg in range(len(lender.page_layout.layer_groups)): beg, end = kit.stale_blocks(mgr, lg, usable) for ordinal in range(-(-usable // DSV4_TPB)): got = None if beg <= ordinal < end else wrong(lg, ordinal) assert got is None, ( - f"block {ordinal} of layer group {lg}, which the cache keeps at the resume " + f"block {ordinal} of layer group {lg}, which attention reads at the resume " f"point {usable}, {NOT_COMPUTED}: {got}" ) resume_bounds = (usable, floor) - assert resume_bounds == expected, ( - f"resume_bounds {resume_bounds}, the rules give {expected}" - ) asked = {} for i, run in enumerate(view.groups): - rows = enumerate(run.block_ordinals.tolist()) + rows = enumerate(run.page_ordinals.tolist()) missing = [ordinal for row, ordinal in rows if not published_row(i, row)] if missing: asked[run.layer_group_index] = missing assert asked == lacking, f"the fetch's rows the publish lacks: {asked}, expected {lacking}" + if not lacking: + assert resume_bounds == (DSV4_END, DSV4_END) + return + assert not floor <= DSV4_END <= usable, ( + f"resume_bounds {resume_bounds} counts the end despite missing kept pages {lacking}" + ) + assert usable < floor, f"missing kept pages still permit a resume: {resume_bounds}" + for lg, ordinals in lacking.items(): + for ordinal in ordinals: + assert wrong(lg, ordinal) is not None, "the missing page already holds source bytes" + + repair_source = kit.published(pub, OTHER, DSV4_PROMPT[:DSV4_END]) + repair_publish = publisher.lend_read(repair_source, 0, DSV4_END, **kit.HOST_NAMED) + repair_publish_view = kit.ready_view(repair_publish, pub) + repair = lender.lend_write(target, 0, DSV4_END, **kit.HOST_NAMED) + repair_view = kit.ready_view(repair, mgr) + + def missing_row(i, row): + run = repair_view.groups[i] + return int(run.page_ordinals[row]) in lacking.get(run.layer_group_index, ()) + + repair_masks = kit.relay(publisher, repair_publish_view, lender, repair_view, missing_row) + assert sum(int(mask.sum()) for mask in repair_masks) == sum(map(len, lacking.values())) + assert return_and_settle(lender, mgr, target, repair, repair_masks) == (DSV4_END, DSV4_END) + repair_publish.release() + for run in repair_view.groups: + for ordinal in run.page_ordinals.tolist(): + assert wrong(run.layer_group_index, ordinal) is None, ( + f"repaired fetch lost block {ordinal} of layer group {run.layer_group_index}" + ) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 951be07f199b..a5398f8fda21 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -79,6 +79,7 @@ def test_resume_restores_page_buffers_and_sparse_metadata_row(sparse: bool) -> N manager.index_mapper = Mock() manager.index_mapper.get_index.return_value = 3 manager._set_page_index_bufs = Mock() + manager._early_freed_index_requests = set() manager.sparse_metadata_batch = Mock() if sparse else None cache = Mock() @@ -90,6 +91,30 @@ def test_resume_restores_page_buffers_and_sparse_metadata_row(sparse: bool) -> N manager.sparse_metadata_batch.add.assert_called_once_with(cache, 3) +@pytest.mark.parametrize("parked", [False, True]) +def test_resume_binds_a_released_slot_s_parked_row_again(parked: bool) -> None: + """A resume of a request whose index slot was released early binds its parked row, if any.""" + manager = object.__new__(KVCacheManagerV2) + manager.index_mapper = Mock() + manager._set_page_index_bufs = Mock() + manager._early_freed_index_requests = {7} + manager.num_pools = 2 + manager.sparse_metadata_batch = Mock() + manager._sharing = Mock() + manager._sharing._parked_row.return_value = ("row 0", "row 1") if parked else None + cache = Mock() + + manager._restore_page_index_bufs(7, cache) + + manager._set_page_index_bufs.assert_not_called() + manager.index_mapper.get_index.assert_not_called() + manager.sparse_metadata_batch.add.assert_not_called() + manager._sharing._parked_row.assert_called_once_with(7) + manager._sharing._on_resume.assert_called_once_with(7) + bound = [args for args, _ in cache.set_base_page_index_buf.call_args_list] + assert bound == ([(0, 0, "row 0"), (0, 1, "row 1")] if parked else []) + + @pytest.mark.parametrize("is_draft", [False, True]) def test_sparse_metadata_publishes_after_preparation(is_draft: bool) -> None: manager = object.__new__(KVCacheManagerV2) @@ -616,6 +641,9 @@ def make_kv_cache_manager( ): mgr = Mock() mgr.kv_connector_manager = None + # No lender attached: no page locks, and the history a window would reach is left as it is. + mgr._holds_page_locks.return_value = False + mgr._lent_history.side_effect = lambda kv_cache, history: history mgr.tokens_per_block = tokens_per_block # A real policy, not an auto-created Mock attribute: the prefix-aware skip # is gated on ALL_REUSABLE (see _skip_pays_off_under_reuse_policy). @@ -5010,6 +5038,91 @@ def test_pending_connector_load_does_not_self_evict_or_report_deadlock() -> None manager.free_resources.assert_not_called() +def test_a_request_holding_page_locks_is_neither_evicted_nor_paused_for_another() -> None: + """Out of pages, a generation request takes no victim among the requests holding page locks.""" + attempts = [] + + def allocate(request): + attempts.append(request.py_request_id) + return len(attempts) > 1 # the first fails; a retry after a victim's release succeeds + + manager = make_kv_cache_manager(can_evict=True, try_allocate_generation_fn=allocate) + manager._holds_page_locks.side_effect = lambda request_id: request_id == 99 + scheduler = make_scheduler(manager, max_num_tokens=100) + locked = make_gen_request(99) + assert not scheduler._is_evictable(locked, set()) + assert not scheduler._is_recompute_pause_candidate(locked, set()) + out = scheduler.schedule_request([make_gen_request(0), locked], set()) + suspended = [call.args[0].py_request_id for call in manager.suspend_request.call_args_list] + taken = ids(out.paused_requests) + ids(out.recompute_paused_requests) + suspended + assert 99 not in taken, f"the locked request taken: {taken}" + assert ids(out.paused_requests) == [0] and suspended == [0] + + +def test_a_request_holding_page_locks_suspends_neither_itself_nor_its_cache() -> None: + """Out of pages with no victim, a request holding page locks stays unscheduled and running.""" + manager = make_kv_cache_manager(can_evict=True, try_allocate_generation_fn=lambda r: False) + manager._holds_page_locks.side_effect = lambda request_id: request_id == 0 + scheduler = make_scheduler(manager, max_num_tokens=100) + locked = make_gen_request(0) + out = scheduler.schedule_request([locked], set()) + assert out.generation_requests == [] and ids(out.paused_requests) != [0] + scheduler._suspend_request(locked) + assert not manager.suspend_request.called + + +def test_a_stall_while_a_lender_locks_pages_is_not_a_deadlock() -> None: + """Passes that reclaim nothing while a lender's locks hold pages are no stall.""" + manager = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, has_cache_tier_below_gpu=False + ) + manager._holds_page_locks.side_effect = lambda request_id: request_id is None + scheduler = make_scheduler(manager, max_num_tokens=1000) + scheduler._DEADLOCK_STALL_ITERS = 3 + requests = [make_ctx_request(0, 100, is_first_context_chunk=False)] + for _ in range(7): + scheduler.schedule_request(requests, set()) + + +def test_a_request_holding_page_locks_in_the_draft_pool_counts_as_locked() -> None: + manager = make_kv_cache_manager(can_evict=True) + draft = make_kv_cache_manager(can_evict=True) + draft.__class__ = KVCacheManagerV2 + draft._holds_page_locks.side_effect = lambda request_id: request_id == 5 + scheduler = make_scheduler(manager, draft_kv_cache_manager=draft) + assert scheduler._holds_page_locks(make_gen_request(5)) + assert not scheduler._holds_page_locks(make_gen_request(6)) + assert not scheduler._is_evictable(make_gen_request(5), set()) + + +@pytest.mark.parametrize("legacy_hook", ["missing", "non_callable", "callable"]) +def test_a_legacy_draft_manager_is_unlocked_regardless_of_page_lock_hook(legacy_hook: str) -> None: + manager = make_kv_cache_manager(can_evict=True) + draft = make_kv_cache_manager(can_evict=True) + if legacy_hook == "missing": + del draft._holds_page_locks + elif legacy_hook == "non_callable": + draft._holds_page_locks = False + scheduler = make_scheduler(manager, draft_kv_cache_manager=draft) + request = make_gen_request(5) + + assert not scheduler._holds_page_locks(request) + manager._holds_page_locks.assert_called_once_with(5) + assert scheduler._is_evictable(request, set()) + if legacy_hook == "callable": + draft._holds_page_locks.assert_not_called() + + +def test_target_page_lock_short_circuits_draft_page_lock_query() -> None: + manager = make_kv_cache_manager(can_evict=True) + manager._holds_page_locks.return_value = True + draft = make_kv_cache_manager(can_evict=True) + scheduler = make_scheduler(manager, draft_kv_cache_manager=draft) + + assert scheduler._holds_page_locks(make_gen_request(5)) + draft._holds_page_locks.assert_not_called() + + def test_parked_connector_load_keeps_kv_pressure_retryable() -> None: manager = make_kv_cache_manager(try_allocate_generation_fn=lambda request: False) connector = Mock() diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py index 143d3a336e5b..dd739791a80b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py @@ -61,6 +61,7 @@ def _make_executor( enable_attention_dp: bool = False, disable_overlap_scheduler: bool = False, kv_cache_transceiver=None, + attached_lender=None, kv_connector_manager=None, is_warmup: bool = False, is_shutdown: bool = False, @@ -130,6 +131,7 @@ def _make_executor( # KV cache manager (resource-manager wrapper). exe.kv_cache_manager = MagicMock() exe.kv_cache_manager.max_beam_width = max_beam_width + exe.kv_cache_manager._sharing = attached_lender exe.kv_cache_manager.impl = MagicMock() exe.kv_cache_manager.impl.need_adjustment = need_adjustment @@ -244,6 +246,10 @@ def test_transceiver_present_returns_false(self): exe = _make_executor(kv_cache_transceiver=MagicMock()) assert PyExecutor._can_pause_for_rebalance(exe) is False + def test_attached_lender_prevents_rebalance(self): + exe = _make_executor(attached_lender=MagicMock()) + assert PyExecutor._can_pause_for_rebalance(exe) is False + def test_connector_present_returns_false(self): """A rebalance reassigns ``slot_id``, and a connector holds addresses derived from it across iterations (``update_state_after_alloc``), so diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index addb326f0a91..38d8d59fc813 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -3107,6 +3107,29 @@ def test_disagg_role_mapper_kinds_default_to_indexed(): } +def test_page_buffer_tensors_default_to_head_major(): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role + from tensorrt_llm._torch.pyexecutor.kv_cache.page_layout import head_major, replicated + + manager = object.__new__(KVCacheManagerV2) + manager.mapping = SimpleNamespace(tp_size=2, tp_rank=0) + manager._page_kv_heads = lambda layer: (2, 4) + + # K/V default to the TRTLLM head-major layout, split by head over TP; the index-key side + # cache is alike on every rank (every shipped index-K, DSA V1 and MiniMax M3, is + # TP-replicated). + assert manager._page_buffer_tensors(0, Role.KEY, 256, 16, False) == head_major(256, 2, 4) + assert manager._page_buffer_tensors(0, Role.INDEX_KEY, 128, 16, False) == replicated(128) + + +def test_page_buffer_tensors_refuse_a_role_the_manager_does_not_declare(): + manager = object.__new__(KVCacheManagerV2) + with pytest.raises(ValueError, match="does not declare how its bytes spread"): + manager._page_buffer_tensors(0, "my_new_cache", 64, 16, False) + + def _index_mapper_capacity_for( *, max_batch_size: int, diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 5fd8ba58086a..55b21ef1ffba 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -4074,6 +4074,10 @@ def test_v2_hybrid_disagg_page_table_preserves_lifecycle_indices(): assert len(page_table.layer_groups) == len(mgr.impl.layer_grouping) assert isinstance(page_table.layer_groups[0], MambaLayerGroup) assert isinstance(page_table.layer_groups[1], AttentionLayerGroup) + # The manager's own page table keeps the same lifecycle order. + groups = mgr.get_page_table().page_layout.layer_groups + assert len(groups) == len(mgr.impl.layer_grouping) + assert [g.kind for g in groups[:2]] == ["state", "attention"] requests = mgr.add_dummy_requests([123], token_nums=[64], is_gen=False) assert len(requests) == 1 @@ -4102,6 +4106,13 @@ def test_v2_hybrid_disagg_page_table_uses_qwen3_next_conv_sections(): assert conv_view.section_bytes == [ dim * d_conv_m1 * conv_elem_size for dim in mgr.conv_section_dims ] + # The manager's own page table declares the same sections, one tensor each. + layout = mgr.get_page_table().page_layout + assert layout.layer_groups[0].kind == "state" + conv = next(b for b in layout.buffers if b.role == str(MambaRole.CONV_STATE)) + assert [t.axes[0].held for t in conv.tensors] == [ + dim * d_conv_m1 * conv_elem_size for dim in mgr.conv_section_dims + ] assert mgr.conv_section_dims == [8, 8, 32] finally: mgr.shutdown() @@ -5998,6 +6009,7 @@ def test_v2_kda_draft_records_replace_the_replay_caches(kda_replay_num_spec): } # A KVCM state buffer of the SSM layers, kept out of disaggregated transfers like the other replay buffers. assert DataRole("kda_state_tok") in mgr.get_disagg_ignored_roles() + assert "kda_state_tok" in mgr._page_untransferred_roles() assert len(mgr.kda_state_tok) == 2 for layer_idx in range(2): layer_cache = mgr.mamba_layer_cache(layer_idx) @@ -6341,6 +6353,7 @@ def test_v2_kda_replay_buffers_are_kvcm_state_buffers(manager_cls): ) ignored = {str(role) for role in mgr.get_disagg_ignored_roles()} assert replay_roles | {"prev_num_accepted_tokens"} <= ignored + assert mgr._page_untransferred_roles() == ignored finally: mgr.shutdown() @@ -6409,6 +6422,7 @@ def test_v2_kda_worker_role_cache(monkeypatch, role, is_disagg, manager_cls): assert not mgr.get_disagg_ignored_roles() & { DataRole(name) for name in _KDA_REPLAY_FIELDS } + assert not mgr._page_untransferred_roles() & set(_KDA_REPLAY_FIELDS) mgr.update_mamba_states( SimpleNamespace(num_seqs=1, num_contexts=1), torch.ones(1, device="cuda", dtype=torch.int32), diff --git a/tests/unittest/_torch/executor/kv_cache/test_spec_recompute_tail.py b/tests/unittest/_torch/executor/kv_cache/test_spec_recompute_tail.py index e8f2ecbf1f62..6eb2eb55edff 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_spec_recompute_tail.py +++ b/tests/unittest/_torch/executor/kv_cache/test_spec_recompute_tail.py @@ -229,6 +229,8 @@ def _reserve(req, *, local_end, offered_end, tail, tokens_per_block=4): mgr.tokens_per_block = tokens_per_block mgr._spec_recompute_tail = tail mgr.is_draft = False + # No lender attached: the history a window would reach is left as it is. + mgr._lent_history.side_effect = lambda kv_cache, history: history KVCacheManagerV2._prepare_connector_prefix_reservation(mgr, req) return connector diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py index 3688bc0ea054..9d39f307a3d3 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py @@ -1116,3 +1116,78 @@ def test_final_admission_rejects_an_invalid_local_allocation() -> None: assert req.py_served_prefix_end == 0 manager.release_unused_connector_reservations(set()) assert connector.releases == [(1, 0, 64)] + + +@pytest.mark.cpu_only +class TestBesideALender: + """The connector loads past a fetch's prefix, no further than a window keeps locked blocks.""" + + class _EveryCache(dict): + def __contains__(self, key): + return True + + @staticmethod + def locking(cap): + """A lender whose locks hold blocks of every cache and let a window reach ``cap`` tokens.""" + every = TestBesideALender._EveryCache + return SimpleNamespace( + _history_cap=lambda kv_cache: cap, + _page_lock_state=lambda: SimpleNamespace( + holding=1, locked_caches=every(), caps_history=True, unreusable={} + ), + ) + + def test_the_connector_is_asked_past_the_prefix_a_fetch_served(self): + connector = FakeConnectorManager(num_matched=32) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=64, capacity=PROMPT_LEN) + kv_cache.history_length = 128 # the fetch's resume raised it + manager.kv_cache_map[req.py_request_id] = kv_cache + req.py_served_prefix_end = 128 + req.set_prepopulated_prompt_len(128, TOKENS_PER_BLOCK) + assert manager._apply_connector_matched_prefix(req) + assert connector.queries == [(0, 128)] + assert req.context_current_position == 160 + assert connector.commits == [(0, 32, False)] + + def test_a_reservation_starts_past_the_prefix_a_fetch_served(self): + connector = FakeConnectorManager(num_matched=64) + connector.prefix_reservations_enabled = True + manager = make_manager(connector) + req = FakeRequest() + cache = FakeKvCache(committed=64, capacity=PROMPT_LEN) + cache.history_length = 128 + manager.kv_cache_map[req.request_id] = cache + req.py_served_prefix_end = 128 + assert manager.prepare_context(req) + assert connector.queries == [(0, 128)] + reservation = connector.get_prefix_reservation(req) + assert (reservation.start, reservation.end) == (128, 192) + assert req.context_current_position == 192 + + @pytest.mark.parametrize("cap,served", [(None, PROMPT_LEN - 1), (100, 100), (40, 64)]) + def test_a_served_prefix_ends_where_a_window_keeps_the_locked_blocks(self, cap, served): + """Under a lock the served end, and so the history, stops where the window keeps them.""" + connector = FakeConnectorManager(num_matched=192) + manager = make_manager(connector) + manager._sharing = self.locking(cap) + req = FakeRequest() + kv_cache = FakeKvCache(committed=64, capacity=PROMPT_LEN) + assert serve(manager, req, kv_cache) is (served > 64) + assert req.context_current_position == served + assert kv_cache.history_length == served + assert connector.commits == [(0, served - 64, False)] + + def test_a_reservation_ends_at_the_whole_blocks_a_window_keeps_under_a_lock(self): + connector = FakeConnectorManager(num_matched=128) + connector.prefix_reservations_enabled = True + manager = make_manager(connector) + manager._sharing = self.locking(100) + req = FakeRequest() + manager.kv_cache_map[req.request_id] = FakeKvCache(committed=0, capacity=0) + assert manager.prepare_context(req) + reservation = connector.get_prefix_reservation(req) + assert (reservation.start, reservation.end) == (0, 96) + assert connector.releases == [(1, 96, 128)] + assert req.context_current_position == 96 diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py index 59746a8a969f..13ec155b3468 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py @@ -20,7 +20,10 @@ """ import gc +import threading +from contextlib import contextmanager from types import SimpleNamespace +from typing import Callable, Iterator import pytest import torch @@ -29,6 +32,7 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import valid_page_slots from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import MemoryKind, RowSelection from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -273,13 +277,12 @@ def test_a_served_cache_is_allocated_off_the_execution_stream(connector): manager.report_batch_to_connector(batch) assert kv_cache.cuda_stream == execution_stream.cuda_stream - with torch.cuda.stream(execution_stream): - torch.cuda._sleep(1 << 31) + with _held_stream(execution_stream) as open_gate: forward_done = torch.cuda.Event() forward_done.record(execution_stream) - manager.wait_allocation_fences([request.py_request_id]) - assert not forward_done.query(), "the fence waited for the execution stream" - execution_stream.synchronize() + manager.wait_allocation_fences([request.py_request_id]) + assert not forward_done.query(), "the fence waited for the execution stream" + open_gate() finally: manager.shutdown() @@ -890,3 +893,101 @@ def test_real_kv_pressure_rejects_reserved_prefix_without_transmission() -> None if blocker is not None: blocker.close() manager.shutdown() + + +def _lender_beside(manager): + from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import StagingOptions, attach_lender + + return attach_lender(manager, scope=b"beside-a-connector", staging=StagingOptions(PROMPT_LEN)) + + +@contextmanager +def _held_stream(stream: torch.cuda.Stream) -> Iterator[Callable[[], None]]: + """Hold a CUDA stream on a host flag until the caller opens it.""" + from cuda.bindings import driver + + flag = torch.zeros(1, dtype=torch.int32, pin_memory=True) + (error,) = driver.cuStreamWaitValue32( + driver.CUstream(stream.cuda_stream), + driver.CUdeviceptr(flag.data_ptr()), + 1, + driver.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_GEQ, + ) + if error == driver.CUresult.CUDA_ERROR_NOT_SUPPORTED: + pytest.skip("this device cannot hold a stream on a host flag") + assert error == driver.CUresult.CUDA_SUCCESS, error + + def open_gate() -> None: + flag.numpy()[0] = 1 + + timed_out = threading.Event() + + def release_after_timeout() -> None: + timed_out.set() + open_gate() + + watchdog = threading.Timer(120.0, release_after_timeout) + watchdog.start() + try: + yield open_gate + finally: + open_gate() + watchdog.cancel() + watchdog.join() + stream.synchronize() + assert not timed_out.is_set(), "the stream stayed held: something waited on it" + + +def _waits_for_side_stream(manager, lend, request, end, rows, ready): + """A device lend gives no view before work queued on the allocation side stream completed.""" + ready_event = None + if ready == "event": + ready_event = torch.cuda.Event() + ready_event.record(manager._stream) + with _held_stream(manager._allocation_stream) as open_gate: + lease = lend( + request, 0, end, memory_kind=MemoryKind.DEVICE, rows=rows, ready_event=ready_event + ) + assert lease.failure is None, lease.failure + manager._stream.synchronize() + assert lease.poll() is None, "a view went out before the work on the allocation side stream" + open_gate() + assert lease.poll() is not None + lease.release() + + +@pytest.mark.parametrize("ready", ["none", "event"]) +@pytest.mark.parametrize("rows", [RowSelection.ALL, RowSelection.NAMED], ids=["all", "named"]) +def test_a_write_in_place_waits_for_the_allocation_side_stream(manager, connector, rows, ready): + """A write in place into a cache allocated on the side stream waits for that stream's work.""" + lender = _lender_beside(manager) + connector.prefix_reservations_enabled = False + request = make_request(request_id=3) + assert schedule(manager, request) + assert request.py_request_id in manager._on_allocation_stream + _waits_for_side_stream(manager, lender.lend_write, request, 2 * TOKENS_PER_BLOCK, rows, ready) + + +@pytest.mark.parametrize("ready", ["none", "event"]) +def test_a_read_in_place_of_a_reused_prefix_waits_for_the_allocation_side_stream( + manager, connector, ready +): + """A read in place of a reused prefix of a cache allocated on the side stream waits for it.""" + lender = _lender_beside(manager) + connector.num_matched = 0 # reuse alone gives the prefix + first = make_request(request_id=4) + assert schedule(manager, first) + batch = run(manager, first) + manager.report_batch_to_connector(batch) + first.context_chunk_size = first.context_remaining_length + first.move_to_next_context_chunk() + manager.update_context_resources(batch) + manager._stream.synchronize() + manager.free_resources(first) + second = make_request(request_id=5) + assert schedule(manager, second) + assert second.py_request_id in manager._on_allocation_stream + reused = manager.kv_cache_map[second.py_request_id].history_length + assert reused >= TOKENS_PER_BLOCK, "the second request reused no block" + end = reused // TOKENS_PER_BLOCK * TOKENS_PER_BLOCK + _waits_for_side_stream(manager, lender.lend_read, second, end, RowSelection.ALL, ready) diff --git a/tests/unittest/_torch/executor/test_layer_type_attention_windows.py b/tests/unittest/_torch/executor/test_layer_type_attention_windows.py index 67f94ca6d15f..6c2449fd0bd9 100644 --- a/tests/unittest/_torch/executor/test_layer_type_attention_windows.py +++ b/tests/unittest/_torch/executor/test_layer_type_attention_windows.py @@ -282,3 +282,46 @@ def test_draft_manager_skips_derivation(monkeypatch): _run_create_kv_cache_manager(monkeypatch, manager_cls, kv_cache_config, is_draft=True) assert captured["kv_cache_config"].max_attention_window is None + + +class _Windows: + """A manager's read-window inputs: one layer group per kept window over ``pp_layers``.""" + + def __init__(self, kept, declared, pp_layers=(2, 3), cache_layers=None): + self._kept = kept + internal = list(range(len(pp_layers))) + grouping = [internal] if len(kept) == 1 else [[layer] for layer in internal] + self.impl = SimpleNamespace(layer_grouping=grouping) + self.pp_layers = list(pp_layers) + self._model_attention_windows = declared + layers = len(pp_layers) if cache_layers is None else cache_layers + self.kv_cache_manager_py_config = SimpleNamespace(layers=[None] * layers) + + def _life_cycle_by_layer_group(self): + return [SimpleNamespace(window_size=window) for window in self._kept] + + +# Internal layers 0 and 1 are model layers 2 and 3, whose declarations decide. +DECLARED = {0: 16, 1: 16, 2: 64, 3: 128} +READ_WINDOWS = { + "second_stage": (_Windows([200, 300], DECLARED), [64, 128]), + "widest_of_the_group": (_Windows([300], {2: 64, 3: 96}), [96]), + "more_cache_layers_than_model_layers": ( + _Windows([200, 300], DECLARED, cache_layers=4), + [200, 300], + ), + "nothing_declared": (_Windows([200, 300], {}), [200, 300]), + "full_attention": (_Windows([None, 300], DECLARED), [None, 128]), + "a_layer_declaring_none": (_Windows([300], {2: 64, 3: None}), [300]), + "a_configured_window_below_the_model_s": (_Windows([256], {2: 4096}, pp_layers=[2]), [256]), + "widest_of_the_group_past_the_kept_window": (_Windows([256], {2: 64, 3: 4096}), [256]), + "non_positive_declarations": (_Windows([200, 300], {2: 0, 3: -1}), [200, 300]), +} + + +@pytest.mark.parametrize("case", list(READ_WINDOWS)) +def test_a_layer_group_reads_the_windows_its_model_layers_declare(case): + """A group reads the widest positive window its model layers declare, at most its kept one.""" + windows, expected = READ_WINDOWS[case] + read = [KVCacheManagerV2._attention_read_window(windows, lg) for lg in range(len(expected))] + assert read == expected diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py index ec448b6e79c7..7b73aebb3b33 100644 --- a/tests/unittest/_torch/executor/test_multimodal_scheduler.py +++ b/tests/unittest/_torch/executor/test_multimodal_scheduler.py @@ -1230,6 +1230,7 @@ def _make_v2_multimodal_scheduler( kv_allocated = {} manager = Mock(spec=KVCacheManagerV2) manager.context_replay_tokens.return_value = None + manager._holds_page_locks.return_value = False manager.tokens_per_block = 10 manager.enable_block_reuse = False manager.enable_joint_kv_cache_reuse = False diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 067c082067a6..f829d9d79256 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -37,6 +37,7 @@ SHUTDOWN_REQUEST_ID, RequestQueueItem, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache_load import KvLoadTracker from tensorrt_llm._torch.pyexecutor.llm_request import ( ExecutorRequest, @@ -3801,6 +3802,58 @@ def test_reset_prefix_cache_clears_target_and_draft_reuse_trees(): stub.draft_kv_cache_manager.reset_reuse_state.assert_called_once_with() +@pytest.mark.parametrize("joint", [False, True]) +@pytest.mark.parametrize("refusing", [None, "target", "draft"]) +def test_reset_prefix_cache_checks_all_lenders_before_clearing_any_tree(joint, refusing): + """A draft lender's refusal leaves both reuse trees and captured routes intact.""" + events = [] + + def manager(name): + cache_manager = object.__new__(KVCacheManagerV2) + + def before_reset(): + events.append((name, "check")) + if name == refusing: + raise RuntimeError("locked cache") + + cache_manager._sharing = types.SimpleNamespace( + _before_reset=before_reset, + _on_reset=lambda: events.append((name, "reset")), + ) + + def clear(): + assert ("target", "check") in events + assert not joint or ("draft", "check") in events + events.append((name, "clear")) + + cache_manager.impl = Mock(clear_reusable_blocks=Mock(side_effect=clear)) + cache_manager.conversation_manager = None + return cache_manager + + executor = object.__new__(PyExecutor) + executor.kv_cache_manager = manager("target") + executor.draft_kv_cache_manager = manager("draft") + executor.enable_joint_kv_cache_reuse = joint + executor.model_engine = types.SimpleNamespace(route_capture=Mock()) + + if refusing == "target" or (joint and refusing == "draft"): + with pytest.raises(RuntimeError, match="locked cache"): + executor.reset_prefix_cache() + executor.kv_cache_manager.impl.clear_reusable_blocks.assert_not_called() + executor.draft_kv_cache_manager.impl.clear_reusable_blocks.assert_not_called() + executor.model_engine.route_capture.clear_shared.assert_not_called() + assert all(action == "check" for _, action in events) + else: + executor.reset_prefix_cache() + executor.kv_cache_manager.impl.clear_reusable_blocks.assert_called_once_with() + if joint: + executor.draft_kv_cache_manager.impl.clear_reusable_blocks.assert_called_once_with() + else: + executor.draft_kv_cache_manager.impl.clear_reusable_blocks.assert_not_called() + assert not any(name == "draft" for name, _ in events) + executor.model_engine.route_capture.clear_shared.assert_called_once_with() + + class TestPendingTransferResponseFlush: def test_rank_local_fatal_error_does_not_issue_adp_response_gather(self): """A lone fatal rank must fail locally rather than desynchronize TP.""" diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py index 91ec477e17fb..83d2a32dc77e 100644 --- a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -105,6 +105,7 @@ def _make_executor( exe.is_shutdown = False exe.drafter = None exe.kv_cache_manager = MagicMock() + exe.kv_cache_manager._sharing = None exe.kv_cache_manager.max_beam_width = 1 exe.kv_cache_manager.impl = MagicMock() exe.kv_cache_manager.impl.need_adjustment = need_adjustment