[ROCm][KVConnector][MoRI-IO] Fix WRITE-mode remote-TP rank collapse (#46332 follow-up) (#47764)
Signed-off-by: avininjamay8 <avininjamay8@users.noreply.github.com> Signed-off-by: avininjamay8 <Avinash.Paul@amd.com> Co-authored-by: avininjamay8 <avininjamay8@users.noreply.github.com> Co-authored-by: avininjamay8 <avpaul@amd.com>
avininjamay8 committed
Jul 28, 2026 at 06:56 UTC
1206891822ca8befe421879a4230ef42a3fc93be
2 files changed
+15
-2
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+9
-1
@@ -477,7 +477,15 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
477
remote_port=int(remote_handshake_port),
478
remote_handshake_port=int(remote_handshake_port),
479
remote_notify_port=int(remote_notify_port),
480
- tp_size=kv_transfer_params.get("tp_size", 1),
480
+ # Remote peer TP degree (used as remote_tp_size downstream). The
481
+ # proxy advertises it under "remote_tp_size"; #46332 read "tp_size"
482
+ # which is absent on WRITE producer requests -> defaulted to 1 ->
483
+ # rank collapse. Read the right key; 0 == unknown (== homogeneous).
484
+ tp_size=int(
485
+ kv_transfer_params.get("remote_tp_size")
486
+ or kv_transfer_params.get("tp_size")
487
+ or 0
488
+ ),
489
remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
490
remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0),
491
)
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+6
-1
@@ -1351,6 +1351,9 @@ class MoRIIOConnectorWorker:
1351
return {remote_agent_name}
1352
1353
def _remote_tp_rank(self, remote_tp_size: int) -> int:
1354
+ # 0/unknown remote TP == homogeneous (avoids collapsing all ranks to 0).
1355
+ if remote_tp_size == 0:
1356
+ remote_tp_size = self.world_size
1357
return get_moriio_remote_tp_rank(self.tp_rank, self.world_size, remote_tp_size)
1358
1359
def _background_moriio_handshake(
@@ -2145,7 +2148,9 @@ class MoRIIOConnectorWorker:
2148
validate_moriio_heterogeneous_tp_kv_heads(
2149
local_tp_size=self.world_size,
2150
remote_tp_size=(
2148
- remote_tp_size if remote_tp_size is not None else self.world_size
2151
+ remote_tp_size
2152
+ if remote_tp_size and remote_tp_size > 0
2153
+ else self.world_size
2154
),
2155
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
2156
is_mla=self._is_mla_cache_layer(layer_name),