add ops

torch_mlu_ops-v1.3.2/csrc/kernels/copy_blocks.mluh

@@ -0,0 +1,37 @@
+/*************************************************************************
+ * Copyright (C) [2023-2024] by Cambricon, Inc.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+ * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+ * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ *************************************************************************/
+#ifndef CSRC_KERNELS_COPY_BLOCKS_MLUH_
+#define CSRC_KERNELS_COPY_BLOCKS_MLUH_
+
+#include <vector>
+#include "cnnl.h"
+#include "kernel_utils.h"
+
+namespace tmo {
+/**
+ * @brief Perform copy_blocks operation.
+ * @param queue: The queue for mlu.
+ * @param key_caches: Output/Input. Pointer to the MLU memory that stores the key_caches
+ * vector<tensor> which has shape [num_layers<num_blocks, num_heads, block_size, head_size>].
+ * @param value_caches: Output/Input. Pointer to the MLU memory that stores the value_caches
+ * vector<tensor> which has shape [num_layers<num_blocks, num_heads, block_size, head_size>].
+ * @param block_mapping_vec: block_mapping vector.
+ * @param block_size_in_bytes: one block data size.
+ */
+KernelStatus invokeCopyBlocksKernel(const cnrtQueue_t queue,
+                                    const std::vector<void *> &key_caches,
+                                    const std::vector<void *> &value_caches,
+                                    const std::vector<int32_t> &block_mapping_vec,
+                                    const size_t block_size_in_bytes);
+}  // namespace tmo
+
+#endif  // CSRC_KERNELS_COPY_BLOCKS_MLUH_