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 2, 18:15 UTC
[ROCm] mori: add InterNodeV1LL inter-node kernel selection via VLLM_MORI_INTERNODE_KERNEL (#41751)
Signed-off-by: jatseng-ai <jatseng@amd.com>
jatseng-ai committed
May 27, 2026 at 09:33 UTC
05c50c721e07035dd36169fa8bb2825a7b952555
7 files changed
+36
-14
tests/kernels/moe/modular_kernel_tools/common.py
+1
-1
@@ -249,7 +249,7 @@ class Config:
249
250
def needs_mori(self):
251
info = prepare_finalize_info(self.prepare_finalize_type)
252
- return info.backend == "mori"
252
+ return info.backend in ("mori_high_throughput", "mori_low_latency")
253
254
def all2all_backend(self):
255
info = prepare_finalize_info(self.prepare_finalize_type)
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+1
-1
@@ -232,7 +232,7 @@ if has_mori():
232
standard_format,
233
fp8_types,
234
blocked_quantization_support=True,
235
- backend="mori",
235
+ backend="mori_high_throughput",
236
supports_apply_weight_on_input=False,
237
)
238
tests/kernels/moe/test_moe_layer.py
+5
-3
@@ -92,7 +92,7 @@ PARALLEL_COMBOS = [
92
BACKENDS = ["allgather_reducescatter"]
93
94
if has_mori():
95
- BACKENDS += ["mori"]
95
+ BACKENDS += ["mori_high_throughput", "mori_low_latency"]
96
97
if has_flashinfer_nvlink_two_sided():
98
BACKENDS += ["flashinfer_nvlink_two_sided"]
@@ -118,7 +118,8 @@ QUANT_METHODS = [
118
# fmt: off
119
BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
120
"allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
121
- "mori": {None, "fp8", "modelopt_fp8"},
121
+ "mori_high_throughput": {None, "fp8", "modelopt_fp8"},
122
+ "mori_low_latency": {None, "fp8", "modelopt_fp8"},
123
"flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
124
"flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501
125
"deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
@@ -129,7 +130,8 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
130
# Map from backend -> (DP/EP support, DP support, TP support, SP support)
131
BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool, bool]] = {
132
"allgather_reducescatter": (True, True, True, True),
132
- "mori": (True, False, False, True),
133
+ "mori_high_throughput": (True, False, False, True),
134
+ "mori_low_latency": (True, False, False, True),
135
"flashinfer_nvlink_two_sided": (False, True, False, False),
136
"flashinfer_nvlink_one_sided": (False, True, False, False),
137
"deepep_low_latency": (True, False, False, True),
vllm/config/parallel.py
+6
-3
@@ -42,7 +42,8 @@ All2AllBackend = Literal[
42
"pplx",
43
"deepep_high_throughput",
44
"deepep_low_latency",
45
- "mori",
45
+ "mori_high_throughput",
46
+ "mori_low_latency",
47
"nixl_ep",
48
"allgather_reducescatter",
49
"flashinfer_all2allv", # temporary alias for flashinfer_nvlink_two_sided
@@ -177,7 +178,8 @@ class ParallelConfig:
178
- "allgather_reducescatter": All2all based on allgather and reducescatter
179
- "deepep_high_throughput": Use deepep high-throughput kernels
180
- "deepep_low_latency": Use deepep low-latency kernels
180
- - "mori": Use mori kernels
181
+ - "mori_high_throughput": MoRI EP with InterNodeV1 for multi-node
182
+ - "mori_low_latency": MoRI EP with InterNodeV1LL for multi-node
183
- "nixl_ep": Use nixl-ep kernels
184
- "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl
185
- "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels"""
@@ -621,7 +623,8 @@ class ParallelConfig:
623
"allgather_reducescatter",
624
"deepep_high_throughput",
625
"deepep_low_latency",
624
- "mori",
626
+ "mori_high_throughput",
627
+ "mori_low_latency",
628
"nixl_ep",
629
)
630
and self.enable_expert_parallel
vllm/distributed/device_communicators/all2all.py
+12
-3
@@ -764,14 +764,19 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
764
765
766
class MoriAll2AllManager(All2AllManagerBase):
767
- def __init__(self, cpu_group):
767
+ def __init__(self, cpu_group, all2all_backend: str):
768
assert has_mori(), (
769
"MoRI kernels not found. Please follow https://github.com/ROCm/mori/blob/main/README.md"
770
" to install MoRI kernels."
771
) # noqa
772
+ assert all2all_backend in (
773
+ "mori_high_throughput",
774
+ "mori_low_latency",
775
+ ), f"unsupported MoRI all2all backend: {all2all_backend!r}"
776
import mori
777
778
super().__init__(cpu_group)
779
+ self._all2all_backend = all2all_backend
780
self.handle_cache = Cache()
781
782
torch._C._distributed_c10d._register_process_group("mori", cpu_group)
@@ -805,8 +810,12 @@ class MoriAll2AllManager(All2AllManagerBase):
810
warp_num_per_block = 16
811
block_num = 80
812
else:
808
- # multi node
809
- kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1
813
+ # Multi-node: kernel follows --all2all-backend (mirrors deepep_* split).
814
+ # mori_low_latency → InterNodeV1LL; mori_high_throughput → V1.
815
+ if self._all2all_backend == "mori_low_latency":
816
+ kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1LL
817
+ else:
818
+ kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1
819
if on_gfx942():
820
warp_num_per_block = 16
821
block_num = 32
vllm/distributed/device_communicators/cuda_communicator.py
+7
-2
@@ -137,10 +137,15 @@ class CudaCommunicator(DeviceCommunicatorBase):
137
self.all2all_manager = DeepEPLLAll2AllManager(
138
self.cpu_group, tcp_store_group
139
)
140
- elif self.all2all_backend == "mori":
140
+ elif self.all2all_backend in (
141
+ "mori_high_throughput",
142
+ "mori_low_latency",
143
+ ):
144
from .all2all import MoriAll2AllManager
145
143
- self.all2all_manager = MoriAll2AllManager(self.cpu_group)
146
+ self.all2all_manager = MoriAll2AllManager(
147
+ self.cpu_group, self.all2all_backend
148
+ )
149
elif self.all2all_backend == "nixl_ep":
150
from .all2all import NixlEPAll2AllManager
151
vllm/model_executor/layers/fused_moe/config.py
+4
-1
@@ -1091,7 +1091,10 @@ class FusedMoEParallelConfig:
1091
1092
@property
1093
def use_mori_kernels(self):
1094
- return self.use_all2all_kernels and self.all2all_backend == "mori"
1094
+ return self.use_all2all_kernels and self.all2all_backend in (
1095
+ "mori_high_throughput",
1096
+ "mori_low_latency",
1097
+ )
1098
1099
@property
1100
def use_nixl_ep_kernels(self):