@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5a468ff7c7
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 10, 19:00 UTC

[MoE Refactor] Split `invoke_fused_moe_kernel` (#31050)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

Yongye Zhu committed Jan 2, 2026 at 13:47 UTC 5a468ff7c7463a9ed4e6353481a6c0dfb2bfa100
1 file changed +266 -102
vllm/model_executor/layers/fused_moe/fused_moe.py
+266 -102
@@ -541,11 +541,70 @@ def fused_moe_kernel(
541 tl.store(c_ptrs, accumulator, mask=c_mask)
542
543
544 -def invoke_fused_moe_kernel(
544 +# NOTE(zyongye): we can remove all the wna16 kernel
545 +# once we drop off sm75 support
546 +def invoke_fused_moe_wna16_cuda_kernel(
547 + A: torch.Tensor,
548 + B: torch.Tensor,
549 + C: torch.Tensor,
550 + B_scale: torch.Tensor | None,
551 + B_zp: torch.Tensor | None,
552 + topk_weights: torch.Tensor | None,
553 + sorted_token_ids: torch.Tensor,
554 + expert_ids: torch.Tensor,
555 + num_tokens_post_padded: torch.Tensor,
556 + mul_routed_weight: bool,
557 + top_k: int,
558 + config: dict[str, Any],
559 + block_shape: list[int],
560 +):
561 + assert B_scale is not None and B_scale.ndim == 3
562 + assert B_zp is None or B_zp.ndim == 3
563 + assert block_shape is None or block_shape[0] == 0
564 +
565 + M = A.size(0)
566 + num_tokens = M * top_k
567 + bit = 4
568 +
569 + config = config.copy()
570 + config.update(
571 + get_moe_wna16_block_config(
572 + config=config,
573 + use_moe_wna16_cuda=True,
574 + num_valid_tokens=num_tokens,
575 + size_k=A.size(1),
576 + size_n=B.size(1),
577 + num_experts=B.size(1),
578 + group_size=block_shape[1],
579 + real_top_k=top_k,
580 + block_size_m=config["BLOCK_SIZE_M"],
581 + )
582 + )
583 +
584 + ops.moe_wna16_gemm(
585 + A,
586 + C,
587 + B,
588 + B_scale,
589 + B_zp,
590 + topk_weights if mul_routed_weight else None,
591 + sorted_token_ids,
592 + expert_ids,
593 + num_tokens_post_padded,
594 + top_k,
595 + config["BLOCK_SIZE_M"],
596 + config["BLOCK_SIZE_N"],
597 + config["BLOCK_SIZE_K"],
598 + bit,
599 + )
600 +
601 +
602 +# NOTE(zyongye): we can remove all the wna16 kernel
603 +# once we drop off sm75 support
604 +def invoke_fused_moe_wna16_triton_kernel(
605 A: torch.Tensor,
606 B: torch.Tensor,
607 C: torch.Tensor,
548 - A_scale: torch.Tensor | None,
608 B_scale: torch.Tensor | None,
609 B_zp: torch.Tensor | None,
610 topk_weights: torch.Tensor | None,
@@ -556,6 +615,96 @@ def invoke_fused_moe_kernel(
615 top_k: int,
616 config: dict[str, Any],
617 compute_type: tl.dtype,
618 + use_int8_w8a16: bool,
619 + use_int4_w4a16: bool,
620 + block_shape: list[int],
621 +):
622 + assert B_scale is not None and B_scale.ndim == 3
623 + assert B_zp is None or B_zp.ndim == 3
624 + assert block_shape is None or block_shape[0] == 0
625 +
626 + M = A.size(0)
627 + num_tokens = M * top_k
628 +
629 + EM = sorted_token_ids.size(0)
630 + if A.size(0) < config["BLOCK_SIZE_M"]:
631 + # optimize for small batch_size.
632 + # We assume that top_ids of each token is unique,
633 + # so num_valid_experts <= batch_size <= BLOCK_SIZE_M,
634 + # and we can skip some invalid blocks.
635 + EM = min(sorted_token_ids.size(0), A.size(0) * top_k * config["BLOCK_SIZE_M"])
636 + grid = lambda META: (
637 + triton.cdiv(EM, META["BLOCK_SIZE_M"])
638 + * triton.cdiv(B.size(1), META["BLOCK_SIZE_N"]),
639 + )
640 + config = config.copy()
641 + config.update(
642 + get_moe_wna16_block_config(
643 + config=config,
644 + use_moe_wna16_cuda=False,
645 + num_valid_tokens=num_tokens,
646 + size_k=A.size(1),
647 + size_n=B.size(1),
648 + num_experts=B.size(1),
649 + group_size=block_shape[1],
650 + real_top_k=top_k,
651 + block_size_m=config["BLOCK_SIZE_M"],
652 + )
653 + )
654 +
655 + fused_moe_kernel_gptq_awq[grid](
656 + A,
657 + B,
658 + C,
659 + B_scale,
660 + B_zp,
661 + topk_weights,
662 + sorted_token_ids,
663 + expert_ids,
664 + num_tokens_post_padded,
665 + B.size(1),
666 + A.size(1),
667 + EM,
668 + num_tokens,
669 + A.stride(0),
670 + A.stride(1),
671 + B.stride(0),
672 + B.stride(2),
673 + B.stride(1),
674 + C.stride(1),
675 + C.stride(2),
676 + B_scale.stride(0),
677 + B_scale.stride(2),
678 + B_scale.stride(1),
679 + B_zp.stride(0) if B_zp is not None else 0,
680 + B_zp.stride(2) if B_zp is not None else 0,
681 + B_zp.stride(1) if B_zp is not None else 0,
682 + block_k_diviable=A.size(1) % config["BLOCK_SIZE_K"] == 0,
683 + group_size=block_shape[1],
684 + MUL_ROUTED_WEIGHT=mul_routed_weight,
685 + top_k=top_k,
686 + compute_type=compute_type,
687 + has_zp=B_zp is not None,
688 + use_int4_w4a16=use_int4_w4a16,
689 + use_int8_w8a16=use_int8_w8a16,
690 + **config,
691 + )
692 +
693 +
694 +def invoke_fused_moe_triton_kernel(
695 + A: torch.Tensor,
696 + B: torch.Tensor,
697 + C: torch.Tensor,
698 + A_scale: torch.Tensor | None,
699 + B_scale: torch.Tensor | None,
700 + topk_weights: torch.Tensor | None,
701 + sorted_token_ids: torch.Tensor,
702 + expert_ids: torch.Tensor,
703 + num_tokens_post_padded: torch.Tensor,
704 + mul_routed_weight: bool,
705 + top_k: int,
706 + config: dict[str, Any],
707 + compute_type: tl.dtype,
708 use_fp8_w8a8: bool,
709 use_int8_w8a8: bool,
710 use_int8_w8a16: bool,
@@ -563,7 +712,7 @@ def invoke_fused_moe_kernel(
712 per_channel_quant: bool,
713 block_shape: list[int] | None = None,
714 B_bias: torch.Tensor | None = None,
566 -) -> None:
715 +):
716 assert topk_weights is not None or not mul_routed_weight
717 assert topk_weights is None or topk_weights.stride(1) == 1
718 assert sorted_token_ids.stride(0) == 1
@@ -576,7 +725,6 @@ def invoke_fused_moe_kernel(
725 assert block_shape is None or triton.cdiv(
726 B.size(-1), block_shape[1]
727 ) == B_scale.size(-1)
579 -
728 elif use_int8_w8a16 or use_int4_w4a16:
729 assert B_scale is not None
730 assert block_shape is None or block_shape[0] == 0
@@ -599,13 +747,90 @@ def invoke_fused_moe_kernel(
747 * triton.cdiv(B.size(1), META["BLOCK_SIZE_N"]),
748 )
749 HAS_BIAS = B_bias is not None
602 - if (
603 - (use_int8_w8a16 or use_int4_w4a16)
604 - and block_shape is not None
605 - and block_shape[1] > 0
750 +
751 + config = config.copy()
752 + config["SPLIT_K"] = 1
753 + BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K")
754 + if block_shape is not None:
755 + BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1]))
756 + fused_moe_kernel[grid](
757 + A,
758 + B,
759 + C,
760 + B_bias,
761 + A_scale,
762 + B_scale,
763 + topk_weights,
764 + sorted_token_ids,
765 + expert_ids,
766 + num_tokens_post_padded,
767 + B.size(1),
768 + B.size(2),
769 + EM,
770 + num_tokens,
771 + A.stride(0),
772 + A.stride(1),
773 + B.stride(0),
774 + B.stride(2),
775 + B.stride(1),
776 + C.stride(1),
777 + C.stride(2),
778 + A_scale.stride(0) if A_scale is not None and A_scale.ndim == 2 else 0,
779 + A_scale.stride(1) if A_scale is not None and A_scale.ndim == 2 else 0,
780 + B_scale.stride(0) if B_scale is not None and B_scale.ndim >= 2 else 0,
781 + B_scale.stride(2) if B_scale is not None and B_scale.ndim == 3 else 0,
782 + B_scale.stride(1) if B_scale is not None and B_scale.ndim >= 2 else 0,
783 + B_bias.stride(0) if B_bias is not None else 0,
784 + B_bias.stride(1) if B_bias is not None else 0,
785 + 0 if block_shape is None else block_shape[0],
786 + 0 if block_shape is None else block_shape[1],
787 + MUL_ROUTED_WEIGHT=mul_routed_weight,
788 + top_k=top_k,
789 + compute_type=compute_type,
790 + use_fp8_w8a8=use_fp8_w8a8,
791 + use_int8_w8a8=use_int8_w8a8,
792 + use_int8_w8a16=use_int8_w8a16,
793 + per_channel_quant=per_channel_quant,
794 + HAS_BIAS=HAS_BIAS,
795 + BLOCK_SIZE_K=BLOCK_SIZE_K,
796 + **config,
797 + )
798 +
799 +
800 +def dispatch_fused_moe_kernel(
801 + A: torch.Tensor,
802 + B: torch.Tensor,
803 + C: torch.Tensor,
804 + A_scale: torch.Tensor | None,
805 + B_scale: torch.Tensor | None,
806 + B_zp: torch.Tensor | None,
807 + topk_weights: torch.Tensor | None,
808 + sorted_token_ids: torch.Tensor,
809 + expert_ids: torch.Tensor,
810 + num_tokens_post_padded: torch.Tensor,
811 + mul_routed_weight: bool,
812 + top_k: int,
813 + config: dict[str, Any],
814 + compute_type: tl.dtype,
815 + use_fp8_w8a8: bool,
816 + use_int8_w8a8: bool,
817 + use_int8_w8a16: bool,
818 + use_int4_w4a16: bool,
819 + per_channel_quant: bool,
820 + block_shape: list[int] | None = None,
821 + B_bias: torch.Tensor | None = None,
822 +) -> None:
823 + assert topk_weights is not None or not mul_routed_weight
824 + assert topk_weights is None or topk_weights.stride(1) == 1
825 + assert sorted_token_ids.stride(0) == 1
826 +
827 + M = A.size(0)
828 + num_tokens = M * top_k
829 +
830 + if (use_int8_w8a16 or use_int4_w4a16) and (
831 + block_shape is not None and block_shape[1] > 0
832 ):
607 - assert B_scale is not None and B_scale.ndim == 3
608 - assert B_zp is None or B_zp.ndim == 3
833 + assert B_bias is None
834
835 use_moe_wna16_cuda = should_moe_wna16_use_cuda(
836 num_valid_tokens=num_tokens,
@@ -613,41 +838,25 @@ def invoke_fused_moe_kernel(
838 num_experts=B.size(0),
839 bit=4 if use_int4_w4a16 else 8,
840 )
616 - config = config.copy()
617 - config.update(
618 - get_moe_wna16_block_config(
619 - config=config,
620 - use_moe_wna16_cuda=use_moe_wna16_cuda,
621 - num_valid_tokens=num_tokens,
622 - size_k=A.size(1),
623 - size_n=B.size(1),
624 - num_experts=B.size(1),
625 - group_size=block_shape[1],
626 - real_top_k=top_k,
627 - block_size_m=config["BLOCK_SIZE_M"],
628 - )
629 - )
841
842 if use_moe_wna16_cuda:
632 - bit = 4 if use_int4_w4a16 else 8
633 - ops.moe_wna16_gemm(
843 + invoke_fused_moe_wna16_cuda_kernel(
844 A,
635 - C,
845 B,
846 + C,
847 B_scale,
848 B_zp,
639 - topk_weights if mul_routed_weight else None,
849 + topk_weights,
850 sorted_token_ids,
851 expert_ids,
852 num_tokens_post_padded,
853 + mul_routed_weight,
854 top_k,
644 - config["BLOCK_SIZE_M"],
645 - config["BLOCK_SIZE_N"],
646 - config["BLOCK_SIZE_K"],
647 - bit,
855 + config,
856 + block_shape,
857 )
858 return
650 - fused_moe_kernel_gptq_awq[grid](
859 + invoke_fused_moe_wna16_triton_kernel(
860 A,
861 B,
862 C,
@@ -657,80 +866,37 @@ def invoke_fused_moe_kernel(
866 sorted_token_ids,
867 expert_ids,
868 num_tokens_post_padded,
660 - B.size(1),
661 - A.size(1),
662 - EM,
663 - num_tokens,
664 - A.stride(0),
665 - A.stride(1),
666 - B.stride(0),
667 - B.stride(2),
668 - B.stride(1),
669 - C.stride(1),
670 - C.stride(2),
671 - B_scale.stride(0),
672 - B_scale.stride(2),
673 - B_scale.stride(1),
674 - B_zp.stride(0) if B_zp is not None else 0,
675 - B_zp.stride(2) if B_zp is not None else 0,
676 - B_zp.stride(1) if B_zp is not None else 0,
677 - block_k_diviable=A.size(1) % config["BLOCK_SIZE_K"] == 0,
678 - group_size=block_shape[1],
679 - MUL_ROUTED_WEIGHT=mul_routed_weight,
680 - top_k=top_k,
681 - compute_type=compute_type,
682 - has_zp=B_zp is not None,
683 - use_int4_w4a16=use_int4_w4a16,
684 - use_int8_w8a16=use_int8_w8a16,
685 - **config,
869 + mul_routed_weight,
870 + top_k,
871 + config,
872 + compute_type,
873 + use_int8_w8a16,
874 + use_int4_w4a16,
875 + block_shape,
876 )
877 +
878 else:
688 - config = config.copy()
689 - config["SPLIT_K"] = 1
690 - BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K")
691 - if block_shape is not None:
692 - BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1]))
693 - fused_moe_kernel[grid](
879 + invoke_fused_moe_triton_kernel(
880 A,
881 B,
882 C,
697 - B_bias,
883 A_scale,
884 B_scale,
885 topk_weights,
886 sorted_token_ids,
887 expert_ids,
888 num_tokens_post_padded,
704 - B.size(1),
705 - B.size(2),
706 - EM,
707 - num_tokens,
708 - A.stride(0),
709 - A.stride(1),
710 - B.stride(0),
711 - B.stride(2),
712 - B.stride(1),
713 - C.stride(1),
714 - C.stride(2),
715 - A_scale.stride(0) if A_scale is not None and A_scale.ndim == 2 else 0,
716 - A_scale.stride(1) if A_scale is not None and A_scale.ndim == 2 else 0,
717 - B_scale.stride(0) if B_scale is not None and B_scale.ndim >= 2 else 0,
718 - B_scale.stride(2) if B_scale is not None and B_scale.ndim == 3 else 0,
719 - B_scale.stride(1) if B_scale is not None and B_scale.ndim >= 2 else 0,
720 - B_bias.stride(0) if B_bias is not None else 0,
721 - B_bias.stride(1) if B_bias is not None else 0,
722 - 0 if block_shape is None else block_shape[0],
723 - 0 if block_shape is None else block_shape[1],
724 - MUL_ROUTED_WEIGHT=mul_routed_weight,
725 - top_k=top_k,
726 - compute_type=compute_type,
727 - use_fp8_w8a8=use_fp8_w8a8,
728 - use_int8_w8a8=use_int8_w8a8,
729 - use_int8_w8a16=use_int8_w8a16,
730 - per_channel_quant=per_channel_quant,
731 - HAS_BIAS=HAS_BIAS,
732 - BLOCK_SIZE_K=BLOCK_SIZE_K,
733 - **config,
889 + mul_routed_weight,
890 + top_k,
891 + config,
892 + compute_type,
893 + use_fp8_w8a8,
894 + use_int8_w8a8,
895 + use_int8_w8a16,
896 + use_int4_w4a16,
897 + per_channel_quant,
898 + block_shape,
899 + B_bias,
900 )
901
902
@@ -1997,7 +2163,7 @@ def fused_experts_impl(
2163 ignore_invalid_experts=True,
2164 )
2165
2000 - invoke_fused_moe_kernel(
2166 + dispatch_fused_moe_kernel(
2167 qcurr_hidden_states,
2168 w1,
2169 intermediate_cache1,
@@ -2056,7 +2222,7 @@ def fused_experts_impl(
2222 if expert_map is not None:
2223 intermediate_cache3.zero_()
2224
2059 - invoke_fused_moe_kernel(
2225 + dispatch_fused_moe_kernel(
2226 qintermediate_cache2,
2227 w2,
2228 intermediate_cache3,
@@ -2207,13 +2373,12 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2373 topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map
2374 )
2375
2210 - invoke_fused_moe_kernel(
2376 + invoke_fused_moe_triton_kernel(
2377 hidden_states,
2378 w1,
2379 intermediate_cache1,
2380 a1q_scale,
2381 self.w1_scale,
2216 - self.w1_zp,
2382 None, # topk_weights
2383 sorted_token_ids,
2384 expert_ids,
@@ -2245,13 +2410,12 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2410 self.block_shape,
2411 )
2412
2248 - invoke_fused_moe_kernel(
2413 + invoke_fused_moe_triton_kernel(
2414 qintermediate_cache2,
2415 w2,
2416 intermediate_cache3,
2417 a2q_scale,
2418 self.w2_scale,
2254 - self.w2_zp,
2419 topk_weights,
2420 sorted_token_ids,
2421 expert_ids,