@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3f8a874065

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

[Kernels] Enable FlashInfer FP8 Blockscale on SM90 (for TEP DSR1) (#27134)

Signed-off-by: Duncan Moss <djm.moss@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Duncan Moss committed Nov 14, 2025 at 08:02 UTC 3f8a8740656f2c0079b9e2b1623a0758a61104af

4 files changed +179 -68

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+22 -1

@@ -57,6 +57,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

57 tp_rank: int = 0,

58 tp_size: int = 1,

59 use_dp: bool = False,

60 + use_deepseek_fp8_block_scale: bool = False,

61 ):

62 super().__init__(quant_config)

63 assert quant_config.quant_dtype in ("nvfp4", torch.float8_e4m3fn, None), (

@@ -69,6 +70,10 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

70 self.tp_size = tp_size

71 self.out_dtype = out_dtype

72 self.use_dp = use_dp

73 + # Enables DeepSeek-style FP8 block-scale path:

74 + # - pass per-block weight scales to the kernel

75 + # - skip input activation quantization (kernel applies scaling)

76 + self.use_deepseek_fp8_block_scale = use_deepseek_fp8_block_scale

78 @property

79 def activation_formats(

@@ -147,7 +152,12 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

152 "Only activation silu is supported in FlashInferExperts"

153 )

154

150 - if self.quant_dtype == torch.float8_e4m3fn:

155 + # Select quantization metadata based on FP8 format/path

156 + if (

157 + self.quant_dtype == torch.float8_e4m3fn

158 + and not self.use_deepseek_fp8_block_scale

159 + ):

160 + # FP8 per-tensor path: use global alphas/scales; do not pass input_sf

161 quant_scales = [

162 self.g1_alphas,

163 self.a2_gscale,

@@ -176,6 +186,15 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

186 # FlashInfer API requires weight to be long for nvfp4

187 fc1_expert_weights = w1.view(torch.long)

188 fc2_expert_weights = w2.view(torch.long)

189 + elif self.use_deepseek_fp8_block_scale:

190 + # FP8 block-scale path: provide block-scale weights, omit a1q_scale

191 + quant_scales = [

192 + self.w1_scale,

193 + self.w2_scale,

194 + ]

195 + a1q_scale = None

196 + fc1_expert_weights = w1

197 + fc2_expert_weights = w2

198 else:

199 quant_scales = None

200 a1q_scale = None

@@ -196,6 +215,8 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

215 ep_size=self.ep_size,

216 ep_rank=self.ep_rank,

217 output=output,

218 + # Informs FlashInfer to use the block-scale decoding path when True

219 + use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale,

220 )

221

222

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+97 -50

@@ -28,11 +28,15 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

28 self,

29 use_dp: bool,

30 num_dispatchers: int = 1,

31 + use_deepseek_fp8_block_scale: bool = False,

32 ):

33 super().__init__()

34 self.num_dispatchers_ = num_dispatchers

35 self.use_dp = use_dp

36 self.local_tokens = None

37 + # Toggle for DeepSeek-style FP8 block-scale path where activations are

38 + # not quantized here and weight block scales are consumed by the kernel.

39 + self.use_deepseek_fp8_block_scale = use_deepseek_fp8_block_scale

41 @property

42 def activation_format(self) -> mk.FusedMoEActivationFormat:

@@ -73,8 +77,9 @@ class FlashInferAllToAllMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFina

77 self,

78 use_dp: bool,

79 num_dispatchers: int = 1,

80 + use_deepseek_fp8_block_scale: bool = False,

81 ):

77 - super().__init__(use_dp, num_dispatchers)

82 + super().__init__(use_dp, num_dispatchers, use_deepseek_fp8_block_scale)

83 self.alltoall_info = None

85 # Initialize all2all_manager only for DP case

@@ -97,15 +102,19 @@ class FlashInferAllToAllMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFina

102 )

103

104 if not self.use_dp:

100 - # Non-DP case: standard quantization

