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

[MoE] Nvfp4 Masked Gemm: Add flashinfer grouped_gemm_nt_masked (#25990)

Signed-off-by: Shu Wang. <shuw@nvidia.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

Shu Wang committed Nov 19, 2025 at 15:29 UTC 613abb50d5715ba693ee9d5b727e8385b98e7185
10 files changed +1062 -33
.buildkite/test-pipeline.yaml
+1
@@ -921,6 +921,7 @@ steps:
921 - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
922 - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
923 - pytest -v -s tests/kernels/moe/test_flashinfer.py
924 + - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
925
926 - label: Blackwell Fusion and Compile Tests # 30 min
927 timeout_in_minutes: 40
tests/kernels/moe/test_cutedsl_moe.py new
+582
@@ -0,0 +1,582 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import pytest
6 +
7 +from vllm.platforms import current_platform
8 +
9 +if not current_platform.has_device_capability(100):
10 + pytest.skip(
11 + reason="Nvfp4 Requires compute capability of 10 or above.",
12 + allow_module_level=True,
13 + )
14 +
15 +import torch
16 +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.flashinfer_cutedsl_moe import (
21 + flashinfer_cutedsl_moe_masked,
22 +)
23 +from vllm.utils.flashinfer import (
24 + flashinfer_cutedsl_grouped_gemm_nt_masked as cutedsl_gmm_masked,
25 +)
26 +from vllm.utils.flashinfer import (
27 + scaled_fp4_grouped_quantize,
28 +)
29 +
30 +kE2M1ToFloat = torch.tensor(
31 + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32
32 +)
33 +
34 +FLOAT8_E4M3_MAX = 448.0
35 +FLOAT4_E2M1_MAX = 6.0
36 +
37 +
38 +def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
39 + m_tiles = (m + 128 - 1) // 128
40 + f = block_size * 4
41 + k_tiles = (k + f - 1) // f
42 + tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
43 + tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
44 + out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
45 + return out[0:m, 0:k]
46 +
47 +
48 +def dequantize_nvfp4_to_dtype(
49 + tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16
50 +):
51 + """Dequantize the fp4 tensor back to high precision."""
52 + # Two fp4 values are packed into one uint8.
53 + assert tensor_fp4.dtype == torch.uint8
54 + m, packed_k = tensor_fp4.shape
55 + k = packed_k * 2
56 + tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
57 + tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
58 + tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
59 + tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
60 + tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
61 +
62 + # scale the tensor
63 + out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
64 + return out.to(dtype=dtype)
65 +
66 +
67 +def break_fp4_bytes(a, dtype):
68 + assert a.dtype == torch.uint8
69 + m, n = a.shape
70 +
71 + # Vectorized nibble processing
72 + a_flat = a.flatten()
73 + high = (a_flat & 0xF0) >> 4 # Upper nibbles
74 + low = a_flat & 0x0F # Lower nibbles
75 +
76 + # Combine nibbles for batch processing
77 + combined = torch.stack((low, high), dim=1).flatten()
78 +
79 + # Vectorized sign and magnitude extraction
80 + signs = (combined & 0x08).to(torch.bool) # Sign bits
81 + abs_vals = (combined & 0x07).to(torch.long) # Magnitude indices
82 +
83 + # Device-aware lookup and sign application
84 + kE2M1 = kE2M1ToFloat.to(device=a.device)
85 + values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)
86 +
87 + # Reshape to final form
88 + return values.reshape(m, n * 2).to(dtype=dtype)
89 +
90 +
91 +def generate_balanced_routing(
92 + hidden_states: torch.Tensor, num_experts: int, top_k: int
93 +):
94 + """
95 + Generate routing weights and topk indices such that every expert is active.
96 + Returns routing_weights, topk_idx
97 + """
98 +
99 + num_tokens, hidden_dim = hidden_states.shape
100 + # num_tokens = batch_size * seq_len
101 +
102 + # First, assign at least one token per expert
103 + tokens_per_expert = torch.arange(num_tokens) % num_experts
104 + tokens_per_expert = tokens_per_expert[torch.randperm(num_tokens)] # shuffle
105 +
106 + # Each token has top_k experts — start with one guaranteed expert
107 + topk_idx = torch.full((num_tokens, top_k), -1, dtype=torch.long)
108 + topk_idx[:, 0] = tokens_per_expert
109 +
110 + # For remaining top_k - 1 experts, pick randomly (allowing repeats)
111 + if top_k > 1:
112 + random_choices = torch.randint(0, num_experts, (num_tokens, top_k - 1))
113 + topk_idx[:, 1:] = random_choices
114 +
115 + # Normalize routing weights so each token's weights sum to 1
116 + routing_weights = torch.rand(num_tokens, top_k)
117 + routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
118 +
119 + # Reshape back if needed
120 + routing_weights = routing_weights.view(num_tokens, top_k)
121 + topk_idx = topk_idx.view(num_tokens, top_k)
122 +
123 + return routing_weights, topk_idx
124 +
125 +
126 +def prepare_inputs(
127 + hidden_states: torch.Tensor,
128 + router_logits: torch.Tensor,
129 + num_experts: int,
130 + topk: int,
131 +):
132 + routing_weights, topk_idx = generate_balanced_routing(
133 + router_logits, num_experts, topk
134 + )
135 +
136 + masked_m = []
137 + for i in range(num_experts):
138 + mask = topk_idx.view(-1) == i
139 + masked_m.append(mask.sum())
140 +
141 + masked_m = torch.tensor(masked_m, dtype=torch.int32)
142 + # Intialize the hidden_states_3d with ones instead of empty to avoid nan
143 + # issue.
144 + hidden_states_3d = torch.ones(
145 + (num_experts, max(masked_m), hidden_states.shape[1]), dtype=hidden_states.dtype
146 + )
147 + for i in range(num_experts):
148 + hidden_states_3d[i, : masked_m[i], :] = hidden_states[topk_idx.view(-1) == i]
149 +
150 + return hidden_states_3d, masked_m, topk_idx, routing_weights
151 +
152 +
153 +MNK_FACTORS = [
154 + (2, 1024, 1024),
155 + (2, 1024, 1536),
156 + (2, 3072, 1024),
157 + (2, 3072, 1536),
158 + (64, 1024, 1024),
159 + (64, 1024, 1536),
160 + (64, 3072, 1024),
161 + (64, 2048, 1024),
162 + (224, 1024, 1024),
163 + (224, 1024, 1536),
164 +]
165 +
166 +
167 +# Reference implementation of torch_moe
168 +def torch_moe(a, w1, w2, score, topk, expert_map):
169 + B, D = a.shape
170 + a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
171 + out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
172 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
173 + topk_weight, topk_ids = torch.topk(score, topk)
174 + topk_weight = topk_weight.view(-1)
175 + topk_ids = topk_ids.view(-1)
176 + if expert_map is not None:
177 + topk_ids = expert_map[topk_ids]
178 + for i in range(w1.shape[0]):
179 + mask = topk_ids == i
180 + if mask.sum():
181 + out[mask] = SiluAndMul()(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
182 + 0, 1
183 + )
184 + return (
185 + out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
186 + ).sum(dim=1)
187 +
188 +
189 +def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids):
190 + B, D = a.shape
191 + a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
192 + out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
193 +
194 + topk_weight = topk_weight.view(-1)
195 + topk_ids = topk_ids.view(-1)
196 +
197 + for i in range(w1.shape[0]):
198 + mask = topk_ids == i
199 + if mask.sum():
200 + m = w1[i].shape[0]
201 + assert m % 2 == 0
202 + # Note: w1 and w3 are swapped!
203 + w3_expert, w1_expert = w1[i][m // 2 :, :], w1[i][: m // 2, :]
204 + inter = F.silu(a[mask] @ w1_expert.t()) * (a[mask] @ w3_expert.t())
205 + inter_gs = torch.tensor(1.0).cuda()
206 + inter_q, inter_blockscale = fp4_quantize(inter, inter_gs)
207 + inter = dequantize_nvfp4_to_dtype(
208 + inter_q,
209 + inter_blockscale,
210 + inter_gs,
211 + dtype=inter.dtype,
212 + device=inter.device,
213 + block_size=16,
214 + ).cuda()
215 + out[mask] = inter @ w2[i].transpose(0, 1)
216 + return (
217 + out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
218 + ).sum(dim=1)
219 +
220 +
221 +def grouped_gemm_ref(
222 + hidden_states_expanded: torch.Tensor,
223 + hidden_states_3d: torch.Tensor,
224 + weights: torch.Tensor,
225 + topk_idx: torch.Tensor,
226 + masked_m: torch.Tensor,
227 + B: int,
228 + topk: int,
229 + num_experts: int,
230 + *,
231 + block_size: int = 16,
232 +) -> torch.Tensor:
233 + """
234 + Computes the reference grouped GEMM (fp4 quantized per-expert loop),
235 + computes flashinfer grouped GEMM (for scale consistency),
236 + and returns ONLY the repacked reference output: out_ref.
237 +
238 + Returns:
239 + out_ref: Tensor [num_experts, max_m, n_out]
240 + """
241 + device_hs = hidden_states_expanded.device
242 + device_w = weights.device
243 + out_dtype = weights.dtype
244 + n_out = weights.shape[1]
245 +
246 + # Flattened reference output (B*topk, n_out)
247 + out = torch.zeros((B * topk, n_out), dtype=out_dtype, device=device_w)
248 +
249 + # Per-expert reference compute loop
250 + for i in range(num_experts):
251 + mask = topk_idx.view(-1) == i
252 + if mask.any():
253 + lhs = hidden_states_expanded[mask]
254 + rhs = weights[i]
255 +
256 + a_amax = lhs.abs().max().to(torch.float32).to(device_hs)
257 + b_amax = rhs.abs().max().to(torch.float32).to(device_w)
258 +
259 + a_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / a_amax
260 + b_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax
261 +
262 + lhsq, lhsq_sf = fp4_quantize(lhs, a_gs)
263 + rhsq, rhsq_sf = fp4_quantize(rhs, b_gs)
264 +
265 + lhs_in_dtype = dequantize_nvfp4_to_dtype(
266 + lhsq,
267 + lhsq_sf,
268 + a_gs,
269 + dtype=lhs.dtype,
270 + device=device_hs,
271 + block_size=block_size,
272 + )
273 + rhs_in_dtype = dequantize_nvfp4_to_dtype(
274 + rhsq,
275 + rhsq_sf,
276 + b_gs,
277 + dtype=rhs.dtype,
278 + device=device_w,
279 + block_size=block_size,
280 + )
281 +
282 + out[mask] = lhs_in_dtype @ rhs_in_dtype.t()
283 +
284 + # Determine per-expert max_m
285 + max_m_val = int(masked_m.max().item())
286 +
287 + # Repack into [num_experts, max_m, n_out]
288 + out_ref = torch.zeros(
289 + (num_experts, max_m_val, n_out),
290 + dtype=out.dtype,
291 + device=out.device,
292 + )
293 + expert_slot = [0] * num_experts
294 +
295 + for i, expert_id in enumerate(topk_idx.view(-1).tolist()):
296 + slot = expert_slot[expert_id]
297 + if slot < max_m_val:
298 + out_ref[expert_id, slot, :] = out[i]
299 + expert_slot[expert_id] += 1
300 + else:
301 + raise IndexError(
302 + f"Expert {expert_id} exceeded max slots ({max_m_val}). "
303 + "Increase max_m or check masked_m."
304 + )
305 +
306 + return out_ref
307 +
308 +
309 +def flashinfer_cutedsl_grouped_gemm_nt_masked(
310 + hidden_states: torch.Tensor, # 3d
311 + input_global_scale: torch.Tensor, # (l,)
312 + weights: torch.Tensor,
313 + w_global_scale: torch.Tensor, # (l,)
314 + masked_m: torch.Tensor,
315 +):
316 + # hidden_states: [l, m, k]
317 + # weights: [l, n, k]
318 + aq, aq_sf = scaled_fp4_grouped_quantize(
319 + hidden_states,
320 + masked_m.to(hidden_states.device),
321 + input_global_scale,
322 + )
323 + num_experts, n, k = weights.shape
324 + bq, bq_sf = scaled_fp4_grouped_quantize(
325 + weights,
326 + torch.full((num_experts,), n, device=weights.device, dtype=torch.int32),
327 + w_global_scale,
328 + )
329 +
330 + out = torch.zeros(
331 + (num_experts, max(masked_m), n), dtype=weights.dtype, device=aq.device
332 + )
333 + out = out.permute(1, 2, 0) # requirement of kernel
334 + sf_vec_size = 16
335 + ab_dtype = "float4_e2m1fn"
336 + sf_dtype = "float8_e4m3fn"
337 + c_dtype = "bfloat16"
338 + alpha = 1.0 / (input_global_scale * w_global_scale).to(out.dtype).view(
339 + 1, 1, num_experts
340 + )
341 +
342 + def get_cute_dtype(input: torch.Tensor) -> str:
343 + if input.dtype == torch.bfloat16:
344 + return "bfloat16"
345 + elif input.dtype == torch.float16:
346 + return "float16"
347 + elif input.dtype == torch.float32:
348 + return "float32"
349 + else:
350 + raise ValueError(f"Unsupported cute dtype {input.dtype}")
351 +
352 + cutedsl_gmm_masked(
353 + (aq, aq_sf),
354 + (bq, bq_sf),
355 + out,
356 + masked_m.to(aq.device),
357 + ab_dtype=ab_dtype,
358 + sf_dtype=sf_dtype,
359 + c_dtype=c_dtype,
360 + sf_vec_size=sf_vec_size,
361 + alpha=alpha,
362 + alpha_dtype=get_cute_dtype(alpha),
363 + )
364 +
365 + return out
366 +
367 +
368 +@pytest.mark.parametrize("bs, hidden_dim, inter_dim", [(2, 128, 256), (16, 128, 512)])
369 +@pytest.mark.parametrize("topk", [1, 2, 4])
370 +@torch.inference_mode()
371 +def test_flashinfer_cutedsl_moe_masked(
372 + bs: int, hidden_dim: int, inter_dim: int, topk: int
373 +):
374 + torch.manual_seed(42)
375 + device = "cuda"
376 + num_experts = 8
377 + hidden_states = (
378 + torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device) / 5.0
379 + )
380 + w1 = (
381 + torch.randn(
382 + num_experts, 2 * inter_dim, hidden_dim, dtype=torch.bfloat16, device=device
383 + )
384 + / 10.0
385 + )
386 + w2 = (
387 + torch.randn(
388 + num_experts, hidden_dim, inter_dim, dtype=torch.bfloat16, device=device
389 + )
390 + / 10.0
391 + )
392 + router_logits = torch.randn(bs, num_experts, dtype=torch.float32)
393 +
394 + hidden_states_expanded = (
395 + hidden_states.view(bs, -1, hidden_dim)
396 + .repeat(1, topk, 1)
397 + .reshape(-1, hidden_dim)
398 + )
399 + hidden_states_3d, masked_m, topk_idx, routing_weights = prepare_inputs(
400 + hidden_states_expanded, router_logits, num_experts, topk
401 + )
402 +
403 + w1_amax = w1.abs().amax(dim=(1, 2)).to(torch.float32).to(w1.device)
404 + w2_amax = w2.abs().amax(dim=(1, 2)).to(torch.float32).to(w2.device)
405 + input_global_scale = torch.ones(
406 + (num_experts,), dtype=torch.float32, device=hidden_states.device
407 + )
408 +
409 + w1_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
410 + w2_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
411 + a2_global_scale = torch.ones(
412 + (num_experts,), dtype=torch.float32, device=hidden_states.device
413 + ) # assume intermediate scale is 1.0
414 +
415 + w1_fp4, w1_blockscale = scaled_fp4_grouped_quantize(
416 + w1,
417 + torch.ones(num_experts, dtype=torch.int32, device=w1.device) * 2 * inter_dim,
418 + w1_global_scale,
419 + )
420 + w2_fp4, w2_blockscale = scaled_fp4_grouped_quantize(
421 + w2,
422 + torch.ones(num_experts, dtype=torch.int32, device=w2.device) * hidden_dim,
423 + w2_global_scale,
424 + )
425 +
426 + w1_alpha = 1.0 / (input_global_scale * w1_global_scale)
427 + w2_alpha = 1.0 / (a2_global_scale * w2_global_scale)
428 +
429 + out = torch.empty_like(hidden_states_3d)
430 + # Note: the 1st dim shouldn't be bs
431 + wk = torch.empty(
432 + num_experts,
433 + hidden_states_3d.shape[1],
434 + inter_dim * 2,
435 + dtype=hidden_states_3d.dtype,
436 + device=hidden_states.device,
437 + )
438 + flashinfer_cutedsl_moe_masked(
439 + hidden_states_3d.to(hidden_states.device),
440 + input_global_scale,
441 + w1_fp4.permute(2, 0, 1),
442 + w1_blockscale,
443 + w1_alpha,
444 + w2_fp4.permute(2, 0, 1),
445 + a2_global_scale,
446 + w2_blockscale,
447 + w2_alpha,
448 + masked_m.to(hidden_states.device),
449 + wk,
450 + out,
451 + )
452 +
453 + # reference
454 + a_fp4, a_scale_interleaved = fp4_quantize(hidden_states, input_global_scale)
455 + a_in_dtype = dequantize_nvfp4_to_dtype(
456 + a_fp4,
457 + a_scale_interleaved,
458 + input_global_scale,
459 + dtype=hidden_states.dtype,
460 + device=hidden_states.device,
461 + block_size=16,
462 + )
463 + w1_d = torch.empty(
464 + (num_experts, 2 * inter_dim, hidden_dim), device=w1.device, dtype=w1.dtype
465 + )
466 + w2_d = torch.empty(
467 + (num_experts, hidden_dim, inter_dim), device=w2.device, dtype=w2.dtype
468 + )
469 +
470 + for idx in range(0, num_experts):
471 + w1_fp4_sliced, w1_blockscale_sliced = fp4_quantize(
472 + w1[idx], w1_global_scale[idx]
473 + )
474 + w2_fp4_sliced, w2_blockscale_sliced = fp4_quantize(
475 + w2[idx], w2_global_scale[idx]
476 + )
477 + w1_d[idx] = dequantize_nvfp4_to_dtype(
478 + w1_fp4_sliced,
479 + w1_blockscale_sliced,
480 + w1_global_scale[idx],
481 + dtype=w1.dtype,
482 + device=w1.device,
483 + block_size=16,
484 + )
485 + w2_d[idx] = dequantize_nvfp4_to_dtype(
486 + w2_fp4_sliced,
487 + w2_blockscale_sliced,
488 + w2_global_scale[idx],
489 + dtype=w2.dtype,
490 + device=w2.device,
491 + block_size=16,
492 + )
493 +
494 + ref_output = torch_moe_nvfp4(
495 + a_in_dtype,
496 + w1_d,
497 + w2_d,
498 + topk,
499 + routing_weights.to(a_in_dtype.device),
500 + topk_idx.to(a_in_dtype.device),
501 + )
502 + out_weighted = torch.zeros_like(ref_output, device=out.device, dtype=out.dtype)
503 +
504 + positions = torch.nonzero(masked_m[topk_idx], as_tuple=False)
505 + rows, cols = positions[:, 0], positions[:, 1]
506 + experts = topk_idx[rows, cols]
507 + for i in range(num_experts):
508 + mask = experts == i
509 + if mask.any():
510 + idx = torch.nonzero(mask, as_tuple=False).squeeze(-1)
511 + r, c = rows[idx], cols[idx]
512 + out_weighted[r] += out[i, : len(r), :] * routing_weights[r, c].to(
513 + out.device
514 + ).unsqueeze(-1)
515 + torch.testing.assert_close(
516 + out_weighted.cpu(), ref_output.cpu(), atol=2e-1, rtol=2e-1
517 + )
518 +
519 +
520 +@pytest.mark.parametrize(
521 + "bs, hidden_dim, inter_dim, topk", [(2, 128, 256, 2), (16, 128, 512, 5)]
522 +)
523 +@torch.inference_mode()
524 +def test_grouped_gemm_nt_masked(
525 + bs: int, hidden_dim: int, inter_dim: int, topk: int
526 +) -> None:
527 + torch.manual_seed(42)
528 + B = bs
529 + D = hidden_dim
530 + N = inter_dim
531 + # CuteDSL group gemm has issue when not all experts are active.
532 + # i.e. masked = [2, 3, 0, 0, 1] where the 2nd and 3rd experts are inactive
533 + # see https://github.com/flashinfer-ai/flashinfer/issues/1856
534 + num_experts = bs
535 + hidden_states = torch.randn(B, D, dtype=torch.bfloat16, device="cuda")
536 + weights = torch.randn(num_experts, N, D, dtype=torch.bfloat16, device="cuda")
537 + router_logits = torch.randn(B, num_experts, dtype=torch.float32)
538 +
539 + hidden_states_expanded = (
540 + hidden_states.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
541 + )
542 + hidden_states_3d, masked_m, topk_idx, _ = prepare_inputs(
543 + hidden_states_expanded, router_logits, num_experts, topk
544 + )
545 +
546 + a_amax = (
547 + hidden_states_3d.abs()
548 + .amax(dim=(1, 2))
549 + .to(torch.float32)
550 + .to(hidden_states.device)
551 + )
552 + b_amax = weights.abs().amax(dim=(1, 2)).to(torch.float32).to(weights.device)
553 + a_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / a_amax
554 + b_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax
555 + out_flashinfer = flashinfer_cutedsl_grouped_gemm_nt_masked(
556 + hidden_states_3d.to(hidden_states.device), a_gs, weights, b_gs, masked_m
557 + )
558 + # reference
559 + out_ref = grouped_gemm_ref(
560 + hidden_states_expanded=hidden_states_expanded,
561 + hidden_states_3d=hidden_states_3d,
562 + weights=weights,
563 + topk_idx=topk_idx,
564 + masked_m=masked_m,
565 + B=B,
566 + topk=topk,
567 + num_experts=num_experts,
568 + )
569 + # Note: just to compare the masked position due to cutedsl may write nan
570 + # into unmasked position.
571 + for i in range(num_experts):
572 + torch.testing.assert_close(
573 + out_flashinfer.permute(2, 0, 1)[i, : masked_m[i]],
574 + out_ref.to(out_flashinfer.device)[i, : masked_m[i]],
575 + atol=1e-1,
576 + rtol=1e-1,
577 + )
578 +
579 +
580 +if __name__ == "__main__":
581 + test_flashinfer_cutedsl_moe_masked(16, 128, 512, 4)
582 + test_grouped_gemm_nt_masked(16, 128, 512, 4)
vllm/envs.py
+6 -2
@@ -157,7 +157,9 @@ if TYPE_CHECKING:
157 VLLM_USE_FLASHINFER_MOE_FP16: bool = False
158 VLLM_USE_FLASHINFER_MOE_FP8: bool = False
159 VLLM_USE_FLASHINFER_MOE_FP4: bool = False
160 - VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency"] = "latency"
160 + VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency", "masked_gemm"] = (
161 + "latency"
162 + )
163 VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024
164 VLLM_XGRAMMAR_CACHE_MB: int = 0
165 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
@@ -1238,7 +1240,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1240 # - "latency":
1241 # Uses TensorRT-LLM kernels optimized for low-latency inference.
1242 "VLLM_FLASHINFER_MOE_BACKEND": env_with_choices(
1241 - "VLLM_FLASHINFER_MOE_BACKEND", "latency", ["throughput", "latency"]
1243 + "VLLM_FLASHINFER_MOE_BACKEND",
1244 + "latency",
1245 + ["throughput", "latency", "masked_gemm"],
1246 ),
1247 # Control the workspace buffer size for the FlashInfer backend.
1248 "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int(
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+14 -2
@@ -6,6 +6,7 @@ import deep_ep
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 FusedMoEQuantConfig
12 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -27,6 +28,8 @@ logger = init_logger(__name__)
28 DEEPEP_QUANT_BLOCK_SIZE = 128
29 DEEPEP_QUANT_BLOCK_SHAPE = [DEEPEP_QUANT_BLOCK_SIZE, DEEPEP_QUANT_BLOCK_SIZE]
30
31 +logger = init_logger(__name__)
32 +
33
34 def dequant_fp8(
35 expert_x_fp8: torch.Tensor, expert_x_scales: torch.Tensor
@@ -187,16 +190,25 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
190
191 # TODO (varun): Optimization - Use a batched version of quant
192 x = x.view((-1, hidden_dim))
193 + q_dtype = quant_config.quant_dtype
194 +
195 + if envs.VLLM_FLASHINFER_MOE_BACKEND == "masked_gemm":
196 + logger.info_once(
197 + "Skip quantization when using FlashInfer CUTEDSL(masked_gemm) "
198 + "for ModelOptNvFp4FusedMoE."
199 + )
200 + q_dtype = None
201 +
202 x, x_scales = moe_kernel_quantize_input(
203 x,
204 quant_config.a1_scale,
193 - quant_config.quant_dtype,
205 + q_dtype,
206 quant_config.per_act_token_quant,
207 quant_config.block_shape,
208 )
209 x = x.view((num_experts, -1, hidden_dim))
210
199 - if quant_config.quant_dtype is not None:
211 + if q_dtype is not None:
212 assert x_scales is not None
213 x_scales = normalize_batched_scales_shape(x_scales, num_experts)
214
vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py new
+346
@@ -0,0 +1,346 @@
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.logger import init_logger
8 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
10 + TopKWeightAndReduceDelegate,
11 +)
12 +from vllm.utils.flashinfer import (
13 + flashinfer_cutedsl_grouped_gemm_nt_masked,
14 + has_flashinfer_cutedsl_grouped_gemm_nt_masked,
15 + scaled_fp4_grouped_quantize,
16 + silu_and_mul_scaled_nvfp4_experts_quantize,
17 +)
18 +
19 +logger = init_logger(__name__)
20 +
21 +
22 +def is_valid_flashinfer_cutedsl_fused_moe(
23 + hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor
24 +) -> bool:
25 + """
26 + Check if the given problem size is supported by the FlashInfer CuteDSL MoE
27 + kernel.
28 + """
29 + if not has_flashinfer_cutedsl_grouped_gemm_nt_masked():
30 + logger.debug_once(
31 + "FlashInferCuteDSLExperts disabled: "
32 + "flashinfer_cutedsl_fused_moe not available."
33 + )
34 + return False
35 + # Data type checks
36 + if (
37 + w1.dtype != torch.uint8
38 + or w2.dtype != torch.uint8
39 + or hidden_states.dtype not in [torch.float32, torch.float16, torch.bfloat16]
40 + ):
41 + logger.debug_once(
42 + "FlashInferCuteDSLExperts disabled: w1/w2 must be torch.uint8 "
43 + f"(got w1={w1.dtype}, w2={w2.dtype}), hidden_states must be "
44 + f"float32, float16, or bfloat16 (got {hidden_states.dtype})."
45 + )
46 + return False
47 + return True
48 +
49 +
50 +class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):
51 + def __init__(
52 + self,
53 + out_dtype: torch.dtype,
54 + quant_config: FusedMoEQuantConfig,
55 + ):
56 + super().__init__(quant_config)
57 + assert quant_config.quant_dtype == "nvfp4", (
58 + "Only nvfp4 quantization are currently supported."
59 + )
60 + self.out_dtype = out_dtype
61 +
62 + @property
63 + def activation_formats(
64 + self,
65 + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
66 + return (
67 + mk.FusedMoEActivationFormat.BatchedExperts,
68 + mk.FusedMoEActivationFormat.BatchedExperts,
69 + )
70 +
71 + def supports_expert_map(self) -> bool:
72 + return False
73 +
74 + def supports_chunking(self) -> bool:
75 + # This refers to TP chunking; DP chunking is handled separately.
76 + # TODO(shuw@nvidia.com): Set to False to be consistent with
77 + # batched_deep_gemm_moe
78 + return False
79 +
80 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
81 + # Let PrepareAndFinalize::finalize() decide the impl.
82 + return TopKWeightAndReduceDelegate()
83 +
84 + def workspace_shapes(
85 + self,
86 + M: int,
87 + N: int,
88 + K: int,
89 + topk: int,
90 + global_num_experts: int,
91 + local_num_experts: int,
92 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
93 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
94 + # We use global_num_experts due to how moe_align_block_size handles
95 + # expert_maps.
96 + """
97 + Compute the shapes for the temporary and final outputs of the two gemms
98 + and activation in the fused expert function. Since the gemms are
99 + independent, the workspace for the first gemm can be shared with the
100 + workspace for the last gemm.
101 +
102 + Returns a tuple of:
103 + - workspace13 shape tuple: must be large enough to hold the
104 + result of either expert gemm.
105 + - workspace2 shape tuple: must be large enough to hold the
106 + result of the activation function.
107 + - output shape tuple: must be exact size of the final gemm output.
108 + - Workspace type: The dtype to use for the workspace tensors.
109 + - Note: in order for activation chunking to work, the first dimension
110 + of each tuple must be the number of tokens.
111 + """
112 + output_shape = (local_num_experts, M, K)
113 + workspace2 = (local_num_experts, M, N)
114 + workspace1 = output_shape
115 + return (workspace1, workspace2, output_shape)
116 +
117 + def apply(
118 + self,
119 + output: torch.Tensor,
120 + hidden_states: torch.Tensor,
121 + w1: torch.Tensor,
122 + w2: torch.Tensor,
123 + topk_weights: torch.Tensor,
124 + topk_ids: torch.Tensor,
125 + activation: str,
126 + global_num_experts: int,
127 + expert_map: torch.Tensor | None,
128 + a1q_scale: torch.Tensor | None,
129 + a2_scale: torch.Tensor | None, # Not used
130 + workspace13: torch.Tensor | None,
131 + workspace2: torch.Tensor | None,
132 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
133 + apply_router_weight_on_input: bool | None,
134 + ):
135 + assert self.quant_dtype == "nvfp4", (
136 + "Only nvfp4 quantization are currently supported."
137 + )
138 + # Ensure w1_scale and w2_scale are not None before calling view
139 + assert self.w1_scale is not None and self.w2_scale is not None, (
140 + "w1_scale and w2_scale must not be None for FlashInferExperts"
141 + )
142 + assert expert_tokens_meta is not None
143 + expert_num_tokens = expert_tokens_meta.expert_num_tokens
144 + assert hidden_states.ndim == 3
145 + assert self.w1_scale.ndim == 3
146 + assert self.w2_scale.ndim == 3
147 + flashinfer_cutedsl_moe_masked(
148 + hidden_states=hidden_states,
149 + input_global_scale=self.a1_gscale,
150 + w1=w1,
151 + w1_blockscale=self.w1_scale,
152 + w1_alpha=self.g1_alphas,
153 + w2=w2,
154 + a2_global_scale=self.a2_gscale,
155 + w2_blockscale=self.w2_scale,
156 + w2_alpha=self.g2_alphas,
157 + masked_m=expert_num_tokens,
158 + workspace=workspace2,
159 + out=output,
160 + )
161 +
162 +
163 +def get_cute_dtype(input: torch.Tensor) -> str:
164 + if input.dtype == torch.bfloat16:
165 + return "bfloat16"
166 + elif input.dtype == torch.float16:
167 + return "float16"
168 + elif input.dtype == torch.float32:
169 + return "float32"
170 + else:
171 + raise ValueError(f"Unsupported cute dtype {input.dtype}")
172 +
173 +
174 +def flashinfer_cutedsl_moe_masked(
175 + hidden_states: torch.Tensor,
176 + input_global_scale: torch.Tensor,
177 + w1: torch.Tensor,
178 + w1_blockscale: torch.Tensor,
179 + w1_alpha,
180 + w2: torch.Tensor,
181 + a2_global_scale: torch.Tensor,
182 + w2_blockscale: torch.Tensor,
183 + w2_alpha,
184 + masked_m: torch.Tensor,
185 + workspace: torch.Tensor,
186 + out: torch.Tensor,
187 +):
188 + """
189 + Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL
190 + kernels.
191 +
192 + Args:
193 + hidden_states (torch.Tensor): [num_experts, m, k], bf16
194 + input_global_scale (torch.Tensor): (l,)
195 + w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8
196 + w1_blockscale (torch.Tensor): blockscale factors, e4m3,
197 + w1_alpha (torch.Tensor): (l,)
198 + w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8
199 + a2_global_scale (torch.Tensor): (l,)
200 + w2_blockscale (torch.Tensor): blockscale factors, e4m3,
201 + w2_alpha (torch.Tensor): (l,)
202 + masked_m (torch.Tensor): Masked dimension indices
203 + workspace (torch.Tensor): For gateup_output
204 +
205 + Notes:
206 + - Assumes max(masked_m) <= m.
207 + """
208 +
209 + # === Assertions on dtypes ===
210 + assert input_global_scale.dtype == torch.float32, (
211 + f"input_global_scale must be float32, got {input_global_scale.dtype}"
212 + )
213 + assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}"
214 + assert w1_blockscale.dtype == torch.float8_e4m3fn, (
215 + f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}"
216 + )
217 + assert w1_alpha.dtype == torch.float32, (
218 + f"w1_alpha must be float32, got {w1_alpha.dtype}"
219 + )
220 + assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}"
221 + assert a2_global_scale.dtype == torch.float32, (
222 + f"a2_global_scale must be float32, got {a2_global_scale.dtype}"
223 + )
224 + assert w2_blockscale.dtype == torch.float8_e4m3fn, (
225 + f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}"
226 + )
227 + assert w2_alpha.dtype == torch.float32, (
228 + f"w2_alpha must be float32, got {w2_alpha.dtype}"
229 + )
230 +
231 + # === Assertions on shapes ===
232 + n = w2.shape[-1] * 2 # intermediate dimension
233 + num_experts, m, k = hidden_states.shape
234 +
235 + assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}"
236 + assert w1.shape[-1] * 2 == k, (
237 + f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}"
238 + )
239 + assert w2.shape[-2:] == (
240 + k,
241 + n // 2,
242 + ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}"
243 +
244 + assert input_global_scale.shape == (num_experts,), (
245 + f"input_global_scale must be (l,), got {input_global_scale.shape}"
246 + )
247 + assert w1_alpha.shape == (num_experts,), (
248 + f"w1_alpha must be (l,), got {w1_alpha.shape}"
249 + )
250 + assert a2_global_scale.shape == (num_experts,), (
251 + f"a2_global_scale must be (l,), got {a2_global_scale.shape}"
252 + )
253 + assert w2_alpha.shape == (num_experts,), (
254 + f"w2_alpha must be (l,), got {w2_alpha.shape}"
255 + )
256 +
257 + aq, aq_sf = scaled_fp4_grouped_quantize(
258 + hidden_states,
259 + masked_m,
260 + input_global_scale,
261 + )
262 +
263 + workspace = workspace.permute(1, 2, 0) # requirement of kernel
264 + sf_vec_size = 16
265 + assert aq_sf.dtype == torch.float8_e4m3fn
266 + assert aq.dtype == torch.uint8
267 + ab_dtype = "float4_e2m1fn"
268 + sf_dtype = "float8_e4m3fn"
269 +
270 + c_dtype = get_cute_dtype(hidden_states)
271 +
272 + # Gemm1
273 + flashinfer_cutedsl_grouped_gemm_nt_masked(
274 + (aq, aq_sf),
275 + (w1.permute(1, 2, 0), w1_blockscale),
276 + workspace,
277 + masked_m,
278 + ab_dtype=ab_dtype,
279 + sf_dtype=sf_dtype,
280 + c_dtype=c_dtype,
281 + sf_vec_size=sf_vec_size,
282 + alpha=w1_alpha.view(1, 1, num_experts),
283 + alpha_dtype=get_cute_dtype(w1_alpha),
284 + ) # in logical [m, n, l]
285 +
286 + # SILU and quantization
287 + diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize(
288 + workspace.permute(2, 0, 1),
289 + masked_m,
290 + a2_global_scale,
291 + )
292 +
293 + # Gemm2
294 + out = out.permute(1, 2, 0) # requirement of kernel
295 + flashinfer_cutedsl_grouped_gemm_nt_masked(
296 + (diq, diq_sf),
297 + (w2.permute(1, 2, 0), w2_blockscale),
298 + out,
299 + masked_m,
300 + ab_dtype=ab_dtype,
301 + sf_dtype=sf_dtype,
302 + c_dtype=c_dtype,
303 + sf_vec_size=sf_vec_size,
304 + alpha=w2_alpha.view(1, 1, num_experts),
305 + alpha_dtype=get_cute_dtype(w2_alpha),
306 + ) # in logical [m, k, l]
307 + out = out.permute(2, 0, 1)
308 +
309 +
310 +def flashinfer_cutedsl_moe_fp4(
311 + hidden_states: torch.Tensor,
312 + w1: torch.Tensor,
313 + w2: torch.Tensor,
314 + topk_weights: torch.Tensor,
315 + topk_ids: torch.Tensor,
316 + quant_config: FusedMoEQuantConfig,
317 + inplace: bool = False,
318 + activation: str = "silu",
319 + global_num_experts: int = -1,
320 + expert_map: torch.Tensor | None = None,
321 + apply_router_weight_on_input: bool = False,
322 +) -> torch.Tensor:
323 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
324 + create_flashinfer_prepare_finalize,
325 + )
326 +
327 + fused_experts = mk.FusedMoEModularKernel(
328 + create_flashinfer_prepare_finalize(use_dp=False), # could be swapped later
329 + FlashInferCuteDSLExperts(
330 + out_dtype=hidden_states.dtype,
331 + quant_config=quant_config,
332 + ),
333 + )
334 +
335 + return fused_experts(
336 + hidden_states=hidden_states,
337 + w1=w1,
338 + w2=w2,
339 + topk_weights=topk_weights,
340 + topk_ids=topk_ids,
341 + inplace=inplace,
342 + activation=activation,
343 + global_num_experts=global_num_experts,
344 + expert_map=expert_map,
345 + apply_router_weight_on_input=apply_router_weight_on_input,
346 + )
vllm/model_executor/layers/quantization/modelopt.py
+21 -9
@@ -1468,7 +1468,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1468 gemm1_weight = layer.w13_weight.data
1469 gemm1_weight_scale = layer.w13_weight_scale.data
1470
1471 - if self.allow_flashinfer:
1471 + if (
1472 + self.allow_flashinfer
1473 + and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
1474 + ):
1475 gemm1_weight, gemm1_weight_scale = reorder_w1w3_to_w3w1(
1476 gemm1_weight, gemm1_weight_scale, dim=-2
1477 )
@@ -1746,17 +1749,26 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1749 workspace=layer.workspace,
1750 )
1751
1749 - elif (
1750 - self.allow_flashinfer
1751 - and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
1752 - ):
1753 - from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
1754 - flashinfer_cutlass_moe_fp4,
1752 + elif self.allow_flashinfer:
1753 + assert self.flashinfer_moe_backend in (
1754 + FlashinferMoeBackend.CUTLASS,
1755 + FlashinferMoeBackend.CUTEDSL,
1756 )
1757 + if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
1758 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
1759 + flashinfer_cutlass_moe_fp4,
1760 + )
1761
1757 - assert self.moe_quant_config is not None
1762 + flashinfer_fn_moe_fp4 = flashinfer_cutlass_moe_fp4
1763 + else:
1764 + from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import ( # noqa: E501
1765 + flashinfer_cutedsl_moe_fp4,
1766 + )
1767 +
1768 + flashinfer_fn_moe_fp4 = flashinfer_cutedsl_moe_fp4
1769
1759 - return flashinfer_cutlass_moe_fp4(
1770 + assert self.moe_quant_config is not None
1771 + return flashinfer_fn_moe_fp4(
1772 hidden_states=x,
1773 w1=layer.w13_weight,
1774 w2=layer.w2_weight,
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+33 -10
@@ -10,6 +10,9 @@ from vllm.model_executor.layers.fused_moe.config import (
10 FusedMoEConfig,
11 FusedMoEQuantConfig,
12 )
13 +from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import (
14 + FlashInferCuteDSLExperts,
15 +)
16 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
17 FlashInferExperts,
18 )
@@ -17,10 +20,14 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize im
20 create_flashinfer_prepare_finalize,
21 )
22 from vllm.platforms import current_platform
20 -from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
23 +from vllm.utils.flashinfer import (
24 + has_flashinfer_cutedsl_grouped_gemm_nt_masked,
25 + has_flashinfer_cutlass_fused_moe,
26 +)
27
28 __all__ = [
29 "is_flashinfer_fp4_cutlass_moe_available",
30 + "is_flashinfer_fp4_cutedsl_moe_available",
31 "reorder_w1w3_to_w3w1",
32 "build_flashinfer_fp4_cutlass_moe_prepare_finalize",
33 ]
@@ -36,6 +43,16 @@ def is_flashinfer_fp4_cutlass_moe_available() -> bool:
43 )
44
45
46 +def is_flashinfer_fp4_cutedsl_moe_available() -> bool:
47 + """Return ``True`` when FlashInfer CUTEDSL NV-FP4 kernels can be used."""
48 + return (
49 + envs.VLLM_USE_FLASHINFER_MOE_FP4
50 + and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
51 + and current_platform.is_cuda()
52 + and current_platform.is_device_capability(100)
53 + )
54 +
55 +
56 def reorder_w1w3_to_w3w1(
57 weight: torch.Tensor, scale: torch.Tensor, dim: int = -2
58 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -72,15 +89,21 @@ def select_nvfp4_gemm_impl(
89 """Return a GEMM *experts* implementation for NV-FP4 fused-MoE layers"""
90
91 if allow_flashinfer:
75 - return FlashInferExperts(
76 - out_dtype=moe.in_dtype,
77 - quant_config=moe_quant_config,
78 - ep_rank=moe.moe_parallel_config.ep_rank,
79 - ep_size=moe.moe_parallel_config.ep_size,
80 - tp_rank=moe.moe_parallel_config.tp_rank,
81 - tp_size=moe.moe_parallel_config.tp_size,
82 - use_dp=moe.moe_parallel_config.dp_size > 1,
83 - )
92 + if envs.VLLM_FLASHINFER_MOE_BACKEND == "masked_gemm":
93 + return FlashInferCuteDSLExperts(
94 + out_dtype=moe.in_dtype,
95 + quant_config=moe_quant_config,
96 + )
97 + elif envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput":
98 + return FlashInferExperts(
99 + out_dtype=moe.in_dtype,
100 + quant_config=moe_quant_config,
101 + ep_rank=moe.moe_parallel_config.ep_rank,
102 + ep_size=moe.moe_parallel_config.ep_size,
103 + tp_rank=moe.moe_parallel_config.tp_rank,
104 + tp_size=moe.moe_parallel_config.tp_size,
105 + use_dp=moe.moe_parallel_config.dp_size > 1,
106 + )
107
108 # native cutlass experts currently don't support DP; TP case won't call this
109 raise ValueError(
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+12 -9
@@ -25,6 +25,7 @@ logger = init_logger(__name__)
25 class FlashinferMoeBackend(Enum):
26 TENSORRT_LLM = "TensorRT-LLM"
27 CUTLASS = "CUTLASS"
28 + CUTEDSL = "CUTEDSL"
29
30
31 def calculate_tile_tokens_dim(num_tokens, top_k, num_experts):
@@ -273,19 +274,21 @@ def flashinfer_cutlass_moe_fp8(
274
275
276 def get_flashinfer_moe_backend() -> FlashinferMoeBackend:
277 + backend_map = {
278 + "throughput": FlashinferMoeBackend.CUTLASS,
279 + "latency": FlashinferMoeBackend.TENSORRT_LLM,
280 + "masked_gemm": FlashinferMoeBackend.CUTEDSL,
281 + }
282 +
283 flashinfer_moe_backend = envs.VLLM_FLASHINFER_MOE_BACKEND
277 - # Prefer CUTLASS on SM90 to cover both SM90/SM100 generations
278 - if flashinfer_moe_backend == "throughput" or current_platform.is_device_capability(
279 - 90
280 - ):
284 + if flashinfer_moe_backend in backend_map:
285 + return backend_map[flashinfer_moe_backend]
286 + elif current_platform.is_device_capability(90):
287 return FlashinferMoeBackend.CUTLASS
282 - elif flashinfer_moe_backend == "latency":
283 - return FlashinferMoeBackend.TENSORRT_LLM
288
285 - allowed_backends = ["throughput", "latency"]
289 raise ValueError(
287 - f"Unknown flashinfer moe backend: {flashinfer_moe_backend}"
288 - f" expected one of {allowed_backends}"
290 + f"Unknown flashinfer moe backend: {flashinfer_moe_backend!r}. "
291 + f"Expected one of {list(backend_map.keys())}."
292 )
293
294
vllm/model_executor/layers/quantization/utils/nvfp4_moe_support.py
+5 -1
@@ -5,6 +5,7 @@ from dataclasses import dataclass
5 import vllm.envs as envs
6 from vllm.logger import init_logger
7 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
8 + is_flashinfer_fp4_cutedsl_moe_available,
9 is_flashinfer_fp4_cutlass_moe_available,
10 )
11 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
@@ -32,7 +33,10 @@ def detect_nvfp4_moe_support(class_name: str = "") -> NvFp4Support:
33 """Detect platform support for NV-FP4 fused-MoE path"""
34 cutlass_supported = cutlass_fp4_supported()
35
35 - allow_flashinfer = cutlass_supported and is_flashinfer_fp4_cutlass_moe_available()
36 + allow_flashinfer = cutlass_supported and (
37 + is_flashinfer_fp4_cutlass_moe_available()
38 + or is_flashinfer_fp4_cutedsl_moe_available()
39 + )
40
41 if allow_flashinfer:
42 _logger.info_once(
vllm/utils/flashinfer.py
+42
@@ -114,7 +114,17 @@ flashinfer_trtllm_fp8_per_tensor_scale_moe = _lazy_import_wrapper(
114 flashinfer_cutlass_fused_moe = _lazy_import_wrapper(
115 "flashinfer.fused_moe", "cutlass_fused_moe"
116 )
117 +flashinfer_cutedsl_grouped_gemm_nt_masked = _lazy_import_wrapper(
118 + "flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"
119 +)
120 flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")
121 +nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize")
122 +silu_and_mul_scaled_nvfp4_experts_quantize = _lazy_import_wrapper(
123 + "flashinfer", "silu_and_mul_scaled_nvfp4_experts_quantize"
124 +)
125 +scaled_fp4_grouped_quantize = _lazy_import_wrapper(
126 + "flashinfer", "scaled_fp4_grouped_quantize"
127 +)
128 nvfp4_block_scale_interleave = _lazy_import_wrapper(
129 "flashinfer", "nvfp4_block_scale_interleave"
130 )
@@ -166,6 +176,14 @@ def has_flashinfer_moe() -> bool:
176 )
177
178
179 +@functools.cache
180 +def has_flashinfer_cutedsl() -> bool:
181 + """Return ``True`` if FlashInfer cutedsl module is available."""
182 + return (
183 + has_flashinfer() and importlib.util.find_spec("flashinfer.cute_dsl") is not None
184 + )
185 +
186 +
187 @functools.cache
188 def has_flashinfer_cutlass_fused_moe() -> bool:
189 """Return `True` if FlashInfer CUTLASS fused MoE is available."""
@@ -187,6 +205,26 @@ def has_flashinfer_cutlass_fused_moe() -> bool:
205 return True
206
207
208 +@functools.cache
209 +def has_flashinfer_cutedsl_grouped_gemm_nt_masked() -> bool:
210 + """Return ``True`` if FlashInfer CUTLASS fused MoE is available."""
211 + if not has_flashinfer_cutedsl():
212 + return False
213 +
214 + # Check if all required functions are available
215 + required_functions = [
216 + ("flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"),
217 + ("flashinfer", "scaled_fp4_grouped_quantize"),
218 + ("flashinfer", "silu_and_scaled_nvfp4_experts_quantize"),
219 + ]
220 +
221 + for module_name, attr_name in required_functions:
222 + mod = _get_submodule(module_name)
223 + if not mod or not hasattr(mod, attr_name):
224 + return False
225 + return True
226 +
227 +
228 @functools.cache
229 def has_nvidia_artifactory() -> bool:
230 """Return `True` if NVIDIA's artifactory is accessible.
@@ -472,7 +510,10 @@ __all__ = [
510 "has_flashinfer",
511 "flashinfer_trtllm_fp8_block_scale_moe",
512 "flashinfer_cutlass_fused_moe",
513 + "flashinfer_cutedsl_grouped_gemm_nt_masked",
514 "flashinfer_fp4_quantize",
515 + "silu_and_mul_scaled_nvfp4_experts_quantize",
516 + "scaled_fp4_grouped_quantize",
517 "nvfp4_block_scale_interleave",
518 "trtllm_fp4_block_scale_moe",
519 "autotune",
@@ -480,6 +521,7 @@ __all__ = [
521 "has_flashinfer_comm",
522 "has_flashinfer_all2all",
523 "has_flashinfer_cutlass_fused_moe",
524 + "has_flashinfer_cutedsl_grouped_gemm_nt_masked",
525 "has_nvidia_artifactory",
526 "supports_trtllm_attention",
527 "can_use_trtllm_attention",