[Perf] Optimize moe permute by pre-allocate buffer, 9~14% kernel performance improvement

benchmarks/kernels/benchmark_moe_permute_unpermute.py

@@ -10,6 +10,7 @@
 
 from vllm.model_executor.layers.fused_moe import fused_topk
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
+    MoEPermuteScratch,
     moe_permute,
     moe_unpermute,
 )
@@ -54,6 +55,15 @@ def benchmark_permute(
     topk_weights, topk_ids, token_expert_indices = fused_topk(
         qhidden_states, input_gating, topk, False
     )
+    scratch = MoEPermuteScratch(
+        max_num_tokens=num_tokens,
+        topk=topk,
+        num_experts=num_experts,
+        num_local_experts=num_experts,
+        device=qhidden_states.device,
+        hidden_size=hidden_size,
+        hidden_dtype=qhidden_states.dtype,
+    )
 
     def prepare(i: int):
         input_gating.copy_(gating_output[i])
@@ -65,6 +75,7 @@ def run():
             topk_ids=topk_ids,
             n_expert=num_experts,
             expert_map=None,
+            scratch=scratch,
         )
 
     # JIT compilation & warmup
@@ -123,6 +134,15 @@ def benchmark_unpermute(
     topk_weights, topk_ids, token_expert_indices = fused_topk(
         qhidden_states, input_gating, topk, False
     )
+    scratch = MoEPermuteScratch(
+        max_num_tokens=num_tokens,
+        topk=topk,
+        num_experts=num_experts,
+        num_local_experts=num_experts,
+        device=qhidden_states.device,
+        hidden_size=hidden_size,
+        hidden_dtype=qhidden_states.dtype,
+    )
 
     def prepare():
         (
@@ -137,6 +157,7 @@ def prepare():
             topk_ids=topk_ids,
             n_expert=num_experts,
             expert_map=None,
+            scratch=scratch,
         )
         # convert to fp16/bf16 as gemm output
         return (

csrc/moe/moe_ops.h

@@ -62,6 +62,9 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
 
 bool moe_permute_unpermute_supported();
 
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t num_experts);
+
 void shuffle_rows(const torch::Tensor& input_tensor,
                   const torch::Tensor& dst2src_map,
                   torch::Tensor& output_tensor);

csrc/moe/moe_permute_unpermute_op.cu

@@ -8,7 +8,34 @@
 // moe_permute kernels require at least CUDA 12.0
 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
 
-void moe_permute(
+namespace {
+
+torch::Tensor maybe_allocate_tensor(
+    const std::optional<torch::Tensor>& maybe_tensor,
+    at::IntArrayRef expected_sizes, torch::ScalarType dtype, c10::Device device,
+    char const* name) {
+  auto expected_numel = c10::multiply_integers(expected_sizes);
+  if (maybe_tensor.has_value()) {
+    auto tensor = maybe_tensor.value();
+    TORCH_CHECK(tensor.device() == device, name, " must be on the same device");
+    TORCH_CHECK(tensor.scalar_type() == dtype, name, " has incorrect dtype");
+    TORCH_CHECK(tensor.numel() >= expected_numel, name,
+                " is too small for the requested shape");
+    auto flat_tensor = tensor.reshape({tensor.numel()});
+    return flat_tensor.narrow(0, 0, expected_numel).view(expected_sizes);
+  }
+  return torch::empty(expected_sizes, torch::dtype(dtype).device(device));
+}
+
+}  // namespace
+
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t n_expert) {
+  return static_cast<int64_t>(
+      CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
+}
+
+void moe_permute_impl(
     const torch::Tensor& input,                      // [n_token, hidden]
     const torch::Tensor& topk_ids,                   // [n_token, topk]
     const torch::Tensor& token_expert_indices,       // [n_token, topk]
@@ -17,7 +44,11 @@ void moe_permute(
     torch::Tensor& permuted_input,             // [permuted_size, hidden]
     torch::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
     torch::Tensor& inv_permuted_idx,           // [n_token, topk]
-    torch::Tensor& permuted_idx) {             // [permute_size]
+    torch::Tensor& permuted_idx,               // [permute_size]
+    const std::optional<torch::Tensor>& maybe_sort_workspace,
+    const std::optional<torch::Tensor>& maybe_permuted_experts_id,
+    const std::optional<torch::Tensor>& maybe_sorted_row_idx,
+    const std::optional<torch::Tensor>& maybe_topk_ids_for_sort) {
   TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
               "expert_first_token_offset must be int64");
   TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
@@ -27,49 +58,49 @@ void moe_permute(
   TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
               "inv_permuted_idx must be int32");
   TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
-              "expert_first_token_offset shape != n_local_expert+1")
+              "expert_first_token_offset shape != n_local_expert+1");
   TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
               "token_expert_indices shape must be same as inv_permuted_idx");
