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

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 13, 08:15 UTC

[Performance][B200] Fix deepgemm prologue (#27897)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Varun Sundar Rabindranath committed Nov 12, 2025 at 16:13 UTC 74a9a9faad062e6fe90f7d238b5fad7e2f3da237

6 files changed +163 -48

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

@@ -232,6 +232,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

232 """

233 super().__init__(quant_config)

234 assert self.block_shape == get_mk_alignment_for_contiguous_layout()

235 + assert self.quant_config.use_fp8_w8a8

236 self.max_num_tokens = max_num_tokens

237 self.num_dispatchers = num_dispatchers

238

@@ -250,6 +251,12 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

251 def supports_expert_map(self) -> bool:

252 return False

253

254 + def supports_packed_ue8m0_act_scales(self) -> bool:

255 + """

256 + DeepGemm supports packed ue8m0 activation scales format in devices == sm100

257 + """

258 + return current_platform.is_device_capability(100)

259 +

260 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

261 # Let PrepareAndFinalize::finalize() decide the impl.

262 return TopKWeightAndReduceDelegate()

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+29

@@ -6,6 +6,7 @@ import deep_ep

6 import torch

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

9 +from vllm.logger import init_logger

10 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

11 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

12 TopKWeightAndReduceDelegate,

@@ -20,6 +21,8 @@ from vllm.v1.worker.ubatching import (

21 dbo_maybe_run_recv_hook,

22 )

24 +logger = init_logger(__name__)

25 +

26 # DeepEP kernels quantize dispatch inputs in 128 element chunks.

27 DEEPEP_QUANT_BLOCK_SIZE = 128

28 DEEPEP_QUANT_BLOCK_SHAPE = [DEEPEP_QUANT_BLOCK_SIZE, DEEPEP_QUANT_BLOCK_SIZE]

@@ -94,6 +97,29 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

97 self.handles: list[tuple | None] = [None, None]

98 self.num_dispatchers_ = num_dispatchers

100 + # We don't have enough information to determine if we should dispatch

101 + # activation scales in a packed ue8m0 format during object construction

102 + # time. This setting is handled by post_init_setup.

103 + self.use_ue8m0_dispatch = False

104 +

105 + def post_init_setup(self, fused_experts: mk.FusedMoEPermuteExpertsUnpermute):

106 + if not fused_experts.supports_packed_ue8m0_act_scales():

107 + # Early exit.

108 + return

109 +

110 + if self.use_fp8_dispatch:

111 + logger.debug_once(

112 + "Update DeepEPLLPrepareFinalize to do packed ue8m0 scales dispatch."

113 + )

114 + self.use_ue8m0_dispatch = True

115 + else:

116 + logger.warning_once(

117 + "DeepEPLLPrepareAndFinalize is setup to dispatch raw/unquantized "

118 + f"activations despite ({fused_experts.__class__.__name__}) being able "

119 + "to support quantized activations.",

120 + scope="local",

121 + )

122 +

123 def num_dispatchers(self) -> int:

124 return self.num_dispatchers_

125

@@ -206,6 +232,9 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

232 self.max_tokens_per_rank,

233 num_experts,

234 use_fp8=self.use_fp8_dispatch,

235 + # round_scale needs to be set to dispatch in ue8m0

236 + round_scale=self.use_ue8m0_dispatch,

237 + use_ue8m0=self.use_ue8m0_dispatch,

238 async_finish=False,

239 return_recv_hook=True,

240 )

vllm/model_executor/layers/fused_moe/modular_kernel.py

+25

@@ -149,6 +149,15 @@ class FusedMoEPrepareAndFinalize(ABC):

149 described above.

150 """

151

152 + def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"):

153 + """

154 + Initialize FusedMoEPrepareAndFinalize settings that depend on

155 + FusedMoEPermuteExpertsUnpermute experts object.

156 + The FusedMoEPrepareAndFinalize implementations that have such

157 + dependencies may choose to override this function.

158 + """

159 + return

160 +

161 @abstractmethod

162 def prepare(

163 self,

@@ -503,6 +512,13 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

512 """

513 raise NotImplementedError

514

515 + def supports_packed_ue8m0_act_scales(self) -> bool:

516 + """

517 + A flag indicating whether or not this class can process packed ue8m0

518 + activation scales.

519 + """

520 + return False

521 +

522 def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:

523 """

524 Workspace type: The dtype to use for the workspace tensors.

@@ -698,6 +714,8 @@ class FusedMoEModularKernel(torch.nn.Module):

714 self.prepare_finalize = prepare_finalize

715 self.fused_experts = fused_experts

716 self.shared_experts = shared_experts

717 +

718 + self._post_init_setup()

719 assert (

720 prepare_finalize.activation_format == fused_experts.activation_formats[0]

721 ), (

@@ -707,6 +725,13 @@ class FusedMoEModularKernel(torch.nn.Module):

725 f"{fused_experts.activation_formats[0]}"

726 )

727

728 + def _post_init_setup(self):

729 + """

730 + Resolve any leftover setup dependencies between self.prepare_finalize

731 + and self.fused_experts here.

732 + """

733 + self.prepare_finalize.post_init_setup(self.fused_experts)

734 +

735 def supports_expert_map(self) -> bool:

736 """

737 A flag indicating whether or not this class supports expert maps.

vllm/model_executor/layers/quantization/fp8.py

+34 -40

@@ -60,11 +60,10 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

60 create_fp8_input_scale,

61 create_fp8_scale_parameter,

62 create_fp8_weight_parameter,

63 - expert_weight_is_col_major,

63 + deepgemm_post_process_fp8_weight_block,

64 maybe_post_process_fp8_weight_block,

65 process_fp8_weight_block_strategy,

66 process_fp8_weight_tensor_strategy,

67 - requant_weight_ue8m0_inplace,

67 validate_fp8_block_shape,

68 )

69 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (

@@ -94,7 +93,6 @@ from vllm.model_executor.utils import set_weight_attrs

93 from vllm.platforms import current_platform

94 from vllm.scalar_type import scalar_types

95 from vllm.utils.deep_gemm import (

97 - get_col_major_tma_aligned_tensor,

96 is_deep_gemm_e8m0_used,

97 is_deep_gemm_supported,

98 )

@@ -846,15 +844,31 @@ class Fp8MoEMethod(FusedMoEMethodBase):

844

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

846 # it ahead of time for performance reasons.

849 - if self.allow_deep_gemm and not is_deep_gemm_e8m0_used():

850 - if expert_weight_is_col_major(layer.w13_weight_scale_inv):

851 - layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(

852 - layer.w13_weight_scale_inv

847 + if self.allow_deep_gemm:

848 + dg_w13_weight, dg_w13_weight_scale_inv = (

849 + deepgemm_post_process_fp8_weight_block(

850 + wq=layer.w13_weight.data,

851 + ws=layer.w13_weight_scale_inv.data,

852 + quant_block_shape=tuple(layer.weight_block_size),

853 + use_e8m0=is_deep_gemm_e8m0_used(),

854 )

854 - if expert_weight_is_col_major(layer.w2_weight_scale_inv):

855 - layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(

856 - layer.w2_weight_scale_inv

855 + )

856 + dg_w2_weight, dg_w2_weight_scale_inv = (

857 + deepgemm_post_process_fp8_weight_block(

858 + wq=layer.w2_weight.data,

859 + ws=layer.w2_weight_scale_inv.data,

860 + quant_block_shape=tuple(layer.weight_block_size),

861 + use_e8m0=is_deep_gemm_e8m0_used(),

862 )

863 + )

864 + layer.w13_weight = Parameter(dg_w13_weight, requires_grad=False)

865 + layer.w13_weight_scale_inv = Parameter(

866 + dg_w13_weight_scale_inv, requires_grad=False

867 + )

868 + layer.w2_weight = Parameter(dg_w2_weight, requires_grad=False)

869 + layer.w2_weight_scale_inv = Parameter(

870 + dg_w2_weight_scale_inv, requires_grad=False

871 + )

872

873 # If checkpoint is fp16, quantize in place.

874 elif not self.quant_config.is_checkpoint_fp8_serialized:

@@ -990,31 +1004,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1004 del layer.w13_input_scale

1005 del layer.w2_input_scale

1006

993 - if is_deep_gemm_e8m0_used() and self.block_quant:

994 - assert layer.weight_block_size is not None

995 - # Re-quantise the expert weights so their scales are UE8M0.

996 - block_sz = tuple(layer.weight_block_size)

997 - requant_weight_ue8m0_inplace(

998 - layer.w13_weight.data,

999 - layer.w13_weight_scale_inv.data,

1000 - block_sz,

1001 - )

1002 - requant_weight_ue8m0_inplace(

1003 - layer.w2_weight.data,

1004 - layer.w2_weight_scale_inv.data,

1005 - block_sz,

1006 - )

1007 -

1008 - # Ensure column-major TMA alignment expected by DeepGEMM.

1009 - if expert_weight_is_col_major(layer.w13_weight_scale_inv):

1010 - layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(

1011 - layer.w13_weight_scale_inv

1012 - )

1013 - if expert_weight_is_col_major(layer.w2_weight_scale_inv):

1014 - layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(

1015 - layer.w2_weight_scale_inv

1016 - )

1017 -

1007 def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

1008 if (

1009 self.rocm_aiter_moe_enabled

@@ -1037,7 +1026,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1026 layer: torch.nn.Module,

1027 ) -> FusedMoEPermuteExpertsUnpermute:

1028 from vllm.model_executor.layers.fused_moe import (

1040 - BatchedTritonOrDeepGemmExperts,

1029 + BatchedDeepGemmExperts,

1030 + BatchedTritonExperts,

1031 TritonOrDeepGemmExperts,

1032 )

1033

@@ -1053,20 +1043,24 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1043 ):

1044 max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()

1045 assert max_num_tokens_per_rank is not None

1046 +

1047 + experts_impl = (

1048 + BatchedDeepGemmExperts if self.allow_deep_gemm else BatchedTritonExperts

1049 + )

1050 logger.debug(

1057 - "BatchedTritonOrDeepGemmExperts(%s): "

1058 - "max_tokens_per_rank=%s, block_size=%s, per_act_token=%s",

1051 + "%s(%s): max_tokens_per_rank=%s, block_size=%s, per_act_token=%s",

1052 + experts_impl.__name__,

1053 self.__class__.__name__,

1054 max_num_tokens_per_rank,

1055 self.weight_block_size,

1056 False,

1057 )

1064 - return BatchedTritonOrDeepGemmExperts(

1058 + return experts_impl(

1059 max_num_tokens=max_num_tokens_per_rank,

1060 num_dispatchers=prepare_finalize.num_dispatchers(),

1061 quant_config=self.moe_quant_config,

1068 - allow_deep_gemm=self.allow_deep_gemm,

1062 )

1063 +

1064 elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

1065 experts = select_cutlass_fp8_gemm_impl(

1066 self.moe,

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

+53 -4

@@ -34,6 +34,7 @@ from vllm.utils.deep_gemm import (

34 is_deep_gemm_e8m0_used,

35 is_deep_gemm_supported,

36 should_use_deepgemm_for_fp8_linear,

37 + transform_sf_into_required_layout,

38 )

39 from vllm.utils.torch_utils import direct_register_custom_op

@@ -929,6 +930,50 @@ def requant_weight_ue8m0_inplace(

930 s_old.copy_(s_requant)

931

932

933 +def deepgemm_post_process_fp8_weight_block(

934 + wq: torch.Tensor, ws: torch.Tensor, quant_block_shape: tuple[int], use_e8m0: bool

935 +) -> tuple[torch.Tensor, torch.Tensor]:

936 + assert wq.dtype == torch.float8_e4m3fn, (

937 + "Expected quantized tensor dtype "

938 + f"to be torch.float8_e4m3fn, got {wq.dtype} instead."

939 + )

940 + assert ws.dtype == torch.float32, (

941 + f"Expected tensor scales dtype to be torch.float32, got {ws.dtype} instead"

942 + )

943 +

944 + if use_e8m0:

945 + requant_weight_ue8m0_inplace(wq, ws, block_size=quant_block_shape)

946 +

947 + original_ndim = wq.ndim

948 + if wq.ndim == 2:

949 + assert ws.ndim == 2

950 + wq = wq.unsqueeze(0)

951 + ws = ws.unsqueeze(0)

952 +

953 + # From https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/utils/layout.hpp#L46

954 + recipe = (1, 128, 128)

955 +

956 + # Ref : https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/apis/gemm.hpp

957 + # DeepGemm uses the `transform_sf_into_required_layout` function to

958 + # represent scales in the correct format.

959 + dg_ws = transform_sf_into_required_layout(

960 + sf=ws,

961 + mn=wq.size(1),

962 + k=wq.size(2),

963 + recipe=recipe,

964 + num_groups=wq.size(0),

965 + # is the scale factors for A in (Refers to the argument A in A @ B).

966 + # Weights are B.

967 + is_sfa=False,

968 + )

969 +

970 + if original_ndim == 2:

971 + wq = wq.squeeze(0)

972 + dg_ws = dg_ws.squeeze(0)

973 +

974 + return wq, dg_ws

975 +

976 +

977 def _maybe_pad_fp8_weight(weight: torch.Tensor) -> torch.Tensor:

978 """Pad the weight tensor. This is an optimization on ROCm platform, which

979 can benefit from tensors located far enough from one another in memory"""

@@ -1141,11 +1186,15 @@ def maybe_post_process_fp8_weight_block(layer: torch.nn.Module):

1186 should_use_deepgemm = should_use_deepgemm_for_fp8_linear(

1187 layer.orig_dtype, layer.weight

1188 )

1144 - if is_deep_gemm_e8m0_used() and should_use_deepgemm:

1145 - block_sz = tuple(layer.weight_block_size)

1146 - requant_weight_ue8m0_inplace(

1147 - layer.weight.data, layer.weight_scale.data, block_sz

1189 + if should_use_deepgemm:

1190 + dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block(

1191 + wq=layer.weight.data,

1192 + ws=layer.weight_scale.data,

1193 + quant_block_shape=tuple(layer.weight_block_size),

1194 + use_e8m0=is_deep_gemm_e8m0_used(),

1195 )

1196 + layer.weight = torch.nn.Parameter(dg_weight, requires_grad=False)

1197 + layer.weight_scale = torch.nn.Parameter(dg_weight_scale, requires_grad=False)

1198

1199

1200 def expert_weight_is_col_major(x: torch.Tensor) -> bool:

vllm/utils/deep_gemm.py

+15 -4

@@ -49,10 +49,6 @@ def is_deep_gemm_e8m0_used() -> bool:

49 logger.info_once("DeepGEMM E8M0 disabled: _fp8_gemm_nt_impl not found")

50 return False

52 - if envs.VLLM_USE_FLASHINFER_MOE_FP8:

53 - logger.info_once("DeepGEMM E8M0 disabled: FlashInfer MOE is enabled.")

54 - return False

55 -

52 if envs.VLLM_USE_DEEP_GEMM_E8M0:

53 logger.info_once("DeepGEMM E8M0 enabled on current platform.")

54 return True

@@ -77,6 +73,7 @@ _fp8_paged_mqa_logits_impl: Callable[..., Any] | None = None

73 _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None

74 _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None

75 _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None

76 +_transform_sf_into_required_layout_impl: Callable[..., Any] | None = None

79 def _lazy_init() -> None:

@@ -86,6 +83,7 @@ def _lazy_init() -> None:

83 global _get_paged_mqa_logits_metadata_impl

84 global _get_mn_major_tma_aligned_tensor_impl

85 global _get_mk_alignment_for_contiguous_layout_impl

86 + global _transform_sf_into_required_layout_impl

87 # fast path

88 if (

89 _fp8_gemm_nt_impl is not None

@@ -95,6 +93,7 @@ def _lazy_init() -> None:

93 or _fp8_paged_mqa_logits_impl is not None

94 or _get_paged_mqa_logits_metadata_impl is not None

95 or _get_mk_alignment_for_contiguous_layout_impl is not None

96 + or _transform_sf_into_required_layout_impl is not None

97 ):

98 return

@@ -124,6 +123,9 @@ def _lazy_init() -> None:

123 _get_mk_alignment_for_contiguous_layout_impl = getattr(

124 _dg, "get_mk_alignment_for_contiguous_layout", None

125 )

126 + _transform_sf_into_required_layout_impl = getattr(

127 + _dg, "transform_sf_into_required_layout", None

128 + )

129

130

131 def get_num_sms() -> int:

@@ -179,6 +181,15 @@ def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):

181 )

182

183

184 +def transform_sf_into_required_layout(*args, **kwargs):

185 + _lazy_init()

186 + if _transform_sf_into_required_layout_impl is None:

187 + return _missing(*args, **kwargs)

188 + return _transform_sf_into_required_layout_impl(

189 + *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs

190 + )

191 +

192 +

193 def fp8_mqa_logits(

194 q: torch.Tensor,

195 kv: tuple[torch.Tensor, torch.Tensor],