727c41f3fd · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any, Optional

6

7 import torch

8 from torch.nn import Module

9 -from torch.nn.parameter import Parameter

9 from torch.utils._python_dispatch import TorchDispatchMode

10

11 import vllm.envs as envs

@@ -728,6 +727,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):

727 self.quant_config = quant_config

728 self.weight_block_size = self.quant_config.weight_block_size

729 self.block_quant: bool = self.weight_block_size is not None

730 + self.weight_scale_name = (

731 + "weight_scale_inv" if self.block_quant else "weight_scale"

732 + )

733 self.fp8_backend = get_fp8_moe_backend(

734 self.block_quant, layer.moe_parallel_config, self.moe.is_lora_enabled

735 )

@@ -832,38 +834,28 @@ class Fp8MoEMethod(FusedMoEMethodBase):

834

835 # WEIGHT_SCALES

836 if not self.block_quant:

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

836 - # They will be combined to a single scale after weight loading.

837 - w13_weight_scale = torch.nn.Parameter(

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

839 - )

840 - w2_weight_scale = torch.nn.Parameter(

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

842 - )

843 - layer.register_parameter("w13_weight_scale", w13_weight_scale)

844 - layer.register_parameter("w2_weight_scale", w2_weight_scale)

837 + # For per-tensor quant, the scales are per expert and weight.

838 + w13_scale_data = torch.ones(num_experts, 2, dtype=torch.float32)

839 + w2_scale_data = torch.ones(num_experts, dtype=torch.float32)

840 else:

846 - w13_weight_scale = torch.nn.Parameter(

847 - torch.ones(

848 - num_experts,

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

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

851 - dtype=torch.float32,

852 - ),

853 - requires_grad=False,

841 + # For block quant, the scales are per block (typically 128x128).

842 + w13_scale_data = torch.ones(

843 + num_experts,

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

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

846 + dtype=torch.float32,

847 )

855 - w2_weight_scale = torch.nn.Parameter(

856 - torch.ones(

857 - num_experts,

858 - (hidden_size + block_n - 1) // block_n,

859 - (intermediate_size_per_partition + block_k - 1) // block_k,

860 - dtype=torch.float32,

861 - ),

862 - requires_grad=False,

848 + w2_scale_data = torch.ones(

849 + num_experts,

850 + (hidden_size + block_n - 1) // block_n,

851 + (intermediate_size_per_partition + block_k - 1) // block_k,

852 + dtype=torch.float32,

853 )

864 - layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)

865 - layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)

866 - assert self.quant_config.activation_scheme == "dynamic"

854 + w13_weight_scale = torch.nn.Parameter(w13_scale_data, requires_grad=False)

855 + w2_weight_scale = torch.nn.Parameter(w2_scale_data, requires_grad=False)

856 + # Note: name is weight_scale for tensor, weight_scale_inv for block.

857 + layer.register_parameter(f"w13_{self.weight_scale_name}", w13_weight_scale)

858 + layer.register_parameter(f"w2_{self.weight_scale_name}", w2_weight_scale)

859

860 # Add the quantization method used (per tensor/grouped/channel)

861 # to ensure the weight scales are loaded in properly

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

869

870 # INPUT_SCALES

871 if self.quant_config.activation_scheme == "static":

872 + assert not self.block_quant

873 w13_input_scale = torch.nn.Parameter(

874 torch.ones(num_experts, dtype=torch.float32), requires_grad=False

875 )

@@ -893,158 +886,60 @@ class Fp8MoEMethod(FusedMoEMethodBase):

886 layer.w13_input_scale = None

887 layer.w2_input_scale = None

888

896 - def process_weights_after_loading(self, layer: Module) -> None:

897 - if getattr(layer, "_already_called_process_weights_after_loading", False):

898 - return

899 -

900 - # TODO (rob): refactor block quant into separate class.

901 - if self.block_quant:

902 - assert self.quant_config.activation_scheme == "dynamic"

