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, 11:15 UTC
[Perf] Optimize `moe_permute` for CUTLASS FP8 (#32892)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 28, 2026 at 13:15 UTC
c4e744dbd41f23ee7fd554a86cc1bf516082552d
3 files changed
+47
-44
csrc/moe/moe_permute_unpermute_op.cu
+24
-9
@@ -73,25 +73,40 @@ void moe_permute(
73
get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
74
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
75
76
+ // DeepGEMM: use getMIndices kernel to compute
77
+ // 1) align_expert_first_token_offset (aligned prefix offsets)
78
+ // 2) m_indices (expert id for each aligned row)
79
+ // eg. expert0: 3, expert1: 5, expert2: 2 tokens respectively
80
+ // expert_first_token_offset = [0, 3, 8, 10], align_block_size = 4
81
+ // expert0: 3->4, expert1: 5->8, expert2: 2->4
82
+ // align_expert_first_token_offset = [0, 4, 12, 16]
83
+ // so m_indices = [0,0,0,0, 1,1,1,1,1,1,1,1, 2,2,2,2]
84
+ torch::Tensor align_expert_first_token_offset;
85
+ const int64_t* aligned_expert_first_token_offset_ptr = nullptr;
86
+ if (align_block_size.has_value()) {
87
+ align_expert_first_token_offset =
88
+ torch::zeros_like(expert_first_token_offset);
89
+ getMIndices(get_ptr<int64_t>(expert_first_token_offset),
90
+ get_ptr<int64_t>(align_expert_first_token_offset),
91
+ get_ptr<int>(m_indices), n_local_expert, align_block_size_value,
92
+ stream);
93
+ aligned_expert_first_token_offset_ptr =
94
+ get_ptr<int64_t>(align_expert_first_token_offset);
95
+ }
96
+
97
// dispatch expandInputRowsKernelLauncher
98
MOE_DISPATCH(input.scalar_type(), [&] {
99
expandInputRowsKernelLauncher<scalar_t>(
100
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
101
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
102
get_ptr<int>(inv_permuted_idx), get_ptr<int>(permuted_idx),
82
- get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
83
- n_hidden, topk, n_local_expert, align_block_size_value, stream);
103
+ get_ptr<int64_t>(expert_first_token_offset),
104
+ aligned_expert_first_token_offset_ptr, n_token, valid_num_ptr, n_hidden,
105
+ topk, n_local_expert, align_block_size_value, stream);
106
});
107
86
- // get m_indices and update expert_first_token_offset with align block
108
// this is only required for DeepGemm and not required for CUTLASS group gemm
109
if (align_block_size.has_value()) {
89
- auto align_expert_first_token_offset =
90
- torch::zeros_like(expert_first_token_offset);
91
- getMIndices(get_ptr<int64_t>(expert_first_token_offset),
92
- get_ptr<int64_t>(align_expert_first_token_offset),
93
- get_ptr<int>(m_indices), n_local_expert, align_block_size_value,
94
- stream);
110
expert_first_token_offset.copy_(align_expert_first_token_offset);
111
}
112
}
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+2
-1
@@ -60,7 +60,8 @@ void expandInputRowsKernelLauncher(
60
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
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* expert_first_token_offset, int64_t const num_rows,
63
+ int64_t const* expert_first_token_offset,
64
+ int64_t const* aligned_expert_first_token_offset, int64_t const num_rows,
65
int64_t const* num_valid_tokens_ptr, int64_t const cols, int const k,
66
int num_local_experts, const int& align_block_size, cudaStream_t stream);
67
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl
+21
-34
@@ -5,7 +5,8 @@ __global__ void expandInputRowsKernel(
5
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
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* expert_first_token_offset, int64_t const num_rows,
8
+ int64_t const* expert_first_token_offset,
9
+ int64_t const* aligned_expert_first_token_offset, int64_t const num_rows,
10
int64_t const* num_dest_rows, int64_t const cols, int64_t k,
11
int num_local_experts, int align_block_size) {
12
// Reverse permutation map.
@@ -18,35 +19,22 @@ __global__ void expandInputRowsKernel(
19
expanded_dest_row_to_expanded_source_row[expanded_dest_row];
20
int expert_id = sorted_experts[expanded_dest_row];
21
21
- extern __shared__ int64_t smem_expert_first_token_offset[];
22
if constexpr (ALIGN_BLOCK_SIZE) {
23
- // load g2s
24
- for (int idx = threadIdx.x; idx < num_local_experts + 1;
25
- idx += blockDim.x) {
26
- smem_expert_first_token_offset[idx] =
27
- __ldg(expert_first_token_offset + idx);
23
+ // convert (unaligned) expanded_dest_row -> aligned expanded_dest_row.
24
+ // aligned_expert_first_token_offset[e] provides the aligned prefix start
25
+ // for expert e. For non-local experts we map to the end (total aligned M).
26
+ int64_t aligned_base = 0;
27
+ int64_t token_offset_in_expert = 0;
28
+ if (expert_id >= num_local_experts) {
29
+ aligned_base =
30
+ __ldg(aligned_expert_first_token_offset + num_local_experts);
31
+ token_offset_in_expert = 0;
32
+ } else {
33
+ aligned_base = __ldg(aligned_expert_first_token_offset + expert_id);
34
+ token_offset_in_expert =
35
+ expanded_dest_row - __ldg(expert_first_token_offset + expert_id);
36
}
29
- __syncthreads();
30
- int lane_idx = threadIdx.x & 31;
31
-
32
- if (lane_idx == 0) {
33
- // set token_offset_in_expert = 0 if this expert is not local expert
34
- int token_offset_in_expert =
35
- expert_id >= num_local_experts
36
- ? 0
37
- : expanded_dest_row - smem_expert_first_token_offset[expert_id];
38
- int64_t accumulate_align_offset = 0;
39
-#pragma unroll 1
40
- for (int eidx = 1; eidx <= min(expert_id, num_local_experts); eidx++) {
41
- auto n_token_in_expert = smem_expert_first_token_offset[eidx] -
42
- smem_expert_first_token_offset[eidx - 1];
43
- accumulate_align_offset += (n_token_in_expert + align_block_size - 1) /
44
- align_block_size * align_block_size;
45
- }
46
- expanded_dest_row = accumulate_align_offset + token_offset_in_expert;
47
- }
48
- // lane0 shuffle broadcast align_expanded_dest_row
49
- expanded_dest_row = __shfl_sync(0xffffffff, expanded_dest_row, 0);
37
+ expanded_dest_row = aligned_base + token_offset_in_expert;
38
}
39
40
if (threadIdx.x == 0) {
@@ -88,7 +76,8 @@ void expandInputRowsKernelLauncher(
76
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
77
int const* expanded_dest_row_to_expanded_source_row,
78
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
91
- int64_t* expert_first_token_offset, int64_t const num_rows,
79
+ int64_t const* expert_first_token_offset,
80
+ int64_t const* aligned_expert_first_token_offset, int64_t const num_rows,
81
int64_t const* num_valid_tokens_ptr, int64_t const cols, int const k,
82
int num_local_experts, const int& align_block_size, cudaStream_t stream) {
83
int64_t const blocks = num_rows * k;
@@ -104,14 +93,12 @@ void expandInputRowsKernelLauncher(
93
bool is_align_block_size = align_block_size != -1;
94
auto func = func_map[is_check_skip][is_align_block_size];
95
107
- int64_t smem_size = sizeof(int64_t) * (num_local_experts + 1);
108
-
109
- func<<<blocks, threads, smem_size, stream>>>(
96
+ func<<<blocks, threads, 0, stream>>>(
97
unpermuted_input, permuted_output, sorted_experts,
98
expanded_dest_row_to_expanded_source_row,
99
expanded_source_row_to_expanded_dest_row, permuted_idx,
113
- expert_first_token_offset, num_rows, num_valid_tokens_ptr, cols, k,
114
- num_local_experts, align_block_size);
100
+ expert_first_token_offset, aligned_expert_first_token_offset, num_rows,
101
+ num_valid_tokens_ptr, cols, k, num_local_experts, align_block_size);
102
}
103
104
template <class T, class U>