@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2e693f48e7

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

[Perf] Add TRTLLM FP8 MoE Modular Kernel (#36307)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

Wei Zhao committed Mar 12, 2026 at 10:32 UTC 2e693f48e7bd6fa621c8ce2c753ae76360793a04

3 files changed +230 -108

tests/kernels/moe/test_flashinfer.py

+2 -2

@@ -19,7 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import (

19 fp8_w8a8_moe_quant_config,

20 )

21 from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (

22 - TrtLlmFp8Experts,

22 + TrtLlmFp8ExpertsMonolithic,

23 )

24 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

25 FlashInferExperts,

@@ -247,7 +247,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

247 allow_new_interface=True,

248 use_monolithic=True,

249 ),

250 - TrtLlmFp8Experts(

250 + TrtLlmFp8ExpertsMonolithic(

251 moe_config=td.layer.moe,

252 quant_config=quant_config,

253 ),

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

+176 -55

@@ -4,6 +4,7 @@

4 import torch

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

7 +from vllm.logger import init_logger

8 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

9 from vllm.model_executor.layers.fused_moe.config import (

10 FusedMoEConfig,

@@ -11,6 +12,9 @@ from vllm.model_executor.layers.fused_moe.config import (

12 FusedMoEQuantConfig,

13 RoutingMethodType,

14 )

15 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

16 + TopKWeightAndReduceNoOP,

17 +)

18 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

19 activation_to_flashinfer_int,

20 )

@@ -22,10 +26,13 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

26 )

27 from vllm.platforms import current_platform

29 +logger = init_logger(__name__)

30 +

26 -class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic):

32 +class TrtLlmFp8ExpertsBase:

33 """

28 - Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface.

34 + Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic

35 + interfaces.

36 """

38 def __init__(

@@ -33,8 +40,6 @@ class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic):

40 moe_config: FusedMoEConfig,

41 quant_config: FusedMoEQuantConfig,

42 ):

36 - super().__init__(moe_config, quant_config)

37 -

43 self.routing_method_type = moe_config.routing_method

44 self.topk = moe_config.experts_per_token

45 self.intermediate_size_per_partition = (

@@ -44,6 +49,173 @@ class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic):

49 self.local_num_experts = moe_config.num_local_experts

50 self.ep_rank = moe_config.moe_parallel_config.ep_rank

52 + self.quant_config = quant_config

53 +

54 + @staticmethod

55 + def activation_format() -> mk.FusedMoEActivationFormat:

56 + return mk.FusedMoEActivationFormat.Standard

57 +

58 + @staticmethod

59 + def _supports_current_device() -> bool:

60 + """Supports only Blackwell-family GPUs."""

61 + p = current_platform

62 + # Add check flashinfer trtllm is available

63 + return p.is_cuda() and p.is_device_capability_family(100)

64 +

65 + @staticmethod

66 + def _supports_no_act_and_mul() -> bool:

67 + """Does not support non-gated MoE (i.e. Nanotron-3-Nano)."""

68 + return True

69 +

70 + @staticmethod

71 + def _supports_activation(activation: MoEActivation) -> bool:

72 + """Supports only SiLU and RELU^2 non-gated activation."""

73 + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

74 +

75 + @staticmethod

76 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

77 + """Monolithic kernel so only use with naive DP/EP and TP."""

78 + return (

79 + not moe_parallel_config.use_all2all_kernels

80 + or moe_parallel_config.use_naive_all2all_kernels

81 + ) and not moe_parallel_config.enable_eplb

82 +

83 + @staticmethod

84 + def _supports_router_logits_dtype(

85 + router_logits_dtype: torch.dtype | None,

86 + routing_method: RoutingMethodType,

87 + ) -> bool:

88 + """

89 + The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.

90 + Only DeepSeekV3 routing supports float32 router_logits (which is converted

91 + internally in the kernel).

92 + """

93 + if router_logits_dtype == torch.float32:

94 + # Only DeepSeekV3 routing handles float32 logits

95 + # https://github.com/flashinfer-ai/flashinfer/issues/2469

96 + return routing_method == RoutingMethodType.DeepSeekV3

97 + return True

98 +

99 + def supports_chunking(self) -> bool:

100 + return False

101 +

102 + def supports_expert_map(self) -> bool:

103 + return False

104 +

105 +

106 +class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):

107 + """

108 + Fp8 TRTLLM-Gen MoE kernels. Supports modular interface.

109 + """

110 +

111 + @staticmethod

112 + def _supports_quant_scheme(

113 + weight_key: QuantKey | None,

114 + activation_key: QuantKey | None,

115 + ) -> bool:

116 + """Supports Fp8 block."""

