From 39353807bbda0bd05e1c071f2bc0e03319536430 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 04:12:44 -0700 Subject: [PATCH 01/25] [None][test] Gate the gpt_oss modeling_v2 target under disaggregated serving Context server on one GPU, generation server on another, KV moved by the NIXL Python transceiver, both workers forced to the target with TRTLLM_MODELING_V2=require. The target is unchanged: each side only ever sees one of the two batch shapes the aggregated case alternates between. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../defs/accuracy/test_modeling_v2_gpt_oss.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py b/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py index 2b033d506a9c..cf8cd5fc12e2 100644 --- a/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py +++ b/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py @@ -35,6 +35,7 @@ from ..conftest import llm_models_root from .accuracy_core import GSM8K, LlmapiAccuracyTestHarness +from .test_disaggregated_serving import launch_disaggregated_llm, run_accuracy_test # The targets assert their own SM at construction: certification is per GPU # architecture, and a receipt from another one says nothing here. @@ -87,3 +88,65 @@ def test_gsm8k(self, mocker): ) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm, extra_evaluator_kwargs=self.extra_evaluator_kwargs) + + @skip_not_sm103 + @pytest.mark.skip_less_device(2) + def test_gsm8k_disagg(self, mocker): + """The same target on both sides of a disaggregated deployment. + + Context server on one GPU, generation server on another, KV moved by + the NIXL Python transceiver. The target itself is unchanged: the context + server only ever builds prefill batches and the generation server only + decode batches, so each side settles on one of the two targets that the + aggregated case already alternates between. + + ``modeling_v2`` is an ``LLM`` argument, so it goes into each server's + config: ``trtllm-serve`` reads its ``LLM`` arguments from a YAML the + harness writes. The validate switch is an environment variable and + travels through ``extra_env``. + """ + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 8192) + mocker.patch.dict(GSM8K.EVALUATE_KWARGS, {"scores_filter": "exact_match,flexible-extract"}) + transceiver = { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + "max_tokens_in_buffer": 4096, + } + kv_cache = { + "max_attention_window": [128, 32768], + "enable_block_reuse": False, + "free_gpu_memory_fraction": 0.5, + } + ctx_server_config = { + "modeling_v2": "require", + "disable_overlap_scheduler": True, + "cache_transceiver_config": transceiver, + "kv_cache_config": kv_cache, + "tensor_parallel_size": 1, + } + gen_server_config = { + "modeling_v2": "require", + "disable_overlap_scheduler": False, + "cache_transceiver_config": transceiver, + "kv_cache_config": kv_cache, + "tensor_parallel_size": 1, + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": {"num_instances": 1}, + "generation_servers": {"num_instances": 1}, + } + with launch_disaggregated_llm( + disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + extra_env={MODELING_V2_VALIDATE_ENV: "1"}, + ) as llm: + run_accuracy_test( + llm, + self.MODEL_NAME, + test_sets=["GSM8K"], + extra_evaluator_kwargs={GSM8K: self.extra_evaluator_kwargs}, + ) From b4b6f27335b86dd4e593ea6030e3cc24947c8e8b Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 05:15:15 -0700 Subject: [PATCH 02/25] [None][doc] Design spec for model-owned disaggregated KV transfer planning Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- ...26-10-10-disagg-transfer-planner-design.md | 221 ++++++++++++++++++ 1 file changed, 221 insertions(+) create mode 100644 docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md diff --git a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md new file mode 100644 index 000000000000..2d61c5db5ecb --- /dev/null +++ b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md @@ -0,0 +1,221 @@ +# Disaggregated KV transfer: model-owned transfer planning + +Date: 2026-10-10 +Status: draft for review +Scope: `tensorrt_llm/_torch/disaggregation/native/` (transfer v2, Python/NIXL runtime) + +## Problem + +The transfer v2 sender decides *what bytes go where* in one shared code path: +`Sender._build_kv_write_meta` drives `PeerRegistrar`, which dispatches to two +hard-wired policies (`AttentionPolicy`, `MambaPolicy`) keyed by a closed +`CacheKind`/`MapperKind` enum set. Model knowledge leaks into the shared +module at several points: + +| Site | What it knows | +|------|---------------| +| `native/transfer.py:1762-1777` | `CacheKind.STATE` branch, `MambaPolicy.is_paired`, helix CP strided slicing | +| `native/transfer.py:3134-3181` | `_fanin_bounce_safe` (duplicate-head, replicated views), `_get_mamba_slot` | +| `native/transfer.py:3262-3273` | `mamba_receiver_payload_bytes` for bounce sizing | +| `native/transfer.py:3381-3396` | receiver-side `validate_peer_compatible` of both policies | +| `native/peer.py` | `_POLICY_CLASSES` closed dict; `get_peer_overlap` computes attention `duplicate_head_factor` | +| `native/rank_info.py` | `AttentionInfo` is the only model-side description a rank can publish | +| `resource/page.py:37` | `MapperKind` closed enum | + +A model whose cache has a new structure (a new recurrent state layout, a new +sharding of attention state, a sparse index pool with its own replication +rule) cannot support disaggregation without editing this shared module, and +every such edit has to preserve the invariants of every other model that +already flows through it. + +## Idea + +Apply the modeling_v2 principle: a target is one concrete implementation of +one concrete topology, so the target owns its transfer plan for that +topology. The framework keeps everything that is about *moving* bytes and +*tracking* a transfer; the model owns everything that is about *which* +bytes pair with which. + +The current shared implementation is not deleted. It becomes +`DefaultTransferPlanner`, the planner every v1 model and every modeling_v2 +target that does not override it uses. This mirrors the shared ops in the +modeling_v2 catalog. + +## Boundary + +### Model-owned: `TransferPlanner` + +New module `tensorrt_llm/_torch/disaggregation/base/planner.py`. + +```python +@dataclass +class PeerTargets: + ranks: list[int] # peer flat ranks this rank exchanges with + uniform_writers: bool # every writer contributes equal bytes (receiver fan-in bounce split) + +@dataclass +class TransferPlan: + src_ptrs: np.ndarray # int64, device addresses on this rank + dst_ptrs: np.ndarray # int64, device addresses on the peer rank + sizes: np.ndarray # int64, bytes per (src, dst) pair + # An empty plan is valid: this rank writes nothing to this peer but still + # completes the handshake, as duplicate-head non-owners do today. + +class TransferPlanner(ABC): + planner_id: str # must match on both sides of a transfer + + @abstractmethod + def describe(self) -> bytes: + """Opaque self-description carried in RankInfo.planner_blob. Only this + planner class decodes it (it runs on both sides).""" + + @abstractmethod + def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: + """Raise PeerIncompatibleError with a field-level reason. Replaces + PeerRegistrar._check_peer_compatible and both validate_peer_compatible.""" + + @abstractmethod + def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> PeerTargets: + """Which peer ranks this rank exchanges with for one DP group of the peer. + Replaces get_peer_overlap + should_send_pool + _fanin_bounce_safe.""" + + @abstractmethod + def plan( + self, + self_ri: RankInfo, + peer_ri: RankInfo, + src_block_ids_per_lg: list[np.ndarray], + dst_block_ids_per_lg: list[np.ndarray], + token_range: TokenRange, + prompt_len: int, + dst_start_token: Optional[int], + ) -> TransferPlan: + """Sender side. Replaces the body of _build_kv_write_meta.""" + + @abstractmethod + def receiver_payload_bytes( + self, self_ri: RankInfo, peer_ri: RankInfo, dst_block_ids_per_lg: list[np.ndarray] + ) -> Optional[int]: + """Receiver side, bytes beyond the paged blocks that ride the same write + (recurrent state today). None = cannot size, receiver disables bounce. + Replaces _get_mamba_slot + mamba_receiver_payload_bytes.""" +``` + +Threading contract is unchanged from today: `plan` is called where +`_build_kv_write_meta` is called now, `targets` and `receiver_payload_bytes` +where `dispatch_task` calls their predecessors. Planner-internal caches +(peer overlap, mappers) follow the same ownership as `PeerRegistrar`'s caches +do now. + +### Framework-owned (unchanged responsibilities) + +- `WriteMeta` assembly from a `TransferPlan` plus bookkeeping + (`unique_rid`, `slice_id`, `bounce_dst_base`, `expected_transfers = len(targets.ranks)`). +- NIXL agent requests, bounce buffers (Python and C++ staging), aux transfer. +- Task lifecycle: `_LogicalOutcomes`, physical ownership, in-doubt polling, + retirement, cancellation, ZMQ messaging, pipelined chunk numbering. +- Gen-first ADP broadcast: the union of `targets(...)` over all peer DP ranks + stays in `Receiver.dispatch_task`; it is a protocol quirk, not a model one. +- Peer bookkeeping: peer `RankInfo` cache, `KVRegionExtractorV1` per peer, + aux transfer layout cache, sender endpoints. These move out of + `PeerRegistrar` into a framework `PeerRegistry`; `PeerRegistrar`'s + remaining content (policies, mappers, overlap, pool mapping) becomes the + internals of `DefaultTransferPlanner`. + +### Wire format: `RankInfo` + +`RankInfo` stays the envelope. Changes: + +- Add `planner_id: str` and `planner_blob: bytes`. +- Add `enable_attention_dp: bool` as a top-level topology field; it is + consumed by `tp_size_per_dp_group`, which is topology, not attention. +- Move `attention: AttentionInfo` into `DefaultTransferPlanner.describe()`. + Nothing outside `AttentionPolicy` reads it once `enable_attention_dp` is + lifted out. +- `page_table` stays top-level: the bounce layer and the region extractor + consume it, and the default planner needs it too. A custom planner + receives it like any other envelope field. +- `aux_meta` stays top-level (aux transfer is framework-owned). + +Framework check at `PeerRegistry.register` and `Receiver._get_sender_info`: +`peer_ri.planner_id != self.planner.planner_id` raises +`PeerIncompatibleError` before `check_peer` runs. This is the one +cross-model invariant the framework enforces: both sides run the same +planner class. Blob format compatibility across wheel versions is the +planner's responsibility, exactly as `AttentionInfo` is today. + +### Mount point + +The planner comes from the model class, not the KV cache manager: + +```python +class PretrainedModel/DecoderModelForCausalLM: + def build_transfer_planner(self, kv_cache_manager) -> TransferPlanner: + return DefaultTransferPlanner(kv_cache_manager) +``` + +A modeling_v2 target overrides `build_transfer_planner` in its `modeling.py`, +beside `REQUIRED_TRTLLM_OPS`, and keeps its planner in its own target +directory. Reason for the model over the cache manager: a plan is a function +of the parallel topology *and* the cache layout, and the target is the one +object that already owns both. The cache manager builds the page table and +knows nothing about which peer topologies the target supports. + +Plumbing: `create_kv_cache_transceiver` (`pyexecutor/_util.py:4232`) gains a +`planner` argument built from `model_engine.model`; +`KvCacheTransceiverV2.__init__` passes it into `TransferWorkerConfig`; +`TransferWorker` hands it to `Sender`, `Receiver` and `PeerRegistry`. + +## What moves where + +| From | To | +|------|----| +| `Sender._build_kv_write_meta` body (`transfer.py:1710-1836`) | `DefaultTransferPlanner.plan`; `_build_kv_write_meta` shrinks to plan -> `WriteMeta` | +| `Sender._pair_ordinals` | `DefaultTransferPlanner` (a custom planner may import it as a helper) | +| `Receiver._fanin_bounce_safe`, `_get_mamba_slot`, `mamba_receiver_payload_bytes` call | `DefaultTransferPlanner.targets(...).uniform_writers` / `receiver_payload_bytes` | +| `Receiver._get_sender_info` policy validation (`3381-3396`) | `planner.check_peer` | +| `PeerRegistrar` policies, `_kv_map_cache`, `_overlap_cache`, `_lg_pool_mapping_cache`, `get_pool_mapping`, `get_kv_map`, `get_peer_overlap`, `should_send_pool`, `_owns_tp_fan_in` | `DefaultTransferPlanner` (file `native/default_planner.py`) | +| `PeerRegistrar` peer `RankInfo` cache, extractor cache, aux layout cache, `should_send_aux` | framework `PeerRegistry` (stays in `native/peer.py`) | +| `mixers/attention`, `mixers/ssm` | unchanged, now imported only by `default_planner.py` | + +After the move, `native/transfer.py` imports nothing from `mixers/` and +mentions neither `CacheKind` nor `MapperKind`. That is the grep gate. + +## Non-goals + +- The C++ transceiver and the `UCX`/`MPI` backends. +- Aux (draft tokens, logits) transfer. +- Bounce internals and the pipelined-chunk protocol. +- Changing any mapper or policy behavior. Phase 1 is a pure extraction with + bit-identical transfer plans. +- Writing the first target-owned planner. This spec creates the hook; the + first custom planner is its own piece of work with its own ctx/gen target + pair. + +## Testing + +- Existing `tests/unittest/disaggregated/test_peer.py`, + `test_kda_mamba_transfer.py`, `test_kv_transfer.py` and `test_rank_info.py` + keep passing with `PeerRegistrar(...)` construction retargeted to + `DefaultTransferPlanner` / `PeerRegistry`. No behavior change means no + golden change. +- One new contract test (`tests/unittest/disaggregated/test_transfer_planner.py`): + a toy `TransferPlanner` (fixed ranks, fixed ptr arrays, `describe()` + returning a marker) driven through `Sender._build_kv_write_meta` and + `Receiver.dispatch_task` with a stub agent that records the `WriteMeta` it + is handed. Asserts the plan arrives unmodified and that a `planner_id` + mismatch fails the request with `PeerIncompatibleError` rather than the + worker. The stub agent is new and lives in the test module; no existing + fixture drives `Sender`/`Receiver` directly today. +- Grep gate on `native/transfer.py` (no `mixers`, `CacheKind`, `MapperKind`), + enforced in the same test module. +- Integration: the existing disagg L0 tests on the native backend are the + parity gate for Phase 1. + +## Phasing + +1. Extraction: `TransferPlanner` protocol, `DefaultTransferPlanner`, + `PeerRegistry` split, `RankInfo` field changes, model mount point, + plumbing. Zero behavior change. +2. First target-owned planner in a modeling_v2 target directory (separate + spec, picks the model and the ctx/gen target pair). From eb47a8be5943bdbfbc009e437645ab18bb859b3a Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 05:22:57 -0700 Subject: [PATCH 03/25] [None][doc] Implementation plan for model-owned disaggregated transfer planning Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../2026-10-10-disagg-transfer-planner.md | 1077 +++++++++++++++++ 1 file changed, 1077 insertions(+) create mode 100644 docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md diff --git a/docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md b/docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md new file mode 100644 index 000000000000..8f5fbe6ea419 --- /dev/null +++ b/docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md @@ -0,0 +1,1077 @@ +# Model-owned disaggregated KV transfer planning — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Put every "which bytes pair with which" decision of the transfer v2 sender/receiver behind a `TransferPlanner` interface that a model class can override, with the current shared logic becoming `DefaultTransferPlanner`, and `native/transfer.py` no longer knowing about attention, Mamba, `CacheKind` or `MapperKind`. + +**Architecture:** A new `TransferPlanner` ABC (`base/planner.py`) with `check_peer / targets / plan / receiver_payload_bytes / describe`. `DefaultTransferPlanner` (`native/default_planner.py`) receives the policy, mapper, overlap and pool-matching engine that lives in `PeerRegistrar` today plus the body of `Sender._build_kv_write_meta`, `Receiver._fanin_bounce_safe`, `Receiver._get_mamba_slot`. `PeerRegistrar` keeps only framework bookkeeping (peer `RankInfo` cache, extractors, aux layout cache, aux ownership) and holds the planner. `Sender`/`Receiver` call the planner through `self._planner`. The planner is produced by a factory the model class exposes (`build_transfer_planner`), plumbed `_util.py -> create_kv_cache_transceiver -> KvCacheTransceiverV2 -> TransferWorkerConfig`. Zero behavior change: every existing transfer plan is bit-identical. + +**Tech Stack:** Python 3.12, numpy, msgpack (RankInfo wire), pytest. No GPU for the unit tasks; the final task runs the GPU disagg gate on a GB300 node. + +**Spec:** `docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md` + +**Deviations from the spec (decided while planning, all in the direction of less churn):** +1. The framework bookkeeping class keeps its name `PeerRegistrar` (spec: `PeerRegistry`). ~10 test files mock `sender._registrar` / `receiver._registrar`; renaming buys nothing. +2. `RankInfo.attention` stays where it is and `DefaultTransferPlanner.describe()` returns `b""`. The spec moved `AttentionInfo` into the blob for cleanliness; that touches both policies, `tp_size_per_dp_group` and every RankInfo-building test for no functional gain. A custom planner does not need `attention`; it can leave it `None` and use `planner_blob`. Revisit only if a real planner needs it. +3. The model hook is a factory with the inputs the worker actually has: `build_transfer_planner(kv_cache_manager, self_rank_info, self_extractor) -> TransferPlanner`, not `build_transfer_planner(kv_cache_manager)`. `RankInfo` and the extractor are built by `TransferWorker` from the cache manager; the default planner consumes them, and duplicating that derivation in every planner would be wrong. +4. `PeerOverlap` moves with the engine into `default_planner.py` (it is the default planner's internal overlap record). Tests import it from there. +5. `TransferPlanner.forget_peer(peer_ri)` is added (default no-op) so `PeerRegistrar.unregister` can drop planner-side caches. It is cache hygiene the current `unregister` already does for `_kv_map_cache`/`_lg_pool_mapping_cache`. + +## Global Constraints + +- NVIDIA SPDX copyright header (2026) on every new file; bump the year on modified files whose header is older. +- `git commit -s` (DCO); no AI attribution, no co-authors. +- Zero behavior change in Phase 1: every existing unit test under `tests/unittest/disaggregated/` keeps passing with only construction/import changes; the GPU gate `TestModelingV2GptOss120bSm103Tp1::test_gsm8k_disagg` and `TestGPTOSS::test_kv_cache_v2_nixl_python` must pass. +- After the refactor `native/transfer.py` contains none of the strings `mixers`, `CacheKind`, `MapperKind`, `MambaPolicy`, `AttentionPolicy` (grep gate, enforced by a test). +- Both peers must run the same planner class: `planner_id` mismatch -> `PeerIncompatibleError`, failing that request/peer, never the worker. +- Python-only change. Run unit tests inside the container with `PYTHONPATH=` and main's venv first on `PATH` (see `/scratch/.../users/fredw/mv2-disagg/run_disagg_test.sh`); `srun --mpi=none` for anything that spawns MPI. + +## Review Focus + +1. **Empty plan from a non-owning rank.** A rank whose `plan()` returns three empty arrays (duplicate-head non-owner, PP stage with no overlap) must still produce a `WriteMeta` and complete the handshake, exactly like today's `_build_kv_write_meta` with no fragments. Pinned in Task 4 (`test_build_kv_write_meta_empty_plan_still_delivers`). +2. **`receiver_payload_bytes` returning `None`.** Receiver must disable bounce (not crash, not reserve a wrongly sized region) when the planner cannot size the extra payload. Pinned in Task 5 (`test_dispatch_task_payload_none_disables_bounce`). +3. **Gen-first ADP broadcast union.** With `ctx_dp_rank is None` the receiver unions `targets()` over all peer DP ranks and must set `uniform_writers=False` on the union so fan-in bounce is never attempted. Pinned by the existing ownership regression test that covers broadcast (adapted in Task 5) plus an assertion added there. +4. **planner_id mismatch on the receiver's first contact.** `_get_sender_info` must raise `PeerIncompatibleError` before `REGISTER_RANK_INFO` is sent, and cache the diagnosis. Pinned in Task 6 (`test_receiver_rejects_planner_id_mismatch`). +5. **Old peer wire format (no `planner_id`).** A `RankInfo` serialized without the two new fields must still deserialize (defaults `""`/`b""`) so a mixed-version pair fails with the explicit mismatch error rather than a msgpack `TypeError`. Pinned in Task 2 (`test_rank_info_from_bytes_without_planner_fields`). + +--- + +### Task 1: `TransferPlanner` protocol module + +**Files:** +- Create: `tensorrt_llm/_torch/disaggregation/base/planner.py` +- Modify: `tensorrt_llm/_torch/disaggregation/native/transfer.py:195-203` (replace the `PeerIncompatibleError` class body with an import) +- Test: `tests/unittest/disaggregated/test_transfer_planner.py` (new) + +**Interfaces:** +- Produces: + ```python + class PeerIncompatibleError(ValueError) + @dataclass class PeerTargets: ranks: list[int]; uniform_writers: bool + @dataclass class TransferPlan: src_ptrs: np.ndarray; dst_ptrs: np.ndarray; sizes: np.ndarray + @classmethod empty() -> TransferPlan + class TransferPlanner(ABC): + planner_id: str (class attribute, must be overridden) + def describe(self) -> bytes + def check_peer(self, self_ri, peer_ri) -> None # raises PeerIncompatibleError + def targets(self, self_ri, peer_ri, peer_dp_rank: int) -> PeerTargets + def plan(self, self_ri, peer_ri, src_block_ids_per_lg, dst_block_ids_per_lg, + token_range, prompt_len, dst_start_token) -> TransferPlan + def receiver_payload_bytes(self, self_ri, peer_ri, dst_block_ids_per_lg) -> Optional[int] + def forget_peer(self, peer_ri) -> None # default no-op + ``` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unittest/disaggregated/test_transfer_planner.py +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Contract tests for the model-owned transfer planner seam.""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.disaggregation.base.planner import ( + PeerIncompatibleError, + PeerTargets, + TransferPlan, + TransferPlanner, +) + + +def test_transfer_plan_empty_is_three_int64_arrays(): + plan = TransferPlan.empty() + for arr in (plan.src_ptrs, plan.dst_ptrs, plan.sizes): + assert arr.dtype == np.int64 and arr.size == 0 + + +def test_transfer_planner_is_abstract(): + with pytest.raises(TypeError): + TransferPlanner() # type: ignore[abstract] + + +def test_peer_incompatible_error_is_value_error(): + assert issubclass(PeerIncompatibleError, ValueError) + assert PeerTargets(ranks=[0], uniform_writers=True).ranks == [0] +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run (inside container, from `/tests/unittest`): `python -m pytest disaggregated/test_transfer_planner.py -q` +Expected: FAIL with `ModuleNotFoundError: tensorrt_llm._torch.disaggregation.base.planner` + +- [ ] **Step 3: Write the module** + +```python +# tensorrt_llm/_torch/disaggregation/base/planner.py +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The seam between the transfer framework and a model's KV-transfer plan. + +The framework (``native/transfer.py``) moves bytes and tracks transfers. Which +bytes of this rank pair with which bytes of a peer rank is a property of the +model's cache layout and the two parallel topologies, so a model class may +supply its own :class:`TransferPlanner`; ``native/default_planner.py`` is the +one every model gets unless it overrides ``build_transfer_planner``. + +Both sides of a transfer must run the same planner class (``planner_id``): +the framework enforces that and nothing else about the plan. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional + +import numpy as np + +if TYPE_CHECKING: + from tensorrt_llm._torch.disaggregation.base.backend import TokenRange + from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + + +class PeerIncompatibleError(ValueError): + """A peer cannot exchange KV with this rank. Fails the peer, never the worker.""" + + +@dataclass +class PeerTargets: + """Which peer ranks this rank exchanges with for one peer DP group.""" + + ranks: List[int] + # Every writer contributes an equal share of the receiver's region, so the + # receiver may split one fan-in bounce reservation ``total // len(ranks)``. + uniform_writers: bool = False + + +@dataclass +class TransferPlan: + """Device addresses to copy, pairwise. ``sizes[i]`` bytes from ``src_ptrs[i]`` to ``dst_ptrs[i]``. + + An empty plan is valid: this rank writes nothing to this peer but still + completes the handshake (a duplicate-head non-owner, a PP stage with no + overlapping layers). + """ + + src_ptrs: np.ndarray + dst_ptrs: np.ndarray + sizes: np.ndarray + + @classmethod + def empty(cls) -> "TransferPlan": + e = np.array([], dtype=np.int64) + return cls(src_ptrs=e, dst_ptrs=e.copy(), sizes=e.copy()) + + +class TransferPlanner(ABC): + """Model-owned KV transfer planning. One instance per rank, shared by Sender and Receiver. + + Threading is the framework's: ``plan`` runs where the sender builds write + metadata, ``targets``/``receiver_payload_bytes``/``check_peer`` run on the + receiver's dispatch path and both listeners. Keep per-peer caches keyed by + ``(peer_ri.instance_name, peer_ri.instance_rank)`` and drop them in + ``forget_peer``. + """ + + #: Both sides must agree. The framework compares it before ``check_peer``. + planner_id: str = "" + + def describe(self) -> bytes: + """Opaque self-description carried in ``RankInfo.planner_blob``; only this class decodes it.""" + return b"" + + @abstractmethod + def check_peer(self, self_ri: "RankInfo", peer_ri: "RankInfo") -> None: + """Raise :class:`PeerIncompatibleError` with a field-level reason, else return.""" + + @abstractmethod + def targets(self, self_ri: "RankInfo", peer_ri: "RankInfo", peer_dp_rank: int) -> PeerTargets: + """Peer flat ranks this rank exchanges with for the peer's DP group ``peer_dp_rank``.""" + + @abstractmethod + def plan( + self, + self_ri: "RankInfo", + peer_ri: "RankInfo", + src_block_ids_per_lg: List[np.ndarray], + dst_block_ids_per_lg: List[np.ndarray], + token_range: "TokenRange", + prompt_len: int, + dst_start_token: Optional[int], + ) -> TransferPlan: + """Sender side: the bytes this rank writes into ``peer_ri`` for one chunk.""" + + @abstractmethod + def receiver_payload_bytes( + self, self_ri: "RankInfo", peer_ri: "RankInfo", dst_block_ids_per_lg: List[np.ndarray] + ) -> Optional[int]: + """Receiver side: bytes beyond the paged blocks that ride the same write. + + ``None`` means the planner cannot size them; the receiver then disables + its bounce reservation for this request. + """ + + def forget_peer(self, peer_ri: "RankInfo") -> None: + """Drop any per-peer cache. Called by ``PeerRegistrar.unregister``.""" + return None +``` + +In `native/transfer.py` replace the class + +```python +class PeerIncompatibleError(ValueError): + ... +``` + +with + +```python +from tensorrt_llm._torch.disaggregation.base.planner import PeerIncompatibleError # noqa: F401 (re-exported) +``` + +placed with the other imports (keep the name importable from `transfer` because tests and `transceiver.py` import it from there — check with `grep -rn "PeerIncompatibleError" tensorrt_llm tests`). + +- [ ] **Step 4: Run tests** + +Run: `python -m pytest disaggregated/test_transfer_planner.py -q` → PASS (3). Then `python -c "from tensorrt_llm._torch.disaggregation.native.transfer import PeerIncompatibleError"` → ok. + +- [ ] **Step 5: Commit** + +```bash +git add tensorrt_llm/_torch/disaggregation/base/planner.py tensorrt_llm/_torch/disaggregation/native/transfer.py tests/unittest/disaggregated/test_transfer_planner.py +git commit -s -m "[None][refactor] Add the TransferPlanner seam for disaggregated KV transfer" +``` + +--- + +### Task 2: `RankInfo` carries `planner_id` / `planner_blob` + +**Files:** +- Modify: `tensorrt_llm/_torch/disaggregation/native/rank_info.py` (dataclass fields after `device_id`) +- Test: `tests/unittest/disaggregated/test_rank_info.py` (append) + +**Interfaces:** +- Produces: `RankInfo.planner_id: str = ""`, `RankInfo.planner_blob: bytes = b""`; round-trips through `to_bytes`/`from_bytes`; absent on the wire -> defaults. + +- [ ] **Step 1: Write the failing tests** (append to `test_rank_info.py`; reuse that file's existing RankInfo builder helper if there is one, else build a minimal `RankInfo` as below) + +```python +def _minimal_rank_info(**kw): + from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + base = dict( + instance_name="ctx", instance_rank=0, tp_size=1, tp_rank=0, pp_size=1, pp_rank=0, + layer_num_per_pp=[2], sender_endpoints=[], self_endpoint="", transfer_engine_info=b"", + ) + base.update(kw) + return RankInfo(**base) + + +def test_rank_info_planner_fields_round_trip(): + from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + ri = _minimal_rank_info(planner_id="toy", planner_blob=b"\x00\x01") + back = RankInfo.from_bytes(ri.to_bytes()) + assert back.planner_id == "toy" and back.planner_blob == b"\x00\x01" + + +def test_rank_info_from_bytes_without_planner_fields(): + import msgpack + from dataclasses import asdict + from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + data = asdict(_minimal_rank_info()) + for k in ("planner_id", "planner_blob", "attention", "aux_meta", "page_table"): + data.pop(k, None) + back = RankInfo.from_bytes(msgpack.packb(data)) + assert back.planner_id == "" and back.planner_blob == b"" +``` + +- [ ] **Step 2: Run** `python -m pytest disaggregated/test_rank_info.py -q -k planner` → FAIL (`unexpected keyword argument 'planner_id'`). + +- [ ] **Step 3: Implement** — in the `RankInfo` dataclass add after `device_id: int = 0`: + +```python + # Which TransferPlanner class produced this rank's plan, and that planner's + # opaque self-description. The framework compares planner_id across peers; + # only the planner decodes planner_blob. + planner_id: str = "" + planner_blob: bytes = b"" +``` + +`to_bytes` (uses `asdict`) and `from_bytes` (`cls(**unpacked)`) need no change; msgpack encodes `bytes` natively. + +- [ ] **Step 4: Run** `python -m pytest disaggregated/test_rank_info.py -q` → PASS. + +- [ ] **Step 5: Commit** `git commit -s -m "[None][refactor] Carry planner_id and planner_blob on RankInfo"` + +--- + +### Task 3: `DefaultTransferPlanner` owns the engine; `PeerRegistrar` keeps bookkeeping + +**Files:** +- Create: `tensorrt_llm/_torch/disaggregation/native/default_planner.py` +- Modify: `tensorrt_llm/_torch/disaggregation/native/peer.py` (becomes ~120 lines) +- Modify tests (construction only): `tests/unittest/disaggregated/test_peer.py`, `test_pool_matching.py`, `test_extractor.py:350-392`, `test_kda_mamba_transfer.py:62,289`, `test_bounce.py:229-262`, `tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py:33,1945` + +**Interfaces:** +- Produces: + ```python + # native/default_planner.py + LGPoolKey = Tuple[int, int] + @dataclass class PeerOverlap (unchanged fields) + class DefaultTransferPlanner(TransferPlanner): + planner_id = "trtllm.default" + def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1) + # engine (moved verbatim from PeerRegistrar): + get_pool_mapping(peer_ri) -> Dict[LGPoolKey, LGPoolKey] + get_kv_map(peer_ri, self_pool_key, peer_pool_key) -> RegionMapperBase + get_peer_overlap(peer_ri, peer_dp_rank) -> PeerOverlap + should_send_pool(peer_overlap, peer_ri, layer_group_id, pool_idx) -> bool + _owns_tp_fan_in, _find_overlap, _get_buffers_per_layer, _get_policy, _POLICY_CLASSES + peer_extractor(peer_ri) -> KVRegionExtractorV1 # planner-private cache + # moved from transfer.py: + @staticmethod _pair_ordinals(src_block_ids, dst_block_ids, dst_start_block=0) + @staticmethod _fanin_bounce_safe(overlap, peer_ri, receiver_page_table) -> bool + @staticmethod _get_mamba_slot(page_table, dst_block_ids_per_lg) -> Optional[int] + # protocol: + check_peer, targets, plan, receiver_payload_bytes, forget_peer + # native/peer.py + class PeerRegistrar: + def __init__(self, self_rank_info, self_extractor, planner: TransferPlanner) + planner (property), self_rank_info, self_extractor, register, unregister, + get_peer_rank_info, peer_extractor(peer_name, peer_rank), get_aux_transfer_layout, + cache_aux_transfer_layout, should_send_aux, _unique_key + ``` +- Consumes: Task 1 types. + +- [ ] **Step 1: Create `default_planner.py` by moving code.** Start the file with the NVIDIA header and these imports (superset of what moves): + +```python +from __future__ import annotations + +from collections import Counter +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np + +from tensorrt_llm import logger +from tensorrt_llm._torch.disaggregation.base.backend import TokenRange +from tensorrt_llm._torch.disaggregation.base.planner import ( + PeerIncompatibleError, + PeerTargets, + TransferPlan, + TransferPlanner, +) +from tensorrt_llm._torch.disaggregation.base.region import RegionMapperBase +from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( + MambaPolicy, + mamba_receiver_payload_bytes, +) +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.page import ( + CacheKind, + KVCachePageTable, + MapperKind, + PoolView, +) +from tensorrt_llm._torch.disaggregation.resource.utils import ( + get_layer_byte_ranges, + get_layer_to_layer_group, + get_pool_view_global_layer_ids, +) +``` + +Then move, **unchanged in body**, from `peer.py`: `LGPoolKey`, `PeerOverlap`, and from `PeerRegistrar`: `_POLICY_CLASSES`, `get_pool_mapping`, `_get_policy`, `get_kv_map`, `_get_buffers_per_layer`, `_find_overlap`, `get_peer_overlap`, `_owns_tp_fan_in`, `should_send_pool`. Rename every `"PeerRegistrar.` prefix inside error strings to `"DefaultTransferPlanner.`. The class header and constructor: + +```python +class DefaultTransferPlanner(TransferPlanner): + """The planner every model gets unless it overrides ``build_transfer_planner``. + + Attention (PAGED) and recurrent-state (STATE) layer groups are planned by + ``AttentionPolicy`` and ``MambaPolicy``; pools are matched by role and + layer overlap; TP/PP/CP overlap decides which peer ranks receive what. + """ + + planner_id = "trtllm.default" + + _POLICY_CLASSES = { + CacheKind.PAGED: AttentionPolicy, + CacheKind.STATE: MambaPolicy, + } + + def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1): + self._ri = self_rank_info + self._self_ext_cache = self_extractor + self._policies: Dict[CacheKind, Union[AttentionPolicy, MambaPolicy]] = {} + self._kv_map_cache: Dict[tuple, RegionMapperBase] = {} + self._overlap_cache: Dict[str, PeerOverlap] = {} + self._lg_pool_mapping_cache: Dict[str, Dict[LGPoolKey, LGPoolKey]] = {} + self._peer_ext_cache: Dict[str, KVRegionExtractorV1] = {} + + @staticmethod + def _unique_key(name: str, rank: int) -> str: + return name + str(rank) + + def peer_extractor(self, peer_ri: RankInfo) -> KVRegionExtractorV1: + key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) + ext = self._peer_ext_cache.get(key) + if ext is None: + ext = KVRegionExtractorV1(peer_ri.page_table) + self._peer_ext_cache[key] = ext + return ext + + def forget_peer(self, peer_ri: RankInfo) -> None: + key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) + self._peer_ext_cache.pop(key, None) + self._lg_pool_mapping_cache.pop(key, None) + for k in [k for k in self._kv_map_cache if k[0] == key]: + del self._kv_map_cache[k] +``` + +(`get_peer_overlap` caches under `_unique_key(instance_name, peer_dp_rank)`; leave that as is — it is per instance, not per rank.) + +- [ ] **Step 2: Add the moved statics.** Move `Sender._pair_ordinals` (transfer.py:1683-1708) verbatim as `DefaultTransferPlanner._pair_ordinals`. Move `Receiver._fanin_bounce_safe` (3134-3170) verbatim as a `@staticmethod` with the same signature. Replace `Receiver._get_mamba_slot` with: + +```python + @staticmethod + def _get_mamba_slot( + page_table: Optional[KVCachePageTable], dst_block_ids_per_lg: List[np.ndarray] + ) -> Optional[int]: + """The receiver slot of the first STATE layer group, or None when there is none.""" + if page_table is None: + return None + for lg_idx, lg in enumerate(page_table.layer_groups): + if lg.kind == CacheKind.STATE: + block_ids = dst_block_ids_per_lg[lg_idx] + if block_ids.size > 0: + return int(block_ids[0]) + return None +``` + +- [ ] **Step 3: Implement the protocol methods.** + +```python + # ---- TransferPlanner protocol ------------------------------------------------- + + def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: + if not self._get_policy(CacheKind.PAGED).check_peer_compatible(peer_ri): + raise PeerIncompatibleError( + "DefaultTransferPlanner.check_peer: attention geometry differs from the peer" + ) + self_page_table = self._self_ext_cache.page_table + try: + MambaPolicy.validate_peer_compatible( + self._ri, peer_ri, self_page_table, peer_ri.page_table + ) + AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) + except ValueError as e: + raise PeerIncompatibleError(str(e)) from e + + self_layers = sum(self._ri.layer_num_per_pp) + peer_layers = sum(peer_ri.layer_num_per_pp) + if self_layers != peer_layers: + # One side may hold speculative (MTP) layers the other lacks; pool + # matching transfers only the layers both hold. + logger.warning( + "DefaultTransferPlanner: layer count differs " + f"(local={self_layers}, peer={peer_layers}), allowing partial layer transfer." + ) + self._warn_nhd_head_mismatch(peer_ri) + + def _warn_nhd_head_mismatch(self, peer_ri: RankInfo) -> None: + # body = the NHD warning block of today's PeerRegistrar.register (peer.py:97-112), + # using self._self_ext_cache.page_table and self.get_pool_mapping(peer_ri) + ... + + def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> PeerTargets: + overlap = self.get_peer_overlap(peer_ri, peer_dp_rank) + uniform = ( + self_ri.cp_size == 1 + and peer_ri.cp_size == 1 + and self._fanin_bounce_safe(overlap, peer_ri, self._self_ext_cache.page_table) + ) + return PeerTargets(ranks=list(overlap.ranks), uniform_writers=uniform) + + def plan( + self, + self_ri: RankInfo, + peer_ri: RankInfo, + src_block_ids_per_lg: List[np.ndarray], + dst_block_ids_per_lg: List[np.ndarray], + token_range: TokenRange, + prompt_len: int, + dst_start_token: Optional[int], + ) -> TransferPlan: + is_partial = token_range.start > 0 or token_range.end < prompt_len + if is_partial and (self_ri.cp_size > 1 or peer_ri.cp_size > 1): + raise ValueError( + "enable_pipelined_transfer is not supported with context parallelism " + f"(sender cp_size={self_ri.cp_size}, receiver cp_size={peer_ri.cp_size})" + ) + targets = self.get_peer_overlap(peer_ri, peer_ri.dp_rank) + extractor = self._self_ext_cache + peer_extractor = self.peer_extractor(peer_ri) + pool_mapping = self.get_pool_mapping(peer_ri) + + src_frag_parts: list[np.ndarray] = [] + dst_frag_parts: list[np.ndarray] = [] + size_specs: list[tuple[int, int]] = [] + # <<< paste the per-pool loop of Sender._build_kv_write_meta (transfer.py:1750-1800) + # verbatim, with these renames: + # self._registrar.should_send_pool(...) -> self.should_send_pool(...) + # self._registrar.self_rank_info -> self_ri + # req_info.dst_start_token -> dst_start_token + # src_block_ids_per_groups / dst_block_ids_per_groups -> src_block_ids_per_lg / dst_block_ids_per_lg + # Sender._pair_ordinals -> self._pair_ordinals + # self._registrar.get_kv_map(...) -> self.get_kv_map(...) + if not src_frag_parts: + return TransferPlan.empty() + counts, values = zip(*size_specs) + return TransferPlan( + src_ptrs=np.concatenate(src_frag_parts), + dst_ptrs=np.concatenate(dst_frag_parts), + sizes=np.repeat(np.array(values, dtype=np.int64), counts), + ) + + def receiver_payload_bytes( + self, self_ri: RankInfo, peer_ri: RankInfo, dst_block_ids_per_lg: List[np.ndarray] + ) -> Optional[int]: + receiver_pt = self._self_ext_cache.page_table + dst_slot = self._get_mamba_slot(receiver_pt, dst_block_ids_per_lg) + if dst_slot is None: + return 0 + if peer_ri.page_table is None: + return None # cannot size the recurrent-state payload + return mamba_receiver_payload_bytes( + sender_page_table=peer_ri.page_table, + receiver_page_table=receiver_pt, + dst_slot=dst_slot, + ) +``` + +Note on `plan` vs today: today `_build_kv_write_meta` also records perf-timer marks and `expected_transfers`; those stay in the framework (Task 4). + +- [ ] **Step 4: Shrink `peer.py` to the bookkeeping class.** + +```python +class PeerRegistrar: + """Framework-side peer bookkeeping: who the peers are, their page-table extractors, + the aux transfer layouts, and aux ownership. Planning is the planner's.""" + + def __init__( + self, + self_rank_info: RankInfo, + self_extractor: KVRegionExtractorV1, + planner: TransferPlanner, + ): + self._ri = self_rank_info + self._self_ext_cache = self_extractor + self._planner = planner + self._peer_ri_cache: Dict[str, RankInfo] = {} + self._peer_ext_cache: Dict[str, KVRegionExtractorV1] = {} + self._aux_transfer_layout_cache: Dict[str, AuxTransferLayout] = {} + + @property + def planner(self) -> TransferPlanner: + return self._planner + + def register(self, peer_name: str, peer_rank: int, peer_ri: RankInfo): + if peer_ri.planner_id != self._planner.planner_id: + raise PeerIncompatibleError( + f"PeerRegistrar.register: peer {peer_name} (rank={peer_rank}) plans transfers " + f"with '{peer_ri.planner_id}', this rank with '{self._planner.planner_id}'; " + "both sides of a disaggregated deployment must run the same model planner" + ) + self._planner.check_peer(self._ri, peer_ri) + key = self._unique_key(peer_name, peer_rank) + self._peer_ri_cache[key] = peer_ri + self._aux_transfer_layout_cache.pop(key, None) + self._peer_ext_cache[key] = KVRegionExtractorV1(peer_ri.page_table) + + def unregister(self, peer_name: str, peer_rank: int): + key = self._unique_key(peer_name, peer_rank) + peer_ri = self._peer_ri_cache.pop(key, None) + self._peer_ext_cache.pop(key, None) + self._aux_transfer_layout_cache.pop(key, None) + if peer_ri is not None: + self._planner.forget_peer(peer_ri) +``` + +Keep verbatim: `peer_extractor`, `self_extractor`, `get_peer_rank_info`, `get_aux_transfer_layout`, `cache_aux_transfer_layout`, `self_rank_info`, `_unique_key`, `should_send_aux`. Delete everything else (it moved). Imports left: `Dict`, `AuxTransferLayout`, `PeerIncompatibleError`, `TransferPlanner`, `RankInfo`, `KVRegionExtractorV1`. The old `register` raised `ValueError("... is incompatible with local rank.")` when `_check_peer_compatible` returned False; `check_peer` now raises `PeerIncompatibleError` (a `ValueError`) — callers that catch `ValueError` keep working. + +- [ ] **Step 5: Adapt the direct tests (construction only).** + - `test_peer.py:38`: `from tensorrt_llm._torch.disaggregation.native.default_planner import PeerOverlap, DefaultTransferPlanner`; replace `PeerRegistrar(` with `DefaultTransferPlanner(` and `PeerRegistrar._get_buffers_per_layer` with `DefaultTransferPlanner._get_buffers_per_layer`. Any test that calls `reg.register(...)` and expects `ValueError` keeps passing (`PeerIncompatibleError` is one) — if it calls `register` it needs `PeerRegistrar(ri, ext, DefaultTransferPlanner(ri, ext))`; grep `register(` in the file to find those. + - `test_pool_matching.py:38,179,361`: same substitution. + - `test_extractor.py:350,365`: same. + - `test_kda_mamba_transfer.py:62,289`: `registrar = PeerRegistrar(ctx_ri, ext, DefaultTransferPlanner(ctx_ri, ext))` (this test exercises `register`). + - `test_deepseek_v4_cache_manager.py:33,1945`: it calls `register` then `get_pool_mapping`/`get_kv_map`: build `planner = DefaultTransferPlanner(gen_rank_info, ext)`, `PeerRegistrar(gen_rank_info, ext, planner).register("ctx", 0, ctx_rank_info)`, then call `planner.get_pool_mapping(...)` / `planner.get_kv_map(...)`. + - `test_bounce.py:239`: `safe = DefaultTransferPlanner._fanin_bounce_safe` (import from default_planner; drop the `tfr` importorskip if unused). + - Also move the `PeerOverlap` import anywhere else: `grep -rn "PeerOverlap" tests tensorrt_llm`. + +- [ ] **Step 6: Run** `python -m pytest disaggregated/test_peer.py disaggregated/test_pool_matching.py disaggregated/test_extractor.py disaggregated/test_bounce.py disaggregated/test_rank_info.py disaggregated/test_transfer_planner.py -q` → PASS. (`test_kda_mamba_transfer.py` and the deepseek_v4 test need a GPU; run them in Task 7.) `transfer.py` is still importing `PeerRegistrar` and calling the removed methods — it is not touched until Task 4, so do **not** run transfer-level tests yet. + +- [ ] **Step 7: Commit** `git commit -s -m "[None][refactor] Move transfer planning out of PeerRegistrar into DefaultTransferPlanner"` + +--- + +### Task 4: `Sender` plans through the planner + +**Files:** +- Modify: `tensorrt_llm/_torch/disaggregation/native/transfer.py` — `Sender.__init__`, `_build_kv_write_meta`, `_build_aux_write_meta:1841`, readiness sites `1088`, `2228`, `2242`; delete `_pair_ordinals`. +- Modify tests: `tests/unittest/disaggregated/test_chunked_transfer.py:150-232`, `test_aux_buffer_registration.py:55-80`, `test_unsubmitted_session.py:45-70`, and any other `sender._registrar.get_peer_overlap` site (`grep -rln "get_peer_overlap" tests/unittest/disaggregated`). +- Test: `tests/unittest/disaggregated/test_transfer_planner.py` (add `test_build_kv_write_meta_uses_planner_plan`, `test_build_kv_write_meta_empty_plan_still_delivers`) + +**Interfaces:** +- Consumes: `PeerRegistrar.planner`, `TransferPlanner.targets/plan`, `TransferPlan`. +- Produces: `Sender._planner` attribute; `Sender._expected_transfers(peer_ri) -> int`. + +- [ ] **Step 1: Write the failing tests** (append to `test_transfer_planner.py`) + +```python +from types import SimpleNamespace +from unittest.mock import Mock + +from tensorrt_llm._torch.disaggregation.base.backend import Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.native.transfer import KVSendTask, RecvReqInfo, Sender +from tensorrt_llm._torch.disaggregation.resource.page import CacheKind +from tensorrt_llm.llmapi.disagg_utils import DisaggregatedParams # adjust to the module test_chunked_transfer.py uses + + +class _ToyPlanner(TransferPlanner): + planner_id = "toy" + + def __init__(self, plan: TransferPlan, ranks=(0,)): + self._plan, self._ranks, self.calls = plan, list(ranks), [] + + def check_peer(self, self_ri, peer_ri): + return None + + def targets(self, self_ri, peer_ri, peer_dp_rank): + return PeerTargets(ranks=list(self._ranks), uniform_writers=True) + + def plan(self, self_ri, peer_ri, src, dst, token_range, prompt_len, dst_start_token): + self.calls.append((src, dst, token_range, prompt_len, dst_start_token)) + return self._plan + + def receiver_payload_bytes(self, self_ri, peer_ri, dst): + return 0 + + +def _sender_with(planner): + peer_ri = SimpleNamespace( + instance_name="gen", instance_rank=3, dp_rank=0, device_id=1, cp_size=1, + self_endpoint="tcp://gen:1", page_table=None, + ) + sender = Sender.__new__(Sender) + sender._registrar = Mock() + sender._registrar.self_rank_info = SimpleNamespace(cp_size=1, instance_name="ctx", instance_rank=0) + sender._registrar.get_peer_rank_info.return_value = peer_ri + sender._planner = planner + return sender + + +def _task_and_req(): + ids = np.array([5, 6], dtype=np.int64) + task = KVSendTask( + Chunk(block_ids_per_layer_groups=[ids], kind_per_layer_group=[CacheKind.PAGED], + token_range=TokenRange(0, 64), is_last=True), + DisaggregatedParams(), # mirror test_chunked_transfer._make_params() + slice_id=0, prompt_len=64, + ) + req = RecvReqInfo(sender_req_id=1, instance_name="gen", instance_rank=3, + block_ids_per_layer_groups=[np.array([9, 10], dtype=np.int64)], + unique_rid=1, dst_start_token=32, slice_id=0) + return task, req + + +def test_build_kv_write_meta_uses_planner_plan(): + plan = TransferPlan(src_ptrs=np.array([100, 200], dtype=np.int64), + dst_ptrs=np.array([300, 400], dtype=np.int64), + sizes=np.array([16, 16], dtype=np.int64)) + planner = _ToyPlanner(plan, ranks=(3, 4)) + sender = _sender_with(planner) + task, req = _task_and_req() + wm = sender._build_kv_write_meta(task, req) + assert wm.src_ptrs is plan.src_ptrs and wm.dst_ptrs is plan.dst_ptrs and wm.sizes is plan.sizes + assert wm.expected_transfers == 2 and task.expected_transfers == 2 + assert wm.dst_device_id == 1 and wm.peer_rank == 3 + (src, dst, tr, plen, dst_start), = planner.calls + assert src is task._chunk.block_ids_per_layer_groups and dst is req.block_ids_per_layer_groups + assert (tr.start, tr.end, plen, dst_start) == (0, 64, 64, 32) + + +def test_build_kv_write_meta_empty_plan_still_delivers(): + sender = _sender_with(_ToyPlanner(TransferPlan.empty())) + task, req = _task_and_req() + wm = sender._build_kv_write_meta(task, req) + assert wm.sizes.size == 0 and wm.expected_transfers == 1 and wm.is_last_slice is True +``` + +(Use the same `_make_params()`/`TokenRange` helpers as `test_chunked_transfer.py` — copy them if they are module-private.) + +- [ ] **Step 2: Run** `python -m pytest disaggregated/test_transfer_planner.py -q -k build_kv` → FAIL (`AttributeError: _planner` / old body calls `_registrar.get_peer_overlap`). + +- [ ] **Step 3: Implement in `transfer.py`.** + +`Sender.__init__`: after `self._registrar = peer_registrar` add `self._planner = peer_registrar.planner`. Add a helper right after `_get_session`: + +```python + def _expected_transfers(self, peer_ri: RankInfo) -> int: + """How many writes the peer's DP group expects from this instance.""" + return len( + self._planner.targets(self._registrar.self_rank_info, peer_ri, peer_ri.dp_rank).ranks + ) +``` + +Replace the four `len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks)` with `self._expected_transfers(peer_ri)` (lines ~1088, 1841, 2228, 2242). Replace `_build_kv_write_meta` with: + +```python + @nvtx_range("_build_kv_write_meta") + def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> WriteMeta: + peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) + self_ri = self._registrar.self_rank_info + timer = task._perf_timer + if timer: + timer.record_prepare_args_start(peer_ri.instance_rank) + expected_transfers = self._expected_transfers(peer_ri) + # Recorded on the task as well: the count on the write meta settles the task, while the + # handle has to answer how many writes are still out after one of them has already failed. + task.expected_transfers = expected_transfers + + plan = self._planner.plan( + self_ri, + peer_ri, + task._chunk.block_ids_per_layer_groups, + req_info.block_ids_per_layer_groups, + task._chunk.token_range, + task._prompt_len, + req_info.dst_start_token, + ) + + if timer: + timer.record_prepare_args_end(peer_ri.instance_rank) + timer.record_transfer_sizes( + peer_ri.instance_rank, int(plan.sizes.sum()), plan.dst_ptrs.size + ) + + return WriteMeta( + task=task, + src_ptrs=plan.src_ptrs, + dst_ptrs=plan.dst_ptrs, + sizes=plan.sizes, + dst_device_id=peer_ri.device_id, + expected_transfers=expected_transfers, + peer_name=peer_ri.instance_name + str(peer_ri.instance_rank), + peer_rank=peer_ri.instance_rank, + peer_endpoint=peer_ri.self_endpoint, + unique_rid=task._unique_rid, + slice_id=task.slice_id, + receiver_slice_id=req_info.slice_id if req_info.slice_id is not None else 0, + is_last_slice=task._chunk.is_last, + bounce_dst_base=req_info.bounce_dst_base, + ) +``` + +Delete `Sender._pair_ordinals`. + +- [ ] **Step 4: Adapt sender-side mock fixtures.** + - `test_chunked_transfer.py` `_make_projection_sender()`: keep `registrar` for `self_rank_info`/`get_peer_rank_info`; build the planner as the real engine with mocked internals so the ordinal-pairing tests still exercise `DefaultTransferPlanner.plan`: + ```python + planner = DefaultTransferPlanner.__new__(DefaultTransferPlanner) + planner._self_ext_cache = extractor + planner.peer_extractor = Mock(return_value=extractor) + planner.get_peer_overlap = Mock(return_value=SimpleNamespace(ranks=[0])) + planner.should_send_pool = Mock(return_value=True) + planner.get_pool_mapping = Mock(return_value={(0, 0): (0, 0), (1, 0): (1, 0)}) + planner.get_kv_map = Mock(return_value=mapper) + sender._planner = planner + ``` + and replace `registrar.self_rank_info = SimpleNamespace(cp_size=1, cp_rank=0)` with one that also has `dp_rank=0`; `registrar.get_peer_overlap`/`should_send_kv`/`get_pool_mapping`/`peer_extractor`/`get_kv_map` lines go away. `expected_transfers` now comes from `planner.targets` → since `targets` is the real method calling the mocked `get_peer_overlap` and `_fanin_bounce_safe(...)` with `self_rank_info.cp_size` — give `peer_ri` a `cp_size=1` (it has one) and the SimpleNamespace overlap `duplicate_head_factor=1, overlap_pp_size=1`. + - `test_aux_buffer_registration.py:69-71`: `sender = SimpleNamespace(_registrar=registrar, _planner=Mock(), _expected_transfers=lambda ri: 1)` — or set `sender._planner.targets.return_value = PeerTargets(ranks=[0])` and call `Sender._build_aux_write_meta(sender, ...)`; since `_expected_transfers` is a bound method on `Sender`, pass it explicitly as above. + - `test_unsubmitted_session.py:63-64` and every other `sender._registrar.get_peer_overlap.return_value = SimpleNamespace(ranks=[0])`: replace with `sender._planner = Mock(); sender._planner.targets.return_value = PeerTargets(ranks=[0])`. + - Find all: `grep -rn "get_peer_overlap\|should_send_kv\|get_pool_mapping\|get_kv_map" tests/unittest/disaggregated | grep -v test_peer.py | grep -v test_pool_matching.py`. + +- [ ] **Step 5: Run** `python -m pytest disaggregated/test_transfer_planner.py disaggregated/test_chunked_transfer.py disaggregated/test_aux_buffer_registration.py disaggregated/test_unsubmitted_session.py -q` → PASS. + +- [ ] **Step 6: Commit** `git commit -s -m "[None][refactor] Sender builds KV write metadata from the planner's TransferPlan"` + +--- + +### Task 5: `Receiver` dispatches through the planner; `transfer.py` sheds model knowledge + +**Files:** +- Modify: `tensorrt_llm/_torch/disaggregation/native/transfer.py` — imports (lines 65-70, 81), `Receiver.__init__`, delete `_fanin_bounce_safe` and `_get_mamba_slot`, rewrite the middle of `dispatch_task` (3215-3275), `_get_sender_info` (3381-3403). +- Modify tests: `tests/unittest/disaggregated/test_transfer_ownership_regressions.py` (3 receiver fixtures, lines ~602, ~699, ~872), any other `receiver._registrar.get_peer_overlap` site. +- Test: `test_transfer_planner.py` (add `test_transfer_module_has_no_model_knowledge`, `test_dispatch_task_payload_none_disables_bounce`) + +**Interfaces:** +- Consumes: `PeerTargets.uniform_writers`, `TransferPlanner.receiver_payload_bytes/check_peer`, `RankInfo.planner_id`. +- Produces: `Receiver._planner`. + +- [ ] **Step 1: Write the failing tests** + +```python +import inspect +import tensorrt_llm._torch.disaggregation.native.transfer as transfer_mod + + +def test_transfer_module_has_no_model_knowledge(): + src = inspect.getsource(transfer_mod) + for forbidden in ("mixers", "CacheKind", "MapperKind", "MambaPolicy", "AttentionPolicy"): + assert forbidden not in src, f"native/transfer.py mentions {forbidden}" +``` + +For `test_dispatch_task_payload_none_disables_bounce`, copy the smallest receiver fixture from `test_transfer_ownership_regressions.py` (the ~860-900 one: single writer, `_enforce_physical_ownership=False`), give it `receiver._planner = Mock()` with `targets.return_value = PeerTargets(ranks=[0], uniform_writers=True)` and `receiver_payload_bytes.return_value = None`, a `receiver._bounce = Mock()`, and assert after `receiver.dispatch_task(task)` that `receiver._bounce.reserve.assert_not_called()` and `receiver._request_sender_data` was called once. + +- [ ] **Step 2: Run** `python -m pytest disaggregated/test_transfer_planner.py -q -k "no_model_knowledge or payload_none"` → FAIL. + +- [ ] **Step 3: Implement.** + +Imports: delete `AttentionPolicy`, `MambaPolicy`, `mamba_receiver_payload_bytes`, `CacheKind`, `MapperKind` imports (keep `KVCachePageTable` only if still referenced — grep after the edit; add `from tensorrt_llm._torch.disaggregation.base.planner import PeerIncompatibleError, PeerTargets`). Delete `Receiver._fanin_bounce_safe` and `Receiver._get_mamba_slot`. + +`Receiver.__init__`: after `self._registrar = peer_registrar` add `self._planner = peer_registrar.planner`. + +`dispatch_task`, replace from `if sender_dp_rank is not None:` (first occurrence, ~3215) through the `bounced = ...` statement with: + +```python + self_ri = self._registrar.self_rank_info + if sender_dp_rank is not None: + # Normal path: ctx_dp_rank is known, send to overlapping ranks. + peer_targets = self._planner.targets(self_ri, peer_infos, sender_dp_rank) + dp0_targets = peer_targets + else: + # Gen-first with ADP: ctx_dp_rank unknown — broadcast REQUEST_DATA + # to ALL ctx sender ranks so every DP group receives it. The real + # writer count exceeds the per-group count, so the union is never + # a uniform writer set. + dp0_targets = self._planner.targets(self_ri, peer_infos, 0) + all_ranks_set: set[int] = set(dp0_targets.ranks) + for dp in range(1, peer_infos.dp_size): + all_ranks_set.update(self._planner.targets(self_ri, peer_infos, dp).ranks) + all_ranks = list(all_ranks_set) + logger.debug( + f"Receiver.dispatch_task: ADP broadcast path, dp_size={peer_infos.dp_size}, " + f"all_ranks={all_ranks}" + ) + peer_targets = PeerTargets(ranks=all_ranks, uniform_writers=False) + + # expected_transfers reflects one DP group even under broadcast: only + # one group will actually process the context request. + task.expected_transfers = len(dp0_targets.ranks) + # TP fan-in splits ONE region equally, so allow it only when the planner says every + # writer contributes the same share; never under ADP broadcast. + allow_bounce = task.expected_transfers == 1 or ( + sender_dp_rank is not None and peer_targets.uniform_writers + ) + # Bytes beyond the paged blocks that ride the same coalesced write + # (recurrent state today) size the bounce region too; None = unknown. + extra_bytes = self._planner.receiver_payload_bytes( + self_ri, peer_infos, receiver_req.block_ids_per_layer_groups + ) + if extra_bytes is None: + allow_bounce = False + extra_bytes = 0 + bounced = allow_bounce and self._bounce.reserve( + receiver_req, task.expected_transfers, extra_bytes=extra_bytes + ) +``` + +and rename every later `peer_overlap.ranks` in the method to `peer_targets.ranks`. + +`_get_sender_info`: replace the `try: MambaPolicy.validate_peer_compatible(...) ... except ValueError as e:` block with: + +```python + try: + if sender_info.planner_id != self._planner.planner_id: + raise PeerIncompatibleError( + f"context peer plans transfers with '{sender_info.planner_id}', " + f"this rank with '{self._planner.planner_id}'" + ) + self._planner.check_peer(self._registrar.self_rank_info, sender_info) + except PeerIncompatibleError as e: + msg = ( + f"context peer at '{info_endpoint}' is incompatible: {e} " + "(cached: all further requests to this context peer fail " + "fast; restart this generation worker to re-validate, " + "e.g. after redeploying a compatible server on the same " + "endpoint)" + ) + self._incompatible_peers[info_endpoint] = msg + raise PeerIncompatibleError(msg) from e +``` + +(Keep the explanatory comment above it, reworded: the sender-side gate is `PeerRegistrar.register` in the listener thread.) + +- [ ] **Step 4: Adapt receiver fixtures.** In each `receiver._registrar = SimpleNamespace(get_peer_overlap=..., self_extractor=..., self_rank_info=...)`: drop `get_peer_overlap`/`self_extractor`, keep `self_rank_info`, and add `receiver._planner = SimpleNamespace(targets=Mock(side_effect=lambda _s, _p, dp: PeerTargets(ranks=overlaps[dp].ranks, uniform_writers=True)), receiver_payload_bytes=Mock(return_value=0))` (single-overlap fixtures: `Mock(return_value=PeerTargets(ranks=[...], uniform_writers=True))`). In the broadcast fixture (`overlaps = {0: ..., 1: ...}`) add the assertion that `receiver._bounce.reserve` was called with `expected_transfers == 2` only if it already asserts bounce behavior; otherwise assert `allow_bounce` indirectly via `reserve.assert_not_called()` when `ctx_dp_rank is None` (Review Focus 3). + +- [ ] **Step 5: Run** `python -m pytest disaggregated -q -x --deselect disaggregated/test_kv_transfer.py --deselect disaggregated/test_kv_transfer_mp.py --ignore=disaggregated/test_kda_mamba_transfer.py` (GPU files excluded) → PASS, including `test_transfer_module_has_no_model_knowledge`. + +- [ ] **Step 6: Commit** `git commit -s -m "[None][refactor] Receiver dispatches through the planner; transfer.py no longer knows cache kinds"` + +--- + +### Task 6: Planner factory plumbing and the model mount point + +**Files:** +- Modify: `tensorrt_llm/_torch/disaggregation/native/transfer.py` — `TransferWorkerConfig` (+`planner_factory`), `TransferWorker._setup_peer_infrastructure` +- Modify: `tensorrt_llm/_torch/disaggregation/transceiver.py:237-300` (`KvCacheTransceiverV2.__init__` +`planner_factory`) +- Modify: `tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py:144-220` (`create_kv_cache_transceiver` +`model=None`) +- Modify: `tensorrt_llm/_torch/pyexecutor/_util.py:4232` (pass `model=model_engine.model`) +- Modify: `tensorrt_llm/_torch/models/modeling_utils.py:486+` (`DecoderModelForCausalLM.build_transfer_planner`) +- Test: `test_transfer_planner.py` (add `test_registrar_rejects_planner_id_mismatch`, `test_receiver_rejects_planner_id_mismatch`, `test_worker_uses_model_planner_factory`) + +**Interfaces:** +- Produces: + ```python + TransferPlannerFactory = Callable[[KVCacheManager, RankInfo, KVRegionExtractorV1], TransferPlanner] + TransferWorkerConfig.planner_factory: Optional[TransferPlannerFactory] = None + KvCacheTransceiverV2(mapping, dist, kv_cache_manager, cache_transceiver_config, planner_factory=None) + create_kv_cache_transceiver(..., mamba_cache_manager=None, model=None) + DecoderModelForCausalLM.build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor) -> TransferPlanner + ``` + +- [ ] **Step 1: Write the failing tests** + +```python +def test_registrar_rejects_planner_id_mismatch(): + from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar + self_ri = _minimal_rank_info(planner_id="toy") # reuse helper from test_rank_info or inline + peer_ri = _minimal_rank_info(instance_name="gen", planner_id="trtllm.default") + reg = PeerRegistrar(self_ri, Mock(), _ToyPlanner(TransferPlan.empty())) + with pytest.raises(PeerIncompatibleError, match="toy"): + reg.register("gen", 0, peer_ri) + assert reg._peer_ri_cache == {} + + +def test_receiver_rejects_planner_id_mismatch(monkeypatch): + from tensorrt_llm._torch.disaggregation.native.transfer import Receiver + recv = Receiver.__new__(Receiver) + recv._incompatible_peers = {} + recv._sender_ep_instance_map = {} + recv._registrar = SimpleNamespace(self_rank_info=_minimal_rank_info(planner_id="toy")) + recv._planner = _ToyPlanner(TransferPlan.empty()) + recv._should_register_peer = lambda params: True + sender_info = _minimal_rank_info(instance_name="ctx", planner_id="trtllm.default") + fake_messenger = Mock(); fake_messenger.receive.return_value = [sender_info.to_bytes()] + monkeypatch.setattr(transfer_mod, "ZMQMessenger", Mock(return_value=fake_messenger)) + recv._get_or_connect_dealer = Mock() + params = SimpleNamespace(ctx_info_endpoint="tcp://ctx:9", ctx_dp_rank=0) # match _extract_info_endpoint + with pytest.raises(PeerIncompatibleError, match="trtllm.default"): + recv._get_sender_info(params) + recv._get_or_connect_dealer.assert_not_called() # REGISTER_RANK_INFO never sent + assert "tcp://ctx:9" in recv._incompatible_peers # cached + + +def test_worker_uses_model_planner_factory(): + from tensorrt_llm._torch.disaggregation.native.transfer import TransferWorker + worker = TransferWorker.__new__(TransferWorker) + worker._rank_info = _minimal_rank_info() + made = _ToyPlanner(TransferPlan.empty()) + seen = {} + def factory(kvm, ri, ext): + seen["args"] = (kvm, ri, ext); return made + worker._config = SimpleNamespace(planner_factory=factory) + kvm = Mock(); kvm.mapping.rank = 1 + monkeypatch_ext = Mock() + with pytest.MonkeyPatch.context() as mp: + mp.setattr(transfer_mod, "KVRegionExtractorV1", Mock(return_value=monkeypatch_ext)) + mp.setattr(transfer_mod, "RankInfoServer", Mock()) + worker._setup_peer_infrastructure(kvm) + assert worker._peer_registrar.planner is made + assert seen["args"] == (kvm, worker._rank_info, monkeypatch_ext) + assert worker._rank_info.planner_id == "toy" and worker._rank_info.planner_blob == b"" +``` + +- [ ] **Step 2: Run** → FAIL (`TransferWorkerConfig` has no `planner_factory`; `_setup_peer_infrastructure` builds `PeerRegistrar(ri, ext)` with two args). + +- [ ] **Step 3: Implement.** + +`transfer.py`: + +```python +TransferPlannerFactory = Callable[[KVCacheManager, RankInfo, KVRegionExtractorV1], TransferPlanner] + +@dataclass +class TransferWorkerConfig: + ... + # Model-owned planner. None -> DefaultTransferPlanner. + planner_factory: Optional[TransferPlannerFactory] = None +``` + +```python + def _setup_peer_infrastructure(self, kvm: KVCacheManager): + self._rank_info_server = RankInfoServer(self._rank_info) if kvm.mapping.rank == 0 else None + self._kv_extractor = KVRegionExtractorV1(kvm) + factory = self._config.planner_factory or ( + lambda _kvm, ri, ext: DefaultTransferPlanner(ri, ext) + ) + self._planner = factory(kvm, self._rank_info, self._kv_extractor) + # The envelope carries who planned and what they want the peer's planner to know. + self._rank_info.planner_id = self._planner.planner_id + self._rank_info.planner_blob = self._planner.describe() + self._peer_registrar = PeerRegistrar(self._rank_info, self._kv_extractor, self._planner) +``` + +(`RankInfoServer` holds the object reference and serializes on request, as it already does for `self_endpoint`/`transfer_engine_info` set later.) Import `DefaultTransferPlanner` from `native.default_planner` and `TransferPlanner` from `base.planner`. + +`transceiver.py`: `KvCacheTransceiverV2.__init__(..., cache_transceiver_config, planner_factory=None)` → `TransferWorkerConfig(..., planner_factory=planner_factory)`. + +`kv_cache_transceiver.py`: `create_kv_cache_transceiver(..., mamba_cache_manager=None, model=None)`; in the Python branch: + +```python + planner_factory = getattr(model, "build_transfer_planner", None) + return KvCacheTransceiverV2(mapping, dist, kv_cache_manager, + cache_transceiver_config, + planner_factory=planner_factory) +``` + +`_util.py:4232`: add `model=model_engine.model`. + +`modeling_utils.py`, in `DecoderModelForCausalLM`: + +```python + def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): + """The disaggregated KV transfer plan for this model on this rank. + + A model whose cache layout the default planner does not understand + returns its own :class:`TransferPlanner` here; both the context and + the generation deployment of that model then run the same planner. + """ + from tensorrt_llm._torch.disaggregation.native.default_planner import ( + DefaultTransferPlanner, + ) + + return DefaultTransferPlanner(self_rank_info, self_extractor) +``` + +(Local import: `modeling_utils` must not import the disaggregation package at module load.) + +- [ ] **Step 4: Run** `python -m pytest disaggregated -q --ignore=disaggregated/test_kv_transfer.py --ignore=disaggregated/test_kv_transfer_mp.py --ignore=disaggregated/test_kda_mamba_transfer.py` → PASS; `python -m pytest llmapi/test_llm_args.py -q -k transceiver` and `api_stability` are unaffected (no LLM-arg change) — run `python -m pytest api_stability -q` once to be sure. + +- [ ] **Step 5: Commit** `git commit -s -m "[None][feat] Let a model class supply its disaggregated TransferPlanner"` + +--- + +### Task 7: GPU parity gate + +**Files:** none (verification only). + +- [ ] **Step 1: Holder job.** `cd /scratch/.../users/fredw/mv2-disagg && sbatch --parsable --qos=interactive alloc.sbatch` (writes `alloc.jobid`); wait for RUNNING. + +- [ ] **Step 2: GPU unit tests** via `in_alloc.sh` with a runner that sets the same env as `run_disagg_test.sh` but `cd tests/unittest`: + `python -m pytest disaggregated/test_kv_transfer.py disaggregated/test_kda_mamba_transfer.py disaggregated/test_kv_transfer_mp.py _torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py -q -x -p no:cacheprovider` → PASS (these exercise `TransferWorker` end-to-end with the default planner, including Mamba/KDA and the DeepSeek-V4 indexer pools). + +- [ ] **Step 3: End-to-end gate**: `./in_alloc.sh run_disagg_test.sh "accuracy/test_modeling_v2_gpt_oss.py::TestModelingV2GptOss120bSm103Tp1::test_gsm8k_disagg"` → PASSED, GSM8K within tolerance of 90.3 (baseline on this branch before the refactor: 89.84). Then `"accuracy/test_disaggregated_serving.py::TestGPTOSS::test_kv_cache_v2_nixl_python"` (stock model, both cache-manager variants; needs 4 GPUs) → PASSED. Grep both logs for `planner_id` errors (none expected). + +- [ ] **Step 4: Release the node** (`scancel $(cat alloc.jobid)`), record results in the PR description draft (not in code comments). + +--- + +## Self-review notes + +- Spec coverage: protocol (T1), wire fields (T2), DefaultTransferPlanner + registrar split (T3), sender (T4), receiver + grep gate (T5), factory/mount/plumbing + planner_id test (T6), parity (T7). Spec items consciously not done are listed under *Deviations* at the top. +- Names used consistently: `PeerTargets.ranks/uniform_writers`, `TransferPlan.src_ptrs/dst_ptrs/sizes`, `TransferPlanner.check_peer/targets/plan/receiver_payload_bytes/describe/forget_peer/planner_id`, `PeerRegistrar(self_rank_info, self_extractor, planner)` + `.planner`, `DefaultTransferPlanner(self_rank_info, self_extractor)`, `TransferWorkerConfig.planner_factory`, `build_transfer_planner(kv_cache_manager, self_rank_info, self_extractor)`. +- Review Focus items each have a pinned test: 1→T4, 2→T5, 3→T5 (broadcast fixture assertion), 4→T6, 5→T2. From 9c4d45f0577da374b5a5e9cf9202c7b6b71371d7 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 05:28:27 -0700 Subject: [PATCH 04/25] [None][refactor] Add the TransferPlanner seam for disaggregated KV transfer Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../_torch/disaggregation/base/planner.py | 126 ++++++++++++++++++ .../_torch/disaggregation/native/transfer.py | 13 +- .../disaggregated/test_transfer_planner.py | 41 ++++++ 3 files changed, 170 insertions(+), 10 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/base/planner.py create mode 100644 tests/unittest/disaggregated/test_transfer_planner.py diff --git a/tensorrt_llm/_torch/disaggregation/base/planner.py b/tensorrt_llm/_torch/disaggregation/base/planner.py new file mode 100644 index 000000000000..b820a6fac5ce --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/base/planner.py @@ -0,0 +1,126 @@ +# 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 seam between the transfer framework and a model's KV-transfer plan. + +The framework (``native/transfer.py``) moves bytes and tracks transfers. Which +bytes of this rank pair with which bytes of a peer rank is a property of the +model's cache layout and the two parallel topologies, so a model class may +supply its own :class:`TransferPlanner`; ``native/default_planner.py`` is the +one every model gets unless it overrides ``build_transfer_planner``. + +Both sides of a transfer must run the same planner class (``planner_id``): +the framework enforces that and nothing else about the plan. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional + +import numpy as np + +if TYPE_CHECKING: + from tensorrt_llm._torch.disaggregation.base.backend import TokenRange + from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + + +class PeerIncompatibleError(ValueError): + """A peer cannot exchange KV with this rank. Fails the peer, never the worker.""" + + +@dataclass +class PeerTargets: + """Which peer ranks this rank exchanges with for one peer DP group.""" + + ranks: List[int] + # Every writer contributes an equal share of the receiver's region, so the + # receiver may split one fan-in bounce reservation ``total // len(ranks)``. + uniform_writers: bool = False + + +@dataclass +class TransferPlan: + """Device addresses to copy, pairwise: ``sizes[i]`` bytes from ``src_ptrs[i]`` to ``dst_ptrs[i]``. + + An empty plan is valid: this rank writes nothing to this peer but still + completes the handshake (a duplicate-head non-owner, a PP stage with no + overlapping layers). + """ + + src_ptrs: np.ndarray + dst_ptrs: np.ndarray + sizes: np.ndarray + + @classmethod + def empty(cls) -> "TransferPlan": + return cls( + src_ptrs=np.array([], dtype=np.int64), + dst_ptrs=np.array([], dtype=np.int64), + sizes=np.array([], dtype=np.int64), + ) + + +class TransferPlanner(ABC): + """Model-owned KV transfer planning. One instance per rank, shared by Sender and Receiver. + + Threading is the framework's: ``plan`` runs where the sender builds write + metadata; ``targets``, ``receiver_payload_bytes`` and ``check_peer`` run on + the receiver's dispatch path and on both listeners. Keep per-peer caches + keyed by ``(peer_ri.instance_name, peer_ri.instance_rank)`` and drop them + in ``forget_peer``. + """ + + #: Both sides must agree. The framework compares it before ``check_peer``. + planner_id: str = "" + + def describe(self) -> bytes: + """Opaque self-description carried in ``RankInfo.planner_blob``; only this class decodes it.""" + return b"" + + @abstractmethod + def check_peer(self, self_ri: "RankInfo", peer_ri: "RankInfo") -> None: + """Raise :class:`PeerIncompatibleError` with a field-level reason, else return.""" + + @abstractmethod + def targets(self, self_ri: "RankInfo", peer_ri: "RankInfo", peer_dp_rank: int) -> PeerTargets: + """Peer flat ranks this rank exchanges with for the peer's DP group ``peer_dp_rank``.""" + + @abstractmethod + def plan( + self, + self_ri: "RankInfo", + peer_ri: "RankInfo", + src_block_ids_per_lg: List[np.ndarray], + dst_block_ids_per_lg: List[np.ndarray], + token_range: "TokenRange", + prompt_len: int, + dst_start_token: Optional[int], + ) -> TransferPlan: + """Sender side: the bytes this rank writes into ``peer_ri`` for one chunk.""" + + @abstractmethod + def receiver_payload_bytes( + self, self_ri: "RankInfo", peer_ri: "RankInfo", dst_block_ids_per_lg: List[np.ndarray] + ) -> Optional[int]: + """Receiver side: bytes beyond the paged blocks that ride the same write. + + ``None`` means the planner cannot size them; the receiver then disables + its bounce reservation for this request. + """ + + def forget_peer(self, peer_ri: "RankInfo") -> None: + """Drop any per-peer cache. Called by ``PeerRegistrar.unregister``.""" + return None diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 56c61cf1ab59..d7b6f3ceec0e 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -68,6 +68,9 @@ mamba_receiver_payload_bytes, ) from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar +from tensorrt_llm._torch.disaggregation.base.planner import ( # noqa: F401 (re-exported) + PeerIncompatibleError, +) from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.native.retirement import ( @@ -192,16 +195,6 @@ class MessageType: SESSION_QUIESCED = b"SESSION_QUIESCED" -class PeerIncompatibleError(ValueError): - """A context peer failed the KV/recurrent-state layout compatibility check. - - Subclasses ValueError so existing ``except ValueError`` handlers still - match, while letting dispatch_task catch peer incompatibility narrowly and - fail only the requests targeting that peer instead of crashing the - executor loop. - """ - - class TaskStatus(Enum): INIT = "INIT" TRANSFERRING = "TRANSFERRING" diff --git a/tests/unittest/disaggregated/test_transfer_planner.py b/tests/unittest/disaggregated/test_transfer_planner.py new file mode 100644 index 000000000000..39e2b0211b31 --- /dev/null +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -0,0 +1,41 @@ +# 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. +"""Contract tests for the model-owned transfer planner seam.""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.disaggregation.base.planner import ( + PeerIncompatibleError, + PeerTargets, + TransferPlan, + TransferPlanner, +) + + +def test_transfer_plan_empty_is_three_int64_arrays(): + plan = TransferPlan.empty() + for arr in (plan.src_ptrs, plan.dst_ptrs, plan.sizes): + assert arr.dtype == np.int64 and arr.size == 0 + + +def test_transfer_planner_is_abstract(): + with pytest.raises(TypeError): + TransferPlanner() # type: ignore[abstract] + + +def test_peer_incompatible_error_is_value_error(): + assert issubclass(PeerIncompatibleError, ValueError) + assert PeerTargets(ranks=[0], uniform_writers=True).ranks == [0] From 6ba0a9255a7688161c25374d301bac543030d89f Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 05:32:24 -0700 Subject: [PATCH 05/25] [None][refactor] Carry planner_id and planner_blob on RankInfo Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../_torch/disaggregation/native/rank_info.py | 5 +++ .../unittest/disaggregated/test_rank_info.py | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 87731d303204..b786ca97273a 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -44,6 +44,11 @@ class RankInfo: cp_size: int = 1 cp_rank: int = 0 device_id: int = 0 + # Which TransferPlanner class produced this rank's plan, and that planner's + # opaque self-description. The framework compares planner_id across peers; + # only the planner decodes planner_blob. + planner_id: str = "" + planner_blob: bytes = b"" attention: Optional[AttentionInfo] = None aux_meta: Optional[AuxBufferMeta] = None diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index a1d25e01c3f3..50fd1c0c27e4 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -203,3 +203,41 @@ def test_rank_info_represents_cache_element_bytes( restored = RankInfo.from_bytes(rank_info.to_bytes()) assert restored.attention.element_bytes == expected_element_bytes assert isinstance(restored.attention.element_bytes, expected_type) + + +def _minimal_rank_info(**overrides) -> RankInfo: + fields = dict( + instance_name="ctx", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + self_endpoint="", + transfer_engine_info=b"", + ) + fields.update(overrides) + return RankInfo(**fields) + + +def test_rank_info_planner_fields_round_trip(): + ri = _minimal_rank_info(planner_id="toy", planner_blob=b"\x00\x01") + back = RankInfo.from_bytes(ri.to_bytes()) + assert back.planner_id == "toy" + assert back.planner_blob == b"\x00\x01" + + +def test_rank_info_from_bytes_without_planner_fields(): + """A peer serialized before planner_id existed still deserializes, so the + mismatch surfaces as the explicit PeerIncompatibleError and not a TypeError.""" + import msgpack + from dataclasses import asdict + + data = asdict(_minimal_rank_info()) + for key in ("planner_id", "planner_blob", "attention", "aux_meta", "page_table"): + data.pop(key, None) + back = RankInfo.from_bytes(msgpack.packb(data)) + assert back.planner_id == "" + assert back.planner_blob == b"" From 9153ea9c5685be3d4cde27779187493003a0bc67 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 05:41:39 -0700 Subject: [PATCH 06/25] [None][refactor] Move transfer planning out of PeerRegistrar into DefaultTransferPlanner PeerRegistrar keeps peer bookkeeping (rank infos, extractors, aux layouts, aux ownership) and gates registration on the planner: same planner_id on both sides, then the planner's own compatibility check. Pool matching, mapper construction, topology overlap, send ownership, ordinal pairing, the fan-in uniformity gate and recurrent-state payload sizing move to DefaultTransferPlanner behind the TransferPlanner protocol. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../disaggregation/native/default_planner.py | 802 ++++++++++++++++++ .../_torch/disaggregation/native/peer.py | 554 +----------- .../_torch/disaggregation/native/transfer.py | 3 +- .../test_deepseek_v4_cache_manager.py | 10 +- tests/unittest/disaggregated/test_bounce.py | 4 +- .../unittest/disaggregated/test_extractor.py | 10 +- .../disaggregated/test_kda_mamba_transfer.py | 5 +- tests/unittest/disaggregated/test_peer.py | 69 +- .../disaggregated/test_pool_matching.py | 8 +- 9 files changed, 901 insertions(+), 564 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/native/default_planner.py diff --git a/tensorrt_llm/_torch/disaggregation/native/default_planner.py b/tensorrt_llm/_torch/disaggregation/native/default_planner.py new file mode 100644 index 000000000000..5ee4b895154b --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/native/default_planner.py @@ -0,0 +1,802 @@ +# 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 transfer planner every model gets unless it supplies its own. + +Attention (PAGED) and recurrent-state (STATE) layer groups are planned by +``AttentionPolicy`` and ``MambaPolicy``; pools are matched by role and layer +overlap; TP/PP/CP overlap decides which peer ranks receive what. All of it is +behind the :class:`TransferPlanner` protocol, so ``native/transfer.py`` only +ever sees ranks, pointer arrays and byte counts. +""" + +from __future__ import annotations + +from collections import Counter +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np + +from tensorrt_llm import logger +from tensorrt_llm._utils import nvtx_range +from tensorrt_llm._torch.disaggregation.base.backend import TokenRange +from tensorrt_llm._torch.disaggregation.base.planner import ( + PeerIncompatibleError, + PeerTargets, + TransferPlan, + TransferPlanner, +) +from tensorrt_llm._torch.disaggregation.base.region import RegionMapperBase +from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( + MambaPolicy, + mamba_receiver_payload_bytes, +) +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.page import ( + CacheKind, + KVCachePageTable, + MapperKind, + PoolView, +) +from tensorrt_llm._torch.disaggregation.resource.utils import ( + get_layer_byte_ranges, + get_layer_to_layer_group, + get_pool_view_global_layer_ids, +) + +# Type alias for (lg_idx, pool_idx) pair +LGPoolKey = Tuple[int, int] + + +@dataclass +class PeerOverlap: + overlap_pp_size: int = 0 + overlap_tp_size: int = 0 + overlap_cp_size: int = 0 + duplicate_head_factor: int = 1 + peer_duplicate_head_factor: int = 1 + ranks: List[int] = field(default_factory=list) + + +class DefaultTransferPlanner(TransferPlanner): + """Policy-driven planning over the page tables both peers publish. + + Per-peer state (pool mapping, mappers, peer extractor) is cached by peer + key and dropped in :meth:`forget_peer`; the TP/PP/CP overlap is cached per + peer *instance* and DP group because every rank of a DP group shares it. + """ + + planner_id = "trtllm.default" + + # Registry: CacheKind -> PolicyClass. Add new layer types here. The + # policy owns every mapper kind of its life cycle and picks the mapper + # per view in build_mapper. + _POLICY_CLASSES = { + CacheKind.PAGED: AttentionPolicy, + CacheKind.STATE: MambaPolicy, + } + + def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1): + self._ri = self_rank_info + self._self_ext_cache = self_extractor + self._policies: Dict[CacheKind, Union[AttentionPolicy, MambaPolicy]] = {} + self._kv_map_cache: Dict[ + tuple, RegionMapperBase + ] = {} # key: (peer_key, self_lg_pool_key, peer_lg_pool_key) + self._overlap_cache: Dict[str, PeerOverlap] = {} + self._lg_pool_mapping_cache: Dict[ + str, Dict[LGPoolKey, LGPoolKey] + ] = {} # peer_key -> {(self_lg, self_pi) -> (peer_lg, peer_pi)} + self._peer_ext_cache: Dict[str, KVRegionExtractorV1] = {} + + @staticmethod + def _unique_key(name: str, rank: int) -> str: + return name + str(rank) + + def peer_extractor(self, peer_ri: RankInfo) -> KVRegionExtractorV1: + key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) + ext = self._peer_ext_cache.get(key) + if ext is None: + ext = KVRegionExtractorV1(peer_ri.page_table) + self._peer_ext_cache[key] = ext + return ext + + # ---- TransferPlanner protocol ------------------------------------------------- + + def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: + if self_ri.bounded_replay_on_generation != peer_ri.bounded_replay_on_generation: + raise PeerIncompatibleError( + "DefaultTransferPlanner.check_peer: bounded_replay_on_generation must match " + f"between peers (local={self_ri.bounded_replay_on_generation}, " + f"peer={peer_ri.bounded_replay_on_generation})" + ) + if not self._get_policy(CacheKind.PAGED).check_peer_compatible(peer_ri): + raise PeerIncompatibleError( + "DefaultTransferPlanner.check_peer: attention geometry differs from the peer " + f"(peer={peer_ri.instance_name}, rank={peer_ri.instance_rank})" + ) + # Recurrent-state (Mamba/KDA) layout gate. Raises with a field-level + # diagnostic instead of returning False, so the precise mismatch + # reaches the caller. + self_page_table = self._self_ext_cache.page_table + try: + MambaPolicy.validate_peer_compatible( + self._ri, + peer_ri, + self_page_table, + peer_ri.page_table, + ) + AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) + except ValueError as e: + raise PeerIncompatibleError(str(e)) from e + + self_layers = sum(self._ri.layer_num_per_pp) + peer_layers = sum(peer_ri.layer_num_per_pp) + if self_layers != peer_layers: + # Allow mismatch when one side has speculative (e.g. MTP) layers + # that the other side doesn't. The pool_mapping logic will only + # transfer layers that exist on both sides. + logger.warning( + "DefaultTransferPlanner: layer count differs " + f"(local={self_layers}, peer={peer_layers}), " + "allowing partial layer transfer." + ) + self._warn_nhd_head_mismatch(peer_ri) + + def _warn_nhd_head_mismatch(self, peer_ri: RankInfo) -> None: + head_match, _ = self._get_policy(CacheKind.PAGED).head_match(peer_ri) + if not head_match: + self_page_table = self._self_ext_cache.page_table + nhd_fragments_per_token = sum( + len( + self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].buffer_entries + ) + for layer_group_id, pool_idx in self.get_pool_mapping(peer_ri) + if self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].mapper_kind + == MapperKind.NHD + ) + if nhd_fragments_per_token: + local_heads = self._ri.attention.kv_heads_per_rank + peer_heads = peer_ri.attention.kv_heads_per_rank + logger.warning_once( + "NHD head-mismatched disaggregated KV transfer has no " + "contiguous staging path and will emit approximately " + f"{nhd_fragments_per_token} NIXL descriptors per transferred " + "token per peer, excluding block-level replicated pools " + f"(local_kv_heads={local_heads}, peer_kv_heads={peer_heads}). " + "Long-context TEP/DEP transfers may have high latency.", + key=( + "native-nhd-head-mismatch-" + f"{local_heads}-{peer_heads}-{nhd_fragments_per_token}" + ), + ) + + def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> PeerTargets: + overlap = self.get_peer_overlap(peer_ri, peer_dp_rank) + # Helix writers contribute unequal shares (one KV owner, one state + # sender, the rest empty), so CP on either side rules out the equal + # fan-in split as well. + uniform = ( + self_ri.cp_size == 1 + and peer_ri.cp_size == 1 + and self._fanin_bounce_safe(overlap, peer_ri, self._self_ext_cache.page_table) + ) + return PeerTargets(ranks=list(overlap.ranks), uniform_writers=uniform) + + def plan( + self, + self_ri: RankInfo, + peer_ri: RankInfo, + src_block_ids_per_lg: List[np.ndarray], + dst_block_ids_per_lg: List[np.ndarray], + token_range: TokenRange, + prompt_len: int, + dst_start_token: Optional[int], + ) -> TransferPlan: + # Only a piece short of the whole prompt runs the chunk projection below, and that is the + # part with no answer under CP; a sole piece covering everything takes the untouched path. + is_partial = token_range.start > 0 or token_range.end < prompt_len + if is_partial and (self_ri.cp_size > 1 or peer_ri.cp_size > 1): + raise ValueError( + "enable_pipelined_transfer is not supported with context parallelism " + f"(sender cp_size={self_ri.cp_size}, receiver cp_size={peer_ri.cp_size})" + ) + targets = self.get_peer_overlap(peer_ri, peer_ri.dp_rank) + + # Aggregate fragment pointers from all matching pool pairs. + # Each pool pair produces one or more region pairs (rp), each containing + # a numpy array of src/dst pointers and a uniform bytes_per_region. + src_frag_parts: list[np.ndarray] = [] + dst_frag_parts: list[np.ndarray] = [] + # Instead of calling np.full() per region pair to build a size array and + # then np.concatenate() all of them, we record (count, bytes_per_region) + # tuples and construct the final sizes array with a single np.repeat(). + # For 48k+ items this avoids many small allocations in the hot loop. + size_specs: list[tuple[int, int]] = [] + extractor = self._self_ext_cache + peer_extractor = self.peer_extractor(peer_ri) + pool_mapping = self.get_pool_mapping(peer_ri) + + # Aggregate fragments from all matching pools using numpy concatenation. + # Send ownership is per pool: replicated pools elect one fan-in + # owner, sharded pools keep head-duplication routing. + for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): + if not self.should_send_pool(targets, peer_ri, self_lg, self_pi): + continue + + lg_info = extractor.page_table.layer_groups[self_lg] + src_block_ids = src_block_ids_per_lg[self_lg] + dst_block_ids = dst_block_ids_per_lg[peer_lg] + + # Extract regions: STATE = per-layer ptrs, PAGED = per-block ptrs + if lg_info.kind == CacheKind.STATE: + if src_block_ids.size == 0 or dst_block_ids.size == 0: + continue + if not MambaPolicy.is_paired(self_ri, peer_ri): + continue + src_region = extractor.extract_slot(int(src_block_ids[0]), self_lg, self_pi) + dst_region = peer_extractor.extract_slot(int(dst_block_ids[0]), peer_lg, peer_pi) + else: + src_block_ids = np.asarray(src_block_ids, dtype=np.int64) + dst_block_ids = np.asarray(dst_block_ids, dtype=np.int64) + if peer_ri.cp_size > 1 and self_ri.cp_size == 1: + # Helix: the receiver owns global blocks [cp_rank::cp_size] + # (same protocol as partition_context_for_helix), so its + # table is the strided subset of ours; block reuse is + # rejected under helix. + src_block_ids = src_block_ids[peer_ri.cp_rank :: peer_ri.cp_size] + tpb = extractor.page_table.tokens_per_block + dst_start_block = ( + dst_start_token // tpb if dst_start_token is not None else 0 + ) + # Both tables are positional (index == block ordinal, -1 == not + # held), so pairing is a plain intersection: SWA eviction, prefix + # reuse, and pipelined chunk bounds are all already encoded as + # holes by whichever side owns that state. + src_block_ids, dst_block_ids = self._pair_ordinals( + src_block_ids, dst_block_ids, dst_start_block=dst_start_block + ) + if src_block_ids.size == 0: + continue + src_region = extractor.extract(src_block_ids, self_lg, self_pi) + dst_region = peer_extractor.extract(dst_block_ids, peer_lg, peer_pi) + + # Map and collect fragments (same for all layer types) + mapper = self.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) + region_pair = mapper.map(src_region, dst_region) + region_pairs = region_pair if isinstance(region_pair, list) else [region_pair] + for rp in region_pairs: + src_frag_parts.append(rp.src.memory.ptrs) + dst_frag_parts.append(rp.dst.memory.ptrs) + size_specs.append((rp.src.memory.ptrs.size, rp.src.memory.bytes_per_region)) + + if not src_frag_parts: + return TransferPlan.empty() + # Build the sizes array in one shot: np.repeat expands each + # bytes_per_region value by its count, e.g.: + # values=[4096, 8192], counts=[100, 200] + # -> [4096]*100 ++ [8192]*200 + counts, values = zip(*size_specs) + return TransferPlan( + src_ptrs=np.concatenate(src_frag_parts), + dst_ptrs=np.concatenate(dst_frag_parts), + sizes=np.repeat(np.array(values, dtype=np.int64), counts), + ) + + def receiver_payload_bytes( + self, self_ri: RankInfo, peer_ri: RankInfo, dst_block_ids_per_lg: List[np.ndarray] + ) -> Optional[int]: + receiver_pt = self._self_ext_cache.page_table + dst_slot = self._get_mamba_slot(receiver_pt, dst_block_ids_per_lg) + if dst_slot is None: + return 0 + if peer_ri.page_table is None: + return None # cannot size the recurrent-state payload + return mamba_receiver_payload_bytes( + sender_page_table=peer_ri.page_table, + receiver_page_table=receiver_pt, + dst_slot=dst_slot, + ) + + def forget_peer(self, peer_ri: RankInfo) -> None: + key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) + self._peer_ext_cache.pop(key, None) + self._lg_pool_mapping_cache.pop(key, None) + for k in [k for k in self._kv_map_cache if k[0] == key]: + del self._kv_map_cache[k] + + # ---- helpers moved from the transfer framework --------------------------------- + + @staticmethod + def _get_mamba_slot( + page_table: Optional[KVCachePageTable], dst_block_ids_per_lg: List[np.ndarray] + ) -> Optional[int]: + """The receiver slot of the first STATE layer group, or None when there is none.""" + if page_table is None: + return None + for lg_idx, lg in enumerate(page_table.layer_groups): + if lg.kind == CacheKind.STATE: + block_ids = dst_block_ids_per_lg[lg_idx] + if block_ids.size > 0: + return int(block_ids[0]) + return None + + @staticmethod + def _pair_ordinals( + src_block_ids: np.ndarray, + dst_block_ids: np.ndarray, + dst_start_block: int = 0, + ) -> tuple[np.ndarray, np.ndarray]: + """Pair src/dst positional block tables ordinal by ordinal. + + Both tables are indexed by block ordinal (see Chunk); an entry is -1 + where that side has nothing to offer/accept. A block moves only where + both sides hold it, so a sender that has not evicted yet, a receiver + that already reused a prefix, a receiver with a wider speculative + window, or a sender chunk that blanks everything outside its range all + fall out of the same intersection without any per-case arithmetic. + Ordinals below ``dst_start_block`` are excluded as well. + """ + if src_block_ids.size != dst_block_ids.size: + raise ValueError( + f"positional block tables differ in length: src={src_block_ids.size} " + f"dst={dst_block_ids.size}; both sides must span the same prompt" + ) + mask = (src_block_ids >= 0) & (dst_block_ids >= 0) + if dst_start_block > 0: + mask[:dst_start_block] = False + idx = np.flatnonzero(mask) + return src_block_ids[idx], dst_block_ids[idx] + + @nvtx_range("_build_kv_write_meta") + + @staticmethod + def _fanin_bounce_safe( + overlap: PeerOverlap, + peer_ri: RankInfo, + receiver_page_table: Optional[KVCachePageTable], + ) -> bool: + """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. + The split assumes every writer contributes the same size, which holds when: + * duplicate_head_factor == 1 -- else some ranks don't send KV yet still count in + expected_transfers, so the live writers overflow their slots; + * the PP layer split is even -- a single PP stage (overlap_pp_size <= 1) is trivially fine; + for PP fan-in, every overlapping stage must hold the same number of layers + (peer_ri.layer_num_per_pp all-equal) or per-writer sizes differ. If that full per-stage + list isn't available (shorter than the fan-in degree), be conservative and fall back. + Otherwise fall back to the per-fragment path (correct, just not coalesced). + Equal layer count means equal bytes only when the per-block sizes match; reserve() rejects + the mismatched case, so this only needs the count to split evenly.""" + if overlap.duplicate_head_factor != 1: + return False + if overlap.overlap_pp_size > 1: + lpp = getattr(peer_ri, "layer_num_per_pp", None) + if not lpp or len(lpp) < overlap.overlap_pp_size or len(set(lpp)) != 1: + return False + # Replicated pools (e.g. MiniMax M3 index-key) are sent by one elected + # fan-in owner only, so with multiple writers their contributions + # differ in size and the equal split is invalid. Inspect both endpoints: + # a masked PP stage may advertise no replicated view even though another + # stage owns one that is visible in the receiver's page table. + if len(overlap.ranks) > 1: + for page_table in (peer_ri.page_table, receiver_page_table): + if page_table is None: + continue + for layer_group in page_table.layer_groups: + for pool_view in getattr(layer_group, "pool_views", ()): + if pool_view.mapper_kind == MapperKind.REPLICATED: + return False + return True + + # ---- pool matching, mappers and topology overlap ------------------------------ + + def get_pool_mapping(self, peer_ri: RankInfo) -> Dict[LGPoolKey, LGPoolKey]: + """Get mapping from (self_lg_idx, self_pool_idx) -> (peer_lg_idx, peer_pool_idx). + + Two-step matching: + 1. Find peer layer_group via layer_to_layer_group (global_layer_id -> lg_idx). + 2. Within the matched peer layer_group, find the unique peer pool whose + ``PoolView.pool_role`` equals self's and whose global_layer_ids + overlap. + + Layer-overlap is required: a peer pool with the same pool_role but + zero layer overlap with self is *not* a match — the two pools cover + disjoint layers and have nothing to transfer. + + A self layer group never matches multiple peer layer groups, so the + result is one peer pool per self pool. Layer groups partition each + rank's layers by attention/life-cycle class, which both sides derive + from the same model config; PP only changes which layers overlap (the + fan-out across peer PP ranks is handled by calling this method once + per peer rank). Step 1 raises if this invariant is ever violated. + """ + key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) + if key in self._lg_pool_mapping_cache: + return self._lg_pool_mapping_cache[key] + + mapping: Dict[LGPoolKey, LGPoolKey] = {} + self_pt = self._self_ext_cache.page_table + peer_pt = peer_ri.page_table + + if self_pt is None or peer_pt is None: + self._lg_pool_mapping_cache[key] = mapping + return mapping + if not self_pt.layer_groups or not peer_pt.layer_groups: + self._lg_pool_mapping_cache[key] = mapping + return mapping + + # Build kind-specific peer layer-to-group mappings lazily. In hybrid + # models (e.g. Qwen3Next) a global_layer_id may appear in both a + # PAGED and a STATE layer group, so the lookup must be scoped to + # the same CacheKind as self's layer group. + _peer_l2g_cache: Dict[CacheKind, Dict[int, int]] = {} + + def _peer_l2g(kind: CacheKind) -> Dict[int, int]: + if kind not in _peer_l2g_cache: + _peer_l2g_cache[kind] = get_layer_to_layer_group(peer_pt, kind) + return _peer_l2g_cache[kind] + + for self_lg_idx, self_lg in enumerate(self_pt.layer_groups): + if self_lg.kind not in (CacheKind.PAGED, CacheKind.STATE): + continue + for self_pi, self_pv in enumerate(self_lg.pool_views): + # Every view carries buffer_entries, so a view's exact layer + # set always comes from its entries (a view may cover a + # subset of the LG when V2 splits an LG into multiple pools + # by buffer-size class, or when a role class exists only on + # some layers, e.g. sparse-layer index-K). + pv_global_ids = get_pool_view_global_layer_ids(self_pv, self_lg) + if not pv_global_ids: + continue + + # Step 1: find the peer layer_group via overlapping global_layer_ids. + # A self layer group (hence each of its pool views) never matches + # multiple peer layer groups: layer groups partition a rank's + # layers by attention/life-cycle class, both sides derive that + # class from the same model config, and global ids are + # PP-invariant. So PP only changes WHICH layers overlap — layers + # the peer doesn't hold are a legal skip (PP slices, one-sided + # MTP layers) — never how many peer LGs they land in; the PP + # fan-out is handled by per-peer-rank calls of this method. A + # multi-LG hit therefore means the two peers group layers + # differently (unsupported topology), and we fail loudly instead + # of silently transferring only the first LG's overlap. + peer_layer_to_group = _peer_l2g(self_lg.kind) + peer_lg_indices = { + peer_layer_to_group[g] for g in pv_global_ids if g in peer_layer_to_group + } + if not peer_lg_indices: + continue + if len(peer_lg_indices) > 1: + raise ValueError( + "DefaultTransferPlanner.get_pool_mapping: pool view " + f"(lg={self_lg_idx}, pool={self_pi}) spans multiple peer " + f"layer groups {sorted(peer_lg_indices)}; mismatched layer " + "grouping between peers is not supported" + ) + peer_lg_idx = next(iter(peer_lg_indices)) + peer_lg = peer_pt.layer_groups[peer_lg_idx] + + # Step 2: pick the first peer pool with the same pool_role + # whose layers overlap self's (zero-overlap pools cover + # disjoint layers — nothing to transfer). + # + # Uniqueness assumption: at most one peer pool can match on + # both ``pool_role`` (frozenset equality) and layer overlap. + # We do *not* assume ``pool_role`` is unique within a peer LG + # — V2 may split an LG into multiple same-role pools by + # buffer-size class (e.g. VSWA). What we rely on is that both + # peers run the same pool-grouping logic, so for every self_pv + # there is exactly one peer pool with the same role *and* an + # overlapping layer set; other same-role peer pools cover + # disjoint layers and fall out via the overlap filter. + self_layer_set = set(pv_global_ids) + matched_peer_pi = None + for peer_pi, peer_pv in enumerate(peer_lg.pool_views): + if peer_pv.pool_role != self_pv.pool_role: + continue + peer_global_ids = get_pool_view_global_layer_ids(peer_pv, peer_lg) + if not set(peer_global_ids) & self_layer_set: + continue + if peer_pv.mapper_kind != self_pv.mapper_kind: + raise ValueError( + "DefaultTransferPlanner.get_pool_mapping: incompatible mapper " + f"kinds for pool role {sorted(self_pv.pool_role)} " + f"(local={self_pv.mapper_kind.name}, " + f"peer={peer_pv.mapper_kind.name}, peer_pool={peer_pi})" + ) + matched_peer_pi = peer_pi + break + + if matched_peer_pi is not None: + mapping[(self_lg_idx, self_pi)] = (peer_lg_idx, matched_peer_pi) + + self._lg_pool_mapping_cache[key] = mapping + return mapping + + def _get_policy(self, kind: CacheKind) -> Union[AttentionPolicy, MambaPolicy]: + """Return the policy for a CacheKind (lazily instantiated).""" + if kind not in self._policies: + self._policies[kind] = self._POLICY_CLASSES[kind](self._ri) + return self._policies[kind] + + def get_kv_map( + self, + peer_ri: RankInfo, + self_pool_key: LGPoolKey, + peer_pool_key: LGPoolKey, + ) -> RegionMapperBase: + """Get mapper for a specific pool pair. + + Args: + peer_ri: Peer rank info. + self_pool_key: (self_lg_idx, self_pool_idx). + peer_pool_key: (peer_lg_idx, peer_pool_idx). + """ + peer_key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) + cache_key = (peer_key, self_pool_key, peer_pool_key) + if cache_key in self._kv_map_cache: + return self._kv_map_cache[cache_key] + + self_pt = self._self_ext_cache.page_table + peer_pt = peer_ri.page_table + assert self_pt is not None + assert peer_pt is not None + self_lg_idx, self_pi = self_pool_key + peer_lg_idx, peer_pi = peer_pool_key + self_lg = self_pt.layer_groups[self_lg_idx] + peer_lg = peer_pt.layer_groups[peer_lg_idx] + self_pv = self_lg.pool_views[self_pi] + peer_pv = peer_lg.pool_views[peer_pi] + + if self_pv.mapper_kind != peer_pv.mapper_kind: + raise ValueError( + "DefaultTransferPlanner.get_kv_map: incompatible mapper kinds " + f"(local={self_pv.mapper_kind.name}, peer={peer_pv.mapper_kind.name})" + ) + + # Every view is entries-driven: resolve the overlap layers to + # slot-relative byte offsets on each side from the views' buffer + # entries. Layer selection is explicit, so mappers never assume a + # uniform layer stride (other role classes may interleave), and no + # convention about global-id/byte-offset ordering is needed. + self_global_ids = get_pool_view_global_layer_ids(self_pv, self_lg) + peer_global_ids = get_pool_view_global_layer_ids(peer_pv, peer_lg) + # Iterate the overlap in self's physical slot order (not sorted by + # global id) so that layers whose regions are contiguous on both + # sides stay adjacent in the offset arrays and the mappers can merge + # them into one fragment even when global-id order diverges from the + # physical layout. Order only affects run merging (a perf property): + # each layer's byte offset is looked up explicitly below, so any + # iteration order transfers correct bytes. + overlap = set(self_global_ids) & set(peer_global_ids) + overlapping_layers = [gid for gid in self_global_ids if gid in overlap] + + self_starts, self_bytes_per_layer = get_layer_byte_ranges(self_pv) + peer_starts, peer_bytes_per_layer = get_layer_byte_ranges(peer_pv) + self_g2l = {ll.global_layer_id: ll.local_layer_id for ll in self_lg.local_layers} + peer_g2l = {ll.global_layer_id: ll.local_layer_id for ll in peer_lg.local_layers} + self_layer_offsets = np.array( + [self_starts[self_g2l[gid]] for gid in overlapping_layers], dtype=np.int64 + ) + peer_layer_offsets = np.array( + [peer_starts[peer_g2l[gid]] for gid in overlapping_layers], dtype=np.int64 + ) + # Per-layer buffer count (K and V are separate buffers within a + # layer's region); head-mismatch mappers slice heads inside each. + self_buffers_per_layer = self._get_buffers_per_layer( + self_pv, + layer_group_id=self_lg_idx, + pool_idx=self_pi, + ) + peer_buffers_per_layer = self._get_buffers_per_layer( + peer_pv, + layer_group_id=peer_lg_idx, + pool_idx=peer_pi, + ) + + # Polymorphic dispatch: attention vs mamba policy + policy = self._get_policy(self_lg.kind) + + # For mamba, compute ptr-array layer offsets for partial PP overlap. + # extract_slot returns ptrs for ALL local_layer_ids in sorted order; + # the mapper must slice only the overlapping subset. + extra_kwargs = {} + if self_lg.kind == CacheKind.STATE: + # Positions index the ptr array ``extract_slot`` builds, which is + # ordered by view offset. + self_gid_to_pos = {gid: i for i, gid in enumerate(self_global_ids)} + peer_gid_to_pos = {gid: i for i, gid in enumerate(peer_global_ids)} + self_overlap_positions = [self_gid_to_pos[gid] for gid in overlapping_layers] + peer_overlap_positions = [peer_gid_to_pos[gid] for gid in overlapping_layers] + # Under contiguous PP partitioning, overlapping layers form a + # contiguous block in the ptrs array. + extra_kwargs["src_layer_off"] = ( + self_overlap_positions[0] if self_overlap_positions else 0 + ) + extra_kwargs["dst_layer_off"] = ( + peer_overlap_positions[0] if peer_overlap_positions else 0 + ) + + mapper = policy.build_mapper( + peer_ri=peer_ri, + mapper_kind=self_pv.mapper_kind, + self_layer_offsets=self_layer_offsets, + peer_layer_offsets=peer_layer_offsets, + self_bytes_per_layer=self_bytes_per_layer, + peer_bytes_per_layer=peer_bytes_per_layer, + self_buffers_per_layer=self_buffers_per_layer, + peer_buffers_per_layer=peer_buffers_per_layer, + self_lg=self_lg, + peer_lg=peer_lg, + self_pv=self_pv, + peer_pv=peer_pv, + **extra_kwargs, + ) + + self._kv_map_cache[cache_key] = mapper + return mapper + + @staticmethod + def _get_buffers_per_layer( + pool_view: PoolView, + *, + layer_group_id: int, + pool_idx: int, + ) -> int: + """Per-layer buffer count of a view (e.g. K+V -> 2, key-only -> 1). + + Views are bucketed per (layer group, pool, mapper kind) at page-table + build time, so every layer in a view carries the same role set and + hence the same entry count — a skewed distribution should never occur. + Still verify it per layer rather than via total-count divisibility: + e.g. 1 + 3 entries over two layers passes ``total % layers == 0`` yet + would make head-slicing mappers split every layer at wrong offsets. + """ + entries = pool_view.buffer_entries + if len(entries) == 0: + return 1 + counts = Counter(int(e["local_layer_id"]) for e in entries) + distinct = set(counts.values()) + if len(distinct) != 1: + raise ValueError( + "PoolView buffer entries are not evenly distributed across layers: " + f"layer_group={layer_group_id}, pool={pool_idx}, " + f"per-layer entry counts={sorted(counts.items())}" + ) + return distinct.pop() + + @staticmethod + def _find_overlap(self_val, peer_val, self_rank, peer_rank=None): + if self_val <= peer_val: + overlap = peer_val // self_val + start = self_rank * overlap + (peer_rank * peer_val if peer_rank is not None else 0) + end = start + overlap + else: + ratio = self_val // peer_val + start = (self_rank // ratio) + (peer_rank * peer_val if peer_rank is not None else 0) + overlap = 1 + end = start + overlap + + return overlap, start, end + + def get_peer_overlap(self, peer_rank_info: RankInfo, peer_dp_rank: int) -> PeerOverlap: + peer_ri = peer_rank_info + key = self._unique_key(peer_ri.instance_name, peer_dp_rank) + if key in self._overlap_cache: + return self._overlap_cache[key] + + # compute pp overlap and target layers + self_start_layer = sum(self._ri.layer_num_per_pp[: self._ri.pp_rank]) + self_end_layer = self_start_layer + self._ri.layer_num_per_pp[self._ri.pp_rank] + + pre = 0 + tgt_pp_ranks: List[int] = [] + for p in range(peer_ri.pp_size): + peer_start_layer = pre + peer_end_layer = peer_start_layer + peer_ri.layer_num_per_pp[p] + if self_start_layer < peer_end_layer and self_end_layer > peer_start_layer: + tgt_pp_ranks.append(p) + pre += peer_ri.layer_num_per_pp[p] + + if tgt_pp_ranks == []: + targets = PeerOverlap() + self._overlap_cache[key] = targets + return targets + + peer_start_pp = tgt_pp_ranks[0] + overlap_pp_size = len(tgt_pp_ranks) + peer_end_pp = peer_start_pp + overlap_pp_size + + self_tp_per_dp = self._ri.tp_size_per_dp_group + peer_tp_per_dp = peer_ri.tp_size_per_dp_group + self_tp_rank_in_dp = self._ri.tp_rank % self_tp_per_dp + + overlap_tp_size, peer_start_tp, peer_end_tp = self._find_overlap( + self_tp_per_dp, peer_tp_per_dp, self_tp_rank_in_dp, peer_dp_rank + ) + overlap_cp_size, peer_start_cp, peer_end_cp = self._find_overlap( + self._ri.cp_size, peer_ri.cp_size, self._ri.cp_rank + ) + + ranks: List[int] = [] + for pp in range(peer_start_pp, peer_end_pp): + for tp in range(peer_start_tp, peer_end_tp): + for cp in range(peer_start_cp, peer_end_cp): + # CP-minor flat rank (ppRank*(TP*CP) + tpRank*CP + + # cpRank), matching Mapping and the C++ transceiver; a + # TP-minor formula mis-routes when tp > 1 and cp > 1. + # Loop nesting mirrors the layout so ranks come out in + # ascending order. + ranks.append(pp * peer_ri.tp_size * peer_ri.cp_size + tp * peer_ri.cp_size + cp) + + dup_head, peer_dup_head = self._get_policy(CacheKind.PAGED).duplicate_head_factors(peer_ri) + + targets = PeerOverlap( + overlap_pp_size=overlap_pp_size, + overlap_tp_size=overlap_tp_size, + overlap_cp_size=overlap_cp_size, + duplicate_head_factor=dup_head, + peer_duplicate_head_factor=peer_dup_head, + ranks=ranks, + ) + self._overlap_cache[key] = targets + return targets + + def _owns_tp_fan_in(self, peer_rank_info: RankInfo) -> bool: + """Elect one owner when replicated bytes fan in across TP ranks. + + A peer with fewer TP shards receives identical replicated data from + several local ranks. Rotate the elected owner by the destination's + DP rank (mirroring ``should_send_kv``'s head-duplication pairing) so + that with a multi-DP-group generation side the extra replicated + traffic spreads across local ranks instead of always landing on the + first rank of each fan-in group. + """ + ratio = max( + 1, + self._ri.tp_size_per_dp_group // peer_rank_info.tp_size_per_dp_group, + ) + self_tp_rank = self._ri.tp_rank % self._ri.tp_size_per_dp_group + return self_tp_rank % ratio == peer_rank_info.dp_rank % ratio + + def should_send_pool( + self, + peer_overlap: PeerOverlap, + peer_rank_info: RankInfo, + layer_group_id: int, + pool_idx: int, + ) -> bool: + """Return whether this rank owns the transfer of one view pair. + + ``pool_idx`` indexes the layer group's ``pool_views`` list (one view + per role class; several views may share a physical pool). Each view + is kind-homogeneous, so ownership is a single per-view decision: + replicated views use one sender per fan-in group (the policy returns + None for them), sharded views retain head-duplication routing. + + For mamba layer groups, ownership is based on mamba's own TP routing + (independent of attention's duplicate-head logic): + - If mamba_tp == 1 (attention_dp enabled): fan-in election (one sender) + - Otherwise: always send (each rank owns unique TP-sharded state) + """ + layer_group = self._self_ext_cache.page_table.layer_groups[layer_group_id] + pool_view = layer_group.pool_views[pool_idx] + # Delegate to the policy's should_send. None means fan-in election. + policy = self._get_policy(layer_group.kind) + result = policy.should_send(peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind) + return self._owns_tp_fan_in(peer_rank_info) if result is None else result diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index d8d37658195b..9420690a085b 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -12,104 +12,56 @@ # 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. +"""Framework-side peer bookkeeping for the native transfer runtime. -from collections import Counter -from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple, Union +Who the peers are, their page-table extractors, the aux transfer layouts and +aux ownership live here. How KV bytes pair between this rank and a peer is the +:class:`TransferPlanner`'s business; the registrar only gates registration on +it and hands it out to the sender and receiver. +""" -import numpy as np +from typing import Dict, Optional -from tensorrt_llm import logger -from tensorrt_llm._torch.disaggregation.base.region import RegionMapperBase +from tensorrt_llm._torch.disaggregation.base.planner import ( + PeerIncompatibleError, + TransferPlanner, +) from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxTransferLayout -from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy -from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, MapperKind, PoolView -from tensorrt_llm._torch.disaggregation.resource.utils import ( - get_layer_byte_ranges, - get_layer_to_layer_group, - get_pool_view_global_layer_ids, -) - -# Type alias for (lg_idx, pool_idx) pair -LGPoolKey = Tuple[int, int] - - -@dataclass -class PeerOverlap: - overlap_pp_size: int = 0 - overlap_tp_size: int = 0 - overlap_cp_size: int = 0 - duplicate_head_factor: int = 1 - peer_duplicate_head_factor: int = 1 - ranks: List[int] = field(default_factory=list) class PeerRegistrar: - # Registry: CacheKind -> PolicyClass. Add new layer types here. The - # policy owns every mapper kind of its life cycle and picks the mapper - # per view in build_mapper. - _POLICY_CLASSES = { - CacheKind.PAGED: AttentionPolicy, - CacheKind.STATE: MambaPolicy, - } - - def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1): + def __init__( + self, + self_rank_info: RankInfo, + self_extractor: KVRegionExtractorV1, + planner: TransferPlanner, + ): self._ri = self_rank_info - self._policies: Dict[CacheKind, Union[AttentionPolicy, MambaPolicy]] = {} - self._peer_ri_cache: Dict[str, RankInfo] = {} - self._kv_map_cache: Dict[ - tuple, RegionMapperBase - ] = {} # key: (peer_key, self_lg_pool_key, peer_lg_pool_key) self._self_ext_cache = self_extractor + self._planner = planner + self._peer_ri_cache: Dict[str, RankInfo] = {} self._peer_ext_cache: Dict[str, KVRegionExtractorV1] = {} - self._overlap_cache: Dict[str, PeerOverlap] = {} self._aux_transfer_layout_cache: Dict[str, AuxTransferLayout] = {} - self._lg_pool_mapping_cache: Dict[ - str, Dict[LGPoolKey, LGPoolKey] - ] = {} # peer_key -> {(self_lg, self_pi) -> (peer_lg, peer_pi)} + + @property + def planner(self) -> TransferPlanner: + return self._planner def register(self, peer_name: str, peer_rank: int, peer_ri: RankInfo): - assert self._self_ext_cache is not None - if not self._check_peer_compatible(peer_ri): - raise ValueError( - f"PeerRegistrar.register: peer {peer_name} (rank={peer_rank}) is incompatible with local rank." + """Admit a peer. Raises :class:`PeerIncompatibleError` (a ``ValueError``) otherwise.""" + if peer_ri.planner_id != self._planner.planner_id: + raise PeerIncompatibleError( + f"PeerRegistrar.register: peer {peer_name} (rank={peer_rank}) plans transfers " + f"with '{peer_ri.planner_id}', this rank with '{self._planner.planner_id}'; " + "both sides of a disaggregated deployment must run the same model planner" ) + self._planner.check_peer(self._ri, peer_ri) key = self._unique_key(peer_name, peer_rank) self._peer_ri_cache[key] = peer_ri self._aux_transfer_layout_cache.pop(key, None) - peer_ri = self.get_peer_rank_info(peer_name, peer_rank) - extractor = KVRegionExtractorV1(peer_ri.page_table) - self._peer_ext_cache[key] = extractor - - head_match, _ = self._get_policy(CacheKind.PAGED).head_match(peer_ri) - if not head_match: - self_page_table = self._self_ext_cache.page_table - nhd_fragments_per_token = sum( - len( - self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].buffer_entries - ) - for layer_group_id, pool_idx in self.get_pool_mapping(peer_ri) - if self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].mapper_kind - == MapperKind.NHD - ) - if nhd_fragments_per_token: - local_heads = self._ri.attention.kv_heads_per_rank - peer_heads = peer_ri.attention.kv_heads_per_rank - logger.warning_once( - "NHD head-mismatched disaggregated KV transfer has no " - "contiguous staging path and will emit approximately " - f"{nhd_fragments_per_token} NIXL descriptors per transferred " - "token per peer, excluding block-level replicated pools " - f"(local_kv_heads={local_heads}, peer_kv_heads={peer_heads}). " - "Long-context TEP/DEP transfers may have high latency.", - key=( - "native-nhd-head-mismatch-" - f"{local_heads}-{peer_heads}-{nhd_fragments_per_token}" - ), - ) + self._peer_ext_cache[key] = KVRegionExtractorV1(peer_ri.page_table) def peer_extractor(self, peer_name: str, peer_rank: int) -> KVRegionExtractorV1: return self._peer_ext_cache[self._unique_key(peer_name, peer_rank)] @@ -121,17 +73,11 @@ def self_extractor(self) -> KVRegionExtractorV1: def unregister(self, peer_name: str, peer_rank: int): key = self._unique_key(peer_name, peer_rank) - if key in self._peer_ri_cache: - del self._peer_ri_cache[key] - if key in self._peer_ext_cache: - del self._peer_ext_cache[key] + peer_ri = self._peer_ri_cache.pop(key, None) + self._peer_ext_cache.pop(key, None) self._aux_transfer_layout_cache.pop(key, None) - # Clean up kv_map_cache entries for this peer - keys_to_remove = [k for k in self._kv_map_cache if k[0] == key] - for k in keys_to_remove: - del self._kv_map_cache[k] - if key in self._lg_pool_mapping_cache: - del self._lg_pool_mapping_cache[key] + if peer_ri is not None: + self._planner.forget_peer(peer_ri) def get_peer_rank_info(self, peer_name: str, peer_rank: int): return self._peer_ri_cache[self._unique_key(peer_name, peer_rank)] @@ -162,436 +108,6 @@ def validate_bounded_replay_compatible(self, peer_ri: RankInfo) -> None: f"peer={peer_ri.bounded_replay_on_generation})." ) - def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: - self.validate_bounded_replay_compatible(peer_ri) - if not self._get_policy(CacheKind.PAGED).check_peer_compatible(peer_ri): - return False - - # Recurrent-state (Mamba/KDA) layout gate. Raises ValueError with a - # field-level diagnostic instead of returning False, so the precise - # mismatch reaches the caller of register(). - self_page_table = ( - self._self_ext_cache.page_table if self._self_ext_cache is not None else None - ) - MambaPolicy.validate_peer_compatible( - self._ri, - peer_ri, - self_page_table, - peer_ri.page_table, - ) - AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) - - self_layers = sum(self._ri.layer_num_per_pp) - peer_layers = sum(peer_ri.layer_num_per_pp) - if self_layers != peer_layers: - # Allow mismatch when one side has speculative (e.g. MTP) layers - # that the other side doesn't. The pool_mapping logic will only - # transfer layers that exist on both sides. - logger.warning( - "PeerRegistrar: layer count differs " - f"(local={self_layers}, peer={peer_layers}), " - "allowing partial layer transfer." - ) - - return True - - def get_pool_mapping(self, peer_ri: RankInfo) -> Dict[LGPoolKey, LGPoolKey]: - """Get mapping from (self_lg_idx, self_pool_idx) -> (peer_lg_idx, peer_pool_idx). - - Two-step matching: - 1. Find peer layer_group via layer_to_layer_group (global_layer_id -> lg_idx). - 2. Within the matched peer layer_group, find the unique peer pool whose - ``PoolView.pool_role`` equals self's and whose global_layer_ids - overlap. - - Layer-overlap is required: a peer pool with the same pool_role but - zero layer overlap with self is *not* a match — the two pools cover - disjoint layers and have nothing to transfer. - - A self layer group never matches multiple peer layer groups, so the - result is one peer pool per self pool. Layer groups partition each - rank's layers by attention/life-cycle class, which both sides derive - from the same model config; PP only changes which layers overlap (the - fan-out across peer PP ranks is handled by calling this method once - per peer rank). Step 1 raises if this invariant is ever violated. - """ - key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) - if key in self._lg_pool_mapping_cache: - return self._lg_pool_mapping_cache[key] - - mapping: Dict[LGPoolKey, LGPoolKey] = {} - self_pt = self._self_ext_cache.page_table - peer_pt = peer_ri.page_table - - if self_pt is None or peer_pt is None: - self._lg_pool_mapping_cache[key] = mapping - return mapping - if not self_pt.layer_groups or not peer_pt.layer_groups: - self._lg_pool_mapping_cache[key] = mapping - return mapping - - # Build kind-specific peer layer-to-group mappings lazily. In hybrid - # models (e.g. Qwen3Next) a global_layer_id may appear in both a - # PAGED and a STATE layer group, so the lookup must be scoped to - # the same CacheKind as self's layer group. - _peer_l2g_cache: Dict[CacheKind, Dict[int, int]] = {} - - def _peer_l2g(kind: CacheKind) -> Dict[int, int]: - if kind not in _peer_l2g_cache: - _peer_l2g_cache[kind] = get_layer_to_layer_group(peer_pt, kind) - return _peer_l2g_cache[kind] - - for self_lg_idx, self_lg in enumerate(self_pt.layer_groups): - if self_lg.kind not in (CacheKind.PAGED, CacheKind.STATE): - continue - for self_pi, self_pv in enumerate(self_lg.pool_views): - # Every view carries buffer_entries, so a view's exact layer - # set always comes from its entries (a view may cover a - # subset of the LG when V2 splits an LG into multiple pools - # by buffer-size class, or when a role class exists only on - # some layers, e.g. sparse-layer index-K). - pv_global_ids = get_pool_view_global_layer_ids(self_pv, self_lg) - if not pv_global_ids: - continue - - # Step 1: find the peer layer_group via overlapping global_layer_ids. - # A self layer group (hence each of its pool views) never matches - # multiple peer layer groups: layer groups partition a rank's - # layers by attention/life-cycle class, both sides derive that - # class from the same model config, and global ids are - # PP-invariant. So PP only changes WHICH layers overlap — layers - # the peer doesn't hold are a legal skip (PP slices, one-sided - # MTP layers) — never how many peer LGs they land in; the PP - # fan-out is handled by per-peer-rank calls of this method. A - # multi-LG hit therefore means the two peers group layers - # differently (unsupported topology), and we fail loudly instead - # of silently transferring only the first LG's overlap. - peer_layer_to_group = _peer_l2g(self_lg.kind) - peer_lg_indices = { - peer_layer_to_group[g] for g in pv_global_ids if g in peer_layer_to_group - } - if not peer_lg_indices: - continue - if len(peer_lg_indices) > 1: - raise ValueError( - "PeerRegistrar.get_pool_mapping: pool view " - f"(lg={self_lg_idx}, pool={self_pi}) spans multiple peer " - f"layer groups {sorted(peer_lg_indices)}; mismatched layer " - "grouping between peers is not supported" - ) - peer_lg_idx = next(iter(peer_lg_indices)) - peer_lg = peer_pt.layer_groups[peer_lg_idx] - - # Step 2: pick the first peer pool with the same pool_role - # whose layers overlap self's (zero-overlap pools cover - # disjoint layers — nothing to transfer). - # - # Uniqueness assumption: at most one peer pool can match on - # both ``pool_role`` (frozenset equality) and layer overlap. - # We do *not* assume ``pool_role`` is unique within a peer LG - # — V2 may split an LG into multiple same-role pools by - # buffer-size class (e.g. VSWA). What we rely on is that both - # peers run the same pool-grouping logic, so for every self_pv - # there is exactly one peer pool with the same role *and* an - # overlapping layer set; other same-role peer pools cover - # disjoint layers and fall out via the overlap filter. - self_layer_set = set(pv_global_ids) - matched_peer_pi = None - for peer_pi, peer_pv in enumerate(peer_lg.pool_views): - if peer_pv.pool_role != self_pv.pool_role: - continue - peer_global_ids = get_pool_view_global_layer_ids(peer_pv, peer_lg) - if not set(peer_global_ids) & self_layer_set: - continue - if peer_pv.mapper_kind != self_pv.mapper_kind: - raise ValueError( - "PeerRegistrar.get_pool_mapping: incompatible mapper " - f"kinds for pool role {sorted(self_pv.pool_role)} " - f"(local={self_pv.mapper_kind.name}, " - f"peer={peer_pv.mapper_kind.name}, peer_pool={peer_pi})" - ) - matched_peer_pi = peer_pi - break - - if matched_peer_pi is not None: - mapping[(self_lg_idx, self_pi)] = (peer_lg_idx, matched_peer_pi) - - self._lg_pool_mapping_cache[key] = mapping - return mapping - - def _get_policy(self, kind: CacheKind) -> Union[AttentionPolicy, MambaPolicy]: - """Return the policy for a CacheKind (lazily instantiated).""" - if kind not in self._policies: - self._policies[kind] = self._POLICY_CLASSES[kind](self._ri) - return self._policies[kind] - - def get_kv_map( - self, - peer_ri: RankInfo, - self_pool_key: LGPoolKey, - peer_pool_key: LGPoolKey, - ) -> RegionMapperBase: - """Get mapper for a specific pool pair. - - Args: - peer_ri: Peer rank info. - self_pool_key: (self_lg_idx, self_pool_idx). - peer_pool_key: (peer_lg_idx, peer_pool_idx). - """ - peer_key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) - cache_key = (peer_key, self_pool_key, peer_pool_key) - if cache_key in self._kv_map_cache: - return self._kv_map_cache[cache_key] - - self_pt = self._self_ext_cache.page_table - peer_pt = peer_ri.page_table - assert self_pt is not None - assert peer_pt is not None - self_lg_idx, self_pi = self_pool_key - peer_lg_idx, peer_pi = peer_pool_key - self_lg = self_pt.layer_groups[self_lg_idx] - peer_lg = peer_pt.layer_groups[peer_lg_idx] - self_pv = self_lg.pool_views[self_pi] - peer_pv = peer_lg.pool_views[peer_pi] - - if self_pv.mapper_kind != peer_pv.mapper_kind: - raise ValueError( - "PeerRegistrar.get_kv_map: incompatible mapper kinds " - f"(local={self_pv.mapper_kind.name}, peer={peer_pv.mapper_kind.name})" - ) - - # Every view is entries-driven: resolve the overlap layers to - # slot-relative byte offsets on each side from the views' buffer - # entries. Layer selection is explicit, so mappers never assume a - # uniform layer stride (other role classes may interleave), and no - # convention about global-id/byte-offset ordering is needed. - self_global_ids = get_pool_view_global_layer_ids(self_pv, self_lg) - peer_global_ids = get_pool_view_global_layer_ids(peer_pv, peer_lg) - # Iterate the overlap in self's physical slot order (not sorted by - # global id) so that layers whose regions are contiguous on both - # sides stay adjacent in the offset arrays and the mappers can merge - # them into one fragment even when global-id order diverges from the - # physical layout. Order only affects run merging (a perf property): - # each layer's byte offset is looked up explicitly below, so any - # iteration order transfers correct bytes. - overlap = set(self_global_ids) & set(peer_global_ids) - overlapping_layers = [gid for gid in self_global_ids if gid in overlap] - - self_starts, self_bytes_per_layer = get_layer_byte_ranges(self_pv) - peer_starts, peer_bytes_per_layer = get_layer_byte_ranges(peer_pv) - self_g2l = {ll.global_layer_id: ll.local_layer_id for ll in self_lg.local_layers} - peer_g2l = {ll.global_layer_id: ll.local_layer_id for ll in peer_lg.local_layers} - self_layer_offsets = np.array( - [self_starts[self_g2l[gid]] for gid in overlapping_layers], dtype=np.int64 - ) - peer_layer_offsets = np.array( - [peer_starts[peer_g2l[gid]] for gid in overlapping_layers], dtype=np.int64 - ) - # Per-layer buffer count (K and V are separate buffers within a - # layer's region); head-mismatch mappers slice heads inside each. - self_buffers_per_layer = self._get_buffers_per_layer( - self_pv, - layer_group_id=self_lg_idx, - pool_idx=self_pi, - ) - peer_buffers_per_layer = self._get_buffers_per_layer( - peer_pv, - layer_group_id=peer_lg_idx, - pool_idx=peer_pi, - ) - - # Polymorphic dispatch: attention vs mamba policy - policy = self._get_policy(self_lg.kind) - - # For mamba, compute ptr-array layer offsets for partial PP overlap. - # extract_slot returns ptrs for ALL local_layer_ids in sorted order; - # the mapper must slice only the overlapping subset. - extra_kwargs = {} - if self_lg.kind == CacheKind.STATE: - # Positions index the ptr array ``extract_slot`` builds, which is - # ordered by view offset. - self_gid_to_pos = {gid: i for i, gid in enumerate(self_global_ids)} - peer_gid_to_pos = {gid: i for i, gid in enumerate(peer_global_ids)} - self_overlap_positions = [self_gid_to_pos[gid] for gid in overlapping_layers] - peer_overlap_positions = [peer_gid_to_pos[gid] for gid in overlapping_layers] - # Under contiguous PP partitioning, overlapping layers form a - # contiguous block in the ptrs array. - extra_kwargs["src_layer_off"] = ( - self_overlap_positions[0] if self_overlap_positions else 0 - ) - extra_kwargs["dst_layer_off"] = ( - peer_overlap_positions[0] if peer_overlap_positions else 0 - ) - - mapper = policy.build_mapper( - peer_ri=peer_ri, - mapper_kind=self_pv.mapper_kind, - self_layer_offsets=self_layer_offsets, - peer_layer_offsets=peer_layer_offsets, - self_bytes_per_layer=self_bytes_per_layer, - peer_bytes_per_layer=peer_bytes_per_layer, - self_buffers_per_layer=self_buffers_per_layer, - peer_buffers_per_layer=peer_buffers_per_layer, - self_lg=self_lg, - peer_lg=peer_lg, - self_pv=self_pv, - peer_pv=peer_pv, - **extra_kwargs, - ) - - self._kv_map_cache[cache_key] = mapper - return mapper - - @staticmethod - def _get_buffers_per_layer( - pool_view: PoolView, - *, - layer_group_id: int, - pool_idx: int, - ) -> int: - """Per-layer buffer count of a view (e.g. K+V -> 2, key-only -> 1). - - Views are bucketed per (layer group, pool, mapper kind) at page-table - build time, so every layer in a view carries the same role set and - hence the same entry count — a skewed distribution should never occur. - Still verify it per layer rather than via total-count divisibility: - e.g. 1 + 3 entries over two layers passes ``total % layers == 0`` yet - would make head-slicing mappers split every layer at wrong offsets. - """ - entries = pool_view.buffer_entries - if len(entries) == 0: - return 1 - counts = Counter(int(e["local_layer_id"]) for e in entries) - distinct = set(counts.values()) - if len(distinct) != 1: - raise ValueError( - "PoolView buffer entries are not evenly distributed across layers: " - f"layer_group={layer_group_id}, pool={pool_idx}, " - f"per-layer entry counts={sorted(counts.items())}" - ) - return distinct.pop() - - @staticmethod - def _find_overlap(self_val, peer_val, self_rank, peer_rank=None): - if self_val <= peer_val: - overlap = peer_val // self_val - start = self_rank * overlap + (peer_rank * peer_val if peer_rank is not None else 0) - end = start + overlap - else: - ratio = self_val // peer_val - start = (self_rank // ratio) + (peer_rank * peer_val if peer_rank is not None else 0) - overlap = 1 - end = start + overlap - - return overlap, start, end - - def get_peer_overlap(self, peer_rank_info: RankInfo, peer_dp_rank: int) -> PeerOverlap: - peer_ri = peer_rank_info - key = self._unique_key(peer_ri.instance_name, peer_dp_rank) - if key in self._overlap_cache: - return self._overlap_cache[key] - - # compute pp overlap and target layers - self_start_layer = sum(self._ri.layer_num_per_pp[: self._ri.pp_rank]) - self_end_layer = self_start_layer + self._ri.layer_num_per_pp[self._ri.pp_rank] - - pre = 0 - tgt_pp_ranks: List[int] = [] - for p in range(peer_ri.pp_size): - peer_start_layer = pre - peer_end_layer = peer_start_layer + peer_ri.layer_num_per_pp[p] - if self_start_layer < peer_end_layer and self_end_layer > peer_start_layer: - tgt_pp_ranks.append(p) - pre += peer_ri.layer_num_per_pp[p] - - if tgt_pp_ranks == []: - targets = PeerOverlap() - self._overlap_cache[key] = targets - return targets - - peer_start_pp = tgt_pp_ranks[0] - overlap_pp_size = len(tgt_pp_ranks) - peer_end_pp = peer_start_pp + overlap_pp_size - - self_tp_per_dp = self._ri.tp_size_per_dp_group - peer_tp_per_dp = peer_ri.tp_size_per_dp_group - self_tp_rank_in_dp = self._ri.tp_rank % self_tp_per_dp - - overlap_tp_size, peer_start_tp, peer_end_tp = self._find_overlap( - self_tp_per_dp, peer_tp_per_dp, self_tp_rank_in_dp, peer_dp_rank - ) - overlap_cp_size, peer_start_cp, peer_end_cp = self._find_overlap( - self._ri.cp_size, peer_ri.cp_size, self._ri.cp_rank - ) - - ranks: List[int] = [] - for pp in range(peer_start_pp, peer_end_pp): - for tp in range(peer_start_tp, peer_end_tp): - for cp in range(peer_start_cp, peer_end_cp): - # CP-minor flat rank (ppRank*(TP*CP) + tpRank*CP + - # cpRank), matching Mapping and the C++ transceiver; a - # TP-minor formula mis-routes when tp > 1 and cp > 1. - # Loop nesting mirrors the layout so ranks come out in - # ascending order. - ranks.append(pp * peer_ri.tp_size * peer_ri.cp_size + tp * peer_ri.cp_size + cp) - - dup_head, peer_dup_head = self._get_policy(CacheKind.PAGED).duplicate_head_factors(peer_ri) - - targets = PeerOverlap( - overlap_pp_size=overlap_pp_size, - overlap_tp_size=overlap_tp_size, - overlap_cp_size=overlap_cp_size, - duplicate_head_factor=dup_head, - peer_duplicate_head_factor=peer_dup_head, - ranks=ranks, - ) - self._overlap_cache[key] = targets - return targets - - def _owns_tp_fan_in(self, peer_rank_info: RankInfo) -> bool: - """Elect one owner when replicated bytes fan in across TP ranks. - - A peer with fewer TP shards receives identical replicated data from - several local ranks. Rotate the elected owner by the destination's - DP rank (mirroring ``should_send_kv``'s head-duplication pairing) so - that with a multi-DP-group generation side the extra replicated - traffic spreads across local ranks instead of always landing on the - first rank of each fan-in group. - """ - ratio = max( - 1, - self._ri.tp_size_per_dp_group // peer_rank_info.tp_size_per_dp_group, - ) - self_tp_rank = self._ri.tp_rank % self._ri.tp_size_per_dp_group - return self_tp_rank % ratio == peer_rank_info.dp_rank % ratio - - def should_send_pool( - self, - peer_overlap: PeerOverlap, - peer_rank_info: RankInfo, - layer_group_id: int, - pool_idx: int, - ) -> bool: - """Return whether this rank owns the transfer of one view pair. - - ``pool_idx`` indexes the layer group's ``pool_views`` list (one view - per role class; several views may share a physical pool). Each view - is kind-homogeneous, so ownership is a single per-view decision: - replicated views use one sender per fan-in group (the policy returns - None for them), sharded views retain head-duplication routing. - - For mamba layer groups, ownership is based on mamba's own TP routing - (independent of attention's duplicate-head logic): - - If mamba_tp == 1 (attention_dp enabled): fan-in election (one sender) - - Otherwise: always send (each rank owns unique TP-sharded state) - """ - layer_group = self._self_ext_cache.page_table.layer_groups[layer_group_id] - pool_view = layer_group.pool_views[pool_idx] - # Delegate to the policy's should_send. None means fan-in election. - policy = self._get_policy(layer_group.kind) - result = policy.should_send(peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind) - return self._owns_tp_fan_in(peer_rank_info) if result is None else result - def should_send_aux(self, peer_rank_info: RankInfo) -> bool: # to ensure the transfer aux is not duplicated diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index d7b6f3ceec0e..50ed8f5da416 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -67,7 +67,8 @@ MambaPolicy, mamba_receiver_payload_bytes, ) -from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar +from tensorrt_llm._torch.disaggregation.native.default_planner import PeerOverlap +from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.base.planner import ( # noqa: F401 (re-exported) PeerIncompatibleError, ) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index 4362fb47665f..c5eabf0e469d 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -30,6 +30,7 @@ DeepseekV4AttentionType, compress_ratio_has_attention, ) +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( @@ -1968,13 +1969,16 @@ def test_disagg_pool_mapping_prefers_exact_pool_view_layer_set(self): try: ctx_rank_info = RankInfo.from_kv_cache_manager("ctx", ctx_cache_manager, 0) gen_rank_info = RankInfo.from_kv_cache_manager("gen", gen_cache_manager, 0) - registrar = PeerRegistrar(gen_rank_info, KVRegionExtractorV1(gen_cache_manager)) + gen_ext = KVRegionExtractorV1(gen_cache_manager) + planner = DefaultTransferPlanner(gen_rank_info, gen_ext) + registrar = PeerRegistrar(gen_rank_info, gen_ext, planner) + ctx_rank_info.planner_id = planner.planner_id registrar.register("ctx", 0, ctx_rank_info) - pool_mapping = registrar.get_pool_mapping(ctx_rank_info) + pool_mapping = planner.get_pool_mapping(ctx_rank_info) assert pool_mapping for self_pool_key, peer_pool_key in pool_mapping.items(): - registrar.get_kv_map(ctx_rank_info, self_pool_key, peer_pool_key) + planner.get_kv_map(ctx_rank_info, self_pool_key, peer_pool_key) finally: ctx_cache_manager.shutdown() gen_cache_manager.shutdown() diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 70b761b1326f..32206c7bbdae 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -233,10 +233,10 @@ def test_fanin_bounce_safe_gate() -> None: some ranks don't send KV yet count in expected_transfers) must fall back to the per-fragment path. """ - tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.resource.page import MapperKind - safe = tfr.Receiver._fanin_bounce_safe + safe = DefaultTransferPlanner._fanin_bounce_safe def ov(dup: int, pp: int, ranks: tuple[int, ...] = (0,)) -> SimpleNamespace: return SimpleNamespace(duplicate_head_factor=dup, overlap_pp_size=pp, ranks=list(ranks)) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ec0a638524cb..fdb8d1bd10ba 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -347,6 +347,7 @@ def test_v1_dsa_indexer_replicated_transfer_across_pp( import torch from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ReplicatedMapper + from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo @@ -362,7 +363,9 @@ def test_v1_dsa_indexer_replicated_transfer_across_pp( try: ctx_extractor = KVRegionExtractorV1(ctx) ctx_ri = RankInfo.from_kv_cache_manager("ctx", ctx, device_id=0) - registrar = PeerRegistrar(ctx_ri, ctx_extractor) + registrar = PeerRegistrar( + ctx_ri, ctx_extractor, DefaultTransferPlanner(ctx_ri, ctx_extractor) + ) # Fill ctx indexer pool with position-dependent bytes; zero gen pools. ctx_pool_tensor = ctx.impl.get_indexer_k_cache_pool() @@ -381,12 +384,13 @@ def test_v1_dsa_indexer_replicated_transfer_across_pp( for gen_pp_rank, gen in enumerate(gens): gen_ri = RankInfo.from_kv_cache_manager("gen", gen, device_id=0) + gen_ri.planner_id = DefaultTransferPlanner.planner_id registrar.register("gen", gen_ri.instance_rank, gen_ri) - mapping = registrar.get_pool_mapping(gen_ri) + mapping = registrar.planner.get_pool_mapping(gen_ri) # KV pool and indexer pool each match 1:1. assert mapping == {(0, 0): (0, 0), (0, 1): (0, 1)} - mapper = registrar.get_kv_map(gen_ri, (0, 1), (0, 1)) + mapper = registrar.planner.get_kv_map(gen_ri, (0, 1), (0, 1)) assert isinstance(mapper, ReplicatedMapper) gen_extractor = registrar.peer_extractor("gen", gen_ri.instance_rank) diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 9494849f23fd..0feb34f20b9a 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -59,6 +59,7 @@ MambaPolicy, mamba_receiver_payload_bytes, ) +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( @@ -286,7 +287,9 @@ def test_kda_hetero_tp_sharded_accepted(): MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) # And via the registrar entry point used at runtime. - registrar = PeerRegistrar(ctx_ri, KVRegionExtractorV1(ctx_pt)) + ctx_ext = KVRegionExtractorV1(ctx_pt) + registrar = PeerRegistrar(ctx_ri, ctx_ext, DefaultTransferPlanner(ctx_ri, ctx_ext)) + gen_ri.planner_id = DefaultTransferPlanner.planner_id registrar.register("kda_gen", 1, gen_ri) finally: for mgr in ctx_mgrs + gen_mgrs: diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 7ff5f5e72e7c..c00784ea839f 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -16,7 +16,7 @@ import numpy as np import pytest -import tensorrt_llm._torch.disaggregation.native.peer as peer_module +import tensorrt_llm._torch.disaggregation.native.default_planner as default_planner_module from tensorrt_llm._torch.disaggregation.base.region import MemRegionGroup, SpecRegion from tensorrt_llm._torch.disaggregation.native.auxiliary import ( AuxBufferMeta, @@ -35,7 +35,11 @@ MambaHeadMatchMapper, MambaPolicy, ) -from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar +from tensorrt_llm._torch.disaggregation.native.default_planner import ( + DefaultTransferPlanner, + PeerOverlap, +) +from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import ( @@ -183,6 +187,7 @@ def make_rankinfo( cp_size=cp_size, cp_rank=cp_rank, device_id=0, + planner_id=DefaultTransferPlanner.planner_id, layer_num_per_pp=layer_num_per_pp, self_endpoint="", transfer_engine_info=b"", @@ -203,7 +208,7 @@ def make_rankinfo( def _make_peer_registrar_and_peer_ri(self_ri, peer_ri): pt = make_page_table() real_kv_extractor = KVRegionExtractorV1(pt) - reg = PeerRegistrar(self_ri, real_kv_extractor) + reg = PeerRegistrar(self_ri, real_kv_extractor, DefaultTransferPlanner(self_ri, real_kv_extractor)) if peer_ri.page_table is None: peer_ri = make_rankinfo( instance_name=peer_ri.instance_name, @@ -250,7 +255,7 @@ def test_basic_overlap(): layer_num_per_pp=[1, 1], ) reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) - overlap = reg.get_peer_overlap(peer_ri, peer_dp_rank=0) + overlap = reg.planner.get_peer_overlap(peer_ri, peer_dp_rank=0) assert overlap.overlap_pp_size == 2 assert overlap.overlap_tp_size == 1 assert overlap.ranks == [0, 1] @@ -278,7 +283,7 @@ def test_no_overlap(): layer_num_per_pp=[2], ) reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) - overlap = reg.get_peer_overlap(peer_ri, 0) + overlap = reg.planner.get_peer_overlap(peer_ri, 0) assert overlap.overlap_pp_size == 0 assert overlap.ranks == [] @@ -307,7 +312,7 @@ def test_attention_free_stage_has_no_head_duplication( ) reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) - overlap = reg.get_peer_overlap(peer_ri, peer_dp_rank=0) + overlap = reg.planner.get_peer_overlap(peer_ri, peer_dp_rank=0) assert overlap.duplicate_head_factor == 1 assert overlap.peer_duplicate_head_factor == 1 @@ -336,7 +341,7 @@ def test_pp_ratio_peer_smaller(): layer_num_per_pp=[3], ) reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) - overlap = reg.get_peer_overlap(peer_ri, 0) + overlap = reg.planner.get_peer_overlap(peer_ri, 0) assert overlap.overlap_pp_size > 0 assert all(r >= 0 for r in overlap.ranks) @@ -349,7 +354,7 @@ def test_tp_overlap(): "peer", tp_size=4, tp_rank=0, pp_size=1, pp_rank=0, cp_size=1, cp_rank=0 ) reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) - overlap = reg.get_peer_overlap(peer_ri, 0) + overlap = reg.planner.get_peer_overlap(peer_ri, 0) assert overlap.overlap_tp_size in [1, 2] assert all(isinstance(r, int) for r in overlap.ranks) @@ -362,7 +367,7 @@ def test_cp_overlap(): "peer", cp_size=4, cp_rank=0, pp_size=1, pp_rank=0, tp_size=1, tp_rank=0 ) reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) - overlap = reg.get_peer_overlap(peer_ri, 0) + overlap = reg.planner.get_peer_overlap(peer_ri, 0) assert overlap.overlap_cp_size in [1, 2] assert all(isinstance(r, int) for r in overlap.ranks) @@ -389,7 +394,7 @@ def test_multiple_overlap(): layer_num_per_pp=[1, 1, 1, 1], ) reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) - overlap = reg.get_peer_overlap(peer_ri, peer_dp_rank=0) + overlap = reg.planner.get_peer_overlap(peer_ri, peer_dp_rank=0) expected_overlap = PeerOverlap( overlap_pp_size=2, overlap_tp_size=2, @@ -404,7 +409,9 @@ def test_multiple_overlap(): def _make_peer_registrar(self_rankinfo): pt = self_rankinfo.page_table if self_rankinfo.page_table else make_page_table() real_kv_extractor = KVRegionExtractorV1(pt) - reg = PeerRegistrar(self_rankinfo, real_kv_extractor) + reg = PeerRegistrar( + self_rankinfo, real_kv_extractor, DefaultTransferPlanner(self_rankinfo, real_kv_extractor) + ) return reg @@ -564,7 +571,7 @@ def test_peer_registrar_get_kv_map_identity(): page_table=make_page_table(), ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) # Full contiguous overlap: one whole-region fragment per block. assert isinstance(mapper, IntactMapper) pair = mapper.map( @@ -593,7 +600,7 @@ def test_peer_registrar_get_kv_map_head_match(): page_table=peer_pt, ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) # Partial layer overlap with matching heads: a single shifted fragment. assert isinstance(mapper, IntactMapper) pair = mapper.map( @@ -625,7 +632,7 @@ def test_peer_registrar_get_kv_map_head_mismatch(): page_table=make_page_table(block_bytes=[2048], mamba_tp=1), ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, HNDHeadMismatchMapper) @@ -650,7 +657,7 @@ def test_peer_registrar_get_kv_map_uses_physical_offset_order(): page_table=peer_pt, ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, IntactMapper) pair = mapper.map( SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=1024)), @@ -684,7 +691,7 @@ def test_peer_registrar_get_kv_map_non_contiguous_overlap_splits_fragments(): page_table=peer_pt, ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, IntactMapper) pairs = mapper.map( SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=1536)), @@ -718,7 +725,7 @@ def test_peer_registrar_get_kv_map_merges_non_monotonic_contiguous_layout(): page_table=peer_pt, ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, IntactMapper) pair = mapper.map( SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=1024)), @@ -832,7 +839,7 @@ def test_replicated_mapper_ignores_kv_head_mismatch(): layer_num_per_pp=[1], page_table=peer_pt, ) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, ReplicatedMapper) @@ -867,7 +874,7 @@ def test_replicated_pool_owner_rotates_by_destination_dp_rank(peer_dp_rank): layer_num_per_pp=[1], ) reg = _make_peer_registrar(self_ri) - ownership.append(reg.should_send_pool(overlap, peer_ri, 0, 0)) + ownership.append(reg.planner.should_send_pool(overlap, peer_ri, 0, 0)) # ratio = self_tp(8) / peer_tp_per_dp(1) = 8: exactly one owner, at the # slot selected by the destination dp rank. @@ -1043,7 +1050,7 @@ def test_get_buffers_per_layer_rejects_non_uniform_nhd_entries(): ) with pytest.raises(ValueError, match="layer_group=3, pool=4"): - PeerRegistrar._get_buffers_per_layer( + DefaultTransferPlanner._get_buffers_per_layer( pool_view, layer_group_id=3, pool_idx=4, @@ -1071,7 +1078,7 @@ def test_non_uniform_view_geometry_rejected_for_all_kinds(): reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) with pytest.raises(ValueError, match="not uniform"): - reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) def test_peer_registrar_allows_byte_aligned_subbyte_head_mismatch(): @@ -1095,7 +1102,7 @@ def test_peer_registrar_allows_byte_aligned_subbyte_head_mismatch(): reg = _make_peer_registrar(self_ri) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, HNDHeadMismatchMapper) @@ -1147,7 +1154,7 @@ def test_peer_registrar_dispatches_nhd_mapper(): reg = _make_peer_registrar(self_ri) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, NHDHeadMismatchMapper) @@ -1166,7 +1173,7 @@ def test_peer_registrar_warns_for_nhd_head_mismatch(monkeypatch): ) warnings = [] monkeypatch.setattr( - peer_module.logger, + default_planner_module.logger, "warning_once", lambda *message, key: warnings.append((" ".join(map(str, message)), key)), ) @@ -1200,7 +1207,7 @@ def test_peer_registrar_nhd_head_match_uses_intact_mapper(): reg = _make_peer_registrar(self_ri) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + mapper = reg.planner.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, IntactMapper) assert not isinstance(mapper, NHDHeadMismatchMapper) @@ -1311,8 +1318,8 @@ def _add_replicated_side_view(page_table, *, bytes_per_layer=8): def test_policy_dispatch_keys_on_cache_kind_only(): reg = _make_peer_registrar(make_rankinfo(instance_name="local")) - assert isinstance(reg._get_policy(CacheKind.PAGED), AttentionPolicy) - assert isinstance(reg._get_policy(CacheKind.STATE), MambaPolicy) + assert isinstance(reg.planner._get_policy(CacheKind.PAGED), AttentionPolicy) + assert isinstance(reg.planner._get_policy(CacheKind.STATE), MambaPolicy) def test_mamba_replicated_side_view_elects_one_sender(): @@ -1334,8 +1341,8 @@ def test_mamba_replicated_side_view_elects_one_sender(): instance_name="local", tp_size=2, tp_rank=tp_rank, page_table=self_pt ) reg = _make_peer_registrar(self_ri) - conv_ownership.append(reg.should_send_pool(overlap, peer_ri, 1, 0)) - side_ownership.append(reg.should_send_pool(overlap, peer_ri, 1, side_idx)) + conv_ownership.append(reg.planner.should_send_pool(overlap, peer_ri, 1, 0)) + side_ownership.append(reg.planner.should_send_pool(overlap, peer_ri, 1, side_idx)) assert conv_ownership == [True, True] assert side_ownership == [True, False] @@ -1355,9 +1362,9 @@ def test_mamba_replicated_side_view_copies_whole_layers_under_tp_mismatch(): peer_ri = make_rankinfo(instance_name="peer", tp_size=1, page_table=peer_pt) reg = _make_peer_registrar(self_ri) - mapper = reg.get_kv_map(peer_ri, (1, self_side), (1, peer_side)) + mapper = reg.planner.get_kv_map(peer_ri, (1, self_side), (1, peer_side)) assert isinstance(mapper, MambaHeadMatchMapper) - conv_mapper = reg.get_kv_map(peer_ri, (1, 0), (1, 0)) + conv_mapper = reg.planner.get_kv_map(peer_ri, (1, 0), (1, 0)) assert isinstance(conv_mapper, ConvStateMismatchMapper) diff --git a/tests/unittest/disaggregated/test_pool_matching.py b/tests/unittest/disaggregated/test_pool_matching.py index bcaa8225e03a..beb7f88531da 100644 --- a/tests/unittest/disaggregated/test_pool_matching.py +++ b/tests/unittest/disaggregated/test_pool_matching.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Golden tests for ``PeerRegistrar.get_pool_mapping``. +"""Golden tests for ``DefaultTransferPlanner.get_pool_mapping``. Locks the current (pre-refactor) behavior of disagg pool matching across the representative scenarios that the refactor must keep working: @@ -35,7 +35,7 @@ import pytest from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo -from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import ( @@ -176,7 +176,7 @@ def _rank_info( def _registrar(self_pt, **kwargs): self_ri = _rank_info(name="self", page_table=self_pt, **kwargs) - return PeerRegistrar(self_ri, KVRegionExtractorV1(self_pt)) + return DefaultTransferPlanner(self_ri, KVRegionExtractorV1(self_pt)) def _kv_pool_view(pool_idx, local_layer_ids): @@ -358,7 +358,7 @@ def test_skewed_buffer_entries_per_layer_raises(): """ view = _pool_view(0, [(0, "key"), (1, "key"), (1, "key"), (1, "key")]) with pytest.raises(ValueError, match="not evenly distributed"): - PeerRegistrar._get_buffers_per_layer(view, layer_group_id=0, pool_idx=0) + DefaultTransferPlanner._get_buffers_per_layer(view, layer_group_id=0, pool_idx=0) def test_duplicate_global_layer_id_across_lgs_raises(): From c8f15d811182105ebb7c560021990a0a41b55385 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 05:46:15 -0700 Subject: [PATCH 07/25] [None][refactor] Sender builds KV write metadata from the planner's TransferPlan Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../_torch/disaggregation/base/planner.py | 9 +- .../disaggregation/native/default_planner.py | 5 +- .../_torch/disaggregation/native/transfer.py | 163 ++++-------------- .../test_aux_buffer_registration.py | 3 +- .../disaggregated/test_chunked_transfer.py | 29 +++- .../disaggregated/test_transfer_planner.py | 117 +++++++++++++ .../disaggregated/test_unsubmitted_session.py | 3 +- 7 files changed, 181 insertions(+), 148 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/base/planner.py b/tensorrt_llm/_torch/disaggregation/base/planner.py index b820a6fac5ce..e0dea1e3954b 100644 --- a/tensorrt_llm/_torch/disaggregation/base/planner.py +++ b/tensorrt_llm/_torch/disaggregation/base/planner.py @@ -28,7 +28,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, AbstractSet, List, Optional, Tuple import numpy as np @@ -108,8 +108,13 @@ def plan( token_range: "TokenRange", prompt_len: int, dst_start_token: Optional[int], + excluded_pool_views: AbstractSet[Tuple[int, int]], ) -> TransferPlan: - """Sender side: the bytes this rank writes into ``peer_ri`` for one chunk.""" + """Sender side: the bytes this rank writes into ``peer_ri`` for one chunk. + + ``excluded_pool_views`` are ``(layer_group, pool_view)`` pairs the cache + manager keeps out of this chunk; the plan writes nothing from them. + """ @abstractmethod def receiver_payload_bytes( diff --git a/tensorrt_llm/_torch/disaggregation/native/default_planner.py b/tensorrt_llm/_torch/disaggregation/native/default_planner.py index 5ee4b895154b..c062a8c7d6b6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/default_planner.py +++ b/tensorrt_llm/_torch/disaggregation/native/default_planner.py @@ -25,7 +25,7 @@ from collections import Counter from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple, Union +from typing import AbstractSet, Dict, List, Optional, Tuple, Union import numpy as np @@ -206,6 +206,7 @@ def plan( token_range: TokenRange, prompt_len: int, dst_start_token: Optional[int], + excluded_pool_views: AbstractSet[Tuple[int, int]], ) -> TransferPlan: # Only a piece short of the whole prompt runs the chunk projection below, and that is the # part with no answer under CP; a sole piece covering everything takes the untouched path. @@ -235,6 +236,8 @@ def plan( # Send ownership is per pool: replicated pools elect one fan-in # owner, sharded pools keep head-duplication routing. for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): + if (self_lg, self_pi) in excluded_pool_views: + continue if not self.should_send_pool(targets, peer_ri, self_lg, self_pi): continue diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 50ed8f5da416..7da13939e56d 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -921,6 +921,7 @@ def __init__( retirement_watchdog: Optional[RetirementWatchdog] = None, ) -> None: self._registrar = peer_registrar + self._planner = peer_registrar.planner self._device_id = peer_registrar.self_rank_info.device_id self._agent = agent self._bounce = bounce @@ -1079,7 +1080,7 @@ def setup_session(self, tx_session: "TxSession"): peer_ri = self._registrar.get_peer_rank_info( req_info.instance_name, req_info.instance_rank ) - expected_count = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks) + expected_count = self._expected_transfers(peer_ri) if self._is_req_ready(unique_rid, expected_count): with tx_session.lock: tx_session.receiver_ready = True @@ -1096,6 +1097,12 @@ def _get_session(self, unique_rid: Optional[int]) -> Optional["TxSession"]: return session return session_entry + def _expected_transfers(self, peer_ri: RankInfo) -> int: + """How many writes the peer's DP group expects from this instance.""" + return len( + self._planner.targets(self._registrar.self_rank_info, peer_ri, peer_ri.dp_rank).ranks + ) + def _begin_task_operation(self, task: SendTaskBase, peer_rank: int) -> Optional[bool]: """Return True for admission, None for a duplicate, or False for rejection.""" if not task.begin_physical_operation(peer_rank): @@ -1673,153 +1680,43 @@ def _deliver_aux_to_agent(self, write_meta: WriteMeta): f"AUX send perf logging failed for request {write_meta.unique_rid}: {error}" ) - @staticmethod - def _pair_ordinals( - src_block_ids: np.ndarray, - dst_block_ids: np.ndarray, - dst_start_block: int = 0, - ) -> tuple[np.ndarray, np.ndarray]: - """Pair src/dst positional block tables ordinal by ordinal. - - Both tables are indexed by block ordinal (see Chunk); an entry is -1 - where that side has nothing to offer/accept. A block moves only where - both sides hold it, so a sender that has not evicted yet, a receiver - that already reused a prefix, a receiver with a wider speculative - window, or a sender chunk that blanks everything outside its range all - fall out of the same intersection without any per-case arithmetic. - Ordinals below ``dst_start_block`` are excluded as well. - """ - if src_block_ids.size != dst_block_ids.size: - raise ValueError( - f"positional block tables differ in length: src={src_block_ids.size} " - f"dst={dst_block_ids.size}; both sides must span the same prompt" - ) - mask = (src_block_ids >= 0) & (dst_block_ids >= 0) - if dst_start_block > 0: - mask[:dst_start_block] = False - idx = np.flatnonzero(mask) - return src_block_ids[idx], dst_block_ids[idx] - @nvtx_range("_build_kv_write_meta") def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> WriteMeta: peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) self_ri = self._registrar.self_rank_info - token_range = task._chunk.token_range - # Only a piece short of the whole prompt runs the chunk projection below, and that is the - # part with no answer under CP; a sole piece covering everything takes the untouched path. - is_partial = token_range.start > 0 or token_range.end < task._prompt_len - if is_partial and (self_ri.cp_size > 1 or peer_ri.cp_size > 1): - raise ValueError( - "enable_pipelined_transfer is not supported with context parallelism " - f"(sender cp_size={self_ri.cp_size}, receiver cp_size={peer_ri.cp_size})" - ) timer = task._perf_timer if timer: timer.record_prepare_args_start(peer_ri.instance_rank) - targets = self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank) - expected_transfers = len(targets.ranks) + expected_transfers = self._expected_transfers(peer_ri) # Recorded on the task as well: the count on the write meta settles the task, while the # handle has to answer how many writes are still out after one of them has already failed. task.expected_transfers = expected_transfers - # Aggregate fragment pointers from all matching pool pairs. - # Each pool pair produces one or more region pairs (rp), each containing - # a numpy array of src/dst pointers and a uniform bytes_per_region. - src_frag_parts: list[np.ndarray] = [] - dst_frag_parts: list[np.ndarray] = [] - # Instead of calling np.full() per region pair to build a size array and - # then np.concatenate() all of them, we record (count, bytes_per_region) - # tuples and construct the final sizes array with a single np.repeat(). - # For 48k+ items this avoids many small allocations in the hot loop. - size_specs: list[tuple[int, int]] = [] - dst_device_id = peer_ri.device_id - extractor = self._registrar.self_extractor - peer_extractor = self._registrar.peer_extractor( - peer_ri.instance_name, peer_ri.instance_rank + # Which of this rank's bytes land where on the peer is the model's planner's call; + # an empty plan still goes out so the peer's expected-writer count settles. + plan = self._planner.plan( + self_ri, + peer_ri, + task._chunk.block_ids_per_layer_groups, + req_info.block_ids_per_layer_groups, + task._chunk.token_range, + task._prompt_len, + req_info.dst_start_token, + task._chunk.excluded_pool_views, ) - pool_mapping = self._registrar.get_pool_mapping(peer_ri) - dst_block_ids_per_groups = req_info.block_ids_per_layer_groups - src_block_ids_per_groups = task._chunk.block_ids_per_layer_groups - - # Aggregate fragments from all matching pools using numpy concatenation. - # Send ownership is per pool: replicated pools elect one fan-in - # owner, sharded pools keep head-duplication routing. - for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): - if (self_lg, self_pi) in task._chunk.excluded_pool_views: - continue - if not self._registrar.should_send_pool(targets, peer_ri, self_lg, self_pi): - continue - - lg_info = extractor.page_table.layer_groups[self_lg] - src_block_ids = src_block_ids_per_groups[self_lg] - dst_block_ids = dst_block_ids_per_groups[peer_lg] - - # Extract regions: STATE = per-layer ptrs, PAGED = per-block ptrs - if lg_info.kind == CacheKind.STATE: - if src_block_ids.size == 0 or dst_block_ids.size == 0: - continue - if not MambaPolicy.is_paired(self._registrar.self_rank_info, peer_ri): - continue - src_region = extractor.extract_slot(int(src_block_ids[0]), self_lg, self_pi) - dst_region = peer_extractor.extract_slot(int(dst_block_ids[0]), peer_lg, peer_pi) - else: - src_block_ids = np.asarray(src_block_ids, dtype=np.int64) - dst_block_ids = np.asarray(dst_block_ids, dtype=np.int64) - if peer_ri.cp_size > 1 and self_ri.cp_size == 1: - # Helix: the receiver owns global blocks [cp_rank::cp_size] - # (same protocol as partition_context_for_helix), so its - # table is the strided subset of ours; block reuse is - # rejected under helix. - src_block_ids = src_block_ids[peer_ri.cp_rank :: peer_ri.cp_size] - tpb = extractor.page_table.tokens_per_block - dst_start_block = ( - req_info.dst_start_token // tpb if req_info.dst_start_token is not None else 0 - ) - # Both tables are positional (index == block ordinal, -1 == not - # held), so pairing is a plain intersection: SWA eviction, prefix - # reuse, and pipelined chunk bounds are all already encoded as - # holes by whichever side owns that state. - src_block_ids, dst_block_ids = Sender._pair_ordinals( - src_block_ids, dst_block_ids, dst_start_block=dst_start_block - ) - if src_block_ids.size == 0: - continue - src_region = extractor.extract(src_block_ids, self_lg, self_pi) - dst_region = peer_extractor.extract(dst_block_ids, peer_lg, peer_pi) - - # Map and collect fragments (same for all layer types) - mapper = self._registrar.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) - region_pair = mapper.map(src_region, dst_region) - region_pairs = region_pair if isinstance(region_pair, list) else [region_pair] - for rp in region_pairs: - src_frag_parts.append(rp.src.memory.ptrs) - dst_frag_parts.append(rp.dst.memory.ptrs) - size_specs.append((rp.src.memory.ptrs.size, rp.src.memory.bytes_per_region)) - - if src_frag_parts: - src_frags = np.concatenate(src_frag_parts) - dst_frags = np.concatenate(dst_frag_parts) - # Build the kv_sizes array in one shot: np.repeat expands each - # bytes_per_region value by its count, e.g.: - # values=[4096, 8192], counts=[100, 200] - # → [4096]*100 ++ [8192]*200 - counts, values = zip(*size_specs) - kv_sizes = np.repeat(np.array(values, dtype=np.int64), counts) - else: - src_frags = np.array([], dtype=np.int64) - dst_frags = np.array([], dtype=np.int64) - kv_sizes = np.array([], dtype=np.int64) if timer: timer.record_prepare_args_end(peer_ri.instance_rank) - timer.record_transfer_sizes(peer_ri.instance_rank, int(kv_sizes.sum()), dst_frags.size) + timer.record_transfer_sizes( + peer_ri.instance_rank, int(plan.sizes.sum()), plan.dst_ptrs.size + ) return WriteMeta( task=task, - src_ptrs=src_frags, - dst_ptrs=dst_frags, - sizes=kv_sizes, - dst_device_id=dst_device_id, + src_ptrs=plan.src_ptrs, + dst_ptrs=plan.dst_ptrs, + sizes=plan.sizes, + dst_device_id=peer_ri.device_id, expected_transfers=expected_transfers, peer_name=peer_ri.instance_name + str(peer_ri.instance_rank), peer_rank=peer_ri.instance_rank, @@ -1836,7 +1733,7 @@ def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> Wri timer = task._perf_timer if timer: timer.record_prepare_args_start(peer_ri.instance_rank) - expected_transfers = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks) + expected_transfers = self._expected_transfers(peer_ri) src_ptrs = np.array([], dtype=np.int64) dst_ptrs = np.array([], dtype=np.int64) @@ -2223,7 +2120,7 @@ def _save_peer_req_info(self, peer_transfer_req_info: RecvReqInfo): req_info = peer_transfer_req_info self._add_req_info(req_info.unique_rid, req_info.instance_rank, req_info) peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) - expected_transfers = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks) + expected_transfers = self._expected_transfers(peer_ri) if self._is_req_ready(req_info.unique_rid, expected_transfers): session = self._get_session(req_info.unique_rid) if session is not None and not session.receiver_ready: @@ -2237,7 +2134,7 @@ def has_all_peer_req_infos(self, unique_rid: int) -> bool: def _has_all_peer_req_infos(self, req_info: RecvReqInfo) -> bool: peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) - expected_transfers = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks) + expected_transfers = self._expected_transfers(peer_ri) return self._is_req_ready(req_info.unique_rid, expected_transfers) def clear_session(self, unique_rid: int): diff --git a/tests/unittest/disaggregated/test_aux_buffer_registration.py b/tests/unittest/disaggregated/test_aux_buffer_registration.py index 956d2c261ca6..198129ca8d72 100644 --- a/tests/unittest/disaggregated/test_aux_buffer_registration.py +++ b/tests/unittest/disaggregated/test_aux_buffer_registration.py @@ -67,10 +67,9 @@ def _build_aux_write_meta( registrar = Mock() registrar.self_rank_info = SimpleNamespace(aux_meta=src_buffer.meta) registrar.get_peer_rank_info.return_value = peer_rank_info - registrar.get_peer_overlap.return_value = SimpleNamespace(ranks=[0]) registrar.should_send_aux.return_value = True registrar.get_aux_transfer_layout.return_value = None - sender = SimpleNamespace(_registrar=registrar) + sender = SimpleNamespace(_registrar=registrar, _expected_transfers=lambda peer_ri: 1) task = SimpleNamespace(_perf_timer=None, _slot=src_slot, _unique_rid=7) req_info = SimpleNamespace( instance_name="peer", instance_rank=1, aux_slot=dst_slot, unique_rid=7 diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index d75c51a6ca77..d0f49ce489e1 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -27,6 +27,7 @@ from tensorrt_llm import DisaggregatedParams from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.transfer import ( AgentResult, KVSendTask, @@ -174,20 +175,30 @@ def _make_projection_sender() -> Sender: ) registrar = MagicMock() - registrar.self_rank_info = SimpleNamespace(cp_size=1, cp_rank=0) + registrar.self_rank_info = SimpleNamespace(cp_size=1, cp_rank=0, dp_rank=0) registrar.self_extractor = extractor registrar.get_peer_rank_info.return_value = peer_ri - registrar.get_peer_overlap.return_value = SimpleNamespace(ranks=[0]) - registrar.should_send_kv.return_value = True - registrar.get_pool_mapping.return_value = { - (0, 0): (0, 0), - (1, 0): (1, 0), - } - registrar.peer_extractor.return_value = extractor - registrar.get_kv_map.return_value = mapper + + # The real planner's plan() over mocked pool matching / mappers, so the + # ordinal pairing under test is the production code path. + planner = DefaultTransferPlanner.__new__(DefaultTransferPlanner) + planner._self_ext_cache = extractor + planner.peer_extractor = MagicMock(return_value=extractor) + planner.get_peer_overlap = MagicMock( + return_value=SimpleNamespace(ranks=[0], duplicate_head_factor=1, overlap_pp_size=1) + ) + planner.should_send_pool = MagicMock(return_value=True) + planner.get_pool_mapping = MagicMock( + return_value={ + (0, 0): (0, 0), + (1, 0): (1, 0), + } + ) + planner.get_kv_map = MagicMock(return_value=mapper) sender = Sender.__new__(Sender) sender._registrar = registrar + sender._planner = planner return sender diff --git a/tests/unittest/disaggregated/test_transfer_planner.py b/tests/unittest/disaggregated/test_transfer_planner.py index 39e2b0211b31..b2e9f0d2bf8b 100644 --- a/tests/unittest/disaggregated/test_transfer_planner.py +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -39,3 +39,120 @@ def test_transfer_planner_is_abstract(): def test_peer_incompatible_error_is_value_error(): assert issubclass(PeerIncompatibleError, ValueError) assert PeerTargets(ranks=[0], uniform_writers=True).ranks == [0] + + +# --------------------------------------------------------------------------- # +# Sender drives the planner +# --------------------------------------------------------------------------- # +from types import SimpleNamespace # noqa: E402 +from unittest.mock import Mock # noqa: E402 + +from tensorrt_llm import DisaggregatedParams # noqa: E402 +from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, TokenRange # noqa: E402 +from tensorrt_llm._torch.disaggregation.native.transfer import ( # noqa: E402 + KVSendTask, + RecvReqInfo, + Sender, +) + + +class _ToyPlanner(TransferPlanner): + """Fixed ranks and a fixed plan; records what the framework handed it.""" + + planner_id = "toy" + + def __init__(self, plan: TransferPlan, ranks=(0,)): + self._plan = plan + self._ranks = list(ranks) + self.calls: list = [] + + def check_peer(self, self_ri, peer_ri): + return None + + def targets(self, self_ri, peer_ri, peer_dp_rank): + return PeerTargets(ranks=list(self._ranks), uniform_writers=True) + + def plan( + self, self_ri, peer_ri, src, dst, token_range, prompt_len, dst_start_token, excluded + ): + self.calls.append((src, dst, token_range, prompt_len, dst_start_token, excluded)) + return self._plan + + def receiver_payload_bytes(self, self_ri, peer_ri, dst): + return 0 + + +def _sender_with(planner: TransferPlanner) -> Sender: + peer_ri = SimpleNamespace( + instance_name="gen", + instance_rank=3, + dp_rank=0, + device_id=1, + cp_size=1, + self_endpoint="tcp://gen:1", + page_table=None, + ) + sender = Sender.__new__(Sender) + sender._registrar = Mock() + sender._registrar.self_rank_info = SimpleNamespace( + cp_size=1, instance_name="ctx", instance_rank=0 + ) + sender._registrar.get_peer_rank_info.return_value = peer_ri + sender._planner = planner + return sender + + +def _task_and_req(): + task = KVSendTask( + Chunk( + block_ids_per_layer_groups=[np.array([5, 6], dtype=np.int64)], + kind_per_layer_group=[CacheKind.PAGED], + token_range=TokenRange(0, 64), + is_last=True, + ), + DisaggregatedParams(disagg_request_id=1), + slice_id=0, + prompt_len=64, + ) + req = RecvReqInfo( + sender_req_id=1, + instance_name="gen", + instance_rank=3, + block_ids_per_layer_groups=[np.array([9, 10], dtype=np.int64)], + unique_rid=1, + dst_start_token=32, + slice_id=0, + ) + return task, req + + +def test_build_kv_write_meta_uses_planner_plan(): + plan = TransferPlan( + src_ptrs=np.array([100, 200], dtype=np.int64), + dst_ptrs=np.array([300, 400], dtype=np.int64), + sizes=np.array([16, 16], dtype=np.int64), + ) + planner = _ToyPlanner(plan, ranks=(3, 4)) + sender = _sender_with(planner) + task, req = _task_and_req() + + wm = sender._build_kv_write_meta(task, req) + + assert wm.src_ptrs is plan.src_ptrs and wm.dst_ptrs is plan.dst_ptrs and wm.sizes is plan.sizes + assert wm.expected_transfers == 2 and task.expected_transfers == 2 + assert wm.dst_device_id == 1 and wm.peer_rank == 3 and wm.peer_endpoint == "tcp://gen:1" + ((src, dst, token_range, prompt_len, dst_start, excluded),) = planner.calls + assert src is task._chunk.block_ids_per_layer_groups + assert dst is req.block_ids_per_layer_groups + assert excluded is task._chunk.excluded_pool_views + assert (token_range.start, token_range.end, prompt_len, dst_start) == (0, 64, 64, 32) + + +def test_build_kv_write_meta_empty_plan_still_delivers(): + sender = _sender_with(_ToyPlanner(TransferPlan.empty())) + task, req = _task_and_req() + + wm = sender._build_kv_write_meta(task, req) + + assert wm.sizes.size == 0 and wm.src_ptrs.size == 0 and wm.dst_ptrs.size == 0 + assert wm.expected_transfers == 1 and wm.is_last_slice is True diff --git a/tests/unittest/disaggregated/test_unsubmitted_session.py b/tests/unittest/disaggregated/test_unsubmitted_session.py index 09eb9bfa2296..8ce7cd150328 100644 --- a/tests/unittest/disaggregated/test_unsubmitted_session.py +++ b/tests/unittest/disaggregated/test_unsubmitted_session.py @@ -61,7 +61,8 @@ def _sender(*, ownership: bool = True) -> transfer_mod.Sender: sender._registrar.get_peer_rank_info.return_value = SimpleNamespace( instance_name="gen", instance_rank=0, dp_rank=0, self_endpoint="tcp://gen:1234" ) - sender._registrar.get_peer_overlap.return_value = SimpleNamespace(ranks=[0]) + sender._planner = Mock() + sender._planner.targets.return_value = SimpleNamespace(ranks=[0], uniform_writers=True) return sender From bec115d9a5891370a4751c9bd4d02a0e9d2f787d Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 06:29:52 -0700 Subject: [PATCH 08/25] [None][refactor] Receiver dispatches through the planner; transfer.py no longer knows cache kinds The receiver asks the planner which peer ranks to request from, whether those writers are uniform enough for a fan-in bounce split, how many extra bytes ride along, and whether a newly met context peer is compatible. The attention/Mamba policies, CacheKind and MapperKind are gone from the transfer module; a test greps for them. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../_torch/disaggregation/native/peer.py | 9 - .../_torch/disaggregation/native/transfer.py | 175 +++++------------- .../disaggregated/test_cache_reuse_adapter.py | 24 +-- .../test_transfer_ownership_regressions.py | 31 +++- .../disaggregated/test_transfer_planner.py | 58 ++++++ 5 files changed, 140 insertions(+), 157 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 9420690a085b..29eee2e13d72 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -99,15 +99,6 @@ def self_rank_info(self) -> RankInfo: def _unique_key(self, name: str, rank: int) -> str: return name + str(rank) - def validate_bounded_replay_compatible(self, peer_ri: RankInfo) -> None: - """Require the same generation replay mode on both peers.""" - if self._ri.bounded_replay_on_generation != peer_ri.bounded_replay_on_generation: - raise ValueError( - "PeerRegistrar: bounded_replay_on_generation must match between peers " - f"(local={self._ri.bounded_replay_on_generation}, " - f"peer={peer_ri.bounded_replay_on_generation})." - ) - def should_send_aux(self, peer_rank_info: RankInfo) -> bool: # to ensure the transfer aux is not duplicated diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 7da13939e56d..9c0b48b5a125 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -62,16 +62,11 @@ get_non_empty_aux_indices, ) 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 ( - MambaPolicy, - mamba_receiver_payload_bytes, -) -from tensorrt_llm._torch.disaggregation.native.default_planner import PeerOverlap -from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.base.planner import ( # noqa: F401 (re-exported) PeerIncompatibleError, + PeerTargets, ) +from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.native.retirement import ( @@ -82,7 +77,6 @@ from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager @@ -2871,6 +2865,7 @@ def __init__( retirement_watchdog: Optional[RetirementWatchdog] = None, ): self._registrar = peer_registrar + self._planner = peer_registrar.planner self._agent = agent self._bounce = bounce # Internal component gate only. Production wiring is added after the @@ -3023,56 +3018,6 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: slice_id=task.slice_id, ) - @staticmethod - def _fanin_bounce_safe( - overlap: PeerOverlap, - peer_ri: RankInfo, - receiver_page_table: Optional[KVCachePageTable], - ) -> bool: - """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. - The split assumes every writer contributes the same size, which holds when: - * duplicate_head_factor == 1 -- else some ranks don't send KV yet still count in - expected_transfers, so the live writers overflow their slots; - * the PP layer split is even -- a single PP stage (overlap_pp_size <= 1) is trivially fine; - for PP fan-in, every overlapping stage must hold the same number of layers - (peer_ri.layer_num_per_pp all-equal) or per-writer sizes differ. If that full per-stage - list isn't available (shorter than the fan-in degree), be conservative and fall back. - Otherwise fall back to the per-fragment path (correct, just not coalesced). - Equal layer count means equal bytes only when the per-block sizes match; reserve() rejects - the mismatched case, so this only needs the count to split evenly.""" - if overlap.duplicate_head_factor != 1: - return False - if overlap.overlap_pp_size > 1: - lpp = getattr(peer_ri, "layer_num_per_pp", None) - if not lpp or len(lpp) < overlap.overlap_pp_size or len(set(lpp)) != 1: - return False - # Replicated pools (e.g. MiniMax M3 index-key) are sent by one elected - # fan-in owner only, so with multiple writers their contributions - # differ in size and the equal split is invalid. Inspect both endpoints: - # a masked PP stage may advertise no replicated view even though another - # stage owns one that is visible in the receiver's page table. - if len(overlap.ranks) > 1: - for page_table in (peer_ri.page_table, receiver_page_table): - if page_table is None: - continue - for layer_group in page_table.layer_groups: - for pool_view in getattr(layer_group, "pool_views", ()): - if pool_view.mapper_kind == MapperKind.REPLICATED: - return False - return True - - def _get_mamba_slot(self, req_info: RecvReqInfo) -> Optional[int]: - """Extract the mamba slot index from block_ids_per_layer_groups, or None.""" - pt = self._registrar.self_extractor.page_table - if pt is None: - return None - for lg_idx, lg in enumerate(pt.layer_groups): - if lg.kind == CacheKind.STATE: - block_ids = req_info.block_ids_per_layer_groups[lg_idx] - if block_ids.size > 0: - return int(block_ids[0]) - return None - def dispatch_task(self, task: KVRecvTask) -> None: params = task._params logger.debug( @@ -3107,64 +3052,43 @@ def dispatch_task(self, task: KVRecvTask) -> None: task.fail(e) return + self_ri = self._registrar.self_rank_info if sender_dp_rank is not None: # Normal path: ctx_dp_rank is known, send to overlapping ranks. - peer_overlap = self._registrar.get_peer_overlap(peer_infos, sender_dp_rank) + peer_targets = self._planner.targets(self_ri, peer_infos, sender_dp_rank) + dp0_targets = peer_targets else: # Gen-first with ADP: ctx_dp_rank unknown — broadcast REQUEST_DATA - # to ALL ctx sender ranks so every DP group receives it. - # get_peer_overlap returns ranks for one DP group (topology is - # symmetric), so use dp_rank=0 as representative. - dp_size = peer_infos.dp_size - dp0_overlap = self._registrar.get_peer_overlap(peer_infos, 0) - # Union of overlapping ranks across all DP groups for broadcast (deduplicated) - all_ranks_set: set[int] = set(dp0_overlap.ranks) - for dp in range(1, dp_size): - all_ranks_set.update(self._registrar.get_peer_overlap(peer_infos, dp).ranks) + # to ALL ctx sender ranks so every DP group receives it. The real + # writer count exceeds the per-group count, so the union is never + # a uniform writer set. + dp0_targets = self._planner.targets(self_ri, peer_infos, 0) + all_ranks_set: set[int] = set(dp0_targets.ranks) + for dp in range(1, peer_infos.dp_size): + all_ranks_set.update(self._planner.targets(self_ri, peer_infos, dp).ranks) all_ranks = list(all_ranks_set) logger.debug( - f"Receiver.dispatch_task: ADP broadcast path, dp_size={dp_size}, " + f"Receiver.dispatch_task: ADP broadcast path, dp_size={peer_infos.dp_size}, " f"all_ranks={all_ranks}" ) - peer_overlap = type(dp0_overlap)(ranks=all_ranks) + peer_targets = PeerTargets(ranks=all_ranks, uniform_writers=False) - # In gen-first ADP broadcast, peer_overlap contains the union of all DP - # groups, but expected_transfers should reflect per-DP-group count since - # only one DP group will actually process the context request. - if sender_dp_rank is not None: - task.expected_transfers = len(peer_overlap.ranks) - else: - task.expected_transfers = len(dp0_overlap.ranks) - # TP fan-in splits ONE region equally, so allow it only for a uniform writer set: - # _fanin_bounce_safe() (TP-by-head / even-PP), and never under ADP broadcast (sender_dp_rank - # None), where the real writer count exceeds expected_transfers and would overflow the slot. - topo_overlap = peer_overlap if sender_dp_rank is not None else dp0_overlap - # Helix writers contribute unequal shares (one KV owner, one state - # sender, the rest empty), so the fan-in equal split would overflow - # into neighboring reservations; use the per-fragment path. - cp_involved = self._registrar.self_rank_info.cp_size > 1 or peer_infos.cp_size > 1 + # expected_transfers reflects one DP group even under broadcast: only + # one group will actually process the context request. + task.expected_transfers = len(dp0_targets.ranks) + # TP fan-in splits ONE region equally, so allow it only when the planner says every + # writer contributes the same share, and never under ADP broadcast. allow_bounce = task.expected_transfers == 1 or ( - sender_dp_rank is not None - and not cp_involved - and self._fanin_bounce_safe( - topo_overlap, - peer_infos, - self._registrar.self_extractor.page_table, - ) + sender_dp_rank is not None and peer_targets.uniform_writers ) - # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks, - # so the bounce region must be sized for those bytes too. - extra_bytes = 0 - mamba_dst_slot = self._get_mamba_slot(receiver_req) - if mamba_dst_slot is not None: - if peer_infos.page_table is None: - allow_bounce = False # cannot size the recurrent-state payload - else: - extra_bytes = mamba_receiver_payload_bytes( - sender_page_table=peer_infos.page_table, - receiver_page_table=self._registrar.self_extractor.page_table, - dst_slot=mamba_dst_slot, - ) + # Bytes beyond the paged blocks that ride the same coalesced write + # (recurrent state today) size the bounce region too; None = unknown. + extra_bytes = self._planner.receiver_payload_bytes( + self_ri, peer_infos, receiver_req.block_ids_per_layer_groups + ) + if extra_bytes is None: + allow_bounce = False + extra_bytes = 0 bounced = allow_bounce and self._bounce.reserve( receiver_req, task.expected_transfers, extra_bytes=extra_bytes ) @@ -3174,7 +3098,7 @@ def dispatch_task(self, task: KVRecvTask) -> None: f"dispatch_task: RxSession {task._unique_rid} not found; " "session may have been closed before dispatch" ) - sender_endpoints = {peer_infos.sender_endpoints[rank] for rank in peer_overlap.ranks} + sender_endpoints = {peer_infos.sender_endpoints[rank] for rank in peer_targets.ranks} # Fan-in: each sender gets its own sub-region base (writers must not overwrite); else serialize once. fanin_bounce = bounced and task.expected_transfers > 1 key = (receiver_req.unique_rid, receiver_req.slice_id) @@ -3183,7 +3107,7 @@ def dispatch_task(self, task: KVRecvTask) -> None: session.mark_transferring(task.slice_id) # Cache sender endpoints so cancel() can send CANCEL_SESSION to them. session._sender_endpoints.update(sender_endpoints) - for i, rank in enumerate(peer_overlap.ranks): + for i, rank in enumerate(peer_targets.ranks): if task._perf_timer is not None: task._perf_timer.record_task_start(rank) if fanin_bounce: @@ -3192,7 +3116,7 @@ def dispatch_task(self, task: KVRecvTask) -> None: self._request_sender_data(peer_infos.sender_endpoints[rank], receiver_req_bytes) return - peer_ranks = list(peer_overlap.ranks) + peer_ranks = list(peer_targets.ranks) serialized_requests: list[tuple[int, bytes]] = [] for i, rank in enumerate(peer_ranks): if fanin_bounce: @@ -3271,27 +3195,22 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: finally: messenger.stop() - # Peer compatibility gates on the receiver side. - # The sender-side check (PeerRegistrar.register) runs in the - # sender's listener thread, where exceptions are only logged, so - # reject here — before REGISTER_RANK_INFO is even sent, so no - # dealers are connected and no partial registration happens for - # the bad peer. The failure is converted to PeerIncompatibleError - # (handled in dispatch_task) so only requests targeting this peer - # fail, and cached so later requests fail fast. + # Receiver-side gate. The sender-side check (PeerRegistrar.register) + # runs in the sender's listener thread, where exceptions are only + # logged, so reject here — before REGISTER_RANK_INFO is even sent, + # so no dealers are connected and no partial registration happens + # for the bad peer. Only requests targeting this peer fail + # (PeerIncompatibleError is handled in dispatch_task), and the + # verdict is cached so later requests fail fast. try: - self._registrar.validate_bounded_replay_compatible(sender_info) - MambaPolicy.validate_peer_compatible( - self._registrar.self_rank_info, - sender_info, - self._registrar.self_extractor.page_table, - sender_info.page_table, - ) - AttentionPolicy.validate_peer_compatible( - self._registrar.self_extractor.page_table, - sender_info.page_table, - ) - except ValueError as e: + if sender_info.planner_id != self._planner.planner_id: + raise PeerIncompatibleError( + f"context peer plans transfers with '{sender_info.planner_id}', " + f"this rank with '{self._planner.planner_id}'; both sides of a " + "disaggregated deployment must run the same model planner" + ) + self._planner.check_peer(self._registrar.self_rank_info, sender_info) + except PeerIncompatibleError as e: msg = ( f"context peer at '{info_endpoint}' is incompatible: {e} " "(cached: all further requests to this context peer fail " diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index 2a8ca61c646a..d4854b09d5d2 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -12,7 +12,7 @@ # 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. -"""Tests for CacheReuseAdapter, the positional _create_chunk, and Sender ordinal pairing.""" +"""Tests for CacheReuseAdapter, the positional _create_chunk, and DefaultTransferPlanner ordinal pairing.""" from types import SimpleNamespace from unittest.mock import MagicMock @@ -20,7 +20,7 @@ import numpy as np import pytest -from tensorrt_llm._torch.disaggregation.native.transfer import Sender +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.resource.cache_reuse import ( CacheReuseAdapter, _CacheReuseAdapterV1, @@ -37,7 +37,7 @@ def _arr(values): # --------------------------------------------------------------------------- -# Sender._pair_ordinals: the whole alignment protocol is one intersection. +# DefaultTransferPlanner._pair_ordinals: the whole alignment protocol is one intersection. # --------------------------------------------------------------------------- @@ -50,32 +50,32 @@ class TestPairOrdinals: """ def test_identity(self): - src, dst = Sender._pair_ordinals(_arr([10, 11, 12]), _arr([20, 21, 22])) + src, dst = DefaultTransferPlanner._pair_ordinals(_arr([10, 11, 12]), _arr([20, 21, 22])) np.testing.assert_array_equal(src, [10, 11, 12]) np.testing.assert_array_equal(dst, [20, 21, 22]) def test_receiver_prefix_reuse_is_a_dst_hole(self): - src, dst = Sender._pair_ordinals(_arr([10, 11, 12, 13]), _arr([-1, -1, 22, 23])) + src, dst = DefaultTransferPlanner._pair_ordinals(_arr([10, 11, 12, 13]), _arr([-1, -1, 22, 23])) np.testing.assert_array_equal(src, [12, 13]) np.testing.assert_array_equal(dst, [22, 23]) def test_sender_chunk_bounds_are_src_holes(self): - src, dst = Sender._pair_ordinals(_arr([-1, 11, 12, -1]), _arr([20, 21, 22, 23])) + src, dst = DefaultTransferPlanner._pair_ordinals(_arr([-1, 11, 12, -1]), _arr([20, 21, 22, 23])) np.testing.assert_array_equal(src, [11, 12]) np.testing.assert_array_equal(dst, [21, 22]) def test_holes_on_both_sides_take_the_intersection(self): - src, dst = Sender._pair_ordinals(_arr([10, -1, 12, 13]), _arr([-1, 21, 22, -1])) + src, dst = DefaultTransferPlanner._pair_ordinals(_arr([10, -1, 12, 13]), _arr([-1, 21, 22, -1])) np.testing.assert_array_equal(src, [12]) np.testing.assert_array_equal(dst, [22]) def test_no_overlap_is_empty(self): - src, dst = Sender._pair_ordinals(_arr([10, 11, -1]), _arr([-1, -1, 22])) + src, dst = DefaultTransferPlanner._pair_ordinals(_arr([10, 11, -1]), _arr([-1, -1, 22])) assert src.size == 0 assert dst.size == 0 def test_dst_start_block_masks_the_head(self): - src, dst = Sender._pair_ordinals( + src, dst = DefaultTransferPlanner._pair_ordinals( _arr([10, 11, 12, 13]), _arr([20, 21, 22, 23]), dst_start_block=2 ) np.testing.assert_array_equal(src, [12, 13]) @@ -83,7 +83,7 @@ def test_dst_start_block_masks_the_head(self): def test_length_mismatch_raises(self): with pytest.raises(ValueError, match="differ in length"): - Sender._pair_ordinals(_arr([10, 11]), _arr([20, 21, 22])) + DefaultTransferPlanner._pair_ordinals(_arr([10, 11]), _arr([20, 21, 22])) def test_sender_not_yet_evicted_but_receiver_has(self): # Regression (DeepSeek-V4 SWA, prompt_len == 2 blocks, window == 1 @@ -91,14 +91,14 @@ def test_sender_not_yet_evicted_but_receiver_has(self): # advanced, while gen pre-declared history_length=prompt_len and never # allocated ordinal 0. Pairing by ordinal copies X1 -> G1; the old # length-based suffix arithmetic copied X0 -> G1. - src, dst = Sender._pair_ordinals(_arr([100, 101]), _arr([-1, 201])) + src, dst = DefaultTransferPlanner._pair_ordinals(_arr([100, 101]), _arr([-1, 201])) np.testing.assert_array_equal(src, [101]) np.testing.assert_array_equal(dst, [201]) def test_receiver_with_wider_speculative_window(self): # Only gen runs spec decoding, so its SWA window keeps one more block. # That block has no source; it simply is not written. - src, dst = Sender._pair_ordinals(_arr([-1, -1, 12, 13]), _arr([-1, 21, 22, 23])) + src, dst = DefaultTransferPlanner._pair_ordinals(_arr([-1, -1, 12, 13]), _arr([-1, 21, 22, 23])) np.testing.assert_array_equal(src, [12, 13]) np.testing.assert_array_equal(dst, [22, 23]) diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index cfa6213af139..5d1ca77a0a0c 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -28,6 +28,7 @@ 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.planner import PeerTargets from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult from tensorrt_llm._torch.disaggregation.native.bounce.core import TransferContext @@ -586,10 +587,12 @@ def test_gen_first_no_retry_adp_count_seal_waits_for_one_writer_group() -> None: receiver._sessions = {} receiver._pre_cancelled_rids = {} receiver._bounce = _BounceProbe() + receiver._bounce.reserve = Mock(wraps=receiver._bounce.reserve) receiver._enforce_physical_ownership = True receiver._shutdown = False receiver_req = SimpleNamespace( unique_rid=rid, + block_ids_per_layer_groups=[], slice_id=0, bounce_dst_base=None, to_bytes=Mock(return_value=b"receiver-request"), @@ -600,14 +603,20 @@ def test_gen_first_no_retry_adp_count_seal_waits_for_one_writer_group() -> None: 1: SimpleNamespace(ranks=[2, 3]), } receiver._registrar = SimpleNamespace( - get_peer_overlap=Mock(side_effect=lambda _peer, dp_rank: overlaps[dp_rank]), - self_extractor=SimpleNamespace(page_table=None), self_rank_info=SimpleNamespace( cp_size=1, instance_name="gen", instance_rank=0, ), ) + receiver._planner = SimpleNamespace( + targets=Mock( + side_effect=lambda _self_ri, _peer_ri, dp_rank: PeerTargets( + ranks=overlaps[dp_rank].ranks, uniform_writers=True + ) + ), + receiver_payload_bytes=Mock(return_value=0), + ) receiver._get_sender_info = Mock( return_value=SimpleNamespace( sender_endpoints={rank: f"tcp://sender-{rank}" for rank in range(4)}, @@ -632,6 +641,8 @@ def test_gen_first_no_retry_adp_count_seal_waits_for_one_writer_group() -> None: session.dispatch_prepared_receive(task) assert task.expected_transfers == 2 + # The broadcast union is never a uniform writer set: fan-in bounce must not be attempted. + receiver._bounce.reserve.assert_not_called() assert {call.args[0] for call in receiver._request_sender_data.call_args_list} == { f"tcp://sender-{rank}" for rank in range(4) } @@ -684,6 +695,7 @@ def test_partial_bounced_publication_waits_for_queued_writer_success( receiver._dealers = {} receiver_req = SimpleNamespace( unique_rid=rid, + block_ids_per_layer_groups=[], slice_id=0, mamba_state_index=None, bounce_dst_base=None, @@ -696,10 +708,12 @@ def test_partial_bounced_publication_waits_for_queued_writer_success( overlap_pp_size=1, ) receiver._registrar = SimpleNamespace( - get_peer_overlap=Mock(return_value=overlap), - self_extractor=SimpleNamespace(page_table=None), self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0, cp_size=1), ) + receiver._planner = SimpleNamespace( + targets=Mock(return_value=PeerTargets(ranks=overlap.ranks, uniform_writers=True)), + receiver_payload_bytes=Mock(return_value=0), + ) receiver._get_sender_info = Mock( return_value=SimpleNamespace( sender_endpoints={0: "tcp://sender-0", 1: "tcp://sender-1"}, @@ -861,6 +875,7 @@ def test_cancel_after_publication_cannot_overtake_request_data( receiver._build_recv_req_info = Mock( return_value=SimpleNamespace( unique_rid=rid, + block_ids_per_layer_groups=[], slice_id=0, mamba_state_index=None, bounce_dst_base=None, @@ -868,10 +883,10 @@ def test_cancel_after_publication_cannot_overtake_request_data( ) ) overlap = SimpleNamespace(ranks=[0]) - receiver._registrar = SimpleNamespace( - get_peer_overlap=Mock(return_value=overlap), - self_extractor=SimpleNamespace(page_table=None), - self_rank_info=SimpleNamespace(cp_size=1), + receiver._registrar = SimpleNamespace(self_rank_info=SimpleNamespace(cp_size=1)) + receiver._planner = SimpleNamespace( + targets=Mock(return_value=PeerTargets(ranks=overlap.ranks, uniform_writers=True)), + receiver_payload_bytes=Mock(return_value=0), ) receiver._get_sender_info = Mock( return_value=SimpleNamespace( diff --git a/tests/unittest/disaggregated/test_transfer_planner.py b/tests/unittest/disaggregated/test_transfer_planner.py index b2e9f0d2bf8b..c0114d987bc7 100644 --- a/tests/unittest/disaggregated/test_transfer_planner.py +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -156,3 +156,61 @@ def test_build_kv_write_meta_empty_plan_still_delivers(): assert wm.sizes.size == 0 and wm.src_ptrs.size == 0 and wm.dst_ptrs.size == 0 assert wm.expected_transfers == 1 and wm.is_last_slice is True + + +# --------------------------------------------------------------------------- # +# Receiver drives the planner; transfer.py knows no cache kinds +# --------------------------------------------------------------------------- # +import inspect # noqa: E402 + +import tensorrt_llm._torch.disaggregation.native.transfer as transfer_mod # noqa: E402 +from tensorrt_llm._torch.disaggregation.native.transfer import Receiver # noqa: E402 + + +def test_transfer_module_has_no_model_knowledge(): + src = inspect.getsource(transfer_mod) + for forbidden in ("mixers", "CacheKind", "MapperKind", "MambaPolicy", "AttentionPolicy"): + assert forbidden not in src, f"native/transfer.py mentions {forbidden}" + + +def _receiver_with(planner) -> tuple[Receiver, SimpleNamespace]: + recv = Receiver.__new__(Receiver) + recv._enforce_physical_ownership = False + recv._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(cp_size=1, instance_name="gen", instance_rank=0) + ) + recv._planner = planner + recv._bounce = Mock() + req = SimpleNamespace( + unique_rid=7, + slice_id=0, + block_ids_per_layer_groups=[np.array([1], dtype=np.int64)], + bounce_dst_base=None, + to_bytes=Mock(return_value=b"req"), + ) + recv._build_recv_req_info = Mock(return_value=req) + recv._get_sender_info = Mock( + return_value=SimpleNamespace( + sender_endpoints={0: "tcp://s0"}, page_table=None, dp_size=1, cp_size=1 + ) + ) + session = SimpleNamespace(mark_transferring=Mock(), _sender_endpoints=set()) + recv._get_session = Mock(return_value=session) + recv._request_sender_data = Mock() + task = SimpleNamespace( + _params=SimpleNamespace(ctx_dp_rank=0), _unique_rid=7, slice_id=0, _perf_timer=None + ) + return recv, task + + +def test_dispatch_task_payload_none_disables_bounce(): + planner = Mock() + planner.targets.return_value = PeerTargets(ranks=[0], uniform_writers=True) + planner.receiver_payload_bytes.return_value = None + recv, task = _receiver_with(planner) + + recv.dispatch_task(task) + + recv._bounce.reserve.assert_not_called() + recv._request_sender_data.assert_called_once_with("tcp://s0", b"req") + assert task.expected_transfers == 1 From 81bdecfc1104cad177490b912e88a785c63edc92 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 08:33:28 -0700 Subject: [PATCH 09/25] [None][feat] Let a model class supply its disaggregated TransferPlanner TransferWorker builds the planner through a factory and stamps planner_id and planner_blob on the rank's RankInfo; DecoderModelForCausalLM provides the default factory, and a modeling_v2 target overrides it in its own directory. Plumbed from the executor through create_kv_cache_transceiver and KvCacheTransceiverV2. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../disaggregation/kv_cache_transceiver.py | 10 +- .../_torch/disaggregation/native/transfer.py | 19 +++- .../_torch/disaggregation/transceiver.py | 3 + tensorrt_llm/_torch/models/modeling_utils.py | 12 +++ tensorrt_llm/_torch/pyexecutor/_util.py | 3 +- .../disaggregated/test_kv_transfer.py | 5 +- .../disaggregated/test_transfer_planner.py | 102 ++++++++++++++++++ 7 files changed, 148 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py index e93b8a819c59..83a17040fc03 100644 --- a/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py @@ -148,7 +148,8 @@ def create_kv_cache_transceiver( kv_cache_manager: KVCacheManager, attention_type: AttentionTypeCpp, cache_transceiver_config: CacheTransceiverConfig, - mamba_cache_manager: Optional[BaseMambaCacheManager] = None + mamba_cache_manager: Optional[BaseMambaCacheManager] = None, + model: Optional[object] = None, ) -> Optional["KvCacheTransceiver"]: resolve_cache_transceiver_config(cache_transceiver_config) if cache_transceiver_config is None or cache_transceiver_config.backend is None: @@ -224,8 +225,13 @@ def create_kv_cache_transceiver( KvCacheTransceiverV2 logger.info("Using KvCacheTransceiverV2") # MixedMambaHybridCacheManager contains both the KV and Mamba pools. + # The model class may plan its own KV transfers (see + # DecoderModelForCausalLM.build_transfer_planner); anything without + # the hook gets the default planner. + planner_factory = getattr(model, "build_transfer_planner", None) return KvCacheTransceiverV2(mapping, dist, kv_cache_manager, - cache_transceiver_config) + cache_transceiver_config, + planner_factory=planner_factory) # Explicit legacy runtime: transceiver_runtime is None or "CPP". return BindKvCacheTransceiver(mapping, dist, kv_cache_manager, diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 9c0b48b5a125..864ad30402b7 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -65,7 +65,9 @@ from tensorrt_llm._torch.disaggregation.base.planner import ( # noqa: F401 (re-exported) PeerIncompatibleError, PeerTargets, + TransferPlanner, ) +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo @@ -4323,6 +4325,12 @@ class TransferWorkerConfig: # trailing _BYTES, lowercased); dict > env > default. Ignored when # agent_buffer_size_mb == 0. agent_bounce_params: Optional[Dict[str, str]] = None + # Builds the model's TransferPlanner from what the worker derives from the + # cache manager. None -> DefaultTransferPlanner. + planner_factory: Optional["TransferPlannerFactory"] = None + + +TransferPlannerFactory = Callable[[KVCacheManager, RankInfo, KVRegionExtractorV1], TransferPlanner] class TransferWorker: @@ -4411,7 +4419,16 @@ def sweep_stale_req_infos(self): def _setup_peer_infrastructure(self, kvm: KVCacheManager): self._rank_info_server = RankInfoServer(self._rank_info) if kvm.mapping.rank == 0 else None self._kv_extractor = KVRegionExtractorV1(kvm) - self._peer_registrar = PeerRegistrar(self._rank_info, self._kv_extractor) + factory = self._config.planner_factory or ( + lambda _kvm, ri, ext: DefaultTransferPlanner(ri, ext) + ) + self._planner = factory(kvm, self._rank_info, self._kv_extractor) + # The envelope says who planned this rank's transfers and carries what + # that planner wants the peer's planner to know. RankInfoServer holds + # the object and serializes on request, so setting it here is in time. + self._rank_info.planner_id = self._planner.planner_id + self._rank_info.planner_blob = self._planner.describe() + self._peer_registrar = PeerRegistrar(self._rank_info, self._kv_extractor, self._planner) def _setup_transfer_engine(self): torch.cuda.set_device(self._config.device_id) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 3c9e4f245db7..6116fa5447e6 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -48,6 +48,7 @@ from tensorrt_llm._torch.disaggregation.native.retirement import QuiescenceFatalEvent from tensorrt_llm._torch.disaggregation.native.transfer import ( RxSession, + TransferPlannerFactory, TransferWorker, TransferWorkerConfig, ) @@ -241,6 +242,7 @@ def __init__( dist: Distributed, kv_cache_manager: KVCacheManager, cache_transceiver_config: CacheTransceiverConfig, + planner_factory: Optional["TransferPlannerFactory"] = None, ): self._shutdown_lock = threading.Lock() self._shutdown_complete = False @@ -317,6 +319,7 @@ def __init__( else 0 ), agent_bounce_params=cache_transceiver_config.agent_bounce_params, + planner_factory=planner_factory, ) ) if enforce_physical_ownership: diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 0a6060b47288..b1e0a339b4fa 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -801,6 +801,18 @@ def cache_derived_state(self) -> None: None. """ + def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): + """The disaggregated KV transfer plan for this model on this rank. + + A model whose cache layout the default planner does not understand + returns its own ``TransferPlanner`` here; the context and the + generation deployment of that model then run the same planner. + """ + from tensorrt_llm._torch.disaggregation.native.default_planner import \ + DefaultTransferPlanner + + return DefaultTransferPlanner(self_rank_info, self_extractor) + def post_load_weights(self) -> None: """Run the default staged post-load hook sequence. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7aedaf8a81e6..f21f28a739b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -4304,7 +4304,8 @@ def create_py_executor_instance( kv_cache_transceiver = create_kv_cache_transceiver( mapping, dist, kv_cache_manager, attention_type, - cache_transceiver_config, mamba_cache_manager) + cache_transceiver_config, mamba_cache_manager, + model=model_engine.model) waiting_queue_policy = (scheduler_config.waiting_queue_policy if scheduler_config is not None else diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 88d4623289d8..539a0bd1eedc 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -42,6 +42,7 @@ import torch import tensorrt_llm +import tensorrt_llm._torch.disaggregation.native.default_planner as default_planner_mod import tensorrt_llm._torch.disaggregation.native.transfer as transfer_mod import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm @@ -1943,7 +1944,7 @@ def test_incompatible_peer_fails_only_affected_requests( ) -> None: """An incompatible context peer must fail only the requests targeting it. - When peer validation rejects a peer during first + When the planner rejects a peer during first registration, the receiver must (a) not raise out of receive() — the session fails through the normal transfer-error path (WaitResult.FAILED) with the diagnostic preserved; (b) cache the incompatibility so later @@ -2017,7 +2018,7 @@ def raiser(*args, **kwargs): try: with pytest.MonkeyPatch.context() as mp: - mp.setattr(transfer_mod.MambaPolicy, "validate_peer_compatible", staticmethod(raiser)) + mp.setattr(default_planner_mod.MambaPolicy, "validate_peer_compatible", staticmethod(raiser)) # (a) First request to the incompatible peer: receive() must not # raise; the session fails with the diagnostic preserved. diff --git a/tests/unittest/disaggregated/test_transfer_planner.py b/tests/unittest/disaggregated/test_transfer_planner.py index c0114d987bc7..28380b6281c4 100644 --- a/tests/unittest/disaggregated/test_transfer_planner.py +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -214,3 +214,105 @@ def test_dispatch_task_payload_none_disables_bounce(): recv._bounce.reserve.assert_not_called() recv._request_sender_data.assert_called_once_with("tcp://s0", b"req") assert task.expected_transfers == 1 + + +# --------------------------------------------------------------------------- # +# Same planner on both sides; the worker takes the model's planner +# --------------------------------------------------------------------------- # +from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar # noqa: E402 +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo # noqa: E402 +from tensorrt_llm._torch.disaggregation.native.transfer import TransferWorker # noqa: E402 + + +def _minimal_rank_info(**overrides) -> RankInfo: + fields = dict( + instance_name="ctx", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + self_endpoint="", + transfer_engine_info=b"", + ) + fields.update(overrides) + return RankInfo(**fields) + + +def test_registrar_rejects_planner_id_mismatch(): + self_ri = _minimal_rank_info(planner_id="toy") + peer_ri = _minimal_rank_info(instance_name="gen", planner_id="trtllm.default") + reg = PeerRegistrar(self_ri, Mock(), _ToyPlanner(TransferPlan.empty())) + + with pytest.raises(PeerIncompatibleError, match="toy"): + reg.register("gen", 0, peer_ri) + with pytest.raises(KeyError): + reg.get_peer_rank_info("gen", 0) + + +def test_receiver_rejects_planner_id_mismatch(monkeypatch): + recv = Receiver.__new__(Receiver) + recv._incompatible_peers = {} + recv._sender_ep_instance_map = {} + recv._registrar = SimpleNamespace(self_rank_info=_minimal_rank_info(planner_id="toy")) + recv._planner = _ToyPlanner(TransferPlan.empty()) + sender_info = _minimal_rank_info( + instance_name="ctx", planner_id="trtllm.default", sender_endpoints=["tcp://ctx:1"] + ) + messenger = Mock() + messenger.receive.return_value = [sender_info.to_bytes()] + monkeypatch.setattr(transfer_mod, "ZMQMessenger", Mock(return_value=messenger)) + recv._get_or_connect_dealer = Mock() + params = SimpleNamespace(ctx_info_endpoint="tcp://ctx:9") + + with pytest.raises(PeerIncompatibleError, match="trtllm.default"): + recv._get_sender_info(params) + + recv._get_or_connect_dealer.assert_not_called() # REGISTER_RANK_INFO never sent + assert "tcp://ctx:9" in recv._incompatible_peers # verdict cached + with pytest.raises(PeerIncompatibleError): # and fails fast next time + recv._get_sender_info(params) + + +def test_worker_uses_model_planner_factory(monkeypatch): + worker = TransferWorker.__new__(TransferWorker) + worker._rank_info = _minimal_rank_info() + made = _ToyPlanner(TransferPlan.empty()) + seen = {} + + def factory(kvm, ri, ext): + seen["args"] = (kvm, ri, ext) + return made + + worker._config = SimpleNamespace(planner_factory=factory) + kvm = Mock() + kvm.mapping.rank = 1 + extractor = Mock() + monkeypatch.setattr(transfer_mod, "KVRegionExtractorV1", Mock(return_value=extractor)) + monkeypatch.setattr(transfer_mod, "RankInfoServer", Mock()) + + worker._setup_peer_infrastructure(kvm) + + assert worker._peer_registrar.planner is made + assert seen["args"] == (kvm, worker._rank_info, extractor) + assert worker._rank_info.planner_id == "toy" + assert worker._rank_info.planner_blob == b"" + + +def test_worker_defaults_to_default_planner(monkeypatch): + from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner + + worker = TransferWorker.__new__(TransferWorker) + worker._rank_info = _minimal_rank_info() + worker._config = SimpleNamespace(planner_factory=None) + kvm = Mock() + kvm.mapping.rank = 0 + monkeypatch.setattr(transfer_mod, "KVRegionExtractorV1", Mock(return_value=Mock())) + monkeypatch.setattr(transfer_mod, "RankInfoServer", Mock()) + + worker._setup_peer_infrastructure(kvm) + + assert isinstance(worker._peer_registrar.planner, DefaultTransferPlanner) + assert worker._rank_info.planner_id == DefaultTransferPlanner.planner_id From 15b9efe8e0950b770bd306594a8294a239af5289 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 09:16:51 -0700 Subject: [PATCH 10/25] [None][fix] Keep the default planner's statics static, wrap every peer-vetting error, cache targets A decorator left on _fanin_bounce_safe would have made it a bound method under a profiler and failed every transfer; check_peer now converts any ValueError raised while vetting a peer into PeerIncompatibleError so the receiver fails only that peer and caches the verdict; targets() is cached per peer DP group because the sender asks on every request. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../2026-10-10-disagg-transfer-planner.md | 6 ++ .../disaggregation/native/default_planner.py | 77 +++++++++-------- .../native/mixers/attention/peer.py | 2 +- .../disaggregated/test_chunked_transfer.py | 1 + .../disaggregated/test_transfer_planner.py | 84 +++++++++++++++++++ 5 files changed, 136 insertions(+), 34 deletions(-) diff --git a/docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md b/docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md index 8f5fbe6ea419..603a2149485d 100644 --- a/docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md +++ b/docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md @@ -16,6 +16,12 @@ 3. The model hook is a factory with the inputs the worker actually has: `build_transfer_planner(kv_cache_manager, self_rank_info, self_extractor) -> TransferPlanner`, not `build_transfer_planner(kv_cache_manager)`. `RankInfo` and the extractor are built by `TransferWorker` from the cache manager; the default planner consumes them, and duplicating that derivation in every planner would be wrong. 4. `PeerOverlap` moves with the engine into `default_planner.py` (it is the default planner's internal overlap record). Tests import it from there. 5. `TransferPlanner.forget_peer(peer_ri)` is added (default no-op) so `PeerRegistrar.unregister` can drop planner-side caches. It is cache hygiene the current `unregister` already does for `_kv_map_cache`/`_lg_pool_mapping_cache`. +6. The receiver's first-contact gate now runs the planner's full `check_peer` (attention + geometry, both layout validators, layer-count and NHD warnings), where before it ran only + the two layout validators and left geometry mismatches to the sender's listener thread, + whose logged error left the generation request hanging until its timeout. A mismatched + pair now fails fast on the generation side with a cached diagnostic. This is a deliberate + behavior improvement, recorded here because Phase 1 is otherwise zero-change. ## Global Constraints diff --git a/tensorrt_llm/_torch/disaggregation/native/default_planner.py b/tensorrt_llm/_torch/disaggregation/native/default_planner.py index c062a8c7d6b6..8fadee58c29d 100644 --- a/tensorrt_llm/_torch/disaggregation/native/default_planner.py +++ b/tensorrt_llm/_torch/disaggregation/native/default_planner.py @@ -30,7 +30,6 @@ import numpy as np from tensorrt_llm import logger -from tensorrt_llm._utils import nvtx_range from tensorrt_llm._torch.disaggregation.base.backend import TokenRange from tensorrt_llm._torch.disaggregation.base.planner import ( PeerIncompatibleError, @@ -98,6 +97,7 @@ def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1 tuple, RegionMapperBase ] = {} # key: (peer_key, self_lg_pool_key, peer_lg_pool_key) self._overlap_cache: Dict[str, PeerOverlap] = {} + self._targets_cache: Dict[str, PeerTargets] = {} self._lg_pool_mapping_cache: Dict[ str, Dict[LGPoolKey, LGPoolKey] ] = {} # peer_key -> {(self_lg, self_pi) -> (peer_lg, peer_pi)} @@ -118,22 +118,22 @@ def peer_extractor(self, peer_ri: RankInfo) -> KVRegionExtractorV1: # ---- TransferPlanner protocol ------------------------------------------------- def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: - if self_ri.bounded_replay_on_generation != peer_ri.bounded_replay_on_generation: - raise PeerIncompatibleError( - "DefaultTransferPlanner.check_peer: bounded_replay_on_generation must match " - f"between peers (local={self_ri.bounded_replay_on_generation}, " - f"peer={peer_ri.bounded_replay_on_generation})" - ) - if not self._get_policy(CacheKind.PAGED).check_peer_compatible(peer_ri): - raise PeerIncompatibleError( - "DefaultTransferPlanner.check_peer: attention geometry differs from the peer " - f"(peer={peer_ri.instance_name}, rank={peer_ri.instance_rank})" - ) - # Recurrent-state (Mamba/KDA) layout gate. Raises with a field-level - # diagnostic instead of returning False, so the precise mismatch - # reaches the caller. - self_page_table = self._self_ext_cache.page_table + """Vet a peer. Every layout or geometry mismatch surfaces as PeerIncompatibleError + so the caller fails only that peer's requests and caches the verdict.""" try: + if self_ri.bounded_replay_on_generation != peer_ri.bounded_replay_on_generation: + raise ValueError( + "bounded_replay_on_generation must match between peers " + f"(local={self_ri.bounded_replay_on_generation}, " + f"peer={peer_ri.bounded_replay_on_generation})" + ) + if not self._get_policy(CacheKind.PAGED).check_peer_compatible(peer_ri): + raise ValueError( + "attention geometry differs from the peer " + f"(peer={peer_ri.instance_name}, rank={peer_ri.instance_rank})" + ) + # Recurrent-state (Mamba/KDA) layout gate, with a field-level diagnostic. + self_page_table = self._self_ext_cache.page_table MambaPolicy.validate_peer_compatible( self._ri, peer_ri, @@ -141,21 +141,25 @@ def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) + + self_layers = sum(self._ri.layer_num_per_pp) + peer_layers = sum(peer_ri.layer_num_per_pp) + if self_layers != peer_layers: + # Allow mismatch when one side has speculative (e.g. MTP) layers + # that the other side doesn't. The pool_mapping logic will only + # transfer layers that exist on both sides. + logger.warning( + "DefaultTransferPlanner: layer count differs " + f"(local={self_layers}, peer={peer_layers}), " + "allowing partial layer transfer." + ) + # Pool matching runs here too; a peer that groups layers differently + # raises from it and must be refused like any other mismatch. + self._warn_nhd_head_mismatch(peer_ri) + except PeerIncompatibleError: + raise except ValueError as e: - raise PeerIncompatibleError(str(e)) from e - - self_layers = sum(self._ri.layer_num_per_pp) - peer_layers = sum(peer_ri.layer_num_per_pp) - if self_layers != peer_layers: - # Allow mismatch when one side has speculative (e.g. MTP) layers - # that the other side doesn't. The pool_mapping logic will only - # transfer layers that exist on both sides. - logger.warning( - "DefaultTransferPlanner: layer count differs " - f"(local={self_layers}, peer={peer_layers}), " - "allowing partial layer transfer." - ) - self._warn_nhd_head_mismatch(peer_ri) + raise PeerIncompatibleError(f"DefaultTransferPlanner.check_peer: {e}") from e def _warn_nhd_head_mismatch(self, peer_ri: RankInfo) -> None: head_match, _ = self._get_policy(CacheKind.PAGED).head_match(peer_ri) @@ -186,6 +190,13 @@ def _warn_nhd_head_mismatch(self, peer_ri: RankInfo) -> None: ) def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> PeerTargets: + # The sender asks on every request; the fan-in gate walks both page + # tables, so the answer is kept per (peer instance, DP group) like the + # overlap it derives from. + key = self._unique_key(peer_ri.instance_name, peer_dp_rank) + cached = self._targets_cache.get(key) + if cached is not None: + return cached overlap = self.get_peer_overlap(peer_ri, peer_dp_rank) # Helix writers contribute unequal shares (one KV owner, one state # sender, the rest empty), so CP on either side rules out the equal @@ -195,7 +206,9 @@ def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> Pe and peer_ri.cp_size == 1 and self._fanin_bounce_safe(overlap, peer_ri, self._self_ext_cache.page_table) ) - return PeerTargets(ranks=list(overlap.ranks), uniform_writers=uniform) + targets = PeerTargets(ranks=list(overlap.ranks), uniform_writers=uniform) + self._targets_cache[key] = targets + return targets def plan( self, @@ -365,8 +378,6 @@ def _pair_ordinals( idx = np.flatnonzero(mask) return src_block_ids[idx], dst_block_ids[idx] - @nvtx_range("_build_kv_write_meta") - @staticmethod def _fanin_bounce_safe( overlap: PeerOverlap, diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index f141cce93ff4..3092eeda94f3 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -127,7 +127,7 @@ class ReplicatedMapper(IntactMapper): indexer K), so no head slicing applies. Layer selection under partial PP overlap happens through the per-layer offsets, and fan-in routing (one owning sender per destination) is decided upstream by - ``PeerRegistrar.should_send_pool``. + ``DefaultTransferPlanner.should_send_pool``. """ def __init__( diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index d0f49ce489e1..95195a6416ca 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -183,6 +183,7 @@ def _make_projection_sender() -> Sender: # ordinal pairing under test is the production code path. planner = DefaultTransferPlanner.__new__(DefaultTransferPlanner) planner._self_ext_cache = extractor + planner._targets_cache = {} planner.peer_extractor = MagicMock(return_value=extractor) planner.get_peer_overlap = MagicMock( return_value=SimpleNamespace(ranks=[0], duplicate_head_factor=1, overlap_pp_size=1) diff --git a/tests/unittest/disaggregated/test_transfer_planner.py b/tests/unittest/disaggregated/test_transfer_planner.py index 28380b6281c4..842011e50b33 100644 --- a/tests/unittest/disaggregated/test_transfer_planner.py +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -316,3 +316,87 @@ def test_worker_defaults_to_default_planner(monkeypatch): assert isinstance(worker._peer_registrar.planner, DefaultTransferPlanner) assert worker._rank_info.planner_id == DefaultTransferPlanner.planner_id + + +# --------------------------------------------------------------------------- # +# DefaultTransferPlanner contract pins (from the whole-branch review) +# --------------------------------------------------------------------------- # +from tensorrt_llm._torch.disaggregation.native.default_planner import ( # noqa: E402 + DefaultTransferPlanner, +) + + +def test_default_planner_statics_survive_a_wrapping_nvtx_range(): + """With a profiler attached, nvtx_range wraps the function it decorates. A decorator + accidentally stacked on a @staticmethod then turns the gate into a bound method and + every targets() call fails. Import the module with a wrapping nvtx_range (a fresh + interpreter, so this test cannot disturb the others) and check the statics stay static.""" + import subprocess + import sys + + code = """ +import functools +import tensorrt_llm._utils as u +def wrapping_nvtx_range(*_a, **_k): + def deco(func): + @functools.wraps(func) + def inner(*args, **kwargs): + return func(*args, **kwargs) + return inner + return deco +u.nvtx_range = wrapping_nvtx_range +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner as P +for name in ("_fanin_bounce_safe", "_pair_ordinals", "_get_mamba_slot", "_get_buffers_per_layer", "_find_overlap"): + assert isinstance(P.__dict__[name], staticmethod), name +print("STATICS_OK") +""" + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, timeout=600 + ) + assert "STATICS_OK" in result.stdout, result.stdout[-2000:] + result.stderr[-4000:] + + +def _default_planner_pair(): + """A default planner and a compatible peer over the small real page tables test_peer builds.""" + from tests.unittest.disaggregated.test_peer import make_page_table, make_rankinfo + from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 + + self_ri = make_rankinfo(instance_name="gen", tp_size=1, tp_rank=0, page_table=make_page_table()) + peer_ri = make_rankinfo( + instance_name="ctx", instance_rank=0, tp_size=1, tp_rank=0, page_table=make_page_table() + ) + planner = DefaultTransferPlanner(self_ri, KVRegionExtractorV1(self_ri.page_table)) + return planner, self_ri, peer_ri + + +def test_default_planner_check_peer_wraps_every_value_error(monkeypatch): + """Any ValueError raised while vetting a peer must surface as PeerIncompatibleError, + so the receiver fails only that peer's requests and caches the verdict.""" + planner, self_ri, peer_ri = _default_planner_pair() + # A head-count mismatch is what sends check_peer into pool matching. + peer_ri.attention.kv_heads_per_rank = 1 + + def boom(_peer_ri): + raise ValueError("pool view spans multiple peer layer groups") + + monkeypatch.setattr(planner, "get_pool_mapping", boom) + with pytest.raises(PeerIncompatibleError, match="multiple peer layer groups"): + planner.check_peer(self_ri, peer_ri) + + +def test_default_planner_targets_is_cached_per_peer_dp_group(monkeypatch): + """The sender asks for targets() on every request; the fan-in gate walks both page + tables, so the answer must be computed once per (peer instance, DP group).""" + planner, self_ri, peer_ri = _default_planner_pair() + calls = [] + real = DefaultTransferPlanner._fanin_bounce_safe + + def counting(*args, **kwargs): + calls.append(1) + return real(*args, **kwargs) + + monkeypatch.setattr(planner, "_fanin_bounce_safe", counting) + first = planner.targets(self_ri, peer_ri, 0) + second = planner.targets(self_ri, peer_ri, 0) + assert first == second and first.ranks == [0] + assert len(calls) == 1 From 265197f5ffeeb3298ad749fcd43def32660562e8 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 09:20:08 -0700 Subject: [PATCH 11/25] [None][chore] Apply the repo formatters to the transfer-planner change Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../disaggregation/kv_cache_transceiver.py | 4 +++- .../disaggregation/native/default_planner.py | 4 +--- .../_torch/disaggregation/native/peer.py | 5 +---- .../_torch/disaggregation/native/transfer.py | 12 +++++----- tensorrt_llm/_torch/models/modeling_utils.py | 3 ++- tensorrt_llm/_torch/pyexecutor/_util.py | 11 ++++++---- .../disaggregated/test_cache_reuse_adapter.py | 16 ++++++++++---- .../unittest/disaggregated/test_extractor.py | 2 +- .../disaggregated/test_kda_mamba_transfer.py | 2 +- .../disaggregated/test_kv_transfer.py | 4 +++- tests/unittest/disaggregated/test_peer.py | 12 +++++----- .../disaggregated/test_pool_matching.py | 2 +- .../unittest/disaggregated/test_rank_info.py | 10 ++++++--- .../test_transfer_ownership_regressions.py | 2 +- .../disaggregated/test_transfer_planner.py | 22 +++++++++++++------ 15 files changed, 68 insertions(+), 43 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py index 83a17040fc03..d1096d745ef8 100644 --- a/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py @@ -229,7 +229,9 @@ def create_kv_cache_transceiver( # DecoderModelForCausalLM.build_transfer_planner); anything without # the hook gets the default planner. planner_factory = getattr(model, "build_transfer_planner", None) - return KvCacheTransceiverV2(mapping, dist, kv_cache_manager, + return KvCacheTransceiverV2(mapping, + dist, + kv_cache_manager, cache_transceiver_config, planner_factory=planner_factory) diff --git a/tensorrt_llm/_torch/disaggregation/native/default_planner.py b/tensorrt_llm/_torch/disaggregation/native/default_planner.py index 8fadee58c29d..0e265928aeec 100644 --- a/tensorrt_llm/_torch/disaggregation/native/default_planner.py +++ b/tensorrt_llm/_torch/disaggregation/native/default_planner.py @@ -276,9 +276,7 @@ def plan( # rejected under helix. src_block_ids = src_block_ids[peer_ri.cp_rank :: peer_ri.cp_size] tpb = extractor.page_table.tokens_per_block - dst_start_block = ( - dst_start_token // tpb if dst_start_token is not None else 0 - ) + dst_start_block = dst_start_token // tpb if dst_start_token is not None else 0 # Both tables are positional (index == block ordinal, -1 == not # held), so pairing is a plain intersection: SWA eviction, prefix # reuse, and pipelined chunk bounds are all already encoded as diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 29eee2e13d72..e7fed70f3503 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -22,10 +22,7 @@ from typing import Dict, Optional -from tensorrt_llm._torch.disaggregation.base.planner import ( - PeerIncompatibleError, - TransferPlanner, -) +from tensorrt_llm._torch.disaggregation.base.planner import PeerIncompatibleError, TransferPlanner from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxTransferLayout from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 864ad30402b7..67e7fb3c3af1 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -47,6 +47,11 @@ TransferOp, TransferRequest, ) +from tensorrt_llm._torch.disaggregation.base.planner import ( # noqa: F401 (re-exported) + PeerIncompatibleError, + PeerTargets, + TransferPlanner, +) from tensorrt_llm._torch.disaggregation.base.transfer import ( ReceiverBase, RxSessionBase, @@ -61,13 +66,8 @@ build_aux_transfer_layout, get_non_empty_aux_indices, ) -from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message -from tensorrt_llm._torch.disaggregation.base.planner import ( # noqa: F401 (re-exported) - PeerIncompatibleError, - PeerTargets, - TransferPlanner, -) from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner +from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index b1e0a339b4fa..5d2407575b15 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -801,7 +801,8 @@ def cache_derived_state(self) -> None: None. """ - def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): + def build_transfer_planner(self, kv_cache_manager, self_rank_info, + self_extractor): """The disaggregated KV transfer plan for this model on this rank. A model whose cache layout the default planner does not understand diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f21f28a739b3..b886139efe15 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -4302,10 +4302,13 @@ def create_py_executor_instance( if isinstance(kv_cache_manager, BaseMambaCacheManager): mamba_cache_manager = kv_cache_manager - kv_cache_transceiver = create_kv_cache_transceiver( - mapping, dist, kv_cache_manager, attention_type, - cache_transceiver_config, mamba_cache_manager, - model=model_engine.model) + kv_cache_transceiver = create_kv_cache_transceiver(mapping, + dist, + kv_cache_manager, + attention_type, + cache_transceiver_config, + mamba_cache_manager, + model=model_engine.model) waiting_queue_policy = (scheduler_config.waiting_queue_policy if scheduler_config is not None else diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index d4854b09d5d2..2b683e16b849 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -55,17 +55,23 @@ def test_identity(self): np.testing.assert_array_equal(dst, [20, 21, 22]) def test_receiver_prefix_reuse_is_a_dst_hole(self): - src, dst = DefaultTransferPlanner._pair_ordinals(_arr([10, 11, 12, 13]), _arr([-1, -1, 22, 23])) + src, dst = DefaultTransferPlanner._pair_ordinals( + _arr([10, 11, 12, 13]), _arr([-1, -1, 22, 23]) + ) np.testing.assert_array_equal(src, [12, 13]) np.testing.assert_array_equal(dst, [22, 23]) def test_sender_chunk_bounds_are_src_holes(self): - src, dst = DefaultTransferPlanner._pair_ordinals(_arr([-1, 11, 12, -1]), _arr([20, 21, 22, 23])) + src, dst = DefaultTransferPlanner._pair_ordinals( + _arr([-1, 11, 12, -1]), _arr([20, 21, 22, 23]) + ) np.testing.assert_array_equal(src, [11, 12]) np.testing.assert_array_equal(dst, [21, 22]) def test_holes_on_both_sides_take_the_intersection(self): - src, dst = DefaultTransferPlanner._pair_ordinals(_arr([10, -1, 12, 13]), _arr([-1, 21, 22, -1])) + src, dst = DefaultTransferPlanner._pair_ordinals( + _arr([10, -1, 12, 13]), _arr([-1, 21, 22, -1]) + ) np.testing.assert_array_equal(src, [12]) np.testing.assert_array_equal(dst, [22]) @@ -98,7 +104,9 @@ def test_sender_not_yet_evicted_but_receiver_has(self): def test_receiver_with_wider_speculative_window(self): # Only gen runs spec decoding, so its SWA window keeps one more block. # That block has no source; it simply is not written. - src, dst = DefaultTransferPlanner._pair_ordinals(_arr([-1, -1, 12, 13]), _arr([-1, 21, 22, 23])) + src, dst = DefaultTransferPlanner._pair_ordinals( + _arr([-1, -1, 12, 13]), _arr([-1, 21, 22, 23]) + ) np.testing.assert_array_equal(src, [12, 13]) np.testing.assert_array_equal(dst, [22, 23]) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index fdb8d1bd10ba..2e5685a57cce 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -346,8 +346,8 @@ def test_v1_dsa_indexer_replicated_transfer_across_pp( """ import torch - from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ReplicatedMapper from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner + from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ReplicatedMapper from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 0feb34f20b9a..a441f53e181c 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -55,11 +55,11 @@ import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( MambaPolicy, mamba_receiver_payload_bytes, ) -from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 539a0bd1eedc..cfb87bc65097 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -2018,7 +2018,9 @@ def raiser(*args, **kwargs): try: with pytest.MonkeyPatch.context() as mp: - mp.setattr(default_planner_mod.MambaPolicy, "validate_peer_compatible", staticmethod(raiser)) + mp.setattr( + default_planner_mod.MambaPolicy, "validate_peer_compatible", staticmethod(raiser) + ) # (a) First request to the incompatible peer: receive() must not # raise; the session fails with the diagnostic preserved. diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index c00784ea839f..1477375efc23 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -22,6 +22,10 @@ AuxBufferMeta, build_aux_transfer_layout, ) +from tensorrt_llm._torch.disaggregation.native.default_planner import ( + DefaultTransferPlanner, + PeerOverlap, +) from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ( AttentionPolicy, HNDHeadMismatchMapper, @@ -35,10 +39,6 @@ MambaHeadMatchMapper, MambaPolicy, ) -from tensorrt_llm._torch.disaggregation.native.default_planner import ( - DefaultTransferPlanner, - PeerOverlap, -) from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 @@ -208,7 +208,9 @@ def make_rankinfo( def _make_peer_registrar_and_peer_ri(self_ri, peer_ri): pt = make_page_table() real_kv_extractor = KVRegionExtractorV1(pt) - reg = PeerRegistrar(self_ri, real_kv_extractor, DefaultTransferPlanner(self_ri, real_kv_extractor)) + reg = PeerRegistrar( + self_ri, real_kv_extractor, DefaultTransferPlanner(self_ri, real_kv_extractor) + ) if peer_ri.page_table is None: peer_ri = make_rankinfo( instance_name=peer_ri.instance_name, diff --git a/tests/unittest/disaggregated/test_pool_matching.py b/tests/unittest/disaggregated/test_pool_matching.py index beb7f88531da..3b316bc9cf66 100644 --- a/tests/unittest/disaggregated/test_pool_matching.py +++ b/tests/unittest/disaggregated/test_pool_matching.py @@ -34,8 +34,8 @@ import numpy as np import pytest -from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner +from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import ( diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index 50fd1c0c27e4..cc50877dd9b9 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -230,11 +230,15 @@ def test_rank_info_planner_fields_round_trip(): def test_rank_info_from_bytes_without_planner_fields(): - """A peer serialized before planner_id existed still deserializes, so the - mismatch surfaces as the explicit PeerIncompatibleError and not a TypeError.""" - import msgpack + """An old-wire peer without planner fields still deserializes. + + The planner mismatch then surfaces as the explicit PeerIncompatibleError + rather than a msgpack TypeError. + """ from dataclasses import asdict + import msgpack + data = asdict(_minimal_rank_info()) for key in ("planner_id", "planner_blob", "attention", "aux_meta", "page_table"): data.pop(key, None) diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index 5d1ca77a0a0c..1e4471840186 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -28,8 +28,8 @@ 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.planner import PeerTargets from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base.planner import PeerTargets 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.transfer import ( diff --git a/tests/unittest/disaggregated/test_transfer_planner.py b/tests/unittest/disaggregated/test_transfer_planner.py index 842011e50b33..829f27326f4c 100644 --- a/tests/unittest/disaggregated/test_transfer_planner.py +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -327,10 +327,13 @@ def test_worker_defaults_to_default_planner(monkeypatch): def test_default_planner_statics_survive_a_wrapping_nvtx_range(): - """With a profiler attached, nvtx_range wraps the function it decorates. A decorator + """The planner's static helpers stay static under a wrapping nvtx_range. + + With a profiler attached, nvtx_range wraps the function it decorates. A decorator accidentally stacked on a @staticmethod then turns the gate into a bound method and every targets() call fails. Import the module with a wrapping nvtx_range (a fresh - interpreter, so this test cannot disturb the others) and check the statics stay static.""" + interpreter, so this test cannot disturb the others) and check the statics stay static. + """ import subprocess import sys @@ -358,8 +361,8 @@ def inner(*args, **kwargs): def _default_planner_pair(): """A default planner and a compatible peer over the small real page tables test_peer builds.""" - from tests.unittest.disaggregated.test_peer import make_page_table, make_rankinfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 + from tests.unittest.disaggregated.test_peer import make_page_table, make_rankinfo self_ri = make_rankinfo(instance_name="gen", tp_size=1, tp_rank=0, page_table=make_page_table()) peer_ri = make_rankinfo( @@ -370,8 +373,10 @@ def _default_planner_pair(): def test_default_planner_check_peer_wraps_every_value_error(monkeypatch): - """Any ValueError raised while vetting a peer must surface as PeerIncompatibleError, - so the receiver fails only that peer's requests and caches the verdict.""" + """Every ValueError raised while vetting a peer becomes PeerIncompatibleError. + + That way the receiver fails only that peer's requests and caches the verdict. + """ planner, self_ri, peer_ri = _default_planner_pair() # A head-count mismatch is what sends check_peer into pool matching. peer_ri.attention.kv_heads_per_rank = 1 @@ -385,8 +390,11 @@ def boom(_peer_ri): def test_default_planner_targets_is_cached_per_peer_dp_group(monkeypatch): - """The sender asks for targets() on every request; the fan-in gate walks both page - tables, so the answer must be computed once per (peer instance, DP group).""" + """targets() is computed once per (peer instance, DP group). + + The sender asks on every request and the fan-in gate walks both page tables, + so the answer must be cached. + """ planner, self_ri, peer_ri = _default_planner_pair() calls = [] real = DefaultTransferPlanner._fanin_bounce_safe From aad9ca5b7e019b1776b7ea79607b666f1367428a Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:07:46 -0700 Subject: [PATCH 12/25] [None][doc] Plan for the DeepSeek-R1 dep4 target-owned transfer planner Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- ...10-10-deepseek-r1-dep4-transfer-planner.md | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md diff --git a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md new file mode 100644 index 000000000000..9bed9bb05c1c --- /dev/null +++ b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md @@ -0,0 +1,66 @@ +# DeepSeek-R1 dep4 target-owned TransferPlanner — Phase 2 plan + +> **For agentic workers:** REQUIRED SUB-SKILL: superpowers:executing-plans. Steps use `- [ ]`. + +**Goal:** The first modeling_v2 target that plans its own disaggregated KV transfer: the +`r1_0528_nvfp4__sm_103__dep4` target ships a `TransferPlanner` in its own directory, overrides +`build_transfer_planner`, and a ctx dep4 -> gen dep4 deployment across two 4-GPU nodes runs on it. + +**Design (agreed in chat, 2026-10-10):** +- Under ADP dep4 every ctx rank pairs with exactly one gen rank (same dp_rank; flat rank == dp_rank + since pp=cp=1), MLA is one K buffer per layer (`kv_factor=1`), FP8 KV, one PAGED layer group over + one physical pool. The target planner therefore needs no pool matching, no policy, no mapper: + `targets` = `[peer_dp_rank]`, `plan` = ordinal pairing then `base + block*slot_stride`, merging + runs of blocks that are consecutive on both sides into one fragment (the default planner emits + one fragment per block). +- `check_peer` is the target declaring the only topology it supports: peer `dp_size == 4`, + `pp_size == cp_size == 1`, `enable_attention_dp`, one PAGED group with one pool, equal + `slot_bytes`, `tokens_per_block` and layer count. Anything else -> `PeerIncompatibleError`. +- `planner_id = "deepseek_r1_0528_nvfp4_sm103_dep4"`; `receiver_payload_bytes` = 0; `describe` = b"". +- Correctness proof is a CPU parity test against `DefaultTransferPlanner` on a synthesized R1 + page table: split the target planner's merged runs back to slot granularity and the fragment + multiset must equal the default planner's. +- GPU gate: two nodes (ctx on A, gen on B), NIXL over UCX tcp/eth0, `TRTLLM_MODELING_V2=require`, + GSM8K against the `deepseek-ai/DeepSeek-R1-0528` NVFP4/FP8 reference (94.24). The in-repo + accuracy test is the standard localhost 8-GPU form for CI; the two-node launcher is a dev script + outside the repo because the harness only knows localhost. + +**Spec:** `docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md` (Phase 2 section) +and the agreed design above. + +## Global Constraints +- NVIDIA header on new files; `git commit -s`; no AI attribution. +- Target isolation (`test_targets_do_not_share_files`): the planner file imports the catalog or + framework packages absolutely, never another target. +- Zero change to `native/` and to the default planner. + +## Review Focus +1. Non-owning / empty pairing: a chunk whose ordinals the peer does not hold must yield + `TransferPlan.empty()` (pinned in the parity test with an all-holes dst table). +2. Merged runs must stop at any discontinuity on *either* side (pinned: src consecutive, dst not). +3. `dst_start_token` masking must match the default planner (pinned in parity). +4. A peer with a second pool (e.g. MTP layer in its own pool) must be refused, not silently + half-transferred (pinned in `test_check_peer_refuses_*`). +5. The target planner must never be selected for the gen side when the ctx side runs the default + planner: framework `planner_id` gate (already pinned in Phase 1). + +--- + +### Task 1: planner + parity test +Files: Create `tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py`; +Modify `.../modeling.py` (override `build_transfer_planner`); +Test: `tests/unittest/_torch/modeling_v2/models/test_modeling_v2_deepseek_transfer_planner.py`. +Steps: write the parity + refusal tests, see them fail on import, implement, see them pass, commit. + +### Task 2: in-repo accuracy gate (8 GPUs, localhost form) +Files: `tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py` += `test_gsm8k_disagg` +(ctx dep4 + gen dep4, NIXL PYTHON, require mode via extra_env, DuckLLM args patched with +`kv_cache_quant_algo=FP8` so the R1 reference entry matches). Not runnable on 4-GPU nodes; it is +the CI form of what Task 3 runs by hand. + +### Task 3: two-node dev launcher + GSM8K +Files (outside repo, `/scratch/.../users/fredw/mv2-disagg/twonode/`): `alloc2.sbatch` +(2 nodes x 4 GPUs, keepalive on both), `serve_ctx.sh`, `serve_gen.sh`, `disagg_server.sh`, +`client_gsm8k.py` (copy of the harness client closure around `run_accuracy_test`). +Verify: both workers log `planner_id=deepseek_r1_0528_nvfp4_sm103_dep4`... (grep "Resolved" + +transceiver logs), GSM8K within tolerance of 94.24, no `PeerIncompatibleError`. From 6f4057a63fd90c6d2a95ab7a89b78ecf9d645ed8 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:14:18 -0700 Subject: [PATCH 13/25] [None][feat] DeepSeek-R1 dep4 target plans its own disaggregated KV transfer Under attention-DP dep4 every context rank has one generation twin and MLA keeps a block's latent K in one contiguous slot, so the target's planner is ordinal pairing plus slot arithmetic, merging consecutive blocks into one fragment. It accepts only the dep4-to-dep4 pair. A CPU parity test proves the plan covers exactly the bytes the default planner would move. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../r1_0528_nvfp4__sm_103__dep4/modeling.py | 5 + .../r1_0528_nvfp4__sm_103__dep4/transfer.py | 162 +++++++++++++++ ...t_modeling_v2_deepseek_transfer_planner.py | 189 ++++++++++++++++++ 3 files changed, 356 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_transfer_planner.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py index f84c25f519f1..04a20e8e30a4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py @@ -50,6 +50,7 @@ from .ops.moe.fused_moe import fused_moe from .ops.moe.noaux_tc_op import noaux_tc_op from .ops.quantization.fp4_quantize import fp4_quantize +from .transfer import DeepseekR1Dep4TransferPlanner _CACHED_CTX_FIELDS = ( "enable_context_mla_with_cached_kv", @@ -1787,6 +1788,10 @@ def __init__(self, model_config: ModelConfig): self.draft_model = DraftModel(self.model, self.lm_head, self.logits_processor) self.spec_worker = get_spec_worker(self.spec_config, model_config, model_config.mapping) + def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): + """This target's own disaggregated KV plan: dep4 context to dep4 generation.""" + return DeepseekR1Dep4TransferPlanner(self_rank_info, self_extractor) + def forward( self, attn_metadata: AttentionMetadata, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py new file mode 100644 index 000000000000..f2bc6f6fab79 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py @@ -0,0 +1,162 @@ +# 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. +"""Disaggregated KV transfer plan for the dep4 DeepSeek-R1 target. + +Under attention-DP dep4 every context rank has exactly one generation twin, the +rank with the same dp_rank. MLA keeps one latent-K region per layer in one +physical pool, so a block's bytes are one contiguous slot. The plan is then +plain arithmetic: pair block ordinals both sides hold, address the slots, and +copy. Blocks that are consecutive on both sides move as one fragment. + +This planner accepts exactly one peer shape -- the same target on four ranks -- +and refuses anything else; the generic default planner stays untouched for the +models that need pool matching, head re-sharding or recurrent state. +""" + +from __future__ import annotations + +from typing import AbstractSet, List, Optional, Tuple + +import numpy as np + +from tensorrt_llm._torch.disaggregation.base.backend import TokenRange +from tensorrt_llm._torch.disaggregation.base.planner import ( + PeerIncompatibleError, + PeerTargets, + TransferPlan, + TransferPlanner, +) +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.page import ( + CacheKind, + KVCachePageTable, + PhysicalPool, +) +from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool + +_DP_SIZE = 4 + + +def _sole_pool(page_table: Optional[KVCachePageTable], side: str) -> PhysicalPool: + """The one PAGED pool this target keeps its latent KV in, or a refusal.""" + if page_table is None: + raise PeerIncompatibleError(f"{side}: no page table published") + if len(page_table.layer_groups) != 1 or page_table.layer_groups[0].kind != CacheKind.PAGED: + raise PeerIncompatibleError( + f"{side}: expected exactly one PAGED layer group, " + f"got {[lg.kind.name for lg in page_table.layer_groups]}" + ) + lg = page_table.layer_groups[0] + pools = page_table.pool_groups[lg.pool_group_idx].pools + if len(lg.pool_views) != 1 or len(pools) != 1: + raise PeerIncompatibleError( + f"{side}: expected one pool for the latent KV, " + f"got {len(lg.pool_views)} views over {len(pools)} pools" + ) + return get_physical_pool(page_table, 0, lg.pool_views[0].pool_idx) + + +class DeepseekR1Dep4TransferPlanner(TransferPlanner): + planner_id = "deepseek_r1_0528_nvfp4_sm103_dep4" + + def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1): + self._ri = self_rank_info + self._page_table = self_extractor.page_table + self._pool = _sole_pool(self._page_table, "this rank") + + # ---- TransferPlanner protocol ------------------------------------------------- + + def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: + for ri, side in ((self_ri, "this rank"), (peer_ri, "peer")): + if ri.dp_size != _DP_SIZE or ri.tp_size != _DP_SIZE: + raise PeerIncompatibleError( + f"{side}: dp_size={ri.dp_size} tp_size={ri.tp_size}; this target runs dep{_DP_SIZE} only" + ) + if ri.pp_size != 1 or ri.cp_size != 1: + raise PeerIncompatibleError( + f"{side}: pp_size={ri.pp_size} cp_size={ri.cp_size}; this target runs pp=1 cp=1 only" + ) + if ri.attention is None or not ri.attention.enable_attention_dp: + raise PeerIncompatibleError(f"{side}: attention DP is off") + peer_pool = _sole_pool(peer_ri.page_table, "peer") + self_layers = len(self._page_table.layer_groups[0].local_layers) + peer_layers = len(peer_ri.page_table.layer_groups[0].local_layers) + if self_layers != peer_layers: + raise PeerIncompatibleError( + f"layer count differs: this rank {self_layers}, peer {peer_layers}" + ) + if peer_ri.page_table.tokens_per_block != self._page_table.tokens_per_block: + raise PeerIncompatibleError( + f"tokens_per_block differs: this rank {self._page_table.tokens_per_block}, " + f"peer {peer_ri.page_table.tokens_per_block}" + ) + if peer_pool.slot_bytes != self._pool.slot_bytes: + raise PeerIncompatibleError( + f"latent KV slot differs: this rank {self._pool.slot_bytes} bytes, " + f"peer {peer_pool.slot_bytes} bytes" + ) + + def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> PeerTargets: + # pp = cp = 1 and tp == dp, so the peer's flat rank is its dp_rank, and + # one writer per destination means the fan-in split is trivially uniform. + return PeerTargets(ranks=[peer_dp_rank], uniform_writers=True) + + def plan( + self, + self_ri: RankInfo, + peer_ri: RankInfo, + src_block_ids_per_lg: List[np.ndarray], + dst_block_ids_per_lg: List[np.ndarray], + token_range: TokenRange, + prompt_len: int, + dst_start_token: Optional[int], + excluded_pool_views: AbstractSet[Tuple[int, int]], + ) -> TransferPlan: + del token_range, prompt_len # cp == 1: a partial piece needs no projection + if (0, 0) in excluded_pool_views: # the sole latent pool is layer group 0, view 0 + return TransferPlan.empty() + tpb = self._page_table.tokens_per_block + dst_start_block = dst_start_token // tpb if dst_start_token is not None else 0 + src_ids, dst_ids = DefaultTransferPlanner._pair_ordinals( + np.asarray(src_block_ids_per_lg[0], dtype=np.int64), + np.asarray(dst_block_ids_per_lg[0], dtype=np.int64), + dst_start_block=dst_start_block, + ) + if src_ids.size == 0: + return TransferPlan.empty() + + src_pool, dst_pool = self._pool, _sole_pool(peer_ri.page_table, "peer") + slot = src_pool.slot_bytes + # Consecutive slots are consecutive bytes only when the stride is the slot. + mergeable = src_pool.slot_stride_bytes == slot and dst_pool.slot_stride_bytes == slot + if mergeable and src_ids.size > 1: + breaks = np.flatnonzero((np.diff(src_ids) != 1) | (np.diff(dst_ids) != 1)) + 1 + starts = np.concatenate(([0], breaks)) + lengths = np.diff(np.concatenate((starts, [src_ids.size]))) + else: + starts = np.arange(src_ids.size) + lengths = np.ones(src_ids.size, dtype=np.int64) + return TransferPlan( + src_ptrs=src_pool.base_address + src_ids[starts] * src_pool.slot_stride_bytes, + dst_ptrs=dst_pool.base_address + dst_ids[starts] * dst_pool.slot_stride_bytes, + sizes=lengths.astype(np.int64) * slot, + ) + + def receiver_payload_bytes( + self, self_ri: RankInfo, peer_ri: RankInfo, dst_block_ids_per_lg: List[np.ndarray] + ) -> Optional[int]: + return 0 # no recurrent state rides along with the paged blocks diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_transfer_planner.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_transfer_planner.py new file mode 100644 index 000000000000..490147e6d733 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_transfer_planner.py @@ -0,0 +1,189 @@ +# 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 dep4 DeepSeek-R1 target plans its own KV transfer; prove it equals the default plan.""" + +import numpy as np +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.transfer import ( # noqa: E501 + DeepseekR1Dep4TransferPlanner, +) +from tensorrt_llm._torch.disaggregation.base import TokenRange +from tensorrt_llm._torch.disaggregation.base.planner import PeerIncompatibleError, TransferPlan +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner +from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.page import ( + BUFFER_ENTRY_DTYPE, + AttentionLayerGroup, + KVCachePageTable, + LocalLayer, + PhysicalPool, + PhysicalPoolGroup, + PoolView, +) + +pytestmark = pytest.mark.cpu_only + +NUM_LAYERS = 6 +LAYER_BYTES = 512 # one latent-K region per layer per block +SLOT_BYTES = NUM_LAYERS * LAYER_BYTES +TPB = 16 + + +def _mla_page_table(base: int, *, num_layers: int = NUM_LAYERS, extra_pool: bool = False): + """One PAGED group, one K-only pool: the shape RankInfo.from_kv_cache_manager builds for MLA.""" + entries = np.array( + [(i, i * LAYER_BYTES, LAYER_BYTES) for i in range(num_layers)], dtype=BUFFER_ENTRY_DTYPE + ) + views = [PoolView(pool_idx=0, buffer_entries=entries)] + pools = [PhysicalPool(base_address=base, slot_bytes=num_layers * LAYER_BYTES, num_slots=256)] + if extra_pool: + views.append(PoolView(pool_idx=1, buffer_entries=entries, pool_role=frozenset({"mtp"}))) + pools.append( + PhysicalPool(base_address=base + (1 << 30), slot_bytes=SLOT_BYTES, num_slots=256) + ) + lg = AttentionLayerGroup( + pool_group_idx=0, + kv_head_num_per_rank=1, + sliding_window_size=None, + local_layers=[LocalLayer(local_layer_id=i, global_layer_id=i) for i in range(num_layers)], + pool_views=views, + ) + return KVCachePageTable( + tokens_per_block=TPB, layer_groups=[lg], pool_groups=[PhysicalPoolGroup(pools=pools)] + ) + + +def _dep4_rank_info(name: str, dp_rank: int, page_table, *, dp_size: int = 4, pp_size: int = 1): + return RankInfo( + instance_name=name, + instance_rank=dp_rank, + tp_size=dp_size, + tp_rank=dp_rank, + pp_size=pp_size, + pp_rank=0, + dp_size=dp_size, + dp_rank=dp_rank, + layer_num_per_pp=[len(page_table.layer_groups[0].local_layers)], + sender_endpoints=[], + self_endpoint="", + transfer_engine_info=b"", + attention=AttentionInfo( + kv_heads_per_rank=1, + tokens_per_block=TPB, + dims_per_head=512, + element_bytes=1, + enable_attention_dp=True, + is_mla=True, + ), + page_table=page_table, + ) + + +def _pair(dp_rank: int = 2): + self_pt, peer_pt = _mla_page_table(0x1000_0000), _mla_page_table(0x9000_0000) + self_ri = _dep4_rank_info("ctx", dp_rank, self_pt) + peer_ri = _dep4_rank_info("gen", dp_rank, peer_pt) + ext = KVRegionExtractorV1(self_pt) + return ( + DeepseekR1Dep4TransferPlanner(self_ri, ext), + DefaultTransferPlanner(self_ri, ext), + self_ri, + peer_ri, + ) + + +def _fragments_at_slot_granularity(plan: TransferPlan) -> set: + """Split merged runs into per-slot pieces so plans can be compared as sets.""" + out = set() + for s, d, n in zip(plan.src_ptrs.tolist(), plan.dst_ptrs.tolist(), plan.sizes.tolist()): + assert n % SLOT_BYTES == 0 + for k in range(n // SLOT_BYTES): + out.add((s + k * SLOT_BYTES, d + k * SLOT_BYTES, SLOT_BYTES)) + return out + + +def _tables(src, dst): + return [np.array(src, dtype=np.int64)], [np.array(dst, dtype=np.int64)] + + +@pytest.mark.parametrize( + "src,dst,dst_start_token", + [ + ([3, 4, 5, 6], [10, 11, 12, 13], None), # all consecutive both sides + ([3, 4, 9, 10], [10, 11, 12, 13], None), # src breaks a run + ([3, 4, 5, 6], [10, 11, 20, 21], None), # dst breaks a run + ([3, -1, 5, 6], [10, 11, -1, 13], None), # holes on both sides + ([3, 4, 5, 6], [10, 11, 12, 13], 2 * TPB), # receiver prefix reuse + ([3, 4, 5, 6], [-1, -1, -1, -1], None), # peer holds nothing + ], + ids=["contiguous", "src_break", "dst_break", "holes", "dst_start", "nothing_to_send"], +) +def test_plan_matches_default_planner_byte_for_byte(src, dst, dst_start_token): + target, default, self_ri, peer_ri = _pair() + s, d = _tables(src, dst) + tr = TokenRange(0, len(src) * TPB) + + ours = target.plan(self_ri, peer_ri, s, d, tr, len(src) * TPB, dst_start_token, frozenset()) + ref = default.plan(self_ri, peer_ri, s, d, tr, len(src) * TPB, dst_start_token, frozenset()) + + assert _fragments_at_slot_granularity(ours) == _fragments_at_slot_granularity(ref) + assert int(ours.sizes.sum()) == int(ref.sizes.sum()) + + +def test_contiguous_blocks_merge_into_one_fragment(): + target, _, self_ri, peer_ri = _pair() + s, d = _tables([3, 4, 5, 6], [10, 11, 12, 13]) + plan = target.plan(self_ri, peer_ri, s, d, TokenRange(0, 4 * TPB), 4 * TPB, None, frozenset()) + assert plan.sizes.tolist() == [4 * SLOT_BYTES] + + +def test_targets_is_the_matching_dp_rank_only(): + target, default, self_ri, peer_ri = _pair(dp_rank=2) + ours, ref = target.targets(self_ri, peer_ri, 2), default.targets(self_ri, peer_ri, 2) + assert ours.ranks == ref.ranks == [2] + assert ours.uniform_writers is True + + +def test_check_peer_accepts_the_dep4_twin_and_refuses_everything_else(): + target, _, self_ri, peer_ri = _pair() + target.check_peer(self_ri, peer_ri) # the supported pair + with pytest.raises(PeerIncompatibleError, match="dp_size"): + target.check_peer( + self_ri, _dep4_rank_info("gen", 0, _mla_page_table(0x9000_0000), dp_size=2) + ) + with pytest.raises(PeerIncompatibleError, match="pool"): + target.check_peer( + self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000, extra_pool=True)) + ) + with pytest.raises(PeerIncompatibleError, match="layer"): + target.check_peer( + self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000, num_layers=5)) + ) + + +def test_target_class_builds_this_planner(): + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4 import ( # noqa: E501 + modeling, + ) + + cls = modeling.ModelingV2DeepseekR10528Nvfp4Sm103Dep4 + self_pt = _mla_page_table(0x1000_0000) + ri = _dep4_rank_info("ctx", 0, self_pt) + planner = cls.build_transfer_planner(object(), None, ri, KVRegionExtractorV1(self_pt)) + assert isinstance(planner, DeepseekR1Dep4TransferPlanner) + assert planner.planner_id == "deepseek_r1_0528_nvfp4_sm103_dep4" From 0dab135e305b0937390ab3d1a9ee02a07a7645b3 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:15:10 -0700 Subject: [PATCH 14/25] [None][test] Gate the DeepSeek-R1 dep4 target under disaggregated serving Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../accuracy/test_modeling_v2_deepseek_v3.py | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py b/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py index a2ef52bc341c..a278b2525a8e 100644 --- a/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py +++ b/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py @@ -37,6 +37,7 @@ from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV from tensorrt_llm._utils import get_sm_version from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MTPDecodingConfig +from tensorrt_llm.quantization import QuantAlgo from ..conftest import llm_models_root from .accuracy_core import ( @@ -45,6 +46,7 @@ assert_acceptance_length, compute_acceptance_length, ) +from .test_disaggregated_serving import launch_disaggregated_llm, run_accuracy_test # The targets assert their own SM at construction: certification is per GPU # architecture, and a receipt from another one says nothing here. @@ -194,3 +196,50 @@ def test_mtp3_acceptance(self, mode, mocker): acceptance_length = compute_acceptance_length(llm) print(f"[AL] {mode} acceptance_length = {acceptance_length:.3f}") assert_acceptance_length(self.ACCEPTANCE_KEY, acceptance_length) + + @skip_not_sm103 + @pytest.mark.skip_less_device(8) + def test_gsm8k_disagg(self, mocker): + """dep4 context server to dep4 generation server, KV planned by the target itself. + + The target ships its own ``TransferPlanner``; both workers run it, so the + framework's planner_id handshake passes and every KV block moves by the + target's slot arithmetic rather than the generic pool matching. ``require`` + goes into each server's config, since ``modeling_v2`` is an ``LLM`` argument. + """ + mocker.patch.dict(GSM8K.EVALUATE_KWARGS, _SCORES_FILTER) + transceiver = { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + "max_tokens_in_buffer": 4096, + } + common = { + "modeling_v2": "require", + "tensor_parallel_size": 4, + "moe_expert_parallel_size": 4, + "enable_attention_dp": True, + "max_batch_size": 16, + "max_seq_len": 8192, + "max_num_tokens": 8192, + "cache_transceiver_config": transceiver, + "kv_cache_config": {"free_gpu_memory_fraction": 0.6}, + } + ctx_server_config = {**common, "disable_overlap_scheduler": True} + gen_server_config = {**common, "disable_overlap_scheduler": False} + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": {"num_instances": 1}, + "generation_servers": {"num_instances": 1}, + } + with launch_disaggregated_llm( + disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + extra_env={MODELING_V2_VALIDATE_ENV: "1"}, + ) as llm: + # The harness infers only the weight quantization from the path; the + # R1 reference entry is keyed on the FP8 KV cache as well. + llm.args.quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + run_accuracy_test(llm, self.MODEL_NAME, test_sets=["GSM8K"]) From 341d260ee31acc08b087667accf52372cc965aa6 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 11:10:16 -0700 Subject: [PATCH 15/25] [None][chore] Log which TransferPlanner a transfer worker runs Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- tensorrt_llm/_torch/disaggregation/native/transfer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 67e7fb3c3af1..2ed4a8433d16 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -4428,6 +4428,10 @@ def _setup_peer_infrastructure(self, kvm: KVCacheManager): # the object and serializes on request, so setting it here is in time. self._rank_info.planner_id = self._planner.planner_id self._rank_info.planner_blob = self._planner.describe() + logger.info( + f"TransferWorker: KV transfer planned by {type(self._planner).__name__} " + f"(planner_id={self._planner.planner_id!r})" + ) self._peer_registrar = PeerRegistrar(self._rank_info, self._kv_extractor, self._planner) def _setup_transfer_engine(self): From 159018c6b99ee469c063566228656dfc8ccc460b Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 10 Oct 2026 11:38:42 -0700 Subject: [PATCH 16/25] [None][doc] Record the two-node DeepSeek-R1 dep4 planner gate results Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../2026-10-10-deepseek-r1-dep4-transfer-planner.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md index 9bed9bb05c1c..1c7c8645d0ec 100644 --- a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md +++ b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md @@ -64,3 +64,16 @@ Files (outside repo, `/scratch/.../users/fredw/mv2-disagg/twonode/`): `alloc2.sb `client_gsm8k.py` (copy of the harness client closure around `run_accuracy_test`). Verify: both workers log `planner_id=deepseek_r1_0528_nvfp4_sm103_dep4`... (grep "Resolved" + transceiver logs), GSM8K within tolerance of 94.24, no `PeerIncompatibleError`. + +## Results (2026-10-10) + +- Task 1: planner + parity test committed (9a9b56e485); 33 tests green including the modeling_v2 + claims and target-contract suites. +- Task 2: in-repo 8-GPU gate committed (048bd64c25); collected, not runnable on 4-GPU nodes. +- Task 3: two nodes (ctx dep4 on nvl72d202-T12, gen dep4 on nvl72d202-T16, NIXL over UCX tcp/eth0, + static worker URLs, `TRTLLM_MODELING_V2=require`). Both workers log + `KV transfer planned by DeepseekR1Dep4TransferPlanner`, no default-planner fallback, no + `PeerIncompatibleError`. GSM8K flexible-extract **95.22** against the 94.24 reference, PASSED; + 1319 requests in 6.5 min. Launcher findings (loopback-only HTTP discovery, `--server_role` + crash in static mode, `UCX_TLS=^ib` requirement) are recorded in the dev scripts and memory, + not in the repo. Logs: `/scratch/fsw/portfolios/coreai/projects/coreai_comparch_trtllm/users/fredw/mv2-disagg/twonode/logs/run_20261010T182505Z`. From d3a22f6fa91dbba2437f46c69606c96a7b569587 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sun, 11 Oct 2026 01:42:11 -0700 Subject: [PATCH 17/25] [None][refactor] Bind disaggregated KV planners to the targets they cover, not per target A planner owns a compatibility set of (context target, generation target) pairs; a target is one member. The target used to name its planner by overriding build_transfer_planner, which capped each planner at one target. Now the planner declares COVERS, modeling_v2/transfers resolves the planner from the target id, AutoModelForCausalLM installs the factory on the instance, and describe() publishes the target id so the peer's planner can refuse uncovered pairs. The DeepSeek dep4 planner moves out of the target directory as DeepseekMlaLatentTransferPlanner; the target's modeling.py carries no transfer code. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- ...10-10-deepseek-r1-dep4-transfer-planner.md | 22 ++++ ...26-10-10-disagg-transfer-planner-design.md | 27 ++++- .../_experimental/modeling_v2/__init__.py | 2 + .../modeling_v2/_router_index.py | 17 +++ .../r1_0528_nvfp4__sm_103__dep4/modeling.py | 5 - .../modeling_v2/transfers/__init__.py | 71 ++++++++++++ .../deepseek_mla_latent.py} | 66 ++++++----- tensorrt_llm/_torch/models/modeling_auto.py | 8 ++ .../test_deepseek_mla_latent.py} | 43 +++---- .../transfers/test_transfers_binding.py | 105 ++++++++++++++++++ 10 files changed, 312 insertions(+), 54 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/transfers/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/{models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py => transfers/deepseek_mla_latent.py} (72%) rename tests/unittest/_torch/modeling_v2/{models/deepseek_v3/test_modeling_v2_deepseek_transfer_planner.py => transfers/test_deepseek_mla_latent.py} (83%) create mode 100644 tests/unittest/_torch/modeling_v2/transfers/test_transfers_binding.py diff --git a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md index 1c7c8645d0ec..27533e908ddc 100644 --- a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md +++ b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md @@ -77,3 +77,25 @@ transceiver logs), GSM8K within tolerance of 94.24, no `PeerIncompatibleError`. 1319 requests in 6.5 min. Launcher findings (loopback-only HTTP discovery, `--server_role` crash in static mode, `UCX_TLS=^ib` requirement) are recorded in the dev scripts and memory, not in the repo. Logs: `/scratch/fsw/portfolios/coreai/projects/coreai_comparch_trtllm/users/fredw/mv2-disagg/twonode/logs/run_20261010T182505Z`. + +## Redesign of the binding (after review) + +The target-overrides-`build_transfer_planner` binding was rejected in review: one planner +must be able to serve several targets, including pairs such as ctx dep4 -> gen dep8 and +targets from different cores. Implemented instead: + +- Planner moved to `modeling_v2/transfers/deepseek_mla_latent.py` as + `DeepseekMlaLatentTransferPlanner`, `planner_id = "modeling_v2.deepseek_mla_latent"`, with a + `COVERS` matrix of `(ctx target id, gen target id)` pairs; today only + `(deepseek_v3/r1_0528_nvfp4__sm_103__dep4, same)`. `describe()` publishes the own target id and + `check_peer` refuses any pair outside the matrix; the dp/tp/pp/cp checks are gone because the + target id already pins the topology, the layout checks (one pool, layer count, + `tokens_per_block`, `slot_bytes`) stay because the arithmetic depends on them. +- `modeling_v2/transfers/__init__.py`: `PLANNERS`, `planner_for_target`, `bind_transfer_planner`; + `target_id_of` lives in `_router_index` so `AutoModelForCausalLM` can ask it cheaply. +- `AutoModelForCausalLM.from_config` binds after construction; the target's `modeling.py` is back + to zero transfer-related lines. +- Tests: `tests/unittest/_torch/modeling_v2/transfers/` (binding direction, lookup, duplicate + coverage, bind; planner parity/merge/targets/pair refusal/describe). 74 tests green with the + claims, routing and disaggregation planner suites. The two-node GSM8K gate was run before the + redesign; the plan arithmetic is unchanged and the parity test still pins it. diff --git a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md index 2d61c5db5ecb..b9e8570ef1bb 100644 --- a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md +++ b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md @@ -154,12 +154,31 @@ class PretrainedModel/DecoderModelForCausalLM: return DefaultTransferPlanner(kv_cache_manager) ``` -A modeling_v2 target overrides `build_transfer_planner` in its `modeling.py`, -beside `REQUIRED_TRTLLM_OPS`, and keeps its planner in its own target -directory. Reason for the model over the cache manager: a plan is a function +The default is the method above. A modeling_v2 target does **not** override +it: a planner is the owner of a compatibility set, the (context target, +generation target) pairs it can move bytes between, and a target is one +member of such a set. Letting the member name the set would cap every set at +one member. So the binding runs the other way: + +- A planner class declares `COVERS`, a frozenset of `(ctx_target_id, + gen_target_id)` pairs. Target ids are `/`, the + directory being the target's identity (`____`). +- `modeling_v2/transfers/` holds every planner and the one lookup, + `planner_for_target(target_id)`: zero planners -> default plan, two -> + configuration error. +- `AutoModelForCausalLM.from_config` asks that lookup for the target it just + built and installs the covering planner's factory as the instance's + `build_transfer_planner`. The framework reads the hook and nothing else; + neither the target nor the disaggregation code names a planner. +- `describe()` publishes the rank's own target id; the peer's planner reads + it from `planner_blob` and refuses any pair its matrix lacks. Equal + `planner_id` on both sides therefore means "we belong to the same + compatibility set", not "we are the same target". + +Reason for the model over the cache manager still holds: a plan is a function of the parallel topology *and* the cache layout, and the target is the one object that already owns both. The cache manager builds the page table and -knows nothing about which peer topologies the target supports. +knows nothing about which peer topologies are supported. Plumbing: `create_kv_cache_transceiver` (`pyexecutor/_util.py:4232`) gains a `planner` argument built from `model_engine.model`; diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index c5140e7c912b..f81814ff79fb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -33,6 +33,7 @@ NoTarget, modeling_v2_resolve, step_contract_enabled, + target_id_of, ) __all__ = [ @@ -42,4 +43,5 @@ "NoTarget", "modeling_v2_resolve", "step_contract_enabled", + "target_id_of", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index cdb288f7220f..3f7a0a6e706e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -101,6 +101,23 @@ class NoTarget(Exception): """ +def target_id_of(model_cls: type) -> Optional[str]: + """``/`` for a class defined in a modeling_v2 target, else None. + + The directory name is a target's identity (``____``); + the family prefix keeps two families' checkpoints from colliding. Cheap on + purpose: ``AutoModelForCausalLM`` asks this for every model it builds. + """ + module = getattr(model_cls, "__module__", None) or "" + prefix = f"{_PACKAGE}.models." + if not module.startswith(prefix): + return None + parts = module[len(prefix) :].split(".") + if len(parts) != 3: # family.directory.modeling + return None + return f"{parts[0]}/{parts[1]}" + + def routing_module(arch: str): """Import and return the routing module for ``arch``, or None.""" name = MODELING_V2_ROUTERS.get(arch) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py index 04a20e8e30a4..f84c25f519f1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py @@ -50,7 +50,6 @@ from .ops.moe.fused_moe import fused_moe from .ops.moe.noaux_tc_op import noaux_tc_op from .ops.quantization.fp4_quantize import fp4_quantize -from .transfer import DeepseekR1Dep4TransferPlanner _CACHED_CTX_FIELDS = ( "enable_context_mla_with_cached_kv", @@ -1788,10 +1787,6 @@ def __init__(self, model_config: ModelConfig): self.draft_model = DraftModel(self.model, self.lm_head, self.logits_processor) self.spec_worker = get_spec_worker(self.spec_config, model_config, model_config.mapping) - def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): - """This target's own disaggregated KV plan: dep4 context to dep4 generation.""" - return DeepseekR1Dep4TransferPlanner(self_rank_info, self_extractor) - def forward( self, attn_metadata: AttentionMetadata, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/__init__.py new file mode 100644 index 000000000000..412f156867ca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/__init__.py @@ -0,0 +1,71 @@ +# 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. +"""Which disaggregated KV transfer planner serves a modeling_v2 target. + +A planner owns a compatibility set: the (context target, generation target) +pairs it can move bytes between, declared on the class as ``COVERS``. Targets +declare nothing. This module is the one place that knows every planner, and it +answers one question: given the target a model instance was built from, which +planner covers it? None means the engine keeps the framework's default plan. + +Binding happens where the model is constructed (``AutoModelForCausalLM``), +by installing the covering planner's factory as the instance's +``build_transfer_planner``; the framework reads that hook and nothing else. +""" + +from __future__ import annotations + +from typing import Optional, Type + +from tensorrt_llm._torch.disaggregation.base.planner import TransferPlanner + +from .._router_index import target_id_of +from .deepseek_mla_latent import DeepseekMlaLatentTransferPlanner + +#: Every planner a target may be served by. Order carries no meaning: a target +#: covered by two planners is a configuration error, not a priority question. +PLANNERS: tuple = (DeepseekMlaLatentTransferPlanner,) + + +def planner_for_target(target_id: str) -> Optional[Type[TransferPlanner]]: + """The planner whose matrix mentions ``target_id`` on either side, or None.""" + found = [p for p in PLANNERS if any(target_id in pair for pair in p.COVERS)] + if len(found) > 1: + names = ", ".join(p.__name__ for p in found) + raise ValueError(f"target {target_id!r} is covered by more than one planner: {names}") + return found[0] if found else None + + +def bind_transfer_planner(model, model_cls: type) -> Optional[Type[TransferPlanner]]: + """Install the covering planner's factory on ``model``; return the planner or None. + + The factory has the ``build_transfer_planner`` signature the framework + calls, closing over the target id so the planner can publish it. + """ + target_id = target_id_of(model_cls) + if target_id is None: + return None + planner_cls = planner_for_target(target_id) + if planner_cls is None: + return None + + def build_transfer_planner(kv_cache_manager, self_rank_info, self_extractor): + return planner_cls(self_rank_info, self_extractor, self_target=target_id) + + model.build_transfer_planner = build_transfer_planner + return planner_cls + + +__all__ = ["PLANNERS", "bind_transfer_planner", "planner_for_target", "target_id_of"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py similarity index 72% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py index f2bc6f6fab79..a60cf733844e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/transfer.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py @@ -12,22 +12,26 @@ # 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. -"""Disaggregated KV transfer plan for the dep4 DeepSeek-R1 target. +"""Disaggregated KV transfer plan for DeepSeek targets that keep MLA latent KV in one pool. -Under attention-DP dep4 every context rank has exactly one generation twin, the -rank with the same dp_rank. MLA keeps one latent-K region per layer in one -physical pool, so a block's bytes are one contiguous slot. The plan is then +Under attention-DP every context rank has a generation twin with the same +dp_rank. MLA keeps one latent-K region per layer in one physical pool, so a +block's bytes for all local layers are one contiguous slot. The plan is then plain arithmetic: pair block ordinals both sides hold, address the slots, and copy. Blocks that are consecutive on both sides move as one fragment. -This planner accepts exactly one peer shape -- the same target on four ranks -- -and refuses anything else; the generic default planner stays untouched for the -models that need pool matching, head re-sharding or recurrent state. +Which deployments may talk to each other is the ``COVERS`` matrix below: pairs +of (context target, generation target) ids. A target never names this planner; +``modeling_v2.transfers`` finds the planner from the target id, and ``describe`` +publishes that id so the peer's planner can look the pair up too. Anything not +in the matrix is refused at registration; the generic default planner stays +untouched for the models that need pool matching, head re-sharding or +recurrent state. """ from __future__ import annotations -from typing import AbstractSet, List, Optional, Tuple +from typing import AbstractSet, FrozenSet, List, Optional, Tuple import numpy as np @@ -48,7 +52,7 @@ ) from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool -_DP_SIZE = 4 +_R1_DEP4 = "deepseek_v3/r1_0528_nvfp4__sm_103__dep4" def _sole_pool(page_table: Optional[KVCachePageTable], side: str) -> PhysicalPool: @@ -70,28 +74,39 @@ def _sole_pool(page_table: Optional[KVCachePageTable], side: str) -> PhysicalPoo return get_physical_pool(page_table, 0, lg.pool_views[0].pool_idx) -class DeepseekR1Dep4TransferPlanner(TransferPlanner): - planner_id = "deepseek_r1_0528_nvfp4_sm103_dep4" +class DeepseekMlaLatentTransferPlanner(TransferPlanner): + planner_id = "modeling_v2.deepseek_mla_latent" - def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1): + #: (context target, generation target) pairs this planner knows how to move + #: bytes between. The same pair is listed once; the planner serves both + #: roles, so a pair is read from the sender's side or the receiver's side. + COVERS: FrozenSet[Tuple[str, str]] = frozenset({(_R1_DEP4, _R1_DEP4)}) + + def __init__( + self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1, *, self_target: str + ): self._ri = self_rank_info + self._target = self_target self._page_table = self_extractor.page_table self._pool = _sole_pool(self._page_table, "this rank") + def describe(self) -> bytes: + return self._target.encode() + + def _pair(self, peer_ri: RankInfo) -> Tuple[str, str]: + """The (self, peer) target pair, or a refusal if the matrix lacks it.""" + peer_target = peer_ri.planner_blob.decode() + if (self._target, peer_target) in self.COVERS or (peer_target, self._target) in self.COVERS: + return self._target, peer_target + raise PeerIncompatibleError( + f"no transfer plan for the pair ({self._target!r}, {peer_target!r}); " + f"covered pairs: {sorted(self.COVERS)}" + ) + # ---- TransferPlanner protocol ------------------------------------------------- def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: - for ri, side in ((self_ri, "this rank"), (peer_ri, "peer")): - if ri.dp_size != _DP_SIZE or ri.tp_size != _DP_SIZE: - raise PeerIncompatibleError( - f"{side}: dp_size={ri.dp_size} tp_size={ri.tp_size}; this target runs dep{_DP_SIZE} only" - ) - if ri.pp_size != 1 or ri.cp_size != 1: - raise PeerIncompatibleError( - f"{side}: pp_size={ri.pp_size} cp_size={ri.cp_size}; this target runs pp=1 cp=1 only" - ) - if ri.attention is None or not ri.attention.enable_attention_dp: - raise PeerIncompatibleError(f"{side}: attention DP is off") + self._pair(peer_ri) peer_pool = _sole_pool(peer_ri.page_table, "peer") self_layers = len(self._page_table.layer_groups[0].local_layers) peer_layers = len(peer_ri.page_table.layer_groups[0].local_layers) @@ -111,8 +126,9 @@ def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: ) def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> PeerTargets: - # pp = cp = 1 and tp == dp, so the peer's flat rank is its dp_rank, and - # one writer per destination means the fan-in split is trivially uniform. + # Every covered pair is attention-DP with pp = cp = 1 and tp == dp, so + # the peer's flat rank is its dp_rank, and one writer per destination + # means the fan-in split is trivially uniform. return PeerTargets(ranks=[peer_dp_rank], uniform_writers=True) def plan( diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 0592317d293a..5d6f3eec2abd 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from typing import Generic, Optional, Type +from .._experimental.modeling_v2 import target_id_of from ..model_config import ModelConfig from ..utils import model_extra_attrs from .modeling_utils import (DecoderModelForCausalLM, TConfig, TModel, @@ -77,4 +78,11 @@ def from_config( with model_extra_attrs(extra_attrs): model = cls(config) model.extra_attrs = extra_attrs + # A modeling_v2 target may be served by a planner that covers it (see + # modeling_v2.transfers). Resolved here, where the target is known, + # so neither the target nor the disaggregation code names the other. + if target_id_of(cls) is not None: + from .._experimental.modeling_v2.transfers import \ + bind_transfer_planner + bind_transfer_planner(model, cls) return model diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_transfer_planner.py b/tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_transfer_planner.py rename to tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py index 490147e6d733..e187cb4d9368 100644 --- a/tests/unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_transfer_planner.py +++ b/tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py @@ -12,13 +12,13 @@ # 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 dep4 DeepSeek-R1 target plans its own KV transfer; prove it equals the default plan.""" +"""The DeepSeek MLA-latent planner must produce the default planner's bytes for the pairs it covers.""" import numpy as np import pytest -from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.transfer import ( # noqa: E501 - DeepseekR1Dep4TransferPlanner, +from tensorrt_llm._torch._experimental.modeling_v2.transfers.deepseek_mla_latent import ( + DeepseekMlaLatentTransferPlanner, ) from tensorrt_llm._torch.disaggregation.base import TokenRange from tensorrt_llm._torch.disaggregation.base.planner import PeerIncompatibleError, TransferPlan @@ -42,6 +42,7 @@ LAYER_BYTES = 512 # one latent-K region per layer per block SLOT_BYTES = NUM_LAYERS * LAYER_BYTES TPB = 16 +DEP4 = "deepseek_v3/r1_0528_nvfp4__sm_103__dep4" def _mla_page_table(base: int, *, num_layers: int = NUM_LAYERS, extra_pool: bool = False): @@ -68,8 +69,12 @@ def _mla_page_table(base: int, *, num_layers: int = NUM_LAYERS, extra_pool: bool ) -def _dep4_rank_info(name: str, dp_rank: int, page_table, *, dp_size: int = 4, pp_size: int = 1): +def _dep4_rank_info( + name: str, dp_rank: int, page_table, *, dp_size: int = 4, pp_size: int = 1, target: str = DEP4 +): return RankInfo( + planner_id=DeepseekMlaLatentTransferPlanner.planner_id, + planner_blob=target.encode(), instance_name=name, instance_rank=dp_rank, tp_size=dp_size, @@ -100,7 +105,7 @@ def _pair(dp_rank: int = 2): peer_ri = _dep4_rank_info("gen", dp_rank, peer_pt) ext = KVRegionExtractorV1(self_pt) return ( - DeepseekR1Dep4TransferPlanner(self_ri, ext), + DeepseekMlaLatentTransferPlanner(self_ri, ext, self_target=DEP4), DefaultTransferPlanner(self_ri, ext), self_ri, peer_ri, @@ -159,12 +164,17 @@ def test_targets_is_the_matching_dp_rank_only(): assert ours.uniform_writers is True -def test_check_peer_accepts_the_dep4_twin_and_refuses_everything_else(): +def test_check_peer_accepts_a_covered_pair_and_refuses_everything_else(): target, _, self_ri, peer_ri = _pair() - target.check_peer(self_ri, peer_ri) # the supported pair - with pytest.raises(PeerIncompatibleError, match="dp_size"): + target.check_peer(self_ri, peer_ri) # (dep4, dep4) is in the matrix + with pytest.raises(PeerIncompatibleError, match="pair"): target.check_peer( - self_ri, _dep4_rank_info("gen", 0, _mla_page_table(0x9000_0000), dp_size=2) + self_ri, + _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000), target="deepseek_v3/other"), + ) + with pytest.raises(PeerIncompatibleError, match="pair"): + target.check_peer( + self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000), target="") ) with pytest.raises(PeerIncompatibleError, match="pool"): target.check_peer( @@ -176,14 +186,7 @@ def test_check_peer_accepts_the_dep4_twin_and_refuses_everything_else(): ) -def test_target_class_builds_this_planner(): - from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4 import ( # noqa: E501 - modeling, - ) - - cls = modeling.ModelingV2DeepseekR10528Nvfp4Sm103Dep4 - self_pt = _mla_page_table(0x1000_0000) - ri = _dep4_rank_info("ctx", 0, self_pt) - planner = cls.build_transfer_planner(object(), None, ri, KVRegionExtractorV1(self_pt)) - assert isinstance(planner, DeepseekR1Dep4TransferPlanner) - assert planner.planner_id == "deepseek_r1_0528_nvfp4_sm103_dep4" +def test_describe_publishes_own_target_id(): + target, _, _, _ = _pair() + assert target.describe() == DEP4.encode() + assert target.planner_id == "modeling_v2.deepseek_mla_latent" diff --git a/tests/unittest/_torch/modeling_v2/transfers/test_transfers_binding.py b/tests/unittest/_torch/modeling_v2/transfers/test_transfers_binding.py new file mode 100644 index 000000000000..3f2c6100ddfa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/transfers/test_transfers_binding.py @@ -0,0 +1,105 @@ +# 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. +"""A planner declares the targets it covers; a target knows nothing about planners.""" + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2 import transfers +from tensorrt_llm._torch._experimental.modeling_v2.transfers.deepseek_mla_latent import ( + DeepseekMlaLatentTransferPlanner, +) + +pytestmark = pytest.mark.cpu_only + +DEEPSEEK_DEP4 = "deepseek_v3/r1_0528_nvfp4__sm_103__dep4" +GPT_OSS_TP1 = "gpt_oss/gpt_oss_120b__sm_103__tp1" + + +def _target_cls(dotted_dir: str): + family, name = dotted_dir.split("/") + module = __import__( + f"tensorrt_llm._torch._experimental.modeling_v2.models.{family}.{name}.modeling", + fromlist=["_"], + ) + return next( + v + for k, v in vars(module).items() + if k.startswith("ModelingV2") + and isinstance(v, type) + and v.__module__ == module.__name__ + and "Target" not in k + ) + + +def test_target_id_is_family_slash_directory(): + assert transfers.target_id_of(_target_cls(DEEPSEEK_DEP4)) == DEEPSEEK_DEP4 + assert transfers.target_id_of(_target_cls(GPT_OSS_TP1)) == GPT_OSS_TP1 + assert transfers.target_id_of(dict) is None # not a modeling_v2 target + + +def test_planner_is_found_by_the_target_it_covers_and_nothing_else(): + assert transfers.planner_for_target(DEEPSEEK_DEP4) is DeepseekMlaLatentTransferPlanner + assert transfers.planner_for_target(GPT_OSS_TP1) is None + assert transfers.planner_for_target("nobody/such__sm_0__tp1") is None + + +def test_no_target_source_names_a_planner(): + """The binding direction: planners list targets, never the reverse.""" + from pathlib import Path + + root = Path(transfers.__file__).resolve().parents[1] / "models" + hits = [ + p.relative_to(root) + for p in root.rglob("*.py") + if "transfer_planner" in p.read_text() or "TransferPlanner" in p.read_text() + ] + assert hits == [], hits + + +def test_two_planners_covering_one_target_is_an_error(monkeypatch): + class Twin(DeepseekMlaLatentTransferPlanner): + pass + + monkeypatch.setattr(transfers, "PLANNERS", (DeepseekMlaLatentTransferPlanner, Twin)) + with pytest.raises(ValueError, match=DEEPSEEK_DEP4): + transfers.planner_for_target(DEEPSEEK_DEP4) + + +class _StubModel: + def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): + return "default" + + +def test_bind_leaves_an_uncovered_model_on_the_default_hook(): + stub = _StubModel() + assert transfers.bind_transfer_planner(stub, _target_cls(GPT_OSS_TP1)) is None + assert stub.build_transfer_planner(None, None, None) == "default" + assert transfers.bind_transfer_planner(stub, dict) is None + + +def test_bind_installs_a_factory_that_knows_its_own_target(monkeypatch): + seen = {} + + def fake_init(self, self_rank_info, self_extractor, *, self_target): + seen["args"] = (self_rank_info, self_extractor, self_target) + + monkeypatch.setattr(DeepseekMlaLatentTransferPlanner, "__init__", fake_init) + stub = _StubModel() + assert transfers.bind_transfer_planner(stub, _target_cls(DEEPSEEK_DEP4)) is ( + DeepseekMlaLatentTransferPlanner + ) + planner = stub.build_transfer_planner("kvm", "ri", "ext") + assert isinstance(planner, DeepseekMlaLatentTransferPlanner) + assert seen["args"] == ("ri", "ext", DEEPSEEK_DEP4) From da9b201335c80217182d8719dbe9bc5426611488 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sun, 11 Oct 2026 01:58:50 -0700 Subject: [PATCH 18/25] [None][doc] Record the two-node re-gate of the planner rebinding Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md index 27533e908ddc..0a249031682f 100644 --- a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md +++ b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md @@ -97,5 +97,9 @@ targets from different cores. Implemented instead: to zero transfer-related lines. - Tests: `tests/unittest/_torch/modeling_v2/transfers/` (binding direction, lookup, duplicate coverage, bind; planner parity/merge/targets/pair refusal/describe). 74 tests green with the - claims, routing and disaggregation planner suites. The two-node GSM8K gate was run before the - redesign; the plan arithmetic is unchanged and the parity test still pins it. + claims, routing and disaggregation planner suites. +- Two-node re-gate with the new binding (ctx dep4 on nvl72d185-T13, gen dep4 on nvl72d185-T16, + same launcher): all four ranks log + `KV transfer planned by DeepseekMlaLatentTransferPlanner (planner_id='modeling_v2.deepseek_mla_latent')`, + no `PeerIncompatibleError`; GSM8K flexible-extract **95.07** against the 94.24 reference, PASSED. + Logs: `mv2-disagg/twonode/logs/run_20261011T084450Z`. From 037c5868fe8a3f39c9769716f64ee5148f2e3164 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sun, 11 Oct 2026 11:50:07 -0700 Subject: [PATCH 19/25] [None][refactor] Hand TransferPlanner the request descriptions whole plan() took the fields of the sender's Chunk and the receiver's RecvReqInfo as separate positional arguments, so any per-request fact added to either type later would not reach a planner, and the drop would be silent. plan() now takes (self_ri, peer_ri, chunk, req_info, prompt_len) and receiver_payload_bytes() takes (self_ri, peer_ri, req_info). RecvReqInfo moves from native/transfer.py to base/backend.py beside Chunk: it is the receiver's wire message, the protocol type the planner contract names. The sender's KVSendTask is deliberately not passed; it carries transfer progress, counters and timers a planner must not touch. DefaultTransferPlanner and DeepseekMlaLatentTransferPlanner read the same values from the two objects; plans are unchanged. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- ...10-10-deepseek-r1-dep4-transfer-planner.md | 25 ++++++++ ...26-10-10-disagg-transfer-planner-design.md | 19 ++++-- .../transfers/deepseek_mla_latent.py | 22 ++++--- .../_torch/disaggregation/base/__init__.py | 2 + .../_torch/disaggregation/base/backend.py | 51 ++++++++++++++++ .../_torch/disaggregation/base/planner.py | 20 +++---- .../disaggregation/native/default_planner.py | 21 +++---- .../_torch/disaggregation/native/transfer.py | 60 +------------------ .../transfers/test_deepseek_mla_latent.py | 32 +++++++--- .../disaggregated/test_transfer_planner.py | 28 ++++----- 10 files changed, 163 insertions(+), 117 deletions(-) diff --git a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md index 0a249031682f..88d36dc11451 100644 --- a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md +++ b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md @@ -103,3 +103,28 @@ targets from different cores. Implemented instead: `KV transfer planned by DeepseekMlaLatentTransferPlanner (planner_id='modeling_v2.deepseek_mla_latent')`, no `PeerIncompatibleError`; GSM8K flexible-extract **95.07** against the 94.24 reference, PASSED. Logs: `mv2-disagg/twonode/logs/run_20261011T084450Z`. + +## Interface hardening (after review) + +Reviewing DeepSeek-V4.1 (PR #19727, merged 2026-10-11) against this branch showed +that `plan()` took a copy of `Chunk`'s field list as positional arguments. main +added `Chunk.excluded_pool_views` (remote-tail / bounded-replay SWA views the +generation side rebuilds locally) and filters on it inside the sender's pool +loop, which this branch moved into `DefaultTransferPlanner.plan`. With the old +signature that field could not reach the planner, and a textual rebase would +have merged cleanly while dropping the filter. + +`plan(self_ri, peer_ri, chunk, req_info, prompt_len)` and +`receiver_payload_bytes(self_ri, peer_ri, req_info)` now receive the two request +descriptions whole. `RecvReqInfo` moved from `native/transfer.py` to +`base/backend.py` next to `Chunk`, since it is the wire protocol type the planner +contract depends on. The sender's `KVSendTask` is deliberately not passed: it +carries transfer progress, counters and timers. + +Still to do at rebase time, not here: honor `chunk.excluded_pool_views` in +`DefaultTransferPlanner.plan` and fold `validate_bounded_replay_compatible` into +`check_peer`; both land on the planner, not on `transfer.py`. + +Verified: 292 unit tests green (planner, chunked transfer, ownership regressions, +unsubmitted session, peer, modeling_v2 transfers/claims/routing/target +contract/op bind). diff --git a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md index b9e8570ef1bb..2bc76d2ff3d1 100644 --- a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md +++ b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md @@ -84,23 +84,32 @@ class TransferPlanner(ABC): self, self_ri: RankInfo, peer_ri: RankInfo, - src_block_ids_per_lg: list[np.ndarray], - dst_block_ids_per_lg: list[np.ndarray], - token_range: TokenRange, + chunk: Chunk, # sender's description of the piece (base/backend.py) + req_info: RecvReqInfo, # receiver's description of where it lands (base/backend.py) prompt_len: int, - dst_start_token: Optional[int], ) -> TransferPlan: """Sender side. Replaces the body of _build_kv_write_meta.""" @abstractmethod def receiver_payload_bytes( - self, self_ri: RankInfo, peer_ri: RankInfo, dst_block_ids_per_lg: list[np.ndarray] + self, self_ri: RankInfo, peer_ri: RankInfo, req_info: RecvReqInfo ) -> Optional[int]: """Receiver side, bytes beyond the paged blocks that ride the same write (recurrent state today). None = cannot size, receiver disables bounce. Replaces _get_mamba_slot + mamba_receiver_payload_bytes.""" ``` +`plan` and `receiver_payload_bytes` take the two request descriptions whole, +never fields picked out of them. `Chunk` is what the sender holds for the piece +(block tables per layer group, token range, and whatever per-request facts the +transceiver records later, e.g. pool views a request excludes); `RecvReqInfo` is +the receiver's wire message (its block tables, `dst_start_token`, bounce +address, aux slot). Both live in `base/backend.py`, so the planner package +depends on protocol types only, never on `native/transfer.py`'s task objects +(`KVSendTask` carries transfer progress and timers a planner must not touch). +A new per-request fact is added to one of these two types and every planner +sees it; the interface does not change and nothing silently stops flowing. + Threading contract is unchanged from today: `plan` is called where `_build_kv_write_meta` is called now, `targets` and `receiver_payload_bytes` where `dispatch_task` calls their predecessors. Planner-internal caches diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py b/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py index a60cf733844e..eda6f71fa68d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py @@ -31,11 +31,11 @@ from __future__ import annotations -from typing import AbstractSet, FrozenSet, List, Optional, Tuple +from typing import FrozenSet, Optional, Tuple import numpy as np -from tensorrt_llm._torch.disaggregation.base.backend import TokenRange +from tensorrt_llm._torch.disaggregation.base.backend import Chunk, RecvReqInfo from tensorrt_llm._torch.disaggregation.base.planner import ( PeerIncompatibleError, PeerTargets, @@ -135,21 +135,19 @@ def plan( self, self_ri: RankInfo, peer_ri: RankInfo, - src_block_ids_per_lg: List[np.ndarray], - dst_block_ids_per_lg: List[np.ndarray], - token_range: TokenRange, + chunk: Chunk, + req_info: RecvReqInfo, prompt_len: int, - dst_start_token: Optional[int], - excluded_pool_views: AbstractSet[Tuple[int, int]], ) -> TransferPlan: - del token_range, prompt_len # cp == 1: a partial piece needs no projection - if (0, 0) in excluded_pool_views: # the sole latent pool is layer group 0, view 0 + del prompt_len # cp == 1: a partial piece needs no projection + if (0, 0) in chunk.excluded_pool_views: # the sole latent pool is layer group 0, view 0 return TransferPlan.empty() tpb = self._page_table.tokens_per_block + dst_start_token = req_info.dst_start_token dst_start_block = dst_start_token // tpb if dst_start_token is not None else 0 src_ids, dst_ids = DefaultTransferPlanner._pair_ordinals( - np.asarray(src_block_ids_per_lg[0], dtype=np.int64), - np.asarray(dst_block_ids_per_lg[0], dtype=np.int64), + np.asarray(chunk.block_ids_per_layer_groups[0], dtype=np.int64), + np.asarray(req_info.block_ids_per_layer_groups[0], dtype=np.int64), dst_start_block=dst_start_block, ) if src_ids.size == 0: @@ -173,6 +171,6 @@ def plan( ) def receiver_payload_bytes( - self, self_ri: RankInfo, peer_ri: RankInfo, dst_block_ids_per_lg: List[np.ndarray] + self, self_ri: RankInfo, peer_ri: RankInfo, req_info: RecvReqInfo ) -> Optional[int]: return 0 # no recurrent state rides along with the paged blocks diff --git a/tensorrt_llm/_torch/disaggregation/base/__init__.py b/tensorrt_llm/_torch/disaggregation/base/__init__.py index 1d3e9dd767d3..d3dcdfac8fad 100644 --- a/tensorrt_llm/_torch/disaggregation/base/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/base/__init__.py @@ -29,6 +29,7 @@ Fetches, Outcome, Publishes, + RecvReqInfo, TokenRange, ) @@ -43,5 +44,6 @@ "Fetches", "Outcome", "Publishes", + "RecvReqInfo", "TokenRange", ] diff --git a/tensorrt_llm/_torch/disaggregation/base/backend.py b/tensorrt_llm/_torch/disaggregation/base/backend.py index 3e5c500826a1..2fcbc58c0b4e 100644 --- a/tensorrt_llm/_torch/disaggregation/base/backend.py +++ b/tensorrt_llm/_torch/disaggregation/base/backend.py @@ -20,6 +20,7 @@ from enum import IntEnum from typing import List, Optional, Protocol, Set, Tuple, Union, runtime_checkable +import msgpack import numpy as np # TODO: @@ -38,6 +39,7 @@ "Fetches", "Outcome", "Publishes", + "RecvReqInfo", "TokenRange", ] @@ -114,6 +116,55 @@ def __post_init__(self): ) +@dataclass +class RecvReqInfo: + """The receiver's description of one request, sent to the sender over the wire. + + The receiver-side counterpart of :class:`Chunk`: everything the sender's planner may + need about where the bytes land. New per-request receiver facts belong here, so they + reach the planner without changing its interface. + """ + + sender_req_id: int + instance_name: str + instance_rank: int + # Positional pool slots per layer group (see Chunk): entry i is the + # receiver's slot for block ordinal i, or -1 where it accepts nothing. + block_ids_per_layer_groups: list[np.ndarray] + unique_rid: int + # Optional extra lower bound (block-aligned token offset) below which the + # receiver accepts nothing, on top of the -1 holes already in its table. + dst_start_token: Optional[int] = None + aux_slot: Optional[int] = None + slice_id: Optional[int] = None + bounce_dst_base: Optional[int] = None + + def to_bytes(self) -> bytes: + return msgpack.packb( + { + "sender_req_id": self.sender_req_id, + "instance_name": self.instance_name, + "instance_rank": self.instance_rank, + "block_ids_per_layer_groups": [ + arr.tobytes() for arr in self.block_ids_per_layer_groups + ], + "unique_rid": self.unique_rid, + "dst_start_token": self.dst_start_token, + "aux_slot": self.aux_slot, + "slice_id": self.slice_id, + "bounce_dst_base": self.bounce_dst_base, + } + ) + + @classmethod + def from_bytes(cls, data: bytes) -> "RecvReqInfo": + d = msgpack.unpackb(data, raw=False) + d["block_ids_per_layer_groups"] = [ + np.frombuffer(b, dtype=np.int64).copy() for b in d["block_ids_per_layer_groups"] + ] + return cls(**d) + + @dataclass class CacheExtent: """One ask: which content, and where this rank's blocks for it are. diff --git a/tensorrt_llm/_torch/disaggregation/base/planner.py b/tensorrt_llm/_torch/disaggregation/base/planner.py index e0dea1e3954b..00945f1bf21c 100644 --- a/tensorrt_llm/_torch/disaggregation/base/planner.py +++ b/tensorrt_llm/_torch/disaggregation/base/planner.py @@ -28,12 +28,12 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, AbstractSet, List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional import numpy as np if TYPE_CHECKING: - from tensorrt_llm._torch.disaggregation.base.backend import TokenRange + from tensorrt_llm._torch.disaggregation.base.backend import Chunk, RecvReqInfo from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo @@ -103,22 +103,22 @@ def plan( self, self_ri: "RankInfo", peer_ri: "RankInfo", - src_block_ids_per_lg: List[np.ndarray], - dst_block_ids_per_lg: List[np.ndarray], - token_range: "TokenRange", + chunk: "Chunk", + req_info: "RecvReqInfo", prompt_len: int, - dst_start_token: Optional[int], - excluded_pool_views: AbstractSet[Tuple[int, int]], ) -> TransferPlan: """Sender side: the bytes this rank writes into ``peer_ri`` for one chunk. - ``excluded_pool_views`` are ``(layer_group, pool_view)`` pairs the cache - manager keeps out of this chunk; the plan writes nothing from them. + ``chunk`` is this rank's description of what it holds for the piece (block + tables per layer group, token range), ``req_info`` the receiver's description + of where it lands (its block tables, a lower bound below which it accepts + nothing). Both arrive whole so a planner honors every per-request fact the + framework records, present or future, without a signature change. """ @abstractmethod def receiver_payload_bytes( - self, self_ri: "RankInfo", peer_ri: "RankInfo", dst_block_ids_per_lg: List[np.ndarray] + self, self_ri: "RankInfo", peer_ri: "RankInfo", req_info: "RecvReqInfo" ) -> Optional[int]: """Receiver side: bytes beyond the paged blocks that ride the same write. diff --git a/tensorrt_llm/_torch/disaggregation/native/default_planner.py b/tensorrt_llm/_torch/disaggregation/native/default_planner.py index 0e265928aeec..5526efaf4539 100644 --- a/tensorrt_llm/_torch/disaggregation/native/default_planner.py +++ b/tensorrt_llm/_torch/disaggregation/native/default_planner.py @@ -25,12 +25,12 @@ from collections import Counter from dataclasses import dataclass, field -from typing import AbstractSet, Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import numpy as np from tensorrt_llm import logger -from tensorrt_llm._torch.disaggregation.base.backend import TokenRange +from tensorrt_llm._torch.disaggregation.base.backend import Chunk, RecvReqInfo from tensorrt_llm._torch.disaggregation.base.planner import ( PeerIncompatibleError, PeerTargets, @@ -214,13 +214,14 @@ def plan( self, self_ri: RankInfo, peer_ri: RankInfo, - src_block_ids_per_lg: List[np.ndarray], - dst_block_ids_per_lg: List[np.ndarray], - token_range: TokenRange, + chunk: Chunk, + req_info: RecvReqInfo, prompt_len: int, - dst_start_token: Optional[int], - excluded_pool_views: AbstractSet[Tuple[int, int]], ) -> TransferPlan: + src_block_ids_per_lg = chunk.block_ids_per_layer_groups + dst_block_ids_per_lg = req_info.block_ids_per_layer_groups + token_range = chunk.token_range + dst_start_token = req_info.dst_start_token # Only a piece short of the whole prompt runs the chunk projection below, and that is the # part with no answer under CP; a sole piece covering everything takes the untouched path. is_partial = token_range.start > 0 or token_range.end < prompt_len @@ -249,7 +250,7 @@ def plan( # Send ownership is per pool: replicated pools elect one fan-in # owner, sharded pools keep head-duplication routing. for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): - if (self_lg, self_pi) in excluded_pool_views: + if (self_lg, self_pi) in chunk.excluded_pool_views: continue if not self.should_send_pool(targets, peer_ri, self_lg, self_pi): continue @@ -312,10 +313,10 @@ def plan( ) def receiver_payload_bytes( - self, self_ri: RankInfo, peer_ri: RankInfo, dst_block_ids_per_lg: List[np.ndarray] + self, self_ri: RankInfo, peer_ri: RankInfo, req_info: RecvReqInfo ) -> Optional[int]: receiver_pt = self._self_ext_cache.page_table - dst_slot = self._get_mamba_slot(receiver_pt, dst_block_ids_per_lg) + dst_slot = self._get_mamba_slot(receiver_pt, req_info.block_ids_per_layer_groups) if dst_slot is None: return 0 if peer_ri.page_table is None: diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 2ed4a8433d16..a66e7727ec9f 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -27,7 +27,6 @@ from enum import Enum from typing import TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List, Optional, Union -import msgpack import numpy as np import torch @@ -38,7 +37,7 @@ import tensorrt_llm.bindings from tensorrt_llm import logger -from tensorrt_llm._torch.disaggregation.base import Chunk +from tensorrt_llm._torch.disaggregation.base import Chunk, RecvReqInfo from tensorrt_llm._torch.disaggregation.base.agent import ( BaseTransferAgent, MemoryDescs, @@ -107,48 +106,6 @@ class _TransferNotSubmittedError(RuntimeError): """The sender rejected an operation before the backend could access memory.""" -@dataclass -class RecvReqInfo: - sender_req_id: int - instance_name: str - instance_rank: int - # Positional pool slots per layer group (see Chunk): entry i is the - # receiver's slot for block ordinal i, or -1 where it accepts nothing. - block_ids_per_layer_groups: list[np.ndarray] - unique_rid: int - # Optional extra lower bound (block-aligned token offset) below which the - # receiver accepts nothing, on top of the -1 holes already in its table. - dst_start_token: Optional[int] = None - aux_slot: Optional[int] = None - slice_id: Optional[int] = None - bounce_dst_base: Optional[int] = None - - def to_bytes(self) -> bytes: - return msgpack.packb( - { - "sender_req_id": self.sender_req_id, - "instance_name": self.instance_name, - "instance_rank": self.instance_rank, - "block_ids_per_layer_groups": [ - arr.tobytes() for arr in self.block_ids_per_layer_groups - ], - "unique_rid": self.unique_rid, - "dst_start_token": self.dst_start_token, - "aux_slot": self.aux_slot, - "slice_id": self.slice_id, - "bounce_dst_base": self.bounce_dst_base, - } - ) - - @classmethod - def from_bytes(cls, data: bytes) -> "RecvReqInfo": - d = msgpack.unpackb(data, raw=False) - d["block_ids_per_layer_groups"] = [ - np.frombuffer(b, dtype=np.int64).copy() for b in d["block_ids_per_layer_groups"] - ] - return cls(**d) - - @dataclass class ReadMeta: unique_rid: int @@ -1690,16 +1647,7 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write # Which of this rank's bytes land where on the peer is the model's planner's call; # an empty plan still goes out so the peer's expected-writer count settles. - plan = self._planner.plan( - self_ri, - peer_ri, - task._chunk.block_ids_per_layer_groups, - req_info.block_ids_per_layer_groups, - task._chunk.token_range, - task._prompt_len, - req_info.dst_start_token, - task._chunk.excluded_pool_views, - ) + plan = self._planner.plan(self_ri, peer_ri, task._chunk, req_info, task._prompt_len) if timer: timer.record_prepare_args_end(peer_ri.instance_rank) @@ -3085,9 +3033,7 @@ def dispatch_task(self, task: KVRecvTask) -> None: ) # Bytes beyond the paged blocks that ride the same coalesced write # (recurrent state today) size the bounce region too; None = unknown. - extra_bytes = self._planner.receiver_payload_bytes( - self_ri, peer_infos, receiver_req.block_ids_per_layer_groups - ) + extra_bytes = self._planner.receiver_payload_bytes(self_ri, peer_infos, receiver_req) if extra_bytes is None: allow_bounce = False extra_bytes = 0 diff --git a/tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py b/tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py index e187cb4d9368..3707730dc760 100644 --- a/tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py +++ b/tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py @@ -20,7 +20,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.transfers.deepseek_mla_latent import ( DeepseekMlaLatentTransferPlanner, ) -from tensorrt_llm._torch.disaggregation.base import TokenRange +from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, RecvReqInfo, TokenRange from tensorrt_llm._torch.disaggregation.base.planner import PeerIncompatibleError, TransferPlan from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo @@ -122,8 +122,24 @@ def _fragments_at_slot_granularity(plan: TransferPlan) -> set: return out -def _tables(src, dst): - return [np.array(src, dtype=np.int64)], [np.array(dst, dtype=np.int64)] +def _request(src, dst, dst_start_token=None): + """Sender-side chunk and receiver-side request for one single-layer-group transfer.""" + prompt_len = len(src) * TPB + chunk = Chunk( + block_ids_per_layer_groups=[np.array(src, dtype=np.int64)], + kind_per_layer_group=[CacheKind.PAGED], + token_range=TokenRange(0, prompt_len), + is_last=True, + ) + req_info = RecvReqInfo( + sender_req_id=1, + instance_name="gen", + instance_rank=2, + block_ids_per_layer_groups=[np.array(dst, dtype=np.int64)], + unique_rid=1, + dst_start_token=dst_start_token, + ) + return chunk, req_info, prompt_len @pytest.mark.parametrize( @@ -140,11 +156,10 @@ def _tables(src, dst): ) def test_plan_matches_default_planner_byte_for_byte(src, dst, dst_start_token): target, default, self_ri, peer_ri = _pair() - s, d = _tables(src, dst) - tr = TokenRange(0, len(src) * TPB) + chunk, req_info, prompt_len = _request(src, dst, dst_start_token) - ours = target.plan(self_ri, peer_ri, s, d, tr, len(src) * TPB, dst_start_token, frozenset()) - ref = default.plan(self_ri, peer_ri, s, d, tr, len(src) * TPB, dst_start_token, frozenset()) + ours = target.plan(self_ri, peer_ri, chunk, req_info, prompt_len) + ref = default.plan(self_ri, peer_ri, chunk, req_info, prompt_len) assert _fragments_at_slot_granularity(ours) == _fragments_at_slot_granularity(ref) assert int(ours.sizes.sum()) == int(ref.sizes.sum()) @@ -152,8 +167,7 @@ def test_plan_matches_default_planner_byte_for_byte(src, dst, dst_start_token): def test_contiguous_blocks_merge_into_one_fragment(): target, _, self_ri, peer_ri = _pair() - s, d = _tables([3, 4, 5, 6], [10, 11, 12, 13]) - plan = target.plan(self_ri, peer_ri, s, d, TokenRange(0, 4 * TPB), 4 * TPB, None, frozenset()) + plan = target.plan(self_ri, peer_ri, *_request([3, 4, 5, 6], [10, 11, 12, 13])) assert plan.sizes.tolist() == [4 * SLOT_BYTES] diff --git a/tests/unittest/disaggregated/test_transfer_planner.py b/tests/unittest/disaggregated/test_transfer_planner.py index 829f27326f4c..559b57aa279a 100644 --- a/tests/unittest/disaggregated/test_transfer_planner.py +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -48,12 +48,13 @@ def test_peer_incompatible_error_is_value_error(): from unittest.mock import Mock # noqa: E402 from tensorrt_llm import DisaggregatedParams # noqa: E402 -from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, TokenRange # noqa: E402 -from tensorrt_llm._torch.disaggregation.native.transfer import ( # noqa: E402 - KVSendTask, +from tensorrt_llm._torch.disaggregation.base import ( # noqa: E402 + CacheKind, + Chunk, RecvReqInfo, - Sender, + TokenRange, ) +from tensorrt_llm._torch.disaggregation.native.transfer import KVSendTask, Sender # noqa: E402 class _ToyPlanner(TransferPlanner): @@ -72,13 +73,11 @@ def check_peer(self, self_ri, peer_ri): def targets(self, self_ri, peer_ri, peer_dp_rank): return PeerTargets(ranks=list(self._ranks), uniform_writers=True) - def plan( - self, self_ri, peer_ri, src, dst, token_range, prompt_len, dst_start_token, excluded - ): - self.calls.append((src, dst, token_range, prompt_len, dst_start_token, excluded)) + def plan(self, self_ri, peer_ri, chunk, req_info, prompt_len): + self.calls.append((chunk, req_info, prompt_len)) return self._plan - def receiver_payload_bytes(self, self_ri, peer_ri, dst): + def receiver_payload_bytes(self, self_ri, peer_ri, req_info): return 0 @@ -141,11 +140,9 @@ def test_build_kv_write_meta_uses_planner_plan(): assert wm.src_ptrs is plan.src_ptrs and wm.dst_ptrs is plan.dst_ptrs and wm.sizes is plan.sizes assert wm.expected_transfers == 2 and task.expected_transfers == 2 assert wm.dst_device_id == 1 and wm.peer_rank == 3 and wm.peer_endpoint == "tcp://gen:1" - ((src, dst, token_range, prompt_len, dst_start, excluded),) = planner.calls - assert src is task._chunk.block_ids_per_layer_groups - assert dst is req.block_ids_per_layer_groups - assert excluded is task._chunk.excluded_pool_views - assert (token_range.start, token_range.end, prompt_len, dst_start) == (0, 64, 64, 32) + # The planner gets both request descriptions whole, not fields picked out of them. + ((chunk, req_info, prompt_len),) = planner.calls + assert chunk is task._chunk and req_info is req and prompt_len == 64 def test_build_kv_write_meta_empty_plan_still_delivers(): @@ -214,6 +211,9 @@ def test_dispatch_task_payload_none_disables_bounce(): recv._bounce.reserve.assert_not_called() recv._request_sender_data.assert_called_once_with("tcp://s0", b"req") assert task.expected_transfers == 1 + # The receiver hands its request description to the planner whole. + (_self_ri, _peer_ri, req_info), _ = planner.receiver_payload_bytes.call_args + assert req_info is recv._build_recv_req_info.return_value # --------------------------------------------------------------------------- # From b7958a452fe118a7563fdbee8483dad06c127983 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:03:24 -0700 Subject: [PATCH 20/25] [None][refactor] Identify transfer-planner deployments by architecture and topology, not by modeling_v2 target A planner's compatibility matrix named deployments by modeling_v2 target id, and the binding looked only at modeling_v2 classes. A context worker running the built-in DeepseekV3ForCausalLM therefore got the default planner while a generation worker on the modeling_v2 dep4 target got DeepseekMlaLatentTransferPlanner, and the framework's planner_id check refused the pair at registration. The bytes were compatible all along: the MLA latent KV layout is the cache manager's, not the implementation's. Deployment (architecture plus tp/pp/cp/moe_ep/moe_tp/attention_dp) is now the identity. It is read off the ModelConfig for every model AutoModelForCausalLM builds, so a built-in model and a modeling_v2 target on the same topology bind to the same planner and publish the same description, and either may feed the other. COVERS pairs, describe() and the pair lookup use Deployment; the planner package moves from modeling_v2/transfers to disaggregation/planners since it was never a modeling_v2 property; target_id_of goes away with its last caller. A consequence worth naming: a pure built-in DeepseekV3 dep4 disaggregated deployment now also runs on DeepseekMlaLatentTransferPlanner instead of the default planner. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../_experimental/modeling_v2/__init__.py | 2 - .../modeling_v2/_router_index.py | 17 --- .../planners}/__init__.py | 49 ++++--- .../planners}/deepseek_mla_latent.py | 57 +++++---- .../disaggregation/planners/deployment.py | 75 +++++++++++ tensorrt_llm/_torch/models/modeling_auto.py | 14 +- .../test_deepseek_mla_latent_planner.py} | 34 +++-- .../disaggregation/test_planner_binding.py | 121 ++++++++++++++++++ .../transfers/test_transfers_binding.py | 105 --------------- 9 files changed, 286 insertions(+), 188 deletions(-) rename tensorrt_llm/_torch/{_experimental/modeling_v2/transfers => disaggregation/planners}/__init__.py (51%) rename tensorrt_llm/_torch/{_experimental/modeling_v2/transfers => disaggregation/planners}/deepseek_mla_latent.py (76%) create mode 100644 tensorrt_llm/_torch/disaggregation/planners/deployment.py rename tests/unittest/_torch/{modeling_v2/transfers/test_deepseek_mla_latent.py => disaggregation/test_deepseek_mla_latent_planner.py} (89%) create mode 100644 tests/unittest/_torch/disaggregation/test_planner_binding.py delete mode 100644 tests/unittest/_torch/modeling_v2/transfers/test_transfers_binding.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index f81814ff79fb..c5140e7c912b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -33,7 +33,6 @@ NoTarget, modeling_v2_resolve, step_contract_enabled, - target_id_of, ) __all__ = [ @@ -43,5 +42,4 @@ "NoTarget", "modeling_v2_resolve", "step_contract_enabled", - "target_id_of", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 3f7a0a6e706e..cdb288f7220f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -101,23 +101,6 @@ class NoTarget(Exception): """ -def target_id_of(model_cls: type) -> Optional[str]: - """``/`` for a class defined in a modeling_v2 target, else None. - - The directory name is a target's identity (``____``); - the family prefix keeps two families' checkpoints from colliding. Cheap on - purpose: ``AutoModelForCausalLM`` asks this for every model it builds. - """ - module = getattr(model_cls, "__module__", None) or "" - prefix = f"{_PACKAGE}.models." - if not module.startswith(prefix): - return None - parts = module[len(prefix) :].split(".") - if len(parts) != 3: # family.directory.modeling - return None - return f"{parts[0]}/{parts[1]}" - - def routing_module(arch: str): """Import and return the routing module for ``arch``, or None.""" name = MODELING_V2_ROUTERS.get(arch) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/__init__.py b/tensorrt_llm/_torch/disaggregation/planners/__init__.py similarity index 51% rename from tensorrt_llm/_torch/_experimental/modeling_v2/transfers/__init__.py rename to tensorrt_llm/_torch/disaggregation/planners/__init__.py index 412f156867ca..51ecf0e4219d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/planners/__init__.py @@ -12,13 +12,16 @@ # 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. -"""Which disaggregated KV transfer planner serves a modeling_v2 target. +"""Which disaggregated KV transfer planner serves a deployment. -A planner owns a compatibility set: the (context target, generation target) -pairs it can move bytes between, declared on the class as ``COVERS``. Targets -declare nothing. This module is the one place that knows every planner, and it -answers one question: given the target a model instance was built from, which -planner covers it? None means the engine keeps the framework's default plan. +A planner owns a compatibility set: the (context deployment, generation +deployment) pairs it can move bytes between, declared on the class as +``COVERS``. A :class:`Deployment` is the architecture plus the parallel +topology, so one entry covers the built-in model and any modeling_v2 target +alike. Models declare nothing. This module is the one place that knows every +planner, and it answers one question: given the deployment a model instance +was built for, which planner covers it? None means the engine keeps the +framework's default plan. Binding happens where the model is constructed (``AutoModelForCausalLM``), by installing the covering planner's factory as the instance's @@ -27,45 +30,49 @@ from __future__ import annotations -from typing import Optional, Type +from typing import TYPE_CHECKING, Optional, Type from tensorrt_llm._torch.disaggregation.base.planner import TransferPlanner -from .._router_index import target_id_of from .deepseek_mla_latent import DeepseekMlaLatentTransferPlanner +from .deployment import Deployment -#: Every planner a target may be served by. Order carries no meaning: a target -#: covered by two planners is a configuration error, not a priority question. +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + +#: Every planner a deployment may be served by. Order carries no meaning: a +#: deployment covered by two planners is a configuration error, not a priority +#: question. PLANNERS: tuple = (DeepseekMlaLatentTransferPlanner,) -def planner_for_target(target_id: str) -> Optional[Type[TransferPlanner]]: - """The planner whose matrix mentions ``target_id`` on either side, or None.""" - found = [p for p in PLANNERS if any(target_id in pair for pair in p.COVERS)] +def planner_for(deployment: Deployment) -> Optional[Type[TransferPlanner]]: + """The planner whose matrix mentions ``deployment`` on either side, or None.""" + found = [p for p in PLANNERS if any(deployment in pair for pair in p.COVERS)] if len(found) > 1: names = ", ".join(p.__name__ for p in found) - raise ValueError(f"target {target_id!r} is covered by more than one planner: {names}") + raise ValueError(f"{deployment} is covered by more than one planner: {names}") return found[0] if found else None -def bind_transfer_planner(model, model_cls: type) -> Optional[Type[TransferPlanner]]: +def bind_transfer_planner(model, model_config: "ModelConfig") -> Optional[Type[TransferPlanner]]: """Install the covering planner's factory on ``model``; return the planner or None. The factory has the ``build_transfer_planner`` signature the framework - calls, closing over the target id so the planner can publish it. + calls, closing over the deployment so the planner can publish it. """ - target_id = target_id_of(model_cls) - if target_id is None: + deployment = Deployment.of(model_config) + if deployment is None: return None - planner_cls = planner_for_target(target_id) + planner_cls = planner_for(deployment) if planner_cls is None: return None def build_transfer_planner(kv_cache_manager, self_rank_info, self_extractor): - return planner_cls(self_rank_info, self_extractor, self_target=target_id) + return planner_cls(self_rank_info, self_extractor, self_deployment=deployment) model.build_transfer_planner = build_transfer_planner return planner_cls -__all__ = ["PLANNERS", "bind_transfer_planner", "planner_for_target", "target_id_of"] +__all__ = ["PLANNERS", "Deployment", "bind_transfer_planner", "planner_for"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py similarity index 76% rename from tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py rename to tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py index eda6f71fa68d..d4425992f991 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/transfers/deepseek_mla_latent.py +++ b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py @@ -12,7 +12,7 @@ # 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. -"""Disaggregated KV transfer plan for DeepSeek targets that keep MLA latent KV in one pool. +"""Disaggregated KV transfer plan for DeepSeek deployments that keep MLA latent KV in one pool. Under attention-DP every context rank has a generation twin with the same dp_rank. MLA keeps one latent-K region per layer in one physical pool, so a @@ -21,12 +21,14 @@ copy. Blocks that are consecutive on both sides move as one fragment. Which deployments may talk to each other is the ``COVERS`` matrix below: pairs -of (context target, generation target) ids. A target never names this planner; -``modeling_v2.transfers`` finds the planner from the target id, and ``describe`` -publishes that id so the peer's planner can look the pair up too. Anything not -in the matrix is refused at registration; the generic default planner stays -untouched for the models that need pool matching, head re-sharding or -recurrent state. +of (context deployment, generation deployment). A deployment is the +architecture plus its parallel topology, so the built-in DeepseekV3 model and a +modeling_v2 target on the same topology are the same deployment and may feed +each other. A model never names this planner; ``disaggregation.planners`` +finds the planner from the deployment, and ``describe`` publishes that +deployment so the peer's planner can look the pair up too. Anything not in the +matrix is refused at registration; the generic default planner stays untouched +for the models that need pool matching, head re-sharding or recurrent state. """ from __future__ import annotations @@ -52,7 +54,12 @@ ) from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool -_R1_DEP4 = "deepseek_v3/r1_0528_nvfp4__sm_103__dep4" +from .deployment import Deployment + +# Attention-DP over 4 ranks, experts sharded 4 ways, attention replicated. +_DEEPSEEK_DEP4 = Deployment( + architecture="DeepseekV3ForCausalLM", tp=4, pp=1, cp=1, moe_ep=4, moe_tp=1, attention_dp=True +) def _sole_pool(page_table: Optional[KVCachePageTable], side: str) -> PhysicalPool: @@ -75,32 +82,36 @@ def _sole_pool(page_table: Optional[KVCachePageTable], side: str) -> PhysicalPoo class DeepseekMlaLatentTransferPlanner(TransferPlanner): - planner_id = "modeling_v2.deepseek_mla_latent" + planner_id = "trtllm.deepseek_mla_latent" - #: (context target, generation target) pairs this planner knows how to move - #: bytes between. The same pair is listed once; the planner serves both - #: roles, so a pair is read from the sender's side or the receiver's side. - COVERS: FrozenSet[Tuple[str, str]] = frozenset({(_R1_DEP4, _R1_DEP4)}) + #: (context deployment, generation deployment) pairs this planner knows how + #: to move bytes between. The same pair is listed once; the planner serves + #: both roles, so a pair is read from the sender's side or the receiver's. + COVERS: FrozenSet[Tuple[Deployment, Deployment]] = frozenset({(_DEEPSEEK_DEP4, _DEEPSEEK_DEP4)}) def __init__( - self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1, *, self_target: str + self, + self_rank_info: RankInfo, + self_extractor: KVRegionExtractorV1, + *, + self_deployment: Deployment, ): self._ri = self_rank_info - self._target = self_target + self._deployment = self_deployment self._page_table = self_extractor.page_table self._pool = _sole_pool(self._page_table, "this rank") def describe(self) -> bytes: - return self._target.encode() + return self._deployment.to_bytes() - def _pair(self, peer_ri: RankInfo) -> Tuple[str, str]: - """The (self, peer) target pair, or a refusal if the matrix lacks it.""" - peer_target = peer_ri.planner_blob.decode() - if (self._target, peer_target) in self.COVERS or (peer_target, self._target) in self.COVERS: - return self._target, peer_target + def _pair(self, peer_ri: RankInfo) -> Tuple[Deployment, Deployment]: + """The (self, peer) deployment pair, or a refusal if the matrix lacks it.""" + peer = Deployment.from_bytes(peer_ri.planner_blob) + if (self._deployment, peer) in self.COVERS or (peer, self._deployment) in self.COVERS: + return self._deployment, peer raise PeerIncompatibleError( - f"no transfer plan for the pair ({self._target!r}, {peer_target!r}); " - f"covered pairs: {sorted(self.COVERS)}" + f"no transfer plan for the pair ({self._deployment}, {peer}); " + f"covered pairs: {[tuple(map(str, pair)) for pair in self.COVERS]}" ) # ---- TransferPlanner protocol ------------------------------------------------- diff --git a/tensorrt_llm/_torch/disaggregation/planners/deployment.py b/tensorrt_llm/_torch/disaggregation/planners/deployment.py new file mode 100644 index 000000000000..af172ea0a41a --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/planners/deployment.py @@ -0,0 +1,75 @@ +# 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 a transfer planner identifies a deployment by. + +A KV transfer plan depends on the model architecture and on how each side is +sharded, not on which implementation computed the KV: a built-in +``DeepseekV3ForCausalLM`` and a modeling_v2 target for the same checkpoint on +the same topology write the same bytes into the same pools. ``Deployment`` +therefore names exactly those two things, so a planner's compatibility matrix +covers both implementations at once and a context on one can feed a +generation on the other. +""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + + +@dataclass(frozen=True) +class Deployment: + """The architecture of a model and the parallel topology it runs under.""" + + architecture: str + tp: int + pp: int + cp: int + moe_ep: int + moe_tp: int + attention_dp: bool + + @classmethod + def of(cls, model_config: "ModelConfig") -> Optional["Deployment"]: + """The deployment ``model_config`` describes, or None when it names no architecture.""" + architectures = getattr(model_config.pretrained_config, "architectures", None) + if not architectures: + return None + m = model_config.mapping + return cls( + architecture=architectures[0], + tp=m.tp_size, + pp=m.pp_size, + cp=m.cp_size, + moe_ep=m.moe_ep_size, + moe_tp=m.moe_tp_size, + attention_dp=bool(m.enable_attention_dp), + ) + + def to_bytes(self) -> bytes: + """Wire form for ``RankInfo.planner_blob``.""" + return json.dumps(asdict(self), sort_keys=True).encode() + + @classmethod + def from_bytes(cls, blob: bytes) -> Optional["Deployment"]: + """Inverse of :meth:`to_bytes`; None for anything that is not one.""" + try: + return cls(**json.loads(blob.decode())) + except (ValueError, TypeError, UnicodeDecodeError): + return None diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 5d6f3eec2abd..d7f5fb5ae909 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 from typing import Generic, Optional, Type -from .._experimental.modeling_v2 import target_id_of from ..model_config import ModelConfig from ..utils import model_extra_attrs from .modeling_utils import (DecoderModelForCausalLM, TConfig, TModel, @@ -78,11 +77,10 @@ def from_config( with model_extra_attrs(extra_attrs): model = cls(config) model.extra_attrs = extra_attrs - # A modeling_v2 target may be served by a planner that covers it (see - # modeling_v2.transfers). Resolved here, where the target is known, - # so neither the target nor the disaggregation code names the other. - if target_id_of(cls) is not None: - from .._experimental.modeling_v2.transfers import \ - bind_transfer_planner - bind_transfer_planner(model, cls) + # A KV transfer planner may cover this deployment (architecture and + # parallel topology; see disaggregation.planners). Resolved here, where + # the config is known and before the transceiver asks for the hook, so + # neither the model nor the disaggregation code names the other. + from ..disaggregation.planners import bind_transfer_planner + bind_transfer_planner(model, config) return model diff --git a/tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py b/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py similarity index 89% rename from tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py rename to tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py index 3707730dc760..fd34378c196f 100644 --- a/tests/unittest/_torch/modeling_v2/transfers/test_deepseek_mla_latent.py +++ b/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py @@ -14,17 +14,20 @@ # limitations under the License. """The DeepSeek MLA-latent planner must produce the default planner's bytes for the pairs it covers.""" +from dataclasses import replace + import numpy as np import pytest -from tensorrt_llm._torch._experimental.modeling_v2.transfers.deepseek_mla_latent import ( - DeepseekMlaLatentTransferPlanner, -) from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, RecvReqInfo, TokenRange from tensorrt_llm._torch.disaggregation.base.planner import PeerIncompatibleError, TransferPlan from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.planners.deepseek_mla_latent import ( + DeepseekMlaLatentTransferPlanner, +) +from tensorrt_llm._torch.disaggregation.planners.deployment import Deployment from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import ( BUFFER_ENTRY_DTYPE, @@ -42,7 +45,9 @@ LAYER_BYTES = 512 # one latent-K region per layer per block SLOT_BYTES = NUM_LAYERS * LAYER_BYTES TPB = 16 -DEP4 = "deepseek_v3/r1_0528_nvfp4__sm_103__dep4" +DEP4 = Deployment( + architecture="DeepseekV3ForCausalLM", tp=4, pp=1, cp=1, moe_ep=4, moe_tp=1, attention_dp=True +) def _mla_page_table(base: int, *, num_layers: int = NUM_LAYERS, extra_pool: bool = False): @@ -70,11 +75,11 @@ def _mla_page_table(base: int, *, num_layers: int = NUM_LAYERS, extra_pool: bool def _dep4_rank_info( - name: str, dp_rank: int, page_table, *, dp_size: int = 4, pp_size: int = 1, target: str = DEP4 + name: str, dp_rank: int, page_table, *, dp_size: int = 4, pp_size: int = 1, blob: bytes = None ): return RankInfo( planner_id=DeepseekMlaLatentTransferPlanner.planner_id, - planner_blob=target.encode(), + planner_blob=DEP4.to_bytes() if blob is None else blob, instance_name=name, instance_rank=dp_rank, tp_size=dp_size, @@ -105,7 +110,7 @@ def _pair(dp_rank: int = 2): peer_ri = _dep4_rank_info("gen", dp_rank, peer_pt) ext = KVRegionExtractorV1(self_pt) return ( - DeepseekMlaLatentTransferPlanner(self_ri, ext, self_target=DEP4), + DeepseekMlaLatentTransferPlanner(self_ri, ext, self_deployment=DEP4), DefaultTransferPlanner(self_ri, ext), self_ri, peer_ri, @@ -184,11 +189,16 @@ def test_check_peer_accepts_a_covered_pair_and_refuses_everything_else(): with pytest.raises(PeerIncompatibleError, match="pair"): target.check_peer( self_ri, - _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000), target="deepseek_v3/other"), + _dep4_rank_info( + "gen", + 2, + _mla_page_table(0x9000_0000), + blob=replace(DEP4, tp=8, moe_ep=8).to_bytes(), + ), ) with pytest.raises(PeerIncompatibleError, match="pair"): target.check_peer( - self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000), target="") + self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000), blob=b"") ) with pytest.raises(PeerIncompatibleError, match="pool"): target.check_peer( @@ -200,7 +210,7 @@ def test_check_peer_accepts_a_covered_pair_and_refuses_everything_else(): ) -def test_describe_publishes_own_target_id(): +def test_describe_publishes_own_deployment(): target, _, _, _ = _pair() - assert target.describe() == DEP4.encode() - assert target.planner_id == "modeling_v2.deepseek_mla_latent" + assert Deployment.from_bytes(target.describe()) == DEP4 + assert target.planner_id == "trtllm.deepseek_mla_latent" diff --git a/tests/unittest/_torch/disaggregation/test_planner_binding.py b/tests/unittest/_torch/disaggregation/test_planner_binding.py new file mode 100644 index 000000000000..2d3dfe0e1cb2 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/test_planner_binding.py @@ -0,0 +1,121 @@ +# 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. +"""A planner covers deployments; a deployment is the architecture plus its parallel +topology, which a built-in model and a modeling_v2 target describe identically.""" + +from dataclasses import replace +from pathlib import Path + +import pytest +from transformers import PretrainedConfig + +import tensorrt_llm._torch._experimental.modeling_v2 as modeling_v2 +from tensorrt_llm._torch.disaggregation import planners +from tensorrt_llm._torch.disaggregation.planners.deepseek_mla_latent import ( + DeepseekMlaLatentTransferPlanner, +) +from tensorrt_llm._torch.disaggregation.planners.deployment import Deployment +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + +DEEPSEEK_DEP4 = Deployment( + architecture="DeepseekV3ForCausalLM", tp=4, pp=1, cp=1, moe_ep=4, moe_tp=1, attention_dp=True +) +GPT_OSS_TP1 = Deployment( + architecture="GptOssForCausalLM", tp=1, pp=1, cp=1, moe_ep=1, moe_tp=1, attention_dp=False +) +_DEP4_MAPPING = dict( + world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=True +) + + +def _model_config(architecture: str, **mapping_kwargs) -> ModelConfig: + return ModelConfig( + pretrained_config=PretrainedConfig(architectures=[architecture]), + mapping=Mapping(**mapping_kwargs), + ) + + +def test_deployment_reads_the_architecture_and_the_parallel_topology(): + assert Deployment.of(_model_config("DeepseekV3ForCausalLM", **_DEP4_MAPPING)) == DEEPSEEK_DEP4 + assert Deployment.of(_model_config("GptOssForCausalLM")) == GPT_OSS_TP1 + assert Deployment.of(ModelConfig(pretrained_config=PretrainedConfig())) is None + + +def test_deployment_survives_the_wire(): + assert Deployment.from_bytes(DEEPSEEK_DEP4.to_bytes()) == DEEPSEEK_DEP4 + assert Deployment.from_bytes(b"") is None + assert Deployment.from_bytes(b"not json") is None + + +def test_planner_is_found_by_the_deployment_it_covers_and_nothing_else(): + assert planners.planner_for(DEEPSEEK_DEP4) is DeepseekMlaLatentTransferPlanner + assert planners.planner_for(GPT_OSS_TP1) is None + # The same checkpoint on another topology is another deployment. + assert planners.planner_for(replace(DEEPSEEK_DEP4, tp=8, moe_ep=8)) is None + + +def test_no_target_source_names_a_planner(): + """The binding direction: planners list deployments, never the reverse.""" + root = Path(modeling_v2.__file__).resolve().parent / "models" + hits = [ + p.relative_to(root) + for p in root.rglob("*.py") + if "transfer_planner" in p.read_text() or "TransferPlanner" in p.read_text() + ] + assert hits == [], hits + + +def test_two_planners_covering_one_deployment_is_an_error(monkeypatch): + class Twin(DeepseekMlaLatentTransferPlanner): + pass + + monkeypatch.setattr(planners, "PLANNERS", (DeepseekMlaLatentTransferPlanner, Twin)) + with pytest.raises(ValueError, match="DeepseekV3ForCausalLM"): + planners.planner_for(DEEPSEEK_DEP4) + + +class _StubModel: + def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): + return "default" + + +def test_bind_leaves_an_uncovered_model_on_the_default_hook(): + stub = _StubModel() + assert planners.bind_transfer_planner(stub, _model_config("GptOssForCausalLM")) is None + assert stub.build_transfer_planner(None, None, None) == "default" + no_arch = ModelConfig(pretrained_config=PretrainedConfig()) + assert planners.bind_transfer_planner(stub, no_arch) is None + + +def test_bind_installs_a_factory_that_knows_its_own_deployment(monkeypatch): + """The class the model was built from is not an input: a built-in + DeepseekV3ForCausalLM and a modeling_v2 target on the same topology bind + to the same planner and publish the same description, which is what lets + a context on one implementation feed a generation on the other.""" + seen = {} + + def fake_init(self, self_rank_info, self_extractor, *, self_deployment): + seen["args"] = (self_rank_info, self_extractor, self_deployment) + + monkeypatch.setattr(DeepseekMlaLatentTransferPlanner, "__init__", fake_init) + stub = _StubModel() + config = _model_config("DeepseekV3ForCausalLM", **_DEP4_MAPPING) + assert planners.bind_transfer_planner(stub, config) is DeepseekMlaLatentTransferPlanner + planner = stub.build_transfer_planner("kvm", "ri", "ext") + assert isinstance(planner, DeepseekMlaLatentTransferPlanner) + assert seen["args"] == ("ri", "ext", DEEPSEEK_DEP4) diff --git a/tests/unittest/_torch/modeling_v2/transfers/test_transfers_binding.py b/tests/unittest/_torch/modeling_v2/transfers/test_transfers_binding.py deleted file mode 100644 index 3f2c6100ddfa..000000000000 --- a/tests/unittest/_torch/modeling_v2/transfers/test_transfers_binding.py +++ /dev/null @@ -1,105 +0,0 @@ -# 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. -"""A planner declares the targets it covers; a target knows nothing about planners.""" - -import pytest - -from tensorrt_llm._torch._experimental.modeling_v2 import transfers -from tensorrt_llm._torch._experimental.modeling_v2.transfers.deepseek_mla_latent import ( - DeepseekMlaLatentTransferPlanner, -) - -pytestmark = pytest.mark.cpu_only - -DEEPSEEK_DEP4 = "deepseek_v3/r1_0528_nvfp4__sm_103__dep4" -GPT_OSS_TP1 = "gpt_oss/gpt_oss_120b__sm_103__tp1" - - -def _target_cls(dotted_dir: str): - family, name = dotted_dir.split("/") - module = __import__( - f"tensorrt_llm._torch._experimental.modeling_v2.models.{family}.{name}.modeling", - fromlist=["_"], - ) - return next( - v - for k, v in vars(module).items() - if k.startswith("ModelingV2") - and isinstance(v, type) - and v.__module__ == module.__name__ - and "Target" not in k - ) - - -def test_target_id_is_family_slash_directory(): - assert transfers.target_id_of(_target_cls(DEEPSEEK_DEP4)) == DEEPSEEK_DEP4 - assert transfers.target_id_of(_target_cls(GPT_OSS_TP1)) == GPT_OSS_TP1 - assert transfers.target_id_of(dict) is None # not a modeling_v2 target - - -def test_planner_is_found_by_the_target_it_covers_and_nothing_else(): - assert transfers.planner_for_target(DEEPSEEK_DEP4) is DeepseekMlaLatentTransferPlanner - assert transfers.planner_for_target(GPT_OSS_TP1) is None - assert transfers.planner_for_target("nobody/such__sm_0__tp1") is None - - -def test_no_target_source_names_a_planner(): - """The binding direction: planners list targets, never the reverse.""" - from pathlib import Path - - root = Path(transfers.__file__).resolve().parents[1] / "models" - hits = [ - p.relative_to(root) - for p in root.rglob("*.py") - if "transfer_planner" in p.read_text() or "TransferPlanner" in p.read_text() - ] - assert hits == [], hits - - -def test_two_planners_covering_one_target_is_an_error(monkeypatch): - class Twin(DeepseekMlaLatentTransferPlanner): - pass - - monkeypatch.setattr(transfers, "PLANNERS", (DeepseekMlaLatentTransferPlanner, Twin)) - with pytest.raises(ValueError, match=DEEPSEEK_DEP4): - transfers.planner_for_target(DEEPSEEK_DEP4) - - -class _StubModel: - def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): - return "default" - - -def test_bind_leaves_an_uncovered_model_on_the_default_hook(): - stub = _StubModel() - assert transfers.bind_transfer_planner(stub, _target_cls(GPT_OSS_TP1)) is None - assert stub.build_transfer_planner(None, None, None) == "default" - assert transfers.bind_transfer_planner(stub, dict) is None - - -def test_bind_installs_a_factory_that_knows_its_own_target(monkeypatch): - seen = {} - - def fake_init(self, self_rank_info, self_extractor, *, self_target): - seen["args"] = (self_rank_info, self_extractor, self_target) - - monkeypatch.setattr(DeepseekMlaLatentTransferPlanner, "__init__", fake_init) - stub = _StubModel() - assert transfers.bind_transfer_planner(stub, _target_cls(DEEPSEEK_DEP4)) is ( - DeepseekMlaLatentTransferPlanner - ) - planner = stub.build_transfer_planner("kvm", "ri", "ext") - assert isinstance(planner, DeepseekMlaLatentTransferPlanner) - assert seen["args"] == ("ri", "ext", DEEPSEEK_DEP4) From d2df0400096911f342e72484ea7a0de62740dfd0 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:32:40 -0700 Subject: [PATCH 21/25] [None][refactor] Let each transfer planner judge its deployments with supports_deployment Replace the Deployment dataclass and the COVERS pair matrix with one class method per planner, supports_deployment(architecture, mapping) -> bool. The planner package asks every planner and binds the one that says yes. The matrix encoded the full topology (tp, moe_ep, moe_tp) into an identity that nothing consumed: plan() reads geometry from the page table and RankInfo already carries the peer's parallel sizes. It also made coverage an enumeration, so dep4 -> dep8 needed a new entry although the MLA latent KV is not sharded and the plan is the same for any width. The predicate states the real condition: a DeepSeek-V3 architecture under attention-DP with pp = cp = 1. The peer no longer re-validates the deployment in check_peer. The framework only connects ranks with the same planner_id, and the only way to get this planner_id is to have passed supports_deployment at bind time, so the peer has already been judged by the same function. describe() falls back to the base class's empty blob; check_peer keeps the latent geometry comparison (pool shape, layer count, tokens_per_block, slot bytes), which is what decides whether the bytes line up. Signed-off-by: Fred Wei Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../disaggregation/planners/__init__.py | 42 +++++----- .../planners/deepseek_mla_latent.py | 72 +++++++---------- .../disaggregation/planners/deployment.py | 75 ------------------ tensorrt_llm/_torch/models/modeling_auto.py | 2 +- .../test_deepseek_mla_latent_planner.py | 39 ++------- .../disaggregation/test_planner_binding.py | 79 ++++++++----------- 6 files changed, 87 insertions(+), 222 deletions(-) delete mode 100644 tensorrt_llm/_torch/disaggregation/planners/deployment.py diff --git a/tensorrt_llm/_torch/disaggregation/planners/__init__.py b/tensorrt_llm/_torch/disaggregation/planners/__init__.py index 51ecf0e4219d..77bde4b95ee4 100644 --- a/tensorrt_llm/_torch/disaggregation/planners/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/planners/__init__.py @@ -14,14 +14,13 @@ # limitations under the License. """Which disaggregated KV transfer planner serves a deployment. -A planner owns a compatibility set: the (context deployment, generation -deployment) pairs it can move bytes between, declared on the class as -``COVERS``. A :class:`Deployment` is the architecture plus the parallel -topology, so one entry covers the built-in model and any modeling_v2 target -alike. Models declare nothing. This module is the one place that knows every -planner, and it answers one question: given the deployment a model instance -was built for, which planner covers it? None means the engine keeps the -framework's default plan. +Each planner class judges deployments itself through +``supports_deployment(architecture, mapping)``: the model architecture and the +parallel topology are the whole input, so the built-in model and any +modeling_v2 target of the same architecture are judged alike. Models declare +nothing. This module is the one place that knows every planner, and it answers +one question: which planner supports the deployment a model instance was built +for? None means the engine keeps the framework's default plan. Binding happens where the model is constructed (``AutoModelForCausalLM``), by installing the covering planner's factory as the instance's @@ -35,44 +34,43 @@ from tensorrt_llm._torch.disaggregation.base.planner import TransferPlanner from .deepseek_mla_latent import DeepseekMlaLatentTransferPlanner -from .deployment import Deployment if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.mapping import Mapping #: Every planner a deployment may be served by. Order carries no meaning: a -#: deployment covered by two planners is a configuration error, not a priority +#: deployment two planners support is a configuration error, not a priority #: question. PLANNERS: tuple = (DeepseekMlaLatentTransferPlanner,) -def planner_for(deployment: Deployment) -> Optional[Type[TransferPlanner]]: - """The planner whose matrix mentions ``deployment`` on either side, or None.""" - found = [p for p in PLANNERS if any(deployment in pair for pair in p.COVERS)] +def planner_for(architecture: str, mapping: "Mapping") -> Optional[Type[TransferPlanner]]: + """The planner that supports ``architecture`` sharded as ``mapping`` says, or None.""" + found = [p for p in PLANNERS if p.supports_deployment(architecture, mapping)] if len(found) > 1: names = ", ".join(p.__name__ for p in found) - raise ValueError(f"{deployment} is covered by more than one planner: {names}") + raise ValueError(f"{architecture} is supported by more than one planner: {names}") return found[0] if found else None def bind_transfer_planner(model, model_config: "ModelConfig") -> Optional[Type[TransferPlanner]]: - """Install the covering planner's factory on ``model``; return the planner or None. + """Install the supporting planner's factory on ``model``; return the planner or None. - The factory has the ``build_transfer_planner`` signature the framework - calls, closing over the deployment so the planner can publish it. + The factory has the ``build_transfer_planner`` signature the framework calls. """ - deployment = Deployment.of(model_config) - if deployment is None: + architectures = getattr(model_config.pretrained_config, "architectures", None) + if not architectures: return None - planner_cls = planner_for(deployment) + planner_cls = planner_for(architectures[0], model_config.mapping) if planner_cls is None: return None def build_transfer_planner(kv_cache_manager, self_rank_info, self_extractor): - return planner_cls(self_rank_info, self_extractor, self_deployment=deployment) + return planner_cls(self_rank_info, self_extractor) model.build_transfer_planner = build_transfer_planner return planner_cls -__all__ = ["PLANNERS", "Deployment", "bind_transfer_planner", "planner_for"] +__all__ = ["PLANNERS", "bind_transfer_planner", "planner_for"] diff --git a/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py index d4425992f991..a5b820e7a197 100644 --- a/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py +++ b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py @@ -20,20 +20,22 @@ plain arithmetic: pair block ordinals both sides hold, address the slots, and copy. Blocks that are consecutive on both sides move as one fragment. -Which deployments may talk to each other is the ``COVERS`` matrix below: pairs -of (context deployment, generation deployment). A deployment is the -architecture plus its parallel topology, so the built-in DeepseekV3 model and a -modeling_v2 target on the same topology are the same deployment and may feed -each other. A model never names this planner; ``disaggregation.planners`` -finds the planner from the deployment, and ``describe`` publishes that -deployment so the peer's planner can look the pair up too. Anything not in the -matrix is refused at registration; the generic default planner stays untouched -for the models that need pool matching, head re-sharding or recurrent state. +``supports_deployment`` says which deployments this plan holds for: a +DeepSeek-V3 architecture under attention-DP with no pipeline or context +parallelism. The latent KV is never sharded across ranks, so the width of +either side does not matter and a dep4 context may feed a dep8 generation. +Which implementation computed the KV does not matter either: the built-in +``DeepseekV3ForCausalLM`` and a modeling_v2 target write the same bytes. The +framework only connects ranks that bound the same ``planner_id``, so a peer +has already passed the same check for itself; ``check_peer`` then compares the +latent geometry both sides published. The generic default planner stays +untouched for the models that need pool matching, head re-sharding or +recurrent state. """ from __future__ import annotations -from typing import FrozenSet, Optional, Tuple +from typing import TYPE_CHECKING, Optional import numpy as np @@ -54,12 +56,8 @@ ) from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool -from .deployment import Deployment - -# Attention-DP over 4 ranks, experts sharded 4 ways, attention replicated. -_DEEPSEEK_DEP4 = Deployment( - architecture="DeepseekV3ForCausalLM", tp=4, pp=1, cp=1, moe_ep=4, moe_tp=1, attention_dp=True -) +if TYPE_CHECKING: + from tensorrt_llm.mapping import Mapping def _sole_pool(page_table: Optional[KVCachePageTable], side: str) -> PhysicalPool: @@ -84,40 +82,24 @@ def _sole_pool(page_table: Optional[KVCachePageTable], side: str) -> PhysicalPoo class DeepseekMlaLatentTransferPlanner(TransferPlanner): planner_id = "trtllm.deepseek_mla_latent" - #: (context deployment, generation deployment) pairs this planner knows how - #: to move bytes between. The same pair is listed once; the planner serves - #: both roles, so a pair is read from the sender's side or the receiver's. - COVERS: FrozenSet[Tuple[Deployment, Deployment]] = frozenset({(_DEEPSEEK_DEP4, _DEEPSEEK_DEP4)}) + @classmethod + def supports_deployment(cls, architecture: str, mapping: "Mapping") -> bool: + """Whether this plan holds for ``architecture`` sharded as ``mapping`` says.""" + return ( + architecture == "DeepseekV3ForCausalLM" + and mapping.pp_size == 1 + and mapping.cp_size == 1 + and bool(mapping.enable_attention_dp) + ) - def __init__( - self, - self_rank_info: RankInfo, - self_extractor: KVRegionExtractorV1, - *, - self_deployment: Deployment, - ): + def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1): self._ri = self_rank_info - self._deployment = self_deployment self._page_table = self_extractor.page_table self._pool = _sole_pool(self._page_table, "this rank") - def describe(self) -> bytes: - return self._deployment.to_bytes() - - def _pair(self, peer_ri: RankInfo) -> Tuple[Deployment, Deployment]: - """The (self, peer) deployment pair, or a refusal if the matrix lacks it.""" - peer = Deployment.from_bytes(peer_ri.planner_blob) - if (self._deployment, peer) in self.COVERS or (peer, self._deployment) in self.COVERS: - return self._deployment, peer - raise PeerIncompatibleError( - f"no transfer plan for the pair ({self._deployment}, {peer}); " - f"covered pairs: {[tuple(map(str, pair)) for pair in self.COVERS]}" - ) - # ---- TransferPlanner protocol ------------------------------------------------- def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: - self._pair(peer_ri) peer_pool = _sole_pool(peer_ri.page_table, "peer") self_layers = len(self._page_table.layer_groups[0].local_layers) peer_layers = len(peer_ri.page_table.layer_groups[0].local_layers) @@ -137,9 +119,9 @@ def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: ) def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> PeerTargets: - # Every covered pair is attention-DP with pp = cp = 1 and tp == dp, so - # the peer's flat rank is its dp_rank, and one writer per destination - # means the fan-in split is trivially uniform. + # Every supported deployment is attention-DP with pp = cp = 1 and + # tp == dp, so the peer's flat rank is its dp_rank, and one writer per + # destination means the fan-in split is trivially uniform. return PeerTargets(ranks=[peer_dp_rank], uniform_writers=True) def plan( diff --git a/tensorrt_llm/_torch/disaggregation/planners/deployment.py b/tensorrt_llm/_torch/disaggregation/planners/deployment.py deleted file mode 100644 index af172ea0a41a..000000000000 --- a/tensorrt_llm/_torch/disaggregation/planners/deployment.py +++ /dev/null @@ -1,75 +0,0 @@ -# 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 a transfer planner identifies a deployment by. - -A KV transfer plan depends on the model architecture and on how each side is -sharded, not on which implementation computed the KV: a built-in -``DeepseekV3ForCausalLM`` and a modeling_v2 target for the same checkpoint on -the same topology write the same bytes into the same pools. ``Deployment`` -therefore names exactly those two things, so a planner's compatibility matrix -covers both implementations at once and a context on one can feed a -generation on the other. -""" - -from __future__ import annotations - -import json -from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Optional - -if TYPE_CHECKING: - from tensorrt_llm._torch.model_config import ModelConfig - - -@dataclass(frozen=True) -class Deployment: - """The architecture of a model and the parallel topology it runs under.""" - - architecture: str - tp: int - pp: int - cp: int - moe_ep: int - moe_tp: int - attention_dp: bool - - @classmethod - def of(cls, model_config: "ModelConfig") -> Optional["Deployment"]: - """The deployment ``model_config`` describes, or None when it names no architecture.""" - architectures = getattr(model_config.pretrained_config, "architectures", None) - if not architectures: - return None - m = model_config.mapping - return cls( - architecture=architectures[0], - tp=m.tp_size, - pp=m.pp_size, - cp=m.cp_size, - moe_ep=m.moe_ep_size, - moe_tp=m.moe_tp_size, - attention_dp=bool(m.enable_attention_dp), - ) - - def to_bytes(self) -> bytes: - """Wire form for ``RankInfo.planner_blob``.""" - return json.dumps(asdict(self), sort_keys=True).encode() - - @classmethod - def from_bytes(cls, blob: bytes) -> Optional["Deployment"]: - """Inverse of :meth:`to_bytes`; None for anything that is not one.""" - try: - return cls(**json.loads(blob.decode())) - except (ValueError, TypeError, UnicodeDecodeError): - return None diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index d7f5fb5ae909..e606c963502c 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -77,7 +77,7 @@ def from_config( with model_extra_attrs(extra_attrs): model = cls(config) model.extra_attrs = extra_attrs - # A KV transfer planner may cover this deployment (architecture and + # A KV transfer planner may support this deployment (architecture and # parallel topology; see disaggregation.planners). Resolved here, where # the config is known and before the transceiver asks for the hook, so # neither the model nor the disaggregation code names the other. diff --git a/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py b/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py index fd34378c196f..7642eee85836 100644 --- a/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py +++ b/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py @@ -12,9 +12,7 @@ # 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 DeepSeek MLA-latent planner must produce the default planner's bytes for the pairs it covers.""" - -from dataclasses import replace +"""The DeepSeek MLA-latent planner must produce the default planner's bytes for the deployments it supports.""" import numpy as np import pytest @@ -27,7 +25,6 @@ from tensorrt_llm._torch.disaggregation.planners.deepseek_mla_latent import ( DeepseekMlaLatentTransferPlanner, ) -from tensorrt_llm._torch.disaggregation.planners.deployment import Deployment from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import ( BUFFER_ENTRY_DTYPE, @@ -45,9 +42,6 @@ LAYER_BYTES = 512 # one latent-K region per layer per block SLOT_BYTES = NUM_LAYERS * LAYER_BYTES TPB = 16 -DEP4 = Deployment( - architecture="DeepseekV3ForCausalLM", tp=4, pp=1, cp=1, moe_ep=4, moe_tp=1, attention_dp=True -) def _mla_page_table(base: int, *, num_layers: int = NUM_LAYERS, extra_pool: bool = False): @@ -74,12 +68,9 @@ def _mla_page_table(base: int, *, num_layers: int = NUM_LAYERS, extra_pool: bool ) -def _dep4_rank_info( - name: str, dp_rank: int, page_table, *, dp_size: int = 4, pp_size: int = 1, blob: bytes = None -): +def _dep4_rank_info(name: str, dp_rank: int, page_table, *, dp_size: int = 4, pp_size: int = 1): return RankInfo( planner_id=DeepseekMlaLatentTransferPlanner.planner_id, - planner_blob=DEP4.to_bytes() if blob is None else blob, instance_name=name, instance_rank=dp_rank, tp_size=dp_size, @@ -110,7 +101,7 @@ def _pair(dp_rank: int = 2): peer_ri = _dep4_rank_info("gen", dp_rank, peer_pt) ext = KVRegionExtractorV1(self_pt) return ( - DeepseekMlaLatentTransferPlanner(self_ri, ext, self_deployment=DEP4), + DeepseekMlaLatentTransferPlanner(self_ri, ext), DefaultTransferPlanner(self_ri, ext), self_ri, peer_ri, @@ -183,23 +174,9 @@ def test_targets_is_the_matching_dp_rank_only(): assert ours.uniform_writers is True -def test_check_peer_accepts_a_covered_pair_and_refuses_everything_else(): +def test_check_peer_refuses_a_peer_whose_latent_geometry_differs(): target, _, self_ri, peer_ri = _pair() - target.check_peer(self_ri, peer_ri) # (dep4, dep4) is in the matrix - with pytest.raises(PeerIncompatibleError, match="pair"): - target.check_peer( - self_ri, - _dep4_rank_info( - "gen", - 2, - _mla_page_table(0x9000_0000), - blob=replace(DEP4, tp=8, moe_ep=8).to_bytes(), - ), - ) - with pytest.raises(PeerIncompatibleError, match="pair"): - target.check_peer( - self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000), blob=b"") - ) + target.check_peer(self_ri, peer_ri) with pytest.raises(PeerIncompatibleError, match="pool"): target.check_peer( self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000, extra_pool=True)) @@ -208,9 +185,3 @@ def test_check_peer_accepts_a_covered_pair_and_refuses_everything_else(): target.check_peer( self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000, num_layers=5)) ) - - -def test_describe_publishes_own_deployment(): - target, _, _, _ = _pair() - assert Deployment.from_bytes(target.describe()) == DEP4 - assert target.planner_id == "trtllm.deepseek_mla_latent" diff --git a/tests/unittest/_torch/disaggregation/test_planner_binding.py b/tests/unittest/_torch/disaggregation/test_planner_binding.py index 2d3dfe0e1cb2..6e21e9ae5eca 100644 --- a/tests/unittest/_torch/disaggregation/test_planner_binding.py +++ b/tests/unittest/_torch/disaggregation/test_planner_binding.py @@ -12,10 +12,10 @@ # 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. -"""A planner covers deployments; a deployment is the architecture plus its parallel -topology, which a built-in model and a modeling_v2 target describe identically.""" +"""A planner says which deployments it supports; binding asks every planner and +installs the one that says yes. The architecture and the parallel topology are +the only inputs, so a built-in model and a modeling_v2 target bind alike.""" -from dataclasses import replace from pathlib import Path import pytest @@ -26,51 +26,42 @@ from tensorrt_llm._torch.disaggregation.planners.deepseek_mla_latent import ( DeepseekMlaLatentTransferPlanner, ) -from tensorrt_llm._torch.disaggregation.planners.deployment import Deployment from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.mapping import Mapping pytestmark = pytest.mark.cpu_only -DEEPSEEK_DEP4 = Deployment( - architecture="DeepseekV3ForCausalLM", tp=4, pp=1, cp=1, moe_ep=4, moe_tp=1, attention_dp=True -) -GPT_OSS_TP1 = Deployment( - architecture="GptOssForCausalLM", tp=1, pp=1, cp=1, moe_ep=1, moe_tp=1, attention_dp=False -) -_DEP4_MAPPING = dict( - world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=True +DEEPSEEK = "DeepseekV3ForCausalLM" +DEP4 = Mapping(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=True) +DEP8 = Mapping(world_size=8, tp_size=8, moe_ep_size=8, moe_tp_size=1, enable_attention_dp=True) +TP4 = Mapping(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=False) +DEP4_PP2 = Mapping( + world_size=8, tp_size=4, pp_size=2, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=True ) -def _model_config(architecture: str, **mapping_kwargs) -> ModelConfig: +def _model_config(architecture: str, mapping: Mapping) -> ModelConfig: return ModelConfig( - pretrained_config=PretrainedConfig(architectures=[architecture]), - mapping=Mapping(**mapping_kwargs), + pretrained_config=PretrainedConfig(architectures=[architecture]), mapping=mapping ) -def test_deployment_reads_the_architecture_and_the_parallel_topology(): - assert Deployment.of(_model_config("DeepseekV3ForCausalLM", **_DEP4_MAPPING)) == DEEPSEEK_DEP4 - assert Deployment.of(_model_config("GptOssForCausalLM")) == GPT_OSS_TP1 - assert Deployment.of(ModelConfig(pretrained_config=PretrainedConfig())) is None - +def test_mla_latent_planner_supports_attention_dp_deepseek_of_any_width(): + supports = DeepseekMlaLatentTransferPlanner.supports_deployment + assert supports(DEEPSEEK, DEP4) + assert supports(DEEPSEEK, DEP8) + assert not supports("GptOssForCausalLM", DEP4) + assert not supports(DEEPSEEK, TP4) + assert not supports(DEEPSEEK, DEP4_PP2) -def test_deployment_survives_the_wire(): - assert Deployment.from_bytes(DEEPSEEK_DEP4.to_bytes()) == DEEPSEEK_DEP4 - assert Deployment.from_bytes(b"") is None - assert Deployment.from_bytes(b"not json") is None - -def test_planner_is_found_by_the_deployment_it_covers_and_nothing_else(): - assert planners.planner_for(DEEPSEEK_DEP4) is DeepseekMlaLatentTransferPlanner - assert planners.planner_for(GPT_OSS_TP1) is None - # The same checkpoint on another topology is another deployment. - assert planners.planner_for(replace(DEEPSEEK_DEP4, tp=8, moe_ep=8)) is None +def test_planner_is_found_by_the_deployment_it_supports_and_nothing_else(): + assert planners.planner_for(DEEPSEEK, DEP4) is DeepseekMlaLatentTransferPlanner + assert planners.planner_for("GptOssForCausalLM", Mapping()) is None def test_no_target_source_names_a_planner(): - """The binding direction: planners list deployments, never the reverse.""" + """The binding direction: planners judge deployments, never the reverse.""" root = Path(modeling_v2.__file__).resolve().parent / "models" hits = [ p.relative_to(root) @@ -80,13 +71,13 @@ def test_no_target_source_names_a_planner(): assert hits == [], hits -def test_two_planners_covering_one_deployment_is_an_error(monkeypatch): +def test_two_planners_supporting_one_deployment_is_an_error(monkeypatch): class Twin(DeepseekMlaLatentTransferPlanner): pass monkeypatch.setattr(planners, "PLANNERS", (DeepseekMlaLatentTransferPlanner, Twin)) - with pytest.raises(ValueError, match="DeepseekV3ForCausalLM"): - planners.planner_for(DEEPSEEK_DEP4) + with pytest.raises(ValueError, match=DEEPSEEK): + planners.planner_for(DEEPSEEK, DEP4) class _StubModel: @@ -94,28 +85,26 @@ def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extracto return "default" -def test_bind_leaves_an_uncovered_model_on_the_default_hook(): +def test_bind_leaves_an_unsupported_model_on_the_default_hook(): stub = _StubModel() - assert planners.bind_transfer_planner(stub, _model_config("GptOssForCausalLM")) is None + assert ( + planners.bind_transfer_planner(stub, _model_config("GptOssForCausalLM", Mapping())) is None + ) assert stub.build_transfer_planner(None, None, None) == "default" no_arch = ModelConfig(pretrained_config=PretrainedConfig()) assert planners.bind_transfer_planner(stub, no_arch) is None -def test_bind_installs_a_factory_that_knows_its_own_deployment(monkeypatch): - """The class the model was built from is not an input: a built-in - DeepseekV3ForCausalLM and a modeling_v2 target on the same topology bind - to the same planner and publish the same description, which is what lets - a context on one implementation feed a generation on the other.""" +def test_bind_installs_the_supporting_planner_as_the_factory(monkeypatch): seen = {} - def fake_init(self, self_rank_info, self_extractor, *, self_deployment): - seen["args"] = (self_rank_info, self_extractor, self_deployment) + def fake_init(self, self_rank_info, self_extractor): + seen["args"] = (self_rank_info, self_extractor) monkeypatch.setattr(DeepseekMlaLatentTransferPlanner, "__init__", fake_init) stub = _StubModel() - config = _model_config("DeepseekV3ForCausalLM", **_DEP4_MAPPING) + config = _model_config(DEEPSEEK, DEP4) assert planners.bind_transfer_planner(stub, config) is DeepseekMlaLatentTransferPlanner planner = stub.build_transfer_planner("kvm", "ri", "ext") assert isinstance(planner, DeepseekMlaLatentTransferPlanner) - assert seen["args"] == ("ri", "ext", DEEPSEEK_DEP4) + assert seen["args"] == ("ri", "ext") From e0e8792f7832e150fb62dbfec270dfd8cc241e3f Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 12 Oct 2026 20:00:32 -0700 Subject: [PATCH 22/25] [None][feat] Select the KV transfer planner by name through TRTLLM_KV_TRANSFER_PLANNER The planner is now chosen by the user, not by the framework. Setting TRTLLM_KV_TRANSFER_PLANNER to a planner_id binds that planner to the model; unset, every model keeps the default plan. PLANNERS becomes the map from planner_id to class, and the auto-selection step (planner_for, the "two planners support one deployment" error) goes away with it. supports_deployment is now a check rather than a selection criterion: a named planner that does not support the architecture and topology the model is built for fails binding with the reason, as does an unknown name. Both sides of a disaggregated pair set the same name; the framework already refuses to connect ranks whose planner_id differs. Signed-off-by: Fred Wei Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../disaggregation/planners/__init__.py | 70 ++++++++++-------- .../planners/deepseek_mla_latent.py | 3 +- tensorrt_llm/_torch/models/modeling_auto.py | 8 +-- .../disaggregation/test_planner_binding.py | 71 +++++++++++-------- 4 files changed, 87 insertions(+), 65 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/planners/__init__.py b/tensorrt_llm/_torch/disaggregation/planners/__init__.py index 77bde4b95ee4..22d956a98f2d 100644 --- a/tensorrt_llm/_torch/disaggregation/planners/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/planners/__init__.py @@ -12,24 +12,25 @@ # 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. -"""Which disaggregated KV transfer planner serves a deployment. +"""The KV transfer planner the user asks for, bound to the model it will serve. -Each planner class judges deployments itself through -``supports_deployment(architecture, mapping)``: the model architecture and the -parallel topology are the whole input, so the built-in model and any -modeling_v2 target of the same architecture are judged alike. Models declare -nothing. This module is the one place that knows every planner, and it answers -one question: which planner supports the deployment a model instance was built -for? None means the engine keeps the framework's default plan. +The user names a planner by setting ``TRTLLM_KV_TRANSFER_PLANNER`` to its +``planner_id``; ``PLANNERS`` maps every id to its class. Unset, every model +keeps the framework's default plan. The named planner must support the +deployment, the model architecture under the parallel topology it is built +for, or binding fails with a clear reason. Both sides of a disaggregated pair +must name the same planner; the framework refuses to connect ranks whose +``planner_id`` differs. Binding happens where the model is constructed (``AutoModelForCausalLM``), -by installing the covering planner's factory as the instance's +by installing the planner's factory as the instance's ``build_transfer_planner``; the framework reads that hook and nothing else. """ from __future__ import annotations -from typing import TYPE_CHECKING, Optional, Type +import os +from typing import TYPE_CHECKING, Dict, Optional, Type from tensorrt_llm._torch.disaggregation.base.planner import TransferPlanner @@ -37,34 +38,43 @@ if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm.mapping import Mapping -#: Every planner a deployment may be served by. Order carries no meaning: a -#: deployment two planners support is a configuration error, not a priority -#: question. -PLANNERS: tuple = (DeepseekMlaLatentTransferPlanner,) +PLANNER_ENV_VAR = "TRTLLM_KV_TRANSFER_PLANNER" - -def planner_for(architecture: str, mapping: "Mapping") -> Optional[Type[TransferPlanner]]: - """The planner that supports ``architecture`` sharded as ``mapping`` says, or None.""" - found = [p for p in PLANNERS if p.supports_deployment(architecture, mapping)] - if len(found) > 1: - names = ", ".join(p.__name__ for p in found) - raise ValueError(f"{architecture} is supported by more than one planner: {names}") - return found[0] if found else None +#: Every planner a user may name, keyed by its ``planner_id``. +PLANNERS: Dict[str, Type[TransferPlanner]] = { + p.planner_id: p for p in (DeepseekMlaLatentTransferPlanner,) +} def bind_transfer_planner(model, model_config: "ModelConfig") -> Optional[Type[TransferPlanner]]: - """Install the supporting planner's factory on ``model``; return the planner or None. + """Install the planner named by the environment on ``model``; return it, or None if none is named. - The factory has the ``build_transfer_planner`` signature the framework calls. + Raises ``ValueError`` when the name is unknown or the planner does not + support this deployment. """ - architectures = getattr(model_config.pretrained_config, "architectures", None) - if not architectures: + key = os.environ.get(PLANNER_ENV_VAR, "") + if not key: return None - planner_cls = planner_for(architectures[0], model_config.mapping) + planner_cls = PLANNERS.get(key) if planner_cls is None: - return None + raise ValueError( + f"{PLANNER_ENV_VAR}={key!r} names no KV transfer planner; " + f"known planners: {sorted(PLANNERS)}" + ) + architectures = getattr(model_config.pretrained_config, "architectures", None) + if not architectures: + raise ValueError( + f"{PLANNER_ENV_VAR}={key!r} needs a model architecture, " + "but the pretrained config names none" + ) + mapping = model_config.mapping + if not planner_cls.supports_deployment(architectures[0], mapping): + raise ValueError( + f"{planner_cls.__name__} does not support {architectures[0]} with " + f"tp={mapping.tp_size} pp={mapping.pp_size} cp={mapping.cp_size} " + f"attention_dp={bool(mapping.enable_attention_dp)}" + ) def build_transfer_planner(kv_cache_manager, self_rank_info, self_extractor): return planner_cls(self_rank_info, self_extractor) @@ -73,4 +83,4 @@ def build_transfer_planner(kv_cache_manager, self_rank_info, self_extractor): return planner_cls -__all__ = ["PLANNERS", "bind_transfer_planner", "planner_for"] +__all__ = ["PLANNER_ENV_VAR", "PLANNERS", "bind_transfer_planner"] diff --git a/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py index a5b820e7a197..959f484bbb22 100644 --- a/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py +++ b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py @@ -20,7 +20,8 @@ plain arithmetic: pair block ordinals both sides hold, address the slots, and copy. Blocks that are consecutive on both sides move as one fragment. -``supports_deployment`` says which deployments this plan holds for: a +The user selects this planner by name (see ``disaggregation.planners``); +``supports_deployment`` then checks that the plan holds for the deployment: a DeepSeek-V3 architecture under attention-DP with no pipeline or context parallelism. The latent KV is never sharded across ranks, so the width of either side does not matter and a dep4 context may feed a dep8 generation. diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index e606c963502c..cf719704dba7 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -77,10 +77,10 @@ def from_config( with model_extra_attrs(extra_attrs): model = cls(config) model.extra_attrs = extra_attrs - # A KV transfer planner may support this deployment (architecture and - # parallel topology; see disaggregation.planners). Resolved here, where - # the config is known and before the transceiver asks for the hook, so - # neither the model nor the disaggregation code names the other. + # The user may name a KV transfer planner for this deployment (see + # disaggregation.planners). Bound here, where the config is known and + # before the transceiver asks for the hook, so neither the model nor + # the disaggregation code names the other. from ..disaggregation.planners import bind_transfer_planner bind_transfer_planner(model, config) return model diff --git a/tests/unittest/_torch/disaggregation/test_planner_binding.py b/tests/unittest/_torch/disaggregation/test_planner_binding.py index 6e21e9ae5eca..caa54411d442 100644 --- a/tests/unittest/_torch/disaggregation/test_planner_binding.py +++ b/tests/unittest/_torch/disaggregation/test_planner_binding.py @@ -12,9 +12,8 @@ # 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. -"""A planner says which deployments it supports; binding asks every planner and -installs the one that says yes. The architecture and the parallel topology are -the only inputs, so a built-in model and a modeling_v2 target bind alike.""" +"""The user names a KV transfer planner through ``TRTLLM_KV_TRANSFER_PLANNER``; the +named planner checks that it supports the deployment and is bound to the model.""" from pathlib import Path @@ -23,6 +22,7 @@ import tensorrt_llm._torch._experimental.modeling_v2 as modeling_v2 from tensorrt_llm._torch.disaggregation import planners +from tensorrt_llm._torch.disaggregation.planners import PLANNER_ENV_VAR from tensorrt_llm._torch.disaggregation.planners.deepseek_mla_latent import ( DeepseekMlaLatentTransferPlanner, ) @@ -32,6 +32,7 @@ pytestmark = pytest.mark.cpu_only DEEPSEEK = "DeepseekV3ForCausalLM" +MLA_LATENT = DeepseekMlaLatentTransferPlanner.planner_id DEP4 = Mapping(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=True) DEP8 = Mapping(world_size=8, tp_size=8, moe_ep_size=8, moe_tp_size=1, enable_attention_dp=True) TP4 = Mapping(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=False) @@ -46,6 +47,11 @@ def _model_config(architecture: str, mapping: Mapping) -> ModelConfig: ) +class _StubModel: + def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): + return "default" + + def test_mla_latent_planner_supports_attention_dp_deepseek_of_any_width(): supports = DeepseekMlaLatentTransferPlanner.supports_deployment assert supports(DEEPSEEK, DEP4) @@ -55,13 +61,13 @@ def test_mla_latent_planner_supports_attention_dp_deepseek_of_any_width(): assert not supports(DEEPSEEK, DEP4_PP2) -def test_planner_is_found_by_the_deployment_it_supports_and_nothing_else(): - assert planners.planner_for(DEEPSEEK, DEP4) is DeepseekMlaLatentTransferPlanner - assert planners.planner_for("GptOssForCausalLM", Mapping()) is None +def test_every_planner_is_keyed_by_its_planner_id(): + assert planners.PLANNERS == {MLA_LATENT: DeepseekMlaLatentTransferPlanner} + assert PLANNER_ENV_VAR == "TRTLLM_KV_TRANSFER_PLANNER" def test_no_target_source_names_a_planner(): - """The binding direction: planners judge deployments, never the reverse.""" + """Planners are chosen by the user, never by a model's source.""" root = Path(modeling_v2.__file__).resolve().parent / "models" hits = [ p.relative_to(root) @@ -71,40 +77,45 @@ def test_no_target_source_names_a_planner(): assert hits == [], hits -def test_two_planners_supporting_one_deployment_is_an_error(monkeypatch): - class Twin(DeepseekMlaLatentTransferPlanner): - pass - - monkeypatch.setattr(planners, "PLANNERS", (DeepseekMlaLatentTransferPlanner, Twin)) - with pytest.raises(ValueError, match=DEEPSEEK): - planners.planner_for(DEEPSEEK, DEP4) - - -class _StubModel: - def build_transfer_planner(self, kv_cache_manager, self_rank_info, self_extractor): - return "default" - - -def test_bind_leaves_an_unsupported_model_on_the_default_hook(): +def test_without_the_variable_every_model_keeps_the_default_hook(monkeypatch): + monkeypatch.delenv(PLANNER_ENV_VAR, raising=False) stub = _StubModel() - assert ( - planners.bind_transfer_planner(stub, _model_config("GptOssForCausalLM", Mapping())) is None - ) + assert planners.bind_transfer_planner(stub, _model_config(DEEPSEEK, DEP4)) is None assert stub.build_transfer_planner(None, None, None) == "default" - no_arch = ModelConfig(pretrained_config=PretrainedConfig()) - assert planners.bind_transfer_planner(stub, no_arch) is None + monkeypatch.setenv(PLANNER_ENV_VAR, "") + assert planners.bind_transfer_planner(stub, _model_config(DEEPSEEK, DEP4)) is None -def test_bind_installs_the_supporting_planner_as_the_factory(monkeypatch): +def test_the_named_planner_is_installed_as_the_factory(monkeypatch): seen = {} def fake_init(self, self_rank_info, self_extractor): seen["args"] = (self_rank_info, self_extractor) monkeypatch.setattr(DeepseekMlaLatentTransferPlanner, "__init__", fake_init) + monkeypatch.setenv(PLANNER_ENV_VAR, MLA_LATENT) stub = _StubModel() - config = _model_config(DEEPSEEK, DEP4) - assert planners.bind_transfer_planner(stub, config) is DeepseekMlaLatentTransferPlanner + assert ( + planners.bind_transfer_planner(stub, _model_config(DEEPSEEK, DEP4)) + is DeepseekMlaLatentTransferPlanner + ) planner = stub.build_transfer_planner("kvm", "ri", "ext") assert isinstance(planner, DeepseekMlaLatentTransferPlanner) assert seen["args"] == ("ri", "ext") + + +def test_an_unknown_key_is_refused_with_the_known_keys(monkeypatch): + monkeypatch.setenv(PLANNER_ENV_VAR, "trtllm.no_such_planner") + with pytest.raises(ValueError, match=MLA_LATENT): + planners.bind_transfer_planner(_StubModel(), _model_config(DEEPSEEK, DEP4)) + + +def test_a_planner_that_does_not_support_the_deployment_is_refused(monkeypatch): + monkeypatch.setenv(PLANNER_ENV_VAR, MLA_LATENT) + with pytest.raises(ValueError, match="GptOssForCausalLM"): + planners.bind_transfer_planner(_StubModel(), _model_config("GptOssForCausalLM", DEP4)) + with pytest.raises(ValueError, match=DEEPSEEK): + planners.bind_transfer_planner(_StubModel(), _model_config(DEEPSEEK, TP4)) + no_arch = ModelConfig(pretrained_config=PretrainedConfig()) + with pytest.raises(ValueError, match="architecture"): + planners.bind_transfer_planner(_StubModel(), no_arch) From 9a471672a12d8d5e0aed93df01287c9c91ae74db Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 12 Oct 2026 20:07:47 -0700 Subject: [PATCH 23/25] [None][doc] Document TRTLLM_KV_TRANSFER_PLANNER and bring the planner design notes up to date Signed-off-by: Fred Wei Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- docs/source/features/disagg-serving.md | 2 + ...10-10-deepseek-r1-dep4-transfer-planner.md | 50 ++++++++++--------- ...26-10-10-disagg-transfer-planner-design.md | 38 +++++++------- 3 files changed, 48 insertions(+), 42 deletions(-) diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index cc583ac31d07..afe65fe1d200 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -380,6 +380,8 @@ TRT-LLM uses some environment variables to control the behavior of disaggregated * `TRTLLM_NIXL_KVCACHE_BACKEND`: Selects the transport NIXL itself uses. Valid values are `UCX` (default) and `LIBFABRIC`; an unsupported value logs a warning and falls back to `UCX`. `LIBFABRIC` additionally requires a NIXL build carrying the libfabric plugin — see the [disaggregated serving examples](source:examples/disaggregated/README.md). +* `TRTLLM_KV_TRANSFER_PLANNER`: Names a model-specific KV transfer planner by its `planner_id`, for example `trtllm.deepseek_mla_latent` for DeepSeek-V3 architectures under attention DP with no pipeline or context parallelism. Unset (the default), every model uses the generic transfer plan. The named planner checks at start-up that it supports the model architecture and parallel topology and fails with the reason otherwise. Context and generation servers must set the same value; the transceiver refuses to connect workers whose planners differ. Requires `transceiver_runtime: PYTHON`. + * `TRTLLM_GPU_KEEPALIVE`: If set to `1`, a generation worker that is waiting at the benchmark fill gate (`TLLM_BENCHMARK_REQ_QUEUES_SIZE`) keeps a resident warp on every SM in ~100 ms chunks instead of idling through the wait, so GPU-activity metrics do not read idle while the context tier fills it. The work is drained when the gate opens and never overlaps a forward pass. The default value is `0`. There are some other useful environment variables that may help when encountering failures or performance issues. diff --git a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md index 88d36dc11451..64a1df731df1 100644 --- a/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md +++ b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md @@ -82,27 +82,25 @@ transceiver logs), GSM8K within tolerance of 94.24, no `PeerIncompatibleError`. The target-overrides-`build_transfer_planner` binding was rejected in review: one planner must be able to serve several targets, including pairs such as ctx dep4 -> gen dep8 and -targets from different cores. Implemented instead: - -- Planner moved to `modeling_v2/transfers/deepseek_mla_latent.py` as - `DeepseekMlaLatentTransferPlanner`, `planner_id = "modeling_v2.deepseek_mla_latent"`, with a - `COVERS` matrix of `(ctx target id, gen target id)` pairs; today only - `(deepseek_v3/r1_0528_nvfp4__sm_103__dep4, same)`. `describe()` publishes the own target id and - `check_peer` refuses any pair outside the matrix; the dp/tp/pp/cp checks are gone because the - target id already pins the topology, the layout checks (one pool, layer count, - `tokens_per_block`, `slot_bytes`) stay because the arithmetic depends on them. -- `modeling_v2/transfers/__init__.py`: `PLANNERS`, `planner_for_target`, `bind_transfer_planner`; - `target_id_of` lives in `_router_index` so `AutoModelForCausalLM` can ask it cheaply. -- `AutoModelForCausalLM.from_config` binds after construction; the target's `modeling.py` is back - to zero transfer-related lines. -- Tests: `tests/unittest/_torch/modeling_v2/transfers/` (binding direction, lookup, duplicate - coverage, bind; planner parity/merge/targets/pair refusal/describe). 74 tests green with the - claims, routing and disaggregation planner suites. -- Two-node re-gate with the new binding (ctx dep4 on nvl72d185-T13, gen dep4 on nvl72d185-T16, - same launcher): all four ranks log - `KV transfer planned by DeepseekMlaLatentTransferPlanner (planner_id='modeling_v2.deepseek_mla_latent')`, - no `PeerIncompatibleError`; GSM8K flexible-extract **95.07** against the 94.24 reference, PASSED. - Logs: `mv2-disagg/twonode/logs/run_20261011T084450Z`. +targets from different cores. A `COVERS` matrix of target-id pairs, then of +architecture-plus-topology pairs, followed and was dropped in turn: it encoded a topology +nothing consumed and made coverage an enumeration. Implemented instead: + +- Planner at `disaggregation/planners/deepseek_mla_latent.py` as + `DeepseekMlaLatentTransferPlanner`, `planner_id = "trtllm.deepseek_mla_latent"`, with + `supports_deployment(architecture, mapping)`: `DeepseekV3ForCausalLM`, attention DP, + pp = cp = 1, any width. `check_peer` keeps the layout checks (one pool, layer count, + `tokens_per_block`, `slot_bytes`) because the arithmetic depends on them; `describe()` is + the base class's empty blob. +- `disaggregation/planners/__init__.py`: `PLANNER_ENV_VAR = "TRTLLM_KV_TRANSFER_PLANNER"`, + `PLANNERS` (planner_id -> class), `bind_transfer_planner`. The user names the planner; the + framework never picks one. +- `AutoModelForCausalLM.from_config` binds after construction for every model; the target's + `modeling.py` has zero transfer-related lines. +- Tests: `tests/unittest/_torch/disaggregation/test_planner_binding.py` (supports_deployment, + PLANNERS keys, binding direction, unset/empty variable, named planner installed, unknown name, + unsupported deployment) and `test_deepseek_mla_latent_planner.py` (parity, merge, targets, + geometry refusal). ## Interface hardening (after review) @@ -121,9 +119,13 @@ descriptions whole. `RecvReqInfo` moved from `native/transfer.py` to contract depends on. The sender's `KVSendTask` is deliberately not passed: it carries transfer progress, counters and timers. -Still to do at rebase time, not here: honor `chunk.excluded_pool_views` in -`DefaultTransferPlanner.plan` and fold `validate_bounded_replay_compatible` into -`check_peer`; both land on the planner, not on `transfer.py`. +Done at the rebase onto main: `DefaultTransferPlanner.plan` skips the +`chunk.excluded_pool_views` main added (the DeepSeek planner returns an empty plan +when its sole pool is excluded), and the `bounded_replay_on_generation` peer check +main added to `PeerRegistrar` lives in `DefaultTransferPlanner.check_peer`; both +sit on the planner, not on `transfer.py`. The disaggregated gates pass +`modeling_v2: require` through each server's config, since main made the switch +an `LLM` argument. Verified: 292 unit tests green (planner, chunked transfer, ownership regressions, unsubmitted session, peer, modeling_v2 transfers/claims/routing/target diff --git a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md index 2bc76d2ff3d1..329163b07cc1 100644 --- a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md +++ b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md @@ -164,25 +164,27 @@ class PretrainedModel/DecoderModelForCausalLM: ``` The default is the method above. A modeling_v2 target does **not** override -it: a planner is the owner of a compatibility set, the (context target, -generation target) pairs it can move bytes between, and a target is one -member of such a set. Letting the member name the set would cap every set at -one member. So the binding runs the other way: - -- A planner class declares `COVERS`, a frozenset of `(ctx_target_id, - gen_target_id)` pairs. Target ids are `/`, the - directory being the target's identity (`____`). -- `modeling_v2/transfers/` holds every planner and the one lookup, - `planner_for_target(target_id)`: zero planners -> default plan, two -> - configuration error. -- `AutoModelForCausalLM.from_config` asks that lookup for the target it just - built and installs the covering planner's factory as the instance's +it, and the framework does not pick a planner either: the user does. + +- `TRTLLM_KV_TRANSFER_PLANNER=` names the planner. Unset, every + model keeps the default plan. +- `disaggregation/planners/` holds every model-specific planner and + `PLANNERS`, the map from `planner_id` to class. Each planner class answers + `supports_deployment(architecture, mapping) -> bool`: the model architecture + and the parallel topology are the whole input, so a built-in model and a + modeling_v2 target of the same architecture are judged alike, and any + width that satisfies the planner's conditions passes (dep4 may feed dep8). +- `AutoModelForCausalLM.from_config` calls `bind_transfer_planner`, which + looks the name up, checks `supports_deployment` for the model just built + (an unknown name or an unsupported deployment is a start-up error with the + reason) and installs the planner's factory as the instance's `build_transfer_planner`. The framework reads the hook and nothing else; - neither the target nor the disaggregation code names a planner. -- `describe()` publishes the rank's own target id; the peer's planner reads - it from `planner_blob` and refuses any pair its matrix lacks. Equal - `planner_id` on both sides therefore means "we belong to the same - compatibility set", not "we are the same target". + neither the model nor the disaggregation code names a planner. +- `describe()` is the base class's empty blob. The framework connects only + ranks with equal `planner_id`, and the only way to hold a `planner_id` is + to have passed the same `supports_deployment` at bind time, so the peer is + already judged; `check_peer` compares the cache geometry both sides + published (pool shape, layer count, `tokens_per_block`, `slot_bytes`). Reason for the model over the cache manager still holds: a plan is a function of the parallel topology *and* the cache layout, and the target is the one From 2ce870a537d820042055dc52794cebf29b5487df Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:51:20 -0700 Subject: [PATCH 24/25] [None][test] Build the context-only CSA2 transfer test on the planner seam The test constructs PeerRegistrar and reads the peer overlap and pool mapping from it. Those now live on DefaultTransferPlanner, the registrar takes the planner, and a peer's RankInfo carries the planner id the transfer worker stamps on it, so the test does the same. Signed-off-by: Fred Wei Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../test_context_only_cache_transfer.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/unittest/disaggregated/test_context_only_cache_transfer.py b/tests/unittest/disaggregated/test_context_only_cache_transfer.py index 4c80b135ed7b..3da322100ec7 100644 --- a/tests/unittest/disaggregated/test_context_only_cache_transfer.py +++ b/tests/unittest/disaggregated/test_context_only_cache_transfer.py @@ -13,6 +13,7 @@ CSA2CacheRole, ) from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout +from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.native.transfer import RecvReqInfo, Sender @@ -160,21 +161,25 @@ def test_context_only_physical_cache_transfer_preserves_logical_owners( "generation", target, torch.cuda.current_device() ).to_bytes() ) - registrar = PeerRegistrar(source_info, KVRegionExtractorV1(source)) + source_extractor = KVRegionExtractorV1(source) + planner = DefaultTransferPlanner(source_info, source_extractor) + registrar = PeerRegistrar(source_info, source_extractor, planner) for info, dp_rank in ((source_info, source_dp_rank), (target_info, target_dp_rank)): assert info.instance_rank == info.tp_rank == info.dp_rank == (dp_rank or 0) assert info.tp_size == info.dp_size == (4 if dp_rank is not None else 1) assert info.tp_size_per_dp_group == 1 assert info.attention.enable_attention_dp == (dp_rank is not None) + # The transfer worker stamps the planner on the rank info it publishes. + target_info.planner_id = DefaultTransferPlanner.planner_id registrar.register("generation", target_info.instance_rank, target_info) - assert registrar.get_peer_overlap(target_info, target_info.dp_rank).ranks == [ + assert planner.get_peer_overlap(target_info, target_info.dp_rank).ranks == [ target_info.instance_rank ] - receiver_registrar = PeerRegistrar(target_info, KVRegionExtractorV1(target)) - assert receiver_registrar.get_peer_overlap(source_info, source_info.dp_rank).ranks == [ + receiver_planner = DefaultTransferPlanner(target_info, KVRegionExtractorV1(target)) + assert receiver_planner.get_peer_overlap(source_info, source_info.dp_rank).ranks == [ source_info.instance_rank ] - pool_mapping = registrar.get_pool_mapping(target_info) + pool_mapping = planner.get_pool_mapping(target_info) matched_global_ids = set() for (src_group, src_view), (dst_group, dst_view) in pool_mapping.items(): source_group = source_info.page_table.layer_groups[src_group] @@ -206,6 +211,7 @@ def test_context_only_physical_cache_transfer_preserves_logical_owners( assert len(target_slice.excluded_pool_views) == 20 sender = Sender.__new__(Sender) sender._registrar = registrar + sender._planner = planner task = SimpleNamespace( _chunk=source_slice, _prompt_len=length, From 3472a11dd65b60db65837c9b5c4b655afdd472fe Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Tue, 13 Oct 2026 08:21:23 -0700 Subject: [PATCH 25/25] [None][feat] Split the transfer planner into a cache describer and a row pairer The planner's plan(self_ri, peer_ri, chunk, req_info, prompt_len) mixed two kinds of knowledge: what one rank's cache holds for a request (its page table and slot tables), and how this rank's layout pairs with a peer's. Split them so each half can later move to the layer that owns it. CacheDescriber.rows(page_table, slot_table_per_layer_group, excluded_views) is single-sided: one LayerGroupRows per pool view, carrying block ordinals, slot addresses, the row size and the ordinal span. DefaultCacheDescriber is slot arithmetic off the page table. Its describer_id travels on RankInfo and is compared at registration and at the receiver gate like planner_id. TransferPlanner.pair(self_ri, peer_ri, self_rows, peer_rows, *, token_range, prompt_len, dst_start_token) is two-sided: pool matching, ownership, ordinal intersection and the mappers over the rows it is given. It no longer reads slot tables or addresses memory, and no longer keeps a peer extractor. plan_transfer() composes the two for one piece; the sender calls only that. _pair_ordinals becomes _pair_rows over row ordinals, with the helix stride and the span check inside. STATE rows are the request's one slot; _state_region expands it to per-layer pointers for the mamba mappers. The DeepSeek MLA latent planner pairs the default describer's rows and merges rows that are consecutive in memory on both sides. Signed-off-by: Fred Wei Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- ...26-10-10-disagg-transfer-planner-design.md | 77 ++++++-- .../_torch/disaggregation/base/planner.py | 136 +++++++++++-- .../native/default_describer.py | 63 ++++++ .../disaggregation/native/default_planner.py | 179 ++++++++++-------- .../_torch/disaggregation/native/peer.py | 6 + .../_torch/disaggregation/native/rank_info.py | 2 + .../_torch/disaggregation/native/transfer.py | 42 +++- .../planners/deepseek_mla_latent.py | 56 +++--- .../test_deepseek_mla_latent_planner.py | 16 +- .../disaggregation/test_default_describer.py | 112 +++++++++++ .../disaggregated/test_cache_reuse_adapter.py | 75 +++++--- .../disaggregated/test_chunked_transfer.py | 39 ++-- .../test_context_only_cache_transfer.py | 2 + .../unittest/disaggregated/test_rank_info.py | 13 +- .../disaggregated/test_transfer_planner.py | 74 ++++++-- 15 files changed, 693 insertions(+), 199 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/native/default_describer.py create mode 100644 tests/unittest/_torch/disaggregation/test_default_describer.py diff --git a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md index 329163b07cc1..996f46618cae 100644 --- a/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md +++ b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md @@ -61,6 +61,33 @@ class TransferPlan: # An empty plan is valid: this rank writes nothing to this peer but still # completes the handshake, as duplicate-head non-owners do today. +@dataclass(frozen=True) +class LayerGroupRows: + """The rows one rank holds for a request in one pool view of one layer group. + A row is a whole pool slot: block ordinal block_ordinals[i] starts at device + address addresses[i] and is row_bytes long. span_blocks is how many ordinals + the request covers on this side. A STATE group is one row, ordinal 0.""" + layer_group_index: int + view_index: int + kind: CacheKind + block_ordinals: np.ndarray + addresses: np.ndarray + row_bytes: int + span_blocks: int + +class CacheDescriber(ABC): + describer_id: str # must match on both sides of a transfer + + @abstractmethod + def rows( + self, + page_table: KVCachePageTable, + slot_table_per_layer_group: Sequence[np.ndarray], # positional, -1 = not held + excluded_views: AbstractSet[tuple[int, int]], # (layer group, view) left out + ) -> list[LayerGroupRows]: + """Single-sided: what this rank's cache holds for a request, off its own + page table. Knows the model's cache layout, never the peer.""" + class TransferPlanner(ABC): planner_id: str # must match on both sides of a transfer @@ -80,15 +107,19 @@ class TransferPlanner(ABC): Replaces get_peer_overlap + should_send_pool + _fanin_bounce_safe.""" @abstractmethod - def plan( + def pair( self, self_ri: RankInfo, peer_ri: RankInfo, - chunk: Chunk, # sender's description of the piece (base/backend.py) - req_info: RecvReqInfo, # receiver's description of where it lands (base/backend.py) + self_rows: Sequence[LayerGroupRows], # what this rank holds for the piece + peer_rows: Sequence[LayerGroupRows], # what the receiver holds + *, + token_range: TokenRange, # the piece within the prompt prompt_len: int, + dst_start_token: Optional[int], # receiver accepts nothing below this ) -> TransferPlan: - """Sender side. Replaces the body of _build_kv_write_meta.""" + """Sender side, two-sided: pair the rows. Replaces the body of + _build_kv_write_meta.""" @abstractmethod def receiver_payload_bytes( @@ -97,18 +128,36 @@ class TransferPlanner(ABC): """Receiver side, bytes beyond the paged blocks that ride the same write (recurrent state today). None = cannot size, receiver disables bounce. Replaces _get_mamba_slot + mamba_receiver_payload_bytes.""" + +def plan_transfer(describer, planner, self_ri, peer_ri, chunk, req_info, prompt_len) -> TransferPlan: + """Framework composition for one piece: describe this side from chunk, describe + the receiver from its wire message (req_info) against its published page table, + then planner.pair(...).""" ``` -`plan` and `receiver_payload_bytes` take the two request descriptions whole, -never fields picked out of them. `Chunk` is what the sender holds for the piece -(block tables per layer group, token range, and whatever per-request facts the -transceiver records later, e.g. pool views a request excludes); `RecvReqInfo` is -the receiver's wire message (its block tables, `dst_start_token`, bounce -address, aux slot). Both live in `base/backend.py`, so the planner package -depends on protocol types only, never on `native/transfer.py`'s task objects -(`KVSendTask` carries transfer progress and timers a planner must not touch). -A new per-request fact is added to one of these two types and every planner -sees it; the interface does not change and nothing silently stops flowing. +The split follows where the knowledge lives. The describer looks at one side +only: it reads this rank's page table and a request's positional slot tables +(`Chunk.block_ids_per_layer_groups`, `RecvReqInfo.block_ids_per_layer_groups`) +and says which whole slots exist for the request and where. That is the model's +cache layout, nothing about the peer, so it is the half a KV cache manager +could export on its own. The planner looks at both sides: pool matching, head +re-sharding, who writes, and the ordinal intersection that decides which rows +move. It only ever consumes rows, so it never addresses memory itself and +never reads a slot table. + +`RecvReqInfo` stays the receiver's wire message (its block tables, +`dst_start_token`, bounce address, aux slot); the sender rebuilds the +receiver's rows from it with the same describer. `receiver_payload_bytes` +still takes it whole. Both request types live in `base/backend.py`, so the +planner package depends on protocol types only, never on +`native/transfer.py`'s task objects. + +The framework compares `describer_id` and `planner_id` across peers before +anything else; both travel on `RankInfo`. Today every model runs +`DefaultCacheDescriber` (`native/default_describer.py`): one row per held +slot at `base_address + slot * slot_stride_bytes`. There is no selection hook +for a describer yet because no model needs one; when one does, it is chosen +where the planner is chosen. Threading contract is unchanged from today: `plan` is called where `_build_kv_write_meta` is called now, `targets` and `receiver_payload_bytes` diff --git a/tensorrt_llm/_torch/disaggregation/base/planner.py b/tensorrt_llm/_torch/disaggregation/base/planner.py index 00945f1bf21c..0812cbcf3447 100644 --- a/tensorrt_llm/_torch/disaggregation/base/planner.py +++ b/tensorrt_llm/_torch/disaggregation/base/planner.py @@ -14,27 +14,36 @@ # limitations under the License. """The seam between the transfer framework and a model's KV-transfer plan. -The framework (``native/transfer.py``) moves bytes and tracks transfers. Which -bytes of this rank pair with which bytes of a peer rank is a property of the -model's cache layout and the two parallel topologies, so a model class may -supply its own :class:`TransferPlanner`; ``native/default_planner.py`` is the -one every model gets unless it overrides ``build_transfer_planner``. - -Both sides of a transfer must run the same planner class (``planner_id``): -the framework enforces that and nothing else about the plan. +The framework (``native/transfer.py``) moves bytes and tracks transfers. What +a model contributes splits in two, each behind its own interface: + +* :class:`CacheDescriber` looks at one side only. Given this rank's page table + and what a request holds, it says which rows (whole pool slots) exist for the + request and where they are. It knows the model's cache layout, not the peer. +* :class:`TransferPlanner` looks at both sides. Given the rows each side + described, it decides which bytes of this rank pair with which bytes of the + peer: pool matching, head re-sharding, who writes. + +:func:`plan_transfer` composes the two for one piece of one request. Both +sides of a transfer must run the same describer and the same planner +(``describer_id``, ``planner_id``): the framework enforces that and nothing +else about the plan. """ from __future__ import annotations from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, AbstractSet, List, Optional, Sequence, Tuple import numpy as np +from .backend import CacheKind + if TYPE_CHECKING: - from tensorrt_llm._torch.disaggregation.base.backend import Chunk, RecvReqInfo + from tensorrt_llm._torch.disaggregation.base.backend import Chunk, RecvReqInfo, TokenRange from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable class PeerIncompatibleError(ValueError): @@ -73,10 +82,65 @@ def empty(cls) -> "TransferPlan": ) +@dataclass(frozen=True) +class LayerGroupRows: + """The rows one rank holds for a request in one pool view of one layer group. + + A row is a whole pool slot. Row ``i`` holds block ordinal ``block_ordinals[i]`` + (the block's position in the prompt) and starts at device address + ``addresses[i]``; every row is ``row_bytes`` long. Ordinals this side does not + hold are simply absent. ``span_blocks`` is how many ordinals the request + covers on this side, so two sides can be checked to describe the same prompt. + + A STATE (recurrent) group holds one checkpoint for the request: one row, + ordinal 0, at the slot's address. How a slot splits into layers and heads is + the page table's business, which the planner's mappers read. + """ + + layer_group_index: int + view_index: int + kind: CacheKind + block_ordinals: np.ndarray # int64, strictly increasing + addresses: np.ndarray # int64, one per row + row_bytes: int + span_blocks: int + + def __len__(self) -> int: + return int(self.block_ordinals.size) + + +class CacheDescriber(ABC): + """How one rank's cache holds a request. One instance per rank; knows no peer. + + This is the single-sided half of planning: it reads the model's page table + and a request's positional slot tables and produces :class:`LayerGroupRows`. + The sender runs it twice per piece, over its own table and, from the + receiver's wire message, over the receiver's. + """ + + #: Both sides must agree. The framework compares it before any registration. + describer_id: str = "" + + @abstractmethod + def rows( + self, + page_table: "KVCachePageTable", + slot_table_per_layer_group: Sequence[np.ndarray], + excluded_views: AbstractSet[Tuple[int, int]], + ) -> List[LayerGroupRows]: + """One entry per pool view the request may transfer, in page-table order. + + ``slot_table_per_layer_group[g]`` is positional: entry ``i`` is the pool slot + holding block ordinal ``i``, or -1 where this side holds nothing (see + ``Chunk``). Views named in ``excluded_views`` as ``(layer_group, view)`` get + no entry at all. + """ + + class TransferPlanner(ABC): """Model-owned KV transfer planning. One instance per rank, shared by Sender and Receiver. - Threading is the framework's: ``plan`` runs where the sender builds write + Threading is the framework's: ``pair`` runs where the sender builds write metadata; ``targets``, ``receiver_payload_bytes`` and ``check_peer`` run on the receiver's dispatch path and on both listeners. Keep per-peer caches keyed by ``(peer_ri.instance_name, peer_ri.instance_rank)`` and drop them @@ -99,21 +163,24 @@ def targets(self, self_ri: "RankInfo", peer_ri: "RankInfo", peer_dp_rank: int) - """Peer flat ranks this rank exchanges with for the peer's DP group ``peer_dp_rank``.""" @abstractmethod - def plan( + def pair( self, self_ri: "RankInfo", peer_ri: "RankInfo", - chunk: "Chunk", - req_info: "RecvReqInfo", + self_rows: Sequence[LayerGroupRows], + peer_rows: Sequence[LayerGroupRows], + *, + token_range: "TokenRange", prompt_len: int, + dst_start_token: Optional[int], ) -> TransferPlan: - """Sender side: the bytes this rank writes into ``peer_ri`` for one chunk. + """Sender side: the bytes this rank writes into ``peer_ri`` for one piece. - ``chunk`` is this rank's description of what it holds for the piece (block - tables per layer group, token range), ``req_info`` the receiver's description - of where it lands (its block tables, a lower bound below which it accepts - nothing). Both arrive whole so a planner honors every per-request fact the - framework records, present or future, without a signature change. + ``self_rows`` is what this rank holds for the piece, ``peer_rows`` what the + receiver holds, both as the describer produced them. ``token_range`` is the + piece within the prompt of ``prompt_len`` tokens; ``dst_start_token`` is a + lower bound below which the receiver accepts nothing, on top of the rows + it already left out. """ @abstractmethod @@ -129,3 +196,32 @@ def receiver_payload_bytes( def forget_peer(self, peer_ri: "RankInfo") -> None: """Drop any per-peer cache. Called by ``PeerRegistrar.unregister``.""" return None + + +def plan_transfer( + describer: CacheDescriber, + planner: TransferPlanner, + self_ri: "RankInfo", + peer_ri: "RankInfo", + chunk: "Chunk", + req_info: "RecvReqInfo", + prompt_len: int, +) -> TransferPlan: + """The sender's plan for one piece: describe both sides, then pair them. + + The receiver's rows are rebuilt here from the slot tables in its wire + message and its published page table. + """ + self_rows = describer.rows( + self_ri.page_table, chunk.block_ids_per_layer_groups, chunk.excluded_pool_views + ) + peer_rows = describer.rows(peer_ri.page_table, req_info.block_ids_per_layer_groups, frozenset()) + return planner.pair( + self_ri, + peer_ri, + self_rows, + peer_rows, + token_range=chunk.token_range, + prompt_len=prompt_len, + dst_start_token=req_info.dst_start_token, + ) diff --git a/tensorrt_llm/_torch/disaggregation/native/default_describer.py b/tensorrt_llm/_torch/disaggregation/native/default_describer.py new file mode 100644 index 000000000000..5d59df4533f9 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/native/default_describer.py @@ -0,0 +1,63 @@ +# 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 cache describer every model gets: one row per held slot, addressed by slot arithmetic.""" + +from __future__ import annotations + +from typing import AbstractSet, List, Sequence, Tuple + +import numpy as np + +from tensorrt_llm._torch.disaggregation.base.planner import CacheDescriber, LayerGroupRows +from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable +from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool + + +class DefaultCacheDescriber(CacheDescriber): + """Rows straight off the page table: every view of every layer group, one row per slot held. + + A PAGED row is a block's slot; a STATE row is the request's one state slot. + Both are ``base_address + slot * slot_stride_bytes``, ``slot_bytes`` long. + """ + + describer_id = "trtllm.default" + + def rows( + self, + page_table: KVCachePageTable, + slot_table_per_layer_group: Sequence[np.ndarray], + excluded_views: AbstractSet[Tuple[int, int]], + ) -> List[LayerGroupRows]: + out: List[LayerGroupRows] = [] + for lg_idx, lg in enumerate(page_table.layer_groups): + table = np.asarray(slot_table_per_layer_group[lg_idx], dtype=np.int64) + held = np.flatnonzero(table >= 0) + for view_idx, pool_view in enumerate(lg.pool_views): + if (lg_idx, view_idx) in excluded_views: + continue + pool = get_physical_pool(page_table, lg_idx, pool_view.pool_idx) + assert pool.slot_stride_bytes is not None + out.append( + LayerGroupRows( + layer_group_index=lg_idx, + view_index=view_idx, + kind=lg.kind, + block_ordinals=held, + addresses=pool.base_address + pool.slot_stride_bytes * table[held], + row_bytes=pool.slot_bytes, + span_blocks=int(table.size), + ) + ) + return out diff --git a/tensorrt_llm/_torch/disaggregation/native/default_planner.py b/tensorrt_llm/_torch/disaggregation/native/default_planner.py index 5526efaf4539..75b2794fff44 100644 --- a/tensorrt_llm/_torch/disaggregation/native/default_planner.py +++ b/tensorrt_llm/_torch/disaggregation/native/default_planner.py @@ -16,28 +16,34 @@ Attention (PAGED) and recurrent-state (STATE) layer groups are planned by ``AttentionPolicy`` and ``MambaPolicy``; pools are matched by role and layer -overlap; TP/PP/CP overlap decides which peer ranks receive what. All of it is -behind the :class:`TransferPlanner` protocol, so ``native/transfer.py`` only -ever sees ranks, pointer arrays and byte counts. +overlap; TP/PP/CP overlap decides which peer ranks receive what. The rows each +side holds come in from the cache describer; this class only pairs them. All of +it is behind the :class:`TransferPlanner` protocol, so ``native/transfer.py`` +only ever sees ranks, pointer arrays and byte counts. """ from __future__ import annotations from collections import Counter from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Sequence, Tuple, Union import numpy as np from tensorrt_llm import logger -from tensorrt_llm._torch.disaggregation.base.backend import Chunk, RecvReqInfo +from tensorrt_llm._torch.disaggregation.base.backend import RecvReqInfo, TokenRange from tensorrt_llm._torch.disaggregation.base.planner import ( + LayerGroupRows, PeerIncompatibleError, PeerTargets, TransferPlan, TransferPlanner, ) -from tensorrt_llm._torch.disaggregation.base.region import RegionMapperBase +from tensorrt_llm._torch.disaggregation.base.region import ( + MemRegionGroup, + RegionMapperBase, + SpecRegion, +) from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( MambaPolicy, @@ -101,20 +107,13 @@ def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1 self._lg_pool_mapping_cache: Dict[ str, Dict[LGPoolKey, LGPoolKey] ] = {} # peer_key -> {(self_lg, self_pi) -> (peer_lg, peer_pi)} - self._peer_ext_cache: Dict[str, KVRegionExtractorV1] = {} + # (peer_key or "", layer group, view) -> per-layer offsets within a STATE slot + self._state_layout_cache: Dict[Tuple[str, int, int], Tuple[np.ndarray, int]] = {} @staticmethod def _unique_key(name: str, rank: int) -> str: return name + str(rank) - def peer_extractor(self, peer_ri: RankInfo) -> KVRegionExtractorV1: - key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) - ext = self._peer_ext_cache.get(key) - if ext is None: - ext = KVRegionExtractorV1(peer_ri.page_table) - self._peer_ext_cache[key] = ext - return ext - # ---- TransferPlanner protocol ------------------------------------------------- def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: @@ -210,18 +209,17 @@ def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> Pe self._targets_cache[key] = targets return targets - def plan( + def pair( self, self_ri: RankInfo, peer_ri: RankInfo, - chunk: Chunk, - req_info: RecvReqInfo, + self_rows: Sequence[LayerGroupRows], + peer_rows: Sequence[LayerGroupRows], + *, + token_range: TokenRange, prompt_len: int, + dst_start_token: Optional[int], ) -> TransferPlan: - src_block_ids_per_lg = chunk.block_ids_per_layer_groups - dst_block_ids_per_lg = req_info.block_ids_per_layer_groups - token_range = chunk.token_range - dst_start_token = req_info.dst_start_token # Only a piece short of the whole prompt runs the chunk projection below, and that is the # part with no answer under CP; a sole piece covering everything takes the untouched path. is_partial = token_range.start > 0 or token_range.end < prompt_len @@ -231,6 +229,18 @@ def plan( f"(sender cp_size={self_ri.cp_size}, receiver cp_size={peer_ri.cp_size})" ) targets = self.get_peer_overlap(peer_ri, peer_ri.dp_rank) + self_by_view = {(r.layer_group_index, r.view_index): r for r in self_rows} + peer_by_view = {(r.layer_group_index, r.view_index): r for r in peer_rows} + tpb = self._self_ext_cache.page_table.tokens_per_block + dst_start_block = dst_start_token // tpb if dst_start_token is not None else 0 + # Helix: the receiver owns global blocks [cp_rank::cp_size] (same protocol as + # partition_context_for_helix), so its ordinals are the strided subset of ours; + # block reuse is rejected under helix. + helix = ( + (peer_ri.cp_rank, peer_ri.cp_size) + if peer_ri.cp_size > 1 and self_ri.cp_size == 1 + else None + ) # Aggregate fragment pointers from all matching pool pairs. # Each pool pair produces one or more region pairs (rp), each containing @@ -242,53 +252,39 @@ def plan( # tuples and construct the final sizes array with a single np.repeat(). # For 48k+ items this avoids many small allocations in the hot loop. size_specs: list[tuple[int, int]] = [] - extractor = self._self_ext_cache - peer_extractor = self.peer_extractor(peer_ri) pool_mapping = self.get_pool_mapping(peer_ri) # Aggregate fragments from all matching pools using numpy concatenation. # Send ownership is per pool: replicated pools elect one fan-in # owner, sharded pools keep head-duplication routing. for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): - if (self_lg, self_pi) in chunk.excluded_pool_views: + src = self_by_view.get((self_lg, self_pi)) + if src is None: # a view the describer left out for this piece continue + dst = peer_by_view[(peer_lg, peer_pi)] if not self.should_send_pool(targets, peer_ri, self_lg, self_pi): continue - lg_info = extractor.page_table.layer_groups[self_lg] - src_block_ids = src_block_ids_per_lg[self_lg] - dst_block_ids = dst_block_ids_per_lg[peer_lg] - - # Extract regions: STATE = per-layer ptrs, PAGED = per-block ptrs - if lg_info.kind == CacheKind.STATE: - if src_block_ids.size == 0 or dst_block_ids.size == 0: + # Build regions: STATE = per-layer ptrs of the one slot, PAGED = per-block ptrs + if src.kind == CacheKind.STATE: + if len(src) == 0 or len(dst) == 0: continue if not MambaPolicy.is_paired(self_ri, peer_ri): continue - src_region = extractor.extract_slot(int(src_block_ids[0]), self_lg, self_pi) - dst_region = peer_extractor.extract_slot(int(dst_block_ids[0]), peer_lg, peer_pi) + src_region = self._state_region("", self_ri.page_table, src) + dst_region = self._state_region( + self._unique_key(peer_ri.instance_name, peer_ri.instance_rank), + peer_ri.page_table, + dst, + ) else: - src_block_ids = np.asarray(src_block_ids, dtype=np.int64) - dst_block_ids = np.asarray(dst_block_ids, dtype=np.int64) - if peer_ri.cp_size > 1 and self_ri.cp_size == 1: - # Helix: the receiver owns global blocks [cp_rank::cp_size] - # (same protocol as partition_context_for_helix), so its - # table is the strided subset of ours; block reuse is - # rejected under helix. - src_block_ids = src_block_ids[peer_ri.cp_rank :: peer_ri.cp_size] - tpb = extractor.page_table.tokens_per_block - dst_start_block = dst_start_token // tpb if dst_start_token is not None else 0 - # Both tables are positional (index == block ordinal, -1 == not - # held), so pairing is a plain intersection: SWA eviction, prefix - # reuse, and pipelined chunk bounds are all already encoded as - # holes by whichever side owns that state. - src_block_ids, dst_block_ids = self._pair_ordinals( - src_block_ids, dst_block_ids, dst_start_block=dst_start_block + src_sel, dst_sel = self._pair_rows( + src, dst, dst_start_block=dst_start_block, helix=helix ) - if src_block_ids.size == 0: + if src_sel.size == 0: continue - src_region = extractor.extract(src_block_ids, self_lg, self_pi) - dst_region = peer_extractor.extract(dst_block_ids, peer_lg, peer_pi) + src_region = SpecRegion(MemRegionGroup(src.addresses[src_sel], src.row_bytes)) + dst_region = SpecRegion(MemRegionGroup(dst.addresses[dst_sel], dst.row_bytes)) # Map and collect fragments (same for all layer types) mapper = self.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) @@ -312,6 +308,24 @@ def plan( sizes=np.repeat(np.array(values, dtype=np.int64), counts), ) + def _state_region( + self, peer_key: str, page_table: KVCachePageTable, rows: LayerGroupRows + ) -> SpecRegion: + """Per-layer pointers of a STATE row: ``slot address + view offset of each layer``. + + View offsets are explicit because a recurrent role may exist on only a + subset of layers or be interleaved with another role in a V2 pool. + """ + key = (peer_key, rows.layer_group_index, rows.view_index) + layout = self._state_layout_cache.get(key) + if layout is None: + pv = page_table.layer_groups[rows.layer_group_index].pool_views[rows.view_index] + layer_offsets, bytes_per_layer = get_layer_byte_ranges(pv) + layout = (np.array(sorted(layer_offsets.values()), dtype=np.int64), bytes_per_layer) + self._state_layout_cache[key] = layout + ordered_offsets, bytes_per_layer = layout + return SpecRegion(MemRegionGroup(ordered_offsets + int(rows.addresses[0]), bytes_per_layer)) + def receiver_payload_bytes( self, self_ri: RankInfo, peer_ri: RankInfo, req_info: RecvReqInfo ) -> Optional[int]: @@ -329,10 +343,11 @@ def receiver_payload_bytes( def forget_peer(self, peer_ri: RankInfo) -> None: key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) - self._peer_ext_cache.pop(key, None) self._lg_pool_mapping_cache.pop(key, None) for k in [k for k in self._kv_map_cache if k[0] == key]: del self._kv_map_cache[k] + for k in [k for k in self._state_layout_cache if k[0] == key]: + del self._state_layout_cache[k] # ---- helpers moved from the transfer framework --------------------------------- @@ -351,31 +366,43 @@ def _get_mamba_slot( return None @staticmethod - def _pair_ordinals( - src_block_ids: np.ndarray, - dst_block_ids: np.ndarray, + def _pair_rows( + src: LayerGroupRows, + dst: LayerGroupRows, + *, dst_start_block: int = 0, + helix: Optional[Tuple[int, int]] = None, ) -> tuple[np.ndarray, np.ndarray]: - """Pair src/dst positional block tables ordinal by ordinal. - - Both tables are indexed by block ordinal (see Chunk); an entry is -1 - where that side has nothing to offer/accept. A block moves only where - both sides hold it, so a sender that has not evicted yet, a receiver - that already reused a prefix, a receiver with a wider speculative - window, or a sender chunk that blanks everything outside its range all - fall out of the same intersection without any per-case arithmetic. - Ordinals below ``dst_start_block`` are excluded as well. + """Pair the rows both sides hold, ordinal by ordinal; returns row indices into each side. + + Rows are keyed by block ordinal, and a side lists only the ordinals it + holds, so a block moves exactly where both sides hold it: a sender that + has not evicted yet, a receiver that already reused a prefix, a receiver + with a wider speculative window, or a sender chunk that blanks everything + outside its range all fall out of the same intersection without any + per-case arithmetic. Ordinals below ``dst_start_block`` are excluded as + well. Under ``helix=(cp_rank, cp_size)`` the receiver holds every + ``cp_size``-th global ordinal from ``cp_rank``, numbered from 0. """ - if src_block_ids.size != dst_block_ids.size: + src_ordinals = src.block_ordinals + src_index = np.arange(src_ordinals.size) + src_span = src.span_blocks + if helix is not None: + cp_rank, cp_size = helix + keep = (src_ordinals >= cp_rank) & ((src_ordinals - cp_rank) % cp_size == 0) + src_index = src_index[keep] + src_ordinals = (src_ordinals[keep] - cp_rank) // cp_size + src_span = len(range(cp_rank, src_span, cp_size)) + if src_span != dst.span_blocks: raise ValueError( - f"positional block tables differ in length: src={src_block_ids.size} " - f"dst={dst_block_ids.size}; both sides must span the same prompt" + f"positional block tables differ in length: src={src_span} " + f"dst={dst.span_blocks}; both sides must span the same prompt" ) - mask = (src_block_ids >= 0) & (dst_block_ids >= 0) - if dst_start_block > 0: - mask[:dst_start_block] = False - idx = np.flatnonzero(mask) - return src_block_ids[idx], dst_block_ids[idx] + common, src_pos, dst_pos = np.intersect1d( + src_ordinals, dst.block_ordinals, assume_unique=True, return_indices=True + ) + accepted = common >= dst_start_block + return src_index[src_pos[accepted]], dst_pos[accepted] @staticmethod def _fanin_bounce_safe( @@ -626,11 +653,11 @@ def get_kv_map( policy = self._get_policy(self_lg.kind) # For mamba, compute ptr-array layer offsets for partial PP overlap. - # extract_slot returns ptrs for ALL local_layer_ids in sorted order; + # _state_region returns ptrs for ALL local_layer_ids in sorted order; # the mapper must slice only the overlapping subset. extra_kwargs = {} if self_lg.kind == CacheKind.STATE: - # Positions index the ptr array ``extract_slot`` builds, which is + # Positions index the ptr array ``_state_region`` builds, which is # ordered by view offset. self_gid_to_pos = {gid: i for i, gid in enumerate(self_global_ids)} peer_gid_to_pos = {gid: i for i, gid in enumerate(peer_global_ids)} diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index e7fed70f3503..4a013bba7345 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -54,6 +54,12 @@ def register(self, peer_name: str, peer_rank: int, peer_ri: RankInfo): f"with '{peer_ri.planner_id}', this rank with '{self._planner.planner_id}'; " "both sides of a disaggregated deployment must run the same model planner" ) + if peer_ri.describer_id != self._ri.describer_id: + raise PeerIncompatibleError( + f"PeerRegistrar.register: peer {peer_name} (rank={peer_rank}) describes its cache " + f"with '{peer_ri.describer_id}', this rank with '{self._ri.describer_id}'; " + "both sides of a disaggregated deployment must run the same cache describer" + ) self._planner.check_peer(self._ri, peer_ri) key = self._unique_key(peer_name, peer_rank) self._peer_ri_cache[key] = peer_ri diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index b786ca97273a..52f9bba1004b 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -49,6 +49,8 @@ class RankInfo: # only the planner decodes planner_blob. planner_id: str = "" planner_blob: bytes = b"" + # Which CacheDescriber class described this rank's rows; compared across peers too. + describer_id: str = "" attention: Optional[AttentionInfo] = None aux_meta: Optional[AuxBufferMeta] = None diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index a66e7727ec9f..ab354c5c9d47 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -47,9 +47,11 @@ TransferRequest, ) from tensorrt_llm._torch.disaggregation.base.planner import ( # noqa: F401 (re-exported) + CacheDescriber, PeerIncompatibleError, PeerTargets, TransferPlanner, + plan_transfer, ) from tensorrt_llm._torch.disaggregation.base.transfer import ( ReceiverBase, @@ -65,6 +67,7 @@ build_aux_transfer_layout, get_non_empty_aux_indices, ) +from tensorrt_llm._torch.disaggregation.native.default_describer import DefaultCacheDescriber from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar @@ -869,12 +872,14 @@ def __init__( self, peer_registrar: PeerRegistrar, agent: BaseTransferAgent, + describer: CacheDescriber, bounce=None, enforce_physical_ownership: bool = False, retirement_watchdog: Optional[RetirementWatchdog] = None, ) -> None: self._registrar = peer_registrar self._planner = peer_registrar.planner + self._describer = describer self._device_id = peer_registrar.self_rank_info.device_id self._agent = agent self._bounce = bounce @@ -1645,9 +1650,18 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write # handle has to answer how many writes are still out after one of them has already failed. task.expected_transfers = expected_transfers - # Which of this rank's bytes land where on the peer is the model's planner's call; - # an empty plan still goes out so the peer's expected-writer count settles. - plan = self._planner.plan(self_ri, peer_ri, task._chunk, req_info, task._prompt_len) + # Which of this rank's bytes land where on the peer is the model's call: the describer + # says what each side holds, the planner pairs it. An empty plan still goes out so the + # peer's expected-writer count settles. + plan = plan_transfer( + self._describer, + self._planner, + self_ri, + peer_ri, + task._chunk, + req_info, + task._prompt_len, + ) if timer: timer.record_prepare_args_end(peer_ri.instance_rank) @@ -3157,7 +3171,14 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: f"this rank with '{self._planner.planner_id}'; both sides of a " "disaggregated deployment must run the same model planner" ) - self._planner.check_peer(self._registrar.self_rank_info, sender_info) + self_ri = self._registrar.self_rank_info + if sender_info.describer_id != self_ri.describer_id: + raise PeerIncompatibleError( + f"context peer describes its cache with '{sender_info.describer_id}', " + f"this rank with '{self_ri.describer_id}'; both sides of a " + "disaggregated deployment must run the same cache describer" + ) + self._planner.check_peer(self_ri, sender_info) except PeerIncompatibleError as e: msg = ( f"context peer at '{info_endpoint}' is incompatible: {e} " @@ -4369,14 +4390,18 @@ def _setup_peer_infrastructure(self, kvm: KVCacheManager): lambda _kvm, ri, ext: DefaultTransferPlanner(ri, ext) ) self._planner = factory(kvm, self._rank_info, self._kv_extractor) - # The envelope says who planned this rank's transfers and carries what - # that planner wants the peer's planner to know. RankInfoServer holds - # the object and serializes on request, so setting it here is in time. + self._describer = DefaultCacheDescriber() + # The envelope says who described this rank's cache and who planned its + # transfers, and carries what the planner wants the peer's planner to + # know. RankInfoServer holds the object and serializes on request, so + # setting it here is in time. + self._rank_info.describer_id = self._describer.describer_id self._rank_info.planner_id = self._planner.planner_id self._rank_info.planner_blob = self._planner.describe() logger.info( f"TransferWorker: KV transfer planned by {type(self._planner).__name__} " - f"(planner_id={self._planner.planner_id!r})" + f"(planner_id={self._planner.planner_id!r}, " + f"describer_id={self._describer.describer_id!r})" ) self._peer_registrar = PeerRegistrar(self._rank_info, self._kv_extractor, self._planner) @@ -4423,6 +4448,7 @@ def _setup_transfer_engine(self): self._sender = Sender( self._peer_registrar, self._agent, + self._describer, bounce=self._bounce, enforce_physical_ownership=self._config.enforce_physical_ownership, retirement_watchdog=self._retirement_watchdog, diff --git a/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py index 959f484bbb22..bb0dcdbaac1d 100644 --- a/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py +++ b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py @@ -17,8 +17,9 @@ Under attention-DP every context rank has a generation twin with the same dp_rank. MLA keeps one latent-K region per layer in one physical pool, so a block's bytes for all local layers are one contiguous slot. The plan is then -plain arithmetic: pair block ordinals both sides hold, address the slots, and -copy. Blocks that are consecutive on both sides move as one fragment. +plain arithmetic: pair the rows both sides hold by block ordinal and copy +slot for slot. Rows that are consecutive in memory on both sides move as one +fragment. The user selects this planner by name (see ``disaggregation.planners``); ``supports_deployment`` then checks that the plan holds for the deployment: a @@ -36,12 +37,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, Sequence import numpy as np -from tensorrt_llm._torch.disaggregation.base.backend import Chunk, RecvReqInfo +from tensorrt_llm._torch.disaggregation.base.backend import RecvReqInfo, TokenRange from tensorrt_llm._torch.disaggregation.base.planner import ( + LayerGroupRows, PeerIncompatibleError, PeerTargets, TransferPlan, @@ -125,42 +127,44 @@ def targets(self, self_ri: RankInfo, peer_ri: RankInfo, peer_dp_rank: int) -> Pe # destination means the fan-in split is trivially uniform. return PeerTargets(ranks=[peer_dp_rank], uniform_writers=True) - def plan( + def pair( self, self_ri: RankInfo, peer_ri: RankInfo, - chunk: Chunk, - req_info: RecvReqInfo, + self_rows: Sequence[LayerGroupRows], + peer_rows: Sequence[LayerGroupRows], + *, + token_range: TokenRange, prompt_len: int, + dst_start_token: Optional[int], ) -> TransferPlan: - del prompt_len # cp == 1: a partial piece needs no projection - if (0, 0) in chunk.excluded_pool_views: # the sole latent pool is layer group 0, view 0 + del token_range, prompt_len # cp == 1: a partial piece needs no projection + # The sole latent pool is layer group 0, view 0; the describer leaves it + # out of a piece that excludes it. + src = next((r for r in self_rows if (r.layer_group_index, r.view_index) == (0, 0)), None) + if src is None: return TransferPlan.empty() + dst = next(r for r in peer_rows if (r.layer_group_index, r.view_index) == (0, 0)) tpb = self._page_table.tokens_per_block - dst_start_token = req_info.dst_start_token dst_start_block = dst_start_token // tpb if dst_start_token is not None else 0 - src_ids, dst_ids = DefaultTransferPlanner._pair_ordinals( - np.asarray(chunk.block_ids_per_layer_groups[0], dtype=np.int64), - np.asarray(req_info.block_ids_per_layer_groups[0], dtype=np.int64), - dst_start_block=dst_start_block, + src_sel, dst_sel = DefaultTransferPlanner._pair_rows( + src, dst, dst_start_block=dst_start_block ) - if src_ids.size == 0: + if src_sel.size == 0: return TransferPlan.empty() - src_pool, dst_pool = self._pool, _sole_pool(peer_ri.page_table, "peer") - slot = src_pool.slot_bytes - # Consecutive slots are consecutive bytes only when the stride is the slot. - mergeable = src_pool.slot_stride_bytes == slot and dst_pool.slot_stride_bytes == slot - if mergeable and src_ids.size > 1: - breaks = np.flatnonzero((np.diff(src_ids) != 1) | (np.diff(dst_ids) != 1)) + 1 + src_ptrs, dst_ptrs, slot = src.addresses[src_sel], dst.addresses[dst_sel], src.row_bytes + # Consecutive rows are consecutive bytes only when they sit one slot apart on both sides. + if src_sel.size > 1: + breaks = np.flatnonzero((np.diff(src_ptrs) != slot) | (np.diff(dst_ptrs) != slot)) + 1 starts = np.concatenate(([0], breaks)) - lengths = np.diff(np.concatenate((starts, [src_ids.size]))) + lengths = np.diff(np.concatenate((starts, [src_sel.size]))) else: - starts = np.arange(src_ids.size) - lengths = np.ones(src_ids.size, dtype=np.int64) + starts = np.zeros(1, dtype=np.int64) + lengths = np.ones(1, dtype=np.int64) return TransferPlan( - src_ptrs=src_pool.base_address + src_ids[starts] * src_pool.slot_stride_bytes, - dst_ptrs=dst_pool.base_address + dst_ids[starts] * dst_pool.slot_stride_bytes, + src_ptrs=src_ptrs[starts], + dst_ptrs=dst_ptrs[starts], sizes=lengths.astype(np.int64) * slot, ) diff --git a/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py b/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py index 7642eee85836..0767a16e8577 100644 --- a/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py +++ b/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py @@ -18,7 +18,12 @@ import pytest from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, RecvReqInfo, TokenRange -from tensorrt_llm._torch.disaggregation.base.planner import PeerIncompatibleError, TransferPlan +from tensorrt_llm._torch.disaggregation.base.planner import ( + PeerIncompatibleError, + TransferPlan, + plan_transfer, +) +from tensorrt_llm._torch.disaggregation.native.default_describer import DefaultCacheDescriber from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo @@ -154,8 +159,9 @@ def test_plan_matches_default_planner_byte_for_byte(src, dst, dst_start_token): target, default, self_ri, peer_ri = _pair() chunk, req_info, prompt_len = _request(src, dst, dst_start_token) - ours = target.plan(self_ri, peer_ri, chunk, req_info, prompt_len) - ref = default.plan(self_ri, peer_ri, chunk, req_info, prompt_len) + describer = DefaultCacheDescriber() + ours = plan_transfer(describer, target, self_ri, peer_ri, chunk, req_info, prompt_len) + ref = plan_transfer(describer, default, self_ri, peer_ri, chunk, req_info, prompt_len) assert _fragments_at_slot_granularity(ours) == _fragments_at_slot_granularity(ref) assert int(ours.sizes.sum()) == int(ref.sizes.sum()) @@ -163,7 +169,9 @@ def test_plan_matches_default_planner_byte_for_byte(src, dst, dst_start_token): def test_contiguous_blocks_merge_into_one_fragment(): target, _, self_ri, peer_ri = _pair() - plan = target.plan(self_ri, peer_ri, *_request([3, 4, 5, 6], [10, 11, 12, 13])) + plan = plan_transfer( + DefaultCacheDescriber(), target, self_ri, peer_ri, *_request([3, 4, 5, 6], [10, 11, 12, 13]) + ) assert plan.sizes.tolist() == [4 * SLOT_BYTES] diff --git a/tests/unittest/_torch/disaggregation/test_default_describer.py b/tests/unittest/_torch/disaggregation/test_default_describer.py new file mode 100644 index 000000000000..ddf25118539d --- /dev/null +++ b/tests/unittest/_torch/disaggregation/test_default_describer.py @@ -0,0 +1,112 @@ +# 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 default describer turns a request's positional slot tables into rows off the page table.""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.disaggregation.base import CacheKind +from tensorrt_llm._torch.disaggregation.native.default_describer import DefaultCacheDescriber +from tensorrt_llm._torch.disaggregation.resource.page import ( + BUFFER_ENTRY_DTYPE, + AttentionLayerGroup, + KVCachePageTable, + LocalLayer, + MambaLayerGroup, + PhysicalPool, + PhysicalPoolGroup, + PoolView, +) + +pytestmark = pytest.mark.cpu_only + +KV_SLOT = 1024 +STATE_SLOT = 4096 +STATE_STRIDE = 8192 # the state pool pads every slot + + +def _page_table() -> KVCachePageTable: + """A PAGED group with two views over two pools, and a STATE group with one padded pool.""" + entries = np.array([(0, 0, KV_SLOT)], dtype=BUFFER_ENTRY_DTYPE) + paged = AttentionLayerGroup( + pool_group_idx=0, + kv_head_num_per_rank=1, + local_layers=[LocalLayer(local_layer_id=0, global_layer_id=0)], + pool_views=[ + PoolView(pool_idx=0, buffer_entries=entries), + PoolView(pool_idx=1, buffer_entries=entries, pool_role=frozenset({"index"})), + ], + ) + state = MambaLayerGroup( + pool_group_idx=1, + local_layers=[LocalLayer(local_layer_id=0, global_layer_id=1)], + pool_views=[PoolView(pool_idx=0, buffer_entries=entries)], + ) + return KVCachePageTable( + tokens_per_block=16, + layer_groups=[paged, state], + pool_groups=[ + PhysicalPoolGroup( + pools=[ + PhysicalPool(base_address=0x1000, slot_bytes=KV_SLOT, num_slots=64), + PhysicalPool(base_address=0x9000, slot_bytes=KV_SLOT, num_slots=64), + ] + ), + PhysicalPoolGroup( + pools=[ + PhysicalPool( + base_address=0x20000, + slot_bytes=STATE_SLOT, + num_slots=8, + slot_stride_bytes=STATE_STRIDE, + ) + ] + ), + ], + ) + + +def _tables(paged, state): + return [np.array(paged, dtype=np.int64), np.array(state, dtype=np.int64)] + + +def test_rows_hold_only_the_ordinals_the_table_holds_and_keep_the_span(): + rows = DefaultCacheDescriber().rows(_page_table(), _tables([7, -1, 9, -1], [3]), frozenset()) + + assert [(r.layer_group_index, r.view_index, r.kind) for r in rows] == [ + (0, 0, CacheKind.PAGED), + (0, 1, CacheKind.PAGED), + (1, 0, CacheKind.STATE), + ] + kv, index, state = rows + assert kv.block_ordinals.tolist() == [0, 2] and kv.span_blocks == 4 and len(kv) == 2 + assert kv.addresses.tolist() == [0x1000 + 7 * KV_SLOT, 0x1000 + 9 * KV_SLOT] + assert kv.row_bytes == KV_SLOT + # The second view reads the same slot table against its own pool. + assert index.block_ordinals.tolist() == [0, 2] + assert index.addresses.tolist() == [0x9000 + 7 * KV_SLOT, 0x9000 + 9 * KV_SLOT] + # A STATE group is the request's one slot, addressed by the pool's stride. + assert state.block_ordinals.tolist() == [0] and state.span_blocks == 1 + assert state.addresses.tolist() == [0x20000 + 3 * STATE_STRIDE] + assert state.row_bytes == STATE_SLOT + + +def test_excluded_views_get_no_rows_and_empty_tables_keep_their_span(): + rows = DefaultCacheDescriber().rows(_page_table(), _tables([-1, -1], []), {(0, 1)}) + + assert [(r.layer_group_index, r.view_index) for r in rows] == [(0, 0), (1, 0)] + kv, state = rows + assert len(kv) == 0 and kv.span_blocks == 2 + assert len(state) == 0 and state.span_blocks == 0 diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index 2b683e16b849..b5d69d4c83f2 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -12,7 +12,7 @@ # 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. -"""Tests for CacheReuseAdapter, the positional _create_chunk, and DefaultTransferPlanner ordinal pairing.""" +"""Tests for CacheReuseAdapter, the positional _create_chunk, and DefaultTransferPlanner row pairing.""" from types import SimpleNamespace from unittest.mock import MagicMock @@ -20,6 +20,8 @@ import numpy as np import pytest +from tensorrt_llm._torch.disaggregation.base import CacheKind +from tensorrt_llm._torch.disaggregation.base.planner import LayerGroupRows from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.resource.cache_reuse import ( CacheReuseAdapter, @@ -36,60 +38,74 @@ def _arr(values): return np.array(values, dtype=np.int64) +def _rows(table) -> LayerGroupRows: + """Rows off a positional slot table (index == block ordinal, -1 == not held), addressed by slot.""" + table = _arr(table) + held = np.flatnonzero(table >= 0) + return LayerGroupRows( + layer_group_index=0, + view_index=0, + kind=CacheKind.PAGED, + block_ordinals=held, + addresses=table[held], + row_bytes=1, + span_blocks=int(table.size), + ) + + +def _pair(src_table, dst_table, **kw): + """The slots paired by DefaultTransferPlanner._pair_rows, as (src slots, dst slots).""" + src, dst = _rows(src_table), _rows(dst_table) + src_sel, dst_sel = DefaultTransferPlanner._pair_rows(src, dst, **kw) + return src.addresses[src_sel], dst.addresses[dst_sel] + + # --------------------------------------------------------------------------- -# DefaultTransferPlanner._pair_ordinals: the whole alignment protocol is one intersection. +# DefaultTransferPlanner._pair_rows: the whole alignment protocol is one intersection. # --------------------------------------------------------------------------- -class TestPairOrdinals: - """Both tables are positional (index == block ordinal, -1 == not held). +class TestPairRows: + """Each side lists the ordinals it holds; a block moves only where both sides hold it. - The sender moves a block only where both sides hold it. Every previously - special-cased situation (SWA eviction, receiver prefix reuse, wider receiver - window, pipelined chunk bounds) is just a different hole pattern. + Every previously special-cased situation (SWA eviction, receiver prefix + reuse, wider receiver window, pipelined chunk bounds) is just a different + hole pattern. """ def test_identity(self): - src, dst = DefaultTransferPlanner._pair_ordinals(_arr([10, 11, 12]), _arr([20, 21, 22])) + src, dst = _pair(_arr([10, 11, 12]), _arr([20, 21, 22])) np.testing.assert_array_equal(src, [10, 11, 12]) np.testing.assert_array_equal(dst, [20, 21, 22]) def test_receiver_prefix_reuse_is_a_dst_hole(self): - src, dst = DefaultTransferPlanner._pair_ordinals( - _arr([10, 11, 12, 13]), _arr([-1, -1, 22, 23]) - ) + src, dst = _pair(_arr([10, 11, 12, 13]), _arr([-1, -1, 22, 23])) np.testing.assert_array_equal(src, [12, 13]) np.testing.assert_array_equal(dst, [22, 23]) def test_sender_chunk_bounds_are_src_holes(self): - src, dst = DefaultTransferPlanner._pair_ordinals( - _arr([-1, 11, 12, -1]), _arr([20, 21, 22, 23]) - ) + src, dst = _pair(_arr([-1, 11, 12, -1]), _arr([20, 21, 22, 23])) np.testing.assert_array_equal(src, [11, 12]) np.testing.assert_array_equal(dst, [21, 22]) def test_holes_on_both_sides_take_the_intersection(self): - src, dst = DefaultTransferPlanner._pair_ordinals( - _arr([10, -1, 12, 13]), _arr([-1, 21, 22, -1]) - ) + src, dst = _pair(_arr([10, -1, 12, 13]), _arr([-1, 21, 22, -1])) np.testing.assert_array_equal(src, [12]) np.testing.assert_array_equal(dst, [22]) def test_no_overlap_is_empty(self): - src, dst = DefaultTransferPlanner._pair_ordinals(_arr([10, 11, -1]), _arr([-1, -1, 22])) + src, dst = _pair(_arr([10, 11, -1]), _arr([-1, -1, 22])) assert src.size == 0 assert dst.size == 0 def test_dst_start_block_masks_the_head(self): - src, dst = DefaultTransferPlanner._pair_ordinals( - _arr([10, 11, 12, 13]), _arr([20, 21, 22, 23]), dst_start_block=2 - ) + src, dst = _pair(_arr([10, 11, 12, 13]), _arr([20, 21, 22, 23]), dst_start_block=2) np.testing.assert_array_equal(src, [12, 13]) np.testing.assert_array_equal(dst, [22, 23]) def test_length_mismatch_raises(self): with pytest.raises(ValueError, match="differ in length"): - DefaultTransferPlanner._pair_ordinals(_arr([10, 11]), _arr([20, 21, 22])) + _pair(_arr([10, 11]), _arr([20, 21, 22])) def test_sender_not_yet_evicted_but_receiver_has(self): # Regression (DeepSeek-V4 SWA, prompt_len == 2 blocks, window == 1 @@ -97,19 +113,26 @@ def test_sender_not_yet_evicted_but_receiver_has(self): # advanced, while gen pre-declared history_length=prompt_len and never # allocated ordinal 0. Pairing by ordinal copies X1 -> G1; the old # length-based suffix arithmetic copied X0 -> G1. - src, dst = DefaultTransferPlanner._pair_ordinals(_arr([100, 101]), _arr([-1, 201])) + src, dst = _pair(_arr([100, 101]), _arr([-1, 201])) np.testing.assert_array_equal(src, [101]) np.testing.assert_array_equal(dst, [201]) def test_receiver_with_wider_speculative_window(self): # Only gen runs spec decoding, so its SWA window keeps one more block. # That block has no source; it simply is not written. - src, dst = DefaultTransferPlanner._pair_ordinals( - _arr([-1, -1, 12, 13]), _arr([-1, 21, 22, 23]) - ) + src, dst = _pair(_arr([-1, -1, 12, 13]), _arr([-1, 21, 22, 23])) np.testing.assert_array_equal(src, [12, 13]) np.testing.assert_array_equal(dst, [22, 23]) + def test_helix_receiver_holds_every_cp_size_th_ordinal(self): + # cp_size=2: receiver cp_rank=1 owns global ordinals 1, 3, 5 as its ordinals 0, 1, 2. + src, dst = _pair(_arr([10, 11, 12, 13, 14, 15]), _arr([20, -1, 22]), helix=(1, 2)) + np.testing.assert_array_equal(src, [11, 15]) + np.testing.assert_array_equal(dst, [20, 22]) + # The strided span must still match the receiver's table. + with pytest.raises(ValueError, match="differ in length"): + _pair(_arr([10, 11, 12, 13, 14]), _arr([20, 21, 22]), helix=(1, 2)) + # --------------------------------------------------------------------------- # _create_chunk: positional table over ceil(prompt_len / tpb) ordinals. diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 95195a6416ca..f9bbbad1d3c0 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -26,6 +26,7 @@ from tensorrt_llm import DisaggregatedParams from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base.planner import CacheDescriber, LayerGroupRows from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.transfer import ( @@ -149,6 +150,7 @@ def _make_projection_sender() -> Sender: self_endpoint="tcp://decode:0", cp_size=1, cp_rank=0, + page_table=None, ) extractor = MagicMock() @@ -159,14 +161,27 @@ def _make_projection_sender() -> Sender: SimpleNamespace(kind=CacheKind.PAGED, sliding_window_size=_PROJECTION_PROMPT_TOKENS), ], ) - # extract(region_ids, layer_group_id, pool_idx) - the caller passes the - # layer-group and pool indices positionally, so swallow them. - extractor.extract.side_effect = lambda block_ids, *_, **__: SimpleNamespace( - memory=SimpleNamespace( - ptrs=np.asarray(block_ids, dtype=np.int64), - bytes_per_region=1, - ) - ) + + class _SlotAddressedDescriber(CacheDescriber): + """One view per layer group; a row's address is its slot id, one byte long.""" + + def rows(self, page_table, slot_table_per_layer_group, excluded_views): + out = [] + for lg_idx, table in enumerate(slot_table_per_layer_group): + table = np.asarray(table, dtype=np.int64) + held = np.flatnonzero(table >= 0) + out.append( + LayerGroupRows( + layer_group_index=lg_idx, + view_index=0, + kind=CacheKind.PAGED, + block_ordinals=held, + addresses=table[held], + row_bytes=1, + span_blocks=int(table.size), + ) + ) + return out mapper = MagicMock() mapper.map.side_effect = lambda src_region, dst_region: SimpleNamespace( @@ -175,16 +190,15 @@ def _make_projection_sender() -> Sender: ) registrar = MagicMock() - registrar.self_rank_info = SimpleNamespace(cp_size=1, cp_rank=0, dp_rank=0) + registrar.self_rank_info = SimpleNamespace(cp_size=1, cp_rank=0, dp_rank=0, page_table=None) registrar.self_extractor = extractor registrar.get_peer_rank_info.return_value = peer_ri - # The real planner's plan() over mocked pool matching / mappers, so the + # The real planner's pair() over mocked pool matching / mappers, so the # ordinal pairing under test is the production code path. planner = DefaultTransferPlanner.__new__(DefaultTransferPlanner) planner._self_ext_cache = extractor planner._targets_cache = {} - planner.peer_extractor = MagicMock(return_value=extractor) planner.get_peer_overlap = MagicMock( return_value=SimpleNamespace(ranks=[0], duplicate_head_factor=1, overlap_pp_size=1) ) @@ -200,6 +214,7 @@ def _make_projection_sender() -> Sender: sender = Sender.__new__(Sender) sender._registrar = registrar sender._planner = planner + sender._describer = _SlotAddressedDescriber() return sender @@ -314,7 +329,7 @@ def _projection_task_for(start_block: int, end_block: int) -> KVSendTask: def test_context_parallelism_is_refused_only_for_a_piece_short_of_the_prompt(): """The guard only fires for a piece short of the whole prompt; a sole whole-prompt piece passes.""" sender = _make_projection_sender() - sender._registrar.self_rank_info = SimpleNamespace(cp_size=2, cp_rank=0) + sender._registrar.self_rank_info = SimpleNamespace(cp_size=2, cp_rank=0, page_table=None) sender._build_kv_write_meta(_projection_task_for(0, 8), _make_projection_req_info()) with pytest.raises(ValueError, match="context parallelism"): diff --git a/tests/unittest/disaggregated/test_context_only_cache_transfer.py b/tests/unittest/disaggregated/test_context_only_cache_transfer.py index 3da322100ec7..019cd7c3db36 100644 --- a/tests/unittest/disaggregated/test_context_only_cache_transfer.py +++ b/tests/unittest/disaggregated/test_context_only_cache_transfer.py @@ -13,6 +13,7 @@ CSA2CacheRole, ) from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout +from tensorrt_llm._torch.disaggregation.native.default_describer import DefaultCacheDescriber from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo @@ -212,6 +213,7 @@ def test_context_only_physical_cache_transfer_preserves_logical_owners( sender = Sender.__new__(Sender) sender._registrar = registrar sender._planner = planner + sender._describer = DefaultCacheDescriber() task = SimpleNamespace( _chunk=source_slice, _prompt_len=length, diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index cc50877dd9b9..5008d28e77e0 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -223,10 +223,11 @@ def _minimal_rank_info(**overrides) -> RankInfo: def test_rank_info_planner_fields_round_trip(): - ri = _minimal_rank_info(planner_id="toy", planner_blob=b"\x00\x01") + ri = _minimal_rank_info(planner_id="toy", planner_blob=b"\x00\x01", describer_id="rows") back = RankInfo.from_bytes(ri.to_bytes()) assert back.planner_id == "toy" assert back.planner_blob == b"\x00\x01" + assert back.describer_id == "rows" def test_rank_info_from_bytes_without_planner_fields(): @@ -240,8 +241,16 @@ def test_rank_info_from_bytes_without_planner_fields(): import msgpack data = asdict(_minimal_rank_info()) - for key in ("planner_id", "planner_blob", "attention", "aux_meta", "page_table"): + for key in ( + "planner_id", + "planner_blob", + "describer_id", + "attention", + "aux_meta", + "page_table", + ): data.pop(key, None) back = RankInfo.from_bytes(msgpack.packb(data)) assert back.planner_id == "" assert back.planner_blob == b"" + assert back.describer_id == "" diff --git a/tests/unittest/disaggregated/test_transfer_planner.py b/tests/unittest/disaggregated/test_transfer_planner.py index 559b57aa279a..76327d3f9161 100644 --- a/tests/unittest/disaggregated/test_transfer_planner.py +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -18,6 +18,8 @@ import pytest from tensorrt_llm._torch.disaggregation.base.planner import ( + CacheDescriber, + LayerGroupRows, PeerIncompatibleError, PeerTargets, TransferPlan, @@ -34,6 +36,8 @@ def test_transfer_plan_empty_is_three_int64_arrays(): def test_transfer_planner_is_abstract(): with pytest.raises(TypeError): TransferPlanner() # type: ignore[abstract] + with pytest.raises(TypeError): + CacheDescriber() # type: ignore[abstract] def test_peer_incompatible_error_is_value_error(): @@ -73,15 +77,42 @@ def check_peer(self, self_ri, peer_ri): def targets(self, self_ri, peer_ri, peer_dp_rank): return PeerTargets(ranks=list(self._ranks), uniform_writers=True) - def plan(self, self_ri, peer_ri, chunk, req_info, prompt_len): - self.calls.append((chunk, req_info, prompt_len)) + def pair( + self, self_ri, peer_ri, self_rows, peer_rows, *, token_range, prompt_len, dst_start_token + ): + self.calls.append((self_rows, peer_rows, token_range, prompt_len, dst_start_token)) return self._plan def receiver_payload_bytes(self, self_ri, peer_ri, req_info): return 0 -def _sender_with(planner: TransferPlanner) -> Sender: +class _ToyDescriber(CacheDescriber): + """One row per held slot at address == slot; records what the framework handed it.""" + + describer_id = "toy" + + def __init__(self): + self.calls: list = [] + + def rows(self, page_table, slot_table_per_layer_group, excluded_views): + self.calls.append((page_table, slot_table_per_layer_group, excluded_views)) + table = np.asarray(slot_table_per_layer_group[0], dtype=np.int64) + held = np.flatnonzero(table >= 0) + return [ + LayerGroupRows( + layer_group_index=0, + view_index=0, + kind=CacheKind.PAGED, + block_ordinals=held, + addresses=table[held], + row_bytes=1, + span_blocks=int(table.size), + ) + ] + + +def _sender_with(planner: TransferPlanner, describer: CacheDescriber | None = None) -> Sender: peer_ri = SimpleNamespace( instance_name="gen", instance_rank=3, @@ -89,15 +120,16 @@ def _sender_with(planner: TransferPlanner) -> Sender: device_id=1, cp_size=1, self_endpoint="tcp://gen:1", - page_table=None, + page_table="peer page table", ) sender = Sender.__new__(Sender) sender._registrar = Mock() sender._registrar.self_rank_info = SimpleNamespace( - cp_size=1, instance_name="ctx", instance_rank=0 + cp_size=1, instance_name="ctx", instance_rank=0, page_table="self page table" ) sender._registrar.get_peer_rank_info.return_value = peer_ri sender._planner = planner + sender._describer = describer or _ToyDescriber() return sender @@ -108,6 +140,7 @@ def _task_and_req(): kind_per_layer_group=[CacheKind.PAGED], token_range=TokenRange(0, 64), is_last=True, + excluded_pool_views={(0, 1)}, ), DisaggregatedParams(disagg_request_id=1), slice_id=0, @@ -125,14 +158,15 @@ def _task_and_req(): return task, req -def test_build_kv_write_meta_uses_planner_plan(): +def test_build_kv_write_meta_describes_both_sides_then_pairs(): plan = TransferPlan( src_ptrs=np.array([100, 200], dtype=np.int64), dst_ptrs=np.array([300, 400], dtype=np.int64), sizes=np.array([16, 16], dtype=np.int64), ) planner = _ToyPlanner(plan, ranks=(3, 4)) - sender = _sender_with(planner) + describer = _ToyDescriber() + sender = _sender_with(planner, describer) task, req = _task_and_req() wm = sender._build_kv_write_meta(task, req) @@ -140,9 +174,16 @@ def test_build_kv_write_meta_uses_planner_plan(): assert wm.src_ptrs is plan.src_ptrs and wm.dst_ptrs is plan.dst_ptrs and wm.sizes is plan.sizes assert wm.expected_transfers == 2 and task.expected_transfers == 2 assert wm.dst_device_id == 1 and wm.peer_rank == 3 and wm.peer_endpoint == "tcp://gen:1" - # The planner gets both request descriptions whole, not fields picked out of them. - ((chunk, req_info, prompt_len),) = planner.calls - assert chunk is task._chunk and req_info is req and prompt_len == 64 + # The describer sees each side's page table and slot tables; only this side excludes views. + (self_pt, self_tables, self_excluded), (peer_pt, peer_tables, peer_excluded) = describer.calls + assert self_pt == "self page table" and self_tables is task._chunk.block_ids_per_layer_groups + assert self_excluded == {(0, 1)} + assert peer_pt == "peer page table" and peer_tables is req.block_ids_per_layer_groups + assert not peer_excluded + # The planner pairs the rows the describer produced, plus the request's bounds. + ((self_rows, peer_rows, token_range, prompt_len, dst_start_token),) = planner.calls + assert self_rows[0].addresses.tolist() == [5, 6] and peer_rows[0].addresses.tolist() == [9, 10] + assert token_range == task._chunk.token_range and prompt_len == 64 and dst_start_token == 32 def test_build_kv_write_meta_empty_plan_still_delivers(): @@ -252,6 +293,17 @@ def test_registrar_rejects_planner_id_mismatch(): reg.get_peer_rank_info("gen", 0) +def test_registrar_rejects_describer_id_mismatch(): + self_ri = _minimal_rank_info(planner_id="toy", describer_id="toy") + peer_ri = _minimal_rank_info( + instance_name="gen", planner_id="toy", describer_id="trtllm.default" + ) + reg = PeerRegistrar(self_ri, Mock(), _ToyPlanner(TransferPlan.empty())) + + with pytest.raises(PeerIncompatibleError, match="describes its cache"): + reg.register("gen", 0, peer_ri) + + def test_receiver_rejects_planner_id_mismatch(monkeypatch): recv = Receiver.__new__(Receiver) recv._incompatible_peers = {} @@ -349,7 +401,7 @@ def inner(*args, **kwargs): return deco u.nvtx_range = wrapping_nvtx_range from tensorrt_llm._torch.disaggregation.native.default_planner import DefaultTransferPlanner as P -for name in ("_fanin_bounce_safe", "_pair_ordinals", "_get_mamba_slot", "_get_buffers_per_layer", "_find_overlap"): +for name in ("_fanin_bounce_safe", "_pair_rows", "_get_mamba_slot", "_get_buffers_per_layer", "_find_overlap"): assert isinstance(P.__dict__[name], staticmethod), name print("STATICS_OK") """