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, 23:45 UTC
Fix CUDA permute/unpermute for use with DeepGemm Moe (#17934)
Signed-off-by: Caleb_Du <Caleb_Du@zju.edu.cn>
Caleb_Du committed
Jul 27, 2025 at 22:08 UTC
57c22e57f989b466a46a990243bb7f072a668b7f
8 files changed
+236
-209
benchmarks/kernels/benchmark_moe_permute_unpermute.py
+43
-33
@@ -8,12 +8,13 @@ import ray
8
import torch
9
from transformers import AutoConfig
10
11
-from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
11
+from vllm.model_executor.layers.fused_moe.fused_moe import *
12
+from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
13
_moe_permute,
14
_moe_unpermute_and_reduce,
15
+ moe_permute,
16
+ moe_unpermute,
17
)
15
-from vllm.model_executor.layers.fused_moe.fused_moe import *
16
-from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import *
18
from vllm.model_executor.layers.fused_moe.utils import _fp8_quantize
19
from vllm.platforms import current_platform
20
from vllm.utils import FlexibleArgumentParser
@@ -63,18 +64,19 @@ def benchmark_permute(
64
65
def run():
66
if use_customized_permute:
66
- (permuted_hidden_states, first_token_off, inv_perm_idx, m_indices) = (
67
- moe_permute(
68
- qhidden_states,
69
- topk_weights=topk_weights,
70
- topk_ids=topk_ids,
71
- token_expert_indices=token_expert_indices,
72
- topk=topk,
73
- n_expert=num_experts,
74
- n_local_expert=num_experts,
75
- expert_map=None,
76
- align_block_size=align_block_size,
77
- )
67
+ (
68
+ permuted_hidden_states,
69
+ a1q_scale,
70
+ first_token_off,
71
+ inv_perm_idx,
72
+ m_indices,
73
+ ) = moe_permute(
74
+ qhidden_states,
75
+ a1q_scale=None,
76
+ topk_ids=topk_ids,
77
+ n_expert=num_experts,
78
+ expert_map=None,
79
+ align_block_size=align_block_size,
80
)
81
else:
82
(
@@ -150,18 +152,19 @@ def benchmark_unpermute(
152
153
def prepare():
154
if use_customized_permute:
153
- (permuted_hidden_states, first_token_off, inv_perm_idx, m_indices) = (
154
- moe_permute(
155
- qhidden_states,
156
- topk_weights=topk_weights,
157
- topk_ids=topk_ids,
158
- token_expert_indices=token_expert_indices,
159
- topk=topk,
160
- n_expert=num_experts,
161
- n_local_expert=num_experts,
162
- expert_map=None,
163
- align_block_size=align_block_size,
164
- )
155
+ (
156
+ permuted_hidden_states,
157
+ a1q_scale,
158
+ first_token_off,
159
+ inv_perm_idx,
160
+ m_indices,
161
+ ) = moe_permute(
162
+ qhidden_states,
163
+ a1q_scale=None,
164
+ topk_ids=topk_ids,
165
+ n_expert=num_experts,
166
+ expert_map=None,
167
+ align_block_size=align_block_size,
168
)
169
# convert to fp16/bf16 as gemm output
170
return (
@@ -191,16 +194,19 @@ def benchmark_unpermute(
194
195
def run(input: tuple):
196
if use_customized_permute:
194
- (permuted_hidden_states, first_token_off, inv_perm_idx, m_indices) = input
197
+ (
198
+ permuted_hidden_states,
199
+ first_token_off,
200
+ inv_perm_idx,
201
+ m_indices,
202
+ ) = input
203
+ output = torch.empty_like(hidden_states)
204
moe_unpermute(
205
+ output,
206
permuted_hidden_states,
207
topk_weights,
198
- topk_ids,
208
inv_perm_idx,
209
first_token_off,
201
- topk,
202
- num_experts,
203
- num_experts,
210
)
211
else:
212
(
@@ -211,7 +217,11 @@ def benchmark_unpermute(
217
inv_perm,
218
) = input
219
_moe_unpermute_and_reduce(
214
- output_hidden_states, permuted_hidden_states, inv_perm, topk_weights
220
+ output_hidden_states,
221
+ permuted_hidden_states,
222
+ inv_perm,
223
+ topk_weights,
224
+ True,
225
)
226
227
# JIT compilation & warmup
csrc/moe/moe_permute_unpermute_op.cu
+35
-38
@@ -10,32 +10,28 @@
10
11
void moe_permute(
12
const torch::Tensor& input, // [n_token, hidden]
13
- const torch::Tensor& topk_weights, //[n_token, topk]
14
- torch::Tensor& topk_ids, // [n_token, topk]
13
+ const torch::Tensor& topk_ids, // [n_token, topk]
14
const torch::Tensor& token_expert_indices, // [n_token, topk]
15
const std::optional<torch::Tensor>& expert_map, // [n_expert]
16
int64_t n_expert, int64_t n_local_expert, int64_t topk,
17
const std::optional<int64_t>& align_block_size,
19
- torch::Tensor&
20
- permuted_input, // [topk * n_token/align_block_size_m, hidden]
18
+ torch::Tensor& permuted_input, // [permuted_size, hidden]
19
torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
22
- torch::Tensor& src_row_id2dst_row_id_map, // [n_token, topk]
20
+ torch::Tensor& inv_permuted_idx, // [n_token, topk]
21
+ torch::Tensor& permuted_idx, // [permute_size]
22
torch::Tensor& m_indices) { // [align_expand_m]
24
- TORCH_CHECK(topk_weights.scalar_type() == at::ScalarType::Float,
25
- "topk_weights must be float32");
23
TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
24
"expert_first_token_offset must be int64");
25
TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
26
"topk_ids must be int32");
27
TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int,
28
"token_expert_indices must be int32");
32
- TORCH_CHECK(src_row_id2dst_row_id_map.scalar_type() == at::ScalarType::Int,
33
- "src_row_id2dst_row_id_map must be int32");
29
+ TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
30
+ "inv_permuted_idx must be int32");
31
TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
32
"expert_first_token_offset shape != n_local_expert+1")
36
- TORCH_CHECK(
37
- src_row_id2dst_row_id_map.sizes() == token_expert_indices.sizes(),
38
- "token_expert_indices shape must be same as src_row_id2dst_row_id_map");
33
+ TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
34
+ "token_expert_indices shape must be same as inv_permuted_idx");
35
auto n_token = input.sizes()[0];
36
auto n_hidden = input.sizes()[1];
37
auto align_block_size_value =
@@ -46,8 +42,9 @@ void moe_permute(
42
auto sort_workspace = torch::empty(
43
{sorter_size},
44
torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
45
+ auto copy_topk_ids = topk_ids.clone(); // copy topk_ids for preprocess
46
auto permuted_experts_id = torch::empty_like(topk_ids);
50
- auto dst_row_id2src_row_id_map = torch::empty_like(src_row_id2dst_row_id_map);
47
+ auto sorted_row_idx = torch::empty_like(inv_permuted_idx);
48
auto align_expert_first_token_offset =
49
torch::zeros_like(expert_first_token_offset);
50
@@ -67,24 +64,22 @@ void moe_permute(
64
const int* expert_map_ptr = get_ptr<int>(expert_map.value());
65
valid_num_ptr =
66
get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
70
- preprocessTopkIdLauncher(get_ptr<int>(topk_ids), n_token * topk,
67
+ preprocessTopkIdLauncher(get_ptr<int>(copy_topk_ids), n_token * topk,
68
expert_map_ptr, n_expert, stream);
69
}
70
// expert sort topk expert id and scan expert id get expert_first_token_offset
74
- sortAndScanExpert(get_ptr<int>(topk_ids), get_ptr<int>(token_expert_indices),
75
- get_ptr<int>(permuted_experts_id),
76
- get_ptr<int>(dst_row_id2src_row_id_map),
77
- get_ptr<int64_t>(expert_first_token_offset), n_token,
78
- n_expert, n_local_expert, topk, sorter,
79
- get_ptr<int>(sort_workspace), stream);
71
+ sortAndScanExpert(
72
+ get_ptr<int>(copy_topk_ids), get_ptr<int>(token_expert_indices),
73
+ get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
74
+ get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
75
+ n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
76
77
// dispatch expandInputRowsKernelLauncher
78
MOE_DISPATCH(input.scalar_type(), [&] {
79
expandInputRowsKernelLauncher<scalar_t>(
80
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
85
- get_ptr<float>(topk_weights), get_ptr<int>(permuted_experts_id),
86
- get_ptr<int>(dst_row_id2src_row_id_map),
87
- get_ptr<int>(src_row_id2dst_row_id_map),
81
+ get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
82
+ get_ptr<int>(inv_permuted_idx), get_ptr<int>(permuted_idx),
83
get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
84
n_hidden, topk, n_local_expert, align_block_size_value, stream);
85
});
@@ -101,32 +96,34 @@ void moe_permute(
96
}
97
98
void moe_unpermute(
104
- const torch::Tensor& permuted_hidden_states, // [n_token * topk, hidden]
105
- const torch::Tensor& topk_weights, //[n_token, topk]
106
- const torch::Tensor& topk_ids, // [n_token, topk]
107
- const torch::Tensor& src_row_id2dst_row_id_map, // [n_token, topk]
108
- const torch::Tensor& expert_first_token_offset, // [n_local_expert+1]
109
- int64_t n_expert, int64_t n_local_expert, int64_t topk,
99
+ const torch::Tensor& permuted_hidden_states, // [n_token * topk, hidden]
100
+ const torch::Tensor& topk_weights, // [n_token, topk]
101
+ const torch::Tensor& inv_permuted_idx, // [n_token, topk]
102
+ const std::optional<torch::Tensor>&
103
+ expert_first_token_offset, // [n_local_expert+1]
104
+ int64_t topk,
105
torch::Tensor& hidden_states // [n_token, hidden]
106
) {
112
- TORCH_CHECK(src_row_id2dst_row_id_map.sizes() == topk_ids.sizes(),
113
- "topk_ids shape must be same as src_row_id2dst_row_id_map");
114
- TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
115
- "topk_ids must be int32");
107
TORCH_CHECK(
108
permuted_hidden_states.scalar_type() == hidden_states.scalar_type(),
118
- "topk_ids dtype must be same as src_row_id2dst_row_id_map");
109
+ "permuted_hidden_states dtype must be same as hidden_states");
110
auto n_token = hidden_states.size(0);
111
auto n_hidden = hidden_states.size(1);
112
auto stream = at::cuda::getCurrentCUDAStream().stream();
122
- const int64_t* valid_ptr =
123
- get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
113
+
114
+ int64_t const* valid_ptr = nullptr;
115
+ if (expert_first_token_offset.has_value()) {
116
+ int n_local_expert = expert_first_token_offset.value().size(0) - 1;
117
+ valid_ptr =
118
+ get_ptr<int64_t>(expert_first_token_offset.value()) + n_local_expert;
119
+ }
120
+
121
MOE_DISPATCH(hidden_states.scalar_type(), [&] {
122
finalizeMoeRoutingKernelLauncher<scalar_t, scalar_t>(
123
get_ptr<scalar_t>(permuted_hidden_states),
124
get_ptr<scalar_t>(hidden_states), get_ptr<float>(topk_weights),
128
- get_ptr<int>(src_row_id2dst_row_id_map), get_ptr<int>(topk_ids),
129
- n_token, n_hidden, topk, valid_ptr, stream);
125
+ get_ptr<int>(inv_permuted_idx), n_token, n_hidden, topk, valid_ptr,
126
+ stream);
127
});
128
}
129
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+1
-1
@@ -177,7 +177,7 @@ __global__ void getMIndicesKernel(int64_t* expert_first_token_offset,
177
int tidx = threadIdx.x;
178
extern __shared__ int64_t smem_expert_first_token_offset[];
179
for (int i = tidx; i <= num_local_expert; i += blockDim.x) {
180
- smem_expert_first_token_offset[tidx] = __ldg(expert_first_token_offset + i);
180
+ smem_expert_first_token_offset[i] = __ldg(expert_first_token_offset + i);
181
}
182
__syncthreads();
183
auto last_token_offset = smem_expert_first_token_offset[eidx + 1];
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+4
-16
@@ -57,31 +57,19 @@ void sortAndScanExpert(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,
61
- const float* unpermuted_scales, int* sorted_experts,
60
+ T const* unpermuted_input, T* permuted_output, int* sorted_experts,
61
int const* expanded_dest_row_to_expanded_source_row,
63
- int* expanded_source_row_to_expanded_dest_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,
64
int64_t const* num_valid_tokens_ptr, int64_t const cols, int const k,
65
int num_local_experts, const int& align_block_size, cudaStream_t stream);
66
68
-// Final kernel to unpermute and scale
69
-// This kernel unpermutes the original data, does the k-way reduction and
70
-// performs the final skip connection.
71
-template <typename T, typename OutputType, bool CHECK_SKIPPED>
72
-__global__ void finalizeMoeRoutingKernel(
73
- T const* expanded_permuted_rows, OutputType* reduced_unpermuted_output,
74
- float const* scales, int const* expanded_source_row_to_expanded_dest_row,
75
- int const* expert_for_source_row, int64_t const orig_cols, int64_t const k,
76
- int64_t const* num_valid_ptr);
77
-
67
template <class T, class OutputType>
68
void finalizeMoeRoutingKernelLauncher(
69
T const* expanded_permuted_rows, OutputType* reduced_unpermuted_output,
70
float const* scales, int const* expanded_source_row_to_expanded_dest_row,
82
- int const* expert_for_source_row, int64_t const num_rows,
83
- int64_t const cols, int64_t const k, int64_t const* num_valid_ptr,
84
- cudaStream_t stream);
71
+ int64_t const num_rows, int64_t const cols, int64_t const k,
72
+ int64_t const* num_valid_ptr, cudaStream_t stream);
73
74
void preprocessTopkIdLauncher(int* topk_id_ptr, int size,
75
const int* expert_map_ptr, int num_experts,
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl
+18
-21
@@ -2,10 +2,9 @@
2
3
template <typename T, bool CHECK_SKIPPED, bool ALIGN_BLOCK_SIZE>
4
__global__ void expandInputRowsKernel(
5
- T const* unpermuted_input, T* permuted_output,
6
- const float* unpermuted_scales, int* sorted_experts,
5
+ T const* unpermuted_input, T* permuted_output, int* sorted_experts,
6
int const* expanded_dest_row_to_expanded_source_row,
8
- int* expanded_source_row_to_expanded_dest_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,
9
int64_t const* num_dest_rows, int64_t const cols, int64_t k,
10
int num_local_experts, int align_block_size) {
@@ -54,6 +53,10 @@ __global__ void expandInputRowsKernel(
53
assert(expanded_dest_row <= INT32_MAX);
54
expanded_source_row_to_expanded_dest_row[expanded_source_row] =
55
static_cast<int>(expanded_dest_row);
56
+ // skip non local expert token
57
+ if (!CHECK_SKIPPED || blockIdx.x < *num_dest_rows) {
58
+ permuted_idx[expanded_dest_row] = expanded_source_row;
59
+ }
60
}
61
62
if (!CHECK_SKIPPED || blockIdx.x < *num_dest_rows) {
@@ -62,7 +65,7 @@ __global__ void expandInputRowsKernel(
65
using DataElem = cutlass::Array<T, ELEM_PER_THREAD>;
66
67
// Duplicate and permute rows
65
- int64_t const source_row = expanded_source_row % num_rows;
68
+ int64_t const source_row = expanded_source_row / k;
69
70
auto const* source_row_ptr =
71
reinterpret_cast<DataElem const*>(unpermuted_input + source_row * cols);
@@ -82,10 +85,9 @@ __global__ void expandInputRowsKernel(
85
86
template <typename T>
87
void expandInputRowsKernelLauncher(
85
- T const* unpermuted_input, T* permuted_output,
86
- const float* unpermuted_scales, int* sorted_experts,
88
+ T const* unpermuted_input, T* permuted_output, int* sorted_experts,
89
int const* expanded_dest_row_to_expanded_source_row,
88
- int* expanded_source_row_to_expanded_dest_row,
90
+ int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
91
int64_t* expert_first_token_offset, int64_t const num_rows,
92
int64_t const* num_valid_tokens_ptr, int64_t const cols, int const k,
93
int num_local_experts, const int& align_block_size, cudaStream_t stream) {
@@ -105,11 +107,11 @@ void expandInputRowsKernelLauncher(
107
int64_t smem_size = sizeof(int64_t) * (num_local_experts + 1);
108
109
func<<<blocks, threads, smem_size, stream>>>(
108
- unpermuted_input, permuted_output, unpermuted_scales, sorted_experts,
110
+ unpermuted_input, permuted_output, sorted_experts,
111
expanded_dest_row_to_expanded_source_row,
110
- expanded_source_row_to_expanded_dest_row, expert_first_token_offset,
111
- num_rows, num_valid_tokens_ptr, cols, k, num_local_experts,
112
- align_block_size);
112
+ 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);
115
}
116
117
template <class T, class U>
@@ -128,11 +130,9 @@ template <typename T, typename OutputType, bool CHECK_SKIPPED>
130
__global__ void finalizeMoeRoutingKernel(
131
T const* expanded_permuted_rows, OutputType* reduced_unpermuted_output,
132
float const* scales, int const* expanded_source_row_to_expanded_dest_row,
131
- int const* expert_for_source_row, int64_t const orig_cols, int64_t const k,
132
- int64_t const* num_valid_ptr) {
133
+ int64_t const orig_cols, int64_t const k, int64_t const* num_valid_ptr) {
134
assert(orig_cols % 4 == 0);
135
int64_t const original_row = blockIdx.x;
135
- int64_t const num_rows = gridDim.x;
136
auto const offset = original_row * orig_cols;
137
OutputType* reduced_row_ptr = reduced_unpermuted_output + offset;
138
int64_t const num_valid = *num_valid_ptr;
@@ -159,14 +159,13 @@ __global__ void finalizeMoeRoutingKernel(
159
ComputeElem thread_output;
160
thread_output.fill(0);
161
for (int k_idx = 0; k_idx < k; ++k_idx) {
162
- int64_t const expanded_original_row = original_row + k_idx * num_rows;
162
+ int64_t const expanded_original_row = original_row * k + k_idx;
163
int64_t const expanded_permuted_row =
164
expanded_source_row_to_expanded_dest_row[expanded_original_row];
165
166
int64_t const k_offset = original_row * k + k_idx;
167
float const row_scale = scales[k_offset];
168
169
- // Check after row_rescale has accumulated
169
if (CHECK_SKIPPED && expanded_permuted_row >= num_valid) {
170
continue;
171
}
@@ -189,9 +188,8 @@ template <class T, class OutputType>
188
void finalizeMoeRoutingKernelLauncher(
189
T const* expanded_permuted_rows, OutputType* reduced_unpermuted_output,
190
float const* scales, int const* expanded_source_row_to_expanded_dest_row,
192
- int const* expert_for_source_row, int64_t const num_rows,
193
- int64_t const cols, int64_t const k, int64_t const* num_valid_ptr,
194
- cudaStream_t stream) {
191
+ int64_t const num_rows, int64_t const cols, int64_t const k,
192
+ int64_t const* num_valid_ptr, cudaStream_t stream) {
193
int64_t const blocks = num_rows;
194
int64_t const threads = 256;
195
bool const check_finished = num_valid_ptr != nullptr;
@@ -201,6 +199,5 @@ void finalizeMoeRoutingKernelLauncher(
199
auto* const kernel = func_map[check_finished];
200
kernel<<<blocks, threads, 0, stream>>>(
201
expanded_permuted_rows, reduced_unpermuted_output, scales,
204
- expanded_source_row_to_expanded_dest_row, expert_for_source_row, cols, k,
205
- num_valid_ptr);
202
+ expanded_source_row_to_expanded_dest_row, cols, k, num_valid_ptr);
203
}
csrc/moe/torch_bindings.cpp
+5
-6
@@ -56,18 +56,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
56
" -> Tensor");
57
58
m.def(
59
- "moe_permute(Tensor input, Tensor topk_weight, Tensor! topk_ids,"
59
+ "moe_permute(Tensor input, Tensor topk_ids,"
60
"Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
61
"int n_local_expert,"
62
"int topk, int? align_block_size,Tensor! permuted_input, Tensor! "
63
- "expert_first_token_offset, Tensor! src_row_id2dst_row_id_map, Tensor! "
64
- "m_indices)->()");
63
+ "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
64
+ "permuted_idx, Tensor! m_indices)->()");
65
66
m.def(
67
"moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
68
- "Tensor topk_ids,Tensor src_row_id2dst_row_id_map, Tensor "
69
- "expert_first_token_offset, int n_expert, int n_local_expert,int "
70
- "topk, Tensor! hidden_states)->()");
68
+ "Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
69
+ "int topk, Tensor! hidden_states)->()");
70
71
m.def("moe_permute_unpermute_supported() -> bool");
72
m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
tests/kernels/moe/test_moe_permute_unpermute.py
+82
-50
@@ -17,28 +17,34 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
17
moe_permute, moe_permute_unpermute_supported, moe_unpermute)
18
from vllm.platforms import current_platform
19
20
-NUM_EXPERTS = [16, 64]
20
+NUM_EXPERTS = [16, 64, 256]
21
TOP_KS = [2, 4, 6, 8]
22
EP_SIZE = [1, 4, 16]
23
current_platform.seed_everything(0)
24
25
26
-def torch_permute(hidden_states: torch.Tensor,
27
- topk_ids: torch.Tensor,
28
- token_expert_indices: torch.Tensor,
29
- topk: int,
30
- n_expert: int,
31
- n_local_expert: int,
32
- start_expert: int,
33
- expert_map: Optional[torch.Tensor] = None,
34
- align_block_size: Optional[int] = None,
35
- fill_invalid_expert: int = -1) -> list[torch.Tensor]:
26
+def torch_permute(
27
+ hidden_states: torch.Tensor,
28
+ topk_ids: torch.Tensor,
29
+ # token_expert_indices: torch.Tensor,
30
+ topk: int,
31
+ n_expert: int,
32
+ n_local_expert: int,
33
+ start_expert: int,
34
+ expert_map: Optional[torch.Tensor] = None,
35
+ align_block_size: Optional[int] = None,
36
+ fill_invalid_expert: int = -1) -> list[torch.Tensor]:
37
n_token, n_hidden = hidden_states.shape[0], hidden_states.shape[1]
38
if expert_map is not None:
39
is_local_expert = (expert_map[topk_ids] != -1)
40
not_local_expert = (expert_map[topk_ids] == -1)
41
topk_ids = is_local_expert * (
42
topk_ids - start_expert) + not_local_expert * (topk_ids + n_expert)
43
+ token_expert_indices = torch.arange(0,
44
+ n_token * topk,
45
+ dtype=torch.int32,
46
+ device=hidden_states.device).reshape(
47
+ (n_token, topk))
48
49
sorted_topk_ids, sorted_indices = torch.sort(topk_ids.flatten(),
50
stable=True)
@@ -59,8 +65,8 @@ def torch_permute(hidden_states: torch.Tensor,
65
valid_row_idx = []
66
if align_block_size is None:
67
62
- permuted_hidden_states = hidden_states[dst_row_id2src_row_id_map %
63
- n_token, ...]
68
+ permuted_hidden_states = hidden_states[dst_row_id2src_row_id_map //
69
+ topk, ...]
70
permuted_row_size = permuted_hidden_states.shape[0]
71
m_indices = torch.empty(permuted_row_size,
72
device="cuda",
@@ -73,14 +79,21 @@ def torch_permute(hidden_states: torch.Tensor,
79
0, n_token * topk, device="cuda",
80
dtype=torch.int32)[src2dst_idx].reshape((n_token, topk))
81
valid_row_idx += [i for i in range(expert_first_token_offset[-1])]
82
+ dst_row_id2src_row_id_map[
83
+ expert_first_token_offset[-1]:] = n_token * topk
84
return [
85
permuted_hidden_states, expert_first_token_offset,
78
- src_row_id2dst_row_id_map, m_indices, valid_row_idx
86
+ src_row_id2dst_row_id_map, dst_row_id2src_row_id_map, m_indices,
87
+ valid_row_idx
88
]
89
else:
90
permuted_row_size = (topk * n_token + n_expert *
91
(align_block_size - 1) + align_block_size -
92
1) // align_block_size * align_block_size
93
+ permuted_idx = torch.full((permuted_row_size, ),
94
+ n_token * topk,
95
+ dtype=torch.int32,
96
+ device=hidden_states.device)
97
permuted_hidden_states = torch.empty((permuted_row_size, n_hidden),
98
device="cuda",
99
dtype=hidden_states.dtype)
@@ -105,13 +118,16 @@ def torch_permute(hidden_states: torch.Tensor,
118
align_first_token_offset = align_expert_first_token_offset[i - 1]
119
align_last_token_offset = align_expert_first_token_offset[i]
120
dst_row_id2src_row_id_in_expert = dst_row_id2src_row_id_map[
108
- first_token_offset:first_token_offset +
109
- n_token_in_expert] % n_token
121
+ first_token_offset:first_token_offset + n_token_in_expert]
122
# store token in current expert with align_first_token_offset
123
permuted_hidden_states[align_first_token_offset:\
124
align_first_token_offset+n_token_in_expert,\
125
...] = hidden_states[\
114
- dst_row_id2src_row_id_in_expert, ...]
126
+ dst_row_id2src_row_id_in_expert // topk,\
127
+ ...]
128
+ permuted_idx[align_first_token_offset:\
129
+ align_first_token_offset+\
130
+ n_token_in_expert] = dst_row_id2src_row_id_in_expert
131
# set current expert m_indices
132
m_indices[align_first_token_offset:align_last_token_offset] = i - 1
133
valid_row_idx += [
@@ -135,7 +151,7 @@ def torch_permute(hidden_states: torch.Tensor,
151
src2dst_idx].reshape((n_token, topk))
152
return [
153
permuted_hidden_states, align_expert_first_token_offset,
138
- align_src_row_id2dst_row_id, m_indices, valid_row_idx
154
+ align_src_row_id2dst_row_id, permuted_idx, m_indices, valid_row_idx
155
]
156
157
@@ -146,15 +162,18 @@ def torch_unpermute(permuted_hidden_states: torch.Tensor,
162
valid_row_idx: torch.Tensor, topk: int,
163
n_expert: int) -> torch.Tensor:
164
# ignore invalid row
165
+ n_hidden = permuted_hidden_states.shape[1]
166
mask = torch.zeros(permuted_hidden_states.shape[0],
167
dtype=bool,
168
device="cuda")
169
mask[valid_row_idx] = True
170
permuted_hidden_states[~mask] = 0
154
- idx = src_row_id2dst_row_id_map.flatten()[
155
- token_expert_indices.flatten()].reshape(token_expert_indices.shape)
156
- output = permuted_hidden_states[idx, ...] * topk_weights[..., None]
157
- output = output.sum(dim=1).to(permuted_hidden_states.dtype)
171
+
172
+ permuted_hidden_states = permuted_hidden_states[
173
+ src_row_id2dst_row_id_map.flatten(), ...]
174
+ permuted_hidden_states = permuted_hidden_states.view(-1, topk, n_hidden)
175
+ output = (permuted_hidden_states * topk_weights.unsqueeze(2)).sum(1).to(
176
+ permuted_hidden_states.dtype)
177
return output
178
179
@@ -184,43 +203,56 @@ def test_moe_permute_unpermute(n_token: int, n_hidden: int, topk: int,
203
gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype)
204
topk_weights, topk_ids, token_expert_indices = fused_topk(
205
hidden_states, gating_output, topk, False)
187
- gold0, gold1, gold2, gold3, valid_row_idx = torch_permute(
188
- hidden_states,
189
- topk_ids,
190
- token_expert_indices,
191
- topk,
192
- n_expert,
193
- n_local_expert,
194
- start_expert,
195
- expert_map=expert_map,
196
- align_block_size=align_block_size,
197
- fill_invalid_expert=fill_invalid_expert)
198
-
199
- result0, result1, result2, result3 = moe_permute(
200
- hidden_states, topk_weights, topk_ids, token_expert_indices, topk,
201
- n_expert, n_local_expert, expert_map, align_block_size,
202
- fill_invalid_expert)
206
+ (gold_permuted_hidden_states, gold_expert_first_token_offset,
207
+ gold_inv_permuted_idx, gold_permuted_idx, gold_m_indices,
208
+ valid_row_idx) = torch_permute(
209
+ hidden_states,
210
+ topk_ids,
211
+ # token_expert_indices,
212
+ topk,
213
+ n_expert,
214
+ n_local_expert,
215
+ start_expert,
216
+ expert_map=expert_map,
217
+ align_block_size=align_block_size,
218
+ fill_invalid_expert=fill_invalid_expert)
219
+
220
+ (permuted_hidden_states, _, expert_first_token_offset, inv_permuted_idx,
221
+ m_indices) = moe_permute(hidden_states=hidden_states,
222
+ a1q_scale=None,
223
+ topk_ids=topk_ids,
224
+ n_expert=n_expert,
225
+ n_local_expert=n_local_expert,
226
+ expert_map=expert_map,
227
+ align_block_size=align_block_size,
228
+ fill_invalid_expert=fill_invalid_expert)
229
230
# check expert_first_token_offset
205
- torch.testing.assert_close(gold1, result1, atol=0, rtol=0)
231
+ torch.testing.assert_close(gold_expert_first_token_offset,
232
+ expert_first_token_offset,
233
+ atol=0,
234
+ rtol=0)
235
# check src_row_id2dst_row_id_map
207
- torch.testing.assert_close(gold2, result2, atol=0, rtol=0)
236
+ torch.testing.assert_close(gold_inv_permuted_idx.flatten(),
237
+ inv_permuted_idx,
238
+ atol=0,
239
+ rtol=0)
240
# check mindice
209
- torch.testing.assert_close(gold3, result3, atol=0, rtol=0)
241
+ torch.testing.assert_close(gold_m_indices, m_indices, atol=0, rtol=0)
242
# check permuted_hidden_states, only valid token
211
- torch.testing.assert_close(gold0[valid_row_idx],
212
- result0[valid_row_idx],
243
+ torch.testing.assert_close(gold_permuted_hidden_states[valid_row_idx],
244
+ permuted_hidden_states[valid_row_idx],
245
atol=0,
246
rtol=0)
215
-
247
# add a random tensor to simulate group gemm
217
- result0 = 0.5 * result0 + torch.randn_like(result0)
248
+ result0 = 0.5 * permuted_hidden_states + torch.randn_like(
249
+ permuted_hidden_states)
250
+ result4 = torch.empty_like(hidden_states)
251
+ moe_unpermute(result4, result0, topk_weights, inv_permuted_idx,
252
+ expert_first_token_offset)
253
219
- result4 = moe_unpermute(result0, topk_weights, topk_ids, result2, result1,
220
- topk, n_expert, n_local_expert)
254
gold4 = torch_unpermute(result0, topk_weights, topk_ids,
222
- token_expert_indices, result2, valid_row_idx, topk,
223
- n_local_expert)
224
-
255
+ token_expert_indices, inv_permuted_idx,
256
+ valid_row_idx, topk, n_local_expert)
257
# check unpermuted hidden
258
torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0)
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+48
-44
@@ -76,43 +76,43 @@ def _moe_unpermute_and_reduce(
76
77
def moe_permute(
78
hidden_states: torch.Tensor,
79
- topk_weights: torch.Tensor,
79
+ a1q_scale: Optional[torch.Tensor],
80
topk_ids: torch.Tensor,
81
- token_expert_indices: torch.Tensor,
82
- topk: int,
81
n_expert: int,
84
- n_local_expert: int,
82
+ n_local_expert: int = -1,
83
expert_map: Optional[torch.Tensor] = None,
84
align_block_size: Optional[int] = None,
85
fill_invalid_expert: int = -1
88
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
86
+) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor,
87
+ torch.Tensor]:
88
"""
89
This function expands and permutes activation to gather uncontinuous tokens
90
for each expert.
91
Parameters:
92
- hidden_states (torch.Tensor): The input tensor to the MoE layer.
94
- - topk_weights (torch.Tensor): topk expert route weight for each token.
93
+ - a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states
94
- topk_ids (torch.Tensor): topk expert route id for each token.
96
- - token_expert_indices (torch.Tensor): indice for expanded hidden.
97
- - topk (int): The number of top-k experts to select.
95
- n_expert (int): The number of expert.
96
- n_local_expert (int): The number of expert in current EP rank.
97
- expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
101
- from the global expert space to the local expert space of the expert
98
+ from the global expert space to the local expert space of the expert
99
parallel shard.
100
- align_block_size (Optional[int]): align group gemm block size for deepgemm
101
- fill_invalid_expert(int): fill expert id in m_indices for invalid expert
102
to workaround DeepGemm unsupported -1 in m_indices
103
Returns:
104
- permuted_hidden_states (torch.Tensor): permuted activation.
105
+ - a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states
106
- expert_first_token_offset (torch.Tensor): offset of the first token
107
of each expert for standard grouped gemm. if enable 'align_block_size'
108
expert_first_token_offset will align up to 'align_block_size'.
111
- - src_row_id2dst_row_id_map (torch.Tensor): idx map for moe_unpermute.
109
+ - inv_permuted_idx (torch.Tensor): idx map for moe_unpermute.
110
+ - permuted_idx (torch.Tensor): idx map from hidden to permuted_hidden.
111
- m_indices: m_indices for grouped gemm in deepgemm,`m_indices[i]` records
112
the group which the j-th row of the LHS belong to.`
113
"""
114
n_token, n_hidden = hidden_states.size()
115
+ topk = topk_ids.size(1)
116
assert (n_hidden * hidden_states.element_size()
117
) % 16 == 0, "permue kernel need hidden dim align to 16B"
118
permuted_row_size = n_token * topk
@@ -120,12 +120,19 @@ def moe_permute(
120
permuted_row_size = (permuted_row_size + n_expert *
121
(align_block_size - 1) + align_block_size -
122
1) // align_block_size * align_block_size
123
-
123
+ if n_local_expert == -1:
124
+ n_local_expert = n_expert
125
permuted_hidden_states = torch.empty(
126
(permuted_row_size, n_hidden),
127
dtype=hidden_states.dtype,
128
device=hidden_states.device,
129
)
130
+ token_expert_indices = torch.arange(0,
131
+ n_token * topk,
132
+ dtype=torch.int32,
133
+ device=hidden_states.device).reshape(
134
+ (n_token, topk))
135
+
136
m_indices = torch.full((permuted_row_size, ),
137
fill_invalid_expert,
138
dtype=torch.int32,
@@ -133,57 +140,54 @@ def moe_permute(
140
expert_first_token_offset = torch.empty(n_local_expert + 1,
141
dtype=torch.int64,
142
device=hidden_states.device)
136
- src_row_id2dst_row_id_map = torch.empty((n_token, topk),
137
- dtype=torch.int32,
138
- device=hidden_states.device)
139
- torch.ops._moe_C.moe_permute(hidden_states, topk_weights, topk_ids,
140
- token_expert_indices, expert_map, n_expert,
141
- n_local_expert, topk, align_block_size,
142
- permuted_hidden_states,
143
- expert_first_token_offset,
144
- src_row_id2dst_row_id_map, m_indices)
145
- return (permuted_hidden_states, expert_first_token_offset,
146
- src_row_id2dst_row_id_map, m_indices)
143
+ permuted_idx = torch.full((permuted_row_size, ),
144
+ n_token * topk,
145
+ dtype=torch.int32,
146
+ device=hidden_states.device)
147
+ inv_permuted_idx = torch.empty((n_token, topk),
148
+ dtype=torch.int32,
149
+ device=hidden_states.device)
150
+ topk_ids = topk_ids.to(torch.int32)
151
+ torch.ops._moe_C.moe_permute(hidden_states, topk_ids, token_expert_indices,
152
+ expert_map, n_expert, n_local_expert, topk,
153
+ align_block_size, permuted_hidden_states,
154
+ expert_first_token_offset, inv_permuted_idx,
155
+ permuted_idx, m_indices)
156
+ if a1q_scale is not None:
157
+ a1q_scale = a1q_scale[permuted_idx.clamp(max=n_token * topk - 1) //
158
+ topk]
159
+ return (permuted_hidden_states, a1q_scale, expert_first_token_offset,
160
+ inv_permuted_idx.flatten(), m_indices)
161
162
163
def moe_unpermute(
164
+ out: torch.Tensor,
165
permuted_hidden_states: torch.Tensor,
166
topk_weights: torch.Tensor,
152
- topk_ids: torch.Tensor,
153
- src_row_id2dst_row_id_map: torch.Tensor,
154
- expert_first_token_offset: torch.Tensor,
155
- topk: int,
156
- n_expert: int,
157
- n_local_expert: int,
158
-) -> torch.Tensor:
167
+ inv_permuted_idx: torch.Tensor,
168
+ expert_first_token_offset: Optional[torch.Tensor] = None,
169
+) -> None:
170
"""
171
This function expands and permutes activation to gathering uncontinuous
172
tokens for each expert.
173
Parameters:
174
+ - out (torch.Tensor): output tensor
175
- permuted_hidden_states (torch.Tensor): permuted activation.
176
- topk_weights (torch.Tensor): topk expert route weight for each token.
165
- - topk_ids (torch.Tensor): topk expert route id for each token.
166
- - expert_first_token_offset (torch.Tensor): offset of the first token
167
- of each expert for grouped gemm.
168
- - topk (int): The number of top-k experts to select.
169
- - n_expert (int): The number of expert.
170
- - n_local_expert (int): The number of expert in current EP rank.
177
+ - inv_permuted_idx (torch.Tensor): row idx map for moe_unpermute.
178
+ - expert_first_token_offset (Optional[torch.Tensor]): offset of the first
179
+ token of each expert for grouped gemm.
180
Returns:
181
- hidden_states (torch.Tensor): The reduced and unpermuted activation
182
tensor.
183
"""
175
- n_token, n_hidden = topk_weights.size(0), permuted_hidden_states.size(-1)
184
+ topk = topk_weights.size(1)
185
+ n_hidden = permuted_hidden_states.size(-1)
186
assert (n_hidden * permuted_hidden_states.element_size()
187
) % 16 == 0, "unpermue kernel need hidden dim align to 16B"
178
- hidden_states = torch.empty((n_token, n_hidden),
179
- dtype=permuted_hidden_states.dtype,
180
- device=permuted_hidden_states.device)
181
-
188
torch.ops._moe_C.moe_unpermute(permuted_hidden_states, topk_weights,
183
- topk_ids, src_row_id2dst_row_id_map,
184
- expert_first_token_offset, n_expert,
185
- n_local_expert, topk, hidden_states)
186
- return hidden_states
189
+ inv_permuted_idx, expert_first_token_offset,
190
+ topk, out)
191
192
193
def moe_permute_unpermute_supported():