From 4f3fca728899f5d0e99d3d6f4b934a95f6b8e0cf Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Sun, 9 Aug 2026 12:36:00 +0800 Subject: [PATCH 01/10] refactor(ops): migrate paged attention decode to canonical InfiniOps --- .../paged_attention_infiniops.cc | 116 +++++++++++++++++- test/infinicore/ops/paged_attention.py | 8 +- xmake.lua | 39 ++++-- 3 files changed, 147 insertions(+), 16 deletions(-) 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/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..6f208fd6f 100644 --- a/xmake.lua +++ b/xmake.lua @@ -356,26 +356,38 @@ 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 end local skipped_ops = { paged_attention_infinilm = true, paged_attention_prefill_infinilm = true } - local filtered = {} + local selected = {} + local selected_set = {} + local with_linked_flash_attention = 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_attention = true + end + end + end + + if with_linked_flash_attention 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, ",") + + return table.concat(selected, ","), with_linked_flash_attention end local function get_infiniops_backend_cmake_arg() @@ -419,7 +431,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_attention = configure_infiniops_ops(os.getenv("INFINI_OPS_OPS")) + if with_linked_flash_attention 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 +742,10 @@ 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 + add_defines("ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE") + end add_links("infiniops") add_rpathdirs(INFINI_ROOT .. "/lib") on_load(function (target) From bf24faa6b54c0d9410310465054693070492217c Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Sun, 9 Aug 2026 23:15:21 +0800 Subject: [PATCH 02/10] refactor(ops): migrate paged attention prefill to canonical InfiniOps --- .../mha_varlen_flashattn.cc | 169 +++++++++++++++++- submodules/InfiniOps | 2 +- xmake.lua | 33 +++- 3 files changed, 191 insertions(+), 13 deletions(-) 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) From 98432df0f0923742b260a1a8f06dc88224f68040 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 10 Aug 2026 11:39:32 +0800 Subject: [PATCH 03/10] refactor(ops): remove InfiniLM-suffixed InfiniOps calls --- .../ops/kv_caching/kv_caching_infiniops.cc | 50 ------------- .../paged_attention_infiniops.cc | 66 ++++++++--------- .../paged_attention_prefill_infiniops.cc | 68 ----------------- .../random_sample/random_sample_infiniops.cc | 42 ----------- src/infinicore/ops/rope/rope_infiniops.cc | 73 ------------------- xmake.lua | 29 +------- 6 files changed, 34 insertions(+), 294 deletions(-) delete mode 100644 src/infinicore/ops/kv_caching/kv_caching_infiniops.cc delete mode 100644 src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc delete mode 100644 src/infinicore/ops/random_sample/random_sample_infiniops.cc delete mode 100644 src/infinicore/ops/rope/rope_infiniops.cc diff --git a/src/infinicore/ops/kv_caching/kv_caching_infiniops.cc b/src/infinicore/ops/kv_caching/kv_caching_infiniops.cc deleted file mode 100644 index 5a2774f0d..000000000 --- a/src/infinicore/ops/kv_caching/kv_caching_infiniops.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/kv_caching.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/kv_caching_infinilm.h" - -namespace infinicore::op::kv_caching_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -struct PlannedMeta { - TensorMeta k_cache, v_cache, k, v, past_kv_lengths; - graph::GraphTensor k_cache_tensor, v_cache_tensor, k_tensor, v_tensor, past_kv_lengths_tensor; -}; -} // namespace - -void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &past_kv_lengths) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, past_kv_lengths); - return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(past_kv_lengths), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(past_kv_lengths)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - infini::ops::KvCachingInfinilm::Call( - handle, - config, - planned->k.tensor(planned->k_tensor), - planned->v.tensor(planned->v_tensor), - planned->past_kv_lengths.tensor(planned->past_kv_lengths_tensor), - planned->k_cache.tensor(planned->k_cache_tensor), - planned->v_cache.tensor(planned->v_cache_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::cleanup_dispatcher(), &cleanup); - return true; -}(); -} // namespace infinicore::op::kv_caching_impl::infiniops -#endif diff --git a/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc b/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc index c8e8b7697..8e37a3ac7 100644 --- a/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc +++ b/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc @@ -4,27 +4,28 @@ #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::infiniop { +void *plan(Tensor out, + const Tensor &q, + const Tensor &k_cache, + const Tensor &v_cache, + const Tensor &block_tables, + const Tensor &cache_lens, + std::optional alibi_slopes, + float scale); +void run(void *planned_meta); +void cleanup(void **planned_meta_ptr); +} // namespace infinicore::op::paged_attention_impl::infiniop + namespace infinicore::op::paged_attention_impl::infiniops { namespace { using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -constexpr std::size_t kMaxPagedAttentionSplits = 8; - -std::size_t WorkspaceSizeInBytes(const Tensor &q) { - return kMaxPagedAttentionSplits - * static_cast(q->size(0)) - * static_cast(q->size(1)) - * static_cast(q->size(2) + 2) - * sizeof(float); -} - bool canUseFlashAttention(const Tensor &out, const Tensor &q, const Tensor &k_cache, @@ -77,12 +78,12 @@ bool canUseFlashAttention(const Tensor &out, } struct PlannedMeta { - TensorMeta out, q, k_cache, v_cache, block_tables, cache_lens; TensorMeta flash_out, flash_q, flash_k_cache, flash_v_cache; + TensorMeta block_tables, cache_lens; std::optional alibi_slopes; - 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; + void *fallback_meta; bool use_flash_attention; float scale; }; @@ -107,30 +108,30 @@ void *plan(Tensor out, 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}); + void *fallback_meta = use_flash_attention + ? nullptr + : paged_attention_impl::infiniop::plan( + out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes, scale); 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), + TensorMeta(block_tables), TensorMeta(cache_lens), alibi_slopes ? std::optional{TensorMeta(*alibi_slopes)} : std::nullopt, - 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, + fallback_meta, use_flash_attention, scale}; } void run(void *planned_meta) { auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - 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) { + 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 cache_lens{ @@ -168,21 +169,16 @@ void run(void *planned_meta) { } #endif - infini::ops::PagedAttentionInfinilm::Call( - handle, - config, - planned->q.tensor(planned->q_tensor), - planned->k_cache.tensor(planned->k_cache_tensor), - planned->v_cache.tensor(planned->v_cache_tensor), - planned->block_tables.tensor(planned->block_tables_tensor), - planned->cache_lens.tensor(planned->cache_lens_tensor), - planned->alibi_slopes ? std::optional{planned->alibi_slopes->tensor(planned->alibi_slopes_tensor.value()->data())} : std::nullopt, - planned->scale, - planned->out.tensor(planned->out_tensor)); + INFINICORE_ASSERT(planned->fallback_meta != nullptr); + paged_attention_impl::infiniop::run(planned->fallback_meta); } void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); + auto planned = *reinterpret_cast(planned_meta_ptr); + if (planned->fallback_meta != nullptr) { + paged_attention_impl::infiniop::cleanup(&planned->fallback_meta); + } + delete planned; *planned_meta_ptr = nullptr; } diff --git a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc b/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc deleted file mode 100644 index bb842a6c6..000000000 --- a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include "infinicore/ops/paged_attention_prefill.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/paged_attention_prefill_infinilm.h" - -#include - -namespace infinicore::op::paged_attention_prefill_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor out, - Tensor q, - Tensor k_cache, - Tensor v_cache, - Tensor block_tables, - Tensor kv_lens, - Tensor cum_seqlens_q, - std::optional alibi_slopes, - float scale) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q); - if (alibi_slopes) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *alibi_slopes); - } - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta out_meta(out); - TensorMeta q_meta(q); - TensorMeta k_cache_meta(k_cache); - TensorMeta v_cache_meta(v_cache); - TensorMeta block_tables_meta(block_tables); - TensorMeta kv_lens_meta(kv_lens); - TensorMeta cum_seqlens_q_meta(cum_seqlens_q); - std::optional alibi_slopes_meta; - if (alibi_slopes) { - alibi_slopes_meta.emplace(*alibi_slopes); - } - - infini::ops::PagedAttentionPrefillInfinilm::Call( - handle, - config, - q_meta.tensor(q), - k_cache_meta.tensor(k_cache), - v_cache_meta.tensor(v_cache), - block_tables_meta.tensor(block_tables), - kv_lens_meta.tensor(kv_lens), - cum_seqlens_q_meta.tensor(cum_seqlens_q), - alibi_slopes_meta ? std::optional{alibi_slopes_meta->tensor(*alibi_slopes)} : std::nullopt, - scale, - out_meta.tensor(out)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(PagedAttentionPrefill::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::paged_attention_prefill_impl::infiniops -#endif diff --git a/src/infinicore/ops/random_sample/random_sample_infiniops.cc b/src/infinicore/ops/random_sample/random_sample_infiniops.cc deleted file mode 100644 index 5146b1443..000000000 --- a/src/infinicore/ops/random_sample/random_sample_infiniops.cc +++ /dev/null @@ -1,42 +0,0 @@ -#include "infinicore/ops/random_sample.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/random_sample_infinilm.h" - -namespace infinicore::op::random_sample_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor indices, Tensor logits, float random_val, float topp, int topk, float temperature) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(indices->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(indices, logits); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta indices_meta(indices); - TensorMeta logits_meta(logits); - infini::ops::RandomSampleInfinilm::Call( - handle, - config, - logits_meta.tensor(logits), - random_val, - topp, - static_cast(topk), - temperature, - indices_meta.tensor(indices)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(RandomSample::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::random_sample_impl::infiniops -#endif diff --git a/src/infinicore/ops/rope/rope_infiniops.cc b/src/infinicore/ops/rope/rope_infiniops.cc deleted file mode 100644 index 045a5aca3..000000000 --- a/src/infinicore/ops/rope/rope_infiniops.cc +++ /dev/null @@ -1,73 +0,0 @@ -#include "infinicore/ops/rope.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/rotary_embedding_infinilm.h" - -#include - -namespace infinicore::op::rope_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -struct PlannedMeta { - TensorMeta x_out, x, pos, sin, cos; - graph::GraphTensor x_out_tensor, x_tensor, pos_tensor, sin_tensor, cos_tensor; - bool is_neox; -}; - -bool toInfiniOpsIsNeox(infinicore::nn::RoPE::Algo algo) { - switch (algo) { - case infinicore::nn::RoPE::Algo::GPT_J: - return true; - case infinicore::nn::RoPE::Algo::GPT_NEOX: - return false; - default: - throw std::runtime_error("Unsupported RoPE algorithm"); - } -} -} // namespace - -void *plan(Tensor x_out, - const Tensor &x, - const Tensor &pos, - const Tensor &sin, - const Tensor &cos, - infinicore::nn::RoPE::Algo algo) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(x_out->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x_out, x, pos, sin, cos); - return new PlannedMeta{ - TensorMeta(x_out), TensorMeta(x), TensorMeta(pos), TensorMeta(sin), TensorMeta(cos), - graph::GraphTensor(x_out), graph::GraphTensor(x), graph::GraphTensor(pos), graph::GraphTensor(sin), graph::GraphTensor(cos), - toInfiniOpsIsNeox(algo)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - infini::ops::RotaryEmbeddingInfinilm::Call( - handle, - config, - planned->x.tensor(planned->x_tensor), - planned->pos.tensor(planned->pos_tensor), - planned->sin.tensor(planned->sin_tensor), - planned->cos.tensor(planned->cos_tensor), - planned->is_neox, - planned->x_out.tensor(planned->x_out_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(RoPE::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(RoPE::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(RoPE::cleanup_dispatcher(), &cleanup); - return true; -}(); -} // namespace infinicore::op::rope_impl::infiniops -#endif diff --git a/xmake.lua b/xmake.lua index f41d7bc6d..de559e0ec 100644 --- a/xmake.lua +++ b/xmake.lua @@ -361,44 +361,22 @@ local function configure_infiniops_ops(infiniops_ops) return infiniops_ops, false, false end - local skipped_ops = { - paged_attention_infinilm = true, - paged_attention_prefill_infinilm = true - } 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 (has_config("nv-gpu") or not skipped_ops[op]) then + if #op > 0 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 + if has_config("nv-gpu") and 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 + if has_config("nv-gpu") and 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 - - 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 @@ -1021,7 +999,6 @@ target("infinicore_cpp_api") end if has_config("infiniops") and not has_config("nv-gpu") then remove_files("src/infinicore/ops/paged_attention/paged_attention_infiniops.cc") - remove_files("src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc") end if has_config("mutual-awareness") then add_files("src/infinicore/analyzer/*.cc") From 01727462426cd0a1230d0d140afd89931ef21c6b Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 10 Aug 2026 23:48:24 +0800 Subject: [PATCH 04/10] refactor(ops): use canonical KV-cache attention --- .../ops/mha_kvcache/mha_kvcache_flashattn.cc | 151 +++++++++++++++++- 1 file changed, 149 insertions(+), 2 deletions(-) diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc index 50cded8ab..47fbbed4a 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc @@ -1,6 +1,15 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) #include "infinicore/ops/mha_kvcache.hpp" +#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE +#include "../infiniops_impl.hpp" + +#include "base/flash_attn_with_kvcache.h" + +#include +#include +#endif + #include "infinicore/adaptor/flash_attention_adaptor.hpp" #include @@ -18,11 +27,86 @@ #endif namespace infinicore::op::mha_kvcache_impl::flashattn { +namespace { + +#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +bool canUseInfiniOps(const Tensor &out, + const Tensor &q, + const Tensor &k_cache, + const Tensor &v_cache, + const Tensor &seqlens_k, + const Tensor &block_table, + const std::optional &alibi_slopes) { + const auto dtype = q->dtype(); + if (out->device().getType() != Device::Type::NVIDIA + || q->ndim() != 4 + || out->ndim() != 4 + || k_cache->ndim() != 4 + || v_cache->ndim() != 4 + || q->size(1) != 1 + || k_cache->shape() != v_cache->shape() + || out->shape() != q->shape() + || (dtype != DataType::F16 && dtype != DataType::BF16) + || out->dtype() != dtype + || k_cache->dtype() != dtype + || v_cache->dtype() != dtype + || q->size(0) == 0 + || q->size(2) == 0 + || k_cache->size(1) == 0 + || k_cache->size(2) == 0 + || q->size(2) % k_cache->size(2) != 0 + || q->size(3) == 0 + || q->size(3) > 256 + || q->size(3) % 8 != 0 + || q->size(3) != k_cache->size(3) + || q->stride(3) != 1 + || out->stride(3) != 1 + || k_cache->stride(3) != 1 + || v_cache->stride(3) != 1 + || seqlens_k->ndim() != 1 + || seqlens_k->size(0) != q->size(0) + || seqlens_k->dtype() != DataType::I32 + || !seqlens_k->is_contiguous() + || block_table->ndim() != 2 + || block_table->size(0) != q->size(0) + || block_table->dtype() != DataType::I32 + || !block_table->is_contiguous() + || k_cache->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() != out->device() + || (alibi_slopes.value()->ndim() == 1 + && alibi_slopes.value()->size(0) != q->size(2)) + || (alibi_slopes.value()->ndim() == 2 + && (alibi_slopes.value()->size(0) != q->size(0) + || alibi_slopes.value()->size(1) != q->size(2))))) { + return false; + } + + return true; +} +#endif + +} // namespace struct PlannedMeta { graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; std::optional alibi_slopes; float scale; +#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE + bool use_infiniops{false}; + std::optional infiniops_out, infiniops_q, infiniops_k_cache, + infiniops_v_cache, infiniops_seqlens_k, infiniops_block_table; + std::optional infiniops_alibi_slopes; +#endif }; void *plan(Tensor out, @@ -33,7 +117,7 @@ void *plan(Tensor out, const Tensor &block_table, std::optional alibi_slopes, float scale) { - return new PlannedMeta{ + auto *planned = new PlannedMeta{ graph::GraphTensor(out), graph::GraphTensor(q), graph::GraphTensor(k_cache), @@ -42,14 +126,77 @@ void *plan(Tensor out, graph::GraphTensor(block_table), alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, scale}; + +#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE + planned->use_infiniops = canUseInfiniOps( + out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes); + if (planned->use_infiniops) { + planned->infiniops_out.emplace(out); + planned->infiniops_q.emplace(q); + planned->infiniops_k_cache.emplace(k_cache); + planned->infiniops_v_cache.emplace(v_cache); + planned->infiniops_seqlens_k.emplace(seqlens_k); + planned->infiniops_block_table.emplace(block_table); + if (alibi_slopes) { + planned->infiniops_alibi_slopes.emplace(*alibi_slopes); + } + } +#endif + + return planned; } void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + +#ifdef ENABLE_INFINIOPS_LINKED_FLASH_ATTN_WITH_KVCACHE + 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 cache_seqlens{ + p->infiniops_seqlens_k->tensor(p->seqlens_k)}; + const std::optional block_table{ + p->infiniops_block_table->tensor(p->block_table)}; + const std::optional alibi_slopes = p->alibi_slopes + ? std::optional{p->infiniops_alibi_slopes->tensor(*p->alibi_slopes)} + : std::nullopt; + + infini::ops::FlashAttnWithKvcache::Call( + handle, + config, + p->infiniops_q->tensor(p->q), + p->infiniops_k_cache->tensor(p->k_cache), + p->infiniops_v_cache->tensor(p->v_cache), + no_tensor, + no_tensor, + no_tensor, + no_tensor, + cache_seqlens, + no_tensor, + no_tensor, + block_table, + alibi_slopes, + std::optional{p->scale}, + true, + std::vector{-1, -1}, + 0.0, + true, + std::int64_t{0}, + false, + p->infiniops_out->tensor(p->out), + no_tensor); + return; + } +#endif + #if defined(ENABLE_FLASH_ATTN) #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); // Paged KV caches must be contiguous for flash-attn; avoid extra copies for q/metadata when already dense. const bool out_need_copy_back = !p->out->is_contiguous(); From a8006d1bf628fbacf66671b1cb00a04b5a1f522a Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 10 Aug 2026 23:49:16 +0800 Subject: [PATCH 05/10] refactor(nn): use canonical RotaryEmbedding for RoPE --- include/infinicore/nn/rope.hpp | 3 +- include/infinicore/ops.hpp | 1 + include/infinicore/ops/rotary_embedding.hpp | 23 +++++ src/infinicore/nn/rope.cc | 55 ++++++++++++ .../ops/rotary_embedding/rotary_embedding.cc | 51 +++++++++++ .../rotary_embedding_infiniops.cc | 84 +++++++++++++++++++ 6 files changed, 216 insertions(+), 1 deletion(-) create mode 100644 include/infinicore/ops/rotary_embedding.hpp create mode 100644 src/infinicore/ops/rotary_embedding/rotary_embedding.cc create mode 100644 src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc diff --git a/include/infinicore/nn/rope.hpp b/include/infinicore/nn/rope.hpp index eaeba8712..37f350678 100644 --- a/include/infinicore/nn/rope.hpp +++ b/include/infinicore/nn/rope.hpp @@ -94,9 +94,10 @@ class RoPE : public Module { std::string extra_repr() const; protected: - // Buffers (sin and cos cache tables) - not exposed in state_dict + // RoPE cache tables are not exposed in state_dict. INFINICORE_NN_BUFFER(sin_cache); INFINICORE_NN_BUFFER(cos_cache); + INFINICORE_NN_BUFFER(cos_sin_cache); private: void initialize_cache(); diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index b960e4a49..d9bf1ad11 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -73,6 +73,7 @@ #include "ops/relu.hpp" #include "ops/rms_norm.hpp" #include "ops/rope.hpp" +#include "ops/rotary_embedding.hpp" #include "ops/rot.hpp" #include "ops/rotg.hpp" #include "ops/rotm.hpp" diff --git a/include/infinicore/ops/rotary_embedding.hpp b/include/infinicore/ops/rotary_embedding.hpp new file mode 100644 index 000000000..6c9841f1e --- /dev/null +++ b/include/infinicore/ops/rotary_embedding.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(RotaryEmbedding, const Tensor &, Tensor, std::optional, const Tensor &, int64_t, bool, int64_t, bool); + +void rotary_embedding_(const Tensor &positions, + Tensor query, + std::optional key, + const Tensor &cos_sin_cache, + int64_t head_size, + bool is_neox, + int64_t rope_dim_offset = 0, + bool inverse = false); + +} // namespace infinicore::op diff --git a/src/infinicore/nn/rope.cc b/src/infinicore/nn/rope.cc index 115e162e2..5e30c7e80 100644 --- a/src/infinicore/nn/rope.cc +++ b/src/infinicore/nn/rope.cc @@ -3,6 +3,7 @@ #include "../utils.hpp" #include "infinicore/ops/mrope.hpp" #include "infinicore/ops/rope.hpp" +#include "infinicore/ops/rotary_embedding.hpp" #include #include #include @@ -12,6 +13,24 @@ #include namespace infinicore::nn { +namespace { + +template +std::vector combine_cos_sin_cache(const std::vector &cos, + const std::vector &sin, + size_t max_seq_len, + size_t cache_dim) { + std::vector combined(max_seq_len * cache_dim * 2); + for (size_t pos = 0; pos < max_seq_len; ++pos) { + std::copy_n(cos.begin() + pos * cache_dim, cache_dim, + combined.begin() + pos * cache_dim * 2); + std::copy_n(sin.begin() + pos * cache_dim, cache_dim, + combined.begin() + (pos * 2 + 1) * cache_dim); + } + return combined; +} + +} // namespace RoPE::RoPE(size_t head_dim, size_t rotary_dim, @@ -58,6 +77,12 @@ void RoPE::initialize_cache() { INFINICORE_NN_BUFFER_INIT(sin_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); INFINICORE_NN_BUFFER_INIT(cos_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); +#ifdef ENABLE_INFINIOPS_API + if (device_.getType() == Device::Type::NVIDIA && !mrope_section_) { + INFINICORE_NN_BUFFER_INIT(cos_sin_cache, ({max_seq_len_, rotary_dim_}, dtype_, device_)); + } +#endif + // Pre-compute sin and cos values // Frequency generation always uses GPT-J style (theta^(-2j/rotary_dim)). // The rotation algorithm (algo_) controls how dimensions are paired in the kernel. @@ -95,6 +120,11 @@ void RoPE::initialize_cache() { auto cos_f32_cpu = Tensor::from_blob(cos_data.data(), {max_seq_len_, cache_dim}, DataType::F32, cpu_device); sin_cache_->copy_from(sin_f32_cpu); cos_cache_->copy_from(cos_f32_cpu); + if (cos_sin_cache_) { + auto combined = combine_cos_sin_cache(cos_data, sin_data, max_seq_len_, cache_dim); + auto combined_cpu = Tensor::from_blob(combined.data(), {max_seq_len_, rotary_dim_}, DataType::F32, cpu_device); + cos_sin_cache_->copy_from(combined_cpu); + } } else if (dtype_ == DataType::BF16) { // Convert F32 to BF16 using the same conversion as Python's ml_dtypes.bfloat16 // This uses round-to-nearest-even (matching _f32_to_bf16 implementation) @@ -112,6 +142,11 @@ void RoPE::initialize_cache() { // copy_from handles cross-device copying to target device sin_cache_->copy_from(sin_bf16_cpu); cos_cache_->copy_from(cos_bf16_cpu); + if (cos_sin_cache_) { + auto combined = combine_cos_sin_cache(cos_bf16_data, sin_bf16_data, max_seq_len_, cache_dim); + auto combined_cpu = Tensor::from_blob(combined.data(), {max_seq_len_, rotary_dim_}, DataType::BF16, cpu_device); + cos_sin_cache_->copy_from(combined_cpu); + } } else if (dtype_ == DataType::F16) { // Convert F32 to F16 std::vector sin_f16_data(max_seq_len_ * cache_dim); @@ -127,6 +162,11 @@ void RoPE::initialize_cache() { sin_cache_->copy_from(sin_f16_cpu); cos_cache_->copy_from(cos_f16_cpu); + if (cos_sin_cache_) { + auto combined = combine_cos_sin_cache(cos_f16_data, sin_f16_data, max_seq_len_, cache_dim); + auto combined_cpu = Tensor::from_blob(combined.data(), {max_seq_len_, rotary_dim_}, DataType::F16, cpu_device); + cos_sin_cache_->copy_from(combined_cpu); + } } else { throw std::runtime_error( "RoPE cache dtype conversion not yet supported for dtype: " @@ -148,6 +188,21 @@ Tensor RoPE::forward(const Tensor &x, const Tensor &pos, bool in_place) const { } } +#ifdef ENABLE_INFINIOPS_API + if (cos_sin_cache_) { + if (!in_place) { + y->copy_from(x); + } + op::rotary_embedding_(pos, + y, + std::nullopt, + cos_sin_cache_, + static_cast(head_dim_), + algo_ == Algo::GPT_NEOX); + return y; + } +#endif + size_t ndim = x->ndim(); op::rope_(y->narrow({{ndim - 1, 0, rotary_dim_}}), x->narrow({{ndim - 1, 0, rotary_dim_}}), diff --git a/src/infinicore/ops/rotary_embedding/rotary_embedding.cc b/src/infinicore/ops/rotary_embedding/rotary_embedding.cc new file mode 100644 index 000000000..3fe6e3ad4 --- /dev/null +++ b/src/infinicore/ops/rotary_embedding/rotary_embedding.cc @@ -0,0 +1,51 @@ +#include "infinicore/ops/rotary_embedding.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RotaryEmbedding); + +RotaryEmbedding::RotaryEmbedding(const Tensor &positions, + Tensor query, + std::optional key, + const Tensor &cos_sin_cache, + int64_t head_size, + bool is_neox, + int64_t rope_dim_offset, + bool inverse) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(positions, query, cos_sin_cache); + if (key) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(query, *key); + } + INFINICORE_GRAPH_OP_DISPATCH( + query->device().getType(), positions, query, key, cos_sin_cache, + head_size, is_neox, rope_dim_offset, inverse); +} + +void RotaryEmbedding::execute(const Tensor &positions, + Tensor query, + std::optional key, + const Tensor &cos_sin_cache, + int64_t head_size, + bool is_neox, + int64_t rope_dim_offset, + bool inverse) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + RotaryEmbedding, positions, query, key, cos_sin_cache, head_size, + is_neox, rope_dim_offset, inverse); +} + +void rotary_embedding_(const Tensor &positions, + Tensor query, + std::optional key, + const Tensor &cos_sin_cache, + int64_t head_size, + bool is_neox, + int64_t rope_dim_offset, + bool inverse) { + RotaryEmbedding::execute(positions, query, key, cos_sin_cache, head_size, + is_neox, rope_dim_offset, inverse); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc new file mode 100644 index 000000000..86ea05e9e --- /dev/null +++ b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc @@ -0,0 +1,84 @@ +#include "infinicore/ops/rotary_embedding.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/rotary_embedding.h" + +namespace infinicore::op::rotary_embedding_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta positions, query, cos_sin_cache; + std::optional key; + graph::GraphTensor positions_tensor, query_tensor, cos_sin_cache_tensor; + std::optional key_tensor; + int64_t head_size; + bool is_neox; + int64_t rope_dim_offset; + bool inverse; +}; + +} // namespace + +void *plan(const Tensor &positions, + Tensor query, + std::optional key, + const Tensor &cos_sin_cache, + int64_t head_size, + bool is_neox, + int64_t rope_dim_offset, + bool inverse) { + INFINICORE_ASSERT(query->device().getType() == Device::Type::NVIDIA); + return new PlannedMeta{ + TensorMeta(positions), + TensorMeta(query), + TensorMeta(cos_sin_cache), + key ? std::optional{TensorMeta(*key)} : std::nullopt, + graph::GraphTensor(positions), + graph::GraphTensor(query), + graph::GraphTensor(cos_sin_cache), + key ? std::optional{graph::GraphTensor(*key)} : std::nullopt, + head_size, + is_neox, + rope_dim_offset, + inverse}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + const std::optional key = planned->key + ? std::optional{planned->key->tensor(*planned->key_tensor)} + : std::nullopt; + infini::ops::RotaryEmbedding::Call( + handle, + config, + planned->positions.tensor(planned->positions_tensor), + planned->query.tensor(planned->query_tensor), + key, + planned->cos_sin_cache.tensor(planned->cos_sin_cache_tensor), + planned->head_size, + planned->is_neox, + planned->rope_dim_offset, + planned->inverse); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + RotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::NVIDIA, &plan); + RotaryEmbedding::run_dispatcher().registerDevice(Device::Type::NVIDIA, &run); + RotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::NVIDIA, &cleanup); + return true; +}(); + +} // namespace infinicore::op::rotary_embedding_impl::infiniops +#endif From e5fbef93d4613dfc974b06511322ae707c1274e9 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 10 Aug 2026 23:52:56 +0800 Subject: [PATCH 06/10] refactor(ops): use Argmax for greedy sampling --- .../ops/random_sample/random_sample.cc | 45 +++++++++++++++++++ submodules/InfiniOps | 2 +- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/src/infinicore/ops/random_sample/random_sample.cc b/src/infinicore/ops/random_sample/random_sample.cc index bb0b41ce8..318f70bee 100644 --- a/src/infinicore/ops/random_sample/random_sample.cc +++ b/src/infinicore/ops/random_sample/random_sample.cc @@ -2,7 +2,46 @@ #include "../../utils.hpp" +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/argmax.h" +#endif + namespace infinicore::op { +namespace { + +#ifdef ENABLE_INFINIOPS_API +bool tryGreedyWithInfiniOps(Tensor indices, Tensor logits, int topk) { + const auto dtype = logits->dtype(); + if (logits->device().getType() != Device::Type::NVIDIA + || topk != 1 + || logits->ndim() != 1 + || logits->numel() == 0 + || !logits->is_contiguous() + || (dtype != DataType::F16 && dtype != DataType::BF16 && dtype != DataType::F32) + || indices->numel() != 1 + || indices->dtype() != DataType::I64 + || !indices->is_contiguous()) { + return false; + } + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + const std::optional no_dim; + infini::ops::Argmax::Call( + handle, + config, + infiniops::TensorMeta(logits).tensor(logits), + no_dim, + false, + infiniops::TensorMeta(indices).tensor(indices)); + return true; +} +#endif + +} // namespace common::OpDispatcher &RandomSample::dispatcher() { static common::OpDispatcher dispatcher_; @@ -14,6 +53,12 @@ void RandomSample::execute( float random_val, float topp, int topk, float temperature) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(indices, logits); infinicore::context::setDevice(logits->device()); +#ifdef ENABLE_INFINIOPS_API + if (tryGreedyWithInfiniOps(indices, logits, topk)) { + return; + } +#endif + dispatcher().lookup(logits->device().getType())( indices, logits, random_val, topp, topk, temperature); } diff --git a/submodules/InfiniOps b/submodules/InfiniOps index 1c865aea5..8875fb5a0 160000 --- a/submodules/InfiniOps +++ b/submodules/InfiniOps @@ -1 +1 @@ -Subproject commit 1c865aea58a6af8bbfdc67b76a4ed0ea8d1c167d +Subproject commit 8875fb5a02630f61481f9ab6669a7ff111292ff5 From 44ef940f31a4e93d9a31f6bec4c73a6a92f2402e Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 01:17:48 +0800 Subject: [PATCH 07/10] refactor(ops): use canonical FusedAddRmsNorm --- .../add_rms_norm/add_rms_norm_infiniops.cc | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc b/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc index f03cdec78..3ca8bf2ab 100644 --- a/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc +++ b/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc @@ -3,7 +3,8 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/add_rms_norm.h" +#include "base/copy.h" +#include "base/fused_add_rms_norm.h" #include @@ -16,6 +17,7 @@ struct PlannedMeta { TensorMeta out, residual, a, b, weight; graph::GraphTensor out_tensor, residual_tensor, a_tensor, b_tensor, weight_tensor; float epsilon; + bool copy_input, copy_residual; }; } // namespace @@ -35,7 +37,9 @@ void *plan(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const graph::GraphTensor(a), graph::GraphTensor(b), graph::GraphTensor(weight), - epsilon}; + epsilon, + out->data() != a->data(), + residual->data() != b->data()}; } void run(void *planned_meta) { @@ -45,15 +49,23 @@ void run(void *planned_meta) { handle.set_stream(context::getStream()); infini::ops::Config config; - infini::ops::AddRmsNorm::Call( + auto out = planned->out.tensor(planned->out_tensor); + auto residual = planned->residual.tensor(planned->residual_tensor); + + if (planned->copy_input) { + infini::ops::Copy::Call(handle, config, planned->a.tensor(planned->a_tensor), false, out); + } + if (planned->copy_residual) { + infini::ops::Copy::Call(handle, config, planned->b.tensor(planned->b_tensor), false, residual); + } + + infini::ops::FusedAddRmsNorm::Call( handle, config, - planned->a.tensor(planned->a_tensor), - planned->b.tensor(planned->b_tensor), - planned->weight.tensor(planned->weight_tensor), - std::optional{planned->epsilon}, - planned->out.tensor(planned->out_tensor), - planned->residual.tensor(planned->residual_tensor)); + out, + residual, + std::optional{planned->weight.tensor(planned->weight_tensor)}, + planned->epsilon); } void cleanup(void **planned_meta_ptr) { From 2f4973b96602c51585dfcfc6dce740af7446d511 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 01:18:08 +0800 Subject: [PATCH 08/10] refactor(ops): use canonical SiluAndMul --- src/infinicore/ops/swiglu/swiglu_infiniops.cc | 42 ++++++++++++++++--- 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/src/infinicore/ops/swiglu/swiglu_infiniops.cc b/src/infinicore/ops/swiglu/swiglu_infiniops.cc index 5b99a8eef..feacfd3fc 100644 --- a/src/infinicore/ops/swiglu/swiglu_infiniops.cc +++ b/src/infinicore/ops/swiglu/swiglu_infiniops.cc @@ -3,7 +3,8 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/swiglu.h" +#include "base/copy.h" +#include "base/silu_and_mul.h" namespace infinicore::op::swiglu_impl::infiniops { namespace { @@ -11,8 +12,8 @@ namespace { using TensorMeta = ::infinicore::op::infiniops::TensorMeta; struct PlannedMeta { - TensorMeta c, a, b; - graph::GraphTensor c_tensor, a_tensor, b_tensor; + TensorMeta c, a, b, packed, gate, up; + graph::GraphTensor c_tensor, a_tensor, b_tensor, packed_tensor, gate_tensor, up_tensor; }; } // namespace @@ -20,14 +21,32 @@ struct PlannedMeta { void *plan(Tensor c, const Tensor &a, const Tensor &b) { INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().getType())); INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + INFINICORE_ASSERT(c->shape() == a->shape()); + INFINICORE_ASSERT(a->shape() == b->shape()); + INFINICORE_ASSERT(c->dtype() == a->dtype()); + INFINICORE_ASSERT(a->dtype() == b->dtype()); + INFINICORE_ASSERT(!a->shape().empty()); + + auto packed_shape = a->shape(); + packed_shape.back() *= 2; + auto packed = Tensor::empty(packed_shape, a->dtype(), a->device()); + auto hidden_size = a->size(a->ndim() - 1); + auto gate = packed->narrow({{packed->ndim() - 1, 0, hidden_size}}); + auto up = packed->narrow({{packed->ndim() - 1, hidden_size, hidden_size}}); return new PlannedMeta{ TensorMeta(c), TensorMeta(a), TensorMeta(b), + TensorMeta(packed), + TensorMeta(gate), + TensorMeta(up), graph::GraphTensor(c), graph::GraphTensor(a), - graph::GraphTensor(b)}; + graph::GraphTensor(b), + graph::GraphTensor(packed), + graph::GraphTensor(gate), + graph::GraphTensor(up)}; } void run(void *planned_meta) { @@ -37,11 +56,22 @@ void run(void *planned_meta) { handle.set_stream(context::getStream()); infini::ops::Config config; - infini::ops::Swiglu::Call( + infini::ops::Copy::Call( handle, config, - planned->a.tensor(planned->a_tensor), planned->b.tensor(planned->b_tensor), + false, + planned->gate.tensor(planned->gate_tensor)); + infini::ops::Copy::Call( + handle, + config, + planned->a.tensor(planned->a_tensor), + false, + planned->up.tensor(planned->up_tensor)); + infini::ops::SiluAndMul::Call( + handle, + config, + planned->packed.tensor(planned->packed_tensor), planned->c.tensor(planned->c_tensor)); } From 6609f20367f2549fbe5d5713bfbad609d5a855ec Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 01:39:13 +0800 Subject: [PATCH 09/10] style(ops): format rotary embedding with clang-format 16 --- .../ops/rotary_embedding/rotary_embedding.cc | 14 +++++++------- .../rotary_embedding/rotary_embedding_infiniops.cc | 4 ++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/infinicore/ops/rotary_embedding/rotary_embedding.cc b/src/infinicore/ops/rotary_embedding/rotary_embedding.cc index 3fe6e3ad4..3856b7db4 100644 --- a/src/infinicore/ops/rotary_embedding/rotary_embedding.cc +++ b/src/infinicore/ops/rotary_embedding/rotary_embedding.cc @@ -7,13 +7,13 @@ namespace infinicore::op { INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RotaryEmbedding); RotaryEmbedding::RotaryEmbedding(const Tensor &positions, - Tensor query, - std::optional key, - const Tensor &cos_sin_cache, - int64_t head_size, - bool is_neox, - int64_t rope_dim_offset, - bool inverse) { + Tensor query, + std::optional key, + const Tensor &cos_sin_cache, + int64_t head_size, + bool is_neox, + int64_t rope_dim_offset, + bool inverse) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(positions, query, cos_sin_cache); if (key) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(query, *key); diff --git a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc index 86ea05e9e..79330f0e0 100644 --- a/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc +++ b/src/infinicore/ops/rotary_embedding/rotary_embedding_infiniops.cc @@ -53,8 +53,8 @@ void run(void *planned_meta) { handle.set_stream(context::getStream()); infini::ops::Config config; const std::optional key = planned->key - ? std::optional{planned->key->tensor(*planned->key_tensor)} - : std::nullopt; + ? std::optional{planned->key->tensor(*planned->key_tensor)} + : std::nullopt; infini::ops::RotaryEmbedding::Call( handle, config, From a69708c0ae2e4e2d1c5fe415a7aed7dbd5b61430 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 02:14:43 +0800 Subject: [PATCH 10/10] refactor(ops): use PyTorch Argmax provider --- src/infinicore/ops/random_sample/random_sample.cc | 1 + submodules/InfiniOps | 2 +- xmake.lua | 4 ++++ 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/infinicore/ops/random_sample/random_sample.cc b/src/infinicore/ops/random_sample/random_sample.cc index 318f70bee..22c9783c0 100644 --- a/src/infinicore/ops/random_sample/random_sample.cc +++ b/src/infinicore/ops/random_sample/random_sample.cc @@ -29,6 +29,7 @@ bool tryGreedyWithInfiniOps(Tensor indices, Tensor logits, int topk) { infini::ops::Handle handle; handle.set_stream(context::getStream()); infini::ops::Config config; + config.set_implementation_index(8); const std::optional no_dim; infini::ops::Argmax::Call( handle, diff --git a/submodules/InfiniOps b/submodules/InfiniOps index 8875fb5a0..47c1c4969 160000 --- a/submodules/InfiniOps +++ b/submodules/InfiniOps @@ -1 +1 @@ -Subproject commit 8875fb5a02630f61481f9ab6669a7ff111292ff5 +Subproject commit 47c1c4969c4312b4f15fc73928c72e4c262f79ce diff --git a/xmake.lua b/xmake.lua index de559e0ec..2ae1de1c1 100644 --- a/xmake.lua +++ b/xmake.lua @@ -417,6 +417,10 @@ local function build_infiniops_external(xmake_os) "-DGENERATE_PYTHON_BINDINGS=OFF", "-DCMAKE_BUILD_TYPE=Release" } + if has_config("nv-gpu") then + table.insert(cmake_config_args, "-DWITH_TORCH=ON") + table.insert(cmake_config_args, "-DINFINI_OPS_TORCH_OPS=argmax") + end if has_config("iluvatar-gpu") and has_config("aten") then table.insert(cmake_config_args, "-DTORCH_CXX11_ABI=0") table.insert(cmake_config_args, "-DCMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=0")