@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3e887d2e0c
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 15, 12:15 UTC

permute/unpermute kernel for moe optimization (#14568)

Signed-off-by: Caleb_Du <Caleb_Du@zju.edu.cn>

Caleb_Du committed May 3, 2025 at 02:31 UTC 3e887d2e0c1fcd65efbfe02db8a824c761fe4d41
19 files changed +1474 -28
CMakeLists.txt
+13 -1
@@ -15,7 +15,6 @@ project(vllm_extensions LANGUAGES CXX)
15
16 # CUDA by default, can be overridden by using -DVLLM_TARGET_DEVICE=... (used by setup.py)
17 set(VLLM_TARGET_DEVICE "cuda" CACHE STRING "Target device backend for vLLM")
18 -
18 message(STATUS "Build type: ${CMAKE_BUILD_TYPE}")
19 message(STATUS "Target device: ${VLLM_TARGET_DEVICE}")
20
@@ -682,6 +681,17 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
681 endif()
682 endif()
683
684 +if(VLLM_GPU_LANG STREQUAL "CUDA")
685 + set(MOE_PERMUTE_SRC
686 + "csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu"
687 + "csrc/moe/moe_permute_unpermute_op.cu")
688 +
689 + set_gencode_flags_for_srcs(
690 + SRCS "${MARLIN_PERMUTE_SRC}"
691 + CUDA_ARCHS "${MOE_PERMUTE_ARCHS}")
692 +
693 + list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}")
694 +endif()
695 message(STATUS "Enabling moe extension.")
696 define_gpu_extension_target(
697 _moe_C
@@ -690,6 +700,8 @@ define_gpu_extension_target(
700 SOURCES ${VLLM_MOE_EXT_SRC}
701 COMPILE_FLAGS ${VLLM_GPU_FLAGS}
702 ARCHITECTURES ${VLLM_GPU_ARCHES}
703 + INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
704 + INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
705 USE_SABI 3
706 WITH_SOABI)
707
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py
+2 -1
@@ -90,7 +90,8 @@ def bench_run(results: list[benchmark.Measurement], model: str,
90
91 score = torch.randn((m, num_experts), device="cuda", dtype=dtype)
92
93 - topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
93 + topk_weights, topk_ids, token_expert_indices = fused_topk(
94 + a, score, topk, renormalize=False)
95
96 def run_triton_moe(a: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
97 topk_weights: torch.Tensor, topk_ids: torch.Tensor,
benchmarks/kernels/benchmark_moe.py
+2 -2
@@ -115,8 +115,8 @@ def benchmark_config(config: BenchmarkConfig,
115 from vllm.model_executor.layers.fused_moe import override_config
116 with override_config(config):
117 if use_deep_gemm:
118 - topk_weights, topk_ids = fused_topk(x, input_gating, topk,
119 - False)
118 + topk_weights, topk_ids, token_expert_indices = fused_topk(
119 + x, input_gating, topk, False)
120 return fused_experts(
121 x,
122 w1,
benchmarks/kernels/benchmark_moe_permute_unpermute.py new
+349
@@ -0,0 +1,349 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import argparse
4 +from typing import Any, TypedDict
5 +
6 +import ray
7 +import torch
8 +from transformers import AutoConfig
9 +
10 +from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
11 + _moe_permute, _moe_unpermute_and_reduce)
12 +from vllm.model_executor.layers.fused_moe.fused_moe import *
13 +from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import *
14 +from vllm.model_executor.layers.fused_moe.utils import _fp8_quantize
15 +from vllm.platforms import current_platform
16 +from vllm.utils import FlexibleArgumentParser
17 +
18 +FP8_DTYPE = current_platform.fp8_dtype()
19 +
20 +
21 +class BenchmarkConfig(TypedDict):
22 + BLOCK_SIZE_M: int
23 + BLOCK_SIZE_N: int
24 + BLOCK_SIZE_K: int
25 + GROUP_SIZE_M: int
26 + num_warps: int
27 + num_stages: int
28 +
29 +
30 +def benchmark_permute(num_tokens: int,
31 + num_experts: int,
32 + hidden_size: int,
33 + topk: int,
34 + dtype: torch.dtype,
35 + use_fp8_w8a8: bool,
36 + use_int8_w8a16: bool,
37 + num_iters: int = 100,
38 + use_customized_permute: bool = False) -> float:
39 + # init_dtype = torch.float16 if use_fp8_w8a8 else dtype
40 + hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
41 + # output_hidden_states = torch.empty_like(hidden_states)
42 + if use_fp8_w8a8:
43 + align_block_size = 128 # deepgemm needs 128 m aligned block
44 + qhidden_states, scale = _fp8_quantize(hidden_states, None, None)
45 + else:
46 + align_block_size = None
47 + qhidden_states = hidden_states
48 +
49 + gating_output = torch.randn(num_iters,
50 + num_tokens,
51 + num_experts,
52 + dtype=torch.float32)
53 +
54 + input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32)
55 + topk_weights, topk_ids, token_expert_indices = fused_topk(
56 + qhidden_states, input_gating, topk, False)
57 +
58 + def prepare(i: int):
59 + input_gating.copy_(gating_output[i])
60 +
61 + def run():
62 + if use_customized_permute:
63 + (permuted_hidden_states, first_token_off, inv_perm_idx,
64 + m_indices) = moe_permute(
65 + qhidden_states,
66 + topk_weights=topk_weights,
67 + topk_ids=topk_ids,
68 + token_expert_indices=token_expert_indices,
69 + topk=topk,
70 + n_expert=num_experts,
71 + n_local_expert=num_experts,
72 + expert_map=None,
73 + align_block_size=align_block_size,
74 + )
75 + else:
76 + (permuted_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
77 + inv_perm) = _moe_permute(qhidden_states, None, topk_ids,
78 + num_experts, None, align_block_size)
79 +
80 + # JIT compilation & warmup
81 + run()
82 + torch.cuda.synchronize()
83 +
84 + # Capture 10 invocations with CUDA graph
85 + graph = torch.cuda.CUDAGraph()
86 + with torch.cuda.graph(graph):
87 + for _ in range(10):
88 + run()
89 + torch.cuda.synchronize()
90 +
91 + # Warmup
92 + for _ in range(5):
93 + graph.replay()
94 + torch.cuda.synchronize()
95 +
96 + start_event = torch.cuda.Event(enable_timing=True)
97 + end_event = torch.cuda.Event(enable_timing=True)
98 +
99 + latencies: list[float] = []
100 + for i in range(num_iters):
101 + prepare(i)
102 + torch.cuda.synchronize()
103 +
104 + start_event.record()
105 + graph.replay()
106 + end_event.record()
107 + end_event.synchronize()
108 + latencies.append(start_event.elapsed_time(end_event))
109 + avg = sum(latencies) / (num_iters * 10) * 1000 # us
110 + graph.reset()
111 + return avg
112 +
113 +
114 +def benchmark_unpermute(num_tokens: int,
115 + num_experts: int,
116 + hidden_size: int,
117 + topk: int,
118 + dtype: torch.dtype,
119 + use_fp8_w8a8: bool,
120 + use_int8_w8a16: bool,
121 + num_iters: int = 100,
122 + use_customized_permute: bool = False) -> float:
123 + # init_dtype = torch.float16 if use_fp8_w8a8 else dtype
124 + hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
125 + output_hidden_states = torch.empty_like(hidden_states)
126 + if use_fp8_w8a8:
127 + align_block_size = 128 # deepgemm needs 128 m aligned block
128 + qhidden_states, scale = _fp8_quantize(hidden_states, None, None)
129 + else:
130 + align_block_size = None
131 + qhidden_states = hidden_states
132 +
133 + input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32)
134 +
135 + topk_weights, topk_ids, token_expert_indices = fused_topk(
136 + qhidden_states, input_gating, topk, False)
137 +
138 + def prepare():
139 + if use_customized_permute:
140 + (permuted_hidden_states, first_token_off, inv_perm_idx,
141 + m_indices) = moe_permute(
142 + qhidden_states,
143 + topk_weights=topk_weights,
144 + topk_ids=topk_ids,
145 + token_expert_indices=token_expert_indices,
146 + topk=topk,
147 + n_expert=num_experts,
148 + n_local_expert=num_experts,
149 + expert_map=None,
150 + align_block_size=align_block_size,
151 + )
152 + # convert to fp16/bf16 as gemm output
153 + return (permuted_hidden_states.to(dtype), first_token_off,
154 + inv_perm_idx, m_indices)
155 + else:
156 + (permuted_qhidden_states, a1q_scale, sorted_token_ids, expert_ids,
157 + inv_perm) = _moe_permute(qhidden_states, None, topk_ids,
158 + num_experts, None, align_block_size)
159 + # convert to fp16/bf16 as gemm output
160 + return (permuted_qhidden_states.to(dtype), a1q_scale,
161 + sorted_token_ids, expert_ids, inv_perm)
162 +
163 + def run(input: tuple):
164 + if use_customized_permute:
165 + (permuted_hidden_states, first_token_off, inv_perm_idx,
166 + m_indices) = input
167 + moe_unpermute(permuted_hidden_states, topk_weights, topk_ids,
168 + inv_perm_idx, first_token_off, topk, num_experts,
169 + num_experts)
170 + else:
171 + (permuted_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
172 + inv_perm) = input
173 + _moe_unpermute_and_reduce(output_hidden_states,
174 + permuted_hidden_states, inv_perm,
175 + topk_weights)
176 +
177 + # JIT compilation & warmup
178 + input = prepare()
179 + run(input)
180 + torch.cuda.synchronize()
181 +
182 + # Capture 10 invocations with CUDA graph
183 + graph = torch.cuda.CUDAGraph()
184 + with torch.cuda.graph(graph):
185 + for _ in range(10):
186 + run(input)
187 + torch.cuda.synchronize()
188 +
189 + # Warmup
190 + for _ in range(5):
191 + graph.replay()
192 + torch.cuda.synchronize()
193 +
194 + start_event = torch.cuda.Event(enable_timing=True)
195 + end_event = torch.cuda.Event(enable_timing=True)
196 +
197 + latencies: list[float] = []
198 + for i in range(num_iters):
199 + torch.cuda.synchronize()
200 + start_event.record()
201 + graph.replay()
202 + end_event.record()
203 + end_event.synchronize()
204 + latencies.append(start_event.elapsed_time(end_event))
205 + avg = sum(latencies) / (num_iters * 10) * 1000 # us
206 + graph.reset()
207 + return avg
208 +
209 +
210 +@ray.remote(num_gpus=1)
211 +class BenchmarkWorker:
212 +
213 + def __init__(self, seed: int) -> None:
214 + torch.set_default_device("cuda")
215 + current_platform.seed_everything(seed)
216 + self.seed = seed
217 + # Get the device ID to allocate tensors and kernels
218 + # on the respective GPU. This is required for Ray to work
219 + # correctly with multi-GPU tuning on the ROCm platform.
220 + self.device_id = int(ray.get_gpu_ids()[0])
221 +
222 + def benchmark(
223 + self,
224 + num_tokens: int,
225 + num_experts: int,
226 + hidden_size: int,
227 + topk: int,
228 + dtype: torch.dtype,
229 + use_fp8_w8a8: bool,
230 + use_int8_w8a16: bool,
231 + use_customized_permute: bool = False,
232 + ) -> tuple[dict[str, int], float]:
233 + current_platform.seed_everything(self.seed)
234 +
235 + permute_time = benchmark_permute(
236 + num_tokens,
237 + num_experts,
238 + hidden_size,
239 + topk,
240 + dtype,
241 + use_fp8_w8a8,
242 + use_int8_w8a16,
243 + num_iters=100,
244 + use_customized_permute=use_customized_permute)
245 + unpermute_time = benchmark_unpermute(
246 + num_tokens,
247 + num_experts,
248 + hidden_size,
249 + topk,
250 + dtype,
251 + use_fp8_w8a8,
252 + use_int8_w8a16,
253 + num_iters=100,
254 + use_customized_permute=use_customized_permute)
255 + return permute_time, unpermute_time
256 +
257 +
258 +def get_weight_block_size_safety(config, default_value=None):
259 +
260 + quantization_config = getattr(config, 'quantization_config', {})
261 + if isinstance(quantization_config, dict):
262 + return quantization_config.get('weight_block_size', default_value)
263 + return default_value
264 +
265 +
266 +def main(args: argparse.Namespace):
267 + print(args)
268 +
269 + config = AutoConfig.from_pretrained(
270 + args.model, trust_remote_code=args.trust_remote_code)
271 + if config.architectures[0] == "DbrxForCausalLM":
272 + E = config.ffn_config.moe_num_experts
273 + topk = config.ffn_config.moe_top_k
274 + elif config.architectures[0] == "JambaForCausalLM":
275 + E = config.num_experts
276 + topk = config.num_experts_per_tok
277 + elif (config.architectures[0] == "DeepseekV3ForCausalLM"
278 + or config.architectures[0] == "DeepseekV2ForCausalLM"):
279 + E = config.n_routed_experts
280 + topk = config.num_experts_per_tok
281 + elif config.architectures[0] in [
282 + "Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM"
283 + ]:
284 + E = config.num_experts
285 + topk = config.num_experts_per_tok
286 +
287 + else:
288 + # Support for llama4
289 + config = config.get_text_config()
290 + # Default: Mixtral.
291 + E = config.num_local_experts
292 + topk = config.num_experts_per_tok
293 +
294 + hidden_size = config.hidden_size
295 + dtype = torch.float16 if current_platform.is_rocm() else config.torch_dtype
296 + use_fp8_w8a8 = args.dtype == "fp8_w8a8"
297 + use_int8_w8a16 = args.dtype == "int8_w8a16"
298 + use_customized_permute = args.use_customized_permute
299 +
300 + if args.batch_size is None:
301 + batch_sizes = [
302 + 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536,
303 + 2048, 3072, 4096
304 + ]
305 + else:
306 + batch_sizes = [args.batch_size]
307 +
308 + ray.init()
309 + num_gpus = int(ray.available_resources()["GPU"])
310 + workers = [BenchmarkWorker.remote(args.seed) for _ in range(num_gpus)]
311 +
312 + def _distribute(method: str, inputs: list[Any]) -> list[Any]:
313 + outputs = []
314 + worker_idx = 0
315 + for input_args in inputs:
316 + worker = workers[worker_idx]
317 + worker_method = getattr(worker, method)
318 + output = worker_method.remote(*input_args)
319 + outputs.append(output)
320 + worker_idx = (worker_idx + 1) % num_gpus
321 + return ray.get(outputs)
322 +
323 + outputs = _distribute(
324 + "benchmark", [(batch_size, E, hidden_size, topk, dtype, use_fp8_w8a8,
325 + use_int8_w8a16, use_customized_permute)
326 + for batch_size in batch_sizes])
327 +
328 + for batch_size, (permute, unpermute) in zip(batch_sizes, outputs):
329 + print(f"Batch size: {batch_size}")
330 + print(f"Permute time: {permute:.2f} us")
331 + print(f"Unpermute time: {unpermute:.2f} us")
332 +
333 +
334 +if __name__ == "__main__":
335 + parser = FlexibleArgumentParser()
336 + parser.add_argument("--model",
337 + type=str,
338 + default="mistralai/Mixtral-8x7B-Instruct-v0.1")
339 + parser.add_argument("--dtype",
340 + type=str,
341 + choices=["auto", "fp8_w8a8", "int8_w8a16"],
342 + default="auto")
343 + parser.add_argument("--use-customized-permute", action="store_true")
344 + parser.add_argument("--seed", type=int, default=0)
345 + parser.add_argument("--batch-size", type=int, required=False)
346 + parser.add_argument("--trust-remote-code", action="store_true")
347 + args = parser.parse_args()
348 +
349 + main(args)
csrc/moe/moe_permute_unpermute_op.cu new
+133
@@ -0,0 +1,133 @@
1 +#include <c10/core/ScalarType.h>
2 +#include <torch/all.h>
3 +#include <ATen/cuda/CUDAContext.h>
4 +#include "permute_unpermute_kernels/moe_permute_unpermute_kernel.h"
5 +#include "permute_unpermute_kernels/dispatch.h"
6 +#include "core/registration.h"
7 +
8 +void moe_permute(
9 + const torch::Tensor& input, // [n_token, hidden]
10 + const torch::Tensor& topk_weights, //[n_token, topk]
11 + torch::Tensor& topk_ids, // [n_token, topk]
12 + const torch::Tensor& token_expert_indicies, // [n_token, topk]
13 + const std::optional<torch::Tensor>& expert_map, // [n_expert]
14 + int64_t n_expert, int64_t n_local_expert, int64_t topk,
15 + const std::optional<int64_t>& align_block_size,
16 + torch::Tensor&
17 + permuted_input, // [topk * n_token/align_block_size_m, hidden]
18 + torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
19 + torch::Tensor& src_row_id2dst_row_id_map, // [n_token, topk]
20 + torch::Tensor& m_indices) { // [align_expand_m]
21 + TORCH_CHECK(topk_weights.scalar_type() == at::ScalarType::Float,
22 + "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_indicies.scalar_type() == at::ScalarType::Int,
28 + "token_expert_indicies must be int32");
29 + TORCH_CHECK(src_row_id2dst_row_id_map.scalar_type() == at::ScalarType::Int,
30 + "src_row_id2dst_row_id_map 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")
33 + TORCH_CHECK(
34 + src_row_id2dst_row_id_map.sizes() == token_expert_indicies.sizes(),
35 + "token_expert_indicies shape must be same as src_row_id2dst_row_id_map");
36 + auto n_token = input.sizes()[0];
37 + auto n_hidden = input.sizes()[1];
38 + auto align_block_size_value =
39 + align_block_size.has_value() ? align_block_size.value() : -1;
40 + auto stream = at::cuda::getCurrentCUDAStream().stream();
41 + const long sorter_size =
42 + CubKeyValueSorter::getWorkspaceSize(n_token * topk, n_expert);
43 + auto sort_workspace = torch::empty(
44 + {sorter_size},
45 + torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
46 + auto permuted_experts_id = torch::empty_like(topk_ids);
47 + auto dst_row_id2src_row_id_map = torch::empty_like(src_row_id2dst_row_id_map);
48 + auto align_expert_first_token_offset =
49 + torch::zeros_like(expert_first_token_offset);
50 +
51 + CubKeyValueSorter sorter{};
52 + int64_t* valid_num_ptr = nullptr;
53 + // pre-process kernel for expert-parallelism:
54 + // no local expert id plus "n_expert" offset for priority to local expert
55 + // map local expert id [n, .., n+n_local_expert-1] to [0, n_local_expert -1]
56 + // For example, 4 expert with ep_size=2. ep_rank=1 owns global expert id
57 + // [2,3] with expert_map[-1, -1, 0, 1], preprocess_topk_id process topk_ids
58 + // and map global expert id [2, 3] to local_expert id [0, 1] and map global
59 + // expert id [0, 1] ( not in ep rank=1) to [4, 5] by plus n_expert. This map
60 + // operation is to make local expert high priority in following sort topk_ids
61 + // and scan local expert_first_token_offset for each ep rank for next group
62 + // gemm.
63 + if (expert_map.has_value()) {
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;
67 + preprocessTopkIdLauncher(get_ptr<int>(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
71 + sortAndScanExpert(get_ptr<int>(topk_ids), get_ptr<int>(token_expert_indicies),
72 + get_ptr<int>(permuted_experts_id),
73 + get_ptr<int>(dst_row_id2src_row_id_map),
74 + get_ptr<int64_t>(expert_first_token_offset), n_token,
75 + n_expert, n_local_expert, topk, sorter,
76 + get_ptr<int>(sort_workspace), stream);
77 +
78 + // dispatch expandInputRowsKernelLauncher
79 + MOE_DISPATCH(input.scalar_type(), [&] {
80 + expandInputRowsKernelLauncher<scalar_t>(
81 + get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
82 + get_ptr<float>(topk_weights), get_ptr<int>(permuted_experts_id),
83 + get_ptr<int>(dst_row_id2src_row_id_map),
84 + get_ptr<int>(src_row_id2dst_row_id_map),
85 + get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
86 + n_hidden, topk, n_local_expert, align_block_size_value, stream);
87 + });
88 +
89 + // get m_indices and update expert_first_token_offset with align block
90 + getMIndices(get_ptr<int64_t>(expert_first_token_offset),
91 + get_ptr<int64_t>(align_expert_first_token_offset),
92 + get_ptr<int>(m_indices), n_local_expert, align_block_size_value,
93 + stream);
94 + if (align_block_size.has_value()) {
95 + // update align_expert_first_token_offset
96 + expert_first_token_offset.copy_(align_expert_first_token_offset);
97 + }
98 +}
99 +
100 +void moe_unpermute(
101 + const torch::Tensor& permuted_hidden_states, // [n_token * topk, hidden]
102 + const torch::Tensor& topk_weights, //[n_token, topk]
103 + const torch::Tensor& topk_ids, // [n_token, topk]
104 + const torch::Tensor& src_row_id2dst_row_id_map, // [n_token, topk]
105 + const torch::Tensor& expert_first_token_offset, // [n_local_expert+1]
106 + int64_t n_expert, int64_t n_local_expert, int64_t topk,
107 + torch::Tensor& hidden_states // [n_token, hidden]
108 +) {
109 + TORCH_CHECK(src_row_id2dst_row_id_map.sizes() == topk_ids.sizes(),
110 + "topk_ids shape must be same as src_row_id2dst_row_id_map");
111 + TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
112 + "topk_ids must be int32");
113 + TORCH_CHECK(
114 + permuted_hidden_states.scalar_type() == hidden_states.scalar_type(),
115 + "topk_ids dtype must be same as src_row_id2dst_row_id_map");
116 + auto n_token = hidden_states.size(0);
117 + auto n_hidden = hidden_states.size(1);
118 + auto stream = at::cuda::getCurrentCUDAStream().stream();
119 + const int64_t* valid_ptr =
120 + get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
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),
125 + get_ptr<int>(src_row_id2dst_row_id_map), get_ptr<int>(topk_ids),
126 + n_token, n_hidden, topk, valid_ptr, stream);
127 + });
128 +}
129 +
130 +TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
131 + m.impl("moe_permute", &moe_permute);
132 + m.impl("moe_unpermute", &moe_unpermute);
133 +}
\ No newline at end of file
csrc/moe/permute_unpermute_kernels/dispatch.h new
+53
@@ -0,0 +1,53 @@
1 +#pragma once
2 +#include <cuda_fp8.h>
3 +#define MOE_SWITCH(TYPE, ...) \
4 + at::ScalarType _st = ::detail::scalar_type(TYPE); \
5 + switch (_st) { \
6 + __VA_ARGS__ \
7 + default: \
8 + TORCH_CHECK(false, "[moe permute]data type dispatch fail!") \
9 + }
10 +
11 +#define MOE_DISPATCH_CASE(enum_type, ...) \
12 + case enum_type: { \
13 + using scalar_t = ScalarType2CudaType<enum_type>::type; \
14 + __VA_ARGS__(); \
15 + break; \
16 + }
17 +#define MOE_DISPATCH_FLOAT_CASE(...) \
18 + MOE_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
19 + MOE_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
20 + MOE_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
21 + MOE_DISPATCH_CASE(at::ScalarType::Float8_e5m2, __VA_ARGS__) \
22 + MOE_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__)
23 +
24 +#define MOE_DISPATCH(TYPE, ...) \
25 + MOE_SWITCH(TYPE, MOE_DISPATCH_FLOAT_CASE(__VA_ARGS__))
26 +
27 +template <at::ScalarType type>
28 +struct ScalarType2CudaType;
29 +
30 +template <>
31 +struct ScalarType2CudaType<at::ScalarType::Float> {
32 + using type = float;
33 +};
34 +template <>
35 +struct ScalarType2CudaType<at::ScalarType::Half> {
36 + using type = half;
37 +};
38 +template <>
39 +struct ScalarType2CudaType<at::ScalarType::BFloat16> {
40 + using type = __nv_bfloat16;
41 +};
42 +
43 +// #if __CUDA_ARCH__ >= 890
44 +// fp8
45 +template <>
46 +struct ScalarType2CudaType<at::ScalarType::Float8_e5m2> {
47 + using type = __nv_fp8_e5m2;
48 +};
49 +template <>
50 +struct ScalarType2CudaType<at::ScalarType::Float8_e4m3fn> {
51 + using type = __nv_fp8_e4m3;
52 +};
53 +// #endif
\ No newline at end of file
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu new
+229
@@ -0,0 +1,229 @@
1 +
2 +#include "moe_permute_unpermute_kernel.h"
3 +
4 +// CubKeyValueSorter definition begin
5 +CubKeyValueSorter::CubKeyValueSorter()
6 + : num_experts_(0), num_bits_(sizeof(int) * 8) {}
7 +
8 +int CubKeyValueSorter::expertsToBits(int num_experts) {
9 + // Max value we represent is V = num_experts + (num_experts - 1) = 2 *
10 + // num_experts - 1 The maximum number of bits is therefore floor(log2(V)) + 1
11 + return static_cast<int>(log2(2 * num_experts - 1)) + 1;
12 +}
13 +
14 +CubKeyValueSorter::CubKeyValueSorter(int const num_experts)
15 + : num_experts_(num_experts), num_bits_(expertsToBits(num_experts)) {}
16 +
17 +void CubKeyValueSorter::updateNumExperts(int const num_experts) {
18 + num_experts_ = num_experts;
19 + num_bits_ = expertsToBits(num_experts);
20 +}
21 +
22 +size_t CubKeyValueSorter::getWorkspaceSize(size_t const num_key_value_pairs,
23 + int const num_experts) {
24 + int num_bits = expertsToBits(num_experts);
25 + size_t required_storage = 0;
26 + int* null_int = nullptr;
27 + cub::DeviceRadixSort::SortPairs(nullptr, required_storage, null_int, null_int,
28 + null_int, null_int, num_key_value_pairs, 0,
29 + num_bits);
30 +
31 + // when num_key_value_pairs, num_experts, num_bits, required_storage = 64,
32 + // 4, 3, 0 The required_storage seems to vary between 0 and 1 for the same
33 + // inputs
34 + if (required_storage == 0) {
35 + required_storage = 1;
36 + }
37 + return required_storage;
38 +}
39 +
40 +void CubKeyValueSorter::run(void* workspace, size_t const workspace_size,
41 + int const* keys_in, int* keys_out,
42 + int const* values_in, int* values_out,
43 + size_t const num_key_value_pairs,
44 + cudaStream_t stream) {
45 + size_t expected_ws_size = getWorkspaceSize(num_key_value_pairs, num_experts_);
46 + size_t actual_ws_size = workspace_size;
47 +
48 + TORCH_CHECK(expected_ws_size <= workspace_size,
49 + "[CubKeyValueSorter::run] The allocated workspace is too small "
50 + "to run this problem.");
51 + cub::DeviceRadixSort::SortPairs(workspace, actual_ws_size, keys_in, keys_out,
52 + values_in, values_out, num_key_value_pairs, 0,
53 + num_bits_, stream);
54 +}
55 +// CubKeyValueSorter definition end
56 +
57 +static inline size_t pad_to_multiple_of_16(size_t const& input) {
58 + static constexpr int ALIGNMENT = 16;
59 + return ALIGNMENT * ((input + ALIGNMENT - 1) / ALIGNMENT);
60 +}
61 +template <class T>
62 +__device__ inline int64_t findTotalEltsLessThanTarget(T const* sorted_indices,
63 + int64_t const arr_length,
64 + T const target) {
65 + int64_t low = 0, high = arr_length - 1, target_location = -1;
66 + while (low <= high) {
67 + int64_t mid = (low + high) / 2;
68 +
69 + if (sorted_indices[mid] >= target) {
70 + high = mid - 1;
71 + } else {
72 + low = mid + 1;
73 + target_location = mid;
74 + }
75 + }
76 + return target_location + 1;
77 +}
78 +
79 +// Calculates the start offset of the tokens for a given expert. The last
80 +// element is the total number of valid tokens
81 +__global__ void computeExpertFirstTokenOffsetKernel(
82 + int const* sorted_experts, int64_t const sorted_experts_len,
83 + int const num_experts, int64_t* expert_first_token_offset) {
84 + // First, compute the global tid. We only need 1 thread per expert.
85 + int const expert = blockIdx.x * blockDim.x + threadIdx.x;
86 +
87 + // Note that expert goes [0, num_experts] (inclusive) because we want a count
88 + // for the total number of active tokens at the end of the scan.
89 + if (expert >= num_experts + 1) {
90 + return;
91 + }
92 + expert_first_token_offset[expert] =
93 + findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, expert);
94 +}
95 +
96 +void computeExpertFirstTokenOffset(int const* sorted_indices,
97 + int const total_indices,
98 + int const num_experts,
99 + int64_t* expert_first_token_offset,
100 + cudaStream_t stream) {
101 + int const num_entries = num_experts + 1;
102 + int const threads = std::min(1024, num_entries);
103 + int const blocks = (num_entries + threads - 1) / threads;
104 +
105 + computeExpertFirstTokenOffsetKernel<<<blocks, threads, 0, stream>>>(
106 + sorted_indices, total_indices, num_experts, expert_first_token_offset);
107 +}
108 +
109 +void sortAndScanExpert(int* expert_for_source_row, const int* source_rows,
110 + int* permuted_experts, int* permuted_rows,
111 + int64_t* expert_first_token_offset, int num_rows,
112 + int num_experts, int num_experts_per_node, int k,
113 + CubKeyValueSorter& sorter, void* sorter_ws,
114 + cudaStream_t stream) {
115 + int64_t const expanded_num_rows = static_cast<int64_t>(k) * num_rows;
116 + // We need to use the full num_experts because that is the sentinel value used
117 + // by topk for disabled experts
118 + sorter.updateNumExperts(num_experts);
119 + size_t const sorter_ws_size_bytes = pad_to_multiple_of_16(
120 + sorter.getWorkspaceSize(expanded_num_rows, num_experts));
121 + sorter.run((void*)sorter_ws, sorter_ws_size_bytes, expert_for_source_row,
122 + permuted_experts, source_rows, permuted_rows, expanded_num_rows,
123 + stream);
124 + computeExpertFirstTokenOffset(permuted_experts, expanded_num_rows,
125 + num_experts_per_node, expert_first_token_offset,
126 + stream);
127 +}
128 +
129 +__global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
130 + const int* expert_map_ptr,
131 + int num_experts) {
132 + auto tidx = threadIdx.x;
133 + auto bidx = blockIdx.x;
134 + auto lidx = tidx & 31;
135 + auto widx = tidx >> 5;
136 + auto warp_count = (blockDim.x + 31) >> 5;
137 + auto offset = bidx * blockDim.x;
138 + auto bound = min(offset + blockDim.x, size);
139 + extern __shared__ int smem_expert_map[];
140 + // store expert_map in smem
141 + for (int i = tidx; i < num_experts; i += blockDim.x) {
142 + smem_expert_map[i] = expert_map_ptr[i];
143 + }
144 + __syncthreads();
145 +
146 + // query global expert id in expert map.
147 + // if global expert id = -1 in exert map, plus n_expert
148 + // else set global expert id = exert map[global expert id]
149 + if (offset + tidx < bound) {
150 + auto topk_id = topk_id_ptr[offset + tidx];
151 + auto local_expert_idx = smem_expert_map[topk_id];
152 + if (local_expert_idx == -1) {
153 + topk_id += num_experts;
154 + } else {
155 + topk_id = local_expert_idx;
156 + }
157 + __syncwarp();
158 + topk_id_ptr[offset + tidx] = topk_id;
159 + }
160 +}
161 +void preprocessTopkIdLauncher(int* topk_id_ptr, int size,
162 + const int* expert_map_ptr, int num_experts,
163 + cudaStream_t stream) {
164 + int block = std::min(size, 1024);
165 + int grid = (size + block - 1) / block;
166 + int smem_size = (num_experts) * sizeof(int);
167 + preprocessTopkIdKernel<<<grid, block, smem_size, stream>>>(
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[tidx] = __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 +}
\ No newline at end of file
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h new
+95
@@ -0,0 +1,95 @@
1 +#pragma once
2 +// reference from tensorrt_llm moe kernel implementation archive in
3 +// https://github.com/BBuf/tensorrt-llm-moe/tree/master
4 +
5 +#include <c10/core/ScalarType.h>
6 +#include <torch/all.h>
7 +#include "dispatch.h"
8 +#include <cub/cub.cuh>
9 +#include <cub/device/device_radix_sort.cuh>
10 +#include <cub/util_type.cuh>
11 +#include "cutlass/numeric_size.h"
12 +#include "cutlass/array.h"
13 +
14 +template <typename T>
15 +inline T* get_ptr(torch::Tensor& t) {
16 + return reinterpret_cast<T*>(t.data_ptr());
17 +}
18 +
19 +template <typename T>
20 +inline const T* get_ptr(const torch::Tensor& t) {
21 + return reinterpret_cast<const T*>(t.data_ptr());
22 +}
23 +
24 +class CubKeyValueSorter {
25 + public:
26 + CubKeyValueSorter();
27 +
28 + CubKeyValueSorter(int const num_experts);
29 +
30 + void updateNumExperts(int const num_experts);
31 +
32 + static size_t getWorkspaceSize(size_t const num_key_value_pairs,
33 + int const num_experts);
34 +
35 + void run(void* workspace, size_t const workspace_size, int const* keys_in,
36 + int* keys_out, int const* values_in, int* values_out,
37 + size_t const num_key_value_pairs, cudaStream_t stream);
38 +
39 + private:
40 + static int expertsToBits(int experts);
41 + int num_experts_;
42 + int num_bits_;
43 +};
44 +
45 +void computeExpertFirstTokenOffset(int const* sorted_indices,
46 + int const total_indices,
47 + int const num_experts,
48 + int64_t* expert_first_token_offset,
49 + cudaStream_t stream);
50 +
51 +void sortAndScanExpert(int* expert_for_source_row, const int* source_rows,
52 + int* permuted_experts, int* permuted_rows,
53 + int64_t* expert_first_token_offset, int num_rows,
54 + int num_experts, int num_experts_per_node, int k,
55 + CubKeyValueSorter& sorter, void* sorter_ws,
56 + cudaStream_t stream);
57 +
58 +template <typename T>
59 +void expandInputRowsKernelLauncher(
60 + T const* unpermuted_input, T* permuted_output,
61 + const float* unpermuted_scales, int* sorted_experts,
62 + int const* expanded_dest_row_to_expanded_source_row,
63 + int* expanded_source_row_to_expanded_dest_row,
64 + int64_t* 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 +
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 +
78 +template <class T, class OutputType>
79 +void finalizeMoeRoutingKernelLauncher(
80 + T const* expanded_permuted_rows, OutputType* reduced_unpermuted_output,
81 + 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);
85 +
86 +void preprocessTopkIdLauncher(int* topk_id_ptr, int size,
87 + const int* expert_map_ptr, int num_experts,
88 + cudaStream_t stream);
89 +
90 +void getMIndices(int64_t* expert_first_token_offset,
91 + int64_t* align_expert_first_token_offset, int* m_indices,
92 + int num_local_expert, const int align_block_size,
93 + cudaStream_t stream);
94 +
95 +#include "moe_permute_unpermute_kernel.inl"
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl new
+211
@@ -0,0 +1,211 @@
1 +#pragma once
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,
7 + int const* expanded_dest_row_to_expanded_source_row,
8 + int* expanded_source_row_to_expanded_dest_row,
9 + int64_t* 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.
13 + // I do this so that later, we can use the source -> dest map to do the k-way
14 + // reduction and unpermuting. I need the reverse map for that reduction to
15 + // allow each threadblock to do 1 k-way reduce without atomics later in MoE. 1
16 + // thread block will be responsible for all k summations.
17 + int64_t expanded_dest_row = blockIdx.x;
18 + int64_t const expanded_source_row =
19 + expanded_dest_row_to_expanded_source_row[expanded_dest_row];
20 + int expert_id = sorted_experts[expanded_dest_row];
21 +
22 + extern __shared__ int64_t smem_expert_first_token_offset[];
23 + int64_t align_expanded_row_accumulate = 0;
24 + if constexpr (ALIGN_BLOCK_SIZE) {
25 + // load g2s
26 + for (int idx = threadIdx.x; idx < num_local_experts + 1;
27 + idx += blockDim.x) {
28 + smem_expert_first_token_offset[idx] =
29 + __ldg(expert_first_token_offset + idx);
30 + }
31 + __syncthreads();
32 + int lane_idx = threadIdx.x & 31;
33 +
34 + if (lane_idx == 0) {
35 + // set token_offset_in_expert = 0 if this expert is not local expert
36 + int token_offset_in_expert =
37 + expert_id >= num_local_experts
38 + ? 0
39 + : expanded_dest_row - smem_expert_first_token_offset[expert_id];
40 + int64_t accumulate_align_offset = 0;
41 +#pragma unroll 1
42 + for (int eidx = 1; eidx <= min(expert_id, num_local_experts); eidx++) {
43 + auto n_token_in_expert = smem_expert_first_token_offset[eidx] -
44 + smem_expert_first_token_offset[eidx - 1];
45 + accumulate_align_offset += (n_token_in_expert + align_block_size - 1) /
46 + align_block_size * align_block_size;
47 + }
48 + expanded_dest_row = accumulate_align_offset + token_offset_in_expert;
49 + }
50 + // lane0 shuffle broadcast align_expanded_dest_row
51 + expanded_dest_row = __shfl_sync(0xffffffff, expanded_dest_row, 0);
52 + }
53 +
54 + if (threadIdx.x == 0) {
55 + assert(expanded_dest_row <= INT32_MAX);
56 + expanded_source_row_to_expanded_dest_row[expanded_source_row] =
57 + static_cast<int>(expanded_dest_row);
58 + }
59 +
60 + if (!CHECK_SKIPPED || blockIdx.x < *num_dest_rows) {
61 + // Load 128-bits per thread
62 + constexpr int64_t ELEM_PER_THREAD = 128 / cutlass::sizeof_bits<T>::value;
63 + using DataElem = cutlass::Array<T, ELEM_PER_THREAD>;
64 +
65 + // Duplicate and permute rows
66 + int64_t const source_k_rank = expanded_source_row / num_rows;
67 + int64_t const source_row = expanded_source_row % num_rows;
68 +
69 + auto const* source_row_ptr =
70 + reinterpret_cast<DataElem const*>(unpermuted_input + source_row * cols);
71 + auto* dest_row_ptr =
72 + reinterpret_cast<DataElem*>(permuted_output + expanded_dest_row * cols);
73 +
74 + int64_t const start_offset = threadIdx.x;
75 + int64_t const stride = blockDim.x;
76 + int64_t const num_elems_in_col = cols / ELEM_PER_THREAD;
77 +
78 + for (int elem_index = start_offset; elem_index < num_elems_in_col;
79 + elem_index += stride) {
80 + dest_row_ptr[elem_index] = source_row_ptr[elem_index];
81 + }
82 + }
83 +}
84 +
85 +template <typename T>
86 +void expandInputRowsKernelLauncher(
87 + T const* unpermuted_input, T* permuted_output,
88 + const float* unpermuted_scales, int* sorted_experts,
89 + int const* expanded_dest_row_to_expanded_source_row,
90 + int* expanded_source_row_to_expanded_dest_row,
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) {
94 + int64_t const blocks = num_rows * k;
95 + int64_t const threads = 256;
96 + using FuncPtr = decltype(&expandInputRowsKernel<T, true, true>);
97 + FuncPtr func_map[2][2] = {
98 + {&expandInputRowsKernel<T, false, false>,
99 + &expandInputRowsKernel<T, false, true>},
100 + {&expandInputRowsKernel<T, true, false>,
101 + &expandInputRowsKernel<T, true, true>},
102 + };
103 + bool is_check_skip = num_valid_tokens_ptr != nullptr;
104 + bool is_align_block_size = align_block_size != -1;
105 + auto func = func_map[is_check_skip][is_align_block_size];
106 +
107 + int64_t smem_size = sizeof(int64_t) * (num_local_experts + 1);
108 +
109 + func<<<blocks, threads, smem_size, stream>>>(
110 + unpermuted_input, permuted_output, unpermuted_scales, sorted_experts,
111 + expanded_dest_row_to_expanded_source_row,
112 + expanded_source_row_to_expanded_dest_row, expert_first_token_offset,
113 + num_rows, num_valid_tokens_ptr, cols, k, num_local_experts,
114 + align_block_size);
115 +}
116 +
117 +template <class T, class U>
118 +__host__ __device__ constexpr static U arrayConvert(T const& input) {
119 + using Type = typename U::Element;
120 + static_assert(T::kElements == U::kElements);
121 + U u;
122 +#pragma unroll
123 + for (int i = 0; i < U::kElements; i++) {
124 + u[i] = static_cast<Type>(input[i]);
125 + }
126 + return u;
127 +}
128 +
129 +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,
133 + int const* expert_for_source_row, int64_t const orig_cols, int64_t const k,
134 + int64_t const* num_valid_ptr) {
135 + assert(orig_cols % 4 == 0);
136 + int64_t const original_row = blockIdx.x;
137 + int64_t const num_rows = gridDim.x;
138 + auto const offset = original_row * orig_cols;
139 + OutputType* reduced_row_ptr = reduced_unpermuted_output + offset;
140 + int64_t const num_valid = *num_valid_ptr;
141 +
142 + // Load 128-bits per thread, according to the smallest data type we read/write
143 + constexpr int64_t FINALIZE_ELEM_PER_THREAD =
144 + 128 / std::min(cutlass::sizeof_bits<OutputType>::value,
145 + cutlass::sizeof_bits<T>::value);
146 +
147 + int64_t const start_offset = threadIdx.x;
148 + int64_t const stride = blockDim.x;
149 + int64_t const num_elems_in_col = orig_cols / FINALIZE_ELEM_PER_THREAD;
150 +
151 + using InputElem = cutlass::Array<T, FINALIZE_ELEM_PER_THREAD>;
152 + using OutputElem = cutlass::Array<OutputType, FINALIZE_ELEM_PER_THREAD>;
153 + using ComputeElem = cutlass::Array<float, FINALIZE_ELEM_PER_THREAD>;
154 + auto const* expanded_permuted_rows_v =
155 + reinterpret_cast<InputElem const*>(expanded_permuted_rows);
156 + auto* reduced_row_ptr_v = reinterpret_cast<OutputElem*>(reduced_row_ptr);
157 +
158 +#pragma unroll
159 + for (int elem_index = start_offset; elem_index < num_elems_in_col;
160 + elem_index += stride) {
161 + ComputeElem thread_output;
162 + thread_output.fill(0);
163 + float row_rescale{0.f};
164 + for (int k_idx = 0; k_idx < k; ++k_idx) {
165 + int64_t const expanded_original_row = original_row + k_idx * num_rows;
166 + int64_t const expanded_permuted_row =
167 + expanded_source_row_to_expanded_dest_row[expanded_original_row];
168 +
169 + int64_t const k_offset = original_row * k + k_idx;
170 + float const row_scale = scales[k_offset];
171 +
172 + // Check after row_rescale has accumulated
173 + if (CHECK_SKIPPED && expanded_permuted_row >= num_valid) {
174 + continue;
175 + }
176 +
177 + auto const* expanded_permuted_rows_row_ptr =
178 + expanded_permuted_rows_v + expanded_permuted_row * num_elems_in_col;
179 +
180 + int64_t const expert_idx = expert_for_source_row[k_offset];
181 +
182 + ComputeElem expert_result = arrayConvert<InputElem, ComputeElem>(
183 + expanded_permuted_rows_row_ptr[elem_index]);
184 + thread_output = thread_output + row_scale * (expert_result);
185 + }
186 +
187 + OutputElem output_elem =
188 + arrayConvert<ComputeElem, OutputElem>(thread_output);
189 + reduced_row_ptr_v[elem_index] = output_elem;
190 + }
191 +}
192 +
193 +template <class T, class OutputType>
194 +void finalizeMoeRoutingKernelLauncher(
195 + T const* expanded_permuted_rows, OutputType* reduced_unpermuted_output,
196 + float const* scales, int const* expanded_source_row_to_expanded_dest_row,
197 + int const* expert_for_source_row, int64_t const num_rows,
198 + int64_t const cols, int64_t const k, int64_t const* num_valid_ptr,
199 + cudaStream_t stream) {
200 + int64_t const blocks = num_rows;
201 + int64_t const threads = 256;
202 + bool const check_finished = num_valid_ptr != nullptr;
203 + using FuncPtr = decltype(&finalizeMoeRoutingKernel<T, OutputType, false>);
204 + FuncPtr func_map[2] = {&finalizeMoeRoutingKernel<T, OutputType, false>,
205 + &finalizeMoeRoutingKernel<T, OutputType, true>};
206 + auto* const kernel = func_map[check_finished];
207 + kernel<<<blocks, threads, 0, stream>>>(
208 + expanded_permuted_rows, reduced_unpermuted_output, scales,
209 + expanded_source_row_to_expanded_dest_row, expert_for_source_row, cols, k,
210 + num_valid_ptr);
211 +}
csrc/moe/torch_bindings.cpp
+22
@@ -53,7 +53,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
53 "int size_m, int size_n, int size_k,"
54 "bool is_full_k, bool use_atomic_add,"
55 "bool use_fp32_reduce, bool is_zp_float) -> Tensor");
56 + m.def(
57 + "marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, "
58 + "Tensor! topk_weights, Tensor! topk_ids, Tensor! b_scales, Tensor! "
59 + "b_zeros, Tensor! g_idx, Tensor! perm, Tensor! workspace, "
60 + "int b_q_type, SymInt size_m, "
61 + "SymInt size_n, SymInt size_k, bool is_k_full, int num_experts, int "
62 + "topk, "
63 + "int moe_block_size, bool replicate_input, bool apply_weights)"
64 + " -> Tensor");
65 +
66 + m.def(
67 + "moe_permute(Tensor input, Tensor topk_weight, Tensor! topk_ids,"
68 + "Tensor token_expert_indicies, Tensor? expert_map, int n_expert,"
69 + "int n_local_expert,"
70 + "int topk, int? align_block_size,Tensor! permuted_input, Tensor! "
71 + "expert_first_token_offset, Tensor! src_row_id2dst_row_id_map, Tensor! "
72 + "m_indices)->()");
73
74 + m.def(
75 + "moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
76 + "Tensor topk_ids,Tensor src_row_id2dst_row_id_map, Tensor "
77 + "expert_first_token_offset, int n_expert, int n_local_expert,int "
78 + "topk, Tensor! hidden_states)->()");
79 // conditionally compiled so impl registration is in source file
80
81 #endif
tests/kernels/moe/test_moe.py
+2 -1
@@ -420,7 +420,8 @@ def test_fused_marlin_moe(
420
421 score = torch.randn((m, e), device="cuda", dtype=dtype)
422
423 - topk_weights, topk_ids = fused_topk(a, score, topk, False)
423 + topk_weights, topk_ids, token_expert_indices = fused_topk(
424 + a, score, topk, False)
425
426 torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, e_map)
427
tests/kernels/moe/test_moe_permute_unpermute.py new
+223
@@ -0,0 +1,223 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Tests for the MOE permute/unpermute kernel
3 +
4 +Run `pytest tests/kernels/test_moe_permute_unpermute.py`.
5 +"""
6 +
7 +from typing import Optional
8 +
9 +import numpy as np
10 +import pytest
11 +import torch
12 +
13 +from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
14 +from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
15 +from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
16 + moe_permute, moe_unpermute)
17 +from vllm.platforms import current_platform
18 +
19 +NUM_EXPERTS = [16, 64]
20 +TOP_KS = [2, 4, 6, 8]
21 +EP_SIZE = [1, 4, 16]
22 +current_platform.seed_everything(0)
23 +
24 +
25 +def torch_permute(hidden_states: torch.Tensor,
26 + topk_ids: torch.Tensor,
27 + token_expert_indices: torch.Tensor,
28 + topk: int,
29 + n_expert: int,
30 + n_local_expert: int,
31 + start_expert: int,
32 + expert_map: Optional[torch.Tensor] = None,
33 + align_block_size: Optional[int] = None,
34 + fill_invalid_expert: int = -1) -> list[torch.Tensor]:
35 + n_token, n_hidden = hidden_states.shape[0], hidden_states.shape[1]
36 + if expert_map is not None:
37 + is_local_expert = (expert_map[topk_ids] != -1)
38 + not_local_expert = (expert_map[topk_ids] == -1)
39 + topk_ids = is_local_expert * (
40 + topk_ids - start_expert) + not_local_expert * (topk_ids + n_expert)
41 +
42 + sorted_topk_ids, sorted_indices = torch.sort(topk_ids.flatten(),
43 + stable=True)
44 + dst_row_id2src_row_id_map = token_expert_indices.flatten()[sorted_indices]
45 +
46 + expert_first_token_offset = torch.zeros(n_local_expert + 1,
47 + dtype=torch.int64,
48 + device="cuda")
49 + idx = 0
50 + for i in range(0, n_local_expert):
51 + cnt = 0
52 + while idx < sorted_topk_ids.numel() and sorted_topk_ids[idx] == i:
53 + cnt += 1
54 + idx += 1
55 + expert_first_token_offset[i + 1] = expert_first_token_offset[i] + cnt
56 +
57 + _, src2dst_idx = torch.sort(dst_row_id2src_row_id_map)
58 + valid_row_idx = []
59 + if align_block_size is None:
60 +
61 + permuted_hidden_states = hidden_states[dst_row_id2src_row_id_map %
62 + n_token, ...]
63 + permuted_row_size = permuted_hidden_states.shape[0]
64 + m_indices = torch.empty(permuted_row_size,
65 + device="cuda",
66 + dtype=torch.int32).fill_(fill_invalid_expert)
67 + for i in range(1, n_local_expert + 1):
68 + first_token_offset = expert_first_token_offset[i - 1]
69 + last_token_offset = expert_first_token_offset[i]
70 + m_indices[first_token_offset:last_token_offset] = i - 1
71 + src_row_id2dst_row_id_map = torch.arange(
72 + 0, n_token * topk, device="cuda",
73 + dtype=torch.int32)[src2dst_idx].reshape((n_token, topk))
74 + valid_row_idx += [i for i in range(expert_first_token_offset[-1])]
75 + return [
76 + permuted_hidden_states, expert_first_token_offset,
77 + src_row_id2dst_row_id_map, m_indices, valid_row_idx
78 + ]
79 + else:
80 + permuted_row_size = (topk * n_token + n_expert *
81 + (align_block_size - 1) + align_block_size -
82 + 1) // align_block_size * align_block_size
83 + permuted_hidden_states = torch.empty((permuted_row_size, n_hidden),
84 + device="cuda",
85 + dtype=hidden_states.dtype)
86 + align_src_row_id2dst_row_id = torch.empty(n_token * topk,
87 + device="cuda",
88 + dtype=torch.int32)
89 + align_expert_first_token_offset = torch.zeros_like(
90 + expert_first_token_offset)
91 + m_indices = torch.empty(permuted_row_size,
92 + device="cuda",
93 + dtype=torch.int32).fill_(fill_invalid_expert)
94 + # get align_permuted_hidden_states,
95 + # valid row_idx and align_expert_first_token_offset
96 + for i in range(1, n_local_expert + 1):
97 + first_token_offset = expert_first_token_offset[i - 1]
98 + last_token_offset = expert_first_token_offset[i]
99 + n_token_in_expert = last_token_offset - first_token_offset
100 + align_expert_first_token_offset[
101 + i] = align_expert_first_token_offset[
102 + i - 1] + (n_token_in_expert + align_block_size -
103 + 1) // align_block_size * align_block_size
104 + align_first_token_offset = align_expert_first_token_offset[i - 1]
105 + align_last_token_offset = align_expert_first_token_offset[i]
106 + dst_row_id2src_row_id_in_expert = dst_row_id2src_row_id_map[
107 + first_token_offset:first_token_offset +
108 + n_token_in_expert] % n_token
109 + # store token in current expert with align_first_token_offset
110 + permuted_hidden_states[align_first_token_offset:\
111 + align_first_token_offset+n_token_in_expert,\
112 + ...] = hidden_states[\
113 + dst_row_id2src_row_id_in_expert, ...]
114 + # set current expert m_indices
115 + m_indices[align_first_token_offset:align_last_token_offset] = i - 1
116 + valid_row_idx += [
117 + i for i in range(align_first_token_offset,
118 + align_first_token_offset + n_token_in_expert)
119 + ]
120 + # get align_src_row_id2dst_row_id
121 + for i in range(n_token * topk):
122 + eid = sorted_topk_ids[i]
123 + if (eid >= n_local_expert):
124 + # check token not in local expert
125 + align_src_row_id2dst_row_id[
126 + i] = align_expert_first_token_offset[-1]
127 + continue
128 + first_token_offset = expert_first_token_offset[eid]
129 + align_first_token_offset = align_expert_first_token_offset[eid]
130 + token_offset = i - first_token_offset
131 + align_src_row_id2dst_row_id[
132 + i] = align_first_token_offset + token_offset
133 + align_src_row_id2dst_row_id = align_src_row_id2dst_row_id[\
134 + src2dst_idx].reshape((n_token, topk))
135 + return [
136 + permuted_hidden_states, align_expert_first_token_offset,
137 + align_src_row_id2dst_row_id, m_indices, valid_row_idx
138 + ]
139 +
140 +
141 +def torch_unpermute(permuted_hidden_states: torch.Tensor,
142 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
143 + token_expert_indices: torch.Tensor,
144 + src_row_id2dst_row_id_map: torch.Tensor,
145 + valid_row_idx: torch.Tensor, topk: int,
146 + n_expert: int) -> torch.Tensor:
147 + # ignore invalid row
148 + mask = torch.zeros(permuted_hidden_states.shape[0],
149 + dtype=bool,
150 + device="cuda")
151 + mask[valid_row_idx] = True
152 + permuted_hidden_states[~mask] = 0
153 + idx = src_row_id2dst_row_id_map.flatten()[
154 + token_expert_indices.flatten()].reshape(token_expert_indices.shape)
155 + output = permuted_hidden_states[idx, ...] * topk_weights[..., None]
156 + output = output.sum(dim=1).to(permuted_hidden_states.dtype)
157 + return output
158 +
159 +
160 +@pytest.mark.parametrize("n_token", [1, 33, 64, 222, 1024, 2048, 3000, 5000])
161 +@pytest.mark.parametrize("n_hidden", [2048, 4096, 7168])
162 +@pytest.mark.parametrize("n_expert", NUM_EXPERTS)
163 +@pytest.mark.parametrize("topk", TOP_KS)
164 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
165 +@pytest.mark.parametrize("ep_size", EP_SIZE)
166 +@pytest.mark.parametrize("align_block_size", [None, 128])
167 +def test_moe_permute_unpermute(n_token: int, n_hidden: int, topk: int,
168 + n_expert: int, ep_size: int, dtype: torch.dtype,
169 + align_block_size: Optional[int]):
170 + fill_invalid_expert = 0
171 + ep_rank = np.random.randint(0, ep_size)
172 + expert_map = None
173 + n_local_expert = n_expert
174 + if (ep_size != 1):
175 + n_local_expert, expert_map = determine_expert_map(
176 + ep_size, ep_rank, n_expert)
177 + expert_map = expert_map.cuda()
178 + start_expert = n_local_expert * ep_rank
179 + current_platform.seed_everything(0)
180 + hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype)
181 + gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype)
182 + topk_weights, topk_ids, token_expert_indices = fused_topk(
183 + hidden_states, gating_output, topk, False)
184 + gold0, gold1, gold2, gold3, valid_row_idx = torch_permute(
185 + hidden_states,
186 + topk_ids,
187 + token_expert_indices,
188 + topk,
189 + n_expert,
190 + n_local_expert,
191 + start_expert,
192 + expert_map=expert_map,
193 + align_block_size=align_block_size,
194 + fill_invalid_expert=fill_invalid_expert)
195 +
196 + result0, result1, result2, result3 = moe_permute(
197 + hidden_states, topk_weights, topk_ids, token_expert_indices, topk,
198 + n_expert, n_local_expert, expert_map, align_block_size,
199 + fill_invalid_expert)
200 +
201 + # check expert_first_token_offset
202 + torch.testing.assert_close(gold1, result1, atol=0, rtol=0)
203 + # check src_row_id2dst_row_id_map
204 + torch.testing.assert_close(gold2, result2, atol=0, rtol=0)
205 + # check mindice
206 + torch.testing.assert_close(gold3, result3, atol=0, rtol=0)
207 + # check permuted_hidden_states, only valid token
208 + torch.testing.assert_close(gold0[valid_row_idx],
209 + result0[valid_row_idx],
210 + atol=0,
211 + rtol=0)
212 +
213 + # add a random tensor to simulate group gemm
214 + result0 = 0.5 * result0 + torch.randn_like(result0)
215 +
216 + result4 = moe_unpermute(result0, topk_weights, topk_ids, result2, result1,
217 + topk, n_expert, n_local_expert)
218 + gold4 = torch_unpermute(result0, topk_weights, topk_ids,
219 + token_expert_indices, result2, valid_row_idx, topk,
220 + n_local_expert)
221 +
222 + # check unpermuted hidden
223 + torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0)
tests/kernels/quantization/test_awq_marlin.py
+2 -1
@@ -84,7 +84,8 @@ def test_fused_marlin_moe_awq(
84
85 score = torch.randn((m, e), device="cuda", dtype=dtype)
86
87 - topk_weights, topk_ids = fused_topk(a, score, topk, False)
87 + topk_weights, topk_ids, token_expert_indices = fused_topk(
88 + a, score, topk, False)
89 marlin_output = torch.ops.vllm.fused_marlin_moe(
90 a,
91 qweight1,
tests/kernels/quantization/test_block_fp8.py
+4 -2
@@ -338,7 +338,8 @@ def deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s, score, topk,
338 M, K = a.shape
339 N = w2.shape[-1]
340
341 - topk_weight, topk_ids = fused_topk(a, score.float(), topk, False)
341 + topk_weight, topk_ids, token_expert_indices = fused_topk(
342 + a, score.float(), topk, False)
343
344 block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
345
@@ -435,7 +436,8 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
436 ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score,
437 topk, block_size)
438
438 - topk_weights, topk_ids = fused_topk(a, score.float(), topk, False)
439 + topk_weights, topk_ids, token_expert_indices = fused_topk(
440 + a, score.float(), topk, False)
441
442 out = deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids)
443
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+2 -2
@@ -71,8 +71,8 @@ def single_marlin_moe(
71 E = w.shape[0]
72 N = w.shape[2] // (num_bits // 2)
73
74 - topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk,
75 - renormalize)
74 + topk_weights, topk_ids, token_expert_indices = fused_topk(
75 + hidden_states, gating_output, topk, renormalize)
76
77 # This might not be an optimal config for a single MMM
78 get_config_func = functools.partial(try_get_optimal_moe_config,
vllm/model_executor/layers/fused_moe/fused_moe.py
+9 -10
@@ -854,7 +854,7 @@ def fused_topk(
854 gating_output: torch.Tensor,
855 topk: int,
856 renormalize: bool,
857 -) -> Tuple[torch.Tensor, torch.Tensor]:
857 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
858 assert hidden_states.shape[0] == gating_output.shape[0], (
859 "Number of tokens mismatch")
860
@@ -868,20 +868,19 @@ def fused_topk(
868 topk,
869 dtype=torch.int32,
870 device=hidden_states.device)
871 - token_expert_indicies = torch.empty(M,
872 - topk,
873 - dtype=torch.int32,
874 - device=hidden_states.device)
871 + token_expert_indices = torch.empty(M,
872 + topk,
873 + dtype=torch.int32,
874 + device=hidden_states.device)
875
876 gating_output_float = gating_output.float() # TODO(woosuk): Optimize this.
877
878 topk_func = dispatch_topk_func()
879 topk_weights, topk_ids = topk_func(topk_weights, topk_ids,
880 - token_expert_indicies,
880 + token_expert_indices,
881 gating_output_float, renormalize)
882
883 - del token_expert_indicies # Not used. Will be used in the future.
884 - return topk_weights, topk_ids
883 + return topk_weights, topk_ids, token_expert_indices
884
885
886 # This is used by the Deepseek-V2 and Deepseek-V3 model
@@ -1510,8 +1509,8 @@ def fused_moe(
1509 topk, renormalize,
1510 num_expert_group, topk_group)
1511 elif custom_routing_function is None:
1513 - topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk,
1514 - renormalize)
1512 + topk_weights, topk_ids, token_expert_indices = fused_topk(
1513 + hidden_states, gating_output, topk, renormalize)
1514 else:
1515 topk_weights, topk_ids = custom_routing_function(
1516 hidden_states, gating_output, topk, renormalize)
vllm/model_executor/layers/fused_moe/layer.py
+5 -4
@@ -801,10 +801,11 @@ class FusedMoE(torch.nn.Module):
801 scoring_func=scoring_func,
802 e_score_correction_bias=e_score_correction_bias)
803 elif custom_routing_function is None:
804 - topk_weights, topk_ids = fused_topk(hidden_states=hidden_states,
805 - gating_output=router_logits,
806 - topk=top_k,
807 - renormalize=renormalize)
804 + topk_weights, topk_ids, token_expert_indices = fused_topk(
805 + hidden_states=hidden_states,
806 + gating_output=router_logits,
807 + topk=top_k,
808 + renormalize=renormalize)
809 else:
810 topk_weights, topk_ids = custom_routing_function(
811 hidden_states=hidden_states,
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py new
+116
@@ -0,0 +1,116 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional, Tuple
3 +
4 +import torch
5 +
6 +
7 +def moe_permute(
8 + hidden_states: torch.Tensor,
9 + topk_weights: torch.Tensor,
10 + topk_ids: torch.Tensor,
11 + token_expert_indices: torch.Tensor,
12 + topk: int,
13 + n_expert: int,
14 + n_local_expert: int,
15 + expert_map: Optional[torch.Tensor] = None,
16 + align_block_size: Optional[int] = None,
17 + fill_invalid_expert: int = -1
18 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
19 + """
20 + This function expands and permutes activation to gather uncontinuous tokens
21 + for each expert.
22 + Parameters:
23 + - hidden_states (torch.Tensor): The input tensor to the MoE layer.
24 + - topk_weights (torch.Tensor): topk expert route weight for each token.
25 + - topk_ids (torch.Tensor): topk expert route id for each token.
26 + - token_expert_indices (torch.Tensor): indice for expanded hidden.
27 + - topk (int): The number of top-k experts to select.
28 + - n_expert (int): The number of expert.
29 + - n_local_expert (int): The number of expert in current EP rank.
30 + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
31 + from the global expert space to the local expert space of the expert
32 + parallel shard.
33 + - align_block_size (Optional[int]): align group gemm block size for deepgemm
34 + - fill_invalid_expert(int): fill expert id in m_indices for invalid expert
35 + to workaround DeepGemm unsupported -1 in m_indices
36 + Returns:
37 + - permuted_hidden_states (torch.Tensor): permuted activation.
38 + - expert_first_token_offset (torch.Tensor): offset of the first token
39 + of each expert for standard grouped gemm. if enable 'align_block_size'
40 + expert_first_token_offset will align up to 'align_block_size'.
41 + - src_row_id2dst_row_id_map (torch.Tensor): idx map for moe_unpermute.
42 + - m_indices: m_indices for grouped gemm in deepgemm,`m_indices[i]` records
43 + the group which the j-th row of the LHS belong to.`
44 + """
45 + n_token, n_hidden = hidden_states.shape
46 + assert (n_hidden * hidden_states.element_size()
47 + ) % 16 == 0, "permue kernel need hidden dim align to 16B"
48 + permuted_row_size = n_token * topk
49 + if align_block_size is not None:
50 + permuted_row_size = (permuted_row_size + n_expert *
51 + (align_block_size - 1) + align_block_size -
52 + 1) // align_block_size * align_block_size
53 +
54 + permuted_hidden_states = torch.empty(
55 + (permuted_row_size, n_hidden),
56 + dtype=hidden_states.dtype,
57 + device=hidden_states.device,
58 + )
59 + m_indices = torch.full((permuted_row_size, ),
60 + fill_invalid_expert,
61 + dtype=torch.int32,
62 + device=hidden_states.device)
63 + expert_first_token_offset = torch.empty(n_local_expert + 1,
64 + dtype=torch.int64,
65 + device=hidden_states.device)
66 + src_row_id2dst_row_id_map = torch.empty((n_token, topk),
67 + dtype=torch.int32,
68 + device=hidden_states.device)
69 + torch.ops._moe_C.moe_permute(hidden_states, topk_weights, topk_ids,
70 + token_expert_indices, expert_map, n_expert,
71 + n_local_expert, topk, align_block_size,
72 + permuted_hidden_states,
73 + expert_first_token_offset,
74 + src_row_id2dst_row_id_map, m_indices)
75 + return (permuted_hidden_states, expert_first_token_offset,
76 + src_row_id2dst_row_id_map, m_indices)
77 +
78 +
79 +def moe_unpermute(
80 + permuted_hidden_states: torch.Tensor,
81 + topk_weights: torch.Tensor,
82 + topk_ids: torch.Tensor,
83 + src_row_id2dst_row_id_map: torch.Tensor,
84 + expert_first_token_offset: torch.Tensor,
85 + topk: int,
86 + n_expert: int,
87 + n_local_expert: int,
88 +) -> torch.Tensor:
89 + """
90 + This function expands and permutes activation to gathering uncontinuous
91 + tokens for each expert.
92 + Parameters:
93 + - permuted_hidden_states (torch.Tensor): permuted activation.
94 + - topk_weights (torch.Tensor): topk expert route weight for each token.
95 + - topk_ids (torch.Tensor): topk expert route id for each token.
96 + - expert_first_token_offset (torch.Tensor): offset of the first token
97 + of each expert for grouped gemm.
98 + - topk (int): The number of top-k experts to select.
99 + - n_expert (int): The number of expert.
100 + - n_local_expert (int): The number of expert in current EP rank.
101 + Returns:
102 + - hidden_states (torch.Tensor): The reduced and unpermuted activation
103 + tensor.
104 + """
105 + n_token, n_hidden = topk_weights.shape[0], permuted_hidden_states.shape[-1]
106 + assert (n_hidden * permuted_hidden_states.element_size()
107 + ) % 16 == 0, "unpermue kernel need hidden dim align to 16B"
108 + hidden_states = torch.empty((n_token, n_hidden),
109 + dtype=permuted_hidden_states.dtype,
110 + device=permuted_hidden_states.device)
111 +
112 + torch.ops._moe_C.moe_unpermute(permuted_hidden_states, topk_weights,
113 + topk_ids, src_row_id2dst_row_id_map,
114 + expert_first_token_offset, n_expert,
115 + n_local_expert, topk, hidden_states)
116 + return hidden_states
vllm/model_executor/models/arctic.py
+2 -4
@@ -175,10 +175,8 @@ class ArcticMoE(nn.Module):
175 # router_logits: (num_tokens, n_experts)
176 router_logits, _ = self.gate(hidden_states)
177 do_normalize = self.top_k > 1
178 - topk_weights, topk_ids = fused_topk(hidden_states,
179 - router_logits,
180 - self.top_k,
181 - renormalize=do_normalize)
178 + topk_weights, topk_ids, token_expert_indices = fused_topk(
179 + hidden_states, router_logits, self.top_k, renormalize=do_normalize)
180 # topk_ids: (num_tokens, k)
181 if self.is_quant:
182 if 2 * num_tokens <= self.num_experts: