[Kernel] GLM5 Router GEMM (#46385)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Jun 24, 2026 at 13:54 UTC
9d6fdc2901df490549f9f73d1f10588a6d16721d
4 files changed
+158
-32
csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu
+49
@@ -286,3 +286,52 @@ template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>(
286
287
template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>(
288
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
289
+
290
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
291
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 6144>(
292
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
293
+
294
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 6144>(
295
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
296
+
297
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 6144>(
298
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
299
+
300
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 6144>(
301
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
302
+
303
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 6144>(
304
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
305
+
306
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 6144>(
307
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
308
+
309
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 6144>(
310
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
311
+
312
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 6144>(
313
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
314
+
315
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 6144>(
316
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
317
+
318
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 6144>(
319
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
320
+
321
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 6144>(
322
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
323
+
324
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 6144>(
325
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
326
+
327
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 6144>(
328
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
329
+
330
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 6144>(
331
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
332
+
333
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 6144>(
334
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
335
+
336
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 6144>(
337
+ __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu
+51
-29
@@ -41,6 +41,7 @@ inline int getSMVersion() {
41
static constexpr int DEFAULT_NUM_EXPERTS = 256;
42
static constexpr int KIMI_K2_NUM_EXPERTS = 384;
43
static constexpr int DEFAULT_HIDDEN_DIM = 7168;
44
+static constexpr int GLM_5_HIDDEN_DIM = 6144;
45
46
template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
47
void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
@@ -121,14 +122,21 @@ void dsv3_router_gemm(
122
123
STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
124
"mat_a and mat_b must have the same hidden_dim");
124
- STD_TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
125
- "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
126
- ", but got hidden_dim=", hidden_dim);
125
+ STD_TORCH_CHECK(
126
+ hidden_dim == DEFAULT_HIDDEN_DIM || hidden_dim == GLM_5_HIDDEN_DIM,
127
+ "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
128
+ " or hidden_dim=", GLM_5_HIDDEN_DIM, ", but got hidden_dim=", hidden_dim);
129
STD_TORCH_CHECK(
130
num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
131
"Expected num_experts=", DEFAULT_NUM_EXPERTS,
132
" or num_experts=", KIMI_K2_NUM_EXPERTS,
133
", but got num_experts=", num_experts);
134
+ // KIMI_K2_NUM_EXPERTS is only instantiated for the default hidden_dim.
135
+ STD_TORCH_CHECK(
136
+ hidden_dim == DEFAULT_HIDDEN_DIM || num_experts == DEFAULT_NUM_EXPERTS,
137
+ "hidden_dim=", GLM_5_HIDDEN_DIM,
138
+ " only supports num_experts=", DEFAULT_NUM_EXPERTS,
139
+ ", but got num_experts=", num_experts);
140
STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
141
"currently num_tokens must be less than or equal to 16 for "
142
"router_gemm");
@@ -148,35 +156,49 @@ void dsv3_router_gemm(
156
157
const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index());
158
159
+ __nv_bfloat16 const* a_ptr =
160
+ reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
161
+ __nv_bfloat16 const* b_ptr =
162
+ reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
163
+
164
if (output.scalar_type() == torch::headeronly::ScalarType::Float) {
152
- if (num_experts == DEFAULT_NUM_EXPERTS) {
153
- LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
154
- unroll_float_output(
155
- num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
156
- reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
157
- reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
158
- } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
159
- LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
160
- unroll_float_output(
161
- num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
162
- reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
163
- reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
165
+ float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
166
+ if (hidden_dim == DEFAULT_HIDDEN_DIM) {
167
+ if (num_experts == DEFAULT_NUM_EXPERTS) {
168
+ LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
169
+ DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
170
+ out_ptr, a_ptr,
171
+ b_ptr, stream);
172
+ } else {
173
+ LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
174
+ DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
175
+ out_ptr, a_ptr,
176
+ b_ptr, stream);
177
+ }
178
+ } else { // GLM_5_HIDDEN_DIM
179
+ LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
180
+ GLM_5_HIDDEN_DIM>::unroll_float_output(num_tokens, out_ptr,
181
+ a_ptr, b_ptr, stream);
182
}
183
} else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
166
- if (num_experts == DEFAULT_NUM_EXPERTS) {
167
- LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
168
- unroll_bf16_output(
169
- num_tokens,
170
- reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
171
- reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
172
- reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
173
- } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
174
- LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
175
- unroll_bf16_output(
176
- num_tokens,
177
- reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
178
- reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
179
- reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
184
+ __nv_bfloat16* out_ptr =
185
+ reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
186
+ if (hidden_dim == DEFAULT_HIDDEN_DIM) {
187
+ if (num_experts == DEFAULT_NUM_EXPERTS) {
188
+ LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
189
+ DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
190
+ out_ptr, a_ptr,
191
+ b_ptr, stream);
192
+ } else {
193
+ LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
194
+ DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
195
+ out_ptr, a_ptr,
196
+ b_ptr, stream);
197
+ }
198
+ } else { // GLM_5_HIDDEN_DIM
199
+ LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
200
+ GLM_5_HIDDEN_DIM>::unroll_bf16_output(num_tokens, out_ptr,
201
+ a_ptr, b_ptr, stream);
202
}
203
}
204
}
csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu
+49
@@ -286,3 +286,52 @@ template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>(
286
287
template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>(
288
float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
289
+
290
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
291
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 6144>(
292
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
293
+
294
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 6144>(
295
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
296
+
297
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 6144>(
298
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
299
+
300
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 6144>(
301
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
302
+
303
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 6144>(
304
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
305
+
306
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 6144>(
307
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
308
+
309
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 6144>(
310
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
311
+
312
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 6144>(
313
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
314
+
315
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 6144>(
316
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
317
+
318
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 6144>(
319
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
320
+
321
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 6144>(
322
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
323
+
324
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 6144>(
325
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
326
+
327
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 6144>(
328
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
329
+
330
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 6144>(
331
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
332
+
333
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 6144>(
334
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
335
+
336
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 6144>(
337
+ float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+9
-3
@@ -14,7 +14,7 @@ from vllm.utils.torch_utils import direct_register_custom_op
14
class GateLinear(ReplicatedLinear):
15
"""MoE gate linear layer with multi-tier GEMM dispatch:
16
17
- 1. DSV3 specialized kernel (SM90+, fp32 out, M<=16, H=7168, E=256/384)
17
+ 1. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
18
2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out,
19
M<=32, H=3072, E=256)
20
3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
@@ -25,9 +25,14 @@ class GateLinear(ReplicatedLinear):
25
method which is only known later).
26
"""
27
28
- # Dimensions supported by the DSV3 specialized kernel
28
+ # Dimensions supported by the DSV3 specialized kernel.
29
+ # Valid (hidden_size, num_experts) combinations:
30
+ # (7168, 256) -> DeepSeek-V3, (7168, 384) -> Kimi-K2,
31
+ # (6144, 256) -> GLM-5
32
DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
30
- DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
33
+ DSV3_SUPPORTED_HIDDEN_SIZES = [7168, 6144]
34
+ # num_experts=384 is only instantiated for hidden_size=7168.
35
+ DSV3_UNSUPPORTED_SHAPES = {(6144, 384)}
36
37
# (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
38
# (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
@@ -71,6 +76,7 @@ class GateLinear(ReplicatedLinear):
76
self.allow_specialized_router_gemm
77
and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS
78
and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
79
+ and (input_size, output_size) not in self.DSV3_UNSUPPORTED_SHAPES
80
)
81
# See https://github.com/vllm-project/vllm/pull/44217
82
# for more details.