@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 63227accf5
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 9, 20:45 UTC

[Kernel] Add topk_sigmoid kernel (#31246)

Signed-off-by: Xin Yang <xyangx@amazon.com>

Xin Yang committed Jan 21, 2026 at 14:49 UTC 63227accf5af941215e5a7fd48c305a15e2c8814
13 files changed +725 -126
benchmarks/kernels/benchmark_fused_topk.py new
+99
@@ -0,0 +1,99 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import itertools
5 +
6 +import torch
7 +
8 +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
9 +from vllm.triton_utils import triton
10 +from vllm.utils.argparse_utils import FlexibleArgumentParser
11 +
12 +num_tokens_range = [2**i for i in range(0, 8, 2)]
13 +num_experts_range = [16, 32, 64, 128, 256, 512]
14 +topk_range = [3, 4]
15 +configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
16 +
17 +
18 +def torch_topk(
19 + gating_output: torch.Tensor,
20 + topk: int,
21 + renormalize: bool,
22 + scoring_func: str = "softmax",
23 +):
24 + if scoring_func == "softmax":
25 + scores = torch.softmax(gating_output.float(), dim=-1)
26 + else:
27 + scores = torch.sigmoid(gating_output.float())
28 + topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
29 +
30 + if renormalize:
31 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
32 +
33 + return topk_weights, topk_ids
34 +
35 +
36 +def get_benchmark(scoring_func):
37 + @triton.testing.perf_report(
38 + triton.testing.Benchmark(
39 + x_names=["num_tokens", "num_experts", "topk"],
40 + x_vals=[list(_) for _ in configs],
41 + line_arg="provider",
42 + line_vals=["torch", "vllm"],
43 + line_names=["Torch", "vLLM"],
44 + styles=[("blue", "-"), ("red", "-")],
45 + ylabel="us",
46 + plot_name=f"fused-topk-perf-{scoring_func}",
47 + args={},
48 + )
49 + )
50 + def benchmark(num_tokens, num_experts, topk, provider):
51 + dtype = torch.bfloat16
52 + hidden_size = 1024
53 + renormalize = True
54 + hidden_states = torch.randn(
55 + (num_tokens, hidden_size), dtype=dtype, device="cuda"
56 + )
57 + gating_output = torch.randn(
58 + (num_tokens, num_experts), dtype=dtype, device="cuda"
59 + )
60 +
61 + quantiles = [0.5, 0.2, 0.8]
62 +
63 + if provider == "torch":
64 + ms, min_ms, max_ms = triton.testing.do_bench(
65 + lambda: torch_topk(
66 + gating_output=gating_output,
67 + topk=topk,
68 + renormalize=renormalize,
69 + scoring_func=scoring_func,
70 + ),
71 + quantiles=quantiles,
72 + )
73 + else:
74 + ms, min_ms, max_ms = triton.testing.do_bench(
75 + lambda: fused_topk(
76 + hidden_states=hidden_states,
77 + gating_output=gating_output,
78 + topk=topk,
79 + renormalize=renormalize,
80 + scoring_func=scoring_func,
81 + ),
82 + quantiles=quantiles,
83 + )
84 +
85 + return 1000 * ms, 1000 * max_ms, 1000 * min_ms
86 +
87 + return benchmark
88 +
89 +
90 +if __name__ == "__main__":
91 + parser = FlexibleArgumentParser(description="Benchmark the MoE topk kernel.")
92 + parser.add_argument("--scoring-func", type=str, default="softmax")
93 + parser.add_argument("--save-path", type=str, default="./configs/fused_topk/")
94 + args = parser.parse_args()
95 +
96 + # Get the benchmark function
97 + benchmark = get_benchmark(args.scoring_func)
98 + # Run performance benchmark
99 + benchmark.run(print_data=True, save_path=args.save_path)
csrc/moe/moe_ops.h
+7 -1
@@ -4,7 +4,13 @@
4
5 void topk_softmax(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
6 torch::Tensor& token_expert_indices,
7 - torch::Tensor& gating_output, bool renormalize);
7 + torch::Tensor& gating_output, bool renormalize,
8 + std::optional<torch::Tensor> bias);
9 +
10 +void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
11 + torch::Tensor& token_expert_indices,
12 + torch::Tensor& gating_output, bool renormalize,
13 + std::optional<torch::Tensor> bias);
14
15 void moe_sum(torch::Tensor& input, torch::Tensor& output);
16
csrc/moe/topk_softmax_kernels.cu
+242 -101
@@ -62,6 +62,12 @@ __device__ __forceinline__ float toFloat(T value) {
62 }
63 }
64
65 +// Scoring function enums
66 +enum ScoringFunc {
67 + SCORING_SOFTMAX = 0, // apply softmax
68 + SCORING_SIGMOID = 1 // apply sigmoid
69 +};
70 +
71 // ====================== Softmax things ===============================
72 // We have our own implementation of softmax here so we can support transposing the output
73 // in the softmax kernel when we extend this module to support expert-choice routing.
@@ -125,6 +131,27 @@ __launch_bounds__(TPB) __global__
131 }
132 }
133
134 +template <int TPB, typename InputType>
135 +__launch_bounds__(TPB) __global__
136 + void moeSigmoid(const InputType* input, const bool* finished, float* output, const int num_cols)
137 +{
138 + const int thread_row_offset = blockIdx.x * num_cols;
139 +
140 + // Don't touch finished rows.
141 + if ((finished != nullptr) && finished[blockIdx.x])
142 + {
143 + return;
144 + }
145 +
146 + for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
147 + {
148 + const int idx = thread_row_offset + ii;
149 + const float val = toFloat(input[idx]);
150 + const float sigmoid_val = 1.0f / (1.0f + __expf(-val));
151 + output[idx] = sigmoid_val;
152 + }
153 +}
154 +
155 template <int TPB, typename IndType>
156 __launch_bounds__(TPB) __global__ void moeTopK(
157 const float* inputs_after_softmax,
@@ -136,7 +163,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
163 const int k,
164 const int start_expert,
165 const int end_expert,
139 - const bool renormalize)
166 + const bool renormalize,
167 + const float* bias)
168 {
169
170 using cub_kvp = cub::KeyValuePair<int, float>;
@@ -162,7 +190,13 @@ __launch_bounds__(TPB) __global__ void moeTopK(
190 {
191 const int idx = thread_read_offset + expert;
192 inp_kvp.key = expert;
165 - inp_kvp.value = inputs_after_softmax[idx];
193 +
194 + // Apply correction bias if provided
195 + if (bias != nullptr) {
196 + inp_kvp.value = inputs_after_softmax[idx] + bias[expert];
197 + } else {
198 + inp_kvp.value = inputs_after_softmax[idx];
199 + }
200
201 for (int prior_k = 0; prior_k < k_idx; ++prior_k)
202 {
@@ -186,12 +220,13 @@ __launch_bounds__(TPB) __global__ void moeTopK(
220 const bool should_process_row = row_is_active && node_uses_expert;
221
222 const int idx = k * block_row + k_idx;
189 - output[idx] = result_kvp.value;
223 + // Return the unbiased scores for output weights
224 + output[idx] = inputs_after_softmax[thread_read_offset + expert];
225 indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
226 assert(indices[idx] >= 0);
227 source_rows[idx] = k_idx * num_rows + block_row;
228 if (renormalize) {
194 - selected_sum += result_kvp.value;
229 + selected_sum += inputs_after_softmax[thread_read_offset + expert];
230 }
231 }
232 __syncthreads();
@@ -225,10 +260,12 @@ __launch_bounds__(TPB) __global__ void moeTopK(
260 2) This implementation assumes k is small, but will work for any k.
261 */
262
228 -template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType, typename InputType = float>
263 +template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType,
264 + typename InputType = float, ScoringFunc SF>
265 __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
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)
266 + void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
267 + int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
268 + const float* bias)
269 {
270 static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
271 std::is_same_v<InputType, __half>,
@@ -353,61 +390,89 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
390 }
391 }
392
356 - // First, we perform a max reduce within the thread. We can do the max in fp16 safely (I think) and just
357 - // convert to float afterwards for the exp + sum reduction.
358 - float thread_max = row_chunk[0];
393 + if constexpr (SF == SCORING_SOFTMAX) {
394 + // First, we perform a max reduce within the thread.
395 + float thread_max = row_chunk[0];
396 #pragma unroll
360 - for (int ii = 1; ii < VPT; ++ii)
361 - {
397 + for (int ii = 1; ii < VPT; ++ii) {
398 thread_max = max(thread_max, row_chunk[ii]);
363 - }
399 + }
400
401 // Now, we find the max within the thread group and distribute among the threads. We use a butterfly reduce.
402 #pragma unroll
367 - for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
368 - {
403 + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
404 + {
405 thread_max = max(thread_max, VLLM_SHFL_XOR_SYNC_WIDTH(thread_max, mask, THREADS_PER_ROW));
370 - }
406 + }
407
372 - // From this point, thread max in all the threads have the max within the row.
373 - // Now, we subtract the max from each element in the thread and take the exp. We also compute the thread local sum.
374 - float row_sum = 0;
408 + // From this point, thread max in all the threads have the max within the row.
409 + // Now, we subtract the max from each element in the thread and take the exp. We also compute the thread local sum.
410 + float row_sum = 0;
411 #pragma unroll
376 - for (int ii = 0; ii < VPT; ++ii)
377 - {
412 + for (int ii = 0; ii < VPT; ++ii)
413 + {
414 row_chunk[ii] = expf(row_chunk[ii] - thread_max);
415 row_sum += row_chunk[ii];
380 - }
416 + }
417
418 // Now, we perform the sum reduce within each thread group. Similar to the max reduce, we use a bufferfly pattern.
419 #pragma unroll
384 - for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
385 - {
420 + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
421 + {
422 row_sum += VLLM_SHFL_XOR_SYNC_WIDTH(row_sum, mask, THREADS_PER_ROW);
387 - }
423 + }
424
389 - // From this point, all threads have the max and the sum for their rows in the thread_max and thread_sum variables
390 - // respectively. Finally, we can scale the rows for the softmax. Technically, for top-k gating we don't need to
391 - // compute the entire softmax row. We can likely look at the maxes and only compute for the top-k values in the row.
392 - // However, this kernel will likely not be a bottle neck and it seems better to closer match torch and find the
393 - // argmax after computing the softmax.
394 - const float reciprocal_row_sum = 1.f / row_sum;
425 + // From this point, all threads have the max and the sum for their rows in the thread_max and thread_sum variables
426 + // respectively. Finally, we can scale the rows for the softmax. Technically, for top-k gating we don't need to
427 + // compute the entire softmax row. We can likely look at the maxes and only compute for the top-k values in the row.
428 + // However, this kernel will likely not be a bottle neck and it seems better to closer match torch and find the
429 + // argmax after computing the softmax.
430 + const float reciprocal_row_sum = 1.f / row_sum;
431
432 #pragma unroll
397 - for (int ii = 0; ii < VPT; ++ii)
398 - {
433 + for (int ii = 0; ii < VPT; ++ii)
434 + {
435 row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
436 + }
437 + } else if constexpr (SF == SCORING_SIGMOID) {
438 +#pragma unroll
439 + for (int ii = 0; ii < VPT; ++ii)
440 + {
441 + row_chunk[ii] = 1.0f / (1.0f + __expf(-row_chunk[ii]));
442 + }
443 + }
444 +
445 + static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
446 +
447 + // If bias is not null, use biased value for selection
448 + float row_chunk_for_choice[VPT];
449 + // Apply correction bias
450 + if (bias != nullptr) {
451 +#pragma unroll
452 + for (int ldg = 0; ldg < LDG_PER_THREAD; ++ldg) {
453 +#pragma unroll
454 + for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
455 + const int expert = first_elt_read_by_thread + ldg * COLS_PER_GROUP_LDG + ii;
456 + float bias_val = expert < NUM_EXPERTS ? bias[expert] : 0.0f;
457 + row_chunk_for_choice[ldg * ELTS_PER_LDG + ii] = row_chunk[ldg * ELTS_PER_LDG + ii] + bias_val;
458 + }
459 + }
460 + } else {
461 +#pragma unroll
462 + for (int ii = 0; ii < VPT; ++ii) {
463 + row_chunk_for_choice[ii] = row_chunk[ii];
464 + }
465 }
466
402 - // Now, softmax_res contains the softmax of the row chunk. Now, I want to find the topk elements in each row, along
467 + // Now, row_chunk contains the softmax / sigmoid of the row chunk. Now, I want to find the topk elements in each row, along
468 // with the max index.
469 int start_col = first_elt_read_by_thread;
405 - static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
470
471 float selected_sum = 0.f;
472 for (int k_idx = 0; k_idx < k; ++k_idx)
473 {
474 // First, each thread does the local argmax
475 + float max_val_for_choice = row_chunk_for_choice[0];
476 float max_val = row_chunk[0];
477 int expert = start_col;
478 #pragma unroll
@@ -416,12 +481,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
481 #pragma unroll
482 for (int ii = 0; ii < ELTS_PER_LDG; ++ii)
483 {
484 + float val_for_choice = row_chunk_for_choice[ldg * ELTS_PER_LDG + ii];
485 float val = row_chunk[ldg * ELTS_PER_LDG + ii];
486
487 // No check on the experts here since columns with the smallest index are processed first and only
488 // updated if > (not >=)
423 - if (val > max_val)
489 + if (val_for_choice > max_val_for_choice)
490 {
491 + max_val_for_choice = val_for_choice;
492 max_val = val;
493 expert = col + ii;
494 }
@@ -434,12 +501,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
501 #pragma unroll
502 for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
503 {
504 + float other_max_for_choice = VLLM_SHFL_XOR_SYNC_WIDTH(max_val_for_choice, mask, THREADS_PER_ROW);
505 float other_max = VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
506 int other_expert = VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
507
508 // We want lower indices to "win" in every thread so we break ties this way
441 - if (other_max > max_val || (other_max == max_val && other_expert < expert))
509 + if (other_max_for_choice > max_val_for_choice || (other_max_for_choice == max_val_for_choice && other_expert < expert))
510 {
511 + max_val_for_choice = other_max_for_choice;
512 max_val = other_max;
513 expert = other_expert;
514 }
@@ -474,7 +543,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
543 {
544 const int offset_for_expert = expert % ELTS_PER_LDG;
545 // Safe to set to any negative value since row_chunk values must be between 0 and 1.
477 - row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
546 + row_chunk_for_choice[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
547 }
548 }
549 }
@@ -508,10 +577,10 @@ struct TopkConstants
577 };
578 } // namespace detail
579
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,
580 +template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
581 +void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
582 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)
583 + const float* bias, cudaStream_t stream)
584 {
585 static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
586 using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
@@ -521,43 +590,51 @@ void topkGatingSoftmaxLauncherHelper(const InputType* input, const bool* finishe
590 const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
591
592 dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
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);
593 + topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
594 + input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias);
595 }
596
597 #ifndef USE_ROCM
529 -#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
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>( \
533 - gating_output, nullptr, topk_weights, topk_indices, token_expert_indices, \
534 - num_tokens, topk, 0, num_experts, renormalize, stream);
598 + #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
599 + static_assert(WARP_SIZE == 32, \
600 + "Unsupported warp size. Only 32 is supported for CUDA"); \
601 + topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, MAX_BYTES, \
602 + IndType, InputType, SF>( \
603 + gating_output, nullptr, topk_weights, topk_indices, \
604 + token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
605 + bias, stream);
606 #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>( \
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>( \
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"); \
607 + #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
608 + if (WARP_SIZE == 64) { \
609 + topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, MAX_BYTES, \
610 + IndType, InputType, SF>( \
611 + gating_output, nullptr, topk_weights, topk_indices, \
612 + token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
613 + bias, stream); \
614 + } else if (WARP_SIZE == 32) { \
615 + topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
616 + IndType, InputType, SF>( \
617 + gating_output, nullptr, topk_weights, topk_indices, \
618 + token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
619 + bias, stream); \
620 + } else { \
621 + assert(false && \
622 + "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
623 }
624 #endif
625
550 -template <typename IndType, typename InputType>
551 -void topkGatingSoftmaxKernelLauncher(
626 +template <typename IndType, typename InputType, ScoringFunc SF>
627 +void topkGatingKernelLauncher(
628 const InputType* gating_output,
629 float* topk_weights,
630 IndType* topk_indices,
631 int* token_expert_indices,
556 - float* softmax_workspace,
632 + float* workspace,
633 const int num_tokens,
634 const int num_experts,
635 const int topk,
636 const bool renormalize,
637 + const float* bias,
638 cudaStream_t stream) {
639 static constexpr int WARPS_PER_TB = 4;
640 static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
@@ -569,64 +646,71 @@ void topkGatingSoftmaxKernelLauncher(
646 #endif
647 switch (num_experts) {
648 case 1:
572 - LAUNCH_SOFTMAX(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
649 + LAUNCH_TOPK(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
650 break;
651 case 2:
575 - LAUNCH_SOFTMAX(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
652 + LAUNCH_TOPK(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
653 break;
654 case 4:
578 - LAUNCH_SOFTMAX(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
655 + LAUNCH_TOPK(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
656 break;
657 case 8:
581 - LAUNCH_SOFTMAX(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
658 + LAUNCH_TOPK(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
659 break;
660 case 16:
584 - LAUNCH_SOFTMAX(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
661 + LAUNCH_TOPK(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
662 break;
663 case 32:
587 - LAUNCH_SOFTMAX(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
664 + LAUNCH_TOPK(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
665 break;
666 case 64:
590 - LAUNCH_SOFTMAX(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
667 + LAUNCH_TOPK(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
668 break;
669 case 128:
593 - LAUNCH_SOFTMAX(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
670 + LAUNCH_TOPK(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
671 break;
672 case 256:
596 - LAUNCH_SOFTMAX(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
673 + LAUNCH_TOPK(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
674 break;
675 case 512:
599 - LAUNCH_SOFTMAX(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
676 + LAUNCH_TOPK(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
677 break;
678 // (CUDA only) support multiples of 64 when num_experts is not power of 2.
679 // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of num_experts,
680 // alternatively we can test 4 bytes loading and enable it in future.
681 #ifndef USE_ROCM
682 case 192:
606 - LAUNCH_SOFTMAX(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
683 + LAUNCH_TOPK(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
684 break;
685 case 320:
609 - LAUNCH_SOFTMAX(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
686 + LAUNCH_TOPK(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
687 break;
688 case 384:
612 - LAUNCH_SOFTMAX(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
689 + LAUNCH_TOPK(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
690 break;
691 case 448:
615 - LAUNCH_SOFTMAX(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
692 + LAUNCH_TOPK(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
693 break;
694 case 576:
618 - LAUNCH_SOFTMAX(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
695 + LAUNCH_TOPK(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
696 break;
697 #endif
698 default: {
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.");
699 + TORCH_CHECK(workspace != nullptr,
700 + "workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
701 static constexpr int TPB = 256;
625 - moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
626 - gating_output, nullptr, softmax_workspace, num_experts);
702 + if constexpr (SF == SCORING_SOFTMAX) {
703 + moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
704 + gating_output, nullptr, workspace, num_experts);
705 + } else if constexpr (SF == SCORING_SIGMOID) {
706 + moeSigmoid<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
707 + gating_output, nullptr, workspace, num_experts);
708 + } else {
709 + TORCH_CHECK(false, "Unsupported scoring func");
710 + }
711 moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
628 - softmax_workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
629 - num_experts, topk, 0, num_experts, renormalize);
712 + workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
713 + num_experts, topk, 0, num_experts, renormalize, bias);
714 }
715 }
716 }
@@ -635,40 +719,55 @@ void topkGatingSoftmaxKernelLauncher(
719 } // namespace vllm
720
721
638 -template<typename ComputeType>
639 -void dispatch_topk_softmax_launch(
722 +template<typename ComputeType, vllm::moe::ScoringFunc SF>
723 +void dispatch_topk_launch(
724 torch::Tensor& gating_output,
725 torch::Tensor& topk_weights,
726 torch::Tensor& topk_indices,
727 torch::Tensor& token_expert_indices,
728 torch::Tensor& softmax_workspace,
645 - int num_tokens, int num_experts, int topk, bool renormalize, cudaStream_t stream)
646 -{
729 + int num_tokens, int num_experts, int topk, bool renormalize,
730 + std::optional<torch::Tensor> bias,
731 + cudaStream_t stream)
732 + {
733 + const float* bias_ptr = nullptr;
734 + if (bias.has_value()) {
735 + const torch::Tensor& bias_tensor = bias.value();
736 + TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "bias tensor must be float32");
737 + TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D");
738 + TORCH_CHECK(bias_tensor.size(0) == num_experts, "bias size mismatch, expected: ", num_experts);
739 + TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous");
740 + bias_ptr = bias_tensor.data_ptr<float>();
741 + }
742 +
743 if (topk_indices.scalar_type() == at::ScalarType::Int) {
648 - vllm::moe::topkGatingSoftmaxKernelLauncher<int, ComputeType>(
744 + vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
745 reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
746 topk_weights.data_ptr<float>(),
747 topk_indices.data_ptr<int>(),
748 token_expert_indices.data_ptr<int>(),
749 softmax_workspace.data_ptr<float>(),
654 - num_tokens, num_experts, topk, renormalize, stream);
750 + num_tokens, num_experts, topk, renormalize,
751 + bias_ptr, stream);
752 } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
656 - vllm::moe::topkGatingSoftmaxKernelLauncher<uint32_t, ComputeType>(
753 + vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
754 reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
755 topk_weights.data_ptr<float>(),
756 topk_indices.data_ptr<uint32_t>(),
757 token_expert_indices.data_ptr<int>(),
758 softmax_workspace.data_ptr<float>(),
662 - num_tokens, num_experts, topk, renormalize, stream);
759 + num_tokens, num_experts, topk, renormalize,
760 + bias_ptr, stream);
761 } else {
762 TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
665 - vllm::moe::topkGatingSoftmaxKernelLauncher<int64_t, ComputeType>(
763 + vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
764 reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
765 topk_weights.data_ptr<float>(),
766 topk_indices.data_ptr<int64_t>(),
767 token_expert_indices.data_ptr<int>(),
768 softmax_workspace.data_ptr<float>(),
671 - num_tokens, num_experts, topk, renormalize, stream);
769 + num_tokens, num_experts, topk, renormalize,
770 + bias_ptr, stream);
771 }
772 }
773
@@ -677,7 +776,8 @@ void topk_softmax(
776 torch::Tensor& topk_indices, // [num_tokens, topk]
777 torch::Tensor& token_expert_indices, // [num_tokens, topk]
778 torch::Tensor& gating_output, // [num_tokens, num_experts]
680 - bool renormalize)
779 + bool renormalize,
780 + std::optional<torch::Tensor> bias)
781 {
782 const int num_experts = gating_output.size(-1);
783 const auto num_tokens = gating_output.numel() / num_experts;
@@ -693,14 +793,55 @@ void topk_softmax(
793 torch::Tensor softmax_workspace = torch::empty({workspace_size}, workspace_options);
794
795 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);
796 + dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
797 + token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
798 + bias, stream);
799 + } else if (gating_output.scalar_type() == at::ScalarType::Half) {
800 + dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
801 + token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
802 + bias, stream);
803 + } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
804 + dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
805 + token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
806 + bias, stream);
807 + } else {
808 + TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
809 + }
810 +}
811 +
812 +void topk_sigmoid(
813 + torch::Tensor& topk_weights, // [num_tokens, topk]
814 + torch::Tensor& topk_indices, // [num_tokens, topk]
815 + torch::Tensor& token_expert_indices, // [num_tokens, topk]
816 + torch::Tensor& gating_output, // [num_tokens, num_experts]
817 + bool renormalize,
818 + std::optional<torch::Tensor> bias)
819 +{
820 + const int num_experts = gating_output.size(-1);
821 + const auto num_tokens = gating_output.numel() / num_experts;
822 + const int topk = topk_weights.size(-1);
823 +
824 + const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
825 + const bool needs_workspace = !is_pow_2 || num_experts > 256;
826 + const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
827 +
828 + const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
829 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
830 + const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float);
831 + torch::Tensor workspace = torch::empty({workspace_size}, workspace_options);
832 +
833 + if (gating_output.scalar_type() == at::ScalarType::Float) {
834 + dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
835 + token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
836 + bias, stream);
837 } 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);
838 + dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
839 + token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
840 + bias, stream);
841 } 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);
842 + dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
843 + token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
844 + bias, stream);
845 } else {
846 TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
847 }
csrc/moe/torch_bindings.cpp
+9 -1
@@ -5,9 +5,17 @@ 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, bool renormalize) -> ()");
8 + "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
9 + "bias) -> ()");
10 m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
11
12 + // Apply topk sigmoid to the gating outputs.
13 + m.def(
14 + "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
15 + "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
16 + "bias) -> ()");
17 + m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
18 +
19 // Calculate the result of moe by summing up the partial results
20 // from all selected experts.
21 m.def("moe_sum(Tensor input, Tensor! output) -> ()");
tests/kernels/moe/test_fused_topk.py new
+137
@@ -0,0 +1,137 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the MoE fused topk kernel
4 +
5 +Run `pytest tests/kernels/moe/test_fused_topk.py`.
6 +"""
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
12 + fused_topk_bias,
13 +)
14 +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
15 +from vllm.platforms import current_platform
16 +
17 +
18 +def torch_topk(
19 + gating_output: torch.Tensor,
20 + topk: int,
21 + renormalize: bool,
22 + e_score_correction_bias: torch.Tensor = None,
23 + scoring_func: str = "softmax",
24 +):
25 + if scoring_func == "softmax":
26 + scores = torch.softmax(gating_output.float(), dim=-1)
27 + else:
28 + assert scoring_func == "sigmoid"
29 + scores = torch.sigmoid(gating_output.float())
30 +
31 + if e_score_correction_bias is not None:
32 + num_experts = gating_output.shape[-1]
33 + scores_for_choice = scores.view(
34 + -1, num_experts
35 + ) + e_score_correction_bias.unsqueeze(0)
36 + _, topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1)
37 + topk_weights = scores.gather(1, topk_ids)
38 + else:
39 + topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
40 +
41 + if renormalize:
42 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
43 +
44 + return topk_weights, topk_ids
45 +
46 +
47 +@pytest.mark.skipif(
48 + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
49 +)
50 +@pytest.mark.parametrize("num_tokens", [1, 33, 56])
51 +@pytest.mark.parametrize("hidden_size", [1024, 2048])
52 +@pytest.mark.parametrize("num_experts", [6, 16])
53 +@pytest.mark.parametrize("topk", [3, 4])
54 +@pytest.mark.parametrize("renormalize", [True, False])
55 +@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
56 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
57 +def test_fused_topk(
58 + num_tokens: int,
59 + hidden_size: int,
60 + num_experts: int,
61 + topk: int,
62 + renormalize: bool,
63 + scoring_func: str,
64 + dtype: torch.dtype,
65 +):
66 + torch.manual_seed(0)
67 + hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
68 + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
69 +
70 + topk_weights_ref, topk_ids_ref = torch_topk(
71 + gating_output=gating_output,
72 + topk=topk,
73 + renormalize=renormalize,
74 + scoring_func=scoring_func,
75 + )
76 +
77 + topk_weights, topk_ids, _ = fused_topk(
78 + hidden_states=hidden_states,
79 + gating_output=gating_output,
80 + topk=topk,
81 + renormalize=renormalize,
82 + scoring_func=scoring_func,
83 + )
84 +
85 + torch.testing.assert_close(
86 + topk_weights_ref.to(torch.float32), topk_weights, atol=1e-2, rtol=1e-2
87 + )
88 + torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0)
89 +
90 +
91 +@pytest.mark.skipif(
92 + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
93 +)
94 +@pytest.mark.parametrize("num_tokens", [1, 33, 56])
95 +@pytest.mark.parametrize("hidden_size", [1024, 2048])
96 +@pytest.mark.parametrize("num_experts", [6, 16])
97 +@pytest.mark.parametrize("topk", [3, 4])
98 +@pytest.mark.parametrize("renormalize", [True, False])
99 +@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
100 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
101 +def test_fused_topk_bias(
102 + num_tokens: int,
103 + hidden_size: int,
104 + num_experts: int,
105 + topk: int,
106 + renormalize: bool,
107 + scoring_func: str,
108 + dtype: torch.dtype,
109 +):
110 + torch.manual_seed(0)
111 + hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
112 + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
113 + e_score_correction_bias = torch.randn(
114 + (num_experts,), dtype=torch.float32, device="cuda"
115 + )
116 +
117 + topk_weights_ref, topk_ids_ref = torch_topk(
118 + gating_output=gating_output,
119 + topk=topk,
120 + renormalize=renormalize,
121 + e_score_correction_bias=e_score_correction_bias,
122 + scoring_func=scoring_func,
123 + )
124 +
125 + topk_weights, topk_ids = fused_topk_bias(
126 + hidden_states=hidden_states,
127 + gating_output=gating_output,
128 + e_score_correction_bias=e_score_correction_bias,
129 + topk=topk,
130 + renormalize=renormalize,
131 + scoring_func=scoring_func,
132 + )
133 +
134 + torch.testing.assert_close(
135 + topk_weights_ref.to(torch.float32), topk_weights, atol=1e-2, rtol=1e-2
136 + )
137 + torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0)
tests/model_executor/test_enabled_custom_ops.py
+17 -3
@@ -18,7 +18,9 @@ from vllm.model_executor.layers.activation import (
18 SiluAndMul,
19 )
20 from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
21 - dispatch_topk_func,
21 + dispatch_topk_sigmoid_func,
22 + dispatch_topk_softmax_func,
23 + vllm_topk_sigmoid,
24 vllm_topk_softmax,
25 )
26 from vllm.model_executor.layers.layernorm import (
@@ -133,8 +135,8 @@ def test_enabled_ops_invalid(env: str):
135 @pytest.mark.parametrize(
136 "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
137 )
136 -def test_topk_dispatch(use_rocm_aiter: bool):
137 - topk_func = dispatch_topk_func(use_rocm_aiter)
138 +def test_topk_softmax_dispatch(use_rocm_aiter: bool):
139 + topk_func = dispatch_topk_softmax_func(use_rocm_aiter)
140
141 if current_platform.is_rocm() and use_rocm_aiter:
142 assert topk_func == rocm_aiter_ops.topk_softmax
@@ -142,6 +144,18 @@ def test_topk_dispatch(use_rocm_aiter: bool):
144 assert topk_func == vllm_topk_softmax
145
146
147 +@pytest.mark.parametrize(
148 + "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
149 +)
150 +def test_topk_sigmoid_dispatch(use_rocm_aiter: bool):
151 + topk_func = dispatch_topk_sigmoid_func(use_rocm_aiter)
152 +
153 + if current_platform.is_rocm() and use_rocm_aiter:
154 + assert topk_func == rocm_aiter_ops.topk_sigmoid
155 + else:
156 + assert topk_func == vllm_topk_sigmoid
157 +
158 +
159 @pytest.mark.parametrize("add_residual", [True, False])
160 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
161 @pytest.mark.parametrize("use_rocm_aiter", [True, False])
vllm/_aiter_ops.py
+39
@@ -200,6 +200,24 @@ def _rocm_aiter_topk_softmax_fake(
200 pass
201
202
203 +def _rocm_aiter_topk_sigmoid_impl(
204 + topk_weights: torch.Tensor,
205 + topk_indices: torch.Tensor,
206 + gating_output: torch.Tensor,
207 +) -> None:
208 + from aiter import topk_sigmoid
209 +
210 + topk_sigmoid(topk_weights, topk_indices, gating_output)
211 +
212 +
213 +def _rocm_aiter_topk_sigmoid_fake(
214 + topk_weights: torch.Tensor,
215 + topk_indices: torch.Tensor,
216 + gating_output: torch.Tensor,
217 +) -> None:
218 + pass
219 +
220 +
221 def _rocm_aiter_biased_grouped_topk_impl(
222 gating_output: torch.Tensor,
223 correction_bias: torch.Tensor,
@@ -985,6 +1003,14 @@ class rocm_aiter_ops:
1003 dispatch_key=current_platform.dispatch_key,
1004 )
1005
1006 + direct_register_custom_op(
1007 + op_name="rocm_aiter_topk_sigmoid",
1008 + op_func=_rocm_aiter_topk_sigmoid_impl,
1009 + mutates_args=["topk_weights", "topk_indices"],
1010 + fake_impl=_rocm_aiter_topk_sigmoid_fake,
1011 + dispatch_key=current_platform.dispatch_key,
1012 + )
1013 +
1014 direct_register_custom_op(
1015 op_name="rocm_aiter_biased_grouped_topk",
1016 op_func=_rocm_aiter_biased_grouped_topk_impl,
@@ -1272,6 +1298,19 @@ class rocm_aiter_ops:
1298 )
1299 return topk_weights, topk_indices
1300
1301 + @staticmethod
1302 + def topk_sigmoid(
1303 + topk_weights: torch.Tensor,
1304 + topk_indices: torch.Tensor,
1305 + token_expert_indices: torch.Tensor,
1306 + gating_output: torch.Tensor,
1307 + renormalize: bool,
1308 + ) -> tuple[torch.Tensor, ...]:
1309 + torch.ops.vllm.rocm_aiter_topk_sigmoid(
1310 + topk_weights, topk_indices, gating_output
1311 + )
1312 + return topk_weights, topk_indices
1313 +
1314 @staticmethod
1315 def biased_grouped_topk(
1316 gating_output: torch.Tensor,
vllm/_custom_ops.py
+25 -1
@@ -2177,9 +2177,33 @@ def topk_softmax(
2177 token_expert_indices: torch.Tensor,
2178 gating_output: torch.Tensor,
2179 renormalize: bool = False,
2180 + e_score_correction_bias: torch.Tensor | None = None,
2181 ) -> None:
2182 torch.ops._moe_C.topk_softmax(
2182 - topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
2183 + topk_weights,
2184 + topk_ids,
2185 + token_expert_indices,
2186 + gating_output,
2187 + renormalize,
2188 + e_score_correction_bias,
2189 + )
2190 +
2191 +
2192 +def topk_sigmoid(
2193 + topk_weights: torch.Tensor,
2194 + topk_ids: torch.Tensor,
2195 + token_expert_indices: torch.Tensor,
2196 + gating_output: torch.Tensor,
2197 + renormalize: bool = False,
2198 + e_score_correction_bias: torch.Tensor | None = None,
2199 +) -> None:
2200 + torch.ops._moe_C.topk_sigmoid(
2201 + topk_weights,
2202 + topk_ids,
2203 + token_expert_indices,
2204 + gating_output,
2205 + renormalize,
2206 + e_score_correction_bias,
2207 )
2208
2209
vllm/model_executor/layers/fused_moe/config.py
+2 -2
@@ -106,14 +106,14 @@ def _quant_flags_to_group_shape(
106 class RoutingMethodType(IntEnum):
107 # Default: Softmax -> TopK
108 Default = (0,)
109 - # Renormalize: TopK -> Softmax
109 + # Renormalize: TopK -> Softmax/Sigmoid
110 Renormalize = (1,)
111 # DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups
112 # -> Top8 experts from the Top4 groups
113 DeepSeekV3 = (2,)
114 # Llama4: Top1 -> Sigmoid
115 Llama4 = (3,)
116 - # RenormalizeNaive: Softmax -> TopK -> Renormalize
116 + # RenormalizeNaive: Softmax/Sigmoid -> TopK -> Renormalize
117 RenormalizeNaive = (4,)
118 # TopK: TopK (no softmax)
119 TopK = (5,)
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+97 -1
@@ -4,6 +4,8 @@ from collections.abc import Callable
4
5 import torch
6
7 +import vllm._custom_ops as ops
8 +from vllm._aiter_ops import rocm_aiter_ops
9 from vllm.distributed.eplb.eplb_state import EplbLayerState
10 from vllm.model_executor.layers.batch_invariant import (
11 vllm_is_batch_invariant,
@@ -12,15 +14,106 @@ from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
14 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
15
16
17 +def vllm_topk_softmax(
18 + topk_weights: torch.Tensor,
19 + topk_indices: torch.Tensor,
20 + token_expert_indices: torch.Tensor,
21 + gating_output: torch.Tensor,
22 + renormalize: bool = False,
23 + e_score_correction_bias: torch.Tensor | None = None,
24 +) -> tuple[torch.Tensor, ...]:
25 + ops.topk_softmax(
26 + topk_weights,
27 + topk_indices,
28 + token_expert_indices,
29 + gating_output,
30 + renormalize,
31 + e_score_correction_bias,
32 + )
33 +
34 + return topk_weights, topk_indices
35 +
36 +
37 +def vllm_topk_sigmoid(
38 + topk_weights: torch.Tensor,
39 + topk_indices: torch.Tensor,
40 + token_expert_indices: torch.Tensor,
41 + gating_output: torch.Tensor,
42 + renormalize: bool = False,
43 + e_score_correction_bias: torch.Tensor | None = None,
44 +) -> tuple[torch.Tensor, ...]:
45 + ops.topk_sigmoid(
46 + topk_weights,
47 + topk_indices,
48 + token_expert_indices,
49 + gating_output,
50 + renormalize,
51 + e_score_correction_bias,
52 + )
53 +
54 + return topk_weights, topk_indices
55 +
56 +
57 def fused_topk_bias(
58 hidden_states: torch.Tensor,
59 gating_output: torch.Tensor,
60 e_score_correction_bias: torch.Tensor,
61 topk: int,
62 renormalize: bool,
63 + scoring_func: str = "softmax",
64 + indices_type: torch.dtype | None = None,
65 ):
66 + if not rocm_aiter_ops.is_fused_moe_enabled():
67 + assert hidden_states.size(0) == gating_output.size(0), (
68 + "Number of tokens mismatch"
69 + )
70 +
71 + M, _ = hidden_states.size()
72 +
73 + topk_weights = torch.empty(
74 + M, topk, dtype=torch.float32, device=hidden_states.device
75 + )
76 + topk_ids = torch.empty(
77 + M,
78 + topk,
79 + dtype=torch.int32 if indices_type is None else indices_type,
80 + device=hidden_states.device,
81 + )
82 + token_expert_indices = torch.empty(
83 + M, topk, dtype=torch.int32, device=hidden_states.device
84 + )
85 +
86 + if scoring_func == "softmax":
87 + topk_weights, topk_ids = vllm_topk_softmax(
88 + topk_weights,
89 + topk_ids,
90 + token_expert_indices,
91 + gating_output,
92 + renormalize,
93 + e_score_correction_bias,
94 + )
95 + return topk_weights, topk_ids
96 + elif scoring_func == "sigmoid":
97 + topk_weights, topk_ids = vllm_topk_sigmoid(
98 + topk_weights,
99 + topk_ids,
100 + token_expert_indices,
101 + gating_output,
102 + renormalize,
103 + e_score_correction_bias,
104 + )
105 + return topk_weights, topk_ids
106 + else:
107 + raise ValueError(f"Unsupported scoring function: {scoring_func}")
108 +
109 n_routed_experts = gating_output.shape[-1]
23 - scores = gating_output.softmax(dim=-1)
110 + if scoring_func == "softmax":
111 + scores = gating_output.softmax(dim=-1)
112 + elif scoring_func == "sigmoid":
113 + scores = gating_output.sigmoid()
114 + else:
115 + raise ValueError(f"Unsupported scoring function: {scoring_func}")
116 +
117 scores_for_choice = scores.view(
118 -1, n_routed_experts
119 ) + e_score_correction_bias.unsqueeze(0)
@@ -43,6 +136,7 @@ class FusedTopKBiasRouter(BaseRouter):
136 global_num_experts: int,
137 eplb_state: EplbLayerState,
138 e_score_correction_bias: torch.Tensor,
139 + scoring_func: str,
140 renormalize: bool = True,
141 routed_scaling_factor: float = 1.0,
142 enable_eplb: bool = False,
@@ -57,6 +151,7 @@ class FusedTopKBiasRouter(BaseRouter):
151 )
152 self.e_score_correction_bias = e_score_correction_bias
153 self.renormalize = renormalize
154 + self.scoring_func = scoring_func
155 self.routed_scaling_factor = routed_scaling_factor
156
157 @property
@@ -80,6 +175,7 @@ class FusedTopKBiasRouter(BaseRouter):
175 e_score_correction_bias=self.e_score_correction_bias.data,
176 topk=self.top_k,
177 renormalize=self.renormalize,
178 + scoring_func=self.scoring_func,
179 )
180
181 if self.routed_scaling_factor != 1.0:
vllm/model_executor/layers/fused_moe/router/fused_topk_router.py
+50 -8
@@ -16,7 +16,7 @@ def vllm_topk_softmax(
16 topk_indices: torch.Tensor,
17 token_expert_indices: torch.Tensor,
18 gating_output: torch.Tensor,
19 - renormalize: bool,
19 + renormalize: bool = False,
20 ) -> tuple[torch.Tensor, ...]:
21 ops.topk_softmax(
22 topk_weights,
@@ -29,7 +29,25 @@ def vllm_topk_softmax(
29 return topk_weights, topk_indices
30
31
32 -def dispatch_topk_func(
32 +def vllm_topk_sigmoid(
33 + topk_weights: torch.Tensor,
34 + topk_indices: torch.Tensor,
35 + token_expert_indices: torch.Tensor,
36 + gating_output: torch.Tensor,
37 + renormalize: bool = False,
38 +) -> tuple[torch.Tensor, ...]:
39 + ops.topk_sigmoid(
40 + topk_weights,
41 + topk_indices,
42 + token_expert_indices,
43 + gating_output,
44 + renormalize,
45 + )
46 +
47 + return topk_weights, topk_indices
48 +
49 +
50 +def dispatch_topk_softmax_func(
51 use_rocm_aiter: bool = False,
52 ) -> Callable[..., tuple[torch.Tensor, ...]]:
53 if use_rocm_aiter:
@@ -37,12 +55,21 @@ def dispatch_topk_func(
55 return vllm_topk_softmax
56
57
58 +def dispatch_topk_sigmoid_func(
59 + use_rocm_aiter: bool = False,
60 +) -> Callable[..., tuple[torch.Tensor, ...]]:
61 + if use_rocm_aiter:
62 + return rocm_aiter_ops.topk_sigmoid
63 + return vllm_topk_sigmoid
64 +
65 +
66 def fused_topk(
67 hidden_states: torch.Tensor,
68 gating_output: torch.Tensor,
69 topk: int,
70 renormalize: bool,
71 indices_type: torch.dtype | None = None,
72 + scoring_func: str = "softmax",
73 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
74 assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
75
@@ -61,12 +88,26 @@ def fused_topk(
88 M, topk, dtype=torch.int32, device=hidden_states.device
89 )
90
64 - topk_func = dispatch_topk_func(use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled())
65 - topk_weights, topk_ids = topk_func(
66 - topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
67 - )
91 + if scoring_func == "softmax":
92 + topk_func = dispatch_topk_softmax_func(
93 + use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled()
94 + )
95 + topk_weights, topk_ids = topk_func(
96 + topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
97 + )
98 +
99 + return topk_weights, topk_ids, token_expert_indices
100 + elif scoring_func == "sigmoid":
101 + topk_func = dispatch_topk_sigmoid_func(
102 + use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled()
103 + )
104 + topk_weights, topk_ids = topk_func(
105 + topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
106 + )
107
69 - return topk_weights, topk_ids, token_expert_indices
108 + return topk_weights, topk_ids, token_expert_indices
109 + else:
110 + raise ValueError(f"Unsupported scoring function: {scoring_func}")
111
112
113 class FusedTopKRouter(BaseRouter):
@@ -82,7 +123,6 @@ class FusedTopKRouter(BaseRouter):
123 enable_eplb: bool = False,
124 indices_type_getter: Callable[[], torch.dtype | None] | None = None,
125 ):
85 - assert scoring_func == "softmax", "FusedTopKRouter only supports softmax."
126 super().__init__(
127 top_k=top_k,
128 global_num_experts=global_num_experts,
@@ -91,6 +131,7 @@ class FusedTopKRouter(BaseRouter):
131 indices_type_getter=indices_type_getter,
132 )
133 self.renormalize = renormalize
134 + self.scoring_func = scoring_func
135
136 @property
137 def routing_method_type(self) -> RoutingMethodType:
@@ -113,6 +154,7 @@ class FusedTopKRouter(BaseRouter):
154 topk=self.top_k,
155 renormalize=self.renormalize,
156 indices_type=indices_type,
157 + scoring_func=self.scoring_func,
158 )
159
160 return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/router_factory.py
+1 -5
@@ -143,17 +143,13 @@ def create_fused_moe_router(
143 router.capture = capture
144 return router
145
146 - if scoring_func != "softmax":
147 - raise ValueError(
148 - "Only softmax scoring function is supported for non-grouped topk."
149 - )
150 -
146 if e_score_correction_bias is not None:
147 router = FusedTopKBiasRouter(
148 top_k=top_k,
149 global_num_experts=global_num_experts,
150 eplb_state=eplb_state,
151 e_score_correction_bias=e_score_correction_bias,
152 + scoring_func=scoring_func,
153 renormalize=renormalize,
154 routed_scaling_factor=routed_scaling_factor,
155 enable_eplb=enable_eplb,
vllm/model_executor/models/minimax_m2.py
-3
@@ -100,9 +100,6 @@ class MiniMaxM2MoE(nn.Module):
100 num_experts=config.num_local_experts,
101 top_k=config.num_experts_per_tok,
102 scoring_func=config.scoring_func,
103 - use_grouped_topk=True,
104 - num_expert_group=1,
105 - topk_group=1,
103 e_score_correction_bias=self.e_score_correction_bias,
104 hidden_size=config.hidden_size,
105 intermediate_size=config.intermediate_size,