[AdvancedCompiler]Topk(cpp wrapper)

ctests/CMakeLists.txt

@@ -21,3 +21,9 @@ add_executable(test_triton_rope test_triton_rope.cpp)
 target_link_libraries(test_triton_rope
   PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
 add_test(NAME test_triton_rope COMMAND test_triton_rope)
+
+
+add_executable(test_triton_special test_triton_special.cpp)
+target_link_libraries(test_triton_special
+  PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
+add_test(NAME test_triton_special COMMAND test_triton_special)

ctests/test_triton_special.cpp

@@ -0,0 +1,40 @@
+#include <gtest/gtest.h>
+#include "c10/util/Logging.h"
+#include "flag_gems/operators.h"
+#include "torch/torch.h"
+
+class topktest
+    : public ::testing::TestWithParam<std::tuple<int64_t, int64_t, int64_t, bool, torch::ScalarType>> {};
+
+TEST_P(topktest, CompareWithPyTorch) {
+  const torch::Device device(torch::kCUDA, 0);
+  auto [batch_size, hiddensize, topk, largest, dtype] = GetParam();
+  auto options = torch::TensorOptions().dtype(dtype).device(device).requires_grad(false);
+  torch::Tensor x = torch::arange(0, hiddensize, options).repeat({batch_size, 1});
+  for (int64_t i = 0; i < batch_size; ++i) {
+    torch::Tensor perm =
+        torch::randperm(hiddensize, torch::TensorOptions().dtype(torch::kLong).device(device));
+    x[i] = x[i].index({perm});
+  }
+  auto x_orig = x.clone();
+
+  auto [out_weight_torch, out_index_torch] = at::topk(x, topk, -1, largest, true);
+  auto [out_weight_triton, out_index_triton] = flag_gems::topk(x, topk, -1, largest, true);
+
+  EXPECT_TRUE(torch::allclose(out_weight_torch, out_weight_triton));
+  EXPECT_TRUE(torch::equal(out_index_torch, out_index_triton));
+}
+
+INSTANTIATE_TEST_SUITE_P(special_op_test,
+                         topktest,
+                         ::testing::Combine(
+                             // batch_size: [4, 8]
+                             ::testing::Values(4, 8),
+                             // hiddensize: [128, 256]
+                             ::testing::Values(128, 256),
+                             // topk: [5]
+                             ::testing::Values(5),
+                             // largest: [true, false]
+                             ::testing::Values(true, false),
+                             // dtype:
+                             ::testing::Values(torch::kFloat32, torch::kFloat16, torch::kBFloat16)));

include/flag_gems/operators.h

@@ -29,4 +29,6 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding(
     const at::Tensor &sin,
     const std::optional<at::Tensor> &position_ids = std::nullopt,
     bool rotary_interleaved = false);
+std::tuple<at::Tensor, at::Tensor> topk(
+    const at::Tensor &x, int64_t k, int64_t dim = -1, bool largest = true, bool sorted = true);
 }  // namespace flag_gems

lib/CMakeLists.txt

@@ -6,7 +6,8 @@ add_library(operators
             mm.cpp
             rms_norm.cpp
             fused_add_rms_norm.cpp
-            rotary_embedding.cpp)
+            rotary_embedding.cpp
+            topk.cpp)
 target_include_directories(operators
   PUBLIC
     $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/include>

lib/topk.cpp

@@ -0,0 +1,99 @@
+#include "flag_gems/operators.h"
+#include "flag_gems/utils.h"
+
+#include <iostream>
+#include <tuple>
+#include "c10/cuda/CUDAStream.h"
+#include "triton_jit/triton_jit_function.h"
+
+namespace flag_gems {
+using namespace triton_jit;
+
+std::tuple<at::Tensor, at::Tensor> topk(
+    const at::Tensor &x, int64_t k, int64_t dim, bool largest, bool sorted) {
+  TORCH_CHECK(x.dim() >= 1, "input tensor must have at least one dimension");
+  TORCH_CHECK(sorted, "currently only support sorted == true");
+  dim = dim < 0 ? x.dim() + dim : dim;
+  TORCH_CHECK(dim == x.dim() - 1, "currently only support topk in last dimension");
+  auto topk_elem_cnt = x.size(dim);
+  int64_t batch_size = 1;
+  for (int i = 0; i < x.dim() - 1; i++) {
+    batch_size *= x.size(i);
+  }
+  bool descending = largest;
+  int64_t chunk_size = (topk_elem_cnt >= 1024) ? 1024 : 256;
+  if (chunk_size < k) {
+    chunk_size = utils::next_power_of_2(k);
+  }
+  int64_t chunk_num = (topk_elem_cnt + chunk_size - 1) / chunk_size;
+  at::Tensor stage1_out = at::empty(batch_size * chunk_num * k, x.options());
+  at::Tensor stage1_out_idx =
+      at::empty(batch_size * chunk_num * k, at::TensorOptions().dtype(torch::kLong).device(x.device()));
+  int64_t stage2_elem_cnt = chunk_num * k;
+  int64_t BLOCK_SIZE = utils::next_power_of_2(stage2_elem_cnt);
+  auto out_shape = x.sizes().vec();
+  out_shape[out_shape.size() - 1] = k;
+  at::Tensor stage2_out = at::empty(out_shape, x.options());
+  at::Tensor stage2_out_idx =
+      at::empty(out_shape, at::TensorOptions().dtype(torch::kLong).device(x.device()));
+  const TritonJITFunction &f1 =
+      TritonJITFunction::getInstance(std::string(utils::get_triton_src_path() / "topk.py"),
+                                     "topk_stage1_kernel");
+  const TritonJITFunction &f2 =
+      TritonJITFunction::getInstance(std::string(utils::get_triton_src_path() / "topk.py"),
+                                     "topk_stage2_kernel");
+  c10::DeviceGuard guard(stage1_out.device());
+  c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();
+  CUstream raw_stream = static_cast<CUstream>(stream.stream());
+  /*
+  def topk_stage1_kernel(y_ptr,
+                         index_ptr,
+                         x_ptr,
+                         k,
+                         N: tl.constexpr,
+                         CHUNK_SIZE: tl.constexpr,
+                         DESCENDING: tl.constexpr):
+*/
+  f1(raw_stream,
+     batch_size,
+     chunk_num,
+     1,
+     /* num_warps */ 8,
+     /* num_stages */ 1,
+     stage1_out,
+     stage1_out_idx,
+     x,
+     k,
+     topk_elem_cnt,
+     chunk_size,
+     descending);
+  /*
+  def topk_stage2_kernel(y_ptr,
+                           index_ptr,
+                           chunk_x,
+                           chunk_index,
+                           sort_dim: tl.constexpr,
+                           k: tl.constexpr,
+                           N: tl.constexpr,
+                           BLOCK_SIZE: tl.constexpr,
+                           DESCENDING: tl.constexpr,):
+*/
+  f2(raw_stream,
+     batch_size,
+     1,
+     1,
+     8,
+     1,
+     stage2_out,
+     stage2_out_idx,
+     stage1_out,
+     stage1_out_idx,
+     dim,
+     k,
+     stage2_elem_cnt,
+     BLOCK_SIZE,
+     descending);
+
+  return std::make_tuple(stage2_out, stage2_out_idx);
+}
+}  // namespace flag_gems

src/flag_gems/csrc/cstub.cpp

@@ -20,6 +20,7 @@ TORCH_LIBRARY(flag_gems, m) {
   m.def(
       "rotary_embedding(Tensor q, Tensor k, Tensor cos, Tensor sin, Tensor? position_ids=None, "
       "bool rotary_interleaved=False) -> (Tensor, Tensor)");  // q and k may be view to other size
+  m.def("topk(Tensor x, SymInt k, int dim, bool largest, bool sorted) -> (Tensor, Tensor)");
 }
 
 TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
@@ -31,5 +32,6 @@ TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
   // Rotary embedding
   m.impl("rotary_embedding", TORCH_FN(rotary_embedding));
   m.impl("rotary_embedding_inplace", TORCH_FN(rotary_embedding_inplace));
+  m.impl("topk", TORCH_FN(topk));
 }
 }  // namespace flag_gems