nixl refactor [2/N]: unify TpKVTopology + HeteroTPTransferConfig into TransferTopology (#39529)

Signed-off-by: Zhanqiu Hu <zhu@redhat.com>

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -631,7 +631,7 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
         mock_thread.return_value.is_alive.return_value = False
 
         worker.use_mla = True
-        worker.kv_topo.is_mla = True
+        worker.transfer_topo.is_mla = True
 
         # MLA cache tensor: shape[-2] is the block size.
         mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16)
@@ -692,9 +692,9 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
         # Override TP rank/size to simulate P TP=2
         prefill_worker.tp_rank = P_TP_RANK
         prefill_worker.tp_size = P_TP_SIZE
-        # Update shared dict so kv_topo sees correct TP size
         prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE
-        prefill_worker.kv_topo.tp_rank = P_TP_RANK
+        prefill_worker.transfer_topo.tp_rank = P_TP_RANK
+        prefill_worker.transfer_topo.tp_size = P_TP_SIZE
 
         prefill_worker.kv_caches_base_addr = [0x1000]
         prefill_worker.block_len_per_layer = [local_block_len]
@@ -714,7 +714,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
         send_meta.ready.set()
 
         # Compute target D ranks using the production code path
-        target_d_ranks = prefill_worker.kv_topo.get_target_remote_ranks(d_tp_size)
+        target_d_ranks = prefill_worker.transfer_topo.handshake_target_ranks(d_tp_size)
 
         mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
         mock_socket.send_multipart = AsyncMock()

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -21,7 +21,7 @@ from vllm import LLM
 from vllm.config import KVTransferConfig, set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     KVOutputAggregator,
-    TpKVTopology,
+    TransferTopology,
     get_current_attn_backend,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1 import nixl
@@ -463,19 +463,20 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         test_shape = self.attn_backends[0].get_kv_cache_shape(
             num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
         )
