diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 45df75eeb3..811bedb084 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -477,7 +477,15 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata): remote_port=int(remote_handshake_port), remote_handshake_port=int(remote_handshake_port), remote_notify_port=int(remote_notify_port), - tp_size=kv_transfer_params.get("tp_size", 1), + # Remote peer TP degree (used as remote_tp_size downstream). The + # proxy advertises it under "remote_tp_size"; #46332 read "tp_size" + # which is absent on WRITE producer requests -> defaulted to 1 -> + # rank collapse. Read the right key; 0 == unknown (== homogeneous). + tp_size=int( + kv_transfer_params.get("remote_tp_size") + or kv_transfer_params.get("tp_size") + or 0 + ), remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0), ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index c478ced221..2a4c028e15 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1351,6 +1351,9 @@ class MoRIIOConnectorWorker: return {remote_agent_name} def _remote_tp_rank(self, remote_tp_size: int) -> int: + # 0/unknown remote TP == homogeneous (avoids collapsing all ranks to 0). + if remote_tp_size == 0: + remote_tp_size = self.world_size return get_moriio_remote_tp_rank(self.tp_rank, self.world_size, remote_tp_size) def _background_moriio_handshake( @@ -2145,7 +2148,9 @@ class MoRIIOConnectorWorker: validate_moriio_heterogeneous_tp_kv_heads( local_tp_size=self.world_size, remote_tp_size=( - remote_tp_size if remote_tp_size is not None else self.world_size + remote_tp_size + if remote_tp_size and remote_tp_size > 0 + else self.world_size ), total_num_kv_heads=self.model_config.get_total_num_kv_heads(), is_mla=self._is_mla_cache_layer(layer_name),