[ROCm] Upgrade MoRI version on rocm dockers required for WideEP DP16 DI CI enablement and fixes for combine API (#56073)
Signed-off-by: Shiksha Patel <shiksha.patel@amd.com> Signed-off-by: lcskrishna <lollachaitanya@gmail.com> Co-authored-by: shikamd123 <shikpate@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Chaitanya Sri Krishna Lolla committed
Sep 29, 2026 at 23:37 UTC
467d81d9a098c5f6695c67832a372a7ad8cd835a
4 files changed
+142
-3
docker/Dockerfile.rock_base
+1
-1
@@ -10,7 +10,7 @@ ARG FA_BRANCH="0e60e394"
10
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
11
ARG AITER_BRANCH="v0.1.23"
12
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
13
-ARG MORI_BRANCH="f95cd5e0ca7b730f92f9ffdfb358484d6965d9f7" # TheRock compatible v1.1.0
13
+ARG MORI_BRANCH="v1.2.3.post1"
14
ARG MORI_REPO="https://github.com/ROCm/mori.git"
15
ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git"
16
ARG ROCM_RUNTIME_COMMIT="25e14349f2606a6b95c4ccd45eb92a1835646cf2"
docker/Dockerfile.rocm_base
+1
-1
@@ -11,7 +11,7 @@ ARG FA_BRANCH="0e60e394"
11
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
12
ARG AITER_BRANCH="v0.1.23"
13
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14
-ARG MORI_BRANCH="v1.1.0"
14
+ARG MORI_BRANCH="v1.2.3.post1"
15
ARG MORI_REPO="https://github.com/ROCm/mori.git"
16
ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git"
17
ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350"
tests/kernels/moe/test_mori_prepare_finalize.py
new
+128
@@ -0,0 +1,128 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import importlib
5
+import sys
6
+import types
7
+from types import SimpleNamespace
8
+
9
+import torch
10
+
11
+
12
+def _describe(tensor):
13
+ if tensor is None:
14
+ return "None (the op was never called)"
15
+ return f"a different tensor of shape {tuple(tensor.shape)}"
16
+
17
+
18
+def _import_mori_prepare_finalize(monkeypatch):
19
+ # This is a contract test for vLLM's adapter; it does not need a real MoRI
20
+ # installation or GPU kernels.
21
+ fake_mori = types.ModuleType("mori")
22
+ fake_mori.ops = SimpleNamespace( # type: ignore[attr-defined]
23
+ EpDispatchCombineOp=object
24
+ )
25
+ monkeypatch.setitem(sys.modules, "mori", fake_mori)
26
+ module_name = "vllm.model_executor.layers.fused_moe.prepare_finalize.mori"
27
+ sys.modules.pop(module_name, None)
28
+ return importlib.import_module(module_name)
29
+
30
+
31
+class _FakeMoriOp:
32
+ def __init__(self):
33
+ self.dispatch_indices = None
34
+ self.combine_indices = None
35
+
36
+ def dispatch(self, a1, weights, scale, indices):
37
+ self.dispatch_indices = indices
38
+ # Model the global dispatched layout seen with EP=16: the experts need
39
+ # these IDs, but passing them to combine violates MoRI's contract.
40
+ dispatch_indices = indices.repeat(16, 1)
41
+ dispatch_a1 = a1.repeat(16, 1)
42
+ dispatch_weights = weights.repeat(16, 1)
43
+ recv_token_num = torch.tensor([dispatch_a1.shape[0]])
44
+ return (
45
+ dispatch_a1,
46
+ dispatch_weights,
47
+ scale,
48
+ dispatch_indices,
49
+ recv_token_num,
50
+ )
51
+
52
+ def combine(self, fused_expert_output, weights, indices):
53
+ self.combine_indices = indices
54
+ return (fused_expert_output,)
55
+
56
+
57
+def test_mori_combine_uses_original_per_rank_topk_ids(monkeypatch):
58
+ mori_module = _import_mori_prepare_finalize(monkeypatch)
59
+ mori_op = _FakeMoriOp()
60
+ adapter = mori_module.MoriPrepareAndFinalize(
61
+ mori_op=mori_op,
62
+ max_tokens_per_rank=16,
63
+ num_dispatchers=1,
64
+ )
65
+
66
+ hidden_states = torch.randn(2, 4)
67
+ router_topk_ids = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32)
68
+ topk_weights = torch.ones(2, 2)
69
+ quant_config = SimpleNamespace(
70
+ is_block_quantized=False,
71
+ is_per_act_token=False,
72
+ )
73
+
74
+ (
75
+ dispatch_a1,
76
+ _,
77
+ _,
78
+ dispatch_ids,
79
+ dispatch_weights,
80
+ ) = adapter.prepare(
81
+ hidden_states,
82
+ topk_weights,
83
+ router_topk_ids,
84
+ num_experts=4,
85
+ expert_map=None,
86
+ apply_router_weight_on_input=False,
87
+ quant_config=quant_config,
88
+ )
89
+
90
+ # Read the shapes before the identity asserts below, which narrow these
91
+ # names to the fake op's Optional attribute type.
92
+ router_shape = tuple(router_topk_ids.shape)
93
+ dispatch_shape = tuple(dispatch_ids.shape)
94
+ expected_rows = router_shape[0] * 16
95
+
96
+ assert dispatch_shape[0] == expected_rows, (
97
+ "prepare() must return the ids dispatch() produced: EP=16 expands "
98
+ f"{router_shape[0]} tokens to {expected_rows} rows, but prepare() "
99
+ f"returned {dispatch_shape[0]} rows"
100
+ )
101
+ assert mori_op.dispatch_indices is router_topk_ids, (
102
+ "dispatch() must be handed this rank's own pre-dispatch routing of "
103
+ f"shape {router_shape}, but it received "
104
+ f"{_describe(mori_op.dispatch_indices)}"
105
+ )
106
+
107
+ output = torch.empty_like(hidden_states)
108
+ adapter.finalize(
109
+ output,
110
+ dispatch_a1,
111
+ dispatch_weights,
112
+ dispatch_ids,
113
+ apply_router_weight_on_input=False,
114
+ weight_and_reduce_impl=SimpleNamespace(),
115
+ )
116
+
117
+ assert mori_op.combine_indices is router_topk_ids, (
118
+ "combine() reduces over this rank's own routing, so it must be handed "
119
+ f"the pre-dispatch tensor of shape {router_shape}, but it received "
120
+ f"{_describe(mori_op.combine_indices)}. finalize() most likely "
121
+ "forwarded its own topk_ids, which the modular kernel has already "
122
+ "rebound to the ids dispatch() returned (ROCm/mori#475)"
123
+ )
124
+ assert mori_op.combine_indices is not dispatch_ids, (
125
+ f"combine() was handed the dispatched ids of shape {dispatch_shape} "
126
+ "instead of this rank's pre-dispatch routing, which silently corrupts "
127
+ "the reduction (ROCm/mori#475)"
128
+ )
vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py
+12
-1
@@ -27,6 +27,7 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
27
self.num_dispatchers_ = num_dispatchers
28
self.max_tokens_per_rank = max_tokens_per_rank
29
self.use_fp8_dispatch = use_fp8_dispatch
30
+ self._dispatch_topk_ids: torch.Tensor | None = None
31
32
@property
33
def activation_format(self) -> mk.FusedMoEActivationFormat:
@@ -92,6 +93,12 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
93
# mori expects one scale slot per token; broadcast.
94
scale = scale.expand(a1.shape[0], 1).contiguous()
95
96
+ # mori's combine() reduces over this rank's own [num_tokens, topk]
97
+ # routing. The modular kernel rebinds topk_ids to the dispatched ids
98
+ # returned below before it calls finalize(), so hold a reference to
99
+ # the pre-dispatch tensor (ROCm/mori#475).
100
+ self._dispatch_topk_ids = topk_ids
101
+
102
(
103
dispatch_a1,
104
dispatch_weights,
@@ -122,9 +129,13 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
129
weight_and_reduce_impl: mk.TopKWeightAndReduce,
130
) -> None:
131
num_token = output.shape[0]
132
+ assert self._dispatch_topk_ids is not None, (
133
+ "finalize() called before prepare(); mori combine() needs the "
134
+ "routing tensor that dispatch() was given"
135
+ )
136
result = self.mori_op.combine(
137
fused_expert_output,
138
None,
128
- topk_ids,
139
+ self._dispatch_topk_ids,
140
)[0]
141
output.copy_(result[:num_token])