A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 03:00 UTC
[ROCm] Add XGMI backend for MoRI Connector (#41753)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
May 22, 2026 at 05:06 UTC
86ccef7d4400a54441057773d8ffb1f61a20af94
3 files changed
+34
-12
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+8
@@ -189,6 +189,7 @@ class MoRIIOConfig:
189
dp_rank: int
190
dp_size: int
191
tp_size: int
192
+ backend: str = "rdma"
193
194
@classmethod
195
def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
@@ -213,6 +214,12 @@ class MoRIIOConfig:
214
dp_size = vllm_config.parallel_config.data_parallel_size
215
tp_size = get_tensor_model_parallel_world_size()
216
port_offset = get_port_offset(dp_rank, tp_rank)
217
+ backend = str(extra_config.get("backend", "rdma")).lower()
218
+ if backend not in ("rdma", "xgmi"):
219
+ raise ValueError(
220
+ f"Invalid MoRIIO backend {backend!r} in kv_connector_extra_config; "
221
+ "must be one of 'rdma' or 'xgmi'."
222
+ )
223
224
return cls(
225
local_ip=get_ip(),
@@ -227,6 +234,7 @@ class MoRIIOConfig:
234
dp_rank=dp_rank,
235
dp_size=dp_size,
236
tp_size=tp_size,
237
+ backend=backend,
238
)
239
240
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+6
-1
@@ -695,7 +695,12 @@ class MoRIIOConnectorWorker:
695
# Agent.
696
self.moriio_wrapper = MoRIIOWrapper(tp_rank=self.tp_rank, dp_rank=self.dp_rank)
697
self.moriio_wrapper.set_moriio_engine(self.moriio_engine)
698
- self.moriio_wrapper.set_backend_type(BackendType.RDMA)
698
+ backend = (
699
+ BackendType.XGMI
700
+ if self.moriio_config.backend == "xgmi"
701
+ else BackendType.RDMA
702
+ )
703
+ self.moriio_wrapper.set_backend_type(backend)
704
self.moriio_wrapper.notify_port = self.moriio_config.notify_port
705
self.local_kv_cache_metadata: list[bytes] = []
706
self.local_kv_cache_size: list[int] = []
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+20
-11
@@ -44,11 +44,13 @@ if TYPE_CHECKING:
44
logger = init_logger(__name__)
45
try:
46
from mori.io import (
47
+ BackendType,
48
EngineDesc,
49
IOEngine,
50
MemoryDesc,
51
PollCqMode,
52
RdmaBackendConfig,
53
+ XgmiBackendConfig,
54
)
55
56
logger.info("MoRIIO is available")
@@ -376,17 +378,24 @@ class MoRIIOWrapper:
378
379
def set_backend_type(self, backend_type):
380
assert self.moriio_engine is not None, "MoRIIO engine must be set first"
379
- qp_per_transfer = envs.VLLM_MORIIO_QP_PER_TRANSFER
380
- post_batch_size = envs.VLLM_MORIIO_POST_BATCH_SIZE
381
- num_worker_threads = envs.VLLM_MORIIO_NUM_WORKERS
382
- poll_mode = PollCqMode.POLLING
383
- rdma_cfg = RdmaBackendConfig(
384
- qp_per_transfer,
385
- post_batch_size,
386
- num_worker_threads,
387
- poll_mode,
388
- )
389
- self.moriio_engine.create_backend(backend_type, rdma_cfg)
381
+ if backend_type == BackendType.XGMI:
382
+ logger.info("Using MoRIIO backend: XGMI")
383
+ self.moriio_engine.create_backend(backend_type, XgmiBackendConfig())
384
+ else:
385
+ logger.info(
386
+ "Using MoRIIO backend: RDMA "
387
+ "(qp_per_transfer=%d, post_batch_size=%d, num_workers=%d)",
388
+ envs.VLLM_MORIIO_QP_PER_TRANSFER,
389
+ envs.VLLM_MORIIO_POST_BATCH_SIZE,
390
+ envs.VLLM_MORIIO_NUM_WORKERS,
391
+ )
392
+ rdma_cfg = RdmaBackendConfig(
393
+ envs.VLLM_MORIIO_QP_PER_TRANSFER,
394
+ envs.VLLM_MORIIO_POST_BATCH_SIZE,
395
+ envs.VLLM_MORIIO_NUM_WORKERS,
396
+ PollCqMode.POLLING,
397
+ )
398
+ self.moriio_engine.create_backend(backend_type, rdma_cfg)
399
400
def get_agent_metadata(self):
401
assert self.moriio_engine is not None, "MoRIIO engine must be set first"