From e868f5c54fbde0ca2b71c86950af94fa6f5ef7fd Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 19 Aug 2026 15:39:18 +0800 Subject: [PATCH] feat: ascend select last token hidden --- .../ascend/select_last_token_hidden_ascend.cc | 225 ++++++++++++++++++ .../ascend/select_last_token_hidden_ascend.h | 8 + .../bang/select_last_token_hidden_bang.mlu | 3 + .../cpu/select_last_token_hidden_cpu.cc | 3 + .../metax/select_last_token_hidden_metax.maca | 3 + .../moore/select_last_token_hidden_moore.mu | 3 + .../nvidia/select_last_token_hidden_nvidia.cu | 3 + .../ops/select_last_token_hidden/operator.cc | 12 + .../select_last_token_hidden.h | 6 + test/infiniop/select_last_token_hidden.py | 1 + 10 files changed, 267 insertions(+) create mode 100644 src/infiniop/ops/select_last_token_hidden/ascend/select_last_token_hidden_ascend.cc create mode 100644 src/infiniop/ops/select_last_token_hidden/ascend/select_last_token_hidden_ascend.h diff --git a/src/infiniop/ops/select_last_token_hidden/ascend/select_last_token_hidden_ascend.cc b/src/infiniop/ops/select_last_token_hidden/ascend/select_last_token_hidden_ascend.cc new file mode 100644 index 000000000..3eb588e86 --- /dev/null +++ b/src/infiniop/ops/select_last_token_hidden/ascend/select_last_token_hidden_ascend.cc @@ -0,0 +1,225 @@ +#include "select_last_token_hidden_ascend.h" +#include "../../../devices/ascend/common_ascend.h" +#include +#include +#include +#include +#include + +namespace op::select_last_token_hidden::ascend { + +struct Descriptor::Opaque { + aclnnTensorDescriptor_t output; + aclnnTensorDescriptor_t hidden_states; + aclnnTensorDescriptor_t offsets; + aclnnTensorDescriptor_t indices; + aclnnScalarDescriptor_t minus_one; + aclnnScalarDescriptor_t alpha; + void *minus_one_value; + void *alpha_value; + void *indices_data; + void *workspace; + uint64_t adds_workspace_size; + uint64_t index_select_workspace_size; + aclOpExecutor *adds_executor; + aclOpExecutor *index_select_executor; + + ~Opaque() { + delete output; + delete hidden_states; + delete offsets; + delete indices; + delete minus_one; + delete alpha; + std::free(minus_one_value); + std::free(alpha_value); + if (indices_data != nullptr) { + aclrtFree(indices_data); + } + if (workspace != nullptr) { + aclrtFree(workspace); + } + aclDestroyAclOpExecutor(adds_executor); + aclDestroyAclOpExecutor(index_select_executor); + } +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t hidden_states_desc, + infiniopTensorDescriptor_t input_offsets_desc) { + const auto output_shape = output_desc->shape(); + const auto hidden_shape = hidden_states_desc->shape(); + const auto offsets_shape = input_offsets_desc->shape(); + + CHECK_OR_RETURN(output_shape.size() == 3 && hidden_shape.size() == 3 && offsets_shape.size() == 1, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(offsets_shape[0] >= 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + const size_t num_requests = offsets_shape[0] - 1; + CHECK_OR_RETURN(output_shape[0] == 1 && output_shape[1] == num_requests + && output_shape[2] == hidden_shape[2], + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(output_desc->isContiguous() && hidden_states_desc->isContiguous() + && input_offsets_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + CHECK_OR_RETURN(input_offsets_desc->dtype() == INFINI_DTYPE_I32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + const auto hidden_dtype = hidden_states_desc->dtype(); + CHECK_OR_RETURN(hidden_dtype == INFINI_DTYPE_F16 || hidden_dtype == INFINI_DTYPE_BF16 + || hidden_dtype == INFINI_DTYPE_F32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(output_desc->dtype() == hidden_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); + + const size_t total_tokens = hidden_shape[0] * hidden_shape[1]; + const size_t hidden_size = hidden_shape[2]; + CHECK_OR_RETURN(total_tokens > 0 && hidden_size > 0, INFINI_STATUS_BAD_TENSOR_SHAPE); + + const auto hidden_acl_dtype = toAclDataType(hidden_dtype); + auto output = new aclnnTensorDescriptor( + hidden_acl_dtype, + {static_cast(num_requests), static_cast(hidden_size)}, + {static_cast(hidden_size), 1}); + auto hidden_states = new aclnnTensorDescriptor( + hidden_acl_dtype, + {static_cast(total_tokens), static_cast(hidden_size)}, + {static_cast(hidden_size), 1}); + auto offsets = new aclnnTensorDescriptor( + ACL_INT32, + {static_cast(num_requests)}, + {1}); + auto indices = new aclnnTensorDescriptor( + ACL_INT32, + {static_cast(num_requests)}, + {1}); + + auto minus_one_value = std::malloc(sizeof(int32_t)); + auto alpha_value = std::malloc(sizeof(int32_t)); + CHECK_OR_RETURN(minus_one_value != nullptr && alpha_value != nullptr, + INFINI_STATUS_INSUFFICIENT_WORKSPACE); + *static_cast(minus_one_value) = -1; + *static_cast(alpha_value) = 1; + auto minus_one = new aclnnScalarDescriptor( + ACL_INT32, minus_one_value, sizeof(int32_t)); + auto alpha = new aclnnScalarDescriptor( + ACL_INT32, alpha_value, sizeof(int32_t)); + + uint64_t adds_workspace_size = 0; + aclOpExecutor *adds_executor = nullptr; + CHECK_ACL(aclnnAddsGetWorkspaceSize( + offsets->tensor, + minus_one->scalar, + alpha->scalar, + indices->tensor, + &adds_workspace_size, + &adds_executor)); + aclSetAclOpExecutorRepeatable(adds_executor); + + uint64_t index_select_workspace_size = 0; + aclOpExecutor *index_select_executor = nullptr; + CHECK_ACL(aclnnIndexSelectGetWorkspaceSize( + hidden_states->tensor, + 0, + indices->tensor, + output->tensor, + &index_select_workspace_size, + &index_select_executor)); + aclSetAclOpExecutorRepeatable(index_select_executor); + + void *indices_data = nullptr; + CHECK_ACL(aclrtMalloc( + &indices_data, + num_requests * sizeof(int32_t), + ACL_MEM_MALLOC_HUGE_FIRST)); + + void *workspace = nullptr; + const uint64_t workspace_size = std::max( + adds_workspace_size, + index_select_workspace_size); + if (workspace_size != 0) { + CHECK_ACL(aclrtMalloc( + &workspace, + workspace_size, + ACL_MEM_MALLOC_HUGE_FIRST)); + } + + auto opaque = new Opaque{ + output, + hidden_states, + offsets, + indices, + minus_one, + alpha, + minus_one_value, + alpha_value, + indices_data, + workspace, + adds_workspace_size, + index_select_workspace_size, + adds_executor, + index_select_executor}; + + auto handle_ascend = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + num_requests, + total_tokens, + hidden_size * infiniSizeOf(hidden_dtype), + opaque, + handle_ascend->device, + handle_ascend->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *output, + const void *hidden_states, + const void *input_offsets, + void *stream) const { + auto offsets_data = const_cast( + static_cast(input_offsets) + 1); + + AclSetTensorAddr( + _opaque->adds_executor, + 0, + _opaque->offsets->tensor, + offsets_data); + AclSetTensorAddr( + _opaque->adds_executor, + 1, + _opaque->indices->tensor, + _opaque->indices_data); + CHECK_ACL(aclnnAdds( + _opaque->workspace, + _opaque->adds_workspace_size, + _opaque->adds_executor, + stream)); + + AclSetTensorAddr( + _opaque->index_select_executor, + 0, + _opaque->hidden_states->tensor, + const_cast(hidden_states)); + AclSetTensorAddr( + _opaque->index_select_executor, + 1, + _opaque->indices->tensor, + _opaque->indices_data); + AclSetTensorAddr( + _opaque->index_select_executor, + 2, + _opaque->output->tensor, + output); + CHECK_ACL(aclnnIndexSelect( + _opaque->workspace, + _opaque->index_select_workspace_size, + _opaque->index_select_executor, + stream)); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::select_last_token_hidden::ascend diff --git a/src/infiniop/ops/select_last_token_hidden/ascend/select_last_token_hidden_ascend.h b/src/infiniop/ops/select_last_token_hidden/ascend/select_last_token_hidden_ascend.h new file mode 100644 index 000000000..36650c71b --- /dev/null +++ b/src/infiniop/ops/select_last_token_hidden/ascend/select_last_token_hidden_ascend.h @@ -0,0 +1,8 @@ +#ifndef __SELECT_LAST_TOKEN_HIDDEN_ASCEND_H__ +#define __SELECT_LAST_TOKEN_HIDDEN_ASCEND_H__ + +#include "../select_last_token_hidden.h" + +DESCRIPTOR(ascend) + +#endif // __SELECT_LAST_TOKEN_HIDDEN_ASCEND_H__ diff --git a/src/infiniop/ops/select_last_token_hidden/bang/select_last_token_hidden_bang.mlu b/src/infiniop/ops/select_last_token_hidden/bang/select_last_token_hidden_bang.mlu index dea71cb30..317bd6528 100644 --- a/src/infiniop/ops/select_last_token_hidden/bang/select_last_token_hidden_bang.mlu +++ b/src/infiniop/ops/select_last_token_hidden/bang/select_last_token_hidden_bang.mlu @@ -60,6 +60,8 @@ __mlu_global__ void selectLastTokenHiddenKernel( namespace op::select_last_token_hidden::bang { +Descriptor::~Descriptor() {} + infiniStatus_t Descriptor::create( infiniopHandle_t handle_, Descriptor **desc_ptr, @@ -107,6 +109,7 @@ infiniStatus_t Descriptor::create( num_requests, total_tokens, hidden_shape[2] * infiniSizeOf(hidden_dtype), + nullptr, handle->device, handle->device_id, internal->getCorePerCluster(), diff --git a/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.cc b/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.cc index d8820f812..f9719553e 100644 --- a/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.cc +++ b/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.cc @@ -8,6 +8,8 @@ namespace op::select_last_token_hidden::cpu { +Descriptor::~Descriptor() {} + infiniStatus_t Descriptor::create( infiniopHandle_t handle, Descriptor **desc_ptr, @@ -42,6 +44,7 @@ infiniStatus_t Descriptor::create( num_requests, total_tokens, hidden_shape[2] * infiniSizeOf(hidden_dtype), + nullptr, handle->device, handle->device_id); return INFINI_STATUS_SUCCESS; diff --git a/src/infiniop/ops/select_last_token_hidden/metax/select_last_token_hidden_metax.maca b/src/infiniop/ops/select_last_token_hidden/metax/select_last_token_hidden_metax.maca index 12eaad603..46c1351b7 100644 --- a/src/infiniop/ops/select_last_token_hidden/metax/select_last_token_hidden_metax.maca +++ b/src/infiniop/ops/select_last_token_hidden/metax/select_last_token_hidden_metax.maca @@ -24,6 +24,8 @@ INFINIOP_METAX_KERNEL selectLastTokenHiddenKernel( namespace op::select_last_token_hidden::metax { +Descriptor::~Descriptor() {} + infiniStatus_t Descriptor::create( infiniopHandle_t handle, Descriptor **desc_ptr, @@ -58,6 +60,7 @@ infiniStatus_t Descriptor::create( num_requests, total_tokens, hidden_shape[2] * infiniSizeOf(hidden_dtype), + nullptr, handle->device, handle->device_id); return INFINI_STATUS_SUCCESS; diff --git a/src/infiniop/ops/select_last_token_hidden/moore/select_last_token_hidden_moore.mu b/src/infiniop/ops/select_last_token_hidden/moore/select_last_token_hidden_moore.mu index f616999c6..370140e49 100644 --- a/src/infiniop/ops/select_last_token_hidden/moore/select_last_token_hidden_moore.mu +++ b/src/infiniop/ops/select_last_token_hidden/moore/select_last_token_hidden_moore.mu @@ -25,6 +25,8 @@ INFINIOP_MOORE_KERNEL selectLastTokenHiddenKernel( namespace op::select_last_token_hidden::moore { +Descriptor::~Descriptor() {} + infiniStatus_t Descriptor::create( infiniopHandle_t handle, Descriptor **desc_ptr, @@ -59,6 +61,7 @@ infiniStatus_t Descriptor::create( num_requests, total_tokens, hidden_shape[2] * infiniSizeOf(hidden_dtype), + nullptr, handle->device, handle->device_id); return INFINI_STATUS_SUCCESS; diff --git a/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cu b/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cu index 1e8db42f1..c6b55514e 100644 --- a/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cu +++ b/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cu @@ -24,6 +24,8 @@ INFINIOP_CUDA_KERNEL selectLastTokenHiddenKernel( namespace op::select_last_token_hidden::nvidia { +Descriptor::~Descriptor() {} + infiniStatus_t Descriptor::create( infiniopHandle_t handle, Descriptor **desc_ptr, @@ -58,6 +60,7 @@ infiniStatus_t Descriptor::create( num_requests, total_tokens, hidden_shape[2] * infiniSizeOf(hidden_dtype), + nullptr, handle->device, handle->device_id); return INFINI_STATUS_SUCCESS; diff --git a/src/infiniop/ops/select_last_token_hidden/operator.cc b/src/infiniop/ops/select_last_token_hidden/operator.cc index f862f849c..b18093fb2 100644 --- a/src/infiniop/ops/select_last_token_hidden/operator.cc +++ b/src/infiniop/ops/select_last_token_hidden/operator.cc @@ -17,6 +17,9 @@ #ifdef ENABLE_CAMBRICON_API #include "bang/select_last_token_hidden_bang.h" #endif +#ifdef ENABLE_ASCEND_API +#include "ascend/select_last_token_hidden_ascend.h" +#endif __INFINI_C infiniStatus_t infiniopCreateSelectLastTokenHiddenDescriptor( infiniopHandle_t handle, @@ -59,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopCreateSelectLastTokenHiddenDescriptor( #endif #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; @@ -106,6 +112,9 @@ __INFINI_C infiniStatus_t infiniopSelectLastTokenHidden( #endif #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; @@ -149,6 +158,9 @@ __INFINI_C infiniStatus_t infiniopDestroySelectLastTokenHiddenDescriptor( #endif #ifdef ENABLE_CAMBRICON_API DESTROY(INFINI_DEVICE_CAMBRICON, bang); +#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/select_last_token_hidden/select_last_token_hidden.h b/src/infiniop/ops/select_last_token_hidden/select_last_token_hidden.h index fe8d4d4b7..846297c24 100644 --- a/src/infiniop/ops/select_last_token_hidden/select_last_token_hidden.h +++ b/src/infiniop/ops/select_last_token_hidden/select_last_token_hidden.h @@ -7,6 +7,8 @@ #define DESCRIPTOR(NAMESPACE) \ namespace op::select_last_token_hidden::NAMESPACE { \ class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ size_t _num_requests; \ size_t _total_tokens; \ size_t _row_bytes; \ @@ -17,11 +19,13 @@ size_t num_requests, \ size_t total_tokens, \ size_t row_bytes, \ + Opaque *opaque, \ infiniDevice_t device_type, \ int device_id, \ size_t kernel_dim_x = 1, \ size_t kernel_dim_y = 1) \ : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), \ _num_requests(num_requests), \ _total_tokens(total_tokens), \ _row_bytes(row_bytes), \ @@ -29,6 +33,8 @@ _kernel_dim_y(kernel_dim_y) {} \ \ public: \ + ~Descriptor(); \ + \ static infiniStatus_t create( \ infiniopHandle_t handle, \ Descriptor **desc_ptr, \ diff --git a/test/infiniop/select_last_token_hidden.py b/test/infiniop/select_last_token_hidden.py index 20b54ae52..b2726d2c5 100644 --- a/test/infiniop/select_last_token_hidden.py +++ b/test/infiniop/select_last_token_hidden.py @@ -43,6 +43,7 @@ _TEST_CASES = [ (3, 17, 65, (0, 2, 8, 17)), (4, 25, 64, (0, 3, 11, 18, 25)), + (3, 10, 63, (0, 2, 6, 10)), (4, 2048, 6144, (0, 512, 1024, 1536, 2048)), ] _TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32]