feat(offload): support native DSv4 checkpoints with LMCache MP

atom/kv_transfer/disaggregation/multi/multi_connector.py

@@ -348,6 +348,13 @@ def close(self) -> None:
 class MultiConnectorScheduler(KVConnectorSchedulerBase):
     """Scheduler-side composite connector."""
 
+    def bind_block_manager(self, block_manager: Any) -> None:
+        """Bind native checkpoint owners before any request can be admitted."""
+        for connector in self._connectors:
+            bind = getattr(connector, "bind_block_manager", None)
+            if callable(bind):
+                bind(block_manager)
+
     def __init__(self, config: Any) -> None:
         self._connectors = _build_subconnectors(config, role="scheduler")
         self.is_producer = any(
@@ -489,6 +496,23 @@ def source_blocks_released(self, seq: Any) -> None:
             if callable(fn):
                 fn(seq)
 
+    def waits_for_transfer_report(self, seq: Any) -> bool:
+        """True only if every sub still deferring this request waits for reports.
+
+        Clock-based reclaim is skipped only when no deferring sub would act on
+        it; a P/D send or an in-process save still gets its abandon path.
+        """
+        deferring = 0
+        for connector in self._connectors:
+            should_defer = getattr(connector, "should_defer_free", None)
+            if not callable(should_defer) or not should_defer(seq):
+                continue
+            deferring += 1
+            callback = getattr(connector, "waits_for_transfer_report", None)
+            if not callable(callback) or callback(seq) is not True:
+                return False
+        return deferring > 0
+
     def abandon_save(self, req_id: Any) -> None:
         # Reclamation of a stalled offload save (see
         # `DenseOffloadConnector.abandon_save`). Only the offload sub tracks
@@ -646,6 +670,25 @@ def protected_block_ids(self, seq: Any):
             protected.update(blocks)
         return frozenset(protected)
 
+    def can_partially_deallocate_state(self, seq: Any) -> bool:
+        """True only if every sub still deferring this request guarantees it.
+
+        Fail-closed like ``protected_block_ids``: releasing the state slot is
+        safe only when no deferring sub still needs the request alive, and a
+        missing method is not approval. ``protected_block_ids`` narrows PAGE
+        blocks only, so it cannot stand in for a sub's state guarantee.
+        """
+        deferring = 0
+        for connector in self._connectors:
+            should_defer = getattr(connector, "should_defer_free", None)
+            if not callable(should_defer) or not should_defer(seq):
+                continue
+            deferring += 1
+            callback = getattr(connector, "can_partially_deallocate_state", None)
+            if not callable(callback) or callback(seq) is not True:
+                return False
+        return deferring > 0
+
     def activate_block_leases(self, seq: Any, block_ids) -> None:
         for connector in self._connectors:
             callback = getattr(connector, "activate_block_leases", None)

atom/kv_transfer/disaggregation/page_region.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Build a `PageRegion` from the tensor that owns its bytes.
+
+Kept out of `types`, which stays a torch-free contract: this is where a
+backend's allocation becomes both halves of a PAGE unit at once.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from atom.kv_transfer.disaggregation.types import KVTransferRegion, PageRegion
+
+
+def page_region(
+    tensor: torch.Tensor,
+    *,
+    semantic_role: str,
+    unit_bytes: int | None = None,
+    total_bytes: int | None = None,
+) -> PageRegion:
+    """The region of ``tensor``'s PAGE blocks and a byte view aliasing them.
+
+    The view is a zero-copy ``uint8 [num_units, 1, unit_bytes]`` alias: the
+    last two dimensions are opaque copy geometry, and LMCache stores the same
+    bytes in the same order whatever shape a block is given.
+
+    ``unit_bytes`` defaults to ``tensor``'s row stride and ``total_bytes`` to
+    all of it. An allocation that holds more than its PAGE blocks (DSV4's
+    planes also hold SLOT rows after them) passes both.
+    """
+    if not tensor.is_contiguous():
+        raise ValueError(f"PAGE tensor {semantic_role!r} must be contiguous")
+    available = tensor.numel() * tensor.element_size()
+    if unit_bytes is None:
+        unit_bytes = tensor.stride(0) * tensor.element_size()
+    if total_bytes is None:
+        total_bytes = available
+    if unit_bytes <= 0 or total_bytes % unit_bytes or total_bytes > available:
+        raise ValueError(
+            f"PAGE tensor {semantic_role!r} has {available} bytes; cannot "
+            f"publish {total_bytes} bytes in units of {unit_bytes}"
+        )
+    # `view`, never `reshape`: a copy would publish bytes nobody writes.
+    byte_view = tensor.view(torch.uint8).view(-1)[:total_bytes]
+    return PageRegion(
+        region=KVTransferRegion(
+            base_addr=tensor.data_ptr(),
+            total_bytes=total_bytes,
+            unit_bytes=unit_bytes,
+            semantic_role=semantic_role,
+        ),
+        view=byte_view.view(-1, 1, unit_bytes),
+    )
+
+
+__all__ = ["page_region"]

atom/kv_transfer/disaggregation/types.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import math
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import Any
@@ -189,11 +190,29 @@ def unit_addr(self, index: int) -> int:
         return self.base_addr + index * self.unit_bytes
 
 
+@dataclass(frozen=True)
+class PageRegion:
+    """One PAGE region and the tensor that aliases exactly its bytes.
+
+    The region is the address form every transport reads (P/D, dense offload);
+    the view is the tensor form LMCache MP registers. They describe the same
+    bytes, so they travel as one value and cannot fall out of step. A producer
+    that only has addresses leaves ``view`` None; LMCache MP then refuses the
+    layout rather than register a partial one. Build one from a tensor with
+    `atom.kv_transfer.disaggregation.page_region.page_region`.
+    """
+
+    region: KVTransferRegion
+    view: Any = None
+
+
 @dataclass
 class KVTransferTensors:
     """Physical PAGE, SLOT, and compressor-only staging region contract.
 
-    ``block_regions`` contain forward-indexed, block-indexed PAGE units.
+    ``pages`` are forward-indexed, block-indexed PAGE units, each with its
+    region and (when published) its tensor view; ``block_regions`` and
+    ``block_tensor_views`` read them back.
     ``swa_block_regions`` is a legacy field name for complete reverse-indexed
     per-request SLOT units, including both compressor state and SWA.
     ``expected_full_slot_region_count`` makes a stateful layout's complete
@@ -213,9 +232,9 @@ class KVTransferTensors:
     protocol can describe replication per region group.
     """
 
-    # Block-indexed PAGE regions, indexed forward by block id.
-    block_regions: list[KVTransferRegion]
-    slot_regions: list[KVTransferRegion]
+    # Block-indexed PAGE regions with their views, indexed forward by block id.
+    pages: list[PageRegion] = field(default_factory=list)
+    slot_regions: list[KVTransferRegion] = field(default_factory=list)
     num_slots: int = 0
     # Optional producer-local -> consumer-global mapping for non-uniform block
     # region layouts. Uniform per-layer groups leave this unset and use the
@@ -231,12 +250,6 @@ class KVTransferTensors:
     scatter_slot: Callable[[int, int], None] | None = None
     # Appended for positional compatibility with existing generic descriptors.
     expected_full_slot_region_count: int | None = None
-    # Zero-copy, physical-order views paired with block regions.
-    # Each view is [num_units, physical_slots_per_unit, opaque_width]; it may
-    # retain its native dtype, and one dim-0 unit must cover exactly the paired
-    # region's unit_bytes. The last two dimensions describe copy geometry, not
-    # a semantic token/head layout.
-    block_tensor_views: list[Any] = field(default_factory=list)
     # Number of TP workers whose complete PAGE layout is byte-identical.
     # ``1`` means no cross-rank deduplication is safe.
     tp_replication_factor: int = 1
@@ -257,17 +270,62 @@ class KVTransferTensors:
     gather_sharded_index: Callable[..., tuple[int, int]] | None = None
     # Appended after the original staging fields for positional compatibility.
     prepare_sharded_index: Callable[..., Any] | None = None
+    # Native PAGE-backed state image contract. The spec names the exact codec
+    # and image geometry; the callback restores that image into an Active SLOT
+    # after a remote retrieve has finished writing its PAGE units. These are
+    # intentionally untyped model-engine objects to keep this module free of
+    # model-engine imports. The PAGE tensor views retain the source allocations;
+    # neither field asks the connector to snapshot a complete live SLOT.
+    paged_state_checkpoint_spec: object | None = None
+    execute_paged_state_copies: Callable[..., None] | None = None
+    # Same declaration for PAGE-backed native state checkpoint images. Kept
+    # separate so future layouts can describe PAGE and recurrent-state
+    # replication independently. Appended for positional compatibility.
+    native_state_tp_replication_factor: int = 1
+    # How many leading `block_regions` make up one native checkpoint PAGE unit.
+    # A draft with a pool of its own appends its regions after the target's;
+    # those are ordinary PAGE KV, never part of the state image. None means
+    # every region. Appended for positional compatibility.
+    paged_state_region_count: int | None = None
     # Scheduler blocks the PAGE regions are addressed in. `init=False` because
     # a backend cannot answer it: `req.block_ids` is the scheduler's id space,
     # and a backend counts in its own page -- a different unit even where it is
     # the same number. Set through `set_block_count`.
     num_blocks: int = field(init=False, default=0)
 
+    @property
+    def block_regions(self) -> tuple[KVTransferRegion, ...]:
+        """Every PAGE region, in publication order (read-only; add `pages`)."""
+        return tuple(page.region for page in self.pages)
+
+    @property
+    def block_tensor_views(self) -> tuple[Any, ...]:
+        """The published views, each a zero-copy byte alias of its region.
+
+        ``uint8 [num_units, 1, unit_bytes]``: the last two dimensions are
+        opaque copy geometry for LMCache MP, not a token/head layout. Empty when
+        the producer published addresses only.
+        """
+        return tuple(page.view for page in self.pages if page.view is not None)
+
+    def merge_pages(self, other: KVTransferTensors) -> None:
+        """Append another contributor's PAGE units (a draft's own pool).
+
+        Replication is a whole-PAGE declaration: one sharded contributor makes
+        the combined object sharded even when this one alone is replicated.
+        """
+        if not other.pages:
+            return
+        self.pages.extend(other.pages)
+        self.tp_replication_factor = math.gcd(
+            self.tp_replication_factor, other.tp_replication_factor
+        )
+
     def set_block_count(self, num_blocks: int) -> None:
         """Fix the block id space, and check every region is in it.
 
-        Called once the last contributor's regions are in the list; a draft
-        appends its own after construction, so this cannot be a `__post_init__`.
+        Called once the last contributor's regions are in; a draft's are merged
+        after construction (`merge_pages`), so this cannot be a `__post_init__`.
 
         A region that does not divide into exactly `num_blocks` units was
         registered in some other unit. Both ends would still agree on