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

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

Support Tensorrt-LLM MoE fp4 for low-latency (#21331)

Signed-off-by: Shu Wang <shuw@nvidia.com> Signed-off-by: Po-Han Huang <pohanh@nvidia.com> Signed-off-by: Shu Wang. <shuw@nvidia.com> Signed-off-by: XIn Li <xinli@nvidia.com> Co-authored-by: XIn Li <xinli@nvidia.com>

Shu Wang committed Aug 7, 2025 at 21:18 UTC a3b9c17b56d09a091e210222d8e1f75cabe65b84

7 files changed +288 -43

vllm/envs.py

+15

@@ -129,6 +129,7 @@ if TYPE_CHECKING:

129 VLLM_SKIP_DEEP_GEMM_WARMUP: bool = False

130 VLLM_USE_FLASHINFER_MOE_FP8: bool = False

131 VLLM_USE_FLASHINFER_MOE_FP4: bool = False

132 + VLLM_FLASHINFER_MOE_BACKEND: str = "throughput"

133 VLLM_XGRAMMAR_CACHE_MB: int = 0

134 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256

135 VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False

@@ -982,6 +983,20 @@ environment_variables: dict[str, Callable[[], Any]] = {

983 "VLLM_ALL2ALL_BACKEND":

984 lambda: os.getenv("VLLM_ALL2ALL_BACKEND", "naive"),

985

986 + # Flashinfer MoE backend for vLLM's fused Mixture-of-Experts support. Both

987 + # require compute capability 10.0 or above.

988 + # Available options:

989 + # - "throughput": [default]

990 + # Uses CUTLASS kernels optimized for high-throughput batch inference.

991 + # - "latency":

992 + # Uses TensorRT-LLM kernels optimized for low-latency inference.

993 + # To set this backend, define the environment variable:

994 + # export VLLM_FLASHINFER_MOE_BACKEND=latency.

995 + # If not set, defaults to "throughput".

996 + "VLLM_FLASHINFER_MOE_BACKEND": lambda: os.getenv(

997 + "VLLM_FLASHINFER_MOE_BACKEND", "throughput"

998 + ),

999 +

1000 # Control the maximum number of tokens per expert supported by the

1001 # NVFP4 MoE CUTLASS Kernel. This value is used to create a buffer for

1002 # the blockscale tensor of activations NVFP4 Quantization.

vllm/model_executor/layers/fused_moe/config.py

+2 -1

@@ -192,7 +192,8 @@ class FusedMoEParallelConfig:

192 @property

193 def use_flashinfer_cutlass_kernels(self):

194 return (envs.VLLM_USE_FLASHINFER_MOE_FP4

195 - and has_flashinfer_cutlass_fused_moe())

195 + and has_flashinfer_cutlass_fused_moe()

196 + and envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput")

197

198 @staticmethod

199 def make(tp_size_: int, dp_size_: int,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+4 -5

@@ -105,7 +105,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

105 detect_nvfp4_moe_support)

106 _nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)

107 self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported

108 - self.allow_flashinfer_cutlass = _nvfp4.allow_flashinfer_cutlass

108 + self.allow_flashinfer = _nvfp4.allow_flashinfer

109 self.use_marlin = _nvfp4.use_marlin

110 self.group_size = 16

111 self.fused_experts = None # type: ignore[assignment]

@@ -212,7 +212,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

212 requires_grad=False)

213

214 # reorder GEMM1 weights and block scales for FlashInfer CUTLASS kernel.

215 - if self.allow_flashinfer_cutlass:

215 + if self.allow_flashinfer:

216 w, s = reorder_w1w3_to_w3w1(layer.w13_weight.data,

217 layer.w13_weight_scale.data,

218 dim=-2)

@@ -266,7 +266,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

266 (layer.w2_input_global_scale), requires_grad=False)

267

268 def maybe_swap_experts_impl(self, moe_parallel_config):

269 - if not self.allow_flashinfer_cutlass:

269 + if not self.allow_flashinfer:

270 return

271 self.fused_experts = build_flashinfer_fp4_cutlass_moe_kernel(

272 moe_parallel_config)

