Repository: vllm-project/vllm
File: tests/kernels/quantization/test_mxfp4_triton_ep.py (Python)

Code before:
        mock_get_config,
        mock_get_backend,
        backend,
        ep_size,
        expected_monolithic,
    ):
        """is_monolithic should be True for TRITON regardless of EP size."""
        mock_get_backend.return_value = backend

        mock_compilation_config = MagicMock()
        mock_compilation_config.max_cudagraph_capture_size = 1024
        mock_vllm_config = MagicMock()
        mock_vllm_config.compilation_config = mock_compilation_config
        mock_get_config.return_value = mock_vllm_config

        moe_config = _make_mock_moe_config(ep_size=ep_size)
        method = Mxfp4MoEMethod(moe_config)

        assert method.is_monolithic == expected_monolithic, (
            f"Expected is_monolithic={expected_monolithic} for "
            f"backend={backend.name}, ep_size={ep_size}, "
            f"but got {method.is_monolithic}."
        )


class TestTritonMoeForwardExpertMap:
    """Test that triton_kernel_moe_forward applies expert_map remapping
    when expert_map is provided (EP active)."""

    @pytest.mark.parametrize("expert_map_present", [False, True])
    def test_routing_path_selection(self, expert_map_present):
        """Verify that the EP-aware routing path is taken when expert_map
        is present, and the legacy_routing path is taken otherwise."""
        # This is a structural test: we mock the routing functions to
        # verify the correct path is exercised.
        mock_expert_map = torch.tensor([0, -1, 1, -1]) if expert_map_present else None

        with (
            patch(
                "vllm.model_executor.layers.fused_moe."
                "gpt_oss_triton_kernels_moe.legacy_routing"
            ) as mock_legacy,
            patch(
                "vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.topk"
            ) as mock_topk,
            patch(
                "vllm.model_executor.layers.fused_moe."
                "gpt_oss_triton_kernels_moe.make_routing_data"
            ) as mock_make_routing,
            patch(
                "vllm.model_executor.layers.fused_moe."

Review comment (suggestion):
@elizabetht  This test doesn't pass for me on H100. Can you fix it by 
```
diff --git a/tests/kernels/quantization/test_mxfp4_triton_ep.py b/tests/kernels/quantization/test_mxfp4_triton_ep.py
index 5ef21810b..7f1e62280 100644
--- a/tests/kernels/quantization/test_mxfp4_triton_ep.py
+++ b/tests/kernels/quantization/test_mxfp4_triton_ep.py
@@ -109,9 +109,12 @@ class TestTritonMoeForwardExpertMap:
     def test_routing_path_selection(self, expert_map_present):
         """Verify that the EP-aware routing path is taken when expert_map
         is present, and the legacy_routing path is taken otherwise."""
+
+        device = "cuda"
+
         # This is a structural test: we mock the routing functions to
         # verify the correct path is exercised.
-        mock_expert_map = torch.tensor([0, -1, 1, -1]) if expert_map_present else None
+        mock_expert_map = torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None
 
         with (
             patch(
@@ -119,7 +122,7 @@ class TestTritonMoeForwardExpertMap:
                 "gpt_oss_triton_kernels_moe.legacy_routing"
             ) as mock_legacy,
             patch(
-                "vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.topk"
+                "triton_kernels.topk.topk"
             ) as mock_topk,
             patch(
                 "vllm.model_executor.layers.fused_moe."
@@ -158,10 +161,10 @@ class TestTritonMoeForwardExpertMap:
 
             mock_fused_experts.return_value = torch.zeros(1, 8)
 
-            hidden = torch.randn(1, 8)
-            w1 = torch.randn(2, 8, 16)
-            w2 = torch.randn(2, 8, 8)
-            logits = torch.randn(1, 4)
+            hidden = torch.randn((1, 8), device = device)
+            w1 = torch.randn((2, 8, 16), device=device)
+            w2 = torch.randn((2, 8, 8), device=device)
+            logits = torch.randn((1, 4), device =device)
 
             triton_kernel_moe_forward(
                 hidden_states=hidden,
```

Code after:
        mock_get_config,
        mock_get_backend,
        backend,
        ep_size,
        expected_monolithic,
    ):
        """is_monolithic should be True for TRITON regardless of EP size."""
        mock_get_backend.return_value = backend

        mock_compilation_config = MagicMock()
        mock_compilation_config.max_cudagraph_capture_size = 1024
        mock_vllm_config = MagicMock()
        mock_vllm_config.compilation_config = mock_compilation_config
        mock_get_config.return_value = mock_vllm_config

        moe_config = _make_mock_moe_config(ep_size=ep_size)
        method = Mxfp4MoEMethod(moe_config)

        assert method.is_monolithic == expected_monolithic, (
            f"Expected is_monolithic={expected_monolithic} for "
            f"backend={backend.name}, ep_size={ep_size}, "
            f"but got {method.is_monolithic}."
        )


class TestTritonMoeForwardExpertMap:
    """Test that triton_kernel_moe_forward applies expert_map remapping
    when expert_map is provided (EP active)."""

    @pytest.mark.parametrize("expert_map_present", [False, True])
    def test_routing_path_selection(self, expert_map_present):
        """Verify that the EP-aware routing path is taken when expert_map
        is present, and the legacy_routing path is taken otherwise."""

        device = "cuda" if torch.cuda.is_available() else "cpu"
        # This is a structural test: we mock the routing functions to
        # verify the correct path is exercised.
        mock_expert_map = (
            torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None
        )

        with (
            patch(
                "vllm.model_executor.layers.fused_moe."
                "gpt_oss_triton_kernels_moe.legacy_routing"
            ) as mock_legacy,
            patch("triton_kernels.topk.topk") as mock_topk,
            patch(
                "vllm.model_executor.layers.fused_moe."
                "gpt_oss_triton_kernels_moe.make_routing_data"
            ) as mock_make_routing,