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