@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fb21d8b4f9

Fork 0

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 2, 23:15 UTC

Add NVFP4 MOE support for Deepseek V4. (#42209)

Signed-off-by: Shiyang Chen <shiychen@nvidia.com>

sychen52 committed May 22, 2026 at 07:21 UTC fb21d8b4f9027f4642637c7bb0acc08c29dce387

9 files changed +217 -17

.buildkite/test_areas/kernels.yaml

+7 -3

@@ -150,9 +150,12 @@ steps:

150 - csrc/quantization/fp4/

151 - csrc/attention/mla/

152 - csrc/quantization/cutlass_w8a8/moe/

153 - - vllm/model_executor/layers/fused_moe/cutlass_moe.py

154 - - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

155 - - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py

153 + - vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py

154 + - vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py

155 + - vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

156 + - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

157 + - vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py

158 + - vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_two_sided.py

159 - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

160 - vllm/v1/attention/backends/flashinfer.py

161 - vllm/v1/attention/backends/mla/cutlass_mla.py

@@ -185,6 +188,7 @@ steps:

188 - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py

189 - pytest -v -s tests/kernels/moe/test_flashinfer.py

190 - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py

191 + - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py

192 - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py

193 # e2e

194 - pytest -v -s tests/models/quantization/test_nvfp4.py

tests/kernels/moe/test_trtllm_nvfp4_moe.py

+76 -7

@@ -22,6 +22,8 @@ from tests.kernels.quantization.nvfp4_utils import (

22 from tests.kernels.utils import torch_moe

23 from vllm import _custom_ops as ops

24 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

25 +from vllm.model_executor.custom_op import CustomOp, op_registry

26 +from vllm.model_executor.layers.activation import SiluAndMulWithClamp

27 from vllm.model_executor.layers.fused_moe import fused_topk

28 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

29 from vllm.model_executor.layers.fused_moe.all2all_utils import (

@@ -58,15 +60,45 @@ MNK_FACTORS = [

60 (64, 704, 4096),

61 ]

63 +_SWIGLU_LIMIT = 0.1

64 +_LARGE_OUTPUT1_SCALE = 32768.0

65 +_CLAMP_OP_NAME = "test_silu_and_mul_with_clamp"

66 +

67 +# Test-only fixed-limit clamp. ``custom_op_name`` makes the class itself

68 +# valid as an ``activation=`` argument to ``torch_moe`` (which only looks

69 +# up ``activation.custom_op_name`` in ``op_registry``), so no

70 +# ``MoEActivation`` enum extension is needed.

71 +if _CLAMP_OP_NAME not in op_registry:

72 +

73 + @CustomOp.register(_CLAMP_OP_NAME)

74 + class _SiluAndMulWithClampTest(SiluAndMulWithClamp):

75 + custom_op_name = _CLAMP_OP_NAME

76 +

77 + def __init__(self, *, compile_native: bool = True) -> None:

78 + super().__init__(_SWIGLU_LIMIT, compile_native=compile_native)

79 +

80 +

81 +SILU_WITH_CLAMP = op_registry[_CLAMP_OP_NAME]

82 +

83 +

84 +ACTIVATION_CASES = [

85 + pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"),

86 + pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"),

87 + pytest.param(

88 + MoEActivation.RELU2_NO_MUL,

89 + MoEActivation.RELU2_NO_MUL,

90 + None,

91 + id="relu2_no_mul",

92 + ),

93 + pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, id="gelu"),

94 +]

95 +

97 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)

98 @pytest.mark.parametrize("e", [128])

99 @pytest.mark.parametrize("topk", [8])

100 @pytest.mark.parametrize("dtype", [torch.bfloat16])

66 -@pytest.mark.parametrize(

67 - "activation",

68 - [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, MoEActivation.GELU],

69 -)

101 +@pytest.mark.parametrize("activation,torch_activation,swiglu_limit", ACTIVATION_CASES)

102 @torch.inference_mode()

103 def test_trtllm_fp4_moe_no_graph(

104 m: int,

@@ -76,6 +108,8 @@ def test_trtllm_fp4_moe_no_graph(

108 topk: int,

109 dtype: torch.dtype,

110 activation: MoEActivation,

111 + torch_activation: MoEActivation | type[SiluAndMulWithClamp],

112 + swiglu_limit: float | None,

113 workspace_init,

114 ):

115 # FlashInfer's trtllm_batched_gemm_runner has no precompiled tile

@@ -113,6 +147,15 @@ def test_trtllm_fp4_moe_no_graph(

147 # Match what oracle/nvfp4.py does for this backend.

148 is_scale_swizzled=False,

149 )

150 + quant_config.gemm1_clamp_limit = swiglu_limit

151 + if swiglu_limit is not None:

152 + assert quant_config.g1_alphas is not None

153 + assert quant_config.a2_gscale is not None

154 + assert torch.all(quant_config.a2_gscale == 1)

155 + # With a2_gscale == 1, g1_alphas is the TRTLLM

156 + # output1_scale_gate_scalar. Make it large enough to catch

157 + # clamp/output-scale coupling in the FlashInfer kernel wrapper.

158 + quant_config.g1_alphas.fill_(_LARGE_OUTPUT1_SCALE)

159

160 score = torch.randn((m, e), device="cuda", dtype=dtype)

161 topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)

@@ -133,6 +176,23 @@ def test_trtllm_fp4_moe_no_graph(

176 max_num_tokens=next_power_of_2(m),

177 )

178

179 + trtllm_inner = TrtLlmNvFp4ExpertsModular(

180 + moe_config=moe_config, quant_config=quant_config

181 + )

182 + # Mimic the production weight-loader path so per-expert tensors that

183 + # are normally precomputed in process_weights_after_loading (g1_scale_c

184 + # and the rescaled gemm1_clamp_limit) get materialized. The test's

185 + # synthetic quant_config has g1_alphas/g2_alphas already at their

186 + # post-fusion values, so we set w13_weight_scale_2 to alias g1_alphas

187 + # (same tensor) and use input_scale=1 to make the in-place

188 + # weight_scale_2 *= input_scale step a no-op.

189 + fake_layer = torch.nn.Module()

190 + fake_layer.w13_weight_scale_2 = quant_config.g1_alphas

191 + fake_layer.w2_weight_scale_2 = quant_config.g2_alphas

192 + fake_layer.w13_input_scale = torch.ones_like(quant_config.g1_alphas)

193 + fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas)