101 - a1q, a1q_scale = moe_kernel_quantize_input(

102 - a1,

103 - quant_config.a1_gscale,

104 - quant_config.quant_dtype,

105 - quant_config.per_act_token_quant,

106 - quant_config.block_shape,

107 - is_fp4_scale_swizzled=not self.use_dp,

108 - )

105 + # Non-DP case: quantize activations unless using block-scale path

106 + if not self.use_deepseek_fp8_block_scale:

107 + a1q, a1q_scale = moe_kernel_quantize_input(

108 + a1,

109 + quant_config.a1_gscale,

110 + quant_config.quant_dtype,

111 + quant_config.per_act_token_quant,

112 + quant_config.block_shape,

113 + is_fp4_scale_swizzled=not self.use_dp,

114 + )

115 + else:

116 + a1q = a1

117 + a1q_scale = None

118 else:

119 # DP case: use FlashInfer AllToAll

120 global_num_tokens_cpu = get_local_sizes()

@@ -122,6 +131,7 @@ class FlashInferAllToAllMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFina

131 top_k,

132 num_experts,

133 quant_config,

134 + use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale,

135 )

136 )

137

@@ -154,8 +164,9 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin

164 self,

165 use_dp: bool,

166 num_dispatchers: int = 1,

167 + use_deepseek_fp8_block_scale: bool = False,

168 ):

158 - super().__init__(use_dp, num_dispatchers)

169 + super().__init__(use_dp, num_dispatchers, use_deepseek_fp8_block_scale)

170

