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

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 10, 02:45 UTC

[Feat] Support non-gated MoE with Marlin, NVFP4 CUTLASS, FP8, INT8, compressed-tensors (#32257)

Signed-off-by: Tomer Natan <tbarnatan@computelab-frontend-8.nvidia.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Tomer Natan <tbarnatan@computelab-frontend-8.nvidia.com> Co-authored-by: mgoin <mgoin64@gmail.com> Co-authored-by: Tomer Natan <tbarnatan@ipp1-1429.ipp1a1.colossus.nvidia.com>

TomerBN-Nvidia committed Jan 16, 2026 at 02:15 UTC c277fbdf31f49c15049fe361ef92a7e34abc7b74

17 files changed +226 -127

tests/kernels/moe/test_cutlass_moe.py

+1 -1

@@ -526,7 +526,7 @@ def test_run_cutlass_moe_fp8(

526 c_strides1 = torch.full((e,), 2 * n, device="cuda", dtype=torch.int64)

527 c_strides2 = torch.full((e,), k, device="cuda", dtype=torch.int64)

528

529 - activation = lambda o, i: torch.ops._C.silu_and_mul(o, i)

529 + activation = "silu"

530 a1q, a1q_scale = moe_kernel_quantize_input(

531 mt.a, mt.a_scale, torch.float8_e4m3fn, per_act_token

532 )

tests/kernels/moe/test_moe.py

+80

@@ -1079,6 +1079,86 @@ def test_fused_marlin_moe_with_bias(m):

1079 torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0)

1080

1081

1082 +@pytest.mark.flaky(reruns=2)

1083 +@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")

1084 +@pytest.mark.parametrize("m", [1, 64, 256])

1085 +@pytest.mark.parametrize("n,k", [(1024, 1024), (2048, 2048)])

1086 +@pytest.mark.parametrize("e,topk", [(8, 2), (64, 4)])

1087 +def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):

1088 + """Test Marlin MoE with non-gated activation (relu2_no_mul).

1089 +

1090 + Non-gated activations like relu2 don't have the gate-up projection pattern,

1091 + so w1 has shape (e, n, k) instead of (e, 2*n, k).

1092 + """

1093 + torch.cuda.manual_seed(42)

1094 +

1095 + group_size = 16 # NVFP4 group size

1096 + is_k_full = True

1097 + quant_type = scalar_types.float4_e2m1f

1098 + dtype = torch.bfloat16

1099 +

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

1101 + # Non-gated: w1 shape is (e, n, k) not (e, 2*n, k)

1102 + w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 10

1103 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10

1104 +

1105 + w1_data = MarlinMoEWeightData.make(

1106 + w=w1,

1107 + quant_type=quant_type,

1108 + group_size=group_size,

1109 + act_order=False,

1110 + )

1111 +

1112 + w2_data = MarlinMoEWeightData.make(

1113 + w=w2,

1114 + quant_type=quant_type,

1115 + group_size=group_size,

1116 + act_order=False,

1117 + )

1118 +

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

1120 +

1121 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)

1122 +

1123 + with set_current_vllm_config(vllm_config):

1124 + torch_output = torch_moe(

1125 + a,

1126 + w1_data.w_ref,

1127 + w2_data.w_ref,

1128 + score,

1129 + topk,

1130 + activation="relu2",

1131 + )

1132 +

1133 + marlin_output = fused_marlin_moe(

1134 + a,

1135 + w1_data.qweight,

1136 + w2_data.qweight,

1137 + None, # bias1

1138 + None, # bias2

1139 + w1_data.scales,

1140 + w2_data.scales,

1141 + score,

1142 + topk_weights,

1143 + topk_ids,

1144 + global_num_experts=e,

1145 + expert_map=None,

1146 + global_scale1=w1_data.global_scale,

1147 + global_scale2=w2_data.global_scale,

1148 + g_idx1=w1_data.g_idx,

1149 + g_idx2=w2_data.g_idx,

1150 + sort_indices1=w1_data.sort_indices,

1151 + sort_indices2=w2_data.sort_indices,

1152 + w1_zeros=w1_data.zeros,

1153 + w2_zeros=w2_data.zeros,

1154 + quant_type_id=quant_type.id,

1155 + is_k_full=is_k_full,

1156 + activation="relu2_no_mul",

1157 + )

1158 +

1159 + torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0)

1160 +

1161 +

1162 @pytest.mark.parametrize("ep_size", [1, 2])

1163 def test_moe_align_block_size_opcheck(ep_size):

1164 num_experts = 4

vllm/envs.py

@@ -1451,6 +1451,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1451 # - "flashinfer-cudnn": use flashinfer cudnn GEMM backend

1452 # - "flashinfer-trtllm": use flashinfer trtllm GEMM backend

1453 # - "flashinfer-cutlass": use flashinfer cutlass GEMM backend

1454 + # - "marlin": use marlin GEMM backend (for GPUs without native FP4 support)

1455 # - <none>: automatically pick an available backend

