@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e8ebbdde83
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 5, 12:15 UTC

[Quantization] Add FlashInfer CuteDSL batched experts backend for NVFP4 MoE (#38251)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io>

Yongye Zhu committed Apr 6, 2026 at 14:57 UTC e8ebbdde8304a8cf89bbd4e101ebdfc25118b125
6 files changed +577 -248
tests/kernels/moe/test_cutedsl_moe.py
+1 -1
@@ -17,7 +17,7 @@ from flashinfer import fp4_quantize
17 from torch.nn import functional as F
18
19 from vllm.model_executor.layers.activation import SiluAndMul
20 -from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
20 +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import ( # noqa: E501
21 flashinfer_cutedsl_moe_masked,
22 )
23 from vllm.utils.flashinfer import (
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py new
+353
@@ -0,0 +1,353 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 +from vllm import envs
8 +from vllm.logger import init_logger
9 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10 +from vllm.model_executor.layers.fused_moe.config import (
11 + FusedMoEConfig,
12 + FusedMoEParallelConfig,
13 + FusedMoEQuantConfig,
14 +)
15 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
16 + TopKWeightAndReduceDelegate,
17 +)
18 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
19 + QuantKey,
20 + kNvfp4Dynamic,
21 + kNvfp4Static,
22 +)
23 +from vllm.platforms import current_platform
24 +from vllm.utils.flashinfer import (
25 + flashinfer_cutedsl_grouped_gemm_nt_masked,
26 + has_flashinfer_cutedsl_grouped_gemm_nt_masked,
27 + scaled_fp4_grouped_quantize,
28 + silu_and_mul_scaled_nvfp4_experts_quantize,
29 +)
30 +
31 +logger = init_logger(__name__)
32 +
33 +
34 +class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
35 + def __init__(
36 + self,
37 + moe_config: FusedMoEConfig,
38 + quant_config: FusedMoEQuantConfig,
39 + max_num_tokens: int,
40 + num_dispatchers: int,
41 + ):
42 + super().__init__(
43 + moe_config=moe_config,
44 + quant_config=quant_config,
45 + max_num_tokens=max_num_tokens,
46 + num_dispatchers=num_dispatchers,
47 + )
48 + assert quant_config.quant_dtype == "nvfp4", (
49 + "Only nvfp4 quantization are currently supported."
50 + )
51 + self.out_dtype = moe_config.in_dtype
52 +
53 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
54 + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
55 + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)
56 +
57 + @staticmethod
58 + def activation_format() -> mk.FusedMoEActivationFormat:
59 + return mk.FusedMoEActivationFormat.BatchedExperts
60 +
61 + @staticmethod
62 + def _supports_current_device() -> bool:
63 + p = current_platform
64 + return (
65 + p.is_cuda()
66 + and p.is_device_capability_family(100)
67 + and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
68 + )
69 +
70 + @staticmethod
71 + def _supports_no_act_and_mul() -> bool:
72 + return False
73 +
74 + @staticmethod
75 + def _supports_quant_scheme(
76 + weight_key: QuantKey | None,
77 + activation_key: QuantKey | None,
78 + ) -> bool:
79 + SUPPORTED_W_A = [
80 + (kNvfp4Static, kNvfp4Dynamic),
81 + ]
82 + return (weight_key, activation_key) in SUPPORTED_W_A
83 +
84 + @staticmethod
85 + def _supports_activation(activation: MoEActivation) -> bool:
86 + return activation == MoEActivation.SILU
87 +
88 + @staticmethod
89 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
90 + return True
91 +
92 + def supports_expert_map(self) -> bool:
93 + return False
94 +
95 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
96 + # Let PrepareAndFinalize::finalize() decide the impl.
97 + return TopKWeightAndReduceDelegate()
98 +
99 + def workspace_shapes(
100 + self,
101 + M: int,
102 + N: int,
103 + K: int,
104 + topk: int,
105 + global_num_experts: int,
106 + local_num_experts: int,
107 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
108 + activation: MoEActivation,
109 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
110 + """
111 + Compute the shapes for the temporary and final outputs of the two gemms
112 + and activation in the fused expert function. Since the gemms are
113 + independent, the workspace for the first gemm can be shared with the
114 + workspace for the last gemm.
115 +
116 + Returns a tuple of:
117 + - workspace13 shape tuple: must be large enough to hold the
118 + result of either expert gemm.
119 + - workspace2 shape tuple: must be large enough to hold the
120 + result of the activation function.
121 + - output shape tuple: must be exact size of the final gemm output.
122 + - Workspace type: The dtype to use for the workspace tensors.
123 + - Note: in order for activation chunking to work, the first dimension
124 + of each tuple must be the number of tokens.
125 + """
126 +
127 + # We use global_num_experts due to how moe_align_block_size handles
128 + # expert_maps.
129 + K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K
130 + output_shape = (local_num_experts, M, K_dim)
131 + workspace2 = (local_num_experts, M, N)
132 + workspace1 = output_shape
133 + return (workspace1, workspace2, output_shape)
134 +
135 + def apply(
136 + self,
137 + output: torch.Tensor,
138 + hidden_states: torch.Tensor,
139 + w1: torch.Tensor,
140 + w2: torch.Tensor,
141 + topk_weights: torch.Tensor,
142 + topk_ids: torch.Tensor,
143 + activation: MoEActivation,
144 + global_num_experts: int,
145 + expert_map: torch.Tensor | None,
146 + a1q_scale: torch.Tensor | None,
147 + a2_scale: torch.Tensor | None, # Not used
148 + workspace13: torch.Tensor | None,
149 + workspace2: torch.Tensor | None,
150 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
151 + apply_router_weight_on_input: bool | None,
152 + ):
153 + assert self.quant_dtype == "nvfp4", (
154 + "Only nvfp4 quantization are currently supported."
155 + )
156 + # Ensure w1_scale and w2_scale are not None before calling view
157 + assert self.w1_scale is not None and self.w2_scale is not None, (
158 + "w1_scale and w2_scale must not be None for FlashInferExperts"
159 + )
160 + assert expert_tokens_meta is not None
161 + expert_num_tokens = expert_tokens_meta.expert_num_tokens
162 + assert hidden_states.ndim == 3
163 + assert self.w1_scale.ndim == 3
164 + assert self.w2_scale.ndim == 3
165 +
166 + input_global_scale = (
167 + None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale
168 + )
169 + flashinfer_hidden_states = (
170 + (hidden_states, a1q_scale)
171 + if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
172 + else hidden_states
173 + )
174 + flashinfer_cutedsl_moe_masked(
175 + hidden_states=flashinfer_hidden_states,
176 + input_global_scale=input_global_scale,
177 + w1=w1,
178 + w1_blockscale=self.w1_scale,
179 + w1_alpha=self.g1_alphas,
180 + w2=w2,
181 + a2_global_scale=self.a2_gscale,
182 + w2_blockscale=self.w2_scale,
183 + w2_alpha=self.g2_alphas,
184 + masked_m=expert_num_tokens,
185 + workspace=workspace2,
186 + out=output,
187 + )
188 +
189 +
190 +def get_cute_dtype(input: torch.Tensor) -> str:
191 + if input.dtype == torch.bfloat16:
192 + return "bfloat16"
193 + elif input.dtype == torch.float16:
194 + return "float16"
195 + elif input.dtype == torch.float32:
196 + return "float32"
197 + else:
198 + raise ValueError(f"Unsupported cute dtype {input.dtype}")
199 +
200 +
201 +def flashinfer_cutedsl_moe_masked(
202 + hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
203 + input_global_scale: torch.Tensor,
204 + w1: torch.Tensor,
205 + w1_blockscale: torch.Tensor,
206 + w1_alpha,
207 + w2: torch.Tensor,
208 + a2_global_scale: torch.Tensor,
209 + w2_blockscale: torch.Tensor,
210 + w2_alpha,
211 + masked_m: torch.Tensor,
212 + workspace: torch.Tensor,
213 + out: torch.Tensor,
214 +):
215 + """
216 + Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL
217 + kernels.
218 +
219 + Args:
220 + hidden_states: Either of the following case
221 + * torch.Tensor: [num_experts, m, k], bf16
222 + * tuple[torch.Tensor, torch.Tensor]: [num_experts, m, k // 2],
223 + uint8, [num_experts, m, k // 16], float8_e4m3fn
224 + input_global_scale (torch.Tensor): (l,)
225 + w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8
226 + w1_blockscale (torch.Tensor): blockscale factors, e4m3,
227 + w1_alpha (torch.Tensor): (l,)
228 + w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8
229 + a2_global_scale (torch.Tensor): (l,)
230 + w2_blockscale (torch.Tensor): blockscale factors, e4m3,
231 + w2_alpha (torch.Tensor): (l,)
232 + masked_m (torch.Tensor): Masked dimension indices
233 + workspace (torch.Tensor): For gateup_output
234 +
235 + Notes:
236 + - Assumes max(masked_m) <= m.
237 + """
238 +
239 + # === Assertions on dtypes ===
240 + assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}"
241 + assert w1_blockscale.dtype == torch.float8_e4m3fn, (
242 + f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}"
243 + )
244 + assert w1_alpha.dtype == torch.float32, (
245 + f"w1_alpha must be float32, got {w1_alpha.dtype}"
246 + )
247 + assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}"
248 + assert a2_global_scale.dtype == torch.float32, (
249 + f"a2_global_scale must be float32, got {a2_global_scale.dtype}"
250 + )
251 + assert w2_blockscale.dtype == torch.float8_e4m3fn, (
252 + f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}"
253 + )
254 + assert w2_alpha.dtype == torch.float32, (
255 + f"w2_alpha must be float32, got {w2_alpha.dtype}"
256 + )
257 +
258 + # === Assertions on shapes ===
259 + n = w2.shape[-1] * 2 # intermediate dimension
260 + if isinstance(hidden_states, tuple):
261 + assert input_global_scale is None, (
262 + "input_global_scale is needed when input needs quant"
263 + )
264 +
265 + aq = hidden_states[0].view(torch.uint8)
266 + aq_sf = hidden_states[1].view(torch.float8_e4m3fn)
267 + # m, k_by_2, num_experts = aq.shape
268 + num_experts, m, k_by_2 = aq.shape
269 + k = k_by_2 * 2
270 + aq = aq.permute(1, 2, 0)
271 + else:
272 + num_experts, m, k = hidden_states.shape
273 +
274 + assert input_global_scale.dtype == torch.float32, (
275 + f"input_global_scale must be float32, got {input_global_scale.dtype}"
276 + )
277 + assert input_global_scale.shape == (num_experts,), (
278 + f"input_global_scale must be (l,), got {input_global_scale.shape}"
279 + )
280 +
281 + aq, aq_sf = scaled_fp4_grouped_quantize(
282 + hidden_states,
283 + masked_m,
284 + input_global_scale,
285 + )
286 +
287 + assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}"
288 + assert w1.shape[-1] * 2 == k, (
289 + f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}"
290 + )
291 + assert w2.shape[-2:] == (
292 + k,
293 + n // 2,
294 + ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}"
295 +
296 + assert w1_alpha.shape == (num_experts,), (
297 + f"w1_alpha must be (l,), got {w1_alpha.shape}"
298 + )
299 + assert a2_global_scale.shape == (num_experts,), (
300 + f"a2_global_scale must be (l,), got {a2_global_scale.shape}"
301 + )
302 + assert w2_alpha.shape == (num_experts,), (
303 + f"w2_alpha must be (l,), got {w2_alpha.shape}"
304 + )
305 +
306 + workspace = workspace.permute(1, 2, 0) # requirement of kernel
307 + sf_vec_size = 16
308 + assert aq_sf.dtype == torch.float8_e4m3fn
309 + assert aq.dtype == torch.uint8
310 + ab_dtype = "float4_e2m1fn"
311 + sf_dtype = "float8_e4m3fn"
312 +
313 + if isinstance(hidden_states, tuple):
314 + c_dtype = "bfloat16"
315 + else:
316 + c_dtype = get_cute_dtype(hidden_states)
317 +
318 + # Gemm1
319 + flashinfer_cutedsl_grouped_gemm_nt_masked(
320 + (aq, aq_sf),
321 + (w1.permute(1, 2, 0), w1_blockscale),
322 + workspace,
323 + masked_m,
324 + ab_dtype=ab_dtype,
325 + sf_dtype=sf_dtype,
326 + c_dtype=c_dtype,
327 + sf_vec_size=sf_vec_size,
328 + alpha=w1_alpha.view(1, 1, num_experts),
329 + alpha_dtype=get_cute_dtype(w1_alpha),
330 + ) # in logical [m, n, l]
331 +
332 + # SILU and quantization
333 + diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize(
334 + workspace.permute(2, 0, 1),
335 + masked_m,
336 + a2_global_scale,
337 + )
338 +
339 + # Gemm2
340 + out = out.permute(1, 2, 0) # requirement of kernel
341 + flashinfer_cutedsl_grouped_gemm_nt_masked(
342 + (diq, diq_sf),
343 + (w2.permute(1, 2, 0), w2_blockscale),
344 + out,
345 + masked_m,
346 + ab_dtype=ab_dtype,
347 + sf_dtype=sf_dtype,
348 + c_dtype=c_dtype,
349 + sf_vec_size=sf_vec_size,
350 + alpha=w2_alpha.view(1, 1, num_experts),
351 + alpha_dtype=get_cute_dtype(w2_alpha),
352 + ) # in logical [m, k, l]
353 + out = out.permute(2, 0, 1)
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+64 -244
@@ -4,8 +4,6 @@
4 import torch
5
6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 -from vllm import envs
8 -from vllm.logger import init_logger
7 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
8 from vllm.model_executor.layers.fused_moe.config import (
9 FusedMoEConfig,
@@ -13,7 +11,7 @@ from vllm.model_executor.layers.fused_moe.config import (
11 FusedMoEQuantConfig,
12 )
13 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
16 - TopKWeightAndReduceDelegate,
14 + TopKWeightAndReduceNoOP,
15 )
16 from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 QuantKey,
@@ -22,33 +20,42 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
20 )
21 from vllm.platforms import current_platform
22 from vllm.utils.flashinfer import (
25 - flashinfer_cutedsl_grouped_gemm_nt_masked,
26 - has_flashinfer_cutedsl_grouped_gemm_nt_masked,
27 - scaled_fp4_grouped_quantize,
28 - silu_and_mul_scaled_nvfp4_experts_quantize,
23 + flashinfer_cute_dsl_fused_moe_nvfp4,
24 + has_flashinfer_cutedsl_moe_nvfp4,
25 )
26
31 -logger = init_logger(__name__)
32 -
27
28 class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
29 + """
30 + CuteDSL NvFP4 MoE experts using the FlashInfer functional API.
31 +
32 + Uses Standard activation format (non-batched). The kernel handles
33 + routing, expert computation, and reduction internally.
34 + Supports expert parallelism natively.
35 + """
36 +
37 def __init__(
38 self,
39 moe_config: FusedMoEConfig,
40 quant_config: FusedMoEQuantConfig,
39 - max_num_tokens: int,
40 - num_dispatchers: int,
41 ):
42 super().__init__(
43 moe_config=moe_config,
44 quant_config=quant_config,
45 - max_num_tokens=max_num_tokens,
46 - num_dispatchers=num_dispatchers,
45 )
46 assert quant_config.quant_dtype == "nvfp4", (
49 - "Only nvfp4 quantization are currently supported."
47 + "Only nvfp4 quantization is currently supported."
48 )
49 self.out_dtype = moe_config.in_dtype
50 + self.hidden_dim = moe_config.hidden_dim
51 + self.intermediate_size_per_partition = (
52 + moe_config.intermediate_size_per_partition
53 + )
54 + self.topk = moe_config.experts_per_token
55 + self.local_num_experts = moe_config.num_local_experts
56 + self.global_num_experts = moe_config.num_experts
57 + self.ep_rank = moe_config.moe_parallel_config.ep_rank
58 + self.local_expert_offset = self.ep_rank * self.local_num_experts
59
60 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
61 layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -56,7 +63,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
63
64 @staticmethod
65 def activation_format() -> mk.FusedMoEActivationFormat:
59 - return mk.FusedMoEActivationFormat.BatchedExperts
66 + return mk.FusedMoEActivationFormat.Standard
67
68 @staticmethod
69 def _supports_current_device() -> bool:
@@ -64,7 +71,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
71 return (
72 p.is_cuda()
73 and p.is_device_capability_family(100)
67 - and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
74 + and has_flashinfer_cutedsl_moe_nvfp4()
75 )
76
77 @staticmethod
@@ -86,15 +93,16 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
93 return activation == MoEActivation.SILU
94
95 @staticmethod
89 - def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
96 + def _supports_parallel_config(
97 + moe_parallel_config: FusedMoEParallelConfig,
98 + ) -> bool:
99 return True
100
101 def supports_expert_map(self) -> bool:
102 return False
103
104 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
96 - # Let PrepareAndFinalize::finalize() decide the impl.
97 - return TopKWeightAndReduceDelegate()
105 + return TopKWeightAndReduceNoOP()
106
107 def workspace_shapes(
108 self,
@@ -107,29 +115,12 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
115 expert_tokens_meta: mk.ExpertTokensMetadata | None,
116 activation: MoEActivation,
117 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
110 - # We use global_num_experts due to how moe_align_block_size handles
111 - # expert_maps.
112 - """
113 - Compute the shapes for the temporary and final outputs of the two gemms
114 - and activation in the fused expert function. Since the gemms are
115 - independent, the workspace for the first gemm can be shared with the
116 - workspace for the last gemm.
117 -
118 - Returns a tuple of:
119 - - workspace13 shape tuple: must be large enough to hold the
120 - result of either expert gemm.
121 - - workspace2 shape tuple: must be large enough to hold the
122 - result of the activation function.
123 - - output shape tuple: must be exact size of the final gemm output.
124 - - Workspace type: The dtype to use for the workspace tensors.
125 - - Note: in order for activation chunking to work, the first dimension
126 - of each tuple must be the number of tokens.
127 - """
128 - K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K
129 - output_shape = (local_num_experts, M, K_dim)
130 - workspace2 = (local_num_experts, M, N)
131 - workspace1 = output_shape
132 - return (workspace1, workspace2, output_shape)
118 + workspace1 = (0,)
119 + workspace2 = (0,)
120 + # K is packed (K//2 for uint8), so output uses hidden_dim.
121 + assert self.hidden_dim == K * 2
122 + output = (M, self.hidden_dim)
123 + return (workspace1, workspace2, output)
124
125 def apply(
126 self,
@@ -143,210 +134,39 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
134 global_num_experts: int,
135 expert_map: torch.Tensor | None,
136 a1q_scale: torch.Tensor | None,
146 - a2_scale: torch.Tensor | None, # Not used
137 + a2_scale: torch.Tensor | None,
138 workspace13: torch.Tensor | None,
139 workspace2: torch.Tensor | None,
140 expert_tokens_meta: mk.ExpertTokensMetadata | None,
141 apply_router_weight_on_input: bool | None,
142 ):
152 - assert self.quant_dtype == "nvfp4", (
153 - "Only nvfp4 quantization are currently supported."
154 - )
155 - # Ensure w1_scale and w2_scale are not None before calling view
156 - assert self.w1_scale is not None and self.w2_scale is not None, (
157 - "w1_scale and w2_scale must not be None for FlashInferExperts"
158 - )
159 - assert expert_tokens_meta is not None
160 - expert_num_tokens = expert_tokens_meta.expert_num_tokens
161 - assert hidden_states.ndim == 3
162 - assert self.w1_scale.ndim == 3
163 - assert self.w2_scale.ndim == 3
164 -
165 - input_global_scale = (
166 - None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale
167 - )
168 - flashinfer_hidden_states = (
169 - (hidden_states, a1q_scale)
170 - if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
171 - else hidden_states
172 - )
173 - flashinfer_cutedsl_moe_masked(
174 - hidden_states=flashinfer_hidden_states,
175 - input_global_scale=input_global_scale,
176 - w1=w1,
177 - w1_blockscale=self.w1_scale,
178 - w1_alpha=self.g1_alphas,
179 - w2=w2,
180 - a2_global_scale=self.a2_gscale,
181 - w2_blockscale=self.w2_scale,
182 - w2_alpha=self.g2_alphas,
183 - masked_m=expert_num_tokens,
184 - workspace=workspace2,
185 - out=output,
186 - )
187 -
188 -
189 -def get_cute_dtype(input: torch.Tensor) -> str:
190 - if input.dtype == torch.bfloat16:
191 - return "bfloat16"
192 - elif input.dtype == torch.float16:
193 - return "float16"
194 - elif input.dtype == torch.float32:
195 - return "float32"
196 - else:
197 - raise ValueError(f"Unsupported cute dtype {input.dtype}")
198 -
199 -
200 -def flashinfer_cutedsl_moe_masked(
201 - hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
202 - input_global_scale: torch.Tensor,
203 - w1: torch.Tensor,
204 - w1_blockscale: torch.Tensor,
205 - w1_alpha,
206 - w2: torch.Tensor,
207 - a2_global_scale: torch.Tensor,
208 - w2_blockscale: torch.Tensor,
209 - w2_alpha,
210 - masked_m: torch.Tensor,
211 - workspace: torch.Tensor,
212 - out: torch.Tensor,
213 -):
214 - """
215 - Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL
216 - kernels.
217 -
218 - Args:
219 - hidden_states: Either of the following case
220 - * torch.Tensor: [num_experts, m, k], bf16
221 - * tuple[torch.Tensor, torch.Tensor]: [num_experts, m, k // 2],
222 - uint8, [num_experts, m, k // 16], float8_e4m3fn
223 - input_global_scale (torch.Tensor): (l,)
224 - w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8
225 - w1_blockscale (torch.Tensor): blockscale factors, e4m3,
226 - w1_alpha (torch.Tensor): (l,)
227 - w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8
228 - a2_global_scale (torch.Tensor): (l,)
229 - w2_blockscale (torch.Tensor): blockscale factors, e4m3,
230 - w2_alpha (torch.Tensor): (l,)
231 - masked_m (torch.Tensor): Masked dimension indices
232 - workspace (torch.Tensor): For gateup_output
233 -
234 - Notes:
235 - - Assumes max(masked_m) <= m.
236 - """
237 -
238 - # === Assertions on dtypes ===
239 - assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}"
240 - assert w1_blockscale.dtype == torch.float8_e4m3fn, (
241 - f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}"
242 - )
243 - assert w1_alpha.dtype == torch.float32, (
244 - f"w1_alpha must be float32, got {w1_alpha.dtype}"
245 - )
246 - assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}"
247 - assert a2_global_scale.dtype == torch.float32, (
248 - f"a2_global_scale must be float32, got {a2_global_scale.dtype}"
249 - )
250 - assert w2_blockscale.dtype == torch.float8_e4m3fn, (
251 - f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}"
252 - )
253 - assert w2_alpha.dtype == torch.float32, (
254 - f"w2_alpha must be float32, got {w2_alpha.dtype}"
255 - )
256 -
257 - # === Assertions on shapes ===
258 - n = w2.shape[-1] * 2 # intermediate dimension
259 - if isinstance(hidden_states, tuple):
260 - assert input_global_scale is None, (
261 - "input_global_scale is needed when input needs quant"
262 - )
263 -
264 - aq = hidden_states[0].view(torch.uint8)
265 - aq_sf = hidden_states[1].view(torch.float8_e4m3fn)
266 - # m, k_by_2, num_experts = aq.shape
267 - num_experts, m, k_by_2 = aq.shape
268 - k = k_by_2 * 2
269 - aq = aq.permute(1, 2, 0)
270 - else:
271 - num_experts, m, k = hidden_states.shape
272 -
273 - assert input_global_scale.dtype == torch.float32, (
274 - f"input_global_scale must be float32, got {input_global_scale.dtype}"
275 - )
276 - assert input_global_scale.shape == (num_experts,), (
277 - f"input_global_scale must be (l,), got {input_global_scale.shape}"
278 - )
279 -
280 - aq, aq_sf = scaled_fp4_grouped_quantize(
281 - hidden_states,
282 - masked_m,
283 - input_global_scale,
284 - )
285 -
286 - assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}"
287 - assert w1.shape[-1] * 2 == k, (
288 - f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}"
289 - )
290 - assert w2.shape[-2:] == (
291 - k,
292 - n // 2,
293 - ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}"
294 -
295 - assert w1_alpha.shape == (num_experts,), (
296 - f"w1_alpha must be (l,), got {w1_alpha.shape}"
297 - )
298 - assert a2_global_scale.shape == (num_experts,), (
299 - f"a2_global_scale must be (l,), got {a2_global_scale.shape}"
300 - )
301 - assert w2_alpha.shape == (num_experts,), (
302 - f"w2_alpha must be (l,), got {w2_alpha.shape}"
303 - )
304 -
305 - workspace = workspace.permute(1, 2, 0) # requirement of kernel
306 - sf_vec_size = 16
307 - assert aq_sf.dtype == torch.float8_e4m3fn
308 - assert aq.dtype == torch.uint8
309 - ab_dtype = "float4_e2m1fn"
310 - sf_dtype = "float8_e4m3fn"
311 -
312 - if isinstance(hidden_states, tuple):
313 - c_dtype = "bfloat16"
314 - else:
315 - c_dtype = get_cute_dtype(hidden_states)
316 -
317 - # Gemm1
318 - flashinfer_cutedsl_grouped_gemm_nt_masked(
319 - (aq, aq_sf),
320 - (w1.permute(1, 2, 0), w1_blockscale),
321 - workspace,
322 - masked_m,
323 - ab_dtype=ab_dtype,
324 - sf_dtype=sf_dtype,
325 - c_dtype=c_dtype,
326 - sf_vec_size=sf_vec_size,
327 - alpha=w1_alpha.view(1, 1, num_experts),
328 - alpha_dtype=get_cute_dtype(w1_alpha),
329 - ) # in logical [m, n, l]
330 -
331 - # SILU and quantization
332 - diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize(
333 - workspace.permute(2, 0, 1),
334 - masked_m,
335 - a2_global_scale,
336 - )
337 -
338 - # Gemm2
339 - out = out.permute(1, 2, 0) # requirement of kernel
340 - flashinfer_cutedsl_grouped_gemm_nt_masked(
341 - (diq, diq_sf),
342 - (w2.permute(1, 2, 0), w2_blockscale),
343 - out,
344 - masked_m,
345 - ab_dtype=ab_dtype,
346 - sf_dtype=sf_dtype,
347 - c_dtype=c_dtype,
348 - sf_vec_size=sf_vec_size,
349 - alpha=w2_alpha.view(1, 1, num_experts),
350 - alpha_dtype=get_cute_dtype(w2_alpha),
351 - ) # in logical [m, k, l]
352 - out = out.permute(2, 0, 1)
143 + assert self.quant_dtype == "nvfp4"
144 + assert a1q_scale is not None
145 + assert self.w1_scale is not None
146 + assert self.w2_scale is not None
147 +
148 + # a1q_scale is (M, K//16) float8_e4m3fn from fp4_quantize.
149 + # The functional API expects x_sf with trailing dim: (M, K//16, 1).
150 + x_sf = a1q_scale.unsqueeze(-1)
151 +
152 + from vllm.utils.flashinfer import _is_fi_autotuning, autotune
153 +
154 + with autotune(_is_fi_autotuning):
155 + flashinfer_cute_dsl_fused_moe_nvfp4(
156 + x=hidden_states,
157 + x_sf=x_sf,
158 + token_selected_experts=topk_ids.to(torch.int32),
159 + token_final_scales=topk_weights.float(),
160 + w1_weight=w1,
161 + w1_weight_sf=self.w1_scale,
162 + w1_alpha=self.g1_alphas,
163 + fc2_input_scale=self.a2_gscale,
164 + w2_weight=w2,
165 + w2_weight_sf=self.w2_scale,
166 + w2_alpha=self.g2_alphas,
167 + num_experts=self.global_num_experts,
168 + top_k=self.topk,
169 + num_local_experts=self.local_num_experts,
170 + local_expert_offset=self.local_expert_offset,
171 + moe_output=output,
172 + )
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+46 -2
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
22 )
23 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
24 prepare_nvfp4_moe_layer_for_fi_or_cutlass,
25 + prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
26 )
27 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
28 FlashinferMoeBackend,
@@ -41,6 +42,7 @@ class NvFp4MoeBackend(Enum):
42 FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
43 FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS"
44 FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL"
45 + FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED"
46 VLLM_CUTLASS = "VLLM_CUTLASS"
47 MARLIN = "MARLIN"
48
@@ -49,6 +51,7 @@ FLASHINFER_NVFP4_MOE_BACKENDS = [
51 NvFp4MoeBackend.FLASHINFER_TRTLLM,
52 NvFp4MoeBackend.FLASHINFER_CUTLASS,
53 NvFp4MoeBackend.FLASHINFER_CUTEDSL,
54 + NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
55 ]
56
57 fi_2_vllm_backend_map: dict[FlashinferMoeBackend, NvFp4MoeBackend] = {
@@ -95,6 +98,13 @@ def backend_to_kernel_cls(
98
99 return [FlashInferCuteDSLExperts]
100
101 + elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED:
102 + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import ( # noqa: E501
103 + FlashInferCuteDSLBatchedExperts,
104 + )
105 +
106 + return [FlashInferCuteDSLBatchedExperts]
107 +
108 elif backend == NvFp4MoeBackend.VLLM_CUTLASS:
109 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
110 CutlassExpertsFp4,
@@ -143,6 +153,7 @@ def select_nvfp4_moe_backend(
153 AVAILABLE_BACKENDS = [
154 NvFp4MoeBackend.FLASHINFER_TRTLLM,
155 NvFp4MoeBackend.FLASHINFER_CUTEDSL,
156 + NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
157 NvFp4MoeBackend.FLASHINFER_CUTLASS,
158 NvFp4MoeBackend.VLLM_CUTLASS,
159 NvFp4MoeBackend.MARLIN,
@@ -198,6 +209,12 @@ def select_nvfp4_moe_backend(
209 runner_backend = config.moe_backend
210 if runner_backend != "auto":
211 requested_backend = map_nvfp4_backend(runner_backend)
212 + # For batched activation format, use batched variant if available.
213 + if (
214 + activation_format == mk.FusedMoEActivationFormat.BatchedExperts
215 + and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
216 + ):
217 + requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
218 return _return_or_raise(
219 requested_backend, config, weight_key, activation_key, activation_format
220 )
@@ -288,7 +305,28 @@ def convert_to_nvfp4_moe_kernel_format(
305 torch.Tensor,
306 torch.Tensor,
307 ]:
291 - if (
308 + if nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
309 + (
310 + w13,
311 + w13_scale,
312 + w13_scale_2,
313 + a13_scale,
314 + w2,
315 + w2_scale,
316 + w2_scale_2,
317 + a2_scale,
318 + ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
319 + layer=layer,
320 + w13=w13,
321 + w13_scale=w13_scale,
322 + w13_scale_2=w13_scale_2,
323 + a13_scale=a13_scale,
324 + w2=w2,
325 + w2_scale=w2_scale,
326 + w2_scale_2=w2_scale_2,
327 + a2_scale=a2_scale,
328 + )
329 + elif (
330 nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS
331 or nvfp4_backend == NvFp4MoeBackend.VLLM_CUTLASS
332 ):
@@ -380,7 +418,13 @@ def make_nvfp4_moe_quant_config(
418 # NOTE(rob): this is a hack until the MoE kernels
419 # create their own quant configs. TRTLLM kernel
420 # does not accept swizzled input quant scales.
383 - is_nvfp4_scale_swizzled=(backend != NvFp4MoeBackend.FLASHINFER_TRTLLM),
421 + is_nvfp4_scale_swizzled=(
422 + backend
423 + not in (
424 + NvFp4MoeBackend.FLASHINFER_TRTLLM,
425 + NvFp4MoeBackend.FLASHINFER_CUTEDSL,
426 + )
427 + ),
428 )
429
430
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+95 -1
@@ -60,6 +60,100 @@ def reorder_w1w3_to_w3w1(
60 )
61
62
63 +def interleave_linear_and_gate(
64 + x: torch.Tensor,
65 + group_size: int = 64,
66 + dim: int = -1,
67 +) -> torch.Tensor:
68 + """Interleave gate and linear weight rows for CuteDSL wrapper."""
69 + sizes = x.size()
70 + dim = dim % x.dim()
71 + assert sizes[dim] % (group_size * 2) == 0, (
72 + f"dim {dim} size {sizes[dim]} must be divisible by {group_size * 2}"
73 + )
74 + prev_sizes = sizes[:dim]
75 + post_sizes = sizes[dim + 1 :]
76 + x = x.view(*prev_sizes, 2, sizes[dim] // (group_size * 2), group_size, *post_sizes)
77 + x = x.transpose(dim, dim + 1).contiguous().view(*sizes)
78 + return x
79 +
80 +
81 +def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
82 + layer: "FusedMoE",
83 + w13: torch.Tensor,
84 + w13_scale: torch.Tensor,
85 + w13_scale_2: torch.Tensor,
86 + a13_scale: torch.Tensor,
87 + w2: torch.Tensor,
88 + w2_scale: torch.Tensor,
89 + w2_scale_2: torch.Tensor,
90 + a2_scale: torch.Tensor,
91 +) -> tuple[
92 + torch.Tensor,
93 + torch.Tensor,
94 + torch.Tensor,
95 + torch.Tensor,
96 + torch.Tensor,
97 + torch.Tensor,
98 + torch.Tensor,
99 + torch.Tensor,
100 +]:
101 + """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend.
102 +
103 + Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper,
104 + and interleaves w13 gate/linear rows.
105 + """
106 + from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
107 +
108 + # Global scaling factors (same as other FlashInfer backends).
109 + num_experts = w13.shape[0]
110 + a13_scale = a13_scale.max().to(torch.float32).expand(num_experts)
111 + a2_scale = a2_scale.max().to(torch.float32).expand(num_experts)
112 +
113 + half = w13.shape[1] // 2
114 + w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1)
115 + w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1)
116 +
117 + # Interleave up/gate rows for w13 weights and scales.
118 + w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)
119 + w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)
120 +
121 + # Convert w13 scale factors: linear → swizzled → MMA layout.
122 + w13_scale = swizzle_blockscale(w13_scale)
123 + E, M_padded, K_sf_padded = w13_scale.shape
124 + w13_scale_flat = w13_scale.reshape(E * M_padded, K_sf_padded)
125 + w13_scale = convert_sf_to_mma_layout(
126 + w13_scale_flat,
127 + m=M_padded,
128 + k=K_sf_padded * 16,
129 + num_groups=E,
130 + sf_vec_size=16,
131 + )
132 +
133 + # Convert w2 scale factors: linear → swizzled → MMA layout.
134 + w2_scale = swizzle_blockscale(w2_scale)
135 + E, M_padded, K_sf_padded = w2_scale.shape
136 + w2_scale_flat = w2_scale.reshape(E * M_padded, K_sf_padded)
137 + w2_scale = convert_sf_to_mma_layout(
138 + w2_scale_flat,
139 + m=M_padded,
140 + k=K_sf_padded * 16,
141 + num_groups=E,
142 + sf_vec_size=16,
143 + )
144 +
145 + return (
146 + w13,
147 + w13_scale,
148 + w13_scale_2,
149 + a13_scale,
150 + w2,
151 + w2_scale,
152 + w2_scale_2,
153 + a2_scale,
154 + )
155 +
156 +
157 def prepare_static_weights_for_trtllm_fp4_moe(
158 # args_dequant,
159 # args,
@@ -221,7 +315,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
315 NvFp4MoeBackend.VLLM_CUTLASS,
316 NvFp4MoeBackend.FLASHINFER_CUTLASS,
317 NvFp4MoeBackend.FLASHINFER_TRTLLM,
224 - NvFp4MoeBackend.FLASHINFER_CUTEDSL,
318 + NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
319 ]
320
321 # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels.
vllm/utils/flashinfer.py
+18
@@ -128,6 +128,12 @@ scaled_fp4_grouped_quantize = _lazy_import_wrapper(
128 nvfp4_block_scale_interleave = _lazy_import_wrapper(
129 "flashinfer.fp4_quantization", "block_scale_interleave"
130 )
131 +flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper(
132 + "flashinfer", "cute_dsl_fused_moe_nvfp4"
133 +)
134 +flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper(
135 + "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout"
136 +)
137 trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
138 "flashinfer", "trtllm_fp4_block_scale_moe"
139 )
@@ -252,6 +258,15 @@ def has_flashinfer_cutedsl_grouped_gemm_nt_masked() -> bool:
258 return True
259
260
261 +@functools.cache
262 +def has_flashinfer_cutedsl_moe_nvfp4() -> bool:
263 + """Return ``True`` if FlashInfer cute_dsl_fused_moe_nvfp4 is available."""
264 + if not has_flashinfer_cutedsl():
265 + return False
266 + mod = _get_submodule("flashinfer")
267 + return mod is not None and hasattr(mod, "cute_dsl_fused_moe_nvfp4")
268 +
269 +
270 @functools.cache
271 def has_nvidia_artifactory() -> bool:
272 """Return `True` if NVIDIA's artifactory is accessible.
@@ -768,6 +783,8 @@ __all__ = [
783 "silu_and_mul_scaled_nvfp4_experts_quantize",
784 "scaled_fp4_grouped_quantize",
785 "nvfp4_block_scale_interleave",
786 + "flashinfer_cute_dsl_fused_moe_nvfp4",
787 + "flashinfer_convert_sf_to_mma_layout",
788 "trtllm_fp4_block_scale_moe",
789 "autotune",
790 "has_flashinfer_moe",
@@ -776,6 +793,7 @@ __all__ = [
793 "has_flashinfer_nvlink_one_sided",
794 "has_flashinfer_cutlass_fused_moe",
795 "has_flashinfer_cutedsl_grouped_gemm_nt_masked",
796 + "has_flashinfer_cutedsl_moe_nvfp4",
797 "has_flashinfer_fp8_blockscale_gemm",
798 "has_nvidia_artifactory",
799 "supports_trtllm_attention",