171 def prepare(

172 self,

@@ -173,22 +184,42 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin

184 if not self.use_dp and quant_config.quant_dtype == "nvfp4":

185 return a1, None, None, topk_ids, topk_weights

186

176 - a1q, a1q_scale = moe_kernel_quantize_input(

177 - a1,

178 - quant_config.a1_gscale,

179 - quant_config.quant_dtype,

180 - quant_config.per_act_token_quant,

181 - quant_config.block_shape,

182 - is_fp4_scale_swizzled=not self.use_dp,

183 - )

187 + if not self.use_deepseek_fp8_block_scale:

188 + a1q, a1q_scale = moe_kernel_quantize_input(

189 + a1,

190 + quant_config.a1_gscale,

191 + quant_config.quant_dtype,

192 + quant_config.per_act_token_quant,

193 + quant_config.block_shape,

194 + is_fp4_scale_swizzled=not self.use_dp,

195 + )

196 + else:

197 + # Block-scale path: pass activations through, omit per-token scales

198 + a1q = a1

199 + a1q_scale = None

200

201 if self.use_dp:

186 - topk_weights, topk_ids, a1q, a1q_scale = get_dp_group().all_gatherv(

187 - [topk_weights, topk_ids, a1q, a1q_scale],

188 - dim=0,

189 - sizes=get_local_sizes(),

190 - )

191 - if quant_config.quant_dtype == "nvfp4":

202 + # Build gather list conditionally - omit a1q_scale if None

203 + # (block-scale path)

204 + gather_list = [topk_weights, topk_ids, a1q]

205 + if a1q_scale is not None:

206 + gather_list.append(a1q_scale)

207 + gathered = get_dp_group().all_gatherv(

208 + gather_list,

209 + dim=0,

210 + sizes=get_local_sizes(),

211 + )

212 + topk_weights, topk_ids, a1q, a1q_scale = gathered

213 + else:

214 + gathered = get_dp_group().all_gatherv(

215 + gather_list,

216 + dim=0,

217 + sizes=get_local_sizes(),

218 + )

219 + topk_weights, topk_ids, a1q = gathered

220 + a1q_scale = None

221 +

222 + if quant_config.quant_dtype == "nvfp4" and a1q_scale is not None:

223 a1q_scale = nvfp4_block_scale_interleave(a1q_scale)

224

225 return a1q, a1q_scale, None, topk_ids, topk_weights

@@ -221,6 +252,7 @@ def flashinfer_alltoall_dispatch(

252 top_k: int,

253 num_experts: int,

254 quant_config: FusedMoEQuantConfig,

255 + use_deepseek_fp8_block_scale: bool = False,

256 ):

257 from flashinfer.comm.trtllm_alltoall import MnnvlMoe

258

@@ -250,30 +282,42 @@ def flashinfer_alltoall_dispatch(

282 )

283 topk_weights = topk_weights.view(dtype=orig_topk_weights_dtype)

284

253 - x, x_sf = moe_kernel_quantize_input(

254 - x,

255 - gs,

256 - quant_config.quant_dtype,

257 - quant_config.per_act_token_quant,

258 - quant_config.block_shape,

259 - is_fp4_scale_swizzled=False, # delay swizzle to after comm

260 - )

261 - x = MnnvlMoe.mnnvl_moe_alltoallv(

262 - x,

263 - alltoall_info,

264 - all2all_manager.workspace_tensor,

265 - ep_rank,

266 - ep_size,

267 - )

285 + if not use_deepseek_fp8_block_scale:

286 + x, x_sf = moe_kernel_quantize_input(

287 + x,

288 + gs,

289 + quant_config.quant_dtype,

290 + quant_config.per_act_token_quant,

291 + quant_config.block_shape,

292 + is_fp4_scale_swizzled=False, # delay swizzle to after comm

293 + )

294 + x = MnnvlMoe.mnnvl_moe_alltoallv(

295 + x,

296 + alltoall_info,

297 + all2all_manager.workspace_tensor,

298 + ep_rank,

299 + ep_size,

300 + )

301

269 - x_sf = MnnvlMoe.mnnvl_moe_alltoallv(

270 - x_sf,

271 - alltoall_info,

272 - all2all_manager.workspace_tensor,

273 - ep_rank,

274 - ep_size,

275 - )

276 - x_sf = nvfp4_block_scale_interleave(x_sf)

302 + x_sf = MnnvlMoe.mnnvl_moe_alltoallv(

303 + x_sf,

304 + alltoall_info,

305 + all2all_manager.workspace_tensor,

306 + ep_rank,

307 + ep_size,

308 + )

309 + if quant_config.quant_dtype == "nvfp4":

310 + x_sf = nvfp4_block_scale_interleave(x_sf)

311 + else:

312 + # Block-scale path: pass activations through without quantization

313 + x_sf = None

314 + x = MnnvlMoe.mnnvl_moe_alltoallv(

315 + x,

316 + alltoall_info,

317 + all2all_manager.workspace_tensor,

318 + ep_rank,

319 + ep_size,

320 + )

321 return alltoall_info, topk_ids, topk_weights, x, x_sf

322

323

@@ -304,6 +348,7 @@ def create_flashinfer_prepare_finalize(

348 use_dp: bool,

349 use_nvfp4: bool = False,

350 enable_alltoallv: bool = False,

351 + use_deepseek_fp8_block_scale: bool = False,

352 ) -> FlashInferCutlassMoEPrepareAndFinalize:

353 """Factory function to create the appropriate FlashInfer implementation."""

354 if use_nvfp4:

@@ -311,5 +356,7 @@ def create_flashinfer_prepare_finalize(

356 return FlashInferAllToAllMoEPrepareAndFinalize(use_dp)

357 else:

358 return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)

314 - # Fp8 only supports AllGather

315 - return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)

359 + # FP8 path currently supported via AllGather; optionally enable block-scale

360 + return FlashInferAllGatherMoEPrepareAndFinalize(

361 + use_dp=use_dp, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale

362 + )

vllm/model_executor/layers/quantization/fp8.py

+36 -12

@@ -3,6 +3,7 @@

4 from collections.abc import Callable

5 from enum import Enum

6 +from functools import partial

7 from typing import TYPE_CHECKING, Any, Optional

9 import torch

@@ -122,10 +123,13 @@ def get_fp8_moe_backend(block_quant: bool) -> Fp8MoeBackend:

123 Select the primary FP8 MoE backend

124 Note: Shape-specific fallbacks may still occur at runtime.

