[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>

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),
         )

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),