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 6, 12:00 UTC
[Perf] Add tuned triton moe config for Qwen3.5 H200, 9.9% E2E throughput improvement (#37340)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 18, 2026 at 14:18 UTC
0ef7f79054b9745e8f683b7881e0b02f1824c047
2 files changed
+180
-9
benchmarks/kernels/benchmark_moe.py
+33
-9
@@ -750,17 +750,20 @@ def get_weight_block_size_safety(config, default_value=None):
750
751
752
def get_model_params(config):
753
- if config.architectures[0] == "DbrxForCausalLM":
753
+ architectures = getattr(config, "architectures", None) or [type(config).__name__]
754
+ architecture = architectures[0]
755
+
756
+ if architecture == "DbrxForCausalLM":
757
E = config.ffn_config.moe_num_experts
758
topk = config.ffn_config.moe_top_k
759
intermediate_size = config.ffn_config.ffn_hidden_size
760
hidden_size = config.hidden_size
758
- elif config.architectures[0] == "JambaForCausalLM":
761
+ elif architecture == "JambaForCausalLM":
762
E = config.num_experts
763
topk = config.num_experts_per_tok
764
intermediate_size = config.intermediate_size
765
hidden_size = config.hidden_size
763
- elif config.architectures[0] in (
766
+ elif architecture in (
767
"DeepseekV2ForCausalLM",
768
"DeepseekV3ForCausalLM",
769
"DeepseekV32ForCausalLM",
@@ -774,7 +777,7 @@ def get_model_params(config):
777
topk = config.num_experts_per_tok
778
intermediate_size = config.moe_intermediate_size
779
hidden_size = config.hidden_size
777
- elif config.architectures[0] in (
780
+ elif architecture in (
781
"Qwen2MoeForCausalLM",
782
"Qwen3MoeForCausalLM",
783
"Qwen3NextForCausalLM",
@@ -783,23 +786,27 @@ def get_model_params(config):
786
topk = config.num_experts_per_tok
787
intermediate_size = config.moe_intermediate_size
788
hidden_size = config.hidden_size
786
- elif config.architectures[0] == "Qwen3VLMoeForConditionalGeneration":
789
+ elif architecture in (
790
+ "Qwen3VLMoeForConditionalGeneration",
791
+ "Qwen3_5MoeForConditionalGeneration",
792
+ "Qwen3_5MoeTextConfig",
793
+ ):
794
text_config = config.get_text_config()
795
E = text_config.num_experts
796
topk = text_config.num_experts_per_tok
797
intermediate_size = text_config.moe_intermediate_size
798
hidden_size = text_config.hidden_size
792
- elif config.architectures[0] == "HunYuanMoEV1ForCausalLM":
799
+ elif architecture == "HunYuanMoEV1ForCausalLM":
800
E = config.num_experts
801
topk = config.moe_topk[0]
802
intermediate_size = config.moe_intermediate_size[0]
803
hidden_size = config.hidden_size
797
- elif config.architectures[0] == "Qwen3OmniMoeForConditionalGeneration":
804
+ elif architecture == "Qwen3OmniMoeForConditionalGeneration":
805
E = config.thinker_config.text_config.num_experts
806
topk = config.thinker_config.text_config.num_experts_per_tok
807
intermediate_size = config.thinker_config.text_config.moe_intermediate_size
808
hidden_size = config.thinker_config.text_config.hidden_size
802
- elif config.architectures[0] == "PixtralForConditionalGeneration":
809
+ elif architecture == "PixtralForConditionalGeneration":
810
# Pixtral can contain different LLM architectures,
811
# recurse to get their parameters
812
return get_model_params(config.get_text_config())
@@ -814,6 +821,23 @@ def get_model_params(config):
821
return E, topk, intermediate_size, hidden_size
822
823
824
+def resolve_dtype(config) -> torch.dtype:
825
+ if current_platform.is_rocm():
826
+ return torch.float16
827
+
828
+ dtype = getattr(config, "dtype", None)
829
+ if dtype is not None:
830
+ return dtype
831
+
832
+ if hasattr(config, "get_text_config"):
833
+ text_config = config.get_text_config()
834
+ dtype = getattr(text_config, "dtype", None)
835
+ if dtype is not None:
836
+ return dtype
837
+
838
+ return torch.bfloat16
839
+
840
+
841
def get_quantization_group_size(config) -> int | None:
842
"""Extract the quantization group size from the HF model config.
843
@@ -861,7 +885,7 @@ def main(args: argparse.Namespace):
885
else:
886
ensure_divisibility(intermediate_size, args.tp_size, "intermediate_size")
887
shard_intermediate_size = 2 * intermediate_size // args.tp_size
864
- dtype = torch.float16 if current_platform.is_rocm() else config.dtype
888
+ dtype = resolve_dtype(config)
889
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
890
use_int8_w8a16 = args.dtype == "int8_w8a16"
891
use_int4_w4a16 = args.dtype == "int4_w4a16"
vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
new
+147
@@ -0,0 +1,147 @@
1
+{
2
+ "triton_version": "3.6.0",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 128,
6
+ "BLOCK_SIZE_K": 128,
7
+ "GROUP_SIZE_M": 1,
8
+ "num_warps": 4,
9
+ "num_stages": 3
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 128,
14
+ "BLOCK_SIZE_K": 128,
15
+ "GROUP_SIZE_M": 16,
16
+ "num_warps": 8,
17
+ "num_stages": 4
18
+ },
19
+ "4": {
20
+ "BLOCK_SIZE_M": 16,
21
+ "BLOCK_SIZE_N": 128,
22
+ "BLOCK_SIZE_K": 256,
23
+ "GROUP_SIZE_M": 16,
24
+ "num_warps": 4,
25
+ "num_stages": 3
26
+ },
27
+ "8": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 128,
30
+ "BLOCK_SIZE_K": 128,
31
+ "GROUP_SIZE_M": 32,
32
+ "num_warps": 4,
33
+ "num_stages": 3
34
+ },
35
+ "16": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 128,
38
+ "BLOCK_SIZE_K": 128,
39
+ "GROUP_SIZE_M": 1,
40
+ "num_warps": 4,
41
+ "num_stages": 3
42
+ },
43
+ "24": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 128,
46
+ "BLOCK_SIZE_K": 256,
47
+ "GROUP_SIZE_M": 16,
48
+ "num_warps": 4,
49
+ "num_stages": 4
50
+ },
51
+ "32": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 128,
54
+ "BLOCK_SIZE_K": 256,
55
+ "GROUP_SIZE_M": 1,
56
+ "num_warps": 4,
57
+ "num_stages": 3
58
+ },
59
+ "48": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 128,
62
+ "BLOCK_SIZE_K": 256,
63
+ "GROUP_SIZE_M": 1,
64
+ "num_warps": 4,
65
+ "num_stages": 3
66
+ },
67
+ "64": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 128,
70
+ "BLOCK_SIZE_K": 256,
71
+ "GROUP_SIZE_M": 1,
72
+ "num_warps": 4,
73
+ "num_stages": 3
74
+ },
75
+ "96": {
76
+ "BLOCK_SIZE_M": 16,
77
+ "BLOCK_SIZE_N": 128,
78
+ "BLOCK_SIZE_K": 256,
79
+ "GROUP_SIZE_M": 32,
80
+ "num_warps": 4,
81
+ "num_stages": 3
82
+ },
83
+ "128": {
84
+ "BLOCK_SIZE_M": 16,
85
+ "BLOCK_SIZE_N": 128,
86
+ "BLOCK_SIZE_K": 256,
87
+ "GROUP_SIZE_M": 16,
88
+ "num_warps": 4,
89
+ "num_stages": 3
90
+ },
91
+ "256": {
92
+ "BLOCK_SIZE_M": 16,
93
+ "BLOCK_SIZE_N": 128,
94
+ "BLOCK_SIZE_K": 128,
95
+ "GROUP_SIZE_M": 1,
96
+ "num_warps": 4,
97
+ "num_stages": 3
98
+ },
99
+ "512": {
100
+ "BLOCK_SIZE_M": 32,
101
+ "BLOCK_SIZE_N": 128,
102
+ "BLOCK_SIZE_K": 128,
103
+ "GROUP_SIZE_M": 32,
104
+ "num_warps": 4,
105
+ "num_stages": 3
106
+ },
107
+ "1024": {
108
+ "BLOCK_SIZE_M": 64,
109
+ "BLOCK_SIZE_N": 128,
110
+ "BLOCK_SIZE_K": 128,
111
+ "GROUP_SIZE_M": 32,
112
+ "num_warps": 4,
113
+ "num_stages": 4
114
+ },
115
+ "1536": {
116
+ "BLOCK_SIZE_M": 64,
117
+ "BLOCK_SIZE_N": 128,
118
+ "BLOCK_SIZE_K": 128,
119
+ "GROUP_SIZE_M": 32,
120
+ "num_warps": 4,
121
+ "num_stages": 4
122
+ },
123
+ "2048": {
124
+ "BLOCK_SIZE_M": 64,
125
+ "BLOCK_SIZE_N": 128,
126
+ "BLOCK_SIZE_K": 256,
127
+ "GROUP_SIZE_M": 16,
128
+ "num_warps": 4,
129
+ "num_stages": 4
130
+ },
131
+ "3072": {
132
+ "BLOCK_SIZE_M": 64,
133
+ "BLOCK_SIZE_N": 128,
134
+ "BLOCK_SIZE_K": 128,
135
+ "GROUP_SIZE_M": 16,
136
+ "num_warps": 4,
137
+ "num_stages": 4
138
+ },
139
+ "4096": {
140
+ "BLOCK_SIZE_M": 64,
141
+ "BLOCK_SIZE_N": 128,
142
+ "BLOCK_SIZE_K": 256,
143
+ "GROUP_SIZE_M": 16,
144
+ "num_warps": 4,
145
+ "num_stages": 4
146
+ }
147
+}