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 6, 15:45 UTC
[Compile] Fix compile warning in `moe_permute` (#36529)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 16, 2026 at 10:16 UTC
e855d380fa59614167362a94e87a21a91f3ab470
3 files changed
+12
-14
csrc/moe/moe_permute_unpermute_op.cu
+3
-4
@@ -73,10 +73,9 @@ void moe_permute(
73
MOE_DISPATCH(input.scalar_type(), [&] {
74
expandInputRowsKernelLauncher<scalar_t>(
75
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
76
- get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
77
- get_ptr<int>(inv_permuted_idx), get_ptr<int>(permuted_idx),
78
- get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
79
- n_hidden, topk, n_local_expert, stream);
76
+ get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
77
+ get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
78
+ n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
79
});
80
}
81
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+1
-1
@@ -57,7 +57,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
57
58
template <typename T>
59
void expandInputRowsKernelLauncher(
60
- T const* unpermuted_input, T* permuted_output, int* sorted_experts,
60
+ T const* unpermuted_input, T* permuted_output,
61
int const* expanded_dest_row_to_expanded_source_row,
62
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
63
int64_t const* expert_first_token_offset, int64_t const num_rows,
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl
+8
-9
@@ -2,7 +2,7 @@
2
3
template <typename T, bool CHECK_SKIPPED>
4
__global__ void expandInputRowsKernel(
5
- T const* unpermuted_input, T* permuted_output, int* sorted_experts,
5
+ T const* unpermuted_input, T* permuted_output,
6
int const* expanded_dest_row_to_expanded_source_row,
7
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
8
int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -16,7 +16,6 @@ __global__ void expandInputRowsKernel(
16
int64_t expanded_dest_row = blockIdx.x;
17
int64_t const expanded_source_row =
18
expanded_dest_row_to_expanded_source_row[expanded_dest_row];
19
- int expert_id = sorted_experts[expanded_dest_row];
19
20
if (threadIdx.x == 0) {
21
assert(expanded_dest_row <= INT32_MAX);
@@ -54,7 +53,7 @@ __global__ void expandInputRowsKernel(
53
54
template <typename T>
55
void expandInputRowsKernelLauncher(
57
- T const* unpermuted_input, T* permuted_output, int* sorted_experts,
56
+ T const* unpermuted_input, T* permuted_output,
57
int const* expanded_dest_row_to_expanded_source_row,
58
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
59
int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -70,12 +69,12 @@ void expandInputRowsKernelLauncher(
69
bool is_check_skip = num_valid_tokens_ptr != nullptr;
70
auto func = func_map[is_check_skip];
71
73
- func<<<blocks, threads, 0, stream>>>(
74
- unpermuted_input, permuted_output, sorted_experts,
75
- expanded_dest_row_to_expanded_source_row,
76
- expanded_source_row_to_expanded_dest_row, permuted_idx,
77
- expert_first_token_offset, num_rows, num_valid_tokens_ptr, cols, k,
78
- num_local_experts);
72
+ func<<<blocks, threads, 0, stream>>>(unpermuted_input, permuted_output,
73
+ expanded_dest_row_to_expanded_source_row,
74
+ expanded_source_row_to_expanded_dest_row,
75
+ permuted_idx, expert_first_token_offset,
76
+ num_rows, num_valid_tokens_ptr, cols, k,
77
+ num_local_experts);
78
}
79
80
template <class T, class U>