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

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 12, 07:15 UTC

[Feat] Support non-gated activations in NVFP4 modelopt path (#29004)

Omer Ullman Argov committed Nov 30, 2025 at 18:02 UTC 39d28108f46ac320bccdf76e84640dbf2bd02bb4

5 files changed +98 -22

tests/kernels/moe/test_flashinfer_moe.py

+17 -7

@@ -16,11 +16,11 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

16 FlashInferExperts,

17 is_valid_flashinfer_cutlass_fused_moe,

18 )

19 +from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import (

20 + create_flashinfer_prepare_finalize,

21 +)

22 from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk

23 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel

21 -from vllm.model_executor.layers.fused_moe.prepare_finalize import (

22 - MoEPrepareAndFinalizeNoEP,

23 -)

24 from vllm.platforms import current_platform

25 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

@@ -48,9 +48,10 @@ MNK_FACTORS = [

48 @pytest.mark.parametrize("e", [40, 64, 256])

49 @pytest.mark.parametrize("topk", [1, 6, 8])

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

51 +@pytest.mark.parametrize("activation", ["silu_and_mul", "relu2"])

52 @torch.inference_mode()

53 def test_flashinfer_fp4_moe_no_graph(

53 - m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype

54 + m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype, activation: str

55 ):

56 current_platform.seed_everything(7)

57 with set_current_vllm_config(

@@ -59,6 +60,7 @@ def test_flashinfer_fp4_moe_no_graph(

60 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10

62 quant_blocksize = 16

63 + is_gated_act = activation == "silu_and_mul"

65 w1_q, w2_q, quant_config = make_test_quant_config(

66 e,

@@ -68,6 +70,7 @@ def test_flashinfer_fp4_moe_no_graph(

70 quant_dtype="nvfp4",

71 block_shape=None,

72 per_act_token_quant=False,

73 + make_gate=is_gated_act,

74 )

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

@@ -76,16 +79,19 @@ def test_flashinfer_fp4_moe_no_graph(

79 assert is_valid_flashinfer_cutlass_fused_moe(a, w1_q, w2_q)

81 flashinfer_experts = FusedMoEModularKernel(

79 - MoEPrepareAndFinalizeNoEP(),

82 + create_flashinfer_prepare_finalize(use_dp=False, use_nvfp4=True),

83 FlashInferExperts(out_dtype=dtype, quant_config=quant_config),

84 )

86 + fi_activation = {"silu_and_mul": "silu", "relu2": "relu2_no_mul"}[activation]

87 +

88 flashinfer_output = flashinfer_experts(

89 hidden_states=a,

90 w1=w1_q,

91 w2=w2_q,

92 topk_weights=topk_weights,

93 topk_ids=topk_ids,

94 + activation=fi_activation,

95 )

97 # Reference check:

@@ -103,7 +109,9 @@ def test_flashinfer_fp4_moe_no_graph(

109 block_size=quant_blocksize,

110 )

111

106 - w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype)

112 + w1_d = torch.empty(

113 + (e, (2 if is_gated_act else 1) * n, k), device="cuda", dtype=dtype

114 + )

115 w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)

116

117 for idx in range(0, e):

@@ -124,7 +132,9 @@ def test_flashinfer_fp4_moe_no_graph(

132 block_size=quant_blocksize,

133 )

134

127 - torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk)

135 + torch_output = torch_moe(

136 + a_in_dtype, w1_d, w2_d, score, topk, activation=activation

137 + )

138

139 torch.testing.assert_close(

140 torch_output, flashinfer_output, atol=1e-1, rtol=1e-1

tests/kernels/moe/utils.py

+10 -1

@@ -264,13 +264,20 @@ def make_test_weights(

264 quant_dtype: torch.dtype | str | None = None,

265 block_shape: list[int] | None = None,

266 per_out_ch_quant: bool = False,

267 + make_gate: bool = True,

268 ) -> tuple[

269 tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None],

270 tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None],

271 ]:

272 return (

273 make_test_weight(

273 - e, 2 * n, k, in_dtype, quant_dtype, block_shape, per_out_ch_quant

274 + e,

275 + (2 if make_gate else 1) * n,

276 + k,

277 + in_dtype,

278 + quant_dtype,

279 + block_shape,

280 + per_out_ch_quant,

281 ),

282 make_test_weight(e, k, n, in_dtype, quant_dtype, block_shape, per_out_ch_quant),

283 )

@@ -297,6 +304,7 @@ def make_test_quant_config(

304 quant_dtype: torch.dtype | str | None = None,

305 per_act_token_quant: bool = False,

306 block_shape: list[int] | None = None,

307 + make_gate: bool = True,

308 ) -> tuple[torch.Tensor, torch.Tensor, FusedMoEQuantConfig]:

309 (_, w1, w1_s, w1_gs), (_, w2, w2_s, w2_gs) = make_test_weights(

310 e,

@@ -306,6 +314,7 @@ def make_test_quant_config(

314 quant_dtype,

315 per_out_ch_quant=per_act_token_quant,

316 block_shape=block_shape,

317 + make_gate=make_gate,

318 )

319

320 # Hacky/trivial scales for nvfp4.

tests/kernels/utils.py

+7 -1

@@ -14,6 +14,7 @@ from torch._prims_common import TensorLikeType

15 from tests.kernels.quant_utils import native_w8a8_block_matmul

16 from vllm.attention.backends.abstract import AttentionType

17 +from vllm.model_executor.custom_op import CustomOp

18 from vllm.model_executor.layers.activation import SiluAndMul

19 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

20 from vllm.utils.torch_utils import make_tensor_with_pad

@@ -839,6 +840,7 @@ def torch_experts(

840 per_act_token_quant=False,

841 block_shape: list[int] | None = None,

842 apply_router_weights_on_input: bool = False,

843 + activation: str = "silu_and_mul",

844 ) -> torch.Tensor:

845 assert (

846 global_num_experts == -1

@@ -881,6 +883,8 @@ def torch_experts(

883

884 f32 = torch.float32

885

886 + act = CustomOp.op_registry[activation]

887 +

888 for i in range(num_experts):

889 mask = topk_ids == i

890 if mask.sum():

@@ -888,7 +892,7 @@ def torch_experts(

892 tmp1 = a[mask] @ w1[i].transpose(0, 1)

893 if b_bias1 is not None:

894 tmp1 = tmp1 + b_bias1[i].view(1, -1).to(tmp1.dtype)

891 - tmp2 = SiluAndMul()(tmp1)

895 + tmp2 = act()(tmp1)

896 out[mask] = tmp2 @ w2[i].transpose(0, 1)

897 if b_bias2 is not None:

898 out[mask] = out[mask] + b_bias2[i].view(1, -1).to(tmp1.dtype)

@@ -969,6 +973,7 @@ def torch_moe(

973 b_bias2: torch.Tensor | None = None,

974 global_num_experts: int = -1,

975 expert_map: torch.Tensor | None = None,

976 + activation: str = "silu_and_mul",

977 ) -> torch.Tensor:

978 score = torch.softmax(score, dim=-1, dtype=torch.float32)

979 topk_weight, topk_ids = torch.topk(score, topk)

@@ -982,6 +987,7 @@ def torch_moe(

987 b_bias1,

988 b_bias2,

989 expert_map,

990 + activation=activation,

991 )

992

993

vllm/model_executor/layers/fused_moe/layer.py

+9 -3

@@ -600,14 +600,20 @@ class FusedMoE(CustomOp):

600 # Avoid circular import

601 from vllm.model_executor.layers.quantization.modelopt import (

602 ModelOptFp8MoEMethod,

603 + ModelOptNvFp4FusedMoE,

604 )

605

606 if not isinstance(

606 - self.quant_method, (UnquantizedFusedMoEMethod, ModelOptFp8MoEMethod)

607 + self.quant_method,

608 + (

609 + UnquantizedFusedMoEMethod,

610 + ModelOptFp8MoEMethod,

611 + ModelOptNvFp4FusedMoE,

612 + ),

613 ):

614 raise NotImplementedError(

615 "is_act_and_mul=False is supported only for unquantized "

610 - "and ModelOpt FP8 moe for now"

616 + ", ModelOpt FP8, and ModelOpt NvFp4 checkpoints"

617 )

618 if not current_platform.is_cuda():

619 raise NotImplementedError(

@@ -1277,7 +1283,7 @@ class FusedMoE(CustomOp):

1283 self._load_combined_w13_weight_scale(

1284 shard_dim=shard_dim,

1285 loaded_weight=loaded_weight,

1280 - param=param,

1286 + param=expert_data,

1287 tp_rank=self.tp_rank,

1288 )

1289 return True if return_success else None

vllm/model_executor/layers/quantization/modelopt.py

+55 -10

@@ -1216,7 +1216,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1216 w13_weight = ModelWeightParameter(

1217 data=torch.empty(

1218 num_experts,

1219 - 2 * intermediate_size_per_partition,

1219 + (2 if self.moe.is_act_and_mul else 1) * intermediate_size_per_partition,

1220 # 2 fp4 items are packed in the input dimension

1221 hidden_size // 2,

1222 dtype=weight_dtype,

@@ -1245,7 +1245,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1245 w13_weight_scale = ModelWeightParameter(

1246 data=torch.empty(

1247 num_experts,

1248 - 2 * intermediate_size_per_partition,

1248 + (2 if self.moe.is_act_and_mul else 1) * intermediate_size_per_partition,

1249 # 2 fp4 items are packed in the input dimension

1250 hidden_size // self.quant_config.group_size,

1251 dtype=weight_scale_dtype,

@@ -1275,7 +1275,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1275 )

1276

1277 w13_weight_scale_2 = PerTensorScaleParameter(

1278 - data=torch.empty(num_experts, 2, dtype=torch.float32),

1278 + data=torch.empty(

1279 + num_experts, 2 if self.moe.is_act_and_mul else 1, dtype=torch.float32

1280 + ),

1281 weight_loader=weight_loader,

1282 )

1283 layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2)

@@ -1296,7 +1298,11 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1298 global_scale_num_experts = global_num_experts if use_global_sf else num_experts

1299

1300 w13_input_scale = PerTensorScaleParameter(

1299 - data=torch.empty(global_scale_num_experts, 2, dtype=torch.float32),

1301 + data=torch.empty(

1302 + global_scale_num_experts,

1303 + 2 if self.moe.is_act_and_mul else 1,

1304 + dtype=torch.float32,

1305 + ),

1306 weight_loader=weight_loader,

1307 )

1308 layer.register_parameter("w13_input_scale", w13_input_scale)

@@ -1312,9 +1318,13 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1318 gemm1_weight = layer.w13_weight.data

1319 gemm1_weight_scale = layer.w13_weight_scale.data

1320

1315 - if self.allow_flashinfer and (

1316 - self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1317 - or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1321 + if (

1322 + self.allow_flashinfer

1323 + and (

1324 + self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1325 + or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1326 + )

1327 + and self.moe.is_act_and_mul

1328 ):

1329 gemm1_weight, gemm1_weight_scale = reorder_w1w3_to_w3w1(

1330 gemm1_weight, gemm1_weight_scale, dim=-2

@@ -1324,7 +1334,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1334 layer.w13_weight_scale = Parameter(gemm1_weight_scale, requires_grad=False)

1335

1336 # Common processing for w13_weight_scale_2

1327 - if not torch.allclose(

1337 + if self.moe.is_act_and_mul and not torch.allclose(

1338 layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]

1339 ):

1340 logger.warning_once(

@@ -1437,11 +1447,39 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1447 w13_blockscale_swizzled, requires_grad=False

1448 )

1449

1450 + w13_weight = layer.w13_weight

1451 + intermediate_size_pad = w13_blockscale_swizzled.size(1) - w13_weight.size(1)

1452 + if intermediate_size_pad:

1453 + # padding gated activations will require to split w1 and w3

1454 + # and pad them individually

1455 + assert not self.moe.is_act_and_mul, (

1456 + "The intermediate size required padding, "

1457 + "but padding is not implemented for gated activations"

1458 + )

1459 +

1460 + layer.w13_weight = Parameter(

1461 + torch.nn.functional.pad(

1462 + w13_weight, (0, 0, 0, intermediate_size_pad)

1463 + ),

1464 + requires_grad=False,

1465 + )

1466 + layer.w2_weight = Parameter(

1467 + torch.nn.functional.pad(

1468 + layer.w2_weight, (0, intermediate_size_pad // 2, 0, 0)

1469 + ),

1470 + requires_grad=False,

1471 + )

1472 + layer.w2_weight_scale = Parameter(

1473 + torch.nn.functional.pad(

1474 + layer.w2_weight_scale, (0, intermediate_size_pad // 16)

1475 + ),

1476 + requires_grad=False,

1477 + )

1478 +

1479 w2_blockscale_swizzled = swizzle_blockscale(layer.w2_weight_scale)

1480 layer.w2_weight_scale = Parameter(

1481 w2_blockscale_swizzled, requires_grad=False

1482 )

1444 - layer.w2_weight = Parameter(layer.w2_weight.data, requires_grad=False)

1483

1484 def get_fused_moe_quant_config(

1485 self, layer: torch.nn.Module

@@ -1484,7 +1522,14 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1522 logical_to_physical_map: torch.Tensor | None = None,

1523 logical_replica_count: torch.Tensor | None = None,

1524 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1487 - assert activation == "silu", "Only SiLU activation is supported."

1525 + if not self.moe.is_act_and_mul:

1526 + assert (

1527 + self.allow_flashinfer

1528 + and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1529 + ), (

1530 + "Non-gated activations are only supported by the"

1531 + " flashinfer CUTLASS backend for modelopt checkpoints"

1532 + )

1533

1534 if (

1535 self.allow_flashinfer