A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 01:00 UTC
[CI] Speed up test_fused_marlin_moe (#40178)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com>
Michael Goin committed
Apr 17, 2026 at 22:26 UTC
48a65ccb02a19f3b9755c13e8c3a4e8d5a368cc7
1 file changed
+47
-38
tests/kernels/moe/test_moe.py
+47
-38
@@ -143,12 +143,14 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [
143
{
144
"a_type": [scalar_types.bfloat16],
145
"b_type": scalar_types.float4_e2m1f,
146
+ "c_type": [scalar_types.bfloat16],
147
"group_blocks": [2],
148
},
149
# MXFP8
150
{
151
"a_type": [scalar_types.bfloat16],
152
"b_type": scalar_types.float8_e4m3fn,
153
+ "c_type": [scalar_types.bfloat16],
154
"group_blocks": [2],
155
},
156
# AWQ-INT4 with INT8 activation
@@ -674,31 +676,35 @@ def test_fused_moe_wn16(
676
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
677
678
679
+MARLIN_MOE_SCENARIOS = [
680
+ # (m, n, k, e, topk, ep_size, act_order, is_k_full)
681
+ # No act_order: is_k_full=True matches usual case (marlin_is_k_full).
682
+ # N>=256 required for Marlin kernel thread config for MXFP8.
683
+ # Single token, small matrices
684
+ (1, 128, 256, 5, 2, 1, False, True),
685
+ # Single token, large matrices
686
+ (1, 1024, 2048, 5, 2, 1, False, True),
687
+ # Unaligned m, small matrices
688
+ (133, 256, 256, 5, 2, 1, False, True),
689
+ # Unaligned m, large matrices
690
+ (133, 1024, 2048, 12, 3, 1, False, True),
691
+ # Aligned batch, small matrices
692
+ (128, 256, 256, 5, 2, 1, False, True),
693
+ # Aligned batch, large matrices
694
+ (128, 1024, 2048, 12, 3, 1, False, True),
695
+ # Expert parallelism
696
+ (64, 1024, 2048, 12, 3, 4, False, True),
697
+ # Act order with is_k_full=True (no tensor parallelism)
698
+ (1, 1024, 2048, 5, 2, 1, True, True),
699
+ # Act order with is_k_full=False (tensor parallelism)
700
+ (133, 256, 256, 5, 2, 1, True, False),
701
+]
702
+
703
+
704
def marlin_moe_generate_valid_test_cases():
705
import itertools
706
680
- m_list = [1, 123, 666]
681
- n_list = [128, 1024]
682
- k_list = [256, 2048]
683
- e_list = [5, 12]
684
- topk_list = [2, 3]
685
- ep_size_list = [1, 4]
686
- act_order_list = [True, False]
687
- is_k_full_list = [True, False]
688
-
689
- all_combinations = itertools.product(
690
- MOE_MARLIN_QUANT_TEST_CONFIGS,
691
- m_list,
692
- n_list,
693
- k_list,
694
- e_list,
695
- topk_list,
696
- ep_size_list,
697
- act_order_list,
698
- is_k_full_list,
699
- )
700
-
701
- def is_invalid(
707
+ def is_valid(
708
a_type,
709
b_type,
710
c_type,
@@ -715,39 +721,42 @@ def marlin_moe_generate_valid_test_cases():
721
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
722
if group_size > 0 and k % group_size != 0:
723
return False
718
-
724
if act_order and group_size in [-1, k, n]:
725
return False
726
if group_size in [k, n]:
727
return False
723
- if not act_order and is_k_full:
728
+ if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
729
return False
725
-
730
return a_type.size_bits < 16 or a_type is c_type
731
732
cases = []
729
- for case in all_combinations:
730
- quant_test_config, m, n, k, _, _, _, act_order, *_ = case
731
- if act_order and not quant_test_config.get("support_act_order", False):
732
- continue
733
-
733
+ for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS:
734
f16_types = [scalar_types.float16]
735
- inner_combinations = itertools.product(
736
- quant_test_config.get("a_type", f16_types),
737
- [quant_test_config["b_type"]],
738
- quant_test_config.get("c_type", f16_types),
739
- quant_test_config["group_blocks"],
735
+ inner_combinations = list(
736
+ itertools.product(
737
+ quant_test_config.get("a_type", f16_types),
738
+ [quant_test_config["b_type"]],
739
+ quant_test_config.get("c_type", f16_types),
740
+ quant_test_config["group_blocks"],
741
+ )
742
)
743
744
+ supports_act_order = quant_test_config.get("support_act_order", False)
745
+
746
for sub_case in inner_combinations:
747
if (
748
sub_case[0] == scalar_types.float8_e4m3fn
749
and current_platform.get_device_capability() not in [89, 120]
750
):
751
continue
748
- args = sub_case + (m, n, k) + case[4:]
749
- if is_invalid(*args):
750
- cases.append(args)
752
+
753
+ for scenario in MARLIN_MOE_SCENARIOS:
754
+ m, n, k, e, topk, ep_size, act_order, is_k_full = scenario
755
+ if act_order and not supports_act_order:
756
+ continue
757
+ args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full)
758
+ if is_valid(*args):
759
+ cases.append(args)
760
return cases
761
762