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

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 15, 03:30 UTC

Support Llama 4 for fused_marlin_moe (#20457)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Jul 4, 2025 at 16:55 UTC 0e3fe896e2a02e7ca74525c1a79f4cd428320eb6

6 files changed +11 -17

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+4 -2

@@ -24,6 +24,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,

24 topk_weights: torch.Tensor,

25 topk_ids: torch.Tensor,

26 quant_type_id: int,

27 + apply_router_weight_on_input: bool = False,

28 global_num_experts: int = -1,

29 expert_map: Optional[torch.Tensor] = None,

30 global_scale1: Optional[torch.Tensor] = None,

@@ -149,7 +150,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,

150 topk_weights,

151 moe_block_size=block_size_m,

152 top_k=topk,

152 - mul_topk_weights=False,

153 + mul_topk_weights=apply_router_weight_on_input,

154 is_ep=expert_map is not None,

155 b_q_type=quant_type,

156 size_m=M,

@@ -182,7 +183,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,

183 topk_weights,

184 moe_block_size=block_size_m,

185 top_k=1,

185 - mul_topk_weights=True,

186 + mul_topk_weights=not apply_router_weight_on_input,

187 is_ep=expert_map is not None,

188 b_q_type=quant_type,

189 size_m=M * topk,

@@ -208,6 +209,7 @@ def fused_marlin_moe_fake(hidden_states: torch.Tensor,

209 topk_weights: torch.Tensor,

210 topk_ids: torch.Tensor,

211 quant_type_id: int,

212 + apply_router_weight_on_input: bool = False,

213 global_num_experts: int = -1,

214 global_scale1: Optional[torch.Tensor] = None,

215 global_scale2: Optional[torch.Tensor] = None,

vllm/model_executor/layers/quantization/awq_marlin.py

+1 -5

@@ -493,11 +493,6 @@ class AWQMoEMethod(FusedMoEMethodBase):

493

494 assert activation == "silu", "Only SiLU activation is supported."

495

496 - if apply_router_weight_on_input:

497 - raise NotImplementedError(

498 - "Apply router weight on input is not supported for"

499 - "fused Marlin MoE method.")

500 -

496 topk_weights, topk_ids = FusedMoE.select_experts(

497 hidden_states=x,

498 router_logits=router_logits,

@@ -520,6 +515,7 @@ class AWQMoEMethod(FusedMoEMethodBase):

515 topk_weights,

516 topk_ids,

517 quant_type_id=self.quant_type.id,

518 + apply_router_weight_on_input=apply_router_weight_on_input,

519 global_num_experts=global_num_experts,

520 expert_map=expert_map,

521 w1_zeros=layer.w13_qzeros,

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

+3 -4

@@ -322,6 +322,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

322 global_scale1=layer.w13_weight_scale_2,

323 global_scale2=layer.w2_weight_scale_2,

324 quant_type_id=scalar_types.float4_e2m1f.id,

325 + apply_router_weight_on_input=apply_router_weight_on_input,

326 global_num_experts=global_num_experts,

327 expert_map=expert_map)

328

@@ -669,8 +670,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

670 if self.use_marlin:

671 assert activation == "silu", (

672 f"{activation} not supported for Marlin MoE.")

672 - assert not apply_router_weight_on_input, (

673 - "Apply router weight on input not supported for Marlin MoE.")

673 return torch.ops.vllm.fused_marlin_moe(

674 x,

675 layer.w13_weight,

@@ -681,6 +680,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

680 topk_weights,

681 topk_ids,

682 quant_type_id=scalar_types.float8_e4m3fn.id,

683 + apply_router_weight_on_input=apply_router_weight_on_input,

684 global_num_experts=global_num_experts,

685 expert_map=expert_map)

686

@@ -1356,8 +1356,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1356

1357 assert activation == "silu", (

1358 f"{activation} not supported for Marlin MoE.")

1359 - assert not apply_router_weight_on_input, (

1360 - "Apply router weight on input not supported for Marlin MoE.")

1359

1360 topk_weights, topk_ids = FusedMoE.select_experts(

1361 hidden_states=x,

@@ -1381,6 +1379,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1379 topk_weights,

1380 topk_ids,

1381 quant_type_id=self.quant_type.id,

1382 + apply_router_weight_on_input=apply_router_weight_on_input,

1383 global_num_experts=global_num_experts,

1384 expert_map=expert_map,

1385 g_idx1=layer.w13_weight_g_idx,

vllm/model_executor/layers/quantization/fp8.py

+1 -2

@@ -889,8 +889,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

889 elif self.use_marlin:

890 assert activation == "silu", (

891 f"{activation} not supported for Marlin MoE.")

892 - assert not apply_router_weight_on_input, (

893 - "Apply router weight on input not supported for Marlin MoE.")

892 return torch.ops.vllm.fused_marlin_moe(

893 x,

894 layer.w13_weight,

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

899 topk_weights,

900 topk_ids,

901 quant_type_id=scalar_types.float8_e4m3fn.id,

902 + apply_router_weight_on_input=apply_router_weight_on_input,

903 global_num_experts=global_num_experts,

904 expert_map=expert_map)

905 else:

vllm/model_executor/layers/quantization/gptq_marlin.py

+1 -4

@@ -645,10 +645,6 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

645 "EPLB not supported for `GPTQMarlinMoEMethod` yet.")

646

647 assert activation == "silu", "Only SiLU activation is supported."

648 - if apply_router_weight_on_input:

649 - raise NotImplementedError(

650 - "Apply router weight on input is not supported for "

651 - "fused Marlin MoE method.")

648

649 topk_weights, topk_ids = FusedMoE.select_experts(

650 hidden_states=x,

@@ -672,6 +668,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

668 topk_weights,

669 topk_ids,

670 quant_type_id=self.quant_type.id,

671 + apply_router_weight_on_input=apply_router_weight_on_input,

672 global_num_experts=global_num_experts,

673 expert_map=expert_map,

674 g_idx1=layer.w13_g_idx,

vllm/model_executor/layers/quantization/modelopt.py

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

700 global_scale1=layer.w13_weight_scale_2,

701 global_scale2=layer.w2_weight_scale_2,

702 quant_type_id=scalar_types.float4_e2m1f.id,

703 + apply_router_weight_on_input=apply_router_weight_on_input,

704 global_num_experts=global_num_experts,

705 expert_map=expert_map)

706