feat(ascend): basic llama operators

src/CMakeLists.txt

@@ -501,6 +501,9 @@ if(WITH_ASCEND)
         # custom kernel implementations (e.g. `RmsNorm` index 1) can call
         # them via the generated launch functions.
         target_compile_definitions(infiniops PUBLIC INFINI_HAS_CUSTOM_KERNELS=1)
+        target_link_libraries(infiniops PRIVATE
+            -Wl,--whole-archive $<TARGET_FILE:no_workspace_kernel> -Wl,--no-whole-archive)
+        add_dependencies(infiniops no_workspace_kernel_build)
     endif()
 endif()
 
@@ -510,8 +513,9 @@ set(INFINI_OPS_SMOKE_BUILD OFF CACHE BOOL
     "Build only the smoke-test operator subset")
 set(_infini_ops_smoke_ops
     add mul relu cast cat gemm matmul linear rms_norm swiglu silu_and_mul
-    causal_softmax causal_softmax_infinilm abs clamp exp)
-set(_infini_ops_smoke_torch_ops abs clamp exp relu)
+    fused_add_rms_norm rotary_embedding causal_softmax
+    causal_softmax_infinilm abs clamp exp argmax)
+set(_infini_ops_smoke_torch_ops abs clamp exp relu argmax)
 
 if(WITH_NVIDIA)
     list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm moe_sum)
@@ -1394,9 +1398,9 @@ if(GENERATE_PYTHON_BINDINGS)
     # Custom `AscendC` kernel objects must be linked directly into ops
     # because the `AscendC` toolchain compiles host stubs with hidden
     # visibility — `libinfiniops.so` cannot re-export those symbols.
-    # The `Operator<..., 1>` template instantiations that call
-    # `aclrtlaunch_*` live in `ops.cc`, so link here with
-    # `--whole-archive` to ensure all launch functions are available.
+    # The Python extension can compile generated instantiations that call
+    # `aclrtlaunch_*` directly, so link here with `--whole-archive`
+    # to ensure all launch functions are available.
     # `$<TARGET_FILE>` works for both real `ascendc_library()` targets and
     # `IMPORTED` targets pointing at a pre-built `.a`.  The
     # `no_workspace_kernel` target is only created inside the

src/native/ascend/custom/rms_norm/op_kernel/rms_norm.cpp

