[None][feat] Add shared KV transfer ownership and progress

tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py

@@ -28,6 +28,11 @@
 if TYPE_CHECKING:
     import torch
 
+    from tensorrt_llm._torch.disaggregation.transfer_manager import (
+        DomainFailure,
+        OperationObservation,
+        TransferManager,
+    )
     from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests
 
 
@@ -131,6 +136,61 @@ def __init__(
         self._unpublished_gen_receives: Dict[
             int, Tuple[LlmRequest, Optional["torch.cuda.Event"]]
         ] = {}
+        self._shared_transfer_manager: Optional["TransferManager"] = None
+        self._shared_updates: List["OperationObservation"] = []
+
+    def attach_shared_transfer_manager(self, manager: "TransferManager") -> None:
+        """Attach the shared holder on its owner thread before submitting work.
+
+        Args:
+            manager: Executor-local holder assembled for a supported profile.
+
+        Raises:
+            RuntimeError: A different holder is already attached.
+        """
+        if (
+            self._shared_transfer_manager is not None
+            and self._shared_transfer_manager is not manager
+        ):
+            raise RuntimeError("a shared transfer manager is already attached")
+        self._shared_transfer_manager = manager
+
+    def has_pending_shared_transfers(self) -> bool:
+        """Return whether shared work still needs owner-loop progress."""
+        manager = self._shared_transfer_manager
+        return manager is not None and manager.pending()
+
+    def take_shared_updates(self) -> Tuple["OperationObservation", ...]:
+        """Drain local observations for orchestration to reconcile across ranks.
+
+        Returns:
+            Operation facts; request associations and readiness belong to orchestration.
+        """
+        updates = tuple(self._shared_updates)
+        self._shared_updates.clear()
+        return updates
+
+    def take_shared_domain_failures(self) -> Tuple["DomainFailure", ...]:
+        """Drain lender-wide stops for orchestration to apply at its recovery boundary."""
+        manager = self._shared_transfer_manager
+        return () if manager is None else manager.take_domain_failures()
+
+    def begin_shared_shutdown(self) -> None:
+        """Close shared admission while keeping the owner loop alive for cleanup."""
+        manager = self._shared_transfer_manager
+        if manager is not None:
+            manager.begin_shutdown(time.monotonic() + manager.shutdown_timeout_s)
+
+    def shared_shutdown_complete(self) -> bool:
+        """Return whether shared resources permit manager teardown."""
+        manager = self._shared_transfer_manager
+        return manager is None or manager.shutdown_status().complete
+
+    def _advance_shared_transfers(self) -> None:
+        """Observe shared work locally without entering native collectives."""
+        manager = self._shared_transfer_manager
+        if manager is not None:
+            self._shared_updates.extend(manager.advance(time.monotonic()))
 
     # -- loop head -----------------------------------------------------------
 
@@ -145,6 +205,7 @@ def handle_errors_synced(self) -> None:
         otherwise the downstream ``tp_gather`` in ``_enqueue_responses``
         deadlocks or leaves peer replicas running.
         """
+        self._advance_shared_transfers()
         pending_ids = self.take_pending_context_failures()
         pending_requests = (
             [req for req in self._registry.active_requests() if get_unique_rid(req) in pending_ids]
@@ -851,13 +912,15 @@ def pace_idle(self) -> None:
         )
         if waiting_on_transfer:
             time.sleep(0.001)
+        elif self.has_pending_shared_transfers():
+            self._shared_transfer_manager.wait_for_progress()
 
 
 class NoopDisaggCoordinator(DisaggTransferCoordinator):
     """Coordinator used when the executor has no KV cache transceiver.
 
-    A null object: no state, no dependencies, no side effects. The loops call
-    every entry point unconditionally, so each one is a no-op here. Transceiver
+    Native entry points are no-ops. An explicitly attached shared holder still
+    progresses without a native transceiver. Transceiver
     enablement is rank-uniform within a collective group, so every rank skips
     the same collectives symmetrically.
     """
@@ -875,7 +938,7 @@ def __init__(self) -> None:
         )
 
     def handle_errors_synced(self) -> None:
-        return None
+        self._advance_shared_transfers()
 
     def admit(self, fitting_gen_init: List[LlmRequest]) -> Tuple[List[LlmRequest], bool]:
         return fitting_gen_init, False
@@ -941,4 +1004,5 @@ def forget_request(self, request_id: int) -> None:
         return None
 
     def pace_idle(self) -> None:
-        return None
+        if self.has_pending_shared_transfers():
+            self._shared_transfer_manager.wait_for_progress()

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.