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 new file mode 100644 index 000000000000..64a1df731df1 --- /dev/null +++ b/docs/superpowers/plans/2026-10-10-deepseek-r1-dep4-transfer-planner.md @@ -0,0 +1,132 @@ +# 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`. + +## 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`. + +## 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. 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) + +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. + +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 +contract/op bind). 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..603a2149485d --- /dev/null +++ b/docs/superpowers/plans/2026-10-10-disagg-transfer-planner.md @@ -0,0 +1,1083 @@ +# 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`. +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 + +- 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. 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..996f46618cae --- /dev/null +++ b/docs/superpowers/specs/2026-10-10-disagg-transfer-planner-design.md @@ -0,0 +1,300 @@ +# 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. + +@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 + + @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 pair( + self, + self_ri: RankInfo, + peer_ri: RankInfo, + 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, two-sided: pair the rows. Replaces the body of + _build_kv_write_meta.""" + + @abstractmethod + def receiver_payload_bytes( + 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.""" + +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(...).""" +``` + +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` +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) +``` + +The default is the method above. A modeling_v2 target does **not** override +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 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 +object that already owns both. The cache manager builds the page table and +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`; +`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). 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 new file mode 100644 index 000000000000..0812cbcf3447 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/base/planner.py @@ -0,0 +1,227 @@ +# 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. 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, 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, TokenRange + from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable + + +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), + ) + + +@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: ``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 + 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 pair( + self, + self_ri: "RankInfo", + peer_ri: "RankInfo", + 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 piece. + + ``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 + def receiver_payload_bytes( + self, self_ri: "RankInfo", peer_ri: "RankInfo", req_info: "RecvReqInfo" + ) -> 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 + + +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/kv_cache_transceiver.py b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py index e93b8a819c59..d1096d745ef8 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,15 @@ def create_kv_cache_transceiver( KvCacheTransceiverV2 logger.info("Using KvCacheTransceiverV2") # MixedMambaHybridCacheManager contains both the KV and Mamba pools. - return KvCacheTransceiverV2(mapping, dist, kv_cache_manager, - cache_transceiver_config) + # 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, + 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/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 new file mode 100644 index 000000000000..75b2794fff44 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/native/default_planner.py @@ -0,0 +1,842 @@ +# 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. 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, Sequence, Tuple, Union + +import numpy as np + +from tensorrt_llm import logger +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 ( + 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, + 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._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)} + # (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) + + # ---- TransferPlanner protocol ------------------------------------------------- + + def check_peer(self, self_ri: RankInfo, peer_ri: RankInfo) -> None: + """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, + 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( + "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(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) + 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: + # 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 + # 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) + ) + targets = PeerTargets(ranks=list(overlap.ranks), uniform_writers=uniform) + self._targets_cache[key] = targets + return targets + + def pair( + self, + self_ri: RankInfo, + peer_ri: RankInfo, + self_rows: Sequence[LayerGroupRows], + peer_rows: Sequence[LayerGroupRows], + *, + 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) + 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 + # 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]] = [] + 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(): + 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 + + # 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 = 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_sel, dst_sel = self._pair_rows( + src, dst, dst_start_block=dst_start_block, helix=helix + ) + if src_sel.size == 0: + continue + 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)) + 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 _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]: + receiver_pt = self._self_ext_cache.page_table + 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: + 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._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 --------------------------------- + + @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_rows( + src: LayerGroupRows, + dst: LayerGroupRows, + *, + dst_start_block: int = 0, + helix: Optional[Tuple[int, int]] = None, + ) -> tuple[np.ndarray, np.ndarray]: + """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. + """ + 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_span} " + f"dst={dst.span_blocks}; both sides must span the same prompt" + ) + 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( + 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. + # _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 ``_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)} + 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/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/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index d8d37658195b..4a013bba7345 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -12,104 +12,59 @@ # 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" + ) + 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 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 +76,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)] @@ -153,445 +102,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 _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/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 87731d303204..52f9bba1004b 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -44,6 +44,13 @@ 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"" + # 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 56c61cf1ab59..ab354c5c9d47 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, @@ -47,6 +46,13 @@ TransferOp, 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, RxSessionBase, @@ -61,13 +67,10 @@ 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.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.peer import PeerOverlap, PeerRegistrar +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 ( @@ -78,7 +81,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 @@ -107,48 +109,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 @@ -192,16 +152,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" @@ -922,11 +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 @@ -1085,7 +1038,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 @@ -1102,6 +1055,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): @@ -1679,153 +1638,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 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, ) - 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, @@ -1842,7 +1691,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) @@ -2229,7 +2078,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: @@ -2243,7 +2092,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): @@ -2980,6 +2829,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 @@ -3132,56 +2982,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( @@ -3216,64 +3016,41 @@ 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) + 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 ) @@ -3283,7 +3060,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) @@ -3292,7 +3069,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: @@ -3301,7 +3078,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: @@ -3380,27 +3157,29 @@ 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_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} " "(cached: all further requests to this context peer fail " @@ -4513,6 +4292,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: @@ -4601,7 +4386,24 @@ 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) + 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"describer_id={self._describer.describer_id!r})" + ) + 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) @@ -4646,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/__init__.py b/tensorrt_llm/_torch/disaggregation/planners/__init__.py new file mode 100644 index 000000000000..22d956a98f2d --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/planners/__init__.py @@ -0,0 +1,86 @@ +# 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 KV transfer planner the user asks for, bound to the model it will serve. + +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 planner's factory as the instance's +``build_transfer_planner``; the framework reads that hook and nothing else. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Dict, Optional, Type + +from tensorrt_llm._torch.disaggregation.base.planner import TransferPlanner + +from .deepseek_mla_latent import DeepseekMlaLatentTransferPlanner + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + +PLANNER_ENV_VAR = "TRTLLM_KV_TRANSFER_PLANNER" + +#: 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 planner named by the environment on ``model``; return it, or None if none is named. + + Raises ``ValueError`` when the name is unknown or the planner does not + support this deployment. + """ + key = os.environ.get(PLANNER_ENV_VAR, "") + if not key: + return None + planner_cls = PLANNERS.get(key) + if planner_cls is 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) + + model.build_transfer_planner = build_transfer_planner + return planner_cls + + +__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 new file mode 100644 index 000000000000..bb0dcdbaac1d --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/planners/deepseek_mla_latent.py @@ -0,0 +1,174 @@ +# 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 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 +block's bytes for all local layers are one contiguous slot. The plan is then +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 +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 TYPE_CHECKING, Optional, Sequence + +import numpy as np + +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.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 + +if TYPE_CHECKING: + from tensorrt_llm.mapping import Mapping + + +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 DeepseekMlaLatentTransferPlanner(TransferPlanner): + planner_id = "trtllm.deepseek_mla_latent" + + @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._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: + 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: + # 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 pair( + self, + self_ri: RankInfo, + peer_ri: RankInfo, + self_rows: Sequence[LayerGroupRows], + peer_rows: Sequence[LayerGroupRows], + *, + token_range: TokenRange, + prompt_len: int, + dst_start_token: Optional[int], + ) -> TransferPlan: + 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_block = dst_start_token // tpb if dst_start_token is not None else 0 + src_sel, dst_sel = DefaultTransferPlanner._pair_rows( + src, dst, dst_start_block=dst_start_block + ) + if src_sel.size == 0: + return TransferPlan.empty() + + 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_sel.size]))) + else: + starts = np.zeros(1, dtype=np.int64) + lengths = np.ones(1, dtype=np.int64) + return TransferPlan( + src_ptrs=src_ptrs[starts], + dst_ptrs=dst_ptrs[starts], + sizes=lengths.astype(np.int64) * slot, + ) + + def receiver_payload_bytes( + 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/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_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 0592317d293a..cf719704dba7 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -77,4 +77,10 @@ def from_config( with model_extra_attrs(extra_attrs): model = cls(config) model.extra_attrs = extra_attrs + # 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/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 0a6060b47288..5d2407575b15 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -801,6 +801,19 @@ 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..b886139efe15 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -4302,9 +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) + 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/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"]) 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}, + ) 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/_torch/disaggregation/test_deepseek_mla_latent_planner.py b/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py new file mode 100644 index 000000000000..0767a16e8577 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/test_deepseek_mla_latent_planner.py @@ -0,0 +1,195 @@ +# 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 DeepSeek MLA-latent planner must produce the default planner's bytes for the deployments it supports.""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, RecvReqInfo, TokenRange +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 +from tensorrt_llm._torch.disaggregation.planners.deepseek_mla_latent import ( + DeepseekMlaLatentTransferPlanner, +) +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( + planner_id=DeepseekMlaLatentTransferPlanner.planner_id, + 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 ( + DeepseekMlaLatentTransferPlanner(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 _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( + "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() + chunk, req_info, prompt_len = _request(src, dst, dst_start_token) + + 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()) + + +def test_contiguous_blocks_merge_into_one_fragment(): + target, _, self_ri, peer_ri = _pair() + 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] + + +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_refuses_a_peer_whose_latent_geometry_differs(): + target, _, self_ri, peer_ri = _pair() + 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)) + ) + with pytest.raises(PeerIncompatibleError, match="layer"): + target.check_peer( + self_ri, _dep4_rank_info("gen", 2, _mla_page_table(0x9000_0000, num_layers=5)) + ) 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/_torch/disaggregation/test_planner_binding.py b/tests/unittest/_torch/disaggregation/test_planner_binding.py new file mode 100644 index 000000000000..caa54411d442 --- /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. +"""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 + +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 import PLANNER_ENV_VAR +from tensorrt_llm._torch.disaggregation.planners.deepseek_mla_latent import ( + DeepseekMlaLatentTransferPlanner, +) +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.mapping import Mapping + +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) +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: Mapping) -> ModelConfig: + return ModelConfig( + pretrained_config=PretrainedConfig(architectures=[architecture]), mapping=mapping + ) + + +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) + assert supports(DEEPSEEK, DEP8) + assert not supports("GptOssForCausalLM", DEP4) + assert not supports(DEEPSEEK, TP4) + assert not supports(DEEPSEEK, DEP4_PP2) + + +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(): + """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) + for p in root.rglob("*.py") + if "transfer_planner" in p.read_text() or "TransferPlanner" in p.read_text() + ] + assert hits == [], hits + + +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(DEEPSEEK, DEP4)) is None + assert stub.build_transfer_planner(None, None, None) == "default" + monkeypatch.setenv(PLANNER_ENV_VAR, "") + assert planners.bind_transfer_planner(stub, _model_config(DEEPSEEK, DEP4)) is None + + +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() + 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) 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_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_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index 2a8ca61c646a..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 Sender ordinal pairing.""" +"""Tests for CacheReuseAdapter, the positional _create_chunk, and DefaultTransferPlanner row pairing.""" from types import SimpleNamespace from unittest.mock import MagicMock @@ -20,7 +20,9 @@ import numpy as np import pytest -from tensorrt_llm._torch.disaggregation.native.transfer import Sender +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, _CacheReuseAdapterV1, @@ -36,54 +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] + + # --------------------------------------------------------------------------- -# Sender._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 = Sender._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 = Sender._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 = Sender._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 = Sender._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 = Sender._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 = Sender._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"): - Sender._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 @@ -91,17 +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 = Sender._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 = Sender._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 d75c51a6ca77..f9bbbad1d3c0 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -26,7 +26,9 @@ 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 ( AgentResult, KVSendTask, @@ -148,6 +150,7 @@ def _make_projection_sender() -> Sender: self_endpoint="tcp://decode:0", cp_size=1, cp_rank=0, + page_table=None, ) extractor = MagicMock() @@ -158,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( @@ -174,20 +190,31 @@ 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, page_table=None) 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 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.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 + sender._describer = _SlotAddressedDescriber() return sender @@ -302,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 4c80b135ed7b..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,8 @@ 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 from tensorrt_llm._torch.disaggregation.native.transfer import RecvReqInfo, Sender @@ -160,21 +162,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 +212,8 @@ 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 + sender._describer = DefaultCacheDescriber() task = SimpleNamespace( _chunk=source_slice, _prompt_len=length, diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ec0a638524cb..2e5685a57cce 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -346,6 +346,7 @@ def test_v1_dsa_indexer_replicated_transfer_across_pp( """ import torch + 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 @@ -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..a441f53e181c 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -55,6 +55,7 @@ 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, @@ -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_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 88d4623289d8..cfb87bc65097 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,9 @@ 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_peer.py b/tests/unittest/disaggregated/test_peer.py index 7ff5f5e72e7c..1477375efc23 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -16,12 +16,16 @@ 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, 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,7 +39,7 @@ MambaHeadMatchMapper, MambaPolicy, ) -from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar +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,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) + 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 +257,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 +285,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 +314,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 +343,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 +356,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 +369,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 +396,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 +411,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 +573,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 +602,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 +634,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 +659,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 +693,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 +727,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 +841,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 +876,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 +1052,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 +1080,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 +1104,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 +1156,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 +1175,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 +1209,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 +1320,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 +1343,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 +1364,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..3b316bc9cf66 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: @@ -34,8 +34,8 @@ import numpy as np import pytest +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.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 ( @@ -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(): diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index a1d25e01c3f3..5008d28e77e0 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -203,3 +203,54 @@ 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", 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(): + """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", + "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_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index cfa6213af139..1e4471840186 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -29,6 +29,7 @@ import tensorrt_llm._torch.disaggregation.transceiver as transceiver_mod from tensorrt_llm import DisaggregatedParams from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange +from tensorrt_llm._torch.disaggregation.base.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 ( @@ -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 new file mode 100644 index 000000000000..76327d3f9161 --- /dev/null +++ b/tests/unittest/disaggregated/test_transfer_planner.py @@ -0,0 +1,462 @@ +# 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 ( + CacheDescriber, + LayerGroupRows, + 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] + with pytest.raises(TypeError): + CacheDescriber() # type: ignore[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 ( # noqa: E402 + CacheKind, + Chunk, + RecvReqInfo, + TokenRange, +) +from tensorrt_llm._torch.disaggregation.native.transfer import KVSendTask, Sender # noqa: E402 + + +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 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 + + +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, + dp_rank=0, + device_id=1, + cp_size=1, + self_endpoint="tcp://gen:1", + 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, 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 + + +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, + excluded_pool_views={(0, 1)}, + ), + 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_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)) + describer = _ToyDescriber() + sender = _sender_with(planner, describer) + 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" + # 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(): + 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 + + +# --------------------------------------------------------------------------- # +# 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 + # 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 + + +# --------------------------------------------------------------------------- # +# 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_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 = {} + 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 + + +# --------------------------------------------------------------------------- # +# 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(): + """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. + """ + 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_rows", "_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 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( + 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): + """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 + + 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): + """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 + + 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 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