[Bugfix] Fix NVFP4/OCP MX MoE emulation (#46254)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Jun 20, 2026 at 23:13 UTC
a346d589f5932d4234bf5bf8718f10e26d187021
5 files changed
+3
-21
.buildkite/test_areas/lm_eval.yaml
+1
@@ -109,6 +109,7 @@ steps:
109
- image-build-amd
110
commands:
111
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
112
+ - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
113
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
114
115
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py
-9
@@ -21,7 +21,6 @@ from vllm.model_executor.layers.fused_moe.config import (
21
FusedMoEQuantConfig,
22
)
23
from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
24
-from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
24
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
25
dequantize_to_dtype,
26
)
@@ -135,14 +134,6 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
134
swizzle=False,
135
)
136
138
- hidden_states, _ = moe_kernel_quantize_input(
139
- A=hidden_states,
140
- A_scale=self.quant_config.a1_gscale,
141
- quant_dtype="nvfp4",
142
- per_act_token_quant=False,
143
- quantization_emulation=True,
144
- )
145
-
137
# Activation quantization/dequantization is deferred to
138
# `moe_kernel_quantize_input` in TritonExperts.apply.
139
super().apply(
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
-11
@@ -21,7 +21,6 @@ from vllm.model_executor.layers.fused_moe.config import (
21
FusedMoEQuantConfig,
22
)
23
from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
24
-from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
24
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
25
from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
26
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -155,16 +154,6 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
154
w2, self.w2_scale_val, hidden_states.dtype
155
)
156
158
- # Apply activation QDQ if needed by the OCP MX scheme
159
- hidden_states, _ = moe_kernel_quantize_input(
160
- A=hidden_states,
161
- A_scale=None,
162
- quant_dtype=self.quant_config.quant_dtype,
163
- per_act_token_quant=False,
164
- ocp_mx_scheme=self.ocp_mx_scheme,
165
- quantization_emulation=True,
166
- )
167
-
157
# Activation quantization/dequantization is deferred to
158
# `moe_kernel_quantize_input` in TritonExperts.apply.
159
super().apply(
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+1
-1
@@ -245,7 +245,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
245
lora_unquantized_hidden_states = hidden_states
246
hidden_states, a1q_scale = moe_kernel_quantize_input(
247
hidden_states,
248
- self.a1_scale,
248
+ self.a1_scale or self.a1_gscale,
249
self.quant_dtype,
250
self.per_act_token_quant,
251
self.block_shape,
vllm/model_executor/layers/fused_moe/utils.py
+1
@@ -296,6 +296,7 @@ def moe_kernel_quantize_input(
296
if not quantization_emulation:
297
return _nvfp4_quantize(A, A_scale, is_sf_swizzled_layout=is_scale_swizzled)
298
else:
299
+ assert A_scale is not None
300
A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
301
return A, None
302
elif quant_dtype == "mxfp4":