diff --git a/include/infinicore/ops/causal_conv1d_ascend_vendor.hpp b/include/infinicore/ops/causal_conv1d_ascend_vendor.hpp new file mode 100644 index 000000000..99e9f6fe0 --- /dev/null +++ b/include/infinicore/ops/causal_conv1d_ascend_vendor.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include "infinicore.h" + +#include +#include +#include + +namespace infinicore::op { + +// Ascend vendor fused causal-conv bridge. The vendor kernel consumes the +// vLLM-Ascend layout directly: x/out [tokens, C], weight [K, C], and state +// [pool, K - 1, C]. It also fuses SiLU and state-cache updates. +INFINICORE_GRAPH_OP_CLASS( + CausalConv1dAscendVendor, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + std::optional, + std::vector, + std::vector, + bool, + bool); + +__export Tensor causal_conv1d_ascend_vendor( + const Tensor &x, + Tensor conv_state, + const Tensor &weight, + std::optional bias, + std::vector query_start_loc, + std::vector cache_indices, + bool fuse_silu, + bool decode); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/matmul_allreduce_add_rmsnorm_ascend.hpp b/include/infinicore/ops/matmul_allreduce_add_rmsnorm_ascend.hpp new file mode 100644 index 000000000..a021229d9 --- /dev/null +++ b/include/infinicore/ops/matmul_allreduce_add_rmsnorm_ascend.hpp @@ -0,0 +1,45 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include "infinicore.h" + +#include +#include + +namespace infinicore::op { + +// vLLM-Ascend vendor bridge: +// add_out = all_reduce(input @ weight^T) + residual +// normalized = rms_norm(add_out, gamma, epsilon) +INFINICORE_GRAPH_OP_CLASS( + MatmulAllReduceAddRmsNormAscend, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + infinicclComm_t, + float); + +__export std::tuple +matmul_allreduce_add_rmsnorm_ascend( + const Tensor &input, + const Tensor &weight, + const Tensor &residual, + const Tensor &gamma, + infinicclComm_t communicator, + float epsilon); + +// vLLM-Ascend vendor bridge: +// add_out = x1 + x2 +// normalized = rms_norm(add_out, gamma, epsilon) +// This directly backs RMSNorm::forward_inplace on Ascend. +__export std::tuple +add_rmsnorm_ascend_vendor( + const Tensor &x1, const Tensor &x2, + const Tensor &gamma, float epsilon); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/matmul_allreduce_ascend.hpp b/include/infinicore/ops/matmul_allreduce_ascend.hpp new file mode 100644 index 000000000..5b873a3af --- /dev/null +++ b/include/infinicore/ops/matmul_allreduce_ascend.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include "infinicore.h" + +#include + +namespace infinicore::op { + +// Ascend CANN MC2 bridge matching torch_npu.npu_mm_all_reduce_base: +// output = all_reduce(input @ weight_transposed). +INFINICORE_GRAPH_OP_CLASS( + MatmulAllReduceAscend, + Tensor, + const Tensor &, + const Tensor &, + infinicclComm_t); + +__export Tensor matmul_allreduce_ascend( + const Tensor &input, + const Tensor &weight_transposed, + infinicclComm_t communicator); + +} // namespace infinicore::op diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc index 32e5c5e9b..bd044fb7d 100644 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ b/src/infinicore/context/allocators/pinnable_block_allocator.cc @@ -47,6 +47,53 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) { std::shared_ptr block; + // The allocator normally keeps free blocks for fast reuse. On memory-constrained + // inference servers, a workload shape change can leave enough cached blocks in + // other size classes to make a new allocation fail. Reclaim only idle, + // non-graph blocks and retry once before surfacing OOM. + auto malloc_with_reclaim = [this](void **ptr, size_t bytes) -> infiniStatus_t { + auto status = infinirtMalloc(ptr, bytes); + if (status == INFINI_STATUS_SUCCESS && *ptr != nullptr) { + return status; + } + + size_t reclaimed_blocks = 0; + size_t reclaimed_bytes = 0; + for (auto &cls : size_classes_) { + for (auto it = cls.free_blocks.begin(); it != cls.free_blocks.end();) { + if (!(*it)->frozen && !(*it)->in_use) { + reclaimed_bytes += (*it)->size; + ++reclaimed_blocks; + INFINICORE_CHECK_ERROR(infinirtFree((*it)->ptr)); + all_blocks_.erase((*it)->ptr); + it = cls.free_blocks.erase(it); + } else { + ++it; + } + } + } + for (auto it = large_blocks_.begin(); it != large_blocks_.end();) { + if (!(*it)->frozen && !(*it)->in_use) { + reclaimed_bytes += (*it)->size; + ++reclaimed_blocks; + INFINICORE_CHECK_ERROR(infinirtFree((*it)->ptr)); + all_blocks_.erase((*it)->ptr); + it = large_blocks_.erase(it); + } else { + ++it; + } + } + if (reclaimed_blocks > 0) { + spdlog::warn("Device allocation of {} bytes failed; reclaimed {} idle blocks ({} bytes) and retrying", + bytes, reclaimed_blocks, reclaimed_bytes); + } + status = infinirtMalloc(ptr, bytes); + if (status == INFINI_STATUS_SUCCESS && *ptr == nullptr) { + return INFINI_STATUS_INTERNAL_ERROR; + } + return status; + }; + // 1. Try size-class allocation for small/medium for (auto &cls : size_classes_) { if (size <= cls.block_size) { @@ -74,7 +121,7 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) { block->in_use = true; block->use_count = 1; - INFINICORE_CHECK_ERROR(infinirtMalloc(&block->ptr, block->size)); + INFINICORE_CHECK_ERROR(malloc_with_reclaim(&block->ptr, block->size)); all_blocks_[block->ptr] = block; return reinterpret_cast(block->ptr); @@ -101,7 +148,7 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) { block->in_use = true; block->use_count = 1; - INFINICORE_CHECK_ERROR(infinirtMalloc(&block->ptr, block->size)); + INFINICORE_CHECK_ERROR(malloc_with_reclaim(&block->ptr, block->size)); large_blocks_.push_back(block); all_blocks_[block->ptr] = block; @@ -166,7 +213,7 @@ void PinnableBlockAllocator::trim() { // Free non-frozen size-class blocks for (auto &cls : size_classes_) { for (auto it = cls.free_blocks.begin(); it != cls.free_blocks.end();) { - if (!(*it)->frozen) { + if (!(*it)->frozen && !(*it)->in_use) { INFINICORE_CHECK_ERROR(infinirtFree((*it)->ptr)); all_blocks_.erase((*it)->ptr); it = cls.free_blocks.erase(it); diff --git a/src/infinicore/nn/rmsnorm.cc b/src/infinicore/nn/rmsnorm.cc index 567d88e49..1e0898102 100644 --- a/src/infinicore/nn/rmsnorm.cc +++ b/src/infinicore/nn/rmsnorm.cc @@ -1,7 +1,13 @@ #include "infinicore/nn/rmsnorm.hpp" #include "infinicore/ops.hpp" +#ifdef ENABLE_ASCEND_API +#include "infinicore/ops/matmul_allreduce_add_rmsnorm_ascend.hpp" +#endif #include +#include +#include #include +#include namespace infinicore::nn { @@ -26,6 +32,20 @@ void RMSNorm::forward_inplace(Tensor &x, Tensor &residual) const { residual = x; x = op::rms_norm(x, weight_, static_cast(eps_)); } else { +#ifdef ENABLE_ASCEND_API + static const bool ascend_vendor_enabled = []() { + const char *value = std::getenv( + "INFINICORE_ASCEND_ADD_RMSNORM_VENDOR"); + return value == nullptr || std::strcmp(value, "0") != 0; + }(); + if (device_.getType() == Device::Type::ASCEND + && ascend_vendor_enabled) { + std::tie(x, residual) = op::add_rmsnorm_ascend_vendor( + x, residual, weight_, + static_cast(eps_)); + return; + } +#endif if (device_.getType() == Device::Type::CPU || device_.getType() == Device::Type::NVIDIA || device_.getType() == Device::Type::ILUVATAR diff --git a/src/infinicore/ops/causal_conv1d/causal_conv1d_ascend_vendor.cc b/src/infinicore/ops/causal_conv1d/causal_conv1d_ascend_vendor.cc new file mode 100644 index 000000000..73fbc83c6 --- /dev/null +++ b/src/infinicore/ops/causal_conv1d/causal_conv1d_ascend_vendor.cc @@ -0,0 +1,430 @@ +#ifdef ENABLE_ASCEND_API + +#include "infinicore/ops/causal_conv1d_ascend_vendor.hpp" +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(CausalConv1dAscendVendor); + +namespace causal_conv1d_ascend_vendor_impl { + +using GetWorkspaceSizeFn = aclnnStatus (*)( + const aclTensor *, + const aclTensor *, + const aclTensor *, + const aclTensor *, + const aclIntArray *, + const aclIntArray *, + const aclIntArray *, + const aclIntArray *, + int64_t, + int64_t, + int64_t, + const aclTensor *, + uint64_t *, + aclOpExecutor **); + +using RunFn = aclnnStatus (*)( + void *, uint64_t, aclOpExecutor *, aclrtStream); + +struct VendorApi { + void *handle = nullptr; + GetWorkspaceSizeFn get_workspace_size = nullptr; + RunFn run = nullptr; + std::string path; +}; + +static void configure_custom_opp_path() { + const char *opp_override = std::getenv("INFINICORE_ASCEND_CUSTOM_OPP_PATH"); + const std::string opp_path = opp_override != nullptr + ? opp_override + : "/vllm-workspace/vllm-ascend/vllm_ascend/" + "_cann_ops_custom/vendors/vllm-ascend"; + const char *current = std::getenv("ASCEND_CUSTOM_OPP_PATH"); + const std::string current_path = current != nullptr ? current : ""; + if (current_path.find(opp_path) == std::string::npos) { + const std::string combined = current_path.empty() + ? opp_path + : opp_path + ":" + current_path; + setenv("ASCEND_CUSTOM_OPP_PATH", combined.c_str(), 1); + } +} + +static VendorApi load_vendor_api() { + configure_custom_opp_path(); + + std::vector candidates; + if (const char *override_path = std::getenv("INFINICORE_ASCEND_CAUSAL_CONV_VENDOR_SO")) { + candidates.emplace_back(override_path); + } + candidates.emplace_back( + "/vllm-workspace/vllm-ascend/vllm_ascend/_cann_ops_custom/" + "vendors/vllm-ascend/op_api/lib/libcust_opapi.so"); + candidates.emplace_back( + "/usr/local/lib/python3.11/site-packages/vllm_ascend/" + "_cann_ops_custom/vendors/vllm-ascend/op_api/lib/libcust_opapi.so"); + + std::string errors; + for (const auto &path : candidates) { + dlerror(); + void *handle = dlopen(path.c_str(), RTLD_NOW | RTLD_GLOBAL); + if (handle == nullptr) { + const char *error = dlerror(); + errors += "\n " + path + ": " + + (error != nullptr ? error : "unknown dlopen error"); + continue; + } + + auto get_workspace_size = reinterpret_cast( + dlsym(handle, "aclnnCausalConv1dGetWorkspaceSize")); + auto run = reinterpret_cast(dlsym(handle, "aclnnCausalConv1d")); + if (get_workspace_size != nullptr && run != nullptr) { + return {handle, get_workspace_size, run, path}; + } + + errors += "\n " + path + ": required symbols are missing"; + dlclose(handle); + } + + throw std::runtime_error( + "[causal_conv1d/ascend/vendor] unable to load vLLM-Ascend " + "libcust_opapi.so. Set INFINICORE_ASCEND_CAUSAL_CONV_VENDOR_SO " + "and INFINICORE_ASCEND_CUSTOM_OPP_PATH if installed elsewhere." + + errors); +} + +static VendorApi &vendor_api() { + static VendorApi api = load_vendor_api(); + return api; +} + +static aclDataType to_acl_dtype(DataType dtype) { + switch (dtype) { + case DataType::F16: + return ACL_FLOAT16; + case DataType::BF16: + return ACL_BF16; + case DataType::F32: + return ACL_FLOAT; + default: + throw std::runtime_error( + "[causal_conv1d/ascend/vendor] only fp16, bf16 and fp32 " + "are supported"); + } +} + +static aclTensor *make_contiguous_acl_tensor( + const Tensor &tensor, + const std::vector &dims) { + std::vector strides(dims.size(), 1); + for (size_t i = dims.size(); i-- > 1;) { + strides[i - 1] = strides[i] * dims[i]; + } + return aclCreateTensor( + dims.data(), + dims.size(), + to_acl_dtype(tensor->dtype()), + strides.data(), + 0, + ACL_FORMAT_ND, + dims.data(), + dims.size(), + const_cast( + reinterpret_cast(tensor->data()))); +} + +struct StreamWorkspace { + void *ptr = nullptr; + uint64_t capacity = 0; + std::vector retired; +}; + +static void *acquire_stream_workspace( + aclrtStream stream, uint64_t bytes) { + if (bytes == 0) { + return nullptr; + } + static auto *mutex = new std::mutex(); + static auto *workspaces = new std::unordered_map(); + std::lock_guard lock(*mutex); + auto &workspace = (*workspaces)[stream]; + if (workspace.capacity < bytes) { + void *new_ptr = nullptr; + auto ret = aclrtMalloc(&new_ptr, bytes, ACL_MEM_MALLOC_HUGE_FIRST); + if (ret != ACL_SUCCESS) { + throw std::runtime_error( + "[causal_conv1d/ascend/vendor] workspace allocation " + "failed: " + + std::to_string(ret)); + } + if (workspace.ptr != nullptr) { + workspace.retired.push_back(workspace.ptr); + } + workspace.ptr = new_ptr; + workspace.capacity = bytes; + } + return workspace.ptr; +} + +struct PlannedMeta { + graph::GraphTensor out; + graph::GraphTensor state; + graph::GraphTensor x; + graph::GraphTensor weight; + std::optional bias; + std::vector query_start_loc; + std::vector cache_indices; + bool fuse_silu; + bool decode; +}; + +void *plan( + Tensor out, + Tensor state, + const Tensor &x, + const Tensor &weight, + std::optional bias, + std::vector query_start_loc, + std::vector cache_indices, + bool fuse_silu, + bool decode) { + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(state), + graph::GraphTensor(x), + graph::GraphTensor(weight), + bias.has_value() + ? std::optional( + graph::GraphTensor(bias.value())) + : std::nullopt, + std::move(query_start_loc), + std::move(cache_indices), + fuse_silu, + decode}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + infinicore::context::setDevice(p->x->device()); + + const auto &x_shape = p->x->shape(); + const auto &weight_shape = p->weight->shape(); + const auto &state_shape = p->state->shape(); + if (x_shape.size() != 3 || x_shape[0] != 1 + || weight_shape.size() != 2 || state_shape.size() != 3) { + throw std::runtime_error( + "[causal_conv1d/ascend/vendor] expected x [1,T,C], " + "weight [K,C], state [pool,K-1,C]"); + } + const int64_t tokens = static_cast(x_shape[1]); + const int64_t channels = static_cast(x_shape[2]); + const int64_t kernel = static_cast(weight_shape[0]); + if (weight_shape[1] != x_shape[2] + || state_shape[1] + 1 != weight_shape[0] + || state_shape[2] != x_shape[2] + || p->query_start_loc.size() != p->cache_indices.size() + 1) { + throw std::runtime_error( + "[causal_conv1d/ascend/vendor] incompatible shapes or " + "metadata lengths"); + } + if (!p->x->is_contiguous() || !p->out->is_contiguous() + || !p->weight->is_contiguous() || !p->state->is_contiguous()) { + throw std::runtime_error( + "[causal_conv1d/ascend/vendor] all tensors must be contiguous"); + } + + Tensor x(p->x); + Tensor out(p->out); + Tensor weight(p->weight); + Tensor state(p->state); + auto *x_acl = make_contiguous_acl_tensor(x, {tokens, channels}); + auto *weight_acl = make_contiguous_acl_tensor(weight, {kernel, channels}); + auto *state_acl = make_contiguous_acl_tensor( + state, + {static_cast(state_shape[0]), + static_cast(state_shape[1]), + channels}); + auto *out_acl = make_contiguous_acl_tensor(out, {tokens, channels}); + aclTensor *bias_acl = nullptr; + if (p->bias.has_value()) { + Tensor bias(p->bias.value()); + bias_acl = make_contiguous_acl_tensor(bias, {channels}); + } + + aclIntArray *query_start_loc_acl = aclCreateIntArray( + p->query_start_loc.data(), p->query_start_loc.size()); + aclIntArray *cache_indices_acl = aclCreateIntArray( + p->cache_indices.data(), p->cache_indices.size()); + + // vLLM-Ascend only accepts initialStateMode in runMode=0 + // (prefill/FN). In decode mode it must be absent. + std::vector initial_state_mode( + p->cache_indices.size(), 0); + aclIntArray *initial_state_mode_acl = nullptr; + if (!p->decode) { + initial_state_mode_acl = aclCreateIntArray( + initial_state_mode.data(), initial_state_mode.size()); + } + + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + auto &api = vendor_api(); + auto ret = api.get_workspace_size( + x_acl, + weight_acl, + bias_acl, + state_acl, + query_start_loc_acl, + cache_indices_acl, + initial_state_mode_acl, + nullptr, + p->fuse_silu ? 1 : 0, + -1, + p->decode ? 1 : 0, + out_acl, + &workspace_size, + &executor); + if (ret != 0) { + const char *message = aclGetRecentErrMsg(); + throw std::runtime_error( + "[causal_conv1d/ascend/vendor] " + "aclnnCausalConv1dGetWorkspaceSize failed: " + + std::to_string(ret) + ", " + + (message != nullptr ? message : "(no ACL error)")); + } + + auto stream = reinterpret_cast(infinicore::context::getStream()); + void *workspace = acquire_stream_workspace(stream, workspace_size); + ret = api.run(workspace, workspace_size, executor, stream); + + aclDestroyTensor(x_acl); + aclDestroyTensor(weight_acl); + aclDestroyTensor(state_acl); + aclDestroyTensor(out_acl); + if (bias_acl != nullptr) { + aclDestroyTensor(bias_acl); + } + aclDestroyIntArray(query_start_loc_acl); + aclDestroyIntArray(cache_indices_acl); + if (initial_state_mode_acl != nullptr) { + aclDestroyIntArray(initial_state_mode_acl); + } + + if (ret != 0) { + const char *message = aclGetRecentErrMsg(); + throw std::runtime_error( + "[causal_conv1d/ascend/vendor] aclnnCausalConv1d failed: " + + std::to_string(ret) + ", " + + (message != nullptr ? message : "(no ACL error)")); + } +} + +void cleanup(void **planned_meta_ptr) { + auto *p = *reinterpret_cast(planned_meta_ptr); + delete p; + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + // This must happen while libinfinicore is loaded, before the first ACLNN + // operator initializes the custom OPP registry. + configure_custom_opp_path(); + CausalConv1dAscendVendor::plan_dispatcher().registerDevice( + Device::Type::ASCEND, &plan); + CausalConv1dAscendVendor::run_dispatcher().registerDevice( + Device::Type::ASCEND, &run); + CausalConv1dAscendVendor::cleanup_dispatcher().registerDevice( + Device::Type::ASCEND, &cleanup); + return true; +}(); + +} // namespace causal_conv1d_ascend_vendor_impl + +CausalConv1dAscendVendor::CausalConv1dAscendVendor( + Tensor out, + Tensor conv_state, + const Tensor &x, + const Tensor &weight, + std::optional bias, + std::vector query_start_loc, + std::vector cache_indices, + bool fuse_silu, + bool decode) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, conv_state, x, weight); + INFINICORE_GRAPH_OP_DISPATCH( + out->device().getType(), + out, + conv_state, + x, + weight, + bias, + std::move(query_start_loc), + std::move(cache_indices), + fuse_silu, + decode); +} + +void CausalConv1dAscendVendor::execute( + Tensor out, + Tensor conv_state, + const Tensor &x, + const Tensor &weight, + std::optional bias, + std::vector query_start_loc, + std::vector cache_indices, + bool fuse_silu, + bool decode) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + CausalConv1dAscendVendor, + out, + conv_state, + x, + weight, + bias, + std::move(query_start_loc), + std::move(cache_indices), + fuse_silu, + decode); +} + +Tensor causal_conv1d_ascend_vendor( + const Tensor &x, + Tensor conv_state, + const Tensor &weight, + std::optional bias, + std::vector query_start_loc, + std::vector cache_indices, + bool fuse_silu, + bool decode) { + auto out = Tensor::empty(x->shape(), x->dtype(), x->device()); + CausalConv1dAscendVendor::execute( + out, + conv_state, + x, + weight, + bias, + std::move(query_start_loc), + std::move(cache_indices), + fuse_silu, + decode); + return out; +} + +} // namespace infinicore::op + +#endif diff --git a/src/infinicore/ops/matmul_allreduce/matmul_allreduce_ascend.cc b/src/infinicore/ops/matmul_allreduce/matmul_allreduce_ascend.cc new file mode 100644 index 000000000..42b95c2f4 --- /dev/null +++ b/src/infinicore/ops/matmul_allreduce/matmul_allreduce_ascend.cc @@ -0,0 +1,303 @@ +#ifdef ENABLE_ASCEND_API + +#include "infinicore/ops/matmul_allreduce_ascend.hpp" +#include "../../../infiniccl/infiniccl_impl.h" +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" + +#include +#include +#include + +extern "C" int HcclGetCommName( + void *communicator, char *communicator_name); + +#include +#include +#include +#include +#include +#include +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MatmulAllReduceAscend); + +namespace matmul_allreduce_ascend_impl { + +static aclDataType to_acl_dtype(DataType dtype) { + switch (dtype) { + case DataType::F16: + return ACL_FLOAT16; + case DataType::BF16: + return ACL_BF16; + default: + throw std::runtime_error( + "[matmul_allreduce/ascend] only fp16 and bf16 are supported"); + } +} + +static aclTensor *make_acl_tensor_2d(const Tensor &tensor) { + const int64_t dims[2] = { + static_cast(tensor->shape()[0]), + static_cast(tensor->shape()[1])}; + const int64_t strides[2] = { + static_cast(tensor->strides()[0]), + static_cast(tensor->strides()[1])}; + return aclCreateTensor( + dims, + 2, + to_acl_dtype(tensor->dtype()), + strides, + 0, + ACL_FORMAT_ND, + dims, + 2, + const_cast( + reinterpret_cast(tensor->data()))); +} + +struct StreamWorkspace { + void *ptr = nullptr; + uint64_t capacity = 0; + std::vector retired; +}; + +static void *acquire_stream_workspace( + aclrtStream stream, uint64_t bytes) { + if (bytes == 0) { + return nullptr; + } + // Rank workers own their stream. Avoid a process-wide mutex on every + // fused row-parallel layer in the decode path. + thread_local std::unordered_map workspaces; + auto &workspace = workspaces[stream]; + if (workspace.capacity < bytes) { + void *new_ptr = nullptr; + auto ret = aclrtMalloc(&new_ptr, bytes, ACL_MEM_MALLOC_HUGE_FIRST); + if (ret != ACL_SUCCESS) { + throw std::runtime_error( + "[matmul_allreduce/ascend] workspace allocation failed: " + + std::to_string(ret)); + } + if (workspace.ptr != nullptr) { + // Calls on one rank stream are ordered. Keep old workspaces alive + // because freeing here would introduce a host synchronization. + workspace.retired.push_back(workspace.ptr); + } + workspace.ptr = new_ptr; + workspace.capacity = bytes; + } + return workspace.ptr; +} + +static const std::string &get_group_name( + infinicclComm_t communicator) { + if (communicator == nullptr || communicator->comm == nullptr) { + throw std::runtime_error( + "[matmul_allreduce/ascend] communicator is null"); + } + // The HCCL communicator name is immutable for the model lifetime. vLLM + // also resolves it once, rather than once per layer invocation. + thread_local std::unordered_map names; + auto found = names.find(communicator); + if (found != names.end()) { + return found->second; + } + char name[128] = {}; + auto ret = HcclGetCommName( + communicator->comm, name); + if (ret != 0) { + throw std::runtime_error( + "[matmul_allreduce/ascend] HcclGetCommName failed: " + + std::to_string(ret)); + } + return names.emplace(communicator, name).first->second; +} + +struct PlannedMeta { + graph::GraphTensor out; + graph::GraphTensor input; + graph::GraphTensor weight; + infinicclComm_t communicator; +}; + +void *plan( + Tensor out, + const Tensor &input, + const Tensor &weight_transposed, + infinicclComm_t communicator) { + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(input), + graph::GraphTensor(weight_transposed), + communicator}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + infinicore::context::setDevice(p->input->device()); + + if (p->input->ndim() != 2 || p->weight->ndim() != 2 + || p->out->ndim() != 2) { + throw std::runtime_error( + "[matmul_allreduce/ascend] expected 2D input, weight and output"); + } + if (p->input->shape()[1] != p->weight->shape()[0] + || p->out->shape()[0] != p->input->shape()[0] + || p->out->shape()[1] != p->weight->shape()[1]) { + throw std::runtime_error( + "[matmul_allreduce/ascend] incompatible matrix shapes"); + } + if (!p->input->is_contiguous() || !p->out->is_contiguous()) { + throw std::runtime_error( + "[matmul_allreduce/ascend] input and output must be contiguous"); + } + if (p->input->dtype() != p->weight->dtype() + || p->input->dtype() != p->out->dtype()) { + throw std::runtime_error( + "[matmul_allreduce/ascend] tensor dtypes must match"); + } + + Tensor input(p->input); + Tensor weight(p->weight); + Tensor out(p->out); + auto *input_acl = make_acl_tensor_2d(input); + auto *weight_acl = make_acl_tensor_2d(weight); + auto *out_acl = make_acl_tensor_2d(out); + if (input_acl == nullptr || weight_acl == nullptr || out_acl == nullptr) { + if (input_acl != nullptr) { + aclDestroyTensor(input_acl); + } + if (weight_acl != nullptr) { + aclDestroyTensor(weight_acl); + } + if (out_acl != nullptr) { + aclDestroyTensor(out_acl); + } + throw std::runtime_error( + "[matmul_allreduce/ascend] aclCreateTensor failed"); + } + + const std::string &group = get_group_name(p->communicator); + static const int64_t comm_turn = []() { + const char *value = std::getenv("INFINICORE_ASCEND_MATMUL_ALLREDUCE_COMM_TURN"); + return value == nullptr ? int64_t{0} + : std::strtoll(value, nullptr, 10); + }(); + static const int64_t stream_mode = []() { + const char *value = std::getenv("INFINICORE_ASCEND_MATMUL_ALLREDUCE_STREAM_MODE"); + return value == nullptr ? int64_t{1} + : std::strtoll(value, nullptr, 10); + }(); + + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + auto ret = aclnnMatmulAllReduceGetWorkspaceSize( + input_acl, + weight_acl, + nullptr, + group.c_str(), + "sum", + comm_turn, + stream_mode, + out_acl, + &workspace_size, + &executor); + if (ret != ACL_SUCCESS) { + const char *message = aclGetRecentErrMsg(); + aclDestroyTensor(input_acl); + aclDestroyTensor(weight_acl); + aclDestroyTensor(out_acl); + throw std::runtime_error( + "[matmul_allreduce/ascend] " + "aclnnMatmulAllReduceGetWorkspaceSize failed: " + + std::to_string(ret) + ", " + + (message != nullptr ? message : "(no ACL error)")); + } + + auto stream = reinterpret_cast(infinicore::context::getStream()); + void *workspace = acquire_stream_workspace(stream, workspace_size); + ret = aclnnMatmulAllReduce( + workspace, workspace_size, executor, stream); + + aclDestroyTensor(input_acl); + aclDestroyTensor(weight_acl); + aclDestroyTensor(out_acl); + + if (ret != ACL_SUCCESS) { + const char *message = aclGetRecentErrMsg(); + throw std::runtime_error( + "[matmul_allreduce/ascend] aclnnMatmulAllReduce failed: " + + std::to_string(ret) + ", " + + (message != nullptr ? message : "(no ACL error)")); + } +} + +void cleanup(void **planned_meta_ptr) { + auto *p = *reinterpret_cast(planned_meta_ptr); + delete p; + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MatmulAllReduceAscend::plan_dispatcher().registerDevice( + Device::Type::ASCEND, &plan); + MatmulAllReduceAscend::run_dispatcher().registerDevice( + Device::Type::ASCEND, &run); + MatmulAllReduceAscend::cleanup_dispatcher().registerDevice( + Device::Type::ASCEND, &cleanup); + return true; +}(); + +} // namespace matmul_allreduce_ascend_impl + +MatmulAllReduceAscend::MatmulAllReduceAscend( + Tensor out, + const Tensor &input, + const Tensor &weight_transposed, + infinicclComm_t communicator) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + out, input, weight_transposed); + INFINICORE_GRAPH_OP_DISPATCH( + out->device().getType(), + out, + input, + weight_transposed, + communicator); +} + +void MatmulAllReduceAscend::execute( + Tensor out, + const Tensor &input, + const Tensor &weight_transposed, + infinicclComm_t communicator) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MatmulAllReduceAscend, + out, + input, + weight_transposed, + communicator); +} + +Tensor matmul_allreduce_ascend( + const Tensor &input, + const Tensor &weight_transposed, + infinicclComm_t communicator) { + if (input->ndim() != 2 || weight_transposed->ndim() != 2) { + throw std::runtime_error( + "[matmul_allreduce/ascend] expected 2D matrices"); + } + auto out = Tensor::empty( + {input->shape()[0], weight_transposed->shape()[1]}, + input->dtype(), + input->device()); + MatmulAllReduceAscend::execute( + out, input, weight_transposed, communicator); + return out; +} + +} // namespace infinicore::op + +#endif diff --git a/src/infinicore/ops/matmul_allreduce_add_rmsnorm/matmul_allreduce_add_rmsnorm_ascend.cc b/src/infinicore/ops/matmul_allreduce_add_rmsnorm/matmul_allreduce_add_rmsnorm_ascend.cc new file mode 100644 index 000000000..a26937376 --- /dev/null +++ b/src/infinicore/ops/matmul_allreduce_add_rmsnorm/matmul_allreduce_add_rmsnorm_ascend.cc @@ -0,0 +1,608 @@ +#ifdef ENABLE_ASCEND_API + +#include "infinicore/ops/matmul_allreduce_add_rmsnorm_ascend.hpp" +#include "../../../infiniccl/infiniccl_impl.h" +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" int HcclGetCommName( + void *communicator, char *communicator_name); + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL( + MatmulAllReduceAddRmsNormAscend); + +namespace matmul_allreduce_add_rmsnorm_ascend_impl { + +using GetWorkspaceSizeFn = aclnnStatus (*)( + const aclTensor *, + const aclTensor *, + const aclTensor *, + const aclTensor *, + char *, + int64_t, + int64_t, + double, + bool, + bool, + const aclTensor *, + const aclTensor *, + uint64_t *, + aclOpExecutor **); + +using RunFn = aclnnStatus (*)( + void *, uint64_t, aclOpExecutor *, aclrtStream); + +using AddRmsNormGetWorkspaceSizeFn = aclnnStatus (*)( + const aclTensor *, + const aclTensor *, + const aclTensor *, + const aclTensor *, + double, + const aclTensor *, + const aclTensor *, + const aclTensor *, + uint64_t *, + aclOpExecutor **); + +struct VendorApi { + void *handle = nullptr; + GetWorkspaceSizeFn get_workspace_size = nullptr; + RunFn run = nullptr; + AddRmsNormGetWorkspaceSizeFn add_rmsnorm_get_workspace_size = nullptr; + RunFn add_rmsnorm_run = nullptr; + std::string path; +}; + +static void configure_custom_opp_path() { + const char *override_path = std::getenv("INFINICORE_ASCEND_CUSTOM_OPP_PATH"); + const std::string opp_path = override_path != nullptr + ? override_path + : "/vllm-workspace/vllm-ascend/vllm_ascend/" + "_cann_ops_custom/vendors/vllm-ascend"; + const char *current = std::getenv("ASCEND_CUSTOM_OPP_PATH"); + const std::string current_path = current != nullptr ? current : ""; + if (current_path.find(opp_path) == std::string::npos) { + const std::string combined = current_path.empty() + ? opp_path + : opp_path + ":" + current_path; + setenv("ASCEND_CUSTOM_OPP_PATH", combined.c_str(), 1); + } +} + +static VendorApi load_vendor_api() { + configure_custom_opp_path(); + std::vector candidates; + if (const char *override_path = std::getenv( + "INFINICORE_ASCEND_MC2_ADD_RMSNORM_VENDOR_SO")) { + candidates.emplace_back(override_path); + } + candidates.emplace_back( + "/vllm-workspace/vllm-ascend/vllm_ascend/_cann_ops_custom/" + "vendors/vllm-ascend/op_api/lib/libcust_opapi.so"); + candidates.emplace_back( + "/usr/local/lib/python3.11/site-packages/vllm_ascend/" + "_cann_ops_custom/vendors/vllm-ascend/op_api/lib/" + "libcust_opapi.so"); + + std::string errors; + for (const auto &path : candidates) { + dlerror(); + void *handle = dlopen(path.c_str(), RTLD_NOW | RTLD_GLOBAL); + if (handle == nullptr) { + const char *error = dlerror(); + errors += "\n " + path + ": " + + (error != nullptr ? error + : "unknown dlopen error"); + continue; + } + auto get_workspace_size = reinterpret_cast(dlsym( + handle, + "aclnnMatmulAllreduceAddRmsnormGetWorkspaceSize")); + auto run = reinterpret_cast(dlsym( + handle, "aclnnMatmulAllreduceAddRmsnorm")); + auto add_rmsnorm_get_workspace_size = reinterpret_cast(dlsym( + handle, "aclnnAddRmsNormBiasGetWorkspaceSize")); + auto add_rmsnorm_run = reinterpret_cast(dlsym( + handle, "aclnnAddRmsNormBias")); + if (get_workspace_size != nullptr && run != nullptr + && add_rmsnorm_get_workspace_size != nullptr + && add_rmsnorm_run != nullptr) { + return {handle, get_workspace_size, run, + add_rmsnorm_get_workspace_size, + add_rmsnorm_run, path}; + } + errors += "\n " + path + ": required symbols are missing"; + dlclose(handle); + } + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] unable to " + "load vLLM-Ascend libcust_opapi.so. Set " + "INFINICORE_ASCEND_MC2_ADD_RMSNORM_VENDOR_SO and " + "INFINICORE_ASCEND_CUSTOM_OPP_PATH if installed elsewhere." + + errors); +} + +static VendorApi &vendor_api() { + static VendorApi api = load_vendor_api(); + return api; +} + +static aclDataType to_acl_dtype(DataType dtype) { + switch (dtype) { + case DataType::F16: + return ACL_FLOAT16; + case DataType::BF16: + return ACL_BF16; + default: + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] only " + "fp16 and bf16 are supported"); + } +} + +static aclTensor *make_acl_tensor(const Tensor &tensor) { + const auto &shape = tensor->shape(); + const auto &strides = tensor->strides(); + std::vector dims(shape.begin(), shape.end()); + std::vector acl_strides( + strides.begin(), strides.end()); + return aclCreateTensor( + dims.data(), + dims.size(), + to_acl_dtype(tensor->dtype()), + acl_strides.data(), + 0, + ACL_FORMAT_ND, + dims.data(), + dims.size(), + const_cast(reinterpret_cast( + tensor->data()))); +} + +struct StreamWorkspace { + void *ptr = nullptr; + uint64_t capacity = 0; + std::vector retired; +}; + +static void *acquire_stream_workspace( + aclrtStream stream, uint64_t bytes) { + if (bytes == 0) { + return nullptr; + } + thread_local std::unordered_map< + aclrtStream, StreamWorkspace> + workspaces; + auto &workspace = workspaces[stream]; + if (workspace.capacity < bytes) { + void *new_ptr = nullptr; + auto ret = aclrtMalloc( + &new_ptr, bytes, ACL_MEM_MALLOC_HUGE_FIRST); + if (ret != ACL_SUCCESS) { + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] " + "workspace allocation failed: " + + std::to_string(ret)); + } + if (workspace.ptr != nullptr) { + workspace.retired.push_back(workspace.ptr); + } + workspace.ptr = new_ptr; + workspace.capacity = bytes; + } + return workspace.ptr; +} + +static void *acquire_stream_rstd_buffer( + aclrtStream stream, uint64_t bytes) { + if (bytes == 0) { + return nullptr; + } + // The vendor op writes rstd asynchronously. Keep a permanent buffer per + // stream instead of returning a temporary Tensor to the allocator before + // the stream has consumed it. + thread_local std::unordered_map< + aclrtStream, StreamWorkspace> + buffers; + auto &buffer = buffers[stream]; + if (buffer.capacity < bytes) { + void *new_ptr = nullptr; + auto ret = aclrtMalloc( + &new_ptr, bytes, ACL_MEM_MALLOC_HUGE_FIRST); + if (ret != ACL_SUCCESS) { + throw std::runtime_error( + "[add_rmsnorm/ascend/vendor] rstd buffer " + "allocation failed: " + + std::to_string(ret)); + } + if (buffer.ptr != nullptr) { + // Do not free an older generation while work on this stream may + // still reference it. Shape growth is rare and bounded. + buffer.retired.push_back(buffer.ptr); + } + buffer.ptr = new_ptr; + buffer.capacity = bytes; + } + return buffer.ptr; +} + +static const std::string &get_group_name( + infinicclComm_t communicator) { + if (communicator == nullptr || communicator->comm == nullptr) { + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] " + "communicator is null"); + } + thread_local std::unordered_map< + infinicclComm_t, std::string> + names; + auto found = names.find(communicator); + if (found != names.end()) { + return found->second; + } + char name[128] = {}; + auto ret = HcclGetCommName(communicator->comm, name); + if (ret != 0) { + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] " + "HcclGetCommName failed: " + + std::to_string(ret)); + } + return names.emplace(communicator, name).first->second; +} + +struct PlannedMeta { + graph::GraphTensor normalized; + graph::GraphTensor add_out; + graph::GraphTensor input; + graph::GraphTensor weight; + graph::GraphTensor residual; + graph::GraphTensor gamma; + infinicclComm_t communicator; + float epsilon; +}; + +void *plan( + Tensor normalized, + Tensor add_out, + const Tensor &input, + const Tensor &weight, + const Tensor &residual, + const Tensor &gamma, + infinicclComm_t communicator, + float epsilon) { + return new PlannedMeta{ + graph::GraphTensor(normalized), + graph::GraphTensor(add_out), + graph::GraphTensor(input), + graph::GraphTensor(weight), + graph::GraphTensor(residual), + graph::GraphTensor(gamma), + communicator, + epsilon}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + infinicore::context::setDevice(p->input->device()); + + if (p->input->ndim() != 2 || p->weight->ndim() != 2 + || p->residual->ndim() != 2 || p->gamma->ndim() != 1 + || p->normalized->ndim() != 2 + || p->add_out->ndim() != 2) { + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] " + "expected 2D matrices and 1D gamma"); + } + const auto rows = p->input->shape()[0]; + const auto out_features = p->weight->shape()[0]; + if (p->input->shape()[1] != p->weight->shape()[1] + || p->residual->shape()[0] != rows + || p->residual->shape()[1] != out_features + || p->gamma->shape()[0] != out_features + || p->normalized->shape() != p->residual->shape() + || p->add_out->shape() != p->residual->shape()) { + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] " + "incompatible tensor shapes"); + } + if (!p->input->is_contiguous() || !p->weight->is_contiguous() + || !p->residual->is_contiguous() + || !p->gamma->is_contiguous() + || !p->normalized->is_contiguous() + || !p->add_out->is_contiguous()) { + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] all " + "tensors must be contiguous"); + } + + Tensor input(p->input); + Tensor weight(p->weight); + Tensor residual(p->residual); + Tensor gamma(p->gamma); + Tensor normalized(p->normalized); + Tensor add_out(p->add_out); + aclTensor *input_acl = make_acl_tensor(input); + aclTensor *weight_acl = make_acl_tensor(weight); + aclTensor *residual_acl = make_acl_tensor(residual); + aclTensor *gamma_acl = make_acl_tensor(gamma); + aclTensor *normalized_acl = make_acl_tensor(normalized); + aclTensor *add_out_acl = make_acl_tensor(add_out); + std::vector tensors{ + input_acl, weight_acl, residual_acl, gamma_acl, + normalized_acl, add_out_acl}; + for (auto *tensor : tensors) { + if (tensor == nullptr) { + for (auto *created : tensors) { + if (created != nullptr) { + aclDestroyTensor(created); + } + } + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] " + "aclCreateTensor failed"); + } + } + + const std::string &group = get_group_name(p->communicator); + std::vector mutable_group(group.begin(), group.end()); + mutable_group.push_back('\0'); + auto &api = vendor_api(); + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + auto ret = api.get_workspace_size( + input_acl, + weight_acl, + residual_acl, + gamma_acl, + mutable_group.data(), + p->communicator->world_size, + p->communicator->rank, + static_cast(p->epsilon), + true, + false, + normalized_acl, + add_out_acl, + &workspace_size, + &executor); + if (ret == ACL_SUCCESS) { + auto stream = reinterpret_cast( + infinicore::context::getStream()); + void *workspace = acquire_stream_workspace(stream, workspace_size); + ret = api.run( + workspace, workspace_size, executor, stream); + } + + for (auto *tensor : tensors) { + aclDestroyTensor(tensor); + } + if (ret != ACL_SUCCESS) { + const char *message = aclGetRecentErrMsg(); + throw std::runtime_error( + "[matmul_allreduce_add_rmsnorm/ascend/vendor] call " + "failed: " + + std::to_string(ret) + ", " + + (message != nullptr ? message : "(no ACL error)")); + } +} + +void cleanup(void **planned_meta_ptr) { + auto *p = *reinterpret_cast( + planned_meta_ptr); + delete p; + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MatmulAllReduceAddRmsNormAscend::plan_dispatcher() + .registerDevice(Device::Type::ASCEND, &plan); + MatmulAllReduceAddRmsNormAscend::run_dispatcher() + .registerDevice(Device::Type::ASCEND, &run); + MatmulAllReduceAddRmsNormAscend::cleanup_dispatcher() + .registerDevice(Device::Type::ASCEND, &cleanup); + return true; +}(); + +} // namespace matmul_allreduce_add_rmsnorm_ascend_impl + +MatmulAllReduceAddRmsNormAscend:: + MatmulAllReduceAddRmsNormAscend( + Tensor normalized, + Tensor add_out, + const Tensor &input, + const Tensor &weight, + const Tensor &residual, + const Tensor &gamma, + infinicclComm_t communicator, + float epsilon) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + normalized, add_out, input, weight, residual, gamma); + INFINICORE_GRAPH_OP_DISPATCH( + normalized->device().getType(), + normalized, + add_out, + input, + weight, + residual, + gamma, + communicator, + epsilon); +} + +void MatmulAllReduceAddRmsNormAscend::execute( + Tensor normalized, + Tensor add_out, + const Tensor &input, + const Tensor &weight, + const Tensor &residual, + const Tensor &gamma, + infinicclComm_t communicator, + float epsilon) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MatmulAllReduceAddRmsNormAscend, + normalized, + add_out, + input, + weight, + residual, + gamma, + communicator, + epsilon); +} + +std::tuple +matmul_allreduce_add_rmsnorm_ascend( + const Tensor &input, + const Tensor &weight, + const Tensor &residual, + const Tensor &gamma, + infinicclComm_t communicator, + float epsilon) { + auto normalized = Tensor::empty( + residual->shape(), residual->dtype(), residual->device()); + auto add_out = Tensor::empty( + residual->shape(), residual->dtype(), residual->device()); + MatmulAllReduceAddRmsNormAscend::execute( + normalized, + add_out, + input, + weight, + residual, + gamma, + communicator, + epsilon); + return {normalized, add_out}; +} + +std::tuple +add_rmsnorm_ascend_vendor( + const Tensor &x1, + const Tensor &x2, + const Tensor &gamma, + float epsilon) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x1, x2, gamma); + infinicore::context::setDevice(x1->device()); + if (x1->shape() != x2->shape() || x1->ndim() < 2 + || gamma->ndim() != 1 + || gamma->shape()[0] != x1->shape().back()) { + throw std::runtime_error( + "[add_rmsnorm/ascend/vendor] incompatible tensor shapes"); + } + if (x1->dtype() != x2->dtype() + || x1->dtype() != gamma->dtype()) { + throw std::runtime_error( + "[add_rmsnorm/ascend/vendor] tensor dtypes must match"); + } + if (!x1->is_contiguous() || !x2->is_contiguous() + || !gamma->is_contiguous()) { + throw std::runtime_error( + "[add_rmsnorm/ascend/vendor] all tensors must be contiguous"); + } + + auto normalized = Tensor::empty( + x1->shape(), x1->dtype(), x1->device()); + auto add_out = Tensor::empty( + x1->shape(), x1->dtype(), x1->device()); + aclTensor *x1_acl = matmul_allreduce_add_rmsnorm_ascend_impl::make_acl_tensor(x1); + aclTensor *x2_acl = matmul_allreduce_add_rmsnorm_ascend_impl::make_acl_tensor(x2); + aclTensor *gamma_acl = matmul_allreduce_add_rmsnorm_ascend_impl::make_acl_tensor(gamma); + aclTensor *normalized_acl = matmul_allreduce_add_rmsnorm_ascend_impl::make_acl_tensor(normalized); + aclTensor *add_out_acl = matmul_allreduce_add_rmsnorm_ascend_impl::make_acl_tensor(add_out); + + std::vector rstd_dims; + rstd_dims.reserve(x1->ndim()); + for (size_t i = 0; i < x1->ndim(); ++i) { + rstd_dims.push_back( + static_cast(x1->shape()[i])); + } + rstd_dims.back() = 1; + std::vector rstd_strides(rstd_dims.size(), 1); + for (ptrdiff_t i = static_cast(rstd_dims.size()) - 2; + i >= 0; --i) { + rstd_strides[i] = rstd_strides[i + 1] * rstd_dims[i + 1]; + } + uint64_t rstd_numel = 1; + for (auto dim : rstd_dims) { + rstd_numel *= static_cast(dim); + } + auto stream = reinterpret_cast( + infinicore::context::getStream()); + void *rstd_ptr = matmul_allreduce_add_rmsnorm_ascend_impl::acquire_stream_rstd_buffer( + stream, rstd_numel * sizeof(float)); + aclTensor *rstd_acl = aclCreateTensor( + rstd_dims.data(), + rstd_dims.size(), + ACL_FLOAT, + rstd_strides.data(), + 0, + ACL_FORMAT_ND, + rstd_dims.data(), + rstd_dims.size(), + rstd_ptr); + + std::vector tensors{ + x1_acl, x2_acl, gamma_acl, normalized_acl, + rstd_acl, add_out_acl}; + for (auto *tensor : tensors) { + if (tensor == nullptr) { + for (auto *created : tensors) { + if (created != nullptr) { + aclDestroyTensor(created); + } + } + throw std::runtime_error( + "[add_rmsnorm/ascend/vendor] aclCreateTensor failed"); + } + } + + auto &api = matmul_allreduce_add_rmsnorm_ascend_impl::vendor_api(); + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + auto ret = api.add_rmsnorm_get_workspace_size( + x1_acl, + x2_acl, + gamma_acl, + nullptr, + static_cast(epsilon), + normalized_acl, + rstd_acl, + add_out_acl, + &workspace_size, + &executor); + if (ret == ACL_SUCCESS) { + void *workspace = matmul_allreduce_add_rmsnorm_ascend_impl:: + acquire_stream_workspace(stream, workspace_size); + ret = api.add_rmsnorm_run( + workspace, workspace_size, executor, stream); + } + + for (auto *tensor : tensors) { + aclDestroyTensor(tensor); + } + if (ret != ACL_SUCCESS) { + const char *message = aclGetRecentErrMsg(); + throw std::runtime_error( + "[add_rmsnorm/ascend/vendor] call failed: " + + std::to_string(ret) + ", " + + (message != nullptr ? message : "(no ACL error)")); + } + return {normalized, add_out}; +} + +} // namespace infinicore::op + +#endif diff --git a/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc b/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc index ead1ca4e2..282f3c82a 100644 --- a/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc +++ b/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc @@ -2,6 +2,7 @@ #include "infinicore/context/context.hpp" #include "infinicore/ops/mha_kvcache.hpp" +#include "infinicore/ops/paged_attention.hpp" #include "native/ascend/workspace_pool_.h" #include @@ -9,8 +10,10 @@ #include #include +#include #include #include +#include #include namespace infinicore::op::mha_kvcache_impl::flashattn_ascend { @@ -38,6 +41,39 @@ host_vector_to_acl_int_array(const std::vector &vec) { return aclCreateIntArray(vec.data(), vec.size()); } +// FIA workspace is scratch storage. Calls submitted to the same ACL stream are +// ordered, so one process-lifetime buffer per stream can be reused without a +// host synchronization. Buffers are intentionally retained until process exit; +// the serving process owns only a small number of streams. +struct StreamWorkspace { + void *ptr = nullptr; + uint64_t capacity = 0; + std::vector retired; +}; + +static void *acquire_stream_workspace(aclrtStream stream, uint64_t bytes) { + if (bytes == 0) { + return nullptr; + } + thread_local std::unordered_map workspaces; + auto &workspace = workspaces[stream]; + if (workspace.capacity < bytes) { + void *new_ptr = nullptr; + auto ret = aclrtMalloc(&new_ptr, bytes, ACL_MEM_MALLOC_HUGE_FIRST); + if (ret != ACL_SUCCESS) { + throw std::runtime_error( + std::string("[mha_kvcache/ascend] cached workspace allocation failed: ") + + std::to_string(ret)); + } + if (workspace.ptr) { + workspace.retired.push_back(workspace.ptr); + } + workspace.ptr = new_ptr; + workspace.capacity = bytes; + } + return workspace.ptr; +} + struct PlannedMeta { graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; std::optional alibi_slopes; @@ -72,9 +108,9 @@ void run(void *planned_meta) { // q/out are BSND [batch, 1, num_heads, head_size] in InfiniCore. For // decode S=1, the same memory can be described to FIA as BNSD // [batch, num_heads, 1, head_size]. - auto q_shape = p->q->shape(); - auto k_shape = p->k_cache->shape(); - auto v_shape = p->v_cache->shape(); + const auto q_shape = p->q->shape(); + const auto k_shape = p->k_cache->shape(); + const auto v_shape = p->v_cache->shape(); if (q_shape.size() != 4 || k_shape.size() != 4 || v_shape.size() != 4) { throw std::runtime_error("[mha_kvcache/ascend] flash attention expects q " @@ -84,9 +120,10 @@ void run(void *planned_meta) { const int64_t batch_size = q_shape[0]; const int64_t num_heads = q_shape[2]; const int64_t head_size = q_shape[3]; + // Ascend paged KV cache is physical BnNBsD. const int64_t num_blocks = k_shape[0]; - const int64_t block_size_val = k_shape[1]; - const int64_t num_kv_heads = k_shape[2]; + const int64_t num_kv_heads = k_shape[1]; + const int64_t block_size_val = k_shape[2]; const int64_t v_head_size = v_shape[3]; if (k_shape[3] != static_cast(head_size)) { @@ -108,18 +145,78 @@ void run(void *planned_meta) { : p->block_table->contiguous(); Tensor out_work = p->out->is_contiguous() ? Tensor(p->out) : p->out->contiguous(); - aclDataType q_dtype = to_acl_dtype(q_work->dtype()); + // FIA V4 supports Qwen3.5's D=256 paged decode with BNSD query and + // contiguous BnBsH cache. Keep the hand-written path as an emergency + // diagnostic fallback only. + const char *disable_fia_head256 = std::getenv("INFINICORE_ASCEND_DISABLE_FIA_HEAD256"); + const bool use_paged_fallback = disable_fia_head256 != nullptr + && std::strcmp(disable_fia_head256, "1") == 0; + if (head_size == 256 && use_paged_fallback) { + auto q_strides = q_work->strides(); + auto k_strides = k_work->strides(); + auto v_strides = v_work->strides(); + auto out_strides = out_work->strides(); + + Tensor q_view = Tensor::strided_from_blob( + const_cast(reinterpret_cast(q_work->data())), + {static_cast(batch_size), + static_cast(num_heads), + static_cast(head_size)}, + {q_strides[0], q_strides[2], q_strides[3]}, + q_work->dtype(), q_work->device()); + Tensor out_view = Tensor::strided_from_blob( + const_cast(reinterpret_cast(out_work->data())), + {static_cast(batch_size), + static_cast(num_heads), + static_cast(v_head_size)}, + {out_strides[0], out_strides[2], out_strides[3]}, + out_work->dtype(), out_work->device()); + Tensor k_view = Tensor::strided_from_blob( + const_cast(reinterpret_cast(k_work->data())), + {static_cast(num_blocks), + static_cast(num_kv_heads), + static_cast(block_size_val), + static_cast(head_size)}, + {k_strides[0], k_strides[1], k_strides[2], k_strides[3]}, + k_work->dtype(), k_work->device()); + Tensor v_view = Tensor::strided_from_blob( + const_cast(reinterpret_cast(v_work->data())), + {static_cast(num_blocks), + static_cast(num_kv_heads), + static_cast(block_size_val), + static_cast(v_head_size)}, + {v_strides[0], v_strides[1], v_strides[2], v_strides[3]}, + v_work->dtype(), v_work->device()); + + Tensor seqlens_k(p->seqlens_k); + paged_attention_( + out_view, q_view, k_view, v_view, bt_work, seqlens_k, + p->alibi_slopes, p->scale); + if (!p->out->is_contiguous()) { + p->out->copy_from(out_work); + } + return; + } - // Read seqlens_k to host + aclDataType q_dtype = to_acl_dtype(q_work->dtype()); auto seqlens_k_shape = p->seqlens_k->shape(); int64_t seqlens_k_len = seqlens_k_shape[0]; std::vector seqlens_k_host(seqlens_k_len); - auto copy_ret = aclrtMemcpy(seqlens_k_host.data(), seqlens_k_len * sizeof(int32_t), - reinterpret_cast(p->seqlens_k->data()), - seqlens_k_len * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); - if (copy_ret != ACL_SUCCESS) { - throw std::runtime_error( - std::string("[mha_kvcache/ascend] copy seqlens_k to host failed: ") + std::to_string(copy_ret)); + + const bool seqlens_on_host = p->seqlens_k->device().getType() == Device::Type::CPU; + if (seqlens_on_host) { + std::memcpy(seqlens_k_host.data(), p->seqlens_k->data(), + seqlens_k_len * sizeof(int32_t)); + } else { + auto copy_ret = aclrtMemcpy( + seqlens_k_host.data(), seqlens_k_len * sizeof(int32_t), + reinterpret_cast(p->seqlens_k->data()), + seqlens_k_len * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); + if (copy_ret != ACL_SUCCESS) { + throw std::runtime_error( + std::string("[mha_kvcache/ascend] copy seqlens_k to host failed: ") + + std::to_string(copy_ret)); + } } // Build actual_seq vectors @@ -140,22 +237,24 @@ void run(void *planned_meta) { q_dims.data(), q_dims.size(), const_cast(reinterpret_cast(q_work->data()))); - // The physical BnBsND cache is contiguous in N and D, so expose it to FIA - // as BnBsH without copying. - std::vector k_dims = {num_blocks, block_size_val, - num_kv_heads * head_size}; - std::vector k_strides = {block_size_val * num_kv_heads * head_size, - num_kv_heads * head_size, 1}; + // Physical BnNBsD avoids flattening N into H and is the preferred Ascend + // paged-attention cache layout. + std::vector k_dims = { + num_blocks, num_kv_heads, block_size_val, head_size}; + std::vector k_strides = { + num_kv_heads * block_size_val * head_size, + block_size_val * head_size, head_size, 1}; aclTensor *k_acl_tensor = aclCreateTensor( k_dims.data(), k_dims.size(), q_dtype, k_strides.data(), 0, ACL_FORMAT_ND, k_dims.data(), k_dims.size(), const_cast(reinterpret_cast(k_work->data()))); aclTensorList *key_acl = aclCreateTensorList(&k_acl_tensor, 1); - std::vector v_dims = {num_blocks, block_size_val, - num_kv_heads * v_head_size}; - std::vector v_strides = {block_size_val * num_kv_heads * v_head_size, - num_kv_heads * v_head_size, 1}; + std::vector v_dims = { + num_blocks, num_kv_heads, block_size_val, v_head_size}; + std::vector v_strides = { + num_kv_heads * block_size_val * v_head_size, + block_size_val * v_head_size, v_head_size, 1}; aclTensor *v_acl_tensor = aclCreateTensor( v_dims.data(), v_dims.size(), q_dtype, v_strides.data(), 0, ACL_FORMAT_ND, v_dims.data(), v_dims.size(), @@ -184,8 +283,18 @@ void run(void *planned_meta) { aclIntArray *seqlens_q_acl = host_vector_to_acl_int_array(actual_seq_q_vec); aclIntArray *seqlens_k_acl = host_vector_to_acl_int_array(actual_seq_k_vec); + aclrtStream stream = static_cast(infinicore::context::getStream()); + static const bool async_all_batch = []() { + const char *value = std::getenv("INFINICORE_ASCEND_FIA_ASYNC_ALL_BATCH"); + return value == nullptr || std::strcmp(value, "0") != 0; + }(); + const bool async_fast_path = seqlens_on_host && (async_all_batch || batch_size <= 2) + && p->q->is_contiguous() + && p->k_cache->is_contiguous() && p->v_cache->is_contiguous() + && p->block_table->is_contiguous() && p->out->is_contiguous(); + // Call CANN API with Paged Attention - // inputLayout="BNSD": query/out=[B,N,S,D], KV cache is BnBsH for paged + // inputLayout="BNSD": query/out=[B,N,S,D], KV cache is BnNBsD for paged // attention. sparse_mode=0: no mask needed for decode (Q_S=1, // IncreFlashAttention branch) uint64_t workspace_size = 0; @@ -240,7 +349,6 @@ void run(void *planned_meta) { + std::to_string(ret) + ", msg: " + (err_msg ? err_msg : "(null)")); } - aclrtStream stream = static_cast(infinicore::context::getStream()); void *workspace = nullptr; if (workspace_size > 0) { workspace = infini::ops::ascend::GetWorkspacePool() diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache.cc index 0c5b3ae8c..0a0dfe90d 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache.cc @@ -13,7 +13,9 @@ MhaKVCache::MhaKVCache(Tensor out, const Tensor &block_table, std::optional alibi_slopes, float scale) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, seqlens_k, block_table); + // Ascend may consume seqlens_k directly from the host to avoid a + // device-to-host synchronization in every attention layer. + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_table); INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); } diff --git a/src/infinicore/ops/multi_head_attention_varlen/ascend/mha_varlen_flashattn_ascend.cc b/src/infinicore/ops/multi_head_attention_varlen/ascend/mha_varlen_flashattn_ascend.cc index 27614a69e..1ffd77cd5 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/ascend/mha_varlen_flashattn_ascend.cc +++ b/src/infinicore/ops/multi_head_attention_varlen/ascend/mha_varlen_flashattn_ascend.cc @@ -2,6 +2,7 @@ #include "infinicore/context/context.hpp" #include "infinicore/ops/mha_varlen.hpp" +#include "infinicore/ops/paged_attention_prefill.hpp" #include "native/ascend/workspace_pool_.h" #include @@ -210,43 +211,90 @@ void run(void *planned_meta) { auto k_shape = p->k->shape(); const int64_t num_heads = q_shape[1]; const int64_t head_size = q_shape[2]; - const int64_t block_size_val = k_shape[1]; - const int64_t num_kv_heads = k_shape[2]; + // Ascend paged KV cache is physical BnNBsD. const int64_t num_blocks = k_shape[0]; + const int64_t num_kv_heads = k_shape[1]; + const int64_t block_size_val = k_shape[2]; Tensor q_work = p->q->is_contiguous() ? Tensor(p->q) : p->q->contiguous(); Tensor k_work = p->k->is_contiguous() ? Tensor(p->k) : p->k->contiguous(); Tensor v_work = p->v->is_contiguous() ? Tensor(p->v) : p->v->contiguous(); Tensor out_work = p->out->is_contiguous() ? Tensor(p->out) : p->out->contiguous(); + aclrtStream stream = static_cast(infinicore::context::getStream()); + const bool use_bnsd = head_size == 256; + int64_t max_q_len = static_cast(q_shape[0]); + Tensor q_acl_storage = q_work; + Tensor out_acl_storage = out_work; + std::vector actual_seq_q_for_acl = actual_seq_q_vec; + + if (use_bnsd) { + max_q_len = 0; + actual_seq_q_for_acl.clear(); + for (int64_t i = 0; i < batch_size; ++i) { + int64_t q_len = cu_q_host[i + 1] - cu_q_host[i]; + actual_seq_q_for_acl.push_back(q_len); + max_q_len = std::max(max_q_len, q_len); + } + + q_acl_storage = Tensor::empty( + {static_cast(batch_size), static_cast(num_heads), + static_cast(max_q_len), static_cast(head_size)}, + q_work->dtype(), q_work->device()); + out_acl_storage = Tensor::empty( + {static_cast(batch_size), static_cast(num_heads), + static_cast(max_q_len), static_cast(head_size)}, + out_work->dtype(), out_work->device()); + + // Pack ragged token-major TND directly into physical BNSD. The + // metadata views are free; copy_from performs one strided device copy + // per request without an intermediate BSND allocation. + for (int64_t i = 0; i < batch_size; ++i) { + const size_t q_begin = static_cast(cu_q_host[i]); + const size_t q_len = static_cast(actual_seq_q_for_acl[i]); + Tensor src = q_work->narrow({{0, q_begin, q_len}}) + ->permute({1, 0, 2}); + Tensor dst = q_acl_storage + ->narrow({{0, static_cast(i), 1}, + {2, 0, q_len}}) + ->squeeze(0); + dst->copy_from(src); + } + } + aclDataType q_dtype = to_acl_dtype(q_work->dtype()); - // Query is already contiguous TND [total_q, num_heads, head_size]. - std::vector q_dims = {static_cast(q_shape[0]), num_heads, - head_size}; - std::vector q_strides = {num_heads * head_size, head_size, 1}; + // TND rejects head_dim=256 on the installed CANN. Pack query once into + // physical BNSD and request BSND output so the consumer keeps its native + // token-major layout. + std::vector q_dims = use_bnsd + ? std::vector{batch_size, num_heads, max_q_len, head_size} + : std::vector{static_cast(q_shape[0]), num_heads, head_size}; + std::vector q_strides = use_bnsd + ? std::vector{num_heads * max_q_len * head_size, + max_q_len * head_size, head_size, 1} + : std::vector{num_heads * head_size, head_size, 1}; aclTensor *query_acl = aclCreateTensor( q_dims.data(), q_dims.size(), q_dtype, q_strides.data(), 0, ACL_FORMAT_ND, q_dims.data(), q_dims.size(), - const_cast(reinterpret_cast(q_work->data()))); - - // The physical BnBsND cache is contiguous in N and D, so expose it to FIA - // as BnBsH without copying. - std::vector k_dims = {num_blocks, block_size_val, - num_kv_heads * head_size}; - std::vector k_strides = {num_kv_heads * block_size_val * head_size, - num_kv_heads * head_size, 1}; + const_cast(reinterpret_cast(q_acl_storage->data()))); + + std::vector k_dims = { + num_blocks, num_kv_heads, block_size_val, head_size}; + std::vector k_strides = { + num_kv_heads * block_size_val * head_size, + block_size_val * head_size, head_size, 1}; aclTensor *k_acl_tensor = aclCreateTensor( k_dims.data(), k_dims.size(), q_dtype, k_strides.data(), 0, ACL_FORMAT_ND, k_dims.data(), k_dims.size(), const_cast(reinterpret_cast(k_work->data()))); aclTensorList *key_acl = aclCreateTensorList(&k_acl_tensor, 1); - // Value uses the same contiguous BnBsH representation. - std::vector v_dims = {num_blocks, block_size_val, - num_kv_heads * head_size}; - std::vector v_strides = {num_kv_heads * block_size_val * head_size, - num_kv_heads * head_size, 1}; + std::vector v_dims = { + num_blocks, num_kv_heads, block_size_val, head_size}; + std::vector v_strides = { + num_kv_heads * block_size_val * head_size, + block_size_val * head_size, head_size, 1}; aclTensor *v_acl_tensor = aclCreateTensor( v_dims.data(), v_dims.size(), q_dtype, v_strides.data(), 0, ACL_FORMAT_ND, v_dims.data(), v_dims.size(), @@ -267,22 +315,29 @@ void run(void *planned_meta) { const_cast(reinterpret_cast(bt_work->data()))); } - // FIA writes directly to contiguous TND output. + // FIA writes directly to the same contiguous output storage. auto out_shape = out_work->shape(); - std::vector out_dims = {static_cast(out_shape[0]), - num_heads, head_size}; - std::vector out_strides = {num_heads * head_size, head_size, 1}; + std::vector out_dims = use_bnsd + ? std::vector{batch_size, num_heads, max_q_len, head_size} + : std::vector{static_cast(out_shape[0]), num_heads, head_size}; + std::vector out_strides = use_bnsd + ? std::vector{num_heads * max_q_len * head_size, + max_q_len * head_size, head_size, 1} + : std::vector{num_heads * head_size, head_size, 1}; aclDataType out_dtype = to_acl_dtype(out_work->dtype()); aclTensor *out_acl = aclCreateTensor( out_dims.data(), out_dims.size(), out_dtype, out_strides.data(), 0, ACL_FORMAT_ND, out_dims.data(), out_dims.size(), - const_cast(reinterpret_cast(out_work->data()))); + const_cast(reinterpret_cast(out_acl_storage->data()))); - aclIntArray *actual_seq_q_acl = host_vector_to_acl_int_array(actual_seq_q_vec); + aclIntArray *actual_seq_q_acl = host_vector_to_acl_int_array(actual_seq_q_for_acl); aclIntArray *actual_seq_k_acl = host_vector_to_acl_int_array(actual_seq_k_vec); - int64_t sparse_mode = 3; // rightDownCausal - aclTensor *atten_mask_acl = create_causal_mask(p); + // A one-token query is decode-equivalent and may attend the whole prefix. + // rightDownCausal is needed only when a request contributes multiple + // query tokens. + const int64_t sparse_mode = use_bnsd && max_q_len == 1 ? 0 : 3; + aclTensor *atten_mask_acl = sparse_mode == 0 ? nullptr : create_causal_mask(p); uint64_t workspace_size = 0; aclOpExecutor *executor = nullptr; @@ -308,8 +363,9 @@ void run(void *planned_meta) { nullptr, // dequantScaleQuery nullptr, // learnableSink num_heads, static_cast(p->scale), 2147483647, 2147483647, - const_cast("TND"), num_kv_heads, sparse_mode, - 0, // innerPrecise: high precision TND path + const_cast(use_bnsd ? "BNSD" : "TND"), + num_kv_heads, sparse_mode, + 0, // innerPrecise: high precision path block_size_val, // blockSize - Paged Attention block size 0, // antiquantMode false, @@ -336,7 +392,6 @@ void run(void *planned_meta) { + std::to_string(ret) + ", msg: " + (err_msg ? err_msg : "(null)")); } - aclrtStream stream = static_cast(infinicore::context::getStream()); void *workspace = nullptr; if (workspace_size > 0) { workspace = infini::ops::ascend::GetWorkspacePool() @@ -346,6 +401,20 @@ void run(void *planned_meta) { ret = aclnnFusedInferAttentionScoreV4(workspace, workspace_size, executor, stream); + if (ret == 0 && use_bnsd) { + // Unpack physical BNSD directly into the caller's ragged TND output. + for (int64_t i = 0; i < batch_size; ++i) { + const size_t q_begin = static_cast(cu_q_host[i]); + const size_t q_len = static_cast(actual_seq_q_for_acl[i]); + Tensor src = out_acl_storage + ->narrow({{0, static_cast(i), 1}, + {2, 0, q_len}}) + ->squeeze(0) + ->permute({1, 0, 2}); + Tensor dst = out_work->narrow({{0, q_begin, q_len}}); + dst->copy_from(src); + } + } // Release aclTensor/aclTensorList/aclIntArray resources aclDestroyTensor(query_acl); diff --git a/src/infiniop/devices/ascend/CMakeLists.txt b/src/infiniop/devices/ascend/CMakeLists.txt index 5229f770c..bd2f4e970 100644 --- a/src/infiniop/devices/ascend/CMakeLists.txt +++ b/src/infiniop/devices/ascend/CMakeLists.txt @@ -31,6 +31,11 @@ ascendc_library(ascend_kernels STATIC ../../ops/paged_caching/ascend/paged_caching_ascend_kernel.cpp ../../ops/paged_attention/ascend/paged_attention_ascend_kernel.cpp ../../ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend_kernel.cpp + ../../ops/causal_conv1d/ascend/causal_conv1d_ascend_kernel.cpp + ../../ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend_kernel.cpp + ../../ops/chunk_gated_delta_rule/ascend/gated_delta_rule_ascend_kernel.cpp + ../../ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_native_kernel.cpp + ../../ops/mrope/ascend/mrope_ascend_kernel.cpp ) target_include_directories(ascend_kernels PRIVATE ../../../../include) diff --git a/src/infiniop/devices/ascend/aclnn_elementwise.h b/src/infiniop/devices/ascend/aclnn_elementwise.h new file mode 100644 index 000000000..4a0509054 --- /dev/null +++ b/src/infiniop/devices/ascend/aclnn_elementwise.h @@ -0,0 +1,34 @@ +#ifndef __INFINIOP_ACLNN_ELEMENTWISE_H__ +#define __INFINIOP_ACLNN_ELEMENTWISE_H__ + +#include "../../operator.h" +#include "../../tensor.h" + +#include + +#define ACLNN_ELEMENTWISE_DESCRIPTOR(OP) \ + namespace op::OP::ascend { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + size_t _workspace_size; \ + \ + Descriptor(Opaque *opaque, size_t workspace_size, \ + infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), _workspace_size(workspace_size) {} \ + \ + public: \ + ~Descriptor(); \ + size_t workspaceSize() const { return _workspace_size; } \ + static infiniStatus_t create( \ + infiniopHandle_t handle, Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t output_desc, \ + std::vector input_descs); \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, void *output, \ + std::vector inputs, void *stream) const; \ + }; \ + } + +#endif // __INFINIOP_ACLNN_ELEMENTWISE_H__ diff --git a/src/infiniop/devices/ascend/aclnn_executor.h b/src/infiniop/devices/ascend/aclnn_executor.h new file mode 100644 index 000000000..086739c8f --- /dev/null +++ b/src/infiniop/devices/ascend/aclnn_executor.h @@ -0,0 +1,67 @@ +#ifndef __INFINIOP_ACLNN_EXECUTOR_H__ +#define __INFINIOP_ACLNN_EXECUTOR_H__ + +#include "common_ascend.h" + +#include +#include + +namespace device::ascend { + +struct AclnnExecutor { + std::vector tensors; + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + + AclnnExecutor() = default; + AclnnExecutor(const AclnnExecutor &) = delete; + AclnnExecutor &operator=(const AclnnExecutor &) = delete; + + ~AclnnExecutor() { + for (auto tensor : tensors) { + delete tensor; + } + if (executor != nullptr) { + aclDestroyAclOpExecutor(executor); + } + } + + void bind(std::initializer_list addresses) const { + size_t index = 0; + for (auto address : addresses) { + AclSetTensorAddr(executor, index, tensors[index]->tensor, address); + ++index; + } + } +}; + +inline infiniStatus_t validateAclnnElementwise( + infiniopTensorDescriptor_t output_desc, + const std::vector &input_descs, + size_t expected_inputs) { + + if (output_desc == nullptr || input_descs.size() != expected_inputs) { + return INFINI_STATUS_BAD_PARAM; + } + if (output_desc->hasBroadcastDim()) { + return INFINI_STATUS_BAD_TENSOR_STRIDES; + } + + for (auto input_desc : input_descs) { + if (input_desc == nullptr) { + return INFINI_STATUS_BAD_PARAM; + } + if (input_desc->dtype() != output_desc->dtype()) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (input_desc->ndim() != output_desc->ndim() + || input_desc->shape() != output_desc->shape()) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + return INFINI_STATUS_SUCCESS; +} + +} // namespace device::ascend + +#endif // __INFINIOP_ACLNN_EXECUTOR_H__ diff --git a/src/infiniop/devices/ascend/common_ascend.cc b/src/infiniop/devices/ascend/common_ascend.cc index d4f35d728..4e4eef8cf 100644 --- a/src/infiniop/devices/ascend/common_ascend.cc +++ b/src/infiniop/devices/ascend/common_ascend.cc @@ -18,7 +18,9 @@ size_t aclnnTensorDescriptor::numel() const { return std::accumulate(shape.begin(), shape.end(), (size_t)1, std::multiplies()); } -aclnnTensorDescriptor::aclnnTensorDescriptor(infiniopTensorDescriptor_t desc, void *data) { +aclnnTensorDescriptor::aclnnTensorDescriptor(infiniopTensorDescriptor_t desc, + void *data, + aclFormat format) { this->ndim = desc->ndim(); this->shape = std::vector(ndim); this->strides = std::vector(ndim); @@ -28,8 +30,7 @@ aclnnTensorDescriptor::aclnnTensorDescriptor(infiniopTensorDescriptor_t desc, vo } this->storageShape = inferStorageShape(this->shape, this->strides); this->dataType = toAclDataType(desc->dtype()); - // TODO: support other formats - this->format = aclFormat::ACL_FORMAT_ND; + this->format = format; this->tensor = aclCreateTensor(this->shape.data(), this->ndim, this->dataType, diff --git a/src/infiniop/devices/ascend/common_ascend.h b/src/infiniop/devices/ascend/common_ascend.h index 4c2e42c9b..fe3edcac3 100644 --- a/src/infiniop/devices/ascend/common_ascend.h +++ b/src/infiniop/devices/ascend/common_ascend.h @@ -35,7 +35,9 @@ struct aclnnTensorDescriptor { aclTensor *tensor; aclnnTensorDescriptor(aclDataType dtype, const std::vector &shape, const std::vector &strides, void *data = nullptr); - aclnnTensorDescriptor(infiniopTensorDescriptor_t y_desc, void *data = nullptr); + aclnnTensorDescriptor(infiniopTensorDescriptor_t y_desc, + void *data = nullptr, + aclFormat format = ACL_FORMAT_ND); ~aclnnTensorDescriptor(); size_t numel() const; diff --git a/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend.cc b/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend.cc new file mode 100644 index 000000000..ec75ec0f0 --- /dev/null +++ b/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend.cc @@ -0,0 +1,61 @@ +#include "causal_conv1d_ascend.h" +#include "../../../devices/ascend/ascend_handle.h" + +namespace op::causal_conv1d::ascend { +struct Descriptor::Opaque { +}; + +Descriptor::~Descriptor() { delete _opaque; } + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t conv_state_desc, + infiniopTensorDescriptor_t final_conv_state_desc, + infiniopTensorDescriptor_t qkv_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t bias_desc, + infiniopTensorDescriptor_t cu_seqlens_desc, + infiniopTensorDescriptor_t initial_state_indices_desc, + infiniopTensorDescriptor_t final_state_indices_desc) { + auto result = CausalConv1dInfo::create( + out_desc, conv_state_desc, final_conv_state_desc, qkv_desc, weight_desc, + bias_desc, cu_seqlens_desc, initial_state_indices_desc, + final_state_indices_desc); + CHECK_RESULT(result); + *desc_ptr = new Descriptor( + new Opaque{}, result.take(), 0, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, void *out, void *conv_state, + void *final_conv_state, const void *qkv, const void *weight, + const void *bias, const void *cu_seqlens, + const void *initial_state_indices, const void *final_state_indices, + void *stream) const { + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + (void)workspace; + return causal_conv1d_kernel_launch( + out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, + initial_state_indices, final_state_indices, _info.data_dtype, + _info.has_bias, _info.has_cu_seqlens, + _info.cu_seqlens_dtype == INFINI_DTYPE_I64, + _info.initial_state_indices_dtype == INFINI_DTYPE_I64, + _info.final_state_indices_dtype == INFINI_DTYPE_I64, + _info.indexed_state_pool, _info.request_count, _info.T, _info.C, + _info.total_tokens, _info.pool_size, + _info.out_strides[0], _info.out_strides[1], _info.out_strides[2], + _info.conv_state_strides[0], _info.conv_state_strides[1], + _info.conv_state_strides[2], + _info.final_conv_state_strides.empty() ? 0 : _info.final_conv_state_strides[0], + _info.final_conv_state_strides.empty() ? 0 : _info.final_conv_state_strides[1], + _info.final_conv_state_strides.empty() ? 0 : _info.final_conv_state_strides[2], + _info.qkv_strides[0], _info.qkv_strides[1], _info.qkv_strides[2], + _info.weight_strides[0], _info.weight_strides[2], + _info.bias_strides.empty() ? 0 : _info.bias_strides[0], stream); +} +} // namespace op::causal_conv1d::ascend diff --git a/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend.h b/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend.h new file mode 100644 index 000000000..8ceaa1f40 --- /dev/null +++ b/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend.h @@ -0,0 +1,22 @@ +#ifndef __CAUSAL_CONV1D_ASCEND_H__ +#define __CAUSAL_CONV1D_ASCEND_H__ +#include "../causal_conv1d.h" +DESCRIPTOR(ascend) +namespace op::causal_conv1d::ascend { +extern "C" infiniStatus_t causal_conv1d_kernel_launch( + void *out, void *conv_state, void *final_conv_state, + const void *qkv, const void *weight, const void *bias, + const void *cu_seqlens, const void *initial_state_indices, + const void *final_state_indices, infiniDtype_t dtype, + bool has_bias, bool has_cu_seqlens, bool cu_seqlens_i64, + bool initial_state_indices_i64, bool final_state_indices_i64, + bool indexed_state_pool, size_t request_count, size_t T, + size_t C, size_t total_tokens, size_t pool_size, + ptrdiff_t out_s0, ptrdiff_t out_s1, ptrdiff_t out_s2, + ptrdiff_t state_s0, ptrdiff_t state_s1, ptrdiff_t state_s2, + ptrdiff_t final_s0, ptrdiff_t final_s1, ptrdiff_t final_s2, + ptrdiff_t qkv_s0, ptrdiff_t qkv_s1, ptrdiff_t qkv_s2, + ptrdiff_t weight_s0, ptrdiff_t weight_s2, ptrdiff_t bias_s0, + void *stream); +} +#endif diff --git a/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend_kernel.cpp b/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend_kernel.cpp new file mode 100644 index 000000000..8acd18f90 --- /dev/null +++ b/src/infiniop/ops/causal_conv1d/ascend/causal_conv1d_ascend_kernel.cpp @@ -0,0 +1,506 @@ +#include "../../../devices/ascend/ascend_kernel_common.h" +#include +using namespace AscendC; + +template +__aicore__ inline float causalDataToFloat(T value) { + if constexpr (std::is_same::value) { + uint32_t bits = static_cast( + *reinterpret_cast(&value)) + << 16; + return *reinterpret_cast(&bits); + } else { + return static_cast(value); + } +} + +template +__aicore__ inline T causalFloatToData(float value) { + if constexpr (std::is_same::value) { + uint32_t bits = *reinterpret_cast(&value); + uint16_t upper = static_cast(bits >> 16); + return *reinterpret_cast(&upper); + } else { + return static_cast(value); + } +} + +__aicore__ inline int64_t causalLoadIndex( + GM_ADDR ptr, bool is_i64, int index, int fallback) { + if (ptr == nullptr) { + return static_cast(fallback); + } + if (is_i64) { + return static_cast( + reinterpret_cast<__gm__ int64_t *>(ptr)[index]); + } + return static_cast( + reinterpret_cast<__gm__ int32_t *>(ptr)[index]); +} + +template +__aicore__ inline float causalLoadHistory( + GlobalTensor &conv_state, GlobalTensor &qkv, + int64_t history_pos, int64_t token_begin, int token_batch, + int channel, ptrdiff_t state_base, ptrdiff_t state_s2, + ptrdiff_t qkv_s0, ptrdiff_t qkv_s1, ptrdiff_t qkv_s2) { + constexpr int64_t STATE_LEN = 3; + if (history_pos < STATE_LEN) { + return causalDataToFloat( + conv_state.GetValue(state_base + history_pos * state_s2)); + } + int64_t token_idx = token_begin + history_pos - STATE_LEN; + ptrdiff_t qkv_offset = static_cast(token_batch) * qkv_s0 + + static_cast(token_idx) * qkv_s1 + + static_cast(channel) * qkv_s2; + return causalDataToFloat(qkv.GetValue(qkv_offset)); +} + +template +class CausalConv1dKernel { +public: + __aicore__ inline void process( + GM_ADDR out_ptr, GM_ADDR conv_state_ptr, GM_ADDR final_state_ptr, + GM_ADDR qkv_ptr, GM_ADDR weight_ptr, GM_ADDR bias_ptr, + GM_ADDR cu_seqlens, GM_ADDR initial_indices, GM_ADDR final_indices, + bool has_bias, bool has_cu, bool cu_i64, bool initial_i64, + bool final_i64, bool indexed_pool, bool update_state_only, bool fuse_state_update, + size_t request_count, size_t T_tokens, + size_t C, size_t total_tokens, size_t pool_size, + ptrdiff_t out_s0, ptrdiff_t out_s1, ptrdiff_t out_s2, + ptrdiff_t state_s0, ptrdiff_t state_s1, ptrdiff_t state_s2, + ptrdiff_t final_s0, ptrdiff_t final_s1, ptrdiff_t final_s2, + ptrdiff_t qkv_s0, ptrdiff_t qkv_s1, ptrdiff_t qkv_s2, + ptrdiff_t weight_s0, ptrdiff_t weight_s2, ptrdiff_t bias_s0) { + GlobalTensor out, conv_state, final_state, qkv, weight, bias; + out.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(out_ptr)); + conv_state.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(conv_state_ptr)); + final_state.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(final_state_ptr)); + qkv.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(qkv_ptr)); + weight.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(weight_ptr)); + bias.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(bias_ptr)); + + const size_t block = GetBlockIdx(); + const size_t block_count = GetBlockNum(); + const bool vector_fast_path = fuse_state_update && !update_state_only + && block_count != 0 && C % block_count == 0 + && ((C / block_count) * sizeof(T)) % BYTE_ALIGN == 0 + && out_s2 == 1 && qkv_s2 == 1 + && state_s1 == 3 && state_s2 == 1 + && weight_s0 == 4 && weight_s2 == 1 + && (!has_bias || bias_s0 == 1); + const size_t fast_tile_len = vector_fast_path ? C / block_count : 1; + const size_t fast_copy_len = alignTileLen(fast_tile_len, BYTE_ALIGN); + const size_t fast_channel_begin = block * fast_tile_len; + + TPipe pipe; + const event_t mte2_to_v = static_cast( + pipe.FetchEventID(HardEvent::MTE2_V)); + const event_t v_to_mte3 = static_cast( + pipe.FetchEventID(HardEvent::V_MTE3)); + const event_t s_to_v = static_cast( + pipe.FetchEventID(HardEvent::S_V)); + const event_t mte2_to_s = static_cast( + pipe.FetchEventID(HardEvent::MTE2_S)); + const event_t s_to_mte3 = static_cast( + pipe.FetchEventID(HardEvent::S_MTE3)); + const event_t v_to_s = static_cast( + pipe.FetchEventID(HardEvent::V_S)); + const event_t mte3_to_v = static_cast( + pipe.FetchEventID(HardEvent::MTE3_V)); + const event_t v_to_mte2 = static_cast( + pipe.FetchEventID(HardEvent::V_MTE2)); + TBuf x_data_buf, weight_data_buf; + TBuf state_raw_buf, weight_raw_buf, qkv_raw_buf; + TBuf bias_data_buf, out_data_buf; + TBuf gather_offsets_buf; + TBuf x_float_buf, weight_float_buf; + TBuf acc_float_buf, tmp_float_buf; + if (vector_fast_path) { + pipe.InitBuffer(x_data_buf, fast_copy_len * sizeof(T)); + pipe.InitBuffer(weight_data_buf, 4 * fast_copy_len * sizeof(T)); + pipe.InitBuffer(state_raw_buf, 3 * fast_copy_len * sizeof(T)); + pipe.InitBuffer(weight_raw_buf, 4 * fast_copy_len * sizeof(T)); + pipe.InitBuffer(qkv_raw_buf, fast_copy_len * sizeof(T)); + pipe.InitBuffer(bias_data_buf, fast_copy_len * sizeof(T)); + pipe.InitBuffer(out_data_buf, fast_copy_len * sizeof(T)); + pipe.InitBuffer(gather_offsets_buf, fast_copy_len * sizeof(uint32_t)); + pipe.InitBuffer(x_float_buf, fast_copy_len * sizeof(float)); + pipe.InitBuffer(weight_float_buf, 4 * fast_copy_len * sizeof(float)); + pipe.InitBuffer(acc_float_buf, fast_copy_len * sizeof(float)); + pipe.InitBuffer(tmp_float_buf, fast_copy_len * sizeof(float)); + + LocalTensor weight_data = weight_data_buf.Get(); + LocalTensor weight_raw = weight_raw_buf.Get(); + LocalTensor weight_float = weight_float_buf.Get(); + LocalTensor gather_offsets = gather_offsets_buf.Get(); + if (weight_s0 != 1) { + DataCopy(weight_raw, + weight[static_cast(fast_channel_begin) * weight_s0], + 4 * fast_copy_len); + SetFlag(mte2_to_v); + WaitFlag(mte2_to_v); + ArithProgression( + gather_offsets.ReinterpretCast(), 0, + static_cast(4 * sizeof(T)), fast_tile_len); + PipeBarrier(); + for (int k = 0; k < 4; ++k) { + if constexpr (std::is_same::value) { + Gather(weight_float[k * fast_copy_len], weight_raw, + gather_offsets, k * sizeof(T), fast_tile_len); + } else { + Gather(weight_data[k * fast_copy_len], weight_raw, + gather_offsets, k * sizeof(T), fast_tile_len); + PipeBarrier(); + Cast(weight_float[k * fast_copy_len], + weight_data[k * fast_copy_len], + RoundMode::CAST_NONE, fast_copy_len); + } + PipeBarrier(); + } + } else { + for (int k = 0; k < 4; ++k) { + const ptrdiff_t weight_offset = k * weight_s2 + static_cast(fast_channel_begin); + if constexpr (std::is_same::value) { + DataCopy(weight_float[k * fast_copy_len], + weight[weight_offset], fast_copy_len); + } else { + DataCopy(weight_data[k * fast_copy_len], + weight[weight_offset], fast_copy_len); + } + SetFlag(mte2_to_v); + WaitFlag(mte2_to_v); + if constexpr (!std::is_same::value) { + Cast(weight_float[k * fast_copy_len], + weight_data[k * fast_copy_len], + RoundMode::CAST_NONE, fast_copy_len); + } + } + } + ArithProgression( + gather_offsets.ReinterpretCast(), 0, + static_cast(3 * sizeof(T)), fast_tile_len); + PipeBarrier(); + } + for (size_t request = 0; request < request_count; ++request) { + int64_t token_begin = 0; + int64_t token_end = static_cast(T_tokens); + int token_batch = static_cast(request); + if (has_cu) { + token_begin = causalLoadIndex(cu_seqlens, cu_i64, request, 0); + token_end = causalLoadIndex(cu_seqlens, cu_i64, request + 1, 0); + token_batch = 0; + if (token_begin < 0 || token_end < token_begin + || token_end > static_cast(total_tokens)) { + return; + } + } + const int64_t request_len = token_end - token_begin; + const int64_t read_slot = indexed_pool + ? causalLoadIndex(initial_indices, initial_i64, request, request) + : static_cast(request); + const int64_t write_slot = indexed_pool && final_indices != nullptr + ? causalLoadIndex(final_indices, final_i64, request, request) + : static_cast(request); + if (read_slot < 0 || write_slot < 0 + || read_slot >= static_cast(pool_size) + || (final_indices != nullptr + && write_slot >= static_cast(pool_size))) { + return; + } + + if (!update_state_only) { + if (vector_fast_path) { + LocalTensor x_data = x_data_buf.Get(); + LocalTensor state_raw = state_raw_buf.Get(); + LocalTensor weight_raw = weight_raw_buf.Get(); + LocalTensor qkv_raw = qkv_raw_buf.Get(); + LocalTensor bias_data = bias_data_buf.Get(); + LocalTensor out_data = out_data_buf.Get(); + LocalTensor x_float = x_float_buf.Get(); + LocalTensor gather_offsets = gather_offsets_buf.Get(); + LocalTensor weight_float = weight_float_buf.Get(); + LocalTensor acc_float = acc_float_buf.Get(); + LocalTensor tmp_float = tmp_float_buf.Get(); + const ptrdiff_t state_base = read_slot * state_s0 + + static_cast(fast_channel_begin) * state_s1; + if (state_s1 != 1) { + DataCopy(state_raw, conv_state[state_base], + 3 * fast_copy_len); + SetFlag(mte2_to_s); + WaitFlag(mte2_to_s); + } + if (has_bias) { + DataCopy(bias_data, + bias[static_cast(fast_channel_begin) * bias_s0], + fast_copy_len); + SetFlag(mte2_to_v); + WaitFlag(mte2_to_v); + } + for (int64_t t = 0; t < request_len; ++t) { + Duplicate(acc_float, 0.0f, fast_copy_len); + for (int k = 0; k < 4; ++k) { + const int64_t history_pos = t + k; + if (history_pos < 3) { + if constexpr (std::is_same::value) { + Gather(x_float, state_raw, gather_offsets, + history_pos * sizeof(T), fast_tile_len); + } else { + Gather(x_data, state_raw, gather_offsets, + history_pos * sizeof(T), fast_tile_len); + } + PipeBarrier(); + } else { + const int64_t token_idx = token_begin + history_pos - 3; + const ptrdiff_t qkv_offset = static_cast(token_batch) * qkv_s0 + + static_cast(token_idx) * qkv_s1 + + static_cast(fast_channel_begin); + if constexpr (std::is_same::value) { + DataCopy(x_float, qkv[qkv_offset], fast_copy_len); + } else { + DataCopy(x_data, qkv[qkv_offset], fast_copy_len); + } + SetFlag(mte2_to_v); + WaitFlag(mte2_to_v); + } + if constexpr (!std::is_same::value) { + Cast(x_float, x_data, RoundMode::CAST_NONE, fast_copy_len); + } + PipeBarrier(); + Mul(tmp_float, x_float, + weight_float[k * fast_copy_len], fast_copy_len); + Add(acc_float, acc_float, tmp_float, fast_copy_len); + SetFlag(v_to_s); + WaitFlag(v_to_s); + SetFlag(v_to_mte2); + WaitFlag(v_to_mte2); + } + if (has_bias) { + if constexpr (std::is_same::value) { + Add(acc_float, acc_float, bias_data, fast_copy_len); + } else { + Cast(tmp_float, bias_data, RoundMode::CAST_NONE, fast_copy_len); + Add(acc_float, acc_float, tmp_float, fast_copy_len); + } + } + const ptrdiff_t out_offset = static_cast(token_batch) * out_s0 + + (token_begin + t) * out_s1 + + static_cast(fast_channel_begin); + if constexpr (std::is_same::value) { + SetFlag(v_to_mte3); + WaitFlag(v_to_mte3); + DataCopy(out[out_offset], acc_float, fast_tile_len); + } else { + Cast(out_data, acc_float, RoundMode::CAST_RINT, fast_copy_len); + SetFlag(v_to_mte3); + WaitFlag(v_to_mte3); + DataCopy(out[out_offset], out_data, fast_tile_len); + } + SetFlag(mte3_to_v); + WaitFlag(mte3_to_v); + } + } else { + for (size_t channel = block; channel < C; channel += block_count) { + const ptrdiff_t state_base = read_slot * state_s0 + + static_cast(channel) * state_s1; + const ptrdiff_t weight_base = static_cast(channel) * weight_s0; + for (int64_t t = 0; t < request_len; ++t) { + float acc = 0.0f; + for (int k = 0; k < 4; ++k) { + const float w = causalDataToFloat( + weight.GetValue(weight_base + k * weight_s2)); + const float x = causalLoadHistory( + conv_state, qkv, t + k, token_begin, token_batch, channel, + state_base, state_s2, qkv_s0, qkv_s1, qkv_s2); + acc += w * x; + } + if (has_bias) { + acc += causalDataToFloat( + bias.GetValue(static_cast(channel) * bias_s0)); + } + const ptrdiff_t out_offset = static_cast(token_batch) * out_s0 + + (token_begin + t) * out_s1 + + static_cast(channel) * out_s2; + out.SetValue(out_offset, causalFloatToData(acc)); + } + } + } + } + if (update_state_only || fuse_state_update) { + // Keep block boundaries on 512-byte cache-line boundaries. This + // prevents different AI cores from flushing overlapping lines of + // the interleaved [C, 3] recurrent-state layout. + constexpr size_t state_align_channels = std::is_same::value ? 128 : 256; + const size_t raw_channels = (C + block_count - 1) / block_count; + const size_t block_channels = ((raw_channels + state_align_channels - 1) / state_align_channels) + * state_align_channels; + const size_t channel_limit = (block + 1) * block_channels; + const size_t channel_end = channel_limit < C ? channel_limit : C; + const bool write_to_pool = final_indices != nullptr; + const bool vector_state_update = vector_fast_path + && (write_to_pool + || (final_s1 == 3 && final_s2 == 1)); + if (vector_state_update) { + LocalTensor state_raw = state_raw_buf.Get(); + LocalTensor qkv_raw = qkv_raw_buf.Get(); + for (int k = 0; k < 3; ++k) { + const int64_t history_pos = request_len + k; + if (history_pos < 3) { + for (size_t i = 0; i < fast_tile_len; ++i) { + state_raw.SetValue( + i * 3 + k, + state_raw.GetValue(i * 3 + history_pos)); + } + } else { + const int64_t token_idx = token_begin + history_pos - 3; + const ptrdiff_t qkv_offset = static_cast(token_batch) * qkv_s0 + + static_cast(token_idx) * qkv_s1 + + static_cast(fast_channel_begin); + DataCopy(qkv_raw, qkv[qkv_offset], fast_copy_len); + SetFlag(mte2_to_s); + WaitFlag(mte2_to_s); + for (size_t i = 0; i < fast_tile_len; ++i) { + state_raw.SetValue( + i * 3 + k, qkv_raw.GetValue(i)); + } + } + } + const ptrdiff_t target_base = write_to_pool + ? write_slot * state_s0 + + static_cast(fast_channel_begin) * state_s1 + : static_cast(request) * final_s0 + + static_cast(fast_channel_begin) * final_s1; + SetFlag(s_to_mte3); + WaitFlag(s_to_mte3); + if (write_to_pool) { + DataCopy(conv_state[target_base], state_raw, + 3 * fast_copy_len); + } else { + DataCopy(final_state[target_base], state_raw, + 3 * fast_copy_len); + } + } else { + for (size_t channel = block * block_channels; channel < channel_end; ++channel) { + const ptrdiff_t state_base = read_slot * state_s0 + + static_cast(channel) * state_s1; + const bool write_to_pool = final_indices != nullptr; + const ptrdiff_t target_base = write_to_pool + ? write_slot * state_s0 + + static_cast(channel) * state_s1 + : static_cast(request) * final_s0 + + static_cast(channel) * final_s1; + const ptrdiff_t target_s2 = write_to_pool ? state_s2 : final_s2; + for (int k = 0; k < 3; ++k) { + const T value = causalFloatToData(causalLoadHistory( + conv_state, qkv, request_len + k, token_begin, token_batch, + channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); + if (write_to_pool) { + conv_state.SetValue(target_base + k * target_s2, value); + } else { + final_state.SetValue(target_base + k * target_s2, value); + } + } + } + } + } + } + if (!update_state_only) { + DataCacheCleanAndInvalid(out); + } + DataCacheCleanAndInvalid(conv_state); + if (final_state_ptr != nullptr) { + DataCacheCleanAndInvalid(final_state); + } + } +}; +#define DEFINE_CAUSAL_CONV1D_KERNEL(NAME, TYPE) \ + __global__ __aicore__ void NAME( \ + GM_ADDR out, GM_ADDR conv_state, GM_ADDR final_state, GM_ADDR qkv, \ + GM_ADDR weight, GM_ADDR bias, GM_ADDR cu, GM_ADDR initial_indices, \ + GM_ADDR final_indices, bool has_bias, bool has_cu, bool cu_i64, \ + bool initial_i64, bool final_i64, bool indexed_pool, \ + bool update_state_only, bool fuse_state_update, \ + size_t request_count, size_t T_tokens, size_t C, \ + size_t total_tokens, size_t pool_size, ptrdiff_t out_s0, \ + ptrdiff_t out_s1, ptrdiff_t out_s2, ptrdiff_t state_s0, \ + ptrdiff_t state_s1, ptrdiff_t state_s2, ptrdiff_t final_s0, \ + ptrdiff_t final_s1, ptrdiff_t final_s2, ptrdiff_t qkv_s0, \ + ptrdiff_t qkv_s1, ptrdiff_t qkv_s2, ptrdiff_t weight_s0, \ + ptrdiff_t weight_s2, ptrdiff_t bias_s0) { \ + CausalConv1dKernel kernel; \ + kernel.process(out, conv_state, final_state, qkv, weight, bias, cu, \ + initial_indices, final_indices, has_bias, has_cu, \ + cu_i64, initial_i64, final_i64, indexed_pool, \ + update_state_only, fuse_state_update, request_count, T_tokens, C, \ + total_tokens, pool_size, \ + out_s0, out_s1, out_s2, state_s0, state_s1, state_s2, \ + final_s0, final_s1, final_s2, qkv_s0, qkv_s1, qkv_s2, \ + weight_s0, weight_s2, bias_s0); \ + } + +DEFINE_CAUSAL_CONV1D_KERNEL(causal_conv1d_half, half) +DEFINE_CAUSAL_CONV1D_KERNEL(causal_conv1d_float, float) +DEFINE_CAUSAL_CONV1D_KERNEL(causal_conv1d_bf16, bfloat16_t) +#undef DEFINE_CAUSAL_CONV1D_KERNEL + +extern "C" infiniStatus_t causal_conv1d_kernel_launch( + void *out, void *conv_state, void *final_state, const void *qkv, + const void *weight, const void *bias, const void *cu, + const void *initial_indices, const void *final_indices, + infiniDtype_t dtype, bool has_bias, bool has_cu, bool cu_i64, + bool initial_i64, bool final_i64, bool indexed_pool, + size_t request_count, size_t T_tokens, size_t C, size_t total_tokens, + size_t pool_size, ptrdiff_t out_s0, ptrdiff_t out_s1, ptrdiff_t out_s2, + ptrdiff_t state_s0, ptrdiff_t state_s1, ptrdiff_t state_s2, + ptrdiff_t final_s0, ptrdiff_t final_s1, ptrdiff_t final_s2, + ptrdiff_t qkv_s0, ptrdiff_t qkv_s1, ptrdiff_t qkv_s2, + ptrdiff_t weight_s0, ptrdiff_t weight_s2, ptrdiff_t bias_s0, + void *stream) { + if (request_count == 0 || C == 0) { + return INFINI_STATUS_SUCCESS; + } + const size_t elem_size = dtype == INFINI_DTYPE_F32 ? sizeof(float) : sizeof(uint16_t); + uint32_t blocks = static_cast(BLOCK_NUM); + while (blocks > 1 + && (C % blocks != 0 + || ((C / blocks) * elem_size) % BYTE_ALIGN != 0)) { + --blocks; + } + const bool fuse_state_update = out_s2 == 1 && qkv_s2 == 1 + && state_s1 == 3 && state_s2 == 1 + && weight_s0 == 4 && weight_s2 == 1 + && (!has_bias || bias_s0 == 1); +#define LAUNCH_CASE(DTYPE, NAME) \ + case DTYPE: \ + NAME<<>>( \ + out, conv_state, final_state, const_cast(qkv), \ + const_cast(weight), const_cast(bias), \ + const_cast(cu), const_cast(initial_indices), \ + const_cast(final_indices), has_bias, has_cu, cu_i64, \ + initial_i64, final_i64, indexed_pool, false, fuse_state_update, \ + request_count, T_tokens, C, total_tokens, pool_size, \ + out_s0, out_s1, out_s2, \ + state_s0, state_s1, state_s2, final_s0, final_s1, final_s2, \ + qkv_s0, qkv_s1, qkv_s2, weight_s0, weight_s2, bias_s0); \ + if (!fuse_state_update) { \ + NAME<<>>( \ + out, conv_state, final_state, const_cast(qkv), \ + const_cast(weight), const_cast(bias), \ + const_cast(cu), const_cast(initial_indices), \ + const_cast(final_indices), has_bias, has_cu, cu_i64, \ + initial_i64, final_i64, indexed_pool, true, false, \ + request_count, T_tokens, C, total_tokens, pool_size, \ + out_s0, out_s1, out_s2, state_s0, state_s1, state_s2, \ + final_s0, final_s1, final_s2, qkv_s0, qkv_s1, qkv_s2, \ + weight_s0, weight_s2, bias_s0); \ + } \ + return INFINI_STATUS_SUCCESS; + switch (dtype) { + LAUNCH_CASE(INFINI_DTYPE_F16, causal_conv1d_half) + LAUNCH_CASE(INFINI_DTYPE_F32, causal_conv1d_float) + LAUNCH_CASE(INFINI_DTYPE_BF16, causal_conv1d_bf16) + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +#undef LAUNCH_CASE +} diff --git a/src/infiniop/ops/causal_conv1d/operator.cc b/src/infiniop/ops/causal_conv1d/operator.cc index 8666d2e62..080ef31a5 100644 --- a/src/infiniop/ops/causal_conv1d/operator.cc +++ b/src/infiniop/ops/causal_conv1d/operator.cc @@ -12,6 +12,9 @@ #include "moore/causal_conv1d_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/causal_conv1d_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateCausalConv1dDescriptor( infiniopHandle_t handle, infiniopCausalConv1dDescriptor_t *desc_ptr, @@ -46,6 +49,9 @@ __INFINI_C infiniStatus_t infiniopCreateCausalConv1dDescriptor( #endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) +#endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -72,6 +78,9 @@ __INFINI_C infiniStatus_t infiniopGetCausalConv1dWorkspaceSize( #endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) +#endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -112,6 +121,9 @@ __INFINI_C infiniStatus_t infiniopCausalConv1d( #endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore) +#endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -138,6 +150,9 @@ __INFINI_C infiniStatus_t infiniopDestroyCausalConv1dDescriptor( #endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) +#endif +#ifdef ENABLE_ASCEND_API + DESTROY(INFINI_DEVICE_ASCEND, ascend) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/chunk_gated_delta_rule/ascend/chunk_gated_delta_rule_ascend.cc b/src/infiniop/ops/chunk_gated_delta_rule/ascend/chunk_gated_delta_rule_ascend.cc new file mode 100644 index 000000000..3d2ec7dfc --- /dev/null +++ b/src/infiniop/ops/chunk_gated_delta_rule/ascend/chunk_gated_delta_rule_ascend.cc @@ -0,0 +1,469 @@ +#include "chunk_gated_delta_rule_ascend.h" +#include "../../../devices/ascend/ascend_handle.h" +#include "../../../devices/ascend/common_ascend.h" +#include "../../recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_native_kernel.h" +#include "gated_delta_rule_ascend_kernel.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace op::chunk_gated_delta_rule::ascend { + +namespace { + +constexpr size_t ACLNN_DIM = 128; +constexpr size_t ACLNN_MAX_SEQ_LEN = 8; +constexpr size_t ACLNN_ALIGNMENT = 512; + +size_t alignAclnn(size_t value) { + return (value + ACLNN_ALIGNMENT - 1) & ~(ACLNN_ALIGNMENT - 1); +} + +size_t reserveAclnn(size_t &cursor, size_t bytes) { + cursor = alignAclnn(cursor); + const size_t result = cursor; + cursor += bytes; + return result; +} + +template +infiniStatus_t copyDeviceVector( + std::vector &dst, const void *src, size_t count) { + std::vector temporary(count); + CHECK_ACL(aclrtMemcpy( + temporary.data(), temporary.size() * sizeof(T), src, + temporary.size() * sizeof(T), ACL_MEMCPY_DEVICE_TO_HOST)); + for (size_t i = 0; i < count; ++i) { + dst[i] = static_cast(temporary[i]); + } + return INFINI_STATUS_SUCCESS; +} + +struct Segment { + size_t offset; + size_t length; + int32_t state_slot; +}; + +} // namespace + +struct Descriptor::Opaque { + struct Call { + aclnnTensorDescriptor_t q = nullptr; + aclnnTensorDescriptor_t k = nullptr; + aclnnTensorDescriptor_t v = nullptr; + aclnnTensorDescriptor_t beta = nullptr; + aclnnTensorDescriptor_t state = nullptr; + aclnnTensorDescriptor_t actual_seq_lengths = nullptr; + aclnnTensorDescriptor_t state_indices = nullptr; + aclnnTensorDescriptor_t g = nullptr; + aclnnTensorDescriptor_t out = nullptr; + aclOpExecutor *executor = nullptr; + uint64_t workspace_size = 0; + + ~Call() { + delete q; + delete k; + delete v; + delete beta; + delete state; + delete actual_seq_lengths; + delete state_indices; + delete g; + delete out; + } + }; + + bool aclnn = false; + std::array>, ACLNN_MAX_SEQ_LEN + 1> calls; + size_t q_offset = 0; + size_t k_offset = 0; + size_t beta_offset = 0; + size_t input_indices_offset = 0; + size_t state_indices_offset = 0; + size_t actual_seq_lengths_offset = 0; + size_t state_staging_offset = 0; + size_t aclnn_workspace_offset = 0; + uint64_t aclnn_workspace_size = 0; + void *buffer = nullptr; + int32_t *host_indices = nullptr; + int32_t *host_lengths = nullptr; + size_t metadata_capacity = 0; + + ~Opaque() { + if (buffer != nullptr) { + aclrtFree(buffer); + } + if (host_indices != nullptr) { + aclrtFreeHost(host_indices); + } + if (host_lengths != nullptr) { + aclrtFreeHost(host_lengths); + } + } +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t initial_state_desc, + infiniopTensorDescriptor_t final_state_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_desc, + infiniopTensorDescriptor_t v_desc, + infiniopTensorDescriptor_t g_desc, + infiniopTensorDescriptor_t beta_desc, + infiniopTensorDescriptor_t cu_seqlens_desc, + infiniopTensorDescriptor_t initial_state_indices_desc, + infiniopTensorDescriptor_t final_state_indices_desc, + bool use_qk_l2norm, + size_t chunk_size) { + + auto result = ChunkGatedDeltaRuleInfo::create( + out_desc, initial_state_desc, final_state_desc, q_desc, k_desc, v_desc, + g_desc, beta_desc, cu_seqlens_desc, initial_state_indices_desc, + final_state_indices_desc, use_qk_l2norm, chunk_size); + CHECK_RESULT(result); + auto info = result.take(); + auto opaque = new Opaque{}; + + const bool state_contiguous = info.initial_state_strides[3] == 1 + && info.initial_state_strides[2] == static_cast(info.Dk) + && info.initial_state_strides[1] + == static_cast(info.Dv * info.Dk) + && info.initial_state_strides[0] + == static_cast(info.Hv * info.Dv * info.Dk); + const bool tensors_contiguous = info.q_strides[1] == static_cast(info.Hk * info.Dk) + && info.q_strides[2] == static_cast(info.Dk) + && info.q_strides[3] == 1 + && info.k_strides[1] == static_cast(info.Hk * info.Dk) + && info.k_strides[2] == static_cast(info.Dk) + && info.k_strides[3] == 1 + && info.v_strides[1] == static_cast(info.Hv * info.Dv) + && info.v_strides[2] == static_cast(info.Dv) + && info.v_strides[3] == 1 + && info.out_strides[1] == static_cast(info.Hv * info.Dv) + && info.out_strides[2] == static_cast(info.Dv) + && info.out_strides[3] == 1 + && info.g_strides[1] == static_cast(info.Hv) + && info.g_strides[2] == 1 + && info.beta_strides[1] == static_cast(info.Hv) + && info.beta_strides[2] == 1; + opaque->aclnn = info.data_dtype == INFINI_DTYPE_BF16 + && info.gate_dtype == INFINI_DTYPE_F32 + && info.use_qk_l2norm + && info.has_cu_seqlens + && info.has_initial_state_indices + && info.has_final_state_indices + && info.Dk == ACLNN_DIM + && info.Dv == ACLNN_DIM + && state_contiguous + && tensors_contiguous; + + size_t workspace_size = info.B * info.Hv * info.Dv * info.Dk + * sizeof(float); + if (opaque->aclnn) { + const int64_t Hk = static_cast(info.Hk); + const int64_t Hv = static_cast(info.Hv); + const int64_t D = static_cast(ACLNN_DIM); + const int64_t pool = static_cast(info.pool_size); + for (size_t length = 1; length <= ACLNN_MAX_SEQ_LEN; ++length) { + const size_t call_count = length == ACLNN_MAX_SEQ_LEN + ? std::max(1, info.total_tokens / ACLNN_MAX_SEQ_LEN + info.B) + : std::max(1, info.B); + auto &pool_calls = opaque->calls[length]; + pool_calls.reserve(call_count); + for (size_t call_index = 0; call_index < call_count; ++call_index) { + auto call = std::make_unique(); + const int64_t L = static_cast(length); + call->q = new aclnnTensorDescriptor( + ACL_BF16, {L, Hk, D}, {Hk * D, D, 1}, nullptr); + call->k = new aclnnTensorDescriptor( + ACL_BF16, {L, Hk, D}, {Hk * D, D, 1}, nullptr); + call->v = new aclnnTensorDescriptor( + ACL_BF16, {L, Hv, D}, {Hv * D, D, 1}, nullptr); + call->beta = new aclnnTensorDescriptor( + ACL_BF16, {L, Hv}, {Hv, 1}, nullptr); + call->state = new aclnnTensorDescriptor( + ACL_BF16, {pool, Hv, D, D}, + {Hv * D * D, D * D, D, 1}, nullptr); + call->actual_seq_lengths = new aclnnTensorDescriptor( + ACL_INT32, {1}, {1}, nullptr); + call->state_indices = new aclnnTensorDescriptor( + ACL_INT32, {L}, {1}, nullptr); + call->g = new aclnnTensorDescriptor( + ACL_FLOAT, {L, Hv}, {Hv, 1}, nullptr); + call->out = new aclnnTensorDescriptor( + ACL_BF16, {L, Hv, D}, {Hv * D, D, 1}, nullptr); + CHECK_ACL(aclnnRecurrentGatedDeltaRuleGetWorkspaceSize( + call->q->tensor, call->k->tensor, call->v->tensor, + call->beta->tensor, call->state->tensor, + call->actual_seq_lengths->tensor, call->state_indices->tensor, + call->g->tensor, nullptr, nullptr, + 1.0f / std::sqrt(static_cast(ACLNN_DIM)), + call->out->tensor, &call->workspace_size, &call->executor)); + CHECK_ACL(aclSetAclOpExecutorRepeatable(call->executor)); + opaque->aclnn_workspace_size = std::max( + opaque->aclnn_workspace_size, call->workspace_size); + pool_calls.push_back(std::move(call)); + } + } + + size_t cursor = 0; + opaque->q_offset = reserveAclnn( + cursor, ACLNN_MAX_SEQ_LEN * info.Hk * ACLNN_DIM + * sizeof(uint16_t)); + opaque->k_offset = reserveAclnn( + cursor, ACLNN_MAX_SEQ_LEN * info.Hk * ACLNN_DIM + * sizeof(uint16_t)); + opaque->beta_offset = reserveAclnn( + cursor, ACLNN_MAX_SEQ_LEN * info.Hv * sizeof(uint16_t)); + opaque->input_indices_offset = reserveAclnn( + cursor, ACLNN_MAX_SEQ_LEN * sizeof(int32_t)); + opaque->state_indices_offset = reserveAclnn( + cursor, ACLNN_MAX_SEQ_LEN * sizeof(int32_t)); + opaque->actual_seq_lengths_offset = reserveAclnn( + cursor, ACLNN_MAX_SEQ_LEN * sizeof(int32_t)); + opaque->state_staging_offset = reserveAclnn(cursor, ACLNN_ALIGNMENT); + opaque->aclnn_workspace_offset = alignAclnn(cursor); + cursor = opaque->aclnn_workspace_offset + + opaque->aclnn_workspace_size; + CHECK_ACL(aclrtMalloc( + &opaque->buffer, cursor, ACL_MEM_MALLOC_HUGE_FIRST)); + opaque->metadata_capacity = info.total_tokens / ACLNN_MAX_SEQ_LEN + info.B; + CHECK_ACL(aclrtMallocHost( + reinterpret_cast(&opaque->host_indices), + opaque->metadata_capacity * ACLNN_MAX_SEQ_LEN * sizeof(int32_t))); + CHECK_ACL(aclrtMallocHost( + reinterpret_cast(&opaque->host_lengths), + opaque->metadata_capacity * sizeof(int32_t))); + workspace_size = std::max(workspace_size, cursor); + } + + *desc_ptr = new Descriptor( + opaque, std::move(info), workspace_size, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *out, + void *initial_state, + void *final_state, + const void *q, + const void *k, + const void *v, + const void *g, + const void *beta, + const void *cu_seqlens, + const void *initial_state_indices, + const void *final_state_indices, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (_info.gate_dtype != INFINI_DTYPE_F32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + if (_opaque->aclnn) { + std::vector cu(_info.B + 1); + std::vector initial_indices(_info.B); + std::vector final_indices(_info.B); + if (_info.cu_seqlens_dtype == INFINI_DTYPE_I64) { + CHECK_STATUS(copyDeviceVector( + cu, cu_seqlens, _info.B + 1)); + } else { + CHECK_STATUS(copyDeviceVector( + cu, cu_seqlens, _info.B + 1)); + } + if (_info.initial_state_indices_dtype == INFINI_DTYPE_I64) { + CHECK_STATUS(copyDeviceVector( + initial_indices, initial_state_indices, _info.B)); + } else { + CHECK_STATUS(copyDeviceVector( + initial_indices, initial_state_indices, _info.B)); + } + if (_info.final_state_indices_dtype == INFINI_DTYPE_I64) { + CHECK_STATUS(copyDeviceVector( + final_indices, final_state_indices, _info.B)); + } else { + CHECK_STATUS(copyDeviceVector( + final_indices, final_state_indices, _info.B)); + } + + std::vector segments; + const size_t state_bytes = _info.Hv * ACLNN_DIM * ACLNN_DIM * sizeof(uint16_t); + auto *state_bytes_ptr = static_cast(initial_state); + for (size_t request = 0; request < _info.B; ++request) { + if (cu[request] < 0 || cu[request + 1] < cu[request] + || cu[request + 1] + > static_cast(_info.total_tokens) + || initial_indices[request] < 0 + || final_indices[request] < 0 + || initial_indices[request] + >= static_cast(_info.pool_size) + || final_indices[request] + >= static_cast(_info.pool_size)) { + return INFINI_STATUS_BAD_PARAM; + } + const size_t source = static_cast(initial_indices[request]); + const size_t destination = static_cast(final_indices[request]); + if (source != destination) { + CHECK_ACL(aclrtMemcpyAsync( + state_bytes_ptr + destination * state_bytes, state_bytes, + state_bytes_ptr + source * state_bytes, state_bytes, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream)); + } + for (size_t offset = static_cast(cu[request]); + offset < static_cast(cu[request + 1]); + offset += ACLNN_MAX_SEQ_LEN) { + const size_t length = std::min( + ACLNN_MAX_SEQ_LEN, + static_cast(cu[request + 1]) - offset); + segments.push_back(Segment{ + offset, length, static_cast(destination)}); + } + } + + if (segments.size() > _opaque->metadata_capacity) { + return INFINI_STATUS_BAD_PARAM; + } + int32_t *host_indices = _opaque->host_indices; + int32_t *host_lengths = _opaque->host_lengths; + for (size_t i = 0; i < segments.size(); ++i) { + std::fill_n( + host_indices + i * ACLNN_MAX_SEQ_LEN, + ACLNN_MAX_SEQ_LEN, segments[i].state_slot); + host_lengths[i] = static_cast(segments[i].length); + } + + auto *base = static_cast(_opaque->buffer); + void *q_normalized = base + _opaque->q_offset; + void *k_normalized = base + _opaque->k_offset; + void *beta_bf16 = base + _opaque->beta_offset; + void *input_indices = base + _opaque->input_indices_offset; + void *state_indices = base + _opaque->state_indices_offset; + void *actual_seq_lengths = base + _opaque->actual_seq_lengths_offset; + void *state_staging = base + _opaque->state_staging_offset; + void *aclnn_workspace = base + _opaque->aclnn_workspace_offset; + + auto add_bytes = [](const void *ptr, size_t bytes) -> void * { + return const_cast( + static_cast(ptr)) + + bytes; + }; + std::array call_cursors{}; + for (size_t i = 0; i < segments.size(); ++i) { + const auto &segment = segments[i]; + const size_t q_bytes = segment.offset * _info.Hk * ACLNN_DIM * sizeof(uint16_t); + const size_t v_bytes = segment.offset * _info.Hv * ACLNN_DIM * sizeof(uint16_t); + const size_t gate_bytes = segment.offset * _info.Hv * sizeof(float); + CHECK_ACL(aclrtMemcpyAsync( + input_indices, + ACLNN_MAX_SEQ_LEN * sizeof(int32_t), + host_indices + i * ACLNN_MAX_SEQ_LEN, + ACLNN_MAX_SEQ_LEN * sizeof(int32_t), + ACL_MEMCPY_HOST_TO_DEVICE, stream)); + + RecurrentGdrNativeParams params{}; + params.B = segment.length; + params.Hk = _info.Hk; + params.Hv = _info.Hv; + params.pool_size = _info.pool_size; + params.initial_indices_i64 = false; + params.final_indices_i64 = false; + params.q_s0 = _info.q_strides[1]; + params.q_s2 = _info.q_strides[2]; + params.k_s0 = _info.k_strides[1]; + params.k_s2 = _info.k_strides[2]; + params.v_s0 = _info.v_strides[1]; + params.v_s2 = _info.v_strides[2]; + params.beta_s0 = _info.beta_strides[1]; + params.beta_s2 = _info.beta_strides[2]; + CHECK_STATUS(recurrent_gdr_native_preprocess_launch( + q_normalized, k_normalized, + add_bytes(v, v_bytes), beta_bf16, state_staging, + actual_seq_lengths, state_indices, + add_bytes(q, q_bytes), add_bytes(k, q_bytes), + add_bytes(v, v_bytes), add_bytes(beta, gate_bytes), + initial_state, input_indices, input_indices, + ¶ms, stream)); + CHECK_ACL(aclrtMemcpyAsync( + state_indices, + segment.length * sizeof(int32_t), + host_indices + i * ACLNN_MAX_SEQ_LEN, + segment.length * sizeof(int32_t), + ACL_MEMCPY_HOST_TO_DEVICE, stream)); + CHECK_ACL(aclrtMemcpyAsync( + actual_seq_lengths, sizeof(int32_t), + host_lengths + i, sizeof(int32_t), + ACL_MEMCPY_HOST_TO_DEVICE, stream)); + + const size_t length = segment.length; + const size_t call_index = call_cursors[length]++; + if (call_index >= _opaque->calls[length].size()) { + return INFINI_STATUS_BAD_PARAM; + } + const auto &call = *_opaque->calls[length][call_index]; + CHECK_ACL(AclSetTensorAddr(call.executor, 0, call.q->tensor, q_normalized)); + CHECK_ACL(AclSetTensorAddr(call.executor, 1, call.k->tensor, k_normalized)); + CHECK_ACL(AclSetTensorAddr(call.executor, 2, call.v->tensor, add_bytes(v, v_bytes))); + CHECK_ACL(AclSetTensorAddr(call.executor, 3, call.beta->tensor, beta_bf16)); + CHECK_ACL(AclSetTensorAddr(call.executor, 4, call.state->tensor, initial_state)); + CHECK_ACL(AclSetTensorAddr(call.executor, 5, call.actual_seq_lengths->tensor, actual_seq_lengths)); + CHECK_ACL(AclSetTensorAddr(call.executor, 6, call.state_indices->tensor, state_indices)); + CHECK_ACL(AclSetTensorAddr(call.executor, 7, call.g->tensor, add_bytes(g, gate_bytes))); + CHECK_ACL(AclSetTensorAddr(call.executor, 8, call.out->tensor, add_bytes(out, v_bytes))); + CHECK_ACL(aclnnRecurrentGatedDeltaRule( + aclnn_workspace, call.workspace_size, call.executor, stream)); + } + CHECK_ACL(aclrtSynchronizeStream(stream)); + return INFINI_STATUS_SUCCESS; + } + + GatedDeltaRuleAscendParams p{}; + p.data_dtype = static_cast(_info.data_dtype); + p.gate_dtype = static_cast(_info.gate_dtype); + p.use_qk_l2norm = _info.use_qk_l2norm; + p.has_cu_seqlens = _info.has_cu_seqlens; + p.cu_seqlens_i64 = _info.cu_seqlens_dtype == INFINI_DTYPE_I64; + p.has_initial_indices = _info.has_initial_state_indices; + p.initial_indices_i64 = _info.initial_state_indices_dtype == INFINI_DTYPE_I64; + p.has_final_indices = _info.has_final_state_indices; + p.final_indices_i64 = _info.final_state_indices_dtype == INFINI_DTYPE_I64; + p.B = _info.B; + p.T = _info.T; + p.total_tokens = _info.total_tokens; + p.Hk = _info.Hk; + p.Hv = _info.Hv; + p.Dk = _info.Dk; + p.Dv = _info.Dv; + p.pool_size = _info.pool_size; + p.value_heads_per_key_head = _info.value_heads_per_key_head; + p.q_scale = 1.0f / std::sqrt(static_cast(_info.Dk)); + for (int i = 0; i < 4; ++i) { + p.out_strides[i] = _info.out_strides[i]; + p.q_strides[i] = _info.q_strides[i]; + p.k_strides[i] = _info.k_strides[i]; + p.v_strides[i] = _info.v_strides[i]; + } + return gated_delta_rule_ascend_kernel_launch( + workspace, out, initial_state, final_state, q, k, v, g, beta, + cu_seqlens, initial_state_indices, final_state_indices, &p, stream); +} + +} // namespace op::chunk_gated_delta_rule::ascend diff --git a/src/infiniop/ops/chunk_gated_delta_rule/ascend/chunk_gated_delta_rule_ascend.h b/src/infiniop/ops/chunk_gated_delta_rule/ascend/chunk_gated_delta_rule_ascend.h new file mode 100644 index 000000000..01eb23aff --- /dev/null +++ b/src/infiniop/ops/chunk_gated_delta_rule/ascend/chunk_gated_delta_rule_ascend.h @@ -0,0 +1,8 @@ +#ifndef __CHUNK_GATED_DELTA_RULE_ASCEND_H__ +#define __CHUNK_GATED_DELTA_RULE_ASCEND_H__ + +#include "../chunk_gated_delta_rule.h" + +DESCRIPTOR(ascend) + +#endif diff --git a/src/infiniop/ops/chunk_gated_delta_rule/ascend/gated_delta_rule_ascend_kernel.cpp b/src/infiniop/ops/chunk_gated_delta_rule/ascend/gated_delta_rule_ascend_kernel.cpp new file mode 100644 index 000000000..050681865 --- /dev/null +++ b/src/infiniop/ops/chunk_gated_delta_rule/ascend/gated_delta_rule_ascend_kernel.cpp @@ -0,0 +1,492 @@ +#include "gated_delta_rule_ascend_kernel.h" +#include "../../../devices/ascend/ascend_kernel_common.h" + +#include + +using namespace AscendC; + +template +__aicore__ inline float gdrToFloat(T value) { + if constexpr (std::is_same::value) { + uint32_t bits = static_cast(*reinterpret_cast(&value)) << 16; + return *reinterpret_cast(&bits); + } else { + return static_cast(value); + } +} + +template +__aicore__ inline T gdrFromFloat(float value) { + if constexpr (std::is_same::value) { + uint32_t bits = *reinterpret_cast(&value); + uint16_t upper = static_cast(bits >> 16); + return *reinterpret_cast(&upper); + } else { + return static_cast(value); + } +} + +__aicore__ inline float gdrExp(float x) { + if (x > 88.0f) { + x = 88.0f; + } else if (x < -87.0f) { + return 0.0f; + } + float scaled = x * 1.4426950408889634f; + int32_t exponent = static_cast(scaled); + if (scaled < static_cast(exponent)) { + --exponent; + } + float r = x - static_cast(exponent) * 0.693145751953125f + - static_cast(exponent) * 1.428606765330187e-6f; + float polynomial = 1.0f + + r * (1.0f + r * (0.5f + r * (0.1666666716f + r * (0.0416666679f + r * (0.0083333338f + r * 0.0013888889f))))); + uint32_t bits = static_cast(exponent + 127) << 23; + return polynomial * *reinterpret_cast(&bits); +} + +__aicore__ inline float gdrLocalSum( + LocalTensor tensor, size_t count) { + float sum = 0.0f; + for (size_t i = 0; i < count; ++i) { + sum += tensor.GetValue(i); + } + return sum; +} + +__aicore__ inline float gdrRsqrt(float x) { + if (x <= 0.0f) { + return 0.0f; + } + float half_x = 0.5f * x; + uint32_t bits = *reinterpret_cast(&x); + bits = 0x5f375a86u - (bits >> 1); + float y = *reinterpret_cast(&bits); + y = y * (1.5f - half_x * y * y); + y = y * (1.5f - half_x * y * y); + y = y * (1.5f - half_x * y * y); + return y; +} + +__aicore__ inline int64_t gdrLoadIndex( + GM_ADDR ptr, bool is_i64, size_t index, int64_t fallback) { + if (ptr == nullptr) { + return fallback; + } + if (is_i64) { + return reinterpret_cast<__gm__ int64_t *>(ptr)[index]; + } + return static_cast( + reinterpret_cast<__gm__ int32_t *>(ptr)[index]); +} + +template +__aicore__ inline void gatedDeltaRuleProcess( + GM_ADDR workspace_ptr, + GM_ADDR out_ptr, + GM_ADDR initial_state_ptr, + GM_ADDR final_state_ptr, + GM_ADDR q_ptr, + GM_ADDR k_ptr, + GM_ADDR v_ptr, + GM_ADDR g_ptr, + GM_ADDR beta_ptr, + GM_ADDR cu_seqlens_ptr, + GM_ADDR initial_indices_ptr, + GM_ADDR final_indices_ptr, + bool use_qk_l2norm, + bool has_cu_seqlens, + bool cu_seqlens_i64, + bool has_initial_indices, + bool initial_indices_i64, + bool has_final_indices, + bool final_indices_i64, + size_t B, + size_t T_tokens, + size_t total_tokens, + size_t Hk, + size_t Hv, + size_t Dk, + size_t Dv, + size_t pool_size, + size_t value_heads_per_key_head, + float q_scale, + ptrdiff_t out_s0, + ptrdiff_t out_s1, + ptrdiff_t out_s2, + ptrdiff_t out_s3, + ptrdiff_t q_s0, + ptrdiff_t q_s1, + ptrdiff_t q_s2, + ptrdiff_t q_s3, + ptrdiff_t k_s0, + ptrdiff_t k_s1, + ptrdiff_t k_s2, + ptrdiff_t k_s3, + ptrdiff_t v_s0, + ptrdiff_t v_s1, + ptrdiff_t v_s2, + ptrdiff_t v_s3) { + + size_t block = GetBlockIdx(); + size_t request = block / Hv; + size_t vh = block % Hv; + if (request >= B) { + return; + } + size_t kh = vh / value_heads_per_key_head; + + GlobalTensor work; + GlobalTensor gate; + GlobalTensor beta; + GlobalTensor out; + GlobalTensor initial_state; + GlobalTensor final_state; + GlobalTensor q; + GlobalTensor k; + GlobalTensor v; + work.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(workspace_ptr)); + gate.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(g_ptr)); + beta.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(beta_ptr)); + out.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(out_ptr)); + initial_state.SetGlobalBuffer( + reinterpret_cast<__gm__ T *>(initial_state_ptr)); + final_state.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(final_state_ptr)); + q.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(q_ptr)); + k.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(k_ptr)); + v.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(v_ptr)); + + int64_t begin = 0; + int64_t end = static_cast(T_tokens); + if (has_cu_seqlens) { + begin = gdrLoadIndex( + cu_seqlens_ptr, cu_seqlens_i64, request, 0); + end = gdrLoadIndex( + cu_seqlens_ptr, cu_seqlens_i64, request + 1, 0); + if (begin < 0 || end < begin + || end > static_cast(total_tokens)) { + return; + } + } + + int64_t read_slot = has_initial_indices + ? gdrLoadIndex( + initial_indices_ptr, initial_indices_i64, + request, static_cast(request)) + : static_cast(request); + int64_t write_slot = has_final_indices + ? gdrLoadIndex( + final_indices_ptr, final_indices_i64, + request, static_cast(request)) + : static_cast(request); + if (read_slot < 0 || read_slot >= static_cast(pool_size) + || (has_final_indices + && (write_slot < 0 + || write_slot >= static_cast(pool_size)))) { + return; + } + + size_t matrix_size = Dv * Dk; + size_t initial_base = (static_cast(read_slot) * Hv + vh) * matrix_size; + + const bool ub_fast_path = Dk == 128 && Dv == 128 + && q_s3 == 1 && k_s3 == 1 && v_s3 == 1 && out_s3 == 1; + if (ub_fast_path) { + constexpr size_t VECTOR_LEN = 128; + constexpr size_t MATRIX_LEN = VECTOR_LEN * VECTOR_LEN; + TPipe pipe; + TBuf state_data_buf, io_data_buf; + TBuf state_float_buf; + TBuf q_float_buf, k_float_buf; + TBuf v_float_buf, out_float_buf; + TBuf tmp_float_buf; + TBuf reduce_work_buf, reduce_result_buf; + pipe.InitBuffer(state_data_buf, MATRIX_LEN * sizeof(T)); + pipe.InitBuffer(io_data_buf, VECTOR_LEN * sizeof(T)); + pipe.InitBuffer(state_float_buf, MATRIX_LEN * sizeof(float)); + pipe.InitBuffer(q_float_buf, VECTOR_LEN * sizeof(float)); + pipe.InitBuffer(k_float_buf, VECTOR_LEN * sizeof(float)); + pipe.InitBuffer(v_float_buf, VECTOR_LEN * sizeof(float)); + pipe.InitBuffer(out_float_buf, VECTOR_LEN * sizeof(float)); + pipe.InitBuffer(tmp_float_buf, VECTOR_LEN * sizeof(float)); + pipe.InitBuffer(reduce_work_buf, VECTOR_LEN * sizeof(float)); + pipe.InitBuffer(reduce_result_buf, 8 * sizeof(float)); + + LocalTensor state_data = state_data_buf.Get(); + LocalTensor io_data = io_data_buf.Get(); + LocalTensor state_float = state_float_buf.Get(); + LocalTensor q_float = q_float_buf.Get(); + LocalTensor k_float = k_float_buf.Get(); + LocalTensor v_float = v_float_buf.Get(); + LocalTensor out_float = out_float_buf.Get(); + LocalTensor tmp_float = tmp_float_buf.Get(); + LocalTensor reduce_work = reduce_work_buf.Get(); + LocalTensor reduce_result = reduce_result_buf.Get(); + + if constexpr (std::is_same::value) { + DataCopy(state_float, initial_state[initial_base], MATRIX_LEN); + } else { + DataCopy(state_data, initial_state[initial_base], MATRIX_LEN); + Cast(state_float, state_data, RoundMode::CAST_NONE, MATRIX_LEN); + } + + for (int64_t token = begin; token < end; ++token) { + size_t token_batch = has_cu_seqlens ? 0 : request; + size_t token_index = static_cast(token); + ptrdiff_t q_base = static_cast(token_batch) * q_s0 + + static_cast(token_index) * q_s1 + + static_cast(kh) * q_s2; + ptrdiff_t k_base = static_cast(token_batch) * k_s0 + + static_cast(token_index) * k_s1 + + static_cast(kh) * k_s2; + ptrdiff_t v_base = static_cast(token_batch) * v_s0 + + static_cast(token_index) * v_s1 + + static_cast(vh) * v_s2; + ptrdiff_t out_base = static_cast(token_batch) * out_s0 + + static_cast(token_index) * out_s1 + + static_cast(vh) * out_s2; + + if constexpr (std::is_same::value) { + DataCopy(q_float, q[q_base], VECTOR_LEN); + DataCopy(k_float, k[k_base], VECTOR_LEN); + DataCopy(v_float, v[v_base], VECTOR_LEN); + } else { + DataCopy(io_data, q[q_base], VECTOR_LEN); + Cast(q_float, io_data, RoundMode::CAST_NONE, VECTOR_LEN); + DataCopy(io_data, k[k_base], VECTOR_LEN); + Cast(k_float, io_data, RoundMode::CAST_NONE, VECTOR_LEN); + DataCopy(io_data, v[v_base], VECTOR_LEN); + Cast(v_float, io_data, RoundMode::CAST_NONE, VECTOR_LEN); + } + + if (use_qk_l2norm) { + Mul(tmp_float, q_float, q_float, VECTOR_LEN); + float q_norm_inv = gdrRsqrt( + gdrLocalSum(tmp_float, VECTOR_LEN)); + Muls(q_float, q_float, q_norm_inv * q_scale, VECTOR_LEN); + + Mul(tmp_float, k_float, k_float, VECTOR_LEN); + float k_norm_inv = gdrRsqrt( + gdrLocalSum(tmp_float, VECTOR_LEN)); + Muls(k_float, k_float, k_norm_inv, VECTOR_LEN); + } else if (q_scale != 1.0f) { + Muls(q_float, q_float, q_scale, VECTOR_LEN); + } + + size_t gate_offset = (token_batch * T_tokens + token_index) * Hv + vh; + float gate_value = gdrExp(gate.GetValue(gate_offset)); + float beta_value = beta.GetValue(gate_offset); + Muls(state_float, state_float, gate_value, MATRIX_LEN); + + for (size_t dv = 0; dv < VECTOR_LEN; ++dv) { + LocalTensor state_row = state_float[dv * VECTOR_LEN]; + Mul(tmp_float, state_row, k_float, VECTOR_LEN); + float kv_memory = gdrLocalSum(tmp_float, VECTOR_LEN); + float delta = (v_float.GetValue(dv) - kv_memory) * beta_value; + Muls(tmp_float, k_float, delta, VECTOR_LEN); + Add(state_row, state_row, tmp_float, VECTOR_LEN); + Mul(tmp_float, state_row, q_float, VECTOR_LEN); + out_float.SetValue( + dv, gdrLocalSum(tmp_float, VECTOR_LEN)); + } + + if constexpr (std::is_same::value) { + DataCopy(out[out_base], out_float, VECTOR_LEN); + } else { + Cast(io_data, out_float, RoundMode::CAST_RINT, VECTOR_LEN); + DataCopy(out[out_base], io_data, VECTOR_LEN); + } + } + + size_t destination_base = has_final_indices + ? (static_cast(write_slot) * Hv + vh) * MATRIX_LEN + : (request * Hv + vh) * MATRIX_LEN; + if constexpr (std::is_same::value) { + if (has_final_indices) { + DataCopy(initial_state[destination_base], state_float, + MATRIX_LEN); + } else { + DataCopy(final_state[destination_base], state_float, + MATRIX_LEN); + } + } else { + Cast(state_data, state_float, RoundMode::CAST_RINT, MATRIX_LEN); + if (has_final_indices) { + DataCopy(initial_state[destination_base], state_data, + MATRIX_LEN); + } else { + DataCopy(final_state[destination_base], state_data, + MATRIX_LEN); + } + } + return; + } + + size_t work_base = (request * Hv + vh) * matrix_size; + for (size_t dv = 0; dv < Dv; ++dv) { + for (size_t dk = 0; dk < Dk; ++dk) { + size_t element = dv * Dk + dk; + work.SetValue( + work_base + element, + gdrToFloat(initial_state.GetValue(initial_base + element))); + } + } + + for (int64_t token = begin; token < end; ++token) { + size_t token_batch = has_cu_seqlens ? 0 : request; + size_t token_index = static_cast(token); + ptrdiff_t q_base = static_cast(token_batch) * q_s0 + + static_cast(token_index) * q_s1 + + static_cast(kh) * q_s2; + ptrdiff_t k_base = static_cast(token_batch) * k_s0 + + static_cast(token_index) * k_s1 + + static_cast(kh) * k_s2; + float q_norm_inv = 1.0f; + float k_norm_inv = 1.0f; + if (use_qk_l2norm) { + float q_sum = 0.0f; + float k_sum = 0.0f; + for (size_t dk = 0; dk < Dk; ++dk) { + float q_value = gdrToFloat(q.GetValue(q_base + dk * q_s3)); + float k_value = gdrToFloat(k.GetValue(k_base + dk * k_s3)); + q_sum += q_value * q_value; + k_sum += k_value * k_value; + } + q_norm_inv = gdrRsqrt(q_sum); + k_norm_inv = gdrRsqrt(k_sum); + } + + size_t gate_offset = (token_batch * T_tokens + token_index) * Hv + vh; + float gate_value = gdrExp(gate.GetValue(gate_offset)); + float beta_value = beta.GetValue(gate_offset); + ptrdiff_t v_base = static_cast(token_batch) * v_s0 + + static_cast(token_index) * v_s1 + + static_cast(vh) * v_s2; + ptrdiff_t out_base = static_cast(token_batch) * out_s0 + + static_cast(token_index) * out_s1 + + static_cast(vh) * out_s2; + + for (size_t dv = 0; dv < Dv; ++dv) { + float kv_memory = 0.0f; + size_t row_base = work_base + dv * Dk; + for (size_t dk = 0; dk < Dk; ++dk) { + float state = work.GetValue(row_base + dk) * gate_value; + work.SetValue(row_base + dk, state); + float key = gdrToFloat(k.GetValue(k_base + dk * k_s3)) * k_norm_inv; + kv_memory += state * key; + } + float value = gdrToFloat(v.GetValue(v_base + dv * v_s3)); + float delta = (value - kv_memory) * beta_value; + float output = 0.0f; + for (size_t dk = 0; dk < Dk; ++dk) { + float key = gdrToFloat(k.GetValue(k_base + dk * k_s3)) * k_norm_inv; + float state = work.GetValue(row_base + dk) + key * delta; + work.SetValue(row_base + dk, state); + float query = gdrToFloat(q.GetValue(q_base + dk * q_s3)) + * q_norm_inv * q_scale; + output += state * query; + } + out.SetValue( + out_base + dv * out_s3, gdrFromFloat(output)); + } + } + + size_t destination_base; + if (has_final_indices) { + destination_base = (static_cast(write_slot) * Hv + vh) * matrix_size; + } else { + destination_base = (request * Hv + vh) * matrix_size; + } + for (size_t element = 0; element < matrix_size; ++element) { + T value = gdrFromFloat(work.GetValue(work_base + element)); + if (has_final_indices) { + initial_state.SetValue(destination_base + element, value); + } else { + final_state.SetValue(destination_base + element, value); + } + } +} + +#define DEFINE_GDR_KERNEL(NAME, TYPE) \ + __global__ __aicore__ void NAME( \ + GM_ADDR workspace, GM_ADDR out, GM_ADDR initial_state, \ + GM_ADDR final_state, GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, \ + GM_ADDR beta, GM_ADDR cu_seqlens, GM_ADDR initial_indices, \ + GM_ADDR final_indices, bool use_qk_l2norm, bool has_cu_seqlens, \ + bool cu_seqlens_i64, bool has_initial_indices, \ + bool initial_indices_i64, bool has_final_indices, \ + bool final_indices_i64, size_t B, size_t T_tokens, \ + size_t total_tokens, size_t Hk, size_t Hv, size_t Dk, size_t Dv, \ + size_t pool_size, size_t value_heads_per_key_head, float q_scale, \ + ptrdiff_t out_s0, ptrdiff_t out_s1, ptrdiff_t out_s2, \ + ptrdiff_t out_s3, ptrdiff_t q_s0, ptrdiff_t q_s1, ptrdiff_t q_s2, \ + ptrdiff_t q_s3, ptrdiff_t k_s0, ptrdiff_t k_s1, ptrdiff_t k_s2, \ + ptrdiff_t k_s3, ptrdiff_t v_s0, ptrdiff_t v_s1, ptrdiff_t v_s2, \ + ptrdiff_t v_s3) { \ + gatedDeltaRuleProcess( \ + workspace, out, initial_state, final_state, q, k, v, g, beta, \ + cu_seqlens, initial_indices, final_indices, use_qk_l2norm, \ + has_cu_seqlens, cu_seqlens_i64, has_initial_indices, \ + initial_indices_i64, has_final_indices, final_indices_i64, B, \ + T_tokens, total_tokens, Hk, Hv, Dk, Dv, pool_size, \ + value_heads_per_key_head, q_scale, out_s0, out_s1, out_s2, \ + out_s3, q_s0, q_s1, q_s2, q_s3, k_s0, k_s1, k_s2, k_s3, \ + v_s0, v_s1, v_s2, v_s3); \ + } + +DEFINE_GDR_KERNEL(gated_delta_rule_half, half) +DEFINE_GDR_KERNEL(gated_delta_rule_bf16, bfloat16_t) +DEFINE_GDR_KERNEL(gated_delta_rule_float, float) +#undef DEFINE_GDR_KERNEL + +extern "C" infiniStatus_t gated_delta_rule_ascend_kernel_launch( + void *workspace, + void *out, + void *initial_state, + void *final_state, + const void *q, + const void *k, + const void *v, + const void *g, + const void *beta, + const void *cu_seqlens, + const void *initial_state_indices, + const void *final_state_indices, + const GatedDeltaRuleAscendParams *p, + void *stream) { + + if (p->B == 0 || p->Hv == 0) { + return INFINI_STATUS_SUCCESS; + } + uint32_t blocks = static_cast(p->B * p->Hv); + +#define LAUNCH_GDR(DTYPE, NAME) \ + case DTYPE: \ + NAME<<>>( \ + workspace, out, initial_state, final_state, \ + const_cast(q), const_cast(k), \ + const_cast(v), const_cast(g), \ + const_cast(beta), const_cast(cu_seqlens), \ + const_cast(initial_state_indices), \ + const_cast(final_state_indices), p->use_qk_l2norm, \ + p->has_cu_seqlens, p->cu_seqlens_i64, p->has_initial_indices, \ + p->initial_indices_i64, p->has_final_indices, \ + p->final_indices_i64, p->B, p->T, p->total_tokens, p->Hk, \ + p->Hv, p->Dk, p->Dv, p->pool_size, \ + p->value_heads_per_key_head, p->q_scale, p->out_strides[0], \ + p->out_strides[1], p->out_strides[2], p->out_strides[3], \ + p->q_strides[0], p->q_strides[1], p->q_strides[2], \ + p->q_strides[3], p->k_strides[0], p->k_strides[1], \ + p->k_strides[2], p->k_strides[3], p->v_strides[0], \ + p->v_strides[1], p->v_strides[2], p->v_strides[3]); \ + return INFINI_STATUS_SUCCESS; + + switch (static_cast(p->data_dtype)) { + LAUNCH_GDR(INFINI_DTYPE_F16, gated_delta_rule_half) + LAUNCH_GDR(INFINI_DTYPE_BF16, gated_delta_rule_bf16) + LAUNCH_GDR(INFINI_DTYPE_F32, gated_delta_rule_float) + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +#undef LAUNCH_GDR +} diff --git a/src/infiniop/ops/chunk_gated_delta_rule/ascend/gated_delta_rule_ascend_kernel.h b/src/infiniop/ops/chunk_gated_delta_rule/ascend/gated_delta_rule_ascend_kernel.h new file mode 100644 index 000000000..0ebb61ee2 --- /dev/null +++ b/src/infiniop/ops/chunk_gated_delta_rule/ascend/gated_delta_rule_ascend_kernel.h @@ -0,0 +1,50 @@ +#ifndef __GATED_DELTA_RULE_ASCEND_KERNEL_H__ +#define __GATED_DELTA_RULE_ASCEND_KERNEL_H__ + +#include "../../../../../include/infinicore.h" +#include +#include + +struct GatedDeltaRuleAscendParams { + int32_t data_dtype; + int32_t gate_dtype; + bool use_qk_l2norm; + bool has_cu_seqlens; + bool cu_seqlens_i64; + bool has_initial_indices; + bool initial_indices_i64; + bool has_final_indices; + bool final_indices_i64; + size_t B; + size_t T; + size_t total_tokens; + size_t Hk; + size_t Hv; + size_t Dk; + size_t Dv; + size_t pool_size; + size_t value_heads_per_key_head; + float q_scale; + ptrdiff_t out_strides[4]; + ptrdiff_t q_strides[4]; + ptrdiff_t k_strides[4]; + ptrdiff_t v_strides[4]; +}; + +extern "C" infiniStatus_t gated_delta_rule_ascend_kernel_launch( + void *workspace, + void *out, + void *initial_state, + void *final_state, + const void *q, + const void *k, + const void *v, + const void *g, + const void *beta, + const void *cu_seqlens, + const void *initial_state_indices, + const void *final_state_indices, + const GatedDeltaRuleAscendParams *params, + void *stream); + +#endif diff --git a/src/infiniop/ops/chunk_gated_delta_rule/operator.cc b/src/infiniop/ops/chunk_gated_delta_rule/operator.cc index 19b5198bf..d4e13ae1c 100644 --- a/src/infiniop/ops/chunk_gated_delta_rule/operator.cc +++ b/src/infiniop/ops/chunk_gated_delta_rule/operator.cc @@ -13,6 +13,9 @@ #ifdef ENABLE_MOORE_API #include "moore/chunk_gated_delta_rule_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/chunk_gated_delta_rule_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateChunkGatedDeltaRuleDescriptor( infiniopHandle_t handle, @@ -53,6 +56,9 @@ __INFINI_C infiniStatus_t infiniopCreateChunkGatedDeltaRuleDescriptor( #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) #endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend) +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -79,6 +85,9 @@ __INFINI_C infiniStatus_t infiniopGetChunkGatedDeltaRuleWorkspaceSize( #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend) +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -112,6 +121,9 @@ __INFINI_C infiniStatus_t infiniopChunkGatedDeltaRule( #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore) #endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend) +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -137,6 +149,9 @@ __INFINI_C infiniStatus_t infiniopDestroyChunkGatedDeltaRuleDescriptor( #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) #endif +#ifdef ENABLE_ASCEND_API + DESTROY(INFINI_DEVICE_ASCEND, ascend) +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend.cc b/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend.cc new file mode 100644 index 000000000..c696fbe97 --- /dev/null +++ b/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend.cc @@ -0,0 +1,71 @@ +#include "fused_gated_delta_net_gating_ascend.h" +#include "../../../devices/ascend/ascend_handle.h" + +namespace op::fused_gated_delta_net_gating::ascend { +extern "C" infiniStatus_t fused_gated_delta_net_gating_kernel_launch( + void *g, void *beta_output, + const void *A_log, const void *a, const void *b, const void *dt_bias, + infiniDtype_t dtype, size_t total, size_t seq_len, size_t hidden, + ptrdiff_t g_s0, ptrdiff_t g_s1, ptrdiff_t g_s2, + ptrdiff_t beta_s0, ptrdiff_t beta_s1, ptrdiff_t beta_s2, + ptrdiff_t A_log_s0, + ptrdiff_t a_s0, ptrdiff_t a_s1, ptrdiff_t a_s2, + ptrdiff_t b_s0, ptrdiff_t b_s1, ptrdiff_t b_s2, + ptrdiff_t dt_bias_s0, float beta, float threshold, void *stream); + +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t g_desc, + infiniopTensorDescriptor_t beta_output_desc, + infiniopTensorDescriptor_t A_log_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t dt_bias_desc, + float beta, + float threshold) { + + auto result = FusedGatedDeltaNetGatingInfo::create( + g_desc, beta_output_desc, A_log_desc, a_desc, b_desc, dt_bias_desc, + beta, threshold); + CHECK_RESULT(result); + + auto handle_ascend = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + new Opaque{}, result.take(), 0, + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *g, + void *beta_output, + const void *A_log, + const void *a, + const void *b, + const void *dt_bias, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + return fused_gated_delta_net_gating_kernel_launch( + g, beta_output, A_log, a, b, dt_bias, + _info.input_dtype, _info.numel(), _info.seq_len, _info.hidden, + _info.g_strides[0], _info.g_strides[1], _info.g_strides[2], + _info.beta_output_strides[0], _info.beta_output_strides[1], + _info.beta_output_strides[2], _info.A_log_strides[0], + _info.a_strides[0], _info.a_strides[1], _info.a_strides[2], + _info.b_strides[0], _info.b_strides[1], _info.b_strides[2], + _info.dt_bias_strides[0], _info.beta, _info.threshold, stream); +} + +} // namespace op::fused_gated_delta_net_gating::ascend diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend.h b/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend.h new file mode 100644 index 000000000..9cfb1b263 --- /dev/null +++ b/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend.h @@ -0,0 +1,8 @@ +#ifndef __FUSED_GATED_DELTA_NET_GATING_ASCEND_H__ +#define __FUSED_GATED_DELTA_NET_GATING_ASCEND_H__ + +#include "../fused_gated_delta_net_gating.h" + +DESCRIPTOR(ascend); + +#endif diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend_kernel.cpp b/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend_kernel.cpp new file mode 100644 index 000000000..8c8ad269c --- /dev/null +++ b/src/infiniop/ops/fused_gated_delta_net_gating/ascend/fused_gated_delta_net_gating_ascend_kernel.cpp @@ -0,0 +1,225 @@ +#include "../../../devices/ascend/ascend_kernel_common.h" + +#include +#include + +using namespace AscendC; + +template +__aicore__ inline float gatingToFloat(T value) { + if constexpr (std::is_same::value) { + uint32_t bits = static_cast( + *reinterpret_cast(&value)) + << 16; + return *reinterpret_cast(&bits); + } else { + return static_cast(value); + } +} + +// AscendC does not expose the host scalar exp/log functions to AICore code. +// These helpers use range reduction and short polynomials instead. +__aicore__ inline float gatingExp(float x) { + if (x > 88.0f) { + x = 88.0f; + } else if (x < -87.0f) { + return 0.0f; + } + + constexpr float inv_ln2 = 1.4426950408889634f; + constexpr float ln2_hi = 0.693145751953125f; + constexpr float ln2_lo = 1.428606765330187e-6f; + float scaled = x * inv_ln2; + int32_t exponent = static_cast(scaled); + if (scaled < static_cast(exponent)) { + --exponent; + } + float r = x - static_cast(exponent) * ln2_hi + - static_cast(exponent) * ln2_lo; + float polynomial = 1.0f + + r * (1.0f + r * (0.5f + r * (0.1666666716f + r * (0.0416666679f + r * (0.0083333338f + r * 0.0013888889f))))); + uint32_t scale_bits = static_cast(exponent + 127) << 23; + float scale = *reinterpret_cast(&scale_bits); + return polynomial * scale; +} + +__aicore__ inline float gatingLog(float x) { + uint32_t bits = *reinterpret_cast(&x); + int32_t exponent = static_cast((bits >> 23) & 0xff) - 127; + bits = (bits & 0x007fffffu) | 0x3f800000u; + float mantissa = *reinterpret_cast(&bits); + + float z = (mantissa - 1.0f) / (mantissa + 1.0f); + float z2 = z * z; + float series = z * (1.0f + z2 * (0.3333333333f + z2 * (0.2f + z2 * (0.1428571429f + z2 * (0.1111111111f + z2 * (0.0909090909f + z2 * 0.0769230769f)))))); + return static_cast(exponent) * 0.6931471805599453f + + 2.0f * series; +} + +__aicore__ inline float gatingSigmoid(float x) { + if (x >= 0.0f) { + float z = gatingExp(-x); + return 1.0f / (1.0f + z); + } + float z = gatingExp(x); + return z / (1.0f + z); +} + +__aicore__ inline float gatingSoftplus(float x, float beta, float threshold) { + float bx = beta * x; + if (bx > threshold) { + return x; + } + if (bx < -20.0f) { + return gatingExp(bx) / beta; + } + return gatingLog(1.0f + gatingExp(bx)) / beta; +} + +template +__aicore__ inline void fused_gated_delta_net_gating_process( + GM_ADDR g_ptr, + GM_ADDR beta_output_ptr, + GM_ADDR A_log_ptr, + GM_ADDR a_ptr, + GM_ADDR b_ptr, + GM_ADDR dt_bias_ptr, + size_t total, + size_t seq_len, + size_t hidden, + ptrdiff_t g_s0, + ptrdiff_t g_s1, + ptrdiff_t g_s2, + ptrdiff_t beta_s0, + ptrdiff_t beta_s1, + ptrdiff_t beta_s2, + ptrdiff_t A_log_s0, + ptrdiff_t a_s0, + ptrdiff_t a_s1, + ptrdiff_t a_s2, + ptrdiff_t b_s0, + ptrdiff_t b_s1, + ptrdiff_t b_s2, + ptrdiff_t dt_bias_s0, + float beta, + float threshold) { + + GlobalTensor g; + GlobalTensor beta_output; + GlobalTensor A_log; + GlobalTensor a; + GlobalTensor b; + GlobalTensor dt_bias; + g.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(g_ptr)); + beta_output.SetGlobalBuffer( + reinterpret_cast<__gm__ float *>(beta_output_ptr)); + A_log.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(A_log_ptr)); + a.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(a_ptr)); + b.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(b_ptr)); + dt_bias.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(dt_bias_ptr)); + + for (size_t linear = 0; linear < total; ++linear) { + size_t h = linear % hidden; + size_t tmp = linear / hidden; + size_t s = tmp % seq_len; + size_t batch = tmp / seq_len; + + ptrdiff_t g_off = static_cast(batch) * g_s0 + + static_cast(s) * g_s1 + + static_cast(h) * g_s2; + ptrdiff_t beta_off = static_cast(batch) * beta_s0 + + static_cast(s) * beta_s1 + + static_cast(h) * beta_s2; + ptrdiff_t a_off = static_cast(batch) * a_s0 + + static_cast(s) * a_s1 + + static_cast(h) * a_s2; + ptrdiff_t b_off = static_cast(batch) * b_s0 + + static_cast(s) * b_s1 + + static_cast(h) * b_s2; + + float x = gatingToFloat(a.GetValue(a_off)) + + gatingToFloat( + dt_bias.GetValue(static_cast(h) * dt_bias_s0)); + float decay = -gatingExp(gatingToFloat( + A_log.GetValue(static_cast(h) * A_log_s0))); + g.SetValue(g_off, decay * gatingSoftplus(x, beta, threshold)); + beta_output.SetValue( + beta_off, gatingSigmoid(gatingToFloat(b.GetValue(b_off)))); + } +} + +#define DEFINE_GATING_KERNEL(NAME, TYPE) \ + __global__ __aicore__ void NAME( \ + GM_ADDR g, GM_ADDR beta_output, GM_ADDR A_log, GM_ADDR a, \ + GM_ADDR b, GM_ADDR dt_bias, size_t total, size_t seq_len, \ + size_t hidden, ptrdiff_t g_s0, ptrdiff_t g_s1, ptrdiff_t g_s2, \ + ptrdiff_t beta_s0, ptrdiff_t beta_s1, ptrdiff_t beta_s2, \ + ptrdiff_t A_log_s0, ptrdiff_t a_s0, ptrdiff_t a_s1, \ + ptrdiff_t a_s2, ptrdiff_t b_s0, ptrdiff_t b_s1, ptrdiff_t b_s2, \ + ptrdiff_t dt_bias_s0, float beta, float threshold) { \ + fused_gated_delta_net_gating_process( \ + g, beta_output, A_log, a, b, dt_bias, total, seq_len, hidden, \ + g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2, A_log_s0, \ + a_s0, a_s1, a_s2, b_s0, b_s1, b_s2, dt_bias_s0, beta, \ + threshold); \ + } + +DEFINE_GATING_KERNEL(fused_gating_half, half) +DEFINE_GATING_KERNEL(fused_gating_float, float) +DEFINE_GATING_KERNEL(fused_gating_bf16, bfloat16_t) +#undef DEFINE_GATING_KERNEL + +extern "C" infiniStatus_t fused_gated_delta_net_gating_kernel_launch( + void *g, + void *beta_output, + const void *A_log, + const void *a, + const void *b, + const void *dt_bias, + infiniDtype_t dtype, + size_t total, + size_t seq_len, + size_t hidden, + ptrdiff_t g_s0, + ptrdiff_t g_s1, + ptrdiff_t g_s2, + ptrdiff_t beta_s0, + ptrdiff_t beta_s1, + ptrdiff_t beta_s2, + ptrdiff_t A_log_s0, + ptrdiff_t a_s0, + ptrdiff_t a_s1, + ptrdiff_t a_s2, + ptrdiff_t b_s0, + ptrdiff_t b_s1, + ptrdiff_t b_s2, + ptrdiff_t dt_bias_s0, + float beta, + float threshold, + void *stream) { + + if (total == 0) { + return INFINI_STATUS_SUCCESS; + } + +#define LAUNCH_GATING(DTYPE, NAME) \ + case DTYPE: \ + NAME<<<1, nullptr, stream>>>( \ + g, beta_output, const_cast(A_log), \ + const_cast(a), const_cast(b), \ + const_cast(dt_bias), total, seq_len, hidden, \ + g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2, A_log_s0, \ + a_s0, a_s1, a_s2, b_s0, b_s1, b_s2, dt_bias_s0, beta, \ + threshold); \ + return INFINI_STATUS_SUCCESS; + + switch (dtype) { + LAUNCH_GATING(INFINI_DTYPE_F16, fused_gating_half) + LAUNCH_GATING(INFINI_DTYPE_BF16, fused_gating_bf16) + LAUNCH_GATING(INFINI_DTYPE_F32, fused_gating_float) + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + +#undef LAUNCH_GATING +} diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc b/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc index 14296e846..d18a12468 100644 --- a/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc +++ b/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc @@ -12,6 +12,9 @@ #include "moore/fused_gated_delta_net_gating_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/fused_gated_delta_net_gating_ascend.h" +#endif __INFINI_C __export infiniStatus_t infiniopCreateFusedGatedDeltaNetGatingDescriptor( infiniopHandle_t handle, @@ -47,6 +50,9 @@ infiniopCreateFusedGatedDeltaNetGatingDescriptor( #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); +#endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif @@ -86,6 +92,9 @@ infiniopGetFusedGatedDeltaNetGatingWorkspaceSize( #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend); +#endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif @@ -133,6 +142,9 @@ infiniopFusedGatedDeltaNetGating( #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); +#endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif @@ -171,6 +183,9 @@ infiniopDestroyFusedGatedDeltaNetGatingDescriptor( #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); +#endif #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif diff --git a/src/infiniop/ops/gelu/ascend/gelu_ascend.cc b/src/infiniop/ops/gelu/ascend/gelu_ascend.cc new file mode 100644 index 000000000..b39816e70 --- /dev/null +++ b/src/infiniop/ops/gelu/ascend/gelu_ascend.cc @@ -0,0 +1,71 @@ +#include "gelu_ascend.h" + +#include "../../../devices/ascend/aclnn_executor.h" +#include +#include + +namespace op::gelu::ascend { + +struct Descriptor::Opaque { + device::ascend::AclnnExecutor op; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + std::vector input_descs) { + + auto status = device::ascend::validateAclnnElementwise(output_desc, input_descs, 1); + if (status != INFINI_STATUS_SUCCESS) { + return status; + } + CHECK_DTYPE(output_desc->dtype(), + INFINI_DTYPE_F16, INFINI_DTYPE_F32, + INFINI_DTYPE_F64, INFINI_DTYPE_BF16); + + auto opaque = std::make_unique(); + opaque->op.tensors = { + new aclnnTensorDescriptor(input_descs[0]), + new aclnnTensorDescriptor(output_desc), + }; + + CHECK_ACL(aclnnGeluGetWorkspaceSize( + opaque->op.tensors[0]->tensor, + opaque->op.tensors[1]->tensor, + &opaque->op.workspace_size, + &opaque->op.executor)); + aclSetAclOpExecutorRepeatable(opaque->op.executor); + + auto handle_ascend = reinterpret_cast(handle); + auto workspace_size = opaque->op.workspace_size; + *desc_ptr = new Descriptor( + opaque.release(), workspace_size, + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *output, std::vector inputs, + void *stream) const { + + if (inputs.size() != 1) { + return INFINI_STATUS_BAD_PARAM; + } + if (workspace_size < workspaceSize()) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + _opaque->op.bind({const_cast(inputs[0]), output}); + CHECK_ACL(aclnnGelu( + workspace, workspace_size, _opaque->op.executor, + static_cast(stream))); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::gelu::ascend diff --git a/src/infiniop/ops/gelu/ascend/gelu_ascend.h b/src/infiniop/ops/gelu/ascend/gelu_ascend.h new file mode 100644 index 000000000..66790f054 --- /dev/null +++ b/src/infiniop/ops/gelu/ascend/gelu_ascend.h @@ -0,0 +1,8 @@ +#ifndef __GELU_ASCEND_H__ +#define __GELU_ASCEND_H__ + +#include "../../../devices/ascend/aclnn_elementwise.h" + +ACLNN_ELEMENTWISE_DESCRIPTOR(gelu) + +#endif // __GELU_ASCEND_H__ diff --git a/src/infiniop/ops/gelu/operator.cc b/src/infiniop/ops/gelu/operator.cc index d07762828..6e002e266 100644 --- a/src/infiniop/ops/gelu/operator.cc +++ b/src/infiniop/ops/gelu/operator.cc @@ -20,6 +20,9 @@ #ifdef ENABLE_CAMBRICON_API #include "bang/gelu_bang.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/gelu_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateGeluDescriptor( infiniopHandle_t handle, @@ -67,6 +70,9 @@ __INFINI_C infiniStatus_t infiniopCreateGeluDescriptor( #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia); #endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -113,6 +119,9 @@ __INFINI_C infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia); #endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -167,6 +176,9 @@ __INFINI_C infiniStatus_t infiniopGelu( #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia); #endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -215,6 +227,9 @@ infiniopDestroyGeluDescriptor(infiniopGeluDescriptor_t desc) { #ifdef ENABLE_ALI_API DELETE(INFINI_DEVICE_ALI, nvidia); #endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/gelutanh/ascend/gelutanh_ascend.cc b/src/infiniop/ops/gelutanh/ascend/gelutanh_ascend.cc new file mode 100644 index 000000000..9adcc2b6b --- /dev/null +++ b/src/infiniop/ops/gelutanh/ascend/gelutanh_ascend.cc @@ -0,0 +1,72 @@ +#include "gelutanh_ascend.h" + +#include "../../../devices/ascend/aclnn_executor.h" +#include +#include + +namespace op::gelutanh::ascend { + +struct Descriptor::Opaque { + device::ascend::AclnnExecutor op; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + std::vector input_descs) { + + auto status = device::ascend::validateAclnnElementwise(output_desc, input_descs, 1); + if (status != INFINI_STATUS_SUCCESS) { + return status; + } + CHECK_DTYPE(output_desc->dtype(), + INFINI_DTYPE_F16, INFINI_DTYPE_F32, + INFINI_DTYPE_F64, INFINI_DTYPE_BF16); + + auto opaque = std::make_unique(); + opaque->op.tensors = { + new aclnnTensorDescriptor(input_descs[0]), + new aclnnTensorDescriptor(output_desc), + }; + + CHECK_ACL(aclnnGeluV2GetWorkspaceSize( + opaque->op.tensors[0]->tensor, + 1, + opaque->op.tensors[1]->tensor, + &opaque->op.workspace_size, + &opaque->op.executor)); + aclSetAclOpExecutorRepeatable(opaque->op.executor); + + auto handle_ascend = reinterpret_cast(handle); + auto workspace_size = opaque->op.workspace_size; + *desc_ptr = new Descriptor( + opaque.release(), workspace_size, + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *output, std::vector inputs, + void *stream) const { + + if (inputs.size() != 1) { + return INFINI_STATUS_BAD_PARAM; + } + if (workspace_size < workspaceSize()) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + _opaque->op.bind({const_cast(inputs[0]), output}); + CHECK_ACL(aclnnGeluV2( + workspace, workspace_size, _opaque->op.executor, + static_cast(stream))); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::gelutanh::ascend diff --git a/src/infiniop/ops/gelutanh/ascend/gelutanh_ascend.h b/src/infiniop/ops/gelutanh/ascend/gelutanh_ascend.h new file mode 100644 index 000000000..7015c1543 --- /dev/null +++ b/src/infiniop/ops/gelutanh/ascend/gelutanh_ascend.h @@ -0,0 +1,8 @@ +#ifndef __GELUTANH_ASCEND_H__ +#define __GELUTANH_ASCEND_H__ + +#include "../../../devices/ascend/aclnn_elementwise.h" + +ACLNN_ELEMENTWISE_DESCRIPTOR(gelutanh) + +#endif // __GELUTANH_ASCEND_H__ diff --git a/src/infiniop/ops/gelutanh/operator.cc b/src/infiniop/ops/gelutanh/operator.cc index 91a213a24..6c7d8d91f 100644 --- a/src/infiniop/ops/gelutanh/operator.cc +++ b/src/infiniop/ops/gelutanh/operator.cc @@ -14,6 +14,9 @@ #ifdef ENABLE_MOORE_API #include "moore/gelutanh_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/gelutanh_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateGeluTanhDescriptor( infiniopHandle_t handle, @@ -50,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateGeluTanhDescriptor( #endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -86,6 +92,9 @@ __INFINI_C infiniStatus_t infiniopGetGeluTanhWorkspaceSize(infiniopGeluTanhDescr #endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) +#endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -128,6 +137,9 @@ __INFINI_C infiniStatus_t infiniopGeluTanh( #endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -164,6 +176,9 @@ __INFINI_C infiniStatus_t infiniopDestroyGeluTanhDescriptor(infiniopGeluTanhDesc #endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/gemm/ascend/gemm_ascend.cc b/src/infiniop/ops/gemm/ascend/gemm_ascend.cc index e8dfb2f0f..5d0662a76 100644 --- a/src/infiniop/ops/gemm/ascend/gemm_ascend.cc +++ b/src/infiniop/ops/gemm/ascend/gemm_ascend.cc @@ -5,6 +5,7 @@ #include #include +#include // Custom hash function for alpha beta pair struct FloatPairHash { @@ -27,6 +28,7 @@ namespace op::gemm::ascend { struct Descriptor::Opaque { aclnnTensorDescriptor_t c, a, b; + bool transa, transb; // cubeMathType // see doc: // https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/80RC3alpha002/apiref/appdevgapi/context/aclnnBatchMatMul.md @@ -64,15 +66,36 @@ infiniStatus_t Descriptor::create( CHECK_RESULT(result); auto info = result.take(); + // aclnnGemm accepts physical row-major tensors plus transpose flags. Passing + // a transposed view directly makes ACLNN materialize that view with a + // Transpose kernel on every invocation, which is especially expensive for + // immutable linear weights. Describe the same storage as a contiguous + // physical matrix and let Cube consume it with the corresponding flag. + auto make_tensor = [](infiniDtype_t dtype, const BlasMatrix &matrix) { + // CANN 8.5.1 produces incorrect results for the F32 transpose-flag path. + bool transposed = dtype != INFINI_DTYPE_F32 + && matrix.row_stride == 1 && matrix.col_stride != 1; + if (transposed) { + return std::make_pair( + new aclnnTensorDescriptor( + toAclDataType(dtype), + {static_cast(matrix.cols), static_cast(matrix.rows)}, + {matrix.col_stride, matrix.row_stride}), + true); + } + return std::make_pair( + new aclnnTensorDescriptor( + toAclDataType(dtype), + {static_cast(matrix.rows), static_cast(matrix.cols)}, + {matrix.row_stride, matrix.col_stride}), + false); + }; + auto c = new aclnnTensorDescriptor(toAclDataType(c_desc->dtype()), {static_cast(info.m), static_cast(info.n)}, {info.c_matrix.row_stride, info.c_matrix.col_stride}); - auto a = new aclnnTensorDescriptor(toAclDataType(a_desc->dtype()), - {static_cast(info.a_matrix.rows), static_cast(info.a_matrix.cols)}, - {info.a_matrix.row_stride, info.a_matrix.col_stride}); - auto b = new aclnnTensorDescriptor(toAclDataType(b_desc->dtype()), - {static_cast(info.b_matrix.rows), static_cast(info.b_matrix.cols)}, - {info.b_matrix.row_stride, info.b_matrix.col_stride}); + auto [a, transa] = make_tensor(a_desc->dtype(), info.a_matrix); + auto [b, transb] = make_tensor(b_desc->dtype(), info.b_matrix); auto tc = c->tensor, ta = a->tensor, @@ -82,10 +105,10 @@ infiniStatus_t Descriptor::create( aclOpExecutor *executor = nullptr; size_t workspace_size = 0; int8_t mt = 1; - CHECK_ACL(aclnnGemmGetWorkspaceSize(ta, tb, tc, 1., 0., 0, 0, tc, mt, &workspace_size, &executor)); + CHECK_ACL(aclnnGemmGetWorkspaceSize(ta, tb, tc, 1., 0., transa, transb, tc, mt, &workspace_size, &executor)); CHECK_ACL(aclSetAclOpExecutorRepeatable(executor)); lookup[std::make_pair(1.0f, 0.0f)] = executor; - CHECK_ACL(aclnnGemmGetWorkspaceSize(ta, tb, tc, 1., 1., 0, 0, tc, mt, &workspace_size, &executor)); + CHECK_ACL(aclnnGemmGetWorkspaceSize(ta, tb, tc, 1., 1., transa, transb, tc, mt, &workspace_size, &executor)); CHECK_ACL(aclSetAclOpExecutorRepeatable(executor)); lookup[std::make_pair(1.0f, 1.0f)] = executor; @@ -95,6 +118,8 @@ infiniStatus_t Descriptor::create( c, a, b, + transa, + transb, mt, std::move(lookup)}, handle->device, handle->device_id); @@ -123,7 +148,7 @@ infiniStatus_t Descriptor::calculate( executor = _opaque->lookup[key]; } else { CHECK_ACL(aclnnGemmGetWorkspaceSize( - ta, tb, tc, alpha, beta, 0, 0, tc, _opaque->mt, + ta, tb, tc, alpha, beta, _opaque->transa, _opaque->transb, tc, _opaque->mt, &workspace_size, &executor)); CHECK_ACL(aclSetAclOpExecutorRepeatable(executor)); _opaque->lookup[key] = executor; diff --git a/src/infiniop/ops/interpolate/operator.cc b/src/infiniop/ops/interpolate/operator.cc index 0197f8cbe..9eeb72950 100644 --- a/src/infiniop/ops/interpolate/operator.cc +++ b/src/infiniop/ops/interpolate/operator.cc @@ -1,6 +1,7 @@ #include "../../operator.h" #include "../../handle.h" #include "infiniop/ops/interpolate.h" +#include #ifdef ENABLE_CPU_API #include "cpu/interpolate_cpu.h" @@ -14,6 +15,9 @@ #ifdef ENABLE_MOORE_API #include "moore/interpolate_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "../upsample_bilinear/ascend/upsample_bilinear_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateInterpolateDescriptor( infiniopHandle_t handle, @@ -57,6 +61,15 @@ __INFINI_C infiniStatus_t infiniopCreateInterpolateDescriptor( #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + case INFINI_DEVICE_ASCEND: + if (mode == nullptr || std::strcmp(mode, "bilinear") != 0) { + return INFINI_STATUS_BAD_PARAM; + } + return op::upsample_bilinear::ascend::Descriptor::create( + handle, reinterpret_cast(desc_ptr), + y_desc, x_desc, align_corners); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -90,6 +103,11 @@ __INFINI_C infiniStatus_t infiniopGetInterpolateWorkspaceSize(infiniopInterpolat #endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_ASCEND_API + case INFINI_DEVICE_ASCEND: + *size = reinterpret_cast(desc)->workspaceSize(); + return INFINI_STATUS_SUCCESS; #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -132,6 +150,11 @@ __INFINI_C infiniStatus_t infiniopInterpolate( #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + case INFINI_DEVICE_ASCEND: + return reinterpret_cast(desc) + ->calculate(workspace, workspace_size, y, x, stream); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -168,6 +191,11 @@ infiniopDestroyInterpolateDescriptor(infiniopInterpolateDescriptor_t desc) { #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + case INFINI_DEVICE_ASCEND: + delete reinterpret_cast(desc); + return INFINI_STATUS_SUCCESS; +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/layer_norm/ascend/layer_norm_ascend.cc b/src/infiniop/ops/layer_norm/ascend/layer_norm_ascend.cc new file mode 100644 index 000000000..2e12096a5 --- /dev/null +++ b/src/infiniop/ops/layer_norm/ascend/layer_norm_ascend.cc @@ -0,0 +1,266 @@ +#include "layer_norm_ascend.h" + +#include "../../../devices/ascend/aclnn_executor.h" +#include +#include + +#include +#include +#include + +namespace op::layer_norm::ascend { + +namespace { + +constexpr size_t ALIGNMENT = 32; + +size_t align_up(size_t value) { + return (value + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT; +} + +std::vector to_i64(const std::vector &values) { + std::vector result; + result.reserve(values.size()); + for (auto value : values) { + result.push_back(static_cast(value)); + } + return result; +} + +std::vector to_i64(const std::vector &values) { + std::vector result; + result.reserve(values.size()); + for (auto value : values) { + result.push_back(static_cast(value)); + } + return result; +} + +std::vector contiguous_strides(const std::vector &shape) { + std::vector strides(shape.size(), 1); + for (size_t i = shape.size(); i > 1; --i) { + strides[i - 2] = strides[i - 1] * shape[i - 1]; + } + return strides; +} + +aclnnTensorDescriptor_t make_tensor( + aclDataType dtype, + const std::vector &shape, + const std::vector &strides) { + return new aclnnTensorDescriptor(dtype, shape, strides, nullptr); +} + +} // namespace + +struct Descriptor::Opaque { + device::ascend::AclnnExecutor standard; + device::ascend::AclnnExecutor affine; + device::ascend::AclnnExecutor reciprocal; + aclIntArray *normalized_shape = nullptr; + size_t acl_workspace_size = 0; + size_t mean_offset = 0; + size_t rstd_offset = 0; + bool bias_exist = false; + + ~Opaque() { + if (normalized_shape != nullptr) { + aclDestroyIntArray(normalized_shape); + } + } +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_standardization_desc, + infiniopTensorDescriptor_t input_std_deviation_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t bias_desc, + float eps) { + + if (output_desc == nullptr || input_standardization_desc == nullptr + || input_std_deviation_desc == nullptr || input_desc == nullptr + || weight_desc == nullptr) { + return INFINI_STATUS_BAD_PARAM; + } + + auto dtype = input_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + if (output_desc->dtype() != dtype + || input_standardization_desc->dtype() != dtype + || input_std_deviation_desc->dtype() != dtype + || weight_desc->dtype() != dtype + || (bias_desc != nullptr && bias_desc->dtype() != dtype)) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + auto result = LayerNormInfo::createLayerNormInfo( + output_desc, input_standardization_desc, input_std_deviation_desc, + input_desc, weight_desc, bias_desc, eps); + CHECK_RESULT(result); + auto info = result.take(); + + auto opaque = std::make_unique(); + opaque->bias_exist = bias_desc != nullptr; + + auto acl_dtype = toAclDataType(dtype); + auto input_shape = to_i64(info.input_shape); + auto keepdim_shape = input_shape; + keepdim_shape.back() = 1; + auto keepdim_strides = contiguous_strides(keepdim_shape); + auto std_deviation_strides = to_i64(info.input_std_deviation_strides); + std_deviation_strides.push_back(1); + auto normalized_size = static_cast(info.normalized_size); + opaque->normalized_shape = aclCreateIntArray(&normalized_size, 1); + if (opaque->normalized_shape == nullptr) { + return INFINI_STATUS_INTERNAL_ERROR; + } + + auto &standard = opaque->standard; + standard.tensors = { + make_tensor(acl_dtype, input_shape, to_i64(info.input_strides)), + make_tensor( + acl_dtype, input_shape, + to_i64(info.input_standardization_strides)), + make_tensor(acl_dtype, keepdim_shape, keepdim_strides), + make_tensor(acl_dtype, keepdim_shape, keepdim_strides), + }; + CHECK_ACL(aclnnLayerNormGetWorkspaceSize( + standard.tensors[0]->tensor, + opaque->normalized_shape, + nullptr, + nullptr, + static_cast(eps), + standard.tensors[1]->tensor, + standard.tensors[2]->tensor, + standard.tensors[3]->tensor, + &standard.workspace_size, + &standard.executor)); + aclSetAclOpExecutorRepeatable(standard.executor); + + auto &affine = opaque->affine; + affine.tensors = { + make_tensor(acl_dtype, input_shape, to_i64(info.input_strides)), + make_tensor( + acl_dtype, {normalized_size}, + to_i64(info.weight_strides)), + }; + if (opaque->bias_exist) { + affine.tensors.push_back(make_tensor( + acl_dtype, {normalized_size}, to_i64(info.bias_strides))); + } + affine.tensors.push_back( + make_tensor(acl_dtype, input_shape, to_i64(info.output_strides))); + affine.tensors.push_back( + make_tensor(acl_dtype, keepdim_shape, keepdim_strides)); + affine.tensors.push_back( + make_tensor(acl_dtype, keepdim_shape, keepdim_strides)); + + size_t affine_output_index = opaque->bias_exist ? 3 : 2; + CHECK_ACL(aclnnLayerNormGetWorkspaceSize( + affine.tensors[0]->tensor, + opaque->normalized_shape, + affine.tensors[1]->tensor, + opaque->bias_exist ? affine.tensors[2]->tensor : nullptr, + static_cast(eps), + affine.tensors[affine_output_index]->tensor, + affine.tensors[affine_output_index + 1]->tensor, + affine.tensors[affine_output_index + 2]->tensor, + &affine.workspace_size, + &affine.executor)); + aclSetAclOpExecutorRepeatable(affine.executor); + + auto &reciprocal = opaque->reciprocal; + reciprocal.tensors = { + make_tensor(acl_dtype, keepdim_shape, keepdim_strides), + make_tensor( + acl_dtype, keepdim_shape, std_deviation_strides), + }; + CHECK_ACL(aclnnReciprocalGetWorkspaceSize( + reciprocal.tensors[0]->tensor, + reciprocal.tensors[1]->tensor, + &reciprocal.workspace_size, + &reciprocal.executor)); + aclSetAclOpExecutorRepeatable(reciprocal.executor); + + opaque->acl_workspace_size = std::max( + standard.workspace_size, + std::max(affine.workspace_size, reciprocal.workspace_size)); + size_t statistic_size = info.othersize * infiniSizeOf(dtype); + opaque->mean_offset = align_up(opaque->acl_workspace_size); + opaque->rstd_offset = align_up(opaque->mean_offset + statistic_size); + size_t workspace_size = opaque->rstd_offset + statistic_size; + + auto handle_ascend = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + dtype, std::move(info), workspace_size, opaque.release(), + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + void *input_standardization, + void *input_std_deviation, + const void *input, + const void *weight, + const void *bias, + void *stream) const { + + if (workspace_size < workspaceSize()) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + void *mean = static_cast(workspace) + _opaque->mean_offset; + void *rstd = static_cast(workspace) + _opaque->rstd_offset; + + _opaque->standard.bind({ + const_cast(input), + input_standardization, + mean, + rstd, + }); + if (_opaque->bias_exist) { + _opaque->affine.bind({ + const_cast(input), + const_cast(weight), + const_cast(bias), + output, + mean, + rstd, + }); + } else { + _opaque->affine.bind({ + const_cast(input), + const_cast(weight), + output, + mean, + rstd, + }); + } + _opaque->reciprocal.bind({rstd, input_std_deviation}); + + auto acl_stream = static_cast(stream); + CHECK_ACL(aclnnLayerNorm( + workspace, _opaque->standard.workspace_size, + _opaque->standard.executor, acl_stream)); + CHECK_ACL(aclnnLayerNorm( + workspace, _opaque->affine.workspace_size, + _opaque->affine.executor, acl_stream)); + CHECK_ACL(aclnnReciprocal( + workspace, _opaque->reciprocal.workspace_size, + _opaque->reciprocal.executor, acl_stream)); + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::layer_norm::ascend diff --git a/src/infiniop/ops/layer_norm/ascend/layer_norm_ascend.h b/src/infiniop/ops/layer_norm/ascend/layer_norm_ascend.h new file mode 100644 index 000000000..cb09390a1 --- /dev/null +++ b/src/infiniop/ops/layer_norm/ascend/layer_norm_ascend.h @@ -0,0 +1,8 @@ +#ifndef __LAYER_NORM_ASCEND_H__ +#define __LAYER_NORM_ASCEND_H__ + +#include "../layer_norm.h" + +DESCRIPTOR(ascend) + +#endif // __LAYER_NORM_ASCEND_H__ diff --git a/src/infiniop/ops/layer_norm/operator.cc b/src/infiniop/ops/layer_norm/operator.cc index 3925e845a..ad5c2a284 100644 --- a/src/infiniop/ops/layer_norm/operator.cc +++ b/src/infiniop/ops/layer_norm/operator.cc @@ -17,6 +17,9 @@ #ifdef ENABLE_CAMBRICON_API #include "bang/layer_norm_bang.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/layer_norm_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateLayerNormDescriptor( infiniopHandle_t handle, @@ -70,6 +73,9 @@ __INFINI_C infiniStatus_t infiniopCreateLayerNormDescriptor( #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -111,6 +117,9 @@ __INFINI_C infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDes #endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -174,6 +183,9 @@ __INFINI_C infiniStatus_t infiniopLayerNorm( #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -219,6 +231,9 @@ infiniopDestroyLayerNormDescriptor(infiniopLayerNormDescriptor_t desc) { #ifdef ENABLE_ILUVATAR_API DELETE(INFINI_DEVICE_ILUVATAR, nvidia); #endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/mrope/ascend/mrope_ascend.cc b/src/infiniop/ops/mrope/ascend/mrope_ascend.cc new file mode 100644 index 000000000..ef93ee36e --- /dev/null +++ b/src/infiniop/ops/mrope/ascend/mrope_ascend.cc @@ -0,0 +1,90 @@ +#include "mrope_ascend.h" +#include "../../../devices/ascend/ascend_handle.h" + +namespace op::mrope::ascend { + +extern "C" infiniStatus_t mrope_ascend_kernel_launch( + void *q_out, void *k_out, const void *q, const void *k, + const void *cos, const void *sin, const void *positions, + infiniDtype_t data_type, bool positions_i64, + size_t num_tokens, size_t num_q_heads, size_t num_kv_heads, + size_t head_size, size_t rotary_dim, size_t half_rotary_dim, + ptrdiff_t q_out_stride_token, ptrdiff_t q_out_stride_head, + ptrdiff_t k_out_stride_token, ptrdiff_t k_out_stride_head, + ptrdiff_t q_stride_token, ptrdiff_t q_stride_head, + ptrdiff_t k_stride_token, ptrdiff_t k_stride_head, + ptrdiff_t cos_stride_position, ptrdiff_t sin_stride_position, + ptrdiff_t positions_stride_axis, ptrdiff_t positions_stride_token, + size_t max_position_embeddings, size_t section_t, + size_t section_h, size_t section_w, + bool positions_has_axes, bool interleaved, void *stream); + +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t q_out_desc, + infiniopTensorDescriptor_t k_out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_desc, + infiniopTensorDescriptor_t cos_desc, + infiniopTensorDescriptor_t sin_desc, + infiniopTensorDescriptor_t positions_desc, + int head_size, + int rotary_dim, + int section_t, + int section_h, + int section_w, + bool interleaved) { + + auto result = MRoPEInfo::create( + q_out_desc, k_out_desc, q_desc, k_desc, cos_desc, sin_desc, + positions_desc, head_size, rotary_dim, section_t, section_h, + section_w, interleaved); + CHECK_RESULT(result); + *desc_ptr = new Descriptor( + result.take(), 0, new Opaque{}, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *q_out, + void *k_out, + const void *q, + const void *k, + const void *cos, + const void *sin, + const void *positions, + void *stream) const { + + (void)workspace; + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (_info.data_type == INFINI_DTYPE_F64) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + return mrope_ascend_kernel_launch( + q_out, k_out, q, k, cos, sin, positions, _info.data_type, + _info.position_type == INFINI_DTYPE_I64, + _info.num_tokens, _info.num_q_heads, _info.num_kv_heads, + _info.head_size, _info.rotary_dim, _info.half_rotary_dim, + _info.q_out_stride_token, _info.q_out_stride_head, + _info.k_out_stride_token, _info.k_out_stride_head, + _info.q_stride_token, _info.q_stride_head, + _info.k_stride_token, _info.k_stride_head, + _info.cos_stride_position, _info.sin_stride_position, + _info.positions_stride_axis, _info.positions_stride_token, + _info.max_position_embeddings, _info.section_t, + _info.section_h, _info.section_w, _info.positions_has_axes, + _info.interleaved, stream); +} + +} // namespace op::mrope::ascend diff --git a/src/infiniop/ops/mrope/ascend/mrope_ascend.h b/src/infiniop/ops/mrope/ascend/mrope_ascend.h new file mode 100644 index 000000000..68ea84c42 --- /dev/null +++ b/src/infiniop/ops/mrope/ascend/mrope_ascend.h @@ -0,0 +1,8 @@ +#ifndef __MROPE_ASCEND_H__ +#define __MROPE_ASCEND_H__ + +#include "../mrope.h" + +DESCRIPTOR(ascend) + +#endif diff --git a/src/infiniop/ops/mrope/ascend/mrope_ascend_kernel.cpp b/src/infiniop/ops/mrope/ascend/mrope_ascend_kernel.cpp new file mode 100644 index 000000000..344ece92b --- /dev/null +++ b/src/infiniop/ops/mrope/ascend/mrope_ascend_kernel.cpp @@ -0,0 +1,228 @@ +#include "../../../devices/ascend/ascend_kernel_common.h" + +#include + +using namespace AscendC; + +template +__aicore__ inline float mropeToFloat(T value) { + if constexpr (std::is_same::value) { + uint32_t bits = static_cast(*reinterpret_cast(&value)) << 16; + return *reinterpret_cast(&bits); + } else { + return static_cast(value); + } +} + +template +__aicore__ inline T mropeFromFloat(float value) { + if constexpr (std::is_same::value) { + uint32_t bits = *reinterpret_cast(&value); + uint16_t upper = static_cast(bits >> 16); + return *reinterpret_cast(&upper); + } else { + return static_cast(value); + } +} + +__aicore__ inline int64_t mropeLoadPosition( + GM_ADDR positions, bool positions_i64, ptrdiff_t offset) { + if (positions_i64) { + return reinterpret_cast<__gm__ int64_t *>(positions)[offset]; + } + return static_cast( + reinterpret_cast<__gm__ int32_t *>(positions)[offset]); +} + +template +__aicore__ inline void mropeRotateHeads( + GlobalTensor &output, + GlobalTensor &input, + GlobalTensor &cos, + GlobalTensor &sin, + GM_ADDR positions, + bool positions_i64, + size_t token, + size_t num_heads, + size_t head_size, + size_t rotary_dim, + size_t half_rotary_dim, + ptrdiff_t output_stride_token, + ptrdiff_t output_stride_head, + ptrdiff_t input_stride_token, + ptrdiff_t input_stride_head, + ptrdiff_t cos_stride_position, + ptrdiff_t sin_stride_position, + ptrdiff_t positions_stride_axis, + ptrdiff_t positions_stride_token, + size_t max_position_embeddings, + size_t section_t, + size_t section_h, + size_t section_w, + bool positions_has_axes, + bool interleaved) { + + for (size_t head = 0; head < num_heads; ++head) { + ptrdiff_t out_base = static_cast(token) * output_stride_token + + static_cast(head) * output_stride_head; + ptrdiff_t in_base = static_cast(token) * input_stride_token + + static_cast(head) * input_stride_head; + for (size_t i = 0; i < half_rotary_dim; ++i) { + size_t axis; + if (interleaved) { + bool h_mask = i % 3 == 1 && i < section_h * 3; + bool w_mask = i % 3 == 2 && i < section_w * 3; + axis = h_mask ? 1 : (w_mask ? 2 : 0); + } else { + axis = i < section_t ? 0 + : (i < section_t + section_h ? 1 : 2); + } + ptrdiff_t position_offset = positions_has_axes + ? static_cast(axis) * positions_stride_axis + + static_cast(token) + * positions_stride_token + : static_cast(token) + * positions_stride_token; + int64_t raw_position = mropeLoadPosition(positions, positions_i64, position_offset); + size_t position = raw_position >= 0 + && static_cast(raw_position) + < max_position_embeddings + ? static_cast(raw_position) + : 0; + float cos_value = mropeToFloat(cos.GetValue( + static_cast(position) * cos_stride_position + i)); + float sin_value = mropeToFloat(sin.GetValue( + static_cast(position) * sin_stride_position + i)); + float x0 = mropeToFloat(input.GetValue(in_base + i)); + float x1 = mropeToFloat(input.GetValue(in_base + i + half_rotary_dim)); + output.SetValue( + out_base + i, + mropeFromFloat(x0 * cos_value - x1 * sin_value)); + output.SetValue( + out_base + i + half_rotary_dim, + mropeFromFloat(x1 * cos_value + x0 * sin_value)); + } + for (size_t i = rotary_dim; i < head_size; ++i) { + output.SetValue(out_base + i, input.GetValue(in_base + i)); + } + } +} + +template +__aicore__ inline void mropeProcess( + GM_ADDR q_out_ptr, GM_ADDR k_out_ptr, GM_ADDR q_ptr, GM_ADDR k_ptr, + GM_ADDR cos_ptr, GM_ADDR sin_ptr, GM_ADDR positions, + bool positions_i64, size_t num_tokens, size_t num_q_heads, + size_t num_kv_heads, size_t head_size, size_t rotary_dim, + size_t half_rotary_dim, ptrdiff_t q_out_stride_token, + ptrdiff_t q_out_stride_head, ptrdiff_t k_out_stride_token, + ptrdiff_t k_out_stride_head, ptrdiff_t q_stride_token, + ptrdiff_t q_stride_head, ptrdiff_t k_stride_token, + ptrdiff_t k_stride_head, ptrdiff_t cos_stride_position, + ptrdiff_t sin_stride_position, ptrdiff_t positions_stride_axis, + ptrdiff_t positions_stride_token, size_t max_position_embeddings, + size_t section_t, size_t section_h, size_t section_w, + bool positions_has_axes, bool interleaved) { + + size_t token = GetBlockIdx(); + if (token >= num_tokens) { + return; + } + GlobalTensor q_out, k_out, q, k, cos, sin; + q_out.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(q_out_ptr)); + k_out.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(k_out_ptr)); + q.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(q_ptr)); + k.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(k_ptr)); + cos.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(cos_ptr)); + sin.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(sin_ptr)); + mropeRotateHeads( + q_out, q, cos, sin, positions, positions_i64, token, num_q_heads, + head_size, rotary_dim, half_rotary_dim, q_out_stride_token, + q_out_stride_head, q_stride_token, q_stride_head, + cos_stride_position, sin_stride_position, positions_stride_axis, + positions_stride_token, max_position_embeddings, section_t, + section_h, section_w, positions_has_axes, interleaved); + mropeRotateHeads( + k_out, k, cos, sin, positions, positions_i64, token, num_kv_heads, + head_size, rotary_dim, half_rotary_dim, k_out_stride_token, + k_out_stride_head, k_stride_token, k_stride_head, + cos_stride_position, sin_stride_position, positions_stride_axis, + positions_stride_token, max_position_embeddings, section_t, + section_h, section_w, positions_has_axes, interleaved); +} + +#define DEFINE_MROPE_KERNEL(NAME, TYPE) \ + __global__ __aicore__ void NAME( \ + GM_ADDR q_out, GM_ADDR k_out, GM_ADDR q, GM_ADDR k, GM_ADDR cos, \ + GM_ADDR sin, GM_ADDR positions, bool positions_i64, \ + size_t num_tokens, size_t num_q_heads, size_t num_kv_heads, \ + size_t head_size, size_t rotary_dim, size_t half_rotary_dim, \ + ptrdiff_t q_out_stride_token, ptrdiff_t q_out_stride_head, \ + ptrdiff_t k_out_stride_token, ptrdiff_t k_out_stride_head, \ + ptrdiff_t q_stride_token, ptrdiff_t q_stride_head, \ + ptrdiff_t k_stride_token, ptrdiff_t k_stride_head, \ + ptrdiff_t cos_stride_position, ptrdiff_t sin_stride_position, \ + ptrdiff_t positions_stride_axis, ptrdiff_t positions_stride_token, \ + size_t max_position_embeddings, size_t section_t, size_t section_h, \ + size_t section_w, bool positions_has_axes, bool interleaved) { \ + mropeProcess( \ + q_out, k_out, q, k, cos, sin, positions, positions_i64, \ + num_tokens, num_q_heads, num_kv_heads, head_size, rotary_dim, \ + half_rotary_dim, q_out_stride_token, q_out_stride_head, \ + k_out_stride_token, k_out_stride_head, q_stride_token, \ + q_stride_head, k_stride_token, k_stride_head, \ + cos_stride_position, sin_stride_position, \ + positions_stride_axis, positions_stride_token, \ + max_position_embeddings, section_t, section_h, section_w, \ + positions_has_axes, interleaved); \ + } + +DEFINE_MROPE_KERNEL(mrope_half, half) +DEFINE_MROPE_KERNEL(mrope_bf16, bfloat16_t) +DEFINE_MROPE_KERNEL(mrope_float, float) +#undef DEFINE_MROPE_KERNEL + +extern "C" infiniStatus_t mrope_ascend_kernel_launch( + void *q_out, void *k_out, const void *q, const void *k, + const void *cos, const void *sin, const void *positions, + infiniDtype_t data_type, bool positions_i64, + size_t num_tokens, size_t num_q_heads, size_t num_kv_heads, + size_t head_size, size_t rotary_dim, size_t half_rotary_dim, + ptrdiff_t q_out_stride_token, ptrdiff_t q_out_stride_head, + ptrdiff_t k_out_stride_token, ptrdiff_t k_out_stride_head, + ptrdiff_t q_stride_token, ptrdiff_t q_stride_head, + ptrdiff_t k_stride_token, ptrdiff_t k_stride_head, + ptrdiff_t cos_stride_position, ptrdiff_t sin_stride_position, + ptrdiff_t positions_stride_axis, ptrdiff_t positions_stride_token, + size_t max_position_embeddings, size_t section_t, + size_t section_h, size_t section_w, + bool positions_has_axes, bool interleaved, void *stream) { + + if (num_tokens == 0) { + return INFINI_STATUS_SUCCESS; + } + uint32_t blocks = static_cast(num_tokens); +#define LAUNCH_MROPE(DTYPE, NAME) \ + case DTYPE: \ + NAME<<>>( \ + q_out, k_out, const_cast(q), const_cast(k), \ + const_cast(cos), const_cast(sin), \ + const_cast(positions), positions_i64, num_tokens, \ + num_q_heads, num_kv_heads, head_size, rotary_dim, \ + half_rotary_dim, q_out_stride_token, q_out_stride_head, \ + k_out_stride_token, k_out_stride_head, q_stride_token, \ + q_stride_head, k_stride_token, k_stride_head, \ + cos_stride_position, sin_stride_position, \ + positions_stride_axis, positions_stride_token, \ + max_position_embeddings, section_t, section_h, section_w, \ + positions_has_axes, interleaved); \ + return INFINI_STATUS_SUCCESS; + switch (data_type) { + LAUNCH_MROPE(INFINI_DTYPE_F16, mrope_half) + LAUNCH_MROPE(INFINI_DTYPE_BF16, mrope_bf16) + LAUNCH_MROPE(INFINI_DTYPE_F32, mrope_float) + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +#undef LAUNCH_MROPE +} diff --git a/src/infiniop/ops/mrope/operator.cc b/src/infiniop/ops/mrope/operator.cc index 2c8ae2a1e..ce8c67bc2 100644 --- a/src/infiniop/ops/mrope/operator.cc +++ b/src/infiniop/ops/mrope/operator.cc @@ -11,6 +11,9 @@ #ifdef ENABLE_MOORE_API #include "moore/mrope_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/mrope_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateMRoPEDescriptor( infiniopHandle_t handle, @@ -58,6 +61,9 @@ __INFINI_C infiniStatus_t infiniopCreateMRoPEDescriptor( #endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -92,6 +98,9 @@ __INFINI_C infiniStatus_t infiniopGetMRoPEWorkspaceSize(infiniopMRoPEDescriptor_ #endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -137,6 +146,9 @@ __INFINI_C infiniStatus_t infiniopMRoPE( #endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -171,6 +183,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMRoPEDescriptor(infiniopMRoPEDescriptor #endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore); +#endif +#ifdef ENABLE_ASCEND_API + DESTROY(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/mul/ascend/mul_ascend.cc b/src/infiniop/ops/mul/ascend/mul_ascend.cc new file mode 100644 index 000000000..30c8a5424 --- /dev/null +++ b/src/infiniop/ops/mul/ascend/mul_ascend.cc @@ -0,0 +1,77 @@ +#include "mul_ascend.h" + +#include "../../../devices/ascend/aclnn_executor.h" +#include +#include + +namespace op::mul::ascend { + +struct Descriptor::Opaque { + device::ascend::AclnnExecutor op; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + std::vector input_descs) { + + auto status = device::ascend::validateAclnnElementwise(output_desc, input_descs, 2); + if (status != INFINI_STATUS_SUCCESS) { + return status; + } + CHECK_DTYPE(output_desc->dtype(), + INFINI_DTYPE_F16, INFINI_DTYPE_F32, + INFINI_DTYPE_F64, INFINI_DTYPE_BF16); + + auto opaque = std::make_unique(); + opaque->op.tensors = { + new aclnnTensorDescriptor(input_descs[0]), + new aclnnTensorDescriptor(input_descs[1]), + new aclnnTensorDescriptor(output_desc), + }; + + CHECK_ACL(aclnnMulGetWorkspaceSize( + opaque->op.tensors[0]->tensor, + opaque->op.tensors[1]->tensor, + opaque->op.tensors[2]->tensor, + &opaque->op.workspace_size, + &opaque->op.executor)); + aclSetAclOpExecutorRepeatable(opaque->op.executor); + + auto handle_ascend = reinterpret_cast(handle); + auto workspace_size = opaque->op.workspace_size; + *desc_ptr = new Descriptor( + opaque.release(), workspace_size, + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *output, std::vector inputs, + void *stream) const { + + if (inputs.size() != 2) { + return INFINI_STATUS_BAD_PARAM; + } + if (workspace_size < workspaceSize()) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + _opaque->op.bind({ + const_cast(inputs[0]), + const_cast(inputs[1]), + output, + }); + CHECK_ACL(aclnnMul( + workspace, workspace_size, _opaque->op.executor, + static_cast(stream))); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::mul::ascend diff --git a/src/infiniop/ops/mul/ascend/mul_ascend.h b/src/infiniop/ops/mul/ascend/mul_ascend.h new file mode 100644 index 000000000..5eae53c95 --- /dev/null +++ b/src/infiniop/ops/mul/ascend/mul_ascend.h @@ -0,0 +1,8 @@ +#ifndef __MUL_ASCEND_H__ +#define __MUL_ASCEND_H__ + +#include "../../../devices/ascend/aclnn_elementwise.h" + +ACLNN_ELEMENTWISE_DESCRIPTOR(mul) + +#endif // __MUL_ASCEND_H__ diff --git a/src/infiniop/ops/mul/operator.cc b/src/infiniop/ops/mul/operator.cc index 33f4e0c88..18948028f 100644 --- a/src/infiniop/ops/mul/operator.cc +++ b/src/infiniop/ops/mul/operator.cc @@ -17,6 +17,9 @@ #ifdef ENABLE_MOORE_API #include "moore/mul_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/mul_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateMulDescriptor( infiniopHandle_t handle, @@ -63,6 +66,9 @@ __INFINI_C infiniStatus_t infiniopCreateMulDescriptor( #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -106,6 +112,9 @@ __INFINI_C infiniStatus_t infiniopGetMulWorkspaceSize(infiniopMulDescriptor_t de #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -158,6 +167,9 @@ __INFINI_C infiniStatus_t infiniopMul( #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -203,6 +215,9 @@ infiniopDestroyMulDescriptor(infiniopMulDescriptor_t desc) { #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/random_sample/ascend/randomsample_aclnn.cc b/src/infiniop/ops/random_sample/ascend/randomsample_aclnn.cc index fbedd48a6..130bdfa48 100644 --- a/src/infiniop/ops/random_sample/ascend/randomsample_aclnn.cc +++ b/src/infiniop/ops/random_sample/ascend/randomsample_aclnn.cc @@ -1,6 +1,9 @@ #include "../../../devices/ascend/common_ascend.h" #include "random_sample_aclnn.h" -#include +#include +#include +#include +#include namespace op::random_sample::ascend { @@ -25,17 +28,12 @@ infiniStatus_t Descriptor::create( auto result = RandomSampleInfo::create(result_desc, probs_desc); CHECK_RESULT(result); CHECK_DTYPE(result->dt_i, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - auto topk_val_dtype = (probs_desc->dtype() == INFINI_DTYPE_BF16 || probs_desc->dtype() == INFINI_DTYPE_F16) - ? INFINI_DTYPE_F32 - : probs_desc->dtype(); - auto workspace_size = utils::align(probs_desc->numel() * infiniSizeOf(topk_val_dtype), 32) - + probs_desc->numel() * infiniSizeOf(INFINI_DTYPE_I64); auto tresult = new aclnnTensorDescriptor(result_desc); auto tprobs = new aclnnTensorDescriptor(probs_desc); *desc_ptr = new Descriptor( result.take(), - workspace_size, + 0, new Opaque{tprobs, tresult}, handle->device, handle->device_id); return INFINI_STATUS_SUCCESS; @@ -45,20 +43,6 @@ size_t Descriptor::minWorkspaceSize() const { return _min_workspace_size; } -extern "C" infiniStatus_t random_sample_kernel_launch( - void *probs, - void *result, - void *topk_val_addr, - void *topk_idx_addr, - float random_val, - float topp, - int topk, - float temperature, - uint64_t n, - infiniDtype_t dt_p, - infiniDtype_t dt_i, - void *stream); - infiniStatus_t Descriptor::calculate( void *workspace, @@ -73,116 +57,129 @@ Descriptor::calculate( if (workspace_size < _min_workspace_size) { return INFINI_STATUS_INSUFFICIENT_WORKSPACE; } - auto topk_ = topk <= (int)_info.n ? topk : (int)_info.n; - bool dosample = topk_ > 1 && temperature != 0.0f && topp != 0.0f && random_val != 0.0f; - auto effective_topk = dosample ? topk_ : 1; - auto topk_shape = std::vector{effective_topk}; - auto topk_stride = std::vector{1}; - - bool use_fp32_topk = (_info.dt_p == INFINI_DTYPE_BF16 || _info.dt_p == INFINI_DTYPE_F16); - - void *probs_for_topk = const_cast(probs); - void *topk_val_addr = workspace; - auto topk_val_bytes = effective_topk * infiniSizeOf(use_fp32_topk ? INFINI_DTYPE_F32 : _info.dt_p); - void *topk_idx_addr = (void *)((uint8_t *)topk_val_addr + utils::align(topk_val_bytes, 32)); - - uint64_t topk_workspace_size = 0; - aclOpExecutor *topk_executor = nullptr; - - if (use_fp32_topk) { - void *probs_fp32; - auto probs_fp32_size = _info.n * infiniSizeOf(INFINI_DTYPE_F32); - CHECK_ACL(aclrtMalloc(&probs_fp32, probs_fp32_size, ACL_MEM_MALLOC_HUGE_FIRST)); - - void *probs_host; - auto probs_host_size = _info.n * infiniSizeOf(_info.dt_p); - CHECK_ACL(aclrtMallocHost(&probs_host, probs_host_size)); - void *probs_fp32_host; - CHECK_ACL(aclrtMallocHost(&probs_fp32_host, _info.n * sizeof(float))); - - CHECK_ACL(aclrtSynchronizeDevice()); - CHECK_ACL(aclrtMemcpy(probs_host, probs_host_size, probs, probs_host_size, ACL_MEMCPY_DEVICE_TO_HOST)); - - auto fp32_ptr = static_cast(probs_fp32_host); - if (_info.dt_p == INFINI_DTYPE_F16) { - auto f16_ptr = static_cast(probs_host); - for (uint64_t i = 0; i < _info.n; i++) { - fp32_ptr[i] = _f16_to_f32(f16_ptr[i]); - } - } else { - auto bf16_ptr = static_cast(probs_host); - for (uint64_t i = 0; i < _info.n; i++) { - fp32_ptr[i] = _bf16_to_f32(bf16_ptr[i]); - } + if (_info.n == 0) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + + // Sampling needs a scalar result on the host. Stage the logits once and do + // stable partial top-k here instead of copying FP32 logits back to the NPU, + // running ACLNN TopK, and launching a second synchronization-prone kernel. + CHECK_ACL(aclrtSynchronizeStream(static_cast(stream))); + void *probs_host = nullptr; + const auto probs_host_size = _info.n * infiniSizeOf(_info.dt_p); + CHECK_ACL(aclrtMallocHost(&probs_host, probs_host_size)); + CHECK_ACL(aclrtMemcpy( + probs_host, probs_host_size, + probs, probs_host_size, + ACL_MEMCPY_DEVICE_TO_HOST)); + + std::vector logits(_info.n); + switch (_info.dt_p) { + case INFINI_DTYPE_F16: { + auto src = static_cast(probs_host); + for (size_t i = 0; i < _info.n; ++i) { + logits[i] = _f16_to_f32(src[i]); + } + break; + } + case INFINI_DTYPE_BF16: { + auto src = static_cast(probs_host); + for (size_t i = 0; i < _info.n; ++i) { + logits[i] = _bf16_to_f32(src[i]); + } + break; + } + case INFINI_DTYPE_F32: { + auto src = static_cast(probs_host); + std::copy(src, src + _info.n, logits.begin()); + break; + } + case INFINI_DTYPE_F64: { + auto src = static_cast(probs_host); + for (size_t i = 0; i < _info.n; ++i) { + logits[i] = static_cast(src[i]); + } + break; + } + default: + CHECK_ACL(aclrtFreeHost(probs_host)); + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + CHECK_ACL(aclrtFreeHost(probs_host)); + + const auto effective_topk = std::min( + static_cast(std::max(topk, 1)), _info.n); + const bool do_sample + = effective_topk > 1 + && temperature != 0.0f + && topp != 0.0f + && random_val != 0.0f; + + std::vector indices(_info.n); + std::iota(indices.begin(), indices.end(), 0); + const auto greater_logit = [&logits](size_t lhs, size_t rhs) { + if (logits[lhs] == logits[rhs]) { + return lhs < rhs; + } + return logits[lhs] > logits[rhs]; + }; + + size_t selected = 0; + if (!do_sample) { + selected = *std::max_element( + indices.begin(), indices.end(), + [&logits](size_t lhs, size_t rhs) { + return logits[lhs] < logits[rhs]; + }); + } else { + std::partial_sort( + indices.begin(), indices.begin() + effective_topk, indices.end(), + greater_logit); + + const double max_logit = logits[indices[0]]; + const double inv_temperature = 1.0 / static_cast(temperature); + const auto weight = [&logits, max_logit, inv_temperature](size_t index) { + return std::exp( + (static_cast(logits[index]) - max_logit) + * inv_temperature); + }; + + double total_mass = 0.0; + for (size_t i = 0; i < _info.n; ++i) { + total_mass += weight(i); + } + double topk_mass = 0.0; + for (size_t i = 0; i < effective_topk; ++i) { + topk_mass += weight(indices[i]); } - CHECK_ACL(aclrtMemcpy(probs_fp32, probs_fp32_size, probs_fp32_host, probs_fp32_size, ACL_MEMCPY_HOST_TO_DEVICE)); + const double limit = static_cast(random_val) + * std::min( + topk_mass, + total_mass * static_cast(topp)); + double cumulative = 0.0; + selected = indices[effective_topk - 1]; + for (size_t i = 0; i < effective_topk; ++i) { + cumulative += weight(indices[i]); + if (limit <= cumulative) { + selected = indices[i]; + break; + } + } + } - CHECK_ACL(aclrtFreeHost(probs_host)); - CHECK_ACL(aclrtFreeHost(probs_fp32_host)); - - int64_t shape = _info.n; - int64_t stride = 1; - auto probs_fp32_desc = new aclnnTensorDescriptor(toAclDataType(INFINI_DTYPE_F32), {shape}, {stride}); - - auto topk_val_fp32_desc = new aclnnTensorDescriptor(toAclDataType(INFINI_DTYPE_F32), topk_shape, topk_stride); - auto topk_idx_desc = new aclnnTensorDescriptor(toAclDataType(INFINI_DTYPE_I64), topk_shape, topk_stride); - - CHECK_ACL(aclnnTopkGetWorkspaceSize(probs_fp32_desc->tensor, - topk_shape[0], - 0, - true, - true, - topk_val_fp32_desc->tensor, - topk_idx_desc->tensor, - &topk_workspace_size, - &topk_executor)); - CHECK_ACL(aclSetAclOpExecutorRepeatable(topk_executor)); - void *topk_workspace; - CHECK_ACL(aclrtMalloc(&topk_workspace, topk_workspace_size, ACL_MEM_MALLOC_HUGE_FIRST)); - AclSetTensorAddr(topk_executor, 0, probs_fp32_desc->tensor, probs_fp32); - AclSetTensorAddr(topk_executor, 1, topk_val_fp32_desc->tensor, topk_val_addr); - AclSetTensorAddr(topk_executor, 2, topk_idx_desc->tensor, topk_idx_addr); - CHECK_ACL(aclnnTopk(topk_workspace, topk_workspace_size, topk_executor, stream)); - CHECK_ACL(aclrtSynchronizeDevice()); - CHECK_ACL(aclrtFree(topk_workspace)); - - delete topk_val_fp32_desc; - delete topk_idx_desc; - delete probs_fp32_desc; - - auto status = random_sample_kernel_launch(probs_fp32, result, topk_val_addr, topk_idx_addr, random_val, topp, effective_topk, temperature, _info.n, INFINI_DTYPE_F32, _info.dt_i, stream); - CHECK_STATUS(status); - CHECK_ACL(aclrtSynchronizeDevice()); - CHECK_ACL(aclrtFree(probs_fp32)); + if (_info.dt_i == INFINI_DTYPE_I32) { + const auto host_result = static_cast(selected); + CHECK_ACL(aclrtMemcpy( + result, sizeof(host_result), + &host_result, sizeof(host_result), + ACL_MEMCPY_HOST_TO_DEVICE)); } else { - auto topk_val = new aclnnTensorDescriptor(toAclDataType(_info.dt_p), topk_shape, topk_stride); - auto topk_idx = new aclnnTensorDescriptor(toAclDataType(INFINI_DTYPE_I64), topk_shape, topk_stride); - - CHECK_ACL(aclnnTopkGetWorkspaceSize(_opaque->probs->tensor, - topk_shape[0], - 0, - true, - true, - topk_val->tensor, - topk_idx->tensor, - &topk_workspace_size, - &topk_executor)); - CHECK_ACL(aclSetAclOpExecutorRepeatable(topk_executor)); - void *topk_workspace; - CHECK_ACL(aclrtMalloc(&topk_workspace, topk_workspace_size, ACL_MEM_MALLOC_HUGE_FIRST)); - AclSetTensorAddr(topk_executor, 0, _opaque->probs->tensor, (void *)probs); - AclSetTensorAddr(topk_executor, 1, topk_val->tensor, topk_val_addr); - AclSetTensorAddr(topk_executor, 2, topk_idx->tensor, topk_idx_addr); - CHECK_ACL(aclnnTopk(topk_workspace, topk_workspace_size, topk_executor, stream)); - CHECK_ACL(aclrtSynchronizeDevice()); - CHECK_ACL(aclrtFree(topk_workspace)); - - auto status = random_sample_kernel_launch(probs_for_topk, result, topk_val_addr, topk_idx_addr, random_val, topp, effective_topk, temperature, _info.n, _info.dt_p, _info.dt_i, stream); - CHECK_STATUS(status); - - delete topk_val; - delete topk_idx; + const auto host_result = static_cast(selected); + CHECK_ACL(aclrtMemcpy( + result, sizeof(host_result), + &host_result, sizeof(host_result), + ACL_MEMCPY_HOST_TO_DEVICE)); } return INFINI_STATUS_SUCCESS; diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_ascend.cc b/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_ascend.cc new file mode 100644 index 000000000..2aaa403ce --- /dev/null +++ b/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_ascend.cc @@ -0,0 +1,462 @@ +#include "recurrent_gated_delta_rule_ascend.h" +#include "../../../devices/ascend/ascend_handle.h" +#include "../../../devices/ascend/common_ascend.h" +#include "../../chunk_gated_delta_rule/ascend/gated_delta_rule_ascend_kernel.h" +#include "recurrent_gated_delta_rule_native_kernel.h" + +#include +#include +#include +#include +#include +#include + +namespace op::recurrent_gated_delta_rule::ascend { + +namespace { + +constexpr size_t NATIVE_DIM = 128; +constexpr size_t NATIVE_ALIGNMENT = 512; + +size_t alignNative(size_t value) { + return (value + NATIVE_ALIGNMENT - 1) & ~(NATIVE_ALIGNMENT - 1); +} + +size_t reserveNative(size_t &cursor, size_t bytes) { + cursor = alignNative(cursor); + size_t result = cursor; + cursor += bytes; + return result; +} + +} // namespace + +struct Descriptor::Opaque { + bool native = false; + aclnnTensorDescriptor_t q = nullptr; + aclnnTensorDescriptor_t k = nullptr; + aclnnTensorDescriptor_t v = nullptr; + aclnnTensorDescriptor_t beta = nullptr; + aclnnTensorDescriptor_t state = nullptr; + aclnnTensorDescriptor_t actual_seq_lengths = nullptr; + aclnnTensorDescriptor_t state_indices = nullptr; + aclnnTensorDescriptor_t g = nullptr; + aclnnTensorDescriptor_t out = nullptr; + aclOpExecutor *executor = nullptr; + void *q_buffer = nullptr; + void *k_buffer = nullptr; + void *v_buffer = nullptr; + void *beta_buffer = nullptr; + void *lengths_buffer = nullptr; + void *indices_buffer = nullptr; + void *native_workspace_buffer = nullptr; + uint64_t native_workspace_size = 0; + std::vector> cached_addresses; + std::vector> cached_descriptors; + std::vector cached_executors; + std::vector cached_workspace_sizes; + size_t q_offset = 0; + size_t k_offset = 0; + size_t v_offset = 0; + size_t beta_offset = 0; + size_t state_offset = 0; + size_t actual_seq_lengths_offset = 0; + size_t state_indices_offset = 0; + size_t native_workspace_offset = 0; + + ~Opaque() { + for (auto &descriptors : cached_descriptors) { + for (auto *descriptor : descriptors) { + delete descriptor; + } + } + delete q; + delete k; + delete v; + delete beta; + delete state; + delete actual_seq_lengths; + delete state_indices; + delete g; + delete out; + if (q_buffer != nullptr) { + aclrtFree(q_buffer); + } + if (k_buffer != nullptr) { + aclrtFree(k_buffer); + } + if (v_buffer != nullptr) { + aclrtFree(v_buffer); + } + if (beta_buffer != nullptr) { + aclrtFree(beta_buffer); + } + if (lengths_buffer != nullptr) { + aclrtFree(lengths_buffer); + } + if (indices_buffer != nullptr) { + aclrtFree(indices_buffer); + } + if (native_workspace_buffer != nullptr) { + aclrtFree(native_workspace_buffer); + } + } +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t initial_state_desc, + infiniopTensorDescriptor_t final_state_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_desc, + infiniopTensorDescriptor_t v_desc, + infiniopTensorDescriptor_t g_desc, + infiniopTensorDescriptor_t beta_desc, + infiniopTensorDescriptor_t initial_state_indices_desc, + infiniopTensorDescriptor_t final_state_indices_desc, + bool use_qk_l2norm) { + + auto result = RecurrentGatedDeltaRuleInfo::create( + out_desc, initial_state_desc, final_state_desc, q_desc, k_desc, v_desc, + g_desc, beta_desc, initial_state_indices_desc, final_state_indices_desc, + use_qk_l2norm); + CHECK_RESULT(result); + auto info = result.take(); + auto opaque = new Opaque{}; + + const bool state_contiguous = info.initial_state_strides[3] == 1 + && info.initial_state_strides[2] == static_cast(info.Dk) + && info.initial_state_strides[1] + == static_cast(info.Dv * info.Dk) + && info.initial_state_strides[0] + == static_cast(info.Hv * info.Dv * info.Dk); + // The native preprocess kernel reads Q/K through their explicit strides + // while normalizing them, so their outer stride need not be contiguous. + const bool tensors_contiguous = info.q_strides[2] == static_cast(info.Dk) + && info.k_strides[2] == static_cast(info.Dk) + && info.v_strides[0] == static_cast(info.Hv * info.Dv) + && info.v_strides[2] == static_cast(info.Dv) + && info.out_strides[0] == static_cast(info.Hv * info.Dv) + && info.out_strides[2] == static_cast(info.Dv) + && info.g_strides[0] == static_cast(info.Hv) + && info.g_strides[2] == 1 + && info.beta_strides[0] == static_cast(info.Hv) + && info.beta_strides[2] == 1; + opaque->native = info.data_dtype == INFINI_DTYPE_BF16 + && info.gate_dtype == INFINI_DTYPE_F32 + && info.use_qk_l2norm + && info.has_initial_state_indices + && info.has_final_state_indices + && info.Dk == NATIVE_DIM + && info.Dv == NATIVE_DIM + && state_contiguous + && tensors_contiguous; + + size_t workspace_size = info.B * info.Hv * info.Dv * info.Dk * sizeof(float); + if (opaque->native) { + const int64_t B = static_cast(info.B); + const int64_t Hk = static_cast(info.Hk); + const int64_t Hv = static_cast(info.Hv); + const int64_t D = static_cast(NATIVE_DIM); + opaque->q = new aclnnTensorDescriptor( + ACL_BF16, {B, Hk, D}, {Hk * D, D, 1}); + opaque->k = new aclnnTensorDescriptor( + ACL_BF16, {B, Hk, D}, {Hk * D, D, 1}); + opaque->v = new aclnnTensorDescriptor( + ACL_BF16, {B, Hv, D}, {Hv * D, D, 1}); + opaque->beta = new aclnnTensorDescriptor( + ACL_BF16, {B, Hv}, {Hv, 1}); + const int64_t pool = static_cast(info.pool_size); + opaque->state = new aclnnTensorDescriptor( + ACL_BF16, {pool, Hv, D, D}, {Hv * D * D, D * D, D, 1}); + opaque->actual_seq_lengths = new aclnnTensorDescriptor( + ACL_INT32, {B}, {1}); + opaque->state_indices = new aclnnTensorDescriptor( + ACL_INT32, {B}, {1}); + opaque->g = new aclnnTensorDescriptor( + ACL_FLOAT, {B, Hv}, {Hv, 1}); + opaque->out = new aclnnTensorDescriptor( + ACL_BF16, {B, Hv, D}, {Hv * D, D, 1}); + + size_t cursor = 0; + opaque->q_offset = reserveNative( + cursor, info.B * info.Hk * NATIVE_DIM * sizeof(uint16_t)); + opaque->k_offset = reserveNative( + cursor, info.B * info.Hk * NATIVE_DIM * sizeof(uint16_t)); + opaque->v_offset = reserveNative( + cursor, info.B * info.Hv * NATIVE_DIM * sizeof(uint16_t)); + opaque->beta_offset = reserveNative( + cursor, info.B * info.Hv * sizeof(uint16_t)); + opaque->state_offset = reserveNative( + cursor, info.B * info.Hv * NATIVE_DIM * NATIVE_DIM + * sizeof(uint16_t)); + opaque->actual_seq_lengths_offset = reserveNative( + cursor, info.B * sizeof(int32_t)); + opaque->state_indices_offset = reserveNative( + cursor, info.B * sizeof(int32_t)); + opaque->native_workspace_offset = alignNative(cursor); + + CHECK_ACL(aclnnRecurrentGatedDeltaRuleGetWorkspaceSize( + opaque->q->tensor, + opaque->k->tensor, + opaque->v->tensor, + opaque->beta->tensor, + opaque->state->tensor, + opaque->actual_seq_lengths->tensor, + opaque->state_indices->tensor, + opaque->g->tensor, + nullptr, + nullptr, + 1.0f / std::sqrt(static_cast(NATIVE_DIM)), + opaque->out->tensor, + &opaque->native_workspace_size, + &opaque->executor)); + CHECK_ACL(aclSetAclOpExecutorRepeatable(opaque->executor)); + CHECK_ACL(aclrtMalloc( + &opaque->q_buffer, info.B * info.Hk * NATIVE_DIM * sizeof(uint16_t), + ACL_MEM_MALLOC_HUGE_FIRST)); + CHECK_ACL(aclrtMalloc( + &opaque->k_buffer, info.B * info.Hk * NATIVE_DIM * sizeof(uint16_t), + ACL_MEM_MALLOC_HUGE_FIRST)); + CHECK_ACL(aclrtMalloc( + &opaque->v_buffer, info.B * info.Hv * NATIVE_DIM * sizeof(uint16_t), + ACL_MEM_MALLOC_HUGE_FIRST)); + CHECK_ACL(aclrtMalloc( + &opaque->beta_buffer, info.B * info.Hv * sizeof(uint16_t), + ACL_MEM_MALLOC_HUGE_FIRST)); + CHECK_ACL(aclrtMalloc( + &opaque->lengths_buffer, info.B * sizeof(int32_t), + ACL_MEM_MALLOC_HUGE_FIRST)); + CHECK_ACL(aclrtMalloc( + &opaque->indices_buffer, info.B * sizeof(int32_t), + ACL_MEM_MALLOC_HUGE_FIRST)); + if (opaque->native_workspace_size > 0) { + CHECK_ACL(aclrtMalloc( + &opaque->native_workspace_buffer, + opaque->native_workspace_size, ACL_MEM_MALLOC_HUGE_FIRST)); + } + workspace_size = opaque->native_workspace_offset + + opaque->native_workspace_size; + } + + *desc_ptr = new Descriptor( + opaque, std::move(info), workspace_size, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *out, + void *initial_state, + void *final_state, + const void *q, + const void *k, + const void *v, + const void *g, + const void *beta, + const void *initial_state_indices, + const void *final_state_indices, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (_info.gate_dtype != INFINI_DTYPE_F32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + if (_opaque->native) { + auto base = static_cast(workspace); + void *q_normalized = _opaque->q_buffer; + void *k_normalized = _opaque->k_buffer; + void *v_contiguous = const_cast(v); + void *beta_bf16 = _opaque->beta_buffer; + void *state_staging = base + _opaque->state_offset; + void *actual_seq_lengths = _opaque->lengths_buffer; + void *state_indices = _opaque->indices_buffer; + void *native_workspace = _opaque->native_workspace_buffer; + + RecurrentGdrNativeParams p{}; + p.B = _info.B; + p.Hk = _info.Hk; + p.Hv = _info.Hv; + p.pool_size = _info.pool_size; + p.initial_indices_i64 = _info.initial_state_indices_dtype == INFINI_DTYPE_I64; + p.final_indices_i64 = _info.final_state_indices_dtype == INFINI_DTYPE_I64; + p.q_s0 = _info.q_strides[0]; + p.q_s2 = _info.q_strides[2]; + p.k_s0 = _info.k_strides[0]; + p.k_s2 = _info.k_strides[2]; + p.v_s0 = _info.v_strides[0]; + p.v_s2 = _info.v_strides[2]; + p.beta_s0 = _info.beta_strides[0]; + p.beta_s2 = _info.beta_strides[2]; + + CHECK_STATUS(recurrent_gdr_native_preprocess_launch( + q_normalized, k_normalized, v_contiguous, beta_bf16, + state_staging, actual_seq_lengths, state_indices, + q, k, v, beta, initial_state, initial_state_indices, + final_state_indices, &p, stream)); + // A single batched ACLNN call avoids one launch per active request. + // Keep an explicit opt-out for CANN regressions or targeted debugging. + static const bool use_batched = []() { + const char *value = std::getenv("INFINICORE_ASCEND_GDR_BATCHED"); + return value == nullptr || std::strcmp(value, "1") == 0; + }(); + if (use_batched) { + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 0, _opaque->q->tensor, q_normalized)); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 1, _opaque->k->tensor, k_normalized)); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 2, _opaque->v->tensor, v_contiguous)); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 3, _opaque->beta->tensor, beta_bf16)); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 4, _opaque->state->tensor, initial_state)); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 5, _opaque->actual_seq_lengths->tensor, + actual_seq_lengths)); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 6, _opaque->state_indices->tensor, + state_indices)); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 7, _opaque->g->tensor, + const_cast(g))); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 8, _opaque->out->tensor, out)); + CHECK_ACL(aclnnRecurrentGatedDeltaRule( + native_workspace, _opaque->native_workspace_size, + _opaque->executor, stream)); + return INFINI_STATUS_SUCCESS; + } + + const int64_t Hk = static_cast(_info.Hk); + const int64_t Hv = static_cast(_info.Hv); + const int64_t D = static_cast(NATIVE_DIM); + const int64_t pool = static_cast(_info.pool_size); + auto add_bytes = [](const void *ptr, size_t bytes) -> void * { + return const_cast(static_cast(ptr)) + bytes; + }; + for (size_t request = 0; request < _info.B; ++request) { + void *q_data = add_bytes( + q_normalized, request * _info.Hk * NATIVE_DIM * sizeof(uint16_t)); + void *k_data = add_bytes( + k_normalized, request * _info.Hk * NATIVE_DIM * sizeof(uint16_t)); + void *v_data = add_bytes( + v_contiguous, request * _info.Hv * NATIVE_DIM * sizeof(uint16_t)); + void *beta_data = add_bytes( + beta_bf16, request * _info.Hv * sizeof(uint16_t)); + void *length_data = add_bytes( + actual_seq_lengths, request * sizeof(int32_t)); + void *index_data = add_bytes(state_indices, request * sizeof(int32_t)); + void *gate_data = add_bytes( + g, request * _info.Hv * sizeof(float)); + void *out_data = add_bytes( + out, request * _info.Hv * NATIVE_DIM * sizeof(uint16_t)); + + std::vector addresses = { + q_data, k_data, v_data, beta_data, initial_state, + length_data, index_data, gate_data, out_data}; + size_t cache_index = 0; + while (cache_index < _opaque->cached_addresses.size() + && _opaque->cached_addresses[cache_index] != addresses) { + ++cache_index; + } + if (cache_index < _opaque->cached_addresses.size()) { + uint64_t cached_workspace_size = _opaque->cached_workspace_sizes[cache_index]; + if (cached_workspace_size > _opaque->native_workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + CHECK_ACL(aclnnRecurrentGatedDeltaRule( + native_workspace, cached_workspace_size, + _opaque->cached_executors[cache_index], stream)); + continue; + } + + auto *q_desc = new aclnnTensorDescriptor( + ACL_BF16, {1, Hk, D}, {Hk * D, D, 1}, q_data); + auto *k_desc = new aclnnTensorDescriptor( + ACL_BF16, {1, Hk, D}, {Hk * D, D, 1}, k_data); + auto *v_desc = new aclnnTensorDescriptor( + ACL_BF16, {1, Hv, D}, {Hv * D, D, 1}, v_data); + auto *beta_desc = new aclnnTensorDescriptor( + ACL_BF16, {1, Hv}, {Hv, 1}, beta_data); + auto *state_desc = new aclnnTensorDescriptor( + ACL_BF16, {pool, Hv, D, D}, + {Hv * D * D, D * D, D, 1}, initial_state); + auto *lengths_desc = new aclnnTensorDescriptor( + ACL_INT32, {1}, {1}, length_data); + auto *indices_desc = new aclnnTensorDescriptor( + ACL_INT32, {1}, {1}, index_data); + auto *gate_desc = new aclnnTensorDescriptor( + ACL_FLOAT, {1, Hv}, {Hv, 1}, gate_data); + auto *out_desc = new aclnnTensorDescriptor( + ACL_BF16, {1, Hv, D}, {Hv * D, D, 1}, out_data); + aclnnTensorDescriptor *request_descs[] = { + q_desc, k_desc, v_desc, beta_desc, state_desc, + lengths_desc, indices_desc, gate_desc, out_desc}; + uint64_t call_workspace_size = 0; + aclOpExecutor *call_executor = nullptr; + CHECK_ACL(aclnnRecurrentGatedDeltaRuleGetWorkspaceSize( + q_desc->tensor, k_desc->tensor, v_desc->tensor, + beta_desc->tensor, state_desc->tensor, lengths_desc->tensor, + indices_desc->tensor, gate_desc->tensor, nullptr, nullptr, + 1.0f / std::sqrt(static_cast(NATIVE_DIM)), + out_desc->tensor, &call_workspace_size, &call_executor)); + if (call_workspace_size > _opaque->native_workspace_size) { + for (auto *descriptor : request_descs) { + delete descriptor; + } + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + CHECK_ACL(aclSetAclOpExecutorRepeatable(call_executor)); + _opaque->cached_addresses.push_back(std::move(addresses)); + _opaque->cached_descriptors.push_back({q_desc, k_desc, v_desc, beta_desc, state_desc, + lengths_desc, indices_desc, gate_desc, out_desc}); + _opaque->cached_executors.push_back(call_executor); + _opaque->cached_workspace_sizes.push_back(call_workspace_size); + CHECK_ACL(aclnnRecurrentGatedDeltaRule( + native_workspace, call_workspace_size, call_executor, stream)); + } + return INFINI_STATUS_SUCCESS; + } + + GatedDeltaRuleAscendParams p{}; + p.data_dtype = static_cast(_info.data_dtype); + p.gate_dtype = static_cast(_info.gate_dtype); + p.use_qk_l2norm = _info.use_qk_l2norm; + p.has_initial_indices = _info.has_initial_state_indices; + p.initial_indices_i64 = _info.initial_state_indices_dtype == INFINI_DTYPE_I64; + p.has_final_indices = _info.has_final_state_indices; + p.final_indices_i64 = _info.final_state_indices_dtype == INFINI_DTYPE_I64; + p.B = _info.B; + p.T = _info.T; + p.total_tokens = _info.T; + p.Hk = _info.Hk; + p.Hv = _info.Hv; + p.Dk = _info.Dk; + p.Dv = _info.Dv; + p.pool_size = _info.pool_size; + p.value_heads_per_key_head = _info.value_heads_per_key_head; + p.q_scale = 1.0f / std::sqrt(static_cast(_info.Dk)); + for (int i = 0; i < 4; ++i) { + p.out_strides[i] = _info.out_strides[i]; + p.q_strides[i] = _info.q_strides[i]; + p.k_strides[i] = _info.k_strides[i]; + p.v_strides[i] = _info.v_strides[i]; + } + return gated_delta_rule_ascend_kernel_launch( + workspace, out, initial_state, final_state, q, k, v, g, beta, + nullptr, initial_state_indices, final_state_indices, &p, stream); +} + +} // namespace op::recurrent_gated_delta_rule::ascend diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_ascend.h b/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_ascend.h new file mode 100644 index 000000000..9d5b7158c --- /dev/null +++ b/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_ascend.h @@ -0,0 +1,8 @@ +#ifndef __RECURRENT_GATED_DELTA_RULE_ASCEND_H__ +#define __RECURRENT_GATED_DELTA_RULE_ASCEND_H__ + +#include "../recurrent_gated_delta_rule.h" + +DESCRIPTOR(ascend) + +#endif diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_native_kernel.cpp b/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_native_kernel.cpp new file mode 100644 index 000000000..97999c7f0 --- /dev/null +++ b/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_native_kernel.cpp @@ -0,0 +1,366 @@ +#include "recurrent_gated_delta_rule_native_kernel.h" +#include "../../../devices/ascend/ascend_kernel_common.h" + +using namespace AscendC; + +namespace { + +constexpr size_t D = 128; +constexpr size_t MATRIX = D * D; +constexpr size_t STATE_TILE = 4096; + +__aicore__ inline float nativeRsqrt(float x) { + if (x <= 0.0f) { + return 0.0f; + } + float scaled = x; + float rescale = 1.0f; + while (scaled > 4.0f) { + scaled *= 0.25f; + rescale *= 0.5f; + } + while (scaled < 1.0f) { + scaled *= 4.0f; + rescale *= 2.0f; + } + float y = 0.75f; + for (int i = 0; i < 6; ++i) { + y *= 1.5f - 0.5f * scaled * y * y; + } + return y * rescale; +} + +__aicore__ inline int64_t loadIndex( + GM_ADDR ptr, bool is_i64, size_t index) { + if (is_i64) { + return reinterpret_cast<__gm__ int64_t *>(ptr)[index]; + } + return static_cast( + reinterpret_cast<__gm__ int32_t *>(ptr)[index]); +} + +__global__ __aicore__ void recurrent_gdr_native_preprocess( + GM_ADDR q_normalized_ptr, + GM_ADDR k_normalized_ptr, + GM_ADDR v_contiguous_ptr, + GM_ADDR beta_bf16_ptr, + GM_ADDR state_staging_ptr, + GM_ADDR actual_seq_lengths_ptr, + GM_ADDR state_indices_ptr, + GM_ADDR q_ptr, + GM_ADDR k_ptr, + GM_ADDR v_ptr, + GM_ADDR beta_ptr, + GM_ADDR state_ptr, + GM_ADDR initial_state_indices_ptr, + GM_ADDR final_state_indices_ptr, + size_t B, + size_t Hk, + size_t Hv, + size_t pool_size, + bool initial_indices_i64, + bool final_indices_i64, + ptrdiff_t q_s0, + ptrdiff_t q_s2, + ptrdiff_t k_s0, + ptrdiff_t k_s2, + ptrdiff_t v_s0, + ptrdiff_t v_s2, + ptrdiff_t beta_s0, + ptrdiff_t beta_s2) { + + size_t block = GetBlockIdx(); + + TPipe pipe; + TQue q_input_queue, k_input_queue, v_input_queue; + TQue q_output_queue, k_output_queue; + TBuf q_float_buf, k_float_buf, state_copy_buf; + pipe.InitBuffer(q_input_queue, 1, D * sizeof(bfloat16_t)); + pipe.InitBuffer(q_output_queue, 1, D * sizeof(bfloat16_t)); + pipe.InitBuffer(k_input_queue, 1, D * sizeof(bfloat16_t)); + pipe.InitBuffer(k_output_queue, 1, D * sizeof(bfloat16_t)); + pipe.InitBuffer(v_input_queue, 1, D * sizeof(bfloat16_t)); + pipe.InitBuffer(q_float_buf, D * sizeof(float)); + pipe.InitBuffer(k_float_buf, D * sizeof(float)); + pipe.InitBuffer(state_copy_buf, STATE_TILE * sizeof(bfloat16_t)); + + GlobalTensor q_normalized, k_normalized; + GlobalTensor v_contiguous, beta_bf16, state_staging; + GlobalTensor q, k, v, state; + GlobalTensor beta; + GlobalTensor actual_seq_lengths, state_indices; + q_normalized.SetGlobalBuffer( + reinterpret_cast<__gm__ bfloat16_t *>(q_normalized_ptr)); + k_normalized.SetGlobalBuffer( + reinterpret_cast<__gm__ bfloat16_t *>(k_normalized_ptr)); + v_contiguous.SetGlobalBuffer( + reinterpret_cast<__gm__ bfloat16_t *>(v_contiguous_ptr)); + beta_bf16.SetGlobalBuffer( + reinterpret_cast<__gm__ bfloat16_t *>(beta_bf16_ptr)); + state_staging.SetGlobalBuffer( + reinterpret_cast<__gm__ bfloat16_t *>(state_staging_ptr)); + q.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(q_ptr)); + k.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(k_ptr)); + v.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(v_ptr)); + beta.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(beta_ptr)); + state.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(state_ptr)); + actual_seq_lengths.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t *>(actual_seq_lengths_ptr)); + state_indices.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t *>(state_indices_ptr)); + + if (block < B * Hk) { + size_t request = block / Hk; + size_t head = block % Hk; + ptrdiff_t q_base = static_cast(request) * q_s0 + + static_cast(head) * q_s2; + ptrdiff_t k_base = static_cast(request) * k_s0 + + static_cast(head) * k_s2; + size_t output_base = block * D; + + LocalTensor q_input = q_input_queue.AllocTensor(); + DataCopy(q_input, q[q_base], D); + q_input_queue.EnQue(q_input); + q_input = q_input_queue.DeQue(); + LocalTensor q_values = q_float_buf.Get(); + Cast(q_values, q_input, AscendC::RoundMode::CAST_NONE, D); + float q_sum = 0.0f; + for (size_t i = 0; i < D; ++i) { + float value = q_values.GetValue(i); + q_sum += value * value; + } + Muls(q_values, q_values, nativeRsqrt(q_sum), D); + LocalTensor q_output = q_output_queue.AllocTensor(); + Cast(q_output, q_values, AscendC::RoundMode::CAST_RINT, D); + q_output_queue.EnQue(q_output); + q_input_queue.FreeTensor(q_input); + q_output = q_output_queue.DeQue(); + DataCopy(q_normalized[output_base], q_output, D); + q_output_queue.FreeTensor(q_output); + + LocalTensor k_input = k_input_queue.AllocTensor(); + DataCopy(k_input, k[k_base], D); + k_input_queue.EnQue(k_input); + k_input = k_input_queue.DeQue(); + LocalTensor k_values = k_float_buf.Get(); + Cast(k_values, k_input, AscendC::RoundMode::CAST_NONE, D); + float k_sum = 0.0f; + for (size_t i = 0; i < D; ++i) { + float value = k_values.GetValue(i); + k_sum += value * value; + } + Muls(k_values, k_values, nativeRsqrt(k_sum), D); + LocalTensor k_output = k_output_queue.AllocTensor(); + Cast(k_output, k_values, AscendC::RoundMode::CAST_RINT, D); + k_output_queue.EnQue(k_output); + k_input_queue.FreeTensor(k_input); + k_output = k_output_queue.DeQue(); + DataCopy(k_normalized[output_base], k_output, D); + k_output_queue.FreeTensor(k_output); + } + + if (block < B * Hv) { + size_t request = block / Hv; + size_t head = block % Hv; + int64_t read_slot = loadIndex( + initial_state_indices_ptr, initial_indices_i64, request); + int64_t write_slot = loadIndex( + final_state_indices_ptr, final_indices_i64, request); + if (read_slot >= 0 && read_slot < static_cast(pool_size) + && write_slot >= 0 + && write_slot < static_cast(pool_size) + && read_slot != write_slot) { + LocalTensor state_local = state_copy_buf.Get(); + size_t source_base = (static_cast(read_slot) * Hv + head) * MATRIX; + size_t destination_base = (static_cast(write_slot) * Hv + head) * MATRIX; + TEventID mte2_event = GetTPipePtr()->FetchEventID(HardEvent::MTE2_S); + TEventID mte3_event = GetTPipePtr()->FetchEventID(HardEvent::MTE3_S); + for (size_t tile = 0; tile < MATRIX; tile += STATE_TILE) { + DataCopy(state_local, state[source_base + tile], STATE_TILE); + SetFlag(mte2_event); + WaitFlag(mte2_event); + DataCopy( + state[destination_base + tile], + state_local, STATE_TILE); + SetFlag(mte3_event); + WaitFlag(mte3_event); + } + } + } +} + +__global__ __aicore__ void recurrent_gdr_native_metadata( + GM_ADDR actual_seq_lengths_ptr, + GM_ADDR state_indices_ptr, + GM_ADDR initial_state_indices_ptr, + size_t B, + bool initial_indices_i64) { + + GlobalTensor actual_seq_lengths, state_indices; + actual_seq_lengths.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t *>(actual_seq_lengths_ptr)); + state_indices.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t *>(state_indices_ptr)); + for (size_t request = 0; request < B; ++request) { + int64_t state_slot = loadIndex( + initial_state_indices_ptr, initial_indices_i64, request); + actual_seq_lengths.SetValue(request, 1); + state_indices.SetValue(request, static_cast(state_slot)); + } +} + +__global__ __aicore__ void recurrent_gdr_native_cast_beta( + GM_ADDR beta_bf16_ptr, + GM_ADDR beta_ptr, + size_t count) { + + size_t copy_len = alignTileLen(count, BYTE_ALIGN); + GlobalTensor beta; + GlobalTensor beta_bf16; + beta.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(beta_ptr)); + beta_bf16.SetGlobalBuffer( + reinterpret_cast<__gm__ bfloat16_t *>(beta_bf16_ptr)); + + TPipe pipe; + TQue input_queue; + TQue output_queue; + pipe.InitBuffer(input_queue, 1, copy_len * sizeof(float)); + pipe.InitBuffer(output_queue, 1, copy_len * sizeof(bfloat16_t)); + + LocalTensor input = input_queue.AllocTensor(); + DataCopy(input, beta, copy_len); + input_queue.EnQue(input); + input = input_queue.DeQue(); + LocalTensor output = output_queue.AllocTensor(); + Cast(output, input, AscendC::RoundMode::CAST_RINT, copy_len); + output_queue.EnQue(output); + input_queue.FreeTensor(input); + output = output_queue.DeQue(); + if (count * sizeof(bfloat16_t) % BYTE_ALIGN != 0) { + DataCopyExtParams params = { + 1, static_cast(count * sizeof(bfloat16_t)), 0, 0, 0}; + DataCopyPad(beta_bf16, output, params); + } else { + DataCopy(beta_bf16, output, count); + } + output_queue.FreeTensor(output); +} + +__global__ __aicore__ void recurrent_gdr_native_commit_state( + GM_ADDR state_ptr, + GM_ADDR state_staging_ptr, + GM_ADDR initial_state_indices_ptr, + GM_ADDR final_state_indices_ptr, + size_t B, + size_t Hv, + size_t pool_size, + bool initial_indices_i64, + bool final_indices_i64) { + + size_t block = GetBlockIdx(); + if (block >= B * Hv) { + return; + } + size_t request = block / Hv; + size_t head = block % Hv; + int64_t read_slot = loadIndex( + initial_state_indices_ptr, initial_indices_i64, request); + int64_t write_slot = loadIndex( + final_state_indices_ptr, final_indices_i64, request); + if (read_slot < 0 || read_slot >= static_cast(pool_size) + || write_slot < 0 || write_slot >= static_cast(pool_size) + || read_slot == write_slot) { + return; + } + + GlobalTensor state, state_staging; + state.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(state_ptr)); + state_staging.SetGlobalBuffer( + reinterpret_cast<__gm__ bfloat16_t *>(state_staging_ptr)); + + TPipe pipe; + TBuf state_buf; + pipe.InitBuffer(state_buf, STATE_TILE * sizeof(bfloat16_t)); + LocalTensor state_local = state_buf.Get(); + size_t staging_base = block * MATRIX; + size_t source_base = (static_cast(read_slot) * Hv + head) * MATRIX; + size_t destination_base = (static_cast(write_slot) * Hv + head) * MATRIX; + TEventID mte2_event = GetTPipePtr()->FetchEventID(HardEvent::MTE2_S); + TEventID mte3_event = GetTPipePtr()->FetchEventID(HardEvent::MTE3_S); + for (size_t tile = 0; tile < MATRIX; tile += STATE_TILE) { + DataCopy(state_local, state[source_base + tile], STATE_TILE); + SetFlag(mte2_event); + WaitFlag(mte2_event); + DataCopy(state[destination_base + tile], state_local, STATE_TILE); + SetFlag(mte3_event); + WaitFlag(mte3_event); + DataCopy(state_local, state_staging[staging_base + tile], STATE_TILE); + SetFlag(mte2_event); + WaitFlag(mte2_event); + DataCopy(state[source_base + tile], state_local, STATE_TILE); + SetFlag(mte3_event); + WaitFlag(mte3_event); + } +} + +} // namespace + +extern "C" infiniStatus_t recurrent_gdr_native_preprocess_launch( + void *q_normalized, + void *k_normalized, + void *v_contiguous, + void *beta_bf16, + void *state_staging, + void *actual_seq_lengths, + void *state_indices, + const void *q, + const void *k, + const void *v, + const void *beta, + const void *state, + const void *initial_state_indices, + const void *final_state_indices, + const RecurrentGdrNativeParams *p, + void *stream) { + + if (p->B == 0) { + return INFINI_STATUS_SUCCESS; + } + uint32_t blocks = static_cast(p->B * (p->Hv > p->Hk ? p->Hv : p->Hk)); + recurrent_gdr_native_preprocess<<>>( + q_normalized, k_normalized, v_contiguous, beta_bf16, state_staging, + actual_seq_lengths, state_indices, const_cast(q), + const_cast(k), const_cast(v), const_cast(beta), + const_cast(state), const_cast(initial_state_indices), + const_cast(final_state_indices), p->B, p->Hk, p->Hv, + p->pool_size, p->initial_indices_i64, p->final_indices_i64, + p->q_s0, p->q_s2, p->k_s0, p->k_s2, p->v_s0, p->v_s2, + p->beta_s0, p->beta_s2); + recurrent_gdr_native_metadata<<<1, nullptr, stream>>>( + actual_seq_lengths, state_indices, + const_cast(final_state_indices), + p->B, p->final_indices_i64); + recurrent_gdr_native_cast_beta<<<1, nullptr, stream>>>( + beta_bf16, const_cast(beta), p->B * p->Hv); + return INFINI_STATUS_SUCCESS; +} + +extern "C" infiniStatus_t recurrent_gdr_native_commit_state_launch( + void *state, + const void *state_staging, + const void *initial_state_indices, + const void *final_state_indices, + const RecurrentGdrNativeParams *p, + void *stream) { + + if (p->B == 0) { + return INFINI_STATUS_SUCCESS; + } + uint32_t blocks = static_cast(p->B * p->Hv); + recurrent_gdr_native_commit_state<<>>( + state, const_cast(state_staging), + const_cast(initial_state_indices), + const_cast(final_state_indices), p->B, p->Hv, p->pool_size, + p->initial_indices_i64, p->final_indices_i64); + return INFINI_STATUS_SUCCESS; +} diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_native_kernel.h b/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_native_kernel.h new file mode 100644 index 000000000..bfe9e0ceb --- /dev/null +++ b/src/infiniop/ops/recurrent_gated_delta_rule/ascend/recurrent_gated_delta_rule_native_kernel.h @@ -0,0 +1,51 @@ +#ifndef __RECURRENT_GATED_DELTA_RULE_NATIVE_KERNEL_H__ +#define __RECURRENT_GATED_DELTA_RULE_NATIVE_KERNEL_H__ + +#include "../../../../../include/infinicore.h" +#include +#include + +struct RecurrentGdrNativeParams { + size_t B; + size_t Hk; + size_t Hv; + size_t pool_size; + bool initial_indices_i64; + bool final_indices_i64; + ptrdiff_t q_s0; + ptrdiff_t q_s2; + ptrdiff_t k_s0; + ptrdiff_t k_s2; + ptrdiff_t v_s0; + ptrdiff_t v_s2; + ptrdiff_t beta_s0; + ptrdiff_t beta_s2; +}; + +extern "C" infiniStatus_t recurrent_gdr_native_preprocess_launch( + void *q_normalized, + void *k_normalized, + void *v_contiguous, + void *beta_bf16, + void *state_staging, + void *actual_seq_lengths, + void *state_indices, + const void *q, + const void *k, + const void *v, + const void *beta, + const void *state, + const void *initial_state_indices, + const void *final_state_indices, + const RecurrentGdrNativeParams *params, + void *stream); + +extern "C" infiniStatus_t recurrent_gdr_native_commit_state_launch( + void *state, + const void *state_staging, + const void *initial_state_indices, + const void *final_state_indices, + const RecurrentGdrNativeParams *params, + void *stream); + +#endif diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc b/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc index fc272fdf9..b15bbebc7 100644 --- a/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc +++ b/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc @@ -13,6 +13,9 @@ #ifdef ENABLE_MOORE_API #include "moore/recurrent_gated_delta_rule_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/recurrent_gated_delta_rule_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateRecurrentGatedDeltaRuleDescriptor( infiniopHandle_t handle, @@ -49,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateRecurrentGatedDeltaRuleDescriptor( #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) #endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend) +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -75,6 +81,9 @@ __INFINI_C infiniStatus_t infiniopGetRecurrentGatedDeltaRuleWorkspaceSize( #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend) +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -109,6 +118,9 @@ __INFINI_C infiniStatus_t infiniopRecurrentGatedDeltaRule( #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore) #endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend) +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -134,6 +146,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRecurrentGatedDeltaRuleDescriptor( #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) #endif +#ifdef ENABLE_ASCEND_API + DESTROY(INFINI_DEVICE_ASCEND, ascend) +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/rms_norm/ascend/rms_norm_aclnn.cc b/src/infiniop/ops/rms_norm/ascend/rms_norm_aclnn.cc index 2e9eb8a25..c0a53d3aa 100644 --- a/src/infiniop/ops/rms_norm/ascend/rms_norm_aclnn.cc +++ b/src/infiniop/ops/rms_norm/ascend/rms_norm_aclnn.cc @@ -1,7 +1,6 @@ #include "rms_norm_aclnn.h" #include "../../../devices/ascend/common_ascend.h" #include -#include extern "C" infiniStatus_t rms_norm_cast_w_launch( void *dst, const void *src, @@ -57,24 +56,17 @@ infiniStatus_t Descriptor::create( auto handle_ascend = reinterpret_cast(handle); - std::vector slice_shape = {static_cast(info.dim())}; - auto slice_stride = std::vector(1, 1); - - aclnnTensorDescriptor_t y = new aclnnTensorDescriptor(toAclDataType(info.atype), slice_shape, slice_stride); - aclnnTensorDescriptor_t x = new aclnnTensorDescriptor(toAclDataType(info.atype), slice_shape, slice_stride); + aclnnTensorDescriptor_t y = new aclnnTensorDescriptor(y_desc); + aclnnTensorDescriptor_t x = new aclnnTensorDescriptor(x_desc); // 仅在跨半精度组合时需要将 w cast 到 atype // (F16 atype + BF16 w, 或 BF16 atype + F16 w) bool needs_cast_w = (info.atype != info.wtype && info.wtype != INFINI_DTYPE_F32); aclnnTensorDescriptor_t w = nullptr; - std::vector w_shape_i64_dbg; - std::vector w_strides_i64_dbg; if (needs_cast_w) { // 规避 constructor #2 的 ndim 内存 corruption 问题 // 先用 constructor #1 从 w_desc 正确构造，再替换 tensor 为正确的 dtype w = new aclnnTensorDescriptor(w_desc); - w_shape_i64_dbg = w->shape; - w_strides_i64_dbg = w->strides; if (w->tensor) { aclDestroyTensor(w->tensor); } @@ -86,8 +78,15 @@ infiniStatus_t Descriptor::create( w = new aclnnTensorDescriptor(w_desc); } - auto rstd_shape = std::vector(1, 1); - auto rstd_strides = std::vector(1, 1); + std::vector rstd_shape; + rstd_shape.reserve(info.ndim() - 1); + for (size_t i = 0; i + 1 < info.ndim(); ++i) { + rstd_shape.push_back(static_cast(info.shape[i])); + } + std::vector rstd_strides(rstd_shape.size(), 1); + for (ptrdiff_t i = static_cast(rstd_shape.size()) - 2; i >= 0; --i) { + rstd_strides[i] = rstd_strides[i + 1] * rstd_shape[i + 1]; + } aclnnTensorDescriptor_t rstd = new aclnnTensorDescriptor(toAclDataType(INFINI_DTYPE_F32), rstd_shape, rstd_strides); size_t workspace_size = 0; @@ -134,51 +133,29 @@ infiniStatus_t Descriptor::calculate( return INFINI_STATUS_INSUFFICIENT_WORKSPACE; } - auto tw = _opaque->w->tensor; - auto tx = _opaque->x->tensor; - auto ty = _opaque->y->tensor; - auto trstd = _opaque->rstd->tensor; - - void *rstdPtr = (void *)((uint8_t *)workspace + _opaque->workspaceSize); + void *rstd_ptr = static_cast(workspace) + _opaque->workspaceSize; void *w_ptr = nullptr; - if (_opaque->needs_cast_w) { - void *cast_w_ptr = (void *)((uint8_t *)workspace + _opaque->cast_w_offset); - void *w_padded_src = (void *)((uint8_t *)workspace + _opaque->w_padded_offset); + void *cast_w_ptr = static_cast(workspace) + _opaque->cast_w_offset; + void *w_padded_src = static_cast(workspace) + _opaque->w_padded_offset; size_t w_bytes = _info.dim() * infiniSizeOf(_info.wtype); - aclrtMemcpyAsync(w_padded_src, _opaque->w_padded_size, (void *)w, w_bytes, - ACL_MEMCPY_DEVICE_TO_DEVICE, (aclrtStream)stream); - rms_norm_cast_w_launch(cast_w_ptr, w_padded_src, _info.wtype, INFINI_DTYPE_F32, _info.dim(), stream); + CHECK_ACL(aclrtMemcpyAsync( + w_padded_src, _opaque->w_padded_size, const_cast(w), w_bytes, + ACL_MEMCPY_DEVICE_TO_DEVICE, static_cast(stream))); + CHECK_STATUS(rms_norm_cast_w_launch( + cast_w_ptr, w_padded_src, _info.wtype, INFINI_DTYPE_F32, + _info.dim(), stream)); w_ptr = cast_w_ptr; } else { - w_ptr = (void *)w; - } - - auto unit = infiniSizeOf(_info.atype); - - AclSetTensorAddr(_opaque->executor, 1, tw, w_ptr); - AclSetTensorAddr(_opaque->executor, 3, trstd, rstdPtr); - - auto ndim = _info.ndim(); - size_t outer = ndim == 2 ? 1 : _info.shape[0]; - size_t inner = ndim == 2 ? _info.shape[0] : _info.shape[1]; - - for (size_t b = 0; b < outer; ++b) { - for (size_t s = 0; s < inner; ++s) { - ptrdiff_t x_offset, y_offset; - if (ndim == 2) { - x_offset = s * _info.x_strides[0]; - y_offset = s * _info.y_strides[0]; - } else { - x_offset = b * _info.x_strides[0] + s * _info.x_strides[1]; - y_offset = b * _info.y_strides[0] + s * _info.y_strides[1]; - } - AclSetTensorAddr(_opaque->executor, 0, tx, ((char *)x) + x_offset * unit); - AclSetTensorAddr(_opaque->executor, 2, ty, ((char *)y) + y_offset * unit); - CHECK_ACL(aclnnRmsNorm(workspace, _opaque->workspaceSize, _opaque->executor, stream)); - } + w_ptr = const_cast(w); } + CHECK_ACL(AclSetTensorAddr(_opaque->executor, 0, _opaque->x->tensor, const_cast(x))); + CHECK_ACL(AclSetTensorAddr(_opaque->executor, 1, _opaque->w->tensor, w_ptr)); + CHECK_ACL(AclSetTensorAddr(_opaque->executor, 2, _opaque->y->tensor, y)); + CHECK_ACL(AclSetTensorAddr(_opaque->executor, 3, _opaque->rstd->tensor, rstd_ptr)); + CHECK_ACL(aclnnRmsNorm( + workspace, _opaque->workspaceSize, _opaque->executor, stream)); return INFINI_STATUS_SUCCESS; } diff --git a/src/infiniop/ops/sigmoid/ascend/sigmoid_ascend.cc b/src/infiniop/ops/sigmoid/ascend/sigmoid_ascend.cc new file mode 100644 index 000000000..520c9fe72 --- /dev/null +++ b/src/infiniop/ops/sigmoid/ascend/sigmoid_ascend.cc @@ -0,0 +1,71 @@ +#include "sigmoid_ascend.h" + +#include "../../../devices/ascend/aclnn_executor.h" +#include +#include + +namespace op::sigmoid::ascend { + +struct Descriptor::Opaque { + device::ascend::AclnnExecutor op; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + std::vector input_descs) { + + auto status = device::ascend::validateAclnnElementwise(output_desc, input_descs, 1); + if (status != INFINI_STATUS_SUCCESS) { + return status; + } + CHECK_DTYPE(output_desc->dtype(), + INFINI_DTYPE_F16, INFINI_DTYPE_F32, + INFINI_DTYPE_F64, INFINI_DTYPE_BF16); + + auto opaque = std::make_unique(); + opaque->op.tensors = { + new aclnnTensorDescriptor(input_descs[0]), + new aclnnTensorDescriptor(output_desc), + }; + + CHECK_ACL(aclnnSigmoidGetWorkspaceSize( + opaque->op.tensors[0]->tensor, + opaque->op.tensors[1]->tensor, + &opaque->op.workspace_size, + &opaque->op.executor)); + aclSetAclOpExecutorRepeatable(opaque->op.executor); + + auto handle_ascend = reinterpret_cast(handle); + auto workspace_size = opaque->op.workspace_size; + *desc_ptr = new Descriptor( + opaque.release(), workspace_size, + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *output, std::vector inputs, + void *stream) const { + + if (inputs.size() != 1) { + return INFINI_STATUS_BAD_PARAM; + } + if (workspace_size < workspaceSize()) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + _opaque->op.bind({const_cast(inputs[0]), output}); + CHECK_ACL(aclnnSigmoid( + workspace, workspace_size, _opaque->op.executor, + static_cast(stream))); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::sigmoid::ascend diff --git a/src/infiniop/ops/sigmoid/ascend/sigmoid_ascend.h b/src/infiniop/ops/sigmoid/ascend/sigmoid_ascend.h new file mode 100644 index 000000000..18425ab3b --- /dev/null +++ b/src/infiniop/ops/sigmoid/ascend/sigmoid_ascend.h @@ -0,0 +1,8 @@ +#ifndef __SIGMOID_ASCEND_H__ +#define __SIGMOID_ASCEND_H__ + +#include "../../../devices/ascend/aclnn_elementwise.h" + +ACLNN_ELEMENTWISE_DESCRIPTOR(sigmoid) + +#endif // __SIGMOID_ASCEND_H__ diff --git a/src/infiniop/ops/sigmoid/operator.cc b/src/infiniop/ops/sigmoid/operator.cc index 0497c0ba7..5eb4656e7 100644 --- a/src/infiniop/ops/sigmoid/operator.cc +++ b/src/infiniop/ops/sigmoid/operator.cc @@ -14,6 +14,9 @@ #ifdef ENABLE_MOORE_API #include "moore/sigmoid_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/sigmoid_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateSigmoidDescriptor( infiniopHandle_t handle, @@ -55,6 +58,9 @@ __INFINI_C infiniStatus_t infiniopCreateSigmoidDescriptor( #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -94,6 +100,9 @@ __INFINI_C infiniStatus_t infiniopGetSigmoidWorkspaceSize(infiniopSigmoidDescrip #endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) +#endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -142,6 +151,9 @@ __INFINI_C infiniStatus_t infiniopSigmoid( #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -184,6 +196,9 @@ infiniopDestroySigmoidDescriptor(infiniopSigmoidDescriptor_t desc) { #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/silu/ascend/silu_ascend.cc b/src/infiniop/ops/silu/ascend/silu_ascend.cc new file mode 100644 index 000000000..140b9bdd8 --- /dev/null +++ b/src/infiniop/ops/silu/ascend/silu_ascend.cc @@ -0,0 +1,129 @@ +#include "silu_ascend.h" + +#include "../../../devices/ascend/aclnn_executor.h" +#include "../../swiglu/ascend/swiglu_ascend.h" +#include +#include +#include +#include + +namespace op::silu::ascend { + +struct Descriptor::Opaque { + std::unique_ptr swiglu_info; + device::ascend::AclnnExecutor sigmoid; + device::ascend::AclnnExecutor mul; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + std::vector input_descs) { + if (false) { + auto result = op::swiglu::ascend::SwigluInfo::create( + output_desc, input_descs[0], input_descs[0]); + CHECK_RESULT(result); + auto opaque = std::make_unique(); + opaque->swiglu_info = std::make_unique( + result.take()); + auto handle_ascend = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + opaque.release(), 0, + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; + } + + auto status = device::ascend::validateAclnnElementwise(output_desc, input_descs, 1); + if (status != INFINI_STATUS_SUCCESS) { + return status; + } + CHECK_DTYPE(output_desc->dtype(), + INFINI_DTYPE_F16, INFINI_DTYPE_F32, + INFINI_DTYPE_F64, INFINI_DTYPE_BF16); + + auto opaque = std::make_unique(); + opaque->sigmoid.tensors = { + new aclnnTensorDescriptor(input_descs[0]), + new aclnnTensorDescriptor(output_desc), + }; + opaque->mul.tensors = { + new aclnnTensorDescriptor(output_desc), + new aclnnTensorDescriptor(input_descs[0]), + new aclnnTensorDescriptor(output_desc), + }; + + CHECK_ACL(aclnnSigmoidGetWorkspaceSize( + opaque->sigmoid.tensors[0]->tensor, + opaque->sigmoid.tensors[1]->tensor, + &opaque->sigmoid.workspace_size, + &opaque->sigmoid.executor)); + aclSetAclOpExecutorRepeatable(opaque->sigmoid.executor); + CHECK_ACL(aclnnMulGetWorkspaceSize( + opaque->mul.tensors[0]->tensor, + opaque->mul.tensors[1]->tensor, + opaque->mul.tensors[2]->tensor, + &opaque->mul.workspace_size, + &opaque->mul.executor)); + aclSetAclOpExecutorRepeatable(opaque->mul.executor); + + auto handle_ascend = reinterpret_cast(handle); + auto workspace_size = std::max( + opaque->sigmoid.workspace_size, + opaque->mul.workspace_size); + *desc_ptr = new Descriptor( + opaque.release(), workspace_size, + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *output, std::vector inputs, + void *stream) const { + + if (inputs.size() != 1) { + return INFINI_STATUS_BAD_PARAM; + } + if (workspace_size < workspaceSize()) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (false) { + const auto &info = *_opaque->swiglu_info; + auto batch = info.ndim == 2 ? 1 : info.shape[0]; + auto seq_len = info.ndim == 2 ? info.shape[0] : info.shape[1]; + auto hidden_size = info.shape[info.ndim - 1]; + auto stride_batch_out = info.ndim == 2 ? 1 : info.c_strides[0]; + auto stride_batch_in = info.ndim == 2 ? 1 : info.a_strides[0]; + auto stride_seq_out = info.ndim == 2 ? info.c_strides[0] : info.c_strides[1]; + auto stride_seq_in = info.ndim == 2 ? info.a_strides[0] : info.a_strides[1]; + return op::swiglu::ascend::swiglu_kernel_launch( + output, + const_cast(inputs[0]), + const_cast(inputs[0]), + info.dtype, batch, seq_len, hidden_size, + stride_batch_out, stride_batch_in, stride_batch_in, + stride_seq_out, stride_seq_in, stride_seq_in, + stream); + } + _opaque->sigmoid.bind({const_cast(inputs[0]), output}); + CHECK_ACL(aclnnSigmoid( + workspace, workspace_size, _opaque->sigmoid.executor, + static_cast(stream))); + + _opaque->mul.bind({ + output, + const_cast(inputs[0]), + output, + }); + CHECK_ACL(aclnnMul( + workspace, workspace_size, _opaque->mul.executor, + static_cast(stream))); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::silu::ascend diff --git a/src/infiniop/ops/silu/ascend/silu_ascend.h b/src/infiniop/ops/silu/ascend/silu_ascend.h new file mode 100644 index 000000000..c87948d09 --- /dev/null +++ b/src/infiniop/ops/silu/ascend/silu_ascend.h @@ -0,0 +1,8 @@ +#ifndef __SILU_ASCEND_H__ +#define __SILU_ASCEND_H__ + +#include "../../../devices/ascend/aclnn_elementwise.h" + +ACLNN_ELEMENTWISE_DESCRIPTOR(silu) + +#endif // __SILU_ASCEND_H__ diff --git a/src/infiniop/ops/silu/operator.cc b/src/infiniop/ops/silu/operator.cc index 602d5d178..502fd0a3c 100644 --- a/src/infiniop/ops/silu/operator.cc +++ b/src/infiniop/ops/silu/operator.cc @@ -14,6 +14,9 @@ #ifdef ENABLE_MOORE_API #include "moore/silu_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/silu_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateSiluDescriptor( infiniopHandle_t handle, @@ -52,6 +55,9 @@ __INFINI_C infiniStatus_t infiniopCreateSiluDescriptor( #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -89,6 +95,9 @@ __INFINI_C infiniStatus_t infiniopGetSiluWorkspaceSize(infiniopSiluDescriptor_t #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -134,6 +143,9 @@ __INFINI_C infiniStatus_t infiniopSilu( #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -173,6 +185,9 @@ infiniopDestroySiluDescriptor(infiniopSiluDescriptor_t desc) { #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); #endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/upsample_bilinear/ascend/upsample_bilinear_ascend.cc b/src/infiniop/ops/upsample_bilinear/ascend/upsample_bilinear_ascend.cc new file mode 100644 index 000000000..1fe3a927a --- /dev/null +++ b/src/infiniop/ops/upsample_bilinear/ascend/upsample_bilinear_ascend.cc @@ -0,0 +1,105 @@ +#include "upsample_bilinear_ascend.h" + +#include "../../../devices/ascend/common_ascend.h" +#include +#include + +namespace op::upsample_bilinear::ascend { + +struct Descriptor::Opaque { + aclnnTensorDescriptor_t input; + aclnnTensorDescriptor_t output; + aclIntArray *output_size; + aclOpExecutor *executor; + + Opaque(aclnnTensorDescriptor_t input_, + aclnnTensorDescriptor_t output_, + aclIntArray *output_size_, + aclOpExecutor *executor_) + : input(input_), output(output_), output_size(output_size_), + executor(executor_) {} + + ~Opaque() { + delete input; + delete output; + if (output_size != nullptr) { + aclDestroyIntArray(output_size); + } + if (executor != nullptr) { + aclDestroyAclOpExecutor(executor); + } + } +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + int align_corners) { + + auto result = UpsampleBilinearInfo::create( + output_desc, input_desc, align_corners); + CHECK_RESULT(result); + auto info = result.take(); + + CHECK_DTYPE(output_desc->dtype(), + INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + + auto input = std::make_unique( + input_desc, nullptr, ACL_FORMAT_NCHW); + auto output = std::make_unique( + output_desc, nullptr, ACL_FORMAT_NCHW); + std::vector output_size_data = { + static_cast(info.h_out()), + static_cast(info.w_out()), + }; + aclIntArray *output_size = aclCreateIntArray( + output_size_data.data(), output_size_data.size()); + if (output_size == nullptr) { + return INFINI_STATUS_INTERNAL_ERROR; + } + + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + auto acl_status = aclnnUpsampleBilinear2dGetWorkspaceSize( + input->tensor, output_size, info.align_corners(), + 0.0, 0.0, output->tensor, + &workspace_size, &executor); + if (acl_status != ACL_SUCCESS) { + GetRecentErrMsg(); + aclDestroyIntArray(output_size); + CHECK_ACL(acl_status); + } + aclSetAclOpExecutorRepeatable(executor); + + auto handle_ascend = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + new Opaque{input.release(), output.release(), output_size, executor}, + std::move(info), workspace_size, + handle_ascend->device, handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *output, const void *input, + void *stream) const { + + if (workspace_size < workspaceSize()) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + AclSetTensorAddr(_opaque->executor, 0, _opaque->input->tensor, + const_cast(input)); + AclSetTensorAddr(_opaque->executor, 1, _opaque->output->tensor, output); + CHECK_ACL(aclnnUpsampleBilinear2d( + workspace, workspace_size, _opaque->executor, + static_cast(stream))); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::upsample_bilinear::ascend diff --git a/src/infiniop/ops/upsample_bilinear/ascend/upsample_bilinear_ascend.h b/src/infiniop/ops/upsample_bilinear/ascend/upsample_bilinear_ascend.h new file mode 100644 index 000000000..814e31d46 --- /dev/null +++ b/src/infiniop/ops/upsample_bilinear/ascend/upsample_bilinear_ascend.h @@ -0,0 +1,8 @@ +#ifndef __UPSAMPLE_BILINEAR_ASCEND_H__ +#define __UPSAMPLE_BILINEAR_ASCEND_H__ + +#include "../upsample_bilinear.h" + +DESCRIPTOR(ascend) + +#endif // __UPSAMPLE_BILINEAR_ASCEND_H__ diff --git a/src/infiniop/ops/upsample_bilinear/operator.cc b/src/infiniop/ops/upsample_bilinear/operator.cc index eb03e6a21..9e253a820 100644 --- a/src/infiniop/ops/upsample_bilinear/operator.cc +++ b/src/infiniop/ops/upsample_bilinear/operator.cc @@ -18,6 +18,10 @@ #include "moore/upsample_bilinear_moore.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/upsample_bilinear_ascend.h" +#endif + extern "C" { // ======================================================================= @@ -60,6 +64,9 @@ __INFINI_C infiniStatus_t infiniopCreateUpsampleBilinearDescriptor( #endif #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ASCEND_API + CREATE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -98,6 +105,9 @@ __INFINI_C infiniStatus_t infiniopGetUpsampleBilinearWorkspaceSize(infiniopUpsam #endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ASCEND_API + GET(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -142,6 +152,9 @@ __INFINI_C infiniStatus_t infiniopUpsampleBilinear( #endif #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ASCEND_API + CALCULATE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -180,6 +193,9 @@ __INFINI_C infiniStatus_t infiniopDestroyUpsampleBilinearDescriptor(infiniopUpsa #endif #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ASCEND_API + DELETE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/test/infinicore/ops/causal_conv1d.py b/test/infinicore/ops/causal_conv1d.py index b3ec8498e..61698b5a9 100644 --- a/test/infinicore/ops/causal_conv1d.py +++ b/test/infinicore/ops/causal_conv1d.py @@ -3,7 +3,6 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -import infinicore import torch from framework import ( BaseOperatorTest, @@ -13,6 +12,7 @@ TestCase, ) +import infinicore # Test cases: # (qkv_shape, qkv_strides, state_shape, weight_shape, bias_shape, @@ -144,6 +144,42 @@ def parse_test_cases(): ) ) + if final_state_indices is not None: + tests.append( + TestCase( + inputs=inputs, + kwargs={}, + output_spec=None, + comparison_target=1, + tolerance=_TOLERANCE_MAP[dtype], + description="CausalConv1d - UPDATED_STATE", + ) + ) + + # Large indexed state update: C=2048 keeps all eight Ascend blocks active + # for FP16/BF16 under the 512-byte alignment rule and catches cache-line + # races that small shapes cannot expose. + for dtype in _TENSOR_DTYPES: + inputs = [ + TensorSpec.from_tensor((1, 1, 2048), None, dtype), + TensorSpec.from_tensor((2, 2048, 3), None, dtype), + TensorSpec.from_tensor((2048, 1, 4), None, dtype), + TensorSpec.from_tensor((2048,), None, dtype), + _manual_i32((0, 1)), + _manual_i32((1,)), + _manual_i32((1,)), + ] + tests.append( + TestCase( + inputs=inputs, + kwargs={}, + output_spec=None, + comparison_target=1, + tolerance=_TOLERANCE_MAP[dtype], + description="CausalConv1d - MULTICORE_UPDATED_STATE", + ) + ) + return tests @@ -177,7 +213,7 @@ def torch_operator(self, qkv, conv_state, weight, *args, **kwargs): ) return torch_causal_conv1d_ref( qkv, - conv_state.clone(), + conv_state, weight, bias=bias, cu_seqlens=cu_seqlens, diff --git a/test/infinicore/ops/mha_kvcache.py b/test/infinicore/ops/mha_kvcache.py index f5be8e05f..0ea587eeb 100644 --- a/test/infinicore/ops/mha_kvcache.py +++ b/test/infinicore/ops/mha_kvcache.py @@ -21,6 +21,8 @@ (2, 1, 4, 4, 64, 256, [7, 250]), (2, 1, 8, 2, 128, 256, [73, 260]), (3, 1, 8, 1, 128, 256, [1, 257, 511]), + # Qwen3.5-27B TP=4 decode shape: N=6, Nkv=1, D=256. + (2, 1, 6, 1, 256, 256, [73, 260]), ] _TOLERANCE_MAP = { @@ -196,6 +198,10 @@ def infinicore_operator( block_table, scale=1.0, ): + if q.device.type == "npu": + # Ascend FIA consumes physical BnNBsD paged cache. + k_cache = k_cache.permute([0, 2, 1, 3]).contiguous() + v_cache = v_cache.permute([0, 2, 1, 3]).contiguous() out = infinicore.mha_kvcache( q, k_cache, diff --git a/test/infinicore/ops/mha_varlen.py b/test/infinicore/ops/mha_varlen.py index 894ff3d17..0cc736832 100644 --- a/test/infinicore/ops/mha_varlen.py +++ b/test/infinicore/ops/mha_varlen.py @@ -22,6 +22,8 @@ (1, 1, 128, 256, [(260, 73), (1, 1)]), (8, 2, 128, 256, [(250,), (7,)]), (8, 2, 128, 256, [(260, 73), (1, 1)]), + (6, 1, 256, 256, [(250,), (7,)]), + (6, 1, 256, 256, [(260, 73), (1, 1)]), ] _MAX_SEQUENCE_LENGTH = 8192 @@ -338,8 +340,11 @@ def infinicore_operator( key = k_cache value = v_cache else: - key = k_cache.permute([0, 2, 1, 3]) - value = v_cache.permute([0, 2, 1, 3]) + if query.device.type == "npu": + key, value = k_cache, v_cache + else: + key = k_cache.permute([0, 2, 1, 3]) + value = v_cache.permute([0, 2, 1, 3]) out = infinicore.mha_varlen( query, key, diff --git a/test/infiniop/gelu.py b/test/infiniop/gelu.py index bb97dd0a2..edfbc9353 100644 --- a/test/infiniop/gelu.py +++ b/test/infiniop/gelu.py @@ -1,24 +1,25 @@ -import torch import ctypes from ctypes import c_uint64 +from enum import Enum, auto + +import torch from libinfiniop import ( LIBINFINIOP, + InfiniDeviceEnum, + InfiniDeviceNames, + InfiniDtype, + InfiniDtypeNames, TestTensor, - get_test_devices, + TestWorkspace, check_error, - test_operator, - get_args, debug, + get_args, + get_test_devices, get_tolerance, - profile_operation, - TestWorkspace, - InfiniDtype, - InfiniDtypeNames, - InfiniDeviceNames, - InfiniDeviceEnum, infiniopOperatorDescriptor_t, + profile_operation, + test_operator, ) -from enum import Enum, auto # ============================================================================== # Configuration (Internal Use Only) @@ -58,14 +59,13 @@ class Inplace(Enum): ] # Data types used for testing -_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32, InfiniDtype.F64] +_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32] # Tolerance map for different data types _TOLERANCE_MAP = { InfiniDtype.BF16: {"atol": 1e-2, "rtol": 1e-2}, InfiniDtype.F16: {"atol": 1e-3, "rtol": 1e-3}, InfiniDtype.F32: {"atol": 1e-5, "rtol": 1e-5}, - InfiniDtype.F64: {"atol": 1e-6, "rtol": 1e-6}, } DEBUG = False @@ -89,7 +89,10 @@ def test( input_stride is not None or output_stride is not None ): return - if device in (InfiniDeviceEnum.CAMBRICON, InfiniDeviceEnum.MOORE) and dtype == InfiniDtype.F64: + if ( + device in (InfiniDeviceEnum.CAMBRICON, InfiniDeviceEnum.MOORE) + and dtype == InfiniDtype.F64 + ): return input = TestTensor(shape, input_stride, dtype, device) diff --git a/test/infiniop/gelutanh.py b/test/infiniop/gelutanh.py index 48cc90f0f..7ac97aa8e 100644 --- a/test/infiniop/gelutanh.py +++ b/test/infiniop/gelutanh.py @@ -49,13 +49,12 @@ class Inplace(Enum): for inplace_item in _INPLACE ] -_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32, InfiniDtype.F64] +_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32] _TOLERANCE_MAP = { InfiniDtype.BF16: {"atol": 2e-2, "rtol": 2e-2}, InfiniDtype.F16: {"atol": 2e-3, "rtol": 2e-3}, InfiniDtype.F32: {"atol": 1e-5, "rtol": 1e-5}, - InfiniDtype.F64: {"atol": 1e-6, "rtol": 1e-6}, } DEBUG = False diff --git a/xmake.lua b/xmake.lua index eecfb4b3e..fc2cacd54 100644 --- a/xmake.lua +++ b/xmake.lua @@ -709,7 +709,7 @@ target("infinicore_cpp_api") add_includedirs(ASCEND_HOME .. "/include/aclnnop") add_includedirs(path.join(os.projectdir(), "submodules/InfiniOps/src")) add_linkdirs(ASCEND_HOME .. "/lib64") - add_links("ascendcl", "nnopbase", "opapi", "runtime") + add_links("ascendcl", "nnopbase", "opapi", "runtime", "dl") add_linkdirs(ASCEND_HOME .. "/../../driver/lib64/driver") add_links("ascend_hal") end