+  auto device = input.device();
   auto n_token = input.sizes()[0];
   auto n_hidden = input.sizes()[1];
+  auto expanded_rows = n_token * topk;
   auto stream = at::cuda::getCurrentCUDAStream().stream();
-  const long sorter_size =
-      CubKeyValueSorter::getWorkspaceSize(n_token * topk, n_expert);
-  auto sort_workspace = torch::empty(
-      {sorter_size},
-      torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
-  torch::Tensor topk_ids_for_sort = topk_ids;
-  auto permuted_experts_id = torch::empty_like(topk_ids);
-  auto sorted_row_idx = torch::empty_like(inv_permuted_idx);
+
+  size_t sorter_size =
+      CubKeyValueSorter::getWorkspaceSize(expanded_rows, n_expert);
+  auto sort_workspace = maybe_allocate_tensor(
+      maybe_sort_workspace, {static_cast<int64_t>(sorter_size)}, torch::kInt8,
+      device, "sort_workspace");
+  auto permuted_experts_id =
+      maybe_allocate_tensor(maybe_permuted_experts_id, topk_ids.sizes(),
+                            at::ScalarType::Int, device, "permuted_experts_id");
+  auto sorted_row_idx =
+      maybe_allocate_tensor(maybe_sorted_row_idx, inv_permuted_idx.sizes(),
+                            at::ScalarType::Int, device, "sorted_row_idx");
 
   CubKeyValueSorter sorter{};
   int64_t* valid_num_ptr = nullptr;
-  // pre-process kernel for expert-parallelism:
-  // no local expert id plus "n_expert" offset for priority to local expert
-  // map local expert id [n, .., n+n_local_expert-1] to [0, n_local_expert -1]
-  // For example, 4 expert with ep_size=2. ep_rank=1 owns global expert id
-  // [2,3] with expert_map[-1, -1, 0, 1], preprocess_topk_id  process topk_ids
-  // and map global expert id [2, 3] to local_expert id [0, 1] and map global
-  // expert id [0, 1] ( not in ep rank=1)  to [4, 5] by plus n_expert. This map
-  // operation is to make local expert high priority in following sort topk_ids
-  // and scan local expert_first_token_offset for each ep rank for next group
-  // gemm.
+  torch::Tensor topk_ids_for_sort = topk_ids;
+
   if (expert_map.has_value()) {
     const int* expert_map_ptr = get_ptr<int>(expert_map.value());
     valid_num_ptr =
         get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
-    topk_ids_for_sort = topk_ids.clone();
+    topk_ids_for_sort =
+        maybe_allocate_tensor(maybe_topk_ids_for_sort, topk_ids.sizes(),
+                              at::ScalarType::Int, device, "topk_ids_for_sort");
+    topk_ids_for_sort.copy_(topk_ids);
     preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
                              expert_map_ptr, n_expert, stream);
   }
