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..c583ab585d95 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,384 @@ 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 new 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 _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. -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``. + """ + 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}") + 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. 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 of dtype ``bool``. + """ -``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(_delivery_mask(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..69b7651a0fe0 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) + + +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], + _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.""" + with pytest.raises(TypeError, match="delivery mask is bool"): + Outcome.for_view(_keyed_view(3), [mask]) + + +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) + 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 + + +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", + [ + (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 +612,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 +630,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 +648,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 +663,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,