@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / be1a85b7a2
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, 23:30 UTC

Revert "[MoE Kernel] Flashinfer nvfp4 cutedsl moe kernel integration" (#38050) (#38169)

Co-authored-by: Zhewen Li <zhewenli@inferact.ai>

Zhewen Li committed Mar 26, 2026 at 07:59 UTC be1a85b7a2929f25c93d469fdd733a3576609e70
6 files changed +248 -577
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_batched_moe import ( # noqa: E501
20 +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
21 flashinfer_cutedsl_moe_masked,
22 )
23 from vllm.utils.flashinfer import (
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py deleted
-353
@@ -1,353 +0,0 @@
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
+244 -64
@@ -4,6 +4,8 @@
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,
@@ -11,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.config import (
13 FusedMoEQuantConfig,
14 )
15 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14 - TopKWeightAndReduceNoOP,
16 + TopKWeightAndReduceDelegate,
17 )
18 from vllm.model_executor.layers.quantization.utils.quant_utils import (
19 QuantKey,
@@ -20,42 +22,33 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
22 )
23 from vllm.platforms import current_platform
24 from vllm.utils.flashinfer import (
23 - flashinfer_cute_dsl_fused_moe_nvfp4,
24 - has_flashinfer_cutedsl_moe_nvfp4,
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
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 - """
33
34 +class FlashInferCuteDSLExperts(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", (
47 - "Only nvfp4 quantization is currently supported."
49 + "Only nvfp4 quantization are currently supported."
50 )
51 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
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)
@@ -63,7 +56,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
56
57 @staticmethod
58 def activation_format() -> mk.FusedMoEActivationFormat:
66 - return mk.FusedMoEActivationFormat.Standard
59 + return mk.FusedMoEActivationFormat.BatchedExperts
60
61 @staticmethod
62 def _supports_current_device() -> bool:
@@ -71,7 +64,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
64 return (
65 p.is_cuda()
66 and p.is_device_capability_family(100)
74 - and has_flashinfer_cutedsl_moe_nvfp4()
67 + and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
68 )
69
70 @staticmethod
@@ -93,16 +86,15 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
86 return activation == MoEActivation.SILU
87
88 @staticmethod
96 - def _supports_parallel_config(
97 - moe_parallel_config: FusedMoEParallelConfig,
98 - ) -> bool:
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:
105 - return TopKWeightAndReduceNoOP()
96 + # Let PrepareAndFinalize::finalize() decide the impl.
97 + return TopKWeightAndReduceDelegate()
98
99 def workspace_shapes(
100 self,
@@ -115,12 +107,29 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
107 expert_tokens_meta: mk.ExpertTokensMetadata | None,
108 activation: MoEActivation,
109 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
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)
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)
133
134 def apply(
135 self,
@@ -134,39 +143,210 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
143 global_num_experts: int,
144 expert_map: torch.Tensor | None,
145 a1q_scale: torch.Tensor | None,
137 - a2_scale: torch.Tensor | None,
146 + a2_scale: torch.Tensor | None, # Not used
147 workspace13: torch.Tensor | None,
148 workspace2: torch.Tensor | None,
149 expert_tokens_meta: mk.ExpertTokensMetadata | None,
150 apply_router_weight_on_input: bool | None,
151 ):
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 - )
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)
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+2 -46
@@ -19,7 +19,6 @@ from vllm.model_executor.layers.fused_moe.config import (
19 )
20 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
21 prepare_nvfp4_moe_layer_for_fi_or_cutlass,
22 - prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
22 )
23 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
24 FlashinferMoeBackend,
@@ -39,7 +38,6 @@ class NvFp4MoeBackend(Enum):
38 FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
39 FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS"
40 FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL"
42 - FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED"
41 VLLM_CUTLASS = "VLLM_CUTLASS"
42 MARLIN = "MARLIN"
43
@@ -48,7 +46,6 @@ FLASHINFER_NVFP4_MOE_BACKENDS = [
46 NvFp4MoeBackend.FLASHINFER_TRTLLM,
47 NvFp4MoeBackend.FLASHINFER_CUTLASS,
48 NvFp4MoeBackend.FLASHINFER_CUTEDSL,
51 - NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
49 ]
50
51 fi_2_vllm_backend_map: dict[FlashinferMoeBackend, NvFp4MoeBackend] = {
@@ -95,13 +92,6 @@ def backend_to_kernel_cls(
92
93 return [FlashInferCuteDSLExperts]
94
98 - elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED:
99 - from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import ( # noqa: E501
100 - FlashInferCuteDSLBatchedExperts,
101 - )
102 -
103 - return [FlashInferCuteDSLBatchedExperts]
104 -
95 elif backend == NvFp4MoeBackend.VLLM_CUTLASS:
96 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
97 CutlassExpertsFp4,
@@ -150,7 +140,6 @@ def select_nvfp4_moe_backend(
140 AVAILABLE_BACKENDS = [
141 NvFp4MoeBackend.FLASHINFER_TRTLLM,
142 NvFp4MoeBackend.FLASHINFER_CUTEDSL,
153 - NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
143 NvFp4MoeBackend.FLASHINFER_CUTLASS,
144 NvFp4MoeBackend.VLLM_CUTLASS,
145 NvFp4MoeBackend.MARLIN,
@@ -206,12 +195,6 @@ def select_nvfp4_moe_backend(
195 runner_backend = config.moe_backend
196 if runner_backend != "auto":
197 requested_backend = map_nvfp4_backend(runner_backend)
209 - # For batched activation format, use batched variant if available.
210 - if (
211 - activation_format == mk.FusedMoEActivationFormat.BatchedExperts
212 - and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
213 - ):
214 - requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
198 return _return_or_raise(
199 requested_backend, config, weight_key, activation_key, activation_format
200 )
@@ -302,28 +285,7 @@ def convert_to_nvfp4_moe_kernel_format(
285 torch.Tensor,
286 torch.Tensor,
287 ]:
305 - if nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
306 - (
307 - w13,
308 - w13_scale,
309 - w13_scale_2,
310 - a13_scale,
311 - w2,
312 - w2_scale,
313 - w2_scale_2,
314 - a2_scale,
315 - ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
316 - layer=layer,
317 - w13=w13,
318 - w13_scale=w13_scale,
319 - w13_scale_2=w13_scale_2,
320 - a13_scale=a13_scale,
321 - w2=w2,
322 - w2_scale=w2_scale,
323 - w2_scale_2=w2_scale_2,
324 - a2_scale=a2_scale,
325 - )
326 - elif (
288 + if (
289 nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS
290 or nvfp4_backend == NvFp4MoeBackend.VLLM_CUTLASS
291 ):
@@ -415,13 +377,7 @@ def make_nvfp4_moe_quant_config(
377 # NOTE(rob): this is a hack until the MoE kernels
378 # create their own quant configs. TRTLLM kernel
379 # does not accept swizzled input quant scales.
418 - is_nvfp4_scale_swizzled=(
419 - backend
420 - not in (
421 - NvFp4MoeBackend.FLASHINFER_TRTLLM,
422 - NvFp4MoeBackend.FLASHINFER_CUTEDSL,
423 - )
424 - ),
380 + is_nvfp4_scale_swizzled=(backend != NvFp4MoeBackend.FLASHINFER_TRTLLM),
381 )
382
383
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+1 -95
@@ -60,100 +60,6 @@ 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 -
63 def prepare_static_weights_for_trtllm_fp4_moe(
64 # args_dequant,
65 # args,
@@ -315,7 +221,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
221 NvFp4MoeBackend.VLLM_CUTLASS,
222 NvFp4MoeBackend.FLASHINFER_CUTLASS,
223 NvFp4MoeBackend.FLASHINFER_TRTLLM,
318 - NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
224 + NvFp4MoeBackend.FLASHINFER_CUTEDSL,
225 ]
226
227 # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels.
vllm/utils/flashinfer.py
-18
@@ -128,12 +128,6 @@ 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 -)
131 trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
132 "flashinfer", "trtllm_fp4_block_scale_moe"
133 )
@@ -257,15 +251,6 @@ def has_flashinfer_cutedsl_grouped_gemm_nt_masked() -> bool:
251 return True
252
253
260 -@functools.cache
261 -def has_flashinfer_cutedsl_moe_nvfp4() -> bool:
262 - """Return ``True`` if FlashInfer cute_dsl_fused_moe_nvfp4 is available."""
263 - if not has_flashinfer_cutedsl():
264 - return False
265 - mod = _get_submodule("flashinfer")
266 - return mod is not None and hasattr(mod, "cute_dsl_fused_moe_nvfp4")
267 -
268 -
254 @functools.cache
255 def has_nvidia_artifactory() -> bool:
256 """Return `True` if NVIDIA's artifactory is accessible.
@@ -782,8 +767,6 @@ __all__ = [
767 "silu_and_mul_scaled_nvfp4_experts_quantize",
768 "scaled_fp4_grouped_quantize",
769 "nvfp4_block_scale_interleave",
785 - "flashinfer_cute_dsl_fused_moe_nvfp4",
786 - "flashinfer_convert_sf_to_mma_layout",
770 "trtllm_fp4_block_scale_moe",
771 "autotune",
772 "has_flashinfer_moe",
@@ -792,7 +775,6 @@ __all__ = [
775 "has_flashinfer_nvlink_one_sided",
776 "has_flashinfer_cutlass_fused_moe",
777 "has_flashinfer_cutedsl_grouped_gemm_nt_masked",
795 - "has_flashinfer_cutedsl_moe_nvfp4",
778 "has_flashinfer_fp8_blockscale_gemm",
779 "has_nvidia_artifactory",
780 "supports_trtllm_attention",