-  // expert sort topk expert id and scan expert id get expert_first_token_offset
+
   sortAndScanExpert(
       get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
       get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
       get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
       n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
 
-  // dispatch expandInputRowsKernelLauncher
   MOE_DISPATCH(input.scalar_type(), [&] {
     expandInputRowsKernelLauncher<scalar_t>(
         get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
@@ -79,6 +110,38 @@ void moe_permute(
   });
 }
 
+void moe_permute(
+    const torch::Tensor& input,                      // [n_token, hidden]
+    const torch::Tensor& topk_ids,                   // [n_token, topk]
+    const torch::Tensor& token_expert_indices,       // [n_token, topk]
+    const std::optional<torch::Tensor>& expert_map,  // [n_expert]
+    int64_t n_expert, int64_t n_local_expert, int64_t topk,
+    torch::Tensor& permuted_input,             // [permuted_size, hidden]
+    torch::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
+    torch::Tensor& inv_permuted_idx,           // [n_token, topk]
+    torch::Tensor& permuted_idx) {             // [permute_size]
+  moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk, permuted_input,
+                   expert_first_token_offset, inv_permuted_idx, permuted_idx,
+                   std::nullopt, std::nullopt, std::nullopt, std::nullopt);
+}
+
+void moe_permute_with_scratch(
+    const torch::Tensor& input, const torch::Tensor& topk_ids,
+    const torch::Tensor& token_expert_indices,
+    const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
+    int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
+    torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
+    torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
+    torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
+    torch::Tensor& topk_ids_for_sort) {
+  moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk, permuted_input,
+                   expert_first_token_offset, inv_permuted_idx, permuted_idx,
+                   sort_workspace, permuted_experts_id, sorted_row_idx,
+                   topk_ids_for_sort);
+}
+
 void moe_unpermute(
     const torch::Tensor& permuted_hidden_states,  // [n_token * topk, hidden]
     const torch::Tensor& topk_weights,            // [n_token, topk]
@@ -169,6 +232,12 @@ void shuffle_rows(const torch::Tensor& input_tensor,
 
 #else
 
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t n_expert) {
+  TORCH_CHECK(
+      false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0");
+}
+
 void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
                  const torch::Tensor& token_expert_indices,
                  const std::optional<torch::Tensor>& expert_map,
@@ -179,6 +248,19 @@ void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
   TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
 }
 
+void moe_permute_with_scratch(
+    const torch::Tensor& input, const torch::Tensor& topk_ids,
+    const torch::Tensor& token_expert_indices,
+    const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
+    int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
+    torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
+    torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
+    torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
+    torch::Tensor& topk_ids_for_sort) {
+  TORCH_CHECK(false,
+              "moe_permute_with_scratch is not supported on CUDA < 12.0");
+}
+
 void moe_unpermute(
     const torch::Tensor& permuted_hidden_states,
     const torch::Tensor& topk_weights, const torch::Tensor& inv_permuted_idx,
@@ -199,5 +281,6 @@ bool moe_permute_unpermute_supported() {
 
 TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
   m.impl("moe_permute", &moe_permute);
+  m.impl("moe_permute_with_scratch", &moe_permute_with_scratch);
   m.impl("moe_unpermute", &moe_unpermute);
 }

csrc/moe/torch_bindings.cpp

@@ -100,13 +100,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
       "permuted_idx)->()");
 
+  m.def(
+      "moe_permute_with_scratch(Tensor input, Tensor topk_ids,"
+      "Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
+      "int n_local_expert,"
+      "int topk, Tensor! permuted_input, Tensor! "
+      "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
+      "permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, "
+      "Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()");
+
   m.def(
       "moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
       "Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
       "int topk, Tensor! hidden_states)->()");
 
   m.def("moe_permute_unpermute_supported() -> bool");
+  m.def(
+      "moe_permute_sort_workspace_size(int num_expanded_rows, int n_expert) -> "
+      "int");
   m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
+  m.impl("moe_permute_sort_workspace_size", &moe_permute_sort_workspace_size);
 
   // Row shuffle for MoE
   m.def(

tests/kernels/moe/test_moe_permute_unpermute.py

@@ -14,6 +14,7 @@
     determine_expert_map,
 )
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
+    MoEPermuteScratch,
     moe_permute,
     moe_permute_unpermute_supported,
     moe_unpermute,
@@ -209,3 +210,79 @@ def test_moe_permute_unpermute(
     )
     # check unpermuted hidden
     torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    n_token = 64
+    n_hidden = 2048
+    n_expert = 16
+    topk = 4
+    hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype)
+    gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype)
+    _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False)
+
+    scratch = MoEPermuteScratch(
+        max_num_tokens=n_token,
+        topk=topk,
+        num_experts=n_expert,
+        num_local_experts=n_expert,
+        device=hidden_states.device,
+        hidden_size=n_hidden,
+        hidden_dtype=hidden_states.dtype,
+    )
+
+    first = moe_permute(
+        hidden_states=hidden_states,
+        a1q_scale=None,
+        topk_ids=topk_ids,
+        n_expert=n_expert,
+        scratch=scratch,
+    )
+    second = moe_permute(
+        hidden_states=hidden_states,
+        a1q_scale=None,
+        topk_ids=topk_ids,
+        n_expert=n_expert,
+        scratch=scratch,
+    )
+
+    (
+        permuted_hidden_states_1,
+        _,
+        expert_first_token_offset_1,
+        inv_permuted_idx_1,
+        permuted_idx_1,
+    ) = first
+    (
+        permuted_hidden_states_2,
+        _,
+        expert_first_token_offset_2,
+        inv_permuted_idx_2,
+        permuted_idx_2,
+    ) = second
+
+    torch.testing.assert_close(permuted_hidden_states_1, permuted_hidden_states_2)
+    torch.testing.assert_close(expert_first_token_offset_1, expert_first_token_offset_2)
+    torch.testing.assert_close(inv_permuted_idx_1, inv_permuted_idx_2)
+    torch.testing.assert_close(permuted_idx_1, permuted_idx_2)
+
+    assert (
+        permuted_hidden_states_1.untyped_storage().data_ptr()
+        == permuted_hidden_states_2.untyped_storage().data_ptr()
+    )
+    assert (
+        expert_first_token_offset_1.untyped_storage().data_ptr()
+        == expert_first_token_offset_2.untyped_storage().data_ptr()
+    )
+    assert (
+        inv_permuted_idx_1.untyped_storage().data_ptr()
+        == scratch.inv_permuted_idx.untyped_storage().data_ptr()
+    )
+    assert (
+        permuted_idx_1.untyped_storage().data_ptr()
+        == scratch.permuted_idx.untyped_storage().data_ptr()
+    )