903 - if current_platform.is_fp8_fnuz():

904 - w13_weight, w13_weight_scale_inv, w13_input_scale = (

905 - normalize_e4m3fn_to_e4m3fnuz(

906 - layer.w13_weight,

907 - layer.w13_weight_scale_inv,

908 - layer.w13_input_scale,

909 - )

910 - )

911 - w2_weight, w2_weight_scale_inv, w2_input_scale = (

912 - normalize_e4m3fn_to_e4m3fnuz(

913 - layer.w2_weight, layer.w2_weight_scale_inv, layer.w2_input_scale

914 - )

915 - )

916 - elif self.flashinfer_moe_backend is not None:

917 - # NOTE: weights have to be swapped since the activation is

918 - # applied on different half for flashinfer vs vllm

919 - w13_weight = swap_w13_to_w31(layer.w13_weight.data)

920 - w13_weight_scale_inv = swap_w13_to_w31(layer.w13_weight_scale_inv.data)

921 - w2_weight = layer.w2_weight.data

922 - w2_weight_scale_inv = layer.w2_weight_scale_inv.data

889 + def _convert_weights_to_kernel_format(

890 + self,

891 + layer: Module,

892 + w13_weight: torch.Tensor,

893 + w2_weight: torch.Tensor,

894 + w13_weight_scale: torch.Tensor,

895 + w2_weight_scale: torch.Tensor,

896 + ) -> None:

897 + if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:

898 + assert self.block_quant

899 + w13_weight, w13_weight_scale = deepgemm_post_process_fp8_weight_block(

900 + wq=w13_weight,

901 + ws=w13_weight_scale,

902 + quant_block_shape=tuple(layer.weight_block_size),

903 + use_e8m0=is_deep_gemm_e8m0_used(),

904 + )

905 + w2_weight, w2_weight_scale = deepgemm_post_process_fp8_weight_block(

906 + wq=w2_weight,

907 + ws=w2_weight_scale,

908 + quant_block_shape=tuple(layer.weight_block_size),

909 + use_e8m0=is_deep_gemm_e8m0_used(),

910 + )

911 + elif self.fp8_backend == Fp8MoeBackend.AITER:

912 + w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(

913 + w13_weight, w2_weight

914 + )

915 + elif self.fp8_backend in [

916 + Fp8MoeBackend.FLASHINFER_CUTLASS,

917 + Fp8MoeBackend.FLASHINFER_TRTLLM,

918 + ]:

919 + w13_weight = swap_w13_to_w31(w13_weight)

920 + if self.block_quant:

921 + w13_weight_scale = swap_w13_to_w31(w13_weight_scale)

922 else:

924 - w13_weight = layer.w13_weight.data

925 - w13_weight_scale_inv = layer.w13_weight_scale_inv.data

926 - w2_weight = layer.w2_weight

927 - w2_weight_scale_inv = layer.w2_weight_scale_inv

928 -

929 - # torch.compile() cannot use Parameter subclasses.

930 - replace_parameter(layer, "w13_weight", w13_weight)

931 - replace_parameter(layer, "w13_weight_scale_inv", w13_weight_scale_inv)

932 - replace_parameter(layer, "w2_weight", w2_weight)

933 - replace_parameter(layer, "w2_weight_scale_inv", w2_weight_scale_inv)

934 - if self.fp8_backend == Fp8MoeBackend.AITER:

935 - # reshaping weights is required for aiter moe kernel.

936 - shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

937 - layer.w13_weight.data, layer.w2_weight.data

938 - )

939 -

940 - replace_parameter(layer, "w13_weight", shuffled_w13)

941 - replace_parameter(layer, "w2_weight", shuffled_w2)

942 -

943 - # DeepGemm scales need to be transposed and aligned. We try to do

944 - # it ahead of time for performance reasons.

945 - if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:

