[ROCm] Add XGMI backend for MoRI Connector (#41753)

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py

@@ -189,6 +189,7 @@ class MoRIIOConfig:
     dp_rank: int
     dp_size: int
     tp_size: int
+    backend: str = "rdma"
 
     @classmethod
     def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
@@ -213,6 +214,12 @@ class MoRIIOConfig:
         dp_size = vllm_config.parallel_config.data_parallel_size
         tp_size = get_tensor_model_parallel_world_size()
         port_offset = get_port_offset(dp_rank, tp_rank)
+        backend = str(extra_config.get("backend", "rdma")).lower()
+        if backend not in ("rdma", "xgmi"):
+            raise ValueError(
+                f"Invalid MoRIIO backend {backend!r} in kv_connector_extra_config; "
+                "must be one of 'rdma' or 'xgmi'."
+            )
 
         return cls(
             local_ip=get_ip(),
@@ -227,6 +234,7 @@ class MoRIIOConfig:
             dp_rank=dp_rank,
             dp_size=dp_size,
             tp_size=tp_size,
+            backend=backend,
         )
 
 

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

@@ -695,7 +695,12 @@ class MoRIIOConnectorWorker:
         # Agent.
         self.moriio_wrapper = MoRIIOWrapper(tp_rank=self.tp_rank, dp_rank=self.dp_rank)
         self.moriio_wrapper.set_moriio_engine(self.moriio_engine)
-        self.moriio_wrapper.set_backend_type(BackendType.RDMA)
+        backend = (
+            BackendType.XGMI
+            if self.moriio_config.backend == "xgmi"
+            else BackendType.RDMA
+        )
+        self.moriio_wrapper.set_backend_type(backend)
         self.moriio_wrapper.notify_port = self.moriio_config.notify_port
         self.local_kv_cache_metadata: list[bytes] = []
         self.local_kv_cache_size: list[int] = []

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py

@@ -44,11 +44,13 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 try:
     from mori.io import (
+        BackendType,
         EngineDesc,
         IOEngine,
         MemoryDesc,
         PollCqMode,
         RdmaBackendConfig,
+        XgmiBackendConfig,
     )
 
     logger.info("MoRIIO is available")
@@ -376,17 +378,24 @@ class MoRIIOWrapper:
 
     def set_backend_type(self, backend_type):
         assert self.moriio_engine is not None, "MoRIIO engine must be set first"
-        qp_per_transfer = envs.VLLM_MORIIO_QP_PER_TRANSFER
-        post_batch_size = envs.VLLM_MORIIO_POST_BATCH_SIZE
-        num_worker_threads = envs.VLLM_MORIIO_NUM_WORKERS
-        poll_mode = PollCqMode.POLLING
-        rdma_cfg = RdmaBackendConfig(
-            qp_per_transfer,
-            post_batch_size,
-            num_worker_threads,
-            poll_mode,
-        )
-        self.moriio_engine.create_backend(backend_type, rdma_cfg)
+        if backend_type == BackendType.XGMI:
+            logger.info("Using MoRIIO backend: XGMI")
+            self.moriio_engine.create_backend(backend_type, XgmiBackendConfig())
+        else:
+            logger.info(
+                "Using MoRIIO backend: RDMA "
+                "(qp_per_transfer=%d, post_batch_size=%d, num_workers=%d)",
+                envs.VLLM_MORIIO_QP_PER_TRANSFER,
+                envs.VLLM_MORIIO_POST_BATCH_SIZE,
+                envs.VLLM_MORIIO_NUM_WORKERS,
+            )
+            rdma_cfg = RdmaBackendConfig(
+                envs.VLLM_MORIIO_QP_PER_TRANSFER,
+                envs.VLLM_MORIIO_POST_BATCH_SIZE,
+                envs.VLLM_MORIIO_NUM_WORKERS,
+                PollCqMode.POLLING,
+            )
+            self.moriio_engine.create_backend(backend_type, rdma_cfg)
 
     def get_agent_metadata(self):
         assert self.moriio_engine is not None, "MoRIIO engine must be set first"