117 + SUPPORTED_W_A = [

118 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),

119 + ]

120 + return (weight_key, activation_key) in SUPPORTED_W_A

121 +

122 + def workspace_shapes(

123 + self,

124 + M: int,

125 + N: int,

126 + K: int,

127 + topk: int,

128 + global_num_experts: int,

129 + local_num_experts: int,

130 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

131 + activation: MoEActivation,

132 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

133 + # The workspaces for this implementation are managed by flashinfer.

134 + workspace1 = (0,)

135 + workspace2 = (0,)

136 + output = (M, K)

137 +

138 + return (workspace1, workspace2, output)

139 +

140 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

141 + return TopKWeightAndReduceNoOP()

142 +

143 + def apply(

144 + self,

145 + output: torch.Tensor,

146 + hidden_states: torch.Tensor,

147 + w1: torch.Tensor,

148 + w2: torch.Tensor,

149 + topk_weights: torch.Tensor,

150 + topk_ids: torch.Tensor,

151 + activation: MoEActivation,

152 + global_num_experts: int,

153 + expert_map: torch.Tensor | None,

154 + a1q_scale: torch.Tensor | None,

155 + a2_scale: torch.Tensor | None,

156 + workspace13: torch.Tensor,

157 + workspace2: torch.Tensor,

158 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

159 + apply_router_weight_on_input: bool,

160 + ):

161 + import flashinfer

162 +

163 + # Pack topk_ids and topk_weights into single tensor

164 + # Format: (expert_id << 16) | (weight_bf16.view(int16))

165 + packed_topk_ids = (topk_ids << 16) | topk_weights.to(torch.bfloat16).view(

166 + torch.int16

167 + )

168 +

169 + # trtllm_fp8_block_scale_routed_moe does not support autotuning

170 + # so skip this kernel during dummy run for autotuning.

171 + import vllm.utils.flashinfer as fi_utils

172 +

173 + if fi_utils._is_fi_autotuning:

174 + return

175 +

176 + assert a1q_scale is not None

177 +

178 + # `trtllm_fp8_block_scale_routed_moe` has a bug and does not write to the

179 + # output tensor in-place so we need to manually copy the result to the

180 + # output tensor

181 + # https://github.com/flashinfer-ai/flashinfer/issues/2703

182 + result = flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(

183 + topk_ids=packed_topk_ids,

184 + routing_bias=None,

185 + hidden_states=hidden_states,

186 + hidden_states_scale=a1q_scale.t().contiguous(), # type: ignore[union-attr]

187 + gemm1_weights=w1,

188 + gemm1_weights_scale=self.quant_config.w1_scale,

189 + gemm2_weights=w2,

190 + gemm2_weights_scale=self.quant_config.w2_scale,

191 + num_experts=global_num_experts,

192 + top_k=self.topk,

193 + n_group=None,

194 + topk_group=None,

195 + intermediate_size=self.intermediate_size_per_partition,

196 + local_expert_offset=self.ep_rank * self.local_num_experts,

197 + local_num_experts=self.local_num_experts,

198 + routed_scaling_factor=None,

199 + routing_method_type=1,

200 + use_shuffled_weight=False,

201 + weight_layout=0,

202 + # output=output,

203 + )

204 + output.copy_(result)

205 +

206 +

207 +class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolithic):

208 + """

209 + Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface.

210 + """

211 +

212 + def __init__(

213 + self,

214 + moe_config: FusedMoEConfig,

215 + quant_config: FusedMoEQuantConfig,

216 + ):

217 + super().__init__(moe_config, quant_config)

218 +

219 # Make additional scales for per-tensor interface.

220 if self.quant_config.is_per_tensor:

221 w1_scale = self.quant_config.w1_scale

@@ -63,22 +235,6 @@ class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic):

235 else torch.ones_like(self._g1_alphas) / self.quant_config.a2_scale

236 )

237

66 - @staticmethod

67 - def activation_format() -> mk.FusedMoEActivationFormat:

68 - return mk.FusedMoEActivationFormat.Standard

69 -

70 - @staticmethod

71 - def _supports_current_device() -> bool:

72 - """Supports only Blackwell-family GPUs."""

73 - p = current_platform

74 - # Add check flashinfer trtllm is available

75 - return p.is_cuda() and p.is_device_capability_family(100)

76 -

77 - @staticmethod

78 - def _supports_no_act_and_mul() -> bool:

79 - """Does not support non-gated MoE (i.e. Nanotron-3-Nano)."""

80 - return True

81 -

238 @staticmethod

239 def _supports_quant_scheme(

240 weight_key: QuantKey | None,

@@ -91,11 +247,6 @@ class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic):

