@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 77c09e1130
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 8, 17:45 UTC

[Refactor] Remove align block size logic in `moe_permute` (#33449)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Feb 6, 2026 at 13:57 UTC 77c09e1130661197ccac2d968a28cd4a557922d5
8 files changed +38 -297
benchmarks/kernels/benchmark_moe_permute_unpermute.py
-6
@@ -44,10 +44,8 @@ def benchmark_permute(
44 hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
45 # output_hidden_states = torch.empty_like(hidden_states)
46 if use_fp8_w8a8:
47 - align_block_size = 128 # deepgemm needs 128 m aligned block
47 qhidden_states, scale = _fp8_quantize(hidden_states, None, None)
48 else:
50 - align_block_size = None
49 qhidden_states = hidden_states
50
51 gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32)
@@ -67,7 +65,6 @@ def benchmark_permute(
65 topk_ids=topk_ids,
66 n_expert=num_experts,
67 expert_map=None,
70 - align_block_size=align_block_size,
68 )
69
70 # JIT compilation & warmup
@@ -117,10 +114,8 @@ def benchmark_unpermute(
114 # init_dtype = torch.float16 if use_fp8_w8a8 else dtype
115 hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
116 if use_fp8_w8a8:
120 - align_block_size = 128 # deepgemm needs 128 m aligned block
117 qhidden_states, scale = _fp8_quantize(hidden_states, None, None)
118 else:
123 - align_block_size = None
119 qhidden_states = hidden_states
120
121 input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32)
@@ -142,7 +137,6 @@ def benchmark_unpermute(
137 topk_ids=topk_ids,
138 n_expert=num_experts,
139 expert_map=None,
145 - align_block_size=align_block_size,
140 )
141 # convert to fp16/bf16 as gemm output
142 return (
csrc/moe/moe_permute_unpermute_op.cu
+5 -39
@@ -14,12 +14,10 @@ void moe_permute(
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,
17 torch::Tensor& permuted_input, // [permuted_size, hidden]
18 torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
19 torch::Tensor& inv_permuted_idx, // [n_token, topk]
21 - torch::Tensor& permuted_idx, // [permute_size]
22 - torch::Tensor& m_indices) { // [align_expand_m]
20 + torch::Tensor& permuted_idx) { // [permute_size]
21 TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
22 "expert_first_token_offset must be int64");
23 TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
@@ -34,8 +32,6 @@ void moe_permute(
32 "token_expert_indices shape must be same as inv_permuted_idx");
33 auto n_token = input.sizes()[0];
34 auto n_hidden = input.sizes()[1];
37 - auto align_block_size_value =
38 - align_block_size.has_value() ? align_block_size.value() : -1;
35 auto stream = at::cuda::getCurrentCUDAStream().stream();
36 const long sorter_size =
37 CubKeyValueSorter::getWorkspaceSize(n_token * topk, n_expert);
@@ -73,42 +69,15 @@ void moe_permute(
69 get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
70 n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
71
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 -
72 // dispatch expandInputRowsKernelLauncher
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),
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);
78 + get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
79 + n_hidden, topk, n_local_expert, stream);
80 });
107 -
108 - // this is only required for DeepGemm and not required for CUTLASS group gemm
109 - if (align_block_size.has_value()) {
110 - expert_first_token_offset.copy_(align_expert_first_token_offset);
111 - }
81 }
82
83 void moe_unpermute(
@@ -201,16 +170,13 @@ void shuffle_rows(const torch::Tensor& input_tensor,
170
171 #else
172
204 -void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_weights,
205 - torch::Tensor& topk_ids,
173 +void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
174 const torch::Tensor& token_expert_indices,
175 const std::optional<torch::Tensor>& expert_map,
176 int64_t n_expert, int64_t n_local_expert, int64_t topk,
209 - const std::optional<int64_t>& align_block_size,
177 torch::Tensor& permuted_input,
178 torch::Tensor& expert_first_token_offset,
212 - torch::Tensor& src_row_id2dst_row_id_map,
213 - torch::Tensor& m_indices) {
179 + torch::Tensor& inv_permuted_idx, torch::Tensor& permuted_idx) {
180 TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
181 }
182
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
-60
@@ -168,64 +168,4 @@ void preprocessTopkIdLauncher(int* topk_id_ptr, int size,
168 topk_id_ptr, size, expert_map_ptr, num_experts);
169 }
170
171 -template <bool ALIGN_BLOCK_SIZE>
172 -__global__ void getMIndicesKernel(int64_t* expert_first_token_offset,
173 - int64_t* align_expert_first_token_offset,
174 - int* m_indices, const int num_local_expert,
175 - const int align_block_size) {
176 - int eidx = blockIdx.x;
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[i] = __ldg(expert_first_token_offset + i);
181 - }
182 - __syncthreads();
183 - auto last_token_offset = smem_expert_first_token_offset[eidx + 1];
184 - auto first_token_offset = smem_expert_first_token_offset[eidx];
185 - int n_token_in_expert = last_token_offset - first_token_offset;
186 -
187 - if constexpr (ALIGN_BLOCK_SIZE) {
188 - n_token_in_expert = (n_token_in_expert + align_block_size - 1) /
189 - align_block_size * align_block_size;
190 - // round up to ALIGN_BLOCK_SIZE
191 - int64_t accumulate_align_offset = 0;
192 - for (int i = 1; i <= eidx + 1; i++) {
193 - int n_token = smem_expert_first_token_offset[i] -
194 - smem_expert_first_token_offset[i - 1];
195 - accumulate_align_offset =
196 - accumulate_align_offset + (n_token + align_block_size - 1) /
197 - align_block_size * align_block_size;
198 - if (i == eidx) {
199 - first_token_offset = accumulate_align_offset;
200 - }
201 - // last block store align_expert_first_token_offset
202 - if (eidx == num_local_expert - 1 && threadIdx.x == 0) {
203 - align_expert_first_token_offset[i] = accumulate_align_offset;
204 - }
205 - }
206 - }
207 - for (int idx = tidx; idx < n_token_in_expert; idx += blockDim.x) {
208 - // update m_indice with expert id
209 - m_indices[first_token_offset + idx] = eidx;
210 - }
211 -}
212 -
213 -void getMIndices(int64_t* expert_first_token_offset,
214 - int64_t* align_expert_first_token_offset, int* m_indices,
215 - int num_local_expert, const int align_block_size,
216 - cudaStream_t stream) {
217 - int block = 256;
218 - int grid = num_local_expert;
219 - int smem_size = sizeof(int64_t) * (num_local_expert + 1);
220 - if (align_block_size == -1) {
221 - getMIndicesKernel<false><<<grid, block, smem_size, stream>>>(
222 - expert_first_token_offset, align_expert_first_token_offset, m_indices,
223 - num_local_expert, align_block_size);
224 - } else {
225 - getMIndicesKernel<true><<<grid, block, smem_size, stream>>>(
226 - expert_first_token_offset, align_expert_first_token_offset, m_indices,
227 - num_local_expert, align_block_size);
228 - }
229 -}
230 -
171 #endif
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+2 -8
@@ -60,10 +60,9 @@ 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 const* expert_first_token_offset,
64 - int64_t const* aligned_expert_first_token_offset, int64_t const num_rows,
63 + int64_t const* expert_first_token_offset, int64_t const num_rows,
64 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);
65 + int num_local_experts, cudaStream_t stream);
66
67 template <class T, class OutputType>
68 void finalizeMoeRoutingKernelLauncher(
@@ -76,9 +75,4 @@ void preprocessTopkIdLauncher(int* topk_id_ptr, int size,
75 const int* expert_map_ptr, int num_experts,
76 cudaStream_t stream);
77
79 -void getMIndices(int64_t* expert_first_token_offset,
80 - int64_t* align_expert_first_token_offset, int* m_indices,
81 - int num_local_expert, const int align_block_size,
82 - cudaStream_t stream);
83 -
78 #include "moe_permute_unpermute_kernel.inl"
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl
+12 -35
@@ -1,14 +1,13 @@
1 #pragma once
2
3 -template <typename T, bool CHECK_SKIPPED, bool ALIGN_BLOCK_SIZE>
3 +template <typename T, bool CHECK_SKIPPED>
4 __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 const* expert_first_token_offset,
9 - int64_t const* aligned_expert_first_token_offset, int64_t const num_rows,
8 + int64_t const* expert_first_token_offset, int64_t const num_rows,
9 int64_t const* num_dest_rows, int64_t const cols, int64_t k,
11 - int num_local_experts, int align_block_size) {
10 + int num_local_experts) {
11 // Reverse permutation map.
12 // I do this so that later, we can use the source -> dest map to do the k-way
13 // reduction and unpermuting. I need the reverse map for that reduction to
@@ -19,24 +18,6 @@ __global__ void expandInputRowsKernel(
18 expanded_dest_row_to_expanded_source_row[expanded_dest_row];
19 int expert_id = sorted_experts[expanded_dest_row];
20
22 - if constexpr (ALIGN_BLOCK_SIZE) {
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 - }
37 - expanded_dest_row = aligned_base + token_offset_in_expert;
38 - }
39 -
21 if (threadIdx.x == 0) {
22 assert(expanded_dest_row <= INT32_MAX);
23 expanded_source_row_to_expanded_dest_row[expanded_source_row] =
@@ -76,29 +57,25 @@ void expandInputRowsKernelLauncher(
57 T const* unpermuted_input, T* permuted_output, int* sorted_experts,
58 int const* expanded_dest_row_to_expanded_source_row,
59 int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
79 - int64_t const* expert_first_token_offset,
80 - int64_t const* aligned_expert_first_token_offset, int64_t const num_rows,
60 + int64_t const* expert_first_token_offset, int64_t const num_rows,
61 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) {
62 + int num_local_experts, cudaStream_t stream) {
63 int64_t const blocks = num_rows * k;
64 int64_t const threads = 256;
85 - using FuncPtr = decltype(&expandInputRowsKernel<T, true, true>);
86 - FuncPtr func_map[2][2] = {
87 - {&expandInputRowsKernel<T, false, false>,
88 - &expandInputRowsKernel<T, false, true>},
89 - {&expandInputRowsKernel<T, true, false>,
90 - &expandInputRowsKernel<T, true, true>},
65 + using FuncPtr = decltype(&expandInputRowsKernel<T, true>);
66 + FuncPtr func_map[2] = {
67 + &expandInputRowsKernel<T, false>,
68 + &expandInputRowsKernel<T, true>,
69 };
70 bool is_check_skip = num_valid_tokens_ptr != nullptr;
93 - bool is_align_block_size = align_block_size != -1;
94 - auto func = func_map[is_check_skip][is_align_block_size];
71 + auto func = func_map[is_check_skip];
72
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,
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);
77 + expert_first_token_offset, num_rows, num_valid_tokens_ptr, cols, k,
78 + num_local_experts);
79 }
80
81 template <class T, class U>
csrc/moe/torch_bindings.cpp
+2 -2
@@ -99,9 +99,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
99 "moe_permute(Tensor input, Tensor topk_ids,"
100 "Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
101 "int n_local_expert,"
102 - "int topk, int? align_block_size,Tensor! permuted_input, Tensor! "
102 + "int topk, Tensor! permuted_input, Tensor! "
103 "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
104 - "permuted_idx, Tensor! m_indices)->()");
104 + "permuted_idx)->()");
105
106 m.def(
107 "moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
tests/kernels/moe/test_moe_permute_unpermute.py
+15 -118
@@ -40,10 +40,8 @@ def torch_permute(
40 n_local_expert: int,
41 start_expert: int,
42 expert_map: torch.Tensor | None = None,
43 - align_block_size: int | None = None,
44 - fill_invalid_expert: int = -1,
43 ) -> list[torch.Tensor]:
46 - n_token, n_hidden = hidden_states.shape[0], hidden_states.shape[1]
44 + n_token = hidden_states.shape[0]
45 if expert_map is not None:
46 is_local_expert = expert_map[topk_ids] != -1
47 not_local_expert = expert_map[topk_ids] == -1
@@ -70,107 +68,19 @@ def torch_permute(
68
69 _, src2dst_idx = torch.sort(dst_row_id2src_row_id_map)
70 valid_row_idx = []
73 - if align_block_size is None:
74 - permuted_hidden_states = hidden_states[dst_row_id2src_row_id_map // topk, ...]
75 - permuted_row_size = permuted_hidden_states.shape[0]
76 - m_indices = torch.empty(
77 - permuted_row_size, device="cuda", dtype=torch.int32
78 - ).fill_(fill_invalid_expert)
79 - for i in range(1, n_local_expert + 1):
80 - first_token_offset = expert_first_token_offset[i - 1]
81 - last_token_offset = expert_first_token_offset[i]
82 - m_indices[first_token_offset:last_token_offset] = i - 1
83 - src_row_id2dst_row_id_map = torch.arange(
84 - 0, n_token * topk, device="cuda", dtype=torch.int32
85 - )[src2dst_idx].reshape((n_token, topk))
86 - valid_row_idx += [i for i in range(expert_first_token_offset[-1])]
87 - dst_row_id2src_row_id_map[expert_first_token_offset[-1] :] = n_token * topk
88 - return [
89 - permuted_hidden_states,
90 - expert_first_token_offset,
91 - src_row_id2dst_row_id_map,
92 - dst_row_id2src_row_id_map,
93 - m_indices,
94 - valid_row_idx,
95 - ]
96 - else:
97 - permuted_row_size = (
98 - (topk * n_token + n_expert * (align_block_size - 1) + align_block_size - 1)
99 - // align_block_size
100 - * align_block_size
101 - )
102 - permuted_idx = torch.full(
103 - (permuted_row_size,),
104 - n_token * topk,
105 - dtype=torch.int32,
106 - device=hidden_states.device,
107 - )
108 - permuted_hidden_states = torch.empty(
109 - (permuted_row_size, n_hidden), device="cuda", dtype=hidden_states.dtype
110 - )
111 - align_src_row_id2dst_row_id = torch.empty(
112 - n_token * topk, device="cuda", dtype=torch.int32
113 - )
114 - align_expert_first_token_offset = torch.zeros_like(expert_first_token_offset)
115 - m_indices = torch.empty(
116 - permuted_row_size, device="cuda", dtype=torch.int32
117 - ).fill_(fill_invalid_expert)
118 - # get align_permuted_hidden_states,
119 - # valid row_idx and align_expert_first_token_offset
120 - for i in range(1, n_local_expert + 1):
121 - first_token_offset = expert_first_token_offset[i - 1]
122 - last_token_offset = expert_first_token_offset[i]
123 - n_token_in_expert = last_token_offset - first_token_offset
124 - align_expert_first_token_offset[i] = (
125 - align_expert_first_token_offset[i - 1]
126 - + (n_token_in_expert + align_block_size - 1)
127 - // align_block_size
128 - * align_block_size
129 - )
130 - align_first_token_offset = align_expert_first_token_offset[i - 1]
131 - align_last_token_offset = align_expert_first_token_offset[i]
132 - dst_row_id2src_row_id_in_expert = dst_row_id2src_row_id_map[
133 - first_token_offset : first_token_offset + n_token_in_expert
134 - ]
135 - # store token in current expert with align_first_token_offset
136 - permuted_hidden_states[
137 - align_first_token_offset : align_first_token_offset + n_token_in_expert,
138 - ...,
139 - ] = hidden_states[dst_row_id2src_row_id_in_expert // topk, ...]
140 - permuted_idx[
141 - align_first_token_offset : align_first_token_offset + n_token_in_expert
142 - ] = dst_row_id2src_row_id_in_expert
143 - # set current expert m_indices
144 - m_indices[align_first_token_offset:align_last_token_offset] = i - 1
145 - valid_row_idx += [
146 - i
147 - for i in range(
148 - align_first_token_offset,
149 - align_first_token_offset + n_token_in_expert,
150 - )
151 - ]
152 - # get align_src_row_id2dst_row_id
153 - for i in range(n_token * topk):
154 - eid = sorted_topk_ids[i]
155 - if eid >= n_local_expert:
156 - # check token not in local expert
157 - align_src_row_id2dst_row_id[i] = align_expert_first_token_offset[-1]
158 - continue
159 - first_token_offset = expert_first_token_offset[eid]
160 - align_first_token_offset = align_expert_first_token_offset[eid]
161 - token_offset = i - first_token_offset
162 - align_src_row_id2dst_row_id[i] = align_first_token_offset + token_offset
163 - align_src_row_id2dst_row_id = align_src_row_id2dst_row_id[src2dst_idx].reshape(
164 - (n_token, topk)
165 - )
166 - return [
167 - permuted_hidden_states,
168 - align_expert_first_token_offset,
169 - align_src_row_id2dst_row_id,
170 - permuted_idx,
171 - m_indices,
172 - valid_row_idx,
173 - ]
71 + permuted_hidden_states = hidden_states[dst_row_id2src_row_id_map // topk, ...]
72 + src_row_id2dst_row_id_map = torch.arange(
73 + 0, n_token * topk, device="cuda", dtype=torch.int32
74 + )[src2dst_idx].reshape((n_token, topk))
75 + valid_row_idx += [i for i in range(expert_first_token_offset[-1])]
76 + dst_row_id2src_row_id_map[expert_first_token_offset[-1] :] = n_token * topk
77 + return [
78 + permuted_hidden_states,
79 + expert_first_token_offset,
80 + src_row_id2dst_row_id_map,
81 + dst_row_id2src_row_id_map,
82 + valid_row_idx,
83 + ]
84
85
86 def torch_unpermute(
@@ -207,7 +117,6 @@ def torch_unpermute(
117 @pytest.mark.parametrize("topk", TOP_KS)
118 @pytest.mark.parametrize("dtype", [torch.bfloat16])
119 @pytest.mark.parametrize("ep_size", EP_SIZE)
210 -@pytest.mark.parametrize("align_block_size", [None, 128])
120 def test_moe_permute_unpermute(
121 n_token: int,
122 n_hidden: int,
@@ -215,11 +124,9 @@ def test_moe_permute_unpermute(
124 n_expert: int,
125 ep_size: int,
126 dtype: torch.dtype,
218 - align_block_size: int | None,
127 ):
128 if not moe_permute_unpermute_supported():
129 pytest.skip("moe_permute_unpermute is not supported on this platform.")
222 - fill_invalid_expert = 0
130 ep_rank = np.random.randint(0, ep_size)
131 expert_map = None
132 n_local_expert = n_expert
@@ -238,7 +145,6 @@ def test_moe_permute_unpermute(
145 gold_expert_first_token_offset,
146 gold_inv_permuted_idx,
147 gold_permuted_idx,
241 - gold_m_indices,
148 valid_row_idx,
149 ) = torch_permute(
150 hidden_states,
@@ -249,8 +155,6 @@ def test_moe_permute_unpermute(
155 n_local_expert,
156 start_expert,
157 expert_map=expert_map,
252 - align_block_size=align_block_size,
253 - fill_invalid_expert=fill_invalid_expert,
158 )
159
160 (
@@ -258,7 +162,7 @@ def test_moe_permute_unpermute(
162 _,
163 expert_first_token_offset,
164 inv_permuted_idx,
261 - m_indices,
165 + _,
166 ) = moe_permute(
167 hidden_states=hidden_states,
168 a1q_scale=None,
@@ -266,8 +170,6 @@ def test_moe_permute_unpermute(
170 n_expert=n_expert,
171 n_local_expert=n_local_expert,
172 expert_map=expert_map,
269 - align_block_size=align_block_size,
270 - fill_invalid_expert=fill_invalid_expert,
173 )
174
175 # check expert_first_token_offset
@@ -278,11 +180,6 @@ def test_moe_permute_unpermute(
180 torch.testing.assert_close(
181 gold_inv_permuted_idx.flatten(), inv_permuted_idx, atol=0, rtol=0
182 )
281 - # check mindice
282 - # current kernel usage assumes deepgemm requires align_block_size
283 - # when it's not provided then we don't compute m_indices (for cutlass)
284 - if align_block_size is not None:
285 - torch.testing.assert_close(gold_m_indices, m_indices, atol=0, rtol=0)
183
184 # check permuted_hidden_states, only valid token
185 torch.testing.assert_close(
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+2 -29
@@ -11,8 +11,6 @@ def moe_permute(
11 n_expert: int,
12 n_local_expert: int = -1,
13 expert_map: torch.Tensor | None = None,
14 - align_block_size: int | None = None,
15 - fill_invalid_expert: int = -1,
14 permuted_hidden_states: torch.Tensor | None = None,
15 ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]:
16 """
@@ -27,9 +25,6 @@ def moe_permute(
25 - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
26 from the global expert space to the local expert space of the expert
27 parallel shard.
30 - - align_block_size (Optional[int]): align group gemm block size for deepgemm
31 - - fill_invalid_expert(int): fill expert id in m_indices for invalid expert
32 - to workaround DeepGemm unsupported -1 in m_indices
28 - permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor.
29 If None, the output tensor will be created in this function.
30 Returns:
@@ -37,12 +32,9 @@ def moe_permute(
32 - a1q_scale (Optional[torch.Tensor]): permuted quant scale for hidden_states
33 if original scale not per-tensor scaling
34 - expert_first_token_offset (torch.Tensor): offset of the first token
40 - of each expert for standard grouped gemm. if enable 'align_block_size'
41 - expert_first_token_offset will align up to 'align_block_size'.
35 + of each expert for standard grouped gemm.
36 - inv_permuted_idx (torch.Tensor): idx map for moe_unpermute.
37 - permuted_idx (torch.Tensor): idx map from hidden to permuted_hidden.
44 - - m_indices: m_indices for grouped gemm in deepgemm,`m_indices[i]` records
45 - the group which the j-th row of the LHS belong to.`
38 """
39 n_token, n_hidden = hidden_states.size()
40 topk = topk_ids.size(1)
@@ -50,17 +42,6 @@ def moe_permute(
42 "permue kernel need hidden dim align to 16B"
43 )
44 permuted_row_size = n_token * topk
53 - if align_block_size is not None:
54 - permuted_row_size = (
55 - (
56 - permuted_row_size
57 - + n_expert * (align_block_size - 1)
58 - + align_block_size
59 - - 1
60 - )
61 - // align_block_size
62 - * align_block_size
63 - )
45 if n_local_expert == -1:
46 n_local_expert = n_expert
47 if permuted_hidden_states is None:
@@ -78,12 +59,6 @@ def moe_permute(
59 0, n_token * topk, dtype=torch.int32, device=hidden_states.device
60 ).reshape((n_token, topk))
61
81 - m_indices = torch.full(
82 - (permuted_row_size,),
83 - fill_invalid_expert,
84 - dtype=torch.int32,
85 - device=hidden_states.device,
86 - )
62 expert_first_token_offset = torch.empty(
63 n_local_expert + 1, dtype=torch.int64, device=hidden_states.device
64 )
@@ -105,12 +80,10 @@ def moe_permute(
80 n_expert,
81 n_local_expert,
82 topk,
108 - align_block_size,
83 permuted_hidden_states,
84 expert_first_token_offset,
85 inv_permuted_idx,
86 permuted_idx,
113 - m_indices,
87 )
88
89 if a1q_scale is not None and a1q_scale.dim() > 1:
@@ -120,7 +93,7 @@ def moe_permute(
93 a1q_scale,
94 expert_first_token_offset,
95 inv_permuted_idx.flatten(),
123 - m_indices,
96 + permuted_idx,
97 )
98
99