From fef55fd26dc633fcc87978f12f1a97e08e351c02 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:21:22 -0700 Subject: [PATCH] [ROCm][CI] Adapt MoE tests to the triton_kernels 3.8 API (#57385) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5 (1M context) --- .../moe/test_gpt_oss_triton_kernels.py | 38 ++++++++++++------- .../moe/test_modular_oai_triton_moe.py | 26 ++++++++----- tests/kernels/moe/utils.py | 24 ++++++++++++ .../layers/fused_moe/oracle/mxfp4.py | 16 +++----- .../layers/quantization/utils/mxfp4_utils.py | 9 +++++ 5 files changed, 79 insertions(+), 34 deletions(-) diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 86028a7feb..c1d83d8349 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/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 diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 47be82dd2a..9665079e37 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/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 ( diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index e762472950..d92e4d6129 100644 --- a/tests/kernels/moe/utils.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index d1c151aca2..7fed4124f5 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/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, diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index a1d0501775..8cc626f386 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/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()