[ROCm][CI] Fall back to lossless Kimi K3 MXFP4 emulation on gfx942 (#50516)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Jul 31, 2026 at 10:33 UTC
5233368daccd3d84293bab7b04bacae1433ea0e4
4 files changed
+34
-5
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
+7
-1
@@ -69,9 +69,15 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
69
self.quantization_emulation = True
70
71
if self.ocp_mx_scheme in {
72
+ OCP_MX_Scheme.w_mxfp4,
73
+ OCP_MX_Scheme.w_mxfp6_e3m2,
74
+ OCP_MX_Scheme.w_mxfp6_e2m3,
75
+ }:
76
+ # Weight-only schemes leave activations unquantized.
77
+ self._quant_dtype = None
78
+ elif self.ocp_mx_scheme in {
79
OCP_MX_Scheme.w_mxfp4_a_mxfp4,
80
}:
74
- # Weight has to be dequantized for mxfp4 emulation.
81
self._quant_dtype = "mxfp4"
82
elif self.ocp_mx_scheme in [
83
OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2,
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+1
@@ -140,6 +140,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
140
MoEActivation.SILU,
141
MoEActivation.GELU,
142
MoEActivation.GELU_TANH,
143
+ MoEActivation.SITU,
144
MoEActivation.SWIGLUOAI,
145
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
146
MoEActivation.SWIGLUSTEP,
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+11
-4
@@ -337,7 +337,10 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
337
backend-level ``is_supported_config`` check filters by device capability).
338
"""
339
if current_platform.is_rocm():
340
- return [Mxfp4MoeBackend.AITER_MXFP4_BF16]
340
+ return [
341
+ Mxfp4MoeBackend.AITER_MXFP4_BF16,
342
+ Mxfp4MoeBackend.EMULATION,
343
+ ]
344
if current_platform.is_xpu():
345
return [Mxfp4MoeBackend.XPU]
346
_AVAILABLE_BACKENDS = [
@@ -1555,8 +1558,12 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1558
w13_bias,
1559
w2_bias,
1560
)
1558
- elif mxfp4_backend == Mxfp4MoeBackend.XPU:
1559
- # No additional transformation needed for XPU backend
1561
+ elif mxfp4_backend in (
1562
+ Mxfp4MoeBackend.XPU,
1563
+ Mxfp4MoeBackend.EMULATION,
1564
+ ):
1565
+ # No additional transformation is needed: XPU consumes the checkpoint
1566
+ # layout directly, while emulation dequantizes that layout at runtime.
1567
return (
1568
w13_weight,
1569
w2_weight,
@@ -1568,7 +1575,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1575
else:
1576
raise ValueError(
1577
f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
1571
- f"Expected TRTLLM, Triton, AITER, or XPU backend."
1578
+ f"Expected TRTLLM, Triton, AITER, XPU, or emulation backend."
1579
)
1580
1581
vllm/model_executor/layers/quantization/mxfp4.py
+15
@@ -15,6 +15,9 @@ from vllm.model_executor.layers.fused_moe import (
15
SharedExperts,
16
)
17
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
18
+from vllm.model_executor.layers.fused_moe.config import (
19
+ mxfp4_w4a16_moe_quant_config,
20
+)
21
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
22
TRITON_BACKENDS,
23
Mxfp4MoeBackend,
@@ -868,6 +871,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
871
w1_scale = layer.w13_weight_scale
872
w2_scale = layer.w2_weight_scale
873
874
+ if self.mxfp4_backend == Mxfp4MoeBackend.EMULATION:
875
+ # Canonical ``mxfp4`` checkpoints are weight-only W4A16. The
876
+ # generic EMULATION config is W4A4, so preserve BF16 activations
877
+ # while the fallback dequantizes only the weights.
878
+ return mxfp4_w4a16_moe_quant_config(
879
+ w1_scale=w1_scale,
880
+ w2_scale=w2_scale,
881
+ w1_bias=w1_bias,
882
+ w2_bias=w2_bias,
883
+ gemm1_clamp_limit=swiglu_limit,
884
+ )
885
+
886
return make_mxfp4_moe_quant_config(
887
mxfp4_backend=self.mxfp4_backend,
888
w1_scale=w1_scale,