diff --git a/tensorrt_llm/_torch/disaggregation/base/backend.py b/tensorrt_llm/_torch/disaggregation/base/backend.py index c583ab585d95..4ba5b1222f26 100644 --- a/tensorrt_llm/_torch/disaggregation/base/backend.py +++ b/tensorrt_llm/_torch/disaggregation/base/backend.py @@ -98,8 +98,8 @@ class CacheKind(IntEnum): @runtime_checkable class SlotRegion(Protocol): - """A contiguous memory region divided into equal-sized slots; a row given by an address - occupies one slot. + """A contiguous memory region divided into equal-sized slots; a row lying in the region occupies + one of its slots. A backend registers only the regions whose :attr:`memory_kind` it can access, each once as ``(address, nbytes)``; rows in a region it has not registered are rows it does not recognize. @@ -141,10 +141,11 @@ def memory_kind(self) -> MemoryKind: class LayerGroupRows(Protocol): """One layer group's rows of a view, described by aligned arrays. - A row is one transfer unit. When addresses are provided, each row occupies one slot in the - region. These members are all a backend reads of a group, apart from those a capability it - declares names. Row addresses need not be consecutive. A backend given rows it does not - recognize rejects the submission with :class:`InvalidSubmission`; it never skips them. + A row is one transfer unit. When slots are provided, row ``i`` occupies slot ``slots[i]`` in + every region :attr:`region_indices` lists. These members are all a backend reads of a group, + apart from those a capability it declares names. Slots need not be consecutive. A backend given + rows it does not recognize rejects the submission with :class:`InvalidSubmission`; it never + skips them. """ @property @@ -152,6 +153,13 @@ def layer_group_index(self) -> int: """The lender's local layer-group index, not the group's position in :attr:`View.groups`.""" ... + @property + def page_ordinals(self) -> np.ndarray: + """One page ordinal per row, ``int64`` of shape ``[rows]``: the request's page the row + holds, as the lender's published page layout numbers pages. Ordinals need not be + consecutive.""" + ... + @property def object_keys(self) -> np.ndarray | None: """One opaque object key per row, ``uint8`` of shape ``[rows, width]`` with the lender's @@ -161,16 +169,16 @@ def object_keys(self) -> np.ndarray | None: ... @property - def addresses(self) -> np.ndarray | None: - """One address per row, ``int64`` of shape ``[rows]``: row ``i`` is the slot at - ``addresses[i]`` in the :class:`SlotRegion` at index :attr:`region_index`. ``None`` exactly - when :attr:`region_index` is.""" + def slots(self) -> np.ndarray | None: + """One slot per row, ``int64`` of shape ``[rows]``: row ``i`` occupies slot ``slots[i]`` in + every :class:`SlotRegion` :attr:`region_indices` lists, its bytes those slots back to back + in the listed order. ``None`` exactly when :attr:`region_indices` is.""" ... @property - def region_index(self) -> int | None: - """The index, among the lender's regions, of the :class:`SlotRegion` the rows sit in, or - ``None``.""" + def region_indices(self) -> tuple[int, ...] | None: + """The indices, among the lender's regions, of the :class:`SlotRegion` objects the rows lie + in, at least one and distinct, in the order a row's bytes run; or ``None``.""" ... def __len__(self) -> int: ... diff --git a/tests/unittest/_torch/disaggregation/test_backend_contract.py b/tests/unittest/_torch/disaggregation/test_backend_contract.py index 69b7651a0fe0..4e0e6b0e1d7e 100644 --- a/tests/unittest/_torch/disaggregation/test_backend_contract.py +++ b/tests/unittest/_torch/disaggregation/test_backend_contract.py @@ -67,13 +67,13 @@ def nbytes(self) -> int: @dataclasses.dataclass(frozen=True, eq=False) class _LayerGroupRows: layer_group_index: int - row_count: int + page_ordinals: np.ndarray object_keys: typing.Optional[np.ndarray] = None - addresses: typing.Optional[np.ndarray] = None - region_index: typing.Optional[int] = None + slots: typing.Optional[np.ndarray] = None + region_indices: typing.Optional[typing.Tuple[int, ...]] = None def __len__(self) -> int: - return self.row_count + return len(self.page_ordinals) @dataclasses.dataclass(frozen=True, eq=False) @@ -86,22 +86,17 @@ def row_masks(self, fill_value: bool = False) -> typing.Tuple[np.ndarray, ...]: 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 + and three in layer group 1. Rows hold pages 0, 1, ... and occupy consecutive slots of region 0; + their keys encode their slots. The arrays are read-only, as a view's metadata is immutable.""" + groups, first_slot = [], 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(" int", + "page_ordinals": "-> np.ndarray", "object_keys": "-> np.ndarray | None", - "addresses": "-> np.ndarray | None", - "region_index": "-> int | None", + "slots": "-> np.ndarray | None", + "region_indices": "-> tuple[int, ...] | None", "__len__": "(self) -> 'int'", }, "SlotRegion": { @@ -246,7 +242,7 @@ 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))) + view = _View((_keyed_view(2).groups[0], _LayerGroupRows(1, np.arange(3)))) assert isinstance(region, SlotRegion) and region.nbytes == 512 assert isinstance(view, View) assert all(isinstance(group, LayerGroupRows) for group in view.groups)