247 ]

248 return (weight_key, activation_key) in SUPPORTED_W_A

249

94 - @staticmethod

95 - def _supports_activation(activation: MoEActivation) -> bool:

96 - """Supports only SiLU and RELU^2 non-gated activation."""

97 - return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

98 -

250 @staticmethod

251 def _supports_routing_method(

252 routing_method: RoutingMethodType,

@@ -123,36 +274,6 @@ class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic):

274 else:

275 raise ValueError("Unsupported quantization scheme.")

276

126 - @staticmethod

127 - def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

128 - """Monolithic kernel so only use with naive DP/EP and TP."""

129 - return (

130 - not moe_parallel_config.use_all2all_kernels

131 - or moe_parallel_config.use_naive_all2all_kernels

132 - ) and not moe_parallel_config.enable_eplb

133 -

134 - @staticmethod

135 - def _supports_router_logits_dtype(

136 - router_logits_dtype: torch.dtype | None,

137 - routing_method: RoutingMethodType,

138 - ) -> bool:

139 - """

140 - The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.

141 - Only DeepSeekV3 routing supports float32 router_logits (which is converted

142 - internally in the kernel).

143 - """

144 - if router_logits_dtype == torch.float32:

145 - # Only DeepSeekV3 routing handles float32 logits

146 - # https://github.com/flashinfer-ai/flashinfer/issues/2469

147 - return routing_method == RoutingMethodType.DeepSeekV3

148 - return True

149 -

150 - def supports_chunking(self) -> bool:

151 - return False

152 -

153 - def supports_expert_map(self) -> bool:

154 - return False

155 -

277 def _apply_per_block(

278 self,

279 hidden_states: torch.Tensor,

vllm/model_executor/layers/fused_moe/oracle/fp8.py

+52 -51

@@ -104,83 +104,84 @@ def _get_priority_backends(

104

105 def backend_to_kernel_cls(

106 backend: Fp8MoeBackend,

107 -) -> type[mk.FusedMoEExperts]:

107 +) -> list[type[mk.FusedMoEExperts]]:

108 if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

109 from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( # noqa: E501

110 - TrtLlmFp8Experts,

110 + TrtLlmFp8ExpertsModular,

111 + TrtLlmFp8ExpertsMonolithic,

112 )

113

113 - return TrtLlmFp8Experts

114 + return [TrtLlmFp8ExpertsMonolithic, TrtLlmFp8ExpertsModular]

115

116 elif backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

117 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

118 FlashInferExperts,

119 )

120

120 - return FlashInferExperts

121 + return [FlashInferExperts]

122

123 elif backend == Fp8MoeBackend.DEEPGEMM:

124 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

125 TritonOrDeepGemmExperts,

126 )

127

127 - return TritonOrDeepGemmExperts

128 + return [TritonOrDeepGemmExperts]

129

130 elif backend == Fp8MoeBackend.BATCHED_DEEPGEMM:

131 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (

132 BatchedDeepGemmExperts,

133 )

134

134 - return BatchedDeepGemmExperts

135 + return [BatchedDeepGemmExperts]

136

137 elif backend == Fp8MoeBackend.MARLIN:

138 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

139 MarlinExperts,

140 )

141

141 - return MarlinExperts

142 + return [MarlinExperts]

143

144 elif backend == Fp8MoeBackend.TRITON:

145 from vllm.model_executor.layers.fused_moe.fused_moe import (

146 TritonExperts,

147 )

148

148 - return TritonExperts

149 + return [TritonExperts]

150

151 elif backend == Fp8MoeBackend.BATCHED_TRITON:

152 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

153 BatchedTritonExperts,

154 )

155

155 - return BatchedTritonExperts

156 + return [BatchedTritonExperts]

157

158 elif backend == Fp8MoeBackend.AITER:

159 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

160 AiterExperts,

161 )

162

162 - return AiterExperts

163 + return [AiterExperts]

164

165 elif backend == Fp8MoeBackend.VLLM_CUTLASS:

166 from vllm.model_executor.layers.fused_moe.triton_cutlass_moe import (

167 TritonOrCutlassExperts,

168 )

169

169 - return TritonOrCutlassExperts

170 + return [TritonOrCutlassExperts]

171

172 elif backend == Fp8MoeBackend.BATCHED_VLLM_CUTLASS:

173 from vllm.model_executor.layers.fused_moe.cutlass_moe import (

174 CutlassBatchedExpertsFp8,

175 )

176

176 - return CutlassBatchedExpertsFp8

177 + return [CutlassBatchedExpertsFp8]

178

179 elif backend == Fp8MoeBackend.XPU:

180 from vllm.model_executor.layers.fused_moe.xpu_fused_moe import (

181 XPUExpertsFp8,

182 )

183

183 - return XPUExpertsFp8

184 + return [XPUExpertsFp8]

185

186 else:

187 raise ValueError(f"Unknown FP8 MoE backend: {backend.value}")

@@ -215,8 +216,9 @@ def select_fp8_moe_backend(

216 Select the primary FP8 MoE backend

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

218 """

219 +

220 if config.is_lora_enabled:

219 - return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON)

221 + return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON)[0]

222

223 # NOTE: the kernels are selected in the following order.

224 AVAILABLE_BACKENDS = _get_priority_backends(config, weight_key, activation_key)

@@ -256,13 +258,13 @@ def select_fp8_moe_backend(

258 activation_key: QuantKey | None,

259 activation_format: mk.FusedMoEActivationFormat,

260 ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:

259 - k_cls = backend_to_kernel_cls(backend)

260 - supported, reason = k_cls.is_supported_config(

261 - k_cls, config, weight_key, activation_key, activation_format

262 - )

263 - if supported:

264 - logger.info_once(_make_log_backend(backend), scope="local")

265 - return backend, k_cls

261 + for k_cls in backend_to_kernel_cls(backend):

262 + supported, reason = k_cls.is_supported_config(

263 + k_cls, config, weight_key, activation_key, activation_format

264 + )

265 + if supported:

266 + logger.info_once(_make_log_backend(backend), scope="local")

267 + return backend, k_cls

268 raise ValueError(_make_log_unsupported(backend, reason))

269

270 # Handle explicit moe_backend from user.

@@ -312,7 +314,7 @@ def select_fp8_moe_backend(

314 raise ValueError(

315 f"FlashInfer MOE backend {fi_backend} does not support FP8 MoE."

316 )

315 - k_cls = backend_to_kernel_cls(backend)

317 + k_cls = backend_to_kernel_cls(backend)[0]

318 return _return_or_raise(

319 backend, config, weight_key, activation_key, activation_format

320 )

@@ -322,23 +324,23 @@ def select_fp8_moe_backend(

324 Fp8MoeBackend.FLASHINFER_TRTLLM,

325 Fp8MoeBackend.FLASHINFER_CUTLASS,

326 ]:

325 - k_cls = backend_to_kernel_cls(backend)

326 - supported, reason = k_cls.is_supported_config(

327 - k_cls,

328 - config,

329 - weight_key,

330 - activation_key,

331 - activation_format,

332 - )

333 -

334 - if supported:

335 - logger.info_once(_make_log_backend(backend), scope="local")

336 - return backend, k_cls

337 - else:

338 - logger.debug_once(

339 - _make_log_unsupported(backend, reason), scope="local"

327 + for k_cls in backend_to_kernel_cls(backend):

328 + supported, reason = k_cls.is_supported_config(

329 + k_cls,

330 + config,

331 + weight_key,

332 + activation_key,

333 + activation_format,

334 )

335

336 + if supported:

337 + logger.info_once(_make_log_backend(backend), scope="local")

338 + return backend, k_cls

339 + else:

340 + logger.debug_once(

341 + _make_log_unsupported(backend, reason), scope="local"

342 + )

343 +

344 raise NotImplementedError(

345 "Found VLLM_USE_FLASHINFER_MOE_FP8=1, but no "

346 "FlashInfer FP8 MoE backend supports the configuration."

@@ -382,20 +384,19 @@ def select_fp8_moe_backend(

384

385 # Select kernels in order of backend.

386 for backend in AVAILABLE_BACKENDS:

385 - k_cls = backend_to_kernel_cls(backend)

386 - supported, reason = k_cls.is_supported_config(

387 - k_cls,

388 - config,

389 - weight_key,

390 - activation_key,

391 - activation_format,

392 - )

393 -

394 - if supported:

395 - logger.info_once(_make_log_backend(backend), scope="local")

396 - return backend, k_cls

397 - else:

398 - logger.debug_once(_make_log_unsupported(backend, reason), scope="local")

387 + for k_cls in backend_to_kernel_cls(backend):

388 + supported, reason = k_cls.is_supported_config(

389 + k_cls,

390 + config,

391 + weight_key,

392 + activation_key,

393 + activation_format,

394 + )

395 + if supported:

396 + logger.info_once(_make_log_backend(backend), scope="local")

397 + return backend, k_cls

398 + else:

399 + logger.debug_once(_make_log_unsupported(backend, reason), scope="local")

400

401 # TODO(rob): per discussion with TPU team, we need a way to register

402 # MoE backends by OOT plugins, rather than having an explicit list