946 - dg_w13_weight, dg_w13_weight_scale_inv = (

947 - deepgemm_post_process_fp8_weight_block(

948 - wq=layer.w13_weight.data,

949 - ws=layer.w13_weight_scale_inv.data,

950 - quant_block_shape=tuple(layer.weight_block_size),

951 - use_e8m0=is_deep_gemm_e8m0_used(),

952 - )

953 - )

954 - dg_w2_weight, dg_w2_weight_scale_inv = (

955 - deepgemm_post_process_fp8_weight_block(

956 - wq=layer.w2_weight.data,

957 - ws=layer.w2_weight_scale_inv.data,

958 - quant_block_shape=tuple(layer.weight_block_size),

959 - use_e8m0=is_deep_gemm_e8m0_used(),

960 - )

961 - )

962 - layer.w13_weight = Parameter(dg_w13_weight, requires_grad=False)

963 - layer.w13_weight_scale_inv = Parameter(

964 - dg_w13_weight_scale_inv, requires_grad=False

965 - )

966 - layer.w2_weight = Parameter(dg_w2_weight, requires_grad=False)

967 - layer.w2_weight_scale_inv = Parameter(

968 - dg_w2_weight_scale_inv, requires_grad=False

969 - )

970 - else:

971 - # Fp8 moe kernels require a single activation scale.

972 - # We take the max of all the scales in case they differ.

973 - if self.quant_config.activation_scheme == "static":

974 - if layer.w13_input_scale is None or layer.w2_input_scale is None:

975 - raise ValueError(

976 - "QuantConfig has static quantization, but found "

977 - "activation scales are None."

978 - )

979 - if not all_close_1d(layer.w13_input_scale) or not all_close_1d(

980 - layer.w2_input_scale

981 - ):

982 - logger.warning_once(

983 - "Found input_scales that are not equal for "

984 - "fp8 MoE layer. Using the maximum across experts "

985 - "for each layer."

986 - )

987 - replace_parameter(layer, "w13_input_scale", layer.w13_input_scale.max())

988 - replace_parameter(layer, "w2_input_scale", layer.w2_input_scale.max())

989 - if current_platform.is_fp8_fnuz():

990 - # Normalize the weights and scales

991 - w13_weight, w13_weight_scale, w13_input_scale = (

992 - normalize_e4m3fn_to_e4m3fnuz(

993 - layer.w13_weight, layer.w13_weight_scale, layer.w13_input_scale

994 - )

995 - )

996 - w2_weight, w2_weight_scale, w2_input_scale = (

997 - normalize_e4m3fn_to_e4m3fnuz(

998 - layer.w2_weight, layer.w2_weight_scale, layer.w2_input_scale

999 - )

1000 - )

1001 - # Reset the parameter

1002 - replace_parameter(layer, "w13_weight", w13_weight)

1003 - replace_parameter(layer, "w13_weight_scale", w13_weight_scale)

1004 - if w13_input_scale is not None:

1005 - replace_parameter(layer, "w13_input_scale", w13_input_scale)

1006 - replace_parameter(layer, "w2_weight", w2_weight)

1007 - replace_parameter(layer, "w2_weight_scale", w2_weight_scale)

1008 - if w2_input_scale is not None:

1009 - replace_parameter(layer, "w2_input_scale", w2_input_scale)

1010 -

1011 - # Fp8 moe kernel needs single weight scale for w13 per expert.

1012 - # We take the max then dequant and requant each expert.

1013 - assert layer.w13_weight_scale is not None

1014 - shard_size = layer.intermediate_size_per_partition

1015 - max_w13_scales = layer.w13_weight_scale.max(dim=1).values

1016 - for expert_id in range(layer.local_num_experts):

1017 - start = 0

1018 - for shard_id in range(2):

1019 - dq_weight = per_tensor_dequantize(

1020 - layer.w13_weight[expert_id][start : start + shard_size, :],

1021 - layer.w13_weight_scale[expert_id][shard_id],

1022 - )

1023 - layer.w13_weight[expert_id][start : start + shard_size, :], _ = (

1024 - ops.scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])

1025 - )

1026 - start += shard_size

