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 14, 05:30 UTC
[Kernel][Performance] Fuse float cast and renormalize to topk softmax kernel (#26717)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com> Signed-off-by: izhuhaoran <izhuhaoran@qq.com>
zhrrr committed
Oct 17, 2025 at 15:30 UTC
75c7ad9918c7fa6d16323db963c19b5f7c45fcfd
5 files changed
+222
-95
csrc/moe/moe_ops.h
+1
-1
@@ -4,7 +4,7 @@
4
5
void topk_softmax(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
6
torch::Tensor& token_expert_indices,
7
- torch::Tensor& gating_output);
7
+ torch::Tensor& gating_output, bool renormalize);
8
9
void moe_sum(torch::Tensor& input, torch::Tensor& output);
10
csrc/moe/topk_softmax_kernels.cu
+216
-87
@@ -16,12 +16,23 @@
16
* See the License for the specific language governing permissions and
17
* limitations under the License.
18
*/
19
+#include <type_traits>
20
#include <torch/all.h>
21
#include <ATen/cuda/CUDAContext.h>
22
#include <c10/cuda/CUDAGuard.h>
23
#include "../cuda_compat.h"
24
#include "../cub_helpers.h"
25
26
+#ifndef USE_ROCM
27
+ #include <cuda_bf16.h>
28
+ #include <cuda_fp16.h>
29
+#else
30
+ #include <hip/hip_bf16.h>
31
+ #include <hip/hip_fp16.h>
32
+ typedef __hip_bfloat16 __nv_bfloat16;
33
+ typedef __hip_bfloat162 __nv_bfloat162;
34
+#endif
35
+
36
#define MAX(a, b) ((a) > (b) ? (a) : (b))
37
#define MIN(a, b) ((a) < (b) ? (a) : (b))
38
@@ -36,16 +47,27 @@ template <
47
/// Alignment requirement in bytes
48
int Alignment = sizeof(T) * N
49
>
39
-class alignas(Alignment) AlignedArray {
40
- float data[N];
50
+struct alignas(Alignment) AlignedArray {
51
+ T data[N];
52
};
53
54
+template <typename T>
55
+__device__ __forceinline__ float toFloat(T value) {
56
+ if constexpr (std::is_same_v<T, float>) {
57
+ return value;
58
+ } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
59
+ return __bfloat162float(value);
60
+ } else if constexpr (std::is_same_v<T, __half>) {
61
+ return __half2float(value);
62
+ }
63
+}
64
+
65
// ====================== Softmax things ===============================
66
// We have our own implementation of softmax here so we can support transposing the output
67
// in the softmax kernel when we extend this module to support expert-choice routing.
46
-template <int TPB>
68
+template <int TPB, typename InputType>
69
__launch_bounds__(TPB) __global__
48
- void moeSoftmax(const float* input, const bool* finished, float* output, const int num_cols)
70
+ void moeSoftmax(const InputType* input, const bool* finished, float* output, const int num_cols)
71
{
72
using BlockReduce = cub::BlockReduce<float, TPB>;
73
__shared__ typename BlockReduce::TempStorage tmpStorage;
@@ -66,7 +88,8 @@ __launch_bounds__(TPB) __global__
88
for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
89
{
90
const int idx = thread_row_offset + ii;
69
- threadData = max(static_cast<float>(input[idx]), threadData);
91
+ const float val = toFloat(input[idx]);
92
+ threadData = max(val, threadData);
93
}
94
95
const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, CubMaxOp());
@@ -81,7 +104,8 @@ __launch_bounds__(TPB) __global__
104
for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
105
{
106
const int idx = thread_row_offset + ii;
84
- threadData += exp((static_cast<float>(input[idx]) - float_max));
107
+ const float val = toFloat(input[idx]);
108
+ threadData += expf(val - float_max);
109
}
110
111
const auto Z = BlockReduce(tmpStorage).Reduce(threadData, CubAddOp());
@@ -95,8 +119,9 @@ __launch_bounds__(TPB) __global__
119
for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
120
{
121
const int idx = thread_row_offset + ii;
98
- const float val = exp((static_cast<float>(input[idx]) - float_max)) * normalizing_factor;
99
- output[idx] = val;
122
+ const float val = toFloat(input[idx]);
123
+ const float softmax_val = expf(val - float_max) * normalizing_factor;
124
+ output[idx] = softmax_val;
125
}
126
}
127
@@ -110,7 +135,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
135
const int num_experts,
136
const int k,
137
const int start_expert,
113
- const int end_expert)
138
+ const int end_expert,
139
+ const bool renormalize)
140
{
141
142
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -125,6 +151,7 @@ __launch_bounds__(TPB) __global__ void moeTopK(
151
152
const bool row_is_active = finished ? !finished[block_row] : true;
153
const int thread_read_offset = blockIdx.x * num_experts;
154
+ float selected_sum = 0.f;
155
for (int k_idx = 0; k_idx < k; ++k_idx)
156
{
157
thread_kvp.key = 0;
@@ -163,9 +190,23 @@ __launch_bounds__(TPB) __global__ void moeTopK(
190
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
191
assert(indices[idx] >= 0);
192
source_rows[idx] = k_idx * num_rows + block_row;
193
+ if (renormalize) {
194
+ selected_sum += result_kvp.value;
195
+ }
196
}
197
__syncthreads();
198
}
199
+
200
+ // Renormalize the k weights for this row to sum to 1, if requested.
201
+ if (renormalize) {
202
+ if (threadIdx.x == 0) {
203
+ const float denom = selected_sum > 0.f ? selected_sum : 1.f;
204
+ for (int k_idx = 0; k_idx < k; ++k_idx) {
205
+ const int idx = k * block_row + k_idx;
206
+ output[idx] = output[idx] / denom;
207
+ }
208
+ }
209
+ }
210
}
211
212
// ====================== TopK softmax things ===============================
@@ -184,21 +225,30 @@ __launch_bounds__(TPB) __global__ void moeTopK(
225
2) This implementation assumes k is small, but will work for any k.
226
*/
227
187
-template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType>
228
+template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType, typename InputType = float>
229
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
189
- void topkGatingSoftmax(const float* input, const bool* finished, float* output, const int num_rows, IndType* indices,
190
- int* source_rows, const int k, const int start_expert, const int end_expert)
230
+ void topkGatingSoftmax(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
231
+ int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize)
232
{
233
+ static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
234
+ std::is_same_v<InputType, __half>,
235
+ "InputType must be float, __nv_bfloat16, or __half");
236
+
237
// We begin by enforcing compile time assertions and setting up compile time constants.
238
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2");
239
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
240
241
// Number of bytes each thread pulls in per load
197
- static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(float);
242
+ static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
243
static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
244
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
245
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
246
247
+ if constexpr (std::is_same_v<InputType, __nv_bfloat16> || std::is_same_v<InputType, __half>) {
248
+ static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
249
+ "ELTS_PER_LDG must be 1 or even for 16-bit conversion");
250
+ }
251
+
252
// Restrictions based on previous section.
253
static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
254
static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
@@ -236,27 +286,71 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
286
287
// We finally start setting up the read pointers for each thread. First, each thread jumps to the start of the
288
// row it will read.
239
- const float* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
289
+ const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
290
291
// Now, we compute the group each thread belong to in order to determine the first column to start loads.
292
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
293
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
244
- const float* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
245
-
246
- // Determine the pointer type to use to read in the data depending on the BYTES_PER_LDG template param. In theory,
247
- // this can support all powers of 2 up to 16.
248
- // NOTE(woosuk): The original implementation uses CUTLASS aligned array here.
249
- // We defined our own aligned array and use it here to avoid the dependency on CUTLASS.
250
- using AccessType = AlignedArray<float, ELTS_PER_LDG>;
294
+ const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
295
296
// Finally, we pull in the data from global mem
297
float row_chunk[VPT];
254
- AccessType* row_chunk_vec_ptr = reinterpret_cast<AccessType*>(&row_chunk);
255
- const AccessType* vec_thread_read_ptr = reinterpret_cast<const AccessType*>(thread_read_ptr);
298
+
299
+ // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert to float
300
+ if constexpr (std::is_same_v<InputType, float>) {
301
+ using VecType = AlignedArray<float, ELTS_PER_LDG>;
302
+ VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
303
+ const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
304
#pragma unroll
257
- for (int ii = 0; ii < LDG_PER_THREAD; ++ii)
258
- {
259
- row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
305
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
306
+ row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
307
+ }
308
+ } else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
309
+ if constexpr (ELTS_PER_LDG >= 2) {
310
+ using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
311
+ float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
312
+ const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
313
+#pragma unroll
314
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
315
+ VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
316
+ int base_idx_f2 = ii * ELTS_PER_LDG / 2;
317
+#pragma unroll
318
+ for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
319
+ row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
320
+ *reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2)
321
+ );
322
+ }
323
+ }
324
+ } else { // ELTS_PER_LDG == 1
325
+#pragma unroll
326
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
327
+ const __nv_bfloat16* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
328
+ row_chunk[ii] = __bfloat162float(*scalar_ptr);
329
+ }
330
+ }
331
+ } else if constexpr (std::is_same_v<InputType, __half>) {
332
+ if constexpr (ELTS_PER_LDG >= 2) {
333
+ using VecType = AlignedArray<__half, ELTS_PER_LDG>;
334
+ float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
335
+ const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
336
+#pragma unroll
337
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
338
+ VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
339
+ int base_idx_f2 = ii * ELTS_PER_LDG / 2;
340
+#pragma unroll
341
+ for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
342
+ row_chunk_f2[base_idx_f2 + jj] = __half22float2(
343
+ *reinterpret_cast<const __half2*>(vec.data + jj * 2)
344
+ );
345
+ }
346
+ }
347
+ } else { // ELTS_PER_LDG == 1
348
+#pragma unroll
349
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
350
+ const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
351
+ row_chunk[ii] = __half2float(*scalar_ptr);
352
+ }
353
+ }
354
}
355
356
// First, we perform a max reduce within the thread. We can do the max in fp16 safely (I think) and just
@@ -310,6 +404,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
404
int start_col = first_elt_read_by_thread;
405
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
406
407
+ float selected_sum = 0.f;
408
for (int k_idx = 0; k_idx < k; ++k_idx)
409
{
410
// First, each thread does the local argmax
@@ -363,6 +458,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
458
output[idx] = max_val;
459
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
460
source_rows[idx] = k_idx * num_rows + thread_row;
461
+ if (renormalize) {
462
+ selected_sum += max_val;
463
+ }
464
}
465
466
// Finally, we clear the value in the thread with the current max if there is another iteration to run.
@@ -380,15 +478,28 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
478
}
479
}
480
}
481
+
482
+ // Renormalize the k weights for this row to sum to 1, if requested.
483
+ if (renormalize) {
484
+ if (thread_group_idx == 0)
485
+ {
486
+ const float denom = selected_sum > 0.f ? selected_sum : 1.f;
487
+ for (int k_idx = 0; k_idx < k; ++k_idx)
488
+ {
489
+ const int idx = k * thread_row + k_idx;
490
+ output[idx] = output[idx] / denom;
491
+ }
492
+ }
493
+ }
494
}
495
496
namespace detail
497
{
498
// Constructs some constants needed to partition the work across threads at compile time.
388
-template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM>
499
+template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename InputType>
500
struct TopkConstants
501
{
391
- static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(float);
502
+ static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
503
static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0, "");
504
static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
505
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
@@ -397,20 +508,21 @@ struct TopkConstants
508
};
509
} // namespace detail
510
400
-template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType>
401
-void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, IndType* indices,
402
- int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, cudaStream_t stream)
511
+template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
512
+void topkGatingSoftmaxLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
513
+ int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
514
+ cudaStream_t stream)
515
{
404
- static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(float) * EXPERTS);
405
- using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM>;
516
+ static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
517
+ using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
518
static constexpr int VPT = Constants::VPT;
519
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
520
const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
521
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
522
523
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
412
- topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM><<<num_blocks, block_dim, 0, stream>>>(
413
- input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert);
524
+ topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType><<<num_blocks, block_dim, 0, stream>>>(
525
+ input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize);
526
}
527
528
#ifndef USE_ROCM
@@ -418,26 +530,26 @@ void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, f
530
static_assert(WARP_SIZE == 32, \
531
"Unsupported warp size. Only 32 is supported for CUDA"); \
532
topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, MAX_BYTES>( \
421
- gating_output, nullptr, topk_weights, topk_indices, \
422
- token_expert_indices, num_tokens, topk, 0, num_experts, stream);
533
+ gating_output, nullptr, topk_weights, topk_indices, token_expert_indices, \
534
+ num_tokens, topk, 0, num_experts, renormalize, stream);
535
#else
536
#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
537
if (WARP_SIZE == 64) { \
538
topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, MAX_BYTES>( \
427
- gating_output, nullptr, topk_weights, topk_indices, \
428
- token_expert_indices, num_tokens, topk, 0, num_experts, stream); \
539
+ gating_output, nullptr, topk_weights, topk_indices, token_expert_indices, \
540
+ num_tokens, topk, 0, num_experts, renormalize, stream); \
541
} else if (WARP_SIZE == 32) { \
542
topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES>( \
431
- gating_output, nullptr, topk_weights, topk_indices, \
432
- token_expert_indices, num_tokens, topk, 0, num_experts, stream); \
543
+ gating_output, nullptr, topk_weights, topk_indices, token_expert_indices, \
544
+ num_tokens, topk, 0, num_experts, renormalize, stream); \
545
} else { \
546
assert(false && "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
547
}
548
#endif
549
438
-template <typename IndType>
550
+template <typename IndType, typename InputType>
551
void topkGatingSoftmaxKernelLauncher(
440
- const float* gating_output,
552
+ const InputType* gating_output,
553
float* topk_weights,
554
IndType* topk_indices,
555
int* token_expert_indices,
@@ -445,11 +557,15 @@ void topkGatingSoftmaxKernelLauncher(
557
const int num_tokens,
558
const int num_experts,
559
const int topk,
560
+ const bool renormalize,
561
cudaStream_t stream) {
562
static constexpr int WARPS_PER_TB = 4;
563
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
564
#ifndef USE_ROCM
452
- static constexpr int BYTES_PER_LDG_MULTIPLE_64 = 8;
565
+ // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
566
+ // elements can be loaded by a warp
567
+ static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
568
+ (std::is_same_v<InputType, __nv_bfloat16> || std::is_same_v<InputType, __half>) ? 4 : 8;
569
#endif
570
switch (num_experts) {
571
case 1:
@@ -506,11 +622,11 @@ void topkGatingSoftmaxKernelLauncher(
622
TORCH_CHECK(softmax_workspace != nullptr,
623
"softmax_workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
624
static constexpr int TPB = 256;
509
- moeSoftmax<TPB><<<num_tokens, TPB, 0, stream>>>(
625
+ moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
626
gating_output, nullptr, softmax_workspace, num_experts);
627
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
628
softmax_workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
513
- num_experts, topk, 0, num_experts);
629
+ num_experts, topk, 0, num_experts, renormalize);
630
}
631
}
632
}
@@ -518,11 +634,50 @@ void topkGatingSoftmaxKernelLauncher(
634
} // namespace moe
635
} // namespace vllm
636
637
+
638
+template<typename ComputeType>
639
+void dispatch_topk_softmax_launch(
640
+ torch::Tensor& gating_output,
641
+ torch::Tensor& topk_weights,
642
+ torch::Tensor& topk_indices,
643
+ torch::Tensor& token_expert_indices,
644
+ torch::Tensor& softmax_workspace,
645
+ int num_tokens, int num_experts, int topk, bool renormalize, cudaStream_t stream)
646
+{
647
+ if (topk_indices.scalar_type() == at::ScalarType::Int) {
648
+ vllm::moe::topkGatingSoftmaxKernelLauncher<int, ComputeType>(
649
+ reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
650
+ topk_weights.data_ptr<float>(),
651
+ topk_indices.data_ptr<int>(),
652
+ token_expert_indices.data_ptr<int>(),
653
+ softmax_workspace.data_ptr<float>(),
654
+ num_tokens, num_experts, topk, renormalize, stream);
655
+ } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
656
+ vllm::moe::topkGatingSoftmaxKernelLauncher<uint32_t, ComputeType>(
657
+ reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
658
+ topk_weights.data_ptr<float>(),
659
+ topk_indices.data_ptr<uint32_t>(),
660
+ token_expert_indices.data_ptr<int>(),
661
+ softmax_workspace.data_ptr<float>(),
662
+ num_tokens, num_experts, topk, renormalize, stream);
663
+ } else {
664
+ TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
665
+ vllm::moe::topkGatingSoftmaxKernelLauncher<int64_t, ComputeType>(
666
+ reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
667
+ topk_weights.data_ptr<float>(),
668
+ topk_indices.data_ptr<int64_t>(),
669
+ token_expert_indices.data_ptr<int>(),
670
+ softmax_workspace.data_ptr<float>(),
671
+ num_tokens, num_experts, topk, renormalize, stream);
672
+ }
673
+}
674
+
675
void topk_softmax(
676
torch::Tensor& topk_weights, // [num_tokens, topk]
677
torch::Tensor& topk_indices, // [num_tokens, topk]
678
torch::Tensor& token_expert_indices, // [num_tokens, topk]
525
- torch::Tensor& gating_output) // [num_tokens, num_experts]
679
+ torch::Tensor& gating_output, // [num_tokens, num_experts]
680
+ bool renormalize)
681
{
682
const int num_experts = gating_output.size(-1);
683
const auto num_tokens = gating_output.numel() / num_experts;
@@ -534,45 +689,19 @@ void topk_softmax(
689
690
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
691
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
537
- torch::Tensor softmax_workspace = torch::empty({workspace_size}, gating_output.options());
538
-
539
- if(topk_indices.scalar_type() == at::ScalarType::Int)
540
- {
541
- vllm::moe::topkGatingSoftmaxKernelLauncher(
542
- gating_output.data_ptr<float>(),
543
- topk_weights.data_ptr<float>(),
544
- topk_indices.data_ptr<int>(),
545
- token_expert_indices.data_ptr<int>(),
546
- softmax_workspace.data_ptr<float>(),
547
- num_tokens,
548
- num_experts,
549
- topk,
550
- stream);
551
- }
552
- else if (topk_indices.scalar_type() == at::ScalarType::UInt32)
553
- {
554
- vllm::moe::topkGatingSoftmaxKernelLauncher(
555
- gating_output.data_ptr<float>(),
556
- topk_weights.data_ptr<float>(),
557
- topk_indices.data_ptr<uint32_t>(),
558
- token_expert_indices.data_ptr<int>(),
559
- softmax_workspace.data_ptr<float>(),
560
- num_tokens,
561
- num_experts,
562
- topk,
563
- stream);
564
- }
565
- else {
566
- TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
567
- vllm::moe::topkGatingSoftmaxKernelLauncher(
568
- gating_output.data_ptr<float>(),
569
- topk_weights.data_ptr<float>(),
570
- topk_indices.data_ptr<int64_t>(),
571
- token_expert_indices.data_ptr<int>(),
572
- softmax_workspace.data_ptr<float>(),
573
- num_tokens,
574
- num_experts,
575
- topk,
576
- stream);
692
+ const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float);
693
+ torch::Tensor softmax_workspace = torch::empty({workspace_size}, workspace_options);
694
+
695
+ if (gating_output.scalar_type() == at::ScalarType::Float) {
696
+ dispatch_topk_softmax_launch<float>(gating_output, topk_weights, topk_indices,
697
+ token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, stream);
698
+ } else if (gating_output.scalar_type() == at::ScalarType::Half) {
699
+ dispatch_topk_softmax_launch<__half>(gating_output, topk_weights, topk_indices,
700
+ token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, stream);
701
+ } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
702
+ dispatch_topk_softmax_launch<__nv_bfloat16>(gating_output, topk_weights, topk_indices,
703
+ token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, stream);
704
+ } else {
705
+ TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
706
}
707
}
csrc/moe/torch_bindings.cpp
+1
-1
@@ -5,7 +5,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
5
// Apply topk softmax to the gating outputs.
6
m.def(
7
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
8
- "token_expert_indices, Tensor gating_output) -> ()");
8
+ "token_expert_indices, Tensor gating_output, bool renormalize) -> ()");
9
m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
10
11
// Calculate the result of moe by summing up the partial results
vllm/_custom_ops.py
+2
-1
@@ -1850,9 +1850,10 @@ def topk_softmax(
1850
topk_ids: torch.Tensor,
1851
token_expert_indices: torch.Tensor,
1852
gating_output: torch.Tensor,
1853
+ renormalize: bool = False,
1854
) -> None:
1855
torch.ops._moe_C.topk_softmax(
1855
- topk_weights, topk_ids, token_expert_indices, gating_output
1856
+ topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
1857
)
1858
1859
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
-5
@@ -1074,9 +1074,8 @@ def vllm_topk_softmax(
1074
topk_indices,
1075
token_expert_indices,
1076
gating_output,
1077
+ renormalize,
1078
)
1078
- if renormalize:
1079
- topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1079
1080
return topk_weights, topk_indices
1081
@@ -1113,11 +1112,9 @@ def fused_topk(
1112
M, topk, dtype=torch.int32, device=hidden_states.device
1113
)
1114
1116
- gating_output_float = gating_output.float() # TODO(woosuk): Optimize this.
1117
-
1115
topk_func = dispatch_topk_func()
1116
topk_weights, topk_ids = topk_func(
1120
- topk_weights, topk_ids, token_expert_indices, gating_output_float, renormalize
1117
+ topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
1118
)
1119
1120
return topk_weights, topk_ids, token_expert_indices