refactor(ops): migrate paged attention prefill 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);

xmake.lua

@@ -358,7 +358,7 @@ local infiniops_external_built = false
 
 local function configure_infiniops_ops(infiniops_ops)
     if not infiniops_ops or #infiniops_ops == 0 then
-        return infiniops_ops, false
+        return infiniops_ops, false, false
     end
 
     local skipped_ops = {
@@ -367,27 +367,39 @@ local function configure_infiniops_ops(infiniops_ops)
     }
     local selected = {}
     local selected_set = {}
-    local with_linked_flash_attention = false
+    local with_linked_flash_attn_with_kvcache = false
+    local with_linked_flash_attn_varlen_func = false
     for _, op in ipairs(infiniops_ops:split("[,;]")) do
         op = op:trim()
         if #op > 0 and (has_config("nv-gpu") or not skipped_ops[op]) then
             table.insert(selected, op)
             selected_set[op] = true
             if has_config("nv-gpu") and (op == "paged_attention_infinilm" or op == "flash_attn_with_kvcache") then
-                with_linked_flash_attention = true
+                with_linked_flash_attn_with_kvcache = true
+            end
+            if has_config("nv-gpu") and (op == "paged_attention_prefill_infinilm" or op == "flash_attn_varlen_func") then
+                with_linked_flash_attn_varlen_func = true
             end
         end
     end
 
-    if with_linked_flash_attention then
+    if with_linked_flash_attn_with_kvcache then
         for _, op in ipairs({"paged_attention_infinilm", "flash_attn_with_kvcache"}) do
             if not selected_set[op] then
                 table.insert(selected, op)
             end
         end
     end
 
-    return table.concat(selected, ","), with_linked_flash_attention
+    if with_linked_flash_attn_varlen_func then
+        for _, op in ipairs({"paged_attention_prefill_infinilm", "flash_attn_varlen_func"}) do
+            if not selected_set[op] then
+                table.insert(selected, op)
+            end
+        end
+    end
+
+    return table.concat(selected, ","), with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func
 end
 
 local function get_infiniops_backend_cmake_arg()
@@ -431,8 +443,8 @@ local function build_infiniops_external(xmake_os)
         table.insert(cmake_config_args, "-DTORCH_CXX11_ABI=0")
         table.insert(cmake_config_args, "-DCMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=0")
     end
-    local infiniops_ops, with_linked_flash_attention = configure_infiniops_ops(os.getenv("INFINI_OPS_OPS"))
-    if with_linked_flash_attention then
+    local infiniops_ops, with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func = configure_infiniops_ops(os.getenv("INFINI_OPS_OPS"))
+    if with_linked_flash_attn_with_kvcache or with_linked_flash_attn_varlen_func then
         table.insert(cmake_config_args, "-DWITH_LINKED=ON")
     end
     if infiniops_ops and #infiniops_ops > 0 then
@@ -742,10 +754,13 @@ target("infinicore_cpp_api")
         end
         add_deps("infiniops_external")
         add_defines("ENABLE_INFINIOPS_API")
-        local _, with_linked_flash_attention = configure_infiniops_ops(os.getenv("INFINI_OPS_OPS"))
-        if with_linked_flash_attention then
+        local _, with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func = configure_infiniops_ops(os.getenv("INFINI_OPS_OPS"))
+        if with_linked_flash_attn_with_kvcache then
             add_defines("ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE")
         end
+        if with_linked_flash_attn_varlen_func then
+            add_defines("ENABLE_INFINIOPS_LINKED_FLASH_ATTN_VARLEN_FUNC")
+        end
         add_links("infiniops")
         add_rpathdirs(INFINI_ROOT .. "/lib")
         on_load(function (target)