-        self.kv_topo = TpKVTopology(
+        self.transfer_topo = TransferTopology(
             tp_rank=self.tp_rank,
+            tp_size=self.world_size,
+            block_size=self.block_size,
             engine_id=self.engine_id,
-            remote_tp_size=self._tp_size,  # shared state
-            remote_block_size=self._block_size,  # shared state
             is_mla=self.use_mla,
+            is_mamba=False,
             total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
             attn_backends=self.attn_backends,
             tensor_shape=test_shape,
         )
 
         self.compat_hash = compute_nixl_compatibility_hash(
-            self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
+            self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks
         )
 
     def _nixl_handshake(
@@ -496,7 +497,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         # Adjust remote block length metadata to satisfy heterogeneous TP
         # invariants enforced during handshake validation.
         remote_block_lens = list(self.block_len_per_layer)
-        tp_ratio = self.kv_topo.tp_ratio(remote_tp_size)
+        tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
         if remote_tp_size > self.world_size:
             # P TP > D TP case, block_len of remote is smaller
             remote_block_lens = [
@@ -731,8 +732,9 @@ class TestNixlHandshake:
             assert set(remote_agents.keys()) == set(range(tp_ratio))
 
             remote_engine_id = worker.REMOTE_ENGINE_ID
-            assert worker._tp_size[remote_engine_id] == remote_tp_size
-            assert -tp_ratio == worker.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
+            remote_info = worker.transfer_topo.get_engine_info(remote_engine_id)
+            assert remote_info.remote_tp_size == remote_tp_size
+            assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size)
             # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks
             assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio
             assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio
@@ -796,7 +798,7 @@ class TestNixlHandshake:
             (conn_p0.connector_worker, conn_p1.connector_worker)
         ):
             worker.world_size = p_tp_size
-            worker.kv_topo.remote_tp_size = {worker.engine_id: p_tp_size}
+            worker.transfer_topo.tp_size = p_tp_size
             worker.tp_rank = rank
             worker.use_mla = True
 
@@ -2337,7 +2339,7 @@ def test_compatibility_hash_validation(
         remote_hash = compute_nixl_compatibility_hash(
             remote_vllm_config,
             decode_worker.backend_name,
-            decode_worker.kv_topo.cross_layers_blocks,
+            decode_worker.transfer_topo.cross_layers_blocks,
         )
 
     prefill_block_size = config_overrides.get("block_size", 16)
@@ -2424,12 +2426,13 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     test_shape = backend.get_kv_cache_shape(
         num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
     )
-    decode_worker.kv_topo = TpKVTopology(
+    decode_worker.transfer_topo = TransferTopology(
         tp_rank=decode_worker.tp_rank,
+        tp_size=decode_worker.world_size,
+        block_size=decode_worker.block_size,
         engine_id=decode_worker.engine_id,
-        remote_tp_size=decode_worker._tp_size,  # shared state
-        remote_block_size=decode_worker._block_size,  # shared state
         is_mla=decode_worker.use_mla,
+        is_mamba=False,
         total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
         attn_backends=[backend],
         tensor_shape=test_shape,
@@ -2438,7 +2441,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     decode_worker.compat_hash = compute_nixl_compatibility_hash(
         decode_worker.vllm_config,
         decode_worker.backend_name,
-        decode_worker.kv_topo.cross_layers_blocks,
+        decode_worker.transfer_topo.cross_layers_blocks,
     )
 
     if error_scenario == "handshake_decode_error":

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -152,13 +152,14 @@ def test_read_blocks_for_req_expands_remote_ids(
 
     remote_engine_id = "remote-engine"
     if has_mamba:
-        worker._mamba_phys_ratio = {remote_engine_id: remote_ratio}
+        worker._physical_blocks_per_logical = {remote_engine_id: remote_ratio}
 
-    # Mock kv_topo: empty remote ranks skips the transfer machinery entirely,
-    # isolating the block-ID expansion logic.
-    worker.kv_topo = MagicMock()
-    worker.kv_topo.get_target_remote_ranks_from_engine_id.return_value = []
-    worker.kv_topo.tp_ratio_from_engine_id.return_value = 1
+    # Mock transfer_topo: empty remote ranks skips the transfer machinery
+    # entirely, isolating the block-ID expansion logic.
+    worker.transfer_topo = MagicMock()
+    worker.transfer_topo.target_remote_ranks.return_value = []
+    worker.transfer_topo.get_engine_info.return_value = MagicMock(remote_tp_size=1)
+    worker.transfer_topo.tp_ratio.return_value = 1
 
     metadata = NixlConnectorMetadata()
     metadata.add_new_req_to_recv(
@@ -317,7 +318,7 @@ def test_get_block_descs_ids_hybrid_ssm():
     worker._has_mamba = True
     worker._is_mamba_group = [False, True]
     worker._physical_blocks_per_logical_kv_block = 1
-    worker._mamba_phys_ratio = {engine_id: 1}
+    worker._physical_blocks_per_logical = {engine_id: 1}
     worker.block_len_per_layer = [100]
     # num_descs = num_regions * num_blocks (no blocks_first doubling)
     worker.num_descs = 2 * num_blocks
@@ -355,7 +356,7 @@ def test_get_block_descs_ids_kernel_block_mismatch():
     worker._has_mamba = True
     worker._is_mamba_group = [False, True]
     worker._physical_blocks_per_logical_kv_block = ratio
-    worker._mamba_phys_ratio = {engine_id: ratio}
+    worker._physical_blocks_per_logical = {engine_id: ratio}
     worker.block_len_per_layer = [100]
     worker.num_descs = 2 * num_blocks  # 800
 
@@ -532,15 +533,15 @@ def test_has_mamba_init(
         ((9216, 524288), 4096, 131),
     ],
 )
-def test_compute_mamba_phys_ratio(ssm_sizes, block_len, expected_ratio):
-    """Verify that compute_mamba_phys_ratio is TP-dependent.
+def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_ratio):
+    """Verify that compute_physical_blocks_per_logical is TP-dependent.
 
     With dimension-sharded Mamba state, the ratio differs across TP sizes
     (e.g. TP=1 → 261, TP=4 → 131 for Nemotron 30B). This is why
-    _mamba_phys_ratio must be stored per-engine.
+    _physical_blocks_per_logical must be stored per-engine.
     """
     from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
-        compute_mamba_phys_ratio,
+        compute_physical_blocks_per_logical,
     )
 
-    assert compute_mamba_phys_ratio(ssm_sizes, block_len) == expected_ratio
+    assert compute_physical_blocks_per_logical(ssm_sizes, block_len) == expected_ratio

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -5,7 +5,7 @@ KV cache helper for store.
 """
 
 from collections.abc import Iterator
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 import torch
@@ -319,31 +319,139 @@ def yield_req_data(
     )
 
 
-@dataclass
-class TpKVTopology:
+def get_current_attn_backends(
+    vllm_config: VllmConfig, layer_names: list[str] | None = None
+) -> list[type[AttentionBackend]]:
+    """Get all distinct attention backends for the given layers.
+
+    Args:
+        vllm_config: The current vLLM configuration.
+        layer_names: Optional list of layer names to scope the lookup.
+            When None, all attention layers are considered.
+
+    Returns:
+        Deduplicated list of attention backend classes.
+    """
+    layer_type = cast(type[Any], AttentionLayerBase)
+    layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
+    if layers:
+        seen: dict[str, type[AttentionBackend]] = {}
+        for layer in layers.values():
+            backend = layer.get_attn_backend()
+            seen[backend.full_cls_name()] = backend
+        return list(seen.values())
+
+    # Fallback for tests, when static_forward_context is empty.
+    logger.debug(
+        "No layers found in the vLLM config. Falling back to default attention backend."
+    )
+    from vllm.v1.attention.selector import get_attn_backend
+
+    return [
+        get_attn_backend(
+            head_size=vllm_config.model_config.get_head_size(),
+            dtype=vllm_config.model_config.dtype,
+            kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+            use_mla=vllm_config.model_config.use_mla,
+        )
+    ]
+
+
+def get_current_attn_backend(
+    vllm_config: VllmConfig, layer_names: list[str] | None = None
+) -> type[AttentionBackend]:
+    """Get the first attention backend for the given layers."""
+    return get_current_attn_backends(vllm_config, layer_names)[0]
+
+
+# ---- Per-engine transfer info ----
+
+
+@dataclass(frozen=True)
+class EngineTransferInfo:
+    """Common per-remote-engine transfer state, computed at handshake.
+
+    Stored per ``engine_id`` inside ``TransferTopology._engines``.
     """
-    Helper class for tensor parallel and KV topology information for
-    mapping between local and remote TP workers.
+
+    remote_tp_size: int
+
+    remote_block_len: int
+    """Block length (bytes)"""
+
+    remote_block_size: int
+    """Tokens per block."""
+
+    remote_physical_blocks_per_logical: int
+    """Physical blocks per logical block."""
+
+
+@dataclass(frozen=True)
+class MambaEngineTransferInfo(EngineTransferInfo):
+    """Extends ``EngineTransferInfo`` with Mamba-hybrid transfer geometry.
+
+    For hybrid SSM+Attention models, FA and Mamba layers may require
+    different numbers of reads from different remote ranks.  This
+    dataclass captures that per-engine transfer plan.
     """
 
+    remote_fa_source_ranks: tuple[int, ...]
+    """Remote ranks carrying unique FA heads for this local rank."""
+
+    remote_all_source_ranks: tuple[int, ...]
+    """All remote ranks this local rank reads from (FA + Mamba)."""
+
+    remote_num_fa_reads: int
+    """Number of distinct remote ranks needed for FA data."""
+
+    remote_num_mamba_reads: int
+    """Number of distinct remote ranks needed for Mamba data."""
+
+    remote_fa_descriptor_bytes: int
+    """Byte size of one FA K (or V) descriptor entry."""
+
+    is_remote_replicated: bool
+    """Whether the remote engine has replicated KV heads
+    (remote_tp_size > total_num_kv_heads)."""
+
+    remote_physical_heads: int
+    """Physical KV heads stored per remote rank."""
+
+
+# ---- Transfer topology ----
+
+
+@dataclass
+class TransferTopology:
+    """Single source of truth for local TP identity and per-engine remote info."""
+
     tp_rank: int
-    remote_tp_size: dict[EngineId, int]
+    tp_size: int
+    block_size: int
+    engine_id: EngineId
     is_mla: bool
+    is_mamba: bool
     total_num_kv_heads: int
     attn_backends: list[type[AttentionBackend]]
-    engine_id: EngineId
-    remote_block_size: dict[EngineId, int]
     tensor_shape: torch.Size | None = None
-    is_mamba: bool = False
 
     def __post_init__(self):
+        self.local_physical_heads = max(1, self.total_num_kv_heads // self.tp_size)
+
+        self._engines: dict[EngineId, EngineTransferInfo] = {}
+        self._fa_source_sets: dict[EngineId, frozenset[int]] = {}
+        self._fa_source_indices: dict[EngineId, dict[int, int]] = {}
+
         # Figure out whether the first dimension of the cache is K/V
-        # or num_blocks. This is used to register the memory regions correctly.
+        # or num_blocks.
         attn_backend = self.attn_backends[0]
         if not self.is_mamba:
             _MOCK_BLOCK_SIZE = 16
             kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape(
-                num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1
+                num_blocks=1,
+                block_size=_MOCK_BLOCK_SIZE,
+                num_kv_heads=1,
+                head_size=1,
             )
             logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
         # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
@@ -358,11 +466,9 @@ class TpKVTopology:
             self._cross_layers_blocks = (
                 len(self.tensor_shape) == len(kv_cache_shape) + 1
             )
-            self.tensor_shape: torch.Size
 
         if self._cross_layers_blocks:
             logger.debug("Using cross-layer KV cache")
-            # prepend layers dimension
             _MOCK_NUM_LAYERS = 80
             kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape
             try:
@@ -372,15 +478,81 @@ class TpKVTopology:
             except (AttributeError, NotImplementedError):
                 assert self.tensor_shape is not None
                 kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
-
-            # In case of cross layers permute kv_cache_shape according to
-            # stride_order to retrieve physical position of block_size
             kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
 
+    # ============================================================
+    # Engine registration
+    # ============================================================
+
+    def register_remote_engine(
+        self,
+        remote_engine_id: EngineId,
+        remote_tp_size: int,
+        remote_block_size: int,
+        remote_block_len: int,
+        remote_physical_blocks_per_logical: int,
+        *,
+        local_block_len: int = 0,
+    ) -> EngineTransferInfo:
+        """Register a remote engine, unifying worker dicts state.
+
+        Only remote engines should be registered here — the local engine's
+        identity (tp_size, block_size, etc.) is set via ``__init__`` params.
+
+        For Mamba models, also computes the Mamba transfer plan and
+        builds the FA source lookup caches.
+
+        Args:
+            local_block_len: Local representative block_len (bytes).
+                Required for Mamba models to compute ``fa_descriptor_bytes``.
+        """
+        assert remote_engine_id != self.engine_id, (
+            f"Cannot register local engine {self.engine_id} as remote. "
+            f"Local identity is set via __init__ params."
+        )
+        if remote_engine_id in self._engines:
+            return self._engines[remote_engine_id]
+        info: EngineTransferInfo
+        if self.is_mamba:
+            info = self._build_mamba_info(
+                remote_tp_size=remote_tp_size,
+                remote_block_size=remote_block_size,
+                remote_block_len=remote_block_len,
+                remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),
+                local_block_len=local_block_len,
+            )
+            assert isinstance(info, MambaEngineTransferInfo)
+            self._fa_source_sets[remote_engine_id] = frozenset(
+                info.remote_fa_source_ranks
+            )
+            self._fa_source_indices[remote_engine_id] = {
+                r: i for i, r in enumerate(info.remote_fa_source_ranks)
+            }
+        else:
+            info = EngineTransferInfo(
+                remote_tp_size=remote_tp_size,
+                remote_block_len=remote_block_len,
+                remote_block_size=remote_block_size,
+                remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),
+            )
+        self._engines[remote_engine_id] = info
+        return info
+
+    def get_engine_info(self, remote_engine_id: EngineId) -> EngineTransferInfo:
+        return self._engines[remote_engine_id]
+
+    # ============================================================
+    # Layout properties
+    # ============================================================
+
     @property
     def is_kv_layout_blocks_first(self) -> bool:
         return self._is_kv_layout_blocks_first
 
+    @property
+    def cross_layers_blocks(self) -> bool:
+        return self._cross_layers_blocks
+
     @property
     def split_k_and_v(self) -> bool:
         # Whether to register regions for K and V separately (when present).
@@ -388,29 +560,16 @@ class TpKVTopology:
             self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first
         )
 
-    @property
-    def tp_size(self) -> int:
-        return self.remote_tp_size[self.engine_id]
+    # ============================================================
+    # Common methods
+    # ============================================================
 
-    @property
-    def block_size(self) -> int:
-        return self.remote_block_size[self.engine_id]
+    def tp_ratio(self, remote_tp_size: int) -> int:
+        """Calculate the tensor parallel ratio between local and remote TP.
 
-    @property
-    def cross_layers_blocks(self) -> bool:
-        return self._cross_layers_blocks
-
-    def tp_ratio(
-        self,
-        remote_tp_size: int,
-    ) -> int:
-        """
-        Calculate the tensor parallel ratio between local and remote TP.
-        We can think of it as the number of local TP workers-per-remote TP
-        workers. Local workers will read from the same remote TP worker in
-        groups of size `tp_ratio`.If remote tp_size > local tp_size, the
-        ratio is flipped (remote_size/local_size) and the returned value is
-        negative.
+        Positive when local_tp >= remote_tp (local workers read from the
+        same remote worker in groups of size ``tp_ratio``).  Negative when
+        remote_tp > local_tp (ratio is flipped).
         """
         if self.tp_size >= remote_tp_size:
             assert self.tp_size % remote_tp_size == 0, (
@@ -418,78 +577,65 @@ class TpKVTopology:
                 f"by remote tensor parallel size {remote_tp_size}."
             )
             return self.tp_size // remote_tp_size
-
         assert remote_tp_size % self.tp_size == 0, (
             f"Remote tensor parallel size {remote_tp_size} is not divisible "
             f"by local tensor parallel size {self.tp_size}."
         )
-        # P TP > D TP case, return the ratio as negative
-        return -remote_tp_size // self.tp_size
+        return -(remote_tp_size // self.tp_size)
 
-    def block_size_ratio(
-        self,
-        remote_block_size: int,
-    ) -> int:
-        """
-        Calculate the block size ratio between local and remote TP.
-        """
+    def block_size_ratio(self, remote_block_size: int) -> int:
+        """Calculate the block size ratio between local and remote."""
         assert self.block_size % remote_block_size == 0, (
             f"Local block size {self.block_size} is not divisible "
             f"by remote block size {remote_block_size} or vice versa."
         )
         return self.block_size // remote_block_size
 
-    def tp_ratio_from_engine_id(
-        self,
-        remote_engine_id: EngineId,
-    ) -> int:
-        remote_tp_size = self.remote_tp_size[remote_engine_id]
-        return self.tp_ratio(remote_tp_size)
-
-    def block_size_ratio_from_engine_id(
-        self,
-        remote_engine_id: EngineId,
-    ) -> int:
-        remote_block_size = self.remote_block_size[remote_engine_id]
-        return self.block_size_ratio(remote_block_size)
-
-    def is_kv_replicated(self, engine_id: EngineId) -> bool:
-        """
-        Whether the KV cache is replicated across TP workers due to the
+    def is_kv_replicated(self, remote_engine_id: EngineId) -> bool:
+        """Whether the KV cache is replicated across TP workers due to the
         number of TP workers being greater than the number of KV heads.
-        When they are equal, each TP rank still owns one distinct KV head,
-        so this is not considered replication.
         """
-        tp_size = self.remote_tp_size[engine_id]
-        return tp_size > self.total_num_kv_heads
+        return self._engines[remote_engine_id].remote_tp_size > self.total_num_kv_heads
 
     def replicates_kv_cache(self, remote_engine_id: EngineId) -> bool:
         # MLA is always replicated as the hidden dim can't be split.
         return self.is_mla or self.is_kv_replicated(remote_engine_id)
 
-    def get_target_remote_ranks(
-        self,
-        remote_tp_size: int,
-    ) -> list[int]:
-        """
-        Get the remote TP rank (on P) that the current local TP rank
-        (on D) will read from. When remote tp_size > local tp_size, we
-        read from multiple remote ranks.
+    @property
+    def local_replicates_kv_cache(self) -> bool:
+        """Whether the local engine's KV cache is replicated."""
+        return self.is_mla or self.tp_size > self.total_num_kv_heads
+
+    def handshake_target_ranks(self, remote_tp_size: int) -> list[int]:
+        """Pre-registration: compute which remote TP ranks to handshake with.
+
+        Pure math based on local/remote TP sizes — does not require
+        the remote engine to be registered yet.
         """
         tp_ratio = self.tp_ratio(remote_tp_size)
         if tp_ratio > 0:
             return [self.tp_rank // tp_ratio]
+        abs_ratio = -tp_ratio
+        return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
 
-        # P TP > D TP case, D reads from |tp_ratio| remote workers.
-        tp_ratio = -tp_ratio
-        return [self.tp_rank * tp_ratio + i for i in range(tp_ratio)]
+    def target_remote_ranks(self, remote_engine_id: EngineId) -> list[int]:
+        """Get the remote TP rank(s) that the current local TP rank will
+        read from.  When remote tp_size > local tp_size, reads from
+        multiple remote ranks.
 
-    def get_target_remote_ranks_from_engine_id(
-        self,
-        remote_engine_id: EngineId,
-    ) -> list[int]:
-        remote_tp_size = self.remote_tp_size[remote_engine_id]
-        return self.get_target_remote_ranks(remote_tp_size)
+        For Mamba models, returns the precomputed ``all_source_ranks``
+        (FA + Mamba union).
+        """
+        info = self._engines[remote_engine_id]
+        if isinstance(info, MambaEngineTransferInfo):
+            return list(info.remote_all_source_ranks)
+
+        tp_ratio = self.tp_ratio(info.remote_tp_size)
+        if tp_ratio > 0:
+            return [self.tp_rank // tp_ratio]
+        # remote TP > local TP: read from |tp_ratio| remote workers
+        abs_ratio = -tp_ratio
+        return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
 
     def get_transfer_cache_regions(
         self, cache: torch.Tensor, layer_spec: "KVCacheSpec"
@@ -498,331 +644,139 @@ class TpKVTopology:
         also accounting for hybrid SSM models specificities.
         """
         if isinstance(layer_spec, MambaSpec):
-            # Register the whole kv cache shared tensor, including SSM/Conv. This is
-            # similar to FI with the difference that SSM/Conv have different sizes
+            # Register the whole kv cache shared tensor, including
+            # SSM/Conv.
             conv, ssm = cache
             return [conv]
 
-        # Check may be hacky but it's matching `_update_hybrid_attention_mamba_layout`.
+        # Check may be hacky but it's matching
+        # `_update_hybrid_attention_mamba_layout`.
         if self.is_mamba and cache.shape[0] == 2:
-            # When MAMBA is present, all backends are blocks first, so that blocks
-            # can be shared between attention layers and mamba layers. Runner
-            # `_update_hybrid_attention_mamba_layout` already adjusted strides
-            # for FlashAttn-like backends so its num_blocks first.
-            # Swap [2<>num_blocks] dims to get required layout for hybrid SSM.
+            # When MAMBA is present, all backends are blocks first, so
+            # that blocks can be shared between attention layers and mamba
+            # layers.  Runner already adjusted strides for FlashAttn-like
+            # backends so its num_blocks first.
+            # Swap [2<>num_blocks] dims for hybrid SSM layout.
             cache = cache.transpose(0, 1)
 
         # Regular case: backends like FA register K/V in separate regions
         return cache if self.split_k_and_v else [cache]
 
+    # ============================================================
+    # Mamba-specific methods
+    # ============================================================
 
-# ---- Mamba-HMA hetero-TP transfer config ----
-#
-# Key insight: with hetero-TP (P_TP > D_TP), FA KV cache may be
-# replicated across P ranks (when P_TP > num_kv_heads), but Mamba
-# conv/SSM state is almost always uniquely sharded per P rank.  So the
-# number of P ranks D must read from can differ between FA and Mamba,
-# and they must be handled separately.
-
-
-def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range:
-    """Physical KV head range stored in a rank's KV cache tensor.
-
-    When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank.
-    When ``tp_size > num_heads``: 1 physical head per rank.  Heads are
-    distributed **contiguously** (matching vLLM's GQA weight partitioning):
-    consecutive ranks share a head before moving to the next one.
-    """
-    if tp_size <= num_heads:
-        assert num_heads % tp_size == 0
-        per_rank = num_heads // tp_size
-        return range(rank * per_rank, (rank + 1) * per_rank)
-    else:
-        h = rank * num_heads // tp_size
-        return range(h, h + 1)
-
-
-def _range_overlap(a: range, b: range) -> range:
-    start = max(a.start, b.start)
-    stop = min(a.stop, b.stop)
-    return range(start, max(start, stop))
-
+    def should_skip_fa(self, remote_engine_id: EngineId, remote_rank: int) -> bool:
+        """Whether to skip FA groups for this remote rank (mamba-only)."""
+        return remote_rank not in self._fa_source_sets[remote_engine_id]
 
-@dataclass
-class HeteroTPTransferConfig:
-    """Precomputed transfer plan for one (D rank, P engine) pair.
-
-    Currently only instantiated for Mamba-HMA (hybrid SSM+Attention) models
-    where FA and mamba require different splitting factors.  Could be extended
-    to other model types that need non-uniform hetero-TP transfer sizing.
-
-    All descriptor sizes are computed here.  The guarantee is:
-        local_entry_size == remote_entry_size   (for NIXL)
-
-    Attributes that start with ``fa_`` concern FlashAttention KV cache.
-    Attributes that start with ``mamba_`` concern Mamba conv/SSM state.
-    """
-
-    # ---- Input parameters (from handshake) ----
-    tp_ratio: int
-    K: int  # total_num_kv_heads (before TP sharding)
-    d_tp: int  # D engine's tensor_parallel_size
-    p_tp: int  # P engine's tensor_parallel_size
-    d_rank: int  # this D worker's TP rank
-    use_mla: bool
-
-    # Per-layer block lengths (bytes, K+V combined for blocks_first).
-    # Uniform across layers for current models.
-    d_block_len: int  # D's block_len_per_layer (representative)
-    p_block_len: int  # P's block_len_per_layer (from handshake)
-    is_blocks_first: bool  # kv_topo.is_kv_layout_blocks_first
-
-    # ---- Derived: computed in __post_init__ ----
-    #
-    # Physical heads per rank (what the KV tensor actually stores)
-    d_physical_heads: int = field(init=False)
-    p_physical_heads: int = field(init=False)
-
-    # How many distinct P ranks D needs for FA data
-    physical_fa_num_reads: int = field(init=False)
-
-    # Which P ranks contribute unique FA heads (ordered by head index)
-    fa_read_targets: list[int] = field(init=False)
-
-    # All P ranks needed for mamba (always abs_tp for tp_ratio < 0)
-    mamba_num_reads: int = field(init=False)
-
-    # All P ranks this D rank communicates with (FA ∪ mamba)
-    transfer_targets: list[int] = field(init=False)
-
-    # FA descriptor entry size (K or V side, for blocks_first layout)
-    # Guaranteed: fa_entry_size is the SAME for local handle AND remote desc.
-    fa_entry_size: int = field(init=False)
-
-    # Replication flags
-    is_d_replicated: bool = field(init=False)
-    is_p_replicated: bool = field(init=False)
-
-    # Pre-built set for fast lookup
-    _fa_target_set: frozenset[int] = field(init=False, repr=False)
-    # Map: P rank → index in fa_read_targets (for head slot offset)
-    _fa_target_index: dict[int, int] = field(init=False, repr=False)
-
-    def __post_init__(self) -> None:
-        K = self.K
-        self.is_d_replicated = self.d_tp > K
-        self.is_p_replicated = self.p_tp > K
-
-        self.d_physical_heads = max(1, K // self.d_tp)
-        self.p_physical_heads = max(1, K // self.p_tp)
-
-        abs_tp = -self.tp_ratio if self.tp_ratio < 0 else 1
-
-        # ---- Mamba range (computed first so FA can prefer ranks in it) ----
-        mamba_range: range | None = None
-        if self.tp_ratio < 0:
-            mamba_range = range(self.d_rank * abs_tp, (self.d_rank + 1) * abs_tp)
-
-        # ---- FA read targets ----
-        if self.use_mla or self.tp_ratio >= 0:
-            self.physical_fa_num_reads = 1
-            self.fa_read_targets = (
-                [0]
-                if self.use_mla
-                # Must match kv_topo.get_target_remote_ranks (d_rank // tp_ratio).
-                else [
-                    self.d_rank // self.tp_ratio if self.tp_ratio > 0 else self.d_rank
-                ]
-            )
-        else:
-            d_needs = _physical_head_range(self.d_tp, K, self.d_rank)
-            # When mamba range exists, prefer P ranks within it so that
-            # FA targets are a subset of mamba transfer_targets (avoids
-            # orphaned FA targets outside the transfer loop).
-            search_range = mamba_range if mamba_range is not None else range(self.p_tp)
-            seen: set[tuple[int, int]] = set()
-            targets: list[int] = []
-            for p in search_range:
-                p_has = _physical_head_range(self.p_tp, K, p)
-                ov = _range_overlap(d_needs, p_has)
-                if len(ov) > 0:
-                    key = (ov.start, ov.stop)
-                    if key not in seen:
-                        seen.add(key)
-                        targets.append(p)
-            if not targets:
-                # Fallback: search globally (should not happen in practice)
-                for p in range(self.p_tp):
-                    p_has = _physical_head_range(self.p_tp, K, p)
-                    ov = _range_overlap(d_needs, p_has)
-                    if len(ov) > 0:
-                        key = (ov.start, ov.stop)
-                        if key not in seen:
-                            seen.add(key)
-                            targets.append(p)
-            self.fa_read_targets = targets
-            self.physical_fa_num_reads = len(targets)
-
-        self._fa_target_set = frozenset(self.fa_read_targets)
-        self._fa_target_index = {r: i for i, r in enumerate(self.fa_read_targets)}
-
-        # ---- Mamba targets ----
-        if mamba_range is not None and abs_tp > self.physical_fa_num_reads:
-            self.mamba_num_reads = abs_tp
-            self.transfer_targets = list(mamba_range)
-        else:
-            self.mamba_num_reads = self.physical_fa_num_reads
-            self.transfer_targets = list(self.fa_read_targets)
-
-        # ---- FA entry size ----
-        # For blocks_first: block_len_per_layer includes K+V; // 2 gives K (or V).
-        # Use min(D, P) because D indexes into P when tp_ratio > 0,
-        # and P is the natural unit when tp_ratio < 0.
-        effective_block_len = min(self.d_block_len, self.p_block_len)
-        if self.is_blocks_first:
-            self.fa_entry_size = effective_block_len // 2
-        else:
-            self.fa_entry_size = effective_block_len
-
-        self._validate()
-
-    def _validate(self) -> None:
-        """Cross-check internal consistency."""
-        if self.is_d_replicated and self.is_p_replicated and self.tp_ratio > 0:
-            logger.info(
-                "Both-replicated hetero-TP: D_TP=%d > P_TP=%d > K=%d. "
-                "Using d_rank // tp_ratio routing with relative head offset.",
-                self.d_tp,
-                self.p_tp,
-                self.K,
-            )
-
-        # FA targets must be a subset of transfer_targets
-        tt_set = set(self.transfer_targets)
-        for t in self.fa_read_targets:
-            if t not in tt_set:
-                logger.error(
-                    "FA target P rank %d is NOT in transfer_targets %s. "
-                    "This will cause missed FA reads!",
-                    t,
-                    self.transfer_targets,
-                )
-
-        # For tp_ratio < 0 with blocks_first: D_K_half / reads should == P_K_half
-        if (
-            self.is_blocks_first
-            and self.tp_ratio < 0
-            and self.physical_fa_num_reads > 0
-        ):
-            d_k_half = self.d_block_len // 2
-            p_k_half = self.p_block_len // 2
-            expected_local = d_k_half // self.physical_fa_num_reads
-            if expected_local != p_k_half:
-                logger.warning(
-                    "FA size mismatch: D_K_half=%d / reads=%d = %d, "
-                    "but P_K_half=%d.  This may indicate a head count or "
-                    "Mamba-HMA inflation inconsistency.",
-                    d_k_half,
-                    self.physical_fa_num_reads,
-                    expected_local,
-                    p_k_half,
-                )
+    def fa_head_slot(self, remote_engine_id: EngineId, remote_rank: int) -> int:
+        """Index into local FA block for this remote rank's head data.
 
-    # ---- Query methods ----
-
-    def should_skip_fa(self, p_rank: int) -> bool:
-        """Whether to skip FA groups for this P rank (mamba-only transfer)."""
-        return p_rank not in self._fa_target_set
-
-    def fa_head_slot(self, p_rank: int) -> int:
-        """Index into D's FA block for this P rank's head data.
-
-        For P ranks in fa_read_targets, returns 0, 1, ..., reads-1.
-        For P ranks NOT in fa_read_targets (replicated duplicates),
-        returns the slot of the matching FA target with the same head.
+        For remote ranks in ``fa_source_ranks``, returns 0, 1, …, reads-1.
+        For ranks NOT in ``fa_source_ranks`` (replicated duplicates),
+        returns the slot of the matching source rank with the same head.
         """
-        if p_rank in self._fa_target_index:
-            return self._fa_target_index[p_rank]
-        # Duplicate head: find which fa_target has the same physical head
-        p_head = _physical_head_range(self.p_tp, self.K, p_rank)
-        for target in self.fa_read_targets:
-            t_head = _physical_head_range(self.p_tp, self.K, target)
-            if _range_overlap(p_head, t_head):
-                return self._fa_target_index[target]
-        return 0  # fallback
-
-    def fa_rank_offset(self, remote_kv_block_len: int) -> int:
-        """Byte offset into P's FA block for this D rank.
-
-        When D is replicated (D_TP > K), multiple D ranks share a head.
-        Computes offset *relative to the target P rank's first head*
-        so it works regardless of how many heads P has.
-        When neither side replicates, falls back to tp_rank % tp_ratio.
-        Returns 0 when D does not index into P's block.
+        fa_index = self._fa_source_indices[remote_engine_id]
+        if remote_rank in fa_index:
+            return fa_index[remote_rank]
+        mamba_info = self._engines[remote_engine_id]
+        assert isinstance(mamba_info, MambaEngineTransferInfo)
+        K = self.total_num_kv_heads
+        remote_tp = mamba_info.remote_tp_size
+        r_head = self._physical_head_range(remote_tp, K, remote_rank)
+        for target in mamba_info.remote_fa_source_ranks:
+            t_head = self._physical_head_range(remote_tp, K, target)
+            if self._range_overlap(r_head, t_head):
+                return fa_index[target]
+        return 0
+
+    def fa_rank_offset(
+        self, remote_engine_id: EngineId, remote_kv_block_len: int
+    ) -> int:
+        """Byte offset into remote FA block for this local rank.
+
+        When local TP is replicated (local_tp > K), multiple local ranks
+        share a head.  Computes offset *relative to the target remote
+        rank's first head* so it works regardless of how many heads the
+        remote has.  Returns 0 when local does not index into remote.
         """
-        if self.use_mla or self.tp_ratio <= 0:
+        mamba_info = self._engines[remote_engine_id]
+        assert isinstance(mamba_info, MambaEngineTransferInfo)
+        tp_ratio = self.tp_ratio(mamba_info.remote_tp_size)
+        if self.is_mla or tp_ratio <= 0:
             return 0
-        if self.is_d_replicated:
-            d_head = self.d_rank * self.K // self.d_tp
-            p_rank = self.fa_read_targets[0]
-            p_start = p_rank * self.K // self.p_tp
-            return (d_head - p_start) * remote_kv_block_len
-        return self.d_rank % self.tp_ratio * remote_kv_block_len
-
-    @property
-    def needs_split_handles(self) -> bool:
-        """Whether per-P-rank split handles are needed.
+        K = self.total_num_kv_heads
+        is_local_replicated = self.tp_size > K
+        if is_local_replicated:
+            local_head = self.tp_rank * K // self.tp_size
+            p_rank = mamba_info.remote_fa_source_ranks[0]
+            p_start = p_rank * K // mamba_info.remote_tp_size
+            return (local_head - p_start) * remote_kv_block_len
+        return self.tp_rank % tp_ratio * remote_kv_block_len
+
+    def needs_split_handles(self, remote_engine_id: EngineId) -> bool:
+        """Whether per-remote-rank split handles are needed.
 
         True when FA and mamba have different read counts, requiring
         different splitting factors in the local handle.
         """
-        return self.tp_ratio < 0 and not self.use_mla and len(self.transfer_targets) > 1
+        mamba_info = self._engines[remote_engine_id]
+        assert isinstance(mamba_info, MambaEngineTransferInfo)
+        tp_ratio = self.tp_ratio(mamba_info.remote_tp_size)
+        return (
+            tp_ratio < 0
+            and not self.is_mla
+            and len(mamba_info.remote_all_source_ranks) > 1
+        )
 
     def compute_split_handle_data(
         self,
+        remote_engine_id: EngineId,
         src_blocks_data: list[tuple[int, int, int]],
         num_fa_descs: int,
         abs_tp: int,
     ) -> list[list[tuple[int, int, int]]]:
-        """Compute per-P-rank (addr, len, tp) triples for Mamba-HMA split handles.
+        """Per-remote-rank (addr, len, dev) triples for Mamba-HMA split
+        handles.
 
         FA descriptors (indices < num_fa_descs) are sliced by
-        ``physical_fa_num_reads``; mamba descriptors are sliced uniformly
+        ``remote_num_fa_reads``; mamba descriptors are sliced uniformly
         by ``abs_tp``.
-
-        Returns one list of triples per transfer target.
         """
+        mamba_info = self._engines[remote_engine_id]
+        assert isinstance(mamba_info, MambaEngineTransferInfo)
         all_handle_data: list[list[tuple[int, int, int]]] = []
-        for p_idx, p_rank in enumerate(self.transfer_targets):
+        for p_idx, p_rank in enumerate(mamba_info.remote_all_source_ranks):
             handle_data: list[tuple[int, int, int]] = []
-            skip_fa = self.should_skip_fa(p_rank)
-            fa_slot = self.fa_head_slot(p_rank) if not skip_fa else 0
-
-            for j, (addr, local_len, tp) in enumerate(src_blocks_data):
+            skip_fa = self.should_skip_fa(remote_engine_id, p_rank)
+            fa_slot = self.fa_head_slot(remote_engine_id, p_rank) if not skip_fa else 0
+            for j, (addr, local_len, dev) in enumerate(src_blocks_data):
                 if j < num_fa_descs:
-                    assert self.physical_fa_num_reads >= 1
-                    fa_chunk = local_len // self.physical_fa_num_reads
-                    handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, tp))
+                    assert mamba_info.remote_num_fa_reads >= 1
+                    fa_chunk = local_len // mamba_info.remote_num_fa_reads
+                    handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, dev))
                 else:
                     mamba_chunk = local_len // abs_tp
-                    handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, tp))
+                    handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, dev))
             all_handle_data.append(handle_data)
         return all_handle_data
 
     def filter_block_ids_for_rank(
         self,
+        remote_engine_id: EngineId,
         remote_rank: int,
         local_ids: BlockIds,
         remote_ids: BlockIds,
         is_mamba_group: list[bool],
     ) -> tuple[BlockIds, BlockIds]:
-        """Zero out FA groups for P ranks outside fa_read_targets.
+        """Zero out FA groups for remote ranks outside ``fa_source_ranks``.
 
         Returns (filtered_local_ids, filtered_remote_ids).  When the
-        remote rank carries FA data for this D rank, returns the inputs
-        unchanged.
+        remote rank carries FA data for this local rank, returns the
+        inputs unchanged.
         """
-        if not self.should_skip_fa(remote_rank):
+        if not self.should_skip_fa(remote_engine_id, remote_rank):
             return local_ids, remote_ids
         num_groups = len(local_ids)
         filtered_local: list[list[int]] = [
@@ -833,108 +787,184 @@ class HeteroTPTransferConfig:
         ]
         return filtered_local, filtered_remote
 
-    def describe(self) -> str:
-        """One-line summary for logging."""
-        return (
-            f"HeteroTPTransferConfig("
-            f"tp_ratio={self.tp_ratio}, K={self.K}, "
-            f"d_tp={self.d_tp}, p_tp={self.p_tp}, d_rank={self.d_rank}, "
-            f"physical_fa_reads={self.physical_fa_num_reads}, "
-            f"mamba_reads={self.mamba_num_reads}, "
-            f"fa_targets={self.fa_read_targets}, "
-            f"transfer_targets={self.transfer_targets}, "
-            f"fa_entry_size={self.fa_entry_size}, "
-            f"d_block_len={self.d_block_len}, p_block_len={self.p_block_len})"
+    def describe(self, remote_engine_id: EngineId) -> str:
+        """One-line summary of transfer config for logging."""
+        info = self._engines[remote_engine_id]
+        base = (
+            f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, "
+            f"K={self.total_num_kv_heads}, "
+            f"local_tp={self.tp_size}, "
+            f"remote_tp={info.remote_tp_size}, "
+            f"local_rank={self.tp_rank}, "
+            f"remote_block_len={info.remote_block_len}"
         )
+        if isinstance(info, MambaEngineTransferInfo):
+            return (
+                f"TransferTopology.mamba({base}, "
+                f"fa_reads={info.remote_num_fa_reads}, "
+                f"mamba_reads={info.remote_num_mamba_reads}, "
+                f"fa_sources={list(info.remote_fa_source_ranks)}, "
+                f"all_sources={list(info.remote_all_source_ranks)}, "
+                f"fa_desc_bytes={info.remote_fa_descriptor_bytes})"
+            )
+        return f"TransferTopology({base})"
+
+    # ============================================================
+    # Private helpers
+    # ============================================================
+    # Mamba-HMA hetero-TP transfer config:
+    # With hetero-TP (P_TP > D_TP), FA KV cache may be replicated across
+    # P ranks (when P_TP > num_kv_heads), but Mamba conv/SSM state is
+    # almost always uniquely sharded per P rank.  So the number of P
+    # ranks D must read from can differ between FA and Mamba, and they
+    # must be handled separately.
+
+    @staticmethod
+    def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range:
+        """Physical KV head range stored in a rank's KV cache tensor.
+
+        When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank.
+        When ``tp_size > num_heads``: 1 physical head per rank.  Heads are
+        distributed **contiguously** (matching vLLM's GQA weight partitioning):
+        consecutive ranks share a head before moving to the next one.
+        """
+        if tp_size <= num_heads:
+            assert num_heads % tp_size == 0
+            per_rank = num_heads // tp_size
+            return range(rank * per_rank, (rank + 1) * per_rank)
+        else:
+            h = rank * num_heads // tp_size
+            return range(h, h + 1)
 
+    @staticmethod
+    def _range_overlap(a: range, b: range) -> range:
+        start = max(a.start, b.start)
+        stop = min(a.stop, b.stop)
+        return range(start, max(start, stop))
 
-def get_current_attn_backends(
-    vllm_config: VllmConfig, layer_names: list[str] | None = None
-) -> list[type[AttentionBackend]]:
-    """Get all distinct attention backends for the given layers.
+    # ============================================================
+    # Private: build Mamba transfer info
+    # ============================================================
 
-    Args:
-        vllm_config: The current vLLM configuration.
-        layer_names: Optional list of layer names to scope the lookup.
-            When None, all attention layers are considered.
+    def _build_mamba_info(
+        self,
+        remote_tp_size: int,
+        remote_block_size: int,
+        remote_block_len: int,
+        remote_physical_blocks_per_logical: int,
+        local_block_len: int,
+    ) -> MambaEngineTransferInfo:
+        """Compute Mamba transfer plan."""
+        K = self.total_num_kv_heads
+        local_tp = self.tp_size
+        local_rank = self.tp_rank
+
+        is_remote_replicated = remote_tp_size > K
+        remote_physical_heads = max(1, K // remote_tp_size)
+
+        if local_tp >= remote_tp_size:
+            assert local_tp % remote_tp_size == 0
+            tp_ratio = local_tp // remote_tp_size
+        else:
+            assert remote_tp_size % local_tp == 0
+            tp_ratio = -(remote_tp_size // local_tp)
 
-    Returns:
-        Deduplicated list of attention backend classes.
-    """
-    layer_type = cast(type[Any], AttentionLayerBase)
-    layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
-    if layers:
-        seen: dict[str, type[AttentionBackend]] = {}
-        for layer in layers.values():
-            backend = layer.get_attn_backend()
-            seen[backend.full_cls_name()] = backend
-        return list(seen.values())
+        abs_tp = -tp_ratio if tp_ratio < 0 else 1
 
-    # Fallback for tests, when static_forward_context is empty.
-    logger.debug(
-        "No layers found in the vLLM config. Falling back to default attention backend."
-    )
-    from vllm.v1.attention.selector import get_attn_backend
+        mamba_range: range | None = None
+        if tp_ratio < 0:
+            mamba_range = range(local_rank * abs_tp, (local_rank + 1) * abs_tp)
 
-    return [
-        get_attn_backend(
-            head_size=vllm_config.model_config.get_head_size(),
-            dtype=vllm_config.model_config.dtype,
-            kv_cache_dtype=vllm_config.cache_config.cache_dtype,
-            use_mla=vllm_config.model_config.use_mla,
-        )
-    ]
+        # ---- FA read targets ----
+        if self.is_mla or tp_ratio >= 0:
+            num_fa_reads = 1
+            fa_source_ranks: list[int] = (
+                [0]
+                if self.is_mla
+                else [local_rank // tp_ratio if tp_ratio > 0 else local_rank]
+            )
+        else:
+            local_needs = self._physical_head_range(local_tp, K, local_rank)
+            search_range = (
+                mamba_range if mamba_range is not None else range(remote_tp_size)
+            )
+            seen: set[tuple[int, int]] = set()
+            fa_source_ranks = []
+            for p in search_range:
+                p_has = self._physical_head_range(remote_tp_size, K, p)
+                ov = self._range_overlap(local_needs, p_has)
+                if len(ov) > 0:
+                    key = (ov.start, ov.stop)
+                    if key not in seen:
+                        seen.add(key)
+                        fa_source_ranks.append(p)
+            if not fa_source_ranks:
+                for p in range(remote_tp_size):
+                    p_has = self._physical_head_range(remote_tp_size, K, p)
+                    ov = self._range_overlap(local_needs, p_has)
+                    if len(ov) > 0:
+                        key = (ov.start, ov.stop)
+                        if key not in seen:
+                            seen.add(key)
+                            fa_source_ranks.append(p)
+            num_fa_reads = len(fa_source_ranks)
 
+        # ---- All source ranks (mamba + FA) ----
+        if mamba_range is not None and abs_tp > num_fa_reads:
+            num_mamba_reads = abs_tp
+            all_source_ranks = list(mamba_range)
+        else:
+            num_mamba_reads = num_fa_reads
+            all_source_ranks = list(fa_source_ranks)
 
-def get_current_attn_backend(
-    vllm_config: VllmConfig, layer_names: list[str] | None = None
-) -> type[AttentionBackend]:
-    """Get the first attention backend for the given layers."""
-    return get_current_attn_backends(vllm_config, layer_names)[0]
+        # ---- FA descriptor bytes ----
+        effective_block_len = min(local_block_len, remote_block_len)
+        if self.is_kv_layout_blocks_first:
+            fa_descriptor_bytes = effective_block_len // 2
+        else:
+            fa_descriptor_bytes = effective_block_len
 
+        # ---- Validation ----
+        is_local_replicated = local_tp > K
+        if is_local_replicated and is_remote_replicated and tp_ratio > 0:
+            logger.info(
+                "Both-replicated hetero-TP: local_tp=%d > remote_tp=%d > K=%d.",
+                local_tp,
+                remote_tp_size,
+                K,
+            )
+        tt_set = set(all_source_ranks)
+        for t in fa_source_ranks:
+            if t not in tt_set:
+                logger.error(
+                    "FA source rank %d NOT in all_source_ranks %s.",
+                    t,
+                    all_source_ranks,
+                )
+        if self.is_kv_layout_blocks_first and tp_ratio < 0 and num_fa_reads > 0:
+            local_k_half = local_block_len // 2
+            remote_k_half = remote_block_len // 2
+            expected = local_k_half // num_fa_reads
+            if expected != remote_k_half:
+                logger.warning(
+                    "FA size mismatch: local_k_half=%d / reads=%d = %d, "
+                    "but remote_k_half=%d.",
+                    local_k_half,
+                    num_fa_reads,
+                    expected,
+                    remote_k_half,
+                )
 
-# TODO (ZhanqiuHu): Consolidate TpKVTopology and HeteroTPTransferConfig
-# into a single engine-agnostic TransferTopology class.
-# 6 of 9 HeteroTPTransferConfig init fields duplicate TpKVTopology data.
-#
-# @dataclass
-# class EngineTransferInfo:
-#     """Per-remote-engine transfer state, computed at handshake."""
-#     p_tp: int
-#     tp_ratio: int
-#     p_block_len: int
-#     block_size: int
-#     # Mamba-specific (None for non-mamba models)
-#     fa_read_targets: list[int] | None = None
-#     transfer_targets: list[int] | None = None
-#     physical_fa_num_reads: int | None = None
-#     mamba_num_reads: int | None = None
-#     fa_entry_size: int | None = None
-#
-# class TransferTopology:
-#     """Single source of truth for TP topology + transfer sizing."""
-#     # Shared (set once at init, replaces duplicate fields)
-#     tp_rank: int          # == TpKVTopology.tp_rank == HeteroTP.d_rank
-#     tp_size: int          # == TpKVTopology.tp_size == HeteroTP.d_tp
-#     total_num_kv_heads: int  # == HeteroTP.K
-#     is_mla: bool          # == HeteroTP.use_mla
-#     is_mamba: bool
-#     is_blocks_first: bool # == HeteroTP.is_blocks_first
-#     d_block_len: int
-#
-#     # Per-engine (populated via register_engine() at handshake)
-#     _engines: dict[EngineId, EngineTransferInfo]
-#
-#     def register_engine(self, engine_id, p_tp, p_block_len, ...): ...
-#
-#     # General (from TpKVTopology)
-#     def tp_ratio(self, engine_id) -> int: ...
-#     def target_remote_ranks(self, engine_id) -> list[int]: ...
-#     def is_kv_replicated(self, engine_id) -> bool: ...
-#
-#     # Mamba-specific (from HeteroTPTransferConfig, gated by is_mamba)
-#     def fa_rank_offset(self, engine_id, block_len) -> int: ...
-#     def physical_fa_num_reads(self, engine_id) -> int: ...
-#     def transfer_targets(self, engine_id) -> list[int]: ...
-#     def should_skip_fa(self, engine_id, p_rank) -> bool: ...
-#     def filter_block_ids_for_rank(self, engine_id, ...) -> ...: ...
+        return MambaEngineTransferInfo(
+            remote_tp_size=remote_tp_size,
+            remote_block_len=remote_block_len,
+            remote_block_size=remote_block_size,
+            remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),
+            remote_fa_source_ranks=tuple(fa_source_ranks),
+            remote_all_source_ranks=tuple(all_source_ranks),
+            remote_num_fa_reads=num_fa_reads,
+            remote_num_mamba_reads=num_mamba_reads,
+            remote_fa_descriptor_bytes=fa_descriptor_bytes,
+            is_remote_replicated=is_remote_replicated,
+            remote_physical_heads=remote_physical_heads,
+        )

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -21,7 +21,7 @@ from vllm import envs
 from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     EngineId,
-    TpKVTopology,
+    TransferTopology,
     get_current_attn_backend,
     get_current_attn_backends,
 )
@@ -764,13 +764,13 @@ class MooncakeConnectorWorker:
         logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
 
         self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
-        self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
-        self.kv_topo = TpKVTopology(
+        self.transfer_topo = TransferTopology(
             tp_rank=self.tp_rank,
+            tp_size=self.tp_size,
+            block_size=self.block_size,
             engine_id=self.engine_id,
-            remote_tp_size=self._tp_size,  # shared state
-            remote_block_size=self._block_size,  # shared state
             is_mla=self.use_mla,
+            is_mamba=False,
             total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
             attn_backends=[backend],
         )
@@ -911,7 +911,7 @@ class MooncakeConnectorWorker:
         self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata
     ):
         pending_reqs: dict[ReqId, SendBlockMeta] = {}
-        remote_tp_ranks = self.kv_topo.get_target_remote_ranks(meta.remote_tp_size)
+        remote_tp_ranks = self.transfer_topo.handshake_target_ranks(meta.remote_tp_size)
         if meta.remote_tp_rank not in remote_tp_ranks:
             # This D worker does not pair with the P worker.
             msg = (
@@ -1256,7 +1256,7 @@ class MooncakeConnectorWorker:
         seen_base_addresses = []
         self.block_len_per_layer = []
 
-        split_k_and_v = self.kv_topo.split_k_and_v
+        split_k_and_v = self.transfer_topo.split_k_and_v
         tensor_size_bytes = None
         for layer_name, cache_or_caches in kv_caches.items():
             cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
@@ -1495,8 +1495,8 @@ class MooncakeConnectorWorker:
         remote_engine_id: EngineId,
         pull_metas: dict[ReqId, PullReqMeta],
     ):
-        remote_tp_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
-            remote_engine_id
+        remote_tp_ranks = self.transfer_topo.handshake_target_ranks(
+            self._tp_size[remote_engine_id]
         )
         count = len(remote_tp_ranks)
         logger.debug(
@@ -1587,7 +1587,7 @@ class MooncakeConnectorWorker:
             )
 
     def _producer_cache_is_replicated(self) -> bool:
-        return self.kv_topo.replicates_kv_cache(self.engine_id)
+        return self.transfer_topo.local_replicates_kv_cache
 
     def _get_transfer_regions(
         self, base_addrs: list[int], block_lens: list[int]
@@ -1595,7 +1595,7 @@ class MooncakeConnectorWorker:
         return _expand_transfer_regions(
             base_addrs=base_addrs,
             block_lens=block_lens,
-            is_kv_layout_blocks_first=self.kv_topo.is_kv_layout_blocks_first,
+            is_kv_layout_blocks_first=self.transfer_topo.is_kv_layout_blocks_first,
         )
 
     def _get_sender_transfer_plan(

vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py

@@ -21,8 +21,8 @@ from vllm import envs
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     BlockIds,
     EngineId,
-    HeteroTPTransferConfig,
-    TpKVTopology,
+    MambaEngineTransferInfo,
+    TransferTopology,
     get_current_attn_backends,
     kv_postprocess_blksize_and_layout_on_receive,
     kv_postprocess_blksize_on_receive,
@@ -49,7 +49,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
     MambaConvSplitInfo,
-    compute_mamba_phys_ratio,
+    compute_physical_blocks_per_logical,
     derive_mamba_conv_split,
 )
 from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
@@ -269,14 +269,12 @@ class NixlConnectorWorker:
         self._registered_descs: list[Any] = []
 
         # ---- Mamba-HMA per-engine state (only used when self._has_mamba) ----
-        # Per-engine transfer config (source of truth for FA/mamba sizing).
-        self._transfer_configs: dict[str, HeteroTPTransferConfig] = {}
-        # NOTE (ZhanqiuHu): _mamba_phys_ratio MUST be per-engine.
-        # compute_mamba_phys_ratio = ceil((conv_bytes + ssm_bytes) / block_len)
+        # NOTE (ZhanqiuHu): _physical_blocks_per_logical MUST be per-engine.
+        # physical_blocks_per_logical = ceil((conv_bytes + ssm_bytes) / block_len)
         # where conv/ssm bytes are per-TP-rank (dimension-sharded).  With
         # heterogeneous TP the per-rank sizes differ, so the ratio differs:
         #   e.g. Nemotron 30B: P(TP=4) → 131, D(TP=1) → 261.
-        self._mamba_phys_ratio: dict[EngineId, int] = {}
+        self._physical_blocks_per_logical: dict[EngineId, int] = {}
 
         # In progress transfers.
         # [req_id -> list[handle]]
@@ -322,10 +320,8 @@ class NixlConnectorWorker:
 
         # lazy initialized in register_kv_caches
         self.compat_hash: str | None = None
-        self.kv_topo: TpKVTopology | None = None
+        self.transfer_topo: TransferTopology | None = None
 
-        self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
-        self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
         # With heterogeneous TP, P must wait for all assigned D TP workers to
         # finish reading before safely freeing the blocks.
         self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
@@ -355,7 +351,6 @@ class NixlConnectorWorker:
                 self.block_size // kernel_block_size
             )
             self.block_size = kernel_block_size
-            self._block_size[self.engine_id] = kernel_block_size
             self.num_blocks *= self._physical_blocks_per_logical_kv_block
 
     def _nixl_handshake(
@@ -384,8 +379,8 @@ class NixlConnectorWorker:
         # Regardless, only handshake with the remote TP rank(s) that current
         # local rank will read from. Note that With homogeneous TP,
         # this happens to be the same single rank_i.
-        assert self.kv_topo is not None
-        p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size)
+        assert self.transfer_topo is not None
+        p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size)
         remote_rank_to_agent_name = {}
         path = make_zmq_path("tcp", host, port)
 
@@ -649,11 +644,11 @@ class NixlConnectorWorker:
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in nixl."""
-        self.kv_topo = TpKVTopology(
+        self.transfer_topo = TransferTopology(
             tp_rank=self.tp_rank,
+            tp_size=self.world_size,
+            block_size=self.block_size,
             engine_id=self.engine_id,
-            remote_tp_size=self._tp_size,  # shared state
-            remote_block_size=self._block_size,  # shared state
             is_mla=self.use_mla,
             total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
             attn_backends=self.attn_backends,
@@ -664,7 +659,7 @@ class NixlConnectorWorker:
             is_mamba=self._has_mamba,
         )
         self.compat_hash = compute_nixl_compatibility_hash(
-            self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
+            self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks
         )
 
         if self.use_host_buffer:
@@ -716,7 +711,7 @@ class NixlConnectorWorker:
             if isinstance(layer_spec, UniformTypeKVCacheSpecs):
                 # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
                 layer_spec = layer_spec.kv_cache_specs[layer_name]
-            cache_list = self.kv_topo.get_transfer_cache_regions(
+            cache_list = self.transfer_topo.get_transfer_cache_regions(
                 cache_or_caches, layer_spec
             )
             # `layer_spec.page_size_bytes` only accounts for logical page_size, that is
@@ -729,7 +724,7 @@ class NixlConnectorWorker:
             )
             # For when registering multiple tensors eg K/V in separate regions.
             physical_page_size = physical_page_size // len(cache_list)
-            if self.kv_topo._cross_layers_blocks:
+            if self.transfer_topo._cross_layers_blocks:
                 # When cross-layers blocks are used, multiply by number of layers
                 physical_page_size = physical_page_size * len(
                     self.kv_cache_config.kv_cache_tensors
@@ -793,7 +788,7 @@ class NixlConnectorWorker:
         self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
         self.num_regions = len(caches_data)
 
-        if self.kv_topo.is_kv_layout_blocks_first:
+        if self.transfer_topo.is_kv_layout_blocks_first:
             # NOTE (NickLucche) When FlashInfer is used, memory is registered
             # with joint KV for each block. This minimizes the overhead in
             # registerMem allowing faster descs queries. In order to be able to
@@ -817,7 +812,7 @@ class NixlConnectorWorker:
         self.dst_num_blocks[self.engine_id] = self.num_blocks
 
         if self._has_mamba:
-            self._mamba_phys_ratio[self.engine_id] = (
+            self._physical_blocks_per_logical[self.engine_id] = (
                 self._physical_blocks_per_logical_kv_block
             )
             logger.info(
@@ -876,11 +871,13 @@ class NixlConnectorWorker:
         conv_offsets = self._conv_decomp.local_conv_offsets
         conv_size, ssm_size = self._mamba_ssm_size
         num_blocks = self._logical_num_blocks * block_size_ratio
-        phys_ratio = self._physical_blocks_per_logical_kv_block
+        physical_per_logical = self._physical_blocks_per_logical_kv_block
 
         result: list[tuple[int, int, int]] = []
         for i, base_addr in enumerate(base_addresses):
-            page_stride = self.block_len_per_layer[i] // block_size_ratio * phys_ratio
+            page_stride = (
+                self.block_len_per_layer[i] // block_size_ratio * physical_per_logical
+            )
             for off, sz in conv_offsets:
                 for blk in range(num_blocks):
                     result.append(
@@ -900,14 +897,14 @@ class NixlConnectorWorker:
     def _build_fa_remote_for_mamba(
         self,
         nixl_agent_meta: NixlAgentMetadata,
-        transfer_cfg: HeteroTPTransferConfig,
         block_size_ratio: int,
-        kv_topo: TpKVTopology,
+        transfer_topo: TransferTopology,
+        remote_engine_id: EngineId,
     ) -> list[tuple[int, int, int]]:
         """Build remote FA descriptors for mamba models.
 
-        Uses transfer_cfg for GQA-aware FA divisor and head-based rank offset
-        instead of the standard uniform tp_ratio split.
+        Uses TransferTopology for GQA-aware FA divisor and head-based rank
+        offset instead of the standard uniform tp_ratio split.
         """
         assert block_size_ratio == 1, (
             "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
@@ -915,7 +912,9 @@ class NixlConnectorWorker:
         )
         # TODO (ZhanqiuHu): unify with register_remote_blocks when Mamba-HMA
         # hetero-TP logic stabilizes.
-        tp_ratio = transfer_cfg.tp_ratio
+        mamba_info = transfer_topo.get_engine_info(remote_engine_id)
+        assert isinstance(mamba_info, MambaEngineTransferInfo)
+        tp_ratio = transfer_topo.tp_ratio(mamba_info.remote_tp_size)
         result: list[tuple[int, int, int]] = []
         for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
             local_block_len = self.get_backend_aware_kv_block_len(
@@ -926,9 +925,11 @@ class NixlConnectorWorker:
                 local_block_len = remote_kv_block_len
 
             if tp_ratio < 0 and not self.use_mla:
-                local_block_len = local_block_len // transfer_cfg.physical_fa_num_reads
+                local_block_len = local_block_len // mamba_info.remote_num_fa_reads
 
-            rank_offset = transfer_cfg.fa_rank_offset(remote_kv_block_len)
+            rank_offset = transfer_topo.fa_rank_offset(
+                remote_engine_id, remote_kv_block_len
+            )
 
             num_blocks = nixl_agent_meta.num_blocks
             page_size = nixl_agent_meta.block_lens[i]
@@ -937,12 +938,12 @@ class NixlConnectorWorker:
                 addr = base_addr + block_offset + rank_offset
                 result.append((addr, local_block_len, nixl_agent_meta.device_id))
 
-            if kv_topo.is_kv_layout_blocks_first:
+            if transfer_topo.is_kv_layout_blocks_first:
                 second_split = self.get_backend_aware_kv_block_len(
                     layer_idx=i, first_split=False, mamba_view=False
                 )
                 if tp_ratio < 0 and not self.use_mla:
-                    second_split = second_split // transfer_cfg.physical_fa_num_reads
+                    second_split = second_split // mamba_info.remote_num_fa_reads
                 for block_id in range(num_blocks):
                     block_offset = block_id * page_size
                     addr = base_addr + block_offset + rank_offset
@@ -981,15 +982,17 @@ class NixlConnectorWorker:
             conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)]
             ssm_read_size = nixl_agent_meta.ssm_sizes[1]
 
-        remote_ratio = self._mamba_phys_ratio[nixl_agent_meta.engine_id]
-        num_blocks = nixl_agent_meta.num_blocks // remote_ratio
+        remote_physical_per_logical = self._physical_blocks_per_logical[
+            nixl_agent_meta.engine_id
+        ]
+        num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical
         device_id = nixl_agent_meta.device_id
 
         result: list[tuple[int, int, int]] = []
         # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
         # block lengths vary across layers (e.g. MLA).
         for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
-            page_stride = nixl_agent_meta.block_lens[i] * remote_ratio
+            page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical
             for off, sz in conv_offsets:
                 for blk in range(num_blocks):
                     result.append((base_addr + blk * page_stride + off, sz, device_id))
@@ -1019,8 +1022,8 @@ class NixlConnectorWorker:
         register another local_xfer_handler using remote block len to ensure
         data copy correctness.
         """
-        assert self.kv_topo is not None
-        kv_topo = self.kv_topo
+        assert self.transfer_topo is not None
+        transfer_topo = self.transfer_topo
 
         block_size_ratio = self.block_size // block_size
         blocks_data: list[tuple[int, int, int]] = []
@@ -1051,7 +1054,7 @@ class NixlConnectorWorker:
                     # (addr, len, device id)
                     blocks_data.append((addr, kv_block_len, self.device_id))
 
-                if kv_topo.is_kv_layout_blocks_first:
+                if transfer_topo.is_kv_layout_blocks_first:
                     second_split = self.get_backend_aware_kv_block_len(
                         layer_idx=i, first_split=False, mamba_view=mamba
                     )
@@ -1153,11 +1156,29 @@ class NixlConnectorWorker:
             )
             return self._remote_agents[engine_id][remote_tp_rank]
 
-        ### Register remote agent metadata
-        if engine_id not in self._tp_size:
-            self._tp_size[engine_id] = remote_tp_size
-        if engine_id not in self._block_size:
-            self._block_size[engine_id] = nixl_agent_meta.block_size
+        ### Register remote engine in TransferTopology (idempotent).
+        assert self.transfer_topo is not None
+        transfer_topo = self.transfer_topo
+        physical_blocks_per_logical = (
+            compute_physical_blocks_per_logical(
+                nixl_agent_meta.ssm_sizes,
+                nixl_agent_meta.block_lens[0],
+            )
+            if self._has_mamba
+            else 1
+        )
+        transfer_topo.register_remote_engine(
+            remote_engine_id=engine_id,
+            remote_tp_size=remote_tp_size,
+            remote_block_size=nixl_agent_meta.block_size,
+            remote_block_len=nixl_agent_meta.block_lens[0],
+            remote_physical_blocks_per_logical=physical_blocks_per_logical,
+            local_block_len=self.block_len_per_layer[0],
+        )
+        if self._has_mamba and engine_id not in self._physical_blocks_per_logical:
+            self._physical_blocks_per_logical[engine_id] = physical_blocks_per_logical
+
+        logger.info("Transfer plan: %s", transfer_topo.describe(engine_id))
 
         remote_agent_name = self.nixl_wrapper.add_remote_agent(
             nixl_agent_meta.agent_metadata
@@ -1170,16 +1191,10 @@ class NixlConnectorWorker:
         # remote:               | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
         # local origin:|          0|          1|          8|         12|
         # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
-        assert self.kv_topo is not None
-        kv_topo = self.kv_topo
-        block_size_ratio = kv_topo.block_size_ratio_from_engine_id(engine_id)
+        block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size)
 
         if engine_id not in self.dst_num_blocks:
             self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
-        if self._has_mamba:
-            self._mamba_phys_ratio[engine_id] = compute_mamba_phys_ratio(
-                nixl_agent_meta.ssm_sizes, nixl_agent_meta.block_lens[0]
-            )
 
         # Keep track of remote agent kv caches base addresses.
         self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
@@ -1189,28 +1204,13 @@ class NixlConnectorWorker:
 
         # This is 1 when P and D `--tensor-parallel-size` match. Otherwise,
         # this is the ratio between the two sizes.
-        tp_ratio = self.kv_topo.tp_ratio_from_engine_id(engine_id)
+        tp_ratio = transfer_topo.tp_ratio(remote_tp_size)
 
         # Handle tp_size>num_kv_heads: replicate KV cache.
         indexes_into_remote = (
-            not self.kv_topo.replicates_kv_cache(engine_id) and tp_ratio > 0
+            not transfer_topo.replicates_kv_cache(engine_id) and tp_ratio > 0
         )
 
-        # Create transfer config (single source of truth for descriptor sizes).
-        if self._has_mamba and engine_id not in self._transfer_configs:
-            self._transfer_configs[engine_id] = HeteroTPTransferConfig(
-                tp_ratio=tp_ratio,
-                K=kv_topo.total_num_kv_heads,
-                d_tp=self.world_size,
-                p_tp=remote_tp_size,
-                d_rank=self.tp_rank,
-                use_mla=self.use_mla,
-                d_block_len=self.block_len_per_layer[0],
-                p_block_len=nixl_agent_meta.block_lens[0],
-                is_blocks_first=kv_topo.is_kv_layout_blocks_first,
-            )
-            logger.info("Created %s", self._transfer_configs[engine_id].describe())
-
         logger.debug(
             "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s",
             engine_id,
@@ -1231,12 +1231,10 @@ class NixlConnectorWorker:
             self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
 
             if self._has_mamba:
-                transfer_cfg = self._transfer_configs.get(engine_id)
-                assert transfer_cfg is not None
-                if transfer_cfg.needs_split_handles:
+                if transfer_topo.needs_split_handles(engine_id):
                     # Mamba-HMA: FA and Mamba use different split factors.
-                    for handle_data in transfer_cfg.compute_split_handle_data(
-                        self.src_blocks_data, self.num_descs, abs_tp
+                    for handle_data in transfer_topo.compute_split_handle_data(
+                        engine_id, self.src_blocks_data, self.num_descs, abs_tp
                     ):
                         descs = self.nixl_wrapper.get_xfer_descs(
                             handle_data, self.nixl_memory_type
@@ -1247,12 +1245,8 @@ class NixlConnectorWorker:
                         self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
 
                     logger.info(
-                        "Mamba-HMA split handles: targets=%s, fa_reads=%s, "
-                        "fa_entry=%s, mamba_reads=%s, num_descs=%s",
-                        transfer_cfg.transfer_targets,
-                        transfer_cfg.physical_fa_num_reads,
-                        transfer_cfg.fa_entry_size,
-                        transfer_cfg.mamba_num_reads,
+                        "Mamba-HMA split handles: %s, num_descs=%s",
+                        transfer_topo.describe(engine_id),
                         self.num_descs,
                     )
             else:
@@ -1321,7 +1315,7 @@ class NixlConnectorWorker:
                         (addr, local_block_len, nixl_agent_meta.device_id)
                     )
 
-                if kv_topo.is_kv_layout_blocks_first:
+                if transfer_topo.is_kv_layout_blocks_first:
                     # With FlashInfer index V separately to allow head splitting.
                     second_split = self.get_backend_aware_kv_block_len(
                         layer_idx=i, first_split=False, mamba_view=mamba
@@ -1360,14 +1354,12 @@ class NixlConnectorWorker:
                 engine_id,
                 remote_tp_rank,
             )
-            transfer_cfg = self._transfer_configs.get(engine_id)
-            assert transfer_cfg is not None
             blocks_data.extend(
                 self._build_fa_remote_for_mamba(
                     nixl_agent_meta,
-                    transfer_cfg,
                     block_size_ratio,
-                    kv_topo,
+                    transfer_topo,
+                    engine_id,
                 )
             )
             blocks_data.extend(
@@ -1403,18 +1395,19 @@ class NixlConnectorWorker:
         """
         remote_engine_id = nixl_agent_meta.engine_id
 
-        assert self._tp_size[remote_engine_id] == remote_tp_size
-        assert self.kv_topo is not None
+        assert self.transfer_topo is not None
+        remote_info = self.transfer_topo.get_engine_info(remote_engine_id)
+        assert remote_info.remote_tp_size == remote_tp_size
 
-        tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
-        block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
-            remote_engine_id
+        tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
+        block_size_ratio = self.transfer_topo.block_size_ratio(
+            nixl_agent_meta.block_size
         )
         # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba.
         # Mamba models can have replicated FA KV with tp_ratio < 0.
         if not self._has_mamba:
             assert not (
-                tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id)
+                tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id)
             )
 
         if self._is_hma_required:
@@ -1467,7 +1460,7 @@ class NixlConnectorWorker:
         if (
             abs(tp_ratio) != 1
             and not self.use_mla
-            and not self.kv_topo.is_kv_replicated(remote_engine_id)
+            and not self.transfer_topo.is_kv_replicated(remote_engine_id)
             and kv_cache_layout != "HND"
             and not self.enable_permute_local_kv
         ):
@@ -1478,7 +1471,7 @@ class NixlConnectorWorker:
 
         # Block len can only vary across layers when using MLA.
         remote_block_len = nixl_agent_meta.block_lens[0]
-        if self.use_mla or self.kv_topo.is_kv_replicated(remote_engine_id):
+        if self.use_mla or self.transfer_topo.is_kv_replicated(remote_engine_id):
             # With replicated KV cache, only the number of blocks can differ.
             # TODO (ZhanqiuHu): For mamba models, validate FA and mamba
             # block_lens separately.
@@ -1594,7 +1587,7 @@ class NixlConnectorWorker:
         if len(self.device_kv_caches) == 0:
             return
         assert block_size_ratio >= 1, "Only nP < nD supported currently."
-        assert self.kv_topo is not None
+        assert self.transfer_topo is not None
         if self.enable_permute_local_kv and block_size_ratio > 1:
             logger.debug(
                 "Post-processing device kv cache on receive by converting "
@@ -1614,7 +1607,7 @@ class NixlConnectorWorker:
                 block_size_ratio,
             )
 
-        split_k_and_v = self.kv_topo.split_k_and_v
+        split_k_and_v = self.transfer_topo.split_k_and_v
 
         for block_ids in block_ids_list:
             indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
@@ -1661,7 +1654,7 @@ class NixlConnectorWorker:
         The scheduler process (via the MultiprocExecutor) will use this output
         to track which workers are done.
         """
-        assert self.kv_topo is not None
+        assert self.transfer_topo is not None
         done_sending = self._get_new_notifs()
         done_recving = self._pop_done_transfers(self._recving_transfers)
 
@@ -1689,8 +1682,9 @@ class NixlConnectorWorker:
                 self.sync_recved_kv_to_device(req_id, meta)
 
             # post processing for heteroblocksize
-            block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
-                meta.remote.engine_id
+            remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id)
+            block_size_ratio = self.transfer_topo.block_size_ratio(
+                remote_info.remote_block_size
             )
             if not self.use_mla and (
                 block_size_ratio > 1 or self.enable_permute_local_kv
@@ -1741,7 +1735,7 @@ class NixlConnectorWorker:
         are reading from the same producer (heterogeneous TP scenario), wait
         for all consumers to be done pulling.
         """
-        assert self.kv_topo is not None
+        assert self.transfer_topo is not None
         notified_req_ids: set[str] = set()
         for notifs in self.nixl_wrapper.get_new_notifs().values():
             for notif in notifs:
@@ -1760,7 +1754,7 @@ class NixlConnectorWorker:
 
                 # NOTE: `tp_ratio` is the opposite when swapping local<>remote
                 n_consumers = int(tp_size)
-                tp_ratio = self.kv_topo.tp_ratio(n_consumers)
+                tp_ratio = self.transfer_topo.tp_ratio(n_consumers)
 
                 # Number of reads *per producer* to wait for.
                 # When remote D TP > local P TP we expect `tp_ratio` reads.
@@ -1901,17 +1895,17 @@ class NixlConnectorWorker:
                 self._reqs_to_send[req_id] = expiration_time
 
     def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
-        assert meta.remote is not None and self.kv_topo is not None
-        remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
-            meta.remote.engine_id
-        )
-        tp_ratio = self.kv_topo.tp_ratio_from_engine_id(meta.remote.engine_id)
+        assert meta.remote is not None and self.transfer_topo is not None
+        engine_id = meta.remote.engine_id
+        remote_ranks = self.transfer_topo.target_remote_ranks(engine_id)
+        remote_info = self.transfer_topo.get_engine_info(engine_id)
+        tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
 
         if self._has_mamba:
             # Expand remote logical → kernel block IDs.
             meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
                 meta.remote.block_ids,
-                self._mamba_phys_ratio[meta.remote.engine_id],
+                self._physical_blocks_per_logical[meta.remote.engine_id],
             )
         else:
             meta.remote.block_ids = self._logical_to_kernel_block_ids(
@@ -1924,7 +1918,7 @@ class NixlConnectorWorker:
                 # the first remote rank (cache is duplicated)..
                 break
 
-            remote_block_size = self.kv_topo.remote_block_size[meta.remote.engine_id]
+            remote_block_size = remote_info.remote_block_size
             logger.debug(
                 "Remote agent %s available, calling _read_blocks"
                 " on remote rank %s with remote block size %s for req %s",
@@ -1955,9 +1949,8 @@ class NixlConnectorWorker:
             remote_ids: BlockIds = meta.remote.block_ids
             if self._has_mamba:
                 # Mamba-HMA: zero out FA groups for P ranks outside fa_read_targets.
-                transfer_cfg = self._transfer_configs.get(meta.remote.engine_id)
-                assert transfer_cfg is not None
-                local_ids, remote_ids = transfer_cfg.filter_block_ids_for_rank(
+                local_ids, remote_ids = self.transfer_topo.filter_block_ids_for_rank(
+                    engine_id,
                     remote_rank,
                     local_ids,
                     remote_ids,
@@ -1999,8 +1992,11 @@ class NixlConnectorWorker:
         Post a READ point-to-point xfer request from a single local worker to
         a single remote worker.
         """
-        assert self.kv_topo is not None
-        block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)
+        assert self.transfer_topo is not None
+        remote_info = self.transfer_topo.get_engine_info(dst_engine_id)
+        block_size_ratio = self.transfer_topo.block_size_ratio(
+            remote_info.remote_block_size
+        )
         if block_size_ratio > 1:
             # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.
             assert not self._is_hma_required
@@ -2190,8 +2186,8 @@ class NixlConnectorWorker:
             # This is like having two "low-level views" of the same storage.
             # `num_fa_descs` offset must be computed per-engine since P and D can
             # have different num_blocks (and thus different FA descs counts).
-            ratio = self._mamba_phys_ratio[engine_id]
-            logical_blocks = num_blocks // ratio
+            physical_per_logical = self._physical_blocks_per_logical[engine_id]
+            logical_blocks = num_blocks // physical_per_logical
             num_fa_descs = self.num_regions * num_blocks
             # 3-read mamba: 4 regions per unique cache tensor (x, B, C, ssm).
             mamba_region_ids = np.arange(len(self.block_len_per_layer) * 4)[:, None]
@@ -2234,21 +2230,22 @@ class NixlConnectorWorker:
         ]
 
     def _logical_to_remote_kernel_block_ids(
-        self, block_ids: BlockIds, remote_ratio: int
+        self, block_ids: BlockIds, remote_physical_per_logical: int
     ) -> BlockIds:
         """Map logical block IDs to physical kernel block IDs on the remote.
 
         Args:
             block_ids: per-group lists of logical block IDs.
-            remote_ratio: remote engine's physical blocks per logical block.
+            remote_physical_per_logical: remote engine's physical blocks
+                per logical block.
 
         Returns:
             Same structure with FA groups expanded (each logical block L
-            becomes kernel blocks [L*remote_ratio .. L*remote_ratio +
-            local_ratio - 1]).  Mamba groups are passed through unchanged.
+            becomes kernel blocks [L*ratio .. L*ratio + local_ratio - 1]).
+            Mamba groups are passed through unchanged.
         """
         local_ratio = self._physical_blocks_per_logical_kv_block
-        if remote_ratio == 1:
+        if remote_physical_per_logical == 1:
             return block_ids
         local_arange = np.arange(local_ratio).reshape(1, -1)
         group_specs = self.kv_cache_config.kv_cache_groups
@@ -2256,7 +2253,7 @@ class NixlConnectorWorker:
         for i, group in enumerate(block_ids):
             if not isinstance(group_specs[i].kv_cache_spec, MambaSpec):
                 arr = np.array(group).reshape(-1, 1)
-                expanded = (arr * remote_ratio + local_arange).flatten()
+                expanded = (arr * remote_physical_per_logical + local_arange).flatten()
                 result.append(expanded.tolist())
             else:
                 # Mamba blocks are 1:1 logical-to-physical (no expansion).
@@ -2296,8 +2293,8 @@ class NixlConnectorWorker:
            +-------------------+         +--------------------+
            |1st_split-2nd_split|         |1st_split-2nd_split |
         """
-        assert self.kv_topo is not None
-        if self.kv_topo.is_kv_layout_blocks_first:
+        assert self.transfer_topo is not None
+        if self.transfer_topo.is_kv_layout_blocks_first:
             # For indexing only half (either just the K or V part).
             if mamba_view:
                 # NOTE (NickLucche) Mamba Opt: this is already skipping the padding so

vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py

@@ -151,7 +151,9 @@ def derive_mamba_conv_split(
     )
 
 
-def compute_mamba_phys_ratio(ssm_sizes: tuple[int, ...], block_len: int) -> int:
+def compute_physical_blocks_per_logical(
+    ssm_sizes: tuple[int, ...], block_len: int
+) -> int:
     """Derive _physical_blocks_per_logical_kv_block from remote metadata.
 
     The remote engine's ratio is not sent directly in the handshake, so we