@@ -200,8 +200,8 @@ class KernelRmsNorm {
 extern "C" __global__ __aicore__ void rms_norm(
     GM_ADDR x, GM_ADDR weight, GM_ADDR y, int64_t totalRows, int64_t dimLength,
     int64_t dimLengthAlign, int64_t formerNum, int64_t formerLength,
-    int64_t tailLength, float eps, int64_t dtypeSize) {
-  if (dtypeSize == 2) {
+    int64_t tailLength, float eps, int64_t isHalf) {
+  if (isHalf != 0) {
     KernelRmsNorm<half> op;
     op.Init(x, weight, y, totalRows, dimLength, dimLengthAlign, formerNum,
             formerLength, tailLength, eps);

src/native/ascend/ops/argmax/kernel.h

@@ -0,0 +1,121 @@
+#ifndef INFINI_OPS_ASCEND_ARGMAX_KERNEL_H_
+#define INFINI_OPS_ASCEND_ARGMAX_KERNEL_H_
+
+#include <algorithm>
+#include <cassert>
+
+#include "acl/acl.h"
+#include "aclnn/aclnn_base.h"
+#include "aclnnop/aclnn_argmax.h"
+#include "aclnnop/aclnn_cast.h"
+#include "base/argmax.h"
+#include "native/ascend/common.h"
+#include "native/ascend/workspace_pool_.h"
+#include "operator.h"
+
+namespace infini::ops {
+
+// Greedy sampling reduction for one contiguous vocabulary vector. CANN
+// ArgMax does not accept BF16, so that dtype is promoted before reduction.
+template <>
+class Operator<Argmax, Device::Type::kAscend> : public Argmax {
+ public:
+  Operator(const Tensor input, const std::optional<int64_t> dim,
+           const bool keepdim, Tensor out)
+      : Argmax(input, dim, keepdim, out),
+        input_cache_(input),
+        out_cache_(out),
+        use_cast_(input.dtype() == DataType::kBFloat16) {
+    assert(input.ndim() == 1 && input.numel() > 0 && input.IsContiguous() &&
+           !dim.has_value() && !keepdim && out.ndim() == 0 &&
+           out.numel() == 1 && out.dtype() == DataType::kInt64 &&
+           (input.dtype() == DataType::kFloat16 ||
+            input.dtype() == DataType::kBFloat16 ||
+            input.dtype() == DataType::kFloat32) &&
+           "Ascend `Argmax` provider 0 supports contiguous 1D float logits, "
+           "no dim, keepdim=false, and a scalar int64 output");
+
+    if (use_cast_) {
+      const auto bytes = input.numel() * sizeof(float);
+      auto ret = aclrtMalloc(&cast_data_, bytes, ACL_MEM_MALLOC_NORMAL_ONLY);
+      assert(ret == ACL_SUCCESS &&
+             "Ascend `Argmax` failed to allocate BF16 promotion buffer");
+      cast_cache_ = ascend::AclTensorCache(
+          {static_cast<int64_t>(input.numel())}, ACL_FLOAT, cast_data_);
+    }
+  }
+
+  ~Operator() {
+    if (!ascend::IsAclRuntimeAlive()) return;
+
+    input_cache_.release();
+    out_cache_.release();
+    if (use_cast_) {
+      cast_cache_.release();
+      aclrtFree(cast_data_);
+    }
+  }
+
+  void operator()(const Tensor input, const std::optional<int64_t> dim,
+                  const bool keepdim, Tensor out) const override {
+    (void)dim;
+    (void)keepdim;
+    auto stream = static_cast<aclrtStream>(stream_);
+    auto t_input = input_cache_.get(const_cast<void*>(input.data()));
+    auto t_out = out_cache_.get(out.data());
+    aclTensor* t_arg_input = t_input;
+
+    if (use_cast_) {
+      auto t_cast = cast_cache_.get(cast_data_);
+      t_arg_input = t_cast;
+      if (!cast_executor_) {
+        auto ret = aclnnCastGetWorkspaceSize(t_input, ACL_FLOAT, t_cast,
+                                             &cast_ws_size_, &cast_executor_);
+        assert(ret == ACL_SUCCESS &&
+               "Ascend `Argmax` BF16 cast workspace query failed");
+        aclSetAclOpExecutorRepeatable(cast_executor_);
+      } else {
+        aclSetInputTensorAddr(cast_executor_, 0, t_input,
+                              const_cast<void*>(input.data()));
+        aclSetOutputTensorAddr(cast_executor_, 0, t_cast, cast_data_);
+      }
+    }
+
+    if (!argmax_executor_) {
+      auto ret = aclnnArgMaxGetWorkspaceSize(
+          t_arg_input, 0, false, t_out, &argmax_ws_size_, &argmax_executor_);
+      assert(ret == ACL_SUCCESS && "Ascend `Argmax` workspace query failed");
+      aclSetAclOpExecutorRepeatable(argmax_executor_);
+    } else {
+      auto arg_input_data = use_cast_ ? cast_data_ : input.data();
+      aclSetInputTensorAddr(argmax_executor_, 0, t_arg_input,
+                            const_cast<void*>(arg_input_data));
+      aclSetOutputTensorAddr(argmax_executor_, 0, t_out, out.data());
+    }
+
+    auto workspace_size = std::max(cast_ws_size_, argmax_ws_size_);
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size);
+    if (use_cast_) {
+      auto ret = aclnnCast(arena.buf, cast_ws_size_, cast_executor_, stream);
+      assert(ret == ACL_SUCCESS && "Ascend `Argmax` BF16 cast failed");
+    }
+    auto ret =
+        aclnnArgMax(arena.buf, argmax_ws_size_, argmax_executor_, stream);
+    assert(ret == ACL_SUCCESS && "Ascend `Argmax` execution failed");
+  }
+
+ private:
+  mutable ascend::AclTensorCache input_cache_;
+  mutable ascend::AclTensorCache out_cache_;
+  mutable ascend::AclTensorCache cast_cache_;
+  bool use_cast_{false};
+  void* cast_data_{nullptr};
+  mutable aclOpExecutor* cast_executor_{nullptr};
+  mutable uint64_t cast_ws_size_{0};
+  mutable aclOpExecutor* argmax_executor_{nullptr};
+  mutable uint64_t argmax_ws_size_{0};
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_ASCEND_ARGMAX_KERNEL_H_

src/native/ascend/ops/causal_softmax/kernel.h

@@ -0,0 +1,120 @@
+#ifndef INFINI_OPS_ASCEND_CAUSAL_SOFTMAX_KERNEL_H_
+#define INFINI_OPS_ASCEND_CAUSAL_SOFTMAX_KERNEL_H_
+
+#include <limits>
+#include <vector>
+
+#include "acl/acl.h"
+#include "aclnn/aclnn_base.h"
+#include "aclnn_copy.h"
+#include "aclnn_masked_fill_scalar.h"
+#include "aclnn_softmax.h"
+#include "base/causal_softmax.h"
+#include "data_type.h"
+#include "native/ascend/common.h"
+#include "native/ascend/workspace_pool_.h"
+#include "operator.h"
+
+namespace infini::ops {
+
+// CANN 8.5 has no single API covering causal-mask-then-softmax. Decompose the
+// operation into a stride-aware copy, masked fill, and last-dimension softmax.
+template <>
+class Operator<CausalSoftmax, Device::Type::kAscend> : public CausalSoftmax {
+ public:
+  Operator(const Tensor input, Tensor out)
+      : CausalSoftmax(input, out), in_cache_(input), out_cache_(out) {
+    temp_size_ = input.numel() * kDataTypeToSize.at(dtype_);
+    Tensor temp_tensor{nullptr, input.shape(), input.dtype(), input.device()};
+    temp_cache_ = ascend::AclTensorCache(temp_tensor);
+
+    // `mask[i][j] = 1` when key position `j` is not visible to query `i`.
+    // Shape `(seq_len, total_seq_len)` broadcasts over leading dimensions.
+    size_t mask_elems = seq_len_ * total_seq_len_;
+    std::vector<uint8_t> mask_host(mask_elems, 0);
+    for (size_t i = 0; i < seq_len_; ++i) {
+      auto vis_end = static_cast<int64_t>(total_seq_len_ - seq_len_ + i);
+      for (auto j = vis_end + 1; j < static_cast<int64_t>(total_seq_len_);
+           ++j) {
+        mask_host[i * total_seq_len_ + j] = 1;
+      }
+    }
+
+    aclrtMalloc(&mask_buf_, mask_elems, ACL_MEM_MALLOC_NORMAL_ONLY);
+    aclrtMemcpy(mask_buf_, mask_elems, mask_host.data(), mask_elems,
+                ACL_MEMCPY_HOST_TO_DEVICE);
+
+    std::vector<int64_t> mshape = {static_cast<int64_t>(seq_len_),
+                                   static_cast<int64_t>(total_seq_len_)};
+    std::vector<int64_t> mstrides = {static_cast<int64_t>(total_seq_len_), 1};
+    mask_tensor_ = aclCreateTensor(mshape.data(), mshape.size(), ACL_BOOL,
+                                   mstrides.data(), 0, ACL_FORMAT_ND,
+                                   mshape.data(), mshape.size(), mask_buf_);
+
+    // `aclCreateScalar` stores the pointer, so the backing value is a member.
+    neg_inf_ = aclCreateScalar(&neg_inf_storage_, ACL_FLOAT);
+  }
+
+  ~Operator() {
+    if (!ascend::IsAclRuntimeAlive()) return;
+
+    if (mask_tensor_) aclDestroyTensor(mask_tensor_);
+    if (mask_buf_) aclrtFree(mask_buf_);
+    if (neg_inf_) aclDestroyScalar(neg_inf_);
+  }
+
+  void operator()(const Tensor input, Tensor out) const override {
+    auto stream = static_cast<aclrtStream>(stream_);
+    auto& temp = ascend::GetWorkspacePool().Ensure(stream, temp_size_, "temp");
+
+    // Descriptors stay owned by the cached operator and outlive all queued
+    // work; only their raw addresses change between invocations.
+    auto t_in = in_cache_.get(const_cast<void*>(input.data()));
+    auto t_out = out_cache_.get(out.data());
+    auto t_temp = temp_cache_.get(temp.buf);
+
+    // CANN consumes these executors even when they are marked repeatable.
+    // Acquire a fresh executor for every transformer layer invocation.
+    aclOpExecutor* copy_exec = nullptr;
+    uint64_t copy_ws = 0;
+    aclnnInplaceCopyGetWorkspaceSize(t_temp, t_in, &copy_ws, &copy_exec);
+    auto& copy_arena = ascend::GetWorkspacePool().Ensure(stream, copy_ws);
+    aclnnInplaceCopy(copy_arena.buf, copy_ws, copy_exec, stream);
+
+    aclOpExecutor* fill_exec = nullptr;
+    uint64_t fill_ws = 0;
+    aclnnInplaceMaskedFillScalarGetWorkspaceSize(t_temp, mask_tensor_, neg_inf_,
+                                                 &fill_ws, &fill_exec);
+    auto& fill_arena = ascend::GetWorkspacePool().Ensure(stream, fill_ws);
+    aclnnInplaceMaskedFillScalar(fill_arena.buf, fill_ws, fill_exec, stream);
+
+    constexpr int64_t kLastDim = -1;
+    aclOpExecutor* softmax_exec = nullptr;
+    uint64_t softmax_ws = 0;
+    aclnnSoftmaxGetWorkspaceSize(t_temp, kLastDim, t_out, &softmax_ws,
+                                 &softmax_exec);
+    auto& softmax_arena = ascend::GetWorkspacePool().Ensure(stream, softmax_ws);
+    aclnnSoftmax(softmax_arena.buf, softmax_ws, softmax_exec, stream);
+  }
+
+ private:
+  mutable ascend::AclTensorCache in_cache_;
+
+  mutable ascend::AclTensorCache out_cache_;
+
+  mutable ascend::AclTensorCache temp_cache_;
+
+  float neg_inf_storage_ = -std::numeric_limits<float>::infinity();
+
+  uint64_t temp_size_ = 0;
+
+  void* mask_buf_ = nullptr;
+
+  aclTensor* mask_tensor_ = nullptr;
+
+  aclScalar* neg_inf_ = nullptr;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/ascend/ops/copy/kernel.h

@@ -0,0 +1,56 @@
+#ifndef INFINI_OPS_ASCEND_COPY_KERNEL_H_
+#define INFINI_OPS_ASCEND_COPY_KERNEL_H_
+
+#include "acl/acl.h"
+#include "aclnn/aclnn_base.h"
+#include "aclnn_copy.h"
+#include "base/copy.h"
+#include "native/ascend/common.h"
+#include "native/ascend/workspace_pool_.h"
+#include "operator.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Copy, Device::Type::kAscend> : public Copy {
+ public:
+  Operator(const Tensor src, const bool non_blocking, Tensor out)
+      : Copy(src, non_blocking, out),
+        in_cache_(BroadcastView(src, out)),
+        out_cache_(out) {}
+
+  void operator()(const Tensor src, const bool /*non_blocking*/,
+                  Tensor out) const override {
+    if (output_size_ == 0) return;
+
+    auto stream = static_cast<aclrtStream>(stream_);
+    auto t_in = in_cache_.get(const_cast<void*>(src.data()));
+    auto t_out = out_cache_.get(out.data());
+
+    // InplaceCopy executors are consumed by CANN even after marking them
+    // repeatable. Reusing the pointer causes a double-free or use-after-free
+    // on a later model layer, so acquire an executor for every invocation.
+    aclOpExecutor* executor = nullptr;
+    uint64_t workspace_size = 0;
+    aclnnInplaceCopyGetWorkspaceSize(t_out, t_in, &workspace_size, &executor);
+
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size);
+    aclnnInplaceCopy(arena.buf, workspace_size, executor, stream);
+  }
+
+ private:
+  static Tensor BroadcastView(const Tensor src, const Tensor out) {
+    return Tensor{const_cast<void*>(src.data()), out.shape(), src.dtype(),
+                  src.device(), BroadcastStrides(src, out)};
+  }
+
+  // Descriptors must outlive the asynchronous ACLNN call. The cached
+  // operator owns them, while each invocation still gets a fresh executor.
+  mutable ascend::AclTensorCache in_cache_;
+
+  mutable ascend::AclTensorCache out_cache_;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_ASCEND_COPY_KERNEL_H_