125 """

125 - # prefer FlashInfer backends when available and enabled on supported GPUs

126 + # Prefer FlashInfer backends on supported GPUs; allow SM90 and SM100.

127 if (

128 current_platform.is_cuda()

128 - and current_platform.is_device_capability(100)

129 + and (

130 + current_platform.is_device_capability(100)

131 + or current_platform.is_device_capability(90)

132 + )

133 and envs.VLLM_USE_FLASHINFER_MOE_FP8

134 and has_flashinfer_moe()

135 ):

@@ -134,14 +138,14 @@ def get_fp8_moe_backend(block_quant: bool) -> Fp8MoeBackend:

138 logger.info_once("Using FlashInfer FP8 MoE TRTLLM backend for SM100")

139 return Fp8MoeBackend.FLASHINFER_TRTLLM

140 else:

137 - if block_quant:

141 + if block_quant and current_platform.is_device_capability(100):

142 raise ValueError(

143 "FlashInfer FP8 MoE throughput backend does not "

144 "support block quantization. Please use "

145 "VLLM_FLASHINFER_MOE_BACKEND=latency "

146 "instead."

147 )

144 - logger.info_once("Using FlashInfer FP8 MoE CUTLASS backend for SM100")

148 + logger.info_once("Using FlashInfer FP8 MoE CUTLASS backend for SM90/SM100")

149 return Fp8MoeBackend.FLASHINFER_CUTLASS

150

151 # weight-only path for older GPUs without native FP8

@@ -641,6 +645,16 @@ class Fp8MoEMethod(FusedMoEMethodBase):

645 self.flashinfer_moe_backend = FlashinferMoeBackend.TENSORRT_LLM

646 elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

647 self.flashinfer_moe_backend = FlashinferMoeBackend.CUTLASS

648 + if self.block_quant:

649 + assert self.weight_block_size == [128, 128], (

650 + f"Only support weight_block_size == [128, 128], "

651 + f"got {self.weight_block_size}"

652 + )

653 + self.flashinfer_moe_fn = partial(

654 + flashinfer_cutlass_moe_fp8,

655 + moe=self.moe,

656 + use_deepseek_fp8_block_scale=self.block_quant,

657 + )

658

659 self.allow_deep_gemm = self.fp8_backend == Fp8MoeBackend.DEEPGEMM

660 self.allow_cutlass_block_scaled_grouped_gemm = (

@@ -1012,8 +1026,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1026 ):

1027 return None

1028 elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

1029 + if self.block_quant:

1030 + assert self.weight_block_size == [128, 128], (

1031 + f"Only support weight_block_size == [128, 128], "

1032 + f"got {self.weight_block_size}"

1033 + )

1034 + # Wire block-scale flag through prepare/finalize when using CUTLASS

1035 prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

1016 - self.moe

1036 + self.moe,

1037 + use_deepseek_fp8_block_scale=self.block_quant,

1038 )

1039 logger.debug_once("%s", prepare_finalize.__class__.__name__)

1040 return prepare_finalize

@@ -1062,9 +1083,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1083 )

1084

1085 elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

1086 + # Select GEMM experts with block-scale when weights are block-quantized

1087 experts = select_cutlass_fp8_gemm_impl(

1088 self.moe,

1089 self.moe_quant_config,

1090 + use_deepseek_fp8_block_scale=self.block_quant,

1091 )

1092 logger.debug_once("Using %s", experts.__class__.__name__)

1093 return experts

@@ -1251,16 +1274,17 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1274 workspace=layer.workspace,

1275 )

1276 elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

1254 - assert not self.block_quant

1255 - assert not renormalize and custom_routing_function is not None

1277 assert activation == "silu", (

1278 f"Expected 'silu' activation but got {activation}"

1279 )

1259 - assert scoring_func == "sigmoid", (

1260 - f"Expected 'sigmoid' scoring func but got {scoring_func}"

1261 - )

1262 -

1263 - result = flashinfer_cutlass_moe_fp8(

1280 + if not self.block_quant:

1281 + assert not renormalize and custom_routing_function is not None

1282 + assert scoring_func == "sigmoid", (

1283 + f"Expected 'sigmoid' scoring func but got {scoring_func}"

1284 + )

1285 + # Delegate to CUTLASS FlashInfer path; function already bound with

1286 + # use_deepseek_fp8_block_scale for block-quant when applicable

1287 + result = self.flashinfer_moe_fn(

1288 x,

1289 layer,

1290 topk_weights,

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

+24 -5

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

17 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501

18 create_flashinfer_prepare_finalize,

19 )

20 +from vllm.platforms import current_platform

22 logger = init_logger(__name__)

@@ -190,17 +191,22 @@ def register_moe_scaling_factors(layer: torch.nn.Module) -> None:

191

192

193 def build_flashinfer_fp8_cutlass_moe_prepare_finalize(

193 - moe: FusedMoEConfig | None,

194 + moe: FusedMoEConfig | None, use_deepseek_fp8_block_scale: bool = False

195 ) -> mk.FusedMoEPrepareAndFinalize:

196 """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""

