InfiniTensor · baominghelly · Aug 12, 2026
diff --git a/src/native/cambricon/cnnl_utils.h b/src/native/cambricon/cnnl_utils.h
@@ -0,0 +1,97 @@
+#ifndef INFINI_OPS_CAMBRICON_CNNL_UTILS_H_
+#define INFINI_OPS_CAMBRICON_CNNL_UTILS_H_
+
+#include <cassert>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "native/cambricon/common.h"
+#include "tensor.h"
+
+#define INFINI_OPS_CNNL_CHECK(call)                                      \
+  do {                                                                   \
+    const auto cnnl_status = (call);                                     \
+    assert(cnnl_status == CNNL_STATUS_SUCCESS && "`" #call "` failed."); \
+    (void)cnnl_status;                                                   \
+  } while (false)
+
+namespace infini::ops::cnnl_utils {
+
+struct HandleDeleter {
+  using pointer = cnnlHandle_t;
+
+  void operator()(pointer handle) const noexcept {
+    if (handle) {
+      (void)cnnlDestroy(handle);
+    }
+  }
+};
+
+using Handle =
+    std::unique_ptr<std::remove_pointer_t<cnnlHandle_t>, HandleDeleter>;
+
+inline Handle CreateHandle() {
+  cnnlHandle_t handle{nullptr};
+  INFINI_OPS_CNNL_CHECK(cnnlCreate(&handle));
+
+  return Handle{handle};
+}
+
+struct TensorDescriptorDeleter {
+  using pointer = cnnlTensorDescriptor_t;
+
+  void operator()(pointer desc) const noexcept {
+    if (desc) {
+      (void)cnnlDestroyTensorDescriptor(desc);
+    }
+  }
+};
+
+using TensorDescriptor =
+    std::unique_ptr<std::remove_pointer_t<cnnlTensorDescriptor_t>,
+                    TensorDescriptorDeleter>;
+
+inline TensorDescriptor CreateTensorDescriptor() {
+  cnnlTensorDescriptor_t desc{nullptr};
+  INFINI_OPS_CNNL_CHECK(cnnlCreateTensorDescriptor(&desc));
+
+  return TensorDescriptor{desc};
+}
+
+inline void SetTensorDescriptor(cnnlTensorDescriptor_t desc, DataType dtype,
+                                const Tensor::Shape& shape,
+                                const Tensor::Strides& strides) {
+  assert(shape.size() == strides.size() &&
+         "`CNNL tensor descriptor` requires matching shape and strides.");
+
+  const auto cnnl_dtype = GetDataType(dtype);
+  assert(cnnl_dtype != CNNL_DTYPE_INVALID &&
+         "`CNNL tensor descriptor` does not support this data type.");
+
+  const auto cnnl_shape =
+      shape.empty() ? std::vector<std::int64_t>{1}
+                    : std::vector<std::int64_t>(shape.begin(), shape.end());
+  const auto cnnl_strides =
+      strides.empty()
+          ? std::vector<std::int64_t>{1}
+          : std::vector<std::int64_t>(strides.begin(), strides.end());
+  const auto ndim = static_cast<int>(cnnl_shape.size());
+
+  INFINI_OPS_CNNL_CHECK(
+      cnnlSetTensorDescriptorEx_v2(desc, CNNL_LAYOUT_ARRAY, cnnl_dtype, ndim,
+                                   cnnl_shape.data(), cnnl_strides.data()));
+}
+
+inline TensorDescriptor MakeTensorDescriptor(DataType dtype,
+                                             const Tensor::Shape& shape,
+                                             const Tensor::Strides& strides) {
+  auto desc = CreateTensorDescriptor();
+  SetTensorDescriptor(desc.get(), dtype, shape, strides);
+
+  return desc;
+}
+
+}  // namespace infini::ops::cnnl_utils
+
+#endif
diff --git a/src/native/cambricon/cnrt_utils.h b/src/native/cambricon/cnrt_utils.h
@@ -0,0 +1,38 @@
+#ifndef INFINI_OPS_CAMBRICON_CNRT_UTILS_H_
+#define INFINI_OPS_CAMBRICON_CNRT_UTILS_H_
+
+#include <cnrt.h>
+
+#include <cassert>
+#include <cstddef>
+#include <memory>
+
+namespace infini::ops::cnrt_utils {
+
+struct DeviceBufferDeleter {
+  using pointer = void*;
+
+  void operator()(pointer buffer) const noexcept {
+    if (buffer) {
+      (void)cnrtFree(buffer);
+    }
+  }
+};
+
+using DeviceBuffer = std::unique_ptr<void, DeviceBufferDeleter>;
+
+inline DeviceBuffer AllocateDeviceBuffer(std::size_t size) {
+  if (size == 0) {
+    return {};
+  }
+
+  void* buffer{nullptr};
+  [[maybe_unused]] const auto status = cnrtMalloc(&buffer, size);
+  assert(status == cnrtSuccess && "`cnrtMalloc` failed.");
+
+  return DeviceBuffer{buffer};
+}
+
+}  // namespace infini::ops::cnrt_utils
+
+#endif
diff --git a/src/native/cambricon/ops/copy/cnnl.h b/src/native/cambricon/ops/copy/cnnl.h
@@ -0,0 +1,80 @@
+#ifndef INFINI_OPS_CAMBRICON_COPY_CNNL_H_
+#define INFINI_OPS_CAMBRICON_COPY_CNNL_H_
+
+#include <algorithm>
+#include <cassert>
+
+#include "base/copy.h"
+#include "native/cambricon/cnnl_utils.h"
+#include "native/cambricon/cnrt_utils.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Copy, Device::Type::kCambricon> : public Copy {
+ public:
+  Operator(const Tensor src, const bool non_blocking, Tensor out)
+      : Copy{src, non_blocking, out} {
+    if (output_size_ == 0) {
+      return;
+    }
+
+    assert(std::all_of(input_strides_.begin(), input_strides_.end(),
+                       [](auto stride) { return stride >= 0; }) &&
+           "`CambriconCopy` does not support negative input strides.");
+    assert(std::all_of(out_strides_.begin(), out_strides_.end(),
+                       [](auto stride) { return stride >= 0; }) &&
+           "`CambriconCopy` does not support negative output strides.");
+
+    cnnl_handle_ = cnnl_utils::CreateHandle();
+    input_desc_ = cnnl_utils::MakeTensorDescriptor(input_type_, input_shape_,
+                                                   input_strides_);
+    out_desc_ =
+        cnnl_utils::MakeTensorDescriptor(out_type_, out_shape_, out_strides_);
+
+    INFINI_OPS_CNNL_CHECK(
+        cnnlGetCopyWorkspaceSize(cnnl_handle_.get(), input_desc_.get(),
+                                 out_desc_.get(), &workspace_size_));
+
+    default_workspace_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_);
+  }
+
+  void operator()(const Tensor src, const bool /*non_blocking*/,
+                  Tensor out) const override {
+    if (output_size_ == 0) {
+      return;
+    }
+
+    INFINI_OPS_CNNL_CHECK(cnnlSetQueue(
+        cnnl_handle_.get(), static_cast<cnrtQueue_t>(stream_ ? stream_ : 0)));
+
+    void* workspace = workspace_ ? workspace_ : default_workspace_.get();
+    const auto workspace_size =
+        workspace_ ? workspace_size_in_bytes_ : workspace_size_;
+    assert(workspace_size >= workspace_size_ &&
+           "`CambriconCopy` requires a sufficiently large workspace.");
+
+    INFINI_OPS_CNNL_CHECK(cnnlCopy_v2(cnnl_handle_.get(), input_desc_.get(),
+                                      src.data(), out_desc_.get(), out.data(),
+                                      workspace, workspace_size_));
+  }
+
+  std::size_t workspace_size_in_bytes() const override {
+    return workspace_size_;
+  }
+
+ private:
+  std::size_t workspace_size_{0};
+
+  cnrt_utils::DeviceBuffer default_workspace_{};
+
+  cnnl_utils::Handle cnnl_handle_{};
+
+  cnnl_utils::TensorDescriptor input_desc_{};
+
+  cnnl_utils::TensorDescriptor out_desc_{};
+};
+
+}  // namespace infini::ops
+
+#endif
diff --git a/tests/test_copy.py b/tests/test_copy.py
@@ -9,6 +9,7 @@
 @pytest.mark.parametrize(
     "input_shape, out_shape, input_strides, out_strides",
     (
+        ((), (), None, None),
         ((100, 100), (100, 100), (1, 100), (100, 1)),
         ((2, 2, 2, 4), (2, 2, 2, 4), (16, 8, 4, 1), (16, 8, 1, 2)),
         ((8, 4, 20, 64), (8, 4, 20, 64), (5120, 64, 256, 1), None),
@@ -59,3 +60,20 @@ def _torch_copy(input, non_blocking, out):
     out.copy_(input, non_blocking=non_blocking)
 
     return out
+
+
+@pytest.mark.parametrize("non_blocking", (False, True))
+@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16))
+def test_copy_preserves_kv_cache_outside_slice(dtype, non_blocking, device):
+    input = torch.randn((2, 3, 3, 4), dtype=dtype, device=device).transpose(1, 2)
+    cache = torch.full((2, 3, 11, 4), 17, dtype=dtype, device=device)
+    out = cache[:, :, 4:7, :]
+
+    before = cache.clone()
+    expected = input.expand_as(out).clone()
+
+    _copy(input, non_blocking, out)
+
+    assert torch.equal(out, expected)
+    assert torch.equal(cache[:, :, :4, :], before[:, :, :4, :])
+    assert torch.equal(cache[:, :, 7:, :], before[:, :, 7:, :])