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 14, 23:45 UTC
[Refactor] Merge Compressed Tensor FP8 `CompressedTensorsW8A8Fp8MoEMethod` and `CompressedTensorsW8A8Fp8MoECutlassMethod` (#21775)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 28, 2025 at 21:47 UTC
48b763d6b5c969024a8a5ae30c2bf9a91e8ac032
1 file changed
+100
-289
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+100
-289
@@ -45,7 +45,6 @@ class GPTQMarlinState(Enum):
45
46
__all__ = [
47
"CompressedTensorsMoEMethod", "CompressedTensorsW8A8Fp8MoEMethod",
48
- "CompressedTensorsW8A8Fp8MoECutlassMethod",
48
"CompressedTensorsW8A8Int8MoEMethod",
49
"CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod",
50
"CompressedTensorsW4A4MoeMethod"
@@ -84,9 +83,8 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
83
elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
84
return CompressedTensorsW4A4MoeMethod()
85
elif (quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
87
- or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)):
88
- return CompressedTensorsW8A8Fp8MoECutlassMethod(quant_config)
89
- elif quant_config._is_fp8_w8a8(weight_quant, input_quant):
86
+ or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
87
+ or quant_config._is_fp8_w8a8(weight_quant, input_quant)):
88
return CompressedTensorsW8A8Fp8MoEMethod(quant_config)
89
elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
90
return CompressedTensorsW8A8Int8MoEMethod(quant_config)
@@ -378,6 +376,14 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
376
377
self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
378
379
+ # cutlass path
380
+ self.is_fp8_w8a8_sm100 = quant_config._is_fp8_w8a8_sm100(
381
+ self.weight_quant, self.input_quant)
382
+ self.use_cutlass = (quant_config._is_fp8_w8a8_sm90(
383
+ self.weight_quant, self.input_quant) or self.is_fp8_w8a8_sm100)
384
+ self.fused_experts = None # type: ignore[assignment]
385
+ self.disable_expert_map = False
386
+
387
def create_weights(self, layer: torch.nn.Module, num_experts: int,
388
hidden_size: int, intermediate_size_per_partition: int,
389
params_dtype: torch.dtype, **extra_weight_attrs):
@@ -558,6 +564,34 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
564
prepare_finalize: FusedMoEPrepareAndFinalize,
565
moe: FusedMoEConfig,
566
) -> FusedMoEPermuteExpertsUnpermute:
567
+ # cutlass path
568
+ if self.use_cutlass:
569
+ from vllm.model_executor.layers.fused_moe import CutlassExpertsFp8
570
+
571
+ use_batched_format = (prepare_finalize.activation_format ==
572
+ FusedMoEActivationFormat.BatchedExperts)
573
+
574
+ num_dispatchers = prepare_finalize.num_dispatchers()
575
+ num_experts = (moe.num_local_experts
576
+ if use_batched_format else moe.num_experts)
577
+
578
+ logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)
579
+
580
+ experts = CutlassExpertsFp8(
581
+ num_experts,
582
+ moe.in_dtype,
583
+ self.input_quant.strategy == QuantizationStrategy.TOKEN,
584
+ self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
585
+ num_dispatchers=num_dispatchers,
586
+ use_batched_format=use_batched_format,
587
+ )
588
+
589
+ self.disable_expert_map = (num_dispatchers > 1
590
+ or not experts.supports_expert_map())
591
+
592
+ return experts
593
+
594
+ # triton path
595
from vllm.model_executor.layers.fused_moe import TritonExperts
596
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
597
BatchedTritonExperts)
@@ -629,6 +663,68 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
663
indices_type=self.topk_indices_dtype,
664
)
665
666
+ # cutlass path
667
+ if self.use_cutlass:
668
+ per_act_token = (
669
+ self.input_quant.strategy == QuantizationStrategy.TOKEN)
670
+ per_channel_quant = (
671
+ self.weight_quant.strategy == QuantizationStrategy.CHANNEL)
672
+
673
+ # small-batch fallback on SM100
674
+ if self.is_fp8_w8a8_sm100 and topk_ids.shape[0] <= 8:
675
+ from vllm.model_executor.layers.fused_moe import fused_experts
676
+ return fused_experts(
677
+ hidden_states=x,
678
+ w1=layer.w13_weight,
679
+ w2=layer.w2_weight,
680
+ topk_weights=topk_weights,
681
+ topk_ids=topk_ids,
682
+ inplace=True,
683
+ activation=activation,
684
+ apply_router_weight_on_input=apply_router_weight_on_input,
685
+ use_fp8_w8a8=True,
686
+ per_channel_quant=per_channel_quant,
687
+ global_num_experts=global_num_experts,
688
+ expert_map=None if self.disable_expert_map else expert_map,
689
+ w1_scale=layer.w13_weight_scale,
690
+ w2_scale=layer.w2_weight_scale,
691
+ a1_scale=layer.w13_input_scale,
692
+ a2_scale=layer.w2_input_scale)
693
+
694
+ if self.fused_experts is None:
695
+ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
696
+ cutlass_moe_fp8)
697
+ return cutlass_moe_fp8(
698
+ x,
699
+ layer.w13_weight,
700
+ layer.w2_weight,
701
+ topk_weights,
702
+ topk_ids,
703
+ per_act_token=per_act_token,
704
+ activation=activation,
705
+ global_num_experts=global_num_experts,
706
+ expert_map=None if self.disable_expert_map else expert_map,
707
+ w1_scale=layer.w13_weight_scale,
708
+ w2_scale=layer.w2_weight_scale,
709
+ a1_scale=layer.w13_input_scale,
710
+ a2_scale=layer.w2_input_scale,
711
+ )
712
+ else:
713
+ return self.fused_experts(
714
+ x,
715
+ layer.w13_weight,
716
+ layer.w2_weight,
717
+ topk_weights,
718
+ topk_ids,
719
+ activation=activation,
720
+ global_num_experts=global_num_experts,
721
+ expert_map=None if self.disable_expert_map else expert_map,
722
+ w1_scale=layer.w13_weight_scale,
723
+ w2_scale=layer.w2_weight_scale,
724
+ a1_scale=layer.w13_input_scale,
725
+ a2_scale=layer.w2_input_scale,
726
+ )
727
+
728
if self.rocm_aiter_moe_enabled:
729
return self.rocm_aiter_fused_experts_func(
730
hidden_states=x,
@@ -685,291 +781,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
781
a2_scale=layer.w2_input_scale)
782
783
688
-class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
689
-
690
- def __init__(
691
- self,
692
- quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501
693
- ):
694
- self.quant_config = quant_config
695
- self.weight_quant = self.quant_config.target_scheme_map["Linear"].get(
696
- "weights")
697
- self.input_quant = self.quant_config.target_scheme_map["Linear"].get(
698
- "input_activations")
699
-
700
- per_tensor = (self.weight_quant.strategy == QuantizationStrategy.TENSOR
701
- and self.input_quant.strategy
702
- == QuantizationStrategy.TENSOR)
703
- per_channel = (
704
- self.weight_quant.strategy == QuantizationStrategy.CHANNEL
705
- and self.input_quant.strategy == QuantizationStrategy.TOKEN)
706
- if not (per_tensor or per_channel):
707
- raise ValueError(
708
- "For FP8 Fused MoE layers, we require per tensor "
709
- "or channelwise, dynamic per token quantization. Found "
710
- f"{self.weight_quant}, {self.input_quant}")
711
-
712
- self.static_input_scales = not self.input_quant.dynamic
713
- if self.static_input_scales and per_channel:
714
- raise ValueError(
715
- "For FP8 Fused MoE layer, we require either per tensor or "
716
- "channelwise, dynamic per token quantization.")
717
-
718
- self.topk_indices_dtype = None
719
- self.fused_experts = None # type: ignore
720
- self.disable_expert_map = False
721
- self.is_fp8_w8a8_sm100 = self.quant_config._is_fp8_w8a8_sm100(
722
- self.weight_quant, self.input_quant)
723
-
724
- def create_weights(self, layer: torch.nn.Module, num_experts: int,
725
- hidden_size: int, intermediate_size_per_partition: int,
726
- params_dtype: torch.dtype, **extra_weight_attrs):
727
-
728
- params_dtype = torch.float8_e4m3fn
729
-
730
- # WEIGHTS
731
- w13_weight = torch.nn.Parameter(torch.empty(
732
- num_experts,
733
- 2 * intermediate_size_per_partition,
734
- hidden_size,
735
- dtype=params_dtype),
736
- requires_grad=False)
737
- layer.register_parameter("w13_weight", w13_weight)
738
- set_weight_attrs(w13_weight, extra_weight_attrs)
739
-
740
- w2_weight = torch.nn.Parameter(torch.empty(
741
- num_experts,
742
- hidden_size,
743
- intermediate_size_per_partition,
744
- dtype=params_dtype),
745
- requires_grad=False)
746
- layer.register_parameter("w2_weight", w2_weight)
747
- set_weight_attrs(w2_weight, extra_weight_attrs)
748
-
749
- # WEIGHT_SCALES
750
- if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
751
- # Allocate 2 scales for w1 and w3 respectively.
752
- # They are combined to a single scale after weight loading.
753
- w13_weight_scale = torch.nn.Parameter(torch.ones(
754
- num_experts, 2, dtype=torch.float32),
755
- requires_grad=False)
756
- layer.register_parameter("w13_weight_scale", w13_weight_scale)
757
- w2_weight_scale = torch.nn.Parameter(torch.ones(
758
- num_experts, dtype=torch.float32),
759
- requires_grad=False)
760
- layer.register_parameter("w2_weight_scale", w2_weight_scale)
761
- # Add PER-TENSOR quantization for FusedMoE.weight_loader.
762
- extra_weight_attrs.update(
763
- {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
764
- set_weight_attrs(w13_weight_scale, extra_weight_attrs)
765
- set_weight_attrs(w2_weight_scale, extra_weight_attrs)
766
-
767
- elif self.weight_quant.strategy == QuantizationStrategy.CHANNEL:
768
- w13_weight_scale = torch.nn.Parameter(torch.ones(
769
- num_experts,
770
- 2 * intermediate_size_per_partition,
771
- 1,
772
- dtype=torch.float32),
773
- requires_grad=False)
774
- layer.register_parameter("w13_weight_scale", w13_weight_scale)
775
- w2_weight_scale = torch.nn.Parameter(torch.ones(
776
- num_experts, hidden_size, 1, dtype=torch.float32),
777
- requires_grad=False)
778
- layer.register_parameter("w2_weight_scale", w2_weight_scale)
779
- # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
780
- extra_weight_attrs.update(
781
- {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value})
782
- set_weight_attrs(w13_weight_scale, extra_weight_attrs)
783
- set_weight_attrs(w2_weight_scale, extra_weight_attrs)
784
-
785
- # INPUT_SCALES
786
- if self.static_input_scales:
787
- w13_input_scale = torch.nn.Parameter(torch.ones(
788
- num_experts, dtype=torch.float32),
789
- requires_grad=False)
790
- layer.register_parameter("w13_input_scale", w13_input_scale)
791
- set_weight_attrs(w13_input_scale, extra_weight_attrs)
792
-
793
- w2_input_scale = torch.nn.Parameter(torch.ones(
794
- num_experts, dtype=torch.float32),
795
- requires_grad=False)
796
- layer.register_parameter("w2_input_scale", w2_input_scale)
797
- set_weight_attrs(w2_input_scale, extra_weight_attrs)
798
- else:
799
- layer.w13_input_scale = None
800
- layer.w2_input_scale = None
801
-
802
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
803
- # Fp8 moe kernels require a single activation scale.
804
- # We take the max of all the scales in case they differ.
805
- if self.static_input_scales:
806
- assert self.input_quant.strategy == QuantizationStrategy.TENSOR
807
- if (layer.w13_input_scale is None or layer.w2_input_scale is None):
808
- raise ValueError(
809
- "QuantConfig has static quantization, but found "
810
- "activation scales are None.")
811
- if (not all_close_1d(layer.w13_input_scale)
812
- or not all_close_1d(layer.w2_input_scale)):
813
- logger.warning_once(
814
- "Found input_scales that are not equal for "
815
- "fp8 MoE layer. Using the maximum across experts "
816
- "for each layer.")
817
- layer.w13_input_scale = torch.nn.Parameter(
818
- layer.w13_input_scale.max(), requires_grad=False)
819
- layer.w2_input_scale = torch.nn.Parameter(
820
- layer.w2_input_scale.max(), requires_grad=False)
821
-
822
- # For Per-TENSOR case, Fp8 moe kernel needs single weight scale
823
- # for w13 per expert. Use max then dequant and requant each expert.
824
- if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
825
- assert layer.w13_weight_scale is not None
826
- shard_size = layer.intermediate_size_per_partition
827
- max_w13_scales = layer.w13_weight_scale.max(dim=1).values
828
- for expert_id in range(layer.local_num_experts):
829
- start = 0
830
- for shard_id in range(2):
831
- dq_weight = per_tensor_dequantize(
832
- layer.w13_weight[expert_id][start:start +
833
- shard_size, :],
834
- layer.w13_weight_scale[expert_id][shard_id])
835
- layer.w13_weight[expert_id][
836
- start:start + shard_size, :], _ = ops.scaled_fp8_quant(
837
- dq_weight, max_w13_scales[expert_id])
838
- start += shard_size
839
- layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
840
- requires_grad=False)
841
-
842
- def select_gemm_impl(
843
- self,
844
- prepare_finalize: FusedMoEPrepareAndFinalize,
845
- moe: FusedMoEConfig,
846
- ) -> FusedMoEPermuteExpertsUnpermute:
847
- from vllm.model_executor.layers.fused_moe import CutlassExpertsFp8
848
-
849
- use_batched_format = (prepare_finalize.activation_format ==
850
- FusedMoEActivationFormat.BatchedExperts)
851
-
852
- num_dispatchers = prepare_finalize.num_dispatchers()
853
-
854
- num_experts = (moe.num_local_experts
855
- if use_batched_format else moe.num_experts)
856
-
857
- logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)
858
-
859
- experts = CutlassExpertsFp8(
860
- num_experts,
861
- moe.in_dtype,
862
- self.input_quant.strategy == QuantizationStrategy.TOKEN,
863
- self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
864
- num_dispatchers=num_dispatchers,
865
- use_batched_format=use_batched_format,
866
- )
867
-
868
- self.disable_expert_map = (num_dispatchers > 1
869
- or not experts.supports_expert_map())
870
-
871
- return experts
872
-
873
- def apply(
874
- self,
875
- layer: torch.nn.Module,
876
- x: torch.Tensor,
877
- router_logits: torch.Tensor,
878
- top_k: int,
879
- renormalize: bool,
880
- use_grouped_topk: bool = False,
881
- topk_group: Optional[int] = None,
882
- num_expert_group: Optional[int] = None,
883
- global_num_experts: int = -1,
884
- expert_map: Optional[torch.Tensor] = None,
885
- custom_routing_function: Optional[Callable] = None,
886
- scoring_func: str = "softmax",
887
- e_score_correction_bias: Optional[torch.Tensor] = None,
888
- apply_router_weight_on_input: bool = False,
889
- activation: str = "silu",
890
- enable_eplb: bool = False,
891
- expert_load_view: Optional[torch.Tensor] = None,
892
- logical_to_physical_map: Optional[torch.Tensor] = None,
893
- logical_replica_count: Optional[torch.Tensor] = None,
894
- ) -> torch.Tensor:
895
- if enable_eplb:
896
- raise NotImplementedError(
897
- "EPLB not supported for "
898
- "`CompressedTensorsW8A8Fp8MoECutlassMethod` yet.")
899
-
900
- topk_weights, topk_ids = FusedMoE.select_experts(
901
- hidden_states=x,
902
- router_logits=router_logits,
903
- use_grouped_topk=use_grouped_topk,
904
- top_k=top_k,
905
- renormalize=renormalize,
906
- topk_group=topk_group,
907
- num_expert_group=num_expert_group,
908
- custom_routing_function=custom_routing_function,
909
- scoring_func=scoring_func,
910
- e_score_correction_bias=e_score_correction_bias)
911
-
912
- per_act_token = (
913
- self.input_quant.strategy == QuantizationStrategy.TOKEN)
914
- per_channel_quant = (
915
- self.weight_quant.strategy == QuantizationStrategy.CHANNEL)
916
- # Triton fused_experts is faster in small batch sizes on SM100.
917
- # Fall back to fused_experts in small batch sizes.
918
- if self.is_fp8_w8a8_sm100 and topk_ids.shape[0] <= 8:
919
- from vllm.model_executor.layers.fused_moe import fused_experts
920
- return fused_experts(
921
- x,
922
- layer.w13_weight,
923
- layer.w2_weight,
924
- topk_weights,
925
- topk_ids,
926
- inplace=True,
927
- activation=activation,
928
- apply_router_weight_on_input=apply_router_weight_on_input,
929
- use_fp8_w8a8=True,
930
- per_channel_quant=per_channel_quant,
931
- global_num_experts=global_num_experts,
932
- expert_map=None if self.disable_expert_map else expert_map,
933
- w1_scale=layer.w13_weight_scale,
934
- w2_scale=layer.w2_weight_scale,
935
- a1_scale=layer.w13_input_scale,
936
- a2_scale=layer.w2_input_scale)
937
- if self.fused_experts is None:
938
- # If no modular kernel is provided, use cutlass_moe_fp8
939
- from vllm.model_executor.layers.fused_moe.cutlass_moe import (
940
- cutlass_moe_fp8)
941
- return cutlass_moe_fp8(
942
- x,
943
- layer.w13_weight,
944
- layer.w2_weight,
945
- topk_weights,
946
- topk_ids,
947
- per_act_token=per_act_token,
948
- activation=activation,
949
- global_num_experts=global_num_experts,
950
- expert_map=None if self.disable_expert_map else expert_map,
951
- w1_scale=layer.w13_weight_scale,
952
- w2_scale=layer.w2_weight_scale,
953
- a1_scale=layer.w13_input_scale,
954
- a2_scale=layer.w2_input_scale,
955
- )
956
- else:
957
- return self.fused_experts(
958
- x,
959
- layer.w13_weight,
960
- layer.w2_weight,
961
- topk_weights,
962
- topk_ids,
963
- activation=activation,
964
- global_num_experts=global_num_experts,
965
- expert_map=None if self.disable_expert_map else expert_map,
966
- w1_scale=layer.w13_weight_scale,
967
- w2_scale=layer.w2_weight_scale,
968
- a1_scale=layer.w13_input_scale,
969
- a2_scale=layer.w2_input_scale,
970
- )
971
-
972
-
784
class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
785
786
def __init__(