[Bugfix][Quantization] Fix OCP MX MoE emulation silently skipping mxfp6 activation QDQ (#52704)
Signed-off-by: xuebwang-amd <xuebwang@amd.com>
xuebwang-amd committed
Aug 20, 2026 at 00:29 UTC
c676232313fce3d607da7c66e41fe8d4739c194d
2 files changed
+195
-23
tests/kernels/moe/test_ocp_mx_moe.py
+148
@@ -13,6 +13,11 @@ from vllm._aiter_ops import (
13
is_aiter_found_and_supported,
14
rocm_aiter_ops,
15
)
16
+from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import (
17
+ activation_quant_dtype,
18
+)
19
+from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
20
+from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme
21
from vllm.platforms import current_platform
22
from vllm.utils.flashinfer import has_flashinfer
23
@@ -1859,3 +1864,146 @@ def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons(
1864
1865
with pytest.raises(NotImplementedError, match="Unsupported reasons"):
1866
mxfp4_oracle.select_mxfp4_moe_backend(moe_config)
1867
+
1868
+
1869
+# Every activation-quantizing OCP MX scheme must map to a `quant_dtype` that
1870
+# `moe_kernel_quantize_input` actually dispatches on. Its final `else` returns
1871
+# the activation untouched, so a name it does not know (e.g. "mxfp6" instead of
1872
+# "mxfp6_e3m2") silently skips the fake-quantization the emulation exists for.
1873
+@pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm")
1874
+@pytest.mark.parametrize("ocp_mx_scheme", list(OCP_MX_Scheme))
1875
+def test_emulation_activation_quant_dtype_is_dispatchable(ocp_mx_scheme):
1876
+ quant_dtype = activation_quant_dtype(ocp_mx_scheme)
1877
+
1878
+ if "_a_" not in ocp_mx_scheme.value:
1879
+ assert quant_dtype is None, "weight-only schemes must not quantize activations"
1880
+ return
1881
+
1882
+ a = torch.randn(64, 128, dtype=torch.bfloat16, device="cuda")
1883
+ a_scale = torch.ones(1, dtype=torch.float32, device="cuda")
1884
+ out, _ = moe_kernel_quantize_input(
1885
+ a, a_scale, quant_dtype, False, None, quantization_emulation=True
1886
+ )
1887
+ assert not torch.equal(out, a), (
1888
+ f"{ocp_mx_scheme.value} -> quant_dtype={quant_dtype!r} left the activation"
1889
+ " unquantized; moe_kernel_quantize_input does not dispatch on it"
1890
+ )
1891
+
1892
+
1893
+@pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm")
1894
+@torch.inference_mode()
1895
+def test_emulation_a_mxfp6_moe_forward_quantizes_activations():
1896
+ """The same property observed through a full MoE forward.
1897
+
1898
+ `w_mxfp4_a_mxfp6_e3m2` and the weight-only `w_mxfp4` differ only in whether
1899
+ activations are fake-quantized, so with identical weights, activations and
1900
+ routing their layer outputs must differ. When the emulation selects a
1901
+ `quant_dtype` `moe_kernel_quantize_input` does not dispatch on, the QDQ is
1902
+ skipped and the two outputs come out bit-identical.
1903
+ """
1904
+ from vllm.config import VllmConfig, set_current_vllm_config
1905
+ from vllm.model_executor.layers.fused_moe import FusedMoEConfig
1906
+ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
1907
+ from vllm.model_executor.layers.fused_moe.config import (
1908
+ FusedMoEParallelConfig,
1909
+ RoutingMethodType,
1910
+ mxfp4_w4a16_moe_quant_config,
1911
+ ocp_mx_moe_quant_config,
1912
+ )
1913
+ from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import (
1914
+ OCP_MXQuantizationEmulationTritonExperts,
1915
+ )
1916
+ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
1917
+ Mxfp4MoeBackend,
1918
+ make_mxfp4_moe_kernel,
1919
+ )
1920
+ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
1921
+ mxfp4_quantize,
1922
+ )
1923
+ from vllm.v1.worker.workspace import init_workspace_manager
1924
+
1925
+ init_workspace_manager(torch.accelerator.current_device_index())
1926
+
1927
+ num_experts, topk = 8, 2
1928
+ hidden_size, intermediate_size, num_tokens = 256, 256, 64
1929
+ dtype, device = torch.bfloat16, "cuda:0"
1930
+
1931
+ torch.manual_seed(0)
1932
+ w13, w13_scale = mxfp4_quantize(
1933
+ torch.randn(
1934
+ num_experts, 2 * intermediate_size, hidden_size, dtype=dtype, device=device
1935
+ )
1936
+ / 8
1937
+ )
1938
+ w2, w2_scale = mxfp4_quantize(
1939
+ torch.randn(
1940
+ num_experts, hidden_size, intermediate_size, dtype=dtype, device=device
1941
+ )
1942
+ / 8
1943
+ )
1944
+ w13, w2 = w13.contiguous(), w2.contiguous()
1945
+ w13_scale, w2_scale = w13_scale.contiguous(), w2_scale.contiguous()
1946
+
1947
+ torch.manual_seed(1)
1948
+ hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
1949
+ topk_weights, topk_ids = torch.topk(
1950
+ torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device),
1951
+ k=topk,
1952
+ dim=-1,
1953
+ )
1954
+ topk_weights = torch.softmax(topk_weights, dim=-1).to(dtype)
1955
+ topk_ids = topk_ids.to(torch.int32)
1956
+
1957
+ def run(quant_config):
1958
+ moe_config = FusedMoEConfig(
1959
+ num_experts=num_experts,
1960
+ experts_per_token=topk,
1961
+ hidden_dim=hidden_size,
1962
+ intermediate_size=intermediate_size,
1963
+ num_local_experts=num_experts,
1964
+ num_logical_experts=num_experts,
1965
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1966
+ activation=MoEActivation.SILU,
1967
+ in_dtype=dtype,
1968
+ device=device,
1969
+ routing_method=RoutingMethodType.Renormalize,
1970
+ )
1971
+ with set_current_vllm_config(VllmConfig()):
1972
+ kernel = make_mxfp4_moe_kernel(
1973
+ moe_quant_config=quant_config,
1974
+ moe_config=moe_config,
1975
+ mxfp4_backend=Mxfp4MoeBackend.EMULATION,
1976
+ experts_cls=OCP_MXQuantizationEmulationTritonExperts,
1977
+ routing_tables=None,
1978
+ )
1979
+ return kernel.apply(
1980
+ hidden_states=hidden_states,
1981
+ w1=w13,
1982
+ w2=w2,
1983
+ topk_weights=topk_weights,
1984
+ topk_ids=topk_ids,
1985
+ activation=MoEActivation.SILU,
1986
+ global_num_experts=num_experts,
1987
+ expert_map=None,
1988
+ apply_router_weight_on_input=False,
1989
+ )
1990
+
1991
+ # A fresh quant config per run: the experts nulls the weight scales on it.
1992
+ a_mxfp6 = ocp_mx_moe_quant_config(
1993
+ quant_dtype="mxfp6_e3m2",
1994
+ weight_dtype="mxfp4",
1995
+ w1_scale=w13_scale,
1996
+ w2_scale=w2_scale,
1997
+ )
1998
+ assert a_mxfp6.ocp_mx_scheme == OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2
1999
+ out_a_mxfp6 = run(a_mxfp6)
2000
+
2001
+ weight_only = mxfp4_w4a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale)
2002
+ assert weight_only.ocp_mx_scheme == OCP_MX_Scheme.w_mxfp4
2003
+ out_weight_only = run(weight_only)
2004
+
2005
+ max_diff = (out_a_mxfp6.float() - out_weight_only.float()).abs().max().item()
2006
+ assert max_diff > 0.0, (
2007
+ "w_mxfp4_a_mxfp6_e3m2 output is bit-identical to weight-only w_mxfp4:"
2008
+ " the emulation never fake-quantized the activations"
2009
+ )
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
+47
-23
@@ -35,6 +35,52 @@ from vllm.utils.import_utils import has_quark
35
logger = init_logger(__name__)
36
37
38
+def activation_quant_dtype(
39
+ ocp_mx_scheme: OCP_MX_Scheme | str,
40
+) -> torch.dtype | str | None:
41
+ """Activation dtype `moe_kernel_quantize_input` should fake-quantize to.
42
+
43
+ Args:
44
+ ocp_mx_scheme: The OCP MX scheme the emulated experts run. Accepts the
45
+ enum member or its string value.
46
+
47
+ Returns:
48
+ A `quant_dtype` `moe_kernel_quantize_input` dispatches on, or None for
49
+ weight-only schemes, which leave activations untouched.
50
+
51
+ Raises:
52
+ NotImplementedError: If the scheme has no emulated activation dtype.
53
+ """
54
+ if ocp_mx_scheme in {
55
+ OCP_MX_Scheme.w_mxfp4,
56
+ OCP_MX_Scheme.w_mxfp6_e3m2,
57
+ OCP_MX_Scheme.w_mxfp6_e2m3,
58
+ }:
59
+ return None
60
+ elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp4_a_mxfp4:
61
+ return "mxfp4"
62
+ elif ocp_mx_scheme in {
63
+ OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2,
64
+ OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2,
65
+ }:
66
+ return "mxfp6_e3m2"
67
+ elif ocp_mx_scheme in {
68
+ OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3,
69
+ OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3,
70
+ }:
71
+ return "mxfp6_e2m3"
72
+ elif ocp_mx_scheme in {
73
+ OCP_MX_Scheme.w_mxfp4_a_fp8,
74
+ OCP_MX_Scheme.w_mxfp6_e3m2_a_fp8,
75
+ OCP_MX_Scheme.w_mxfp6_e2m3_a_fp8,
76
+ }:
77
+ return current_platform.fp8_dtype()
78
+ raise NotImplementedError(
79
+ f"No emulated activation dtype for OCP MX scheme {ocp_mx_scheme}."
80
+ " Please open an issue."
81
+ )
82
+
83
+
84
class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
85
"""
86
Extension of TritonExperts to support emulated OCP MX MoE experts.
@@ -72,29 +118,7 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
118
119
self.quantization_emulation = True
120
75
- if self.ocp_mx_scheme in {
76
- OCP_MX_Scheme.w_mxfp4,
77
- OCP_MX_Scheme.w_mxfp6_e3m2,
78
- OCP_MX_Scheme.w_mxfp6_e2m3,
79
- }:
80
- # Weight-only schemes leave activations unquantized.
81
- self._quant_dtype = None
82
- elif self.ocp_mx_scheme in {
83
- OCP_MX_Scheme.w_mxfp4_a_mxfp4,
84
- }:
85
- self._quant_dtype = "mxfp4"
86
- elif self.ocp_mx_scheme in [
87
- OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2,
88
- OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3,
89
- OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2,
90
- OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3,
91
- ]:
92
- self._quant_dtype = "mxfp6"
93
- elif self.ocp_mx_scheme in [
94
- OCP_MX_Scheme.w_mxfp4_a_fp8,
95
- OCP_MX_Scheme.w_mxfp6_e3m2_a_fp8,
96
- ]:
97
- self._quant_dtype = current_platform.fp8_dtype()
121
+ self._quant_dtype = activation_quant_dtype(self.ocp_mx_scheme)
122
123
@staticmethod
124
def is_supported_config(