From 409f727a842b01500443ce0abf00a8431535044d Mon Sep 17 00:00:00 2001 From: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com> Date: Sat, 17 Oct 2026 02:28:49 -0700 Subject: [PATCH 1/3] [TRTLLM-16186][feat] Define backend contract for KV cache transfers Define, in tensorrt_llm/_torch/disaggregation/base, the contract between a lease holder, which holds leases from the KV cache lender, and a backend that moves the lent rows to or from a peer or a store. base exports Backend, Capabilities, Attempt, Outcome, View, LayerGroupRows, SlotRegion, MemoryKind, Hint, WakeFn, InvalidSubmission and SubmissionRejected, plus TokenRange and CacheKind, which are not part of the contract. Backend replaces Fetches and Publishes. - MemoryKind, in base/memory.py, is the kind of memory a region lies in: HOST is CPU-addressable and DEVICE is the GPU holding the cache. A backend registers only the regions whose SlotRegion.memory_kind it can access, and rows in a region it has not registered are rows it does not recognize. A row given by an address occupies one slot. - A view's groups and their arrays are an immutable snapshot that stays valid after the lease is returned; the lease bounds only the payload's protection. A hint describes only the other end and is an immutable snapshot; which rows this end reads or writes comes from the view alone. - fetch and publish state what the lease holder guarantees about the memory the view locates until the submission's quiesced() is true, and no two rows of a fetch locate overlapping memory. A raise other than the two rejections may have let access to the memory escape, and the lease holder then keeps the lease until shutdown() returns an empty set. quiesced() covers the whole submission, so the leases lent for a backend that may leave a row not proven ended are kept small. - Conclusion and quiescence do not depend on calls to poll() or quiesced(): a backend advances its own work. A backend calls WakeFn only after every change it reports is visible: an outcome or an access end through poll() or quiesced(), a change an optional protocol names where that protocol says. - An instance serves the submissions of one lender's domain, and a process may hold several. A backend's methods and its attempts' belong to one call sequence. - A backend given rows it does not recognize rejects the submission with InvalidSubmission and never skips them. Delivery masks are bool arrays; anything else raises TypeError. Native's per-request types, Chunk, CacheExtent, Delivered, Failed, Cancelled and its own Attempt and Outcome, move unchanged from base/backend.py to native/extent.py; native and transceiver.py change only their imports, and base/transfer.py imports Chunk from native/extent.py for type checking only. The contract tests import nothing of the lender, and they pin that base imports native only for type checking. Signed-off-by: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com> --- .../_torch/disaggregation/base/__init__.py | 42 +- .../_torch/disaggregation/base/backend.py | 491 ++++++++++------ .../_torch/disaggregation/base/memory.py | 28 + .../_torch/disaggregation/base/transfer.py | 5 +- .../_torch/disaggregation/native/extent.py | 205 +++++++ .../_torch/disaggregation/native/fetch.py | 2 +- .../_torch/disaggregation/native/handle.py | 3 +- .../_torch/disaggregation/native/publish.py | 2 +- .../_torch/disaggregation/native/transfer.py | 2 +- .../_torch/disaggregation/transceiver.py | 3 +- .../disaggregation/test_backend_contract.py | 528 ++++++++++++++++-- .../disaggregated/test_chunked_transfer.py | 3 +- .../disaggregated/test_kv_transfer.py | 3 +- .../test_kv_transfer_fail_stop.py | 3 +- .../disaggregated/test_kv_transfer_mp.py | 3 +- .../unittest/disaggregated/test_peer_fetch.py | 7 +- .../disaggregated/test_peer_publish.py | 15 +- .../disaggregated/test_retirement_deadline.py | 3 +- .../disaggregated/test_task_handle.py | 7 +- .../test_transfer_late_settlement.py | 3 +- .../test_transfer_lifecycle_acceptance.py | 3 +- .../test_transfer_ownership_regressions.py | 3 +- .../disaggregated/test_unsubmitted_session.py | 3 +- .../test_verified_reuse_admission.py | 3 +- 24 files changed, 1117 insertions(+), 253 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/base/memory.py create mode 100644 tensorrt_llm/_torch/disaggregation/native/extent.py diff --git a/tensorrt_llm/_torch/disaggregation/base/__init__.py b/tensorrt_llm/_torch/disaggregation/base/__init__.py index 1d3e9dd767d3..c55924e3ab07 100644 --- a/tensorrt_llm/_torch/disaggregation/base/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/base/__init__.py @@ -12,36 +12,44 @@ # 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 cache-transfer contract. +"""The KV cache transfer backend contract. -What this package exports is the contract; what it does not is backend-internal and has to be -imported by file path. +The package exports the contract, plus ``TokenRange`` and ``CacheKind``, which belong to native and +the page table. All of them come from ``backend.py`` except ``MemoryKind``, which ``memory.py`` +defines. The package's other modules are not part of the contract and are imported by their own +module paths, such as ``base.transfer``. """ from .backend import ( Attempt, - CacheExtent, + Backend, CacheKind, - Cancelled, - Chunk, - Delivered, - Failed, - Fetches, + Capabilities, + Hint, + InvalidSubmission, + LayerGroupRows, Outcome, - Publishes, + SlotRegion, + SubmissionRejected, TokenRange, + View, + WakeFn, ) +from .memory import MemoryKind __all__ = [ "Attempt", - "CacheExtent", + "Backend", "CacheKind", - "Cancelled", - "Chunk", - "Delivered", - "Failed", - "Fetches", + "Capabilities", + "Hint", + "InvalidSubmission", + "LayerGroupRows", + "MemoryKind", "Outcome", - "Publishes", + "SlotRegion", + "SubmissionRejected", "TokenRange", + "View", + "WakeFn", ] diff --git a/tensorrt_llm/_torch/disaggregation/base/backend.py b/tensorrt_llm/_torch/disaggregation/base/backend.py index 3e5c500826a1..06824f5a2ce9 100644 --- a/tensorrt_llm/_torch/disaggregation/base/backend.py +++ b/tensorrt_llm/_torch/disaggregation/base/backend.py @@ -12,35 +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. -"""What a cache backend is asked to do, and what it still owes the caller word about.""" +"""The backend contract: what a lease holder hands a backend, and what it gets back. + +The lease holder holds a lease from the KV cache lender and submits its view to a backend. The +caller asks the lease holder for a transfer and supplies its hint. The assembly builds each +backend and hands it the lender's regions, in the lender's order, and its ``WakeFn``. The +obligations stated on each type are part of the contract, checked or not. Methods beyond +:class:`Backend` come as optional protocols, which a backend declares by implementing them. +""" from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass from enum import IntEnum -from typing import List, Optional, Protocol, Set, Tuple, Union, runtime_checkable +from typing import Any, Callable, Protocol, Sequence, runtime_checkable import numpy as np -# TODO: -# 1. Something outside the engine cannot implement this until the file leaves this package. -# 2. Wrapping request ids and endpoints belongs to a later repo-wide typing pass. -# 3. Retrying from a second source needs an (id, epoch) key here and in ``native`` together. +from .memory import MemoryKind __all__ = [ "Attempt", - "CacheExtent", + "Backend", "CacheKind", - "Cancelled", - "Chunk", - "Delivered", - "Failed", - "Fetches", + "Capabilities", + "Hint", + "InvalidSubmission", + "LayerGroupRows", + "MemoryKind", "Outcome", - "Publishes", + "SlotRegion", + "SubmissionRejected", "TokenRange", + "View", + "WakeFn", ] +# Native's token range and layer-group kind, read by its per-request pieces and its page table; +# not part of the backend contract. + @dataclass class TokenRange: @@ -57,225 +67,378 @@ def __post_init__(self): class CacheKind(IntEnum): - """How one layer group's region ids are read, and with them its share of a token range. + """How one layer group's pool slot ids are read, and with them its share of a token range. - PAGED: several block ids covering a span, whose end is an upper bound -- any prefix of it is - valid cache (attention). - STATE: one slot id standing for the request as a whole, whose end is an exact checkpoint -- a - slot rolled to another position is not a prefix of this one (recurrent state). + PAGED: a positional pool-slot table covering a span, whose end is an upper bound -- any prefix + of it is valid cache (attention). + STATE: one pool slot id standing for the request as a whole, whose end is an exact checkpoint + -- a slot rolled to another position is not a prefix of this one (recurrent state). - A backend that cannot tell the two apart reads a slot id as a block id and a checkpoint as a - bound, and neither mistake announces itself. + A backend that cannot tell the two apart reads a state's one slot as a one-entry table and its + checkpoint as a bound, and neither mistake announces itself. """ PAGED = 0 STATE = 1 -@dataclass -class Chunk: - """One piece of a transfer. A monolithic transfer is the sole piece. - - ``block_ids_per_layer_groups[g]`` of a PAGED group is positional: entry ``i`` is the local pool - slot holding block ordinal ``i`` (tokens ``[i * tpb, (i + 1) * tpb)``), or ``-1`` when this side - has nothing to offer/accept there -- SWA-evicted, not yet allocated, already cached on the - receiver, or outside this piece's range. Paged groups carry ``ceil(prompt_len / tpb)`` entries - and beam 0 only; a STATE group carries a single slot. The sender pairs ordinals that are - ``>= 0`` on both sides, so neither side needs to know the other's eviction or reuse state. - - ``token_range`` is the span this piece moves, not a claim that every layer group covers it: a - recurrent-state group carries one slot rather than a span, and a windowed group's table is - mostly holes. - """ +Hint = Any +"""The caller's hint for the other end of a submission (its source or destination), handed to +the backend unread by the lease holder. A hint is an immutable snapshot that the caller never +changes, so the backend may keep a reference to it, for as long as it needs, without copying it. +It describes only the other end: which rows this end reads or writes comes from the view alone.""" - block_ids_per_layer_groups: List[np.ndarray] +WakeFn = Callable[[], None] +"""Handed to a backend by the assembly and called from any thread at least when an outcome lands or +an access ends, and on any further occasion an optional protocol names: thread-safe, reentrant, +never blocks or raises, harmless after shutdown. A backend calls it only after every change it +reports is visible: an outcome or an access end through :meth:`Attempt.poll` or +:meth:`Attempt.quiesced`, a change an optional protocol names where that protocol says.""" - kind_per_layer_group: List[CacheKind] - """How to read each list above, in the same order. - Stated rather than inferred: both kinds are arrays of integers, so nothing in the ids themselves - separates a page from a slot. +@runtime_checkable +class SlotRegion(Protocol): + """A contiguous memory region divided into equal-sized slots; a row given by an address + occupies one slot. + + 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. + Slot ``k`` starts at ``address + k * slot_bytes`` for ``0 <= k < slot_count``; ``nbytes`` + covers all slots, so it is at least ``slot_bytes * slot_count``. A view may select + nonconsecutive slots. + The descriptor and its backing memory remain valid while the backend may access them. """ - token_range: TokenRange # (0, prompt_len) when not cut along tokens - is_last: bool - """Whether this piece ends the series. Getting it wrong ends the transfer early and in silence. + @property + def layout_label(self) -> str: + """An opaque layout label, equal for regions laid out alike; not an allocation identity.""" + ... - Cannot be derived from ``token_range``: a windowed or recurrent-state group never fills the - span, so comparing against the full length lies.""" + @property + def address(self) -> int: ... - excluded_pool_views: Set[Tuple[int, int]] = field(default_factory=set) + @property + def nbytes(self) -> int: ... - def __post_init__(self): - if len(self.kind_per_layer_group) != len(self.block_ids_per_layer_groups): - raise ValueError( - f"{len(self.block_ids_per_layer_groups)} layer groups but " - f"{len(self.kind_per_layer_group)} kinds" - ) + @property + def slot_bytes(self) -> int: + """Bytes of one slot.""" + ... + @property + def slot_count(self) -> int: + """Number of slots, back to back from :attr:`address`.""" + ... -@dataclass -class CacheExtent: - """One ask: which content, and where this rank's blocks for it are. - - One ``name`` covers the request, while ``local``'s blocks are per layer group, positional with - ``-1`` holes (see ``Chunk``). - - Handing one to ``fetch`` or ``publish`` lends it out: a backend may keep it and read it on - another thread. Neither it, its ``Chunk``, nor the arrays inside may be changed while it is - lent -- the transfer would quietly move a different shape, and since a backend is not required - to copy it, nothing would report the change. - - TODO: An ``Attempt`` concluding does not end the loan. It reports a logical outcome, and one - can be reached while writers are still going -- a failure settles the piece while its siblings - write on -- so a backend may still be reading afterwards. Nothing on this surface says when a - caller may write to the extent again; native answers it by holding the only reference until - the session retires. How long this description stays readable and how long the pages it names - stay held are two questions, and neither is answered here. - """ + @property + def memory_kind(self) -> MemoryKind: + """The kind of memory the region lies in. ``DEVICE`` is the GPU of this rank's lender; the + region names no device id.""" + ... - name: int - local: Chunk - """This rank's blocks: the destination when receiving, the source when sending. Same shape, - opposite direction.""" +@runtime_checkable +class LayerGroupRows(Protocol): + """One layer group's rows of a view, described by aligned arrays. - # TODO: Which layer groups a delivery leaves whole is the answering side's to say, not the - # asking side's, and nothing carries it today. + 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. + """ + @property + def layer_group_index(self) -> int: + """The lender's local layer-group index, not the group's position in :attr:`View.groups`.""" + ... -@dataclass(frozen=True) -class Delivered: - """The ask was met. + @property + def object_keys(self) -> np.ndarray | None: + """One opaque object key per row, ``uint8`` of shape ``[rows, width]`` with the lender's + fixed width. Keys are compared only for equality under the lender's identity rules; + they are not payload checksums. ``None`` for rows without keys, which reach only a backend + that declares a capability admitting them.""" + ... - TODO: The evidence is reports, not a fence; the reports travel on a different channel from the - data and the two are not ordered. - """ + @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.""" + ... - token_end: int - """How far this delivery reaches, as an absolute offset rather than a count. + @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 __len__(self) -> int: ... - One value for every layer group, read by each group's kind: a bound for ``PAGED``, a checkpoint - for ``STATE``. - """ - reports_pending: bool = False - """Always false: every report this transfer was owed has arrived. Setting it raises; success - still owed a report is what ``Failed`` and ``Cancelled`` are for. +@runtime_checkable +class View(Protocol): + """The rows one submission moves, all from one lease: at most one :class:`LayerGroupRows` + per layer group. + + Its groups and their members, the rows' metadata, are an immutable snapshot, never changed or + reused: they stay valid for as long as the lease holder keeps the view, after the lease is + returned too. The lease bounds only the protection of the payload, the bytes in the memory + the rows locate. """ - def __post_init__(self): - if self.reports_pending: - raise ValueError("a delivery cannot still be owed a report") + @property + def groups(self) -> Sequence[LayerGroupRows]: ... + def row_masks(self, fill_value: bool = False) -> tuple[np.ndarray, ...]: + """One writable ``bool`` mask of shape ``[len(group)]`` per group, in :attr:`groups` order. -@dataclass(frozen=True) -class Failed: - """The transfer could not be completed.""" + Args: + fill_value: What every entry starts as. + """ + ... - reason: str - reports_pending: bool - """Whether a report this transfer was owed has yet to arrive. Says nothing about hardware.""" +class SubmissionRejected(Exception): + """Refused before work, a peer notification or memory access escaped; also every submission + after :meth:`Backend.shutdown`. Only this type certifies that nothing escaped.""" -@dataclass(frozen=True) -class Cancelled: - """Someone asked for this to stop, and it ended without delivering. - Says nothing about whether the peers have stopped -- nothing here does. A cancelled piece is - normally still owed reports, which is a different question again. +class InvalidSubmission(SubmissionRejected, ValueError): + """Bad arguments or hint, or an undeclared direction; raised before anything escaped.""" + - Cancellation is an act on the whole transfer, not on one piece: there is no way to stop a single - chunk while its siblings continue. +@dataclass(frozen=True) +class Capabilities: + """What a backend implements, fixed for the backend's life. A new capability is a field whose + default means not supported, and rows that need a capability reach only a backend that + declares it. + + Attributes: + fetch: Implements :meth:`Backend.fetch`. + publish: Implements :meth:`Backend.publish`. + cancel: Implements the advisory :meth:`Attempt.cancel`. + + Raises: + ValueError: Neither ``fetch`` nor ``publish`` is set. """ - by_peer: bool = False - """A local cancel is an ordinary end; a peer's is a transfer error, and only the backend knows - which of the two happened.""" + fetch: bool + publish: bool + cancel: bool = False - reports_pending: bool = True + def __post_init__(self): + if not (self.fetch or self.publish): + raise ValueError("a backend implements at least one of fetch and publish") -Outcome = Union[Delivered, Failed, Cancelled] -"""How one delivery ended. ``None`` rather than a member means it has not ended yet. +def _readonly_bool_copy(mask: np.ndarray) -> np.ndarray: + """A read-only copy of ``mask``, which must be a ``bool`` array. -The logical member and its cause are committed at the task/session transition; -polling only observes that decision. ``reports_pending`` can still change, so a -stored outcome carries stale report progress, not a different logical result. + Raises: + TypeError: ``mask`` is not of dtype ``bool``. + """ + copy = np.array(mask) + if copy.dtype != np.bool_: + raise TypeError(f"a delivery mask is bool, not {copy.dtype}") + copy.flags.writeable = False + return copy + + +@dataclass(frozen=True, eq=False) +class Outcome: + """How one submission ended. + + A delivered row is one the destination holds whole and readable: byte for byte as the source + holds it, or as the destination already held it under the row's object key. A + delivered row is final: every write to it has completed and none will follow. A row not + delivered may have been partly written; its memory is reused only once ``quiesced()`` is true. + A miss is nothing delivered and no fault; deliveries and a fault can coexist; a peer's cancel is + a fault, a local cancel is not. An outcome says neither when the request may use the rows, + which the lender decides, nor whether access has ended (``quiesced()``). + + Attributes: + delivered_masks: One read-only ``bool`` mask of shape ``[len(group)]`` per group of the + view, in order, copied from the masks given. + fault: What went wrong, or ``None``. + + Raises: + TypeError: A mask given is not a ``bool`` array. + """ -``reports_pending`` asks one question and only one: is a report this transfer was owed still to -arrive. Receiving waits on the writers' reports, sending on word about its own writes. + delivered_masks: tuple[np.ndarray, ...] + fault: str | None = None -**It is not a release condition, and answering it is not permission to hand memory back.** Reports -and data travel on different channels and are not ordered, so no value of it means the peers have -stopped touching the destination. Whatever eventually answers that is a separate question this -contract does not ask. + def __post_init__(self): + object.__setattr__( + self, + "delivered_masks", + tuple(_readonly_bool_copy(mask) for mask in self.delivered_masks), + ) + + @classmethod + def for_view( + cls, + view: View, + delivered_masks: Sequence[np.ndarray] | None = None, + fault: str | None = None, + ) -> Outcome: + """Returns an outcome for ``view``. + + Args: + view: The view submitted; gives the all-false masks when ``delivered_masks`` is + ``None``. + delivered_masks: One mask per group of ``view``, not checked against it (see + :meth:`check_mask_shapes`), or ``None`` for nothing delivered. + fault: What went wrong, or ``None``. + """ + return cls( + tuple(delivered_masks) if delivered_masks is not None else view.row_masks(False), fault + ) -TODO: A backend that has torn down whatever a report would arrive on answers false -- which is how -this converges today, rather than by the reports arriving. An ``Attempt`` cannot yet settle on its -own evidence. + def check_mask_shapes(self, view: View) -> None: + """Checks the mask count and shapes against ``view``; does not prove delivery. -TODO: Content addressing needs a member for "the source does not hold this", which is worth asking -somewhere else for while a failure is not. -""" + Args: + view: The view the outcome is for. + + Raises: + ValueError: The outcome does not hold one mask of shape ``[len(group)]`` per group of + ``view``. + """ + if len(self.delivered_masks) != len(view.groups): + raise ValueError( + f"outcome has {len(self.delivered_masks)} masks for {len(view.groups)} groups" + ) + for mask, group in zip(self.delivered_masks, view.groups): + if mask.shape != (len(group),): + raise ValueError( + f"layer group {group.layer_group_index}: outcome mask has shape {mask.shape} " + f"for {len(group)} rows" + ) @runtime_checkable class Attempt(Protocol): - """One delivery of one extent.""" + """The handle of one submission. + + Conclusion (:meth:`poll`) and quiescence (:meth:`quiesced`) are independent and come in either + order. Neither depends on these calls: a backend advances its own work, so a lease holder that + calls them only after a wake misses nothing. No method blocks or raises; trouble is a + ``fault``. Every attempt concludes, at the latest when :meth:`Backend.shutdown` forces it; the + lease holder may act on its own deadline before then. The backend alone settles each physical + operation, never an adapter or the lease holder, and :meth:`Backend.shutdown` never stops what + produces the evidence that an operation ended. An attempt's methods are called in its + backend's call sequence (see :class:`Backend`). + """ - def poll(self) -> Optional[Outcome]: - """Non-blocking. ``None`` while the delivery has not reached a conclusion. + def poll(self) -> Outcome | None: + """``None`` until concluded; the same outcome on every call after.""" + ... - A failure is reported as soon as it is known, which may be before every report is in -- - that second question is carried by the outcome itself. - The logical result is stable and does not depend on when the first poll occurs. + def quiesced(self) -> bool: + """True once access to the lent memory has ended, with evidence; its lease may be returned + on it alone. + + A pure read that never advances a transfer. It may turn true late, only on evidence, and + never turns false again. It covers the whole submission: a backend may leave a row not + proven ended, which holds every row's memory, so the leases lent for such a backend are + kept small. """ ... + def cancel(self) -> None: + """Advisory: the backend tries not to start what it has not started, and a started transfer + may still finish. Called only when ``capabilities.cancel`` is true. Never raises, and never + turns :meth:`quiesced` true without evidence; the outcome reads as if not cancelled. Calling + it again, or after the attempt concluded, does nothing.""" + ... + @runtime_checkable -class Fetches(Protocol): - """Somewhere cache can be pulled from. An instance is bound to one request, though it may be a - throwaway wrapper around longer-lived state; the request is not on the wire of this protocol, so - a backend serving an aggregated engine can implement it too.""" +class Backend(Protocol): + """Moves a view's rows between lent memory and a peer or a store. + + An instance serves every submission of every request in its domain, that of one lender, and + itself registers the memory it reads or writes. A process may hold several instances, of one + class or of several, each registering its memory and shutting down on its own; a backend that + allows only one instance has its factory refuse to construct a second. + + The methods of a backend and those of its attempts (:meth:`Attempt.poll`, + :meth:`Attempt.quiesced`, :meth:`Attempt.cancel`) belong to one call sequence: one thread at a + time calls them, the builder, then the executor loop, then the shutdown thread. + :meth:`Attempt.poll` and :meth:`Attempt.quiesced` are also safe against anything the backend + runs concurrently. The lease holder never calls :meth:`shutdown` on one thread while it goes + on calling an attempt's methods on another. + """ + + @property + def capabilities(self) -> Capabilities: ... + + def fetch(self, view: View, source: Hint) -> Attempt: + """Writes the view's rows from ``source``. Returns without waiting on a peer or store. + + The lease holder guarantees that the memory the view's rows locate is writable when it + calls, that no two rows locate overlapping memory, and that nothing but this submission + writes it until the submission's :meth:`Attempt.quiesced` is true. - def fetch(self, extent: CacheExtent, *, src: Optional[str] = None) -> Attempt: - """Start pulling ``extent``, optionally naming which peer to pull it from. + :class:`InvalidSubmission` and :class:`SubmissionRejected` have no side effect; any other + raise, a plain ``ValueError`` included, may have let access to the memory escape, and the + lease holder then keeps the lease until :meth:`shutdown` returns an empty set. Once the + lent memory may be accessed, the call returns the handle and reports any later failure + through it. - An implementation must either honour ``src`` or refuse it. Quietly reading from somewhere - else -- the request's own endpoint, say -- is a disagreement nothing reports. + Args: + view: The rows to write. + source: The caller's hint for where the rows come from. - A submission that fails part way raises rather than answering with a handle, so the caller - stops the request the way it always has. + Returns: + The submission's handle. - TODO: Peers are told one at a time, so a partial failure leaves some already holding the - destination, and raising says nothing about them. Answering with a handle needs a caller - that polls one. + Raises: + InvalidSubmission: Bad arguments or hint, or a direction the backend does not declare. + SubmissionRejected: The backend refuses the submission. """ ... + def publish(self, view: View, destination: Hint) -> Attempt: + """Offers the view's rows to ``destination``. Returns without waiting on a peer or store. -@runtime_checkable -class Publishes(Protocol): - """Somewhere cache can be offered from: the side that generated it, bound to one request. + The lease holder guarantees that the payload the view's rows locate is completely written + and readable when it calls, and stays unchanged until the submission's + :meth:`Attempt.quiesced` is true. - The other direction of ``Fetches``, and a separate protocol because a backend may serve one and - not the other -- a store that only answers reads implements ``Fetches`` alone. - """ + :class:`InvalidSubmission` and :class:`SubmissionRejected` have no side effect; any other + raise, a plain ``ValueError`` included, may have let access to the memory escape, and the + lease holder then keeps the lease until :meth:`shutdown` returns an empty set. Once the + lent memory may be accessed, the call returns the handle and reports any later failure + through it. + + Args: + view: The rows to offer. + destination: The caller's hint for where the rows go. + + Returns: + The submission's handle. + + Raises: + InvalidSubmission: Bad arguments or hint, or a direction the backend does not declare. + SubmissionRejected: The backend refuses the submission. + """ + ... - def publish(self, extent: CacheExtent) -> Attempt: - """Offer ``extent`` to whoever this request's cache is owed to. + def shutdown(self, timeout: float) -> frozenset: + """Stops serving peers and forces an outcome onto every unconcluded attempt. - No counterpart to ``Fetches``'s ``src``: a fetch chooses among sources, while publishing - answers the peers that already asked. Naming a destination is a routing question that - arrives with content addressing, not before. + A forced outcome's fault names the shutdown. Repeatable: each call does what is left. + Registrations are dropped only for memory nothing can access any more. - A submission that fails part way raises rather than answering with a handle, the same way - the other direction does. + Args: + timeout: Seconds the call may take. - Whether the extent is visible to a reader only once all of it has landed is the - implementation's to guarantee, not this protocol's to express: a store that reveals a - half-written extent hands out cache that was never generated. + Returns: + Labels of the backend's choosing for the memory that may still be accessed after this + call: that of unquiesced submissions and that still registered. Empty exactly when + nothing may still be accessed. """ ... diff --git a/tensorrt_llm/_torch/disaggregation/base/memory.py b/tensorrt_llm/_torch/disaggregation/base/memory.py new file mode 100644 index 000000000000..34cd71631661 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/base/memory.py @@ -0,0 +1,28 @@ +# 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 kinds of memory a cache backend can access.""" + +import enum + +__all__ = ["MemoryKind"] + + +class MemoryKind(enum.Enum): + """Memory a backend can access: HOST is CPU-addressable; DEVICE is the GPU holding the cache. + + Each region of memory is of one kind.""" + + HOST = "host" + DEVICE = "device" diff --git a/tensorrt_llm/_torch/disaggregation/base/transfer.py b/tensorrt_llm/_torch/disaggregation/base/transfer.py index 382d2d01d471..cb9a47059370 100644 --- a/tensorrt_llm/_torch/disaggregation/base/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/base/transfer.py @@ -17,12 +17,13 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum -from typing import Optional, cast +from typing import TYPE_CHECKING, Optional, cast from tensorrt_llm import DisaggregatedParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from .backend import Chunk +if TYPE_CHECKING: + from ..native.extent import Chunk class SessionStatus(Enum): diff --git a/tensorrt_llm/_torch/disaggregation/native/extent.py b/tensorrt_llm/_torch/disaggregation/native/extent.py new file mode 100644 index 000000000000..da1f5b7c201c --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/native/extent.py @@ -0,0 +1,205 @@ +# 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. +"""What native's peer-to-peer path is asked to move for one request, and how one piece ended. + +``Outcome`` and ``Attempt`` share their names with the backend contract's in ``base``; the module +tells them apart. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List, Optional, Protocol, Set, Tuple, Union, runtime_checkable + +import numpy as np + +from ..base.backend import CacheKind, TokenRange + +# TODO: +# 1. Native moves cache through these per-request types instead of the backend contract. +# 2. Request ids and endpoints are bare values rather than wrapped types. +# 3. Retrying from a second source needs an (id, epoch) key here and in the sessions together. + +__all__ = [ + "Attempt", + "CacheExtent", + "Cancelled", + "Chunk", + "Delivered", + "Failed", + "Outcome", +] + + +@dataclass +class Chunk: + """One piece of a transfer. A monolithic transfer is the sole piece. + + ``block_ids_per_layer_groups[g]`` of a PAGED group is positional: entry ``i`` is the local pool + slot holding block ordinal ``i`` (tokens ``[i * tpb, (i + 1) * tpb)``), or ``-1`` when this side + has nothing to offer/accept there -- SWA-evicted, not yet allocated, already cached on the + receiver, or outside this piece's range. Paged groups carry ``ceil(prompt_len / tpb)`` entries + and beam 0 only; a STATE group carries a single slot. The sender pairs ordinals that are + ``>= 0`` on both sides, so neither side needs to know the other's eviction or reuse state. + + ``token_range`` is the span this piece moves, not a claim that every layer group covers it: a + recurrent-state group carries one slot rather than a span, and a windowed group's table is + mostly holes. + """ + + block_ids_per_layer_groups: List[np.ndarray] + + kind_per_layer_group: List[CacheKind] + """How to read each list above, in the same order. + + Stated rather than inferred: both kinds are arrays of integers, so nothing in the ids themselves + separates a page from a slot. + """ + + token_range: TokenRange # (0, prompt_len) when not cut along tokens + is_last: bool + """Whether this piece ends the series. Getting it wrong ends the transfer early and in silence. + + Cannot be derived from ``token_range``: a windowed or recurrent-state group never fills the + span, so comparing against the full length lies.""" + + excluded_pool_views: Set[Tuple[int, int]] = field(default_factory=set) + + def __post_init__(self): + if len(self.kind_per_layer_group) != len(self.block_ids_per_layer_groups): + raise ValueError( + f"{len(self.block_ids_per_layer_groups)} layer groups but " + f"{len(self.kind_per_layer_group)} kinds" + ) + + +@dataclass +class CacheExtent: + """One ask: which content, and where this rank's blocks for it are. + + One ``name`` covers the request, while ``local``'s blocks are per layer group, positional with + ``-1`` holes (see ``Chunk``). + + Handing one to ``fetch`` or ``publish`` lends it out: a backend may keep it and read it on + another thread. Neither it, its ``Chunk``, nor the arrays inside may be changed while it is + lent -- the transfer would quietly move a different shape, and since a backend is not required + to copy it, nothing would report the change. + + TODO: An ``Attempt`` concluding does not end the loan. It reports a logical outcome, and one + can be reached while writers are still going -- a failure settles the piece while its siblings + write on -- so a backend may still be reading afterwards. Nothing on this surface says when a + caller may write to the extent again; native answers it by holding the only reference until + the session retires. How long this description stays readable and how long the pages it names + stay held are two questions, and neither is answered here. + """ + + name: int + + local: Chunk + """This rank's blocks: the destination when receiving, the source when sending. Same shape, + opposite direction.""" + + # TODO: Which layer groups a delivery leaves whole is the answering side's to say, not the + # asking side's, and nothing carries it today. + + +@dataclass(frozen=True) +class Delivered: + """The ask was met. + + TODO: The evidence is reports, not a fence; the reports travel on a different channel from the + data and the two are not ordered. + """ + + token_end: int + """How far this delivery reaches, as an absolute offset rather than a count. + + One value for every layer group, read by each group's kind: a bound for ``PAGED``, a checkpoint + for ``STATE``. + """ + + reports_pending: bool = False + """Always false: every report this transfer was owed has arrived. Setting it raises; success + still owed a report is what ``Failed`` and ``Cancelled`` are for. + """ + + def __post_init__(self): + if self.reports_pending: + raise ValueError("a delivery cannot still be owed a report") + + +@dataclass(frozen=True) +class Failed: + """The transfer could not be completed.""" + + reason: str + reports_pending: bool + """Whether a report this transfer was owed has yet to arrive. Says nothing about hardware.""" + + +@dataclass(frozen=True) +class Cancelled: + """Someone asked for this to stop, and it ended without delivering. + + Says nothing about whether the peers have stopped -- nothing here does. A cancelled piece is + normally still owed reports, which is a different question again. + + Cancellation is an act on the whole transfer, not on one piece: there is no way to stop a single + chunk while its siblings continue. + """ + + by_peer: bool = False + """A local cancel is an ordinary end; a peer's is a transfer error, and only the backend knows + which of the two happened.""" + + reports_pending: bool = True + + +Outcome = Union[Delivered, Failed, Cancelled] +"""How one delivery ended. ``None`` rather than a member means it has not ended yet. + +The logical member and its cause are committed at the task/session transition; +polling only observes that decision. ``reports_pending`` can still change, so a +stored outcome carries stale report progress, not a different logical result. + +``reports_pending`` asks one question and only one: is a report this transfer was owed still to +arrive. Receiving waits on the writers' reports, sending on word about its own writes. + +**It is not a release condition, and answering it is not permission to hand memory back.** Reports +and data travel on different channels and are not ordered, so no value of it means the peers have +stopped touching the destination. Whatever eventually answers that is a separate question this +contract does not ask. + +TODO: A backend that has torn down whatever a report would arrive on answers false -- which is how +this converges today, rather than by the reports arriving. An ``Attempt`` cannot yet settle on its +own evidence. + +TODO: Content addressing needs a member for "the source does not hold this", which is worth asking +somewhere else for while a failure is not. +""" + + +@runtime_checkable +class Attempt(Protocol): + """One delivery of one extent.""" + + def poll(self) -> Optional[Outcome]: + """Non-blocking. ``None`` while the delivery has not reached a conclusion. + + A failure is reported as soon as it is known, which may be before every report is in -- + that second question is carried by the outcome itself. + The logical result is stable and does not depend on when the first poll occurs. + """ + ... diff --git a/tensorrt_llm/_torch/disaggregation/native/fetch.py b/tensorrt_llm/_torch/disaggregation/native/fetch.py index c17b283b3108..557c20cc5d13 100644 --- a/tensorrt_llm/_torch/disaggregation/native/fetch.py +++ b/tensorrt_llm/_torch/disaggregation/native/fetch.py @@ -18,10 +18,10 @@ from typing import Optional -from tensorrt_llm._torch.disaggregation.base import Attempt, CacheExtent, Cancelled, Outcome from tensorrt_llm._torch.disaggregation.base.transfer import get_unique_rid from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from .extent import Attempt, CacheExtent, Cancelled, Outcome from .handle import NothingPublished, TaskHandle from .transfer import RxSession, SessionStatus, TransferWorker diff --git a/tensorrt_llm/_torch/disaggregation/native/handle.py b/tensorrt_llm/_torch/disaggregation/native/handle.py index 1f757d3e1dd5..5f8b8545fc55 100644 --- a/tensorrt_llm/_torch/disaggregation/native/handle.py +++ b/tensorrt_llm/_torch/disaggregation/native/handle.py @@ -22,8 +22,7 @@ from typing import Optional -from tensorrt_llm._torch.disaggregation.base import Cancelled, Delivered, Failed, Outcome - +from .extent import Cancelled, Delivered, Failed, Outcome from .transfer import KVRecvTask, KVSendTask, RxSession, SessionStatus, TaskStatus, TxSession diff --git a/tensorrt_llm/_torch/disaggregation/native/publish.py b/tensorrt_llm/_torch/disaggregation/native/publish.py index fc08aa861cfb..6aecf893ff75 100644 --- a/tensorrt_llm/_torch/disaggregation/native/publish.py +++ b/tensorrt_llm/_torch/disaggregation/native/publish.py @@ -16,9 +16,9 @@ from __future__ import annotations -from tensorrt_llm._torch.disaggregation.base import Attempt, CacheExtent from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from .extent import Attempt, CacheExtent from .fetch import refuse_extent_for_another_request from .handle import TaskHandle from .transfer import TxSession diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 56c61cf1ab59..7f5dc0184295 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -38,7 +38,6 @@ import tensorrt_llm.bindings from tensorrt_llm import logger -from tensorrt_llm._torch.disaggregation.base import Chunk from tensorrt_llm._torch.disaggregation.base.agent import ( BaseTransferAgent, MemoryDescs, @@ -61,6 +60,7 @@ build_aux_transfer_layout, get_non_empty_aux_indices, ) +from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 3c9e4f245db7..5b3c842a086a 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -25,7 +25,7 @@ import tensorrt_llm.bindings from tensorrt_llm import logger -from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange from tensorrt_llm._torch.disaggregation.base.agent import use_pure_python_transfer_agent from tensorrt_llm._torch.disaggregation.base.transfer import ( RxSessionBase, @@ -42,6 +42,7 @@ from tensorrt_llm._torch.disaggregation.native.bounce import ( config_from_size as bounce_config_from_size, ) +from tensorrt_llm._torch.disaggregation.native.extent import CacheExtent, Chunk from tensorrt_llm._torch.disaggregation.native.fetch import PeerFetch from tensorrt_llm._torch.disaggregation.native.perf_logger import perf_log_manager from tensorrt_llm._torch.disaggregation.native.publish import PeerPublish diff --git a/tests/unittest/_torch/disaggregation/test_backend_contract.py b/tests/unittest/_torch/disaggregation/test_backend_contract.py index 116670457184..29101cfff6e1 100644 --- a/tests/unittest/_torch/disaggregation/test_backend_contract.py +++ b/tests/unittest/_torch/disaggregation/test_backend_contract.py @@ -1,39 +1,41 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""What the cache-transfer contract promises, pinned. - -These are the invariants the rest of the design leans on: that ``reports_pending`` answers one -question and only one -- whether a report is still owed -- that success cannot be spelled with a -report still outstanding, and that the two things nobody can recompute -- whether a piece ends the -series, and how each group's ids are to be read -- have to be stated rather than defaulted. -""" +"""The backend contract's surface, pinned: its exports, the members and signatures of its +protocols, ``Capabilities`` and ``Outcome``, and a base package that imports nothing of the KV cache +lender. Also what native's per-request types in ``native/extent.py`` promise: ``reports_pending`` +says only whether a report is owed, success cannot leave one owed, ``is_last`` and each group's +kind must be stated, and native's ``Attempt`` needs only ``poll``.""" +import ast import dataclasses +import inspect +import pathlib +import types +import typing import numpy as np import pytest +import tensorrt_llm._torch.disaggregation.base as package +import tensorrt_llm._torch.disaggregation.base.backend as contract_module +import tensorrt_llm._torch.disaggregation.base.memory as memory_module +import tensorrt_llm._torch.disaggregation.native.extent as extent_module from tensorrt_llm._torch.disaggregation.base import ( - Attempt, CacheKind, - Cancelled, - Chunk, - Delivered, - Failed, - Fetches, - Publishes, + LayerGroupRows, + MemoryKind, + Outcome, + SlotRegion, TokenRange, + View, ) +from tensorrt_llm._torch.disaggregation.native.extent import Cancelled, Chunk, Delivered, Failed pytestmark = pytest.mark.cpu_only -def _owes_a_report(outcome) -> bool: - """No conclusion yet, or one that is still owed a report. - - Local to the tests on purpose. It used to be exported from the contract, where the name read - like permission to hand memory back -- which is the one thing it never meant. - """ +def _reports_may_be_pending(outcome) -> bool: + """Reports may still be pending; this says nothing about safe memory release.""" return outcome is None or outcome.reports_pending @@ -48,6 +50,61 @@ def _chunk(**overrides) -> Chunk: return Chunk(**fields) +# Minimal structural implementations. Groups and views compare by identity. +@dataclasses.dataclass(frozen=True) +class _SlotRegion: + layout_label: str + address: int + slot_bytes: int + slot_count: int + memory_kind: MemoryKind + + @property + def nbytes(self) -> int: + return self.slot_bytes * self.slot_count + + +@dataclasses.dataclass(frozen=True, eq=False) +class _LayerGroupRows: + layer_group_index: int + row_count: int + object_keys: typing.Optional[np.ndarray] = None + addresses: typing.Optional[np.ndarray] = None + region_index: typing.Optional[int] = None + + def __len__(self) -> int: + return self.row_count + + +@dataclasses.dataclass(frozen=True, eq=False) +class _View: + groups: typing.Tuple[_LayerGroupRows, ...] + + def row_masks(self, fill_value: bool = False) -> typing.Tuple[np.ndarray, ...]: + return tuple(np.full(len(group), fill_value, dtype=bool) for group in self.groups) + + +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 + 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("". +PROTOCOLS = { + "Backend": { + "capabilities": "-> Capabilities", + "fetch": "(self, view: 'View', source: 'Hint') -> 'Attempt'", + "publish": "(self, view: 'View', destination: 'Hint') -> 'Attempt'", + "shutdown": "(self, timeout: 'float') -> 'frozenset'", + }, + "Attempt": { + "poll": "(self) -> 'Outcome | None'", + "quiesced": "(self) -> 'bool'", + "cancel": "(self) -> 'None'", + }, + "View": { + "groups": "-> Sequence[LayerGroupRows]", + "row_masks": "(self, fill_value: 'bool' = False) -> 'tuple[np.ndarray, ...]'", + }, + "LayerGroupRows": { + "layer_group_index": "-> int", + "object_keys": "-> np.ndarray | None", + "addresses": "-> np.ndarray | None", + "region_index": "-> int | None", + "__len__": "(self) -> 'int'", + }, + "SlotRegion": { + "layout_label": "-> str", + "address": "-> int", + "nbytes": "-> int", + "slot_bytes": "-> int", + "slot_count": "-> int", + "memory_kind": "-> MemoryKind", + }, +} + + +def _protocol_member_signatures(protocol) -> dict: + members = {} + for name, value in vars(protocol).items(): + if name.startswith("_") and name != "__len__": + continue + if isinstance(value, property): + members[name] = f"-> {value.fget.__annotations__['return']}" + else: + members[name] = str(inspect.signature(value)) + return members + + +def test_base_exports_fourteen_names(): + assert len(EXPORTS) == 14 + assert set(contract_module.__all__) == EXPORTS + assert set(package.__all__) == EXPORTS + for name in EXPORTS: + assert getattr(package, name) is getattr(contract_module, name) + assert contract_module.Hint is typing.Any + assert contract_module.WakeFn == typing.Callable[[], None] + + +def test_the_contract_s_memory_kind_is_the_one_its_memory_module_defines(): + """One ``MemoryKind``, defined in ``base/memory.py``, with a member for each kind of memory a + region lies in.""" + assert MemoryKind is memory_module.MemoryKind + assert MemoryKind.__module__ == memory_module.__name__ + assert [kind.name for kind in MemoryKind] == ["HOST", "DEVICE"] + + +@pytest.mark.parametrize("name", sorted(PROTOCOLS)) +def test_each_protocol_has_expected_members_and_signatures(name): + protocol = getattr(contract_module, name) + assert _protocol_member_signatures(protocol) == PROTOCOLS[name] + assert not isinstance(object(), protocol) + + +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))) + assert isinstance(region, SlotRegion) and region.nbytes == 512 + assert isinstance(view, View) + assert all(isinstance(group, LayerGroupRows) for group in view.groups) + masks = view.row_masks(True) + assert [mask.tolist() for mask in masks] == [[True] * 2, [True] * 3] + assert all(mask.flags.writeable for mask in masks) + + +@pytest.mark.parametrize("name", ["SlotRegion", "LayerGroupRows", "View"]) +def test_runtime_protocol_checks_require_every_member(name): + protocol = getattr(contract_module, name) + whole = { + "SlotRegion": _SlotRegion("pool", 4096, 64, 8, MemoryKind.DEVICE), + "LayerGroupRows": _keyed_view(2).groups[0], + "View": _keyed_view(2, 3), + } + complete = { + member: getattr(whole[name], member) for member in _protocol_member_signatures(protocol) + } + assert isinstance(types.SimpleNamespace(**complete), protocol) + for dropped in complete: + rest = {member: value for member, value in complete.items() if member != dropped} + assert not isinstance(types.SimpleNamespace(**rest), protocol), dropped + + +# The constructors and methods of the concrete types a backend builds or calls. +SIGNATURES = { + "Capabilities": "(fetch: 'bool', publish: 'bool', cancel: 'bool' = False) -> None", + "Outcome": "(delivered_masks: 'tuple[np.ndarray, ...]', fault: 'str | None' = None) -> None", + "Outcome.for_view": ( + "(view: 'View', delivered_masks: 'Sequence[np.ndarray] | None' = None," + " fault: 'str | None' = None) -> 'Outcome'" + ), + "Outcome.check_mask_shapes": "(self, view: 'View') -> 'None'", +} + + +@pytest.mark.parametrize("name", sorted(SIGNATURES)) +def test_each_concrete_type_has_its_signature(name): + owner, _, member = name.partition(".") + target = getattr(contract_module, owner) + target = getattr(target, member) if member else target + assert str(inspect.signature(target)) == SIGNATURES[name] + + +def test_capabilities_fields_defaults_and_validation(): + capabilities = contract_module.Capabilities + fields = {f.name: f.default for f in dataclasses.fields(capabilities)} + assert list(fields) == ["fetch", "publish", "cancel"] + assert fields["fetch"] is fields["publish"] is dataclasses.MISSING + assert fields["cancel"] is False + assert capabilities.__dataclass_params__.frozen + caps = capabilities(fetch=True, publish=False) + assert [name for name in dir(caps) if not name.startswith("_")] == sorted(fields) + assert caps.cancel is False + with pytest.raises(dataclasses.FrozenInstanceError): + caps.cancel = True + assert capabilities(fetch=False, publish=True).publish + with pytest.raises(ValueError): + capabilities(fetch=False, publish=False) + + +def test_outcome_fields_and_identity(): + """Equal masks do not make two outcomes one, and an outcome hashes as itself.""" + assert [f.name for f in dataclasses.fields(Outcome)] == ["delivered_masks", "fault"] + assert Outcome.__dataclass_params__.frozen and not Outcome.__dataclass_params__.eq + view = _keyed_view(2) + first, second = Outcome.for_view(view), Outcome.for_view(view) + public = [name for name in dir(first) if not name.startswith("_")] + assert public == ["check_mask_shapes", "delivered_masks", "fault", "for_view"] + assert first == first and first != second + assert len({first, second}) == 2 + with pytest.raises(dataclasses.FrozenInstanceError): + first.fault = "peer reset" + + +def test_outcome_for_view_without_masks_delivers_nothing(): + assert list(inspect.signature(Outcome.for_view).parameters) == [ + "view", + "delivered_masks", + "fault", + ] + view = _keyed_view(2, 0, 3) + outcome = Outcome.for_view(view) + assert [mask.tolist() for mask in outcome.delivered_masks] == [[False] * 2, [], [False] * 3] + assert outcome.fault is None + outcome.check_mask_shapes(view) + assert Outcome.for_view(view, fault="peer reset").fault == "peer reset" + + +def test_outcome_for_view_with_masks_and_a_fault(): + view = _keyed_view(2, 3) + outcome = Outcome.for_view( + view, [np.array([True, False]), np.array([False, True, True])], "peer reset" + ) + assert all(mask.dtype == np.bool_ for mask in outcome.delivered_masks) + assert [mask.tolist() for mask in outcome.delivered_masks] == [ + [True, False], + [False, True, True], + ] + assert outcome.fault == "peer reset" + outcome.check_mask_shapes(view) + + +@pytest.mark.parametrize( + "mask", + [np.array([0, 1, 2]), np.array([0.0, 1.0, 1.0]), [1, 0, 1]], + ids=["int", "float", "list_of_ints"], +) +def test_outcome_refuses_a_mask_that_is_not_bool(mask): + """A status code read as truthiness would turn a failed row into a delivered one.""" + with pytest.raises(TypeError, match="delivery mask is bool"): + Outcome.for_view(_keyed_view(3), [mask]) + + +def test_outcome_masks_are_read_only_copies(): + """A backend may go on writing the masks it gave; the outcome does not change.""" + view = _keyed_view(3) + masks = view.row_masks(True) + outcome = Outcome.for_view(view, masks) + masks[0][0] = False + assert isinstance(outcome.delivered_masks, tuple) + assert outcome.delivered_masks[0].tolist() == [True] * 3 + for mask in outcome.delivered_masks: + assert not mask.flags.writeable + with pytest.raises(ValueError): + mask[0] = False + + +@pytest.mark.parametrize( + "delivered_masks", + [ + (np.zeros(2, dtype=bool),), + (np.zeros(2, dtype=bool), np.zeros(3, dtype=bool), np.zeros(0, dtype=bool)), + (np.zeros(1, dtype=bool), np.zeros(3, dtype=bool)), + (np.zeros(2, dtype=bool), np.zeros(4, dtype=bool)), + (np.zeros((2, 1), dtype=bool), np.zeros(3, dtype=bool)), + (np.zeros(2, dtype=bool), np.zeros((1, 3), dtype=bool)), + ], + ids=["fewer-masks", "more-masks", "short-mask", "long-mask", "2d-mask", "2d-row-mask"], +) +def test_check_mask_shapes_refuses_masks_that_do_not_fit_the_view(delivered_masks): + with pytest.raises(ValueError): + Outcome(delivered_masks).check_mask_shapes(_keyed_view(2, 3)) + + +def test_check_mask_shapes_refuses_a_mask_for_an_empty_group(): + outcome = Outcome((np.zeros(2, dtype=bool), np.zeros(1, dtype=bool))) + with pytest.raises(ValueError): + outcome.check_mask_shapes(_keyed_view(2, 0)) + + +def test_an_invalid_submission_is_a_rejection_and_a_value_error(): + invalid = contract_module.InvalidSubmission + assert contract_module.SubmissionRejected.__bases__ == (Exception,) + assert invalid.__bases__ == (contract_module.SubmissionRejected, ValueError) + assert issubclass(invalid, contract_module.SubmissionRejected) + assert issubclass(invalid, ValueError) + assert not issubclass(contract_module.SubmissionRejected, ValueError) + + +LENDER_MODULE = "tensorrt_llm._torch.pyexecutor.kv_cache.sharing" +NATIVE_PACKAGE = "tensorrt_llm._torch.disaggregation.native" + + +def _type_checking_nodes(tree: ast.AST) -> set: + """The ids of every node under the body of an ``if TYPE_CHECKING:``.""" + ids = set() + for node in ast.walk(tree): + test = getattr(node, "test", None) + if isinstance(node, ast.If) and ( + (isinstance(test, ast.Name) and test.id == "TYPE_CHECKING") + or (isinstance(test, ast.Attribute) and test.attr == "TYPE_CHECKING") + ): + for child in node.body: + ids.update(id(sub) for sub in ast.walk(child)) + return ids + + +def _imported_names(path: str, package_name: str, runtime_only: bool = False) -> list: + """Qualified module and member names imported by ``path``, including deferred imports; with + ``runtime_only``, without those under ``if TYPE_CHECKING:``.""" + with open(path, encoding="utf-8") as source: + tree = ast.parse(source.read()) + skipped = _type_checking_nodes(tree) if runtime_only else set() + found = [] + for node in ast.walk(tree): + if id(node) in skipped: + continue + if isinstance(node, ast.Import): + found.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + module = node.module or "" + if node.level: + anchor = package_name.rsplit(".", node.level - 1)[0] + module = f"{anchor}.{module}" if module else anchor + found.append(module) + found.extend(f"{module}.{alias.name}" for alias in node.names) + return found + + +def _lender_imports(imported_names: list) -> list: + return [ + name + for name in imported_names + if name == LENDER_MODULE or name.startswith(LENDER_MODULE + ".") + ] + + +def _base_modules() -> list: + """Every ``.py`` file under ``base/``, subpackages included, as ``(path, package, label)``; + ``package`` is the package its relative imports resolve in.""" + root = pathlib.Path(package.__file__).parent + modules = [] + for path in sorted(root.rglob("*.py")): + subpackage = ".".join((package.__name__, *path.parent.relative_to(root).parts)) + modules.append((str(path), subpackage, path.relative_to(root).as_posix())) + return modules + + +def test_base_imports_nothing_of_the_lender(): + modules = _base_modules() + assert {"__init__.py", "backend.py", "memory.py"} <= {label for _, _, label in modules} + for path, subpackage, label in modules: + imported = _lender_imports(_imported_names(path, subpackage)) + assert imported == [], f"base/{label} imports {imported}" + + +def test_base_imports_native_only_for_type_checking(): + for path, subpackage, label in _base_modules(): + native = [ + name + for name in _imported_names(path, subpackage, runtime_only=True) + if name == NATIVE_PACKAGE or name.startswith(NATIVE_PACKAGE + ".") + ] + assert native == [], f"base/{label} imports {native} at run time" + + +def test_the_runtime_import_scan_skips_only_type_checking_blocks(tmp_path): + source = tmp_path / "module.py" + source.write_text( + "from typing import TYPE_CHECKING\n" + "if TYPE_CHECKING:\n" + " from ..native.extent import Chunk\n" + "def later():\n" + " from ..native import extent\n" + ) + found = _imported_names(str(source), package.__name__, runtime_only=True) + assert f"{NATIVE_PACKAGE}.extent.Chunk" not in found + assert f"{NATIVE_PACKAGE}.extent" in found + + +def test_the_import_scan_sees_every_kind_of_import(tmp_path): + source = tmp_path / "module.py" + source.write_text( + "from typing import TYPE_CHECKING\n" + "if TYPE_CHECKING:\n" + " from ...pyexecutor.kv_cache.sharing import Hidden\n" + "def later():\n" + " import tensorrt_llm._torch.pyexecutor.kv_cache.sharing._types\n" + " from ...pyexecutor.kv_cache import sharing\n" + ) + found = _lender_imports(_imported_names(str(source), package.__name__)) + assert f"{LENDER_MODULE}.Hidden" in found + assert f"{LENDER_MODULE}._types" in found + assert LENDER_MODULE in found + + # --------------------------------------------------------------------------- # Chunk # --------------------------------------------------------------------------- @@ -127,6 +581,12 @@ def test_cancelled_is_pending_until_told_otherwise(): assert Cancelled().by_peer is False +def test_failed_states_what_is_owed(): + """A failure has no default for ``reports_pending``: nobody can guess it for the caller.""" + with pytest.raises(TypeError): + Failed(reason="x") + + def test_outcomes_are_frozen(): """A stored outcome must not be mutated into a stale claim.""" outcome = Delivered(token_end=4) @@ -139,13 +599,13 @@ def test_outcomes_are_frozen(): # --------------------------------------------------------------------------- -def test_no_conclusion_counts_as_owed(): +def test_no_conclusion_may_still_have_reports_pending(): """A transfer that has not concluded may be queued and about to write.""" - assert _owes_a_report(None) is True + assert _reports_may_be_pending(None) is True -def test_a_delivery_owes_nothing_further(): - assert _owes_a_report(Delivered(token_end=8)) is False +def test_a_delivery_has_no_reports_pending(): + assert _reports_may_be_pending(Delivered(token_end=8)) is False @pytest.mark.parametrize( @@ -157,13 +617,13 @@ def test_a_delivery_owes_nothing_further(): (Cancelled(reports_pending=False), False), ], ) -def test_what_is_owed_does_not_follow_from_the_kind_of_ending(outcome, expected): +def test_reports_pending_does_not_follow_from_the_kind_of_ending(outcome, expected): """How a delivery ended and what it still owes are separate; neither is read off the other.""" - assert _owes_a_report(outcome) is expected + assert _reports_may_be_pending(outcome) is expected # --------------------------------------------------------------------------- -# Protocols +# Native's Attempt # --------------------------------------------------------------------------- @@ -172,15 +632,11 @@ class OnePiece: def poll(self): return None - assert isinstance(OnePiece(), Attempt) - + assert isinstance(OnePiece(), extent_module.Attempt) -def test_reading_and_offering_are_separate_to_implement(): - """A store that only answers reads is asked for nothing it cannot do.""" - class ReadOnly: - def fetch(self, extent, *, src=None): - raise NotImplementedError +def test_something_without_poll_is_not_an_attempt(): + class NoPoll: + pass - assert isinstance(ReadOnly(), Fetches) - assert not isinstance(ReadOnly(), Publishes) + assert not isinstance(NoPoll(), extent_module.Attempt) diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index d75c51a6ca77..8a4f307fb3f6 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -25,8 +25,9 @@ import pytest from tensorrt_llm import DisaggregatedParams -from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult +from tensorrt_llm._torch.disaggregation.native.extent import CacheExtent, Chunk from tensorrt_llm._torch.disaggregation.native.transfer import ( AgentResult, KVSendTask, diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 721b1d062537..28bb58ab9efb 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -55,8 +55,9 @@ else: TRANSFER_AGENT_BINDING_IMPORT_ERROR = None from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams -from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult +from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.transfer import TransferWorker, TransferWorkerConfig from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 diff --git a/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py b/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py index 6ad0e67290ba..29d38828ccc8 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py +++ b/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py @@ -16,9 +16,10 @@ import pytest -from tensorrt_llm._torch.disaggregation.base import Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus from tensorrt_llm._torch.disaggregation.native import transfer as transfer_mod +from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.retirement import ( QuiescenceFatalEvent, RetirementWatchdog, diff --git a/tests/unittest/disaggregated/test_kv_transfer_mp.py b/tests/unittest/disaggregated/test_kv_transfer_mp.py index 0a569cd18642..6df42db04e2a 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_mp.py +++ b/tests/unittest/disaggregated/test_kv_transfer_mp.py @@ -35,8 +35,9 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams -from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus +from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.transfer import TransferWorker, TransferWorkerConfig from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType diff --git a/tests/unittest/disaggregated/test_peer_fetch.py b/tests/unittest/disaggregated/test_peer_fetch.py index 5c8d06aa910a..03941b816b6d 100644 --- a/tests/unittest/disaggregated/test_peer_fetch.py +++ b/tests/unittest/disaggregated/test_peer_fetch.py @@ -17,16 +17,15 @@ import pytest from tensorrt_llm import DisaggregatedParams -from tensorrt_llm._torch.disaggregation.base import ( +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange +from tensorrt_llm._torch.disaggregation.base.transfer import get_unique_rid +from tensorrt_llm._torch.disaggregation.native.extent import ( CacheExtent, - CacheKind, Cancelled, Chunk, Delivered, Failed, - TokenRange, ) -from tensorrt_llm._torch.disaggregation.base.transfer import get_unique_rid from tensorrt_llm._torch.disaggregation.native.fetch import PeerFetch from tensorrt_llm._torch.disaggregation.native.transfer import ( KVRecvTask, diff --git a/tests/unittest/disaggregated/test_peer_publish.py b/tests/unittest/disaggregated/test_peer_publish.py index e1737489d6f3..3bb2ebbe94e5 100644 --- a/tests/unittest/disaggregated/test_peer_publish.py +++ b/tests/unittest/disaggregated/test_peer_publish.py @@ -14,14 +14,8 @@ import pytest from tensorrt_llm import DisaggregatedParams -from tensorrt_llm._torch.disaggregation.base import ( - CacheExtent, - CacheKind, - Chunk, - Delivered, - Publishes, - TokenRange, -) +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange +from tensorrt_llm._torch.disaggregation.native.extent import Attempt, CacheExtent, Chunk, Delivered from tensorrt_llm._torch.disaggregation.native.publish import PeerPublish from tensorrt_llm._torch.disaggregation.native.transfer import TaskStatus, TxSession @@ -65,8 +59,9 @@ def _sending(rid: int = 42) -> TxSession: ) -def test_the_adapter_satisfies_the_contract(): - assert isinstance(PeerPublish(_sending(), _request()), Publishes) +def test_the_adapter_answers_with_an_attempt(): + handle = PeerPublish(_sending(), _request()).publish(_extent()) + assert isinstance(handle, Attempt) def test_each_piece_gets_its_own_handle(): diff --git a/tests/unittest/disaggregated/test_retirement_deadline.py b/tests/unittest/disaggregated/test_retirement_deadline.py index ce530f83c5cc..612c810d97c8 100644 --- a/tests/unittest/disaggregated/test_retirement_deadline.py +++ b/tests/unittest/disaggregated/test_retirement_deadline.py @@ -10,10 +10,11 @@ import numpy as np import pytest -from tensorrt_llm._torch.disaggregation.base import Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus from tensorrt_llm._torch.disaggregation.native import transfer as transfer_mod from tensorrt_llm._torch.disaggregation.native.bounce import NoBounceTransport +from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.retirement import RetirementWatchdog from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle diff --git a/tests/unittest/disaggregated/test_task_handle.py b/tests/unittest/disaggregated/test_task_handle.py index 43f79367cca1..a0a17fe5350c 100644 --- a/tests/unittest/disaggregated/test_task_handle.py +++ b/tests/unittest/disaggregated/test_task_handle.py @@ -27,16 +27,15 @@ import pytest from tensorrt_llm import DisaggregatedParams -from tensorrt_llm._torch.disaggregation.base import ( +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange +from tensorrt_llm._torch.disaggregation.base.transfer import WaitResult +from tensorrt_llm._torch.disaggregation.native.extent import ( CacheExtent, - CacheKind, Cancelled, Chunk, Delivered, Failed, - TokenRange, ) -from tensorrt_llm._torch.disaggregation.base.transfer import WaitResult from tensorrt_llm._torch.disaggregation.native.fetch import PeerFetch from tensorrt_llm._torch.disaggregation.native.handle import TaskHandle from tensorrt_llm._torch.disaggregation.native.transfer import ( diff --git a/tests/unittest/disaggregated/test_transfer_late_settlement.py b/tests/unittest/disaggregated/test_transfer_late_settlement.py index 4dc9c93b1f40..a1321ae6cfab 100644 --- a/tests/unittest/disaggregated/test_transfer_late_settlement.py +++ b/tests/unittest/disaggregated/test_transfer_late_settlement.py @@ -23,8 +23,9 @@ import pytest import tensorrt_llm._torch.disaggregation.native.transfer as transfer_mod -from tensorrt_llm._torch.disaggregation.base import Cancelled, Chunk, Failed, TokenRange +from tensorrt_llm._torch.disaggregation.base import TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus +from tensorrt_llm._torch.disaggregation.native.extent import Cancelled, Chunk, Failed from tensorrt_llm._torch.disaggregation.native.handle import TaskHandle from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle diff --git a/tests/unittest/disaggregated/test_transfer_lifecycle_acceptance.py b/tests/unittest/disaggregated/test_transfer_lifecycle_acceptance.py index a60d8f778a5a..a61b2a86851d 100644 --- a/tests/unittest/disaggregated/test_transfer_lifecycle_acceptance.py +++ b/tests/unittest/disaggregated/test_transfer_lifecycle_acceptance.py @@ -26,9 +26,10 @@ import pytest import tensorrt_llm._torch.disaggregation.transceiver as transceiver_mod -from tensorrt_llm._torch.disaggregation.base import Cancelled, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus from tensorrt_llm._torch.disaggregation.native.bounce.impl import NoBounceTransport +from tensorrt_llm._torch.disaggregation.native.extent import Cancelled, Chunk from tensorrt_llm._torch.disaggregation.native.handle import TaskHandle from tensorrt_llm._torch.disaggregation.native.transfer import AgentResult, KVRecvTask, RxSession from tensorrt_llm._torch.disaggregation.orchestration import coordinator as coordinator_mod diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index cfa6213af139..8fecebd650a1 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -28,9 +28,10 @@ import tensorrt_llm._torch.disaggregation.native.transfer as transfer_mod import tensorrt_llm._torch.disaggregation.transceiver as transceiver_mod from tensorrt_llm import DisaggregatedParams -from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult from tensorrt_llm._torch.disaggregation.native.bounce.core import TransferContext +from tensorrt_llm._torch.disaggregation.native.extent import CacheExtent, Chunk from tensorrt_llm._torch.disaggregation.native.transfer import ( AgentResult, KVRecvTask, diff --git a/tests/unittest/disaggregated/test_unsubmitted_session.py b/tests/unittest/disaggregated/test_unsubmitted_session.py index 09eb9bfa2296..fbe719015676 100644 --- a/tests/unittest/disaggregated/test_unsubmitted_session.py +++ b/tests/unittest/disaggregated/test_unsubmitted_session.py @@ -22,9 +22,10 @@ import numpy as np import pytest -from tensorrt_llm._torch.disaggregation.base import Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus from tensorrt_llm._torch.disaggregation.native import transfer as transfer_mod +from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.retirement import RetirementWatchdog from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle diff --git a/tests/unittest/disaggregated/test_verified_reuse_admission.py b/tests/unittest/disaggregated/test_verified_reuse_admission.py index 1e2a130e0018..8f2373996018 100644 --- a/tests/unittest/disaggregated/test_verified_reuse_admission.py +++ b/tests/unittest/disaggregated/test_verified_reuse_admission.py @@ -36,8 +36,9 @@ import tensorrt_llm._torch.disaggregation.native.bounce as bounce_mod import tensorrt_llm._torch.disaggregation.native.transfer as transfer_mod from tensorrt_llm import DisaggregatedParams -from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base import CacheKind, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus +from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ( IntactMapper, ReplicatedMapper, From 2b24846261f15f3a80e3b8a4596677a61d4b20b0 Mon Sep 17 00:00:00 2001 From: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:39:38 -0700 Subject: [PATCH 2/3] [None][fix] Keep read-only bool delivery masks in Outcome Delivery masks are of dtype bool; anything else raises TypeError. Outcome keeps a read-only bool mask whose type is exactly numpy.ndarray as it is, and its giver writes neither it nor memory it shares afterwards; any other mask is copied. A read-only subclass of numpy.ndarray, such as a masked array, is copied into a plain array, since all() on a masked array skips masked rows. The contract tests add a kept read-only mask beside a copied writable one, a read-only masked array, and read-only int and float masks. Signed-off-by: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com> --- .../_torch/disaggregation/base/backend.py | 15 ++++--- .../disaggregation/test_backend_contract.py | 39 +++++++++++++++++-- 2 files changed, 45 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/base/backend.py b/tensorrt_llm/_torch/disaggregation/base/backend.py index 06824f5a2ce9..c52e6f7ce780 100644 --- a/tensorrt_llm/_torch/disaggregation/base/backend.py +++ b/tensorrt_llm/_torch/disaggregation/base/backend.py @@ -232,12 +232,15 @@ def __post_init__(self): raise ValueError("a backend implements at least one of fetch and publish") -def _readonly_bool_copy(mask: np.ndarray) -> np.ndarray: - """A read-only copy of ``mask``, which must be a ``bool`` array. +def _delivery_mask(mask: np.ndarray) -> np.ndarray: + """``mask`` itself if it is a read-only ``bool`` array whose type is exactly ``numpy.ndarray``; + otherwise a read-only copy of it. Raises: TypeError: ``mask`` is not of dtype ``bool``. """ + if type(mask) is np.ndarray and mask.dtype == np.bool_ and not mask.flags.writeable: + return mask copy = np.array(mask) if copy.dtype != np.bool_: raise TypeError(f"a delivery mask is bool, not {copy.dtype}") @@ -259,11 +262,13 @@ class Outcome: Attributes: delivered_masks: One read-only ``bool`` mask of shape ``[len(group)]`` per group of the - view, in order, copied from the masks given. + view, in order. 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. fault: What went wrong, or ``None``. Raises: - TypeError: A mask given is not a ``bool`` array. + TypeError: A mask given is not of dtype ``bool``. """ delivered_masks: tuple[np.ndarray, ...] @@ -273,7 +278,7 @@ def __post_init__(self): object.__setattr__( self, "delivered_masks", - tuple(_readonly_bool_copy(mask) for mask in self.delivered_masks), + tuple(_delivery_mask(mask) for mask in self.delivered_masks), ) @classmethod diff --git a/tests/unittest/_torch/disaggregation/test_backend_contract.py b/tests/unittest/_torch/disaggregation/test_backend_contract.py index 29101cfff6e1..69b7651a0fe0 100644 --- a/tests/unittest/_torch/disaggregation/test_backend_contract.py +++ b/tests/unittest/_torch/disaggregation/test_backend_contract.py @@ -351,10 +351,21 @@ def test_outcome_for_view_with_masks_and_a_fault(): outcome.check_mask_shapes(view) +def _read_only(array: np.ndarray) -> np.ndarray: + array.flags.writeable = False + return array + + @pytest.mark.parametrize( "mask", - [np.array([0, 1, 2]), np.array([0.0, 1.0, 1.0]), [1, 0, 1]], - ids=["int", "float", "list_of_ints"], + [ + np.array([0, 1, 2]), + np.array([0.0, 1.0, 1.0]), + [1, 0, 1], + _read_only(np.array([0, 1, 2])), + _read_only(np.array([0.0, 1.0, 1.0])), + ], + ids=["int", "float", "list_of_ints", "read_only_int", "read_only_float"], ) def test_outcome_refuses_a_mask_that_is_not_bool(mask): """A status code read as truthiness would turn a failed row into a delivered one.""" @@ -362,8 +373,8 @@ def test_outcome_refuses_a_mask_that_is_not_bool(mask): Outcome.for_view(_keyed_view(3), [mask]) -def test_outcome_masks_are_read_only_copies(): - """A backend may go on writing the masks it gave; the outcome does not change.""" +def test_outcome_copies_writable_masks(): + """A backend may go on writing the writable masks it gave; the outcome does not change.""" view = _keyed_view(3) masks = view.row_masks(True) outcome = Outcome.for_view(view, masks) @@ -376,6 +387,26 @@ def test_outcome_masks_are_read_only_copies(): mask[0] = False +def test_outcome_keeps_read_only_bool_masks(): + """A read-only bool array is kept as given, not copied; a writable one beside it is copied.""" + kept, written = _read_only(np.ones(2, dtype=bool)), np.ones(3, dtype=bool) + outcome = Outcome.for_view(_keyed_view(2, 3), [kept, written]) + assert outcome.delivered_masks[0] is kept + assert outcome.delivered_masks[1] is not written + assert not outcome.delivered_masks[1].flags.writeable + + +def test_outcome_copies_a_read_only_subclass_into_a_plain_array(): + """``all()`` on a masked array skips masked rows, so a failed row there would read as + delivered; the outcome holds a plain array of the data instead.""" + mask = _read_only(np.ma.MaskedArray([True, False], mask=[False, True])) + assert mask.all() + held = Outcome.for_view(_keyed_view(2), [mask]).delivered_masks[0] + assert type(held) is np.ndarray + assert held.tolist() == [True, False] + assert not held.flags.writeable + + @pytest.mark.parametrize( "delivered_masks", [ From 8ac96cf4edf9be251fdf14b29b5384281a1d4e55 Mon Sep 17 00:00:00 2001 From: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:31:49 -0700 Subject: [PATCH 3/3] [None][fix] Promise a new mask from every View.row_masks call A backend can then mark the masks it is given read-only and hand them to Outcome, which keeps a read-only bool mask without copying it. Signed-off-by: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/_torch/disaggregation/base/backend.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/disaggregation/base/backend.py b/tensorrt_llm/_torch/disaggregation/base/backend.py index c52e6f7ce780..c583ab585d95 100644 --- a/tensorrt_llm/_torch/disaggregation/base/backend.py +++ b/tensorrt_llm/_torch/disaggregation/base/backend.py @@ -191,7 +191,8 @@ class View(Protocol): def groups(self) -> Sequence[LayerGroupRows]: ... def row_masks(self, fill_value: bool = False) -> tuple[np.ndarray, ...]: - """One writable ``bool`` mask of shape ``[len(group)]`` per group, in :attr:`groups` order. + """One new writable ``bool`` mask of shape ``[len(group)]`` per group, in :attr:`groups` + order. Args: fill_value: What every entry starts as.