194 + trtllm_inner.process_weights_after_loading(fake_layer)

195 +

196 trtllm_experts = mk.FusedMoEKernel(

197 maybe_make_prepare_finalize(

198 moe=moe_config,

@@ -140,7 +200,7 @@ def test_trtllm_fp4_moe_no_graph(

200 allow_new_interface=True,

201 use_monolithic=False,

202 ),

143 - TrtLlmNvFp4ExpertsModular(moe_config=moe_config, quant_config=quant_config),

203 + trtllm_inner,

204 inplace=False,

205 )

206

@@ -195,7 +255,7 @@ def test_trtllm_fp4_moe_no_graph(

255 )

256

257 torch_output = torch_moe(

198 - a_in_dtype, w1_d, w2_d, score, topk, activation=activation

258 + a_in_dtype, w1_d, w2_d, score, topk, activation=torch_activation

259 )

260

261 torch.testing.assert_close(torch_output, trtllm_output, atol=2e-1, rtol=2e-1)

@@ -203,5 +263,14 @@ def test_trtllm_fp4_moe_no_graph(

263

264 if __name__ == "__main__":

265 test_trtllm_fp4_moe_no_graph(

206 - 64, 704, 4096, 128, 8, torch.bfloat16, MoEActivation.GELU, None

266 + 64,

267 + 704,

268 + 4096,

269 + 128,

270 + 8,

271 + torch.bfloat16,

272 + MoEActivation.GELU,

273 + MoEActivation.GELU,

274 + None,

275 + None,

276 )

vllm/model_executor/layers/fused_moe/config.py

@@ -805,6 +805,7 @@ def nvfp4_moe_quant_config(

805 w1_bias: torch.Tensor | None = None,

806 w2_bias: torch.Tensor | None = None,

807 is_scale_swizzled: bool = True,

808 + gemm1_clamp_limit: float | None = None,

809 ) -> FusedMoEQuantConfig:

810 """

811 Construct a quant config for mxfp4 activations and nvp4 weights.

@@ -823,6 +824,7 @@ def nvfp4_moe_quant_config(

824 per_out_ch_quant=False,

825 block_shape=None,

826 is_scale_swizzled=is_scale_swizzled,

827 + gemm1_clamp_limit=gemm1_clamp_limit,

828 )

829

830

@@ -1271,6 +1273,11 @@ class FusedMoEConfig:

1273 is_act_and_mul: bool = True

1274 is_lora_enabled: bool = False

1275

1276 + # SwiGLU clamp limit. When set, backends that do not implement the clamp

1277 + # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle

1278 + # cannot silently select one and drop the clamp.

1279 + swiglu_limit: float | None = None

1280 +

1281 # This flag is used to disable the inplace optimization

1282 # in MoE kernels. If this flag is True then the kernel

1283 # should not be using inplace. If the flag is false, the

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

+32 -4

@@ -66,6 +66,17 @@ class TrtLlmNvFp4ExpertsBase:

66 else:

67 self.g1_scale_c = self.quant_config.a2_gscale.clone()

69 + if moe_config.is_act_and_mul and quant_config.gemm1_clamp_limit is not None:

70 + device = torch.accelerator.current_device_index()

71 + self.gemm1_clamp_limit = torch.full(

72 + (self.local_num_experts,),

73 + quant_config.gemm1_clamp_limit,

74 + dtype=torch.float32,

75 + device=device,

76 + )

77 + else:

78 + self.gemm1_clamp_limit = None

79 +

80 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

81 layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)

82 layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)

@@ -84,6 +95,20 @@ class TrtLlmNvFp4ExpertsBase:

95 )

96 self.g1_scale_c = layer.g1_scale_c

98 + # Pre-fold the per-expert g1_alphas (= output1_scale_gate_scalar)

99 + # division so the TRTLLM kernel receives the raw-GEMM-space clamp

100 + # directly. g1_alphas is set once here in process_weights_after_loading

101 + # (via the in-place mul above) and never changes again, so this is a

102 + # static, per-expert constant. Register on the layer so EPLB

103 + # rearranges it alongside the other expert tensors.

104 + if self.gemm1_clamp_limit is not None:

105 + gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas

106 + layer.register_parameter(

107 + "gemm1_clamp_limit",

108 + torch.nn.Parameter(gemm1_clamp_limit, requires_grad=False),

109 + )

110 + self.gemm1_clamp_limit = layer.gemm1_clamp_limit

111 +

112 @staticmethod

113 def _supports_current_device() -> bool:

114 """Supports only Blackwell-family GPUs."""

@@ -201,6 +226,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula

226

227 # Pack topk ids and weights into format expected by the kernel.

228 packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)

229 + output1_scale_gate_scalar = self.quant_config.g1_alphas

230

231 # Invoke kernel.

232 flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(

@@ -215,12 +241,12 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula

241 gemm1_bias=None,

242 gemm1_alpha=None,

243 gemm1_beta=None,

218 - gemm1_clamp_limit=None,

244 + gemm1_clamp_limit=self.gemm1_clamp_limit,

245 gemm2_weights=w2,

246 gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),

247 gemm2_bias=None,

248 output1_scale_scalar=self.g1_scale_c,

223 - output1_scale_gate_scalar=self.quant_config.g1_alphas,

249 + output1_scale_gate_scalar=output1_scale_gate_scalar,

250 output2_scale_scalar=self.quant_config.g2_alphas,

251 num_experts=global_num_experts,

252 top_k=self.topk,

@@ -313,6 +339,8 @@ class TrtLlmNvFp4ExpertsMonolithic(

339 if e_score_correction_bias is not None:

340 e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16)

341

342 + output1_scale_gate_scalar = self.quant_config.g1_alphas

343 +

344 # Invoke kernel.

345 # NOTE: Activation padding and output

346 # truncation are handled by the MoE runner's

@@ -328,12 +356,12 @@ class TrtLlmNvFp4ExpertsMonolithic(

356 gemm1_bias=None,

357 gemm1_alpha=None,

358 gemm1_beta=None,

331 - gemm1_clamp_limit=None,

359 + gemm1_clamp_limit=self.gemm1_clamp_limit,

360 gemm2_weights=w2,

361 gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),

362 gemm2_bias=None,

363 output1_scale_scalar=self.g1_scale_c,

336 - output1_scale_gate_scalar=self.quant_config.g1_alphas,

364 + output1_scale_gate_scalar=output1_scale_gate_scalar,

365 output2_scale_scalar=self.quant_config.g2_alphas,

366 num_experts=global_num_experts,

367 top_k=self.topk,

vllm/model_executor/layers/fused_moe/layer.py

@@ -335,6 +335,7 @@ class FusedMoE(PluggableLayer):

335 activation=self.activation,

336 device=vllm_config.device_config.device,

337 routing_method=self.routing_method_type,

338 + swiglu_limit=swiglu_limit,

339 # TODO: in_dtype == out_dtype?

340 disable_inplace=disable_inplace() or shared_experts is not None,

341 )

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

+39 -2

@@ -181,6 +181,15 @@ def select_nvfp4_moe_backend(

181 NvFp4MoeBackend.EMULATION,

182 ]

183

184 + NVFP4_BACKENDS_WITH_CLAMP = {

185 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

186 + }

187 +

188 + if config.swiglu_limit is not None:

189 + AVAILABLE_BACKENDS = [

190 + b for b in AVAILABLE_BACKENDS if b in NVFP4_BACKENDS_WITH_CLAMP

191 + ]

192 +

193 use_batched = config.moe_parallel_config.use_batched_activation_format

194 activation_format = (

195 mk.FusedMoEActivationFormat.BatchedExperts

@@ -234,6 +243,16 @@ def select_nvfp4_moe_backend(

243 and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL

244 ):

245 requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED

246 + if (

247 + config.swiglu_limit is not None

248 + and requested_backend not in NVFP4_BACKENDS_WITH_CLAMP

249 + ):

250 + raise ValueError(

251 + f"Model sets swiglu_limit={config.swiglu_limit}, but the "

252 + f"explicitly requested moe_backend={runner_backend!r} does "

253 + f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm' or "

254 + f"'flashinfer_cutlass' instead."

255 + )

256 return _return_or_raise(

257 requested_backend, config, weight_key, activation_key, activation_format

258 )

@@ -242,17 +261,32 @@ def select_nvfp4_moe_backend(

261 if not envs.VLLM_USE_FLASHINFER_MOE_FP4:

262 # If the user rejects FlashInfer remove those backends.

263 for b in FLASHINFER_NVFP4_MOE_BACKENDS:

245 - AVAILABLE_BACKENDS.remove(b)

264 + if b in AVAILABLE_BACKENDS:

265 + AVAILABLE_BACKENDS.remove(b)

266

267 elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):

268 # If user is explicit about backend, validate it.

269 backend = fi_2_vllm_backend_map[get_flashinfer_moe_backend()]

270 + if (

271 + config.swiglu_limit is not None

272 + and backend not in NVFP4_BACKENDS_WITH_CLAMP

273 + ):

274 + raise ValueError(

275 + f"Model sets swiglu_limit={config.swiglu_limit}, but the "

276 + f"FlashInfer backend selected via VLLM_FLASHINFER_MOE_BACKEND "

277 + f"({backend.value}) does not apply the SwiGLU clamp."

278 + )

279 return _return_or_raise(

280 backend, config, weight_key, activation_key, activation_format

281 )

282 else:

283 # If the user is not explicit about the backend, try each.

255 - for backend in FLASHINFER_NVFP4_MOE_BACKENDS:

284 + fi_backends = [

285 + b

286 + for b in FLASHINFER_NVFP4_MOE_BACKENDS

287 + if config.swiglu_limit is None or b in NVFP4_BACKENDS_WITH_CLAMP

288 + ]

289 + for backend in fi_backends:

290 for k_cls in backend_to_kernel_cls(backend):

291 supported, reason = k_cls.is_supported_config(

292 k_cls,

@@ -439,6 +473,7 @@ def make_nvfp4_moe_quant_config(

473 w2_scale_2: torch.Tensor,

474 a13_scale: torch.Tensor,

475 a2_scale: torch.Tensor,

476 + swiglu_limit: float | None = None,

477 ) -> FusedMoEQuantConfig:

478 if backend == NvFp4MoeBackend.MARLIN:

479 return nvfp4_w4a16_moe_quant_config(

@@ -455,6 +490,7 @@ def make_nvfp4_moe_quant_config(

490 a2_gscale=a2_scale,

491 w1_scale=w13_scale,

492 w2_scale=w2_scale,

493 + gemm1_clamp_limit=swiglu_limit,

494 )

495

496 # Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas.

@@ -478,6 +514,7 @@ def make_nvfp4_moe_quant_config(

514 NvFp4MoeBackend.FLASHINFER_CUTEDSL,

515 )

516 ),

517 + gemm1_clamp_limit=swiglu_limit,

518 )

519

520

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py

@@ -258,6 +258,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

258 w2_scale_2=layer.w2_weight_scale_2,

259 a13_scale=layer.w13_input_scale,

260 a2_scale=layer.w2_input_scale,

261 + swiglu_limit=getattr(layer, "swiglu_limit", None),

262 )

263

264 def apply_monolithic(

vllm/model_executor/layers/quantization/modelopt.py

@@ -1601,6 +1601,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1601 w2_scale_2=layer.w2_weight_scale_2,

1602 a13_scale=layer.w13_input_scale,

1603 a2_scale=layer.w2_input_scale,

1604 + swiglu_limit=getattr(layer, "swiglu_limit", None),

1605 )

1606

1607 @property

vllm/models/deepseek_v4/quant_config.py

+53 -1

@@ -2,6 +2,10 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Quantization config for DeepSeek V4."""

5 +from __future__ import annotations

6 +

7 +from typing import TYPE_CHECKING

8 +

9 from vllm.config import get_current_vllm_config

10 from vllm.model_executor.layers.fused_moe import FusedMoE

11 from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod

@@ -14,6 +18,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

19 _DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8")

21 +if TYPE_CHECKING:

22 + from vllm.model_executor.layers.quantization.modelopt import (

23 + ModelOptNvFp4Config,

24 + )

25 +

27 class DeepseekV4FP8Config(Fp8Config):

28 """FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch.

@@ -38,6 +47,8 @@ class DeepseekV4FP8Config(Fp8Config):

47 def __init__(self, *args, **kwargs):

48 super().__init__(*args, **kwargs)

49 self._resolved_expert_dtype: str | None = None

50 + self._resolved_moe_quant_algo: str | None = None

51 + self._nvfp4_config: ModelOptNvFp4Config | None = None

52 # ``is_scale_e8m0`` is a property that resolves on first read,

53 # by which time the current vllm_config has been set.

@@ -70,6 +81,36 @@ class DeepseekV4FP8Config(Fp8Config):

81 # checkpoints (Flash-Base) store them as float32.

82 return self.expert_dtype == "fp4"

84 + def _resolve_moe_overrides(self) -> None:

85 + if self._resolved_moe_quant_algo is not None:

86 + return

87 + try:

88 + hf_config = get_current_vllm_config().model_config.hf_config

89 + except Exception:

90 + return

91 + quant_cfg = getattr(hf_config, "quantization_config", None) or {}

92 + algo = (quant_cfg.get("moe_quant_algo") or "").upper() or None

93 + self._resolved_moe_quant_algo = algo or ""

94 +

95 + @property

96 + def moe_quant_algo(self) -> str:

97 + self._resolve_moe_overrides()

98 + return self._resolved_moe_quant_algo or ""

99 +

100 + def _get_nvfp4_config(self) -> ModelOptNvFp4Config:

101 + if self._nvfp4_config is None:

102 + from vllm.model_executor.layers.quantization.modelopt import (

103 + ModelOptNvFp4Config,

104 + )

105 +

106 + self._nvfp4_config = ModelOptNvFp4Config(

107 + is_checkpoint_nvfp4_serialized=True,

108 + kv_cache_quant_algo=None,

109 + exclude_modules=[],

110 + group_size=16,

111 + )

112 + return self._nvfp4_config

113 +

114 @classmethod

115 def get_name(cls) -> QuantizationMethods:

116 return "deepseek_v4_fp8"

@@ -97,10 +138,21 @@ class DeepseekV4FP8Config(Fp8Config):

138 ):

139 return UnquantizedFusedMoEMethod(layer.moe_config)

140 if self.expert_dtype == "fp4":

141 + if self.moe_quant_algo == "NVFP4":

142 + from vllm.model_executor.layers.quantization.modelopt import (

143 + ModelOptNvFp4FusedMoE,

144 + )

145 +

146 + return ModelOptNvFp4FusedMoE(

147 + quant_config=self._get_nvfp4_config(),

148 + moe_config=layer.moe_config,

149 + )

150 return Mxfp4MoEMethod(layer.moe_config)

151 # expert_dtype == "fp8": fall through to Fp8Config which

152 # returns Fp8MoEMethod with block-wise float32 scales.

153 return super().get_quant_method(layer, prefix)

154

155 def is_mxfp4_quant(self, prefix, layer):

106 - return isinstance(layer, FusedMoE) and self.expert_dtype == "fp4"

156 + if not isinstance(layer, FusedMoE) or self.expert_dtype != "fp4":

157 + return False

158 + return self.moe_quant_algo != "NVFP4"