diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base index d3ba0bc842..54025dd508 100644 --- a/docker/Dockerfile.rock_base +++ b/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" diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 1c59968005..15e3e2468e 100644 --- a/docker/Dockerfile.rocm_base +++ b/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" diff --git a/tests/kernels/moe/test_mori_prepare_finalize.py b/tests/kernels/moe/test_mori_prepare_finalize.py new file mode 100644 index 0000000000..c5deffe725 --- /dev/null +++ b/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)" + ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py index 75b97c69aa..743235bc7e 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py +++ b/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])