A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 20:00 UTC
[NVFP4] NVFP4 MOE emulation fallback for H100/MI300/MI350, standardize `TritonExperts` usage for OCP MX emulation (#35737)
Signed-off-by: Felix Marty <Felix.Marty@amd.com> Signed-off-by: fxmarty-amd <felmarty@amd.com> Co-authored-by: Kyle Sayers <kylesayrs@gmail.com>
fxmarty-amd committed
Apr 22, 2026 at 17:58 UTC
d622e27d2be9cd4321d75073e4b7ef522204853d
12 files changed
+601
-121
tests/evals/gsm8k/configs/models-mi3xx.txt
+2
@@ -2,3 +2,5 @@ DeepSeek-R1-TP_MI325.yaml
2
DeepSeek-R1-DP_MI325.yaml
3
DeepSeek-V3.2-TP_MI325.yaml
4
DeepSeek-V3.2-DP_MI325.yaml
5
+Qwen3-30B-A3B-NVFP4.yaml
6
+Qwen3.5-35B-A3B-MXFP4-TP2.yaml
\ No newline at end of file
tests/models/quantization/test_nvfp4.py
+23
@@ -120,3 +120,26 @@ def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch):
120
with vllm_runner(model, enforce_eager=eager) as llm:
121
output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)
122
assert output[0][1] == "1 2 3 4 5 6"
123
+
124
+
125
+@pytest.mark.parametrize(
126
+ "model",
127
+ [
128
+ "nvidia/Qwen3-30B-A3B-NVFP4",
129
+ "RedHatAI/Qwen3-30B-A3B-NVFP4",
130
+ ],
131
+)
132
+@pytest.mark.parametrize("backend", ["emulation"])
133
+@pytest.mark.skipif(
134
+ not current_platform.is_rocm(),
135
+ reason="NVFP4 MOE emulation is only useful on AMD Instinct MI3xx",
136
+)
137
+def test_nvfp4_moe(vllm_runner, model, backend, monkeypatch):
138
+ monkeypatch.setenv("VLLM_NVFP4_GEMM_BACKEND", backend)
139
+ with vllm_runner(
140
+ model,
141
+ moe_backend=backend,
142
+ load_format="dummy",
143
+ hf_overrides={"num_hidden_layers": 2},
144
+ ) as llm:
145
+ _ = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)
vllm/config/kernel.py
+5
-1
@@ -115,6 +115,7 @@ MoEBackend = Literal[
115
"flashinfer_cutedsl",
116
"marlin",
117
"aiter",
118
+ "emulation",
119
]
120
121
@@ -142,7 +143,10 @@ class KernelConfig:
143
- "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
144
- "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
145
- "marlin": Use Marlin kernels (weight-only quantization)
145
- - "aiter": Use AMD AITer kernels (ROCm only)"""
146
+ - "aiter": Use AMD AITer kernels (ROCm only)
147
+ - "emulation": use BF16/FP16 GEMM, dequantizing weights and
148
+ running QDQ on activations.
149
+ """
150
151
@field_validator("moe_backend", mode="before")
152
@classmethod
vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py
new
+164
@@ -0,0 +1,164 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+NVFP4 quantization emulation for MoE.
5
+
6
+This file implements NVFP4 emulation for NVFP4 MOE in case the hardware used does not
7
+natively support NVFP4 MOE.
8
+
9
+Weights are dequantized on the fly during each forward, we fall back to calling
10
+`TritonExperts` using BF16, and fake NVFP4 quantize-dequantize
11
+is applied on `a13`, `a2`.
12
+"""
13
+
14
+import torch
15
+
16
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
17
+from vllm.logger import init_logger
18
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19
+from vllm.model_executor.layers.fused_moe.config import (
20
+ FusedMoEConfig,
21
+ FusedMoEQuantConfig,
22
+)
23
+from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
24
+from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
25
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
26
+ dequantize_to_dtype,
27
+)
28
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
29
+ QuantKey,
30
+ kNvfp4Dynamic,
31
+ kNvfp4Static,
32
+)
33
+
34
+logger = init_logger(__name__)
35
+
36
+
37
+class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
38
+ """
39
+ Extension of TritonExperts to support emulated NVFP4 MoE experts.
40
+
41
+ It may be used for NVFP4 models when the device does not have
42
+ native support for this dtype.
43
+ """
44
+
45
+ def __init__(
46
+ self,
47
+ moe_config: FusedMoEConfig,
48
+ quant_config: FusedMoEQuantConfig,
49
+ ):
50
+ super().__init__(moe_config, quant_config)
51
+ logger.warning_once(
52
+ "Using Nvfp4QuantizationEmulationTritonExperts MOE backend. This will"
53
+ " dequantize weights on the fly and may be slower than native"
54
+ " quantized MOE. Consider using a device with native quantization"
55
+ " support (e.g. Nvidia Blackwell) for better performance."
56
+ )
57
+
58
+ # `TritonExperts.apply` expects pre-dequantized weights,
59
+ # which we handle in `apply` below.
60
+ self.w1_scale_val = self.quant_config.w1_scale
61
+ self.w2_scale_val = self.quant_config.w2_scale
62
+
63
+ self.quant_config._w1.scale = None
64
+ self.quant_config._w2.scale = None
65
+
66
+ self.quantization_emulation = True
67
+
68
+ @property
69
+ def quant_dtype(self) -> torch.dtype | str | None:
70
+ return "nvfp4"
71
+
72
+ @property
73
+ def expects_unquantized_inputs(self) -> bool:
74
+ return True
75
+
76
+ @staticmethod
77
+ def _supports_quant_scheme(
78
+ weight_key: QuantKey | None,
79
+ activation_key: QuantKey | None,
80
+ ) -> bool:
81
+ return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)
82
+
83
+ def apply(
84
+ self,
85
+ output: torch.Tensor,
86
+ hidden_states: torch.Tensor,
87
+ w1: torch.Tensor,
88
+ w2: torch.Tensor,
89
+ topk_weights: torch.Tensor,
90
+ topk_ids: torch.Tensor,
91
+ activation: MoEActivation,
92
+ global_num_experts: int,
93
+ expert_map: torch.Tensor | None,
94
+ a1q_scale: torch.Tensor | None,
95
+ a2_scale: torch.Tensor | None,
96
+ workspace13: torch.Tensor,
97
+ workspace2: torch.Tensor,
98
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
99
+ apply_router_weight_on_input: bool,
100
+ ):
101
+ """
102
+ Apply emulated quantized MoE computation.
103
+
104
+ This dequantizes the weights on the fly and calls fused_experts_impl
105
+ with activation quantization support.
106
+ """
107
+ # Dequantize weights if they are quantized
108
+ # For NVFP4, weights are packed in uint8 format
109
+ # w1 shape: [num_experts, 2*intermediate_size, hidden_size//2]
110
+ # w2 shape: [num_experts, hidden_size, intermediate_size//2]
111
+ assert w1.dtype == torch.uint8
112
+ assert w2.dtype == torch.uint8
113
+
114
+ # Dequantize w1 from packed NVFP4 to fp16/bf16
115
+ w13_global_scale = self.quant_config.g1_alphas
116
+
117
+ w1_dequant = dequantize_to_dtype(
118
+ tensor_fp4=w1,
119
+ tensor_sf=self.w1_scale_val,
120
+ global_scale=w13_global_scale,
121
+ dtype=hidden_states.dtype,
122
+ block_size=16,
123
+ swizzle=False,
124
+ )
125
+
126
+ # Dequantize w2 from packed NVFP4 to fp16/bf16
127
+ w2_global_scale = self.quant_config.g2_alphas
128
+
129
+ w2_dequant = dequantize_to_dtype(
130
+ tensor_fp4=w2,
131
+ tensor_sf=self.w2_scale_val,
132
+ global_scale=w2_global_scale,
133
+ dtype=hidden_states.dtype,
134
+ block_size=16,
135
+ swizzle=False,
136
+ )
137
+
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
+
146
+ # Activation quantization/dequantization is deferred to
147
+ # `moe_kernel_quantize_input` in TritonExperts.apply.
148
+ super().apply(
149
+ output=output,
150
+ hidden_states=hidden_states,
151
+ w1=w1_dequant,
152
+ w2=w2_dequant,
153
+ topk_weights=topk_weights,
154
+ topk_ids=topk_ids,
155
+ activation=activation,
156
+ global_num_experts=global_num_experts,
157
+ expert_map=expert_map,
158
+ a1q_scale=None,
159
+ a2_scale=self.quant_config.a2_gscale,
160
+ workspace13=workspace13,
161
+ workspace2=workspace2,
162
+ expert_tokens_meta=expert_tokens_meta,
163
+ apply_router_weight_on_input=apply_router_weight_on_input,
164
+ )
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
new
+186
@@ -0,0 +1,186 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+OCP MX quantization emulation for MoE.
5
+
6
+This file implements OCP MX (MXFP4/MXFP6) emulation for MoE in case the
7
+hardware used does not natively support OCP MX MoE.
8
+
9
+Weights are dequantized on the fly during each forward, we fall back to calling
10
+`TritonExperts` using BF16, and fake OCP MX quantize-dequantize
11
+is applied on activations via `moe_kernel_quantize_input`.
12
+"""
13
+
14
+import torch
15
+
16
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
17
+from vllm.logger import init_logger
18
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19
+from vllm.model_executor.layers.fused_moe.config import (
20
+ FusedMoEConfig,
21
+ FusedMoEQuantConfig,
22
+)
23
+from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
24
+from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
25
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
26
+from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
27
+from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
28
+ OCP_MX_Scheme,
29
+)
30
+
31
+logger = init_logger(__name__)
32
+
33
+
34
+class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
35
+ """
36
+ Extension of TritonExperts to support emulated OCP MX MoE experts.
37
+
38
+ It may be used for OCP MX (MXFP4/MXFP6) models when the device does not
39
+ have native support for these dtypes.
40
+ """
41
+
42
+ def __init__(
43
+ self,
44
+ moe_config: FusedMoEConfig,
45
+ quant_config: FusedMoEQuantConfig,
46
+ ):
47
+ super().__init__(moe_config, quant_config)
48
+ logger.warning_once(
49
+ "Using OCP_MXQuantizationEmulationTritonExperts MOE backend. This"
50
+ " will dequantize weights on the fly and may be slower than native"
51
+ " quantized MOE. Consider using a device with native OCP MX"
52
+ " quantization support for better performance."
53
+ )
54
+
55
+ self.ocp_mx_scheme = quant_config.ocp_mx_scheme
56
+ assert self.ocp_mx_scheme is not None, (
57
+ "ocp_mx_scheme must be set in quant_config for"
58
+ " OCP_MXQuantizationEmulationTritonExperts"
59
+ )
60
+
61
+ # `TritonExperts.apply` expects pre-dequantized weights,
62
+ # which we handle in `apply` below.
63
+ self.w1_scale_val = self.quant_config.w1_scale
64
+ self.w2_scale_val = self.quant_config.w2_scale
65
+
66
+ self.quant_config._w1.scale = None
67
+ self.quant_config._w2.scale = None
68
+
69
+ self.quantization_emulation = True
70
+
71
+ if self.ocp_mx_scheme in {
72
+ OCP_MX_Scheme.w_mxfp4_a_mxfp4,
73
+ }:
74
+ # Weight has to be dequantized for mxfp4 emulation.
75
+ self._quant_dtype = "mxfp4"
76
+ elif self.ocp_mx_scheme in [
77
+ OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2,
78
+ OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3,
79
+ OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2,
80
+ OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3,
81
+ ]:
82
+ self._quant_dtype = "mxfp6"
83
+ elif self.ocp_mx_scheme in [
84
+ OCP_MX_Scheme.w_mxfp4_a_fp8,
85
+ OCP_MX_Scheme.w_mxfp6_e3m2_a_fp8,
86
+ ]:
87
+ # TODO: double check this one
88
+ self._quant_dtype = "mxfp8"
89
+
90
+ @property
91
+ def quant_dtype(self) -> torch.dtype | str | None:
92
+ return self._quant_dtype
93
+
94
+ @property
95
+ def expects_unquantized_inputs(self) -> bool:
96
+ return True
97
+
98
+ @staticmethod
99
+ def _supports_quant_scheme(
100
+ weight_key,
101
+ activation_key,
102
+ ) -> bool:
103
+ # This class is used for emulation only - the oracle selects it
104
+ # directly rather than via quant scheme matching.
105
+ return True
106
+
107
+ def _dequantize_weights(
108
+ self,
109
+ w: torch.Tensor,
110
+ w_scale: torch.Tensor,
111
+ dtype: torch.dtype,
112
+ ) -> torch.Tensor:
113
+ """Dequantize weights based on the OCP MX scheme."""
114
+ if self.ocp_mx_scheme.startswith("w_mxfp4"): # type: ignore[union-attr]
115
+ return dequant_mxfp4(w, w_scale, dtype)
116
+ elif self.ocp_mx_scheme.startswith("w_mxfp6_e3m2"): # type: ignore[union-attr]
117
+ return dequant_mxfp6(w, w_scale, quant_dtype="fp6_e3m2", float_dtype=dtype)
118
+ elif self.ocp_mx_scheme.startswith("w_mxfp6_e2m3"): # type: ignore[union-attr]
119
+ return dequant_mxfp6(w, w_scale, quant_dtype="fp6_e2m3", float_dtype=dtype)
120
+ else:
121
+ raise NotImplementedError(f"Unsupported ocp_mx_scheme={self.ocp_mx_scheme}")
122
+
123
+ def apply(
124
+ self,
125
+ output: torch.Tensor,
126
+ hidden_states: torch.Tensor,
127
+ w1: torch.Tensor,
128
+ w2: torch.Tensor,
129
+ topk_weights: torch.Tensor,
130
+ topk_ids: torch.Tensor,
131
+ activation: MoEActivation,
132
+ global_num_experts: int,
133
+ expert_map: torch.Tensor | None,
134
+ a1q_scale: torch.Tensor | None,
135
+ a2_scale: torch.Tensor | None,
136
+ workspace13: torch.Tensor,
137
+ workspace2: torch.Tensor,
138
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
139
+ apply_router_weight_on_input: bool,
140
+ ):
141
+ """
142
+ Apply emulated quantized MoE computation.
143
+
144
+ This dequantizes the weights on the fly and calls TritonExperts.apply
145
+ with activation quantization support.
146
+ """
147
+ assert w1.dtype == torch.uint8
148
+ assert w2.dtype == torch.uint8
149
+
150
+ # Dequantize w1 and w2 from packed OCP MX format to bf16/fp16
151
+ w1_dequant = self._dequantize_weights(
152
+ w1, self.w1_scale_val, hidden_states.dtype
153
+ )
154
+ w2_dequant = self._dequantize_weights(
155
+ w2, self.w2_scale_val, hidden_states.dtype
156
+ )
157
+
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
+
168
+ # Activation quantization/dequantization is deferred to
169
+ # `moe_kernel_quantize_input` in TritonExperts.apply.
170
+ super().apply(
171
+ output=output,
172
+ hidden_states=hidden_states,
173
+ w1=w1_dequant,
174
+ w2=w2_dequant,
175
+ topk_weights=topk_weights,
176
+ topk_ids=topk_ids,
177
+ activation=activation,
178
+ global_num_experts=global_num_experts,
179
+ expert_map=expert_map,
180
+ a1q_scale=None,
181
+ a2_scale=None,
182
+ workspace13=workspace13,
183
+ workspace2=workspace2,
184
+ expert_tokens_meta=expert_tokens_meta,
185
+ apply_router_weight_on_input=apply_router_weight_on_input,
186
+ )
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+5
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import flashinfer
4
+
5
import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -188,6 +188,8 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
188
expert_tokens_meta: mk.ExpertTokensMetadata | None,
189
apply_router_weight_on_input: bool,
190
):
191
+ import flashinfer
192
+
193
assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
194
assert a1q_scale is not None
195
assert self.quant_config.w1_scale is not None
@@ -306,6 +308,8 @@ class TrtLlmNvFp4ExpertsMonolithic(
308
routed_scaling_factor: float | None = None,
309
topk_group: int | None = None,
310
) -> torch.Tensor:
311
+ import flashinfer
312
+
313
assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
314
assert a1q_scale is not None
315
assert self.quant_config.w1_scale is not None
vllm/model_executor/layers/fused_moe/fused_moe.py
+11
-47
@@ -36,8 +36,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
36
disable_inplace,
37
moe_kernel_quantize_input,
38
)
39
-from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
40
-from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
39
from vllm.model_executor.layers.quantization.utils.quant_utils import (
40
QuantKey,
41
kFp8Dynamic128Sym,
@@ -1708,22 +1706,18 @@ def fused_experts_impl(
1706
w1_bias: torch.Tensor | None = None,
1707
w2_bias: torch.Tensor | None = None,
1708
) -> torch.Tensor:
1709
+ if ocp_mx_scheme is not None:
1710
+ raise NotImplementedError(
1711
+ f"Using ocp_mx_scheme={ocp_mx_scheme} in functional fused_experts call is "
1712
+ "deprecated. Please use OCP_MXQuantizationEmulationTritonExperts."
1713
+ )
1714
+
1715
# Convert string activation to enum for internal use
1716
activation_enum = MoEActivation.from_str(activation)
1717
1718
# Check constraints.
1719
if use_int4_w4a16:
1720
assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch"
1717
- elif ocp_mx_scheme is not None:
1718
- if ocp_mx_scheme.startswith("w_mxfp4"):
1719
- # 16bit activation and fp4x2 packed weight
1720
- assert hidden_states.size(1) == w1.size(2) * 2, "hidden size mismatch"
1721
- elif ocp_mx_scheme.startswith("w_mxfp6"):
1722
- assert hidden_states.size(1) == (w1.size(2) * 4) // 3, (
1723
- "hidden size mismatch"
1724
- )
1725
- else:
1726
- raise NotImplementedError(f"Unsupported ocp_mx_scheme={ocp_mx_scheme}")
1721
else:
1722
assert hidden_states.size(1) == w1.size(2), (
1723
f"Hidden size mismatch {hidden_states.size(1)} != {w1.size(2)}"
@@ -1748,7 +1742,6 @@ def fused_experts_impl(
1742
use_fp8_w8a8=use_fp8_w8a8,
1743
use_int8_w8a16=use_int8_w8a16,
1744
use_int4_w4a16=use_int4_w4a16,
1751
- ocp_mx_scheme=ocp_mx_scheme,
1745
dtype=hidden_states.dtype,
1746
)
1747
@@ -1757,7 +1750,7 @@ def fused_experts_impl(
1750
quant_dtype = _get_config_quant_dtype(
1751
use_fp8_w8a8=use_fp8_w8a8,
1752
use_int8_w8a8=use_int8_w8a8,
1760
- ocp_mx_scheme=ocp_mx_scheme,
1753
+ ocp_mx_scheme=None,
1754
)
1755
1756
get_config_func = functools.partial(
@@ -1802,44 +1795,12 @@ def fused_experts_impl(
1795
1796
out_hidden_states = hidden_states if inplace else torch.empty_like(hidden_states)
1797
1805
- if ocp_mx_scheme is not None:
1806
- # TODO: On platforms for which `current_platform.supports_mx()` is True
1807
- # and for which we have a native OCP mx fused MOE kernel,
1808
- # this dequantization step should not be done.
1809
- if ocp_mx_scheme.startswith("w_mxfp4"):
1810
- # Weight has to be dequantized for mxfp4 emulation.
1811
- w1 = dequant_mxfp4(w1, w1_scale, hidden_states.dtype)
1812
- w1_scale = None
1813
- w2 = dequant_mxfp4(w2, w2_scale, hidden_states.dtype)
1814
- w2_scale = None
1815
- elif ocp_mx_scheme.startswith("w_mxfp6_e3m2"):
1816
- w1 = dequant_mxfp6(
1817
- w1, w1_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
1818
- )
1819
- w1_scale = None
1820
- w2 = dequant_mxfp6(
1821
- w2, w2_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
1822
- )
1823
- w2_scale = None
1824
- elif ocp_mx_scheme.startswith("w_mxfp6_e2m3"):
1825
- w1 = dequant_mxfp6(
1826
- w1, w1_scale, quant_dtype="fp6_e2m3", float_dtype=hidden_states.dtype
1827
- )
1828
- w1_scale = None
1829
- w2 = dequant_mxfp6(
1830
- w2, w2_scale, quant_dtype="fp6_e2m3", float_dtype=hidden_states.dtype
1831
- )
1832
- w2_scale = None
1833
- else:
1834
- raise NotImplementedError(f"Unsupported ocp_mx_scheme={ocp_mx_scheme}")
1835
-
1798
qhidden_states, a1q_scale = moe_kernel_quantize_input(
1799
A=hidden_states,
1800
A_scale=a1_scale,
1801
quant_dtype=quant_dtype,
1802
per_act_token_quant=per_channel_quant,
1803
block_shape=block_shape,
1842
- ocp_mx_scheme=ocp_mx_scheme,
1804
)
1805
1806
sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment(
@@ -1889,7 +1850,6 @@ def fused_experts_impl(
1850
quant_dtype=quant_dtype,
1851
per_act_token_quant=per_channel_quant,
1852
block_shape=block_shape,
1892
- ocp_mx_scheme=ocp_mx_scheme,
1853
)
1854
1855
if expert_map is not None:
@@ -1935,6 +1895,9 @@ class TritonExperts(mk.FusedMoEExpertsModular):
1895
moe_config: FusedMoEConfig,
1896
quant_config: FusedMoEQuantConfig,
1897
):
1898
+ # Whether quantized MOE runs natively, or through
1899
+ # higher-precision + activation QDQ.
1900
+ self.quantization_emulation = False
1901
super().__init__(moe_config, quant_config)
1902
1903
@staticmethod
@@ -2144,6 +2107,7 @@ class TritonExperts(mk.FusedMoEExpertsModular):
2107
self.quant_dtype,
2108
self.per_act_token_quant,
2109
self.block_shape,
2110
+ quantization_emulation=self.quantization_emulation,
2111
)
2112
2113
invoke_fused_moe_triton_kernel(
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+22
@@ -62,6 +62,8 @@ class Mxfp4MoeBackend(Enum):
62
TRITON_UNFUSED = "TRITON_UNFUSED"
63
# XPU
64
XPU = "XPU"
65
+ # Emulation
66
+ EMULATION = "EMULATION"
67
68
69
# Backends that share the same TRTLLM weight format
@@ -143,6 +145,13 @@ def backend_to_kernel_cls(
145
146
return [XPUExpertsMXFp4]
147
148
+ elif backend == Mxfp4MoeBackend.EMULATION:
149
+ from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import (
150
+ OCP_MXQuantizationEmulationTritonExperts,
151
+ )
152
+
153
+ return [OCP_MXQuantizationEmulationTritonExperts]
154
+
155
else:
156
raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}")
157
@@ -158,6 +167,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend:
167
"marlin": Mxfp4MoeBackend.MARLIN,
168
"aiter": Mxfp4MoeBackend.AITER,
169
"xpu": Mxfp4MoeBackend.XPU,
170
+ "emulation": Mxfp4MoeBackend.EMULATION,
171
}
172
if backend := mapping.get(runner_backend):
173
return backend
@@ -181,6 +191,7 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
191
Mxfp4MoeBackend.MARLIN,
192
Mxfp4MoeBackend.BATCHED_MARLIN,
193
Mxfp4MoeBackend.XPU,
194
+ Mxfp4MoeBackend.EMULATION,
195
]
196
return _AVAILABLE_BACKENDS
197
@@ -768,6 +779,17 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
779
w13_bias,
780
w2_bias,
781
)
782
+ elif mxfp4_backend == Mxfp4MoeBackend.EMULATION:
783
+ # No additional transformation needed for emulation backend,
784
+ # weights are dequantized on the fly in the experts class.
785
+ return (
786
+ w13_weight,
787
+ w2_weight,
788
+ w13_weight_scale,
789
+ w2_weight_scale,
790
+ w13_bias,
791
+ w2_bias,
792
+ )
793
else:
794
raise ValueError(
795
f"Unsupported mxfp4_backend: {mxfp4_backend}: "
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+42
@@ -45,6 +45,7 @@ class NvFp4MoeBackend(Enum):
45
FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED"
46
VLLM_CUTLASS = "VLLM_CUTLASS"
47
MARLIN = "MARLIN"
48
+ EMULATION = "EMULATION"
49
50
51
FLASHINFER_NVFP4_MOE_BACKENDS = [
@@ -118,6 +119,12 @@ def backend_to_kernel_cls(
119
)
120
121
return [MarlinExperts]
122
+ elif backend == NvFp4MoeBackend.EMULATION:
123
+ from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import (
124
+ Nvfp4QuantizationEmulationTritonExperts,
125
+ )
126
+
127
+ return [Nvfp4QuantizationEmulationTritonExperts]
128
else:
129
raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}")
130
@@ -130,6 +137,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend:
137
"flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS,
138
"flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL,
139
"marlin": NvFp4MoeBackend.MARLIN,
140
+ "emulation": NvFp4MoeBackend.EMULATION,
141
}
142
if backend := mapping.get(runner_backend):
143
return backend
@@ -157,6 +165,7 @@ def select_nvfp4_moe_backend(
165
NvFp4MoeBackend.FLASHINFER_CUTLASS,
166
NvFp4MoeBackend.VLLM_CUTLASS,
167
NvFp4MoeBackend.MARLIN,
168
+ NvFp4MoeBackend.EMULATION,
169
]
170
171
# NOTE(rob): this is kind of a hack. We need to peak into
@@ -372,6 +381,30 @@ def convert_to_nvfp4_moe_kernel_format(
381
w2_scale_2=w2_scale_2,
382
is_act_and_mul=is_act_and_mul,
383
)
384
+ elif nvfp4_backend == NvFp4MoeBackend.EMULATION:
385
+ if a13_scale is None or a2_scale is None:
386
+ raise ValueError(
387
+ "Activation global scales should not be None, got"
388
+ f" a13_scale={a13_scale}, a2_scale={a2_scale}"
389
+ )
390
+
391
+ if torch.unique(a13_scale).numel() != 1 or torch.unique(a2_scale).numel() != 1:
392
+ logger.warning_once(
393
+ "In NVFP4 linear, the activation global scale for inputs are different"
394
+ " for MOE w13 (gate_up_proj) layer or MOE w2 (down_proj). Using"
395
+ " a13_scale = a13_scale.max() and a2_scale = a2_scale.max()."
396
+ )
397
+
398
+ # 1. We take the max following e.g. quantization/utils/flashinfer_fp4_moe.py.
399
+ # 2. moe_kernel_quantize_input -> ref_nvfp4_quant_dequant
400
+ # use the inverse scale directly (large global scale).
401
+ # NOTE: Before this point, `a13_scale` and `a2_scale` are such that:
402
+ # `FP8_MAX = activation[expert_id].abs().max() * global_scale[expert_id]`,
403
+ # and `global_scale[expert_id]` are small (~1e-4).
404
+ # Taking the largest global scale likely results in overflowing the FP8 range
405
+ # for other experts - other selection strategies may be used.
406
+ a13_scale = 1.0 / a13_scale.max().to(torch.float32)
407
+ a2_scale = 1.0 / a2_scale.max().to(torch.float32)
408
else:
409
raise ValueError(f"Unknown NvFp4 backend for MoE: {nvfp4_backend}")
410
@@ -403,6 +436,15 @@ def make_nvfp4_moe_quant_config(
436
w1_scale=w13_scale,
437
w2_scale=w2_scale,
438
)
439
+ elif backend == NvFp4MoeBackend.EMULATION:
440
+ return nvfp4_moe_quant_config(
441
+ g1_alphas=w13_scale_2,
442
+ g2_alphas=w2_scale_2,
443
+ a1_gscale=a13_scale,
444
+ a2_gscale=a2_scale,
445
+ w1_scale=w13_scale,
446
+ w2_scale=w2_scale,
447
+ )
448
449
# Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas.
450
# The expert's process_weights_after_loading will fuse activation
vllm/model_executor/layers/fused_moe/utils.py
+42
-1
@@ -22,6 +22,9 @@ from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
22
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
23
mxfp8_e4m3_quantize,
24
)
25
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
26
+ ref_nvfp4_quant_dequant,
27
+)
28
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
29
per_tensor_dequantize,
30
)
@@ -253,6 +256,7 @@ def moe_kernel_quantize_input(
256
block_shape: list[int] | None = None,
257
is_fp4_scale_swizzled: bool = True,
258
ocp_mx_scheme: str | None = None,
259
+ quantization_emulation: bool = False,
260
) -> tuple[torch.Tensor, torch.Tensor | None]:
261
# Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation
262
if ocp_mx_scheme is not None:
@@ -274,16 +278,41 @@ def moe_kernel_quantize_input(
278
# activation quantization below.
279
280
if quant_dtype == current_platform.fp8_dtype():
281
+ if quantization_emulation:
282
+ raise NotImplementedError(
283
+ f"moe_kernel_quantize_input does not support quant_dtype={quant_dtype}"
284
+ " MOE quantization emulation. Please open an issue."
285
+ )
286
return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
287
elif quant_dtype == torch.int8:
288
+ if quantization_emulation:
289
+ raise NotImplementedError(
290
+ "moe_kernel_quantize_input does not support quant_dtype=torch.int8"
291
+ " MOE quantization emulation. Please open an issue."
292
+ )
293
return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)
294
elif quant_dtype == "nvfp4":
281
- return _nvfp4_quantize(A, A_scale, is_sf_swizzled_layout=is_fp4_scale_swizzled)
295
+ if not quantization_emulation:
296
+ return _nvfp4_quantize(
297
+ A, A_scale, is_sf_swizzled_layout=is_fp4_scale_swizzled
298
+ )
299
+ else:
300
+ return ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
301
elif quant_dtype == "mxfp4":
302
+ if not quantization_emulation:
303
+ raise NotImplementedError(
304
+ "moe_kernel_quantize_input should not be used for native"
305
+ " quant_dtype='mxfp4' MOE. Please open an issue."
306
+ )
307
return _mxfp4_quantize(A, A_scale, per_act_token_quant, block_shape)
308
elif quant_dtype == "mxfp8":
309
# TODO: `quant_dtype == "mxfp8"` is ambiguous,
310
# should be fp8_e4m3. OCP MX also defines `fp8_e5m2`.
311
+ if quantization_emulation:
312
+ raise NotImplementedError(
313
+ "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE "
314
+ "quantization emulation. Please open an issue."
315
+ )
316
return _mxfp8_e4m3_quantize(
317
A,
318
A_scale,
@@ -292,8 +321,20 @@ def moe_kernel_quantize_input(
321
is_sf_swizzled_layout=is_fp4_scale_swizzled,
322
)
323
elif quant_dtype == "mxfp6_e3m2":
324
+ if not quantization_emulation:
325
+ raise NotImplementedError(
326
+ "moe_kernel_quantize_input should not be used for native "
327
+ " quant_dtype='mxfp6_e3m2'MOE. Please open an issue."
328
+ )
329
+
330
return _mxfp6_e3m2_quantize(A, A_scale, per_act_token_quant, block_shape)
331
elif quant_dtype == "mxfp6_e2m3":
332
+ if not quantization_emulation:
333
+ raise NotImplementedError(
334
+ "moe_kernel_quantize_input should not be used for native"
335
+ " quant_dtype='mxfp6_e2m3' MOE. Please open an issue."
336
+ )
337
+
338
return _mxfp6_e2m3_quantize(A, A_scale, per_act_token_quant, block_shape)
339
else:
340
return A, A_scale
vllm/model_executor/layers/quantization/quark/quark_moe.py
+43
-54
@@ -30,6 +30,7 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_m
30
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
31
TRITON_BACKENDS,
32
Mxfp4MoeBackend,
33
+ backend_to_kernel_cls,
34
convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
35
make_mxfp4_moe_kernel,
36
make_mxfp4_moe_quant_config,
@@ -986,6 +987,8 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
987
f"Please check that the combination is supported in OCP_MX_Scheme."
988
)
989
990
+ # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes,
991
+ # use kernel abstraction for all OCP MX MOE implementations.
992
self.mxfp4_backend: Mxfp4MoeBackend = Mxfp4MoeBackend.NONE
993
self.experts_cls: type[mk.FusedMoEExperts] | None = None
994
self.moe_kernel: mk.FusedMoEKernel | None = None
@@ -994,12 +997,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
997
self.w13_precision_config = None
998
self.w2_precision_config = None
999
997
- if self.ocp_mx_scheme == "w_mxfp4":
998
- self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe)
999
- elif self.ocp_mx_scheme.startswith("w_mxfp4"):
1000
- # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes.
1001
- self.mxfp4_backend = Mxfp4MoeBackend.NONE
1002
-
1000
if self.input_quant is not None:
1001
self.static_input_scales = not self.input_quant.get("is_dynamic")
1002
else:
@@ -1035,6 +1032,18 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1032
self.mxfp4_backend is Mxfp4MoeBackend.NONE or not self.use_rocm_aiter_moe
1033
)
1034
1035
+ if self.ocp_mx_scheme == "w_mxfp4":
1036
+ self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe)
1037
+
1038
+ if self.emulate:
1039
+ # We use the same code path between MXFP4/MXFP6 emulation.
1040
+ self.mxfp4_backend = Mxfp4MoeBackend.EMULATION
1041
+
1042
+ # TODO: Remove `self.mxfp4_backend != Mxfp4MoeBackend.NONE` and make it so that
1043
+ # all MXFP4 backends use the kernel abstraction.
1044
+ if self.mxfp4_backend != Mxfp4MoeBackend.NONE:
1045
+ self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]
1046
+
1047
if self.emulate:
1048
logger.warning_once(
1049
f"The current mode (supports_mx={current_platform.supports_mx()}, "
@@ -1063,7 +1072,12 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1072
act_dtype=act_dtype,
1073
moe_parallel_config=moe_parallel_config,
1074
)
1066
- if self.mxfp4_backend is not None:
1075
+ # In case quantization emulation backend is used, there is no need to apply
1076
+ # MXFP4-specific padding logic as the compute happens in higher precision.
1077
+ if (
1078
+ self.mxfp4_backend is not None
1079
+ and self.mxfp4_backend != Mxfp4MoeBackend.EMULATION
1080
+ ):
1081
hidden_size, intermediate_size_per_partition = (
1082
mxfp4_round_up_hidden_size_and_intermediate_size(
1083
self.mxfp4_backend, hidden_size, intermediate_size_per_partition
@@ -1237,7 +1251,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1251
)
1252
1253
# For w_mxfp4, use oracle functions
1240
- if (
1254
+ if self.emulate or (
1255
self.ocp_mx_scheme == "w_mxfp4"
1256
and self.mxfp4_backend != Mxfp4MoeBackend.NONE
1257
):
@@ -1245,13 +1259,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1259
return
1260
1261
# TODO(bowenbao): gradually migrate to oracles.
1248
- # secondly, process mxfp weights for other schemes
1249
- if self.emulate:
1250
- # Build quant config for emulation path
1251
- self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1252
- torch.accelerator.empty_cache()
1253
- return
1254
-
1262
# Existing AITER path for w_mxfp4_a_mxfp4 and other schemes
1263
from aiter.utility.fp4_utils import e8m0_shuffle
1264
@@ -1345,9 +1352,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1352
self, layer: torch.nn.Module
1353
) -> FusedMoEQuantConfig | None:
1354
# For w_mxfp4 with oracle backend, use oracle function
1348
- if (
1349
- self.ocp_mx_scheme == "w_mxfp4"
1350
- and self.mxfp4_backend != Mxfp4MoeBackend.NONE
1355
+ if self.ocp_mx_scheme == "w_mxfp4" and self.mxfp4_backend not in (
1356
+ Mxfp4MoeBackend.NONE,
1357
+ Mxfp4MoeBackend.EMULATION,
1358
):
1359
w1_scale = layer.w13_weight_scale
1360
w2_scale = layer.w2_weight_scale
@@ -1362,9 +1369,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1369
w2_bias=getattr(layer, "w2_bias", None),
1370
)
1371
1365
- # Existing code for other schemes
1366
- # TODO(bowenbao): kept for emulation fallback, to be refactored into
1367
- # dedicated emulation backend.
1372
+ # Emulation and other schemes
1373
if self.ocp_mx_scheme == "w_mxfp4":
1374
return mxfp4_w4a16_moe_quant_config(
1375
w1_scale=layer.w13_weight_scale,
@@ -1414,7 +1419,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1419
topk_ids: torch.Tensor,
1420
shared_experts_input: torch.Tensor | None,
1421
) -> torch.Tensor:
1417
- # For w_mxfp4 with oracle kernel
1422
+ # For oracle kernel or emulation kernel
1423
if self.moe_kernel is not None:
1424
return self.moe_kernel.apply(
1425
hidden_states=x,
@@ -1429,39 +1434,23 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1434
shared_experts_input=shared_experts_input,
1435
)
1436
1432
- # Existing code for emulation/AITER paths
1433
- if not self.emulate:
1434
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
1435
- rocm_aiter_fused_experts,
1436
- )
1437
-
1438
- return rocm_aiter_fused_experts(
1439
- x,
1440
- layer.w13_weight,
1441
- layer.w2_weight,
1442
- topk_weights=topk_weights,
1443
- topk_ids=topk_ids,
1444
- activation=layer.activation,
1445
- quant_config=self.moe_quant_config,
1446
- moe_config=layer.moe_config,
1447
- expert_map=layer.expert_map,
1448
- )
1449
- else:
1450
- from vllm.model_executor.layers.fused_moe import fused_experts
1437
+ # AITER path
1438
+ # TODO: Refactor this to use modular MOE kernel as well.
1439
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
1440
+ rocm_aiter_fused_experts,
1441
+ )
1442
1452
- return fused_experts(
1453
- x,
1454
- layer.w13_weight,
1455
- layer.w2_weight,
1456
- topk_weights=topk_weights,
1457
- topk_ids=topk_ids,
1458
- inplace=not self.moe.disable_inplace,
1459
- activation=layer.activation,
1460
- global_num_experts=layer.global_num_experts,
1461
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1462
- expert_map=layer.expert_map,
1463
- quant_config=self.moe_quant_config,
1464
- )
1443
+ return rocm_aiter_fused_experts(
1444
+ x,
1445
+ layer.w13_weight,
1446
+ layer.w2_weight,
1447
+ topk_weights=topk_weights,
1448
+ topk_ids=topk_ids,
1449
+ activation=layer.activation,
1450
+ quant_config=self.moe_quant_config,
1451
+ moe_config=layer.moe_config,
1452
+ expert_map=layer.expert_map,
1453
+ )
1454
1455
def apply_monolithic(
1456
self,
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+56
-17
@@ -53,26 +53,52 @@ def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
53
def dequantize_to_dtype(
54
tensor_fp4: torch.Tensor,
55
tensor_sf: torch.Tensor,
56
- global_scale: torch.Tensor | float,
56
+ global_scale: torch.Tensor,
57
dtype: torch.dtype,
58
block_size: int = 16,
59
swizzle: bool | None = True,
60
):
61
- """Dequantize the fp4 tensor back to high precision."""
61
+ """Dequantize the fp4 tensor back to high precision.
62
+
63
+ Supports both 2D and 3D inputs:
64
+ - 2D: [m, packed_k] -> [m, k]
65
+ - 3D: [dim0, m, packed_k] -> [dim0, m, k]
66
+ """
67
# Two fp4 values are packed into one uint8.
68
assert tensor_fp4.dtype == torch.uint8
64
- m, packed_k = tensor_fp4.shape
69
+
70
+ # We handle 3D tensors reshaping them to 2D.
71
+ is_3d = tensor_fp4.ndim == 3
72
+
73
+ if is_3d:
74
+ dim0, m, packed_k = tensor_fp4.shape
75
+ tensor_fp4 = tensor_fp4.reshape(-1, packed_k)
76
+ tensor_sf = tensor_sf.reshape(-1, tensor_sf.shape[-1])
77
+ global_scale = global_scale[:, None, None]
78
+ else:
79
+ m, packed_k = tensor_fp4.shape
80
+
81
k = packed_k * 2
82
tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32)
67
- tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
83
+ tensor_f32 = tensor_f32.reshape(-1, k // block_size, block_size)
84
tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
85
86
if swizzle:
71
- tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
87
+ tensor_sf = convert_swizzled_to_linear( # noqa: E501
88
+ tensor_sf, tensor_f32.size(0), k, block_size
89
+ )
90
+
91
+ if is_3d:
92
+ tensor_sf = tensor_sf.reshape(dim0, m, k // block_size)
93
tensor_sf_dtype = tensor_sf.to(torch.float32) * global_scale
94
95
+ if is_3d:
96
+ tensor_f32 = tensor_f32.reshape(dim0, m, -1, block_size)
97
+
98
# scale the tensor
75
- out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
99
+ out = tensor_f32 * tensor_sf_dtype.unsqueeze(-1)
100
+ out = out.reshape(*out.shape[:-2], -1)
101
+
102
return out.to(dtype)
103
104
@@ -117,6 +143,28 @@ def ref_nvfp4_quant(x, global_scale, block_size):
143
return cast_to_fp4(clipped_x), scale.squeeze(-1)
144
145
146
+def ref_nvfp4_quant_dequant(
147
+ x: torch.Tensor, global_scale: torch.Tensor, block_size: int
148
+) -> tuple[torch.Tensor, None]:
149
+ """
150
+ NVFP4 quantize-dequantize operation.
151
+
152
+ `global_scale` is expected to have a single element.
153
+ """
154
+ x_m, x_k = x.shape
155
+ output_dtype = x.dtype
156
+
157
+ # quantize input to (FP4 and interleaved block scale)
158
+ x_fp4, x_blockscale = ref_nvfp4_quant(x, global_scale, block_size)
159
+
160
+ # dequantize input
161
+ x_fp4 = x_fp4.reshape(x_m, x_k // block_size, block_size)
162
+ x_blockscale = x_blockscale.unsqueeze(-1) / global_scale
163
+ x_dq = (x_fp4 * x_blockscale).reshape(x_m, x_k).to(output_dtype)
164
+
165
+ return x_dq, None
166
+
167
+
168
def run_nvfp4_emulations(
169
x: torch.Tensor,
170
input_global_scale: torch.Tensor,
@@ -125,18 +173,10 @@ def run_nvfp4_emulations(
173
weight_global_scale: torch.Tensor,
174
swizzle: bool | None = True,
175
):
128
- group_size = 16
129
- x_m, x_k = x.shape
176
output_dtype = x.dtype
177
+ group_size = 16
178
132
- # quantize input to (FP4 and interleaved block scale)
133
- x_fp4, x_blockscale = ref_nvfp4_quant(x, input_global_scale, group_size)
134
-
135
- # dequantize input
136
- x_fp4 = x_fp4.reshape(x_m, x_k // group_size, group_size)
137
- x_blockscale = x_blockscale.unsqueeze(-1) / input_global_scale
138
- x_dq = (x_fp4 * x_blockscale).reshape(x_m, x_k).to(output_dtype)
139
- del x_fp4, x_blockscale
179
+ x_dq, _ = ref_nvfp4_quant_dequant(x, input_global_scale, block_size=group_size)
180
181
# dequantize weight
182
w_fp4 = weight.data.view(torch.uint8)
@@ -151,5 +191,4 @@ def run_nvfp4_emulations(
191
192
# matmul
193
out = torch.matmul(x_dq, w_dq.t())
154
- del w_dq, x_dq
194
return out