A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 23:00 UTC
[Perf] Only clone when needed for `moe_permute` (#32273)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 20, 2026 at 11:34 UTC
6c97b9b9b6337d021615ed74f9cf836b65c446d7
3 files changed
+6
-5
csrc/moe/moe_permute_unpermute_op.cu
+4
-3
@@ -42,7 +42,7 @@ 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
45
+ torch::Tensor topk_ids_for_sort = topk_ids;
46
auto permuted_experts_id = torch::empty_like(topk_ids);
47
auto sorted_row_idx = torch::empty_like(inv_permuted_idx);
48
@@ -62,12 +62,13 @@ void moe_permute(
62
const int* expert_map_ptr = get_ptr<int>(expert_map.value());
63
valid_num_ptr =
64
get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
65
- preprocessTopkIdLauncher(get_ptr<int>(copy_topk_ids), n_token * topk,
65
+ topk_ids_for_sort = topk_ids.clone();
66
+ preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
67
expert_map_ptr, n_expert, stream);
68
}
69
// expert sort topk expert id and scan expert id get expert_first_token_offset
70
sortAndScanExpert(
70
- get_ptr<int>(copy_topk_ids), get_ptr<int>(token_expert_indices),
71
+ get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
72
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
73
get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
74
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+1
-1
@@ -109,7 +109,7 @@ void computeExpertFirstTokenOffset(int const* sorted_indices,
109
sorted_indices, total_indices, num_experts, expert_first_token_offset);
110
}
111
112
-void sortAndScanExpert(int* expert_for_source_row, const int* source_rows,
112
+void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
113
int* permuted_experts, int* permuted_rows,
114
int64_t* expert_first_token_offset, int num_rows,
115
int num_experts, int num_experts_per_node, int k,
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+1
-1
@@ -48,7 +48,7 @@ void computeExpertFirstTokenOffset(int const* sorted_indices,
48
int64_t* expert_first_token_offset,
49
cudaStream_t stream);
50
51
-void sortAndScanExpert(int* expert_for_source_row, const int* source_rows,
51
+void sortAndScanExpert(const 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,