[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>
Matt committed
Sep 17, 2026 at 09:21 UTC
fef55fd26dc633fcc87978f12f1a97e08e351c02
5 files changed
+79
-34
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+24
-14
@@ -7,7 +7,7 @@ import torch
7
import torch.nn.functional as F
8
9
from vllm.platforms import current_platform
10
-from vllm.utils.import_utils import has_triton_kernels
10
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
11
12
if not has_triton_kernels():
13
pytest.skip(
@@ -15,9 +15,15 @@ if not has_triton_kernels():
15
allow_module_level=True,
16
)
17
18
-import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
18
import triton_kernels.swiglu
20
-from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
19
+
20
+if get_triton_kernels_version() == "3.8":
21
+ # 3.8: matmul_ogs -> matmul, matmul_ogs_details -> matmul_details
22
+ import triton_kernels.matmul_details.opt_flags as opt_flags
23
+ from triton_kernels.matmul import FlexCtx, PrecisionConfig
24
+else:
25
+ import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
26
+ from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
27
from triton_kernels.numerics import InFlexData
28
from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp
29
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
34
from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
35
triton_kernel_moe_forward,
36
)
37
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs
38
from vllm.utils.math_utils import round_up
39
33
-from .utils import shuffle_weight
40
+from .utils import mxfp4_w_layouts, shuffle_weight
41
42
43
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):
140
141
x_tri = F.pad(x_tri, (0, x_pad, 0, 0), mode="constant", value=0)
142
136
- w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
137
- w_scale_layout, w_scale_layout_opts = (
138
- layout.make_default_matmul_mxfp4_w_scale_layout(
139
- mx_axis=1, num_warps=num_warps
140
- )
141
- )
143
+ (
144
+ w_layout,
145
+ w_layout_opts,
146
+ w_scale_layout,
147
+ w_scale_layout_opts,
148
+ ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps)
149
150
w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1)
151
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):
172
)
173
174
pc1 = PrecisionConfig(
168
- weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
175
+ **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
176
)
177
pc2 = PrecisionConfig(
171
- weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
178
+ **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
179
)
180
181
# tucuate so the rest can run properly
@@ -281,8 +288,6 @@ class Case:
288
@pytest.mark.parametrize("num_token", [2])
289
@pytest.mark.parametrize("tp", [1, 2, 4, 8])
290
def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init):
284
- from triton_kernels.tensor_details import layout
285
-
291
if not hasattr(layout, "make_default_matmul_mxfp4_w_layout"):
292
pytest.skip("make_default_matmul_mxfp4_w_layout not available")
293
@@ -387,6 +392,11 @@ def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk):
392
393
if gptoss_moe.triton_kernels_version == "3.5.1":
394
pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+")
395
+ if gptoss_moe.triton_kernels_version == "3.8":
396
+ pytest.skip(
397
+ "3.8 rebuilds routing from raw ids, so routing_data_from_sparse_topk "
398
+ "just forwards to make_routing_data and the comparison is vacuous"
399
+ )
400
401
from triton_kernels.topk import topk as topk_fn
402
tests/kernels/moe/test_modular_oai_triton_moe.py
+16
-10
@@ -10,7 +10,7 @@ import torch.nn.functional as F
10
11
from tests.utils import wait_for_gpu_memory_to_clear
12
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
13
-from vllm.utils.import_utils import has_triton_kernels
13
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
14
15
if not has_triton_kernels():
16
pytest.skip(
@@ -18,11 +18,14 @@ if not has_triton_kernels():
18
allow_module_level=True,
19
)
20
21
-from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
21
+if get_triton_kernels_version() == "3.8":
22
+ # 3.8: matmul_ogs -> matmul
23
+ from triton_kernels.matmul import FlexCtx, PrecisionConfig
24
+else:
25
+ from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
26
from triton_kernels.numerics import InFlexData
27
from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp
28
from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
25
-from triton_kernels.tensor_details import layout
29
from triton_kernels.testing import assert_close
30
31
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
38
UnfusedOAITritonExperts,
39
)
40
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
41
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs
42
from vllm.platforms import current_platform
43
from vllm.utils.math_utils import round_up
44
from vllm.utils.torch_utils import set_random_seed
45
42
-from .utils import make_dummy_moe_config, shuffle_weight
46
+from .utils import make_dummy_moe_config, mxfp4_w_layouts, shuffle_weight
47
48
MNK = [
49
(1, 512, 384),
@@ -130,10 +134,12 @@ def make_weights(dtype, k, n, e):
134
w2 = w2[..., :n, :k]
135
136
num_warps = 8
133
- w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
134
- w_scale_layout, w_scale_layout_opts = (
135
- layout.make_default_matmul_mxfp4_w_scale_layout(mx_axis=1, num_warps=num_warps)
136
- )
137
+ (
138
+ w_layout,
139
+ w_layout_opts,
140
+ w_scale_layout,
141
+ w_scale_layout_opts,
142
+ ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps)
143
144
w1_tri = convert_layout(wrap_torch_tensor(w1_tri, FP4), w_layout, **w_layout_opts)
145
w1_scale_tri = convert_layout(
@@ -150,10 +156,10 @@ def make_weights(dtype, k, n, e):
156
)
157
158
w1_precision_config = PrecisionConfig(
153
- weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
159
+ **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
160
)
161
w2_precision_config = PrecisionConfig(
156
- weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
162
+ **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
163
)
164
165
return (
tests/kernels/moe/utils.py
+24
@@ -705,3 +705,27 @@ def check_accuracy(a, b, atol, rtol, percent):
705
f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} "
706
f"(threshold: {1 - percent:.4f})"
707
)
708
+
709
+
710
+def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8):
711
+ """Weight/scale layouts for mxfp4 MoE, as (layout, opts) pairs.
712
+
713
+ triton_kernels 3.8 returns layout instances; earlier versions return a
714
+ (layout, opts) tuple.
715
+ """
716
+ from triton_kernels.tensor_details import layout
717
+
718
+ from vllm.utils.import_utils import get_triton_kernels_version
719
+
720
+ if get_triton_kernels_version() == "3.8":
721
+ w = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis)
722
+ s = layout.make_default_matmul_mxfp4_w_scale_layout(
723
+ mx_axis=mx_axis, num_warps=num_warps
724
+ )
725
+ return w, {}, s, {}
726
+
727
+ w, w_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis)
728
+ s, s_opts = layout.make_default_matmul_mxfp4_w_scale_layout(
729
+ mx_axis=mx_axis, num_warps=num_warps
730
+ )
731
+ return w, w_opts, s, s_opts
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+6
-10
@@ -31,7 +31,12 @@ from vllm.model_executor.layers.fused_moe.config import (
31
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
32
swap_w13_to_w31,
33
)
34
-from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
34
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
35
+ _swizzle_mxfp4,
36
+)
37
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
38
+ mx_scale_kwargs as _mx_scale_kwargs,
39
+)
40
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
41
OCP_MX_BLOCK_SIZE,
42
)
@@ -70,15 +75,6 @@ if triton_kernels_version is not None:
75
)
76
77
73
-def _mx_scale_kwargs(scale):
74
- """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size,
75
- 3.5.1/3.6 use weight_scale.
76
- """
77
- if triton_kernels_version == "3.8":
78
- return {"b_mx_scale": scale, "b_microblock_size": 32}
79
- return {"weight_scale": scale}
80
-
81
-
78
def _pack_deepgemm_mxfp4_scales(
79
w13_weight: torch.Tensor,
80
w2_weight: torch.Tensor,
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+9
@@ -36,6 +36,15 @@ def weight_mx_scale(precision_config):
36
return precision_config.weight_scale
37
38
39
+def mx_scale_kwargs(scale):
40
+ """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size,
41
+ 3.5.1/3.6 use weight_scale.
42
+ """
43
+ if get_triton_kernels_version() == "3.8":
44
+ return {"b_mx_scale": scale, "b_microblock_size": 32}
45
+ return {"weight_scale": scale}
46
+
47
+
48
def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
49
"""Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel."""
50
assert has_triton_kernels()