@@ -277,8 +277,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

277 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( # noqa: E501

278 select_nvfp4_gemm_impl)

279

280 - return select_nvfp4_gemm_impl(self.allow_flashinfer_cutlass, moe,

281 - logger)

280 + return select_nvfp4_gemm_impl(self.allow_flashinfer, moe, logger)

281

282 def apply(

283 self,

vllm/model_executor/layers/quantization/modelopt.py

+252 -32

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +from enum import Enum

5 from typing import Any, Callable, Optional, Union

7 import torch

@@ -36,6 +37,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

37 from vllm.model_executor.parameter import (ModelWeightParameter,

38 PerTensorScaleParameter)

39 from vllm.scalar_type import scalar_types

40 +from vllm.utils import next_power_of_2

41 from vllm.utils.flashinfer import has_flashinfer_moe

43 logger = init_logger(__name__)

@@ -44,6 +46,11 @@ QUANT_ALGOS = ["FP8", "NVFP4"]

46 KV_CACHE_QUANT_ALGOS = ["FP8"]

49 +class FlashinferMoeBackend(Enum):

50 + TENSORRT_LLM = "TensorRT-LLM"

51 + CUTLASS = "CUTLASS"

52 +

53 +

54 class ModelOptFp8Config(QuantizationConfig):

55 """Config class for ModelOpt FP8."""

@@ -185,7 +192,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):

192 Args: quant_config: The ModelOpt quantization config.

193 """

194

188 - def __init__(self, quant_config: ModelOptFp8Config):

195 + def __init__(self, quant_config: ModelOptFp8Config) -> None:

196 self.quant_config = quant_config

197 self.fp8_linear = Fp8LinearOp(

198 act_quant_static=True, act_quant_group_shape=GroupShape.PER_TENSOR)

@@ -265,7 +272,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

272 quant_config: The ModelOpt quantization config.

273 """

274

268 - def __init__(self, quant_config: ModelOptFp8Config):

275 + def __init__(self, quant_config: ModelOptFp8Config) -> None:

276 self.quant_config = quant_config

277 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

278 cutlass_fp8_supported)

@@ -670,7 +677,8 @@ class ModelOptNvFp4Config(QuantizationConfig):

677 return cls(is_checkpoint_nvfp4_serialized, kv_cache_quant_algo,

678 exclude_modules, group_size)

679

673 - def is_layer_excluded(self, prefix: str, exclude_modules: list):

680 + def is_layer_excluded(self, prefix: str,

681 + exclude_modules: list[str]) -> bool:

682 import regex as re

683 for pattern in exclude_modules:

684 regex_str = pattern.replace('.', r'\.').replace('*', r'.*')

@@ -714,7 +722,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

722 Args: quant_config: The ModelOpt quantization config.

723 """

724

717 - def __init__(self, quant_config: ModelOptNvFp4Config):

725 + def __init__(self, quant_config: ModelOptNvFp4Config) -> None:

726 self.quant_config = quant_config

727 self.cutlass_nvfp4_supported = cutlass_fp4_supported()

728 self.use_marlin = False

@@ -859,6 +867,16 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

867 return out.view(*output_shape)

868

869

870 +def _get_tile_tokens_dim(num_tokens: int, top_k: int, num_experts: int) -> int:

871 + # Guess tokens per expert assuming perfect expert distribution first.

872 + num_tokens_per_expert = (num_tokens * top_k) // num_experts

873 + # And pad the number to the next power of 2.

874 + tile_tokens_dim = next_power_of_2(num_tokens_per_expert)

875 + # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel.

876 + tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)

877 + return tile_tokens_dim

878 +

879 +

880 class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

881 """

882 MoE Method for FP4 Quantization.

@@ -866,22 +884,40 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

884 quant_config: NVFP4 Quant Config

