| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests that triton_kernel_moe_forward correctly applies expert_map |
| 4 | remapping when expert parallelism (EP) is enabled. |
| 5 | |
| 6 | When expert_map is provided, global expert IDs are remapped to local IDs |
| 7 | via topk + expert_map remap + make_routing_data before building routing |
| 8 | structures, and the expert_map passed downstream to triton_kernel_fused_experts |
| 9 | is None (already applied). |
| 10 | """ |
| 11 | |
| 12 | from unittest.mock import MagicMock, patch |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | |
| 17 | class TestTritonMoeForwardExpertMap: |
| 18 | """Test that triton_kernel_moe_forward applies expert_map remapping |
| 19 | when expert_map is provided (EP active).""" |
| 20 | |
| 21 | def test_expert_map_remap(self): |
| 22 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 23 | mock_expert_map = torch.tensor([0, -1, 1, -1], device=device) |
| 24 | |
| 25 | from vllm.utils.import_utils import import_triton_kernels |
| 26 | |
| 27 | import_triton_kernels() |
| 28 | |
| 29 | mock_routing_data = MagicMock() |
| 30 | mock_gather = MagicMock() |
| 31 | mock_scatter = MagicMock() |
| 32 | |
| 33 | with ( |
| 34 | patch("triton_kernels.topk.topk") as mock_topk, |
| 35 | patch( |
| 36 | "vllm.model_executor.layers.fused_moe.experts." |
| 37 | "gpt_oss_triton_kernels_moe.make_routing_data" |
| 38 | ) as mock_make_routing, |
| 39 | patch( |
| 40 | "vllm.model_executor.layers.fused_moe.experts." |
| 41 | "gpt_oss_triton_kernels_moe.triton_kernel_fused_experts" |
| 42 | ) as mock_fused_experts, |
| 43 | ): |
| 44 | from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 |
| 45 | triton_kernel_moe_forward, |
| 46 | ) |
| 47 | |
| 48 | sparse_result = MagicMock() |
| 49 | sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32) |
| 50 | sparse_result.vals = torch.tensor([[0.6, 0.4]]) |
| 51 | mock_topk.return_value = sparse_result |
| 52 | |
| 53 | mock_make_routing.return_value = ( |
| 54 | mock_routing_data, |
| 55 | mock_gather, |
| 56 | mock_scatter, |
| 57 | ) |
| 58 | |
| 59 | mock_fused_experts.return_value = torch.zeros((1, 8), device=device) |
| 60 | |
| 61 | hidden = torch.randn((1, 8), device=device) |
| 62 | w1 = torch.randn((2, 8, 16), device=device) |
| 63 | w2 = torch.randn((2, 8, 8), device=device) |
| 64 | logits = torch.randn((1, 4), device=device) |
| 65 | |
| 66 | triton_kernel_moe_forward( |
| 67 | hidden_states=hidden, |
| 68 | w1=w1, |
| 69 | w2=w2, |
| 70 | gating_output=logits, |
| 71 | topk=2, |
| 72 | renormalize=True, |
| 73 | expert_map=mock_expert_map, |
| 74 | ) |
| 75 | |
| 76 | mock_topk.assert_called_once() |
| 77 | mock_make_routing.assert_called_once() |
| 78 | |
| 79 | # expert_map should be None in the fused_experts call |
| 80 | # (already applied). |
| 81 | call_kwargs = mock_fused_experts.call_args |
| 82 | assert call_kwargs[1].get("expert_map") is None or (len(call_kwargs[0]) > 0) |