issue/1393 - nsa paged attention and nsa compress paged cache for metax

src/infiniop/ops/nsa_compress_paged_cache/cuda/kernel.cuh

@@ -3,8 +3,13 @@
 
 #include <cstddef>
 #include <cstdint>
+// 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 <cuda_bf16.h>
 #include <cuda_fp16.h>
+#endif
 
 namespace op::nsa_compress_paged_cache::cuda {
 

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__

src/infiniop/ops/nsa_compress_paged_cache/metax/nsa_compress_paged_cache_metax.maca

@@ -0,0 +1,153 @@
+#ifdef ENABLE_METAX_MC_API
+#include <mc_runtime.h>
+#else
+#include <hc_runtime.h>
+#endif
+
+#include <cstdint>
+
+#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 <typename Tindex, typename Tdata>
+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<Tindex, Tdata>(
+        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 <typename Tindex, typename Tdata>
+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<Tindex, Tdata><<<grid, block, 0, stream>>>(
+        static_cast<Tdata *>(k_cmp),
+        static_cast<Tdata *>(v_cmp),
+        static_cast<const Tdata *>(k_cache),
+        static_cast<const Tdata *>(v_cache),
+        static_cast<const Tindex *>(block_tables),
+        static_cast<const Tindex *>(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<device::metax::Handle::Internal> 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<device::metax::Handle *>(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<hcStream_t>(stream_);
+    if (_info.dtype == INFINI_DTYPE_F16) {
+        if (_info.index_dtype == INFINI_DTYPE_I64) {
+            return launchTyped<int64_t, __half>(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream);
+        }
+        if (_info.index_dtype == INFINI_DTYPE_I32) {
+            return launchTyped<int32_t, __half>(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream);
+        }
+        return launchTyped<uint32_t, __half>(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream);
+    }
+    if (_info.index_dtype == INFINI_DTYPE_I64) {
+        return launchTyped<int64_t, __nv_bfloat16>(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream);
+    }
+    if (_info.index_dtype == INFINI_DTYPE_I32) {
+        return launchTyped<int32_t, __nv_bfloat16>(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream);
+    }
+    return launchTyped<uint32_t, __nv_bfloat16>(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, _info, stream);
+}
+
+} // namespace op::nsa_compress_paged_cache::metax

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;

src/infiniop/ops/nsa_paged_attention/cuda/kernel.cuh

@@ -4,8 +4,13 @@
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
+// 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 <cuda_bf16.h>
 #include <cuda_fp16.h>
+#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) {

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__