diff --git a/src/native/cambricon/ops/reshape_and_cache_flash/kernel.h b/src/native/cambricon/ops/reshape_and_cache_flash/kernel.h new file mode 100644 index 000000000..44c34b2b7 --- /dev/null +++ b/src/native/cambricon/ops/reshape_and_cache_flash/kernel.h @@ -0,0 +1,81 @@ +#ifndef INFINI_OPS_CAMBRICON_RESHAPE_AND_CACHE_FLASH_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_RESHAPE_AND_CACHE_FLASH_KERNEL_H_ + +#include + +#include "base/reshape_and_cache_flash.h" +#include "dispatcher.h" +#include "native/cambricon/common.h" +#include "native/cambricon/data_type_.h" + +namespace infini::ops { + +template +void ReshapeAndCacheFlashUnion( + int core_per_cluster, int cluster_count, cnrtQueue_t queue, const void* key, + const void* value, const void* slot_mapping, void* key_cache, + void* value_cache, std::size_t num_tokens, std::size_t num_heads, + std::size_t head_size, std::size_t block_size, + std::ptrdiff_t key_token_stride, std::ptrdiff_t value_token_stride, + std::ptrdiff_t key_head_stride, std::ptrdiff_t value_head_stride, + std::ptrdiff_t key_cache_block_stride, + std::ptrdiff_t value_cache_block_stride, + std::ptrdiff_t key_cache_page_stride, + std::ptrdiff_t value_cache_page_stride, + std::ptrdiff_t key_cache_head_stride, + std::ptrdiff_t value_cache_head_stride); + +template <> +class Operator + : public ReshapeAndCacheFlash { + public: + using ReshapeAndCacheFlash::ReshapeAndCacheFlash; + + Operator(const Tensor key, const Tensor value, const Tensor slot_mapping, + const Tensor k_scale, const Tensor v_scale, + const std::string kv_cache_dtype, Tensor key_cache, + Tensor value_cache) + : ReshapeAndCacheFlash{key, value, slot_mapping, + k_scale, v_scale, kv_cache_dtype, + key_cache, value_cache} { + cnrt_utils::GetLaunchConfig(key.device(), &core_per_cluster_, + &cluster_count_); + } + + void operator()(const Tensor key, const Tensor value, + const Tensor slot_mapping, const Tensor /*k_scale*/, + const Tensor /*v_scale*/, + const std::string /*kv_cache_dtype*/, Tensor key_cache, + Tensor value_cache) const override { + if (num_tokens_ == 0) { + return; + } + + auto queue = static_cast(stream_ ? stream_ : 0); + DispatchFunc< + Device::Type::kCambricon, + List>( + {dtype_}, + [&](auto dtype_tag) { + using T = typename decltype(dtype_tag)::type; + ReshapeAndCacheFlashUnion( + core_per_cluster_, cluster_count_, queue, key.data(), + value.data(), slot_mapping.data(), key_cache.data(), + value_cache.data(), num_tokens_, num_heads_, head_size_, + block_size_, key_token_stride_, value_token_stride_, + key_head_stride_, value_head_stride_, key_cache_block_stride_, + value_cache_block_stride_, key_cache_page_stride_, + value_cache_page_stride_, key_cache_head_stride_, + value_cache_head_stride_); + }, + "CambriconReshapeAndCacheFlash::operator()"); + } + + private: + int core_per_cluster_{0}; + int cluster_count_{0}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/reshape_and_cache_flash/kernel.mlu b/src/native/cambricon/ops/reshape_and_cache_flash/kernel.mlu new file mode 100644 index 000000000..5def4e0b2 --- /dev/null +++ b/src/native/cambricon/ops/reshape_and_cache_flash/kernel.mlu @@ -0,0 +1,97 @@ +#include "kernel.h" + +namespace infini::ops { +namespace reshape_and_cache_flash_detail { + +template +__mlu_global__ void ReshapeAndCacheFlashKernel( + const T* key, const T* value, const int64_t* slot_mapping, T* key_cache, + T* value_cache, std::size_t num_tokens, std::size_t num_heads, + std::size_t head_size, std::size_t block_size, + std::ptrdiff_t key_token_stride, std::ptrdiff_t value_token_stride, + std::ptrdiff_t key_head_stride, std::ptrdiff_t value_head_stride, + std::ptrdiff_t key_cache_block_stride, + std::ptrdiff_t value_cache_block_stride, + std::ptrdiff_t key_cache_page_stride, + std::ptrdiff_t value_cache_page_stride, + std::ptrdiff_t key_cache_head_stride, + std::ptrdiff_t value_cache_head_stride) { + const std::size_t task_count = num_tokens * num_heads; + for (std::size_t task = taskId; task < task_count; task += taskDim) { + const std::size_t token_idx = task / num_heads; + const std::size_t head_idx = task % num_heads; + const int64_t slot = slot_mapping[token_idx]; + if (slot < 0) { + continue; + } + + const auto physical_block = static_cast(slot) / block_size; + const auto block_offset = static_cast(slot) % block_size; + const T* key_src = + key + token_idx * key_token_stride + head_idx * key_head_stride; + const T* value_src = + value + token_idx * value_token_stride + head_idx * value_head_stride; + T* key_dst = key_cache + physical_block * key_cache_block_stride + + block_offset * key_cache_page_stride + + head_idx * key_cache_head_stride; + T* value_dst = value_cache + physical_block * value_cache_block_stride + + block_offset * value_cache_page_stride + + head_idx * value_cache_head_stride; + + for (std::size_t i = 0; i < head_size; ++i) { + key_dst[i] = key_src[i]; + value_dst[i] = value_src[i]; + } + } +} + +} // namespace reshape_and_cache_flash_detail + +template +void ReshapeAndCacheFlashUnion( + int core_per_cluster, int cluster_count, cnrtQueue_t queue, const void* key, + const void* value, const void* slot_mapping, void* key_cache, + void* value_cache, std::size_t num_tokens, std::size_t num_heads, + std::size_t head_size, std::size_t block_size, + std::ptrdiff_t key_token_stride, std::ptrdiff_t value_token_stride, + std::ptrdiff_t key_head_stride, std::ptrdiff_t value_head_stride, + std::ptrdiff_t key_cache_block_stride, + std::ptrdiff_t value_cache_block_stride, + std::ptrdiff_t key_cache_page_stride, + std::ptrdiff_t value_cache_page_stride, + std::ptrdiff_t key_cache_head_stride, + std::ptrdiff_t value_cache_head_stride) { + cnrtDim3_t kernel_dim; + kernel_dim.x = core_per_cluster; + kernel_dim.y = cluster_count; + kernel_dim.z = 1; + + (void)cnrtGetLastError(); + reshape_and_cache_flash_detail::ReshapeAndCacheFlashKernel + <<>>( + reinterpret_cast(key), reinterpret_cast(value), + reinterpret_cast(slot_mapping), + reinterpret_cast(key_cache), reinterpret_cast(value_cache), + num_tokens, num_heads, head_size, block_size, key_token_stride, + value_token_stride, key_head_stride, value_head_stride, + key_cache_block_stride, value_cache_block_stride, + key_cache_page_stride, value_cache_page_stride, key_cache_head_stride, + value_cache_head_stride); + CNRT_CHECK(cnrtGetLastError()); +} + +#define INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(T) \ + template void ReshapeAndCacheFlashUnion( \ + int, int, cnrtQueue_t, const void*, const void*, const void*, void*, \ + void*, std::size_t, std::size_t, std::size_t, std::size_t, \ + std::ptrdiff_t, std::ptrdiff_t, std::ptrdiff_t, std::ptrdiff_t, \ + std::ptrdiff_t, std::ptrdiff_t, std::ptrdiff_t, std::ptrdiff_t, \ + std::ptrdiff_t, std::ptrdiff_t) + +INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(__half); +INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(__bang_bfloat16); +INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(float); + +#undef INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH + +} // namespace infini::ops