From 4edbce98cf714c171f35b8064024a7ce12a74d40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E6=99=A8=20=28Leo=20Cheng=29?= Date: Tue, 14 Jul 2026 13:49:53 +0800 Subject: [PATCH] issue/1393 - add Metax backends for nsa_paged_attention and nsa_compress_paged_cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two operators only had cuda/ and nvidia/ implementations. Each gains metax/*_metax.h + metax/*_metax.maca and an INFINI_DEVICE_METAX branch in operator.cc; the device kernels reuse the shared cuda/kernel.cuh unchanged. That shared header needs two platform guards, keyed on the compiler's own __MACA__ / __MACACC__ rather than on ENABLE_METAX_API, which xmake.lua adds at project scope and would therefore also reach nvcc in a --nv-gpu=y --metax-gpu=y build: - and do not exist on MACA; metax_kernel_common.h provides the equivalents. - #pragma unroll on a runtime-bounded loop is rejected by MACA's clang under -Werror, so MACA joins the existing ILUVATAR/HYGON guard. Signed-off-by: 林晨 (Leo Cheng) --- .../nsa_compress_paged_cache/cuda/kernel.cuh | 5 + .../metax/nsa_compress_paged_cache_metax.h | 8 + .../metax/nsa_compress_paged_cache_metax.maca | 153 +++++++++++++ .../ops/nsa_compress_paged_cache/operator.cc | 16 ++ .../ops/nsa_paged_attention/cuda/kernel.cuh | 13 +- .../metax/nsa_paged_attention_metax.h | 8 + .../metax/nsa_paged_attention_metax.maca | 211 ++++++++++++++++++ .../ops/nsa_paged_attention/operator.cc | 16 ++ 8 files changed, 426 insertions(+), 4 deletions(-) create mode 100644 src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.h create mode 100644 src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.maca create mode 100644 src/infiniop/ops/nsa_paged_attention/metax/nsa_paged_attention_metax.h create mode 100644 src/infiniop/ops/nsa_paged_attention/metax/nsa_paged_attention_metax.maca diff --git a/src/infiniop/ops/nsa_compress_paged_cache/cuda/kernel.cuh b/src/infiniop/ops/nsa_compress_paged_cache/cuda/kernel.cuh index 0b1cacd5d..8dc2a7e71 100644 --- a/src/infiniop/ops/nsa_compress_paged_cache/cuda/kernel.cuh +++ b/src/infiniop/ops/nsa_compress_paged_cache/cuda/kernel.cuh @@ -3,8 +3,13 @@ #include #include +// On MetaX these headers do not exist; the equivalents come from +// devices/metax/metax_kernel_common.h, which the .maca translation unit +// includes before this one. +#if !defined(__MACA__) && !defined(__MACACC__) #include #include +#endif namespace op::nsa_compress_paged_cache::cuda { diff --git a/src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.h b/src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.h new file mode 100644 index 000000000..136fadae6 --- /dev/null +++ b/src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.h @@ -0,0 +1,8 @@ +#ifndef __NSA_COMPRESS_PAGED_CACHE_METAX_H__ +#define __NSA_COMPRESS_PAGED_CACHE_METAX_H__ + +#include "../nsa_compress_paged_cache.h" + +DESCRIPTOR(metax) + +#endif // __NSA_COMPRESS_PAGED_CACHE_METAX_H__ diff --git a/src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.maca b/src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.maca new file mode 100644 index 000000000..f56f55a45 --- /dev/null +++ b/src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.maca @@ -0,0 +1,153 @@ +#ifdef ENABLE_METAX_MC_API +#include +#else +#include +#endif + +#include + +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_kernel_common.h" +#include "nsa_compress_paged_cache_metax.h" + +#include "../cuda/kernel.cuh" + +namespace op::nsa_compress_paged_cache::metax { + +namespace { + +template +INFINIOP_METAX_KERNEL launchNsaCompressPagedCacheHd128( + Tdata *k_cmp, + Tdata *v_cmp, + const Tdata *k_cache, + const Tdata *v_cache, + const Tindex *block_tables, + const Tindex *cache_lens, + size_t max_num_blocks_per_seq, + size_t page_block_size, + size_t subblocks_per_page, + int nsa_block_size, + int update_last_only, + ptrdiff_t k_cmp_block_stride, + ptrdiff_t k_cmp_head_stride, + ptrdiff_t v_cmp_block_stride, + ptrdiff_t v_cmp_head_stride, + ptrdiff_t k_cache_batch_stride, + ptrdiff_t k_cache_head_stride, + ptrdiff_t k_cache_row_stride, + ptrdiff_t v_cache_batch_stride, + ptrdiff_t v_cache_head_stride, + ptrdiff_t v_cache_row_stride, + ptrdiff_t block_table_batch_stride, + ptrdiff_t cache_lens_stride) { + op::nsa_compress_paged_cache::cuda::compressPagedCacheKernel( + k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, max_num_blocks_per_seq, + page_block_size, subblocks_per_page, nsa_block_size, update_last_only, k_cmp_block_stride, k_cmp_head_stride, + v_cmp_block_stride, v_cmp_head_stride, k_cache_batch_stride, k_cache_head_stride, + k_cache_row_stride, v_cache_batch_stride, v_cache_head_stride, v_cache_row_stride, + block_table_batch_stride, cache_lens_stride); +} + +template +infiniStatus_t launchTyped( + void *k_cmp, + void *v_cmp, + const void *k_cache, + const void *v_cache, + const void *block_tables, + const void *cache_lens, + const NsaCompressPagedCacheInfo &info, + hcStream_t stream) { + const size_t max_nsa_blocks = info.update_last_only ? 1 : (info.max_num_blocks_per_seq * info.page_block_size + info.nsa_block_size - 1) / info.nsa_block_size; + dim3 grid(info.num_seqs, max_nsa_blocks, info.num_kv_heads); + dim3 block(128); + launchNsaCompressPagedCacheHd128<<>>( + static_cast(k_cmp), + static_cast(v_cmp), + static_cast(k_cache), + static_cast(v_cache), + static_cast(block_tables), + static_cast(cache_lens), + info.max_num_blocks_per_seq, + info.page_block_size, + info.subblocks_per_page, + info.nsa_block_size, + info.update_last_only, + info.k_cmp_block_stride, + info.k_cmp_head_stride, + info.v_cmp_block_stride, + info.v_cmp_head_stride, + info.k_cache_batch_stride, + info.k_cache_head_stride, + info.k_cache_row_stride, + info.v_cache_batch_stride, + info.v_cache_head_stride, + info.v_cache_row_stride, + info.block_table_batch_stride, + info.cache_lens_stride); + return INFINI_STATUS_SUCCESS; +} + +} // namespace + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t k_cmp_desc, + infiniopTensorDescriptor_t v_cmp_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t v_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t cache_lens_desc, + int nsa_block_size, + int update_last_only) { + auto info_res = NsaCompressPagedCacheInfo::create( + k_cmp_desc, v_cmp_desc, k_cache_desc, v_cache_desc, block_tables_desc, cache_lens_desc, nsa_block_size, update_last_only); + CHECK_RESULT(info_res); + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + info_res.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *k_cmp, + void *v_cmp, + const void *k_cache, + const void *v_cache, + const void *block_tables, + const void *cache_lens, + void *stream_) const { + (void)workspace; + (void)workspace_size; + auto stream = static_cast(stream_); + if (_info.dtype == INFINI_DTYPE_F16) { + if (_info.index_dtype == INFINI_DTYPE_I64) { + return launchTyped(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream); + } + if (_info.index_dtype == INFINI_DTYPE_I32) { + return launchTyped(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream); + } + return launchTyped(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream); + } + if (_info.index_dtype == INFINI_DTYPE_I64) { + return launchTyped(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream); + } + if (_info.index_dtype == INFINI_DTYPE_I32) { + return launchTyped(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream); + } + return launchTyped(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream); +} + +} // namespace op::nsa_compress_paged_cache::metax diff --git a/src/infiniop/ops/nsa_compress_paged_cache/operator.cc b/src/infiniop/ops/nsa_compress_paged_cache/operator.cc index a25311fa3..019ebd431 100644 --- a/src/infiniop/ops/nsa_compress_paged_cache/operator.cc +++ b/src/infiniop/ops/nsa_compress_paged_cache/operator.cc @@ -6,6 +6,10 @@ #include "nvidia/nsa_compress_paged_cache_nvidia.cuh" #endif +#ifdef ENABLE_METAX_API +#include "metax/nsa_compress_paged_cache_metax.h" +#endif + __INFINI_C infiniStatus_t infiniopCreateNsaCompressPagedCacheDescriptor( infiniopHandle_t handle, infiniopNsaCompressPagedCacheDescriptor_t *desc_ptr, @@ -38,6 +42,9 @@ __INFINI_C infiniStatus_t infiniopCreateNsaCompressPagedCacheDescriptor( #endif #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -64,6 +71,9 @@ __INFINI_C infiniStatus_t infiniopGetNsaCompressPagedCacheWorkspaceSize( #endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -98,6 +108,9 @@ __INFINI_C infiniStatus_t infiniopNsaCompressPagedCache( #endif #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -123,6 +136,9 @@ __INFINI_C infiniStatus_t infiniopDestroyNsaCompressPagedCacheDescriptor( #endif #ifdef ENABLE_HYGON_API DESTROY(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + DESTROY(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/nsa_paged_attention/cuda/kernel.cuh b/src/infiniop/ops/nsa_paged_attention/cuda/kernel.cuh index 9866e6f16..78891e814 100644 --- a/src/infiniop/ops/nsa_paged_attention/cuda/kernel.cuh +++ b/src/infiniop/ops/nsa_paged_attention/cuda/kernel.cuh @@ -4,8 +4,13 @@ #include #include #include +// On MetaX these headers do not exist; the equivalents come from +// devices/metax/metax_kernel_common.h, which the .maca translation unit +// includes before this one. +#if !defined(__MACA__) && !defined(__MACACC__) #include #include +#endif namespace op::nsa_paged_attention::cuda { @@ -44,7 +49,7 @@ __device__ inline float blockReduceSum128(float value) { const int lane = threadIdx.x & 31; const int warp = threadIdx.x >> 5; -#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) +#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) && !defined(__MACA__) && !defined(__MACACC__) #pragma unroll #endif for (int offset = 16; offset > 0; offset >>= 1) { @@ -59,7 +64,7 @@ __device__ inline float blockReduceSum128(float value) { float sum = threadIdx.x < 4 ? warp_sums[lane] : 0.0f; if (warp == 0) { -#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) +#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) && !defined(__MACA__) && !defined(__MACACC__) #pragma unroll #endif for (int offset = 2; offset > 0; offset >>= 1) { @@ -132,7 +137,7 @@ __device__ void nsaPagedDecodeHd128Kernel( const int active_select_blocks = max(1, min(select_blocks, kMaxSelectBlocks)); float top_scores[kMaxSelectBlocks]; int top_blocks[kMaxSelectBlocks]; -#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) +#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) && !defined(__MACA__) && !defined(__MACACC__) #pragma unroll #endif for (int i = 0; i < active_select_blocks; ++i) { @@ -177,7 +182,7 @@ __device__ void nsaPagedDecodeHd128Kernel( float sel_acc = 0.0f; float sel_m = -INFINITY; float sel_l = 0.0f; -#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) +#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_HYGON_API) && !defined(__MACA__) && !defined(__MACACC__) #pragma unroll #endif for (int selected = 0; selected < active_select_blocks; ++selected) { diff --git a/src/infiniop/ops/nsa_paged_attention/metax/nsa_paged_attention_metax.h b/src/infiniop/ops/nsa_paged_attention/metax/nsa_paged_attention_metax.h new file mode 100644 index 000000000..401f317e5 --- /dev/null +++ b/src/infiniop/ops/nsa_paged_attention/metax/nsa_paged_attention_metax.h @@ -0,0 +1,8 @@ +#ifndef __NSA_PAGED_ATTENTION_METAX_H__ +#define __NSA_PAGED_ATTENTION_METAX_H__ + +#include "../nsa_paged_attention.h" + +DESCRIPTOR(metax) + +#endif // __NSA_PAGED_ATTENTION_METAX_H__ diff --git a/src/infiniop/ops/nsa_paged_attention/metax/nsa_paged_attention_metax.maca b/src/infiniop/ops/nsa_paged_attention/metax/nsa_paged_attention_metax.maca new file mode 100644 index 000000000..c95e38a6c --- /dev/null +++ b/src/infiniop/ops/nsa_paged_attention/metax/nsa_paged_attention_metax.maca @@ -0,0 +1,211 @@ +#ifdef ENABLE_METAX_MC_API +#include +#else +#include +#endif + +#include +#include + +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_kernel_common.h" +#include "nsa_paged_attention_metax.h" + +#include "../cuda/kernel.cuh" + +namespace op::nsa_paged_attention::metax { + +namespace { + +template +INFINIOP_METAX_KERNEL launchNsaPagedDecodeHd128( + Tdata *out, + const Tdata *q, + const Tdata *k_cmp, + const Tdata *v_cmp, + const Tdata *k_cache, + const Tdata *v_cache, + const Tindex *block_tables, + const Tindex *cache_lens, + const Tgate *gates, + size_t num_heads, + size_t num_kv_heads, + float scale, + size_t max_num_blocks_per_seq, + size_t page_block_size, + size_t subblocks_per_page, + int nsa_block_size, + int window_size, + int select_blocks, + ptrdiff_t q_stride, + ptrdiff_t q_head_stride, + ptrdiff_t k_cmp_block_stride, + ptrdiff_t k_cmp_head_stride, + ptrdiff_t v_cmp_block_stride, + ptrdiff_t v_cmp_head_stride, + ptrdiff_t k_batch_stride, + ptrdiff_t k_head_stride, + ptrdiff_t k_row_stride, + ptrdiff_t v_batch_stride, + ptrdiff_t v_head_stride, + ptrdiff_t v_row_stride, + ptrdiff_t o_stride, + ptrdiff_t o_head_stride, + ptrdiff_t block_table_batch_stride, + ptrdiff_t cache_lens_stride, + ptrdiff_t gates_seq_stride, + ptrdiff_t gates_branch_stride, + ptrdiff_t gates_head_stride) { + op::nsa_paged_attention::cuda::nsaPagedDecodeHd128Kernel( + out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, num_heads, num_kv_heads, scale, + max_num_blocks_per_seq, page_block_size, subblocks_per_page, nsa_block_size, window_size, select_blocks, q_stride, q_head_stride, + k_cmp_block_stride, k_cmp_head_stride, v_cmp_block_stride, v_cmp_head_stride, + k_batch_stride, k_head_stride, k_row_stride, v_batch_stride, v_head_stride, v_row_stride, + o_stride, o_head_stride, block_table_batch_stride, cache_lens_stride, gates_seq_stride, + gates_branch_stride, gates_head_stride); +} + +template +infiniStatus_t launchTyped( + void *out, + const void *q, + const void *k_cmp, + const void *v_cmp, + const void *k_cache, + const void *v_cache, + const void *block_tables, + const void *cache_lens, + const void *gates, + const NsaPagedAttentionInfo &info, + hcStream_t stream) { + dim3 grid(info.num_seqs, info.num_heads); + dim3 block(128); + launchNsaPagedDecodeHd128<<>>( + static_cast(out), + static_cast(q), + static_cast(k_cmp), + static_cast(v_cmp), + static_cast(k_cache), + static_cast(v_cache), + static_cast(block_tables), + static_cast(cache_lens), + static_cast(gates), + info.num_heads, + info.num_kv_heads, + info.scale, + info.max_num_blocks_per_seq, + info.page_block_size, + info.subblocks_per_page, + info.nsa_block_size, + info.window_size, + info.select_blocks, + info.q_stride, + info.q_head_stride, + info.k_cmp_block_stride, + info.k_cmp_head_stride, + info.v_cmp_block_stride, + info.v_cmp_head_stride, + info.k_batch_stride, + info.k_head_stride, + info.k_row_stride, + info.v_batch_stride, + info.v_head_stride, + info.v_row_stride, + info.o_stride, + info.o_head_stride, + info.block_table_batch_stride, + info.cache_lens_stride, + info.gates_seq_stride, + info.gates_branch_stride, + info.gates_head_stride); + return INFINI_STATUS_SUCCESS; +} + +template +infiniStatus_t launchByGate( + void *out, const void *q, const void *k_cmp, const void *v_cmp, const void *k_cache, const void *v_cache, + const void *block_tables, const void *cache_lens, const void *gates, + const NsaPagedAttentionInfo &info, hcStream_t stream) { + if (info.gates_dtype == INFINI_DTYPE_F16) { + return launchTyped(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, info, stream); + } + if (info.gates_dtype == INFINI_DTYPE_BF16) { + return launchTyped(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, info, stream); + } + if (info.gates_dtype == INFINI_DTYPE_F32) { + return launchTyped(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, info, stream); + } + return INFINI_STATUS_BAD_TENSOR_DTYPE; +} + +} // namespace + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_cmp_desc, + infiniopTensorDescriptor_t v_cmp_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t v_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t cache_lens_desc, + infiniopTensorDescriptor_t gates_desc, + float scale, + int nsa_block_size, + int window_size, + int select_blocks) { + auto info_res = NsaPagedAttentionInfo::create(out_desc, q_desc, k_cmp_desc, v_cmp_desc, k_cache_desc, v_cache_desc, + block_tables_desc, cache_lens_desc, gates_desc, scale, nsa_block_size, window_size, select_blocks); + CHECK_RESULT(info_res); + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + info_res.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *out, + const void *q, + const void *k_cmp, + const void *v_cmp, + const void *k_cache, + const void *v_cache, + const void *block_tables, + const void *cache_lens, + const void *gates, + void *stream_) const { + (void)workspace; + (void)workspace_size; + auto stream = static_cast(stream_); + + if (_info.dtype == INFINI_DTYPE_F16) { + if (_info.index_dtype == INFINI_DTYPE_I64) { + return launchByGate(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, _info, stream); + } + if (_info.index_dtype == INFINI_DTYPE_I32) { + return launchByGate(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, _info, stream); + } + return launchByGate(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, _info, stream); + } + if (_info.index_dtype == INFINI_DTYPE_I64) { + return launchByGate(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, _info, stream); + } + if (_info.index_dtype == INFINI_DTYPE_I32) { + return launchByGate(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, _info, stream); + } + return launchByGate(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, _info, stream); +} + +} // namespace op::nsa_paged_attention::metax diff --git a/src/infiniop/ops/nsa_paged_attention/operator.cc b/src/infiniop/ops/nsa_paged_attention/operator.cc index 32bf183b2..c3e8b5b9d 100644 --- a/src/infiniop/ops/nsa_paged_attention/operator.cc +++ b/src/infiniop/ops/nsa_paged_attention/operator.cc @@ -6,6 +6,10 @@ #include "nvidia/nsa_paged_attention_nvidia.cuh" #endif +#ifdef ENABLE_METAX_API +#include "metax/nsa_paged_attention_metax.h" +#endif + __INFINI_C infiniStatus_t infiniopCreateNsaPagedAttentionDescriptor( infiniopHandle_t handle, infiniopNsaPagedAttentionDescriptor_t *desc_ptr, @@ -43,6 +47,9 @@ __INFINI_C infiniStatus_t infiniopCreateNsaPagedAttentionDescriptor( #endif #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -69,6 +76,9 @@ __INFINI_C infiniStatus_t infiniopGetNsaPagedAttentionWorkspaceSize( #endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -106,6 +116,9 @@ __INFINI_C infiniStatus_t infiniopNsaPagedAttention( #endif #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -131,6 +144,9 @@ __INFINI_C infiniStatus_t infiniopDestroyNsaPagedAttentionDescriptor( #endif #ifdef ENABLE_HYGON_API DESTROY(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + DESTROY(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;