feat: add optimized Kunlun rotary embedding backend

CMakeLists.txt

@@ -22,6 +22,7 @@ option(WITH_MARS "Enable Mars backend" OFF)
 option(WITH_CAMBRICON "Enable Cambricon backend" OFF)
 option(WITH_MOORE "Enable Moore backend" OFF)
 option(WITH_ASCEND "Enable Ascend backend" OFF)
+option(WITH_KUNLUN "Enable Kunlun backend" OFF)
 
 option(WITH_TORCH "Enable PyTorch C++ backend" OFF)
 option(WITH_LINKED "Enable linked third-party operator backends" OFF)
@@ -175,7 +176,7 @@ if(AUTO_DETECT_DEVICES)
 
     if(WITH_NVIDIA)
         set(_non_nvidia_gpu_detected FALSE)
-        foreach(_gpu_backend WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_CAMBRICON WITH_MOORE WITH_ASCEND)
+        foreach(_gpu_backend WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_CAMBRICON WITH_MOORE WITH_ASCEND WITH_KUNLUN)
             if(${_gpu_backend})
                 set(_non_nvidia_gpu_detected TRUE)
             endif()
@@ -364,7 +365,7 @@ endif()
 
 # Only one CUDA-like GPU backend can be enabled at a time.
 set(_gpu_backend_count 0)
-foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_MOORE WITH_ASCEND)
+foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_MOORE WITH_ASCEND WITH_KUNLUN)
     if(${_gpu_backend})
         math(EXPR _gpu_backend_count "${_gpu_backend_count} + 1")
     endif()
@@ -633,8 +634,21 @@ if(WITH_ASCEND)
     message(STATUS "Using Ascend from `${ASCEND_HOME}`.")
 endif()
 
+if(WITH_KUNLUN)
+    add_compile_definitions(WITH_KUNLUN=1)
+    set(XRE_ROOT "$ENV{XRE_ROOT}" CACHE PATH "Kunlun XRE SDK root")
+    set(XTDK_ROOT "$ENV{XTDK_ROOT}" CACHE PATH "Kunlun XTDK compiler root")
+    set(KUNLUN_ARCH "xpu3" CACHE STRING "Kunlun device architecture")
+    find_path(KUNLUN_INCLUDE_DIR xpu/runtime.h
+        HINTS "${XRE_ROOT}" PATH_SUFFIXES include REQUIRED)
+    find_library(KUNLUN_RUNTIME_LIB xpurt
+        HINTS "${XRE_ROOT}" PATH_SUFFIXES lib lib64 so REQUIRED)
+    find_program(XTDK_COMPILER clang++ HINTS "${XTDK_ROOT}/bin"
+        NO_DEFAULT_PATH REQUIRED)
+endif()
+
 # If all other platforms are not enabled, CPU is enabled by default.
-if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_THEAD AND NOT WITH_METAX AND NOT WITH_MARS AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND)
+if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_THEAD AND NOT WITH_METAX AND NOT WITH_MARS AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND AND NOT WITH_KUNLUN)
     set(WITH_CPU ON)
 endif()
 

scripts/generate_wrappers.py

@@ -1724,6 +1724,7 @@ def _device_marker_headers(devices):
         "nvidia": "infini/rt/nvidia/device_.h",
         "cambricon": "infini/rt/cambricon/device_.h",
         "ascend": "infini/rt/ascend/device_.h",
+        "kunlun": "infini/rt/kunlun/device_.h",
         "metax": "infini/rt/metax/device_.h",
         "mars": "infini/rt/mars/device_.h",
         "moore": "infini/rt/moore/device_.h",
@@ -2116,7 +2117,7 @@ def _dispatch_gen_batch_size():
         nargs="+",
         default="cpu",
         type=str,
