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

[Core] FlashInfer CUTLASS fused MoE backend (NVFP4) (#20037)

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

Shu Wang committed Jul 17, 2025 at 23:32 UTC c7d8724e7865ba4a54ffdba23ac77eb13d28234b
22 files changed +1095 -271
vllm/_custom_ops.py
+9 -13
@@ -956,11 +956,11 @@ def cutlass_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
956 c_strides, per_act_token, per_out_ch)
957
958
959 -def cutlass_fp4_moe_mm(a_tensors: torch.Tensor, b_tensors: torch.Tensor,
960 - a_scales: torch.Tensor, b_scales: torch.Tensor,
961 - alphas: torch.Tensor, problem_sizes: torch.Tensor,
962 - expert_offsets: torch.Tensor, sf_offsets: torch.Tensor,
963 - out_dtype: torch.dtype, device: torch.device):
959 +def cutlass_fp4_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
960 + b_tensors: torch.Tensor, a_scales: torch.Tensor,
961 + b_scales: torch.Tensor, alphas: torch.Tensor,
962 + problem_sizes: torch.Tensor,
963 + expert_offsets: torch.Tensor, sf_offsets: torch.Tensor):
964 """
965 An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs
966 the gemms for each combination based on the specified problem sizes.
@@ -977,14 +977,10 @@ def cutlass_fp4_moe_mm(a_tensors: torch.Tensor, b_tensors: torch.Tensor,
977 - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
978 MMs used in the fused MoE operation.
979 """
980 - m_topk = a_tensors.shape[0]
981 - n = b_tensors.shape[1]
982 - c_shape = (m_topk, n)
983 - c = torch.empty(c_shape, device=device, dtype=out_dtype)
984 - torch.ops._C.cutlass_fp4_group_mm(c, a_tensors, b_tensors, a_scales,
985 - b_scales, alphas, problem_sizes,
986 - expert_offsets, sf_offsets)
987 - return c.to(out_dtype)
980 + return torch.ops._C.cutlass_fp4_group_mm(out_tensors, a_tensors, b_tensors,
981 + a_scales, b_scales, alphas,
982 + problem_sizes, expert_offsets,
983 + sf_offsets)
984
985
986 # aqlm
vllm/envs.py
+5
@@ -119,6 +119,7 @@ if TYPE_CHECKING:
119 VLLM_TPU_BUCKET_PADDING_GAP: int = 0
120 VLLM_TPU_MOST_MODEL_LEN: Optional[int] = None
121 VLLM_USE_DEEP_GEMM: bool = False
122 + VLLM_USE_FLASHINFER_MOE: bool = False
123 VLLM_XGRAMMAR_CACHE_MB: int = 0
124 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
125 VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
@@ -853,6 +854,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
854 "VLLM_USE_DEEP_GEMM":
855 lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM", "0"))),
856
857 + # Allow use of FlashInfer CUTLASS kernels for fused moe ops.
858 + "VLLM_USE_FLASHINFER_MOE":
859 + lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE", "0"))),
860 +
861 # Control the cache sized used by the xgrammar compiler. The default
862 # of 512 MB should be enough for roughly 1000 JSON schemas.
863 # It can be changed with this variable if needed for some reason.
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+13 -23
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
3 +from typing import Any, Optional
4
5 import torch
6
@@ -255,28 +255,18 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
255 output = (num_experts, max_num_tokens * num_dispatchers, K)
256 return (workspace13, workspace2, output, a.dtype)
257
258 - def apply(
259 - self,
260 - output: torch.Tensor,
261 - hidden_states: torch.Tensor,
262 - w1: torch.Tensor,
263 - w2: torch.Tensor,
264 - topk_weights: torch.Tensor,
265 - topk_ids: torch.Tensor,
266 - activation: str,
267 - global_num_experts: int,
268 - expert_map: Optional[torch.Tensor],
269 - w1_scale: Optional[torch.Tensor],
270 - w2_scale: Optional[torch.Tensor],
271 - w1_zp: Optional[torch.Tensor],
272 - w2_zp: Optional[torch.Tensor],
273 - a1q_scale: Optional[torch.Tensor],
274 - a2_scale: Optional[torch.Tensor],
275 - workspace13: torch.Tensor,
276 - workspace2: torch.Tensor,
277 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
278 - apply_router_weight_on_input: bool,
279 - ):
258 + def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
259 + w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
260 + topk_ids: torch.Tensor, activation: str, global_num_experts: int,
261 + expert_map: Optional[torch.Tensor],
262 + w1_scale: Optional[torch.Tensor],
263 + w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
264 + w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
265 + a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
266 + workspace2: torch.Tensor,
267 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
268 + apply_router_weight_on_input: bool,
269 + extra_expert_args: Optional[dict[str, Any]]):
270 assert expert_tokens_meta is not None
271 expert_num_tokens = expert_tokens_meta.expert_num_tokens
272
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+4 -3
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
3 +from typing import Any, Optional
4
5 import torch
6
@@ -142,7 +142,8 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
142 a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
143 workspace2: torch.Tensor,
144 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
145 - apply_router_weight_on_input: bool):
145 + apply_router_weight_on_input: bool,
146 + extra_expert_args: Optional[dict[str, Any]]):
147 experts = (self.batched_deep_gemm_experts
148 if self.allow_deep_gemm else self.batched_triton_experts)
149 assert experts is not None
@@ -150,4 +151,4 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
151 activation, global_num_experts, expert_map, w1_scale,
152 w2_scale, w1_zp, w2_zp, a1q_scale, a2_scale, workspace13,
153 workspace2, expert_tokens_meta,
153 - apply_router_weight_on_input)
154 + apply_router_weight_on_input, extra_expert_args)
vllm/model_executor/layers/fused_moe/config.py
+16
@@ -15,6 +15,7 @@ from vllm.logger import init_logger
15 from vllm.model_executor.layers.quantization.base_config import (
16 QuantizationConfig)
17 from vllm.utils import cdiv
18 +from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
19
20 logger = init_logger(__name__)
21
@@ -188,6 +189,11 @@ class FusedMoEParallelConfig:
189 return (self.use_all2all_kernels
190 and envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency")
191
192 + @property
193 + def use_flashinfer_cutlass_kernels(self):
194 + return (envs.VLLM_USE_FLASHINFER_MOE
195 + and has_flashinfer_cutlass_fused_moe())
196 +
197 @staticmethod
198 def make(tp_size_: int, dp_size_: int,
199 vllm_parallel_config: ParallelConfig) -> "FusedMoEParallelConfig":
@@ -392,6 +398,10 @@ class FusedMoEConfig:
398 def use_deepep_ll_kernels(self):
399 return self.moe_parallel_config.use_deepep_ll_kernels
400
401 + @property
402 + def use_flashinfer_cutlass_kernels(self):
403 + return self.moe_parallel_config.use_flashinfer_cutlass_kernels
404 +
405 @staticmethod
406 def make(
407 num_experts: int,
@@ -435,6 +445,12 @@ class FusedMoEConfig:
445 if quant_dtype is None and isinstance(quant_config, Fp8Config):
446 quant_dtype = torch.float8_e4m3fn
447
448 + from vllm.model_executor.layers.quantization.modelopt import (
449 + ModelOptNvFp4Config)
450 + if quant_dtype is None and isinstance(quant_config,
451 + ModelOptNvFp4Config):
452 + quant_dtype = torch.uint8
453 +
454 if weight_quant is not None:
455 per_out_ch_quant = (
456 weight_quant.strategy == QuantizationStrategy.CHANNEL)
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+238 -46
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """ CUTLASS based Fused MoE kernels."""
4 -from typing import Callable, Optional
4 +from typing import Any, Callable, Optional
5
6 import torch
7
@@ -14,7 +14,8 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
14 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15 TopKWeightAndReduceDelegate)
16 from vllm.model_executor.layers.fused_moe.utils import (_fp8_quantize,
17 - _resize_cache)
17 + _resize_cache,
18 + extract_required_args)
19 from vllm.scalar_type import scalar_types
20
21 logger = init_logger(__name__)
@@ -298,7 +299,8 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
299 a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
300 workspace2: torch.Tensor,
301 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
301 - apply_router_weight_on_input: bool):
302 + apply_router_weight_on_input: bool,
303 + extra_expert_args: Optional[dict[str, Any]]):
304 assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
305 assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
306
@@ -431,23 +433,28 @@ FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
433 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
434
435
434 -def cutlass_moe_fp4(a: torch.Tensor,
435 - a1_gscale: torch.Tensor,
436 - w1_fp4: torch.Tensor,
437 - w1_blockscale: torch.Tensor,
438 - w1_alphas: torch.Tensor,
439 - a2_gscale: torch.Tensor,
440 - w2_fp4: torch.Tensor,
441 - w2_blockscale: torch.Tensor,
442 - w2_alphas: torch.Tensor,
443 - topk_weights: torch.Tensor,
444 - topk_ids: torch.Tensor,
445 - m: int,
446 - n: int,
447 - k: int,
448 - e: int,
449 - device: torch.device,
450 - apply_router_weight_on_input: bool = False):
436 +def run_cutlass_moe_fp4(
437 + output: torch.Tensor,
438 + a: torch.Tensor,
439 + a1_gscale: torch.Tensor,
440 + w1_fp4: torch.Tensor,
441 + w1_blockscale: torch.Tensor,
442 + w1_alphas: torch.Tensor,
443 + a2_gscale: torch.Tensor,
444 + w2_fp4: torch.Tensor,
445 + w2_blockscale: torch.Tensor,
446 + w2_alphas: torch.Tensor,
447 + topk_weights: torch.Tensor,
448 + topk_ids: torch.Tensor,
449 + workspace13: torch.Tensor,
450 + workspace2: torch.Tensor,
451 + m: int,
452 + n: int,
453 + k: int,
454 + e: int,
455 + device: torch.device,
456 + apply_router_weight_on_input: bool = False,
457 +) -> None:
458 """
459 MoE implementation for FP4 Inputs
460
@@ -487,16 +494,16 @@ def cutlass_moe_fp4(a: torch.Tensor,
494
495 assert (e_w1 == e_w2 and e_w1 == e), ("Number of experts must match",
496 " between weights.")
490 - assert (k_a // 2 == half_k_w1
497 + assert (k_a == half_k_w1 * 2
498 and k == k_w2), ("Hidden size mismatch between a, w1 and w2")
492 - assert (nx2_w1 == n * 2 and half_n_w2 == n // 2), ("mismatch in "
493 - "expected `n`")
499 + assert (nx2_w1 == n * 2 and half_n_w2 * 2 == n), ("mismatch in "
500 + "expected `n`")
501 assert (m == m_a), "input shape mismatch"
502 assert 2 * half_k_w1 == k_w2, "Hidden size mismatch w2 and w1"
503 assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype"
504 assert (topk_weights.size(0) == m and topk_ids.size(0)
505 == m), ("topk must be provided for each row of a")
499 -
506 + topk = topk_ids.size(1)
507 out_dtype = a.dtype
508 num_topk = topk_ids.size(1)
509
@@ -523,7 +530,6 @@ def cutlass_moe_fp4(a: torch.Tensor,
530 blockscale_offsets)
531
532 a = ops.shuffle_rows(a, a_map)
526 -
533 rep_a_fp4, rep_a_blockscale = ops.scaled_fp4_experts_quant(
534 a,
535 a1_gscale,
@@ -531,34 +537,220 @@ def cutlass_moe_fp4(a: torch.Tensor,
537 blockscale_offsets,
538 num_topk,
539 )
534 -
535 - c1 = ops.cutlass_fp4_moe_mm(rep_a_fp4, w1_fp4, rep_a_blockscale,
536 - w1_blockscale, w1_alphas, problem_sizes1,
537 - expert_offsets[:-1], blockscale_offsets[:-1],
538 - out_dtype, device)
540 + c1 = _resize_cache(workspace13, (m * topk, n * 2))
541 + c2 = _resize_cache(workspace2, (m * topk, n))
542 + c3 = _resize_cache(workspace13, (m * topk, k))
543 + ops.cutlass_fp4_moe_mm(c1, rep_a_fp4, w1_fp4, rep_a_blockscale,
544 + w1_blockscale, w1_alphas, problem_sizes1,
545 + expert_offsets[:-1], blockscale_offsets[:-1])
546 del rep_a_fp4, rep_a_blockscale
540 - # hidden size dimension is split to one halfpytho sized tensor.
541 - intermediate = torch.empty((m * num_topk, w1_fp4.size(1) // 2),
542 - device=device,
543 - dtype=out_dtype)
544 -
545 - torch.ops._C.silu_and_mul(intermediate, c1)
546 -
547 + torch.ops._C.silu_and_mul(c2, c1)
548 int_fp4, int_blockscale = ops.scaled_fp4_experts_quant(
548 - intermediate, a2_gscale, expert_offsets, blockscale_offsets, num_topk)
549 + c2, a2_gscale, expert_offsets, blockscale_offsets, num_topk)
550
550 - c2 = ops.cutlass_fp4_moe_mm(int_fp4, w2_fp4, int_blockscale, w2_blockscale,
551 - w2_alphas, problem_sizes2, expert_offsets[:-1],
552 - blockscale_offsets[:-1], out_dtype, device)
551 + ops.cutlass_fp4_moe_mm(c3, int_fp4, w2_fp4, int_blockscale, w2_blockscale,
552 + w2_alphas, problem_sizes2, expert_offsets[:-1],
553 + blockscale_offsets[:-1])
554 del int_fp4, int_blockscale
555
555 - c2 = ops.shuffle_rows(c2, c_map)
556 + c3 = ops.shuffle_rows(c3, c_map)
557 +
558 + assert output.dtype == out_dtype
559 if not apply_router_weight_on_input:
557 - out = (c2.view(m, num_topk, k) *
558 - topk_weights.view(m, num_topk, 1).to(out_dtype)).sum(dim=1)
560 + output.copy_(
561 + (c3.view(m, num_topk, k) *
562 + topk_weights.view(m, num_topk, 1).to(out_dtype)).sum(dim=1),
563 + non_blocking=True)
564 else:
560 - out = c2.view(m, num_topk, k).sum(dim=1)
561 - return out.to(dtype=out_dtype)
565 + output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True)
566 + return
567 +
568 +
569 +class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
570 +
571 + def __init__(
572 + self,
573 + max_experts_per_worker: int,
574 + out_dtype: torch.dtype,
575 + per_act_token_quant: bool,
576 + per_out_ch_quant: bool,
577 + block_shape: Optional[list[int]] = None,
578 + use_batched_format: bool = False,
579 + ):
580 + super().__init__(
581 + FusedMoEQuantConfig(
582 + quant_dtype=torch.uint8,
583 + per_act_token_quant=per_act_token_quant,
584 + per_out_ch_quant=per_out_ch_quant,
585 + block_shape=block_shape,
586 + ))
587 + self.max_experts_per_worker = max_experts_per_worker
588 + self.out_dtype = out_dtype
589 + self.use_batched_format = use_batched_format
590 +
591 + @property
592 + def activation_formats(
593 + self
594 + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
595 + if self.use_batched_format:
596 + return (mk.FusedMoEActivationFormat.BatchedExperts,
597 + mk.FusedMoEActivationFormat.BatchedExperts)
598 + else:
599 + return (mk.FusedMoEActivationFormat.Standard,
600 + mk.FusedMoEActivationFormat.Standard)
601 +
602 + def supports_expert_map(self) -> bool:
603 + return False
604 +
605 + def supports_chunking(self) -> bool:
606 + return True
607 +
608 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
609 + # Let PrepareAndFinalize::finalize() decide the impl.
610 + return TopKWeightAndReduceDelegate()
611 +
612 + def workspace_shapes(
613 + self,
614 + a: torch.Tensor,
615 + aq: torch.Tensor,
616 + M: int,
617 + N: int,
618 + K: int,
619 + topk: int,
620 + global_num_experts: int,
621 + local_num_experts: int,
622 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
623 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
624 + workspace1: tuple[int, ...] = ()
625 + workspace2: tuple[int, ...] = ()
626 + output: tuple[int, ...] = ()
627 + if self.use_batched_format:
628 + padded_M = aq.size(1)
629 + workspace1 = (self.max_experts_per_worker, padded_M, max(N, K))
630 + workspace2 = (self.max_experts_per_worker, padded_M, (N // 2))
631 + output = (self.max_experts_per_worker, padded_M, K)
632 + else:
633 + workspace1 = (M * topk, max(2 * N, K))
634 + workspace2 = (M * topk, N)
635 + output = (M, K)
636 + return (workspace1, workspace2, output,
637 + self.out_dtype if self.out_dtype is not None else a.dtype)
638 +
639 + def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
640 + w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
641 + topk_ids: torch.Tensor, activation: str, global_num_experts: int,
642 + expert_map: Optional[torch.Tensor], w1_scale: torch.Tensor,
643 + w2_scale: torch.Tensor, w1_zp: Optional[torch.Tensor],
644 + w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
645 + a2_scale: torch.Tensor, workspace13: Optional[torch.Tensor],
646 + workspace2: Optional[torch.Tensor],
647 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
648 + apply_router_weight_on_input: bool,
649 + extra_expert_args: Optional[dict[str, Any]]):
650 + required_keys = [
651 + "g1_alphas", "g2_alphas", "a1_gscale", "a2_gscale", "m", "n", "k",
652 + "e", "device"
653 + ]
654 + (g1_alphas, g2_alphas, a1_gscale, a2_gscale, m, n, k, e,
655 + device) = extract_required_args(extra_expert_args, required_keys)
656 + run_cutlass_moe_fp4(
657 + output=output,
658 + a=hidden_states,
659 + a1_gscale=a1_gscale,
660 + w1_fp4=w1,
661 + w1_blockscale=w1_scale,
662 + w1_alphas=g1_alphas,
663 + a2_gscale=a2_gscale,
664 + w2_fp4=w2,
665 + w2_blockscale=w2_scale,
666 + w2_alphas=g2_alphas,
667 + topk_weights=topk_weights,
668 + topk_ids=topk_ids,
669 + workspace13=workspace13,
670 + workspace2=workspace2,
671 + m=m,
672 + n=n,
673 + k=k,
674 + e=e,
675 + device=device,
676 + apply_router_weight_on_input=apply_router_weight_on_input,
677 + )
678 +
679 +
680 +def cutlass_moe_fp4(
681 + a: torch.Tensor,
682 + w1_fp4: torch.Tensor,
683 + w2_fp4: torch.Tensor,
684 + w1_blockscale: torch.Tensor,
685 + w2_blockscale: torch.Tensor,
686 + g1_alphas: torch.Tensor,
687 + g2_alphas: torch.Tensor,
688 + a1_gscale: torch.Tensor,
689 + a2_gscale: torch.Tensor,
690 + topk_weights: torch.Tensor,
691 + topk_ids: torch.Tensor,
692 + m: int,
693 + n: int,
694 + k: int,
695 + e: int,
696 + device: torch.device,
697 + expert_map: Optional[torch.Tensor] = None,
698 + apply_router_weight_on_input: bool = False) -> torch.Tensor:
699 + assert expert_map is None, ("Expert Parallelism / expert_map "
700 + "is currently not supported for "
701 + "ModelOptNvFp4FusedMoE's cutlass_moe_fp4.")
702 + fn = mk.FusedMoEModularKernel(
703 + MoEPrepareAndFinalizeNoEP(),
704 + CutlassExpertsFp4(
705 + max_experts_per_worker=e,
706 + out_dtype=a.dtype,
707 + per_act_token_quant=False,
708 + per_out_ch_quant=False,
709 + use_batched_format=False,
710 + ),
711 + )
712 + extra_expert_args = {
713 + 'g1_alphas': g1_alphas,
714 + 'g2_alphas': g2_alphas,
715 + 'a1_gscale': a1_gscale,
716 + 'a2_gscale': a2_gscale,
717 + 'm': m,
718 + 'n': n,
719 + 'k': k,
720 + 'e': e,
721 + 'device': device,
722 + }
723 +
724 + # NVFP4 requires two levels of quantization, which involves computing some
725 + # scaling factors dynamically. This makes it incompatible with the typical
726 + # prepare -> MoE -> finalize pipeline. Move the quantization logic into the
727 + # MoE body.
728 + extra_prepare_args = {
729 + 'skip_quant': True,
730 + }
731 + # Similar reason as above.
732 + extra_finalize_args = {
733 + 'skip_weight_reduce': True,
734 + }
735 + return fn(
736 + hidden_states=a,
737 + w1=w1_fp4,
738 + w2=w2_fp4,
739 + topk_weights=topk_weights,
740 + topk_ids=topk_ids,
741 + inplace=False,
742 + activation="silu",
743 + global_num_experts=e,
744 + expert_map=None,
745 + w1_scale=w1_blockscale,
746 + w2_scale=w2_blockscale,
747 + a1_scale=None,
748 + a2_scale=None,
749 + apply_router_weight_on_input=apply_router_weight_on_input,
750 + extra_expert_args=extra_expert_args,
751 + extra_prepare_args=extra_prepare_args,
752 + extra_finalize_args=extra_finalize_args,
753 + )
754
755
756 def _valid_cutlass_block_scaled_grouped_gemm(
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+2 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import functools
4 -from typing import Optional
4 +from typing import Any, Optional
5
6 import torch
7
@@ -152,6 +152,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
152 workspace2: torch.Tensor,
153 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
154 apply_router_weight_on_input: bool,
155 + extra_expert_args: Optional[dict[str, Any]],
156 ):
157 assert self.block_shape is not None
158 assert a1q_scale is not None
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+8 -11
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
3 +from typing import Any, Optional
4
5 import deep_ep
6 import torch
@@ -127,16 +127,12 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
127 expert_topk_weights)
128
129 def prepare(
130 - self,
131 - a1: torch.Tensor,
132 - a1_scale: Optional[torch.Tensor],
133 - a2_scale: Optional[torch.Tensor],
134 - topk_weights: torch.Tensor,
135 - topk_ids: torch.Tensor,
136 - num_experts: int,
137 - expert_map: Optional[torch.Tensor],
138 - apply_router_weight_on_input: bool,
130 + self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
131 + a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
132 + topk_ids: torch.Tensor, num_experts: int,
133 + expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
134 quant_config: FusedMoEQuantConfig,
135 + extra_prepare_args: Optional[dict[str, Any]]
136 ) -> tuple[torch.Tensor, Optional[torch.Tensor],
137 Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
138 Optional[torch.Tensor]]:
@@ -191,7 +187,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
187 def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
188 topk_weights: torch.Tensor, topk_ids: torch.Tensor,
189 apply_router_weight_on_input: bool,
194 - weight_and_reduce_impl: mk.TopKWeightAndReduce) -> None:
190 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
191 + extra_finalize_args: Optional[dict[str, Any]]) -> None:
192
193 assert self.handle is not None
194
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+8 -11
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional, Union
3 +from typing import Any, Optional, Union
4
5 import deep_ep
6 import torch
@@ -111,16 +111,12 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
111 return x, x_scales
112
113 def prepare(
114 - self,
115 - a1: torch.Tensor,
116 - a1_scale: Optional[torch.Tensor],
117 - a2_scale: Optional[torch.Tensor],
118 - topk_weights: torch.Tensor,
119 - topk_ids: torch.Tensor,
120 - num_experts: int,
121 - expert_map: Optional[torch.Tensor],
122 - apply_router_weight_on_input: bool,
114 + self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
115 + a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
116 + topk_ids: torch.Tensor, num_experts: int,
117 + expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
118 quant_config: FusedMoEQuantConfig,
119 + extra_prepare_args: Optional[dict[str, Any]]
120 ) -> tuple[torch.Tensor, Optional[torch.Tensor],
121 Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
122 Optional[torch.Tensor]]:
@@ -169,7 +165,8 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
165 def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
166 topk_weights: torch.Tensor, topk_ids: torch.Tensor,
167 apply_router_weight_on_input: bool,
172 - weight_and_reduce_impl: mk.TopKWeightAndReduce) -> None:
168 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
169 + extra_finalize_args: Optional[dict[str, Any]]) -> None:
170 assert isinstance(
171 weight_and_reduce_impl, TopKWeightAndReduceDelegate
172 ), ("Weight application and reduction happens in the combine kernel.")
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py new
+198
@@ -0,0 +1,198 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any, Optional
4 +
5 +import torch
6 +
7 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 +from vllm.logger import init_logger
9 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
11 + TopKWeightAndReduceDelegate)
12 +from vllm.model_executor.layers.fused_moe.utils import extract_required_args
13 +from vllm.utils.flashinfer import (flashinfer_cutlass_fused_moe,
14 + has_flashinfer_cutlass_fused_moe)
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +def is_valid_flashinfer_cutlass_fused_moe(hidden_states: torch.Tensor,
20 + w1: torch.Tensor,
21 + w2: torch.Tensor) -> bool:
22 + """
23 + Check if the given problem size is supported by the FlashInfer CUTLASS MoE
24 + kernel.
25 + """
26 + if not has_flashinfer_cutlass_fused_moe():
27 + logger.debug_once("FlashInferExperts disabled: "
28 + "flashinfer_cutlass_fused_moe not available.")
29 + return False
30 + # Data type checks
31 + if (w1.dtype != torch.uint8 or w2.dtype != torch.uint8
32 + or hidden_states.dtype
33 + not in [torch.float32, torch.float16, torch.bfloat16]):
34 + logger.debug_once(
35 + "FlashInferExperts disabled: w1/w2 must be torch.uint8 "
36 + f"(got w1={w1.dtype}, w2={w2.dtype}), hidden_states must be "
37 + f"float32, float16, or bfloat16 (got {hidden_states.dtype}).")
38 + return False
39 + return True
40 +
41 +
42 +class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
43 +
44 + def __init__(
45 + self,
46 + use_nvfp4_w4a4: bool = False,
47 + use_fp8_w8a8: bool = False,
48 + use_dp: bool = False,
49 + ep_rank: int = 0,
50 + ep_size: int = 1,
51 + tp_rank: int = 0,
52 + tp_size: int = 1,
53 + num_dispatchers: Optional[int] = None,
54 + use_batched_format: bool = False,
55 + ):
56 + super().__init__(
57 + FusedMoEQuantConfig(
58 + quant_dtype=torch.uint8,
59 + per_act_token_quant=False,
60 + block_shape=None,
61 + ))
62 + self.use_nvfp4_w4a4 = use_nvfp4_w4a4
63 + self.use_fp8_w8a8 = use_fp8_w8a8
64 + self.ep_rank = ep_rank
65 + self.ep_size = ep_size
66 + self.tp_rank = tp_rank
67 + self.tp_size = tp_size
68 + self.use_dp = use_dp
69 + assert not use_batched_format or num_dispatchers is not None
70 + self.num_dispatchers = num_dispatchers
71 +
72 + @property
73 + def activation_formats(
74 + self
75 + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
76 + return (mk.FusedMoEActivationFormat.Standard,
77 + mk.FusedMoEActivationFormat.Standard)
78 +
79 + def supports_expert_map(self) -> bool:
80 + return False
81 +
82 + def supports_chunking(self) -> bool:
83 + # This refers to TP chunking; DP chunking is handled separately.
84 + return True
85 +
86 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
87 + # Let PrepareAndFinalize::finalize() decide the impl.
88 + return TopKWeightAndReduceDelegate()
89 +
90 + def workspace_shapes(
91 + self,
92 + a: torch.Tensor,
93 + aq: torch.Tensor,
94 + M: int,
95 + N: int,
96 + K: int,
97 + topk: int,
98 + global_num_experts: int,
99 + local_num_experts: int,
100 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
101 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
102 + # We use global_num_experts due to how moe_align_block_size handles
103 + # expert_maps.
104 + """
105 + Compute the shapes for the temporary and final outputs of the two gemms
106 + and activation in the fused expert function. Since the gemms are
107 + independent, the workspace for the first gemm can be shared with the
108 + workspace for the last gemm.
109 +
110 + Returns a tuple of:
111 + - workspace13 shape tuple: must be large enough to hold the
112 + result of either expert gemm.
113 + - workspace2 shape tuple: must be large enough to hold the
114 + result of the activation function.
115 + - output shape tuple: must be exact size of the final gemm output.
116 + - Workspace type: The dtype to use for the workspace tensors.
117 + - Note: in order for activation chunking to work, the first dimension
118 + of each tuple must be the number of tokens.
119 + """
120 + assert self.use_nvfp4_w4a4 is True, ("Only nvfp4 quantization is "
121 + "currently supported.")
122 + aq_m, aq_n = aq.shape
123 + workspace2 = ()
124 + output_shape = (aq_m, aq_n * 2)
125 + workspace_dtype = a.dtype
126 + workspace1 = output_shape
127 + # The workspace is determined by `aq`, since it comes after any
128 + # potential communication op and is involved in the expert computation.
129 + return (workspace1, workspace2, output_shape, workspace_dtype)
130 +
131 + def apply(
132 + self,
133 + output: torch.Tensor,
134 + hidden_states: torch.Tensor,
135 + w1: torch.Tensor,
136 + w2: torch.Tensor,
137 + topk_weights: torch.Tensor,
138 + topk_ids: torch.Tensor,
139 + activation: str,
140 + global_num_experts: int,
141 + expert_map: Optional[torch.Tensor],
142 + w1_scale: Optional[torch.Tensor],
143 + w2_scale: Optional[torch.Tensor],
144 + w1_zp: Optional[torch.Tensor],
145 + w2_zp: Optional[torch.Tensor],
146 + a1q_scale: Optional[torch.Tensor],
147 + a2_scale: Optional[torch.Tensor], # Not used
148 + workspace13: Optional[torch.Tensor],
149 + workspace2: Optional[torch.Tensor],
150 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
151 + apply_router_weight_on_input: Optional[bool],
152 + extra_expert_args: Optional[dict[str, Any]],
153 + ):
154 + assert extra_expert_args is not None, \
155 + "extra_expert_args must be provided"
156 + required_keys = [
157 + 'g1_alphas', 'g2_alphas', 'a1_gscale', 'a2_gscale', 'out_dtype'
158 + ]
159 +
160 + g1_alphas, g2_alphas, a1_gscale, a2_gscale, out_dtype = (
161 + extract_required_args(extra_expert_args, required_keys))
162 +
163 + # Flashinfer CUTLASS kernel takes scalar global scales,
164 + # min because inv_scale.
165 + assert self.use_nvfp4_w4a4 is True, ("Only nvfp4 quantization is "
166 + "currently supported.")
167 +
168 + # Ensure w1_scale and w2_scale are not None before calling view
169 + assert w1_scale is not None and w2_scale is not None, (
170 + "w1_scale and w2_scale must not "
171 + "be None for FlashInferExperts")
172 +
173 + assert not apply_router_weight_on_input
174 +
175 + quant_scales = [
176 + a1_gscale,
177 + w1_scale.view(torch.int32),
178 + g1_alphas,
179 + a2_gscale,
180 + w2_scale.view(torch.int32),
181 + g2_alphas,
182 + ]
183 + _ = flashinfer_cutlass_fused_moe(
184 + hidden_states,
185 + topk_ids.to(torch.int),
186 + topk_weights,
187 + # FlashInfer API requires weight to be long for nvfp4
188 + w1.view(torch.long),
189 + w2.view(torch.long),
190 + output_dtype=out_dtype,
191 + quant_scales=quant_scales,
192 + input_sf=a1q_scale,
193 + tp_size=self.tp_size,
194 + tp_rank=self.tp_rank,
195 + ep_size=self.ep_size,
196 + ep_rank=self.ep_rank,
197 + output=output,
198 + )
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py new
+114
@@ -0,0 +1,114 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any, Optional
4 +
5 +import torch
6 +
7 +import vllm.envs as envs
8 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 +from vllm.distributed import get_dp_group
10 +from vllm.forward_context import get_forward_context
11 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12 +from vllm.model_executor.layers.fused_moe.utils import (
13 + extract_required_args, moe_kernel_quantize_input)
14 +from vllm.utils.flashinfer import fp4_swizzle_blockscale
15 +
16 +
17 +def get_local_sizes(local_tokens):
18 + cu_sizes = get_forward_context().dp_metadata.cu_tokens_across_dp_cpu
19 + sizes = [cu_sizes[0].item()]
20 + for i in range(1, len(cu_sizes)):
21 + sizes.append((cu_sizes[i] - cu_sizes[i - 1]).item())
22 + max_num_tokens = envs.VLLM_MOE_DP_CHUNK_SIZE
23 + sizes_chunked = [max_num_tokens] * len(sizes)
24 + if local_tokens < max_num_tokens:
25 + # When the number of local tokens is less than max_num_tokens, all other
26 + # ranks will also have fewer than max_num_tokens. The remaining tokens
27 + # are accounted for as residual.
28 + sizes_chunked = [x % max_num_tokens for x in sizes]
29 +
30 + return sizes_chunked
31 +
32 +
33 +class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
34 +
35 + def __init__(
36 + self,
37 + quant_dtype: Optional[torch.dtype] = None,
38 + per_channel_quant: bool = False,
39 + block_shape: Optional[list[int]] = None,
40 + num_dispatchers: int = 1,
41 + ):
42 + super().__init__()
43 + self.per_channel_quant = per_channel_quant
44 + self.block_shape = block_shape
45 + self.quant_dtype = quant_dtype
46 + self.num_dispatchers_ = num_dispatchers
47 +
48 + @property
49 + def activation_format(self) -> mk.FusedMoEActivationFormat:
50 + return mk.FusedMoEActivationFormat.Standard
51 +
52 + def max_num_tokens_per_rank(self) -> Optional[int]:
53 + return None
54 +
55 + def topk_indices_dtype(self) -> Optional[torch.dtype]:
56 + return None
57 +
58 + def num_dispatchers(self) -> int:
59 + return self.num_dispatchers_
60 +
61 + def prepare(
62 + self,
63 + a1: torch.Tensor,
64 + a1_scale: Optional[torch.Tensor], # Not used
65 + a2_scale: Optional[torch.Tensor], # Not used
66 + topk_weights: torch.Tensor,
67 + topk_ids: torch.Tensor,
68 + num_experts: int,
69 + expert_map: Optional[torch.Tensor],
70 + apply_router_weight_on_input: bool,
71 + quant_config: FusedMoEQuantConfig,
72 + extra_prepare_args: Optional[dict[str, Any]]
73 + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
74 + Optional[torch.Tensor], Optional[torch.Tensor]]:
75 +
76 + assert not apply_router_weight_on_input
77 +
78 + (a1_gscale, use_dp, local_tokens) = extract_required_args(
79 + extra_prepare_args, ['a1_gscale', 'use_dp', 'local_tokens'])
80 +
81 + a1q, a1q_scale = moe_kernel_quantize_input(
82 + a1,
83 + a1_gscale,
84 + quant_config.quant_dtype,
85 + self.per_channel_quant,
86 + self.block_shape,
87 + is_fp4_scale_swizzled=not use_dp, # Swizzling after communication
88 + )
89 + if use_dp:
90 + topk_weights, topk_ids, a1q, a1q_scale = \
91 + get_dp_group().all_gatherv([topk_weights, topk_ids, a1q, a1q_scale], # noqa: E501
92 + dim=0,
93 + sizes=get_local_sizes(local_tokens))
94 + a1_m, a1_n = a1q.shape
95 + a1q_scale = fp4_swizzle_blockscale(a1q_scale, a1_m, a1_n * 2)
96 +
97 + return a1q, a1q_scale, None, topk_ids, topk_weights
98 +
99 + def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
100 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
101 + apply_router_weight_on_input: bool,
102 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
103 + extra_finalize_args: Optional[dict[str, Any]]) -> None:
104 +
105 + (use_dp,
106 + local_tokens) = extract_required_args(extra_finalize_args,
107 + ['use_dp', 'local_tokens'])
108 + if use_dp:
109 + fused_expert_output = get_dp_group().reduce_scatterv(
110 + fused_expert_output,
111 + dim=0,
112 + sizes=get_local_sizes(local_tokens),
113 + )
114 + output.copy_(fused_expert_output)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+15 -21
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Fused batched MoE kernel."""
4 -from typing import Optional
4 +from typing import Any, Optional
5
6 import torch
7
@@ -496,16 +496,12 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
496 return self.num_dispatchers_
497
498 def prepare(
499 - self,
500 - a1: torch.Tensor,
501 - a1_scale: Optional[torch.Tensor],
502 - a2_scale: Optional[torch.Tensor],
503 - topk_weights: torch.Tensor,
504 - topk_ids: torch.Tensor,
505 - num_experts: int,
506 - expert_map: Optional[torch.Tensor],
507 - apply_router_weight_on_input: bool,
499 + self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
500 + a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
501 + topk_ids: torch.Tensor, num_experts: int,
502 + expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
503 quant_config: FusedMoEQuantConfig,
504 + extra_prepare_args: Optional[dict[str, Any]]
505 ) -> tuple[torch.Tensor, Optional[torch.Tensor],
506 Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
507 Optional[torch.Tensor]]:
@@ -594,15 +590,11 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
590
591 return b_a1, b_a1_scale, expert_tokens_meta, None, None
592
597 - def finalize(
598 - self,
599 - output: torch.Tensor,
600 - fused_expert_output: torch.Tensor,
601 - topk_weights: torch.Tensor,
602 - topk_ids: torch.Tensor,
603 - apply_router_weight_on_input: bool,
604 - weight_and_reduce_impl: mk.TopKWeightAndReduce,
605 - ) -> None:
593 + def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
594 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
595 + apply_router_weight_on_input: bool,
596 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
597 + extra_finalize_args: Optional[dict[str, Any]]) -> None:
598 if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
599 weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank)
600 weight_and_reduce_impl.apply(
@@ -706,7 +698,8 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
698 a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
699 workspace2: torch.Tensor,
700 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
709 - apply_router_weight_on_input: bool):
701 + apply_router_weight_on_input: bool,
702 + extra_expert_args: Optional[dict[str, Any]]):
703 assert hidden_states.dim() == 3
704 assert expert_tokens_meta is not None
705 expert_num_tokens = expert_tokens_meta.expert_num_tokens
@@ -911,7 +904,8 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
904 a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
905 workspace2: torch.Tensor,
906 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
914 - apply_router_weight_on_input: bool):
907 + apply_router_weight_on_input: bool,
908 + extra_expert_args: Optional[dict[str, Any]]):
909 # Check constraints.
910 if self.use_int4_w4a16:
911 assert hidden_states.size(-1) // 2 == w1.size(2), (
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
@@ -1646,6 +1646,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1646 workspace2: torch.Tensor,
1647 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
1648 apply_router_weight_on_input: bool,
1649 + extra_expert_args: Optional[dict[str, Any]],
1650 ):
1651 # Check constraints.
1652 if self.use_int4_w4a16:
vllm/model_executor/layers/fused_moe/layer.py
+31 -5
@@ -34,6 +34,7 @@ from vllm.model_executor.utils import set_weight_attrs
34 from vllm.platforms import current_platform
35 from vllm.platforms.interface import CpuArchEnum
36 from vllm.utils import direct_register_custom_op, has_deep_ep, has_pplx
37 +from vllm.utils.flashinfer import has_flashinfer
38
39 if current_platform.is_cuda_alike():
40 from .fused_batched_moe import BatchedTritonExperts
@@ -45,6 +46,9 @@ if current_platform.is_cuda_alike():
46 from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize
47 from .deepep_ll_prepare_finalize import (DEEPEP_QUANT_BLOCK_SHAPE,
48 DeepEPLLPrepareAndFinalize)
49 + if has_flashinfer():
50 + from .flashinfer_cutlass_prepare_finalize import (
51 + FlashInferCutlassMoEPrepareAndFinalize)
52 else:
53 fused_experts = None # type: ignore
54 FusedMoEPermuteExpertsUnpermute = None # type: ignore
@@ -99,6 +103,9 @@ class FusedMoEMethodBase(QuantizeMethodBase):
103
104 prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None
105
106 + if moe.use_flashinfer_cutlass_kernels:
107 + prepare_finalize = FlashInferCutlassMoEPrepareAndFinalize(
108 + quant_dtype=moe.quant_dtype, )
109 if moe.use_pplx_kernels:
110 hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
111 moe.max_num_tokens,
@@ -204,6 +211,12 @@ class FusedMoEMethodBase(QuantizeMethodBase):
211 f"{self.__class__.__name__} must select appropriate gemm "
212 "implementation based on the prepare_finalize")
213
214 + def maybe_swap_experts_impl(
215 + self,
216 + moe_parallel_config: FusedMoEParallelConfig,
217 + ):
218 + pass
219 +
220 @abstractmethod
221 def apply(
222 self,
@@ -744,12 +757,15 @@ class FusedMoE(torch.nn.Module):
757 moe_quant_params["intermediate_size_full"] = intermediate_size
758
759 self.quant_method.create_weights(layer=self, **moe_quant_params)
760 + if isinstance(self.quant_method, FusedMoEMethodBase):
761 + self.quant_method.maybe_swap_experts_impl(self.moe_parallel_config)
762
763 # Chunked all2all staging tensor
764 self.batched_hidden_states: Optional[torch.Tensor] = None
765 self.batched_router_logits: Optional[torch.Tensor] = None
766 if (self.moe_parallel_config.use_pplx_kernels
752 - or self.moe_parallel_config.use_deepep_ll_kernels):
767 + or self.moe_parallel_config.use_deepep_ll_kernels
768 + or self.moe_parallel_config.use_flashinfer_cutlass_kernels):
769 self.batched_hidden_states = torch.zeros(
770 (moe.max_num_tokens, self.hidden_size),
771 dtype=moe.in_dtype,
@@ -801,6 +817,10 @@ class FusedMoE(torch.nn.Module):
817 def use_deepep_ll_kernels(self):
818 return self.moe_parallel_config.use_deepep_ll_kernels
819
820 + @property
821 + def use_flashinfer_cutlass_kernels(self):
822 + return self.moe_parallel_config.use_flashinfer_cutlass_kernels
823 +
824 def _load_per_tensor_weight_scale(self, shard_id: str,
825 param: torch.nn.Parameter,
826 loaded_weight: torch.Tensor,
@@ -1402,9 +1422,9 @@ class FusedMoE(torch.nn.Module):
1422 final_hidden_states, non_blocking=True)
1423
1424 ctx = get_forward_context()
1425 + # flashinfer_cutlass_kernels can handle: optional DP + TP/EP
1426 max_tokens_across_dp = ctx.dp_metadata.max_tokens_across_dp_cpu
1427 moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens
1407 -
1428 num_tokens = full_hidden_states.size(0)
1429 for chunk_start_ in range(0, max_tokens_across_dp,
1430 moe_dp_chunk_size_per_rank):
@@ -1424,13 +1444,20 @@ class FusedMoE(torch.nn.Module):
1444 def forward_impl(self, hidden_states: torch.Tensor,
1445 router_logits: torch.Tensor):
1446 assert self.quant_method is not None
1447 + # Route to the chunked forward path using the FlashInfer Cutlass kernel
1448 + # only when data parallelism (DP) is enabled.
1449 + use_flashinfer_cutlass_kernels = (
1450 + self.dp_size > 1
1451 + and self.moe_parallel_config.use_flashinfer_cutlass_kernels)
1452 if (self.moe_parallel_config.use_pplx_kernels
1428 - or self.moe_parallel_config.use_deepep_ll_kernels):
1453 + or self.moe_parallel_config.use_deepep_ll_kernels
1454 + or use_flashinfer_cutlass_kernels):
1455 return self.forward_impl_chunked(hidden_states, router_logits)
1456
1457 do_naive_dispatch_combine: bool = (
1458 self.dp_size > 1
1433 - and not self.moe_parallel_config.use_deepep_ht_kernels)
1459 + and not self.moe_parallel_config.use_deepep_ht_kernels
1460 + and not self.moe_parallel_config.use_flashinfer_cutlass_kernels)
1461 if do_naive_dispatch_combine:
1462 hidden_states, router_logits = get_ep_group().dispatch(
1463 hidden_states, router_logits)
@@ -1460,7 +1487,6 @@ class FusedMoE(torch.nn.Module):
1487
1488 if do_naive_dispatch_combine:
1489 final_hidden_states = get_ep_group().combine(final_hidden_states)
1463 -
1490 if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
1491 # Default set to False. (May have to add shared expert outputs.
1492 final_hidden_states = self.maybe_all_reduce_tensor_model_parallel(
vllm/model_executor/layers/fused_moe/modular_kernel.py
+60 -39
@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
4 from dataclasses import dataclass
5 from enum import Enum
6 from math import prod
7 -from typing import Optional, final
7 +from typing import Any, Optional, final
8
9 import torch
10
@@ -150,16 +150,12 @@ class FusedMoEPrepareAndFinalize(ABC):
150
151 @abstractmethod
152 def prepare(
153 - self,
154 - a1: torch.Tensor,
155 - a1_scale: Optional[torch.Tensor],
156 - a2_scale: Optional[torch.Tensor],
157 - topk_weights: torch.Tensor,
158 - topk_ids: torch.Tensor,
159 - num_experts: int,
160 - expert_map: Optional[torch.Tensor],
161 - apply_router_weight_on_input: bool,
153 + self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
154 + a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
155 + topk_ids: torch.Tensor, num_experts: int,
156 + expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
157 quant_config: FusedMoEQuantConfig,
158 + extra_prepare_args: Optional[dict[str, Any]]
159 ) -> tuple[torch.Tensor, Optional[torch.Tensor],
160 Optional[ExpertTokensMetadata], Optional[torch.Tensor],
161 Optional[torch.Tensor]]:
@@ -190,15 +186,11 @@ class FusedMoEPrepareAndFinalize(ABC):
186 raise NotImplementedError
187
188 @abstractmethod
193 - def finalize(
194 - self,
195 - output: torch.Tensor,
196 - fused_expert_output: torch.Tensor,
197 - topk_weights: torch.Tensor,
198 - topk_ids: torch.Tensor,
199 - apply_router_weight_on_input: bool,
200 - weight_and_reduce_impl: TopKWeightAndReduce,
201 - ) -> None:
189 + def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
190 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
191 + apply_router_weight_on_input: bool,
192 + weight_and_reduce_impl: TopKWeightAndReduce,
193 + extra_finalize_args: Optional[dict[str, Any]]) -> None:
194 """
195 Perform any combine plus apply weights and perform a reduction on the
196 fused experts output.
@@ -376,6 +368,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
368 workspace2: torch.Tensor,
369 expert_tokens_meta: Optional[ExpertTokensMetadata],
370 apply_router_weight_on_input: bool,
371 + extra_expert_args: Optional[dict[str, Any]],
372 ):
373 """
374 This function computes the intermediate result of a Mixture of Experts
@@ -460,21 +453,19 @@ class FusedMoEModularKernel(torch.nn.Module):
453 f"{fused_experts.__class__.__name__}."
454 f"{fused_experts.activation_formats[0]}")
455
463 - def _do_fused_experts(self, fused_out: Optional[torch.Tensor],
464 - a1: torch.Tensor, a1q: torch.Tensor,
465 - w1: torch.Tensor, w2: torch.Tensor,
466 - topk_weights: torch.Tensor, topk_ids: torch.Tensor,
467 - activation: str, global_num_experts: int,
468 - local_num_experts: int,
469 - expert_map: Optional[torch.Tensor],
470 - w1_scale: Optional[torch.Tensor],
471 - w2_scale: Optional[torch.Tensor],
472 - w1_zp: Optional[torch.Tensor],
473 - w2_zp: Optional[torch.Tensor],
474 - a1q_scale: Optional[torch.Tensor],
475 - a2_scale: Optional[torch.Tensor],
476 - expert_tokens_meta: Optional[ExpertTokensMetadata],
477 - apply_router_weight_on_input: bool) -> torch.Tensor:
456 + def _do_fused_experts(
457 + self, fused_out: Optional[torch.Tensor], a1: torch.Tensor,
458 + a1q: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
459 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
460 + activation: str, global_num_experts: int, local_num_experts: int,
461 + expert_map: Optional[torch.Tensor],
462 + w1_scale: Optional[torch.Tensor], w2_scale: Optional[torch.Tensor],
463 + w1_zp: Optional[torch.Tensor], w2_zp: Optional[torch.Tensor],
464 + a1q_scale: Optional[torch.Tensor],
465 + a2_scale: Optional[torch.Tensor],
466 + expert_tokens_meta: Optional[ExpertTokensMetadata],
467 + apply_router_weight_on_input: bool,
468 + extra_expert_args: Optional[dict[str, Any]]) -> torch.Tensor:
469
470 _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
471
@@ -517,7 +508,8 @@ class FusedMoEModularKernel(torch.nn.Module):
508 workspace13=workspace13,
509 workspace2=workspace2,
510 expert_tokens_meta=expert_tokens_meta,
520 - apply_router_weight_on_input=apply_router_weight_on_input)
511 + apply_router_weight_on_input=apply_router_weight_on_input,
512 + extra_expert_args=extra_expert_args)
513
514 return fused_out
515
@@ -541,6 +533,7 @@ class FusedMoEModularKernel(torch.nn.Module):
533 a2_scale: Optional[torch.Tensor],
534 expert_tokens_meta: Optional[ExpertTokensMetadata],
535 apply_router_weight_on_input: bool,
536 + extra_expert_args: Optional[dict[str, Any]],
537 ) -> torch.Tensor:
538
539 _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
@@ -568,7 +561,8 @@ class FusedMoEModularKernel(torch.nn.Module):
561 a1q_scale=a1q_scale,
562 a2_scale=a2_scale,
563 expert_tokens_meta=expert_tokens_meta,
571 - apply_router_weight_on_input=apply_router_weight_on_input)
564 + apply_router_weight_on_input=apply_router_weight_on_input,
565 + extra_expert_args=extra_expert_args)
566
567 # Chunking required case
568 assert num_chunks > 1
@@ -624,6 +618,15 @@ class FusedMoEModularKernel(torch.nn.Module):
618 expert_num_tokens=c_expert_num_tokens,
619 expert_num_tokens_cpu=c_expert_num_tokens_cpu)
620
621 + m = None
622 + if extra_expert_args is not None and 'm' in extra_expert_args:
623 + m = extra_expert_args.get('m')
624 +
625 + if extra_expert_args is not None:
626 + chunked_extra_expert_args = extra_expert_args
627 + else:
628 + chunked_extra_expert_args = {}
629 +
630 for chunk_idx in range(num_chunks):
631 c_a1q, c_a1q_scale, c_a2_scale, c_topk_ids, c_topk_weights = (
632 slice_input_tensors(chunk_idx))
@@ -634,6 +637,11 @@ class FusedMoEModularKernel(torch.nn.Module):
637 expert_tokens_meta, c_topk_ids, local_num_experts,
638 expert_map)
639
640 + s = chunk_idx * CHUNK_SIZE
641 + e = min(s + CHUNK_SIZE, M)
642 +
643 + if m is not None:
644 + chunked_extra_expert_args['m'] = e - s
645 self._do_fused_experts(
646 fused_out=slice_output_tensor(chunk_idx),
647 a1=a1,
@@ -653,7 +661,8 @@ class FusedMoEModularKernel(torch.nn.Module):
661 a1q_scale=c_a1q_scale,
662 a2_scale=c_a2_scale,
663 expert_tokens_meta=c_expert_tokens_meta,
656 - apply_router_weight_on_input=apply_router_weight_on_input)
664 + apply_router_weight_on_input=apply_router_weight_on_input,
665 + extra_expert_args=chunked_extra_expert_args)
666
667 return fused_out
668
@@ -675,6 +684,9 @@ class FusedMoEModularKernel(torch.nn.Module):
684 a1_scale: Optional[torch.Tensor] = None,
685 a2_scale: Optional[torch.Tensor] = None,
686 apply_router_weight_on_input: bool = False,
687 + extra_expert_args: Optional[dict] = None,
688 + extra_prepare_args: Optional[dict] = None,
689 + extra_finalize_args: Optional[dict] = None,
690 ) -> torch.Tensor:
691 """
692 This function computes a Mixture of Experts (MoE) layer using two sets
@@ -707,6 +719,12 @@ class FusedMoEModularKernel(torch.nn.Module):
719 - apply_router_weight_on_input (bool): When true, the topk weights are
720 applied directly on the inputs. This is only applicable when topk is
721 1.
722 + - extra_expert_args (Optional[dict]): Extra keyword arguments to pass to
723 + fused_experts.apply.
724 + - extra_prepare_args (Optional[dict]): Extra keyword arguments to pass
725 + to prepare.
726 + - extra_finalize_args (Optional[dict]): Extra keyword arguments to pass
727 + to finalize.
728
729 Returns:
730 - torch.Tensor: The output tensor after applying the MoE layer.
@@ -730,6 +748,7 @@ class FusedMoEModularKernel(torch.nn.Module):
748 expert_map,
749 apply_router_weight_on_input,
750 self.fused_experts.quant_config,
751 + extra_prepare_args,
752 )
753
754 # Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
@@ -766,11 +785,13 @@ class FusedMoEModularKernel(torch.nn.Module):
785 a1q_scale=a1q_scale,
786 a2_scale=a2_scale,
787 expert_tokens_meta=expert_tokens_meta,
769 - apply_router_weight_on_input=apply_router_weight_on_input)
788 + apply_router_weight_on_input=apply_router_weight_on_input,
789 + extra_expert_args=extra_expert_args)
790
791 self.prepare_finalize.finalize(
792 output, fused_out, topk_weights, topk_ids,
793 apply_router_weight_on_input,
774 - self.fused_experts.finalize_weight_and_reduce_impl())
794 + self.fused_experts.finalize_weight_and_reduce_impl(),
795 + extra_finalize_args)
796
797 return output
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+11 -19
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
3 +from typing import Any, Optional
4
5 import pplx_kernels as pplx
6 import torch
@@ -89,16 +89,12 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
89 return self.num_dispatchers_
90
91 def prepare(
92 - self,
93 - a1: torch.Tensor,
94 - a1_scale: Optional[torch.Tensor],
95 - a2_scale: Optional[torch.Tensor],
96 - topk_weights: torch.Tensor,
97 - topk_ids: torch.Tensor,
98 - num_experts: int,
99 - expert_map: Optional[torch.Tensor],
100 - apply_router_weight_on_input: bool,
92 + self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
93 + a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
94 + topk_ids: torch.Tensor, num_experts: int,
95 + expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
96 quant_config: FusedMoEQuantConfig,
97 + extra_prepare_args: Optional[dict[str, Any]]
98 ) -> tuple[torch.Tensor, Optional[torch.Tensor],
99 Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
100 Optional[torch.Tensor]]:
@@ -217,15 +213,11 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
213
214 return expert_x, expert_x_scale, expert_tokens_meta, None, None
215
220 - def finalize(
221 - self,
222 - output: torch.Tensor,
223 - fused_expert_output: torch.Tensor,
224 - topk_weights: torch.Tensor,
225 - topk_ids: torch.Tensor,
226 - apply_router_weight_on_input: bool,
227 - weight_and_reduce_impl: mk.TopKWeightAndReduce,
228 - ) -> None:
216 + def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
217 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
218 + apply_router_weight_on_input: bool,
219 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
220 + extra_finalize_args: Optional[dict[str, Any]]) -> None:
221 assert isinstance(
222 weight_and_reduce_impl, TopKWeightAndReduceDelegate
223 ), ("Weight application and reduction happens in the combine kernel.")
vllm/model_executor/layers/fused_moe/prepare_finalize.py
+26 -18
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
3 +from typing import Any, Optional
4
5 import torch
6
@@ -38,6 +38,7 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
38 expert_map: Optional[torch.Tensor],
39 apply_router_weight_on_input: bool,
40 quant_config: FusedMoEQuantConfig,
41 + extra_prepare_args: Optional[dict[str, Any]],
42 ) -> tuple[torch.Tensor, Optional[torch.Tensor],
43 Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
44 Optional[torch.Tensor]]:
@@ -48,26 +49,33 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
49 assert topk == 1, \
50 "apply_router_weight_on_input is only implemented for topk=1"
51 a1.mul_(topk_weights.to(a1.dtype))
52 +
53 + if (extra_prepare_args is not None
54 + and extra_prepare_args.get("skip_quant", True)):
55 + # Skip quantization if explicitly requested
56 + return a1, None, None, None, None
57 +
58 a1q, a1q_scale = moe_kernel_quantize_input(
59 a1, a1_scale, quant_config.quant_dtype,
60 quant_config.per_act_token_quant, quant_config.block_shape)
61
62 return a1q, a1q_scale, None, None, None
63
57 - def finalize(
58 - self,
59 - output: torch.Tensor,
60 - fused_expert_output: torch.Tensor,
61 - topk_weights: torch.Tensor,
62 - topk_ids: torch.Tensor,
63 - apply_router_weight_on_input: bool,
64 - weight_and_reduce_impl: mk.TopKWeightAndReduce,
65 - ) -> None:
66 - if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
67 - weight_and_reduce_impl = TopKWeightAndReduceContiguous()
68 - weight_and_reduce_impl.apply(
69 - output=output,
70 - fused_expert_output=fused_expert_output,
71 - topk_weights=topk_weights,
72 - topk_ids=topk_ids,
73 - apply_router_weight_on_input=apply_router_weight_on_input)
64 + def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
65 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
66 + apply_router_weight_on_input: bool,
67 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
68 + extra_finalize_args: Optional[dict[str, Any]]) -> None:
69 + if (extra_finalize_args is not None
70 + and extra_finalize_args.get("skip_weight_reduce", True)):
71 + assert output.shape == fused_expert_output.shape
72 + output.copy_(fused_expert_output)
73 + else:
74 + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
75 + weight_and_reduce_impl = TopKWeightAndReduceContiguous()
76 + weight_and_reduce_impl.apply(
77 + output=output,
78 + fused_expert_output=fused_expert_output,
79 + topk_weights=topk_weights,
80 + topk_ids=topk_ids,
81 + apply_router_weight_on_input=apply_router_weight_on_input)
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+14 -23
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
3 +from typing import Any, Optional
4
5 import torch
6
@@ -119,28 +119,18 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
119 local_num_experts,
120 expert_tokens_meta)
121
122 - def apply(
123 - self,
124 - output: torch.Tensor,
125 - hidden_states: torch.Tensor,
126 - w1: torch.Tensor,
127 - w2: torch.Tensor,
128 - topk_weights: torch.Tensor,
129 - topk_ids: torch.Tensor,
130 - activation: str,
131 - global_num_experts: int,
132 - expert_map: Optional[torch.Tensor],
133 - w1_scale: Optional[torch.Tensor],
134 - w2_scale: Optional[torch.Tensor],
135 - w1_zp: Optional[torch.Tensor],
136 - w2_zp: Optional[torch.Tensor],
137 - a1q_scale: Optional[torch.Tensor],
138 - a2_scale: Optional[torch.Tensor],
139 - workspace13: torch.Tensor,
140 - workspace2: torch.Tensor,
141 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
142 - apply_router_weight_on_input: bool,
143 - ):
122 + def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
123 + w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
124 + topk_ids: torch.Tensor, activation: str, global_num_experts: int,
125 + expert_map: Optional[torch.Tensor],
126 + w1_scale: Optional[torch.Tensor],
127 + w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
128 + w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
129 + a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
130 + workspace2: torch.Tensor,
131 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
132 + apply_router_weight_on_input: bool,
133 + extra_expert_args: Optional[dict[str, Any]]):
134 use_deep_gemm = (self.allow_deep_gemm
135 and (_valid_deep_gemm(hidden_states, w1, w2)
136 or is_blackwell_deep_gemm_used()))
@@ -168,4 +158,5 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
158 workspace2,
159 expert_tokens_meta,
160 apply_router_weight_on_input,
161 + extra_expert_args,
162 )
vllm/model_executor/layers/fused_moe/utils.py
+31 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from math import prod
4 -from typing import Optional, Union
4 +from typing import Any, Optional, Union
5
6 import torch
7
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
15 from vllm.platforms import current_platform
16 from vllm.triton_utils import tl, triton
17 from vllm.utils import cdiv
18 +from vllm.utils.flashinfer import fp4_quantize
19
20
21 @triton.jit
@@ -98,6 +99,16 @@ def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
99 return x.flatten()[:prod(v)].view(*v)
100
101
102 +def _fp4_quantize(
103 + A: torch.Tensor,
104 + A_scale: Optional[torch.Tensor],
105 + is_sf_swizzled_layout: bool,
106 +) -> tuple[torch.Tensor, torch.Tensor]:
107 + return fp4_quantize(A,
108 + A_scale,
109 + is_sf_swizzled_layout=is_sf_swizzled_layout)
110 +
111 +
112 def _fp8_quantize(
113 A: torch.Tensor,
114 A_scale: Optional[torch.Tensor],
@@ -172,11 +183,16 @@ def moe_kernel_quantize_input(
183 quant_dtype: Union[None, torch.dtype, str],
184 per_act_token_quant: bool,
185 block_shape: Optional[list[int]] = None,
186 + is_fp4_scale_swizzled: bool = True,
187 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
188 if quant_dtype == torch.float8_e4m3fn:
189 return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
190 elif quant_dtype == torch.int8:
191 return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)
192 + elif quant_dtype == torch.uint8: # nvfp4
193 + return _fp4_quantize(A,
194 + A_scale,
195 + is_sf_swizzled_layout=is_fp4_scale_swizzled)
196 elif quant_dtype == "mxfp4":
197 return _mxfp4_quantize(A, A_scale, per_act_token_quant, block_shape)
198 else:
@@ -236,3 +252,17 @@ def _validate_scale_shape(
252 assert block_shape is not None
253 expected = (a.shape[0], cdiv(a.shape[1], block_shape[1]))
254 assert a_scale.shape == expected, f"{a_scale.shape} == {expected}"
255 +
256 +
257 +def extract_required_args(
258 + extra_args: Optional[dict[str, Any]],
259 + required_keys: list[str],
260 +) -> tuple[Any, ...]:
261 + if extra_args is None:
262 + raise ValueError("`extra_args` must be provided.")
263 +
264 + missing_keys = [k for k in required_keys if k not in extra_args]
265 + if missing_keys:
266 + raise ValueError(f"Missing keys in `extra_args`: {missing_keys}")
267 +
268 + return tuple(extra_args[k] for k in required_keys)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+5 -5
@@ -339,19 +339,19 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
339 return cutlass_moe_fp4(
340 a=x,
341 w1_fp4=layer.w13_weight,
342 - w1_blockscale=layer.w13_blockscale_swizzled,
343 - w1_alphas=layer.g1_alphas,
342 w2_fp4=layer.w2_weight,
343 + w1_blockscale=layer.w13_blockscale_swizzled,
344 w2_blockscale=layer.w2_blockscale_swizzled,
346 - w2_alphas=layer.g2_alphas,
345 + g1_alphas=layer.g1_alphas,
346 + g2_alphas=layer.g2_alphas,
347 + a1_gscale=layer.w13_input_scale_quant,
348 + a2_gscale=layer.w2_input_scale_quant,
349 topk_weights=topk_weights,
350 topk_ids=topk_ids,
351 m=x.shape[0],
352 n=layer.w2_weight.shape[2] * 2,
353 k=x.shape[1],
354 e=layer.w13_weight.shape[0],
353 - a1_gscale=layer.w13_input_scale_quant,
354 - a2_gscale=layer.w2_input_scale_quant,
355 device=x.device,
356 apply_router_weight_on_input=apply_router_weight_on_input).to(
357 x.dtype)
vllm/model_executor/layers/quantization/modelopt.py
+179 -32
@@ -7,9 +7,15 @@ import torch
7 from torch.nn import Module
8 from torch.nn.parameter import Parameter
9
10 +import vllm.envs as envs
11 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 from vllm._custom_ops import (cutlass_scaled_fp4_mm,
13 cutlass_scaled_mm_supports_fp4, scaled_fp4_quant)
14 +from vllm.distributed import get_ep_group
15 from vllm.logger import init_logger
16 +from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
17 +from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
18 + FlashInferCutlassMoEPrepareAndFinalize)
19 from vllm.model_executor.layers.fused_moe.layer import (
20 FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
21 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
@@ -713,6 +719,18 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
719 self.quant_config = quant_config
720 self.cutlass_nvfp4_supported = cutlass_fp4_supported()
721 self.use_marlin = False
722 + self.allow_flashinfer_cutlass = False
723 +
724 + if envs.VLLM_USE_FLASHINFER_MOE:
725 + if self.cutlass_nvfp4_supported and current_platform.is_cuda() \
726 + and current_platform.is_device_capability(100):
727 + logger.info_once(
728 + "Using FlashInfer kernels for ModelOptNvFp4FusedMoE.")
729 + self.allow_flashinfer_cutlass = True
730 + else:
731 + logger.warning_once(
732 + "Flashinfer CUTLASS Fused MoE not supported "
733 + "or found on the current platform.")
734
735 if not self.cutlass_nvfp4_supported:
736 if is_fp4_marlin_supported():
@@ -722,6 +740,73 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
740 " quantization. Please use Blackwell and"
741 " above.")
742
743 + self.fused_experts = None # type: ignore
744 +
745 + def maybe_swap_experts_impl(
746 + self,
747 + moe_parallel_config: FusedMoEParallelConfig,
748 + ):
749 + if not self.allow_flashinfer_cutlass:
750 + return
751 +
752 + logger.debug_once("FlashInferExperts")
753 + # default to TP/EP case only
754 +
755 + experts_kwargs: dict[str, Any] = {
756 + "use_nvfp4_w4a4": True,
757 + "use_dp": moe_parallel_config.dp_size > 1,
758 + "ep_rank": moe_parallel_config.ep_rank,
759 + "ep_size": moe_parallel_config.ep_size,
760 + "tp_rank": moe_parallel_config.tp_rank,
761 + "tp_size": moe_parallel_config.tp_size,
762 + }
763 +
764 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
765 + FlashInferExperts)
766 + experts = FlashInferExperts(**experts_kwargs)
767 + self.fused_experts = mk.FusedMoEModularKernel(
768 + FlashInferCutlassMoEPrepareAndFinalize(
769 + quant_dtype=torch.uint8,
770 + #meaning 2x e2m1 packed in one, kernel requirement
771 + ),
772 + experts,
773 + )
774 +
775 + # This method update self.fused_experts
776 + # only prepare_finalize is not None call select_gemm_impl
777 + # so when native cutlass fp4, fused_expert is in fuse_moe.py fused_expert
778 + # when it's not called(TP case), we still have 2 kernels to use.
779 + def select_gemm_impl(self, prepare_finalize,
780 + moe) -> mk.FusedMoEPermuteExpertsUnpermute:
781 +
782 + assert moe is not None
783 + assert prepare_finalize is not None
784 + experts = None
785 + all2all_manager = get_ep_group().device_communicator.all2all_manager
786 + assert all2all_manager is not None
787 + if self.allow_flashinfer_cutlass:
788 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
789 + FlashInferExperts)
790 + logger.debug_once("Using FlashInferExperts")
791 + experts = FlashInferExperts(
792 + use_nvfp4_w4a4=True,
793 + use_dp=moe.moe_parallel_config.dp_size > 1,
794 + ep_rank=moe.moe_parallel_config.ep_rank,
795 + ep_size=moe.moe_parallel_config.ep_size,
796 + tp_rank=moe.moe_parallel_config.tp_rank,
797 + tp_size=moe.moe_parallel_config.tp_size,
798 + )
799 + else:
800 + assert moe.dp_size > 1
801 + logger.debug_once("Using CutlassExpertsFp4")
802 + # Currently CutlassExpertsFp4 doesn't support DP
803 + raise ValueError(
804 + "CutlassExpertsFp4 doesn't support DP. "
805 + "Use flashinfer CUTLASS FusedMoE(VLLM_USE_FLASHINFER_MOE)"
806 + " backend instead.")
807 +
808 + return experts
809 +
810 def uses_weight_scale_2_pattern(self) -> bool:
811 """
812 FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales.
@@ -842,8 +927,30 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
927 if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
928
929 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
845 -
930 # GEMM 1
931 + # The FlashInfer Cutlass fused MoE kernel expects the combined weights
932 + # to be ordered as [w3, w1], unlike the standard [w1, w3] layout.
933 + gemm1_weight = layer.w13_weight.data
934 + gemm1_weight_scale = layer.w13_weight_scale.data
935 +
936 + if self.allow_flashinfer_cutlass:
937 + dim = -2
938 + size = gemm1_weight.size(dim)
939 + assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
940 + half = size // 2
941 +
942 + # Reorder weight
943 + w1, w3 = gemm1_weight.split(half, dim=dim)
944 + gemm1_weight = torch.cat([w3, w1], dim=dim).contiguous()
945 +
946 + # Reorder scale
947 + s1, s3 = gemm1_weight_scale.split(half, dim=dim)
948 + gemm1_weight_scale = torch.cat([s3, s1], dim=dim).contiguous()
949 +
950 + layer.w13_weight = Parameter(gemm1_weight, requires_grad=False)
951 + layer.w13_weight_scale = Parameter(gemm1_weight_scale,
952 + requires_grad=False)
953 +
954 if not torch.allclose(layer.w13_weight_scale_2[:, 0],
955 layer.w13_weight_scale_2[:, 1]):
956 logger.warning_once(
@@ -874,9 +981,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
981 layer.w13_input_scale_quant = Parameter(
982 (1 / w13_input_scale).to(torch.float32), requires_grad=False)
983
877 - layer.w13_weight = Parameter(layer.w13_weight.data,
878 - requires_grad=False)
879 -
984 # GEMM 2
985 layer.g2_alphas = Parameter(
986 (layer.w2_input_scale * layer.w2_weight_scale_2).to(torch.float32),
@@ -961,31 +1065,74 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1065 global_num_experts=global_num_experts,
1066 expert_map=expert_map)
1067
964 - assert expert_map is None, ("Expert Parallelism / expert_map "
965 - "is currently not supported for "
966 - "ModelOptNvFp4FusedMoE.")
967 -
968 - from vllm.model_executor.layers.fused_moe.cutlass_moe import (
969 - cutlass_moe_fp4)
970 -
971 - # Cutlass moe takes in activations in BF16/Half precision
972 - # and fp4 quantized weights loaded from the checkpoint
973 - return cutlass_moe_fp4(
974 - a=x,
975 - w1_fp4=layer.w13_weight,
976 - w1_blockscale=layer.w13_blockscale_swizzled,
977 - w1_alphas=layer.g1_alphas,
978 - w2_fp4=layer.w2_weight,
979 - w2_blockscale=layer.w2_blockscale_swizzled,
980 - w2_alphas=layer.g2_alphas,
981 - topk_weights=topk_weights,
982 - topk_ids=topk_ids,
983 - m=x.shape[0],
984 - n=layer.w2_weight.shape[2] * 2,
985 - k=x.shape[1],
986 - e=layer.w13_weight.shape[0],
987 - a1_gscale=layer.w13_input_scale_quant,
988 - a2_gscale=layer.w2_input_scale_quant,
989 - device=x.device,
990 - apply_router_weight_on_input=apply_router_weight_on_input).to(
991 - x.dtype)
1068 + if self.fused_experts is None:
1069 + # If no modular kernel is provided, use cutlass_moe_fp4 for TP case
1070 + # only (no EP).
1071 + from vllm.model_executor.layers.fused_moe.cutlass_moe import (
1072 + cutlass_moe_fp4)
1073 + out = cutlass_moe_fp4(
1074 + a=x,
1075 + w1_fp4=layer.w13_weight,
1076 + w2_fp4=layer.w2_weight,
1077 + w1_blockscale=layer.w13_blockscale_swizzled,
1078 + w2_blockscale=layer.w2_blockscale_swizzled,
1079 + g1_alphas=layer.g1_alphas,
1080 + g2_alphas=layer.g2_alphas,
1081 + a1_gscale=layer.w13_input_scale_quant,
1082 + a2_gscale=layer.w2_input_scale_quant,
1083 + topk_weights=topk_weights,
1084 + topk_ids=topk_ids,
1085 + m=x.shape[0],
1086 + n=layer.w2_weight.shape[2] * 2,
1087 + k=x.shape[1],
1088 + e=layer.w13_weight.shape[0],
1089 + device=x.device,
1090 + expert_map=expert_map,
1091 + apply_router_weight_on_input=apply_router_weight_on_input)
1092 + else:
1093 + # TP or DP case
1094 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
1095 + is_valid_flashinfer_cutlass_fused_moe)
1096 + assert is_valid_flashinfer_cutlass_fused_moe(
1097 + x, layer.w13_weight, layer.w2_weight), (
1098 + "Flashinfer CUTLASS Fused MoE not applicable!")
1099 +
1100 + a1_gscale = torch.min(layer.w13_input_scale_quant)
1101 + a2_gscale = torch.min(layer.w2_input_scale_quant)
1102 + extra_expert_args = {
1103 + 'g1_alphas': layer.g1_alphas,
1104 + 'g2_alphas': layer.g2_alphas,
1105 + 'out_dtype': x.dtype,
1106 + # Avoid confusion with a1_scale and a2_scale
1107 + # where are batch size related.
1108 + 'a1_gscale': a1_gscale,
1109 + 'a2_gscale': a2_gscale,
1110 + }
1111 + extra_prepare_args = {
1112 + 'use_dp': layer.dp_size > 1,
1113 + 'local_tokens': x.shape[0],
1114 + 'a1_gscale': a1_gscale,
1115 + }
1116 + extra_finalize_args = {
1117 + 'use_dp': layer.dp_size > 1,
1118 + 'local_tokens': x.shape[0],
1119 + }
1120 +
1121 + out = self.fused_experts(
1122 + hidden_states=x,
1123 + w1=layer.w13_weight,
1124 + w2=layer.w2_weight,
1125 + topk_weights=topk_weights,
1126 + topk_ids=topk_ids,
1127 + inplace=False, # TODO(shuw): fix later, now output is high prec
1128 + activation=activation,
1129 + global_num_experts=global_num_experts,
1130 + expert_map=expert_map,
1131 + w1_scale=layer.w13_blockscale_swizzled,
1132 + w2_scale=layer.w2_blockscale_swizzled,
1133 + apply_router_weight_on_input=apply_router_weight_on_input,
1134 + extra_expert_args=extra_expert_args,
1135 + extra_prepare_args=extra_prepare_args,
1136 + extra_finalize_args=extra_finalize_args,
1137 + )
1138 + return out
vllm/utils/flashinfer.py new
+107
@@ -0,0 +1,107 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Compatibility wrapper for FlashInfer API changes.
4 +
5 +Users of vLLM should always import **only** these wrappers.
6 +"""
7 +from __future__ import annotations
8 +
9 +import contextlib
10 +import functools
11 +import importlib
12 +import importlib.util
13 +from typing import Any, Callable, NoReturn
14 +
15 +from vllm.logger import init_logger
16 +
17 +logger = init_logger(__name__)
18 +
19 +
20 +@functools.cache
21 +def has_flashinfer() -> bool:
22 + """Return ``True`` if FlashInfer is available."""
23 + # Use find_spec to check if the module exists without importing it
24 + # This avoids potential CUDA initialization side effects
25 + return importlib.util.find_spec("flashinfer") is not None
26 +
27 +
28 +def _missing(*_: Any, **__: Any) -> NoReturn:
29 + """Placeholder for unavailable FlashInfer backend."""
30 + raise RuntimeError(
31 + "FlashInfer backend is not available. Please install the package "
32 + "to enable FlashInfer kernels: "
33 + "https://github.com/flashinfer-ai/flashinfer")
34 +
35 +
36 +def _get_submodule(module_name: str) -> Any | None:
37 + """Safely import a submodule and return it, or None if not available."""
38 + try:
39 + return importlib.import_module(module_name)
40 + except (ImportError, ModuleNotFoundError):
41 + return None
42 +
43 +
44 +# General lazy import wrapper
45 +def _lazy_import_wrapper(module_name: str,
46 + attr_name: str,
47 + fallback_fn: Callable[..., Any] = _missing):
48 + """Create a lazy import wrapper for a specific function."""
49 +
50 + @functools.cache
51 + def _get_impl():
52 + if not has_flashinfer():
53 + return None
54 + mod = _get_submodule(module_name)
55 + return getattr(mod, attr_name, None) if mod else None
56 +
57 + def wrapper(*args, **kwargs):
58 + impl = _get_impl()
59 + if impl is None:
60 + return fallback_fn(*args, **kwargs)
61 + return impl(*args, **kwargs)
62 +
63 + return wrapper
64 +
65 +
66 +# Create lazy wrappers for each function
67 +flashinfer_cutlass_fused_moe = _lazy_import_wrapper("flashinfer.fused_moe",
68 + "cutlass_fused_moe")
69 +fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")
70 +fp4_swizzle_blockscale = _lazy_import_wrapper("flashinfer",
71 + "fp4_swizzle_blockscale")
72 +
73 +# Special case for autotune since it returns a context manager
74 +autotune = _lazy_import_wrapper(
75 + "flashinfer.autotuner",
76 + "autotune",
77 + fallback_fn=lambda *args, **kwargs: contextlib.nullcontext())
78 +
79 +
80 +@functools.cache
81 +def has_flashinfer_cutlass_fused_moe() -> bool:
82 + """Return ``True`` if FlashInfer CUTLASS fused MoE is available."""
83 + if not has_flashinfer():
84 + return False
85 +
86 + # Check if all required functions are available
87 + required_functions = [
88 + ("flashinfer.fused_moe", "cutlass_fused_moe"),
89 + ("flashinfer", "fp4_quantize"),
90 + ("flashinfer", "fp4_swizzle_blockscale"),
91 + ]
92 +
93 + for module_name, attr_name in required_functions:
94 + mod = _get_submodule(module_name)
95 + if not mod or not hasattr(mod, attr_name):
96 + return False
97 + return True
98 +
99 +
100 +__all__ = [
101 + "has_flashinfer",
102 + "has_flashinfer_cutlass_fused_moe",
103 + "flashinfer_cutlass_fused_moe",
104 + "fp4_quantize",
105 + "fp4_swizzle_blockscale",
106 + "autotune",
107 +]