[AMD][OCP MX][CI] Fix tests to not dispatch on `UNFUSED_TRITON` backend on MI300, improve w_mxfp4_a_fp8 emulation support (#46142)
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
fxmarty-amd committed
Jun 23, 2026 at 20:25 UTC
e368415daa2c4a4141904ec9c85e7a0dcaff6160
7 files changed
+74
-17
tests/models/quantization/test_gpt_oss.py
+5
@@ -104,6 +104,11 @@ def test_gpt_oss_attention_quantization(
104
105
model_args = EvaluationConfig(model_name).get_model_args(tp_size)
106
107
+ # Emulation backend on MI300, MI250 is opt-in
108
+ # following https://github.com/vllm-project/vllm/pull/45896
109
+ if not on_gfx950():
110
+ model_args["moe_backend"] = "emulation"
111
+
112
extra_run_kwargs = {
113
"gen_kwargs": {"max_gen_toks": 8000},
114
"apply_chat_template": True,
tests/quantization/test_quark.py
+12
@@ -30,6 +30,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
30
)
31
from vllm.platforms import current_platform
32
33
+if current_platform.is_rocm():
34
+ from vllm.platforms.rocm import on_gfx950
35
+else:
36
+
37
+ def on_gfx950() -> bool:
38
+ return False
39
+
40
+
41
from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch
42
43
# Minimum amd-quark version for MXFP4/OCP_MX tests (single source of truth).
@@ -213,6 +221,10 @@ class AccuracyTestConfig:
221
if model_max_len is not None:
222
model_args["max_model_len"] = model_max_len
223
224
+ # Emulation backend on MI300, MI250 is opt-in following https://github.com/vllm-project/vllm/pull/45896
225
+ if not on_gfx950():
226
+ model_args["moe_backend"] = "emulation"
227
+
228
return model_args
229
230
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
+2
-2
@@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mx
26
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
27
OCP_MX_Scheme,
28
)
29
+from vllm.platforms import current_platform
30
31
logger = init_logger(__name__)
32
@@ -83,8 +84,7 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
84
OCP_MX_Scheme.w_mxfp4_a_fp8,
85
OCP_MX_Scheme.w_mxfp6_e3m2_a_fp8,
86
]:
86
- # TODO: double check this one
87
- self._quant_dtype = "mxfp8"
87
+ self._quant_dtype = current_platform.fp8_dtype()
88
89
@property
90
def quant_dtype(self) -> torch.dtype | str | None:
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+12
-1
@@ -332,12 +332,23 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
332
else:
333
lora_x = hidden_states
334
335
+ # TODO: The fallback to self.a1_scale was added for deferred static
336
+ # activation quantization in https://github.com/vllm-project/vllm/pull/40857.
337
+ # Activation emulation relies solely on `a1q_scale` output of
338
+ # `moe_kernel_quantize_input` - this should be adapted to
339
+ # always solely rely on `a1q_scale`.
340
+ input_scale = (
341
+ a1q_scale
342
+ if self.quantization_emulation
343
+ else (a1q_scale if a1q_scale is not None else self.a1_scale)
344
+ )
345
+
346
def _base_w13_fn():
347
invoke_fused_moe_triton_kernel(
348
hidden_states,
349
w1,
350
intermediate_cache1,
340
- a1q_scale if a1q_scale is not None else self.a1_scale,
351
+ input_scale,
352
self.w1_scale,
353
None, # topk_weights
354
sorted_token_ids,
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+25
-2
@@ -674,6 +674,8 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
674
w2_weight_scale: torch.Tensor,
675
w13_bias: torch.Tensor | None = None,
676
w2_bias: torch.Tensor | None = None,
677
+ w13_input_scale: torch.Tensor | None = None,
678
+ w2_input_scale: torch.Tensor | None = None,
679
_cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None,
680
) -> tuple[
681
torch.Tensor,
@@ -1191,8 +1193,29 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1193
w2_bias,
1194
)
1195
elif mxfp4_backend == Mxfp4MoeBackend.EMULATION:
1194
- # No additional transformation needed for emulation backend,
1195
- # weights are dequantized on the fly in the experts class.
1196
+ w13_has_per_expert_scale = (
1197
+ w13_input_scale is not None
1198
+ and w13_input_scale.ndim == 1
1199
+ and not all_close_1d(w13_input_scale)
1200
+ )
1201
+ w2_has_per_expert_scale = (
1202
+ w2_input_scale is not None
1203
+ and w2_input_scale.ndim == 1
1204
+ and not all_close_1d(w2_input_scale)
1205
+ )
1206
+ if w13_has_per_expert_scale or w2_has_per_expert_scale:
1207
+ logger.warning_once(
1208
+ "Found input_scales that are not equal for OCP MX MoE "
1209
+ "emulation. Using the maximum across experts for each layer."
1210
+ )
1211
+ if w13_input_scale is not None:
1212
+ layer.w13_input_scale = torch.nn.Parameter(
1213
+ w13_input_scale.max().to(torch.float32), requires_grad=False
1214
+ )
1215
+ if w2_input_scale is not None:
1216
+ layer.w2_input_scale = torch.nn.Parameter(
1217
+ w2_input_scale.max().to(torch.float32), requires_grad=False
1218
+ )
1219
return (
1220
w13_weight,
1221
w2_weight,
vllm/model_executor/layers/fused_moe/utils.py
+16
-12
@@ -200,6 +200,16 @@ def _mxfp4_quantize(
200
return A, None
201
202
203
+def _fp8_quantize_dequantize(
204
+ A: torch.Tensor,
205
+ A_scale: torch.Tensor,
206
+):
207
+ qA, qA_scale = ops.scaled_fp8_quant(A, A_scale, use_per_token_if_dynamic=False)
208
+ A = per_tensor_dequantize(qA, qA_scale).to(A.dtype)
209
+
210
+ return A, None
211
+
212
+
213
def _mxfp8_e4m3_quantize(
214
A: torch.Tensor,
215
A_scale: torch.Tensor | None,
@@ -268,23 +278,17 @@ def moe_kernel_quantize_input(
278
# purpose, because there is no native kernel for weight in ocp_mx_scheme
279
# and activation in FP8. The implementation is based on existing
280
# non-emulation ops.
271
- qA, qA_scale = ops.scaled_fp8_quant(
272
- A, A_scale, use_per_token_if_dynamic=False
273
- )
274
- A = per_tensor_dequantize(qA, qA_scale).to(A.dtype)
275
- # After QDQ, we don't need further quantization
276
- return A, None
281
+ # TODO: Remove this `ocp_mx_scheme is not None` block and rely solely
282
+ # on `quantization_emulation`.
283
+ return _fp8_quantize_dequantize(A, A_scale)
284
# else: For other schemes (e.g., *_a_mxfp6_e3m2, *_a_mxfp6_e2m3),
285
# weights are already dequantized, and we proceed with normal
286
# activation quantization below.
280
-
287
if quant_dtype == current_platform.fp8_dtype():
288
if quantization_emulation:
283
- raise NotImplementedError(
284
- f"moe_kernel_quantize_input does not support quant_dtype={quant_dtype}"
285
- " MOE quantization emulation. Please open an issue."
286
- )
287
- return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
289
+ return _fp8_quantize_dequantize(A, A_scale)
290
+ else:
291
+ return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
292
elif quant_dtype == torch.int8:
293
if quantization_emulation:
294
raise NotImplementedError(
vllm/model_executor/layers/quantization/quark/quark_moe.py
+2
@@ -1211,6 +1211,8 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1211
w2_weight_scale=layer.w2_weight_scale,
1212
w13_bias=w13_bias,
1213
w2_bias=w2_bias,
1214
+ w13_input_scale=layer.w13_input_scale,
1215
+ w2_input_scale=layer.w2_input_scale,
1216
)
1217
)
1218