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/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc b/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc index 1f4419129..c8e8b7697 100644 --- a/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc +++ b/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 +#include #include +#include 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 alibi_slopes; - graph::GraphTensor workspace, out_tensor, q_tensor, k_cache_tensor, v_cache_tensor, block_tables_tensor, cache_lens_tensor; + std::optional workspace; + graph::GraphTensor out_tensor, q_tensor, k_cache_tensor, v_cache_tensor, block_tables_tensor, cache_lens_tensor; std::optional alibi_slopes_tensor; + bool use_flash_attention; float scale; }; } // namespace @@ -44,12 +101,21 @@ 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(*alibi_slopes)} : std::nullopt, - graph::GraphTensor(Tensor::empty({WorkspaceSizeInBytes(q)}, DataType::U8, out->device())), + use_flash_attention ? std::nullopt : std::optional{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(*alibi_slopes)} : std::nullopt, + use_flash_attention, scale}; } @@ -57,9 +123,51 @@ void run(void *planned_meta) { auto planned = reinterpret_cast(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 no_tensor; + const std::optional cache_lens{ + planned->cache_lens.tensor(planned->cache_lens_tensor)}; + const std::optional block_tables{ + planned->block_tables.tensor(planned->block_tables_tensor)}; + const std::optional alibi_slopes = planned->alibi_slopes + ? std::optional{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{planned->scale}, + true, + std::vector{-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, 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/test/infinicore/ops/paged_attention.py b/test/infinicore/ops/paged_attention.py index 9809467ef..544178726 100644 --- a/test/infinicore/ops/paged_attention.py +++ b/test/infinicore/ops/paged_attention.py @@ -27,6 +27,7 @@ (3, 8, 8, 128, 16, 1024, False), (3, 8, 8, 64, 16, 1024, False), (8, 64, 8, 128, 16, 2048, False), + (2, 8, 2, 128, 256, 512, False), # Qwen3.6/Qwen3.5 full attention local TP shapes: head_dim=value_dim=256, GQA ratio=6. (1, 24, 4, 256, 16, 32, False), (1, 12, 2, 256, 16, 32, False), @@ -101,6 +102,9 @@ def parse_test_cases(): for dtype in _TENSOR_DTYPES: tolerance = _TOLERANCE_MAP.get(dtype, {"atol": 0, "rtol": 1e-3}) + # The canonical FlashAttention paged path requires int32 indices. + index_dtype = infinicore.int32 if block_size == 256 else infinicore.int64 + # Create typed tensor specs q_spec = TensorSpec.from_tensor(q_shape, None, dtype) k_cache_spec = TensorSpec.from_tensor(k_cache_shape, None, dtype) @@ -109,13 +113,13 @@ def parse_test_cases(): block_tables_shape, init_mode=TensorInitializer.MANUAL, set_tensor=block_tables, - dtype=infinicore.int64, + dtype=index_dtype, ) cache_lens_spec = TensorSpec.from_tensor( cache_lens_shape, init_mode=TensorInitializer.MANUAL, set_tensor=cache_lens_torch, - dtype=infinicore.int64, + dtype=index_dtype, ) # Paged attention operation: returns output tensor diff --git a/xmake.lua b/xmake.lua index eecfb4b3e..f41d7bc6d 100644 --- a/xmake.lua +++ b/xmake.lua @@ -356,26 +356,50 @@ end local infiniops_external_built = false -local function filter_infiniops_ops_for_backend(infiniops_ops) +local function configure_infiniops_ops(infiniops_ops) if not infiniops_ops or #infiniops_ops == 0 then - return infiniops_ops - end - if has_config("nv-gpu") then - return infiniops_ops + return infiniops_ops, false, false end local skipped_ops = { paged_attention_infinilm = true, paged_attention_prefill_infinilm = true } - local filtered = {} + local selected = {} + local selected_set = {} + 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 not skipped_ops[op] then - table.insert(filtered, op) + 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_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_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(filtered, ",") + + 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() @@ -419,7 +443,10 @@ 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 = filter_infiniops_ops_for_backend(os.getenv("INFINI_OPS_OPS")) + 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 table.insert(cmake_config_args, "-DINFINI_OPS_OPS=" .. infiniops_ops) end @@ -727,6 +754,13 @@ target("infinicore_cpp_api") end add_deps("infiniops_external") add_defines("ENABLE_INFINIOPS_API") + 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)