refactor(ops): migrate paged attention to canonical InfiniOps

src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc

@@ -1,6 +1,15 @@
 #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API)
 #include "infinicore/ops/mha_varlen.hpp"
 
+#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_VARLEN_FUNC
+#include "../infiniops_impl.hpp"
+
+#include "base/flash_attn_varlen_func.h"
+
+#include <cstdint>
+#include <vector>
+#endif
+
 #ifdef ENABLE_ATEN
 #include "infinicore/adaptor/aten_adaptor.hpp"
 #include <ATen/ops/scaled_dot_product_attention.h>
@@ -16,13 +25,101 @@
 #include <stdexcept>
 
 namespace infinicore::op::mha_varlen_impl::flashattn {
+namespace {
+
+#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_VARLEN_FUNC
+using TensorMeta = ::infinicore::op::infiniops::TensorMeta;
+
+bool canUseInfiniOps(const Tensor &out,
+                     const Tensor &q,
+                     const Tensor &k,
+                     const Tensor &v,
+                     const Tensor &cum_seqlens_q,
+                     const Tensor &cum_seqlens_k,
+                     const std::optional<Tensor> &block_table,
+                     int max_seqlen_q,
+                     int max_seqlen_k,
+                     const std::optional<Tensor> &alibi_slopes) {
+    const bool paged = block_table.has_value();
+    const auto dtype = q->dtype();
+    if (out->device().getType() != Device::Type::NVIDIA
+        || q->ndim() != 3
+        || out->ndim() != 3
+        || ((paged && (k->ndim() != 4 || v->ndim() != 4))
+            || (!paged && (k->ndim() != 3 || v->ndim() != 3)))
+        || k->shape() != v->shape()
+        || out->shape() != q->shape()
+        || (dtype != DataType::F16 && dtype != DataType::BF16)
+        || out->dtype() != dtype
+        || k->dtype() != dtype
+        || v->dtype() != dtype
+        || q->size(1) == 0
+        || k->size(k->ndim() - 2) == 0
+        || q->size(1) % k->size(k->ndim() - 2) != 0
+        || q->size(2) == 0
+        || q->size(2) > 256
+        || q->size(2) % 8 != 0
+        || q->size(2) != k->size(k->ndim() - 1)
+        || q->stride(2) != 1
+        || out->stride(2) != 1
+        || k->stride(k->ndim() - 1) != 1
+        || v->stride(v->ndim() - 1) != 1
+        || cum_seqlens_q->ndim() != 1
+        || cum_seqlens_k->ndim() != 1
+        || cum_seqlens_q->shape() != cum_seqlens_k->shape()
+        || cum_seqlens_q->numel() < 2
+        || cum_seqlens_q->dtype() != DataType::I32
+        || cum_seqlens_k->dtype() != DataType::I32
+        || !cum_seqlens_q->is_contiguous()
+        || !cum_seqlens_k->is_contiguous()
+        || max_seqlen_q <= 0
+        || max_seqlen_k <= 0) {
+        return false;
+    }
+
+    if (block_table
+        && (block_table.value()->ndim() != 2
+            || block_table.value()->size(0) + 1 != cum_seqlens_q->size(0)
+            || block_table.value()->dtype() != DataType::I32
+            || !block_table.value()->is_contiguous()
+            || k->size(1) % 256 != 0)) {
+        return false;
+    }
+
+    if (alibi_slopes
+        && ((alibi_slopes.value()->ndim() != 1
+             && alibi_slopes.value()->ndim() != 2)
+            || alibi_slopes.value()->dtype() != DataType::F32
+            || !alibi_slopes.value()->is_contiguous()
+            || alibi_slopes.value()->device().getType() != out->device().getType()
+            || alibi_slopes.value()->device().getIndex() != out->device().getIndex()
+            || (alibi_slopes.value()->ndim() == 1
+                && alibi_slopes.value()->size(0) != q->size(1))
+            || (alibi_slopes.value()->ndim() == 2
+                && (alibi_slopes.value()->size(0) + 1
+                        != cum_seqlens_q->size(0)
+                    || alibi_slopes.value()->size(1) != q->size(1))))) {
+        return false;
+    }
+
+    return true;
+}
+#endif
+
+} // namespace
 
 struct PlannedMeta {
     graph::GraphTensor out, q, k, v, cum_seqlens_q, cum_seqlens_k;
     std::optional<graph::GraphTensor> block_table;
     int max_seqlen_q, max_seqlen_k;
     std::optional<graph::GraphTensor> alibi_slopes;
     float scale;
+#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_VARLEN_FUNC
+    bool use_infiniops{false};
+    std::optional<TensorMeta> infiniops_out, infiniops_q, infiniops_k,
+        infiniops_v, infiniops_cum_seqlens_q, infiniops_cum_seqlens_k;
+    std::optional<TensorMeta> infiniops_block_table, infiniops_alibi_slopes;
+#endif
 };
 
 void *plan(Tensor out,
@@ -37,7 +134,7 @@ void *plan(Tensor out,
            std::optional<Tensor> alibi_slopes,
            float scale) {
 
-    return new PlannedMeta{
+    auto planned = new PlannedMeta{
         graph::GraphTensor(out),
         graph::GraphTensor(q),
         graph::GraphTensor(k),
@@ -49,6 +146,28 @@ void *plan(Tensor out,
         max_seqlen_k,
         alibi_slopes ? std::optional<graph::GraphTensor>(graph::GraphTensor(*alibi_slopes)) : std::nullopt,
         scale};
+
+#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_VARLEN_FUNC
+    planned->use_infiniops = canUseInfiniOps(
+        out, q, k, v, cum_seqlens_q, cum_seqlens_k, block_table,
+        max_seqlen_q, max_seqlen_k, alibi_slopes);
+    if (planned->use_infiniops) {
+        planned->infiniops_out.emplace(out);
+        planned->infiniops_q.emplace(q);
+        planned->infiniops_k.emplace(k);
+        planned->infiniops_v.emplace(v);
+        planned->infiniops_cum_seqlens_q.emplace(cum_seqlens_q);
+        planned->infiniops_cum_seqlens_k.emplace(cum_seqlens_k);
+        if (block_table) {
+            planned->infiniops_block_table.emplace(*block_table);
+        }
+        if (alibi_slopes) {
+            planned->infiniops_alibi_slopes.emplace(*alibi_slopes);
+        }
+    }
+#endif
+
+    return planned;
 }
 
 namespace {
@@ -66,14 +185,58 @@ namespace {
 } // namespace
 
 void run(void *planned_meta) {
+    auto *p = reinterpret_cast<PlannedMeta *>(planned_meta);
+
+#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_VARLEN_FUNC
+    if (p->use_infiniops) {
+        infini::ops::Handle handle;
+        handle.set_stream(context::getStream());
+        infini::ops::Config config;
+        config.set_implementation_index(16);
+
+        const std::optional<infini::ops::Tensor> no_tensor;
+        const std::optional<infini::ops::Tensor> block_table = p->block_table
+                                                                 ? std::optional<infini::ops::Tensor>{
+                                                                     p->infiniops_block_table->tensor(*p->block_table)}
+                                                                 : std::nullopt;
+        const std::optional<infini::ops::Tensor> alibi_slopes = p->alibi_slopes
+                                                                  ? std::optional<infini::ops::Tensor>{
+                                                                      p->infiniops_alibi_slopes->tensor(*p->alibi_slopes)}
+                                                                  : std::nullopt;
+
+        infini::ops::FlashAttnVarlenFunc::Call(
+            handle,
+            config,
+            p->infiniops_q->tensor(p->q),
+            p->infiniops_k->tensor(p->k),
+            p->infiniops_v->tensor(p->v),
+            p->infiniops_cum_seqlens_q->tensor(p->cum_seqlens_q),
+            p->infiniops_cum_seqlens_k->tensor(p->cum_seqlens_k),
+            alibi_slopes,
+            block_table,
+            static_cast<std::int64_t>(p->max_seqlen_q),
+            static_cast<std::int64_t>(p->max_seqlen_k),
+            0.0,
+            std::optional<double>{p->scale},
+            true,
+            std::vector<std::int64_t>{-1, -1},
+            0.0,
+            false,
+            false,
+            p->infiniops_out->tensor(p->out),
+            no_tensor,
+            no_tensor);
+        return;
+    }
+#endif
+
 #if !defined(ENABLE_ATEN)
-    (void)planned_meta;
+    (void)p;
     throw std::runtime_error("ATen is not enabled in this build");
 #else
 #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API)
     c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream());
 #endif
-    auto *p = reinterpret_cast<PlannedMeta *>(planned_meta);
 
     auto q = infinicore::adaptor::to_aten_tensor(p->q);
     auto k = infinicore::adaptor::to_aten_tensor(p->k);

src/infinicore/ops/paged_attention/paged_attention_infiniops.cc

@@ -3,10 +3,13 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
+#include "base/flash_attn_with_kvcache.h"
 #include "base/paged_attention_infinilm.h"
 
 #include <cstddef>
+#include <cstdint>
 #include <optional>
+#include <vector>
 
 namespace infinicore::op::paged_attention_impl::infiniops {
 namespace {
@@ -22,11 +25,65 @@ std::size_t WorkspaceSizeInBytes(const Tensor &q) {
          * sizeof(float);
 }
 
+bool canUseFlashAttention(const Tensor &out,
+                          const Tensor &q,
+                          const Tensor &k_cache,
+                          const Tensor &v_cache,
+                          const Tensor &block_tables,
+                          const Tensor &cache_lens) {
+#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE
+    const auto dtype = q->dtype();
+    return out->device().getType() == Device::Type::NVIDIA
+        && q->ndim() == 3
+        && out->ndim() == 3
+        && k_cache->ndim() == 4
+        && v_cache->ndim() == 4
+        && block_tables->ndim() == 2
+        && cache_lens->ndim() == 1
+        && (dtype == DataType::F16 || dtype == DataType::BF16)
+        && out->dtype() == dtype
+        && k_cache->dtype() == dtype
+        && v_cache->dtype() == dtype
+        && k_cache->shape() == v_cache->shape()
+        && out->size(0) == q->size(0)
+        && out->size(1) == q->size(1)
+        && out->size(2) == v_cache->size(3)
+        && q->size(0) == block_tables->size(0)
+        && q->size(0) == cache_lens->size(0)
+        && k_cache->size(1) > 0
+        && q->size(1) % k_cache->size(1) == 0
+        && q->size(2) == k_cache->size(3)
+        && q->size(2) == v_cache->size(3)
+        && q->size(2) <= 256
+        && q->size(2) % 8 == 0
+        && k_cache->size(2) % 256 == 0
+        && q->stride(2) == 1
+        && out->stride(2) == 1
+        && k_cache->stride(3) == 1
+        && v_cache->stride(3) == 1
+        && block_tables->dtype() == DataType::I32
+        && cache_lens->dtype() == DataType::I32
+        && block_tables->is_contiguous()
+        && cache_lens->is_contiguous();
+#else
+    (void)out;
+    (void)q;
+    (void)k_cache;
+    (void)v_cache;
+    (void)block_tables;
+    (void)cache_lens;
+    return false;
+#endif
+}
+
 struct PlannedMeta {
     TensorMeta out, q, k_cache, v_cache, block_tables, cache_lens;
+    TensorMeta flash_out, flash_q, flash_k_cache, flash_v_cache;
     std::optional<TensorMeta> alibi_slopes;
-    graph::GraphTensor workspace, out_tensor, q_tensor, k_cache_tensor, v_cache_tensor, block_tables_tensor, cache_lens_tensor;
+    std::optional<graph::GraphTensor> workspace;
+    graph::GraphTensor out_tensor, q_tensor, k_cache_tensor, v_cache_tensor, block_tables_tensor, cache_lens_tensor;
     std::optional<graph::GraphTensor> alibi_slopes_tensor;
+    bool use_flash_attention;
     float scale;
 };
 } // namespace
@@ -44,22 +101,73 @@ void *plan(Tensor out,
     if (alibi_slopes) {
         INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *alibi_slopes);
     }
+
+    const bool use_flash_attention = canUseFlashAttention(out, q, k_cache, v_cache, block_tables, cache_lens);
+    auto flash_out = out->unsqueeze(1);
+    auto flash_q = q->unsqueeze(1);
+    auto flash_k_cache = k_cache->permute({0, 2, 1, 3});
+    auto flash_v_cache = v_cache->permute({0, 2, 1, 3});
+
     return new PlannedMeta{
         TensorMeta(out), TensorMeta(q), TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(block_tables), TensorMeta(cache_lens),
+        TensorMeta(flash_out), TensorMeta(flash_q), TensorMeta(flash_k_cache), TensorMeta(flash_v_cache),
         alibi_slopes ? std::optional<TensorMeta>{TensorMeta(*alibi_slopes)} : std::nullopt,
-        graph::GraphTensor(Tensor::empty({WorkspaceSizeInBytes(q)}, DataType::U8, out->device())),
+        use_flash_attention ? std::nullopt : std::optional<graph::GraphTensor>{graph::GraphTensor(Tensor::empty({WorkspaceSizeInBytes(q)}, DataType::U8, out->device()))},
         graph::GraphTensor(out), graph::GraphTensor(q), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(block_tables), graph::GraphTensor(cache_lens),
         alibi_slopes ? std::optional<graph::GraphTensor>{graph::GraphTensor(*alibi_slopes)} : std::nullopt,
+        use_flash_attention,
         scale};
 }
 
 void run(void *planned_meta) {
     auto planned = reinterpret_cast<PlannedMeta *>(planned_meta);
     infini::ops::Handle handle;
     handle.set_stream(context::getStream());
-    handle.set_workspace(planned->workspace->data());
-    handle.set_workspace_size_in_bytes(planned->workspace->numel());
+    if (planned->workspace) {
+        handle.set_workspace(planned->workspace.value()->data());
+        handle.set_workspace_size_in_bytes(planned->workspace.value()->numel());
+    }
     infini::ops::Config config;
+
+#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE
+    if (planned->use_flash_attention) {
+        config.set_implementation_index(16);
+        const std::optional<infini::ops::Tensor> no_tensor;
+        const std::optional<infini::ops::Tensor> cache_lens{
+            planned->cache_lens.tensor(planned->cache_lens_tensor)};
+        const std::optional<infini::ops::Tensor> block_tables{
+            planned->block_tables.tensor(planned->block_tables_tensor)};
+        const std::optional<infini::ops::Tensor> alibi_slopes = planned->alibi_slopes
+                                                                  ? std::optional<infini::ops::Tensor>{planned->alibi_slopes->tensor(planned->alibi_slopes_tensor.value()->data())}
+                                                                  : std::nullopt;
+        infini::ops::FlashAttnWithKvcache::Call(
+            handle,
+            config,
+            planned->flash_q.tensor(planned->q_tensor),
+            planned->flash_k_cache.tensor(planned->k_cache_tensor),
+            planned->flash_v_cache.tensor(planned->v_cache_tensor),
+            no_tensor,
+            no_tensor,
+            no_tensor,
+            no_tensor,
+            cache_lens,
+            no_tensor,
+            no_tensor,
+            block_tables,
+            alibi_slopes,
+            std::optional<double>{planned->scale},
+            true,
+            std::vector<std::int64_t>{-1, -1},
+            0.0,
+            true,
+            std::int64_t{0},
+            false,
+            planned->flash_out.tensor(planned->out_tensor),
+            no_tensor);
+        return;
+    }
+#endif
+
     infini::ops::PagedAttentionInfinilm::Call(
         handle,
         config,