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

docker/Dockerfile.rock_base

@@ -10,7 +10,7 @@ ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
 ARG AITER_BRANCH="v0.1.23"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
-ARG MORI_BRANCH="f95cd5e0ca7b730f92f9ffdfb358484d6965d9f7" # TheRock compatible v1.1.0
+ARG MORI_BRANCH="v1.2.3.post1"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
 ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git"
 ARG ROCM_RUNTIME_COMMIT="25e14349f2606a6b95c4ccd45eb92a1835646cf2"

docker/Dockerfile.rocm_base

@@ -11,7 +11,7 @@ ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
 ARG AITER_BRANCH="v0.1.23"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
-ARG MORI_BRANCH="v1.1.0"
+ARG MORI_BRANCH="v1.2.3.post1"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
 ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git"
 ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350"

tests/kernels/moe/test_mori_prepare_finalize.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib
+import sys
+import types
+from types import SimpleNamespace
+
+import torch
+
+
+def _describe(tensor):
+    if tensor is None:
+        return "None (the op was never called)"
+    return f"a different tensor of shape {tuple(tensor.shape)}"
+
+
+def _import_mori_prepare_finalize(monkeypatch):
+    # This is a contract test for vLLM's adapter; it does not need a real MoRI
+    # installation or GPU kernels.
+    fake_mori = types.ModuleType("mori")
+    fake_mori.ops = SimpleNamespace(  # type: ignore[attr-defined]
+        EpDispatchCombineOp=object
+    )
+    monkeypatch.setitem(sys.modules, "mori", fake_mori)
+    module_name = "vllm.model_executor.layers.fused_moe.prepare_finalize.mori"
+    sys.modules.pop(module_name, None)
+    return importlib.import_module(module_name)
+
+
+class _FakeMoriOp:
+    def __init__(self):
+        self.dispatch_indices = None
+        self.combine_indices = None
+
+    def dispatch(self, a1, weights, scale, indices):
+        self.dispatch_indices = indices
+        # Model the global dispatched layout seen with EP=16: the experts need
+        # these IDs, but passing them to combine violates MoRI's contract.
+        dispatch_indices = indices.repeat(16, 1)
+        dispatch_a1 = a1.repeat(16, 1)
+        dispatch_weights = weights.repeat(16, 1)
+        recv_token_num = torch.tensor([dispatch_a1.shape[0]])
+        return (
+            dispatch_a1,
+            dispatch_weights,
+            scale,
+            dispatch_indices,
+            recv_token_num,
+        )
+
+    def combine(self, fused_expert_output, weights, indices):
+        self.combine_indices = indices
+        return (fused_expert_output,)
+
+
+def test_mori_combine_uses_original_per_rank_topk_ids(monkeypatch):
+    mori_module = _import_mori_prepare_finalize(monkeypatch)
+    mori_op = _FakeMoriOp()
+    adapter = mori_module.MoriPrepareAndFinalize(
+        mori_op=mori_op,
+        max_tokens_per_rank=16,
+        num_dispatchers=1,
+    )
+
+    hidden_states = torch.randn(2, 4)
+    router_topk_ids = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32)
+    topk_weights = torch.ones(2, 2)
+    quant_config = SimpleNamespace(
+        is_block_quantized=False,
+        is_per_act_token=False,
+    )
+
+    (
+        dispatch_a1,
+        _,
+        _,
+        dispatch_ids,
+        dispatch_weights,
+    ) = adapter.prepare(
+        hidden_states,
+        topk_weights,
+        router_topk_ids,
+        num_experts=4,
+        expert_map=None,
+        apply_router_weight_on_input=False,
+        quant_config=quant_config,
+    )
+
+    # Read the shapes before the identity asserts below, which narrow these
+    # names to the fake op's Optional attribute type.
+    router_shape = tuple(router_topk_ids.shape)
+    dispatch_shape = tuple(dispatch_ids.shape)
+    expected_rows = router_shape[0] * 16
+
+    assert dispatch_shape[0] == expected_rows, (
+        "prepare() must return the ids dispatch() produced: EP=16 expands "
+        f"{router_shape[0]} tokens to {expected_rows} rows, but prepare() "
+        f"returned {dispatch_shape[0]} rows"
+    )
+    assert mori_op.dispatch_indices is router_topk_ids, (
+        "dispatch() must be handed this rank's own pre-dispatch routing of "
+        f"shape {router_shape}, but it received "
+        f"{_describe(mori_op.dispatch_indices)}"
+    )
+
+    output = torch.empty_like(hidden_states)
+    adapter.finalize(
+        output,
+        dispatch_a1,
+        dispatch_weights,
+        dispatch_ids,
+        apply_router_weight_on_input=False,
+        weight_and_reduce_impl=SimpleNamespace(),
+    )
+
+    assert mori_op.combine_indices is router_topk_ids, (
+        "combine() reduces over this rank's own routing, so it must be handed "
+        f"the pre-dispatch tensor of shape {router_shape}, but it received "
+        f"{_describe(mori_op.combine_indices)}. finalize() most likely "
+        "forwarded its own topk_ids, which the modular kernel has already "
+        "rebound to the ids dispatch() returned (ROCm/mori#475)"
+    )
+    assert mori_op.combine_indices is not dispatch_ids, (
+        f"combine() was handed the dispatched ids of shape {dispatch_shape} "
+        "instead of this rank's pre-dispatch routing, which silently corrupts "
+        "the reduction (ROCm/mori#475)"
+    )

vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py

@@ -27,6 +27,7 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         self.num_dispatchers_ = num_dispatchers
         self.max_tokens_per_rank = max_tokens_per_rank
         self.use_fp8_dispatch = use_fp8_dispatch
+        self._dispatch_topk_ids: torch.Tensor | None = None
 
     @property
     def activation_format(self) -> mk.FusedMoEActivationFormat:
@@ -92,6 +93,12 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
                 # mori expects one scale slot per token; broadcast.
                 scale = scale.expand(a1.shape[0], 1).contiguous()
 
+        # mori's combine() reduces over this rank's own [num_tokens, topk]
+        # routing. The modular kernel rebinds topk_ids to the dispatched ids
+        # returned below before it calls finalize(), so hold a reference to
+        # the pre-dispatch tensor (ROCm/mori#475).
+        self._dispatch_topk_ids = topk_ids
+
         (
             dispatch_a1,
             dispatch_weights,
@@ -122,9 +129,13 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         weight_and_reduce_impl: mk.TopKWeightAndReduce,
     ) -> None:
         num_token = output.shape[0]
+        assert self._dispatch_topk_ids is not None, (
+            "finalize() called before prepare(); mori combine() needs the "
+            "routing tensor that dispatch() was given"
+        )
         result = self.mori_op.combine(
             fused_expert_output,
             None,
-            topk_ids,
+            self._dispatch_topk_ids,
         )[0]
         output.copy_(result[:num_token])