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, 02:15 UTC
Support Llama 4 for cutlass_moe_fp4 (#20453)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 10, 2025 at 04:53 UTC
31b96d1c643c5866dc080b57a71693de1b83cfc6
3 files changed
+80
-74
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+28
-9
@@ -411,13 +411,23 @@ FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
411
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
412
413
414
-def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
415
- w1_fp4: torch.Tensor, w1_blockscale: torch.Tensor,
416
- w1_alphas: torch.Tensor, a2_gscale: torch.Tensor,
417
- w2_fp4: torch.Tensor, w2_blockscale: torch.Tensor,
418
- w2_alphas: torch.Tensor, topk_weights: torch.Tensor,
419
- topk_ids: torch.Tensor, m: int, n: int, k: int, e: int,
420
- device: torch.device):
414
+def cutlass_moe_fp4(a: torch.Tensor,
415
+ a1_gscale: torch.Tensor,
416
+ w1_fp4: torch.Tensor,
417
+ w1_blockscale: torch.Tensor,
418
+ w1_alphas: torch.Tensor,
419
+ a2_gscale: torch.Tensor,
420
+ w2_fp4: torch.Tensor,
421
+ w2_blockscale: torch.Tensor,
422
+ w2_alphas: torch.Tensor,
423
+ topk_weights: torch.Tensor,
424
+ topk_ids: torch.Tensor,
425
+ m: int,
426
+ n: int,
427
+ k: int,
428
+ e: int,
429
+ device: torch.device,
430
+ apply_router_weight_on_input: bool = False):
431
"""
432
MoE implementation for FP4 Inputs
433
@@ -480,6 +490,12 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
490
a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
491
c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
492
493
+ if apply_router_weight_on_input:
494
+ # TODO: this only works for topK=1, will need to update for topK>1
495
+ assert num_topk == 1, \
496
+ "apply_router_weight_on_input is only implemented for topk=1"
497
+ a.mul_(topk_weights.to(out_dtype))
498
+
499
# problem shapes should have [m, n, k]
500
# Note that problem sizes are based on logical number of elements.
501
ops.get_cutlass_moe_mm_data(topk_ids, expert_offsets, problem_sizes1,
@@ -517,8 +533,11 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
533
del int_fp4, int_blockscale
534
535
c2 = ops.shuffle_rows(c2, c_map)
520
- out = (c2.view(m, num_topk, k) *
521
- topk_weights.view(m, num_topk, 1).half()).sum(dim=1)
536
+ if not apply_router_weight_on_input:
537
+ out = (c2.view(m, num_topk, k) *
538
+ topk_weights.view(m, num_topk, 1).to(out_dtype)).sum(dim=1)
539
+ else:
540
+ out = c2.view(m, num_topk, k).sum(dim=1)
541
return out.to(dtype=out_dtype)
542
543
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+20
-20
@@ -295,6 +295,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
295
if enable_eplb:
296
raise NotImplementedError("EPLB not supported for "
297
"`CompressedTensorsW4A4MoeMethod` yet.")
298
+ assert activation == "silu", "Only SiLU activation is supported."
299
300
topk_weights, topk_ids = FusedMoE.select_experts(
301
hidden_states=x,
@@ -326,10 +327,6 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
327
global_num_experts=global_num_experts,
328
expert_map=expert_map)
329
329
- assert activation == "silu", "Only SiLU activation is supported."
330
- assert not apply_router_weight_on_input, (
331
- "Router weight on input is not "
332
- "supported for CompressedTensorsW4A4MoeMethod.")
330
assert expert_map is None, ("Expert Parallelism / expert_map "
331
"is currently not supported for "
332
"CompressedTensorsW4A4MoeMethod.")
@@ -339,22 +336,25 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
336
337
# Cutlass moe takes in activations in BF16/Half precision
338
# and fp4 quantized weights loaded from the checkpoint
342
- return cutlass_moe_fp4(a=x,
343
- w1_fp4=layer.w13_weight,
344
- w1_blockscale=layer.w13_blockscale_swizzled,
345
- w1_alphas=layer.g1_alphas,
346
- w2_fp4=layer.w2_weight,
347
- w2_blockscale=layer.w2_blockscale_swizzled,
348
- w2_alphas=layer.g2_alphas,
349
- topk_weights=topk_weights,
350
- topk_ids=topk_ids,
351
- m=x.shape[0],
352
- n=layer.w2_weight.shape[2] * 2,
353
- k=x.shape[1],
354
- e=layer.w13_weight.shape[0],
355
- a1_gscale=layer.w13_input_scale_quant,
356
- a2_gscale=layer.w2_input_scale_quant,
357
- device=x.device).to(x.dtype)
339
+ return cutlass_moe_fp4(
340
+ a=x,
341
+ w1_fp4=layer.w13_weight,
342
+ w1_blockscale=layer.w13_blockscale_swizzled,
343
+ w1_alphas=layer.g1_alphas,
344
+ w2_fp4=layer.w2_weight,
345
+ w2_blockscale=layer.w2_blockscale_swizzled,
346
+ w2_alphas=layer.g2_alphas,
347
+ topk_weights=topk_weights,
348
+ topk_ids=topk_ids,
349
+ m=x.shape[0],
350
+ n=layer.w2_weight.shape[2] * 2,
351
+ k=x.shape[1],
352
+ e=layer.w13_weight.shape[0],
353
+ a1_gscale=layer.w13_input_scale_quant,
354
+ a2_gscale=layer.w2_input_scale_quant,
355
+ device=x.device,
356
+ apply_router_weight_on_input=apply_router_weight_on_input).to(
357
+ x.dtype)
358
359
360
class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
vllm/model_executor/layers/quantization/modelopt.py
+32
-45
@@ -673,21 +673,21 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
673
if enable_eplb:
674
raise NotImplementedError(
675
"EPLB not supported for `ModelOptNvFp4FusedMoE` yet.")
676
+ assert activation == "silu", "Only SiLU activation is supported."
677
677
- if self.use_marlin:
678
- topk_weights, topk_ids = FusedMoE.select_experts(
679
- hidden_states=x,
680
- router_logits=router_logits,
681
- use_grouped_topk=use_grouped_topk,
682
- top_k=top_k,
683
- renormalize=renormalize,
684
- topk_group=topk_group,
685
- num_expert_group=num_expert_group,
686
- custom_routing_function=custom_routing_function,
687
- scoring_func=scoring_func,
688
- e_score_correction_bias=e_score_correction_bias,
689
- )
678
+ topk_weights, topk_ids = FusedMoE.select_experts(
679
+ hidden_states=x,
680
+ router_logits=router_logits,
681
+ use_grouped_topk=use_grouped_topk,
682
+ top_k=top_k,
683
+ renormalize=renormalize,
684
+ topk_group=topk_group,
685
+ num_expert_group=num_expert_group,
686
+ custom_routing_function=custom_routing_function,
687
+ scoring_func=scoring_func,
688
+ e_score_correction_bias=e_score_correction_bias)
689
690
+ if self.use_marlin:
691
return torch.ops.vllm.fused_marlin_moe(
692
x,
693
layer.w13_weight,
@@ -704,44 +704,31 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
704
global_num_experts=global_num_experts,
705
expert_map=expert_map)
706
707
- assert activation == "silu", "Only SiLU activation is supported."
708
- assert not apply_router_weight_on_input, (
709
- "Router weight on input is not "
710
- "supported for ModelOptNvFp4FusedMoE.")
707
assert expert_map is None, ("Expert Parallelism / expert_map "
708
"is currently not supported for "
709
"ModelOptNvFp4FusedMoE.")
710
715
- topk_weights, topk_ids = FusedMoE.select_experts(
716
- hidden_states=x,
717
- router_logits=router_logits,
718
- use_grouped_topk=use_grouped_topk,
719
- top_k=top_k,
720
- renormalize=renormalize,
721
- topk_group=topk_group,
722
- num_expert_group=num_expert_group,
723
- custom_routing_function=custom_routing_function,
724
- scoring_func=scoring_func,
725
- e_score_correction_bias=e_score_correction_bias)
726
-
711
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
712
cutlass_moe_fp4)
713
714
# Cutlass moe takes in activations in BF16/Half precision
715
# and fp4 quantized weights loaded from the checkpoint
732
- return cutlass_moe_fp4(a=x,
733
- w1_fp4=layer.w13_weight,
734
- w1_blockscale=layer.w13_blockscale_swizzled,
735
- w1_alphas=layer.g1_alphas,
736
- w2_fp4=layer.w2_weight,
737
- w2_blockscale=layer.w2_blockscale_swizzled,
738
- w2_alphas=layer.g2_alphas,
739
- topk_weights=topk_weights,
740
- topk_ids=topk_ids,
741
- m=x.shape[0],
742
- n=layer.w2_weight.shape[2] * 2,
743
- k=x.shape[1],
744
- e=layer.w13_weight.shape[0],
745
- a1_gscale=layer.w13_input_scale_quant,
746
- a2_gscale=layer.w2_input_scale_quant,
747
- device=x.device).to(x.dtype)
716
+ return cutlass_moe_fp4(
717
+ a=x,
718
+ w1_fp4=layer.w13_weight,
719
+ w1_blockscale=layer.w13_blockscale_swizzled,
720
+ w1_alphas=layer.g1_alphas,
721
+ w2_fp4=layer.w2_weight,
722
+ w2_blockscale=layer.w2_blockscale_swizzled,
723
+ w2_alphas=layer.g2_alphas,
724
+ topk_weights=topk_weights,
725
+ topk_ids=topk_ids,
726
+ m=x.shape[0],
727
+ n=layer.w2_weight.shape[2] * 2,
728
+ k=x.shape[1],
729
+ e=layer.w13_weight.shape[0],
730
+ a1_gscale=layer.w13_input_scale_quant,
731
+ a2_gscale=layer.w2_input_scale_quant,
732
+ device=x.device,
733
+ apply_router_weight_on_input=apply_router_weight_on_input).to(
734
+ x.dtype)