-        help="Devices to use. Please pick from `cpu`, `nvidia`, `cambricon`, `ascend`, `metax`, `mars`, `moore`, `iluvatar`, `hygon`, and `thead`. (default: `cpu`)",
+        help="Devices to use. Please pick from `cpu`, `nvidia`, `cambricon`, `ascend`, `metax`, `mars`, `moore`, `iluvatar`, `hygon`, `thead`, and `kunlun`. (default: `cpu`)",
     )
 
     parser.add_argument(

src/CMakeLists.txt

@@ -396,6 +396,23 @@ if(WITH_MOORE)
     list(APPEND DEVICE_LIST "moore")
 endif()
 
+if(WITH_KUNLUN)
+    set(_kunlun_kernel "${CMAKE_CURRENT_SOURCE_DIR}/native/kunlun/ops/rotary_embedding/kernel.xpu")
+    set(KUNLUN_OBJECT "${CMAKE_CURRENT_BINARY_DIR}/rotary_embedding_kunlun.o")
+    add_custom_command(OUTPUT "${KUNLUN_OBJECT}"
+        COMMAND "${XTDK_COMPILER}" --xpu-arch=${KUNLUN_ARCH}
+            -std=c++17 -O2 -fPIC -fno-builtin
+            -I${KUNLUN_INCLUDE_DIR} ${INFINI_RT_INCLUDE_FLAGS}
+            -c "${_kunlun_kernel}" -o "${KUNLUN_OBJECT}"
+        DEPENDS "${_kunlun_kernel}"
+            "${CMAKE_CURRENT_SOURCE_DIR}/native/kunlun/ops/rotary_embedding/launch.h"
+        VERBATIM)
+    target_sources(infiniops PRIVATE "${KUNLUN_OBJECT}")
+    target_include_directories(infiniops PUBLIC "${KUNLUN_INCLUDE_DIR}")
+    target_link_libraries(infiniops PUBLIC "${KUNLUN_RUNTIME_LIB}")
+    list(APPEND DEVICE_LIST "kunlun")
+endif()
+
 if(WITH_CAMBRICON)
     file(GLOB_RECURSE CAMBRICON_MLU_SOURCES CONFIGURE_DEPENDS "native/cambricon/ops/*/*.mlu")
     find_program(CNCC_COMPILER cncc HINTS "${NEUWARE_HOME}/bin" "$ENV{NEUWARE_HOME}/bin" /usr/local/neuware/bin)
@@ -1569,6 +1586,11 @@ if(GENERATE_PYTHON_BINDINGS)
 
     target_link_libraries(ops PRIVATE infiniops)
 
+    # Kunlun dispatch references the host launch stubs in this object.
+    if(WITH_KUNLUN)
+        target_sources(ops PRIVATE "${KUNLUN_OBJECT}")
+    endif()
+
     # Cambricon generated dispatch is compiled into the Python extension and
     # directly references host launch stubs emitted from `.mlu` sources. Link
     # those objects into `ops` too so import does not depend on re-exporting

src/device.h

@@ -16,7 +16,7 @@ using AllDeviceTypes =
     List<Device::Type::kCpu, Device::Type::kNvidia, Device::Type::kCambricon,
          Device::Type::kAscend, Device::Type::kMetax, Device::Type::kMoore,
          Device::Type::kIluvatar, Device::Type::kHygon, Device::Type::kMars,
-         Device::Type::kThead>;
+         Device::Type::kThead, Device::Type::kKunlun>;
 
 template <typename>
 struct ActiveDevicesImpl {

src/native/kunlun/data_type_.h

@@ -0,0 +1,6 @@
+#ifndef INFINI_OPS_KUNLUN_DATA_TYPE__H_
+#define INFINI_OPS_KUNLUN_DATA_TYPE__H_
+
+#include <infini/rt/kunlun/data_type_.h>
+
+#endif

src/native/kunlun/ops/rotary_embedding/kernel.h

@@ -0,0 +1,93 @@
+#ifndef INFINI_OPS_KUNLUN_ROTARY_EMBEDDING_KERNEL_H_
+#define INFINI_OPS_KUNLUN_ROTARY_EMBEDDING_KERNEL_H_
+
+#include <xpu/runtime.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+
+#include "base/rotary_embedding.h"
+#include "dispatcher.h"
+#include "native/kunlun/data_type_.h"
+#include "native/kunlun/ops/rotary_embedding/launch.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<RotaryEmbedding, Device::Type::kKunlun>
+    : public RotaryEmbedding {
+ public:
+  Operator(const Tensor positions, Tensor query, std::optional<Tensor> key,
+           const Tensor cos_sin_cache, int64_t head_size, bool is_neox,
+           int64_t rope_dim_offset = 0, bool inverse = false)
+      : RotaryEmbedding{positions,       query,     key,
+                        cos_sin_cache,   head_size, is_neox,
+                        rope_dim_offset, inverse} {
+    assert(cos_sin_cache.size(0) > 0 && "RoPE requires a nonempty cache");
+    uint64_t available = 0;
+    if (xpu_device_get_attr(&available, XPUATTR_NUM_CLUSTER, device_index_) ==
+            XPU_SUCCESS &&
+        available > 0) {
+      cluster_count_ = static_cast<int>(std::min<uint64_t>(available, 12));
+    }
+  }
+
+  void operator()(const Tensor positions, Tensor query,
+                  std::optional<Tensor> key, const Tensor cos_sin_cache,
+                  int64_t, bool, int64_t, bool) const override {
+    if (num_tokens_ == 0) {
+      return;
+    }
+
+    const auto queue = static_cast<XPUStream>(stream_ ? stream_ : 0);
+    const size_t sequence_length =
+        positions_ndim_ == 2 ? positions_shape_[1] : num_tokens_;
+    const ptrdiff_t query_batch_stride =
+        positions_ndim_ == 2 ? query_strides_[0] : 0;
+    const ptrdiff_t key_batch_stride =
+        positions_ndim_ == 2 && key.has_value() ? key_strides_[0] : 0;
+
+    DispatchFunc<
+        Device::Type::kKunlun,
+        List<DataType::kFloat16, DataType::kBFloat16, DataType::kFloat32>,
+        List<DataType::kFloat16, DataType::kBFloat16, DataType::kFloat32>>(
+        {query_type_, cos_sin_cache_type_},
+        [&](auto query_tag, auto cache_tag) {
+          using T = typename decltype(query_tag)::type;
+          using TCache = typename decltype(cache_tag)::type;
+
+          LaunchRotaryEmbeddingKunlun<T, TCache>(
+              queue, cluster_count_,
+              reinterpret_cast<const int64_t*>(positions.data()),
+              reinterpret_cast<T*>(query.data()),
+              reinterpret_cast<const TCache*>(cos_sin_cache.data()),
+              cos_sin_cache_shape_[0], num_tokens_, sequence_length, num_heads_,
+              rot_dim_, rope_dim_offset_, query_batch_stride,
+              query_token_stride_, query_head_stride_,
+              cos_sin_cache_strides_[0], is_neox_, inverse_);
+
+          if (key.has_value()) {
+            LaunchRotaryEmbeddingKunlun<T, TCache>(
+                queue, cluster_count_,
+                reinterpret_cast<const int64_t*>(positions.data()),
+                reinterpret_cast<T*>(key->data()),
+                reinterpret_cast<const TCache*>(cos_sin_cache.data()),
+                cos_sin_cache_shape_[0], num_tokens_, sequence_length,
+                num_kv_heads_, rot_dim_, rope_dim_offset_, key_batch_stride,
+                key_token_stride_, key_head_stride_, cos_sin_cache_strides_[0],
+                is_neox_, inverse_);
+          }
+        },
+        "KunlunRotaryEmbedding::operator()");
+  }
+
+ private:
+  int cluster_count_{8};
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/kunlun/ops/rotary_embedding/kernel.xpu

@@ -0,0 +1,146 @@
+#include <infini/rt/kunlun/data_type_.h>
+#include <xpu/kernel/xtdk.h>
+#include <xpu/kernel/xtdk_bf16.h>
+#include <xpu/kernel/xtdk_fp16.h>
+
+#include "launch.h"
+
+namespace infini::ops {
+namespace kunlun {
+
+template <typename T>
+struct KernelType {
+  using type = T;
+};
+template <>
+struct KernelType<infini::rt::Float16> {
+  using type = half;
+};
+template <>
+struct KernelType<infini::rt::BFloat16> {
+  using type = bfloat16_t;
+};
+
+template <typename T>
+__device__ float ToFloat(T value) {
+  if constexpr (xpu_std::is_same<T, bfloat16_t>::value) {
+    return __bfloat162float(value);
+  } else {
+    return static_cast<float>(value);
+  }
+}
+
+template <typename T>
+__device__ T FromFloat(float value) {
+  if constexpr (xpu_std::is_same<T, bfloat16_t>::value) {
+    return __float2bfloat16(value);
+  } else {
+    return static_cast<T>(value);
+  }
+}
+
+// XTDK uses 32-bit size_t on device; keep the launch ABI fixed-width.
+template <typename T, typename TCache>
+__global__ void RotaryEmbeddingKunlunKernel(
+    const int64_t* positions, T* data, const TCache* cos_sin_cache,
+    uint64_t cache_rows, uint64_t num_tokens, uint64_t sequence_length,
+    uint64_t num_heads, uint64_t rot_dim, uint64_t rope_dim_offset,
+    int64_t batch_stride, int64_t token_stride, int64_t head_stride,
+    int64_t cache_token_stride, bool is_neox, bool inverse) {
+  // Stripe adjacent workers across clusters to spread small workloads.
+  const uint64_t worker = core_id() * cluster_num() + cluster_id();
+  const uint64_t workers = core_num() * cluster_num();
+  const uint64_t tasks = num_tokens * num_heads;
+  const uint64_t count = tasks / workers + (worker < tasks % workers);
+  const uint64_t first = worker * (tasks / workers) +
+                         (worker < tasks % workers ? worker : tasks % workers);
+  constexpr uint64_t kPairs = 128;
+  __local__ T x[2 * kPairs];
+  __local__ T y[2 * kPairs];
+  __local__ TCache cos[kPairs];
+  __local__ TCache sin[kPairs];
+  __local__ int64_t position[1];
+  const uint64_t pairs = rot_dim / 2;
+
+  for (uint64_t task = first; task < first + count; ++task) {
+    const uint64_t token = task / num_heads;
+    const uint64_t head = task % num_heads;
+    const int64_t offset = (token / sequence_length) * batch_stride +
+                           (token % sequence_length) * token_stride +
+                           head * head_stride + rope_dim_offset;
+    GM2LM(positions + token, position, sizeof(int64_t));
+    // Clamp before narrowing so native int64 positions work during replay.
+    uint64_t row = position[0] < 0 ? 0 : static_cast<uint64_t>(position[0]);
+    if (row >= cache_rows) row = cache_rows - 1;
+    const auto* cache = cos_sin_cache + row * cache_token_stride;
+    for (uint64_t processed = 0; processed < pairs; processed += kPairs) {
+      const uint64_t n =
+          pairs - processed < kPairs ? pairs - processed : kPairs;
+      if (is_neox) {
+        GM2LM(data + offset + processed, x, n * sizeof(T));
+        GM2LM(data + offset + pairs + processed, x + kPairs, n * sizeof(T));
+      } else {
+        GM2LM(data + offset + 2 * processed, x, 2 * n * sizeof(T));
+      }
+      GM2LM(cache + processed, cos, n * sizeof(TCache));
+      GM2LM(cache + pairs + processed, sin, n * sizeof(TCache));
+      for (uint64_t k = 0; k < n; ++k) {
+        const uint64_t i = is_neox ? k : 2 * k;
+        const uint64_t j = is_neox ? k + kPairs : 2 * k + 1;
+        const float a = ToFloat(x[i]);
+        const float b = ToFloat(x[j]);
+        const float c = ToFloat(cos[k]);
+        const float s = inverse ? -ToFloat(sin[k]) : ToFloat(sin[k]);
+        y[i] = FromFloat<T>(a * c - b * s);
+        y[j] = FromFloat<T>(a * s + b * c);
+      }
+      mfence();
+      if (is_neox) {
+        LM2GM(y, data + offset + processed, n * sizeof(T));
+        LM2GM(y + kPairs, data + offset + pairs + processed, n * sizeof(T));
+      } else {
+        LM2GM(y, data + offset + 2 * processed, 2 * n * sizeof(T));
+      }
+    }
+  }
+}
+
+}  // namespace kunlun
+
+template <typename T, typename TCache>
+void LaunchRotaryEmbeddingKunlun(XPUStream stream, int cluster_count,
+                                 const int64_t* positions, T* data,
+                                 const TCache* cos_sin_cache, size_t cache_rows,
+                                 size_t num_tokens, size_t sequence_length,
+                                 size_t num_heads, size_t rot_dim,
+                                 size_t rope_dim_offset, ptrdiff_t batch_stride,
+                                 ptrdiff_t token_stride, ptrdiff_t head_stride,
+                                 ptrdiff_t cache_token_stride, bool is_neox,
+                                 bool inverse) {
+  using K = typename kunlun::KernelType<T>::type;
+  using KCache = typename kunlun::KernelType<TCache>::type;
+  kunlun::RotaryEmbeddingKunlunKernel<K, KCache><<<cluster_count, 64, stream>>>(
+      positions, reinterpret_cast<K*>(data),
+      reinterpret_cast<const KCache*>(cos_sin_cache), cache_rows, num_tokens,
+      sequence_length, num_heads, rot_dim, rope_dim_offset, batch_stride,
+      token_stride, head_stride, cache_token_stride, is_neox, inverse);
+}
+
+#define INSTANTIATE(T, TCACHE)                                           \
+  template void LaunchRotaryEmbeddingKunlun<T, TCACHE>(                  \
+      XPUStream, int, const int64_t*, T*, const TCACHE*, size_t, size_t, \
+      size_t, size_t, size_t, size_t, ptrdiff_t, ptrdiff_t, ptrdiff_t,   \
+      ptrdiff_t, bool, bool)
+
+INSTANTIATE(infini::rt::Float16, infini::rt::Float16);
+INSTANTIATE(infini::rt::Float16, infini::rt::BFloat16);
+INSTANTIATE(infini::rt::Float16, float);
+INSTANTIATE(infini::rt::BFloat16, infini::rt::Float16);
+INSTANTIATE(infini::rt::BFloat16, infini::rt::BFloat16);
+INSTANTIATE(infini::rt::BFloat16, float);
+INSTANTIATE(float, infini::rt::Float16);
+INSTANTIATE(float, infini::rt::BFloat16);
+INSTANTIATE(float, float);
+#undef INSTANTIATE
+
+}  // namespace infini::ops

src/native/kunlun/ops/rotary_embedding/launch.h

@@ -0,0 +1,24 @@
+#ifndef INFINI_OPS_KUNLUN_ROTARY_EMBEDDING_LAUNCH_H_
+#define INFINI_OPS_KUNLUN_ROTARY_EMBEDDING_LAUNCH_H_
+
+#include <xpu/runtime.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace infini::ops {
+
+template <typename T, typename TCache>
+void LaunchRotaryEmbeddingKunlun(XPUStream queue, int cluster_count,
+                                 const int64_t* positions, T* data,
+                                 const TCache* cos_sin_cache, size_t cache_rows,
+                                 size_t num_tokens, size_t sequence_length,
+                                 size_t num_heads, size_t rot_dim,
+                                 size_t rope_dim_offset, ptrdiff_t batch_stride,
+                                 ptrdiff_t token_stride, ptrdiff_t head_stride,
+                                 ptrdiff_t cache_token_stride, bool is_neox,
+                                 bool inverse);
+
+}  // namespace infini::ops
+
+#endif

src/pybind11_utils.h

@@ -204,6 +204,7 @@ inline std::optional<Device::Type> TryDeviceTypeFromString(
       {"iluvatar", Device::Type::kIluvatar},
       {"hygon", Device::Type::kHygon},
       {"thead", Device::Type::kThead},
+      {"kunlun", Device::Type::kKunlun},
   };
 
   auto platform_it{