1027 -

1028 - if self.fp8_backend == Fp8MoeBackend.AITER:

1029 - shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

1030 - layer.w13_weight, layer.w2_weight

1031 - )

1032 -

1033 - replace_parameter(layer, "w13_weight", shuffled_w13)

1034 - replace_parameter(layer, "w2_weight", shuffled_w2)

1035 -

1036 - replace_parameter(layer, "w13_weight_scale", max_w13_scales)

1037 -

1038 - if self.flashinfer_moe_backend is not None:

1039 - # NOTE: weights have to be swapped since the activation is

1040 - # applied on different half for flashinfer vs vllm

1041 - assert not self.block_quant

923 + # TODO(rob): this function is a hack that renames the scaling

924 + # factors in the Module. This is a hack we should clean up.

925 register_moe_scaling_factors(layer)

1043 - w13_weight = swap_w13_to_w31(layer.w13_weight.data)

1044 - if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

926 + if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

927 rotate_flashinfer_fp8_moe_weights(w13_weight, w2_weight)

1046 - layer.w13_weight.data = w13_weight.data

928 + elif self.fp8_backend == Fp8MoeBackend.AITER:

929 + w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(

930 + w13_weight, w2_weight

931 + )

932

933 + # Replace parameters with updated versions. Note that this helper

934 + # function ensures the replacement is compatible with RL weight reloads.

935 + replace_parameter(layer, "w13_weight", w13_weight)

936 + replace_parameter(layer, "w2_weight", w2_weight)

937 + replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_weight_scale)

938 + replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_weight_scale)

939 +

940 + # TODO(rob): we do this after replace_parameter() because

941 + # prepare_moe_fp8_layer_for_marlin uses on the layer's params

942 + # directly. We will refactor this in a follow up PR.

943 if self.fp8_backend == Fp8MoeBackend.MARLIN:

944 prepare_moe_fp8_layer_for_marlin(

945 layer, False, input_dtype=self.marlin_input_dtype

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

948 del layer.w13_input_scale

949 del layer.w2_input_scale

950

951 + def _setup_kernel(self, layer: Module) -> None:

952 + """Setup Modular Kernel for TP Case"""

953 # NOTE(rob): this is a WIP refactor. We are first migrating

954 # all of the kernels in the TP case to use mk. Once this is

955 # done, then we will initialzie the TP case and DP/EP case

@@ -1134,6 +1031,71 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1031 )

1032 self.use_inplace = True

1033

1034 + def process_weights_after_loading(self, layer: Module) -> None:

1035 + if getattr(layer, "_already_called_process_weights_after_loading", False):

1036 + return

1037 +

1038 + # Allow for accessing weights and scales in standard way.

1039 + w13_weight = layer.w13_weight

1040 + w2_weight = layer.w2_weight

1041 + w13_weight_scale = getattr(layer, f"w13_{self.weight_scale_name}")

1042 + w2_weight_scale = getattr(layer, f"w2_{self.weight_scale_name}")

1043 +

1044 + # MI300x and MI325x use FNUZ format for FP8. Convert if needed.

1045 + if current_platform.is_fp8_fnuz():

1046 + w13_weight, w13_weight_scale, layer.w13_input_scale = (

1047 + normalize_e4m3fn_to_e4m3fnuz(

1048 + w13_weight, w13_weight_scale, layer.w13_input_scale

1049 + )

1050 + )

1051 + w2_weight, w2_weight_scale, layer.w2_input_scale = (

1052 + normalize_e4m3fn_to_e4m3fnuz(

1053 + w2_weight, w2_weight_scale, layer.w2_input_scale

1054 + )

1055 + )

1056 +

1057 + # Per tensor kernels require single activation scale. Use the max.

1058 + if self.quant_config.activation_scheme == "static":

1059 + assert not self.block_quant

1060 + assert layer.w13_input_scale is not None

1061 + assert layer.w2_input_scale is not None

