You need to sign in or sign up before continuing.
Unverified Commit dce99862 authored by thatPepe's avatar thatPepe Committed by GitHub
Browse files

Merge pull request #1053 from InfiniTensor/issue/1033xmake

Issue/1033 patch aten and fa adaptations
parents 8d99a8f5 d6e44e84
...@@ -31,15 +31,15 @@ struct InfiniopAttentionDescriptor { ...@@ -31,15 +31,15 @@ struct InfiniopAttentionDescriptor {
float qk_alpha; float qk_alpha;
}; };
__C __export infiniStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t handle, __INFINI_C __export infiniStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t handle,
infiniopAttentionDescriptor_t *desc_ptr, infiniopAttentionDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t out_desc,
infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t q_desc,
infiniopTensorDescriptor_t k_desc, infiniopTensorDescriptor_t k_desc,
infiniopTensorDescriptor_t v_desc, infiniopTensorDescriptor_t v_desc,
infiniopTensorDescriptor_t k_cache_desc, infiniopTensorDescriptor_t k_cache_desc,
infiniopTensorDescriptor_t v_cache_desc, infiniopTensorDescriptor_t v_cache_desc,
size_t pos) { size_t pos) {
if (out_desc->ndim() != 3 || q_desc->ndim() != 3 || k_desc->ndim() != 3 || v_desc->ndim() != 3 || k_cache_desc->ndim() != 3 || v_cache_desc->ndim() != 3) { if (out_desc->ndim() != 3 || q_desc->ndim() != 3 || k_desc->ndim() != 3 || v_desc->ndim() != 3 || k_cache_desc->ndim() != 3 || v_cache_desc->ndim() != 3) {
return INFINI_STATUS_BAD_TENSOR_SHAPE; return INFINI_STATUS_BAD_TENSOR_SHAPE;
} }
...@@ -218,21 +218,21 @@ __C __export infiniStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t h ...@@ -218,21 +218,21 @@ __C __export infiniStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t h
return INFINI_STATUS_SUCCESS; return INFINI_STATUS_SUCCESS;
} }
__C __export infiniStatus_t infiniopGetAttentionWorkspaceSize(infiniopAttentionDescriptor_t desc, size_t *size) { __INFINI_C __export infiniStatus_t infiniopGetAttentionWorkspaceSize(infiniopAttentionDescriptor_t desc, size_t *size) {
*size = ((InfiniopAttentionDescriptor *)desc)->workspace_size; *size = ((InfiniopAttentionDescriptor *)desc)->workspace_size;
return INFINI_STATUS_SUCCESS; return INFINI_STATUS_SUCCESS;
} }
__C __export infiniStatus_t infiniopAttention(infiniopAttentionDescriptor_t desc_, __INFINI_C __export infiniStatus_t infiniopAttention(infiniopAttentionDescriptor_t desc_,
void *workspace_, void *workspace_,
size_t workspace_size_, size_t workspace_size_,
void *out, void *out,
void const *q, void const *q,
void const *k, void const *k,
void const *v, void const *v,
void *k_cache, void *k_cache,
void *v_cache, void *v_cache,
void *stream) { void *stream) {
auto desc = (InfiniopAttentionDescriptor *)desc_; auto desc = (InfiniopAttentionDescriptor *)desc_;
if (workspace_size_ < desc->workspace_size) { if (workspace_size_ < desc->workspace_size) {
return INFINI_STATUS_INSUFFICIENT_WORKSPACE; // STATUS_MEMORY_NOT_ALLOCATED return INFINI_STATUS_INSUFFICIENT_WORKSPACE; // STATUS_MEMORY_NOT_ALLOCATED
...@@ -274,7 +274,7 @@ __C __export infiniStatus_t infiniopAttention(infiniopAttentionDescriptor_t desc ...@@ -274,7 +274,7 @@ __C __export infiniStatus_t infiniopAttention(infiniopAttentionDescriptor_t desc
return INFINI_STATUS_SUCCESS; return INFINI_STATUS_SUCCESS;
} }
__C __export infiniStatus_t infiniopDestroyAttentionDescriptor(infiniopAttentionDescriptor_t desc_) { __INFINI_C __export infiniStatus_t infiniopDestroyAttentionDescriptor(infiniopAttentionDescriptor_t desc_) {
auto desc = (InfiniopAttentionDescriptor *)desc_; auto desc = (InfiniopAttentionDescriptor *)desc_;
if (desc->rearrange_desc_q) { if (desc->rearrange_desc_q) {
CHECK_STATUS(infiniopDestroyRearrangeDescriptor(desc->rearrange_desc_q)); CHECK_STATUS(infiniopDestroyRearrangeDescriptor(desc->rearrange_desc_q));
......
...@@ -24,7 +24,7 @@ ...@@ -24,7 +24,7 @@
#include "moore/causal_softmax_moore.h" #include "moore/causal_softmax_moore.h"
#endif #endif
__C infiniStatus_t infiniopCreateCausalSoftmaxDescriptor( __INFINI_C infiniStatus_t infiniopCreateCausalSoftmaxDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopCausalSoftmaxDescriptor_t *desc_ptr, infiniopCausalSoftmaxDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t y_desc,
...@@ -78,7 +78,7 @@ __C infiniStatus_t infiniopCreateCausalSoftmaxDescriptor( ...@@ -78,7 +78,7 @@ __C infiniStatus_t infiniopCreateCausalSoftmaxDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDescriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDescriptor_t desc, size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
...@@ -125,7 +125,7 @@ __C infiniStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDe ...@@ -125,7 +125,7 @@ __C infiniStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDe
#undef GET #undef GET
} }
__C infiniStatus_t infiniopCausalSoftmax( __INFINI_C infiniStatus_t infiniopCausalSoftmax(
infiniopCausalSoftmaxDescriptor_t desc, infiniopCausalSoftmaxDescriptor_t desc,
void *workspace, size_t workspace_size, void *workspace, size_t workspace_size,
void *y, void *y,
...@@ -177,7 +177,7 @@ __C infiniStatus_t infiniopCausalSoftmax( ...@@ -177,7 +177,7 @@ __C infiniStatus_t infiniopCausalSoftmax(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t infiniopDestroyCausalSoftmaxDescriptor(infiniopCausalSoftmaxDescriptor_t desc) { __INFINI_C infiniStatus_t infiniopDestroyCausalSoftmaxDescriptor(infiniopCausalSoftmaxDescriptor_t desc) {
#define DESTROY(CASE, NAMESPACE) \ #define DESTROY(CASE, NAMESPACE) \
case CASE: \ case CASE: \
......
...@@ -15,7 +15,7 @@ ...@@ -15,7 +15,7 @@
#include "kunlun/clip_kunlun.h" #include "kunlun/clip_kunlun.h"
#endif #endif
__C infiniStatus_t infiniopCreateClipDescriptor( __INFINI_C infiniStatus_t infiniopCreateClipDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopClipDescriptor_t *desc_ptr, infiniopClipDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t y, infiniopTensorDescriptor_t y,
...@@ -62,7 +62,7 @@ __C infiniStatus_t infiniopCreateClipDescriptor( ...@@ -62,7 +62,7 @@ __C infiniStatus_t infiniopCreateClipDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetClipWorkspaceSize(infiniopClipDescriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetClipWorkspaceSize(infiniopClipDescriptor_t desc, size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
...@@ -98,7 +98,7 @@ __C infiniStatus_t infiniopGetClipWorkspaceSize(infiniopClipDescriptor_t desc, s ...@@ -98,7 +98,7 @@ __C infiniStatus_t infiniopGetClipWorkspaceSize(infiniopClipDescriptor_t desc, s
#undef GET #undef GET
} }
__C infiniStatus_t infiniopClip( __INFINI_C infiniStatus_t infiniopClip(
infiniopClipDescriptor_t desc, infiniopClipDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -144,7 +144,7 @@ __C infiniStatus_t infiniopClip( ...@@ -144,7 +144,7 @@ __C infiniStatus_t infiniopClip(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyClipDescriptor(infiniopClipDescriptor_t desc) { infiniopDestroyClipDescriptor(infiniopClipDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -9,16 +9,16 @@ ...@@ -9,16 +9,16 @@
#include "nvidia/conv_nvidia.cuh" #include "nvidia/conv_nvidia.cuh"
#endif #endif
__C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t handle, __INFINI_C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t handle,
infiniopConvDescriptor_t *desc_ptr, infiniopConvDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t y_desc,
infiniopTensorDescriptor_t x_desc, infiniopTensorDescriptor_t x_desc,
infiniopTensorDescriptor_t w_desc, infiniopTensorDescriptor_t w_desc,
infiniopTensorDescriptor_t b_desc, infiniopTensorDescriptor_t b_desc,
void *pads, void *pads,
void *strides, void *strides,
void *dilations, void *dilations,
size_t n) { size_t n) {
#define CREATE(CASE, NAMESPACE) \ #define CREATE(CASE, NAMESPACE) \
case CASE: \ case CASE: \
return op::conv::NAMESPACE::Descriptor::create( \ return op::conv::NAMESPACE::Descriptor::create( \
...@@ -55,7 +55,7 @@ __C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t handle ...@@ -55,7 +55,7 @@ __C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t handle
#undef CREATE #undef CREATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopGetConvWorkspaceSize( infiniopGetConvWorkspaceSize(
infiniopConvDescriptor_t desc, infiniopConvDescriptor_t desc,
size_t *size) { size_t *size) {
...@@ -90,7 +90,7 @@ infiniopGetConvWorkspaceSize( ...@@ -90,7 +90,7 @@ infiniopGetConvWorkspaceSize(
#undef GET #undef GET
} }
__C infiniStatus_t infiniopConv( __INFINI_C infiniStatus_t infiniopConv(
infiniopConvDescriptor_t desc, infiniopConvDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -131,7 +131,7 @@ __C infiniStatus_t infiniopConv( ...@@ -131,7 +131,7 @@ __C infiniStatus_t infiniopConv(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyConvDescriptor(infiniopConvDescriptor_t desc) { infiniopDestroyConvDescriptor(infiniopConvDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
case CASE: \ case CASE: \
......
...@@ -12,7 +12,7 @@ ...@@ -12,7 +12,7 @@
#include "iluvatar/dequantize_w42f16_iluvatar.cuh" #include "iluvatar/dequantize_w42f16_iluvatar.cuh"
#endif #endif
__C infiniStatus_t infiniopCreateDequantizeAWQDescriptor( __INFINI_C infiniStatus_t infiniopCreateDequantizeAWQDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopDequantizeAWQDescriptor_t *desc_ptr, infiniopDequantizeAWQDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t out_desc,
...@@ -54,8 +54,8 @@ __C infiniStatus_t infiniopCreateDequantizeAWQDescriptor( ...@@ -54,8 +54,8 @@ __C infiniStatus_t infiniopCreateDequantizeAWQDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetDequantizeAWQWorkspaceSize(infiniopDequantizeAWQDescriptor_t desc, __INFINI_C infiniStatus_t infiniopGetDequantizeAWQWorkspaceSize(infiniopDequantizeAWQDescriptor_t desc,
size_t *size) { size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
*size = reinterpret_cast<const op::dequantize_awq::NAMESPACE::Descriptor *>(desc)->workspaceSize(); \ *size = reinterpret_cast<const op::dequantize_awq::NAMESPACE::Descriptor *>(desc)->workspaceSize(); \
...@@ -84,7 +84,7 @@ __C infiniStatus_t infiniopGetDequantizeAWQWorkspaceSize(infiniopDequantizeAWQDe ...@@ -84,7 +84,7 @@ __C infiniStatus_t infiniopGetDequantizeAWQWorkspaceSize(infiniopDequantizeAWQDe
#undef GET #undef GET
} }
__C infiniStatus_t infiniopDequantizeAWQ( __INFINI_C infiniStatus_t infiniopDequantizeAWQ(
infiniopDequantizeAWQDescriptor_t desc, infiniopDequantizeAWQDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -123,7 +123,7 @@ __C infiniStatus_t infiniopDequantizeAWQ( ...@@ -123,7 +123,7 @@ __C infiniStatus_t infiniopDequantizeAWQ(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyDequantizeAWQDescriptor(infiniopDequantizeAWQDescriptor_t desc) { infiniopDestroyDequantizeAWQDescriptor(infiniopDequantizeAWQDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -15,7 +15,7 @@ ...@@ -15,7 +15,7 @@
#include "moore/embedding_moore.h" #include "moore/embedding_moore.h"
#endif #endif
__C infiniStatus_t infiniopCreateEmbeddingDescriptor( __INFINI_C infiniStatus_t infiniopCreateEmbeddingDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopEmbeddingDescriptor_t *desc_ptr, infiniopEmbeddingDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t output_desc, infiniopTensorDescriptor_t output_desc,
...@@ -65,7 +65,7 @@ __C infiniStatus_t infiniopCreateEmbeddingDescriptor( ...@@ -65,7 +65,7 @@ __C infiniStatus_t infiniopCreateEmbeddingDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopEmbedding( __INFINI_C infiniStatus_t infiniopEmbedding(
infiniopEmbeddingDescriptor_t desc, infiniopEmbeddingDescriptor_t desc,
void *output, void *output,
const void *input, const void *input,
...@@ -111,7 +111,7 @@ __C infiniStatus_t infiniopEmbedding( ...@@ -111,7 +111,7 @@ __C infiniStatus_t infiniopEmbedding(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t infiniopDestroyEmbeddingDescriptor(infiniopEmbeddingDescriptor_t desc) { __INFINI_C infiniStatus_t infiniopDestroyEmbeddingDescriptor(infiniopEmbeddingDescriptor_t desc) {
#define DESTROY(CASE, NAMESPACE) \ #define DESTROY(CASE, NAMESPACE) \
case CASE: \ case CASE: \
......
...@@ -8,7 +8,7 @@ ...@@ -8,7 +8,7 @@
#endif #endif
#endif #endif
__C infiniStatus_t infiniopCreateFlashAttentionDescriptor( __INFINI_C infiniStatus_t infiniopCreateFlashAttentionDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopFlashAttentionDescriptor_t *desc_ptr, infiniopFlashAttentionDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t out_desc,
...@@ -45,7 +45,7 @@ __C infiniStatus_t infiniopCreateFlashAttentionDescriptor( ...@@ -45,7 +45,7 @@ __C infiniStatus_t infiniopCreateFlashAttentionDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetFlashAttentionWorkspaceSize( __INFINI_C infiniStatus_t infiniopGetFlashAttentionWorkspaceSize(
infiniopFlashAttentionDescriptor_t desc, infiniopFlashAttentionDescriptor_t desc,
size_t *size) { size_t *size) {
...@@ -69,7 +69,7 @@ __C infiniStatus_t infiniopGetFlashAttentionWorkspaceSize( ...@@ -69,7 +69,7 @@ __C infiniStatus_t infiniopGetFlashAttentionWorkspaceSize(
#undef GET_SIZE #undef GET_SIZE
} }
__C infiniStatus_t infiniopFlashAttention( __INFINI_C infiniStatus_t infiniopFlashAttention(
infiniopFlashAttentionDescriptor_t desc, infiniopFlashAttentionDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -99,7 +99,7 @@ __C infiniStatus_t infiniopFlashAttention( ...@@ -99,7 +99,7 @@ __C infiniStatus_t infiniopFlashAttention(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t infiniopDestroyFlashAttentionDescriptor( __INFINI_C infiniStatus_t infiniopDestroyFlashAttentionDescriptor(
infiniopFlashAttentionDescriptor_t desc) { infiniopFlashAttentionDescriptor_t desc) {
#define DESTROY(CASE, NAMESPACE) \ #define DESTROY(CASE, NAMESPACE) \
......
...@@ -15,7 +15,7 @@ ...@@ -15,7 +15,7 @@
#include "kunlun/gelu_kunlun.h" #include "kunlun/gelu_kunlun.h"
#endif #endif
__C infiniStatus_t infiniopCreateGeluDescriptor( __INFINI_C infiniStatus_t infiniopCreateGeluDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopGeluDescriptor_t *desc_ptr, infiniopGeluDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t output_desc, infiniopTensorDescriptor_t output_desc,
...@@ -60,7 +60,7 @@ __C infiniStatus_t infiniopCreateGeluDescriptor( ...@@ -60,7 +60,7 @@ __C infiniStatus_t infiniopCreateGeluDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t desc, size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
...@@ -98,7 +98,7 @@ __C infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t desc, s ...@@ -98,7 +98,7 @@ __C infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t desc, s
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
} }
__C infiniStatus_t infiniopGelu( __INFINI_C infiniStatus_t infiniopGelu(
infiniopGeluDescriptor_t desc, infiniopGeluDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -142,7 +142,7 @@ __C infiniStatus_t infiniopGelu( ...@@ -142,7 +142,7 @@ __C infiniStatus_t infiniopGelu(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyGeluDescriptor(infiniopGeluDescriptor_t desc) { infiniopDestroyGeluDescriptor(infiniopGeluDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -24,7 +24,7 @@ ...@@ -24,7 +24,7 @@
#include "kunlun/gemm_kunlun.h" #include "kunlun/gemm_kunlun.h"
#endif #endif
__C infiniStatus_t infiniopCreateGemmDescriptor( __INFINI_C infiniStatus_t infiniopCreateGemmDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopGemmDescriptor_t *desc_ptr, infiniopGemmDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t c_desc, infiniopTensorDescriptor_t c_desc,
...@@ -84,7 +84,7 @@ __C infiniStatus_t infiniopCreateGemmDescriptor( ...@@ -84,7 +84,7 @@ __C infiniStatus_t infiniopCreateGemmDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopGetGemmWorkspaceSize( infiniopGetGemmWorkspaceSize(
infiniopGemmDescriptor_t desc, infiniopGemmDescriptor_t desc,
size_t *size) { size_t *size) {
...@@ -137,7 +137,7 @@ infiniopGetGemmWorkspaceSize( ...@@ -137,7 +137,7 @@ infiniopGetGemmWorkspaceSize(
#undef GET #undef GET
} }
__C infiniStatus_t infiniopGemm( __INFINI_C infiniStatus_t infiniopGemm(
infiniopGemmDescriptor_t desc, infiniopGemmDescriptor_t desc,
void *workspace, size_t workspace_size, void *workspace, size_t workspace_size,
void *c, void *c,
...@@ -198,7 +198,7 @@ __C infiniStatus_t infiniopGemm( ...@@ -198,7 +198,7 @@ __C infiniStatus_t infiniopGemm(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyGemmDescriptor(infiniopGemmDescriptor_t desc) { infiniopDestroyGemmDescriptor(infiniopGemmDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -9,7 +9,7 @@ ...@@ -9,7 +9,7 @@
#include "metax/kv_caching_metax.h" #include "metax/kv_caching_metax.h"
#endif #endif
__C infiniStatus_t infiniopCreateKVCachingDescriptor( __INFINI_C infiniStatus_t infiniopCreateKVCachingDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopKVCachingDescriptor_t *desc_ptr, infiniopKVCachingDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t k_cache, infiniopTensorDescriptor_t k_cache,
...@@ -57,7 +57,7 @@ __C infiniStatus_t infiniopCreateKVCachingDescriptor( ...@@ -57,7 +57,7 @@ __C infiniStatus_t infiniopCreateKVCachingDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetKVCachingWorkspaceSize( __INFINI_C infiniStatus_t infiniopGetKVCachingWorkspaceSize(
infiniopKVCachingDescriptor_t desc, infiniopKVCachingDescriptor_t desc,
size_t *size) { size_t *size) {
...@@ -95,7 +95,7 @@ __C infiniStatus_t infiniopGetKVCachingWorkspaceSize( ...@@ -95,7 +95,7 @@ __C infiniStatus_t infiniopGetKVCachingWorkspaceSize(
#undef GET_SIZE #undef GET_SIZE
} }
__C infiniStatus_t infiniopKVCaching( __INFINI_C infiniStatus_t infiniopKVCaching(
infiniopKVCachingDescriptor_t desc, infiniopKVCachingDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -139,7 +139,7 @@ __C infiniStatus_t infiniopKVCaching( ...@@ -139,7 +139,7 @@ __C infiniStatus_t infiniopKVCaching(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t infiniopDestroyKVCachingDescriptor( __INFINI_C infiniStatus_t infiniopDestroyKVCachingDescriptor(
infiniopKVCachingDescriptor_t desc) { infiniopKVCachingDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -12,7 +12,7 @@ ...@@ -12,7 +12,7 @@
#include "metax/layer_norm_metax.h" #include "metax/layer_norm_metax.h"
#endif #endif
__C infiniStatus_t infiniopCreateLayerNormDescriptor( __INFINI_C infiniStatus_t infiniopCreateLayerNormDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopLayerNormDescriptor_t *desc_ptr, infiniopLayerNormDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t output_desc, infiniopTensorDescriptor_t output_desc,
...@@ -63,7 +63,7 @@ __C infiniStatus_t infiniopCreateLayerNormDescriptor( ...@@ -63,7 +63,7 @@ __C infiniStatus_t infiniopCreateLayerNormDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDescriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDescriptor_t desc, size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
*size = reinterpret_cast<op::layer_norm::NAMESPACE::Descriptor *>(desc)->workspaceSize(); \ *size = reinterpret_cast<op::layer_norm::NAMESPACE::Descriptor *>(desc)->workspaceSize(); \
...@@ -96,7 +96,7 @@ __C infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDescriptor ...@@ -96,7 +96,7 @@ __C infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDescriptor
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
} }
__C infiniStatus_t infiniopLayerNorm( __INFINI_C infiniStatus_t infiniopLayerNorm(
infiniopLayerNormDescriptor_t desc, infiniopLayerNormDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -149,7 +149,7 @@ __C infiniStatus_t infiniopLayerNorm( ...@@ -149,7 +149,7 @@ __C infiniStatus_t infiniopLayerNorm(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyLayerNormDescriptor(infiniopLayerNormDescriptor_t desc) { infiniopDestroyLayerNormDescriptor(infiniopLayerNormDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -15,7 +15,7 @@ ...@@ -15,7 +15,7 @@
// #include "ascend/logsoftmax_ascend.h" // #include "ascend/logsoftmax_ascend.h"
#endif #endif
__C infiniStatus_t infiniopCreateLogSoftmaxDescriptor( __INFINI_C infiniStatus_t infiniopCreateLogSoftmaxDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopLogSoftmaxDescriptor_t *desc_ptr, infiniopLogSoftmaxDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t y_desc,
...@@ -56,7 +56,7 @@ __C infiniStatus_t infiniopCreateLogSoftmaxDescriptor( ...@@ -56,7 +56,7 @@ __C infiniStatus_t infiniopCreateLogSoftmaxDescriptor(
} }
} }
__C infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
...@@ -90,7 +90,7 @@ __C infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescript ...@@ -90,7 +90,7 @@ __C infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescript
} }
} }
__C infiniStatus_t infiniopLogSoftmax( __INFINI_C infiniStatus_t infiniopLogSoftmax(
infiniopLogSoftmaxDescriptor_t desc, infiniopLogSoftmaxDescriptor_t desc,
void *workspace, size_t workspace_size, void *workspace, size_t workspace_size,
void *y, void *y,
...@@ -129,7 +129,7 @@ __C infiniStatus_t infiniopLogSoftmax( ...@@ -129,7 +129,7 @@ __C infiniStatus_t infiniopLogSoftmax(
} }
} }
__C infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc) { __INFINI_C infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc) {
#define DESTROY(CASE, NAMESPACE) \ #define DESTROY(CASE, NAMESPACE) \
case CASE: \ case CASE: \
......
...@@ -6,7 +6,7 @@ ...@@ -6,7 +6,7 @@
#include "nvidia/lp_norm_nvidia.cuh" #include "nvidia/lp_norm_nvidia.cuh"
#endif #endif
__C infiniStatus_t infiniopCreateLPNormDescriptor( __INFINI_C infiniStatus_t infiniopCreateLPNormDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopLPNormDescriptor_t *desc_ptr, infiniopLPNormDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t output_desc, infiniopTensorDescriptor_t output_desc,
...@@ -47,7 +47,7 @@ __C infiniStatus_t infiniopCreateLPNormDescriptor( ...@@ -47,7 +47,7 @@ __C infiniStatus_t infiniopCreateLPNormDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetLPNormWorkspaceSize(infiniopLPNormDescriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetLPNormWorkspaceSize(infiniopLPNormDescriptor_t desc, size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
*size = reinterpret_cast<op::lp_norm::NAMESPACE::Descriptor *>(desc)->workspaceSize(); \ *size = reinterpret_cast<op::lp_norm::NAMESPACE::Descriptor *>(desc)->workspaceSize(); \
...@@ -75,7 +75,7 @@ __C infiniStatus_t infiniopGetLPNormWorkspaceSize(infiniopLPNormDescriptor_t des ...@@ -75,7 +75,7 @@ __C infiniStatus_t infiniopGetLPNormWorkspaceSize(infiniopLPNormDescriptor_t des
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
} }
__C infiniStatus_t infiniopLPNorm( __INFINI_C infiniStatus_t infiniopLPNorm(
infiniopLPNormDescriptor_t desc, infiniopLPNormDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -114,7 +114,7 @@ __C infiniStatus_t infiniopLPNorm( ...@@ -114,7 +114,7 @@ __C infiniStatus_t infiniopLPNorm(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyLPNormDescriptor(infiniopLPNormDescriptor_t desc) { infiniopDestroyLPNormDescriptor(infiniopLPNormDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -18,7 +18,7 @@ ...@@ -18,7 +18,7 @@
#include "moore/mul_moore.h" #include "moore/mul_moore.h"
#endif #endif
__C infiniStatus_t infiniopCreateMulDescriptor( __INFINI_C infiniStatus_t infiniopCreateMulDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopMulDescriptor_t *desc_ptr, infiniopMulDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t c_desc, infiniopTensorDescriptor_t c_desc,
...@@ -68,7 +68,7 @@ __C infiniStatus_t infiniopCreateMulDescriptor( ...@@ -68,7 +68,7 @@ __C infiniStatus_t infiniopCreateMulDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetMulWorkspaceSize(infiniopMulDescriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetMulWorkspaceSize(infiniopMulDescriptor_t desc, size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
...@@ -109,7 +109,7 @@ __C infiniStatus_t infiniopGetMulWorkspaceSize(infiniopMulDescriptor_t desc, siz ...@@ -109,7 +109,7 @@ __C infiniStatus_t infiniopGetMulWorkspaceSize(infiniopMulDescriptor_t desc, siz
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
} }
__C infiniStatus_t infiniopMul( __INFINI_C infiniStatus_t infiniopMul(
infiniopMulDescriptor_t desc, infiniopMulDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -157,7 +157,7 @@ __C infiniStatus_t infiniopMul( ...@@ -157,7 +157,7 @@ __C infiniStatus_t infiniopMul(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyMulDescriptor(infiniopMulDescriptor_t desc) { infiniopDestroyMulDescriptor(infiniopMulDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -15,7 +15,7 @@ ...@@ -15,7 +15,7 @@
#include "moore/ones_moore.h" #include "moore/ones_moore.h"
#endif #endif
__C infiniStatus_t infiniopCreateOnesDescriptor( __INFINI_C infiniStatus_t infiniopCreateOnesDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopOnesDescriptor_t *desc_ptr, infiniopOnesDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t y_desc,
...@@ -60,7 +60,7 @@ __C infiniStatus_t infiniopCreateOnesDescriptor( ...@@ -60,7 +60,7 @@ __C infiniStatus_t infiniopCreateOnesDescriptor(
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetOnesWorkspaceSize(infiniopOnesDescriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetOnesWorkspaceSize(infiniopOnesDescriptor_t desc, size_t *size) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
...@@ -98,7 +98,7 @@ __C infiniStatus_t infiniopGetOnesWorkspaceSize(infiniopOnesDescriptor_t desc, s ...@@ -98,7 +98,7 @@ __C infiniStatus_t infiniopGetOnesWorkspaceSize(infiniopOnesDescriptor_t desc, s
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
} }
__C infiniStatus_t infiniopOnes( __INFINI_C infiniStatus_t infiniopOnes(
infiniopOnesDescriptor_t desc, infiniopOnesDescriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
...@@ -142,7 +142,7 @@ __C infiniStatus_t infiniopOnes( ...@@ -142,7 +142,7 @@ __C infiniStatus_t infiniopOnes(
#undef CALCULATE #undef CALCULATE
} }
__C infiniStatus_t __INFINI_C infiniStatus_t
infiniopDestroyOnesDescriptor(infiniopOnesDescriptor_t desc) { infiniopDestroyOnesDescriptor(infiniopOnesDescriptor_t desc) {
#define DELETE(CASE, NAMESPACE) \ #define DELETE(CASE, NAMESPACE) \
......
...@@ -12,7 +12,7 @@ ...@@ -12,7 +12,7 @@
#include "metax/paged_attention_metax.h" #include "metax/paged_attention_metax.h"
#endif #endif
__C infiniStatus_t infiniopCreatePagedAttentionDescriptor( __INFINI_C infiniStatus_t infiniopCreatePagedAttentionDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopPagedAttentionDescriptor_t *desc_ptr, infiniopPagedAttentionDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t out_desc,
...@@ -54,7 +54,7 @@ __C infiniStatus_t infiniopCreatePagedAttentionDescriptor( ...@@ -54,7 +54,7 @@ __C infiniStatus_t infiniopCreatePagedAttentionDescriptor(
} }
} }
__C infiniStatus_t infiniopGetPagedAttentionWorkspaceSize( __INFINI_C infiniStatus_t infiniopGetPagedAttentionWorkspaceSize(
infiniopPagedAttentionDescriptor_t desc, infiniopPagedAttentionDescriptor_t desc,
size_t *size) { size_t *size) {
...@@ -84,7 +84,7 @@ __C infiniStatus_t infiniopGetPagedAttentionWorkspaceSize( ...@@ -84,7 +84,7 @@ __C infiniStatus_t infiniopGetPagedAttentionWorkspaceSize(
} }
} }
__C infiniStatus_t infiniopPagedAttention( __INFINI_C infiniStatus_t infiniopPagedAttention(
infiniopPagedAttentionDescriptor_t desc, infiniopPagedAttentionDescriptor_t desc,
void *workspace, size_t workspace_size, void *workspace, size_t workspace_size,
void *out, const void *q, const void *k_cache, const void *v_cache, void *out, const void *q, const void *k_cache, const void *v_cache,
...@@ -118,7 +118,7 @@ __C infiniStatus_t infiniopPagedAttention( ...@@ -118,7 +118,7 @@ __C infiniStatus_t infiniopPagedAttention(
} }
} }
__C infiniStatus_t infiniopDestroyPagedAttentionDescriptor( __INFINI_C infiniStatus_t infiniopDestroyPagedAttentionDescriptor(
infiniopPagedAttentionDescriptor_t desc) { infiniopPagedAttentionDescriptor_t desc) {
#define DESTROY(CASE, NAMESPACE) \ #define DESTROY(CASE, NAMESPACE) \
......
...@@ -12,7 +12,7 @@ ...@@ -12,7 +12,7 @@
#include "moore/paged_attention_prefill_moore.h" #include "moore/paged_attention_prefill_moore.h"
#endif #endif
__C infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor( __INFINI_C infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopPagedAttentionPrefillDescriptor_t *desc_ptr, infiniopPagedAttentionPrefillDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t out_desc,
...@@ -56,7 +56,7 @@ __C infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor( ...@@ -56,7 +56,7 @@ __C infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor(
} }
} }
__C infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize( __INFINI_C infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize(
infiniopPagedAttentionPrefillDescriptor_t desc, infiniopPagedAttentionPrefillDescriptor_t desc,
size_t *size) { size_t *size) {
...@@ -86,7 +86,7 @@ __C infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize( ...@@ -86,7 +86,7 @@ __C infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize(
} }
} }
__C infiniStatus_t infiniopPagedAttentionPrefill( __INFINI_C infiniStatus_t infiniopPagedAttentionPrefill(
infiniopPagedAttentionPrefillDescriptor_t desc, infiniopPagedAttentionPrefillDescriptor_t desc,
void *workspace, size_t workspace_size, void *workspace, size_t workspace_size,
void *out, const void *q, const void *k_cache, const void *v_cache, void *out, const void *q, const void *k_cache, const void *v_cache,
...@@ -123,7 +123,7 @@ __C infiniStatus_t infiniopPagedAttentionPrefill( ...@@ -123,7 +123,7 @@ __C infiniStatus_t infiniopPagedAttentionPrefill(
} }
} }
__C infiniStatus_t infiniopDestroyPagedAttentionPrefillDescriptor( __INFINI_C infiniStatus_t infiniopDestroyPagedAttentionPrefillDescriptor(
infiniopPagedAttentionPrefillDescriptor_t desc) { infiniopPagedAttentionPrefillDescriptor_t desc) {
#define DESTROY(CASE, NAMESPACE) \ #define DESTROY(CASE, NAMESPACE) \
......
...@@ -12,7 +12,7 @@ ...@@ -12,7 +12,7 @@
#include "moore/paged_caching_moore.h" #include "moore/paged_caching_moore.h"
#endif #endif
__C infiniStatus_t infiniopCreatePagedCachingDescriptor( __INFINI_C infiniStatus_t infiniopCreatePagedCachingDescriptor(
infiniopHandle_t handle, infiniopHandle_t handle,
infiniopPagedCachingDescriptor_t *desc_ptr, infiniopPagedCachingDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t k_cache_desc, infiniopTensorDescriptor_t k_cache_desc,
...@@ -49,7 +49,7 @@ __C infiniStatus_t infiniopCreatePagedCachingDescriptor( ...@@ -49,7 +49,7 @@ __C infiniStatus_t infiniopCreatePagedCachingDescriptor(
} }
} }
__C infiniStatus_t infiniopGetPagedCachingWorkspaceSize( __INFINI_C infiniStatus_t infiniopGetPagedCachingWorkspaceSize(
infiniopPagedCachingDescriptor_t desc, infiniopPagedCachingDescriptor_t desc,
size_t *size) { size_t *size) {
...@@ -79,7 +79,7 @@ __C infiniStatus_t infiniopGetPagedCachingWorkspaceSize( ...@@ -79,7 +79,7 @@ __C infiniStatus_t infiniopGetPagedCachingWorkspaceSize(
} }
} }
__C infiniStatus_t infiniopPagedCaching( __INFINI_C infiniStatus_t infiniopPagedCaching(
infiniopPagedCachingDescriptor_t desc, infiniopPagedCachingDescriptor_t desc,
void *workspace, size_t workspace_size, void *workspace, size_t workspace_size,
void *k_cache, void *v_cache, void *k_cache, void *v_cache,
...@@ -113,7 +113,7 @@ __C infiniStatus_t infiniopPagedCaching( ...@@ -113,7 +113,7 @@ __C infiniStatus_t infiniopPagedCaching(
} }
} }
__C infiniStatus_t infiniopDestroyPagedCachingDescriptor( __INFINI_C infiniStatus_t infiniopDestroyPagedCachingDescriptor(
infiniopPagedCachingDescriptor_t desc) { infiniopPagedCachingDescriptor_t desc) {
#define DESTROY(CASE, NAMESPACE) \ #define DESTROY(CASE, NAMESPACE) \
......
...@@ -9,12 +9,12 @@ ...@@ -9,12 +9,12 @@
#include "moore/per_channel_quant_int8_moore.h" #include "moore/per_channel_quant_int8_moore.h"
#endif #endif
__C infiniStatus_t infiniopCreatePerChannelQuantI8Descriptor(infiniopHandle_t handle, __INFINI_C infiniStatus_t infiniopCreatePerChannelQuantI8Descriptor(infiniopHandle_t handle,
infiniopPerChannelQuantI8Descriptor_t *desc_ptr, infiniopPerChannelQuantI8Descriptor_t *desc_ptr,
infiniopTensorDescriptor_t x_packed_desc, infiniopTensorDescriptor_t x_packed_desc,
infiniopTensorDescriptor_t x_scale_desc, infiniopTensorDescriptor_t x_scale_desc,
infiniopTensorDescriptor_t x_zero_desc, infiniopTensorDescriptor_t x_zero_desc,
infiniopTensorDescriptor_t x_desc) { infiniopTensorDescriptor_t x_desc) {
#define CREATE(CASE, NAMESPACE) \ #define CREATE(CASE, NAMESPACE) \
case CASE: \ case CASE: \
return op::per_channel_quant_int8::NAMESPACE::Descriptor::create( \ return op::per_channel_quant_int8::NAMESPACE::Descriptor::create( \
...@@ -40,7 +40,7 @@ __C infiniStatus_t infiniopCreatePerChannelQuantI8Descriptor(infiniopHandle_t ha ...@@ -40,7 +40,7 @@ __C infiniStatus_t infiniopCreatePerChannelQuantI8Descriptor(infiniopHandle_t ha
#undef CREATE #undef CREATE
} }
__C infiniStatus_t infiniopGetPerChannelQuantI8WorkspaceSize(infiniopPerChannelQuantI8Descriptor_t desc, size_t *size) { __INFINI_C infiniStatus_t infiniopGetPerChannelQuantI8WorkspaceSize(infiniopPerChannelQuantI8Descriptor_t desc, size_t *size) {
switch (desc->device_type) { switch (desc->device_type) {
#define GET(CASE, NAMESPACE) \ #define GET(CASE, NAMESPACE) \
case CASE: \ case CASE: \
...@@ -61,14 +61,14 @@ __C infiniStatus_t infiniopGetPerChannelQuantI8WorkspaceSize(infiniopPerChannelQ ...@@ -61,14 +61,14 @@ __C infiniStatus_t infiniopGetPerChannelQuantI8WorkspaceSize(infiniopPerChannelQ
#undef GET #undef GET
} }
__C infiniStatus_t infiniopPerChannelQuantI8(infiniopPerChannelQuantI8Descriptor_t desc, __INFINI_C infiniStatus_t infiniopPerChannelQuantI8(infiniopPerChannelQuantI8Descriptor_t desc,
void *workspace, void *workspace,
size_t workspace_size, size_t workspace_size,
void *x_packed, void *x_packed,
void *x_scale, void *x_scale,
void *x_zero, void *x_zero,
const void *x, const void *x,
void *stream) { void *stream) {
#define QUANT(CASE, NAMESPACE) \ #define QUANT(CASE, NAMESPACE) \
case CASE: \ case CASE: \
return reinterpret_cast<op::per_channel_quant_int8::NAMESPACE::Descriptor *>(desc)->calculate( \ return reinterpret_cast<op::per_channel_quant_int8::NAMESPACE::Descriptor *>(desc)->calculate( \
...@@ -90,7 +90,7 @@ __C infiniStatus_t infiniopPerChannelQuantI8(infiniopPerChannelQuantI8Descriptor ...@@ -90,7 +90,7 @@ __C infiniStatus_t infiniopPerChannelQuantI8(infiniopPerChannelQuantI8Descriptor
#undef QUANT #undef QUANT
} }
__C infiniStatus_t infiniopDestroyPerChannelQuantI8Descriptor(infiniopPerChannelQuantI8Descriptor_t desc) { __INFINI_C infiniStatus_t infiniopDestroyPerChannelQuantI8Descriptor(infiniopPerChannelQuantI8Descriptor_t desc) {
#define DESTROY(CASE, NAMESPACE) \ #define DESTROY(CASE, NAMESPACE) \
case CASE: \ case CASE: \
delete reinterpret_cast<op::per_channel_quant_int8::NAMESPACE::Descriptor *>(desc); \ delete reinterpret_cast<op::per_channel_quant_int8::NAMESPACE::Descriptor *>(desc); \
......
...@@ -16,10 +16,21 @@ static cudaError argMax_( ...@@ -16,10 +16,21 @@ static cudaError argMax_(
void *workspace_ptr, void *workspace_ptr,
size_t &workspace_len, size_t &workspace_len,
cudaStream_t stream) { cudaStream_t stream) {
#if CUDART_VERSION >= 11000 && !defined(ENABLE_QY_API) && !defined(ENABLE_HYGON_API)
// New interface: separate value and index outputs
T *max_value = &kv_pair->value;
int *max_index = &kv_pair->key;
return cub::DeviceReduce::ArgMax(
workspace_ptr, workspace_len,
logits, max_value, max_index, n,
stream);
#else
// Old interface
return cub::DeviceReduce::ArgMax( return cub::DeviceReduce::ArgMax(
workspace_ptr, workspace_len, workspace_ptr, workspace_len,
logits, kv_pair, n, logits, kv_pair, n,
stream); stream);
#endif
} }
template <class Tval, class Tidx> template <class Tval, class Tidx>
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment