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 11, 03:45 UTC
[MoE Refactor][2/N] Use Modular Kernels for Fp8 (#30825)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Dec 19, 2025 at 18:36 UTC
95befecc184778b28c5251d8a2699be8622b683f
1 file changed
+89
-49
vllm/model_executor/layers/quantization/fp8.py
+89
-49
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from enum import Enum
5
-from functools import partial
5
from typing import TYPE_CHECKING, Any, Optional
6
7
import torch
@@ -51,7 +50,6 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
50
FlashinferMoeBackend,
51
apply_flashinfer_per_tensor_scale_fp8,
52
build_flashinfer_fp8_cutlass_moe_prepare_finalize,
54
- flashinfer_cutlass_moe_fp8,
53
get_flashinfer_moe_backend,
54
register_moe_scaling_factors,
55
rotate_flashinfer_fp8_moe_weights,
@@ -728,18 +726,28 @@ class Fp8MoEMethod(FusedMoEMethodBase):
726
self.flashinfer_moe_backend = FlashinferMoeBackend.TENSORRT_LLM
727
elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
728
self.flashinfer_moe_backend = FlashinferMoeBackend.CUTLASS
731
- if self.block_quant:
732
- assert self.weight_block_size == [128, 128], (
733
- f"Only support weight_block_size == [128, 128], "
734
- f"got {self.weight_block_size}"
729
+ if self.block_quant and self.weight_block_size != [128, 128]:
730
+ raise NotImplementedError(
731
+ "FlashInfer CUTLASS FP8 MoE backend only supports block "
732
+ "size [128, 128]."
733
+ )
734
+ if not self.block_quant:
735
+ if layer.renormalize or layer.custom_routing_function is not None:
736
+ raise NotImplementedError(
737
+ "FlashInfer CUTLASS FP8 MoE backend does custom routing "
738
+ f"function or renormalization, but got {layer.renormalize} and "
739
+ f"{layer.custom_routing_function}."
740
+ )
741
+ if layer.scoring_func != "sigmoid":
742
+ raise NotImplementedError(
743
+ "FlashInfer CUTLASS FP8 MoE backend only supports "
744
+ f"'sigmoid' scoring function, but got {layer.scoring_func}."
745
+ )
746
+ if layer.activation != "silu":
747
+ raise NotImplementedError(
748
+ "FlashInfer CUTLASS FP8 MoE backend only supports SiLU "
749
+ "activation function, but got {layer.activation}."
750
)
736
- self.flashinfer_moe_fn = partial(
737
- flashinfer_cutlass_moe_fp8,
738
- moe=self.moe,
739
- use_deepseek_fp8_block_scale=self.block_quant,
740
- )
741
-
742
- self.allow_deep_gemm = self.fp8_backend == Fp8MoeBackend.DEEPGEMM
751
752
def create_weights(
753
self,
@@ -928,7 +936,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
936
937
# DeepGemm scales need to be transposed and aligned. We try to do
938
# it ahead of time for performance reasons.
931
- if self.allow_deep_gemm:
939
+ if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:
940
dg_w13_weight, dg_w13_weight_scale_inv = (
941
deepgemm_post_process_fp8_weight_block(
942
wq=layer.w13_weight.data,
@@ -1039,6 +1047,61 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1047
del layer.w13_input_scale
1048
del layer.w2_input_scale
1049
1050
+ # NOTE(rob): this is a WIP refactor. We are first migrating
1051
+ # all of the kernels in the TP case to use mk. Once this is
1052
+ # done, then we will initialzie the TP case and DP/EP case
1053
+ # via the same code path (i.e. via maybe_init_modular_kernel).
1054
+ # NOTE(rob): in progress migrating all into this format.
1055
+ if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
1056
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
1057
+ FlashInferExperts,
1058
+ )
1059
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
1060
+ FlashInferAllGatherMoEPrepareAndFinalize,
1061
+ )
1062
+
1063
+ config = self.get_fused_moe_quant_config(layer)
1064
+ assert config is not None
1065
+ self.moe_quant_config = config
1066
+
1067
+ self.kernel = mk.FusedMoEModularKernel(
1068
+ FlashInferAllGatherMoEPrepareAndFinalize(
1069
+ use_dp=(self.moe.dp_size > 1),
1070
+ use_deepseek_fp8_block_scale=self.block_quant,
1071
+ ),
1072
+ FlashInferExperts(
1073
+ out_dtype=torch.get_default_dtype(),
1074
+ quant_config=self.moe_quant_config,
1075
+ ep_rank=self.moe.ep_rank,
1076
+ ep_size=self.moe.ep_size,
1077
+ tp_rank=self.moe.tp_rank,
1078
+ tp_size=self.moe.tp_size,
1079
+ use_dp=(self.moe.dp_size > 1),
1080
+ use_deepseek_fp8_block_scale=self.block_quant,
1081
+ ),
1082
+ )
1083
+ self.use_inplace = False
1084
+
1085
+ elif self.fp8_backend in [Fp8MoeBackend.DEEPGEMM, Fp8MoeBackend.TRITON]:
1086
+ from vllm.model_executor.layers.fused_moe import (
1087
+ TritonOrDeepGemmExperts,
1088
+ )
1089
+ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
1090
+ MoEPrepareAndFinalizeNoEP,
1091
+ )
1092
+
1093
+ config = self.get_fused_moe_quant_config(layer)
1094
+ assert config is not None
1095
+ self.moe_quant_config = config
1096
+ self.kernel = mk.FusedMoEModularKernel(
1097
+ MoEPrepareAndFinalizeNoEP(),
1098
+ TritonOrDeepGemmExperts(
1099
+ quant_config=self.moe_quant_config,
1100
+ allow_deep_gemm=(self.fp8_backend == Fp8MoeBackend.DEEPGEMM),
1101
+ ),
1102
+ )
1103
+ self.use_inplace = True
1104
+
1105
def maybe_make_prepare_finalize(
1106
self,
1107
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
@@ -1091,7 +1154,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1154
assert max_num_tokens_per_rank is not None
1155
1156
experts_impl = (
1094
- BatchedDeepGemmExperts if self.allow_deep_gemm else BatchedTritonExperts
1157
+ BatchedDeepGemmExperts
1158
+ if self.fp8_backend == Fp8MoeBackend.DEEPGEMM
1159
+ else BatchedTritonExperts
1160
)
1161
logger.debug(
1162
"%s(%s): max_tokens_per_rank=%s, block_size=%s, per_act_token=%s",
@@ -1126,7 +1191,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1191
)
1192
return TritonOrDeepGemmExperts(
1193
quant_config=self.moe_quant_config,
1129
- allow_deep_gemm=self.allow_deep_gemm,
1194
+ allow_deep_gemm=(self.fp8_backend == Fp8MoeBackend.DEEPGEMM),
1195
)
1196
1197
def get_fused_moe_quant_config(
@@ -1164,6 +1229,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1229
router_logits: torch.Tensor,
1230
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1231
if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
1232
+ # TODO(rob): convert this to MK.
1233
if layer.enable_eplb:
1234
raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.")
1235
assert layer.activation == "silu", (
@@ -1228,6 +1294,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1294
rocm_aiter_fused_experts,
1295
)
1296
1297
+ # TODO(rob): convert this to MK.
1298
result = rocm_aiter_fused_experts(
1299
x,
1300
layer.w13_weight,
@@ -1240,6 +1307,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1307
quant_config=self.moe_quant_config,
1308
)
1309
elif self.use_marlin:
1310
+ # TODO(rob): convert this to MK.
1311
assert layer.activation == "silu", (
1312
f"{layer.activation} not supported for Marlin MoE."
1313
)
@@ -1261,47 +1329,19 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1329
input_dtype=self.marlin_input_dtype,
1330
workspace=layer.workspace,
1331
)
1264
- elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
1265
- assert layer.activation == "silu", (
1266
- f"Expected 'silu' activation but got {layer.activation}"
1267
- )
1268
- if not self.block_quant:
1269
- assert (
1270
- not layer.renormalize and layer.custom_routing_function is not None
1271
- )
1272
- assert layer.scoring_func == "sigmoid", (
1273
- f"Expected 'sigmoid' scoring func but got {layer.scoring_func}"
1274
- )
1275
- # Delegate to CUTLASS FlashInfer path; function already bound with
1276
- # use_deepseek_fp8_block_scale for block-quant when applicable
1277
- result = self.flashinfer_moe_fn(
1332
+ else:
1333
+ result = self.kernel(
1334
x,
1279
- layer,
1335
+ layer.w13_weight,
1336
+ layer.w2_weight,
1337
topk_weights,
1338
topk_ids,
1282
- inplace=False,
1339
+ inplace=self.use_inplace,
1340
activation=layer.activation,
1341
global_num_experts=layer.global_num_experts,
1342
expert_map=layer.expert_map,
1343
apply_router_weight_on_input=layer.apply_router_weight_on_input,
1344
)
1288
- else:
1289
- from vllm.model_executor.layers.fused_moe import fused_experts
1290
-
1291
- result = fused_experts(
1292
- hidden_states=x,
1293
- w1=layer.w13_weight,
1294
- w2=layer.w2_weight,
1295
- topk_weights=topk_weights,
1296
- topk_ids=topk_ids,
1297
- inplace=True,
1298
- activation=layer.activation,
1299
- global_num_experts=layer.global_num_experts,
1300
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1301
- expert_map=layer.expert_map,
1302
- quant_config=self.moe_quant_config,
1303
- allow_deep_gemm=self.allow_deep_gemm,
1304
- )
1345
1346
if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
1347
assert not isinstance(result, tuple), (