@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8435b2e049
Fork 0

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 7, 22:45 UTC

[ModelBash][DSV3] Add TRTLLM DSV3 Router GEMM kernel (6% B1 Speedup) (#34302)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>

Robert Shaw committed Feb 23, 2026 at 09:02 UTC 8435b2e0492525b2ae7361b7f0081febb483bb34
9 files changed +915 -3
CMakeLists.txt
+21
@@ -1101,6 +1101,27 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1101 message(STATUS "Not building Marlin MOE kernels as no compatible archs found"
1102 " in CUDA target architectures")
1103 endif()
1104 +
1105 + # DeepSeek V3 router GEMM kernel - requires SM90+
1106 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
1107 + cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
1108 + else()
1109 + cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
1110 + endif()
1111 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
1112 + set(DSV3_ROUTER_GEMM_SRC
1113 + "csrc/moe/dsv3_router_gemm_entry.cu"
1114 + "csrc/moe/dsv3_router_gemm_float_out.cu"
1115 + "csrc/moe/dsv3_router_gemm_bf16_out.cu")
1116 + set_gencode_flags_for_srcs(
1117 + SRCS "${DSV3_ROUTER_GEMM_SRC}"
1118 + CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
1119 + list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
1120 + message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
1121 + else()
1122 + message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
1123 + " (requires SM90+ and CUDA >= 12.0)")
1124 + endif()
1125 endif()
1126
1127 message(STATUS "Enabling moe extension.")
csrc/moe/dsv3_router_gemm_bf16_out.cu new
+291
@@ -0,0 +1,291 @@
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 <ATen/ATen.h>
22 +#include <ATen/cuda/CUDAContext.h>
23 +
24 +#include <cuda_bf16.h>
25 +#include <cuda_runtime.h>
26 +
27 +#include "dsv3_router_gemm_utils.h"
28 +
29 +// Custom FMA implementation using PTX assembly instructions
30 +__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
31 + float2 const& c) {
32 + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
33 + : "=l"(reinterpret_cast<uint64_t&>(d))
34 + : "l"(reinterpret_cast<uint64_t const&>(a)),
35 + "l"(reinterpret_cast<uint64_t const&>(b)),
36 + "l"(reinterpret_cast<uint64_t const&>(c)));
37 +}
38 +
39 +// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion
40 +template <int VPT>
41 +__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
42 + float* dst) {
43 + __nv_bfloat16* bf16_ptr =
44 + reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
45 +
46 +#pragma unroll
47 + for (int i = 0; i < VPT; i++) {
48 + dst[i] = __bfloat162float(bf16_ptr[i]);
49 + }
50 +}
51 +
52 +template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
53 + int kHiddenDim>
54 +__global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
55 + __nv_bfloat16* out, T const* mat_a, T const* mat_b) {
56 + // Each block handles one expert column
57 + int const n_idx = blockIdx.x;
58 + int const tid = threadIdx.x;
59 + constexpr int kWarpSize = 32;
60 + constexpr int kNumWarps = kBlockSize / kWarpSize;
61 + // Constants for this kernel
62 + constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
63 + constexpr int k_iterations =
64 + kHiddenDim / k_elems_per_k_iteration; // Total K iterations
65 +
66 + // Initialize accumulators for all M rows
67 + float acc[kNumTokens] = {};
68 +
69 + // Shared memory for warp-level reduction
70 + __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps
71 +
72 + // B matrix is in column-major order, so we can directly load a column for the
73 + // n_idx expert
74 + T const* b_col = mat_b + n_idx * kHiddenDim;
75 +
76 + // Pre-compute k_base values for each iteration to help compiler optimize
77 + int k_bases[k_iterations];
78 +#pragma unroll
79 + for (int ki = 0; ki < k_iterations; ki++) {
80 + k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
81 + }
82 +
83 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
84 + asm volatile("griddepcontrol.wait;");
85 +#endif
86 +
87 + // Process the GEMM in chunks
88 + for (int ki = 0; ki < k_iterations; ki++) {
89 + int const k_base = k_bases[ki];
90 +
91 + // Load B matrix values using vector load (8 bf16 values)
92 + uint4 b_vec = *reinterpret_cast<uint4 const*>(b_col + k_base);
93 +
94 + // Convert B values to float
95 + float b_float[VPT];
96 + bf16_uint4_to_float8<VPT>(b_vec, b_float);
97 +
98 +// Process each token
99 +#pragma unroll
100 + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
101 + // Load both rows of A matrix using vector loads
102 + uint4 a_vec = *reinterpret_cast<uint4 const*>(
103 + mat_a + (m_idx * kHiddenDim) + k_base);
104 +
105 + // Convert A values to float
106 + float a_float[VPT];
107 + bf16_uint4_to_float8<VPT>(a_vec, a_float);
108 +
109 +// Process elements in this chunk
110 +#pragma unroll
111 + for (int k = 0; k < VPT; k++) {
112 + float a = a_float[k];
113 + float b = b_float[k];
114 + acc[m_idx] += a * b;
115 + }
116 + }
117 + }
118 +
119 + // Perform warp-level reduction
120 + int const warpSize = 32;
121 + int const warpId = tid / warpSize;
122 + int const laneId = tid % warpSize;
123 +
124 + // Register for warp-level reduction results
125 + float warp_result[kNumTokens];
126 +
127 +#pragma unroll
128 + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
129 + warp_result[m_idx] = acc[m_idx];
130 + }
131 +
132 +// Perform warp-level reduction using optimized butterfly pattern
133 +#pragma unroll
134 + for (int m = 0; m < kNumTokens; m++) {
135 + float sum = warp_result[m];
136 +
137 + // Butterfly reduction pattern
138 + sum += __shfl_xor_sync(0xffffffff, sum, 16);
139 + sum += __shfl_xor_sync(0xffffffff, sum, 8);
140 + sum += __shfl_xor_sync(0xffffffff, sum, 4);
141 + sum += __shfl_xor_sync(0xffffffff, sum, 2);
142 + sum += __shfl_xor_sync(0xffffffff, sum, 1);
143 +
144 + // Only the first thread in each warp stores to shared memory
145 + if (laneId == 0) {
146 + sm_reduction[m][warpId] = sum;
147 + }
148 + }
149 +
150 + __syncthreads();
151 +
152 + // Final reduction across warps (only first thread)
153 + if (tid == 0) {
154 +#pragma unroll
155 + for (int m = 0; m < kNumTokens; m++) {
156 + float final_sum = 0.0f;
157 +
158 +// Sum across the kNumWarps
159 +#pragma unroll
160 + for (int w = 0; w < kNumWarps; w++) {
161 + final_sum += sm_reduction[m][w];
162 + }
163 +
164 + // Write final result
165 + out[m * kNumExperts + n_idx] = __float2bfloat16(final_sum);
166 + }
167 + }
168 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
169 + asm volatile("griddepcontrol.launch_dependents;");
170 +#endif
171 +}
172 +
173 +template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
174 +void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
175 + T const* mat_b, cudaStream_t stream) {
176 + constexpr int VPT = 16 / sizeof(T);
177 + constexpr int kBlockSize = 128;
178 + cudaLaunchConfig_t config;
179 + config.gridDim = kNumExperts;
180 + config.blockDim = kBlockSize;
181 + config.dynamicSmemBytes = 0;
182 + config.stream = stream;
183 + cudaLaunchAttribute attrs[1];
184 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
185 + attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
186 + config.numAttrs = 1;
187 + config.attrs = attrs;
188 + cudaLaunchKernelEx(
189 + &config,
190 + router_gemm_kernel_bf16_output<T, kBlockSize, VPT, kNumTokens,
191 + kNumExperts, kHiddenDim>,
192 + output, mat_a, mat_b);
193 +}
194 +
195 +// Template instantiations for DEFAULT_NUM_EXPERTS experts
196 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 7168>(
197 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
198 +
199 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 7168>(
200 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
201 +
202 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 7168>(
203 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
204 +
205 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 7168>(
206 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
207 +
208 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 7168>(
209 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
210 +
211 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 7168>(
212 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
213 +
214 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 7168>(
215 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
216 +
217 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 7168>(
218 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
219 +
220 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 7168>(
221 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
222 +
223 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 7168>(
224 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
225 +
226 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 7168>(
227 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
228 +
229 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 7168>(
230 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
231 +
232 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 7168>(
233 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
234 +
235 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 7168>(
236 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
237 +
238 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 7168>(
239 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
240 +
241 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 7168>(
242 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
243 +
244 +// Template instantiations for KIMI_K2_NUM_EXPERTS experts
245 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 384, 7168>(
246 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
247 +
248 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 384, 7168>(
249 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
250 +
251 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 384, 7168>(
252 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
253 +
254 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 384, 7168>(
255 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
256 +
257 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 384, 7168>(
258 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
259 +
260 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 384, 7168>(
261 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
262 +
263 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 384, 7168>(
264 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
265 +
266 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 384, 7168>(
267 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
268 +
269 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 384, 7168>(
270 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
271 +
272 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 384, 7168>(
273 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
274 +
275 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 384, 7168>(
276 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
277 +
278 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 384, 7168>(
279 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
280 +
281 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 384, 7168>(
282 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
283 +
284 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 384, 7168>(
285 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
286 +
287 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>(
288 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
289 +
290 +template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>(
291 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
csrc/moe/dsv3_router_gemm_entry.cu new
+163
@@ -0,0 +1,163 @@
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 <ATen/ATen.h>
22 +#include <ATen/cuda/CUDAContext.h>
23 +
24 +#include <cuda_bf16.h>
25 +#include <cuda_runtime.h>
26 +
27 +#include "dsv3_router_gemm_utils.h"
28 +
29 +static constexpr int DEFAULT_NUM_EXPERTS = 256;
30 +static constexpr int KIMI_K2_NUM_EXPERTS = 384;
31 +static constexpr int DEFAULT_HIDDEN_DIM = 7168;
32 +
33 +template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
34 +void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
35 + cudaStream_t stream);
36 +
37 +template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
38 +void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
39 + T const* mat_b, cudaStream_t stream);
40 +
41 +template <int kBegin, int kEnd, int kNumExperts, int kHiddenDim>
42 +struct LoopUnroller {
43 + static void unroll_float_output(int num_tokens, float* output,
44 + __nv_bfloat16 const* input,
45 + __nv_bfloat16 const* weights,
46 + cudaStream_t stream) {
47 + if (num_tokens == kBegin) {
48 + invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts,
49 + kHiddenDim>(output, input, weights, stream);
50 + } else {
51 + LoopUnroller<kBegin + 1, kEnd, kNumExperts,
52 + kHiddenDim>::unroll_float_output(num_tokens, output, input,
53 + weights, stream);
54 + }
55 + }
56 +
57 + static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
58 + __nv_bfloat16 const* input,
59 + __nv_bfloat16 const* weights,
60 + cudaStream_t stream) {
61 + if (num_tokens == kBegin) {
62 + invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts,
63 + kHiddenDim>(output, input, weights, stream);
64 + } else {
65 + LoopUnroller<kBegin + 1, kEnd, kNumExperts,
66 + kHiddenDim>::unroll_bf16_output(num_tokens, output, input,
67 + weights, stream);
68 + }
69 + }
70 +};
71 +
72 +template <int kEnd, int kNumExperts, int kHiddenDim>
73 +struct LoopUnroller<kEnd, kEnd, kNumExperts, kHiddenDim> {
74 + static void unroll_float_output(int num_tokens, float* output,
75 + __nv_bfloat16 const* input,
76 + __nv_bfloat16 const* weights,
77 + cudaStream_t stream) {
78 + if (num_tokens == kEnd) {
79 + invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
80 + output, input, weights, stream);
81 + } else {
82 + throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
83 + }
84 + }
85 +
86 + static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
87 + __nv_bfloat16 const* input,
88 + __nv_bfloat16 const* weights,
89 + cudaStream_t stream) {
90 + if (num_tokens == kEnd) {
91 + invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
92 + output, input, weights, stream);
93 + } else {
94 + throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
95 + }
96 + }
97 +};
98 +
99 +void dsv3_router_gemm(at::Tensor& output, // [num_tokens, num_experts]
100 + const at::Tensor& mat_a, // [num_tokens, hidden_dim]
101 + const at::Tensor& mat_b // [num_experts, hidden_dim]
102 +) {
103 + TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
104 +
105 + const int num_tokens = mat_a.size(0);
106 + const int num_experts = mat_b.size(0);
107 + const int hidden_dim = mat_a.size(1);
108 +
109 + TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
110 + "mat_a and mat_b must have the same hidden_dim");
111 + TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
112 + "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
113 + ", but got hidden_dim=", hidden_dim);
114 + TORCH_CHECK(
115 + num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
116 + "Expected num_experts=", DEFAULT_NUM_EXPERTS,
117 + " or num_experts=", KIMI_K2_NUM_EXPERTS,
118 + ", but got num_experts=", num_experts);
119 + TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
120 + "currently num_tokens must be less than or equal to 16 for "
121 + "router_gemm");
122 + TORCH_CHECK(mat_a.dtype() == at::kBFloat16, "mat_a must be bf16");
123 + TORCH_CHECK(mat_b.dtype() == at::kBFloat16, "mat_b must be bf16");
124 + TORCH_CHECK(output.dtype() == at::kFloat || output.dtype() == at::kBFloat16,
125 + "output must be float32 or bf16");
126 +
127 + auto const sm = getSMVersion();
128 + TORCH_CHECK(sm >= 90 && sm <= 103, "required SM_103 >= CUDA ARCH >= SM_90");
129 +
130 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
131 +
132 + if (output.dtype() == at::kFloat) {
133 + if (num_experts == DEFAULT_NUM_EXPERTS) {
134 + LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
135 + unroll_float_output(
136 + num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
137 + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
138 + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
139 + } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
140 + LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
141 + unroll_float_output(
142 + num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
143 + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
144 + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
145 + }
146 + } else if (output.dtype() == at::kBFloat16) {
147 + if (num_experts == DEFAULT_NUM_EXPERTS) {
148 + LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
149 + unroll_bf16_output(
150 + num_tokens,
151 + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
152 + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
153 + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
154 + } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
155 + LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
156 + unroll_bf16_output(
157 + num_tokens,
158 + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
159 + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
160 + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
161 + }
162 + }
163 +}
csrc/moe/dsv3_router_gemm_float_out.cu new
+291
@@ -0,0 +1,291 @@
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 <ATen/ATen.h>
22 +#include <ATen/cuda/CUDAContext.h>
23 +
24 +#include <cuda_bf16.h>
25 +#include <cuda_runtime.h>
26 +
27 +#include "dsv3_router_gemm_utils.h"
28 +
29 +// Custom FMA implementation using PTX assembly instructions
30 +__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
31 + float2 const& c) {
32 + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
33 + : "=l"(reinterpret_cast<uint64_t&>(d))
34 + : "l"(reinterpret_cast<uint64_t const&>(a)),
35 + "l"(reinterpret_cast<uint64_t const&>(b)),
36 + "l"(reinterpret_cast<uint64_t const&>(c)));
37 +}
38 +
39 +// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion
40 +template <int VPT>
41 +__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
42 + float* dst) {
43 + __nv_bfloat16* bf16_ptr =
44 + reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
45 +
46 +#pragma unroll
47 + for (int i = 0; i < VPT; i++) {
48 + dst[i] = __bfloat162float(bf16_ptr[i]);
49 + }
50 +}
51 +
52 +template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
53 + int kHiddenDim>
54 +__global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
55 + float* out, T const* mat_a, T const* mat_b) {
56 + // Each block handles one expert column
57 + int const n_idx = blockIdx.x;
58 + int const tid = threadIdx.x;
59 + constexpr int kWarpSize = 32;
60 + constexpr int kNumWarps = kBlockSize / kWarpSize;
61 + // Constants for this kernel
62 + constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
63 + constexpr int k_iterations =
64 + kHiddenDim / k_elems_per_k_iteration; // Total K iterations
65 +
66 + // Initialize accumulators for all M rows
67 + float acc[kNumTokens] = {};
68 +
69 + // Shared memory for warp-level reduction
70 + __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps
71 +
72 + // B matrix is in column-major order, so we can directly load a column for the
73 + // n_idx expert
74 + T const* b_col = mat_b + n_idx * kHiddenDim;
75 +
76 + // Pre-compute k_base values for each iteration to help compiler optimize
77 + int k_bases[k_iterations];
78 +#pragma unroll
79 + for (int ki = 0; ki < k_iterations; ki++) {
80 + k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
81 + }
82 +
83 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
84 + asm volatile("griddepcontrol.wait;");
85 +#endif
86 +
87 + // Process the GEMM in chunks
88 + for (int ki = 0; ki < k_iterations; ki++) {
89 + int const k_base = k_bases[ki];
90 +
91 + // Load B matrix values using vector load (8 bf16 values)
92 + uint4 b_vec = *reinterpret_cast<uint4 const*>(b_col + k_base);
93 +
94 + // Convert B values to float
95 + float b_float[VPT];
96 + bf16_uint4_to_float8<VPT>(b_vec, b_float);
97 +
98 +// Process each token
99 +#pragma unroll
100 + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
101 + // Load both rows of A matrix using vector loads
102 + uint4 a_vec = *reinterpret_cast<uint4 const*>(
103 + mat_a + (m_idx * kHiddenDim) + k_base);
104 +
105 + // Convert A values to float
106 + float a_float[VPT];
107 + bf16_uint4_to_float8<VPT>(a_vec, a_float);
108 +
109 +// Process elements in this chunk
110 +#pragma unroll
111 + for (int k = 0; k < VPT; k++) {
112 + float a = a_float[k];
113 + float b = b_float[k];
114 + acc[m_idx] += a * b;
115 + }
116 + }
117 + }
118 +
119 + // Perform warp-level reduction
120 + int const warpSize = 32;
121 + int const warpId = tid / warpSize;
122 + int const laneId = tid % warpSize;
123 +
124 + // Register for warp-level reduction results
125 + float warp_result[kNumTokens];
126 +
127 +#pragma unroll
128 + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
129 + warp_result[m_idx] = acc[m_idx];
130 + }
131 +
132 +// Perform warp-level reduction using optimized butterfly pattern
133 +#pragma unroll
134 + for (int m = 0; m < kNumTokens; m++) {
135 + float sum = warp_result[m];
136 +
137 + // Butterfly reduction pattern
138 + sum += __shfl_xor_sync(0xffffffff, sum, 16);
139 + sum += __shfl_xor_sync(0xffffffff, sum, 8);
140 + sum += __shfl_xor_sync(0xffffffff, sum, 4);
141 + sum += __shfl_xor_sync(0xffffffff, sum, 2);
142 + sum += __shfl_xor_sync(0xffffffff, sum, 1);
143 +
144 + // Only the first thread in each warp stores to shared memory
145 + if (laneId == 0) {
146 + sm_reduction[m][warpId] = sum;
147 + }
148 + }
149 +
150 + __syncthreads();
151 +
152 + // Final reduction across warps (only first thread)
153 + if (tid == 0) {
154 +#pragma unroll
155 + for (int m = 0; m < kNumTokens; m++) {
156 + float final_sum = 0.0f;
157 +
158 +// Sum across the kNumWarps
159 +#pragma unroll
160 + for (int w = 0; w < kNumWarps; w++) {
161 + final_sum += sm_reduction[m][w];
162 + }
163 +
164 + // Write final result
165 + out[m * kNumExperts + n_idx] = final_sum;
166 + }
167 + }
168 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
169 + asm volatile("griddepcontrol.launch_dependents;");
170 +#endif
171 +}
172 +
173 +template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
174 +void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
175 + cudaStream_t stream) {
176 + constexpr int VPT = 16 / sizeof(T);
177 + constexpr int kBlockSize = 128;
178 + cudaLaunchConfig_t config;
179 + config.gridDim = kNumExperts;
180 + config.blockDim = kBlockSize;
181 + config.dynamicSmemBytes = 0;
182 + config.stream = stream;
183 + cudaLaunchAttribute attrs[1];
184 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
185 + attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
186 + config.numAttrs = 1;
187 + config.attrs = attrs;
188 + cudaLaunchKernelEx(
189 + &config,
190 + router_gemm_kernel_float_output<T, kBlockSize, VPT, kNumTokens,
191 + kNumExperts, kHiddenDim>,
192 + output, mat_a, mat_b);
193 +}
194 +
195 +// Template instantiations for DEFAULT_NUM_EXPERTS experts
196 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 7168>(
197 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
198 +
199 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 7168>(
200 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
201 +
202 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 7168>(
203 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
204 +
205 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 7168>(
206 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
207 +
208 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 7168>(
209 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
210 +
211 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 7168>(
212 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
213 +
214 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 7168>(
215 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
216 +
217 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 7168>(
218 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
219 +
220 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 7168>(
221 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
222 +
223 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 7168>(
224 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
225 +
226 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 7168>(
227 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
228 +
229 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 7168>(
230 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
231 +
232 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 7168>(
233 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
234 +
235 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 7168>(
236 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
237 +
238 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 7168>(
239 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
240 +
241 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 7168>(
242 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
243 +
244 +// Template instantiations for KIMI_K2_NUM_EXPERTS experts
245 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 384, 7168>(
246 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
247 +
248 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 384, 7168>(
249 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
250 +
251 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 384, 7168>(
252 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
253 +
254 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 384, 7168>(
255 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
256 +
257 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 384, 7168>(
258 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
259 +
260 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 384, 7168>(
261 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
262 +
263 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 384, 7168>(
264 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
265 +
266 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 384, 7168>(
267 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
268 +
269 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 384, 7168>(
270 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
271 +
272 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 384, 7168>(
273 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
274 +
275 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 384, 7168>(
276 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
277 +
278 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 384, 7168>(
279 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
280 +
281 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 384, 7168>(
282 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
283 +
284 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 384, 7168>(
285 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
286 +
287 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>(
288 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
289 +
290 +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>(
291 + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
csrc/moe/dsv3_router_gemm_utils.h new
+43
@@ -0,0 +1,43 @@
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 +#pragma once
22 +
23 +#include <ATen/cuda/CUDAContext.h>
24 +
25 +#include <cstdlib>
26 +#include <mutex>
27 +
28 +inline int getSMVersion() {
29 + auto* props = at::cuda::getCurrentDeviceProperties();
30 + return props->major * 10 + props->minor;
31 +}
32 +
33 +inline bool getEnvEnablePDL() {
34 + static std::once_flag flag;
35 + static bool enablePDL = false;
36 + std::call_once(flag, [&]() {
37 + if (getSMVersion() >= 90) {
38 + const char* env = std::getenv("TRTLLM_ENABLE_PDL");
39 + enablePDL = env && env[0] == '1' && env[1] == '\0';
40 + }
41 + });
42 + return enablePDL;
43 +}
csrc/moe/moe_ops.h
+12 -1
@@ -55,4 +55,15 @@ bool moe_permute_unpermute_supported();
55
56 void shuffle_rows(const torch::Tensor& input_tensor,
57 const torch::Tensor& dst2src_map,
58 - torch::Tensor& output_tensor);
\ No newline at end of file
58 + torch::Tensor& output_tensor);
59 +
60 +#ifndef USE_ROCM
61 +// DeepSeek V3 optimized router GEMM kernel for SM90+
62 +// Computes output = mat_a @ mat_b.T where:
63 +// mat_a: [num_tokens, hidden_dim] in bf16
64 +// mat_b: [num_experts, hidden_dim] in bf16
65 +// output: [num_tokens, num_experts] in bf16 or fp32
66 +// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
67 +void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
68 + const torch::Tensor& mat_b);
69 +#endif
csrc/moe/torch_bindings.cpp
+4
@@ -124,6 +124,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
124 "routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, "
125 "Tensor)");
126 m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
127 +
128 + // DeepSeek V3 optimized router GEMM for SM90+
129 + m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
130 + m.impl("dsv3_router_gemm", torch::kCUDA, &dsv3_router_gemm);
131 #endif
132 }
133
vllm/_custom_ops.py
+15
@@ -2190,6 +2190,21 @@ def moe_wna16_gemm(
2190 )
2191
2192
2193 +def dsv3_router_gemm(
2194 + hidden_states: torch.Tensor,
2195 + router_weight: torch.Tensor,
2196 + output_dtype: torch.dtype,
2197 +) -> torch.Tensor:
2198 + output = torch.empty(
2199 + hidden_states.shape[0],
2200 + router_weight.shape[0],
2201 + device=hidden_states.device,
2202 + dtype=output_dtype,
2203 + )
2204 + torch.ops._moe_C.dsv3_router_gemm(output, hidden_states, router_weight)
2205 + return output
2206 +
2207 +
2208 def topk_softmax(
2209 topk_weights: torch.Tensor,
2210 topk_ids: torch.Tensor,
vllm/model_executor/models/deepseek_v2.py
+75 -2
@@ -221,6 +221,73 @@ class DeepseekV2MLP(nn.Module):
221 return x
222
223
224 +class DeepSeekV2Gate(ReplicatedLinear):
225 + def __init__(
226 + self,
227 + hidden_size: int,
228 + n_experts: int,
229 + quant_config: QuantizationConfig | None = None,
230 + prefix: str = "",
231 + ):
232 + assert quant_config is None
233 + super().__init__(
234 + hidden_size,
235 + n_experts,
236 + bias=False,
237 + quant_config=quant_config,
238 + prefix=f"{prefix}.gate",
239 + )
240 +
241 + # Unquantized only, will be called "weight".
242 + assert hasattr(self, "weight")
243 + is_hopper_or_blackwell = current_platform.is_device_capability(
244 + (9, 0)
245 + ) or current_platform.is_device_capability_family(100)
246 + SUPPORTED_NUM_EXPERTS = [256, 384]
247 + SUPPORTED_HIDDEN_SIZES = [7168]
248 +
249 + self.allow_dsv3_router_gemm = (
250 + current_platform.is_cuda()
251 + and is_hopper_or_blackwell
252 + and n_experts in SUPPORTED_NUM_EXPERTS
253 + and hidden_size in SUPPORTED_HIDDEN_SIZES
254 + )
255 +
256 + self._out_dtype: torch.dtype | None = None
257 +
258 + def set_out_dtype(self, out_dtype: torch.dtype) -> None:
259 + """
260 + Set out dtype for the router logits. This is needed after
261 + __init__, b/c we need to check if the trtllm kernel is
262 + selected before we decide between bf16 and fp32.
263 + """
264 +
265 + if self._out_dtype is not None:
266 + raise ValueError("out_dtype has already been set")
267 + else:
268 + self._out_dtype = out_dtype
269 +
270 + @property
271 + def out_dtype(self) -> torch.dtype:
272 + if self._out_dtype is None:
273 + raise ValueError("out_dtype has not been set yet")
274 + return self._out_dtype
275 +
276 + def forward(
277 + self,
278 + x: torch.Tensor,
279 + ) -> tuple[torch.Tensor, None]:
280 + """
281 + Use specialized GEMM for low batch size for DSV3 and KIMI.
282 + """
283 + if self.allow_dsv3_router_gemm and x.shape[0] <= 16:
284 + return ops.dsv3_router_gemm(
285 + hidden_states=x, router_weight=self.weight, output_dtype=self.out_dtype
286 + ), None
287 + else:
288 + return super().forward(x)
289 +
290 +
291 class DeepseekV2MoE(nn.Module):
292 def __init__(
293 self,
@@ -249,10 +316,9 @@ class DeepseekV2MoE(nn.Module):
316 "Only silu is supported for now."
317 )
318
252 - self.gate = ReplicatedLinear(
319 + self.gate = DeepSeekV2Gate(
320 config.hidden_size,
321 config.n_routed_experts,
255 - bias=False,
322 quant_config=None,
323 prefix=f"{prefix}.gate",
324 )
@@ -325,6 +391,13 @@ class DeepseekV2MoE(nn.Module):
391 else None,
392 )
393
394 + # NOTE(rob): this is a hack until we finish off the PR for
395 + # merging TRTLLM kernels into the MK framework. Then we can
396 + # query the MonolithicMK for the expected router logits.
397 + self.gate.set_out_dtype(
398 + torch.float32 if self.experts.quant_method.is_monolithic else torch.bfloat16
399 + )
400 +
401 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
402 num_tokens, hidden_dim = hidden_states.shape
403 hidden_states = hidden_states.view(-1, hidden_dim)