885 """

886

869 - def __init__(self, quant_config: ModelOptNvFp4Config):

887 + def __init__(self, quant_config: ModelOptNvFp4Config) -> None:

888 self.quant_config = quant_config

889 from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501

890 detect_nvfp4_moe_support)

891 _nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)

892 self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported

875 - self.allow_flashinfer_cutlass = _nvfp4.allow_flashinfer_cutlass

893 + self.allow_flashinfer = _nvfp4.allow_flashinfer

894 self.use_marlin = _nvfp4.use_marlin

895 + self.flashinfer_moe_backend = None

896 +

897 + if self.allow_flashinfer:

898 + flashinfer_moe_backend = envs.VLLM_FLASHINFER_MOE_BACKEND

899 + if flashinfer_moe_backend == "throughput":

900 + self.flashinfer_moe_backend = FlashinferMoeBackend.CUTLASS

901 + logger.info_once("Using FlashInfer CUTLASS kernels for "

902 + "ModelOptNvFp4FusedMoE.")

903 + elif flashinfer_moe_backend == "latency":

904 + self.flashinfer_moe_backend = FlashinferMoeBackend.TENSORRT_LLM

905 + logger.info_once("Using FlashInfer TensorRT-LLM kernels for "

906 + "ModelOptNvFp4FusedMoE.")

907 + else:

908 + allowed_backends = ["throughput", "latency"]

909 + raise ValueError(

910 + f"Unknown flashinfer moe backend: {flashinfer_moe_backend}"

911 + f" expected one of {allowed_backends}")

912

878 - self.fused_experts = None # type: ignore

913 + self.fused_experts: Optional[

914 + mk.FusedMoEModularKernel] = None # type: ignore[assignment]

915

916 def maybe_swap_experts_impl(

917 self,

918 moe_parallel_config: FusedMoEParallelConfig,

919 ):

884 - if not self.allow_flashinfer_cutlass:

920 + if not self.allow_flashinfer:

921 return

922 self.fused_experts = build_flashinfer_fp4_cutlass_moe_kernel(

923 moe_parallel_config)

@@ -897,8 +933,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

933 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( # noqa: E501

934 select_nvfp4_gemm_impl)

935

900 - return select_nvfp4_gemm_impl(self.allow_flashinfer_cutlass, moe,

901 - logger)

936 + return select_nvfp4_gemm_impl(self.allow_flashinfer, moe, logger)

937

938 def uses_weight_scale_2_pattern(self) -> bool:

939 """

@@ -996,14 +1031,101 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1031 weight_loader=weight_loader)

1032 layer.register_parameter("w2_input_scale", w2_input_scale)

1033

