diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 36f946d8d..efca1f3a4 100644 --- a/src/CMakeLists.txt +++ b/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 $ -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. # `$` 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 diff --git a/src/native/ascend/custom/rms_norm/op_kernel/rms_norm.cpp b/src/native/ascend/custom/rms_norm/op_kernel/rms_norm.cpp index 5c8f4fc67..197301bb3 100644 --- a/src/native/ascend/custom/rms_norm/op_kernel/rms_norm.cpp +++ b/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 op; op.Init(x, weight, y, totalRows, dimLength, dimLengthAlign, formerNum, formerLength, tailLength, eps); diff --git a/src/native/ascend/ops/argmax/kernel.h b/src/native/ascend/ops/argmax/kernel.h new file mode 100644 index 000000000..9db23e0fb --- /dev/null +++ b/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 +#include + +#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 : public Argmax { + public: + Operator(const Tensor input, const std::optional 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(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 dim, + const bool keepdim, Tensor out) const override { + (void)dim; + (void)keepdim; + auto stream = static_cast(stream_); + auto t_input = input_cache_.get(const_cast(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(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(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_ diff --git a/src/native/ascend/ops/causal_softmax/kernel.h b/src/native/ascend/ops/causal_softmax/kernel.h new file mode 100644 index 000000000..2c56677a4 --- /dev/null +++ b/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 +#include + +#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 : 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 mask_host(mask_elems, 0); + for (size_t i = 0; i < seq_len_; ++i) { + auto vis_end = static_cast(total_seq_len_ - seq_len_ + i); + for (auto j = vis_end + 1; j < static_cast(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 mshape = {static_cast(seq_len_), + static_cast(total_seq_len_)}; + std::vector mstrides = {static_cast(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(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(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, ©_ws, ©_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::infinity(); + + uint64_t temp_size_ = 0; + + void* mask_buf_ = nullptr; + + aclTensor* mask_tensor_ = nullptr; + + aclScalar* neg_inf_ = nullptr; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/ascend/ops/copy/kernel.h b/src/native/ascend/ops/copy/kernel.h new file mode 100644 index 000000000..7660a637b --- /dev/null +++ b/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 : 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(stream_); + auto t_in = in_cache_.get(const_cast(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(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_ diff --git a/src/native/ascend/ops/fused_add_rms_norm/kernel.h b/src/native/ascend/ops/fused_add_rms_norm/kernel.h new file mode 100644 index 000000000..c2b1ca710 --- /dev/null +++ b/src/native/ascend/ops/fused_add_rms_norm/kernel.h @@ -0,0 +1,246 @@ +#ifndef INFINI_OPS_ASCEND_FUSED_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_ASCEND_FUSED_ADD_RMS_NORM_KERNEL_H_ + +#include +#include +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_add.h" +#include "aclnn_inplace_add_rms_norm.h" +#include "base/fused_add_rms_norm.h" +#include "base/rms_norm.h" +#include "data_type.h" +#include "native/ascend/common.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" + +namespace infini::ops { +namespace ascend_fused_add_rms_norm_detail { + +inline void* AllocateOnes(Tensor::Size dim, DataType dtype) { + auto element_size = kDataTypeToSize.at(dtype); + auto bytes = dim * element_size; + void* data = nullptr; + auto ret = aclrtMalloc(&data, bytes, ACL_MEM_MALLOC_NORMAL_ONLY); + assert(ret == ACL_SUCCESS && + "`FusedAddRmsNorm` Ascend path failed to allocate unit weight"); + + if (dtype == DataType::kFloat32) { + std::vector host(dim, 1.0f); + ret = + aclrtMemcpy(data, bytes, host.data(), bytes, ACL_MEMCPY_HOST_TO_DEVICE); + } else { + assert((dtype == DataType::kFloat16 || dtype == DataType::kBFloat16) && + "`FusedAddRmsNorm` Ascend path supports fp16, bf16, or fp32"); + auto one_bits = + static_cast(dtype == DataType::kFloat16 ? 0x3c00 : 0x3f80); + std::vector host(dim, one_bits); + ret = + aclrtMemcpy(data, bytes, host.data(), bytes, ACL_MEMCPY_HOST_TO_DEVICE); + } + + assert(ret == ACL_SUCCESS && + "`FusedAddRmsNorm` Ascend path failed to initialize unit weight"); + return data; +} + +} // namespace ascend_fused_add_rms_norm_detail + +// The single-call ACLNN implementation is the default after passing the full +// precision matrix and outperforming the decomposition on all Llama shapes. +template <> +class Operator + : public FusedAddRmsNorm { + public: + Operator(Tensor input, Tensor residual, const std::optional weight, + float epsilon) + : FusedAddRmsNorm(input, residual, weight, epsilon), + shape_(input.shape()), + dtype_(input.dtype()), + element_size_(kDataTypeToSize.at(input.dtype())), + row_bytes_(dim_ * element_size_), + tensor_bytes_(num_tokens_ * row_bytes_), + needs_input_staging_(!input.IsContiguous()), + needs_residual_staging_(!residual.IsContiguous()) { + if (!weight.has_value()) { + unit_weight_data_ = + ascend_fused_add_rms_norm_detail::AllocateOnes(dim_, input.dtype()); + unit_weight_.emplace(unit_weight_data_, Tensor::Shape{dim_}, + input.dtype(), input.device()); + } + + const auto& resolved_weight = weight.has_value() ? *weight : *unit_weight_; + weight_cache_ = ascend::AclTensorCache(resolved_weight); + + std::vector shape(input.shape().begin(), input.shape().end()); + if (needs_input_staging_) { + input_cache_ = + ascend::AclTensorCache(shape, ascend::ToAclDtype(dtype_), nullptr); + } else { + input_cache_ = ascend::AclTensorCache(input); + } + if (needs_residual_staging_) { + residual_cache_ = + ascend::AclTensorCache(shape, ascend::ToAclDtype(dtype_), nullptr); + } else { + residual_cache_ = ascend::AclTensorCache(residual); + } + + rstd_shape_.assign(input.shape().begin(), input.shape().end()); + rstd_shape_.back() = 1; + rstd_size_ = num_tokens_ * sizeof(float); + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + input_cache_.release(); + residual_cache_.release(); + weight_cache_.release(); + if (unit_weight_data_) aclrtFree(unit_weight_data_); + // `rstd_tensor_` remains owned by the repeatable ACLNN executor. + } + + void operator()(Tensor input, Tensor residual, + const std::optional weight, + float epsilon) const override { + auto stream = static_cast(stream_); + void* input_data = input.data(); + void* residual_data = residual.data(); + + if (needs_input_staging_) { + auto& staging = ascend::GetWorkspacePool().Ensure( + stream, tensor_bytes_, "fused_add_rms_norm_input"); + PackRows(staging.buf, input.data(), input_strides_, stream); + input_data = staging.buf; + } + if (needs_residual_staging_) { + auto& staging = ascend::GetWorkspacePool().Ensure( + stream, tensor_bytes_, "fused_add_rms_norm_residual"); + PackRows(staging.buf, residual.data(), residual_strides_, stream); + residual_data = staging.buf; + } + + const void* weight_data = + weight.has_value() ? weight->data() : unit_weight_data_; + auto t_input = input_cache_.get(input_data); + auto t_residual = residual_cache_.get(residual_data); + auto t_weight = weight_cache_.get(const_cast(weight_data)); + + auto& rstd_arena = ascend::GetWorkspacePool().Ensure( + stream, rstd_size_, "fused_add_rms_norm_rstd"); + if (!rstd_tensor_) { + rstd_tensor_ = aclCreateTensor( + rstd_shape_.data(), static_cast(rstd_shape_.size()), + ACL_FLOAT, /*strides=*/nullptr, 0, ACL_FORMAT_ND, rstd_shape_.data(), + static_cast(rstd_shape_.size()), rstd_arena.buf); + } else { + aclSetRawTensorAddr(rstd_tensor_, rstd_arena.buf); + } + + if (!executor_) { + aclnnInplaceAddRmsNormGetWorkspaceSize( + t_input, t_residual, t_weight, static_cast(epsilon), + rstd_tensor_, &ws_size_, &executor_); + aclSetAclOpExecutorRepeatable(executor_); + } else { + aclSetInputTensorAddr(executor_, 0, t_input, input_data); + aclSetInputTensorAddr(executor_, 1, t_residual, residual_data); + aclSetInputTensorAddr(executor_, 2, t_weight, + const_cast(weight_data)); + aclSetOutputTensorAddr(executor_, 0, rstd_tensor_, rstd_arena.buf); + } + + auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_); + aclnnInplaceAddRmsNorm(arena.buf, ws_size_, executor_, stream); + + if (needs_input_staging_) { + UnpackRows(input.data(), input_data, input_strides_, stream); + } + if (needs_residual_staging_) { + UnpackRows(residual.data(), residual_data, residual_strides_, stream); + } + } + + private: + int64_t RowOffset(int64_t row, const Tensor::Strides& strides) const { + int64_t remaining = row; + int64_t offset = 0; + for (int64_t axis = static_cast(shape_.size()) - 2; axis >= 0; + --axis) { + auto coordinate = remaining % static_cast(shape_[axis]); + remaining /= static_cast(shape_[axis]); + offset += coordinate * static_cast(strides[axis]); + } + return offset; + } + + void PackRows(void* dst, const void* src, const Tensor::Strides& strides, + aclrtStream stream) const { + for (int64_t row = 0; row < static_cast(num_tokens_); ++row) { + auto offset = RowOffset(row, strides); + auto ret = aclrtMemcpyAsync( + static_cast(dst) + row * row_bytes_, row_bytes_, + static_cast(src) + offset * element_size_, row_bytes_, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream); + assert(ret == ACL_SUCCESS && + "`FusedAddRmsNorm` Ascend input pack failed"); + } + } + + void UnpackRows(void* dst, const void* src, const Tensor::Strides& strides, + aclrtStream stream) const { + for (int64_t row = 0; row < static_cast(num_tokens_); ++row) { + auto offset = RowOffset(row, strides); + auto ret = aclrtMemcpyAsync( + static_cast(dst) + offset * element_size_, row_bytes_, + static_cast(src) + row * row_bytes_, row_bytes_, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream); + assert(ret == ACL_SUCCESS && + "`FusedAddRmsNorm` Ascend output unpack failed"); + } + } + + Tensor::Shape shape_; + + DataType dtype_; + + uint64_t element_size_{0}; + + uint64_t row_bytes_{0}; + + uint64_t tensor_bytes_{0}; + + bool needs_input_staging_{false}; + + bool needs_residual_staging_{false}; + + void* unit_weight_data_{nullptr}; + + std::optional unit_weight_; + + mutable ascend::AclTensorCache input_cache_; + + mutable ascend::AclTensorCache residual_cache_; + + mutable ascend::AclTensorCache weight_cache_; + + std::vector rstd_shape_; + + uint64_t rstd_size_{0}; + + mutable aclTensor* rstd_tensor_{nullptr}; + + mutable aclOpExecutor* executor_{nullptr}; + + mutable uint64_t ws_size_{0}; +}; + +} // namespace infini::ops + +#include "native/ascend/ops/fused_add_rms_norm/kernel_custom.h" + +#endif diff --git a/src/native/ascend/ops/fused_add_rms_norm/kernel_custom.h b/src/native/ascend/ops/fused_add_rms_norm/kernel_custom.h new file mode 100644 index 000000000..975756d6d --- /dev/null +++ b/src/native/ascend/ops/fused_add_rms_norm/kernel_custom.h @@ -0,0 +1,158 @@ +#ifndef INFINI_OPS_ASCEND_FUSED_ADD_RMS_NORM_KERNEL_CUSTOM_H_ +#define INFINI_OPS_ASCEND_FUSED_ADD_RMS_NORM_KERNEL_CUSTOM_H_ + +#ifdef INFINI_HAS_CUSTOM_KERNELS + +#include +#include +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnnop/aclnn_cast.h" +#include "base/fused_add_rms_norm.h" +#include "native/ascend/common.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" + +extern "C" uint32_t aclrtlaunch_add_rms_norm( + uint32_t block_dim, aclrtStream stream, void* x1, void* x2, void* weight, + void* y, void* x_out, int64_t total_rows, int64_t dim_length, + int64_t dim_length_align, int64_t former_num, int64_t former_length, + int64_t tail_length, float eps, int64_t dtype_size); + +namespace infini::ops { +namespace ascend_fused_add_rms_norm_custom_detail { + +inline void* AllocateFloatOnes(Tensor::Size dim) { + auto bytes = dim * sizeof(float); + void* data = nullptr; + auto ret = aclrtMalloc(&data, bytes, ACL_MEM_MALLOC_NORMAL_ONLY); + assert(ret == ACL_SUCCESS && + "`FusedAddRmsNorm` AscendC unit weight allocation failed"); + + std::vector host(dim, 1.0f); + ret = aclrtMemcpy(data, bytes, host.data(), bytes, ACL_MEMCPY_HOST_TO_DEVICE); + assert(ret == ACL_SUCCESS && + "`FusedAddRmsNorm` AscendC unit weight initialization failed"); + + return data; +} + +} // namespace ascend_fused_add_rms_norm_custom_detail + +// Existing AscendC AddRmsNorm kernel, exposed as implementation index 2 for +// precision and performance comparison with the ACLNN implementations. +template <> +class Operator + : public FusedAddRmsNorm { + public: + Operator(Tensor input, Tensor residual, const std::optional weight, + float epsilon) + : FusedAddRmsNorm(input, residual, weight, epsilon), + dtype_(input.dtype()), + weight_dtype_(weight.has_value() ? weight->dtype() + : DataType::kFloat32) { + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kFloat32) && + "`FusedAddRmsNorm` AscendC kernel supports fp16 and fp32"); + assert(input.IsContiguous() && residual.IsContiguous() && + (!weight.has_value() || weight->IsContiguous()) && + "`FusedAddRmsNorm` AscendC kernel requires contiguous tensors"); + + auto align_elements = 32 / static_cast(kDataTypeToSize.at(dtype_)); + assert(static_cast(dim_) % align_elements == 0 && + "`FusedAddRmsNorm` AscendC kernel requires a 32-byte aligned " + "last dimension"); + assert(dim_ <= 4096 && + "`FusedAddRmsNorm` AscendC kernel exceeds the conservative UB " + "capacity bound; use index 0"); + + total_rows_ = static_cast(input.numel() / dim_); + if (!weight.has_value()) { + weight_fp32_data_ = + ascend_fused_add_rms_norm_custom_detail::AllocateFloatOnes(dim_); + } else if (weight_dtype_ != DataType::kFloat32) { + auto bytes = static_cast(dim_) * sizeof(float); + auto ret = + aclrtMalloc(&weight_fp32_data_, bytes, ACL_MEM_MALLOC_NORMAL_ONLY); + assert(ret == ACL_SUCCESS && + "`FusedAddRmsNorm` AscendC weight allocation failed"); + weight_src_cache_ = + ascend::AclTensorCache({static_cast(dim_)}, + ascend::ToAclDtype(weight_dtype_), nullptr); + weight_dst_cache_ = ascend::AclTensorCache({static_cast(dim_)}, + ACL_FLOAT, weight_fp32_data_); + } + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + weight_src_cache_.release(); + weight_dst_cache_.release(); + if (weight_fp32_data_) aclrtFree(weight_fp32_data_); + } + + void operator()(Tensor input, Tensor residual, + const std::optional weight, + float epsilon) const override { + auto stream = static_cast(stream_); + void* weight_fp32 = weight_fp32_data_; + + if (weight.has_value() && weight_dtype_ == DataType::kFloat32) { + weight_fp32 = const_cast(weight->data()); + } else if (weight.has_value()) { + auto t_src = weight_src_cache_.get(const_cast(weight->data())); + auto t_dst = weight_dst_cache_.get(weight_fp32_data_); + if (!cast_executor_) { + aclnnCastGetWorkspaceSize(t_src, ACL_FLOAT, t_dst, &cast_ws_size_, + &cast_executor_); + aclSetAclOpExecutorRepeatable(cast_executor_); + } else { + aclSetInputTensorAddr(cast_executor_, 0, t_src, + const_cast(weight->data())); + aclSetOutputTensorAddr(cast_executor_, 0, t_dst, weight_fp32_data_); + } + + auto& arena = ascend::GetWorkspacePool().Ensure(stream, cast_ws_size_); + aclnnCast(arena.buf, cast_ws_size_, cast_executor_, stream); + weight_fp32 = weight_fp32_data_; + } + + static constexpr int64_t kMaxBlockDim = 40; + auto used_cores = std::min(total_rows_, kMaxBlockDim); + auto former_length = (total_rows_ + used_cores - 1) / used_cores; + auto tail_length = former_length - 1; + auto former_num = total_rows_ - tail_length * used_cores; + + aclrtlaunch_add_rms_norm( + static_cast(used_cores), stream, input.data(), + residual.data(), weight_fp32, input.data(), residual.data(), + total_rows_, static_cast(dim_), static_cast(dim_), + former_num, former_length, tail_length, epsilon, + static_cast(kDataTypeToSize.at(dtype_))); + } + + private: + DataType dtype_; + + DataType weight_dtype_; + + int64_t total_rows_{0}; + + void* weight_fp32_data_{nullptr}; + + mutable ascend::AclTensorCache weight_src_cache_; + + mutable ascend::AclTensorCache weight_dst_cache_; + + mutable aclOpExecutor* cast_executor_{nullptr}; + + mutable uint64_t cast_ws_size_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_HAS_CUSTOM_KERNELS +#endif // INFINI_OPS_ASCEND_FUSED_ADD_RMS_NORM_KERNEL_CUSTOM_H_ diff --git a/src/native/ascend/ops/rms_norm/kernel.h b/src/native/ascend/ops/rms_norm/kernel.h new file mode 100644 index 000000000..478262583 --- /dev/null +++ b/src/native/ascend/ops/rms_norm/kernel.h @@ -0,0 +1,228 @@ +#ifndef INFINI_OPS_ASCEND_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_ASCEND_RMS_NORM_KERNEL_H_ + +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_rms_norm.h" +#include "aclnnop/aclnn_cast.h" +#include "base/rms_norm.h" +#include "native/ascend/common.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" + +namespace infini::ops { + +template <> +class Operator : public RmsNorm { + public: + Operator(const Tensor input, const Tensor weight, float eps, Tensor out) + : RmsNorm(input, weight, eps, out), weight_cache_(weight) { + assert(ndim_ >= 2 && "`RmsNorm` Ascend path requires rank >= 2"); + assert(input.strides().back() == 1 && out.strides().back() == 1 && + "`RmsNorm` Ascend path requires a contiguous last dimension"); + + needs_input_staging_ = !input.IsContiguous(); + needs_out_staging_ = !out.IsContiguous(); + logical_rows_ = input.numel() / dim_; + element_size_ = kDataTypeToSize.at(input.dtype()); + row_bytes_ = dim_ * element_size_; + tensor_bytes_ = logical_rows_ * row_bytes_; + + std::vector input_shape(input.shape().begin(), + input.shape().end()); + std::vector out_shape(out.shape().begin(), out.shape().end()); + if (needs_input_staging_) { + norm_in_cache_ = ascend::AclTensorCache( + input_shape, ascend::ToAclDtype(input.dtype()), nullptr); + } else { + norm_in_cache_ = ascend::AclTensorCache(input); + } + if (needs_out_staging_) { + norm_out_cache_ = ascend::AclTensorCache( + out_shape, ascend::ToAclDtype(out.dtype()), nullptr); + } else { + norm_out_cache_ = ascend::AclTensorCache(out); + } + + needs_weight_cast_ = + input.dtype() != weight.dtype() && weight.dtype() != DataType::kFloat32; + if (needs_weight_cast_) { + auto fp32_bytes = static_cast(dim_) * sizeof(float); + auto ret = aclrtMalloc(&weight_fp32_data_, fp32_bytes, + ACL_MEM_MALLOC_NORMAL_ONLY); + assert(ret == ACL_SUCCESS && + "`RmsNorm` Ascend path failed to allocate cast weight"); + weight_fp32_cache_ = ascend::AclTensorCache({static_cast(dim_)}, + ACL_FLOAT, weight_fp32_data_); + } + + rstd_shape_.assign(input.shape().begin(), input.shape().end() - 1); + rstd_strides_.resize(rstd_shape_.size()); + int64_t rstd_stride = 1; + for (int64_t axis = static_cast(rstd_shape_.size()) - 1; axis >= 0; + --axis) { + rstd_strides_[axis] = rstd_stride; + rstd_stride *= rstd_shape_[axis]; + } + rstd_size_ = logical_rows_ * sizeof(float); + + // Follow InfiniCore's proven ACLNN path: prepare one repeatable executor + // from stable tensor metadata. Runtime calls only rebind data addresses. + auto t_in = norm_in_cache_.get(nullptr); + auto t_out = norm_out_cache_.get(nullptr); + aclTensor *t_weight; + if (needs_weight_cast_) { + auto t_weight_src = weight_cache_.get(nullptr); + auto t_weight_dst = weight_fp32_cache_.get(weight_fp32_data_); + aclnnCastGetWorkspaceSize(t_weight_src, ACL_FLOAT, t_weight_dst, + &cast_ws_, &cast_exec_); + aclSetAclOpExecutorRepeatable(cast_exec_); + t_weight = t_weight_dst; + } else { + t_weight = weight_cache_.get(nullptr); + } + + rstd_tensor_ = aclCreateTensor( + rstd_shape_.data(), static_cast(rstd_shape_.size()), ACL_FLOAT, + rstd_strides_.data(), 0, ACL_FORMAT_ND, rstd_shape_.data(), + static_cast(rstd_shape_.size()), nullptr); + aclnnRmsNormGetWorkspaceSize(t_in, t_weight, eps, t_out, rstd_tensor_, + &ws_size_, &executor_); + aclSetAclOpExecutorRepeatable(executor_); + total_workspace_size_ = ws_size_ + rstd_size_; + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + norm_in_cache_.release(); + weight_cache_.release(); + norm_out_cache_.release(); + weight_fp32_cache_.release(); + if (weight_fp32_data_) aclrtFree(weight_fp32_data_); + // `rstd_tensor_` leaks with the executor at shutdown (see `64c367c`). + } + + void operator()(const Tensor input, const Tensor weight, float eps, + Tensor out) const override { + auto stream = static_cast(stream_); + void *input_data = const_cast(input.data()); + void *out_data = out.data(); + + if (needs_input_staging_) { + auto &staging = ascend::GetWorkspacePool().Ensure(stream, tensor_bytes_, + "rms_norm_input"); + PackRows(staging.buf, input.data(), input_shape_, input_strides_, stream); + input_data = staging.buf; + } + if (needs_out_staging_) { + auto &staging = ascend::GetWorkspacePool().Ensure(stream, tensor_bytes_, + "rms_norm_output"); + out_data = staging.buf; + } + + auto t_in = norm_in_cache_.get(input_data); + auto t_out = norm_out_cache_.get(out_data); + aclTensor *t_weight; + void *weight_data; + + if (needs_weight_cast_) { + auto t_weight_src = weight_cache_.get(const_cast(weight.data())); + auto t_weight_dst = weight_fp32_cache_.get(weight_fp32_data_); + AclSetTensorAddr(cast_exec_, 0, t_weight_src, + const_cast(weight.data())); + AclSetTensorAddr(cast_exec_, 1, t_weight_dst, weight_fp32_data_); + auto &cast_arena = ascend::GetWorkspacePool().Ensure(stream, cast_ws_); + aclnnCast(cast_arena.buf, cast_ws_, cast_exec_, stream); + t_weight = t_weight_dst; + weight_data = weight_fp32_data_; + } else { + t_weight = weight_cache_.get(const_cast(weight.data())); + weight_data = const_cast(weight.data()); + } + + auto &arena = ascend::GetWorkspacePool().Ensure( + stream, total_workspace_size_, "rms_norm"); + auto *rstd_data = static_cast(arena.buf) + ws_size_; + aclSetRawTensorAddr(rstd_tensor_, rstd_data); + + AclSetTensorAddr(executor_, 0, t_in, input_data); + AclSetTensorAddr(executor_, 1, t_weight, weight_data); + AclSetTensorAddr(executor_, 2, t_out, out_data); + AclSetTensorAddr(executor_, 3, rstd_tensor_, rstd_data); + aclnnRmsNorm(arena.buf, ws_size_, executor_, stream); + + if (needs_out_staging_) { + UnpackRows(out.data(), out_data, out_shape_, out_strides_, stream); + } + } + + private: + int64_t RowOffset(int64_t row, const Tensor::Shape &shape, + const Tensor::Strides &strides) const { + int64_t remaining = row; + int64_t offset = 0; + for (int64_t axis = static_cast(shape.size()) - 2; axis >= 0; + --axis) { + auto coordinate = remaining % static_cast(shape[axis]); + remaining /= static_cast(shape[axis]); + offset += coordinate * static_cast(strides[axis]); + } + return offset; + } + + void PackRows(void *dst, const void *src, const Tensor::Shape &shape, + const Tensor::Strides &strides, aclrtStream stream) const { + for (int64_t row = 0; row < logical_rows_; ++row) { + auto offset = RowOffset(row, shape, strides); + auto ret = aclrtMemcpyAsync( + static_cast(dst) + row * row_bytes_, row_bytes_, + static_cast(src) + offset * element_size_, + row_bytes_, ACL_MEMCPY_DEVICE_TO_DEVICE, stream); + assert(ret == ACL_SUCCESS && "`RmsNorm` input pack failed"); + } + } + + void UnpackRows(void *dst, const void *src, const Tensor::Shape &shape, + const Tensor::Strides &strides, aclrtStream stream) const { + for (int64_t row = 0; row < logical_rows_; ++row) { + auto offset = RowOffset(row, shape, strides); + auto ret = aclrtMemcpyAsync( + static_cast(dst) + offset * element_size_, row_bytes_, + static_cast(src) + row * row_bytes_, row_bytes_, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream); + assert(ret == ACL_SUCCESS && "`RmsNorm` output unpack failed"); + } + } + + mutable ascend::AclTensorCache norm_in_cache_; + mutable ascend::AclTensorCache weight_cache_; + mutable ascend::AclTensorCache norm_out_cache_; + mutable ascend::AclTensorCache weight_fp32_cache_; + bool needs_input_staging_{false}; + bool needs_out_staging_{false}; + bool needs_weight_cast_{false}; + int64_t logical_rows_{0}; + uint64_t element_size_{0}; + uint64_t row_bytes_{0}; + uint64_t tensor_bytes_{0}; + void *weight_fp32_data_{nullptr}; + mutable aclOpExecutor *cast_exec_{nullptr}; + mutable uint64_t cast_ws_{0}; + mutable aclOpExecutor *executor_{nullptr}; + mutable uint64_t ws_size_{0}; + std::vector rstd_shape_; + std::vector rstd_strides_; + uint64_t rstd_size_{0}; + uint64_t total_workspace_size_{0}; + mutable aclTensor *rstd_tensor_{nullptr}; +}; + +} // namespace infini::ops + +#include "native/ascend/ops/rms_norm/kernel_custom.h" + +#endif diff --git a/src/native/ascend/ops/rms_norm/kernel_custom.h b/src/native/ascend/ops/rms_norm/kernel_custom.h new file mode 100644 index 000000000..fcef87d51 --- /dev/null +++ b/src/native/ascend/ops/rms_norm/kernel_custom.h @@ -0,0 +1,170 @@ +#ifndef INFINI_OPS_ASCEND_RMS_NORM_KERNEL_CUSTOM_H_ +#define INFINI_OPS_ASCEND_RMS_NORM_KERNEL_CUSTOM_H_ + +#ifdef INFINI_HAS_CUSTOM_KERNELS + +#include +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnnop/aclnn_cast.h" +#include "base/rms_norm.h" +#include "native/ascend/common.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" + +// Forward-declare the `aclrtlaunch_rms_norm` launch symbol defined by +// the AscendC toolchain from `custom/rms_norm/op_kernel/`. +extern "C" uint32_t aclrtlaunch_rms_norm( + uint32_t block_dim, void* stream, void* input, void* weight, void* out, + int64_t total_rows, int64_t dim_length, int64_t dim_length_align, + int64_t former_num, int64_t former_length, int64_t tail_length, float eps, + int64_t is_half); + +namespace infini::ops { + +// Custom AscendC fused `RmsNorm` kernel (implementation index 1). +// +// A single-kernel implementation that computes `RMSNorm` in one launch, +// avoiding the 5-sub-op decomposition of `aclnnRmsNorm` (index 0). Uses +// `Sqrt` + scalar division instead of `Rsqrt` for higher precision (~1e-7 +// `fp32` error vs ~0.2% with `Rsqrt`). +// +// Select via `implementation_index=1` in Python: +// `infini.ops.rms_norm(input, weight, eps, out, implementation_index=1, +// stream=s)`. +// +// Requirements: +// - Input last dimension must be 32-byte aligned (divisible by 16 for +// `fp16` or 8 for `fp32`). +// - Input, weight, and output must be contiguous. +// - The custom kernel binary must be linked (`BUILD_ASCEND_CUSTOM=ON`). +// +// Use implementation index 0 for other valid layouts and `bf16`. +template <> +class Operator : public RmsNorm { + public: + Operator(const Tensor input, const Tensor weight, float eps, Tensor out) + : RmsNorm(input, weight, eps, out), + dtype_{input.dtype()}, + weight_dtype_{weight.dtype()} { + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || + dtype_ == DataType::kFloat32) && + "`RmsNorm` custom kernel: `input` must be `fp16`, `bf16`, or " + "`fp32`"); + assert((weight_dtype_ == DataType::kFloat16 || + weight_dtype_ == DataType::kBFloat16 || + weight_dtype_ == DataType::kFloat32) && + "`RmsNorm`: `weight` must be `fp16`, `bf16`, or `fp32`"); + + // 32-byte alignment on the last dimension — kernel relies on aligned + // `DataCopyPad` loads/stores. + int64_t align_elems = 32 / static_cast(kDataTypeToSize.at(dtype_)); + dim_length_align_ = + ((static_cast(dim_) + align_elems - 1) / align_elems) * + align_elems; + + assert(dtype_ != DataType::kBFloat16 && + "`RmsNorm` custom kernel does not support `bf16`; use index 0"); + assert(input.IsContiguous() && weight.IsContiguous() && + out.IsContiguous() && + "`RmsNorm` custom kernel requires contiguous tensors"); + assert(static_cast(dim_) == dim_length_align_ && + "`RmsNorm` custom kernel requires a 32-byte aligned last " + "dimension"); + + total_rows_ = static_cast(input.numel() / dim_); + // The custom kernel always reads `weight` as fp32, so fp16 / bf16 + // inputs need a cached `aclnnCast` invocation in `operator()` to + // produce an fp32 shadow buffer on every launch. + if (weight_dtype_ != DataType::kFloat32) { + size_t fp32_bytes = static_cast(dim_) * sizeof(float); + aclrtMalloc(&weight_fp32_data_, fp32_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); + + weight_src_cache_ = + ascend::AclTensorCache({static_cast(dim_)}, + ascend::ToAclDtype(weight_dtype_), nullptr); + weight_dst_cache_ = ascend::AclTensorCache({static_cast(dim_)}, + ACL_FLOAT, weight_fp32_data_); + } + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + // Null cached descriptors — see `AclTensorCache::release()`. + weight_src_cache_.release(); + weight_dst_cache_.release(); + + if (weight_fp32_data_) aclrtFree(weight_fp32_data_); + } + + void operator()(const Tensor input, const Tensor weight, float eps, + Tensor out) const override { + auto stream = static_cast(stream_); + + void* weight_fp32; + + if (weight_dtype_ != DataType::kFloat32) { + auto t_src = weight_src_cache_.get(const_cast(weight.data())); + auto t_dst = weight_dst_cache_.get(weight_fp32_data_); + + if (!cast_exec_) { + aclnnCastGetWorkspaceSize(t_src, ACL_FLOAT, t_dst, &cast_ws_, + &cast_exec_); + aclSetAclOpExecutorRepeatable(cast_exec_); + } else { + aclSetInputTensorAddr(cast_exec_, 0, t_src, + const_cast(weight.data())); + aclSetOutputTensorAddr(cast_exec_, 0, t_dst, weight_fp32_data_); + } + + auto& arena = ascend::GetWorkspacePool().Ensure(stream, cast_ws_); + aclnnCast(arena.buf, cast_ws_, cast_exec_, stream); + weight_fp32 = weight_fp32_data_; + } else { + weight_fp32 = const_cast(weight.data()); + } + + // Block-level tiling. Ascend 910B has 20–40 AIV cores; over-subscribing + // is safe (runtime multiplexes) but wastes one weight load per block. + static constexpr int64_t kMaxBlockDim = 40; + int64_t used_cores = std::min(total_rows_, kMaxBlockDim); + int64_t former_length = (total_rows_ + used_cores - 1) / used_cores; + int64_t tail_length = former_length - 1; + int64_t former_num = total_rows_ - tail_length * used_cores; + uint32_t block_dim = static_cast(used_cores); + + aclrtlaunch_rms_norm(block_dim, stream, const_cast(input.data()), + weight_fp32, out.data(), total_rows_, + static_cast(dim_), dim_length_align_, + former_num, former_length, tail_length, eps, + static_cast(dtype_ == DataType::kFloat16)); + } + + private: + DataType dtype_; + + DataType weight_dtype_; + + int64_t dim_length_align_; + + int64_t total_rows_; + + void* weight_fp32_data_ = nullptr; + + mutable ascend::AclTensorCache weight_src_cache_; + + mutable ascend::AclTensorCache weight_dst_cache_; + + mutable aclOpExecutor* cast_exec_ = nullptr; + + mutable uint64_t cast_ws_ = 0; +}; + +} // namespace infini::ops + +#endif // INFINI_HAS_CUSTOM_KERNELS +#endif // INFINI_OPS_ASCEND_RMS_NORM_KERNEL_CUSTOM_H_ diff --git a/src/native/ascend/ops/rotary_embedding/kernel.h b/src/native/ascend/ops/rotary_embedding/kernel.h new file mode 100644 index 000000000..608e8a550 --- /dev/null +++ b/src/native/ascend/ops/rotary_embedding/kernel.h @@ -0,0 +1,250 @@ +#ifndef INFINI_OPS_ASCEND_ROTARY_EMBEDDING_KERNEL_H_ +#define INFINI_OPS_ASCEND_ROTARY_EMBEDDING_KERNEL_H_ + +#include +#include +#include +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnnop/aclnn_apply_rotary_pos_emb_v2.h" +#include "aclnnop/aclnn_index_select.h" +#include "base/rotary_embedding.h" +#include "native/ascend/common.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" + +namespace infini::ops { + +// Llama-style full-dimension NeoX RoPE. Positions select rows from the packed +// [cos, sin] cache, then CANN rotates query and key together in place. +template <> +class Operator + : public RotaryEmbedding { + public: + Operator(const Tensor positions, Tensor query, std::optional key, + const Tensor cos_sin_cache, int64_t head_size, bool is_neox, + int64_t rope_dim_offset = 0, bool inverse = false) + : RotaryEmbedding(positions, query, key, cos_sin_cache, head_size, + is_neox, rope_dim_offset, inverse), + max_seq_len_(static_cast(cos_sin_cache.size(0))), + element_size_(cos_sin_cache.element_size()) { + assert(key.has_value() && "Ascend `RotaryEmbedding` requires a key tensor"); + assert(is_neox_ && rope_dim_offset_ == 0 && !inverse_ && + rot_dim_ == head_size_ && + "Ascend `RotaryEmbedding` supports full-dimension forward NeoX " + "rotation; use another implementation for partial, inverse, or " + "interleaved rotation"); + assert(cos_sin_cache_type_ == query_type_ && + "Ascend `RotaryEmbedding` requires cache and query dtypes to " + "match"); + assert(query.IsContiguous() && key->IsContiguous() && + cos_sin_cache.IsContiguous() && + "Ascend `RotaryEmbedding` requires contiguous query, key, and " + "cache tensors"); + + const auto num_tokens = static_cast(num_tokens_); + const auto head_dim = head_size_; + const auto acl_dtype = ascend::ToAclDtype(query.dtype()); + const auto table_bytes = + static_cast(max_seq_len_ * head_dim) * element_size_; + auto ret = + aclrtMalloc(&cos_table_data_, table_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); + assert(ret == ACL_SUCCESS && + "Ascend `RotaryEmbedding` failed to allocate cosine table"); + ret = + aclrtMalloc(&sin_table_data_, table_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); + assert(ret == ACL_SUCCESS && + "Ascend `RotaryEmbedding` failed to allocate sine table"); + + UploadCosSinCache(cos_sin_cache); + cos_sin_cache_data_ = cos_sin_cache.data(); + + const auto gathered_bytes = + static_cast(num_tokens * head_dim) * element_size_; + ret = aclrtMalloc(&cos_data_, gathered_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); + assert(ret == ACL_SUCCESS && + "Ascend `RotaryEmbedding` failed to allocate gathered cosine"); + ret = aclrtMalloc(&sin_data_, gathered_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); + assert(ret == ACL_SUCCESS && + "Ascend `RotaryEmbedding` failed to allocate gathered sine"); + + cos_table_cache_ = ascend::AclTensorCache({max_seq_len_, head_dim}, + acl_dtype, cos_table_data_); + sin_table_cache_ = ascend::AclTensorCache({max_seq_len_, head_dim}, + acl_dtype, sin_table_data_); + positions_cache_ = ascend::AclTensorCache( + {num_tokens}, ACL_INT64, const_cast(positions.data())); + cos_out_cache_ = + ascend::AclTensorCache({num_tokens, head_dim}, acl_dtype, cos_data_); + sin_out_cache_ = + ascend::AclTensorCache({num_tokens, head_dim}, acl_dtype, sin_data_); + cos_rotary_cache_ = + ascend::AclTensorCache({num_tokens, 1, head_dim}, acl_dtype, cos_data_); + sin_rotary_cache_ = + ascend::AclTensorCache({num_tokens, 1, head_dim}, acl_dtype, sin_data_); + query_cache_ = ascend::AclTensorCache( + {num_tokens, static_cast(num_heads_), head_dim}, acl_dtype, + query.data()); + key_cache_ = ascend::AclTensorCache( + {num_tokens, static_cast(num_kv_heads_), head_dim}, acl_dtype, + key->data()); + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + cos_table_cache_.release(); + sin_table_cache_.release(); + positions_cache_.release(); + cos_out_cache_.release(); + sin_out_cache_.release(); + cos_rotary_cache_.release(); + sin_rotary_cache_.release(); + query_cache_.release(); + key_cache_.release(); + if (cos_table_data_) aclrtFree(cos_table_data_); + if (sin_table_data_) aclrtFree(sin_table_data_); + if (cos_data_) aclrtFree(cos_data_); + if (sin_data_) aclrtFree(sin_data_); + } + + void operator()(const Tensor positions, Tensor query, + std::optional key, const Tensor cos_sin_cache, + int64_t head_size, bool is_neox, int64_t rope_dim_offset = 0, + bool inverse = false) const override { + assert(key.has_value()); + auto stream = static_cast(stream_); + + if (cos_sin_cache.data() != cos_sin_cache_data_) { + UploadCosSinCache(cos_sin_cache); + cos_sin_cache_data_ = cos_sin_cache.data(); + } + + auto t_cos_table = cos_table_cache_.get(cos_table_data_); + auto t_sin_table = sin_table_cache_.get(sin_table_data_); + auto t_positions = + positions_cache_.get(const_cast(positions.data())); + auto t_cos_out = cos_out_cache_.get(cos_data_); + auto t_sin_out = sin_out_cache_.get(sin_data_); + + if (!cos_index_executor_) { + aclnnIndexSelectGetWorkspaceSize(t_cos_table, 0, t_positions, t_cos_out, + &cos_index_ws_size_, + &cos_index_executor_); + aclSetAclOpExecutorRepeatable(cos_index_executor_); + } else { + aclSetInputTensorAddr(cos_index_executor_, 1, t_positions, + const_cast(positions.data())); + } + + if (!sin_index_executor_) { + aclnnIndexSelectGetWorkspaceSize(t_sin_table, 0, t_positions, t_sin_out, + &sin_index_ws_size_, + &sin_index_executor_); + aclSetAclOpExecutorRepeatable(sin_index_executor_); + } else { + aclSetInputTensorAddr(sin_index_executor_, 1, t_positions, + const_cast(positions.data())); + } + + auto index_ws_size = std::max(cos_index_ws_size_, sin_index_ws_size_); + auto& index_arena = + ascend::GetWorkspacePool().Ensure(stream, index_ws_size); + aclnnIndexSelect(index_arena.buf, cos_index_ws_size_, cos_index_executor_, + stream); + aclnnIndexSelect(index_arena.buf, sin_index_ws_size_, sin_index_executor_, + stream); + + auto t_cos = cos_rotary_cache_.get(cos_data_); + auto t_sin = sin_rotary_cache_.get(sin_data_); + auto t_query = query_cache_.get(query.data()); + auto t_key = key_cache_.get(key->data()); + if (!rotary_executor_) { + aclnnApplyRotaryPosEmbV2GetWorkspaceSize( + t_query, t_key, t_cos, t_sin, 4, const_cast("half"), + &rotary_ws_size_, &rotary_executor_); + aclSetAclOpExecutorRepeatable(rotary_executor_); + } else { + aclSetInputTensorAddr(rotary_executor_, 0, t_query, query.data()); + aclSetInputTensorAddr(rotary_executor_, 1, t_key, key->data()); + aclSetInputTensorAddr(rotary_executor_, 2, t_cos, cos_data_); + aclSetInputTensorAddr(rotary_executor_, 3, t_sin, sin_data_); + } + + auto& rotary_arena = + ascend::GetWorkspacePool().Ensure(stream, rotary_ws_size_); + aclnnApplyRotaryPosEmbV2(rotary_arena.buf, rotary_ws_size_, + rotary_executor_, stream); + } + + private: + void UploadCosSinCache(const Tensor cos_sin_cache) const { + const auto half_dim = head_size_ / 2; + const auto table_bytes = + static_cast(max_seq_len_ * head_size_) * element_size_; + std::vector packed(table_bytes); + std::vector cosine(table_bytes); + std::vector sine(table_bytes); + + auto ret = aclrtMemcpy(packed.data(), table_bytes, cos_sin_cache.data(), + table_bytes, ACL_MEMCPY_DEVICE_TO_HOST); + assert(ret == ACL_SUCCESS && + "Ascend `RotaryEmbedding` failed to read cos/sin cache"); + for (int64_t position = 0; position < max_seq_len_; ++position) { + for (int64_t index = 0; index < half_dim; ++index) { + const auto cos_source = + packed.data() + (position * head_size_ + index) * element_size_; + const auto sin_source = + packed.data() + + (position * head_size_ + half_dim + index) * element_size_; + for (int64_t half = 0; half < 2; ++half) { + const auto destination_index = + position * head_size_ + half * half_dim + index; + std::memcpy(cosine.data() + destination_index * element_size_, + cos_source, element_size_); + std::memcpy(sine.data() + destination_index * element_size_, + sin_source, element_size_); + } + } + } + + ret = aclrtMemcpy(cos_table_data_, table_bytes, cosine.data(), table_bytes, + ACL_MEMCPY_HOST_TO_DEVICE); + assert(ret == ACL_SUCCESS && + "Ascend `RotaryEmbedding` failed to upload cosine table"); + ret = aclrtMemcpy(sin_table_data_, table_bytes, sine.data(), table_bytes, + ACL_MEMCPY_HOST_TO_DEVICE); + assert(ret == ACL_SUCCESS && + "Ascend `RotaryEmbedding` failed to upload sine table"); + } + + int64_t max_seq_len_{0}; + size_t element_size_{0}; + mutable const void* cos_sin_cache_data_{nullptr}; + void* cos_table_data_{nullptr}; + void* sin_table_data_{nullptr}; + void* cos_data_{nullptr}; + void* sin_data_{nullptr}; + mutable ascend::AclTensorCache cos_table_cache_; + mutable ascend::AclTensorCache sin_table_cache_; + mutable ascend::AclTensorCache positions_cache_; + mutable ascend::AclTensorCache cos_out_cache_; + mutable ascend::AclTensorCache sin_out_cache_; + mutable ascend::AclTensorCache cos_rotary_cache_; + mutable ascend::AclTensorCache sin_rotary_cache_; + mutable ascend::AclTensorCache query_cache_; + mutable ascend::AclTensorCache key_cache_; + mutable aclOpExecutor* cos_index_executor_{nullptr}; + mutable uint64_t cos_index_ws_size_{0}; + mutable aclOpExecutor* sin_index_executor_{nullptr}; + mutable uint64_t sin_index_ws_size_{0}; + mutable aclOpExecutor* rotary_executor_{nullptr}; + mutable uint64_t rotary_ws_size_{0}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/ascend/ops/silu_and_mul/kernel.h b/src/native/ascend/ops/silu_and_mul/kernel.h new file mode 100644 index 000000000..84cb39471 --- /dev/null +++ b/src/native/ascend/ops/silu_and_mul/kernel.h @@ -0,0 +1,112 @@ +#ifndef INFINI_OPS_ASCEND_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_ASCEND_SILU_AND_MUL_KERNEL_H_ + +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_copy.h" +#include "aclnnop/aclnn_swi_glu.h" +#include "base/silu_and_mul.h" +#include "native/ascend/common.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" + +namespace infini::ops { + +// `aclnnSwiGlu` splits the last input dimension into `[gate, up]` and +// computes `silu(gate) * up`. CANN writes its output contiguously, so a +// staging tensor is required when the InfiniOps output has padding. +template <> +class Operator : public SiluAndMul { + public: + Operator(const Tensor input, Tensor out) + : SiluAndMul(input, out), input_cache_(input), out_cache_(out) { + needs_copy_ = !is_out_contiguous_; + if (needs_copy_) { + out_staging_size_ = out.numel() * kDataTypeToSize.at(out.dtype()); + } + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + input_cache_.release(); + out_cache_.release(); + if (out_staging_cache_) out_staging_cache_->release(); + } + + void operator()(const Tensor input, Tensor out) const override { + auto t_input = input_cache_.get(const_cast(input.data())); + auto t_out = out_cache_.get(out.data()); + auto stream = static_cast(stream_); + + aclTensor* t_swiglu_out = t_out; + void* swiglu_out_data = out.data(); + + if (needs_copy_) { + auto& staging = ascend::GetWorkspacePool().Ensure( + stream, out_staging_size_, "silu_and_mul_output"); + + if (!out_staging_cache_) { + std::vector out_shape(out_shape_.begin(), out_shape_.end()); + out_staging_cache_.emplace(out_shape, ascend::ToAclDtype(out_type_), + staging.buf); + } + + t_swiglu_out = out_staging_cache_->get(staging.buf); + swiglu_out_data = staging.buf; + } + + if (!swiglu_exec_) { + aclnnSwiGluGetWorkspaceSize(t_input, -1, t_swiglu_out, &swiglu_ws_, + &swiglu_exec_); + aclSetAclOpExecutorRepeatable(swiglu_exec_); + } else { + aclSetInputTensorAddr(swiglu_exec_, 0, t_input, + const_cast(input.data())); + aclSetOutputTensorAddr(swiglu_exec_, 0, t_swiglu_out, swiglu_out_data); + } + + auto& arena = ascend::GetWorkspacePool().Ensure(stream, swiglu_ws_); + aclnnSwiGlu(arena.buf, swiglu_ws_, swiglu_exec_, stream); + + if (needs_copy_) { + if (!copy_exec_) { + aclnnInplaceCopyGetWorkspaceSize(t_out, t_swiglu_out, ©_ws_, + ©_exec_); + aclSetAclOpExecutorRepeatable(copy_exec_); + } else { + aclSetInputTensorAddr(copy_exec_, 0, t_out, out.data()); + aclSetInputTensorAddr(copy_exec_, 1, t_swiglu_out, swiglu_out_data); + } + + auto& copy_arena = ascend::GetWorkspacePool().Ensure(stream, copy_ws_); + aclnnInplaceCopy(copy_arena.buf, copy_ws_, copy_exec_, stream); + } + } + + private: + mutable ascend::AclTensorCache input_cache_; + + mutable ascend::AclTensorCache out_cache_; + + mutable std::optional out_staging_cache_; + + bool needs_copy_{false}; + + uint64_t out_staging_size_{0}; + + mutable aclOpExecutor* swiglu_exec_{nullptr}; + + mutable uint64_t swiglu_ws_{0}; + + mutable aclOpExecutor* copy_exec_{nullptr}; + + mutable uint64_t copy_ws_{0}; +}; + +} // namespace infini::ops + +#endif