[Kernel] Retire the DSv3 router GEMM CUDA kernel (#54040)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Aug 31, 2026 at 15:15 UTC
8e92248f79814bc12a195f5d2529cc245768017c
12 files changed
+18
-1029
CMakeLists.txt
-19
@@ -1417,25 +1417,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1417
message(STATUS "Not building Marlin MOE kernels as no compatible archs found"
1418
" in CUDA target architectures")
1419
endif()
1420
-
1421
- # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0.
1422
- # (fp32_router_gemm has been migrated to _C_stable_libtorch above.)
1423
- cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
1424
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS)
1425
- set(DSV3_ROUTER_GEMM_SRC
1426
- "csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu"
1427
- "csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu"
1428
- "csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu")
1429
- set_gencode_flags_for_srcs(
1430
- SRCS "${DSV3_ROUTER_GEMM_SRC}"
1431
- CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}")
1432
- list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
1433
-
1434
- message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}")
1435
- else()
1436
- message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found"
1437
- " (requires SM90+ and CUDA >= 12.0)")
1438
- endif()
1420
endif()
1421
1422
message(STATUS "Enabling MoE C_stable extension.")
benchmarks/kernels/benchmark_router_gemm.py
+1
-12
@@ -10,10 +10,6 @@ from vllm.transformers_utils.config import get_config
10
from vllm.triton_utils import triton
11
from vllm.utils.argparse_utils import FlexibleArgumentParser
12
13
-# Dimensions supported by the DSV3 specialized kernel
14
-DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
15
-DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
16
-
13
# Dimensions supported by the gpt-oss specialized kernel
14
GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
15
GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
@@ -71,11 +67,6 @@ def get_benchmark(model, max_batch_size, trust_remote_code):
67
is_hopper_or_blackwell = current_platform.is_device_capability(
68
90
69
) or current_platform.is_device_capability_family(100)
74
- allow_dsv3_router_gemm = (
75
- is_hopper_or_blackwell
76
- and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
77
- and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
78
- )
70
allow_gpt_oss_router_gemm = (
71
is_hopper_or_blackwell
72
and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
@@ -116,9 +107,7 @@ def get_benchmark(model, max_batch_size, trust_remote_code):
107
elif provider == "vllm":
108
109
def runner():
119
- if allow_dsv3_router_gemm:
120
- ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
121
- elif allow_fp32_router_gemm:
110
+ if allow_fp32_router_gemm:
111
ops.fp32_router_gemm(mat_a, mat_b)
112
elif allow_gpt_oss_router_gemm:
113
ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
csrc/libtorch_stable/fp32_router_gemm.cu
+2
-32
@@ -4,9 +4,6 @@
4
// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the
5
// supported (E, H) pairs listed at the bottom of this file.
6
// Supports bf16 or fp32 activation; weight is always fp32.
7
-// Adapted from dsv3_router_gemm_float_out.cu.
8
-// (E=256, H=6144) bf16 uses a B300-tuned wide-block geometry; see
9
-// invokeFp32RouterGemm.
7
8
#include <cuda_bf16.h>
9
#include <cuda_runtime.h>
@@ -215,34 +212,10 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
212
float const* mat_b, cudaStream_t stream) {
213
// Geometry tuned on B300 per supported shape, bf16 activation, under a
214
// production-fidelity harness (CUDA-graph replay, per-layer cold weights).
218
- // GLM-5.2 (E=256, H=6144):
219
- // M <= 4 : BS=768, EPB=1 (2.7us vs cast+cuBLAS 8.1us at M=1)
220
- // M in [5, 15]
221
- // or odd : BS=384, EPB=2 (crossover vs BS=768 measured in (4, 8))
222
- // M >= 16, even : BS=192, EPB=2, 2 token groups (M=16 4.79us vs 5.04,
223
- // M=24 5.71 vs 6.38, M=32 6.79 vs 7.72; M=12 loses at
224
- // 0.97x, so the boundary is 16).
215
// Only enabled on the Blackwell family where it was validated; Hopper and
216
// other shapes / fp32 activation keep the legacy geometry.
227
- if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 256 &&
217
+ if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 128 &&
218
kHiddenDim == 6144) {
229
- if (!isBlackwellFamily()) {
230
- launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
231
- output, mat_a, mat_b, stream);
232
- return;
233
- }
234
- if constexpr (kNumTokens <= 4) {
235
- launchFp32RouterGemm<InputT, 768, 1, kNumTokens, kNumExperts, kHiddenDim>(
236
- output, mat_a, mat_b, stream);
237
- } else if constexpr (kNumTokens >= 16 && kNumTokens % 2 == 0) {
238
- launchFp32RouterGemm<InputT, 192, 2, kNumTokens, kNumExperts, kHiddenDim,
239
- 2>(output, mat_a, mat_b, stream);
240
- } else {
241
- launchFp32RouterGemm<InputT, 384, 2, kNumTokens, kNumExperts, kHiddenDim>(
242
- output, mat_a, mat_b, stream);
243
- }
244
- } else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
245
- kNumExperts == 128 && kHiddenDim == 6144) {
219
// MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same
220
// accumulator register cliffs; B300 sweep:
221
// even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x)
@@ -296,8 +269,7 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
269
270
// ---------------------------------------------------------------------------
271
// Explicit instantiations: M=1..32, for both input types, for the supported
299
-// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5], (128, 6144) [MiniMax-M3]
300
-// and (256, 6144) [GLM-5.2].
272
+// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3].
273
// ---------------------------------------------------------------------------
274
275
#define INSTANTIATE(T, M, E, H) \
@@ -342,8 +314,6 @@ INSTANTIATE_ALL(float, 256, 3072)
314
INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
315
INSTANTIATE_ALL(float, 128, 6144)
316
INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
345
-INSTANTIATE_ALL(float, 256, 6144)
346
-INSTANTIATE_ALL(__nv_bfloat16, 256, 6144)
317
318
#undef INSTANTIATE_ALL
319
#undef INSTANTIATE
csrc/libtorch_stable/fp32_router_gemm_entry.cu
+3
-8
@@ -25,12 +25,10 @@ inline int getSMVersion() {
25
static constexpr int FP32_MAX_TOKENS = 32;
26
27
// Supported (hidden_dim, num_experts) pairs (must match the instantiations in
28
-// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3,
29
-// (6144, 256) for GLM-5.2.
28
+// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
29
static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
30
return (hidden_dim == 3072 && num_experts == 256) ||
32
- (hidden_dim == 6144 && num_experts == 128) ||
33
- (hidden_dim == 6144 && num_experts == 256);
31
+ (hidden_dim == 6144 && num_experts == 128);
32
}
33
34
// Forward declarations — 4 template params must match fp32_router_gemm.cu
@@ -79,9 +77,6 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
77
} else if (num_experts == 128 && hidden_dim == 6144) {
78
Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
79
num_tokens, output, mat_a, mat_b, stream);
82
- } else if (num_experts == 256 && hidden_dim == 6144) {
83
- Fp32LoopUnroller<InputT, 256, 6144, 1, FP32_MAX_TOKENS>::unroll(
84
- num_tokens, output, mat_a, mat_b, stream);
80
} else {
81
throw std::invalid_argument(
82
"fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
@@ -116,7 +111,7 @@ void fp32_router_gemm(
111
STD_TORCH_CHECK(
112
fp32_router_gemm_supported(hidden_dim, num_experts),
113
"fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
119
- "(3072, 256), (6144, 128) and (6144, 256)");
114
+ "(3072, 256) and (6144, 128)");
115
STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
116
"fp32_router_gemm: num_tokens must be in [0, 32]");
117
STD_TORCH_CHECK(
csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu
deleted
-337
@@ -1,337 +0,0 @@
1
-/*
2
- * Adapted from SGLang's sgl-kernel implementation, which was adapted from
3
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
4
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
5
- *
6
- * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved.
7
- *
8
- * Licensed under the Apache License, Version 2.0 (the "License");
9
- * you may not use this file except in compliance with the License.
10
- * You may obtain a copy of the License at
11
- *
12
- * http://www.apache.org/licenses/LICENSE-2.0
13
- *
14
- * Unless required by applicable law or agreed to in writing, software
15
- * distributed under the License is distributed on an "AS IS" BASIS,
16
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
- * See the License for the specific language governing permissions and
18
- * limitations under the License.
19
- */
20
-
21
-#include <cstdint>
22
-
23
-#include <cuda_bf16.h>
24
-#include <cuda_runtime.h>
25
-
26
-// Custom FMA implementation using PTX assembly instructions
27
-__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
28
- float2 const& c) {
29
- asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
30
- : "=l"(reinterpret_cast<uint64_t&>(d))
31
- : "l"(reinterpret_cast<uint64_t const&>(a)),
32
- "l"(reinterpret_cast<uint64_t const&>(b)),
33
- "l"(reinterpret_cast<uint64_t const&>(c)));
34
-}
35
-
36
-// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion
37
-template <int VPT>
38
-__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
39
- float* dst) {
40
- __nv_bfloat16* bf16_ptr =
41
- reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
42
-
43
-#pragma unroll
44
- for (int i = 0; i < VPT; i++) {
45
- dst[i] = __bfloat162float(bf16_ptr[i]);
46
- }
47
-}
48
-
49
-template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
50
- int kHiddenDim>
51
-__global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
52
- __nv_bfloat16* out, T const* mat_a, T const* mat_b) {
53
- // Each block handles one expert column
54
- int const n_idx = blockIdx.x;
55
- int const tid = threadIdx.x;
56
- constexpr int kWarpSize = 32;
57
- constexpr int kNumWarps = kBlockSize / kWarpSize;
58
- // Constants for this kernel
59
- constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
60
- constexpr int k_iterations =
61
- kHiddenDim / k_elems_per_k_iteration; // Total K iterations
62
-
63
- // Initialize accumulators for all M rows
64
- float acc[kNumTokens] = {};
65
-
66
- // Shared memory for warp-level reduction
67
- __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps
68
-
69
- // B matrix is in column-major order, so we can directly load a column for the
70
- // n_idx expert
71
- T const* b_col = mat_b + n_idx * kHiddenDim;
72
-
73
- // Pre-compute k_base values for each iteration to help compiler optimize
74
- int k_bases[k_iterations];
75
-#pragma unroll
76
- for (int ki = 0; ki < k_iterations; ki++) {
77
- k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
78
- }
79
-
80
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
81
- cudaGridDependencySynchronize();
82
-#endif
83
-
84
- // Process the GEMM in chunks
85
- for (int ki = 0; ki < k_iterations; ki++) {
86
- int const k_base = k_bases[ki];
87
-
88
- // Load B matrix values using vector load (8 bf16 values)
89
- uint4 b_vec = *reinterpret_cast<uint4 const*>(b_col + k_base);
90
-
91
- // Convert B values to float
92
- float b_float[VPT];
93
- bf16_uint4_to_float8<VPT>(b_vec, b_float);
94
-
95
-// Process each token
96
-#pragma unroll
97
- for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
98
- // Load both rows of A matrix using vector loads
99
- uint4 a_vec = *reinterpret_cast<uint4 const*>(
100
- mat_a + (m_idx * kHiddenDim) + k_base);
101
-
102
- // Convert A values to float
103
- float a_float[VPT];
104
- bf16_uint4_to_float8<VPT>(a_vec, a_float);
105
-
106
-// Process elements in this chunk
107
-#pragma unroll
108
- for (int k = 0; k < VPT; k++) {
109
- float a = a_float[k];
110
- float b = b_float[k];
111
- acc[m_idx] += a * b;
112
- }
113
- }
114
- }
115
-
116
- // Perform warp-level reduction
117
- int const warpSize = 32;
118
- int const warpId = tid / warpSize;
119
- int const laneId = tid % warpSize;
120
-
121
- // Register for warp-level reduction results
122
- float warp_result[kNumTokens];
123
-
124
-#pragma unroll
125
- for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
126
- warp_result[m_idx] = acc[m_idx];
127
- }
128
-
129
-// Perform warp-level reduction using optimized butterfly pattern
130
-#pragma unroll
131
- for (int m = 0; m < kNumTokens; m++) {
132
- float sum = warp_result[m];
133
-
134
- // Butterfly reduction pattern
135
- sum += __shfl_xor_sync(0xffffffff, sum, 16);
136
- sum += __shfl_xor_sync(0xffffffff, sum, 8);
137
- sum += __shfl_xor_sync(0xffffffff, sum, 4);
138
- sum += __shfl_xor_sync(0xffffffff, sum, 2);
139
- sum += __shfl_xor_sync(0xffffffff, sum, 1);
140
-
141
- // Only the first thread in each warp stores to shared memory
142
- if (laneId == 0) {
143
- sm_reduction[m][warpId] = sum;
144
- }
145
- }
146
-
147
- __syncthreads();
148
-
149
- // Final reduction across warps (only first thread)
150
- if (tid == 0) {
151
-#pragma unroll
152
- for (int m = 0; m < kNumTokens; m++) {
153
- float final_sum = 0.0f;
154
-
155
-// Sum across the kNumWarps
156
-#pragma unroll
157
- for (int w = 0; w < kNumWarps; w++) {
158
- final_sum += sm_reduction[m][w];
159
- }
160
-
161
- // Write final result
162
- out[m * kNumExperts + n_idx] = __float2bfloat16(final_sum);
163
- }
164
- }
165
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
166
- cudaTriggerProgrammaticLaunchCompletion();
167
-#endif
168
-}
169
-
170
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
171
-void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
172
- T const* mat_b, cudaStream_t stream) {
173
- constexpr int VPT = 16 / sizeof(T);
174
- constexpr int kBlockSize = 128;
175
- cudaLaunchConfig_t config;
176
- config.gridDim = kNumExperts;
177
- config.blockDim = kBlockSize;
178
- config.dynamicSmemBytes = 0;
179
- config.stream = stream;
180
- cudaLaunchAttribute attrs[1];
181
- attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
182
- attrs[0].val.programmaticStreamSerializationAllowed = 1;
183
- config.numAttrs = 1;
184
- config.attrs = attrs;
185
- cudaLaunchKernelEx(
186
- &config,
187
- router_gemm_kernel_bf16_output<T, kBlockSize, VPT, kNumTokens,
188
- kNumExperts, kHiddenDim>,
189
- output, mat_a, mat_b);
190
-}
191
-
192
-// Template instantiations for DEFAULT_NUM_EXPERTS experts
193
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 7168>(
194
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
195
-
196
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 7168>(
197
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
198
-
199
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 7168>(
200
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
201
-
202
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 7168>(
203
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
204
-
205
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 7168>(
206
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
207
-
208
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 7168>(
209
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
210
-
211
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 7168>(
212
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
213
-
214
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 7168>(
215
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
216
-
217
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 7168>(
218
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
219
-
220
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 7168>(
221
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
222
-
223
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 7168>(
224
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
225
-
226
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 7168>(
227
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
228
-
229
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 7168>(
230
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
231
-
232
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 7168>(
233
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
234
-
235
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 7168>(
236
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
237
-
238
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 7168>(
239
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
240
-
241
-// Template instantiations for KIMI_K2_NUM_EXPERTS experts
242
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 384, 7168>(
243
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
244
-
245
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 384, 7168>(
246
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
247
-
248
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 384, 7168>(
249
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
250
-
251
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 384, 7168>(
252
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
253
-
254
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 384, 7168>(
255
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
256
-
257
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 384, 7168>(
258
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
259
-
260
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 384, 7168>(
261
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
262
-
263
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 384, 7168>(
264
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
265
-
266
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 384, 7168>(
267
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
268
-
269
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 384, 7168>(
270
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
271
-
272
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 384, 7168>(
273
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
274
-
275
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 384, 7168>(
276
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
277
-
278
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 384, 7168>(
279
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
280
-
281
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 384, 7168>(
282
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
283
-
284
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>(
285
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
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
deleted
-210
@@ -1,210 +0,0 @@
1
-/*
2
- * Adapted from SGLang's sgl-kernel implementation, which was adapted from
3
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
4
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
5
- *
6
- * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved.
7
- *
8
- * Licensed under the Apache License, Version 2.0 (the "License");
9
- * you may not use this file except in compliance with the License.
10
- * You may obtain a copy of the License at
11
- *
12
- * http://www.apache.org/licenses/LICENSE-2.0
13
- *
14
- * Unless required by applicable law or agreed to in writing, software
15
- * distributed under the License is distributed on an "AS IS" BASIS,
16
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
- * See the License for the specific language governing permissions and
18
- * limitations under the License.
19
- */
20
-
21
-#include <torch/csrc/stable/library.h>
22
-#include <torch/csrc/stable/tensor.h>
23
-#include <torch/headeronly/core/ScalarType.h>
24
-
25
-#include "libtorch_stable/torch_utils.h"
26
-
27
-#include <cuda_bf16.h>
28
-#include <cuda_runtime.h>
29
-
30
-#include <stdexcept>
31
-
32
-namespace {
33
-
34
-inline int getSMVersion() {
35
- auto* props = get_device_prop();
36
- return props->major * 10 + props->minor;
37
-}
38
-
39
-} // namespace
40
-
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,
48
- cudaStream_t stream);
49
-
50
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
51
-void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
52
- T const* mat_b, cudaStream_t stream);
53
-
54
-template <int kBegin, int kEnd, int kNumExperts, int kHiddenDim>
55
-struct LoopUnroller {
56
- static void unroll_float_output(int num_tokens, float* output,
57
- __nv_bfloat16 const* input,
58
- __nv_bfloat16 const* weights,
59
- cudaStream_t stream) {
60
- if (num_tokens == kBegin) {
61
- invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts,
62
- kHiddenDim>(output, input, weights, stream);
63
- } else {
64
- LoopUnroller<kBegin + 1, kEnd, kNumExperts,
65
- kHiddenDim>::unroll_float_output(num_tokens, output, input,
66
- weights, stream);
67
- }
68
- }
69
-
70
- static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
71
- __nv_bfloat16 const* input,
72
- __nv_bfloat16 const* weights,
73
- cudaStream_t stream) {
74
- if (num_tokens == kBegin) {
75
- invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts,
76
- kHiddenDim>(output, input, weights, stream);
77
- } else {
78
- LoopUnroller<kBegin + 1, kEnd, kNumExperts,
79
- kHiddenDim>::unroll_bf16_output(num_tokens, output, input,
80
- weights, stream);
81
- }
82
- }
83
-};
84
-
85
-template <int kEnd, int kNumExperts, int kHiddenDim>
86
-struct LoopUnroller<kEnd, kEnd, kNumExperts, kHiddenDim> {
87
- static void unroll_float_output(int num_tokens, float* output,
88
- __nv_bfloat16 const* input,
89
- __nv_bfloat16 const* weights,
90
- cudaStream_t stream) {
91
- if (num_tokens == kEnd) {
92
- invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
93
- output, input, weights, stream);
94
- } else {
95
- throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
96
- }
97
- }
98
-
99
- static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
100
- __nv_bfloat16 const* input,
101
- __nv_bfloat16 const* weights,
102
- cudaStream_t stream) {
103
- if (num_tokens == kEnd) {
104
- invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
105
- output, input, weights, stream);
106
- } else {
107
- throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
108
- }
109
- }
110
-};
111
-
112
-void dsv3_router_gemm(
113
- torch::stable::Tensor& output, // [num_tokens, num_experts]
114
- torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim]
115
- torch::stable::Tensor const& mat_b // [num_experts, hidden_dim]
116
-) {
117
- STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
118
-
119
- const int num_tokens = mat_a.size(0);
120
- const int num_experts = mat_b.size(0);
121
- const int hidden_dim = mat_a.size(1);
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");
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");
143
- STD_TORCH_CHECK(
144
- mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
145
- "mat_a must be bf16");
146
- STD_TORCH_CHECK(
147
- mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16,
148
- "mat_b must be bf16");
149
- STD_TORCH_CHECK(
150
- output.scalar_type() == torch::headeronly::ScalarType::Float ||
151
- output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
152
- "output must be float32 or bf16");
153
-
154
- const torch::stable::accelerator::DeviceGuard device_guard(
155
- mat_a.get_device_index());
156
- const int sm = getSMVersion();
157
- STD_TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90");
158
-
159
- const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index());
160
-
161
- __nv_bfloat16 const* a_ptr =
162
- reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
163
- __nv_bfloat16 const* b_ptr =
164
- reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
165
-
166
- if (output.scalar_type() == torch::headeronly::ScalarType::Float) {
167
- float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
168
- if (hidden_dim == DEFAULT_HIDDEN_DIM) {
169
- if (num_experts == DEFAULT_NUM_EXPERTS) {
170
- LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
171
- DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
172
- out_ptr, a_ptr,
173
- b_ptr, stream);
174
- } else {
175
- LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
176
- DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
177
- out_ptr, a_ptr,
178
- b_ptr, stream);
179
- }
180
- } else { // GLM_5_HIDDEN_DIM
181
- LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
182
- GLM_5_HIDDEN_DIM>::unroll_float_output(num_tokens, out_ptr,
183
- a_ptr, b_ptr, stream);
184
- }
185
- } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
186
- __nv_bfloat16* out_ptr =
187
- reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
188
- if (hidden_dim == DEFAULT_HIDDEN_DIM) {
189
- if (num_experts == DEFAULT_NUM_EXPERTS) {
190
- LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
191
- DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
192
- out_ptr, a_ptr,
193
- b_ptr, stream);
194
- } else {
195
- LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
196
- DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
197
- out_ptr, a_ptr,
198
- b_ptr, stream);
199
- }
200
- } else { // GLM_5_HIDDEN_DIM
201
- LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
202
- GLM_5_HIDDEN_DIM>::unroll_bf16_output(num_tokens, out_ptr,
203
- a_ptr, b_ptr, stream);
204
- }
205
- }
206
-}
207
-
208
-STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
209
- m.impl("dsv3_router_gemm", TORCH_BOX(&dsv3_router_gemm));
210
-}
csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu
deleted
-337
@@ -1,337 +0,0 @@
1
-/*
2
- * Adapted from SGLang's sgl-kernel implementation, which was adapted from
3
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
4
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
5
- *
6
- * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved.
7
- *
8
- * Licensed under the Apache License, Version 2.0 (the "License");
9
- * you may not use this file except in compliance with the License.
10
- * You may obtain a copy of the License at
11
- *
12
- * http://www.apache.org/licenses/LICENSE-2.0
13
- *
14
- * Unless required by applicable law or agreed to in writing, software
15
- * distributed under the License is distributed on an "AS IS" BASIS,
16
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
- * See the License for the specific language governing permissions and
18
- * limitations under the License.
19
- */
20
-
21
-#include <cstdint>
22
-
23
-#include <cuda_bf16.h>
24
-#include <cuda_runtime.h>
25
-
26
-// Custom FMA implementation using PTX assembly instructions
27
-__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
28
- float2 const& c) {
29
- asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
30
- : "=l"(reinterpret_cast<uint64_t&>(d))
31
- : "l"(reinterpret_cast<uint64_t const&>(a)),
32
- "l"(reinterpret_cast<uint64_t const&>(b)),
33
- "l"(reinterpret_cast<uint64_t const&>(c)));
34
-}
35
-
36
-// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion
37
-template <int VPT>
38
-__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
39
- float* dst) {
40
- __nv_bfloat16* bf16_ptr =
41
- reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
42
-
43
-#pragma unroll
44
- for (int i = 0; i < VPT; i++) {
45
- dst[i] = __bfloat162float(bf16_ptr[i]);
46
- }
47
-}
48
-
49
-template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
50
- int kHiddenDim>
51
-__global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
52
- float* out, T const* mat_a, T const* mat_b) {
53
- // Each block handles one expert column
54
- int const n_idx = blockIdx.x;
55
- int const tid = threadIdx.x;
56
- constexpr int kWarpSize = 32;
57
- constexpr int kNumWarps = kBlockSize / kWarpSize;
58
- // Constants for this kernel
59
- constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
60
- constexpr int k_iterations =
61
- kHiddenDim / k_elems_per_k_iteration; // Total K iterations
62
-
63
- // Initialize accumulators for all M rows
64
- float acc[kNumTokens] = {};
65
-
66
- // Shared memory for warp-level reduction
67
- __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps
68
-
69
- // B matrix is in column-major order, so we can directly load a column for the
70
- // n_idx expert
71
- T const* b_col = mat_b + n_idx * kHiddenDim;
72
-
73
- // Pre-compute k_base values for each iteration to help compiler optimize
74
- int k_bases[k_iterations];
75
-#pragma unroll
76
- for (int ki = 0; ki < k_iterations; ki++) {
77
- k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
78
- }
79
-
80
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
81
- cudaGridDependencySynchronize();
82
-#endif
83
-
84
- // Process the GEMM in chunks
85
- for (int ki = 0; ki < k_iterations; ki++) {
86
- int const k_base = k_bases[ki];
87
-
88
- // Load B matrix values using vector load (8 bf16 values)
89
- uint4 b_vec = *reinterpret_cast<uint4 const*>(b_col + k_base);
90
-
91
- // Convert B values to float
92
- float b_float[VPT];
93
- bf16_uint4_to_float8<VPT>(b_vec, b_float);
94
-
95
-// Process each token
96
-#pragma unroll
97
- for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
98
- // Load both rows of A matrix using vector loads
99
- uint4 a_vec = *reinterpret_cast<uint4 const*>(
100
- mat_a + (m_idx * kHiddenDim) + k_base);
101
-
102
- // Convert A values to float
103
- float a_float[VPT];
104
- bf16_uint4_to_float8<VPT>(a_vec, a_float);
105
-
106
-// Process elements in this chunk
107
-#pragma unroll
108
- for (int k = 0; k < VPT; k++) {
109
- float a = a_float[k];
110
- float b = b_float[k];
111
- acc[m_idx] += a * b;
112
- }
113
- }
114
- }
115
-
116
- // Perform warp-level reduction
117
- int const warpSize = 32;
118
- int const warpId = tid / warpSize;
119
- int const laneId = tid % warpSize;
120
-
121
- // Register for warp-level reduction results
122
- float warp_result[kNumTokens];
123
-
124
-#pragma unroll
125
- for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
126
- warp_result[m_idx] = acc[m_idx];
127
- }
128
-
129
-// Perform warp-level reduction using optimized butterfly pattern
130
-#pragma unroll
131
- for (int m = 0; m < kNumTokens; m++) {
132
- float sum = warp_result[m];
133
-
134
- // Butterfly reduction pattern
135
- sum += __shfl_xor_sync(0xffffffff, sum, 16);
136
- sum += __shfl_xor_sync(0xffffffff, sum, 8);
137
- sum += __shfl_xor_sync(0xffffffff, sum, 4);
138
- sum += __shfl_xor_sync(0xffffffff, sum, 2);
139
- sum += __shfl_xor_sync(0xffffffff, sum, 1);
140
-
141
- // Only the first thread in each warp stores to shared memory
142
- if (laneId == 0) {
143
- sm_reduction[m][warpId] = sum;
144
- }
145
- }
146
-
147
- __syncthreads();
148
-
149
- // Final reduction across warps (only first thread)
150
- if (tid == 0) {
151
-#pragma unroll
152
- for (int m = 0; m < kNumTokens; m++) {
153
- float final_sum = 0.0f;
154
-
155
-// Sum across the kNumWarps
156
-#pragma unroll
157
- for (int w = 0; w < kNumWarps; w++) {
158
- final_sum += sm_reduction[m][w];
159
- }
160
-
161
- // Write final result
162
- out[m * kNumExperts + n_idx] = final_sum;
163
- }
164
- }
165
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
166
- cudaTriggerProgrammaticLaunchCompletion();
167
-#endif
168
-}
169
-
170
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
171
-void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
172
- cudaStream_t stream) {
173
- constexpr int VPT = 16 / sizeof(T);
174
- constexpr int kBlockSize = 128;
175
- cudaLaunchConfig_t config;
176
- config.gridDim = kNumExperts;
177
- config.blockDim = kBlockSize;
178
- config.dynamicSmemBytes = 0;
179
- config.stream = stream;
180
- cudaLaunchAttribute attrs[1];
181
- attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
182
- attrs[0].val.programmaticStreamSerializationAllowed = 1;
183
- config.numAttrs = 1;
184
- config.attrs = attrs;
185
- cudaLaunchKernelEx(
186
- &config,
187
- router_gemm_kernel_float_output<T, kBlockSize, VPT, kNumTokens,
188
- kNumExperts, kHiddenDim>,
189
- output, mat_a, mat_b);
190
-}
191
-
192
-// Template instantiations for DEFAULT_NUM_EXPERTS experts
193
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 7168>(
194
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
195
-
196
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 7168>(
197
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
198
-
199
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 7168>(
200
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
201
-
202
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 7168>(
203
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
204
-
205
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 7168>(
206
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
207
-
208
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 7168>(
209
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
210
-
211
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 7168>(
212
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
213
-
214
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 7168>(
215
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
216
-
217
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 7168>(
218
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
219
-
220
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 7168>(
221
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
222
-
223
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 7168>(
224
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
225
-
226
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 7168>(
227
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
228
-
229
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 7168>(
230
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
231
-
232
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 7168>(
233
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
234
-
235
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 7168>(
236
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
237
-
238
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 7168>(
239
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
240
-
241
-// Template instantiations for KIMI_K2_NUM_EXPERTS experts
242
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 384, 7168>(
243
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
244
-
245
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 384, 7168>(
246
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
247
-
248
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 384, 7168>(
249
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
250
-
251
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 384, 7168>(
252
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
253
-
254
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 384, 7168>(
255
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
256
-
257
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 384, 7168>(
258
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
259
-
260
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 384, 7168>(
261
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
262
-
263
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 384, 7168>(
264
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
265
-
266
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 384, 7168>(
267
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
268
-
269
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 384, 7168>(
270
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
271
-
272
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 384, 7168>(
273
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
274
-
275
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 384, 7168>(
276
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
277
-
278
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 384, 7168>(
279
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
280
-
281
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 384, 7168>(
282
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
283
-
284
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>(
285
- float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
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);
csrc/libtorch_stable/moe/moe_ops.h
-12
@@ -79,15 +79,3 @@ int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
79
void shuffle_rows(const torch::stable::Tensor& input_tensor,
80
const torch::stable::Tensor& dst2src_map,
81
torch::stable::Tensor& output_tensor);
82
-
83
-#ifndef USE_ROCM
84
-// DeepSeek V3 optimized router GEMM kernel for SM90+
85
-// Computes output = mat_a @ mat_b.T where:
86
-// mat_a: [num_tokens, hidden_dim] in bf16
87
-// mat_b: [num_experts, hidden_dim] in bf16
88
-// output: [num_tokens, num_experts] in bf16 or fp32
89
-// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
90
-void dsv3_router_gemm(torch::stable::Tensor& output,
91
- const torch::stable::Tensor& mat_a,
92
- const torch::stable::Tensor& mat_b);
93
-#endif
csrc/libtorch_stable/moe/torch_bindings.cpp
-4
@@ -126,10 +126,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
126
"topk_group, int topk, bool renormalize, float "
127
"routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, "
128
"Tensor)");
129
-
130
- // DeepSeek V3 optimized router GEMM for SM90+
131
- m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
132
- // conditionally compiled so impl registration is in source file
129
#endif
130
}
131
tests/kernels/test_fp32_router_gemm.py
+2
-3
@@ -3,8 +3,7 @@
3
"""Tests for fp32_router_gemm kernel: activation×weight→fp32.
4
5
Supported (hidden_size, num_experts) pairs:
6
- (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3,
7
- (6144, 256) -> GLM-5.2
6
+ (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
7
8
Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so
9
all tuned geometries (wide-block, experts-per-block, token-group; boundaries
@@ -19,7 +18,7 @@ from vllm._custom_ops import fp32_router_gemm
18
from vllm.platforms import current_platform
19
20
# (hidden_size, num_experts)
22
-SHAPES = [(3072, 256), (6144, 128), (6144, 256)]
21
+SHAPES = [(3072, 256), (6144, 128)]
22
ALL_M = list(range(1, 33))
23
# Absolute tolerance for fp32 kernel vs float64 reference
24
ATOL_FP32 = 2e-4
vllm/_custom_ops.py
-15
@@ -2384,21 +2384,6 @@ def moe_wna16_gemm(
2384
)
2385
2386
2387
-def dsv3_router_gemm(
2388
- hidden_states: torch.Tensor,
2389
- router_weight: torch.Tensor,
2390
- output_dtype: torch.dtype,
2391
-) -> torch.Tensor:
2392
- output = torch.empty(
2393
- hidden_states.shape[0],
2394
- router_weight.shape[0],
2395
- device=hidden_states.device,
2396
- dtype=output_dtype,
2397
- )
2398
- torch.ops._moe_C.dsv3_router_gemm(output, hidden_states, router_weight)
2399
- return output
2400
-
2401
-
2387
def fp32_router_gemm(
2388
hidden_states: torch.Tensor,
2389
router_weight: torch.Tensor,
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+10
-40
@@ -20,30 +20,20 @@ class GateLinear(ReplicatedLinear):
20
21
1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out,
22
K divisible by 8)
23
- 2. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
24
- 3. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32,
25
- (H, E) in {(3072, 256), (6144, 128), (6144, 256)})
26
- 4. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight)
27
- 5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
28
- 6. F.linear via ReplicatedLinear (ultimate fallback)
23
+ 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32,
24
+ (H, E) in {(3072, 256), (6144, 128)})
25
+ 3. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight)
26
+ 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
27
+ 5. F.linear via ReplicatedLinear (ultimate fallback)
28
29
The ``out_dtype`` attribute is mutable and can be set after init
30
(e.g. when the required dtype depends on the expert quantization
31
method which is only known later).
32
"""
33
35
- # Dimensions supported by the DSV3 specialized kernel.
36
- # Valid (hidden_size, num_experts) combinations:
37
- # (7168, 256) -> DeepSeek-V3, (7168, 384) -> Kimi-K2,
38
- # (6144, 256) -> GLM-5
39
- DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
40
- DSV3_SUPPORTED_HIDDEN_SIZES = [7168, 6144]
41
- # num_experts=384 is only instantiated for hidden_size=7168.
42
- DSV3_UNSUPPORTED_SHAPES = {(6144, 384)}
43
-
34
# (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
35
# (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
46
- FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (6144, 256)}
36
+ FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)}
37
FP32_MAX_TOKENS = 32
38
39
def __init__(
@@ -77,18 +67,7 @@ class GateLinear(ReplicatedLinear):
67
)
68
self.out_dtype = out_dtype
69
80
- # DSV3 specialized kernel eligibility (SM90+, exact dims)
70
self.allow_specialized_router_gemm = can_use_specialized_kernels
82
- self.allow_dsv3_router_gemm = (
83
- self.allow_specialized_router_gemm
84
- and self.weight.dtype == torch.bfloat16
85
- and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS
86
- and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
87
- and (input_size, output_size) not in self.DSV3_UNSUPPORTED_SHAPES
88
- )
89
- # See https://github.com/vllm-project/vllm/pull/44217
90
- # for more details.
91
- self._dsv3_max_batch = 16 if is_hopper else 8
71
72
# fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight)
73
vllm_config = get_current_vllm_config_or_none()
@@ -187,16 +166,7 @@ class GateLinear(ReplicatedLinear):
166
output = ll_bf16_gemm(x, self.weight)
167
return output, None
168
190
- # Tier 2: DSV3 specialized kernel (fallback for when cuteDSL unavailable)
191
- if self.allow_dsv3_router_gemm and x.shape[0] <= self._dsv3_max_batch:
192
- output = ops.dsv3_router_gemm(
193
- hidden_states=x,
194
- router_weight=self.weight,
195
- output_dtype=self.out_dtype,
196
- )
197
- return output, None
198
-
199
- # Tier 3: fp32 specialized kernel (H=3072, E=256, M<=32)
169
+ # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32)
170
# Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph
171
# capture does not freeze the runtime num_tokens branch.
172
if self.allow_fp32_router_gemm and x.dtype in (
@@ -208,7 +178,7 @@ class GateLinear(ReplicatedLinear):
178
)
179
return output, None
180
211
- # Tier 4: experimental bf16x3 CuteDSL kernel for fp32 router weights
181
+ # Tier 3: experimental bf16x3 CuteDSL kernel for fp32 router weights
182
if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
183
from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501
184
bf16x3_router_gemm,
@@ -217,12 +187,12 @@ class GateLinear(ReplicatedLinear):
187
output = bf16x3_router_gemm(x, self.weight)
188
return output, None
189
220
- # Tier 5: cuBLAS bf16→fp32
190
+ # Tier 4: cuBLAS bf16→fp32
191
if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
192
output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
193
return output, None
194
225
- # Tier 6: F.linear (ReplicatedLinear)
195
+ # Tier 5: F.linear (ReplicatedLinear)
196
if self.out_dtype is not None and x.dtype != self.weight.dtype:
197
x = x.to(self.weight.dtype)
198
output, output_bias = super().forward(x)