1456 "VLLM_NVFP4_GEMM_BACKEND": env_with_choices(

1457 "VLLM_NVFP4_GEMM_BACKEND",

@@ -1460,6 +1461,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1461 "flashinfer-trtllm",

1462 "flashinfer-cutlass",

1463 "cutlass",

1464 + "marlin",

1465 ],

1466 ),

1467 # Controls garbage collection during CUDA graph capture.

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+26 -18

@@ -2,8 +2,6 @@

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

3 """CUTLASS based Fused MoE kernels."""

5 -from collections.abc import Callable

6 -

5 import torch

7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

@@ -21,7 +19,10 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

19 TopKWeightAndReduceDelegate,

20 TopKWeightAndReduceNoOP,

21 )

24 -from vllm.model_executor.layers.fused_moe.utils import _resize_cache

22 +from vllm.model_executor.layers.fused_moe.utils import (

23 + _resize_cache,

24 + apply_moe_activation,

25 +)

26 from vllm.scalar_type import scalar_types

28 logger = init_logger(__name__)

@@ -33,7 +34,7 @@ def run_cutlass_moe_fp8(

34 w1: torch.Tensor,

35 w2: torch.Tensor,

36 topk_ids: torch.Tensor,

36 - activation_callable: Callable,

37 + activation: str,

38 global_num_experts: int,

39 expert_map: torch.Tensor | None,

40 w1_scale: torch.Tensor | None,

@@ -55,6 +56,7 @@ def run_cutlass_moe_fp8(

56 ):

57 a1q = hidden_states

59 + assert not activation.endswith("_no_mul"), "Only gated activation is supported"

60 assert w1_scale is not None

61 assert w2_scale is not None

62 assert w1.dtype == torch.float8_e4m3fn

@@ -198,7 +200,7 @@ def run_cutlass_moe_fp8(

200 per_out_ch,

201 )

202

201 - activation_callable(act_out, mm1_out)

203 + apply_moe_activation(activation, act_out, mm1_out)

204

205 a2q, a2q_scale = ops.scaled_fp8_quant(

206 act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out

@@ -288,8 +290,6 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

290 if expert_tokens_meta is not None:

291 expert_num_tokens = expert_tokens_meta.expert_num_tokens

292

291 - activation_callable = lambda o, i: self.activation(activation, o, i)

292 -

293 use_batched_format = (

294 self.activation_formats[0] == mk.FusedMoEActivationFormat.BatchedExperts

295 )

@@ -301,7 +301,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

301 w1,

302 w2,

303 topk_ids,

304 - activation_callable,

304 + activation,

305 global_num_experts,

306 expert_map,

307 self.w1_scale,

@@ -436,6 +436,7 @@ def run_cutlass_moe_fp4(

436 w2_alphas: torch.Tensor,

437 topk_weights: torch.Tensor,

438 topk_ids: torch.Tensor,

439 + activation: str,

440 workspace13: torch.Tensor,

441 workspace2: torch.Tensor,

442 m: int,

@@ -544,8 +545,7 @@ def run_cutlass_moe_fp4(

545 num_topk,

546 )

547 c1 = _resize_cache(workspace13, (m * topk, n * 2))

547 - # Note: c2 workspace is no longer needed since SiLU is fused with quantization.

548 - # c3 reuses workspace13 after c1 is consumed.

548 + c2 = _resize_cache(workspace2, (m * topk, n))

549 c3 = _resize_cache(workspace13, (m * topk, k))

550 ops.cutlass_fp4_moe_mm(

551 c1,

@@ -559,10 +559,18 @@ def run_cutlass_moe_fp4(

559 blockscale_offsets[:-1],

560 )

561 del rep_a_fp4, rep_a_blockscale

562 - # Fused SiLU+Mul+NVFP4 quantization

563 - int_fp4, int_blockscale = ops.silu_and_mul_scaled_fp4_experts_quant(

564 - c1, a2_gscale, expert_offsets, blockscale_offsets, num_topk

565 - )

562 + if activation == "silu":

563 + # Fused SiLU+Mul+NVFP4 quantization

564 + # Note: c2 workspace is no longer needed since SiLU is fused with quantization.

565 + # c3 reuses workspace13 after c1 is consumed.

566 + int_fp4, int_blockscale = ops.silu_and_mul_scaled_fp4_experts_quant(

567 + c1, a2_gscale, expert_offsets, blockscale_offsets, num_topk

568 + )

569 + else:

570 + apply_moe_activation(activation, c2, c1)

571 + int_fp4, int_blockscale = ops.scaled_fp4_experts_quant(

572 + c2, a2_gscale, expert_offsets, blockscale_offsets, num_topk

573 + )

574

575 ops.cutlass_fp4_moe_mm(

576 c3,

@@ -693,6 +701,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

701 w2_alphas=self.g2_alphas,

702 topk_weights=topk_weights,

703 topk_ids=topk_ids,

704 + activation=activation,

705 workspace13=workspace13,

706 workspace2=workspace2,

707 m=m,

@@ -711,7 +720,7 @@ def run_cutlass_moe_w4a8_fp8(

720 w1: torch.Tensor,

721 w2: torch.Tensor,

722 topk_ids: torch.Tensor,

714 - activation_callable: Callable,

723 + activation: str,

724 global_num_experts: int,

725 expert_map: torch.Tensor | None,

726 w1_scale: torch.Tensor | None,

@@ -815,7 +824,7 @@ def run_cutlass_moe_w4a8_fp8(

824 s_strides1,

825 )

826

818 - activation_callable(act_out, mm1_out)

827 + apply_moe_activation(activation, act_out, mm1_out)

828

829 a2q, a2q_scale = ops.scaled_fp8_quant(

830 act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out

@@ -936,7 +945,6 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

945 assert self.w2_zp is None, "w2_zp is not supported in CUTLASS MoE"

946

947 expert_num_tokens = None

939 - activation_callable = lambda o, i: self.activation(activation, o, i)

948

949 use_batched_format = (

950 self.activation_formats[0] == mk.FusedMoEActivationFormat.BatchedExperts

@@ -951,7 +959,7 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

959 w1,

960 w2,

961 topk_ids,

954 - activation_callable,

962 + activation,

963 global_num_experts,

964 expert_map,

965 self.w1_scale,

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+16 -24

@@ -17,7 +17,11 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

17 TopKWeightAndReduceDelegate,

18 TopKWeightAndReduceNoOP,

19 )

20 -from vllm.model_executor.layers.fused_moe.utils import _resize_cache, disable_inplace

20 +from vllm.model_executor.layers.fused_moe.utils import (

21 + _resize_cache,

22 + apply_moe_activation,

23 + disable_inplace,

24 +)

25 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

26 marlin_make_workspace_new,

27 marlin_moe_intermediate_size,

@@ -27,21 +31,6 @@ from vllm.platforms import current_platform

31 from vllm.scalar_type import ScalarType, scalar_types

30 -def default_activation_func(

31 - activation: str, output: torch.Tensor, input: torch.Tensor

32 -) -> None:

33 - if activation == "silu":

34 - torch.ops._C.silu_and_mul(output, input)

35 - elif activation == "swigluoai":

36 - # alpha = 1.702, limit = 7.0

37 - torch.ops._C.swigluoai_and_mul(output, input)

38 - else:

39 - raise ValueError(

40 - f"Unsupported activation: {activation}. "

41 - "Only silu and swigluoai activations are supported."

42 - )

43 -

44 -

34 def _fused_marlin_moe(

35 hidden_states: torch.Tensor,

36 w1: torch.Tensor,

@@ -62,7 +51,7 @@ def _fused_marlin_moe(

51 activation: str = "silu",

52 activation_func: Callable[

53 [str, torch.Tensor, torch.Tensor], None

65 - ] = default_activation_func,

54 + ] = apply_moe_activation,

55 input_global_scale1: torch.Tensor | None = None,

56 input_global_scale2: torch.Tensor | None = None,

57 global_scale1: torch.Tensor | None = None,

@@ -83,13 +72,13 @@ def _fused_marlin_moe(

72 assert hidden_states.ndim == 2

73 M, K = hidden_states.size()

74 N = marlin_moe_intermediate_size(w1, w2)

86 -

75 + w13_num_shards = 1 if "no_mul" in activation else 2

76 if workspace is None:

77 workspace = marlin_make_workspace_new(hidden_states.device, 4)

79 if intermediate_cache13 is None:

80 intermediate_cache13 = torch.empty(

92 - (M * num_topk * max(2 * N, K),),

81 + (M * num_topk * max(w13_num_shards * N, K),),

82 device=hidden_states.device,

83 dtype=hidden_states.dtype,

84 )

@@ -101,7 +90,9 @@ def _fused_marlin_moe(

90 dtype=hidden_states.dtype,

91 )

104 - intermediate_cache1 = _resize_cache(intermediate_cache13, (M * num_topk, 2 * N))

93 + intermediate_cache1 = _resize_cache(

94 + intermediate_cache13, (M * num_topk, w13_num_shards * N)

95 + )

97 intermediate_cache3 = _resize_cache(intermediate_cache13, (M * num_topk, K))

@@ -137,16 +128,17 @@ def _fused_marlin_moe(

128 mul_topk_weights=apply_router_weight_on_input,

129 b_q_type=quant_type,

130 size_m=M,

140 - size_n=2 * N,

131 + size_n=w13_num_shards * N,

132 size_k=K,

133 is_k_full=is_k_full,

134 use_atomic_add=False,

135 use_fp32_reduce=True,

136 is_zp_float=False,

137 )

147 -

138 activation_func(

149 - activation, intermediate_cache2, intermediate_cache1.view(-1, 2 * N)

139 + activation,

140 + intermediate_cache2,

141 + intermediate_cache1.view(-1, w13_num_shards * N),

142 )

143

144 if output is None:

@@ -216,7 +208,7 @@ def fused_marlin_moe(

208 activation: str = "silu",

209 activation_func: Callable[

210 [str, torch.Tensor, torch.Tensor], None

219 - ] = default_activation_func,

211 + ] = apply_moe_activation,

212 moe_sum: Callable[[torch.Tensor, torch.Tensor], None] | None = None,

213 expert_map: torch.Tensor | None = None,

214 input_global_scale1: torch.Tensor | None = None,

vllm/model_executor/layers/fused_moe/layer.py

+3 -22

@@ -619,30 +619,11 @@ class FusedMoE(CustomOp):

619 # for heuristic purposes, so it must be initialized first.

620 self.quant_method: FusedMoEMethodBase = _get_quant_method()

621

622 - if not self.moe_config.is_act_and_mul:

623 - # Avoid circular import

624 - from vllm.model_executor.layers.quantization.modelopt import (

625 - ModelOptFp8MoEMethod,

626 - ModelOptNvFp4FusedMoE,

622 + if not self.moe_config.is_act_and_mul and not current_platform.is_cuda():

623 + raise NotImplementedError(

624 + "is_act_and_mul=False is supported only for CUDA for now"

625 )

626

629 - if not isinstance(

630 - self.quant_method,

631 - (

632 - UnquantizedFusedMoEMethod,

633 - ModelOptFp8MoEMethod,

634 - ModelOptNvFp4FusedMoE,

635 - ),

636 - ):

637 - raise NotImplementedError(

638 - "is_act_and_mul=False is supported only for unquantized "

639 - ", ModelOpt FP8, and ModelOpt NvFp4 checkpoints"

640 - )

641 - if not current_platform.is_cuda():

642 - raise NotImplementedError(

643 - "is_act_and_mul=False is supported only for CUDA for now"

644 - )

645 -

627 if self.enable_eplb and not self.quant_method.supports_eplb:

628 # TODO: Add support for additional quantization methods.

629 # The implementation for other quantization methods does not

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

+8 -3

@@ -52,7 +52,7 @@ def select_fp8_moe_backend(

52 block_quant: bool,

53 tp_size: int,

54 with_lora_support: bool,

55 - is_act_and_mul: bool = True,

55 + is_act_and_mul: bool,

56 allow_vllm_cutlass: bool = False,

57 ) -> Fp8MoeBackend:

58 """

@@ -128,7 +128,7 @@ def select_fp8_moe_backend(

128 scope="local",

129 )

130

131 - if use_deep_gemm and moe_use_deep_gemm and block_quant:

131 + if use_deep_gemm and moe_use_deep_gemm and block_quant and is_act_and_mul:

132 if not has_deep_gemm():

133 logger.warning_once(

134 "DeepGEMM backend requested but not available.", scope="local"

@@ -141,7 +141,12 @@ def select_fp8_moe_backend(

141 logger.info_once(_make_log_backend("ROCm AITER"), scope="local")

142 return Fp8MoeBackend.AITER

143

144 - if allow_vllm_cutlass and not block_quant and cutlass_group_gemm_supported():

144 + if (

145 + allow_vllm_cutlass

146 + and not block_quant

147 + and cutlass_group_gemm_supported()

148 + and is_act_and_mul

149 + ):

150 logger.info_once(_make_log_backend("vLLM CUTLASS"), scope="local")

151 return Fp8MoeBackend.VLLM_CUTLASS

152

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

@@ -178,6 +178,7 @@ def convert_to_nvfp4_moe_kernel_format(

178 w2=w2,

179 w2_scale=w2_scale,

180 w2_scale_2=w2_scale_2,

181 + is_act_and_mul=is_act_and_mul,

182 )

183 else:

184 raise ValueError(f"Unknown NvFp4 backend for MoE: {nvfp4_backend}")

vllm/model_executor/layers/fused_moe/utils.py

+2 -1

@@ -367,7 +367,8 @@ def apply_moe_activation(

367 elif activation == GELU_NO_MUL:

368 output.copy_(F.gelu(input))

369 elif activation == RELU2_NO_MUL:

370 - torch.square(F.relu(input), out=output)

370 + F.relu(input, inplace=True)

371 + torch.square(input, out=output)

372 else:

373 raise ValueError(f"Unsupported FusedMoe activation: {activation}")

374

vllm/model_executor/layers/quantization/awq_marlin.py

-2

@@ -764,8 +764,6 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

764 x: torch.Tensor,

765 router_logits: torch.Tensor,

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

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

768 -

767 topk_weights, topk_ids = router.select_experts(

768 hidden_states=x,

769 router_logits=router_logits,

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

+45 -27

@@ -370,12 +370,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

370 layer_name: str | None = None,

371 use_marlin: bool = False,

372 ):

373 - if not moe.is_act_and_mul:

374 - raise ValueError(

375 - "CompressedTensorsW4A4Nvfp4MoEMethod does not yet "

376 - "support non gated MoE models."

377 - )

378 -

373 super().__init__(moe)

374 self.group_size = 16

375 if use_marlin:

@@ -388,6 +382,16 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

382 )

383 else:

384 self.nvfp4_backend = select_nvfp4_moe_backend()

385 +

386 + # TODO: move this type of check into the oracle.

387 + if not self.moe.is_act_and_mul and self.nvfp4_backend not in [

388 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

389 + NvFp4MoeBackend.MARLIN,

390 + ]:

391 + raise NotImplementedError(

392 + "Non-gated activations are only supported by FlashInfer "

393 + f"CUTLASS and Marlin NvFP4 MoE backends, not {self.nvfp4_backend}."

394 + )

395 self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(

396 self.nvfp4_backend

397 )

@@ -404,11 +408,12 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

408 ):

409 layer.num_experts = num_experts

410 layer.params_dtype = params_dtype

411 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1

412

413 w13_weight = torch.nn.Parameter(

414 torch.empty(

415 num_experts,

411 - 2 * intermediate_size_per_partition,

416 + w13_num_shards * intermediate_size_per_partition,

417 # 2 fp4 items are packed in the input dimension

418 hidden_size // 2,

419 requires_grad=False,

@@ -436,7 +441,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

441 w13_weight_scale = torch.nn.Parameter(

442 torch.empty(

443 num_experts,

439 - 2 * intermediate_size_per_partition,

444 + w13_num_shards * intermediate_size_per_partition,

445 # 2 fp4 items are packed in the input dimension

446 hidden_size // self.group_size,

447 dtype=torch.float8_e4m3fn,

@@ -467,7 +472,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

472

473 # Weight Global Scales

474 w13_weight_scale_2 = torch.nn.Parameter(

470 - torch.empty(num_experts, 2, dtype=torch.float32), requires_grad=False

475 + torch.empty(num_experts, w13_num_shards, dtype=torch.float32),

476 + requires_grad=False,

477 )

478 layer.register_parameter("w13_weight_global_scale", w13_weight_scale_2)

479 extra_weight_attrs.update(

@@ -486,7 +492,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

492

493 # Input Global Scales

494 w13_input_scale = torch.nn.Parameter(

489 - torch.empty(num_experts, 2, dtype=torch.float32), requires_grad=False

495 + torch.empty(num_experts, w13_num_shards, dtype=torch.float32),

496 + requires_grad=False,

497 )

498 layer.register_parameter("w13_input_global_scale", w13_input_scale)

499 extra_weight_attrs.update(

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

647 x=x,

648 router_logits=router_logits,

649 top_k=layer.top_k,

650 + activation=layer.activation,

651 global_num_experts=layer.global_num_experts,

652 num_expert_group=layer.num_expert_group,

653 topk_group=layer.topk_group,

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

674 topk_ids=topk_ids,

675 topk_weights=topk_weights,

676 top_k=layer.top_k,

677 + activation=layer.activation,

678 global_num_experts=layer.global_num_experts,

679 )

680 else:

@@ -722,6 +731,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

731 block_quant=self.block_quant,

732 tp_size=moe.tp_size,

733 with_lora_support=moe.is_lora_enabled,

734 + is_act_and_mul=moe.is_act_and_mul,

735 # TODO(rob): enable selecting this externally.

736 allow_vllm_cutlass=True,

737 )

@@ -760,6 +770,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

770 layer.weight_block_size = None

771

772 params_dtype = torch.float8_e4m3fn

773 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1

774

775 if self.block_quant:

776 assert self.weight_block_size is not None

@@ -791,7 +802,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

802 w13_weight = torch.nn.Parameter(

803 torch.empty(

804 num_experts,

794 - 2 * intermediate_size_per_partition,

805 + w13_num_shards * intermediate_size_per_partition,

806 hidden_size,

807 dtype=params_dtype,

808 ),

@@ -814,10 +825,12 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

825

826 # WEIGHT_SCALES

827 if self.weight_quant.strategy == QuantizationStrategy.TENSOR:

817 - # Allocate 2 scales for w1 and w3 respectively.

818 - # They are combined to a single scale after weight loading.

828 + # For gated MoE, allocate 2 scales for w1 and w3 respectively.

829 + # They will be combined to a single scale after weight loading.

830 + # For non-gated MoE, allocate 1 scale for w13.

831 w13_weight_scale = torch.nn.Parameter(

820 - torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False

832 + torch.ones(num_experts, w13_num_shards, dtype=torch.float32),

833 + requires_grad=False,

834 )

835 layer.register_parameter("w13_weight_scale", w13_weight_scale)

836 w2_weight_scale = torch.nn.Parameter(

@@ -835,7 +848,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

848 w13_weight_scale = torch.nn.Parameter(

849 torch.ones(

850 num_experts,

838 - 2 * intermediate_size_per_partition,

851 + w13_num_shards * intermediate_size_per_partition,

852 1,

853 dtype=torch.float32,

854 ),

@@ -858,7 +871,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

871 w13_weight_scale = torch.nn.Parameter(

872 torch.ones(

873 num_experts,

861 - 2 * ((intermediate_size_per_partition + block_n - 1) // block_n),

874 + w13_num_shards

875 + * ((intermediate_size_per_partition + block_n - 1) // block_n),

876 (hidden_size + block_k - 1) // block_k,

877 dtype=torch.float32,

878 ),

@@ -930,11 +944,12 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

944 # Per-tensor kernels use a single scale, for W13, but on disk there

945 # is a separate scale for W1 and W3. Requantize with the max scale.

946 if self.weight_quant.strategy == QuantizationStrategy.TENSOR:

933 - process_fp8_weight_tensor_strategy_moe(

947 + w13, w13_scale = process_fp8_weight_tensor_strategy_moe(

948 w13,

949 w13_scale,

950 shard_size=layer.intermediate_size_per_partition,

951 num_experts=layer.num_local_experts,

952 + is_act_and_mul=self.moe.is_act_and_mul,

953 )

954

955 w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(

@@ -1166,12 +1181,13 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1181 **extra_weight_attrs,

1182 ):

1183 params_dtype = torch.int8

1184 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1

1185

1186 # WEIGHTS

1187 w13_weight = torch.nn.Parameter(

1188 torch.empty(

1189 num_experts,

1174 - 2 * intermediate_size_per_partition,

1190 + w13_num_shards * intermediate_size_per_partition,

1191 hidden_size,

1192 dtype=params_dtype,

1193 ),

@@ -1196,7 +1212,10 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1212 assert self.weight_quant.strategy == QuantizationStrategy.CHANNEL

1213 w13_weight_scale = torch.nn.Parameter(

1214 torch.ones(

1199 - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32

1215 + num_experts,

1216 + w13_num_shards * intermediate_size_per_partition,

1217 + 1,

1218 + dtype=torch.float32,

1219 ),

1220 requires_grad=False,

1221 )

@@ -1296,6 +1315,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1315 **extra_weight_attrs,

1316 ):

1317 intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")

1318 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1

1319

1320 # Will transpose the loaded weight along the

1321 # intermediate and hidden dim sizes. Will

@@ -1307,7 +1327,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1327 torch.empty(

1328 num_experts,

1329 hidden_size // self.packed_factor,

1310 - 2 * intermediate_size_per_partition,

1330 + w13_num_shards * intermediate_size_per_partition,

1331 dtype=torch.int32,

1332 ),

1333 requires_grad=False,

@@ -1352,7 +1372,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1372 torch.ones(

1373 num_experts,

1374 num_groups_w13,

1355 - 2 * intermediate_size_per_partition,

1375 + w13_num_shards * intermediate_size_per_partition,

1376 dtype=params_dtype,

1377 ),

1378 requires_grad=False,

@@ -1600,10 +1620,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1620 x: torch.Tensor,

1621 router_logits: torch.Tensor,

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

1603 - assert layer.activation == "silu", (

1604 - f"{layer.activation} not supported for Marlin MoE."

1605 - )

1606 -

1623 topk_weights, topk_ids = router.select_experts(

1624 hidden_states=x,

1625 router_logits=router_logits,

@@ -1625,6 +1641,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1641 quant_type_id=self.quant_type.id,

1642 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1643 global_num_experts=layer.global_num_experts,

1644 + activation=layer.activation,

1645 expert_map=layer.expert_map,

1646 g_idx1=layer.w13_weight_g_idx,

1647 g_idx2=layer.w2_weight_g_idx,

@@ -1675,11 +1692,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1692 extra_weight_attrs.update(

1693 {"is_transposed": True, "quant_method": self.strategy}

1694 )

1695 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1

1696 w13_weight = torch.nn.Parameter(

1697 torch.empty(

1698 num_experts,

1699 hidden_size // self.packed_factor,

1682 - 2 * intermediate_size_per_partition,

1700 + w13_num_shards * intermediate_size_per_partition,

1701 dtype=torch.int32,

1702 ),

1703 requires_grad=False,

@@ -1712,7 +1730,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1730 torch.ones(

1731 num_experts,

1732 num_groups_w13,

1715 - 2 * intermediate_size_per_partition,

1733 + w13_num_shards * intermediate_size_per_partition,

1734 dtype=params_dtype,

1735 ),

1736 requires_grad=False,

vllm/model_executor/layers/quantization/fp8.py

@@ -637,6 +637,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

637 block_quant=self.block_quant,

638 tp_size=layer.moe_parallel_config.tp_size,

639 with_lora_support=self.moe.is_lora_enabled,

640 + is_act_and_mul=self.moe.is_act_and_mul,

641 )

642

643 if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

vllm/model_executor/layers/quantization/gptq_marlin.py

-2

@@ -900,8 +900,6 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

900 x: torch.Tensor,

901 router_logits: torch.Tensor,

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

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

904 -

903 topk_weights, topk_ids = router.select_experts(

904 hidden_states=x,

905 router_logits=router_logits,

vllm/model_executor/layers/quantization/modelopt.py

+19 -17

@@ -733,6 +733,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

733 block_quant=False,

734 tp_size=moe_config.moe_parallel_config.tp_size,

735 with_lora_support=self.moe.is_lora_enabled,

736 + is_act_and_mul=self.moe.is_act_and_mul,

737 )

738 self.kernel: mk.FusedMoEModularKernel | None = None

739

@@ -789,15 +790,12 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

790 )

791 weight_loader = extra_weight_attrs.get("weight_loader")

792

792 - if self.moe.is_act_and_mul:

793 - w13_up_dim = 2 * intermediate_size_per_partition

794 - else:

795 - w13_up_dim = intermediate_size_per_partition

793 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1

794

795 w13_weight = ModelWeightParameter(

796 data=torch.empty(

797 num_experts,

800 - w13_up_dim,

798 + w13_num_shards * intermediate_size_per_partition,

799 hidden_size,

800 dtype=weight_dtype,

801 ),

@@ -826,7 +824,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

824 # For non-gated MoE, allocate 1 scale for w13.

825 w13_weight_scale = PerTensorScaleParameter(

826 data=torch.full(

829 - (num_experts, 2 if self.moe.is_act_and_mul else 1),

827 + (num_experts, w13_num_shards),

828 1.0,

829 dtype=torch.float32,

830 ),

@@ -1132,6 +1130,9 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

1130 elif envs.VLLM_NVFP4_GEMM_BACKEND == "cutlass":

1131 self.backend = "cutlass"

1132 assert cutlass_fp4_supported(), f"Cutlass is required for {self.backend}"

1133 + elif envs.VLLM_NVFP4_GEMM_BACKEND == "marlin":

1134 + self.backend = "marlin"

1135 + assert is_fp4_marlin_supported(), f"Marlin is required for {self.backend}"

1136

1137 if self.backend == "none":

1138 raise ValueError(

@@ -1337,13 +1338,13 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1338 self.quant_config = quant_config

1339 self.nvfp4_backend = select_nvfp4_moe_backend()

1340 # TODO: move this type of check into the oracle.

1340 - if (

1341 - not self.moe.is_act_and_mul

1342 - and not self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTLASS

1343 - ):

1341 + if not self.moe.is_act_and_mul and self.nvfp4_backend not in [

1342 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

1343 + NvFp4MoeBackend.MARLIN,

1344 + ]:

1345 raise NotImplementedError(

1346 "Non-gated activations are only supported by FlashInfer "

1346 - "CUTLASS NvFP4 MoE backend."

1347 + f"CUTLASS and Marlin NvFP4 MoE backends, not {self.nvfp4_backend}."

1348 )

1349

1350 self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(

@@ -1409,11 +1410,12 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1410 weight_scale_dtype = torch.float8_e4m3fn

1411 weight_loader = extra_weight_attrs.get("weight_loader")

1412 global_num_experts = extra_weight_attrs.get("global_num_experts")

1413 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1

1414 # GEMM 1

1415 w13_weight = ModelWeightParameter(

1416 data=torch.empty(

1417 num_experts,

1416 - (2 if self.moe.is_act_and_mul else 1) * intermediate_size_per_partition,

1418 + w13_num_shards * intermediate_size_per_partition,

1419 # 2 fp4 items are packed in the input dimension

1420 hidden_size // 2,

1421 dtype=weight_dtype,

@@ -1442,7 +1444,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1444 w13_weight_scale = ModelWeightParameter(

1445 data=torch.empty(

1446 num_experts,

1445 - (2 if self.moe.is_act_and_mul else 1) * intermediate_size_per_partition,

1447 + w13_num_shards * intermediate_size_per_partition,

1448 # 2 fp4 items are packed in the input dimension

1449 hidden_size // self.quant_config.group_size,

1450 dtype=weight_scale_dtype,

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

1474 )

1475

1476 w13_weight_scale_2 = PerTensorScaleParameter(

1475 - data=torch.empty(

1476 - num_experts, 2 if self.moe.is_act_and_mul else 1, dtype=torch.float32

1477 - ),

1477 + data=torch.empty(num_experts, w13_num_shards, dtype=torch.float32),

1478 weight_loader=weight_loader,

1479 )

1480 layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2)

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

1495 w13_input_scale = PerTensorScaleParameter(

1496 data=torch.empty(

1497 global_sf_num_experts,

1498 - 2 if self.moe.is_act_and_mul else 1,

1498 + w13_num_shards,

1499 dtype=torch.float32,

1500 ),

1501 weight_loader=weight_loader,

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

1616 x=x,

1617 router_logits=router_logits,

1618 top_k=layer.top_k,

1619 + activation=layer.activation,

1620 global_num_experts=layer.global_num_experts,

1621 num_expert_group=layer.num_expert_group,

1622 topk_group=layer.topk_group,

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

1643 topk_ids=topk_ids,

1644 topk_weights=topk_weights,

1645 top_k=layer.top_k,

1646 + activation=layer.activation,

1647 global_num_experts=layer.global_num_experts,

1648 )

1649 else:

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

+16

@@ -255,6 +255,7 @@ def flashinfer_trtllm_fp4_moe(

255 x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

256 router_logits: torch.Tensor,

257 top_k: int,

258 + activation: str,

259 global_num_experts: int,

260 num_expert_group: int | None,

261 topk_group: int | None,

@@ -269,6 +270,7 @@ def flashinfer_trtllm_fp4_moe(

270 x: Input tensor

271 router_logits: Router logits for expert selection

272 top_k: Number of experts to select per token

273 + activation: Activation function to use

274 global_num_experts: Total number of experts across all ranks

275 num_expert_group: Number of expert groups (for grouped routing)

276 topk_group: Top-k within each group

@@ -282,6 +284,12 @@ def flashinfer_trtllm_fp4_moe(

284

285 from vllm.model_executor.models.llama4 import Llama4MoE

286

287 + # https://github.com/flashinfer-ai/flashinfer/blob/f0277fd1bff90e309e5c19cab36c5dae056d685d/flashinfer/fused_moe/core.py#L2404

288 + assert activation == "silu", (

289 + "Only SiLU activation is supported for FlashInfer TRTLLM FP4 MoE. "

290 + f"{activation} found instead."

291 + )

292 +

293 # Quantize input to FP4

294 if isinstance(x, tuple):

295 hidden_states_fp4, hidden_states_scale_linear_fp4 = x

@@ -352,6 +360,7 @@ def flashinfer_trtllm_fp4_routed_moe(

360 topk_ids: torch.Tensor,

361 topk_weights: torch.Tensor,

362 top_k: int,

363 + activation: str,

364 global_num_experts: int,

365 ) -> torch.Tensor:

366 """

@@ -364,6 +373,7 @@ def flashinfer_trtllm_fp4_routed_moe(

373 x: Input tensor

374 topk_ids: Ids of selected experts

375 top_k: Number of experts to select per token

376 + activation: Activation function to use

377 global_num_experts: Total number of experts across all ranks

378

379 Returns:

@@ -371,6 +381,12 @@ def flashinfer_trtllm_fp4_routed_moe(

381 """

382 import flashinfer

383

384 + # https://github.com/flashinfer-ai/flashinfer/blob/f0277fd1bff90e309e5c19cab36c5dae056d685d/flashinfer/fused_moe/core.py#L2535

385 + assert activation == "silu", (

386 + "Only SiLU activation is supported for FlashInfer TRTLLM FP4 Routed MoE. "

387 + f"{activation} found instead."

388 + )

389 +

390 # Pack top k ids and expert weights into a single int32 tensor, as

391 # required by TRT-LLM

392 packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(

vllm/model_executor/layers/quantization/utils/marlin_utils.py

+1 -8

@@ -233,8 +233,6 @@ def check_moe_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool:

233 intermediate_size_per_partition = layer.intermediate_size_per_partition

234 # apply_router_weight_on_input is not supported for moe marlin

235 supports_router_weight = not layer.apply_router_weight_on_input

236 - # moe marlin requires the activation to be silu

237 - supports_activation = layer.activation == "silu"

236

237 # gate-up: (n, k) = (intermediate_size_per_partition * 2, hidden_size)

238 # down: (n, k) = (hidden_size, intermediate_size_per_partition)

@@ -244,12 +242,7 @@ def check_moe_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool:

242 and intermediate_size_per_partition % max(64, group_size) == 0

243 )

244 supports_group_size = group_size in [-1, 32, 64, 128]

247 - return (

248 - supports_shape

249 - and supports_group_size

250 - and supports_router_weight

251 - and supports_activation

252 - )

245 + return supports_shape and supports_group_size and supports_router_weight

246

247

248 def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tensor):

vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py

+5 -2

@@ -235,6 +235,7 @@ def prepare_nvfp4_moe_layer_for_marlin(

235 w2: torch.Tensor,

236 w2_scale: torch.Tensor,

237 w2_scale_2: torch.Tensor,

238 + is_act_and_mul: bool,

239 ) -> tuple[

240 torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor

241 ]:

@@ -266,8 +267,9 @@ def prepare_nvfp4_moe_layer_for_marlin(

267 # Repack weights to marlin format

268 def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:

269 tensor_list = []

270 + num_shards = 2 if is_act_and_mul else 1

271 if "w13" in name:

270 - size_n, size_k = N * 2, K

272 + size_n, size_k = N * num_shards, K

273 else:

274 size_n, size_k = K, N

275

@@ -300,8 +302,9 @@ def prepare_nvfp4_moe_layer_for_marlin(

302 g_scales = g_scales.to(param_dtype)

303

304 tensor_list = []

305 + num_shards = 2 if is_act_and_mul else 1

306 if "w13" in name:

304 - size_n, size_k = N * 2, K

307 + size_n, size_k = N * num_shards, K

308 else:

309 size_n, size_k = K, N

310