[Core] Enable KimiLinear (KDA/GDN + MLA) PD Separation via NIXL

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -533,12 +533,24 @@ def tp_ratio(self, remote_tp_size: int) -> int:
         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."""
-        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."
+        """Calculate the block size ratio between local and remote.
+
+        Positive when local >= remote (local blocks are larger).
+        Negative when remote > local (remote blocks are larger).
+        """
+        if self.block_size == remote_block_size:
+            return 1
+        if self.block_size > remote_block_size:
+            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}."
+            )
+            return self.block_size // remote_block_size
+        assert remote_block_size % self.block_size == 0, (
+            f"Remote block size {remote_block_size} is not divisible "
+            f"by local block size {self.block_size}."
         )
-        return self.block_size // remote_block_size
+        return -(remote_block_size // self.block_size)
 
     def is_kv_replicated(
         self, remote_engine_id: EngineId, remote_pp_rank: int = 0

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

@@ -95,7 +95,7 @@ def _compute_desc_ids(
     ) -> np.ndarray:
         """Compute NIXL descriptor IDs for given block IDs."""
         num_fa_regions = self.num_regions
-        num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0
+        num_ssm_regions = sum(self._is_ssm_region) * 4 if self._has_mamba else 0
 
         num_blocks = dst_num_blocks
         if block_size_ratio is not None:
@@ -365,6 +365,10 @@ def __init__(
         # Number of NIXL regions. Currently one region per cache
         # (so 1 per layer for MLA, otherwise 2 per layer)
         self.num_regions = 0
+        # Per-region flag: True if the region is an SSM/Mamba layer.
+        # Populated during register_kv_caches; used to route FA descriptors
+        # to attention regions and Mamba descriptors to SSM regions.
+        self._is_ssm_region: list[bool] = []
 
         # nixl_prepped_dlist_handle.
         self.src_xfer_handles_by_block_size: dict[int, int] = {}
@@ -844,6 +848,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         # Enable different block lengths for different layers *only* when MLA is used.
         # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.
         self.block_len_per_layer = list[int]()
+        self._is_ssm_region = list[bool]()
         for layer_name, cache_or_caches in xfer_buffers.items():
             # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
             # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
@@ -898,8 +903,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     "Registering layer %s with cache shape: %s", layer_name, cache.shape
                 )
                 seen_base_addresses.append(base_addr)
+                is_ssm = isinstance(layer_spec, MambaSpec)
+                self._is_ssm_region.append(is_ssm)
                 # Only record non-Mamba page sizes.
-                if isinstance(layer_spec, MambaSpec):
+                if is_ssm:
                     self.block_len_per_layer.append(
                         physical_page_size // self._physical_blocks_per_logical_kv_block
                     )
@@ -941,7 +948,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         assert len(self.block_len_per_layer) == len(seen_base_addresses)
 
         self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
-        self.num_regions = len(caches_data)
+        # FA regions count only attention layers.  SSM/Mamba regions are
+        # served by Mamba descriptors, not FA descriptors, and must be
+        # excluded here so that FA descriptor IDs do not reference them.
+        # (For hybrid MLA+GDN models, building FA descriptors for the GDN
+        # region with virtual K/V split would produce out-of-bounds addresses
+        # under heterogeneous TP, failing NIXL prepXferDlist.)
+        num_attention_base = len(caches_data) - sum(self._is_ssm_region)
+        self.num_regions = num_attention_base
 
         if self.transfer_topo.virtually_split_kv_in_blocks:
             # NOTE (NickLucche) When FlashInfer is used, memory is registered
@@ -1030,6 +1044,12 @@ def _build_mamba_local(
 
         result: list[tuple[int, int, int]] = []
         for i, base_addr in enumerate(base_addresses):
+            # Only build Mamba descriptors for SSM/Mamba regions.
+            # Attention/MLA regions do not contain conv or temporal state;
+            # building Mamba descriptors for them would reference addresses
+            # outside their registered memory, causing prepXferDlist failures.
+            if not self._is_ssm_region[i]:
+                continue
             # Jump one page_size, but ssm page_size may be bigger when kernel
             # locks block size to a specific value (physical_per_logical scale).
             page_stride = (
@@ -1081,6 +1101,10 @@ def _build_mamba_remote(
         # 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):
+            # Only build Mamba descriptors for SSM/Mamba regions.
+            # Attention/MLA regions do not contain conv or temporal state.
+            if i < len(self._is_ssm_region) and not self._is_ssm_region[i]:
+                continue
             page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical
             for off, sz in conv_offsets:
                 for blk in range(num_blocks):
@@ -1106,6 +1130,10 @@ def _build_fa_local(
         num_blocks = self.num_blocks * block_size_ratio
         result: list[tuple[int, int, int]] = []
         for i, base_addr in enumerate(base_addresses):
+            # Only build FA descriptors for attention regions.
+            # SSM/Mamba regions are served by Mamba descriptors.
+            if i < len(self._is_ssm_region) and self._is_ssm_region[i]:
+                continue
             kv_block_len = (
                 self.get_backend_aware_kv_block_len(
                     layer_idx=i, first_split=True, mamba_view=False
@@ -1147,6 +1175,10 @@ def _build_fa_remote(
         num_blocks = nixl_agent_meta.num_blocks
         result: list[tuple[int, int, int]] = []
         for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
+            # Only build FA descriptors for attention regions.
+            # SSM/Mamba regions are served by Mamba descriptors.
+            if i < len(self._is_ssm_region) and self._is_ssm_region[i]:
+                continue
             # Read our whole local region size from remote..
             local_block_len = self.get_backend_aware_kv_block_len(
                 layer_idx=i, first_split=True, mamba_view=False
@@ -1317,7 +1349,32 @@ def add_remote_agent(
         # 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|
-        block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size)
+        # Compute block_size_ratio from actual byte-per-block values so that
+        # heterogeneous TP works even when block_size carries byte values
+        # (e.g. hybrid MLA+GDN models where block_size differs across TP
+        # configs).  _build_fa_remote already uses nixl_agent_meta.block_lens
+        # directly, so this ratio is only needed for handler registration and
+        # descriptor-ID computation.
+        if (
+            self.block_len_per_layer
+            and nixl_agent_meta.block_lens
+            and self.block_len_per_layer[0] != nixl_agent_meta.block_lens[0]
+        ):
+            local_bytes = self.block_len_per_layer[0]
+            remote_bytes = nixl_agent_meta.block_lens[0]
+            if local_bytes > remote_bytes and local_bytes % remote_bytes == 0:
+                block_size_ratio = local_bytes // remote_bytes
+            elif remote_bytes > local_bytes and remote_bytes % local_bytes == 0:
+                block_size_ratio = -(remote_bytes // local_bytes)
+            else:
+                # Non-exact byte division (e.g. hybrid models with
+                # TP-independent MLA component).  Use 1 as fallback;
+                # _build_fa_remote handles bytes via remote block_lens.
+                block_size_ratio = 1
+        else:
+            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
@@ -1342,9 +1399,12 @@ def add_remote_agent(
         plan = self.tp_mappings[engine_id]
 
         ### (Optional) Register local agent memory regions. MLA is not split.
+        # For hybrid MLA+GDN models, SSM state is TP-sharded and must be split
+        # to assemble data from multiple remote P ranks even when MLA
+        # attention is replicated.
         if (
             tp_ratio < 0
-            and not self.use_mla
+            and (not self.use_mla or self._has_mamba)
             and tp_ratio not in self.src_xfer_handles_by_tp_ratio
         ):
             # Remote tp_size > local tp_size: read from multiple remote ranks.
@@ -1425,9 +1485,14 @@ def _validate_remote_agent_handshake(
         assert remote_info.remote_tp_size == remote_tp_size
 
         tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
-        block_size_ratio = self.transfer_topo.block_size_ratio(
-            nixl_agent_meta.block_size
-        )
+        try:
+            block_size_ratio = self.transfer_topo.block_size_ratio(
+                nixl_agent_meta.block_size
+            )
+        except AssertionError:
+            # Heterogeneous TP with non-divisible block sizes (e.g. hybrid
+            # MLA+GDN).  Use 1 as a safe fallback for validation checks.
+            block_size_ratio = 1
         # 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.
         # MLA models do not need to handle kv replication.
@@ -1744,9 +1809,12 @@ def get_finished(self) -> tuple[set[str], set[str]]:
 
             # post processing for heteroblocksize
             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
-            )
+            try:
+                block_size_ratio = self.transfer_topo.block_size_ratio(
+                    remote_info.remote_block_size
+                )
+            except AssertionError:
+                block_size_ratio = 1
             if not self.use_mla and (
                 block_size_ratio > 1 or self.enable_permute_local_kv
             ):
@@ -2054,7 +2122,10 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
         # D may have to perform multiple reads from different remote ranks.
         # MLA opt: when P TP > D TP, only a single read is executed for
         # the first remote rank (cache is duplicated)..
-        if self.use_mla and tp_ratio < 0:
+        # For hybrid MLA+GDN models, SSM state is TP-sharded, so multiple
+        # remote ranks are still needed for the SSM group even when MLA
+        # attention only needs one rank.
+        if self.use_mla and tp_ratio < 0 and not self._has_mamba:
             assert len(read_specs) == 1
 
         for i, spec in enumerate(read_specs):
@@ -2068,17 +2139,23 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
                 req_id,
             )
             # Get side handles.
-            if tp_ratio < 0 and not self.use_mla:
-                assert remote_block_size == self.block_size
+            # For hybrid MLA+GDN with tp_ratio < 0, SSM needs split handles to
+            # assemble data from multiple remote ranks. MLA attention reads
+            # use the full region (replicated, single rank) but the split
+            # handle applies offset 0 + full chunk for FA when fa_num_splits=1.
+            if tp_ratio < 0 and (not self.use_mla or self._has_mamba):
                 # Remote tp_size > local tp_size: we must perform multiple
                 # reads. Get the memory chunk onto which we will write to.
                 local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
             else:
                 # Single read from remote, we write to the whole memory region.
                 # Also handle remote block size different from local block size.
-                local_xfer_side_handle = self.src_xfer_handles_by_block_size[
-                    remote_block_size
-                ]
+                # Use remote block_size handle if registered (block_size_ratio > 1),
+                # otherwise fall back to local block_size handle.
+                local_xfer_side_handle = self.src_xfer_handles_by_block_size.get(
+                    remote_block_size,
+                    self.src_xfer_handles_by_block_size[self.block_size],
+                )
 
             # Destination handle: remote_engine_id -> remote_rank -> handle.
             remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
@@ -2097,11 +2174,19 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
         if self.use_mla and tp_ratio < 0 and read_specs:
             # ..but we still need to notify the other remote ranks that we
             # have the blocks we need so they can update the request state.
+            # Only notify ranks that we actually read from (all_source_ranks),
+            # not all remote ranks — ranks we didn't read from don't track
+            # this request and would log "unrecognized request" errors.
             notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()
             remote_agents = self._remote_agents[meta.remote.engine_id]
-            for rank_to_notify, agent in remote_agents.items():
-                if rank_to_notify != read_specs[0].remote_rank:
-                    self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
+            for rank_to_notify in plan.all_source_ranks:
+                if (
+                    rank_to_notify != read_specs[0].remote_rank
+                    and rank_to_notify in remote_agents
+                ):
+                    self.nixl_wrapper.send_notif(
+                        remote_agents[rank_to_notify], notif_msg=notif_id
+                    )
 
     def _read_blocks(
         self,
@@ -2122,9 +2207,12 @@ def _read_blocks(
         remote_block_ids = read_spec.remote_block_ids
 
         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
-        )
+        try:
+            block_size_ratio = self.transfer_topo.block_size_ratio(
+                remote_info.remote_block_size
+            )
+        except AssertionError:
+            block_size_ratio = 1
         if block_size_ratio > 1:
             # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.
             assert not self._is_hma_required