diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc index a75f4800b..4a5cb7dc9 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc +++ b/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 +#include +#endif + #ifdef ENABLE_ATEN #include "infinicore/adaptor/aten_adaptor.hpp" #include @@ -16,6 +25,88 @@ #include 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 &block_table, + int max_seqlen_q, + int max_seqlen_k, + const std::optional &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; @@ -23,6 +114,12 @@ struct PlannedMeta { int max_seqlen_q, max_seqlen_k; std::optional alibi_slopes; float scale; +#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_VARLEN_FUNC + bool use_infiniops{false}; + std::optional infiniops_out, infiniops_q, infiniops_k, + infiniops_v, infiniops_cum_seqlens_q, infiniops_cum_seqlens_k; + std::optional infiniops_block_table, infiniops_alibi_slopes; +#endif }; void *plan(Tensor out, @@ -37,7 +134,7 @@ void *plan(Tensor out, std::optional 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(*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(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 no_tensor; + const std::optional block_table = p->block_table + ? std::optional{ + p->infiniops_block_table->tensor(*p->block_table)} + : std::nullopt; + const std::optional alibi_slopes = p->alibi_slopes + ? std::optional{ + 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(p->max_seqlen_q), + static_cast(p->max_seqlen_k), + 0.0, + std::optional{p->scale}, + true, + std::vector{-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(planned_meta); auto q = infinicore::adaptor::to_aten_tensor(p->q); auto k = infinicore::adaptor::to_aten_tensor(p->k); diff --git a/submodules/InfiniOps b/submodules/InfiniOps index 21b07ebcf..1c865aea5 160000 --- a/submodules/InfiniOps +++ b/submodules/InfiniOps @@ -1 +1 @@ -Subproject commit 21b07ebcfdb0f993d2f3b672a4e38788e489fb80 +Subproject commit 1c865aea58a6af8bbfdc67b76a4ed0ea8d1c167d diff --git a/xmake.lua b/xmake.lua index 6f208fd6f..f41d7bc6d 100644 --- a/xmake.lua +++ b/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,19 +367,23 @@ 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) @@ -387,7 +391,15 @@ local function configure_infiniops_ops(infiniops_ops) 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)