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 13, 14:30 UTC
[Hardware][AMD][Model] Add Triton MoE tuning support and optimized configs for Qwen3 omni for MI308X (#28373)
Signed-off-by: Xiake Sun <xiake.sun@amd.com> Signed-off-by: Xiake Sun <xisun@amd.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Xiake Sun committed
Nov 10, 2025 at 12:53 UTC
03fa4d3fb3ddfc9ced960d527d075213f32ecf07
2 files changed
+219
-1
benchmarks/kernels/benchmark_moe.py
+6
-1
@@ -211,7 +211,7 @@ def get_rocm_tuning_space(use_fp16):
211
num_warps_range = [1, 2, 4, 8]
212
group_m_range = [1, 4, 8, 16, 32]
213
num_stage_range = [2]
214
- waves_per_eu_range = [0]
214
+ waves_per_eu_range = [0, 1, 2, 4]
215
matrix_instr_nonkdim_range = [16, 32] if use_fp16 else []
216
kpack_range = [1, 2] if use_fp16 else []
217
@@ -616,6 +616,11 @@ def main(args: argparse.Namespace):
616
topk = config.moe_topk[0]
617
intermediate_size = config.moe_intermediate_size[0]
618
hidden_size = config.hidden_size
619
+ elif config.architectures[0] in ["Qwen3OmniMoeForConditionalGeneration"]:
620
+ E = config.thinker_config.text_config.num_experts
621
+ topk = config.thinker_config.text_config.num_experts_per_tok
622
+ intermediate_size = config.thinker_config.text_config.moe_intermediate_size
623
+ hidden_size = config.thinker_config.text_config.hidden_size
624
else:
625
# Support for llama4
626
config = config.get_text_config()
vllm/model_executor/layers/fused_moe/configs/E=128,N=768,device_name=AMD_Instinct_MI308X.json
new
+213
@@ -0,0 +1,213 @@
1
+{
2
+ "triton_version": "3.4.0",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 32,
6
+ "BLOCK_SIZE_K": 64,
7
+ "GROUP_SIZE_M": 1,
8
+ "num_warps": 2,
9
+ "num_stages": 2,
10
+ "waves_per_eu": 2,
11
+ "matrix_instr_nonkdim": 16,
12
+ "kpack": 1
13
+ },
14
+ "2": {
15
+ "BLOCK_SIZE_M": 16,
16
+ "BLOCK_SIZE_N": 32,
17
+ "BLOCK_SIZE_K": 64,
18
+ "GROUP_SIZE_M": 1,
19
+ "num_warps": 2,
20
+ "num_stages": 2,
21
+ "waves_per_eu": 1,
22
+ "matrix_instr_nonkdim": 16,
23
+ "kpack": 2
24
+ },
25
+ "4": {
26
+ "BLOCK_SIZE_M": 16,
27
+ "BLOCK_SIZE_N": 64,
28
+ "BLOCK_SIZE_K": 128,
29
+ "GROUP_SIZE_M": 1,
30
+ "num_warps": 4,
31
+ "num_stages": 2,
32
+ "waves_per_eu": 4,
33
+ "matrix_instr_nonkdim": 16,
34
+ "kpack": 2
35
+ },
36
+ "8": {
37
+ "BLOCK_SIZE_M": 16,
38
+ "BLOCK_SIZE_N": 16,
39
+ "BLOCK_SIZE_K": 64,
40
+ "GROUP_SIZE_M": 1,
41
+ "num_warps": 1,
42
+ "num_stages": 2,
43
+ "waves_per_eu": 4,
44
+ "matrix_instr_nonkdim": 16,
45
+ "kpack": 2
46
+ },
47
+
48
+ "16": {
49
+ "BLOCK_SIZE_M": 16,
50
+ "BLOCK_SIZE_N": 64,
51
+ "BLOCK_SIZE_K": 128,
52
+ "GROUP_SIZE_M": 1,
53
+ "num_warps": 4,
54
+ "num_stages": 2,
55
+ "waves_per_eu": 2,
56
+ "matrix_instr_nonkdim": 16,
57
+ "kpack": 2
58
+ },
59
+ "24": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 32,
62
+ "BLOCK_SIZE_K": 128,
63
+ "GROUP_SIZE_M": 4,
64
+ "num_warps": 2,
65
+ "num_stages": 2,
66
+ "waves_per_eu": 2,
67
+ "matrix_instr_nonkdim": 16,
68
+ "kpack": 2
69
+ },
70
+ "32": {
71
+ "BLOCK_SIZE_M": 16,
72
+ "BLOCK_SIZE_N": 64,
73
+ "BLOCK_SIZE_K": 128,
74
+ "GROUP_SIZE_M": 1,
75
+ "num_warps": 4,
76
+ "num_stages": 2,
77
+ "waves_per_eu": 4,
78
+ "matrix_instr_nonkdim": 16,
79
+ "kpack": 2
80
+ },
81
+ "48": {
82
+ "BLOCK_SIZE_M": 16,
83
+ "BLOCK_SIZE_N": 64,
84
+ "BLOCK_SIZE_K": 128,
85
+ "GROUP_SIZE_M": 1,
86
+ "num_warps": 4,
87
+ "num_stages": 2,
88
+ "waves_per_eu": 4,
89
+ "matrix_instr_nonkdim": 16,
90
+ "kpack": 2
91
+ },
92
+ "64": {
93
+ "BLOCK_SIZE_M": 16,
94
+ "BLOCK_SIZE_N": 64,
95
+ "BLOCK_SIZE_K": 128,
96
+ "GROUP_SIZE_M": 1,
97
+ "num_warps": 4,
98
+ "num_stages": 2,
99
+ "waves_per_eu": 2,
100
+ "matrix_instr_nonkdim": 16,
101
+ "kpack": 2
102
+ },
103
+ "96": {
104
+ "BLOCK_SIZE_M": 16,
105
+ "BLOCK_SIZE_N": 64,
106
+ "BLOCK_SIZE_K": 128,
107
+ "GROUP_SIZE_M": 32,
108
+ "num_warps": 4,
109
+ "num_stages": 2,
110
+ "waves_per_eu": 1,
111
+ "matrix_instr_nonkdim": 16,
112
+ "kpack": 2
113
+ },
114
+ "128": {
115
+ "BLOCK_SIZE_M": 16,
116
+ "BLOCK_SIZE_N": 32,
117
+ "BLOCK_SIZE_K": 64,
118
+ "GROUP_SIZE_M": 4,
119
+ "num_warps": 2,
120
+ "num_stages": 2,
121
+ "waves_per_eu": 1,
122
+ "matrix_instr_nonkdim": 16,
123
+ "kpack": 1
124
+ },
125
+ "256": {
126
+ "BLOCK_SIZE_M": 32,
127
+ "BLOCK_SIZE_N": 32,
128
+ "BLOCK_SIZE_K": 64,
129
+ "GROUP_SIZE_M": 4,
130
+ "num_warps": 2,
131
+ "num_stages": 2,
132
+ "waves_per_eu": 4,
133
+ "matrix_instr_nonkdim": 16,
134
+ "kpack": 2
135
+ },
136
+ "512": {
137
+ "BLOCK_SIZE_M": 64,
138
+ "BLOCK_SIZE_N": 64,
139
+ "BLOCK_SIZE_K": 64,
140
+ "GROUP_SIZE_M": 8,
141
+ "num_warps": 4,
142
+ "num_stages": 2,
143
+ "waves_per_eu": 1,
144
+ "matrix_instr_nonkdim": 16,
145
+ "kpack": 2
146
+ },
147
+ "768": {
148
+ "BLOCK_SIZE_M": 64,
149
+ "BLOCK_SIZE_N": 64,
150
+ "BLOCK_SIZE_K": 128,
151
+ "GROUP_SIZE_M": 1,
152
+ "num_warps": 4,
153
+ "num_stages": 2,
154
+ "waves_per_eu": 2,
155
+ "matrix_instr_nonkdim": 16,
156
+ "kpack": 2
157
+ },
158
+ "1024": {
159
+ "BLOCK_SIZE_M": 64,
160
+ "BLOCK_SIZE_N": 64,
161
+ "BLOCK_SIZE_K": 64,
162
+ "GROUP_SIZE_M": 1,
163
+ "num_warps": 4,
164
+ "num_stages": 2,
165
+ "waves_per_eu": 4,
166
+ "matrix_instr_nonkdim": 16,
167
+ "kpack": 2
168
+ },
169
+ "1536": {
170
+ "BLOCK_SIZE_M": 128,
171
+ "BLOCK_SIZE_N": 128,
172
+ "BLOCK_SIZE_K": 64,
173
+ "GROUP_SIZE_M": 4,
174
+ "num_warps": 4,
175
+ "num_stages": 2,
176
+ "waves_per_eu": 0,
177
+ "matrix_instr_nonkdim": 16,
178
+ "kpack": 1
179
+ },
180
+ "2048": {
181
+ "BLOCK_SIZE_M": 128,
182
+ "BLOCK_SIZE_N": 128,
183
+ "BLOCK_SIZE_K": 64,
184
+ "GROUP_SIZE_M": 1,
185
+ "num_warps": 4,
186
+ "num_stages": 2,
187
+ "waves_per_eu": 0,
188
+ "matrix_instr_nonkdim": 16,
189
+ "kpack": 2
190
+ },
191
+ "3072": {
192
+ "BLOCK_SIZE_M": 128,
193
+ "BLOCK_SIZE_N": 128,
194
+ "BLOCK_SIZE_K": 64,
195
+ "GROUP_SIZE_M": 1,
196
+ "num_warps": 4,
197
+ "num_stages": 2,
198
+ "waves_per_eu": 0,
199
+ "matrix_instr_nonkdim": 16,
200
+ "kpack": 2
201
+ },
202
+ "4096": {
203
+ "BLOCK_SIZE_M": 128,
204
+ "BLOCK_SIZE_N": 128,
205
+ "BLOCK_SIZE_K": 64,
206
+ "GROUP_SIZE_M": 1,
207
+ "num_warps": 4,
208
+ "num_stages": 2,
209
+ "waves_per_eu": 0,
210
+ "matrix_instr_nonkdim": 16,
211
+ "kpack": 2
212
+ }
213
+}