1062 + if not all_close_1d(layer.w13_input_scale) or not all_close_1d(

1063 + layer.w2_input_scale

1064 + ):

1065 + logger.warning_once(

1066 + "Found input_scales that are not equal for "

1067 + "fp8 MoE layer. Using the maximum across experts "

1068 + "for each layer."

1069 + )

1070 + replace_parameter(layer, "w13_input_scale", layer.w13_input_scale.max())

1071 + replace_parameter(layer, "w2_input_scale", layer.w2_input_scale.max())

1072 +

1073 + # Per tensor kernels require single weight scale for w13 per expert, but

1074 + # on disk there is a scale for w1 and w3. Use the max to requantize.

1075 + if not self.block_quant:

1076 + shard_size = layer.intermediate_size_per_partition

1077 + max_w13_scales = w13_weight_scale.max(dim=1).values

1078 + for expert_id in range(layer.local_num_experts):

1079 + start = 0

1080 + for shard_id in range(2):

1081 + dq_weight = per_tensor_dequantize(

1082 + w13_weight[expert_id][start : start + shard_size, :],

1083 + w13_weight_scale[expert_id][shard_id],

1084 + )

1085 + w13_weight[expert_id][start : start + shard_size, :], _ = (

1086 + ops.scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])

1087 + )

1088 + start += shard_size

1089 + w13_weight_scale = max_w13_scales

1090 +

1091 + # Shuffle weights into the runtime format.

1092 + self._convert_weights_to_kernel_format(

1093 + layer, w13_weight, w2_weight, w13_weight_scale, w2_weight_scale

1094 + )

1095 +

1096 + # Setup modular kernel for TP case.

1097 + self._setup_kernel(layer)

1098 +

1099 def maybe_make_prepare_finalize(

1100 self,

1101 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

@@ -1453,32 +1415,26 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1415

1416 # If checkpoint is fp16, quantize in place.

1417 fp8_dtype = current_platform.fp8_dtype()

1456 - w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype)

1457 - w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype)

1418 + w13_weight = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)

1419 + w2_weight = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)

1420

1421 for expert in range(layer.local_num_experts):

1422 w13_weight[expert, :, :], layer.w13_weight_scale[expert] = (

1461 - ops.scaled_fp8_quant(layer.w13_weight.data[expert, :, :])

1423 + ops.scaled_fp8_quant(layer.w13_weight[expert, :, :])

1424 )

1425 w2_weight[expert, :, :], layer.w2_weight_scale[expert] = (

1464 - ops.scaled_fp8_quant(layer.w2_weight.data[expert, :, :])

1426 + ops.scaled_fp8_quant(layer.w2_weight[expert, :, :])

1427 )

1428 replace_parameter(layer, "w13_weight", w13_weight)

1429 replace_parameter(layer, "w2_weight", w2_weight)

1430

1469 - # Reshuffle weights for AITER if needed.

1470 - if self.fp8_backend == Fp8MoeBackend.AITER:

1471 - shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

1472 - layer.w13_weight, layer.w2_weight

1473 - )

1474 - replace_parameter(layer, "w13_weight", shuffled_w13)

1475 - replace_parameter(layer, "w2_weight", shuffled_w2)

1431 + # Shuffle weights into the runtime format.

1432 + self._convert_weights_to_kernel_format(

1433 + layer, w13_weight, w2_weight, layer.w13_weight_scale, layer.w2_weight_scale

1434 + )

1435

1477 - # Rushuffle weights for MARLIN if needed.

1478 - elif self.fp8_backend == Fp8MoeBackend.MARLIN:

1479 - prepare_moe_fp8_layer_for_marlin(

1480 - layer, False, input_dtype=self.marlin_input_dtype

1481 - )

1436 + # Setup modular kernel for TP case.

1437 + self._setup_kernel(layer)

1438

1439

1440 class Fp8KVCacheMethod(BaseKVCacheMethod):

[MoE Refactor][10/N] Cleanup Fp8 Process Weights After Loading (#31169)