@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a38b8af4c3
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

[NVIDIA] Add SM100 Flashinfer Cutlass MoE fp8 backend (#22357)

Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com>

amirkl94 committed Aug 20, 2025 at 01:01 UTC a38b8af4c3647ca615e3f1e334779cda8b1eddd8
6 files changed +612 -138
.buildkite/test-pipeline.yaml
+2
@@ -630,6 +630,7 @@ steps:
630 - vllm/model_executor/layers/fused_moe/cutlass_moe.py
631 - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
632 - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
633 + - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
634 - vllm/v1/attention/backends/flashinfer.py
635 - vllm/compilation/fusion.py
636 - vllm/compilation/fusion_attn.py
@@ -650,6 +651,7 @@ steps:
651 # Fusion
652 - pytest -v -s tests/compile/test_fusion_all_reduce.py
653 - pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern
654 + - pytest -v -s tests/kernels/moe/test_flashinfer.py
655
656 ##### 1 GPU test #####
657 ##### multi gpus test #####
tests/kernels/moe/test_flashinfer.py new
+248
@@ -0,0 +1,248 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import dataclass
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
9 +from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
10 +from vllm.model_executor.layers.fused_moe.layer import FusedMoE
11 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
12 + apply_flashinfer_per_tensor_scale_fp8, flashinfer_cutlass_moe_fp8,
13 + register_moe_scaling_factors, rotate_flashinfer_fp8_moe_weights,
14 + swap_w13_to_w31)
15 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
16 + input_to_float8)
17 +from vllm.model_executor.models.llama4 import Llama4MoE
18 +from vllm.platforms import current_platform
19 +from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
20 +
21 +if not has_flashinfer_cutlass_fused_moe(
22 +) or not current_platform.has_device_capability(100):
23 + pytest.skip("Requires flashinfer_cutlass_fused_moe and nvfp4 support",
24 + allow_module_level=True)
25 +
26 +NUM_EXPERTS = [16]
27 +TOP_KS = [1]
28 +
29 +MNK_FACTORS = [
30 + (256, 8192, 5120),
31 + (256, 4096, 5120),
32 + (127, 8192, 5120),
33 + (127, 4096, 5120),
34 + (10, 8192, 5120),
35 + (10, 4096, 5120),
36 + (1, 8192, 5120),
37 + (1, 4096, 5120),
38 +]
39 +
40 +vllm_config = VllmConfig(parallel_config=ParallelConfig(
41 + pipeline_parallel_size=1))
42 +vllm_config.scheduler_config.max_num_seqs = 128
43 +vllm_config.scheduler_config.max_model_len = 8192
44 +
45 +
46 +def quant_fp8_per_tensor_batches(a):
47 + num_batches = a.size(0)
48 + a_quant = []
49 + a_scales = []
50 +
51 + for i in range(num_batches):
52 + a_fp8, a_global_sf = input_to_float8(a[i])
53 + a_global_sf = 1.0 / a_global_sf
54 + a_quant.append(a_fp8)
55 + a_scales.append(a_global_sf)
56 +
57 + result_a_quant = torch.stack(a_quant)
58 + result_a_scales = torch.stack(a_scales)
59 +
60 + return result_a_quant, result_a_scales
61 +
62 +
63 +@dataclass
64 +class TestData:
65 + hidden_states: torch.Tensor
66 + w13_quantized: torch.Tensor
67 + w2_quantized: torch.Tensor
68 + a1_scale: torch.Tensor
69 + a2_scale: torch.Tensor
70 + w13_weight_scale: torch.Tensor
71 + w2_weight_scale: torch.Tensor
72 + layer: torch.nn.Module
73 +
74 + @staticmethod
75 + def make_moe_tensors_8bit(m: int, k: int, n: int, e: int,
76 + reorder: bool) -> "TestData":
77 + hidden_states = torch.randn(
78 + (m, k), device="cuda", dtype=torch.bfloat16) / 10
79 + w13 = torch.randn((e, 2 * n, k), device="cuda", dtype=torch.bfloat16)
80 + w2 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16)
81 +
82 + # Scale to fp8
83 + _, a1_scale = input_to_float8(hidden_states)
84 + a1_scale = 1.0 / a1_scale
85 + a2_scale = torch.scalar_tensor(1.0).to(device="cuda").to(
86 + dtype=torch.float32)
87 + w13_quantized, w13_weight_scale = quant_fp8_per_tensor_batches(w13)
88 + w2_quantized, w2_weight_scale = quant_fp8_per_tensor_batches(w2)
89 +
90 + layer = torch.nn.Module()
91 + layer.w13_weight = w13_quantized.clone()
92 + layer.w2_weight = w2_quantized.clone()
93 + layer.w13_input_scale = a1_scale
94 + layer.w2_input_scale = a2_scale
95 + layer.w13_weight_scale = w13_weight_scale
96 + layer.w2_weight_scale = w2_weight_scale
97 +
98 + register_moe_scaling_factors(layer)
99 +
100 + # flashinfer expects swapped rows for w13
101 + layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)
102 + if reorder:
103 + rotate_flashinfer_fp8_moe_weights(layer.w13_weight,
104 + layer.w2_weight)
105 + layer.custom_routing_function = Llama4MoE.custom_routing_function
106 + layer.intermediate_size_per_partition = n
107 + layer.ep_rank = 0
108 + layer.local_num_experts = e
109 +
110 + return TestData(
111 + hidden_states=hidden_states,
112 + w13_quantized=w13_quantized,
113 + w2_quantized=w2_quantized,
114 + a1_scale=a1_scale,
115 + a2_scale=a2_scale,
116 + w13_weight_scale=w13_weight_scale,
117 + w2_weight_scale=w2_weight_scale,
118 + layer=layer,
119 + )
120 +
121 +
122 +@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
123 +@pytest.mark.parametrize("e", NUM_EXPERTS)
124 +@pytest.mark.parametrize("topk", TOP_KS)
125 +def test_flashinfer_per_tensor_moe_fp8_no_graph(
126 + m: int,
127 + n: int,
128 + k: int,
129 + e: int,
130 + topk: int,
131 + monkeypatch,
132 +):
133 + current_platform.seed_everything(7)
134 + monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
135 + with set_current_vllm_config(vllm_config):
136 + td = TestData.make_moe_tensors_8bit(m, k, n, e, reorder=True)
137 +
138 + score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
139 + topk_weights, topk_ids = FusedMoE.select_experts(
140 + hidden_states=td.hidden_states,
141 + router_logits=score,
142 + use_grouped_topk=False,
143 + top_k=topk,
144 + renormalize=False,
145 + custom_routing_function=Llama4MoE.custom_routing_function,
146 + scoring_func="softmax")
147 +
148 + output = fused_experts(
149 + td.hidden_states,
150 + td.w13_quantized,
151 + td.w2_quantized,
152 + topk_weights=topk_weights,
153 + topk_ids=topk_ids,
154 + inplace=False,
155 + activation="silu",
156 + use_fp8_w8a8=True,
157 + per_channel_quant=False,
158 + global_num_experts=e,
159 + expert_map=None,
160 + w1_scale=td.w13_weight_scale,
161 + w2_scale=td.w2_weight_scale,
162 + a1_scale=td.a1_scale,
163 + a2_scale=td.a2_scale,
164 + apply_router_weight_on_input=True,
165 + )
166 +
167 + flashinfer_output = apply_flashinfer_per_tensor_scale_fp8(
168 + layer=td.layer,
169 + hidden_states=td.hidden_states,
170 + router_logits=score,
171 + routing_bias=None,
172 + global_num_experts=e,
173 + top_k=topk,
174 + num_expert_group=None,
175 + topk_group=None,
176 + apply_router_weight_on_input=True)
177 +
178 + torch.testing.assert_close(output,
179 + flashinfer_output,
180 + atol=5.5e-2,
181 + rtol=1e-2)
182 +
183 +
184 +@pytest.mark.skip(
185 + "Requires flashinfer version that contains https://github.com/flashinfer-ai/flashinfer/pull/1472"
186 +)
187 +@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
188 +@pytest.mark.parametrize("e", NUM_EXPERTS)
189 +@pytest.mark.parametrize("topk", TOP_KS)
190 +def test_flashinfer_cutlass_moe_fp8_no_graph(
191 + m: int,
192 + n: int,
193 + k: int,
194 + e: int,
195 + topk: int,
196 + monkeypatch,
197 +):
198 + current_platform.seed_everything(7)
199 + monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
200 + with set_current_vllm_config(vllm_config):
201 + td = TestData.make_moe_tensors_8bit(m, k, n, e, reorder=False)
202 +
203 + score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
204 + topk_weights, topk_ids = FusedMoE.select_experts(
205 + hidden_states=td.hidden_states,
206 + router_logits=score,
207 + use_grouped_topk=False,
208 + top_k=topk,
209 + renormalize=False,
210 + custom_routing_function=Llama4MoE.custom_routing_function,
211 + scoring_func="softmax")
212 +
213 + output = fused_experts(
214 + td.hidden_states,
215 + td.w13_quantized,
216 + td.w2_quantized,
217 + topk_weights=topk_weights,
218 + topk_ids=topk_ids,
219 + inplace=False,
220 + activation="silu",
221 + use_fp8_w8a8=True,
222 + per_channel_quant=False,
223 + global_num_experts=e,
224 + expert_map=None,
225 + w1_scale=td.w13_weight_scale,
226 + w2_scale=td.w2_weight_scale,
227 + a1_scale=td.a1_scale,
228 + a2_scale=td.a2_scale,
229 + apply_router_weight_on_input=True,
230 + )
231 +
232 + td.layer.dp_size = 1
233 +
234 + flashinfer_cutlass_output = flashinfer_cutlass_moe_fp8(
235 + td.hidden_states,
236 + td.layer,
237 + topk_weights,
238 + topk_ids,
239 + activation="silu",
240 + global_num_experts=e,
241 + expert_map=None,
242 + apply_router_weight_on_input=True,
243 + )
244 +
245 + torch.testing.assert_close(output,
246 + flashinfer_cutlass_output,
247 + atol=5.5e-2,
248 + rtol=1e-2)
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+33 -22
@@ -61,8 +61,8 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
61 per_act_token_quant=False,
62 block_shape=None,
63 ))
64 - assert quant_dtype == "nvfp4", ("Only nvfp4 quantization is "
65 - "currently supported.")
64 + assert quant_dtype in ("nvfp4", torch.float8_e4m3fn), (
65 + "Only nvfp4,fp8 quantization are currently supported.")
66 self.ep_rank = ep_rank
67 self.ep_size = ep_size
68 self.tp_rank = tp_rank
@@ -122,7 +122,8 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
122 """
123 aq_m, aq_n = aq.shape
124 workspace2 = ()
125 - output_shape = (aq_m, aq_n * 2)
125 + output_shape = (aq_m, aq_n * 2) if self.quant_dtype != \
126 + torch.float8_e4m3fn else (aq_m, aq_n)
127 workspace_dtype = a.dtype
128 workspace1 = output_shape
129 # The workspace is determined by `aq`, since it comes after any
@@ -151,29 +152,39 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
152 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
153 apply_router_weight_on_input: Optional[bool],
154 ):
154 - # Flashinfer CUTLASS kernel takes scalar global scales,
155 - # min because inv_scale.
156 -
157 - # Ensure w1_scale and w2_scale are not None before calling view
158 - assert w1_scale is not None and w2_scale is not None, (
159 - "w1_scale and w2_scale must not "
160 - "be None for FlashInferExperts")
161 -
162 - quant_scales = [
163 - self.a1_gscale,
164 - w1_scale.view(torch.int32),
165 - self.g1_alphas,
166 - self.a2_gscale,
167 - w2_scale.view(torch.int32),
168 - self.g2_alphas,
169 - ]
155 + if self.quant_dtype == torch.float8_e4m3fn:
156 + quant_scales = [
157 + self.g1_alphas, self.a2_gscale, self.g2_alphas, self.a1_gscale
158 + ]
159 +
160 + a1q_scale = None # not passing input_sf in fp8
161 + fc1_expert_weights = w1
162 + fc2_expert_weights = w2
163 + else:
164 + # Ensure w1_scale and w2_scale are not None before calling view
165 + assert w1_scale is not None and w2_scale is not None, (
166 + "w1_scale and w2_scale must not "
167 + "be None for FlashInferExperts")
168 + # Flashinfer CUTLASS kernel takes scalar global scales,
169 + # min because inv_scale.
170 + quant_scales = [
171 + self.a1_gscale,
172 + w1_scale.view(torch.int32),
173 + self.g1_alphas,
174 + self.a2_gscale,
175 + w2_scale.view(torch.int32),
176 + self.g2_alphas,
177 + ]
178 + # FlashInfer API requires weight to be long for nvfp4
179 + fc1_expert_weights = w1.view(torch.long)
180 + fc2_expert_weights = w2.view(torch.long)
181 +
182 _ = flashinfer_cutlass_fused_moe(
183 input=hidden_states,
184 token_selected_experts=topk_ids.to(torch.int),
185 token_final_scales=topk_weights,
174 - # FlashInfer API requires weight to be long for nvfp4
175 - fc1_expert_weights=w1.view(torch.long),
176 - fc2_expert_weights=w2.view(torch.long),
186 + fc1_expert_weights=fc1_expert_weights,
187 + fc2_expert_weights=fc2_expert_weights,
188 output_dtype=self.out_dtype,
189 quant_scales=quant_scales,
190 input_sf=a1q_scale,
vllm/model_executor/layers/quantization/fp8.py
+134 -81
@@ -9,6 +9,7 @@ from torch.nn import Module
9 from torch.nn.parameter import Parameter
10
11 import vllm.envs as envs
12 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13 from vllm import _custom_ops as ops
14 from vllm.distributed import get_tensor_model_parallel_world_size
15 from vllm.logger import init_logger
@@ -23,8 +24,11 @@ from vllm.model_executor.layers.quantization.base_config import (
24 QuantizationConfig, QuantizeMethodBase)
25 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
26 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
26 - apply_flashinfer_per_tensor_scale_fp8, register_moe_scaling_factors,
27 - rotate_flashinfer_fp8_moe_weights, swap_w13_to_w31)
27 + FlashinferMoeBackend, apply_flashinfer_per_tensor_scale_fp8,
28 + build_flashinfer_fp8_cutlass_moe_prepare_finalize,
29 + flashinfer_cutlass_moe_fp8, get_flashinfer_moe_backend,
30 + register_moe_scaling_factors, rotate_flashinfer_fp8_moe_weights,
31 + select_cutlass_fp8_gemm_impl, swap_w13_to_w31)
32 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
33 get_col_major_tma_aligned_tensor, requant_weight_ue8m0_inplace)
34 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
@@ -145,7 +149,7 @@ class Fp8Config(QuantizationConfig):
149 return UnquantizedLinearMethod()
150 return Fp8LinearMethod(self)
151 elif isinstance(layer, FusedMoE):
148 - return Fp8MoEMethod(self, layer.moe_config)
152 + return Fp8MoEMethod(self, layer)
153 elif isinstance(layer, Attention):
154 return Fp8KVCacheMethod(self)
155 return None
@@ -482,16 +486,20 @@ class Fp8MoEMethod(FusedMoEMethodBase):
486 quant_config: The quantization config.
487 """
488
485 - def __init__(self, quant_config: Fp8Config, moe: FusedMoEConfig):
486 - super().__init__(moe)
489 + def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
490 + super().__init__(layer.moe_config)
491 + self.layer = layer
492 self.quant_config = quant_config
493 self.block_quant = self.quant_config.weight_block_size is not None
494
490 - self.flashinfer_moe_enabled = False
495 + self.flashinfer_moe_backend: Optional[FlashinferMoeBackend] = None
496 + self.fused_experts: Optional[
497 + mk.FusedMoEModularKernel] = None # type: ignore
498 if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():
499 + self.flashinfer_moe_backend = get_flashinfer_moe_backend()
500 logger.info_once(
493 - "Using FlashInfer MoE FP8 kernels for Fp8MoEMethod.")
494 - self.flashinfer_moe_enabled = True
501 + f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
502 + )
503 # For GPUs that lack FP8 hardware support, we can leverage the Marlin
504 # kernel for fast weight-only FP8 quantization
505 self.use_marlin = (not current_platform.has_device_capability(89)
@@ -531,6 +539,20 @@ class Fp8MoEMethod(FusedMoEMethodBase):
539 "CutlassBlockScaledGroupedGemm not supported on the current "
540 "platform.")
541
542 + def maybe_make_prepare_finalize(
543 + self,
544 + moe: FusedMoEConfig,
545 + ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
546 + if self.flashinfer_moe_backend != FlashinferMoeBackend.CUTLASS:
547 + return super().maybe_make_prepare_finalize(moe)
548 +
549 + prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
550 + moe,
551 + layer=self.layer,
552 + )
553 + logger.debug_once("%s", prepare_finalize.__class__.__name__)
554 + return prepare_finalize
555 +
556 def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
557 intermediate_size_per_partition: int,
558 params_dtype: torch.dtype, **extra_weight_attrs):
@@ -678,7 +700,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
700 normalize_e4m3fn_to_e4m3fnuz(
701 layer.w2_weight, layer.w2_weight_scale_inv,
702 layer.w2_input_scale)
681 - elif self.flashinfer_moe_enabled:
703 + elif self.flashinfer_moe_backend is not None:
704 # NOTE: weights have to be swapped since the activation is
705 # applied on different half for flashinfer vs vllm
706 w13_weight = swap_w13_to_w31(layer.w13_weight.data)
@@ -686,9 +708,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
708 layer.w13_weight_scale_inv.data)
709 w2_weight = layer.w2_weight.data
710 w2_weight_scale_inv = layer.w2_weight_scale_inv.data
689 - if not self.block_quant:
690 - register_moe_scaling_factors(layer)
691 - rotate_flashinfer_fp8_moe_weights(w13_weight, w2_weight)
711 else:
712 w13_weight = layer.w13_weight.data
713 w13_weight_scale_inv = layer.w13_weight_scale_inv.data
@@ -834,6 +853,17 @@ class Fp8MoEMethod(FusedMoEMethodBase):
853 layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
854 requires_grad=False)
855
856 + if self.flashinfer_moe_backend is not None:
857 + # NOTE: weights have to be swapped since the activation is
858 + # applied on different half for flashinfer vs vllm
859 + assert not self.block_quant
860 + register_moe_scaling_factors(layer)
861 + w13_weight = swap_w13_to_w31(layer.w13_weight.data)
862 + if self.flashinfer_moe_backend == \
863 + FlashinferMoeBackend.TENSORRT_LLM:
864 + rotate_flashinfer_fp8_moe_weights(w13_weight, w2_weight)
865 + layer.w13_weight.data = w13_weight.data
866 +
867 if self.use_marlin:
868 prepare_moe_fp8_layer_for_marlin(layer, False)
869 # Activations not quantized for marlin.
@@ -892,6 +922,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
922 per_act_token_quant=False,
923 allow_deep_gemm=self.allow_deep_gemm,
924 )
925 + elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
926 + experts = select_cutlass_fp8_gemm_impl(
927 + moe,
928 + self.layer,
929 + )
930 + logger.debug_once("Using %s", experts.__class__.__name__)
931 + return experts
932 else:
933 logger.debug(
934 "TritonOrDeepGemmExperts(%s): block_size=%s, per_act_token=%s",
@@ -930,67 +967,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
967 assert logical_to_physical_map is not None
968 assert logical_replica_count is not None
969 assert isinstance(layer, FusedMoE)
933 - if not self.flashinfer_moe_enabled:
934 - topk_weights, topk_ids = FusedMoE.select_experts(
935 - hidden_states=x,
936 - router_logits=router_logits,
937 - use_grouped_topk=use_grouped_topk,
938 - top_k=top_k,
939 - renormalize=renormalize,
940 - topk_group=topk_group,
941 - num_expert_group=num_expert_group,
942 - custom_routing_function=custom_routing_function,
943 - scoring_func=scoring_func,
944 - e_score_correction_bias=e_score_correction_bias,
945 - indices_type=self.topk_indices_dtype,
946 - enable_eplb=enable_eplb,
947 - expert_map=expert_map,
948 - expert_load_view=expert_load_view,
949 - logical_to_physical_map=logical_to_physical_map,
950 - logical_replica_count=logical_replica_count,
951 - )
970
953 - if self.rocm_aiter_moe_enabled:
954 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
955 - rocm_aiter_fused_experts)
956 - return rocm_aiter_fused_experts(
957 - x,
958 - layer.w13_weight,
959 - layer.w2_weight,
960 - topk_weights=topk_weights,
961 - topk_ids=topk_ids,
962 - activation=activation,
963 - use_fp8_w8a8=True,
964 - apply_router_weight_on_input=apply_router_weight_on_input,
965 - w1_scale=(layer.w13_weight_scale_inv
966 - if self.block_quant else layer.w13_weight_scale),
967 - w2_scale=(layer.w2_weight_scale_inv
968 - if self.block_quant else layer.w2_weight_scale),
969 - a1_scale=layer.w13_input_scale,
970 - a2_scale=layer.w2_input_scale,
971 - block_shape=self.quant_config.weight_block_size,
972 - expert_map=expert_map)
973 - elif self.use_marlin:
974 - assert activation == "silu", (
975 - f"{activation} not supported for Marlin MoE.")
976 - return torch.ops.vllm.fused_marlin_moe(
977 - x,
978 - layer.w13_weight,
979 - layer.w2_weight,
980 - None,
981 - None,
982 - layer.w13_weight_scale,
983 - layer.w2_weight_scale,
984 - router_logits,
985 - topk_weights,
986 - topk_ids,
987 - quant_type_id=scalar_types.float8_e4m3fn.id,
988 - apply_router_weight_on_input=apply_router_weight_on_input,
989 - global_num_experts=global_num_experts,
990 - expert_map=expert_map)
991 - elif self.flashinfer_moe_enabled:
992 - assert activation == 'silu'
993 - assert scoring_func == 'sigmoid'
971 + if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
972 + assert activation == 'silu', (
973 + f"Expected 'silu' activation but got {activation}")
974 + assert scoring_func == 'sigmoid', (
975 + f"Expected 'sigmoid' scoring func but got {scoring_func}")
976 if self.block_quant:
977 assert (renormalize and use_grouped_topk
978 and custom_routing_function is None)
@@ -1026,25 +1008,96 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1008 num_expert_group=num_expert_group,
1009 topk_group=topk_group,
1010 apply_router_weight_on_input=apply_router_weight_on_input)
1029 - elif self.fused_experts is not None:
1030 - return self.fused_experts(
1031 - hidden_states=x,
1032 - w1=layer.w13_weight,
1033 - w2=layer.w2_weight,
1011 +
1012 + topk_weights, topk_ids = FusedMoE.select_experts(
1013 + hidden_states=x,
1014 + router_logits=router_logits,
1015 + use_grouped_topk=use_grouped_topk,
1016 + top_k=top_k,
1017 + renormalize=renormalize,
1018 + topk_group=topk_group,
1019 + num_expert_group=num_expert_group,
1020 + custom_routing_function=custom_routing_function,
1021 + scoring_func=scoring_func,
1022 + e_score_correction_bias=e_score_correction_bias,
1023 + indices_type=self.topk_indices_dtype,
1024 + enable_eplb=enable_eplb,
1025 + expert_map=expert_map,
1026 + expert_load_view=expert_load_view,
1027 + logical_to_physical_map=logical_to_physical_map,
1028 + logical_replica_count=logical_replica_count,
1029 + )
1030 +
1031 + if self.rocm_aiter_moe_enabled:
1032 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
1033 + rocm_aiter_fused_experts)
1034 + return rocm_aiter_fused_experts(
1035 + x,
1036 + layer.w13_weight,
1037 + layer.w2_weight,
1038 topk_weights=topk_weights,
1039 topk_ids=topk_ids,
1036 - inplace=True,
1040 activation=activation,
1038 - global_num_experts=global_num_experts,
1041 + use_fp8_w8a8=True,
1042 apply_router_weight_on_input=apply_router_weight_on_input,
1040 - expert_map=expert_map,
1043 w1_scale=(layer.w13_weight_scale_inv
1044 if self.block_quant else layer.w13_weight_scale),
1045 w2_scale=(layer.w2_weight_scale_inv
1046 if self.block_quant else layer.w2_weight_scale),
1047 a1_scale=layer.w13_input_scale,
1048 a2_scale=layer.w2_input_scale,
1047 - )
1049 + block_shape=self.quant_config.weight_block_size,
1050 + expert_map=expert_map)
1051 + elif self.use_marlin:
1052 + assert activation == "silu", (
1053 + f"{activation} not supported for Marlin MoE.")
1054 + return torch.ops.vllm.fused_marlin_moe(
1055 + x,
1056 + layer.w13_weight,
1057 + layer.w2_weight,
1058 + None,
1059 + None,
1060 + layer.w13_weight_scale,
1061 + layer.w2_weight_scale,
1062 + router_logits,
1063 + topk_weights,
1064 + topk_ids,
1065 + quant_type_id=scalar_types.float8_e4m3fn.id,
1066 + apply_router_weight_on_input=apply_router_weight_on_input,
1067 + global_num_experts=global_num_experts,
1068 + expert_map=expert_map)
1069 + elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
1070 + assert self.block_quant is None
1071 + assert (not renormalize and custom_routing_function is not None)
1072 + assert activation == 'silu', (
1073 + f"Expected 'silu' activation but got {activation}")
1074 + assert scoring_func == 'sigmoid', (
1075 + f"Expected 'sigmoid' scoring func but got {scoring_func}")
1076 + if self.fused_experts is not None:
1077 + return self.fused_experts(
1078 + x,
1079 + layer.w13_weight,
1080 + layer.w2_weight,
1081 + topk_weights,
1082 + topk_ids,
1083 + inplace=False,
1084 + activation=activation,
1085 + global_num_experts=global_num_experts,
1086 + expert_map=expert_map,
1087 + apply_router_weight_on_input=apply_router_weight_on_input,
1088 + )
1089 + else:
1090 + return flashinfer_cutlass_moe_fp8(
1091 + x,
1092 + layer,
1093 + topk_weights,
1094 + topk_ids,
1095 + inplace=False,
1096 + activation=activation,
1097 + global_num_experts=global_num_experts,
1098 + expert_map=expert_map,
1099 + apply_router_weight_on_input=apply_router_weight_on_input,
1100 + )
1101 else:
1102 from vllm.model_executor.layers.fused_moe import fused_experts
1103 return fused_experts(
vllm/model_executor/layers/quantization/modelopt.py
+83 -35
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from enum import Enum
4 from typing import Any, Callable, Optional, Union
5
6 import torch
@@ -27,8 +26,11 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
26 build_flashinfer_fp4_cutlass_moe_prepare_finalize, reorder_w1w3_to_w3w1,
27 select_nvfp4_gemm_impl)
28 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
30 - apply_flashinfer_per_tensor_scale_fp8, register_moe_scaling_factors,
31 - rotate_flashinfer_fp8_moe_weights, swap_w13_to_w31)
29 + FlashinferMoeBackend, apply_flashinfer_per_tensor_scale_fp8,
30 + build_flashinfer_fp8_cutlass_moe_prepare_finalize,
31 + flashinfer_cutlass_moe_fp8, get_flashinfer_moe_backend,
32 + register_moe_scaling_factors, rotate_flashinfer_fp8_moe_weights,
33 + select_cutlass_fp8_gemm_impl, swap_w13_to_w31)
34 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
35 apply_fp4_marlin_linear, is_fp4_marlin_supported,
36 prepare_fp4_layer_for_marlin, prepare_moe_fp4_layer_for_marlin)
@@ -49,11 +51,6 @@ QUANT_ALGOS = ["FP8", "NVFP4"]
51 KV_CACHE_QUANT_ALGOS = ["FP8"]
52
53
52 -class FlashinferMoeBackend(Enum):
53 - TENSORRT_LLM = "TensorRT-LLM"
54 - CUTLASS = "CUTLASS"
55 -
56 -
54 class ModelOptFp8Config(QuantizationConfig):
55 """Config class for ModelOpt FP8."""
56
@@ -179,7 +176,7 @@ class ModelOptFp8Config(QuantizationConfig):
176 elif isinstance(layer, Attention):
177 return ModelOptFp8KVCacheMethod(self)
178 elif isinstance(layer, FusedMoE):
182 - return ModelOptFp8MoEMethod(self, layer.moe_config)
179 + return ModelOptFp8MoEMethod(self, layer)
180 return None
181
182
@@ -278,18 +275,49 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
275 def __init__(
276 self,
277 quant_config: ModelOptFp8Config,
281 - moe: FusedMoEConfig,
278 + layer: torch.nn.Module,
279 ) -> None:
283 - super().__init__(moe)
280 + super().__init__(layer.moe_config)
281 + self.layer = layer
282 self.quant_config = quant_config
283 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
284 cutlass_fp8_supported)
285 self.cutlass_fp8_supported = cutlass_fp8_supported()
288 - self.flashinfer_moe_enabled = False
286 + self.flashinfer_moe_backend: Optional[FlashinferMoeBackend] = None
287 + self.fused_experts: Optional[
288 + mk.FusedMoEModularKernel] = None # type: ignore
289 if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():
290 + self.flashinfer_moe_backend = get_flashinfer_moe_backend()
291 logger.info_once(
291 - "Using FlashInfer MoE FP8 kernels for ModelOptFp8MoEMethod.")
292 - self.flashinfer_moe_enabled = True
292 + f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
293 + )
294 +
295 + def maybe_make_prepare_finalize(
296 + self,
297 + moe: FusedMoEConfig,
298 + ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
299 + if self.fused_experts is not None or \
300 + self.flashinfer_moe_backend != FlashinferMoeBackend.CUTLASS:
301 + return super().maybe_make_prepare_finalize(moe)
302 +
303 + prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
304 + moe,
305 + layer=self.layer,
306 + )
307 + logger.debug_once("%s", prepare_finalize.__class__.__name__)
308 + return prepare_finalize
309 +
310 + def select_gemm_impl(
311 + self,
312 + prepare_finalize: mk.FusedMoEPrepareAndFinalize,
313 + moe: FusedMoEConfig,
314 + ) -> mk.FusedMoEPermuteExpertsUnpermute:
315 + experts = select_cutlass_fp8_gemm_impl(
316 + moe,
317 + self.layer,
318 + )
319 + logger.debug_once("Using %s", experts.__class__.__name__)
320 + return experts
321
322 def create_weights(
323 self,
@@ -433,11 +461,12 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
461 layer.w2_input_scale = Parameter(layer.w2_input_scale.max(),
462 requires_grad=False)
463
436 - if self.flashinfer_moe_enabled:
464 + if self.flashinfer_moe_backend is not None:
465 layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)
438 - rotate_flashinfer_fp8_moe_weights(layer.w13_weight,
439 - layer.w2_weight)
466 register_moe_scaling_factors(layer)
467 + if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
468 + rotate_flashinfer_fp8_moe_weights(layer.w13_weight,
469 + layer.w2_weight)
470
471 def apply(
472 self,
@@ -461,14 +490,13 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
490 logical_to_physical_map: Optional[torch.Tensor] = None,
491 logical_replica_count: Optional[torch.Tensor] = None,
492 ) -> torch.Tensor:
464 - assert self.fused_experts is None
465 -
493 if enable_eplb:
494 raise NotImplementedError(
495 "EPLB not supported for `ModelOptFp8MoEMethod` yet.")
496
470 - if self.flashinfer_moe_enabled:
471 - assert activation == 'silu'
497 + if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
498 + assert activation == 'silu', (
499 + f"Expected 'silu' activation but got {activation}")
500 assert not renormalize
501 return apply_flashinfer_per_tensor_scale_fp8(
502 layer=layer,
@@ -495,6 +523,36 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
523 e_score_correction_bias=e_score_correction_bias,
524 indices_type=self.topk_indices_dtype,
525 )
526 +
527 + if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
528 + assert not renormalize
529 + assert activation == 'silu', (
530 + f"Expected 'silu' activation but got {activation}")
531 + if self.fused_experts is not None:
532 + return self.fused_experts(
533 + x,
534 + layer.w13_weight,
535 + layer.w2_weight,
536 + topk_weights,
537 + topk_ids,
538 + inplace=False,
539 + activation=activation,
540 + global_num_experts=global_num_experts,
541 + expert_map=expert_map,
542 + apply_router_weight_on_input=apply_router_weight_on_input,
543 + )
544 + else:
545 + return flashinfer_cutlass_moe_fp8(
546 + x,
547 + layer,
548 + topk_weights,
549 + topk_ids,
550 + inplace=False,
551 + activation=activation,
552 + global_num_experts=global_num_experts,
553 + expert_map=expert_map,
554 + apply_router_weight_on_input=apply_router_weight_on_input,
555 + )
556 from vllm.model_executor.layers.fused_moe.fused_moe import (
557 fused_experts)
558 return fused_experts(
@@ -951,20 +1009,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1009 self.flashinfer_moe_backend = None
1010
1011 if self.allow_flashinfer:
954 - flashinfer_moe_backend = envs.VLLM_FLASHINFER_MOE_BACKEND
955 - if flashinfer_moe_backend == "throughput":
956 - self.flashinfer_moe_backend = FlashinferMoeBackend.CUTLASS
957 - logger.info_once("Using FlashInfer CUTLASS kernels for "
958 - "ModelOptNvFp4FusedMoE.")
959 - elif flashinfer_moe_backend == "latency":
960 - self.flashinfer_moe_backend = FlashinferMoeBackend.TENSORRT_LLM
961 - logger.info_once("Using FlashInfer TensorRT-LLM kernels for "
962 - "ModelOptNvFp4FusedMoE.")
963 - else:
964 - allowed_backends = ["throughput", "latency"]
965 - raise ValueError(
966 - f"Unknown flashinfer moe backend: {flashinfer_moe_backend}"
967 - f" expected one of {allowed_backends}")
1012 + self.flashinfer_moe_backend = get_flashinfer_moe_backend()
1013 + logger.info_once(
1014 + f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
1015 + " for ModelOptNvFp4FusedMoE.")
1016
1017 def maybe_make_prepare_finalize(
1018 self,
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+112
@@ -1,9 +1,26 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from enum import Enum
4 from typing import Optional
5
6 import torch
7
8 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 +from vllm import envs
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
12 +from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
13 + FlashInferExperts)
14 +from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
15 + FlashInferCutlassMoEPrepareAndFinalize)
16 +
17 +logger = init_logger(__name__)
18 +
19 +
20 +class FlashinferMoeBackend(Enum):
21 + TENSORRT_LLM = "TensorRT-LLM"
22 + CUTLASS = "CUTLASS"
23 +
24
25 def calculate_tile_tokens_dim(num_tokens, top_k, num_experts):
26
@@ -144,3 +161,98 @@ def register_moe_scaling_factors(layer: torch.nn.Module) -> None:
161 layer.register_parameter(
162 'output2_scales_scalar',
163 torch.nn.Parameter(output2_scales, requires_grad=False))
164 + layer.register_parameter(
165 + 'w2_input_scale_inv',
166 + torch.nn.Parameter(1.0 / layer.w2_input_scale, requires_grad=False))
167 +
168 +
169 +def build_flashinfer_fp8_cutlass_moe_prepare_finalize(
170 + moe: Optional[FusedMoEConfig],
171 + layer: torch.nn.Module,
172 +) -> mk.FusedMoEPrepareAndFinalize:
173 + """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""
174 + use_dp = moe.moe_parallel_config.dp_size > 1 if moe is not None else False
175 + return FlashInferCutlassMoEPrepareAndFinalize(
176 + use_dp, a1_gscale=layer.w13_input_scale)
177 +
178 +
179 +def select_cutlass_fp8_gemm_impl(
180 + moe: Optional[FusedMoEConfig],
181 + layer: torch.nn.Module,
182 + out_dtype: Optional[torch.dtype] = None,
183 +) -> mk.FusedMoEPermuteExpertsUnpermute:
184 + """Return a GEMM *experts* implementation for fused-MoE layers"""
185 +
186 + from vllm.model_executor.models.llama4 import Llama4MoE
187 + assert layer.custom_routing_function == Llama4MoE.custom_routing_function, \
188 + "FusedMoE flashinfer kernels are only supported for Llama4"
189 +
190 + if moe is not None:
191 + return FlashInferExperts(
192 + g1_alphas=layer.output1_scales_gate_scalar,
193 + g2_alphas=layer.output2_scales_scalar,
194 + a1_gscale=layer.w13_input_scale,
195 + a2_gscale=layer.w2_input_scale_inv,
196 + out_dtype=moe.in_dtype,
197 + quant_dtype=torch.float8_e4m3fn,
198 + ep_rank=moe.moe_parallel_config.ep_rank,
199 + ep_size=moe.moe_parallel_config.ep_size,
200 + tp_rank=moe.moe_parallel_config.tp_rank,
201 + tp_size=moe.moe_parallel_config.tp_size,
202 + )
203 +
204 + assert out_dtype is not None, (
205 + "If moe config is None, out_dtype must be passed")
206 + return FlashInferExperts(
207 + g1_alphas=layer.output1_scales_gate_scalar,
208 + g2_alphas=layer.output2_scales_scalar,
209 + a1_gscale=layer.w13_input_scale,
210 + a2_gscale=layer.w2_input_scale_inv,
211 + out_dtype=out_dtype,
212 + quant_dtype=torch.float8_e4m3fn,
213 + )
214 +
215 +
216 +def flashinfer_cutlass_moe_fp8(
217 + hidden_states: torch.Tensor,
218 + layer: torch.nn.Module,
219 + topk_weights: torch.Tensor,
220 + topk_ids: torch.Tensor,
221 + inplace: bool = False,
222 + activation: str = "silu",
223 + global_num_experts: int = -1,
224 + expert_map: Optional[torch.Tensor] = None,
225 + apply_router_weight_on_input: bool = False,
226 +) -> torch.Tensor:
227 + fused_experts = mk.FusedMoEModularKernel(
228 + build_flashinfer_fp8_cutlass_moe_prepare_finalize(moe=None,
229 + layer=layer),
230 + select_cutlass_fp8_gemm_impl(moe=None,
231 + layer=layer,
232 + out_dtype=hidden_states.dtype))
233 +
234 + return fused_experts(
235 + hidden_states,
236 + layer.w13_weight,
237 + layer.w2_weight,
238 + topk_weights,
239 + topk_ids,
240 + inplace=inplace,
241 + activation=activation,
242 + global_num_experts=global_num_experts,
243 + expert_map=expert_map,
244 + apply_router_weight_on_input=apply_router_weight_on_input,
245 + )
246 +
247 +
248 +def get_flashinfer_moe_backend() -> FlashinferMoeBackend:
249 + flashinfer_moe_backend = envs.VLLM_FLASHINFER_MOE_BACKEND
250 + if flashinfer_moe_backend == "throughput":
251 + return FlashinferMoeBackend.CUTLASS
252 + elif flashinfer_moe_backend == "latency":
253 + return FlashinferMoeBackend.TENSORRT_LLM
254 +
255 + allowed_backends = ["throughput", "latency"]
256 + raise ValueError(
257 + f"Unknown flashinfer moe backend: {flashinfer_moe_backend}"
258 + f" expected one of {allowed_backends}")