@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 074854b24f
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Kernel][B200] `mxfp4` fused cutlass moe (#23696)

Signed-off-by: Duncan Moss <djm.moss@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

Duncan Moss committed Sep 11, 2025 at 14:04 UTC 074854b24f6e0b1e237a004283e1f46d98c0d73c
5 files changed +626 -64
tests/kernels/moe/test_mxfp4_moe.py
+319
@@ -11,6 +11,7 @@ import torch
11 from packaging import version
12
13 from vllm.platforms import current_platform
14 +from vllm.utils.flashinfer import has_flashinfer
15
16 QUARK_MXFP4_AVAILABLE = importlib.util.find_spec(
17 "quark") is not None and version.parse(
@@ -19,6 +20,10 @@ QUARK_MXFP4_AVAILABLE = importlib.util.find_spec(
20 TRTLLM_GEN_MXFP4_AVAILABLE = current_platform.is_cuda(
21 ) and current_platform.is_device_capability(100)
22
23 +HOPPER_MXFP4_BF16_AVAILABLE = (current_platform.is_cuda()
24 + and current_platform.is_device_capability(90)
25 + and has_flashinfer())
26 +
27 if TRTLLM_GEN_MXFP4_AVAILABLE:
28 from flashinfer import (fp4_quantize, mxfp8_quantize,
29 next_positive_power_of_2,
@@ -542,3 +547,317 @@ def test_trtllm_gen_mxfp4_fused_moe(
547 transpose_optimized=transpose_optimized)
548 # relatively loose check since the mxfp4 quantization is less accurate
549 check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8)
550 +
551 +
552 +def _interleave_scales_lastdim_by4(scales: torch.Tensor) -> torch.Tensor:
553 + """Interleave scales on the last dimension by groups of 4, matching
554 + the transformation in mxfp4.py's BF16 (Hopper) path."""
555 + s = scales.to(torch.uint8)
556 + s_shape = s.shape
557 + assert s_shape[-1] % 4 == 0
558 + s = s.reshape(*s_shape[:-1], s_shape[-1] // 4, 4)
559 + # Move the 4-group dimension before the row dimension
560 + permuted = s.permute(0, 2, 1, 3)
561 + # Merge the row dim with the 4-group dim
562 + return permuted.reshape(s_shape[0], s_shape[-1] // 4, s_shape[1] * 4)
563 +
564 +
565 +@pytest.mark.parametrize("topk", [1, 4])
566 +@pytest.mark.parametrize("num_experts", [32])
567 +@pytest.mark.parametrize("num_tokens", [1, 128])
568 +@pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)])
569 +@pytest.mark.parametrize("alpha,beta,limit", [(1.0, 1.0, None),
570 + (1.702, 1.0, 7.0)])
571 +@pytest.mark.skipif(
572 + not HOPPER_MXFP4_BF16_AVAILABLE,
573 + reason="nvidia gpu sm90 and flashinfer are required for this test",
574 +)
575 +def test_flashinfer_cutlass_mxfp4_fused_moe(
576 + topk: int,
577 + num_experts: int,
578 + num_tokens: int,
579 + intermediate_size: int,
580 + hidden_size: int,
581 + alpha: float,
582 + beta: float,
583 + limit: Optional[float],
584 +):
585 + torch.manual_seed(42)
586 + device = "cuda:0"
587 +
588 + # Inputs
589 + hidden_states = torch.randn(num_tokens,
590 + hidden_size,
591 + device=device,
592 + dtype=torch.bfloat16)
593 + # Random MXFP4 weights and scales (uint8), contiguous [w1; w3]
594 + w13_q = torch.randint(
595 + 0,
596 + 256, (num_experts, 2 * intermediate_size, hidden_size // 2),
597 + device=device,
598 + dtype=torch.uint8)
599 + w13_scale = torch.randint(
600 + 118,
601 + 123, (num_experts, 2 * intermediate_size, hidden_size // 32),
602 + device=device,
603 + dtype=torch.uint8)
604 +
605 + w2_q = torch.randint(0,
606 + 256,
607 + (num_experts, hidden_size, intermediate_size // 2),
608 + device=device,
609 + dtype=torch.uint8)
610 + w2_scale = torch.randint(
611 + 118,
612 + 123, (num_experts, hidden_size, intermediate_size // 32),
613 + device=device,
614 + dtype=torch.uint8)
615 + # Bias contiguous [b1; b3]
616 + bias13 = (torch.randn(num_experts,
617 + 2 * intermediate_size,
618 + device=device,
619 + dtype=torch.bfloat16) * 10)
620 + bias2 = (torch.randn(
621 + num_experts, hidden_size, device=device, dtype=torch.bfloat16) * 10)
622 + router_logits = torch.rand(num_tokens,
623 + num_experts,
624 + dtype=torch.float32,
625 + device=device)
626 +
627 + w13_ref = mxfp4_dequantize(w13_q.clone(), w13_scale.clone()).reshape(
628 + num_experts, 2 * intermediate_size, hidden_size)
629 + w2_ref = mxfp4_dequantize(w2_q.clone(), w2_scale.clone()).reshape(
630 + num_experts, hidden_size, intermediate_size)
631 + ref = reference_moe(router_logits.to(torch.float32), topk, num_experts,
632 + hidden_states.to(torch.float32), w13_ref,
633 + bias13.to(torch.float32), w2_ref,
634 + bias2.to(torch.float32), alpha, beta, limit, 'bf16')
635 +
636 + from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
637 +
638 + # Swap halves to arrange as [w3; w1] (kernel expectation)
639 + w1_w, w3_w = torch.chunk(w13_q, 2, dim=1)
640 + w13_q_swapped = torch.cat([w3_w, w1_w], dim=1)
641 +
642 + b1, b3 = torch.chunk(bias13.to(torch.float32), 2, dim=-1)
643 + w13_b = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)
644 +
645 + w1_s, w3_s = torch.chunk(w13_scale, 2, dim=1)
646 + w13_s = torch.cat([w3_s, w1_s], dim=1)
647 + w13_s_inter = _interleave_scales_lastdim_by4(w13_s)
648 + w2_s_inter = _interleave_scales_lastdim_by4(w2_scale)
649 +
650 + routing_weights = torch.nn.functional.softmax(router_logits,
651 + dim=1,
652 + dtype=torch.float32)
653 + token_final_scales, token_selected_experts = torch.topk(routing_weights,
654 + topk,
655 + dim=-1)
656 + token_final_scales = (token_final_scales /
657 + token_final_scales.sum(dim=-1, keepdim=True))
658 + token_selected_experts = token_selected_experts.to(torch.int).contiguous()
659 +
660 + out = torch.empty_like(hidden_states, dtype=torch.bfloat16)
661 + if alpha is not None:
662 + alpha = torch.full((num_experts, ), alpha, device=hidden_states.device)
663 + if beta is not None:
664 + beta = torch.full((num_experts, ), beta, device=hidden_states.device)
665 + if limit is not None:
666 + limit = torch.full((num_experts, ), limit, device=hidden_states.device)
667 +
668 + _ = flashinfer_cutlass_fused_moe(
669 + input=hidden_states,
670 + token_selected_experts=token_selected_experts,
671 + token_final_scales=token_final_scales,
672 + fc1_expert_weights=w13_q_swapped,
673 + fc2_expert_weights=w2_q,
674 + output_dtype=torch.bfloat16,
675 + output=out,
676 + quant_scales=[w13_s_inter.to(torch.uint8),
677 + w2_s_inter.to(torch.uint8)],
678 + fc1_expert_biases=w13_b,
679 + fc2_expert_biases=bias2.to(torch.bfloat16),
680 + swiglu_alpha=alpha,
681 + swiglu_beta=beta,
682 + swiglu_limit=limit,
683 + tp_size=1,
684 + tp_rank=0,
685 + ep_size=1,
686 + ep_rank=0,
687 + use_w4_group_scaling=True,
688 + )
689 +
690 + # Allow some mismatch due to MXFP4 quantization
691 + check_accuracy(ref, out, atol=0, rtol=0.3, percent=0.8)
692 +
693 +
694 +@pytest.mark.parametrize("topk", [1, 4])
695 +@pytest.mark.parametrize("num_experts", [32])
696 +@pytest.mark.parametrize("num_tokens", [1, 128])
697 +@pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)])
698 +@pytest.mark.parametrize("alpha,beta,limit", [(1.0, 1.0, None),
699 + (1.702, 1.0, 7.0)])
700 +@pytest.mark.skipif(
701 + not (current_platform.is_cuda()
702 + and current_platform.is_device_capability(100) and has_flashinfer()),
703 + reason="NVIDIA GPU sm100 and flashinfer are required for this test",
704 +)
705 +def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe(
706 + topk: int,
707 + num_experts: int,
708 + num_tokens: int,
709 + intermediate_size: int,
710 + hidden_size: int,
711 + alpha: Optional[float],
712 + beta: Optional[float],
713 + limit: Optional[float],
714 +):
715 + torch.manual_seed(42)
716 + device = "cuda:0"
717 +
718 + # Inputs
719 + hidden_states = torch.randn(num_tokens,
720 + hidden_size,
721 + device=device,
722 + dtype=torch.bfloat16)
723 + # Float weights in w13 format [w1; w3]
724 + w13 = (torch.randn(num_experts,
725 + 2 * intermediate_size,
726 + hidden_size,
727 + device=device,
728 + dtype=torch.bfloat16) / 10)
729 + w2 = (torch.randn(num_experts,
730 + hidden_size,
731 + intermediate_size,
732 + device=device,
733 + dtype=torch.bfloat16) / 10)
734 + # Bias contiguous [b1; b3]
735 + bias13 = (torch.randn(num_experts,
736 + 2 * intermediate_size,
737 + device=device,
738 + dtype=torch.bfloat16) * 10)
739 + bias2 = (torch.randn(
740 + num_experts, hidden_size, device=device, dtype=torch.bfloat16) * 10)
741 + router_logits = torch.rand(num_tokens,
742 + num_experts,
743 + dtype=torch.float32,
744 + device=device)
745 +
746 + # Quantize weights to MXFP4 per expert (SM100 path)
747 + from flashinfer import mxfp4_quantize
748 +
749 + def quant_mxfp4_batches(a: torch.Tensor, e: int):
750 + qs, sfs = [], []
751 + for i in range(e):
752 + q, sf = mxfp4_quantize(a[i].cuda())
753 + qs.append(q)
754 + sfs.append(sf)
755 + return torch.stack(qs), torch.stack(sfs)
756 +
757 + def dequant_mxfp4_batches(mat_fp4: torch.Tensor,
758 + scale_tensor: torch.Tensor):
759 + num_batches = mat_fp4.size(0)
760 + scale_tensor = scale_tensor.view(num_batches, -1)
761 + from flashinfer import mxfp4_dequantize
762 + return torch.stack([
763 + mxfp4_dequantize(mat_fp4[b, :, :], scale_tensor[b, :])
764 + for b in range(num_batches)
765 + ])
766 +
767 + w13_q, w13_scale = quant_mxfp4_batches(w13, num_experts)
768 + w2_q, w2_scale = quant_mxfp4_batches(w2, num_experts)
769 +
770 + # Reference result using dequantized tensors and reference_moe
771 + w13_ref = dequant_mxfp4_batches(
772 + w13_q.view(torch.uint8),
773 + w13_scale.view(torch.uint8).reshape(-1)).to(torch.float32).reshape(
774 + num_experts, 2 * intermediate_size, hidden_size)
775 + w2_ref = dequant_mxfp4_batches(
776 + w2_q.view(torch.uint8),
777 + w2_scale.view(torch.uint8).reshape(-1)).to(torch.float32).reshape(
778 + num_experts, hidden_size, intermediate_size)
779 +
780 + # Quantize activations for SM100 path and dequantize for reference
781 + hidden_states_q, hidden_states_sf = mxfp8_quantize(hidden_states, True, 32)
782 + # Reference uses BF16 input but quantizes intermediate activation to MXFP8
783 + ref = reference_moe(router_logits.to(torch.float32), topk, num_experts,
784 + hidden_states.to(torch.float32), w13_ref,
785 + bias13.to(torch.float32), w2_ref,
786 + bias2.to(torch.float32), alpha, beta, limit, 'mxfp8')
787 +
788 + # Prepare inputs for FlashInfer CUTLASS fused MoE
789 + from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
790 +
791 + # Swap halves to arrange as [w3; w1] (kernel expectation)
792 + w1_w, w3_w = torch.chunk(w13_q, 2, dim=1)
793 + w13_q_swapped = torch.cat([w3_w, w1_w], dim=1)
794 +
795 + # Swap scales halves to match swapped weights
796 + s1, s3 = torch.chunk(w13_scale, 2, dim=1)
797 + w13_scale_swapped = torch.cat([s3, s1], dim=1)
798 +
799 + b1, b3 = torch.chunk(bias13.to(torch.float32), 2, dim=-1)
800 + w13_b = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)
801 +
802 + # Build routing for kernel
803 + routing_weights = torch.nn.functional.softmax(router_logits,
804 + dim=1,
805 + dtype=torch.float32)
806 + token_final_scales, token_selected_experts = torch.topk(routing_weights,
807 + topk,
808 + dim=-1)
809 + token_final_scales = (token_final_scales /
810 + token_final_scales.sum(dim=-1, keepdim=True))
811 + token_selected_experts = token_selected_experts.to(torch.int).contiguous()
812 +
813 + out = torch.empty_like(hidden_states, dtype=torch.bfloat16)
814 + if alpha is not None:
815 + alpha_t = torch.full((num_experts, ),
816 + alpha,
817 + device=hidden_states.device)
818 + else:
819 + alpha_t = None
820 + if beta is not None:
821 + beta_t = torch.full((num_experts, ), beta, device=hidden_states.device)
822 + else:
823 + beta_t = None
824 + if limit is not None:
825 + limit_t = torch.full((num_experts, ),
826 + limit,
827 + device=hidden_states.device)
828 + else:
829 + limit_t = None
830 +
831 + # Quant scales for SM100 MXFP8+MXFP4 path
832 + fake_input_scale = torch.ones(num_experts, device=device)
833 + quant_scales = [
834 + w13_scale_swapped.view(torch.int32),
835 + fake_input_scale,
836 + w2_scale.view(torch.int32),
837 + fake_input_scale,
838 + ]
839 +
840 + _ = flashinfer_cutlass_fused_moe(
841 + input=hidden_states_q,
842 + token_selected_experts=token_selected_experts,
843 + token_final_scales=token_final_scales,
844 + fc1_expert_weights=w13_q_swapped.contiguous().view(torch.long),
845 + fc2_expert_weights=w2_q.contiguous().view(torch.long),
846 + output_dtype=torch.bfloat16,
847 + output=out,
848 + quant_scales=quant_scales,
849 + fc1_expert_biases=w13_b,
850 + fc2_expert_biases=bias2.to(torch.bfloat16),
851 + swiglu_alpha=alpha_t,
852 + swiglu_beta=beta_t,
853 + swiglu_limit=limit_t,
854 + tp_size=1,
855 + tp_rank=0,
856 + ep_size=1,
857 + ep_rank=0,
858 + use_mxfp8_act_scaling=True,
859 + input_sf=hidden_states_sf,
860 + )
861 +
862 + # Allow some mismatch due to MXFP4 quantization
863 + check_accuracy(ref, out, atol=0, rtol=0.3, percent=0.8)
vllm/envs.py
+12 -1
@@ -166,7 +166,8 @@ if TYPE_CHECKING:
166 VLLM_HAS_FLASHINFER_CUBIN: bool = False
167 VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False
168 VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False
169 - VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
169 + VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False
170 + VLLM_ALLREDUCE_USE_SYMM_MEM: bool = False
171 VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None
172 VLLM_DISABLE_PAD_FOR_CUDAGRAPH: bool = False
173 VLLM_GPT_OSS_USE_CONTAINER_TOOL: bool = False
@@ -1004,6 +1005,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
1005 "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8":
1006 lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "0"))),
1007
1008 + # If set to 1, use the FlashInfer CUTLASS backend for
1009 + # MXFP8 (activation) x MXFP4 (weight) MoE.
1010 + # This is separate from the TRTLLMGEN path controlled by
1011 + # VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8.
1012 + "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS":
1013 + lambda: bool(int(
1014 + os.getenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS", "0")
1015 + )),
1016 +
1017 # If set to 1, use the FlashInfer
1018 # BF16 (activation) x MXFP4 (weight) MoE backend.
1019 "VLLM_USE_FLASHINFER_MOE_MXFP4_BF16":
@@ -1296,6 +1306,7 @@ def compute_hash() -> str:
1306 "VLLM_USE_FLASHINFER_MOE_FP8",
1307 "VLLM_USE_FLASHINFER_MOE_FP4",
1308 "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8",
1309 + "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS",
1310 "VLLM_USE_FLASHINFER_MOE_MXFP4_BF16",
1311 "VLLM_USE_CUDNN_PREFILL",
1312 "VLLM_USE_TRTLLM_ATTENTION",
vllm/model_executor/layers/fused_moe/layer.py
+10 -3
@@ -813,9 +813,16 @@ class FusedMoE(CustomOp):
813
814 # we are padding globally so EP buffer allocation works
815 if quant_config and quant_config.get_name() == "mxfp4":
816 - from vllm.model_executor.layers.quantization.mxfp4 import ( # noqa: E501
817 - should_use_flashinfer_mxfp4)
818 - if current_platform.is_rocm() or should_use_flashinfer_mxfp4():
816 + from vllm.model_executor.layers.quantization.mxfp4 import (
817 + Mxfp4Backend, get_mxfp4_backend)
818 + current_mxfp4_backend = get_mxfp4_backend()
819 + if (current_mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
820 + or current_mxfp4_backend
821 + == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS):
822 + hidden_size = round_up(hidden_size, 128)
823 + elif (current_platform.is_rocm() or current_mxfp4_backend
824 + == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM or
825 + current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
826 hidden_size = round_up(hidden_size, 256)
827
828 # For smuggling this layer into the fused moe custom op
vllm/model_executor/layers/quantization/mxfp4.py
+283 -58
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from enum import Enum
4 from typing import Callable, Optional, Union
5
6 import torch
@@ -33,33 +34,72 @@ from vllm.utils.flashinfer import has_flashinfer
34 logger = init_logger(__name__)
35
36
36 -def _should_use_flashinfer_mxfp4_bf16():
37 - """Determine if FlashInfer MXFP4 BF16 should be used."""
38 - # If explicitly set, respect the setting
39 - if envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16"):
40 - return envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16
41 -
42 - # Enable by default on SM100 if MXFP8 is not explicitly enabled
43 - if (current_platform.is_device_capability(100) and has_flashinfer()
44 - and not envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8")):
45 - logger.info_once(
46 - "Enabling FlashInfer MXFP4 BF16 backend by default for Blackwell. "
47 - "For faster performance, consider setting "
48 - "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8=1, "
49 - "though this may impact accuracy.")
50 - return True
51 -
52 - return False
53 -
54 -
55 -def _should_use_flashinfer_mxfp4_mxfp8():
56 - """Determine if FlashInfer MXFP4 MXFP8 should be used."""
57 - return envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
37 +# enum for mxfp4 backend
38 +class Mxfp4Backend(Enum):
39 + NONE = 0
40 +
41 + # FlashInfer Backend
42 + SM100_FI_MXFP4_MXFP8_TRTLLM = 1
43 + SM100_FI_MXFP4_MXFP8_CUTLASS = 2
44 + SM100_FI_MXFP4_BF16 = 3
45 + SM90_FI_MXFP4_BF16 = 4
46 +
47 + # Marlin Backend
48 + MARLIN = 5
49 +
50 + # Triton Backend
51 + TRITON = 6
52 +
53 +
54 +def get_mxfp4_backend():
55 + # Backend Selection
56 + if current_platform.is_cuda():
57 + if (current_platform.is_device_capability(90) and has_flashinfer()
58 + and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
59 + logger.info_once("Using FlashInfer MXFP4 BF16 backend for SM90")
60 + return Mxfp4Backend.SM90_FI_MXFP4_BF16
61 + elif (current_platform.is_device_capability(100) and has_flashinfer()
62 + and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS):
63 + logger.info_once(
64 + "Using FlashInfer MXFP4 MXFP8 CUTLASS backend for SM100")
65 + return Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
66 + elif (current_platform.is_device_capability(100) and has_flashinfer()
67 + and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8):
68 + logger.info_once(
69 + "Using FlashInfer MXFP4 MXFP8 TRTLLM backend for SM100, "
70 + "for high concurrency throughput workloads consider setting "
71 + "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS=1 for better "
72 + "performance")
73 + return Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
74 + elif current_platform.is_device_capability(100) and has_flashinfer():
75 + logger.info_once(
76 + "Using FlashInfer MXFP4 BF16 backend for SM100, "
77 + "For faster performance on SM100, consider setting "
78 + "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8=1, though this may impact "
79 + "accuracy.")
80 + return Mxfp4Backend.SM100_FI_MXFP4_BF16
81 + elif ((current_platform.is_device_capability(100)
82 + or current_platform.is_device_capability(90))
83 + and not has_flashinfer()):
84 + logger.warning_once(
85 + "MXFP4 MoE is enabled on Hopper/Blackwell but FlashInfer "
86 + "is not available. This may result in degraded performance. "
87 + "Please `pip install vllm[flashinfer]` for best results.")
88
89 + # If FlashInfer is not available, try either Marlin or Triton
90 + if current_platform.get_device_capability(
91 + )[0] < 9 or not has_triton_kernels() or not is_torch_equal_or_newer(
92 + "2.8.0"):
93 + logger.info_once("Using Marlin backend")
94 + return Mxfp4Backend.MARLIN
95 + else:
96 + logger.info_once("Using Triton backend")
97 + return Mxfp4Backend.TRITON
98 + elif current_platform.is_rocm() and has_triton_kernels():
99 + logger.info_once("Using Triton backend")
100 + return Mxfp4Backend.TRITON
101
60 -def should_use_flashinfer_mxfp4():
61 - return (_should_use_flashinfer_mxfp4_mxfp8()
62 - or _should_use_flashinfer_mxfp4_bf16())
102 + return Mxfp4Backend.NONE
103
104
105 class Mxfp4Config(QuantizationConfig):
@@ -113,31 +153,15 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
153 super().__init__(moe)
154 self.topk_indices_dtype = None
155 self.moe = moe
116 - self.use_marlin = self._should_use_marlin()
156 + self.mxfp4_backend = get_mxfp4_backend()
157 self.max_capture_size = get_current_vllm_config(
158 ).compilation_config.max_capture_size
159
120 - if current_platform.is_device_capability(100) and not has_flashinfer():
121 - logger.warning_once(
122 - "MXFP4 MoE is enabled on Blackwell but FlashInfer "
123 - "is not available. This may result in degraded performance. "
124 - "Please `pip install vllm[flashinfer]` for best results.")
160 + assert self.mxfp4_backend != Mxfp4Backend.NONE, (
161 + "No MXFP4 MoE backend (FlashInfer/Marlin/Triton) available."
162 + "Please check your environment and try again.")
163 self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
164
127 - def _should_use_marlin(self):
128 - if envs.VLLM_MXFP4_USE_MARLIN is not None:
129 - return envs.VLLM_MXFP4_USE_MARLIN
130 - if current_platform.is_cuda() and \
131 - not current_platform.is_device_capability(100):
132 - if not current_platform.has_device_capability(90):
133 - # marlin kernel has better performance on ampere
134 - return True
135 - if not has_triton_kernels():
136 - return True
137 - if not is_torch_equal_or_newer("2.8.0"):
138 - return True
139 - return False
140 -
165 def create_weights(self, layer: torch.nn.Module, num_experts: int,
166 hidden_size: int, intermediate_size_per_partition: int,
167 params_dtype: torch.dtype, **extra_weight_attrs):
@@ -157,7 +181,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
181
182 intermediate_size_per_partition_after_pad = \
183 intermediate_size_per_partition
160 - if self.use_marlin:
184 + if self.mxfp4_backend == Mxfp4Backend.MARLIN:
185 # The moe marlin kernel requires that for each linear
186 # n % 256 == 0 and k % 128 == 0.
187 # In gate_up_proj:
@@ -175,16 +199,20 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
199 layer.hidden_size = hidden_size
200 layer.intermediate_size_per_partition = \
201 intermediate_size_per_partition_after_pad
178 - elif should_use_flashinfer_mxfp4():
202 + elif (self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
203 + or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
204 # pad the intermediate size to be a multiple of 2 * mxfp4_block
205 # for to hold non-uniform sharded tensor as well as swizzling
206 # other padding to increase performance
207 intermediate_size_per_partition_after_pad = round_up(
208 intermediate_size_per_partition, 256)
209 hidden_size = round_up(hidden_size, 256)
185 - elif current_platform.is_rocm():
210 + elif current_platform.is_rocm() or (
211 + self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
212 + or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16):
213 intermediate_size_per_partition_after_pad = round_up(
214 intermediate_size_per_partition, 128)
215 + hidden_size = round_up(hidden_size, 128)
216 else:
217 intermediate_size_per_partition_after_pad = round_up(
218 intermediate_size_per_partition, 64)
@@ -264,9 +292,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
292 set_weight_attrs(w2_bias, extra_weight_attrs)
293
294 def process_weights_after_loading(self, layer):
267 - if self.use_marlin:
295 + if self.mxfp4_backend == Mxfp4Backend.MARLIN:
296 prepare_moe_fp4_layer_for_marlin(layer)
269 - elif should_use_flashinfer_mxfp4():
297 + elif (self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
298 + or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
299 from flashinfer.fp4_quantization import (
300 nvfp4_block_scale_interleave)
301 from flashinfer.fused_moe.core import (
@@ -429,7 +458,116 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
458 layer.w2_bias = Parameter(torch.stack(gemm2_bias_shuffled).reshape(
459 self.num_experts, -1),
460 requires_grad=False)
432 - else:
461 + elif (self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
462 + or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16):
463 + layer.gemm1_alpha = Parameter(torch.tensor(
464 + [1.702] * self.num_experts, dtype=torch.float32).cuda(),
465 + requires_grad=False)
466 + layer.gemm1_beta = Parameter(torch.tensor(
467 + [1.0] * self.num_experts, dtype=torch.float32).cuda(),
468 + requires_grad=False)
469 + layer.gemm1_clamp_limit = Parameter(torch.tensor(
470 + [7.0] * self.num_experts, dtype=torch.float32).cuda(),
471 + requires_grad=False)
472 +
473 + sf_block_size = 32 # mxfp4 block size
474 +
475 + # Common shape assertions
476 + assert (layer.w13_weight.dim() == 3
477 + and layer.w13_weight.shape[0] == self.num_experts
478 + and layer.w13_weight.shape[1] == self.intermediate_size * 2
479 + and layer.w13_weight.shape[2] == self.hidden_size // 2)
480 + assert (layer.w13_weight_scale.dim() == 3
481 + and layer.w13_weight_scale.shape[0] == self.num_experts
482 + and layer.w13_weight_scale.shape[1]
483 + == self.intermediate_size * 2
484 + and layer.w13_weight_scale.shape[2]
485 + == self.hidden_size // sf_block_size)
486 + assert (layer.w2_weight.dim() == 3
487 + and layer.w2_weight.shape[0] == self.num_experts
488 + and layer.w2_weight.shape[1] == self.hidden_size and
489 + layer.w2_weight.shape[2] == self.intermediate_size // 2)
490 + assert (layer.w2_weight_scale.dim() == 3
491 + and layer.w2_weight_scale.shape[1] == self.hidden_size
492 + and layer.w2_weight_scale.shape[2]
493 + == self.intermediate_size // sf_block_size)
494 + assert (layer.w13_bias.dim() == 2
495 + and layer.w13_bias.shape[0] == self.num_experts
496 + and layer.w13_bias.shape[1] == self.intermediate_size * 2)
497 + assert (layer.w2_bias.dim() == 2
498 + and layer.w2_bias.shape[0] == self.num_experts
499 + and layer.w2_bias.shape[1] == self.hidden_size)
500 +
501 + # De-interleave and swap for w13 weight, bias, and scales
502 + w13_w = layer.w13_weight.data
503 + gate_w, up_w = w13_w[:, ::2, :], w13_w[:, 1::2, :]
504 + deinterleaved_w13_w = torch.cat([gate_w, up_w], dim=1)
505 + w1_w, w3_w = torch.chunk(deinterleaved_w13_w, 2, dim=1)
506 + w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)
507 +
508 + w13_b = layer.w13_bias.data.to(torch.float32)
509 + gate_b, up_b = w13_b[:, ::2], w13_b[:, 1::2]
510 + deinterleaved_w13_b = torch.cat([gate_b, up_b], dim=1)
511 + b1, b3 = torch.chunk(deinterleaved_w13_b, 2, dim=-1)
512 + w13_bias_swapped = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)
513 +
514 + w13_s = layer.w13_weight_scale.data
515 + gate_s, up_s = w13_s[:, ::2, :], w13_s[:, 1::2, :]
516 + deinterleaved_w13_s = torch.cat([gate_s, up_s], dim=1)
517 + s1, s3 = torch.chunk(deinterleaved_w13_s, 2, dim=1)
518 + w13_scale_swapped = torch.cat([s3, s1], dim=1)
519 +
520 + if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS:
521 + from flashinfer import block_scale_interleave
522 +
523 + orig_shape = w13_scale_swapped.shape
524 + w13_scale_interleaved = block_scale_interleave(
525 + w13_scale_swapped.view(torch.uint8)).reshape(orig_shape)
526 +
527 + w2_s = layer.w2_weight_scale.data
528 + orig_shape = w2_s.shape
529 + w2_scale_interleaved = block_scale_interleave(
530 + w2_s.view(torch.uint8)).reshape(orig_shape)
531 +
532 + layer.w13_weight = Parameter(w13_weight_swapped,
533 + requires_grad=False)
534 + layer.w13_weight_scale = Parameter(w13_scale_interleaved,
535 + requires_grad=False)
536 + layer.w13_bias = Parameter(w13_bias_swapped,
537 + requires_grad=False)
538 + layer.w2_weight_scale = Parameter(w2_scale_interleaved,
539 + requires_grad=False)
540 + elif self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16:
541 +
542 + def _interleave_mxfp4_cutlass_sm90(w):
543 + w_shape = w.shape
544 + w_interleaved = w.reshape(w_shape[0], w_shape[1],
545 + (w_shape[2] // 4), 4)
546 + w_interleaved = w_interleaved.permute(0, 2, 1, 3)
547 + w_interleaved = w_interleaved.reshape(
548 + w_shape[0], w_shape[2] // 4, w_shape[1] * 4)
549 + return w_interleaved
550 +
551 + w31_scales = w13_scale_swapped.to(torch.uint8).view(
552 + torch.uint8)
553 + w31_scales_interleaved = _interleave_mxfp4_cutlass_sm90(
554 + w31_scales)
555 +
556 + w2_weight_scale = layer.w2_weight_scale.data
557 + w2_scales = w2_weight_scale.to(torch.uint8).view(torch.uint8)
558 + w2_scales_interleaved = _interleave_mxfp4_cutlass_sm90(
559 + w2_scales)
560 +
561 + layer.w13_weight = torch.nn.Parameter(torch.cat([w3_w, w1_w],
562 + dim=1),
563 + requires_grad=False)
564 + layer.w13_bias = torch.nn.Parameter(w13_bias_swapped,
565 + requires_grad=False)
566 + layer.w13_weight_scale = torch.nn.Parameter(
567 + w31_scales_interleaved, requires_grad=False)
568 + layer.w2_weight_scale = torch.nn.Parameter(
569 + w2_scales_interleaved, requires_grad=False)
570 + elif self.mxfp4_backend == Mxfp4Backend.TRITON:
571 from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
572
573 w13_bias = layer.w13_bias.to(torch.float32)
@@ -464,6 +602,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
602 layer.w13_weight = None
603 layer.w2_weight = None
604 torch.cuda.empty_cache()
605 + else:
606 + raise ValueError(f"Unsupported backend: {self.mxfp4_backend}")
607
608 def _get_tile_tokens_dim(self, x: torch.Tensor, top_k: int):
609 # Number of tokens in the input tensor.
@@ -500,7 +640,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
640 raise NotImplementedError(
641 "Mxfp4 does not support batched experts format for EP")
642 else:
503 - if should_use_flashinfer_mxfp4():
643 + if (self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
644 + or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
645 # B200 code-path
646 kwargs = {
647 "gemm1_alpha": layer.gemm1_alpha,
@@ -601,7 +742,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
742 if enable_eplb:
743 raise NotImplementedError("EPLB is not supported for mxfp4")
744
604 - if self.use_marlin:
745 + if self.mxfp4_backend == Mxfp4Backend.MARLIN:
746 topk_weights, topk_ids = FusedMoE.select_experts(
747 hidden_states=x,
748 router_logits=router_logits,
@@ -665,16 +806,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
806 logical_replica_count), (
807 "MXFP4 are not supported with this configuration.")
808
668 - if should_use_flashinfer_mxfp4():
669 - from flashinfer import mxfp8_quantize, trtllm_fp4_block_scale_moe
670 - if _should_use_flashinfer_mxfp4_bf16():
809 + if (self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
810 + or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
811 + from flashinfer import trtllm_fp4_block_scale_moe
812 + if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16:
813 assert x.dtype == torch.bfloat16
814 x_quant = x
815 x_scale = None
674 - else:
816 + elif self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM:
817 + from flashinfer import mxfp8_quantize
818 x_quant, x_scale = mxfp8_quantize(x, False) # to mxfp8
819 x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
820 *x.shape[:-1], -1)
821 +
822 trtllm_gen_output = trtllm_fp4_block_scale_moe(
823 router_logits.to(torch.bfloat16),
824 None, # routing_bias
@@ -706,7 +850,86 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
850 tune_max_num_tokens=self.max_capture_size,
851 )[0]
852 return trtllm_gen_output
709 - else:
853 + elif (self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
854 + or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16):
855 + from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
856 +
857 + topk_weights, topk_ids = FusedMoE.select_experts(
858 + hidden_states=x,
859 + router_logits=router_logits,
860 + use_grouped_topk=use_grouped_topk,
861 + top_k=top_k,
862 + renormalize=renormalize,
863 + topk_group=topk_group,
864 + num_expert_group=num_expert_group,
865 + custom_routing_function=custom_routing_function,
866 + scoring_func=scoring_func,
867 + e_score_correction_bias=e_score_correction_bias,
868 + )
869 +
870 + # Backend-specific preparation
871 + if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS:
872 +
873 + from flashinfer import mxfp8_quantize
874 +
875 + x_quant, x_scale = mxfp8_quantize(x, True, 32)
876 +
877 + fake_input_scale = torch.ones(self.num_experts,
878 + device=x.device)
879 + quant_scales = [
880 + layer.w13_weight_scale.contiguous().view(torch.int32),
881 + fake_input_scale,
882 + layer.w2_weight_scale.contiguous().view(torch.int32),
883 + fake_input_scale,
884 + ]
885 +
886 + fi_input = x_quant
887 + extra_kwargs = dict(
888 + use_mxfp8_act_scaling=True,
889 + input_sf=x_scale,
890 + fc1_expert_weights=layer.w13_weight.contiguous().view(
891 + torch.long),
892 + fc2_expert_weights=layer.w2_weight.contiguous().view(
893 + torch.long),
894 + )
895 + elif self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16:
896 + assert x.dtype == torch.bfloat16
897 +
898 + quant_scales = [
899 + layer.w13_weight_scale,
900 + layer.w2_weight_scale,
901 + ]
902 +
903 + fi_input = x
904 + extra_kwargs = dict(
905 + use_w4_group_scaling=True,
906 + fc1_expert_weights=layer.w13_weight,
907 + fc2_expert_weights=layer.w2_weight,
908 + )
909 +
910 + output = torch.empty_like(x, dtype=torch.bfloat16)
911 + _ = flashinfer_cutlass_fused_moe(
912 + input=fi_input,
913 + token_selected_experts=topk_ids.to(torch.int).contiguous(),
914 + token_final_scales=topk_weights,
915 + output_dtype=torch.bfloat16,
916 + output=output,
917 + quant_scales=quant_scales,
918 + fc1_expert_biases=layer.w13_bias,
919 + fc2_expert_biases=layer.w2_bias,
920 + swiglu_alpha=layer.gemm1_alpha,
921 + swiglu_beta=layer.gemm1_beta,
922 + swiglu_limit=layer.gemm1_clamp_limit,
923 + tp_size=self.moe.tp_size,
924 + tp_rank=self.moe.tp_rank,
925 + ep_size=self.moe.ep_size,
926 + ep_rank=self.moe.ep_rank,
927 + tune_max_num_tokens=self.max_capture_size,
928 + **extra_kwargs,
929 + )
930 +
931 + return output
932 + elif self.mxfp4_backend == Mxfp4Backend.TRITON:
933 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501
934 triton_kernel_moe_forward)
935 return triton_kernel_moe_forward(
@@ -724,3 +947,5 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
947 w2_precision=self.w2_precision_config,
948 apply_router_weight_on_input=apply_router_weight_on_input,
949 )
950 + else:
951 + raise ValueError(f"Unsupported backend: {self.mxfp4_backend}")
vllm/model_executor/warmup/kernel_warmup.py
+2 -2
@@ -33,8 +33,8 @@ def kernel_warmup(worker: "Worker"):
33 max_tokens = worker.scheduler_config.max_num_batched_tokens
34 deep_gemm_warmup(model, max_tokens)
35
36 - # FlashInfer kernel autotune for Blackwell (SM 10.0) GPUs
37 - if has_flashinfer() and current_platform.is_device_capability(100):
36 + # FlashInfer autotune for Hopper (SM 9.0) and Blackwell (SM 10.0) GPUs
37 + if has_flashinfer() and current_platform.has_device_capability(90):
38 flashinfer_autotune(worker.model_runner)
39
40 # FlashInfer attention warmup