From 5c0d399b3605793a9ff75a55f14b513590921731 Mon Sep 17 00:00:00 2001 From: zhangxin Date: Sat, 10 Oct 2026 18:58:57 +0800 Subject: [PATCH] feat: add optimized Kunlun rotary embedding backend --- CMakeLists.txt | 20 ++- scripts/generate_wrappers.py | 3 +- src/CMakeLists.txt | 22 +++ src/device.h | 2 +- src/native/kunlun/data_type_.h | 6 + .../kunlun/ops/rotary_embedding/kernel.h | 93 +++++++++++ .../kunlun/ops/rotary_embedding/kernel.xpu | 146 ++++++++++++++++++ .../kunlun/ops/rotary_embedding/launch.h | 24 +++ src/pybind11_utils.h | 1 + src/torch/device_.h | 5 + tests/conftest.py | 3 +- tests/test_rotary_embedding.py | 72 +++++++++ 12 files changed, 391 insertions(+), 6 deletions(-) create mode 100644 src/native/kunlun/data_type_.h create mode 100644 src/native/kunlun/ops/rotary_embedding/kernel.h create mode 100644 src/native/kunlun/ops/rotary_embedding/kernel.xpu create mode 100644 src/native/kunlun/ops/rotary_embedding/launch.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 0d8c3e5b3..9863ba863 100644 --- a/CMakeLists.txt +++ b/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() diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 4b98dfb9b..a839213f1 100644 --- a/scripts/generate_wrappers.py +++ b/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( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index ce9f75770..c90916afa 100644 --- a/src/CMakeLists.txt +++ b/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 diff --git a/src/device.h b/src/device.h index 6773b77ad..4e05bd207 100644 --- a/src/device.h +++ b/src/device.h @@ -16,7 +16,7 @@ using AllDeviceTypes = List; + Device::Type::kThead, Device::Type::kKunlun>; template struct ActiveDevicesImpl { diff --git a/src/native/kunlun/data_type_.h b/src/native/kunlun/data_type_.h new file mode 100644 index 000000000..739ad57a9 --- /dev/null +++ b/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 + +#endif diff --git a/src/native/kunlun/ops/rotary_embedding/kernel.h b/src/native/kunlun/ops/rotary_embedding/kernel.h new file mode 100644 index 000000000..6155a2263 --- /dev/null +++ b/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 + +#include +#include +#include +#include +#include + +#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 + : public RotaryEmbedding { + public: + Operator(const Tensor positions, Tensor query, std::optional 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(std::min(available, 12)); + } + } + + void operator()(const Tensor positions, Tensor query, + std::optional key, const Tensor cos_sin_cache, + int64_t, bool, int64_t, bool) const override { + if (num_tokens_ == 0) { + return; + } + + const auto queue = static_cast(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, + List>( + {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( + queue, cluster_count_, + reinterpret_cast(positions.data()), + reinterpret_cast(query.data()), + reinterpret_cast(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( + queue, cluster_count_, + reinterpret_cast(positions.data()), + reinterpret_cast(key->data()), + reinterpret_cast(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 diff --git a/src/native/kunlun/ops/rotary_embedding/kernel.xpu b/src/native/kunlun/ops/rotary_embedding/kernel.xpu new file mode 100644 index 000000000..25835da82 --- /dev/null +++ b/src/native/kunlun/ops/rotary_embedding/kernel.xpu @@ -0,0 +1,146 @@ +#include +#include +#include +#include + +#include "launch.h" + +namespace infini::ops { +namespace kunlun { + +template +struct KernelType { + using type = T; +}; +template <> +struct KernelType { + using type = half; +}; +template <> +struct KernelType { + using type = bfloat16_t; +}; + +template +__device__ float ToFloat(T value) { + if constexpr (xpu_std::is_same::value) { + return __bfloat162float(value); + } else { + return static_cast(value); + } +} + +template +__device__ T FromFloat(float value) { + if constexpr (xpu_std::is_same::value) { + return __float2bfloat16(value); + } else { + return static_cast(value); + } +} + +// XTDK uses 32-bit size_t on device; keep the launch ABI fixed-width. +template +__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(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(a * c - b * s); + y[j] = FromFloat(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 +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::type; + using KCache = typename kunlun::KernelType::type; + kunlun::RotaryEmbeddingKunlunKernel<<>>( + positions, reinterpret_cast(data), + reinterpret_cast(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( \ + 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 diff --git a/src/native/kunlun/ops/rotary_embedding/launch.h b/src/native/kunlun/ops/rotary_embedding/launch.h new file mode 100644 index 000000000..b4f715f7c --- /dev/null +++ b/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 + +#include +#include + +namespace infini::ops { + +template +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 diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index 4acc97d37..b30e4538a 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -204,6 +204,7 @@ inline std::optional TryDeviceTypeFromString( {"iluvatar", Device::Type::kIluvatar}, {"hygon", Device::Type::kHygon}, {"thead", Device::Type::kThead}, + {"kunlun", Device::Type::kKunlun}, }; auto platform_it{ diff --git a/src/torch/device_.h b/src/torch/device_.h index 3d2090e48..34b151261 100644 --- a/src/torch/device_.h +++ b/src/torch/device_.h @@ -60,6 +60,11 @@ struct TorchDeviceName { static constexpr std::string_view kValue{"cuda"}; }; +template <> +struct TorchDeviceName { + static constexpr std::string_view kValue{"cuda"}; +}; + } // namespace infini::ops::detail #endif diff --git a/tests/conftest.py b/tests/conftest.py index b54f5ed57..f72c54eb6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -25,7 +25,7 @@ def pytest_addoption(parser): "--devices", nargs="+", default=None, - help="Device(s) to test on (e.g., `--devices ascend cpu`). Accepts platform names (`nvidia`, `mars`, `metax`, `iluvatar`, `hygon`, `thead`, `moore`, `cambricon`, `ascend`) or PyTorch device types (`cuda`, `mlu`, `musa`, `npu`). Defaults to all available devices.", + help="Device(s) to test on (e.g., `--devices ascend cpu`). Accepts platform names (`nvidia`, `mars`, `metax`, `iluvatar`, `hygon`, `thead`, `moore`, `cambricon`, `ascend`, `kunlun`) or PyTorch device types (`cuda`, `mlu`, `musa`, `npu`). Defaults to all available devices.", ) parser.addoption( "--report", @@ -199,6 +199,7 @@ def _set_random_seed(seed): "moore": "musa", "cambricon": "mlu", "ascend": "npu", + "kunlun": "cuda", } diff --git a/tests/test_rotary_embedding.py b/tests/test_rotary_embedding.py index c306fa5e7..4bb12a924 100644 --- a/tests/test_rotary_embedding.py +++ b/tests/test_rotary_embedding.py @@ -260,6 +260,78 @@ def test_rotary_embedding_ascend_strided_fused_qkv_view( torch.testing.assert_close(key, expected_key, rtol=rtol, atol=atol) +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16, torch.float32)) +@pytest.mark.parametrize("cache_dtype", (torch.float16, torch.bfloat16, torch.float32)) +@pytest.mark.parametrize("is_neox", (False, True)) +def test_rotary_embedding_kunlun_graph( + dtype, cache_dtype, is_neox, device, implementation_index +): + if torch.device(device).type != "cuda" or not ( + infini.ops.RotaryEmbedding.active_implementation_indices("kunlun") + ): + pytest.skip("requires the Kunlun backend") + # Exercise multiple tiles, partial rotation, batched QKV views, and inverse. + head_size, rot_dim, offset = 520, 512, 2 + fused = torch.randn(2, 4, 8 * head_size, dtype=dtype, device=device) + query = fused[..., : 4 * head_size].view(2, 4, 4, head_size) + key = fused[..., 4 * head_size : 6 * head_size].view(2, 4, 2, head_size) + cache = torch.randn(16, rot_dim + 8, dtype=cache_dtype, device=device)[:, :rot_dim] + positions = torch.zeros((2, 4), dtype=torch.int64, device=device) + + def apply(): + infini.ops.rotary_embedding( + positions, + query, + key, + cache, + head_size, + is_neox, + offset, + True, + implementation_index=implementation_index, + stream=get_stream(query.device), + ) + + apply() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + apply() + for values in ((-1, 0, 15, 16, 2**40, 3, 7, 2**63 - 1), (9, 8, 7, 6, 5, 4, 3, 2)): + original = torch.randn_like(fused) + fused.copy_(original) + positions.copy_( + torch.tensor(values, dtype=torch.int64, device=device).view(2, 4) + ) + expected_query, expected_key = query.cpu(), key.cpu() + _torch_rotary_embedding( + positions.cpu().clamp(0, 15), + expected_query, + expected_key, + cache.cpu(), + head_size, + is_neox, + offset, + True, + ) + graph.replay() + torch.cuda.synchronize() + tolerance = ( + 1e-5 + if dtype == torch.float32 + else (1e-3 if dtype == torch.float16 else 2e-2) + ) + torch.testing.assert_close( + query.cpu(), expected_query, rtol=tolerance, atol=tolerance + ) + torch.testing.assert_close( + key.cpu(), expected_key, rtol=tolerance, atol=tolerance + ) + torch.testing.assert_close( + fused[..., 6 * head_size :], original[..., 6 * head_size :], rtol=0, atol=0 + ) + + def _torch_rotary_embedding( positions, query,