[TRTLLM-17449][feat] Let KVCacheManagerV2 build its page table from layout hooks

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()

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:

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.

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:

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"

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: ...