[PD][Bugfix] Fix NIXL hybrid MLA+mamba heterogeneous TP (#49297)
Signed-off-by: ZeldaHuang <zelda.huanghuang@gmail.com>
Ziming Huang committed
Jul 22, 2026 at 15:11 UTC
7c21548ce38ffe691026edac4eac5a731757283b
2 files changed
+10
-8
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+5
-3
@@ -1598,12 +1598,14 @@ class NixlBaseConnectorWorker:
1598
### (Optional) Register local agent memory regions. MLA is not split.
1599
if (
1600
tp_ratio < 0
1601
- and not self.use_mla
1601
+ and (not self.use_mla or len(plan.all_source_ranks) > 1)
1602
and tp_ratio not in self.src_xfer_handles_by_tp_ratio
1603
):
1604
# Remote tp_size > local tp_size: read from multiple remote ranks.
1605
- # Logically "split" own regions into |tp_ratio| chunks. Mind that
1606
- # we only do this once per remote tp_size (replica-friendly).
1605
+ # Logically "split" own regions into per-source chunks. Hybrid
1606
+ # MLA+SSM also needs this path: MLA is replicated and read once,
1607
+ # while the SSM state is sharded across every remote TP rank.
1608
+ # We only do this once per remote tp_size (replica-friendly).
1609
self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
1610
1611
for handle_data in self._build_local_splits_from_plan(
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+5
-5
@@ -161,9 +161,9 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
161
]
162
163
# D may have to perform multiple reads from different remote ranks.
164
- # MLA opt: when P TP > D TP, only a single read is executed for
165
- # the first remote rank (cache is duplicated)..
166
- if self.use_mla and tp_ratio < 0:
164
+ # Pure MLA reads once because its cache is replicated. Hybrid
165
+ # MLA+SSM still needs one read per SSM source rank.
166
+ if self.use_mla and tp_ratio < 0 and not self._has_mamba:
167
assert len(read_specs) == 1
168
169
for i, spec in enumerate(read_specs):
@@ -177,7 +177,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
177
req_id,
178
)
179
# Get side handles.
180
- if tp_ratio < 0 and not self.use_mla:
180
+ if tp_ratio < 0 and (not self.use_mla or len(read_specs) > 1):
181
assert remote_block_size == self.block_size
182
# Remote tp_size > local tp_size: we must perform multiple
183
# reads. Get the memory chunk onto which we will write to.
@@ -203,7 +203,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
203
remote_xfer_side_handle=remote_xfer_side_handle,
204
)
205
206
- if self.use_mla and tp_ratio < 0 and read_specs:
206
+ if self.use_mla and tp_ratio < 0 and len(read_specs) == 1:
207
# ..but we still need to notify the other remote ranks that we
208
# have the blocks we need so they can update the request state.
209
notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()