[TRTLLM-16186][feat] Lend KV cache pages and recurrent state in place

tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py

@@ -215,7 +215,9 @@ def install_manager_capabilities(manager) -> None:
             "get_disagg_role_mapper_kinds",
             "get_disagg_ignored_roles",
             "rebuild_disagg_imported_cache",
+            "_model_layers_and_tails",
             "get_disagg_global_layer_ids",
+            "get_model_layer_ids",
             "_get_runtime_cache_size_layer_components",
             "_get_generation_request_capacity",
         )
@@ -956,7 +958,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 +977,24 @@ 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 get_model_layer_ids(self, layer_group_id: int) -> list[int]:
+        """The model layer of each internal layer of the layer group; a tail layer belongs to the
+        model layer whose high-precision tail it holds."""
+        return [model_layer for model_layer, _ in self._model_layers_and_tails(layer_group_id)]
+
     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()

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,17 +141,25 @@ 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
     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: ...

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,32 +672,55 @@ 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(
                 f"{where}: object keys must be uint8 [{rows}, width > 0], got {keys.dtype} "
                 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")

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

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.

tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py

@@ -1033,6 +1033,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: