diff --git a/src/infiniop/devices/ascend/CMakeLists.txt b/src/infiniop/devices/ascend/CMakeLists.txt index 5229f770c..5495687ba 100644 --- a/src/infiniop/devices/ascend/CMakeLists.txt +++ b/src/infiniop/devices/ascend/CMakeLists.txt @@ -28,6 +28,7 @@ ascendc_library(ascend_kernels STATIC ../../ops/rope/ascend/rope_ascend_kernel.cpp ../../ops/random_sample/ascend/random_sample_kernel.cpp ../../ops/rms_norm/ascend/cast_kernel.cpp + ../../ops/kv_caching/ascend/kv_caching_ascend_kernel.cpp ../../ops/paged_caching/ascend/paged_caching_ascend_kernel.cpp ../../ops/paged_attention/ascend/paged_attention_ascend_kernel.cpp ../../ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend_kernel.cpp diff --git a/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend.cc b/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend.cc new file mode 100644 index 000000000..d58f45ee4 --- /dev/null +++ b/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend.cc @@ -0,0 +1,53 @@ +#include "kv_caching_ascend.h" +#include "../../../devices/ascend/common_ascend.h" + +namespace op::kv_caching::ascend { + +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t k_cache, + infiniopTensorDescriptor_t v_cache, + infiniopTensorDescriptor_t k, + infiniopTensorDescriptor_t v, + infiniopTensorDescriptor_t past_kv_lengths) { + auto info = KVCachingInfo::createKVCachingInfo(k_cache, v_cache, k, v, past_kv_lengths); + CHECK_RESULT(info); + + auto handle_ascend = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + new Opaque{}, + info.take(), 0, handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + void *stream) const { + (void)workspace; + (void)workspace_size; + + return kv_caching_kernel_launch( + k_cache, v_cache, k, v, past_kv_lengths, + _info.dtype, _info.past_len_dtype, + _info.batch_size, _info.num_kv_heads, _info.seq_len, _info.hidden_dim, + _info.k_cache_strides_0, _info.k_cache_strides_1, _info.k_cache_strides_2, _info.k_cache_strides_3, + _info.v_cache_strides_0, _info.v_cache_strides_1, _info.v_cache_strides_2, _info.v_cache_strides_3, + _info.k_strides_0, _info.k_strides_1, _info.k_strides_2, _info.k_strides_3, + _info.v_strides_0, _info.v_strides_1, _info.v_strides_2, _info.v_strides_3, + stream); +} + +} // namespace op::kv_caching::ascend diff --git a/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend.h b/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend.h new file mode 100644 index 000000000..d737835a6 --- /dev/null +++ b/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend.h @@ -0,0 +1,39 @@ + +#ifndef __KV_CACHING_ASCEND_API_H__ +#define __KV_CACHING_ASCEND_API_H__ + +#include "../kv_caching.h" + +extern "C" infiniStatus_t kv_caching_kernel_launch( + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + infiniDtype_t dtype, + infiniDtype_t past_len_dtype, + size_t batch_size, + size_t num_kv_heads, + size_t seq_len, + size_t hidden_dim, + ptrdiff_t k_cache_strides_0, + ptrdiff_t k_cache_strides_1, + ptrdiff_t k_cache_strides_2, + ptrdiff_t k_cache_strides_3, + ptrdiff_t v_cache_strides_0, + ptrdiff_t v_cache_strides_1, + ptrdiff_t v_cache_strides_2, + ptrdiff_t v_cache_strides_3, + ptrdiff_t k_strides_0, + ptrdiff_t k_strides_1, + ptrdiff_t k_strides_2, + ptrdiff_t k_strides_3, + ptrdiff_t v_strides_0, + ptrdiff_t v_strides_1, + ptrdiff_t v_strides_2, + ptrdiff_t v_strides_3, + void *stream); + +DESCRIPTOR(ascend) + +#endif // __KV_CACHING_ASCEND_API_H__ diff --git a/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend_kernel.cpp b/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend_kernel.cpp new file mode 100644 index 000000000..57fc57419 --- /dev/null +++ b/src/infiniop/ops/kv_caching/ascend/kv_caching_ascend_kernel.cpp @@ -0,0 +1,216 @@ + +#include "../../../devices/ascend/ascend_kernel_common.h" + +using namespace AscendC; + +template +class KVCachingKernel { +public: + __aicore__ inline KVCachingKernel() {} + + __aicore__ inline void init( + GM_ADDR k_cache, + GM_ADDR v_cache, + GM_ADDR k, + GM_ADDR v, + GM_ADDR past_kv_lengths, + size_t batch_size, + size_t num_kv_heads, + size_t seq_len, + size_t hidden_dim, + ptrdiff_t k_cache_strides_0, + ptrdiff_t k_cache_strides_1, + ptrdiff_t k_cache_strides_2, + ptrdiff_t k_cache_strides_3, + ptrdiff_t v_cache_strides_0, + ptrdiff_t v_cache_strides_1, + ptrdiff_t v_cache_strides_2, + ptrdiff_t v_cache_strides_3, + ptrdiff_t k_strides_0, + ptrdiff_t k_strides_1, + ptrdiff_t k_strides_2, + ptrdiff_t k_strides_3, + ptrdiff_t v_strides_0, + ptrdiff_t v_strides_1, + ptrdiff_t v_strides_2, + ptrdiff_t v_strides_3) { + _batch_size = batch_size; + _num_kv_heads = num_kv_heads; + _seq_len = seq_len; + _hidden_dim = hidden_dim; + _k_cache_strides_0 = k_cache_strides_0; + _k_cache_strides_1 = k_cache_strides_1; + _k_cache_strides_2 = k_cache_strides_2; + _k_cache_strides_3 = k_cache_strides_3; + _v_cache_strides_0 = v_cache_strides_0; + _v_cache_strides_1 = v_cache_strides_1; + _v_cache_strides_2 = v_cache_strides_2; + _v_cache_strides_3 = v_cache_strides_3; + _k_strides_0 = k_strides_0; + _k_strides_1 = k_strides_1; + _k_strides_2 = k_strides_2; + _k_strides_3 = k_strides_3; + _v_strides_0 = v_strides_0; + _v_strides_1 = v_strides_1; + _v_strides_2 = v_strides_2; + _v_strides_3 = v_strides_3; + + _k_cache_gm.SetGlobalBuffer(reinterpret_cast<__gm__ Tdata *>(k_cache)); + _v_cache_gm.SetGlobalBuffer(reinterpret_cast<__gm__ Tdata *>(v_cache)); + _k_gm.SetGlobalBuffer(reinterpret_cast<__gm__ Tdata *>(k)); + _v_gm.SetGlobalBuffer(reinterpret_cast<__gm__ Tdata *>(v)); + _past_kv_lengths_gm.SetGlobalBuffer(reinterpret_cast<__gm__ Tidx *>(past_kv_lengths)); + } + + __aicore__ inline void process() { + const size_t work_idx = GetBlockIdx(); + const size_t total = _batch_size * _num_kv_heads * _seq_len; + if (work_idx >= total) { + return; + } + + size_t idx = work_idx; + size_t s = idx % _seq_len; + idx /= _seq_len; + size_t h = idx % _num_kv_heads; + size_t b = idx / _num_kv_heads; + + ptrdiff_t cache_s = static_cast(_past_kv_lengths_gm.GetValue(b)) + static_cast(s); + ptrdiff_t k_cache_base = static_cast(b) * _k_cache_strides_0 + + static_cast(h) * _k_cache_strides_1 + + cache_s * _k_cache_strides_2; + ptrdiff_t v_cache_base = static_cast(b) * _v_cache_strides_0 + + static_cast(h) * _v_cache_strides_1 + + cache_s * _v_cache_strides_2; + ptrdiff_t k_src_base = static_cast(b) * _k_strides_0 + + static_cast(h) * _k_strides_1 + + static_cast(s) * _k_strides_2; + ptrdiff_t v_src_base = static_cast(b) * _v_strides_0 + + static_cast(h) * _v_strides_1 + + static_cast(s) * _v_strides_2; + + for (size_t d = 0; d < _hidden_dim; ++d) { + ptrdiff_t d_offset = static_cast(d); + _k_cache_gm.SetValue(k_cache_base + d_offset * _k_cache_strides_3, + _k_gm.GetValue(k_src_base + d_offset * _k_strides_3)); + _v_cache_gm.SetValue(v_cache_base + d_offset * _v_cache_strides_3, + _v_gm.GetValue(v_src_base + d_offset * _v_strides_3)); + } + } + +private: + GlobalTensor _k_cache_gm; + GlobalTensor _v_cache_gm; + GlobalTensor _k_gm; + GlobalTensor _v_gm; + GlobalTensor _past_kv_lengths_gm; + size_t _batch_size; + size_t _num_kv_heads; + size_t _seq_len; + size_t _hidden_dim; + ptrdiff_t _k_cache_strides_0; + ptrdiff_t _k_cache_strides_1; + ptrdiff_t _k_cache_strides_2; + ptrdiff_t _k_cache_strides_3; + ptrdiff_t _v_cache_strides_0; + ptrdiff_t _v_cache_strides_1; + ptrdiff_t _v_cache_strides_2; + ptrdiff_t _v_cache_strides_3; + ptrdiff_t _k_strides_0; + ptrdiff_t _k_strides_1; + ptrdiff_t _k_strides_2; + ptrdiff_t _k_strides_3; + ptrdiff_t _v_strides_0; + ptrdiff_t _v_strides_1; + ptrdiff_t _v_strides_2; + ptrdiff_t _v_strides_3; +}; + +#define DEFINE_KV_CACHING_KERNEL(KERNEL_NAME, TYPE, IDX_TYPE) \ + extern "C" __global__ __aicore__ void KERNEL_NAME( \ + GM_ADDR k_cache, GM_ADDR v_cache, GM_ADDR k, GM_ADDR v, GM_ADDR past_kv_lengths, \ + size_t batch_size, size_t num_kv_heads, size_t seq_len, size_t hidden_dim, \ + ptrdiff_t k_cache_strides_0, ptrdiff_t k_cache_strides_1, \ + ptrdiff_t k_cache_strides_2, ptrdiff_t k_cache_strides_3, \ + ptrdiff_t v_cache_strides_0, ptrdiff_t v_cache_strides_1, \ + ptrdiff_t v_cache_strides_2, ptrdiff_t v_cache_strides_3, \ + ptrdiff_t k_strides_0, ptrdiff_t k_strides_1, ptrdiff_t k_strides_2, \ + ptrdiff_t k_strides_3, ptrdiff_t v_strides_0, ptrdiff_t v_strides_1, \ + ptrdiff_t v_strides_2, ptrdiff_t v_strides_3) { \ + KVCachingKernel op; \ + op.init(k_cache, v_cache, k, v, past_kv_lengths, batch_size, num_kv_heads, \ + seq_len, hidden_dim, k_cache_strides_0, k_cache_strides_1, \ + k_cache_strides_2, k_cache_strides_3, v_cache_strides_0, \ + v_cache_strides_1, v_cache_strides_2, v_cache_strides_3, \ + k_strides_0, k_strides_1, k_strides_2, k_strides_3, \ + v_strides_0, v_strides_1, v_strides_2, v_strides_3); \ + op.process(); \ + } + +DEFINE_KV_CACHING_KERNEL(kv_caching_kernel_f16_i32, half, int32_t) +DEFINE_KV_CACHING_KERNEL(kv_caching_kernel_f16_i64, half, int64_t) +DEFINE_KV_CACHING_KERNEL(kv_caching_kernel_bf16_i32, bfloat16_t, int32_t) +DEFINE_KV_CACHING_KERNEL(kv_caching_kernel_bf16_i64, bfloat16_t, int64_t) +DEFINE_KV_CACHING_KERNEL(kv_caching_kernel_f32_i32, float, int32_t) +DEFINE_KV_CACHING_KERNEL(kv_caching_kernel_f32_i64, float, int64_t) + +#undef DEFINE_KV_CACHING_KERNEL + +extern "C" infiniStatus_t kv_caching_kernel_launch( + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + infiniDtype_t dtype, + infiniDtype_t past_len_dtype, + size_t batch_size, + size_t num_kv_heads, + size_t seq_len, + size_t hidden_dim, + ptrdiff_t k_cache_strides_0, + ptrdiff_t k_cache_strides_1, + ptrdiff_t k_cache_strides_2, + ptrdiff_t k_cache_strides_3, + ptrdiff_t v_cache_strides_0, + ptrdiff_t v_cache_strides_1, + ptrdiff_t v_cache_strides_2, + ptrdiff_t v_cache_strides_3, + ptrdiff_t k_strides_0, + ptrdiff_t k_strides_1, + ptrdiff_t k_strides_2, + ptrdiff_t k_strides_3, + ptrdiff_t v_strides_0, + ptrdiff_t v_strides_1, + ptrdiff_t v_strides_2, + ptrdiff_t v_strides_3, + void *stream) { + const size_t block_dim = batch_size * num_kv_heads * seq_len; + if (block_dim == 0) { + return INFINI_STATUS_SUCCESS; + } + +#define LAUNCH_KV_CACHING(DTYPE_ENUM, IDX_ENUM, KERNEL_NAME) \ + if (dtype == DTYPE_ENUM && past_len_dtype == IDX_ENUM) { \ + KERNEL_NAME<<>>( \ + k_cache, v_cache, const_cast(k), const_cast(v), \ + const_cast(past_kv_lengths), batch_size, num_kv_heads, \ + seq_len, hidden_dim, k_cache_strides_0, k_cache_strides_1, \ + k_cache_strides_2, k_cache_strides_3, v_cache_strides_0, \ + v_cache_strides_1, v_cache_strides_2, v_cache_strides_3, \ + k_strides_0, k_strides_1, k_strides_2, k_strides_3, \ + v_strides_0, v_strides_1, v_strides_2, v_strides_3); \ + return INFINI_STATUS_SUCCESS; \ + } + + LAUNCH_KV_CACHING(INFINI_DTYPE_F16, INFINI_DTYPE_I32, kv_caching_kernel_f16_i32) + LAUNCH_KV_CACHING(INFINI_DTYPE_F16, INFINI_DTYPE_I64, kv_caching_kernel_f16_i64) + LAUNCH_KV_CACHING(INFINI_DTYPE_BF16, INFINI_DTYPE_I32, kv_caching_kernel_bf16_i32) + LAUNCH_KV_CACHING(INFINI_DTYPE_BF16, INFINI_DTYPE_I64, kv_caching_kernel_bf16_i64) + LAUNCH_KV_CACHING(INFINI_DTYPE_F32, INFINI_DTYPE_I32, kv_caching_kernel_f32_i32) + LAUNCH_KV_CACHING(INFINI_DTYPE_F32, INFINI_DTYPE_I64, kv_caching_kernel_f32_i64) + + return INFINI_STATUS_BAD_TENSOR_DTYPE; + +#undef LAUNCH_KV_CACHING +} diff --git a/src/infiniop/ops/kv_caching/bang/kv_caching_bang.h b/src/infiniop/ops/kv_caching/bang/kv_caching_bang.h new file mode 100644 index 000000000..33a48fad7 --- /dev/null +++ b/src/infiniop/ops/kv_caching/bang/kv_caching_bang.h @@ -0,0 +1,9 @@ + +#ifndef __KV_CACHING_BANG_API_H__ +#define __KV_CACHING_BANG_API_H__ + +#include "../kv_caching.h" + +DESCRIPTOR(bang) + +#endif // __KV_CACHING_BANG_API_H__ diff --git a/src/infiniop/ops/kv_caching/bang/kv_caching_bang.mlu b/src/infiniop/ops/kv_caching/bang/kv_caching_bang.mlu new file mode 100644 index 000000000..709786b42 --- /dev/null +++ b/src/infiniop/ops/kv_caching/bang/kv_caching_bang.mlu @@ -0,0 +1,199 @@ + +#include "../../../devices/bang/common_bang.h" +#include "kv_caching_bang.h" + +namespace { + +template +__mlu_global__ void kvCachingBangKernel( + Tdata *k_cache, + Tdata *v_cache, + const Tdata *k, + const Tdata *v, + const Tidx *past_kv_lengths, + size_t batch_size, + size_t num_kv_heads, + size_t seq_len, + size_t hidden_dim, + ptrdiff_t k_cache_strides_0, + ptrdiff_t k_cache_strides_1, + ptrdiff_t k_cache_strides_2, + ptrdiff_t k_cache_strides_3, + ptrdiff_t v_cache_strides_0, + ptrdiff_t v_cache_strides_1, + ptrdiff_t v_cache_strides_2, + ptrdiff_t v_cache_strides_3, + ptrdiff_t k_strides_0, + ptrdiff_t k_strides_1, + ptrdiff_t k_strides_2, + ptrdiff_t k_strides_3, + ptrdiff_t v_strides_0, + ptrdiff_t v_strides_1, + ptrdiff_t v_strides_2, + ptrdiff_t v_strides_3) { + const size_t task_num = taskDimX * taskDimY; + const bool copy_hidden_contiguous = k_cache_strides_3 == 1 + && v_cache_strides_3 == 1 + && k_strides_3 == 1 + && v_strides_3 == 1; + + if (copy_hidden_contiguous) { + const size_t total_rows = batch_size * num_kv_heads * seq_len; + const size_t bytes = hidden_dim * sizeof(Tdata); + for (size_t linear = taskId; linear < total_rows; linear += task_num) { + size_t idx = linear; + size_t s = idx % seq_len; + idx /= seq_len; + size_t h = idx % num_kv_heads; + size_t b = idx / num_kv_heads; + + ptrdiff_t cache_s = static_cast(past_kv_lengths[b]) + static_cast(s); + ptrdiff_t k_cache_base = static_cast(b) * k_cache_strides_0 + + static_cast(h) * k_cache_strides_1 + + cache_s * k_cache_strides_2; + ptrdiff_t v_cache_base = static_cast(b) * v_cache_strides_0 + + static_cast(h) * v_cache_strides_1 + + cache_s * v_cache_strides_2; + ptrdiff_t k_src_base = static_cast(b) * k_strides_0 + + static_cast(h) * k_strides_1 + + static_cast(s) * k_strides_2; + ptrdiff_t v_src_base = static_cast(b) * v_strides_0 + + static_cast(h) * v_strides_1 + + static_cast(s) * v_strides_2; + + __memcpy(k_cache + k_cache_base, k + k_src_base, bytes, GDRAM2GDRAM); + __memcpy(v_cache + v_cache_base, v + v_src_base, bytes, GDRAM2GDRAM); + } + return; + } + + const size_t total = batch_size * num_kv_heads * seq_len * hidden_dim; + for (size_t linear = taskId; linear < total; linear += task_num) { + size_t idx = linear; + size_t d = idx % hidden_dim; + idx /= hidden_dim; + size_t s = idx % seq_len; + idx /= seq_len; + size_t h = idx % num_kv_heads; + size_t b = idx / num_kv_heads; + + ptrdiff_t cache_s = static_cast(past_kv_lengths[b]) + static_cast(s); + ptrdiff_t k_cache_offset = static_cast(b) * k_cache_strides_0 + + static_cast(h) * k_cache_strides_1 + + cache_s * k_cache_strides_2 + + static_cast(d) * k_cache_strides_3; + ptrdiff_t v_cache_offset = static_cast(b) * v_cache_strides_0 + + static_cast(h) * v_cache_strides_1 + + cache_s * v_cache_strides_2 + + static_cast(d) * v_cache_strides_3; + ptrdiff_t k_src_offset = static_cast(b) * k_strides_0 + + static_cast(h) * k_strides_1 + + static_cast(s) * k_strides_2 + + static_cast(d) * k_strides_3; + ptrdiff_t v_src_offset = static_cast(b) * v_strides_0 + + static_cast(h) * v_strides_1 + + static_cast(s) * v_strides_2 + + static_cast(d) * v_strides_3; + k_cache[k_cache_offset] = k[k_src_offset]; + v_cache[v_cache_offset] = v[v_src_offset]; + } +} + +template +infiniStatus_t launchBangKernel( + cnrtQueue_t queue, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + const op::kv_caching::KVCachingInfo &info, + int core_per_cluster, + int cluster_count) { + size_t total = info.batch_size * info.num_kv_heads * info.seq_len * info.hidden_dim; + if (total == 0) { + return INFINI_STATUS_SUCCESS; + } + cnrtDim3_t kernel_dim; + kernel_dim.x = core_per_cluster; + kernel_dim.y = cluster_count; + kernel_dim.z = 1; + + kvCachingBangKernel<<>>( + static_cast(k_cache), static_cast(v_cache), + static_cast(k), static_cast(v), + static_cast(past_kv_lengths), + info.batch_size, info.num_kv_heads, info.seq_len, info.hidden_dim, + info.k_cache_strides_0, info.k_cache_strides_1, info.k_cache_strides_2, info.k_cache_strides_3, + info.v_cache_strides_0, info.v_cache_strides_1, info.v_cache_strides_2, info.v_cache_strides_3, + info.k_strides_0, info.k_strides_1, info.k_strides_2, info.k_strides_3, + info.v_strides_0, info.v_strides_1, info.v_strides_2, info.v_strides_3); + return INFINI_STATUS_SUCCESS; +} + +} // namespace + +namespace op::kv_caching::bang { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t k_cache, + infiniopTensorDescriptor_t v_cache, + infiniopTensorDescriptor_t k, + infiniopTensorDescriptor_t v, + infiniopTensorDescriptor_t past_kv_lengths) { + auto handle = reinterpret_cast(handle_); + auto info = KVCachingInfo::createKVCachingInfo(k_cache, v_cache, k, v, past_kv_lengths); + CHECK_RESULT(info); + + *desc_ptr = new Descriptor( + new Opaque{static_cast(handle)->internal()}, + info.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + void *stream) const { + (void)workspace; + (void)workspace_size; + auto queue = reinterpret_cast(stream); + const int core_per_cluster = _opaque->internal->getCorePerCluster(); + const int cluster_count = _opaque->internal->getClusterCount(); + +#define LAUNCH_WITH_IDX(TDATA) \ + if (_info.past_len_dtype == INFINI_DTYPE_I32) { \ + return launchBangKernel(queue, k_cache, v_cache, k, v, past_kv_lengths, _info, core_per_cluster, cluster_count); \ + } \ + return launchBangKernel(queue, k_cache, v_cache, k, v, past_kv_lengths, _info, core_per_cluster, cluster_count) + + switch (_info.dtype) { + case INFINI_DTYPE_F16: + LAUNCH_WITH_IDX(half); + case INFINI_DTYPE_BF16: + LAUNCH_WITH_IDX(bfloat16_t); + case INFINI_DTYPE_F32: + LAUNCH_WITH_IDX(float); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + +#undef LAUNCH_WITH_IDX +} + +} // namespace op::kv_caching::bang diff --git a/src/infiniop/ops/kv_caching/cpu/kv_caching_cpu.cc b/src/infiniop/ops/kv_caching/cpu/kv_caching_cpu.cc new file mode 100644 index 000000000..a6e52d643 --- /dev/null +++ b/src/infiniop/ops/kv_caching/cpu/kv_caching_cpu.cc @@ -0,0 +1,126 @@ +#include "kv_caching_cpu.h" +#include "../../../devices/cpu/common_cpu.h" + +namespace op::kv_caching::cpu { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t k_cache, + infiniopTensorDescriptor_t v_cache, + infiniopTensorDescriptor_t k, + infiniopTensorDescriptor_t v, + infiniopTensorDescriptor_t past_kv_lengths) { + + auto info = KVCachingInfo::createKVCachingInfo(k_cache, v_cache, k, v, past_kv_lengths); + CHECK_RESULT(info); + + *desc_ptr = new Descriptor( + nullptr, + info.take(), + 0, + handle->device, + handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +template +static void calculateImpl( + const KVCachingInfo &info, + Tdata *k_cache, + Tdata *v_cache, + const Tdata *k, + const Tdata *v, + const Tidx *past_kv_lengths) { + + const auto batch_size = static_cast(info.batch_size); + const auto num_kv_heads = static_cast(info.num_kv_heads); + const auto seq_len = static_cast(info.seq_len); + const auto hidden_dim = static_cast(info.hidden_dim); + + const ptrdiff_t total = batch_size * num_kv_heads * seq_len * hidden_dim; + +#ifdef ENABLE_OMP +#pragma omp parallel for schedule(static) +#endif + for (ptrdiff_t linear = 0; linear < total; ++linear) { + ptrdiff_t idx = linear; + const ptrdiff_t d = idx % hidden_dim; + idx /= hidden_dim; + const ptrdiff_t s = idx % seq_len; + idx /= seq_len; + const ptrdiff_t h = idx % num_kv_heads; + const ptrdiff_t b = idx / num_kv_heads; + + const ptrdiff_t cache_s = static_cast(past_kv_lengths[b]) + s; + + const ptrdiff_t k_cache_offset = b * info.k_cache_strides_0 + + h * info.k_cache_strides_1 + + cache_s * info.k_cache_strides_2 + + d * info.k_cache_strides_3; + const ptrdiff_t v_cache_offset = b * info.v_cache_strides_0 + + h * info.v_cache_strides_1 + + cache_s * info.v_cache_strides_2 + + d * info.v_cache_strides_3; + const ptrdiff_t k_offset = b * info.k_strides_0 + + h * info.k_strides_1 + + s * info.k_strides_2 + + d * info.k_strides_3; + const ptrdiff_t v_offset = b * info.v_strides_0 + + h * info.v_strides_1 + + s * info.v_strides_2 + + d * info.v_strides_3; + + k_cache[k_cache_offset] = k[k_offset]; + v_cache[v_cache_offset] = v[v_offset]; + } +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + void *stream) const { + +#define DISPATCH_PAST_LEN(TDATA) \ + if (_info.past_len_dtype == INFINI_DTYPE_I32) { \ + calculateImpl(_info, \ + reinterpret_cast(k_cache), \ + reinterpret_cast(v_cache), \ + reinterpret_cast(k), \ + reinterpret_cast(v), \ + reinterpret_cast(past_kv_lengths)); \ + return INFINI_STATUS_SUCCESS; \ + } \ + if (_info.past_len_dtype == INFINI_DTYPE_I64) { \ + calculateImpl(_info, \ + reinterpret_cast(k_cache), \ + reinterpret_cast(v_cache), \ + reinterpret_cast(k), \ + reinterpret_cast(v), \ + reinterpret_cast(past_kv_lengths)); \ + return INFINI_STATUS_SUCCESS; \ + } \ + return INFINI_STATUS_BAD_TENSOR_DTYPE + + switch (_info.dtype) { + case INFINI_DTYPE_F16: + DISPATCH_PAST_LEN(fp16_t); + case INFINI_DTYPE_BF16: + DISPATCH_PAST_LEN(bf16_t); + case INFINI_DTYPE_F32: + DISPATCH_PAST_LEN(float); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + +#undef DISPATCH_PAST_LEN +} + +} // namespace op::kv_caching::cpu diff --git a/src/infiniop/ops/kv_caching/cpu/kv_caching_cpu.h b/src/infiniop/ops/kv_caching/cpu/kv_caching_cpu.h new file mode 100644 index 000000000..7441719c4 --- /dev/null +++ b/src/infiniop/ops/kv_caching/cpu/kv_caching_cpu.h @@ -0,0 +1,8 @@ +#ifndef __KV_CACHING_CPU_H__ +#define __KV_CACHING_CPU_H__ + +#include "../kv_caching.h" + +DESCRIPTOR(cpu) + +#endif // __KV_CACHING_CPU_H__ diff --git a/src/infiniop/ops/kv_caching/kunlun/kv_caching_kunlun.h b/src/infiniop/ops/kv_caching/kunlun/kv_caching_kunlun.h new file mode 100644 index 000000000..a7cd5ea9d --- /dev/null +++ b/src/infiniop/ops/kv_caching/kunlun/kv_caching_kunlun.h @@ -0,0 +1,9 @@ + +#ifndef __KV_CACHING_KUNLUN_API_H__ +#define __KV_CACHING_KUNLUN_API_H__ + +#include "../kv_caching.h" + +DESCRIPTOR(kunlun) + +#endif // __KV_CACHING_KUNLUN_API_H__ diff --git a/src/infiniop/ops/kv_caching/kunlun/kv_caching_kunlun.xpu b/src/infiniop/ops/kv_caching/kunlun/kv_caching_kunlun.xpu new file mode 100644 index 000000000..6ce1d3394 --- /dev/null +++ b/src/infiniop/ops/kv_caching/kunlun/kv_caching_kunlun.xpu @@ -0,0 +1,172 @@ + +#include "../../../devices/kunlun/kunlun_common.h" +#include "../../../devices/kunlun/kunlun_handle.h" +#include "../../../devices/kunlun/kunlun_kernel_common.h" +#include "kv_caching_kunlun.h" + +namespace { + +template +__global__ void kvCachingKunlunKernel( + Tmem *k_cache, + Tmem *v_cache, + const Tmem *k, + const Tmem *v, + const Tidx *past_kv_lengths, + uint32_t batch_size, + uint32_t num_kv_heads, + uint32_t seq_len, + uint32_t hidden_dim, + int32_t k_cache_strides_0, + int32_t k_cache_strides_1, + int32_t k_cache_strides_2, + int32_t k_cache_strides_3, + int32_t v_cache_strides_0, + int32_t v_cache_strides_1, + int32_t v_cache_strides_2, + int32_t v_cache_strides_3, + int32_t k_strides_0, + int32_t k_strides_1, + int32_t k_strides_2, + int32_t k_strides_3, + int32_t v_strides_0, + int32_t v_strides_1, + int32_t v_strides_2, + int32_t v_strides_3) { + int cid = core_id(); + int ncores = core_num(); + if (cid >= ncores) { + return; + } + uint32_t thread_id = static_cast(ncores * cluster_id() + cid); + uint32_t nthreads = static_cast(ncores * cluster_num()); + uint32_t total = batch_size * num_kv_heads * seq_len * hidden_dim; + + for (uint32_t linear = thread_id; linear < total; linear += nthreads) { + uint32_t idx = linear; + uint32_t d = idx % hidden_dim; + idx /= hidden_dim; + uint32_t s = idx % seq_len; + idx /= seq_len; + uint32_t h = idx % num_kv_heads; + uint32_t b = idx / num_kv_heads; + + int32_t cache_s = static_cast(past_kv_lengths[b]) + static_cast(s); + int32_t k_cache_offset = static_cast(b) * k_cache_strides_0 + + static_cast(h) * k_cache_strides_1 + + cache_s * k_cache_strides_2 + + static_cast(d) * k_cache_strides_3; + int32_t v_cache_offset = static_cast(b) * v_cache_strides_0 + + static_cast(h) * v_cache_strides_1 + + cache_s * v_cache_strides_2 + + static_cast(d) * v_cache_strides_3; + int32_t k_src_offset = static_cast(b) * k_strides_0 + + static_cast(h) * k_strides_1 + + static_cast(s) * k_strides_2 + + static_cast(d) * k_strides_3; + int32_t v_src_offset = static_cast(b) * v_strides_0 + + static_cast(h) * v_strides_1 + + static_cast(s) * v_strides_2 + + static_cast(d) * v_strides_3; + k_cache[k_cache_offset] = k[k_src_offset]; + v_cache[v_cache_offset] = v[v_src_offset]; + } +} + +} // namespace + +namespace op::kv_caching::kunlun { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t k_cache, + infiniopTensorDescriptor_t v_cache, + infiniopTensorDescriptor_t k, + infiniopTensorDescriptor_t v, + infiniopTensorDescriptor_t past_kv_lengths) { + auto handle = reinterpret_cast(handle_); + auto info = KVCachingInfo::createKVCachingInfo(k_cache, v_cache, k, v, past_kv_lengths); + CHECK_RESULT(info); + + *desc_ptr = new Descriptor( + new Opaque{handle->internal()}, + info.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +template +infiniStatus_t launchKunlunKernel( + kunlunStream_t stream, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + const KVCachingInfo &info) { + size_t total = info.batch_size * info.num_kv_heads * info.seq_len * info.hidden_dim; + if (total == 0) { + return INFINI_STATUS_SUCCESS; + } + kvCachingKunlunKernel + <<<12, 64, stream>>>( + reinterpret_cast<__global_ptr__ Tmem *>(k_cache), + reinterpret_cast<__global_ptr__ Tmem *>(v_cache), + reinterpret_cast<__global_ptr__ const Tmem *>(k), + reinterpret_cast<__global_ptr__ const Tmem *>(v), + reinterpret_cast<__global_ptr__ const Tidx *>(past_kv_lengths), + static_cast(info.batch_size), static_cast(info.num_kv_heads), + static_cast(info.seq_len), static_cast(info.hidden_dim), + static_cast(info.k_cache_strides_0), static_cast(info.k_cache_strides_1), + static_cast(info.k_cache_strides_2), static_cast(info.k_cache_strides_3), + static_cast(info.v_cache_strides_0), static_cast(info.v_cache_strides_1), + static_cast(info.v_cache_strides_2), static_cast(info.v_cache_strides_3), + static_cast(info.k_strides_0), static_cast(info.k_strides_1), + static_cast(info.k_strides_2), static_cast(info.k_strides_3), + static_cast(info.v_strides_0), static_cast(info.v_strides_1), + static_cast(info.v_strides_2), static_cast(info.v_strides_3)); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + void *stream) const { + (void)workspace; + (void)workspace_size; + auto kunlun_stream = reinterpret_cast(stream); + +#define LAUNCH_WITH_IDX(TMEM) \ + if (_info.past_len_dtype == INFINI_DTYPE_I32) { \ + return launchKunlunKernel(kunlun_stream, k_cache, v_cache, k, v, past_kv_lengths, _info); \ + } \ + return launchKunlunKernel(kunlun_stream, k_cache, v_cache, k, v, past_kv_lengths, _info) + + switch (_info.dtype) { + case INFINI_DTYPE_F32: + LAUNCH_WITH_IDX(uint32_t); + case INFINI_DTYPE_F16: + LAUNCH_WITH_IDX(uint16_t); + case INFINI_DTYPE_BF16: + LAUNCH_WITH_IDX(uint16_t); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + +#undef LAUNCH_WITH_IDX +} + +} // namespace op::kv_caching::kunlun diff --git a/src/infiniop/ops/kv_caching/moore/kv_caching_moore.h b/src/infiniop/ops/kv_caching/moore/kv_caching_moore.h new file mode 100644 index 000000000..caba162a4 --- /dev/null +++ b/src/infiniop/ops/kv_caching/moore/kv_caching_moore.h @@ -0,0 +1,9 @@ + +#ifndef __KV_CACHING_MOORE_API_H__ +#define __KV_CACHING_MOORE_API_H__ + +#include "../kv_caching.h" + +DESCRIPTOR(moore) + +#endif // __KV_CACHING_MOORE_API_H__ diff --git a/src/infiniop/ops/kv_caching/moore/kv_caching_moore.mu b/src/infiniop/ops/kv_caching/moore/kv_caching_moore.mu new file mode 100644 index 000000000..0563d8b62 --- /dev/null +++ b/src/infiniop/ops/kv_caching/moore/kv_caching_moore.mu @@ -0,0 +1,160 @@ + +#include "../../../devices/moore/moore_common.h" +#include "kv_caching_moore.h" + +namespace { + +template +__global__ void kvCachingMooreKernel( + Tdata *k_cache, + Tdata *v_cache, + const Tdata *k, + const Tdata *v, + const Tidx *past_kv_lengths, + size_t batch_size, + size_t num_kv_heads, + size_t seq_len, + size_t hidden_dim, + ptrdiff_t k_cache_strides_0, + ptrdiff_t k_cache_strides_1, + ptrdiff_t k_cache_strides_2, + ptrdiff_t k_cache_strides_3, + ptrdiff_t v_cache_strides_0, + ptrdiff_t v_cache_strides_1, + ptrdiff_t v_cache_strides_2, + ptrdiff_t v_cache_strides_3, + ptrdiff_t k_strides_0, + ptrdiff_t k_strides_1, + ptrdiff_t k_strides_2, + ptrdiff_t k_strides_3, + ptrdiff_t v_strides_0, + ptrdiff_t v_strides_1, + ptrdiff_t v_strides_2, + ptrdiff_t v_strides_3) { + + size_t tid = blockIdx.x * blockDim.x + threadIdx.x; + size_t total = batch_size * num_kv_heads * seq_len * hidden_dim; + size_t grid_size = blockDim.x * gridDim.x; + + for (size_t linear = tid; linear < total; linear += grid_size) { + size_t idx = linear; + size_t d = idx % hidden_dim; + idx /= hidden_dim; + size_t s = idx % seq_len; + idx /= seq_len; + size_t h = idx % num_kv_heads; + size_t b = idx / num_kv_heads; + + ptrdiff_t cache_s = static_cast(past_kv_lengths[b]) + static_cast(s); + ptrdiff_t k_cache_offset = static_cast(b) * k_cache_strides_0 + + static_cast(h) * k_cache_strides_1 + + cache_s * k_cache_strides_2 + + static_cast(d) * k_cache_strides_3; + ptrdiff_t v_cache_offset = static_cast(b) * v_cache_strides_0 + + static_cast(h) * v_cache_strides_1 + + cache_s * v_cache_strides_2 + + static_cast(d) * v_cache_strides_3; + ptrdiff_t k_src_offset = static_cast(b) * k_strides_0 + + static_cast(h) * k_strides_1 + + static_cast(s) * k_strides_2 + + static_cast(d) * k_strides_3; + ptrdiff_t v_src_offset = static_cast(b) * v_strides_0 + + static_cast(h) * v_strides_1 + + static_cast(s) * v_strides_2 + + static_cast(d) * v_strides_3; + k_cache[k_cache_offset] = k[k_src_offset]; + v_cache[v_cache_offset] = v[v_src_offset]; + } +} + +template +infiniStatus_t launchMooreKernel( + musaStream_t stream, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + const op::kv_caching::KVCachingInfo &info) { + constexpr unsigned int block_size = 256; + size_t total = info.batch_size * info.num_kv_heads * info.seq_len * info.hidden_dim; + if (total == 0) { + return INFINI_STATUS_SUCCESS; + } + size_t num_blocks = (total + block_size - 1) / block_size; + kvCachingMooreKernel<<>>( + static_cast(k_cache), static_cast(v_cache), + static_cast(k), static_cast(v), + static_cast(past_kv_lengths), + info.batch_size, info.num_kv_heads, info.seq_len, info.hidden_dim, + info.k_cache_strides_0, info.k_cache_strides_1, info.k_cache_strides_2, info.k_cache_strides_3, + info.v_cache_strides_0, info.v_cache_strides_1, info.v_cache_strides_2, info.v_cache_strides_3, + info.k_strides_0, info.k_strides_1, info.k_strides_2, info.k_strides_3, + info.v_strides_0, info.v_strides_1, info.v_strides_2, info.v_strides_3); + return INFINI_STATUS_SUCCESS; +} + +} // namespace + +namespace op::kv_caching::moore { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t k_cache, + infiniopTensorDescriptor_t v_cache, + infiniopTensorDescriptor_t k, + infiniopTensorDescriptor_t v, + infiniopTensorDescriptor_t past_kv_lengths) { + auto handle = reinterpret_cast(handle_); + auto info = KVCachingInfo::createKVCachingInfo(k_cache, v_cache, k, v, past_kv_lengths); + CHECK_RESULT(info); + + *desc_ptr = new Descriptor( + new Opaque{handle->internal()}, + info.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + void *stream) const { + (void)workspace; + (void)workspace_size; + auto musa_stream = reinterpret_cast(stream); + +#define LAUNCH_WITH_IDX(TDATA) \ + if (_info.past_len_dtype == INFINI_DTYPE_I32) { \ + return launchMooreKernel(musa_stream, k_cache, v_cache, k, v, past_kv_lengths, _info); \ + } \ + return launchMooreKernel(musa_stream, k_cache, v_cache, k, v, past_kv_lengths, _info) + + switch (_info.dtype) { + case INFINI_DTYPE_F16: + LAUNCH_WITH_IDX(half); + case INFINI_DTYPE_BF16: + LAUNCH_WITH_IDX(uint16_t); + case INFINI_DTYPE_F32: + LAUNCH_WITH_IDX(float); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + +#undef LAUNCH_WITH_IDX +} + +} // namespace op::kv_caching::moore diff --git a/src/infiniop/ops/kv_caching/operator.cc b/src/infiniop/ops/kv_caching/operator.cc index e83bb1240..87d8e585e 100644 --- a/src/infiniop/ops/kv_caching/operator.cc +++ b/src/infiniop/ops/kv_caching/operator.cc @@ -2,12 +2,27 @@ #include "../../handle.h" #include "infiniop/ops/kv_caching.h" +#if defined(ENABLE_CPU_API) +#include "cpu/kv_caching_cpu.h" +#endif #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/kv_caching_nvidia.cuh" #endif #if defined(ENABLE_METAX_API) #include "metax/kv_caching_metax.h" #endif +#if defined(ENABLE_MOORE_API) +#include "moore/kv_caching_moore.h" +#endif +#if defined(ENABLE_CAMBRICON_API) +#include "bang/kv_caching_bang.h" +#endif +#if defined(ENABLE_ASCEND_API) +#include "ascend/kv_caching_ascend.h" +#endif +#if defined(ENABLE_KUNLUN_API) +#include "kunlun/kv_caching_kunlun.h" +#endif __INFINI_C infiniStatus_t infiniopCreateKVCachingDescriptor( infiniopHandle_t handle, @@ -31,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateKVCachingDescriptor( switch (handle->device) { +#ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); +#endif #ifdef ENABLE_NVIDIA_API CREATE(INFINI_DEVICE_NVIDIA, nvidia); #endif @@ -49,6 +67,18 @@ __INFINI_C infiniStatus_t infiniopCreateKVCachingDescriptor( #if defined(ENABLE_METAX_API) CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MOORE_API + CREATE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_CAMBRICON_API + CREATE(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); +#endif +#ifdef ENABLE_KUNLUN_API + CREATE(INFINI_DEVICE_KUNLUN, kunlun); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -69,6 +99,9 @@ __INFINI_C infiniStatus_t infiniopGetKVCachingWorkspaceSize( switch (desc->device_type) { +#ifdef ENABLE_CPU_API + GET_SIZE(INFINI_DEVICE_CPU, cpu); +#endif #ifdef ENABLE_NVIDIA_API GET_SIZE(INFINI_DEVICE_NVIDIA, nvidia); #endif @@ -87,6 +120,18 @@ __INFINI_C infiniStatus_t infiniopGetKVCachingWorkspaceSize( #if defined(ENABLE_METAX_API) GET_SIZE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MOORE_API + GET_SIZE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_CAMBRICON_API + GET_SIZE(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_ASCEND_API + GET_SIZE(INFINI_DEVICE_ASCEND, ascend); +#endif +#ifdef ENABLE_KUNLUN_API + GET_SIZE(INFINI_DEVICE_KUNLUN, kunlun); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -113,6 +158,9 @@ __INFINI_C infiniStatus_t infiniopKVCaching( switch (desc->device_type) { +#ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); +#endif #ifdef ENABLE_NVIDIA_API CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); #endif @@ -131,6 +179,18 @@ __INFINI_C infiniStatus_t infiniopKVCaching( #if defined(ENABLE_METAX_API) CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MOORE_API + CALCULATE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_CAMBRICON_API + CALCULATE(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); +#endif +#ifdef ENABLE_KUNLUN_API + CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -149,6 +209,9 @@ __INFINI_C infiniStatus_t infiniopDestroyKVCachingDescriptor( switch (desc->device_type) { +#ifdef ENABLE_CPU_API + DELETE(INFINI_DEVICE_CPU, cpu); +#endif #ifdef ENABLE_NVIDIA_API DELETE(INFINI_DEVICE_NVIDIA, nvidia); #endif @@ -167,6 +230,18 @@ __INFINI_C infiniStatus_t infiniopDestroyKVCachingDescriptor( #if defined(ENABLE_METAX_API) DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MOORE_API + DELETE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_CAMBRICON_API + DELETE(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); +#endif +#ifdef ENABLE_KUNLUN_API + DELETE(INFINI_DEVICE_KUNLUN, kunlun); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;