mindspore 2.4.1__cp311-cp311-manylinux1_x86_64.whl → 2.4.10__cp311-cp311-manylinux1_x86_64.whl

mindspore/lib/plugin/ascend/custom_ascendc_910b/op_impl/ai_core/tbe/custom_ascendc_910b_impl/dynamic/decoder_kv_cache.py

@@ -0,0 +1,215 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+"""
+Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
+"""
+
+import os, sys
+import ctypes
+import json
+import shutil
+from tbe.common.platform import get_soc_spec
+from tbe.common.utils import para_check
+from tbe.tikcpp import compile_op, replay_op, check_op_cap, generalize_op_params, get_code_channel, OpInfo
+from tbe.tikcpp.compile_op import CommonUtility, AscendCLogLevel
+from tbe.common.buildcfg import get_default_build_config
+from impl.util.platform_adapter import tbe_register
+from tbe.common.buildcfg import get_current_build_config
+PYF_PATH = os.path.dirname(os.path.realpath(__file__))
+
+DTYPE_MAP = {"float32": ["DT_FLOAT", "float"],
+    "float16": ["DT_FLOAT16", "half"],
+    "int8": ["DT_INT8", "int8_t"],
+    "int16": ["DT_INT16", "int16_t"],
+    "int32": ["DT_INT32", "int32_t"],
+    "int64": ["DT_INT64", "int64_t"],
+    "uint1": ["DT_UINT1", "uint8_t"],
+    "uint8": ["DT_UINT8", "uint8_t"],
+    "uint16": ["DT_UINT16", "uint16_t"],
+    "uint32": ["DT_UINT32", "uint32_t"],
+    "uint64": ["DT_UINT64", "uint64_t"],
+    "bool": ["DT_BOOL", "bool"],
+    "double": ["DT_DOUBLE", "double"],
+    "dual": ["DT_DUAL", "unknown"],
+    "dual_sub_int8": ["DT_DUAL_SUB_INT8", "unknown"],
+    "dual_sub_uint8": ["DT_DUAL_SUB_UINT8", "unknown"],
+    "string": ["DT_STRING", "unknown"],
+    "complex32": ["DT_COMPLEX32", "unknown"],
+    "complex64": ["DT_COMPLEX64", "unknown"],
+    "complex128": ["DT_COMPLEX128", "unknown"],
+    "qint8": ["DT_QINT8", "unknown"],
+    "qint16": ["DT_QINT16", "unknown"],
+    "qint32": ["DT_QINT32", "unknown"],
+    "quint8": ["DT_QUINT8", "unknown"],
+    "quint16": ["DT_QUINT16", "unknown"],
+    "resource": ["DT_RESOURCE", "unknown"],
+    "string_ref": ["DT_STRING_REF", "unknown"],
+    "int4": ["DT_INT4", "int4b_t"],
+    "bfloat16": ["DT_BF16", "bfloat16_t"]}
+
+def add_dtype_fmt_option_single(x, x_n, is_ref: bool = False):
+    options = []
+    x_fmt = x.get("format")
+    x_dtype = x.get("dtype")
+    x_n_in_kernel = x_n + '_REF' if is_ref else x_n
+    options.append("-DDTYPE_{n}={t}".format(n=x_n_in_kernel, t=DTYPE_MAP.get(x_dtype)[1]))
+    options.append("-DORIG_DTYPE_{n}={ot}".format(n=x_n_in_kernel, ot=DTYPE_MAP.get(x_dtype)[0]))
+    options.append("-DFORMAT_{n}=FORMAT_{f}".format(n=x_n_in_kernel, f=x_fmt))
+    return options
+
+def get_dtype_fmt_options(__inputs__, __outputs__):
+    options = []
+    input_names = ['cache', 'update', 'valid_seq_len', 'batch_index', 'seq_len_axis', 'new_max_seq_len', 'cur_max_seq_len']
+    output_names = ['out']
+    unique_param_name_set = set()
+    for idx, x in enumerate(__inputs__):
+        if x is None:
+            continue
+        x_n = input_names[idx].upper()
+        unique_param_name_set.add(x_n)
+        options += add_dtype_fmt_option_single(x, x_n)
+
+    for idx, x in enumerate(__outputs__):
+        if x is None:
+            continue
+        x_n = output_names[idx].upper()
+        if x_n in unique_param_name_set:
+            options += add_dtype_fmt_option_single(x, x_n, True)
+        else:
+            options += add_dtype_fmt_option_single(x, x_n)
+    return options
+
+def load_dso(so_path):
+    try:
+        ctypes.CDLL(so_path)
+    except OSError as error :
+        CommonUtility.print_compile_log("", error, AscendCLogLevel.LOG_ERROR)
+        raise RuntimeError("cannot open %s" %(so_path))
+    else:
+        msg = "load so succ " + so_path
+        CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO)
+
+def get_shortsoc_compile_option(compile_option_list: list, shortsoc:str):
+    compile_options = []
+    if shortsoc in compile_option_list:
+        compile_options.extend(compile_option_list[shortsoc])
+    if '__ALLSOC__' in compile_option_list:
+        compile_options.extend(compile_option_list['__ALLSOC__'])
+    return compile_options
+
+def get_kernel_source(src_file, dir_snake, dir_ex):
+    src_ex = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, src_file)
+    if os.path.exists(src_ex):
+        return src_ex
+    src = os.environ.get('BUILD_KERNEL_SRC')
+    if src and os.path.exists(src):
+        return src
+    src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, src_file)
+    if os.path.exists(src):
+        return src
+    src = os.path.join(PYF_PATH, src_file)
+    if os.path.exists(src):
+        return src
+    src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, dir_snake + ".cpp")
+    if os.path.exists(src):
+        return src
+    src = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, dir_ex + ".cpp")
+    if os.path.exists(src):
+        return src
+    src = os.path.join(PYF_PATH, "..", "ascendc", os.path.splitext(src_file)[0], src_file)
+    if os.path.exists(src):
+        return src
+    return src_ex
+
+def _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_):
+    __inputs__ = []
+    for arg in [cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__]:
+        if arg != None:
+            if isinstance(arg, (list, tuple)):
+                if len(arg) == 0:
+                    continue
+                __inputs__.append(arg[0])
+            else:
+                __inputs__.append(arg)
+        else:
+            __inputs__.append(arg)
+    __outputs__ = []
+    for arg in [out_out_]:
+        if arg != None:
+            if isinstance(arg, (list, tuple)):
+                if len(arg) == 0:
+                    continue
+                __outputs__.append(arg[0])
+            else:
+                __outputs__.append(arg)
+        else:
+            __outputs__.append(arg)
+    __attrs__ = []
+    return __inputs__, __outputs__, __attrs__
+
+@tbe_register.register_operator("DecoderKvCache", trans_bool_to_s8=False)
+@para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
+def decoder_kv_cache(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, kernel_name="decoder_kv_cache", impl_mode=""):
+    # do ascendc build step
+    if get_current_build_config("enable_op_prebuild"):
+        return
+    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+    options = get_dtype_fmt_options(__inputs__, __outputs__)
+    options += ["-x", "cce"]
+    bisheng = os.environ.get('BISHENG_REAL_PATH')
+    if bisheng is None:
+        bisheng = shutil.which("bisheng")
+    if bisheng != None:
+        bisheng_path = os.path.dirname(bisheng)
+        tikcpp_path = os.path.realpath(os.path.join(bisheng_path, "..", "..", "tikcpp"))
+    else:
+        tikcpp_path = os.path.realpath("/usr/local/Ascend/latest/compiler/tikcpp")
+    options.append("-I" + tikcpp_path)
+    options.append("-I" + os.path.join(tikcpp_path, "..", "..", "include"))
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw"))
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "impl"))
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "interface"))
+    options.append("-I" + os.path.join(PYF_PATH, "..", "ascendc", "common"))
+    if impl_mode == "high_performance":
+        options.append("-DHIGH_PERFORMANCE=1")
+    elif impl_mode == "high_precision":
+        options.append("-DHIGH_PRECISION=1")
+    if get_current_build_config("enable_deterministic_mode") == 1:
+        options.append("-DDETERMINISTIC_MODE=1")
+    else:
+        options.append("-DDETERMINISTIC_MODE=0")
+
+    custom_compile_options = {},
+    custom_all_compile_options = {},
+    soc_version = get_soc_spec("SOC_VERSION")
+    soc_short = get_soc_spec("SHORT_SOC_VERSION").lower()
+    custom_compile_options_soc = get_shortsoc_compile_option(custom_compile_options[0], soc_short)
+    custom_all_compile_options_soc = get_shortsoc_compile_option(custom_all_compile_options[0], soc_short)
+    options += custom_all_compile_options_soc
+    options += custom_compile_options_soc
+
+    origin_func_name = "decoder_kv_cache"
+    ascendc_src_dir_ex = "decoder_kv_cache"
+    ascendc_src_dir = "decoder_kv_cache"
+    ascendc_src_file = "decoder_kv_cache.cpp"
+    src = get_kernel_source(ascendc_src_file, ascendc_src_dir, ascendc_src_dir_ex)
+
+    msg = "start compile Acend C Operator DecoderKvCache, kernel name is " + kernel_name
+    CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO)
+    op_type = "DecoderKvCache"
+    code_channel = get_code_channel(src, kernel_name, op_type, options)
+    op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\
+        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__], origin_outputs = [out_out_],\
+                param_type_dynamic = False, mc2_ctx = [], param_type_list = ['required', 'required', 'required', 'required', 'required', 'required', 'required', 'required'], init_value_list = [None],\
+                output_shape_depend_on_compute = [])
+    compile_op(src, origin_func_name, op_info, options, code_channel, '{}')
+
+def op_select_format(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+    result = check_op_cap("op_select_format", "DecoderKvCache", __inputs__, __outputs__, __attrs__)
+    return result.decode("utf-8")
+
+def get_op_specific_info(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+    result = check_op_cap("get_op_specific_info", "DecoderKvCache", __inputs__, __outputs__, __attrs__)
+    return result.decode("utf-8")

mindspore/lib/plugin/ascend/custom_ascendc_910b/op_impl/ai_core/tbe/custom_ascendc_910b_impl/dynamic/prompt_kv_cache.cpp

@@ -0,0 +1,274 @@
+/**
+ * Copyright 2023 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "kernel_operator.h"
+#include "kernel_utils.h"
+using namespace AscendC;
+
+namespace {
+constexpr int64_t kBufferNum = 1;
+const int64_t kDivisor = 4;
+static __aicore__ inline int64_t CeilRound(int64_t value, int64_t divisor) {
+  if (divisor == 0) {
+    return 0;
+  }
+  return (value + divisor - 1) / divisor * divisor;
+}
+}  // namespace
+
+template <typename T>
+class KernelPromptKvCache {
+ public:
+  __aicore__ inline KernelPromptKvCache() {}
+
+  __aicore__ inline void GetIndex(GM_ADDR batch_index, GM_ADDR valid_seq_len) {
+    int64_t batch_index_ub_num = CeilRound(ub_, kDivisor);
+    int64_t valid_seq_len_ub_num = CeilRound(ub_, kDivisor);
+    batch_index_gm_.SetGlobalBuffer((__gm__ int64_t *)batch_index, batch_index_ub_num);
+
+    int64_t total_num = batch_index_ub_num + valid_seq_len_ub_num;
+    pipe_.InitBuffer(index_queue_, 1, total_num * sizeof(int64_t));
+    batch_index_tensor_ = index_queue_.AllocTensor<int64_t>();
+    DataCopy(batch_index_tensor_, batch_index_gm_, batch_index_ub_num);
+
+    valid_seq_len_gm_.SetGlobalBuffer((__gm__ int64_t *)valid_seq_len, valid_seq_len_ub_num);
+    valid_seq_len_tensor_ = batch_index_tensor_[batch_index_ub_num];
+    DataCopy(valid_seq_len_tensor_, valid_seq_len_gm_, valid_seq_len_ub_num);
+  }
+
+  __aicore__ inline void UpdateCache(GM_ADDR cache, GM_ADDR update) {
+    pipe_.InitBuffer(update_queue_, kBufferNum, former_block_us_ * d_ * sizeof(T));
+    for (int64_t i = 0; i < each_core_bs_num_; ++i) {
+      int64_t bh_idx = core_idx_ * former_each_core_bs_num_ + i;
+      int64_t ub_idx = bh_idx / h_;
+      int64_t h_idx = bh_idx % h_;
+      pipe_barrier((pipe_t)PIPE_ALL);
+      int64_t cache_b_idx = batch_index_tensor_.GetValue(ub_idx);
+      int64_t s_idx = valid_seq_len_tensor_.GetValue(ub_idx);
+      if (cache_b_idx < 0 || cache_b_idx >= b_) {
+        continue;
+      }
+      if (s_idx < 0 || s_idx + us_ > s_) {
+        continue;
+      }
+
+      for (int64_t j = 0; j < split_us_; ++j) {
+        int64_t u_block_len;
+        if (j == split_us_ - 1) {
+          u_block_len = tail_block_us_ * d_;
+        } else {
+          u_block_len = former_block_us_ * d_;
+        }
+        LocalTensor<T> update_in_local_tensor = update_queue_.AllocTensor<T>();
+        update_gm_.SetGlobalBuffer(
+          (__gm__ T *)update + ub_idx * update_b_stride_ + h_idx * update_h_stride_ + j * former_block_us_ * d_,
+          u_block_len);
+        out_gm_.SetGlobalBuffer((__gm__ T *)cache + cache_b_idx * cache_b_stride_ + h_idx * cache_h_stride_ +
+                                  s_idx * d_ + j * former_block_us_ * d_,
+                                u_block_len);
+        pipe_barrier((pipe_t)PIPE_ALL);
+        DataCopy(update_in_local_tensor, update_gm_, u_block_len);
+        update_queue_.EnQue(update_in_local_tensor);
+        LocalTensor<T> update_in_local_tensor_out = update_queue_.DeQue<T>();
+        pipe_barrier((pipe_t)PIPE_ALL);
+        DataCopy(out_gm_, update_in_local_tensor_out, u_block_len);
+        update_queue_.FreeTensor(update_in_local_tensor_out);
+      }
+    }
+  }
+
+  __aicore__ inline void InitParam(GM_ADDR tiling) {
+    GET_TILING_DATA(tiling_data, tiling);
+    core_num_ = tiling_data.core_num;
+    b_ = tiling_data.b;
+    h_ = tiling_data.h;
+    s_ = tiling_data.s;
+    d_ = tiling_data.d;
+    ub_ = tiling_data.ub;
+    us_ = tiling_data.us;
+    former_each_core_bs_num_ = tiling_data.former_each_core_bs_num;
+    tail_each_core_bs_num_ = tiling_data.tail_each_core_bs_num;
+    split_us_ = tiling_data.split_us;
+    former_block_us_ = tiling_data.former_block_us;
+    tail_block_us_ = tiling_data.tail_block_us;
+  }
+
+  __aicore__ inline void Process(GM_ADDR cache, GM_ADDR update, GM_ADDR valid_seq_len, GM_ADDR batch_index,
+                                 GM_ADDR seq_len_axis, GM_ADDR new_max_seq_len, GM_ADDR cur_max_seq_len,
+                                 GM_ADDR tiling_data) {
+    core_idx_ = GetBlockIdx();
+    InitParam(tiling_data);
+    if (core_idx_ >= core_num_) {
+      return;
+    }
+    if (g_coreType == AIC) {
+      return;
+    }
+    if (core_idx_ != core_num_ - 1) {
+      each_core_bs_num_ = former_each_core_bs_num_;
+    } else {
+      each_core_bs_num_ = tail_each_core_bs_num_;
+    }
+
+    GetIndex(batch_index, valid_seq_len);
+
+    cache_h_stride_ = s_ * d_;
+    cache_b_stride_ = h_ * cache_h_stride_;
+
+    update_h_stride_ = us_ * d_;
+    update_b_stride_ = h_ * update_h_stride_;
+
+    UpdateCache(cache, update);
+    index_queue_.FreeTensor(batch_index_tensor_);
+  }
+
+ private:
+  // gm
+  GlobalTensor<T> update_gm_;
+  GlobalTensor<int64_t> valid_seq_len_gm_;
+  GlobalTensor<int64_t> batch_index_gm_;
+  GlobalTensor<int64_t> new_max_seq_len_gm_;
+  GlobalTensor<T> out_gm_;
+
+  // local gm
+  LocalTensor<int64_t> valid_seq_len_tensor_;
+  LocalTensor<int64_t> batch_index_tensor_;
+
+  TPipe pipe_;
+  TQue<QuePosition::VECIN, 1> update_queue_;
+  TQue<QuePosition::VECIN, 1> index_queue_;
+  TQue<QuePosition::VECIN, 1> new_max_seq_len_queue_;
+
+  int64_t core_idx_ = 0;
+  int64_t core_num_ = 0;
+  int64_t each_core_bs_num_ = 0;
+  int64_t former_each_core_bs_num_ = 0;
+  int64_t tail_each_core_bs_num_ = 0;
+  int64_t b_ = 0;
+  int64_t h_ = 0;
+  int64_t s_ = 0;
+  int64_t d_ = 0;
+  int64_t ub_ = 0;
+  int64_t us_ = 0;
+  int64_t split_us_ = 0;
+  int64_t former_block_us_ = 0;
+  int64_t tail_block_us_ = 0;
+
+  int64_t cache_b_stride_ = 0;
+  int64_t cache_h_stride_ = 0;
+  int64_t update_b_stride_ = 0;
+  int64_t update_h_stride_ = 0;
+};
+
+template <typename T>
+class KernelPromptKvCacheCopyAll {
+ public:
+  __aicore__ inline KernelPromptKvCacheCopyAll() {}
+
+  __aicore__ inline void InitParam(GM_ADDR tiling) {
+    GET_TILING_DATA(tiling_data, tiling);
+    core_num_ = tiling_data.core_num;
+    former_each_core_bs_num_ = tiling_data.former_each_core_bs_num;
+    tail_each_core_bs_num_ = tiling_data.tail_each_core_bs_num;
+    split_us_ = tiling_data.split_us;
+    former_block_us_ = tiling_data.former_block_us;
+    tail_block_us_ = tiling_data.tail_block_us;
+  }
+
+  __aicore__ inline void Process(GM_ADDR cache, GM_ADDR update, GM_ADDR valid_seq_len, GM_ADDR batch_index,
+                                 GM_ADDR seq_len_axis, GM_ADDR new_max_seq_len, GM_ADDR cur_max_seq_len,
+                                 GM_ADDR tiling_data) {
+    core_idx_ = GetBlockIdx();
+    InitParam(tiling_data);
+    if (core_idx_ >= core_num_) {
+      return;
+    }
+    if (g_coreType == AIC) {
+      return;
+    }
+    if (core_idx_ != core_num_ - 1) {
+      each_core_bs_num_ = former_each_core_bs_num_;
+    } else {
+      each_core_bs_num_ = tail_each_core_bs_num_;
+    }
+    pipe_barrier((pipe_t)PIPE_ALL);
+
+    pipe_.InitBuffer(update_queue_, kBufferNum, former_block_us_ * sizeof(T));
+    for (int64_t i = 0; i < split_us_; ++i) {
+      int64_t u_block_len;
+      if (i == split_us_ - 1) {
+        u_block_len = tail_block_us_;
+      } else {
+        u_block_len = former_block_us_;
+      }
+      LocalTensor<T> update_in_local_tensor = update_queue_.AllocTensor<T>();
+      update_gm_.SetGlobalBuffer((__gm__ T *)update + core_idx_ * former_each_core_bs_num_ + i * former_block_us_,
+                                 u_block_len);
+      out_gm_.SetGlobalBuffer((__gm__ T *)cache + core_idx_ * former_each_core_bs_num_ + i * former_block_us_,
+                              u_block_len);
+      pipe_barrier((pipe_t)PIPE_ALL);
+      DataCopy(update_in_local_tensor, update_gm_, u_block_len);
+      update_queue_.EnQue(update_in_local_tensor);
+      LocalTensor<T> update_in_local_tensor_out = update_queue_.DeQue<T>();
+      pipe_barrier((pipe_t)PIPE_ALL);
+      DataCopy(out_gm_, update_in_local_tensor_out, u_block_len);
+      update_queue_.FreeTensor(update_in_local_tensor_out);
+    }
+  }
+
+ private:
+  // gm
+  GlobalTensor<T> update_gm_;
+  GlobalTensor<T> out_gm_;
+
+  TPipe pipe_;
+  TQue<QuePosition::VECIN, 1> update_queue_;
+
+  int64_t core_idx_ = 0;
+  int64_t core_num_ = 0;
+  int64_t each_core_bs_num_ = 0;
+  int64_t former_each_core_bs_num_ = 0;
+  int64_t tail_each_core_bs_num_ = 0;
+
+  int64_t split_us_ = 0;
+  int64_t former_block_us_ = 0;
+  int64_t tail_block_us_ = 0;
+};
+
+extern "C" __global__ __aicore__ void prompt_kv_cache(GM_ADDR cache, GM_ADDR update, GM_ADDR valid_seq_len,
+                                                      GM_ADDR batch_index, GM_ADDR seq_len_axis,
+                                                      GM_ADDR new_max_seq_len, GM_ADDR cur_max_seq_len, GM_ADDR out,
+                                                      GM_ADDR workspace, GM_ADDR tiling) {
+  if (TILING_KEY_IS(1)) {
+    KernelPromptKvCache<int8_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, tiling);
+  } else if (TILING_KEY_IS(2)) {
+    KernelPromptKvCache<int16_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, tiling);
+  } else if (TILING_KEY_IS(4)) {
+    KernelPromptKvCache<int32_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, tiling);
+  } else if (TILING_KEY_IS(10)) {
+    KernelPromptKvCacheCopyAll<int8_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, tiling);
+  } else if (TILING_KEY_IS(20)) {
+    KernelPromptKvCacheCopyAll<int16_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, tiling);
+  } else if (TILING_KEY_IS(40)) {
+    KernelPromptKvCacheCopyAll<int32_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, tiling);
+  }
+}