[ROCm][CI] Adapt MoE tests to the triton_kernels 3.8 API (#57385)

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

tests/kernels/moe/test_gpt_oss_triton_kernels.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn.functional as F
 
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
 
 if not has_triton_kernels():
     pytest.skip(
@@ -15,9 +15,15 @@ if not has_triton_kernels():
         allow_module_level=True,
     )
 
-import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
 import triton_kernels.swiglu
-from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
+
+if get_triton_kernels_version() == "3.8":
+    # 3.8: matmul_ogs -> matmul, matmul_ogs_details -> matmul_details
+    import triton_kernels.matmul_details.opt_flags as opt_flags
+    from triton_kernels.matmul import FlexCtx, PrecisionConfig
+else:
+    import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
+    from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
 from triton_kernels.numerics import InFlexData
 from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp
 from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
@@ -28,9 +34,10 @@ from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_co
 from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
     triton_kernel_moe_forward,
 )
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs
 from vllm.utils.math_utils import round_up
 
-from .utils import shuffle_weight
+from .utils import mxfp4_w_layouts, shuffle_weight
 
 
 def deshuffle(w: torch.Tensor):
@@ -133,12 +140,12 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int):
 
         x_tri = F.pad(x_tri, (0, x_pad, 0, 0), mode="constant", value=0)
 
-        w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
-        w_scale_layout, w_scale_layout_opts = (
-            layout.make_default_matmul_mxfp4_w_scale_layout(
-                mx_axis=1, num_warps=num_warps
-            )
-        )
+        (
+            w_layout,
+            w_layout_opts,
+            w_scale_layout,
+            w_scale_layout_opts,
+        ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps)
 
         w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1)
         w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, torch.bfloat16, axis=1)
@@ -165,10 +172,10 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int):
         )
 
         pc1 = PrecisionConfig(
-            weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
+            **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
         )
         pc2 = PrecisionConfig(
-            weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
+            **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
         )
 
         # tucuate so the rest can run properly
@@ -281,8 +288,6 @@ class Case:
 @pytest.mark.parametrize("num_token", [2])
 @pytest.mark.parametrize("tp", [1, 2, 4, 8])
 def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init):
-    from triton_kernels.tensor_details import layout
-
     if not hasattr(layout, "make_default_matmul_mxfp4_w_layout"):
         pytest.skip("make_default_matmul_mxfp4_w_layout not available")
 
@@ -387,6 +392,11 @@ def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk):
 
     if gptoss_moe.triton_kernels_version == "3.5.1":
         pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+")
+    if gptoss_moe.triton_kernels_version == "3.8":
+        pytest.skip(
+            "3.8 rebuilds routing from raw ids, so routing_data_from_sparse_topk "
+            "just forwards to make_routing_data and the comparison is vacuous"
+        )
 
     from triton_kernels.topk import topk as topk_fn
 

tests/kernels/moe/test_modular_oai_triton_moe.py

@@ -10,7 +10,7 @@ import torch.nn.functional as F
 
 from tests.utils import wait_for_gpu_memory_to_clear
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
-from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
 
 if not has_triton_kernels():
     pytest.skip(
@@ -18,11 +18,14 @@ if not has_triton_kernels():
         allow_module_level=True,
     )
 
-from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
+if get_triton_kernels_version() == "3.8":
+    # 3.8: matmul_ogs -> matmul
+    from triton_kernels.matmul import FlexCtx, PrecisionConfig
+else:
+    from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
 from triton_kernels.numerics import InFlexData
 from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp
 from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
-from triton_kernels.tensor_details import layout
 from triton_kernels.testing import assert_close
 
 from vllm.config import VllmConfig, set_current_vllm_config
@@ -35,11 +38,12 @@ from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe imp
     UnfusedOAITritonExperts,
 )
 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import round_up
 from vllm.utils.torch_utils import set_random_seed
 
-from .utils import make_dummy_moe_config, shuffle_weight
+from .utils import make_dummy_moe_config, mxfp4_w_layouts, shuffle_weight
 
 MNK = [
     (1, 512, 384),
@@ -130,10 +134,12 @@ def make_weights(dtype, k, n, e):
     w2 = w2[..., :n, :k]
 
     num_warps = 8
-    w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
-    w_scale_layout, w_scale_layout_opts = (
-        layout.make_default_matmul_mxfp4_w_scale_layout(mx_axis=1, num_warps=num_warps)
-    )
+    (
+        w_layout,
+        w_layout_opts,
+        w_scale_layout,
+        w_scale_layout_opts,
+    ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps)
 
     w1_tri = convert_layout(wrap_torch_tensor(w1_tri, FP4), w_layout, **w_layout_opts)
     w1_scale_tri = convert_layout(
@@ -150,10 +156,10 @@ def make_weights(dtype, k, n, e):
     )
 
     w1_precision_config = PrecisionConfig(
-        weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
+        **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
     )
     w2_precision_config = PrecisionConfig(
-        weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
+        **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
     )
 
     return (

tests/kernels/moe/utils.py

@@ -705,3 +705,27 @@ def check_accuracy(a, b, atol, rtol, percent):
             f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} "
             f"(threshold: {1 - percent:.4f})"
         )
+
+
+def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8):
+    """Weight/scale layouts for mxfp4 MoE, as (layout, opts) pairs.
+
+    triton_kernels 3.8 returns layout instances; earlier versions return a
+    (layout, opts) tuple.
+    """
+    from triton_kernels.tensor_details import layout
+
+    from vllm.utils.import_utils import get_triton_kernels_version
+
+    if get_triton_kernels_version() == "3.8":
+        w = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis)
+        s = layout.make_default_matmul_mxfp4_w_scale_layout(
+            mx_axis=mx_axis, num_warps=num_warps
+        )
+        return w, {}, s, {}
+
+    w, w_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis)
+    s, s_opts = layout.make_default_matmul_mxfp4_w_scale_layout(
+        mx_axis=mx_axis, num_warps=num_warps
+    )
+    return w, w_opts, s, s_opts

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

@@ -31,7 +31,12 @@ from vllm.model_executor.layers.fused_moe.config import (
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     swap_w13_to_w31,
 )
-from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+    _swizzle_mxfp4,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+    mx_scale_kwargs as _mx_scale_kwargs,
+)
 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
     OCP_MX_BLOCK_SIZE,
 )
@@ -70,15 +75,6 @@ if triton_kernels_version is not None:
         )
 
 
-def _mx_scale_kwargs(scale):
-    """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size,
-    3.5.1/3.6 use weight_scale.
-    """
-    if triton_kernels_version == "3.8":
-        return {"b_mx_scale": scale, "b_microblock_size": 32}
-    return {"weight_scale": scale}
-
-
 def _pack_deepgemm_mxfp4_scales(
     w13_weight: torch.Tensor,
     w2_weight: torch.Tensor,

vllm/model_executor/layers/quantization/utils/mxfp4_utils.py

@@ -36,6 +36,15 @@ def weight_mx_scale(precision_config):
     return precision_config.weight_scale
 
 
+def mx_scale_kwargs(scale):
+    """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size,
+    3.5.1/3.6 use weight_scale.
+    """
+    if get_triton_kernels_version() == "3.8":
+        return {"b_mx_scale": scale, "b_microblock_size": 32}
+    return {"weight_scale": scale}
+
+
 def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
     """Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel."""
     assert has_triton_kernels()