1034 + def prepare_static_weight_layouts_for_trtllm_moe(

1035 + self,

1036 + gemm1_weights: torch.Tensor,

1037 + gemm2_weights: torch.Tensor,

1038 + gemm1_scales_linear_fp4_bytes: torch.Tensor,

1039 + gemm2_scales_linear_fp4_bytes: torch.Tensor,

1040 + hidden_size: int,

1041 + intermediate_size: int,

1042 + num_experts: int,

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

1044 + """Prepare quantized weights for kernel (done offline with weights)."""

1045 + from flashinfer import (reorder_rows_for_gated_act_gemm,

1046 + shuffle_matrix_a, shuffle_matrix_sf_a)

1047 + epilogue_tile_m = 128 # FIXME: this depends on the kernel internals

1048 +

1049 + # Convert quantized weights to proper formats

1050 + gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(

1051 + num_experts, 2 * intermediate_size, hidden_size // 2) # packed fp4

1052 + gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(

1053 + torch.float8_e4m3fn).reshape(num_experts, 2 * intermediate_size,

1054 + hidden_size //

1055 + 16) # fp8 scaling factors

1056 +

1057 + gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(

1058 + num_experts, hidden_size, intermediate_size // 2) # packed fp4

1059 + gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view(

1060 + torch.float8_e4m3fn).reshape(num_experts, hidden_size,

1061 + intermediate_size //

1062 + 16) # fp8 scaling factors

1063 +

1064 + # Reorder rows of W1 and scales for fused gated activation

1065 + gemm1_weights_fp4_interleaved = []

1066 + gemm1_scales_fp4_interleaved = []

1067 + for i in range(num_experts):

1068 + gemm1_weights_fp4_interleaved.append(

1069 + reorder_rows_for_gated_act_gemm(gemm1_weights_fp4[i].clone()))

1070 + gemm1_scales_fp4_interleaved.append(

1071 + reorder_rows_for_gated_act_gemm(

1072 + gemm1_scales_linear_fp4[i].clone()))

1073 +

1074 + # Stack weights and scales for all experts

1075 + gemm1_weights_fp4_interleaved = torch.stack(

1076 + gemm1_weights_fp4_interleaved).reshape(num_experts,

1077 + 2 * intermediate_size,

1078 + hidden_size // 2)

1079 + gemm1_scales_fp4_interleaved = torch.stack(

1080 + gemm1_scales_fp4_interleaved).reshape(num_experts,

1081 + 2 * intermediate_size,

1082 + hidden_size // 16)

1083 +

1084 + # Shuffle weights and scaling factors for transposed mma output

1085 + gemm1_weights_fp4_shuffled = []

1086 + gemm1_scales_fp4_shuffled = []

1087 + gemm2_weights_fp4_shuffled = []

1088 + gemm2_scales_fp4_shuffled = []

1089 + for i in range(num_experts):

1090 + gemm1_weights_fp4_shuffled.append(

1091 + shuffle_matrix_a(

1092 + gemm1_weights_fp4_interleaved[i].view(torch.uint8),

1093 + epilogue_tile_m))

1094 + gemm1_scales_fp4_shuffled.append(

1095 + shuffle_matrix_sf_a(

1096 + gemm1_scales_fp4_interleaved[i].view(torch.uint8),

1097 + epilogue_tile_m))

1098 +

1099 + gemm2_weights_fp4_shuffled.append(

1100 + shuffle_matrix_a(gemm2_weights_fp4[i].view(torch.uint8),

1101 + epilogue_tile_m))

1102 + gemm2_scales_fp4_shuffled.append(

1103 + shuffle_matrix_sf_a(

1104 + gemm2_scales_linear_fp4[i].view(torch.uint8),

1105 + epilogue_tile_m))

1106 +

1107 + # Stack weights for all experts

1108 + gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)

1109 + gemm1_scales_fp4_shuffled = (

1110 + torch.stack(gemm1_scales_fp4_shuffled).view(

1111 + torch.float8_e4m3fn).reshape(num_experts,

1112 + 2 * intermediate_size,

1113 + hidden_size // 16))

1114 +

1115 + gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)

1116 + gemm2_scales_fp4_shuffled = (

1117 + torch.stack(gemm2_scales_fp4_shuffled).view(

1118 + torch.float8_e4m3fn).reshape(num_experts, hidden_size,

1119 + intermediate_size // 16))

1120 + return (gemm1_weights_fp4_shuffled, gemm1_scales_fp4_shuffled,

1121 + gemm2_weights_fp4_shuffled, gemm2_scales_fp4_shuffled)

1122 +

1123 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

1000 - # GEMM 1

1001 - # The FlashInfer Cutlass fused MoE kernel expects the combined weights

1002 - # to be ordered as [w3, w1], unlike the standard [w1, w3] layout.

1124 + # GEMM 1 processing

1125 gemm1_weight = layer.w13_weight.data

1126 gemm1_weight_scale = layer.w13_weight_scale.data

1127

1006 - if self.allow_flashinfer_cutlass:

1128 + if self.allow_flashinfer:

1129 gemm1_weight, gemm1_weight_scale = reorder_w1w3_to_w3w1(

1130 gemm1_weight, gemm1_weight_scale, dim=-2)

1131

@@ -1011,6 +1133,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1133 layer.w13_weight_scale = Parameter(gemm1_weight_scale,

1134 requires_grad=False)

1135

1136 + # Common processing for w13_weight_scale_2

1137 if not torch.allclose(layer.w13_weight_scale_2[:, 0],

1138 layer.w13_weight_scale_2[:, 1]):

1139 logger.warning_once(

@@ -1021,26 +1144,18 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1144 layer.w13_weight_scale_2 = Parameter(w13_weight_scale_2,

1145 requires_grad=False)

1146

1147 + # Common processing for input scales and alphas

1148 w13_input_scale = layer.w13_input_scale.max(dim=1).values.to(

1149 torch.float32)

1150 layer.g1_alphas = Parameter(

1151 (w13_input_scale * w13_weight_scale_2).to(torch.float32),

1152 requires_grad=False)

1153

1030 - assert (layer.w13_weight_scale.shape[2] % 16 == 0), (

1031 - "Expected weight_scale.dim(1) to be divisible by 16")

1032 - assert (layer.w13_weight_scale.dtype == torch.float8_e4m3fn), (

1033 - "Weight Blockscale must be represented as FP8-E4M3")

1034 - w13_blockscale_swizzled = swizzle_blockscale(layer.w13_weight_scale)

1035 -

1036 - layer.w13_blockscale_swizzled = Parameter(w13_blockscale_swizzled,

1037 - requires_grad=False)

1038 -

1154 # This is for quantization, so we need to invert it.

1155 layer.w13_input_scale_quant = Parameter(

1156 (1 / w13_input_scale).to(torch.float32), requires_grad=False)

1157

1043 - # GEMM 2

1158 + # GEMM 2 processing

1159 layer.g2_alphas = Parameter(

1160 (layer.w2_input_scale * layer.w2_weight_scale_2).to(torch.float32),

1161 requires_grad=False)

@@ -1049,15 +1164,63 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1164 layer.w2_input_scale_quant = Parameter(

1165 (1 / layer.w2_input_scale).to(torch.float32), requires_grad=False)

1166

1052 - assert (layer.w2_weight_scale.shape[2] % 16 == 0), (

1053 - "Expected weight_scale.dim(1) to be divisible by 16")

1054 - assert (layer.w2_weight_scale.dtype == torch.float8_e4m3fn), (

1055 - "Weight Blockscale must be represented as FP8-E4M3")

1056 - w2_blockscale_swizzled = swizzle_blockscale(layer.w2_weight_scale)

1167 + # TensorRT-LLM specific processing

1168 + if self.allow_flashinfer and \

1169 + self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

1170 + # Prepare static weights for TRT-LLM kernel

1171 + (gemm1_weights_fp4_shuffled, gemm1_scales_fp4_shuffled,

1172 + gemm2_weights_fp4_shuffled, gemm2_scales_fp4_shuffled

1173 + ) = self.prepare_static_weight_layouts_for_trtllm_moe(

1174 + layer.w13_weight,

1175 + layer.w2_weight,

1176 + layer.w13_weight_scale,

1177 + layer.w2_weight_scale,

1178 + layer.w2_weight.size(-2), # hidden_size

1179 + layer.w13_weight.size(-2) // 2, # intermediate_size

1180 + layer.w13_weight.size(0), # num_experts

1181 + )

1182 +

1183 + layer.gemm1_weights_fp4_shuffled = Parameter(

1184 + gemm1_weights_fp4_shuffled, requires_grad=False)

1185 + layer.gemm2_weights_fp4_shuffled = Parameter(

1186 + gemm2_weights_fp4_shuffled, requires_grad=False)

1187 + layer.gemm1_scales_fp4_shuffled = Parameter(

1188 + gemm1_scales_fp4_shuffled, requires_grad=False)

1189 + layer.gemm2_scales_fp4_shuffled = Parameter(

1190 + gemm2_scales_fp4_shuffled, requires_grad=False)

1191 +

1192 + # Additional parameter needed for TRT-LLM

1193 + layer.g1_scale_c = Parameter(

1194 + (layer.w2_input_scale_quant * layer.g1_alphas).to(

1195 + torch.float32),

1196 + requires_grad=False,

1197 + )

1198

1058 - layer.w2_blockscale_swizzled = Parameter(w2_blockscale_swizzled,

1059 - requires_grad=False)

1060 - layer.w2_weight = Parameter(layer.w2_weight.data, requires_grad=False)

1199 + # Clean up weights that won't be used by TRT-LLM

1200 + del layer.w2_weight

1201 + del layer.w2_weight_scale

1202 + del layer.w13_weight

1203 + del layer.w13_weight_scale

1204 + else:

1205 + # Non-TRT-LLM processing (Cutlass or non-flashinfer)

1206 + assert (layer.w13_weight_scale.shape[2] % 16 == 0), (

1207 + "Expected weight_scale.dim(1) to be divisible by 16")

1208 + assert (layer.w13_weight_scale.dtype == torch.float8_e4m3fn), (

1209 + "Weight Blockscale must be represented as FP8-E4M3")

1210 + w13_blockscale_swizzled = swizzle_blockscale(

1211 + layer.w13_weight_scale)

1212 + layer.w13_blockscale_swizzled = Parameter(w13_blockscale_swizzled,

1213 + requires_grad=False)

1214 +

1215 + assert (layer.w2_weight_scale.shape[2] % 16 == 0), (

1216 + "Expected weight_scale.dim(1) to be divisible by 16")

1217 + assert (layer.w2_weight_scale.dtype == torch.float8_e4m3fn), (

1218 + "Weight Blockscale must be represented as FP8-E4M3")

1219 + w2_blockscale_swizzled = swizzle_blockscale(layer.w2_weight_scale)

1220 + layer.w2_blockscale_swizzled = Parameter(w2_blockscale_swizzled,

1221 + requires_grad=False)

1222 + layer.w2_weight = Parameter(layer.w2_weight.data,

1223 + requires_grad=False)

1224

1225 if self.use_marlin:

1226 prepare_moe_fp4_layer_for_marlin(layer)

@@ -1095,6 +1258,60 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1258 "EPLB not supported for `ModelOptNvFp4FusedMoE` yet.")

1259 assert activation == "silu", "Only SiLU activation is supported."

1260

1261 + if self.allow_flashinfer and \

1262 + self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

1263 + import flashinfer

1264 +

1265 + from vllm.model_executor.models.llama4 import Llama4MoE

1266 +

1267 + a1_gscale = layer.w13_input_scale_quant

1268 + (hidden_states_fp4,

1269 + hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(

1270 + x,

1271 + a1_gscale,

1272 + is_sf_swizzled_layout=False,

1273 + )

1274 + use_llama4_routing = \

1275 + custom_routing_function is Llama4MoE.custom_routing_function

1276 + routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3

1277 + if use_llama4_routing:

1278 + routing_method_type = flashinfer.RoutingMethodType.Llama4

1279 + out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(

1280 + routing_logits=router_logits

1281 + if use_llama4_routing else router_logits.to(torch.float32),

1282 + routing_bias=e_score_correction_bias,

1283 + hidden_states=hidden_states_fp4,

1284 + hidden_states_scale=hidden_states_scale_linear_fp4.view(

1285 + torch.float8_e4m3fn).flatten(),

1286 + gemm1_weights=layer.gemm1_weights_fp4_shuffled.data,

1287 + gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.data.view(

1288 + torch.float8_e4m3fn),

1289 + gemm1_bias=None,

1290 + gemm1_alpha=None,

1291 + gemm1_beta=None,

1292 + gemm1_clamp_limit=None,

1293 + gemm2_weights=layer.gemm2_weights_fp4_shuffled.data,

1294 + gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.data.view(

1295 + torch.float8_e4m3fn),

1296 + gemm2_bias=None,

1297 + output1_scale_scalar=layer.g1_scale_c.data,

1298 + output1_scale_gate_scalar=layer.g1_alphas.data,

1299 + output2_scale_scalar=layer.g2_alphas.data,

1300 + num_experts=global_num_experts,

1301 + top_k=top_k,

1302 + n_group=num_expert_group,

1303 + topk_group=topk_group,

1304 + intermediate_size=layer.intermediate_size_per_partition,

1305 + local_expert_offset=layer.ep_rank * layer.local_num_experts,

1306 + local_num_experts=layer.local_num_experts,

1307 + routed_scaling_factor=None,

1308 + tile_tokens_dim=_get_tile_tokens_dim(x.shape[0], top_k,

1309 + layer.local_num_experts),

1310 + routing_method_type=routing_method_type,

1311 + do_finalize=True,

1312 + )[0]

1313 + return out

1314 +

1315 topk_weights, topk_ids = FusedMoE.select_experts(

1316 hidden_states=x,

1317 router_logits=router_logits,

@@ -1149,6 +1366,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1366 expert_map=expert_map,

1367 apply_router_weight_on_input=apply_router_weight_on_input)

1368 else:

1369 + assert self.allow_flashinfer and \

1370 + self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1371 out = flashinfer_fp4_cutlass_moe_forward(

1372 self.fused_experts,

1373 layer,

@@ -1160,4 +1379,5 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1379 expert_map=expert_map,

1380 apply_router_weight_on_input=apply_router_weight_on_input,

1381 )

1382 +

1383 return out

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

+9 -3

@@ -126,7 +126,7 @@ def flashinfer_fp4_cutlass_moe_forward(

126

127

128 def select_nvfp4_gemm_impl(

129 - allow_flashinfer_cutlass: bool,

129 + allow_flashinfer: bool,

130 moe, # FusedMoEConfig

131 logger):

132 """Return a GEMM *experts* implementation for NV-FP4 fused-MoE layers"""

@@ -137,8 +137,14 @@ def select_nvfp4_gemm_impl(

137 all2all_manager = get_ep_group().device_communicator.all2all_manager

138 assert all2all_manager is not None

139

140 - if allow_flashinfer_cutlass:

141 - logger.debug_once("Using FlashInferExperts")

140 + if allow_flashinfer:

141 + flashinfer_backend = envs.VLLM_FLASHINFER_MOE_BACKEND

142 + if flashinfer_backend != "throughput":

143 + raise ValueError(

144 + f"Only throughput backend is supported for FlashInferExperts, "

145 + f"but got {flashinfer_backend}.")

146 + logger.debug_once(

147 + "Initializing FlashInferExperts with throughput backend.")

148 return FlashInferExperts(

149 use_nvfp4_w4a4=True,

150 use_dp=moe.moe_parallel_config.dp_size > 1,

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

+2 -2

@@ -21,7 +21,7 @@ class NvFp4Support:

21 """Result container for NV-FP4 capability probing."""

23 cutlass_supported: bool

24 - allow_flashinfer_cutlass: bool

24 + allow_flashinfer: bool

25 use_marlin: bool

@@ -54,6 +54,6 @@ def detect_nvfp4_moe_support(class_name: str = "") -> NvFp4Support:

55 return NvFp4Support(

56 cutlass_supported=cutlass_supported,

57 - allow_flashinfer_cutlass=allow_flashinfer,

57 + allow_flashinfer=allow_flashinfer,

58 use_marlin=use_marlin,

59 )

vllm/utils/flashinfer.py

@@ -86,6 +86,8 @@ flashinfer_cutlass_fused_moe = _lazy_import_wrapper("flashinfer.fused_moe",

86 fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")

87 nvfp4_block_scale_interleave = _lazy_import_wrapper(

88 "flashinfer", "nvfp4_block_scale_interleave")

89 +trtllm_fp4_block_scale_moe = _lazy_import_wrapper(

90 + "flashinfer", "trtllm_fp4_block_scale_moe")

92 # Special case for autotune since it returns a context manager

93 autotune = _lazy_import_wrapper(

@@ -112,6 +114,7 @@ def has_flashinfer_cutlass_fused_moe() -> bool:

114 ("flashinfer.fused_moe", "cutlass_fused_moe"),

115 ("flashinfer", "fp4_quantize"),

116 ("flashinfer", "nvfp4_block_scale_interleave"),

117 + ("flashinfer.fused_moe", "trtllm_fp4_block_scale_moe"),

118 ]

119

120 for module_name, attr_name in required_functions:

@@ -188,6 +191,7 @@ __all__ = [

191 "flashinfer_cutlass_fused_moe",

192 "fp4_quantize",

193 "nvfp4_block_scale_interleave",

194 + "trtllm_fp4_block_scale_moe",

195 "autotune",

196 "has_flashinfer_moe",

197 "has_flashinfer_cutlass_fused_moe",