197 use_dp = moe.moe_parallel_config.dp_size > 1 if moe is not None else False

197 - return create_flashinfer_prepare_finalize(use_dp)

198 + # Propagate block-scale flag so prepare/finalize can skip act quantization

199 + # and inform the kernel to consume per-block weight scales.

200 + return create_flashinfer_prepare_finalize(

201 + use_dp, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale

202 + )

203

204

205 def select_cutlass_fp8_gemm_impl(

206 moe: FusedMoEConfig | None,

207 quant_config: FusedMoEQuantConfig,

208 out_dtype: torch.dtype | None = None,

209 + use_deepseek_fp8_block_scale: bool = False,

210 ) -> mk.FusedMoEPermuteExpertsUnpermute:

211 """Return a GEMM *experts* implementation for fused-MoE layers"""

212

@@ -212,12 +218,14 @@ def select_cutlass_fp8_gemm_impl(

218 ep_size=moe.moe_parallel_config.ep_size,

219 tp_rank=moe.moe_parallel_config.tp_rank,

220 tp_size=moe.moe_parallel_config.tp_size,

221 + use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,

222 )

223

224 assert out_dtype is not None, "If moe config is None, out_dtype must be passed"

225 return FlashInferExperts(

226 out_dtype=out_dtype,

227 quant_config=quant_config,

228 + use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,

229 )

230

231

@@ -231,14 +239,22 @@ def flashinfer_cutlass_moe_fp8(

239 global_num_experts: int = -1,

240 expert_map: torch.Tensor | None = None,

241 apply_router_weight_on_input: bool = False,

242 + use_deepseek_fp8_block_scale: bool = False,

243 + moe: FusedMoEConfig | None = None,

244 ) -> torch.Tensor:

245 quant_config = layer.quant_method.get_fused_moe_quant_config(layer)

246 assert quant_config is not None

247

248 + # Construct modular kernel with block-scale support when requested.

249 fused_experts = mk.FusedMoEModularKernel(

239 - build_flashinfer_fp8_cutlass_moe_prepare_finalize(moe=None),

250 + build_flashinfer_fp8_cutlass_moe_prepare_finalize(

251 + moe=moe, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale

252 + ),

253 select_cutlass_fp8_gemm_impl(

241 - moe=None, quant_config=quant_config, out_dtype=hidden_states.dtype

254 + moe=moe,

255 + quant_config=quant_config,

256 + out_dtype=hidden_states.dtype,

257 + use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,

258 ),

259 )

260

@@ -258,7 +274,10 @@ def flashinfer_cutlass_moe_fp8(

274

275 def get_flashinfer_moe_backend() -> FlashinferMoeBackend:

276 flashinfer_moe_backend = envs.VLLM_FLASHINFER_MOE_BACKEND

261 - if flashinfer_moe_backend == "throughput":

277 + # Prefer CUTLASS on SM90 to cover both SM90/SM100 generations

278 + if flashinfer_moe_backend == "throughput" or current_platform.is_device_capability(

279 + 90

280 + ):

281 return FlashinferMoeBackend.CUTLASS

282 elif flashinfer_moe_backend == "latency":

283 return FlashinferMoeBackend.TENSORRT_LLM