From f8679ea325aafa728c5fd6253b59e67210d48f77 Mon Sep 17 00:00:00 2001 From: Jingbo-gao Date: Thu, 16 Jul 2026 21:23:52 +0800 Subject: [PATCH 1/3] adapt flash attention to Ascend backend --- include/infinicore/adaptor/aten_adaptor.hpp | 5 + .../adaptor/flash_attention_adaptor.hpp | 9 +- python/infinicore/__init__.py | 6 + python/infinicore/device.py | 1 + python/infinicore/ops/ascend_flash_attn.py | 208 +++++++++++ .../ops/mha_kvcache/mha_kvcache_flashattn.cc | 7 +- .../mha_kvcache_flashattn_ascend.cc | 313 ++++++++++++++++ .../mha_varlen_flashattn.cc | 8 +- .../mha_varlen_flashattn_ascend.cc | 344 ++++++++++++++++++ src/infiniop/ops/flash_attention/operator.cc | 3 + xmake.lua | 42 +++ xmake/ascend.lua | 1 + 12 files changed, 942 insertions(+), 5 deletions(-) create mode 100644 python/infinicore/ops/ascend_flash_attn.py create mode 100644 src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc create mode 100644 src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn_ascend.cc diff --git a/include/infinicore/adaptor/aten_adaptor.hpp b/include/infinicore/adaptor/aten_adaptor.hpp index 74b053dc5..6ce3b426c 100644 --- a/include/infinicore/adaptor/aten_adaptor.hpp +++ b/include/infinicore/adaptor/aten_adaptor.hpp @@ -50,6 +50,11 @@ inline at::Device to_at_device(const Device &device) { else if (device.getType() == Device::Type::MOORE) { return at::Device(at::DeviceType::PrivateUse1, device.getIndex()); } +#endif +#if defined(ENABLE_ASCEND_FLASH_ATTN) + else if (device.getType() == Device::Type::ASCEND) { + return at::Device(at::DeviceType::PrivateUse1, device.getIndex()); + } #endif else { throw std::runtime_error("Unsupported device type for ATen"); diff --git a/include/infinicore/adaptor/flash_attention_adaptor.hpp b/include/infinicore/adaptor/flash_attention_adaptor.hpp index 2c195cf0f..b1e778e55 100644 --- a/include/infinicore/adaptor/flash_attention_adaptor.hpp +++ b/include/infinicore/adaptor/flash_attention_adaptor.hpp @@ -5,7 +5,12 @@ // NVIDIA flash-attn-nvidia.so uses namespace flash. The pip/MetaX flash_attn_2_cuda extension // exports the same entry points at global scope (no namespace), matching FLASH_NAMESPACE builds // where the namespace is empty. -#if !defined(ENABLE_METAX_API) +// +// Ascend (aclnn C API path) does NOT use the flash:: namespace at all — the aclnn kernels are +// called directly from the dedicated *._ascend.cc implementation files, so this header is only +// included by the NVIDIA/MetaX/QY code paths. We still guard the namespace below so that +// existing code compiles unchanged when ENABLE_ASCEND_FLASH_ATTN is defined. +#if !defined(ENABLE_METAX_API) && !defined(ENABLE_ASCEND_FLASH_ATTN) namespace flash { #endif std::vector @@ -136,7 +141,7 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size #endif ); -#if !defined(ENABLE_METAX_API) +#if !defined(ENABLE_METAX_API) && !defined(ENABLE_ASCEND_FLASH_ATTN) } // namespace flash #endif #endif // ENABLE_FLASH_ATTN diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index 808995640..cba4aed1e 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -113,6 +113,10 @@ moore_mate_flash_attn_prefill, ) from infinicore.ops.mrope import mrope +from infinicore.ops.ascend_flash_attn import ( + ascend_flash_attn_decode, + ascend_flash_attn_prefill, +) from infinicore.ops.mul import mul from infinicore.ops.mul_scalar import mul_scalar from infinicore.ops.narrow import narrow @@ -292,6 +296,8 @@ "var_mean", "moore_mate_flash_attn_prefill", "moore_mate_flash_attn_decode", + "ascend_flash_attn_prefill", + "ascend_flash_attn_decode", "var", "topk", "all", diff --git a/python/infinicore/device.py b/python/infinicore/device.py index 8858a1235..3607da183 100644 --- a/python/infinicore/device.py +++ b/python/infinicore/device.py @@ -1,3 +1,4 @@ +import torch from infinicore.lib import _infinicore _infinicore_2_python_dict = {} diff --git a/python/infinicore/ops/ascend_flash_attn.py b/python/infinicore/ops/ascend_flash_attn.py new file mode 100644 index 000000000..01379c866 --- /dev/null +++ b/python/infinicore/ops/ascend_flash_attn.py @@ -0,0 +1,208 @@ +""" +Paged Flash-Attention wrapper backed by Ascend CANN aclnnFusedInferAttentionScore. + +Runtime requirements: + - torch_npu (Ascend PyTorch adapter) + - CANN toolkit (provides libopapi.so, libascendcl.so, etc.) + +Provides two entry points: + - ascend_flash_attn_decode: decode with layout (num_blocks, block_size, num_kv_heads, head_size) + - ascend_flash_attn_prefill: variable-length prefill (used by mha_varlen) + +NOTE: The actual compute is done by the C++ *._ascend.cc files that call +aclnnFusedInferAttentionScore directly. This Python module exists as a +convenience wrapper for users who prefer a pure-Python calling convention +(e.g. for testing or benchmarking) and mirrors the moore_mate_flash_attn.py +interface. +""" + +import torch + +try: + import torch_npu # noqa: F401 - register NPU backend + + _NPU_AVAILABLE = True +except ImportError: + _NPU_AVAILABLE = False + + +def is_available() -> bool: + return _NPU_AVAILABLE + + +def _check_available(): + if not _NPU_AVAILABLE: + raise RuntimeError( + "torch_npu is not installed. " + "Please install torch_npu and CANN toolkit first." + ) + + +@torch.inference_mode() +def ascend_flash_attn_decode( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + scale: float, + block_size: int, + max_seq_len: int, +) -> torch.Tensor: + """ + Decode entry point for Ascend NPU. + + This is a thin wrapper that validates inputs and delegates to + infinicore's C++ mha_kvcache op (which calls + aclnnFusedInferAttentionScore internally when ENABLE_ASCEND_FLASH_ATTN + is defined at build time). + + Args: + q: (num_seqs, num_heads, head_size) on NPU + k_cache: (num_blocks, block_size, num_kv_heads, head_size) on NPU + v_cache: (num_blocks, block_size, num_kv_heads, head_size) on NPU + block_tables: (num_seqs, max_blocks_per_seq) int32 on NPU + seq_lens: (num_seqs,) int32 on NPU + scale: softmax scale (typically 1/sqrt(head_size)) + block_size: paged KV block size + max_seq_len: maximum sequence length (reserved, not used by aclnn) + + Returns: + out: (num_seqs, num_heads, head_size) on NPU + """ + _check_available() + + import infinicore + + num_seqs, num_heads, head_size = q.shape + + q_4d = q.unsqueeze(1) # (num_seqs, 1, num_heads, head_size) + + out_4d = infinicore.mha_kvcache( + q=infinicore.from_torch(q_4d), + k_cache=infinicore.from_torch(k_cache), + v_cache=infinicore.from_torch(v_cache), + seqlens_k=infinicore.from_torch(seq_lens), + block_table=infinicore.from_torch(block_tables), + alibi_slopes=None, + scale=scale, + ) + + return out_4d.squeeze(1) # (num_seqs, num_heads, head_size) + + +@torch.inference_mode() +def ascend_flash_attn_prefill( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + block_tables: torch.Tensor, + scale: float, + max_seqlen_q: int, + max_seqlen_k: int, + block_size: int, + causal: bool = True, +) -> torch.Tensor: + """ + Variable-length prefill entry point for Ascend NPU. + + Delegates to infinicore's C++ mha_varlen op (which calls + aclnnFusedInferAttentionScore internally). + + Args: + q: (total_q, num_heads, head_size) on NPU + k_cache: (num_blocks, block_size, num_kv_heads, head_size) on NPU + v_cache: (num_blocks, block_size, num_kv_heads, head_size) on NPU + cu_seqlens_q: (batch_size + 1,) int32 on NPU + cu_seqlens_k: (batch_size + 1,) int32 on NPU + block_tables: (batch_size, max_blocks_per_seq) int32 on NPU + scale: softmax scale + max_seqlen_q: max query sequence length in the batch + max_seqlen_k: max key sequence length in the batch + block_size: paged KV block size + causal: whether to apply causal mask + + Returns: + out: (total_q, num_heads, head_size) on NPU + """ + _check_available() + + import infinicore + + out = infinicore.mha_varlen( + q=infinicore.from_torch(q), + k=infinicore.from_torch(k_cache), + v=infinicore.from_torch(v_cache), + cum_seqlens_q=infinicore.from_torch(cu_seqlens_q), + cum_seqlens_k=infinicore.from_torch(cu_seqlens_k), + block_table=infinicore.from_torch(block_tables), + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + alibi_slopes=None, + scale=scale, + ) + + return out + + +if __name__ == "__main__": + if not is_available(): + raise SystemExit("torch_npu not available, please install torch_npu first.") + + device = torch.device("npu:0") + + num_seqs, num_heads, num_kv_heads = 2, 8, 2 + head_size, block_size, max_seq_len = 128, 16, 64 + num_blocks = 32 + + q = torch.randn(num_seqs, num_heads, head_size, dtype=torch.float16, device=device) + k_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, + dtype=torch.float16, device=device, + ) + v_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, + dtype=torch.float16, device=device, + ) + block_tables = torch.zeros(num_seqs, 4, dtype=torch.int32, device=device) + block_tables[0, 0] = 0 + block_tables[1, 0] = 1 + seq_lens = torch.tensor([32, 48], dtype=torch.int32, device=device) + + out = ascend_flash_attn_decode( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_tables=block_tables, + seq_lens=seq_lens, + scale=head_size**-0.5, + block_size=block_size, + max_seq_len=max_seq_len, + ) + torch.npu.synchronize() + assert list(out.shape) == list(q.shape), f"Expected {q.shape}, got {out.shape}" + print("ascend_flash_attn_decode test passed") + + # --- Prefill test --- + total_q = 64 # sum of sequence lengths + cu_seqlens_q = torch.tensor([0, 32, 64], dtype=torch.int32, device=device) + cu_seqlens_k = torch.tensor([0, 32, 64], dtype=torch.int32, device=device) + q_prefill = torch.randn(total_q, num_heads, head_size, dtype=torch.float16, device=device) + out_prefill = ascend_flash_attn_prefill( + q=q_prefill, + k_cache=k_cache, + v_cache=v_cache, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + block_tables=block_tables, + scale=head_size**-0.5, + max_seqlen_q=32, + max_seqlen_k=32, + block_size=block_size, + causal=True, + ) + torch.npu.synchronize() + assert list(out_prefill.shape) == list(q_prefill.shape), f"Expected {q_prefill.shape}, got {out_prefill.shape}" + print("ascend_flash_attn_prefill test passed") diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc index 0167c17df..fbed186bc 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc @@ -44,7 +44,12 @@ void *plan(Tensor out, } void run(void *planned_meta) { -#ifdef ENABLE_FLASH_ATTN +#if defined(ENABLE_ASCEND_FLASH_ATTN) + // Ascend path: delegate to the dedicated _ascend.cc implementation. + // This should never be reached because the Ascend implementation registers + // its own plan/run/cleanup for Device::Type::ASCEND. If it does, throw. + throw std::runtime_error("[mha_kvcache/flashattn] Ascend path should use flashattn_ascend impl"); +#elif 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 diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc new file mode 100644 index 000000000..da9273cf0 --- /dev/null +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc @@ -0,0 +1,313 @@ +#if defined(ENABLE_ASCEND_FLASH_ATTN) + +#include "infinicore/context/context.hpp" +#include "infinicore/ops/mha_kvcache.hpp" + + +#include +#include +#include + +#include +#include +#include +#include + +namespace infinicore::op::mha_kvcache_impl::flashattn_ascend { + +static aclDataType to_acl_dtype(DataType dtype) { + switch (dtype) { + case DataType::F16: + return ACL_FLOAT16; + case DataType::BF16: + return ACL_BF16; + case DataType::F32: + return ACL_FLOAT; + case DataType::I32: + return ACL_INT32; + case DataType::I64: + return ACL_INT64; + default: + throw std::runtime_error( + "[mha_kvcache/ascend] Unsupported dtype for aclTensor"); + } +} + +static aclIntArray * +host_vector_to_acl_int_array(const std::vector &vec) { + return aclCreateIntArray(vec.data(), vec.size()); +} + +struct PlannedMeta { + graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; + std::optional alibi_slopes; + float scale; +}; + +void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, + const Tensor &v_cache, const Tensor &seqlens_k, + const Tensor &block_table, std::optional alibi_slopes, + float scale) { + return new PlannedMeta{graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k_cache), + graph::GraphTensor(v_cache), + graph::GraphTensor(seqlens_k), + graph::GraphTensor(block_table), + alibi_slopes ? std::optional( + graph::GraphTensor(*alibi_slopes)) + : std::nullopt, + scale}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + infinicore::context::setDevice(p->q->device()); + + if (p->alibi_slopes.has_value()) { + throw std::runtime_error("[mha_kvcache/ascend] ALiBi not supported by " + "aclnnFusedInferAttentionScore"); + } + + // q/out are BSND [batch, 1, num_heads, head_size] in InfiniCore. For + // decode S=1, the same memory can be described to FIA as BNSD + // [batch, num_heads, 1, head_size]. + auto q_shape = p->q->shape(); + auto k_shape = p->k_cache->shape(); + auto v_shape = p->v_cache->shape(); + + if (q_shape.size() != 4 || k_shape.size() != 4 || v_shape.size() != 4) { + throw std::runtime_error("[mha_kvcache/ascend] flash attention expects q " + "and KV cache to be 4D tensors"); + } + + const int64_t batch_size = q_shape[0]; + const int64_t num_heads = q_shape[2]; + const int64_t head_size = q_shape[3]; + const int64_t num_blocks = k_shape[0]; + const int64_t block_size_val = k_shape[1]; + const int64_t num_kv_heads = k_shape[2]; + const int64_t v_head_size = v_shape[3]; + + if (k_shape[3] != static_cast(head_size)) { + throw std::runtime_error( + "[mha_kvcache/ascend] k_cache head_size does not match q head_size"); + } + if (v_shape[0] != k_shape[0] || v_shape[1] != k_shape[1] || + v_shape[2] != k_shape[2]) { + throw std::runtime_error( + "[mha_kvcache/ascend] k_cache and v_cache shapes are incompatible"); + } + + Tensor q_work = p->q->is_contiguous() ? Tensor(p->q) : p->q->contiguous(); + Tensor k_work = p->k_cache->is_contiguous() ? Tensor(p->k_cache) + : p->k_cache->contiguous(); + Tensor v_work = p->v_cache->is_contiguous() ? Tensor(p->v_cache) + : p->v_cache->contiguous(); + Tensor bt_work = p->block_table->is_contiguous() + ? Tensor(p->block_table) + : p->block_table->contiguous(); + Tensor out_work = + p->out->is_contiguous() ? Tensor(p->out) : p->out->contiguous(); + + aclDataType q_dtype = to_acl_dtype(q_work->dtype()); + + // Read seqlens_k to host + auto seqlens_k_shape = p->seqlens_k->shape(); + int64_t seqlens_k_len = seqlens_k_shape[0]; + std::vector seqlens_k_host(seqlens_k_len); + auto copy_ret = + aclrtMemcpy(seqlens_k_host.data(), seqlens_k_len * sizeof(int32_t), + reinterpret_cast(p->seqlens_k->data()), + seqlens_k_len * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); + if (copy_ret != ACL_SUCCESS) { + throw std::runtime_error( + std::string("[mha_kvcache/ascend] copy seqlens_k to host failed: ") + + std::to_string(copy_ret)); + } + + // Build actual_seq vectors + std::vector actual_seq_q_vec(batch_size, + 1); // decode: always 1 query token + std::vector actual_seq_k_vec; + actual_seq_k_vec.reserve(batch_size); + for (int64_t i = 0; i < batch_size; ++i) { + actual_seq_k_vec.push_back(seqlens_k_host[i]); + } + + // BNSD [batch, num_heads, 1, head_size], viewed from BSND memory. + std::vector q_dims = {batch_size, num_heads, 1, head_size}; + std::vector q_strides = {num_heads * head_size, head_size, head_size, + 1}; + aclTensor *query_acl = aclCreateTensor( + q_dims.data(), q_dims.size(), q_dtype, q_strides.data(), 0, ACL_FORMAT_ND, + q_dims.data(), q_dims.size(), + const_cast(reinterpret_cast(q_work->data()))); + + // The physical BnBsND cache is contiguous in N and D, so expose it to FIA + // as BnBsH without copying. + std::vector k_dims = {num_blocks, block_size_val, + num_kv_heads * head_size}; + std::vector k_strides = {block_size_val * num_kv_heads * head_size, + num_kv_heads * head_size, 1}; + aclTensor *k_acl_tensor = aclCreateTensor( + k_dims.data(), k_dims.size(), q_dtype, k_strides.data(), 0, ACL_FORMAT_ND, + k_dims.data(), k_dims.size(), + const_cast(reinterpret_cast(k_work->data()))); + aclTensorList *key_acl = aclCreateTensorList(&k_acl_tensor, 1); + + std::vector v_dims = {num_blocks, block_size_val, + num_kv_heads * v_head_size}; + std::vector v_strides = {block_size_val * num_kv_heads * v_head_size, + num_kv_heads * v_head_size, 1}; + aclTensor *v_acl_tensor = aclCreateTensor( + v_dims.data(), v_dims.size(), q_dtype, v_strides.data(), 0, ACL_FORMAT_ND, + v_dims.data(), v_dims.size(), + const_cast(reinterpret_cast(v_work->data()))); + aclTensorList *value_acl = aclCreateTensorList(&v_acl_tensor, 1); + + // Block table: [batch, max_blocks_per_seq] INT32 on device + auto bt_shape = bt_work->shape(); + std::vector bt_dims = {bt_shape[0], bt_shape[1]}; + std::vector bt_strides = {bt_shape[1], 1}; + aclTensor *block_table_acl = aclCreateTensor( + bt_dims.data(), bt_dims.size(), ACL_INT32, bt_strides.data(), 0, + ACL_FORMAT_ND, bt_dims.data(), bt_dims.size(), + const_cast(reinterpret_cast(bt_work->data()))); + + // BNSD [batch, num_heads, 1, head_size], viewed from BSND memory. + std::vector out_dims = {batch_size, num_heads, 1, head_size}; + std::vector out_strides = {num_heads * head_size, head_size, + head_size, 1}; + aclDataType out_dtype = to_acl_dtype(out_work->dtype()); + aclTensor *out_acl = aclCreateTensor( + out_dims.data(), out_dims.size(), out_dtype, out_strides.data(), 0, + ACL_FORMAT_ND, out_dims.data(), out_dims.size(), + const_cast(reinterpret_cast(out_work->data()))); + + aclIntArray *seqlens_q_acl = host_vector_to_acl_int_array(actual_seq_q_vec); + aclIntArray *seqlens_k_acl = host_vector_to_acl_int_array(actual_seq_k_vec); + + // Call CANN API with Paged Attention + // inputLayout="BNSD": query/out=[B,N,S,D], KV cache is BnBsH for paged + // attention. sparse_mode=0: no mask needed for decode (Q_S=1, + // IncreFlashAttention branch) + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + + aclnnStatus ret = aclnnFusedInferAttentionScoreV4GetWorkspaceSize( + query_acl, key_acl, value_acl, + nullptr, // pseShift + nullptr, // atten_mask (not needed for decode) + seqlens_q_acl, seqlens_k_acl, nullptr, nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, + block_table_acl, // blockTable - Paged Attention + nullptr, // queryPaddingSize + nullptr, // kvPaddingSize + nullptr, // keyAntiquantScale + nullptr, // keyAntiquantOffset + nullptr, // valueAntiquantScale + nullptr, // valueAntiquantOffset + nullptr, // keySharedPrefix + nullptr, // valueSharedPrefix + nullptr, // actualSharedPrefixLen + nullptr, // queryRope + nullptr, // keyRope + nullptr, // keyRopeAntiquantScale + nullptr, // dequantScaleQuery + nullptr, // learnableSink + num_heads, static_cast(p->scale), 2147483647, 2147483647, + const_cast("BNSD"), num_kv_heads, + 0, // sparse_mode=0 (no mask for decode) + 0, // innerPrecise + block_size_val, // blockSize - Paged Attention block size + 0, // antiquantMode + false, + 0, // keyAntiquantMode + 0, // valueAntiquantMode + 0, // queryQuantMode + out_acl, nullptr, &workspace_size, &executor); + + if (ret != 0) { + aclDestroyTensor(query_acl); + aclDestroyTensorList(key_acl); + aclDestroyTensorList(value_acl); + aclDestroyTensor(block_table_acl); + aclDestroyTensor(out_acl); + aclDestroyIntArray(seqlens_q_acl); + aclDestroyIntArray(seqlens_k_acl); + const char *err_msg = aclGetRecentErrMsg(); + throw std::runtime_error( + std::string( + "[mha_kvcache/ascend] " + "aclnnFusedInferAttentionScoreV4GetWorkspaceSize failed: ") + + std::to_string(ret) + ", msg: " + (err_msg ? err_msg : "(null)")); + } + + void *workspace = nullptr; + if (workspace_size > 0) { + aclError alloc_ret = + aclrtMalloc(&workspace, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST); + if (alloc_ret != ACL_SUCCESS) { + aclDestroyTensor(query_acl); + aclDestroyTensorList(key_acl); + aclDestroyTensorList(value_acl); + aclDestroyTensor(block_table_acl); + aclDestroyTensor(out_acl); + aclDestroyIntArray(seqlens_q_acl); + aclDestroyIntArray(seqlens_k_acl); + throw std::runtime_error( + std::string("[mha_kvcache/ascend] aclrtMalloc workspace failed: ") + + std::to_string(alloc_ret)); + } + } + + aclrtStream stream = + static_cast(infinicore::context::getStream()); + ret = aclnnFusedInferAttentionScoreV4(workspace, workspace_size, executor, + stream); + aclrtSynchronizeStream(stream); + + if (workspace) { + aclrtFree(workspace); + } + + // Release aclTensor/aclTensorList/aclIntArray resources + aclDestroyTensor(query_acl); + aclDestroyTensorList(key_acl); + aclDestroyTensorList(value_acl); + aclDestroyTensor(block_table_acl); + aclDestroyTensor(out_acl); + aclDestroyIntArray(seqlens_q_acl); + aclDestroyIntArray(seqlens_k_acl); + + if (ret != 0) { + const char *err_msg = aclGetRecentErrMsg(); + throw std::runtime_error( + std::string( + "[mha_kvcache/ascend] aclnnFusedInferAttentionScoreV4 failed: ") + + std::to_string(ret) + ", msg: " + (err_msg ? err_msg : "(null)")); + } + + // Copy back if out was not contiguous + if (!p->out->is_contiguous()) { + p->out->copy_from(out_work); + } +} + +void cleanup(void **planned_meta_ptr) { + auto *p = *reinterpret_cast(planned_meta_ptr); + delete p; + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MhaKVCache::plan_dispatcher().registerDevice(Device::Type::ASCEND, &plan); + MhaKVCache::run_dispatcher().registerDevice(Device::Type::ASCEND, &run); + MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::ASCEND, + &cleanup); + return true; +}(); + +} // namespace infinicore::op::mha_kvcache_impl::flashattn_ascend +#endif 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 4a233bf5b..ebc3b0893 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 @@ -65,7 +65,10 @@ namespace { } // namespace void run(void *planned_meta) { -#ifndef ENABLE_ATEN +#if defined(ENABLE_ASCEND_FLASH_ATTN) + (void)planned_meta; + throw std::runtime_error("[mha_varlen/flashattn] Ascend path should use flashattn_ascend impl"); +#elif !defined(ENABLE_ATEN) (void)planned_meta; throw std::runtime_error("ATen is not enabled in this build"); #else @@ -120,7 +123,8 @@ void run(void *planned_meta) { std::nullopt, 0.0, true, - std::optional(static_cast(p->scale))); + static_cast(p->scale), + num_heads != num_kv_heads); out_work.copy_(result.permute({0, 2, 1, 3}).reshape({q.size(0), num_heads, value_dim})); if (out_need_copy_back) { p->out->copy_from(out_work_ic); diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn_ascend.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn_ascend.cc new file mode 100644 index 000000000..8bd51839b --- /dev/null +++ b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn_ascend.cc @@ -0,0 +1,344 @@ +#if defined(ENABLE_ASCEND_FLASH_ATTN) + +#include "infinicore/context/context.hpp" +#include "infinicore/ops/mha_varlen.hpp" + + +#include +#include +#include + +#include +#include +#include +#include + +namespace infinicore::op::mha_varlen_impl::flashattn_ascend { + +static aclDataType to_acl_dtype(DataType dtype) { + switch (dtype) { + case DataType::F16: + return ACL_FLOAT16; + case DataType::BF16: + return ACL_BF16; + case DataType::F32: + return ACL_FLOAT; + case DataType::I32: + return ACL_INT32; + case DataType::I64: + return ACL_INT64; + default: + throw std::runtime_error( + "[mha_varlen/ascend] Unsupported dtype for aclTensor"); + } +} + +static aclIntArray * +host_vector_to_acl_int_array(const std::vector &vec) { + return aclCreateIntArray(vec.data(), vec.size()); +} + +struct PlannedMeta { + graph::GraphTensor out, q, k, v, cum_seqlens_q, cum_seqlens_k; + std::optional block_table; + int max_seqlen_q, max_seqlen_k; + std::optional alibi_slopes; + float scale; + // Per-device causal mask (each PlannedMeta gets its own) + void *mask_dev = nullptr; + aclTensor *mask_acl = nullptr; +}; + +static aclTensor *create_causal_mask(PlannedMeta *p) { + if (p->mask_acl != nullptr) { + return p->mask_acl; + } + + int64_t mask_size = 2048 * 2048; + + std::vector mask_host(mask_size, 0); + for (int64_t i = 0; i < 2048; ++i) { + for (int64_t j = 0; j < 2048; ++j) { + if (j > i) { + mask_host[i * 2048 + j] = 1; + } + } + } + + aclrtMalloc(&p->mask_dev, mask_size, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMemcpy(p->mask_dev, mask_size, mask_host.data(), mask_size, + ACL_MEMCPY_HOST_TO_DEVICE); + + std::vector mask_dims = {2048, 2048}; + std::vector mask_strides = {2048, 1}; + p->mask_acl = aclCreateTensor( + mask_dims.data(), mask_dims.size(), ACL_UINT8, mask_strides.data(), 0, + ACL_FORMAT_ND, mask_dims.data(), mask_dims.size(), p->mask_dev); + return p->mask_acl; +} + +void *plan(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, + const Tensor &cum_seqlens_q, const Tensor &cum_seqlens_k, + std::optional block_table, int max_seqlen_q, + int max_seqlen_k, std::optional alibi_slopes, float scale) { + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k), + graph::GraphTensor(v), + graph::GraphTensor(cum_seqlens_q), + graph::GraphTensor(cum_seqlens_k), + block_table + ? std::optional(graph::GraphTensor(*block_table)) + : std::nullopt, + max_seqlen_q, + max_seqlen_k, + alibi_slopes + ? std::optional(graph::GraphTensor(*alibi_slopes)) + : std::nullopt, + scale}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + infinicore::context::setDevice(p->q->device()); + + if (p->alibi_slopes.has_value()) { + throw std::runtime_error("[mha_varlen/ascend] ALiBi not supported by " + "aclnnFusedInferAttentionScore"); + } + + auto cu_q_tensor = p->cum_seqlens_q; + auto cu_k_tensor = p->cum_seqlens_k; + auto cu_q_shape = cu_q_tensor->shape(); + auto cu_k_shape = cu_k_tensor->shape(); + int64_t cu_q_len = cu_q_shape[0]; + int64_t cu_k_len = cu_k_shape[0]; + int64_t batch_size = cu_q_len - 1; + + std::vector cu_q_host(cu_q_len); + std::vector cu_k_host(cu_k_len); + aclrtMemcpy(cu_q_host.data(), cu_q_len * sizeof(int32_t), + reinterpret_cast(cu_q_tensor->data()), + cu_q_len * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); + aclrtMemcpy(cu_k_host.data(), cu_k_len * sizeof(int32_t), + reinterpret_cast(cu_k_tensor->data()), + cu_k_len * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); + + std::vector actual_seq_q_vec; + std::vector actual_seq_k_vec; + for (int64_t i = 0; i < batch_size; ++i) { + int64_t k_len = cu_k_host[i + 1] - cu_k_host[i]; + // TND query uses cumulative sequence lengths. Paged KV cache keeps + // actualSeqLengthsKv in per-batch mode. + actual_seq_q_vec.push_back(cu_q_host[i + 1]); + actual_seq_k_vec.push_back(k_len); + } + + auto q_shape = p->q->shape(); + auto k_shape = p->k->shape(); + const int64_t num_heads = q_shape[1]; + const int64_t head_size = q_shape[2]; + const int64_t block_size_val = k_shape[1]; + const int64_t num_kv_heads = k_shape[2]; + const int64_t num_blocks = k_shape[0]; + + Tensor q_work = p->q->is_contiguous() ? Tensor(p->q) : p->q->contiguous(); + Tensor k_work = p->k->is_contiguous() ? Tensor(p->k) : p->k->contiguous(); + Tensor v_work = p->v->is_contiguous() ? Tensor(p->v) : p->v->contiguous(); + Tensor out_work = + p->out->is_contiguous() ? Tensor(p->out) : p->out->contiguous(); + + aclDataType q_dtype = to_acl_dtype(q_work->dtype()); + + // Query is already contiguous TND [total_q, num_heads, head_size]. + std::vector q_dims = {static_cast(q_shape[0]), num_heads, + head_size}; + std::vector q_strides = {num_heads * head_size, head_size, 1}; + aclTensor *query_acl = aclCreateTensor( + q_dims.data(), q_dims.size(), q_dtype, q_strides.data(), 0, ACL_FORMAT_ND, + q_dims.data(), q_dims.size(), + const_cast(reinterpret_cast(q_work->data()))); + + // The physical BnBsND cache is contiguous in N and D, so expose it to FIA + // as BnBsH without copying. + std::vector k_dims = {num_blocks, block_size_val, + num_kv_heads * head_size}; + std::vector k_strides = {num_kv_heads * block_size_val * head_size, + num_kv_heads * head_size, 1}; + aclTensor *k_acl_tensor = aclCreateTensor( + k_dims.data(), k_dims.size(), q_dtype, k_strides.data(), 0, ACL_FORMAT_ND, + k_dims.data(), k_dims.size(), + const_cast(reinterpret_cast(k_work->data()))); + aclTensorList *key_acl = aclCreateTensorList(&k_acl_tensor, 1); + + // Value uses the same contiguous BnBsH representation. + std::vector v_dims = {num_blocks, block_size_val, + num_kv_heads * head_size}; + std::vector v_strides = {num_kv_heads * block_size_val * head_size, + num_kv_heads * head_size, 1}; + aclTensor *v_acl_tensor = aclCreateTensor( + v_dims.data(), v_dims.size(), q_dtype, v_strides.data(), 0, ACL_FORMAT_ND, + v_dims.data(), v_dims.size(), + const_cast(reinterpret_cast(v_work->data()))); + aclTensorList *value_acl = aclCreateTensorList(&v_acl_tensor, 1); + + // Block table: [batch, max_blocks_per_seq] INT32 on device + aclTensor *block_table_acl = nullptr; + if (p->block_table.has_value()) { + auto &bt = p->block_table.value(); + Tensor bt_work = bt->is_contiguous() ? Tensor(bt) : bt->contiguous(); + auto bt_shape = bt_work->shape(); + std::vector bt_dims = {bt_shape[0], bt_shape[1]}; + std::vector bt_strides = {bt_shape[1], 1}; + block_table_acl = aclCreateTensor( + bt_dims.data(), bt_dims.size(), ACL_INT32, bt_strides.data(), 0, + ACL_FORMAT_ND, bt_dims.data(), bt_dims.size(), + const_cast(reinterpret_cast(bt_work->data()))); + } + + // FIA writes directly to contiguous TND output. + auto out_shape = out_work->shape(); + std::vector out_dims = {static_cast(out_shape[0]), + num_heads, head_size}; + std::vector out_strides = {num_heads * head_size, head_size, 1}; + aclDataType out_dtype = to_acl_dtype(out_work->dtype()); + aclTensor *out_acl = aclCreateTensor( + out_dims.data(), out_dims.size(), out_dtype, out_strides.data(), 0, + ACL_FORMAT_ND, out_dims.data(), out_dims.size(), + const_cast(reinterpret_cast(out_work->data()))); + + aclIntArray *actual_seq_q_acl = + host_vector_to_acl_int_array(actual_seq_q_vec); + aclIntArray *actual_seq_k_acl = + host_vector_to_acl_int_array(actual_seq_k_vec); + + int64_t sparse_mode = 3; // rightDownCausal + aclTensor *atten_mask_acl = create_causal_mask(p); + + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + + aclnnStatus ret = aclnnFusedInferAttentionScoreV4GetWorkspaceSize( + query_acl, key_acl, value_acl, + nullptr, // pseShift + atten_mask_acl, actual_seq_q_acl, actual_seq_k_acl, nullptr, nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, + block_table_acl, // blockTable - Paged Attention + nullptr, // queryPaddingSize + nullptr, // kvPaddingSize + nullptr, // keyAntiquantScale + nullptr, // keyAntiquantOffset + nullptr, // valueAntiquantScale + nullptr, // valueAntiquantOffset + nullptr, // keySharedPrefix + nullptr, // valueSharedPrefix + nullptr, // actualSharedPrefixLen + nullptr, // queryRope + nullptr, // keyRope + nullptr, // keyRopeAntiquantScale + nullptr, // dequantScaleQuery + nullptr, // learnableSink + num_heads, static_cast(p->scale), 2147483647, 2147483647, + const_cast("TND"), num_kv_heads, sparse_mode, + 0, // innerPrecise: high precision TND path + block_size_val, // blockSize - Paged Attention block size + 0, // antiquantMode + false, + 0, // keyAntiquantMode + 0, // valueAntiquantMode + 0, // queryQuantMode + out_acl, nullptr, &workspace_size, &executor); + + if (ret != 0) { + aclDestroyTensor(query_acl); + aclDestroyTensorList(key_acl); + aclDestroyTensorList(value_acl); + if (block_table_acl) + aclDestroyTensor(block_table_acl); + aclDestroyTensor(out_acl); + aclDestroyIntArray(actual_seq_q_acl); + aclDestroyIntArray(actual_seq_k_acl); + const char *err_msg = aclGetRecentErrMsg(); + throw std::runtime_error( + std::string( + "[mha_varlen/ascend] " + "aclnnFusedInferAttentionScoreV4GetWorkspaceSize failed: ") + + std::to_string(ret) + ", msg: " + (err_msg ? err_msg : "(null)")); + } + + void *workspace = nullptr; + if (workspace_size > 0) { + aclError alloc_ret = + aclrtMalloc(&workspace, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST); + if (alloc_ret != ACL_SUCCESS) { + aclDestroyTensor(query_acl); + aclDestroyTensorList(key_acl); + aclDestroyTensorList(value_acl); + if (block_table_acl) + aclDestroyTensor(block_table_acl); + aclDestroyTensor(out_acl); + aclDestroyIntArray(actual_seq_q_acl); + aclDestroyIntArray(actual_seq_k_acl); + throw std::runtime_error( + std::string("[mha_varlen/ascend] aclrtMalloc workspace failed: ") + + std::to_string(alloc_ret)); + } + } + + aclrtStream stream = + static_cast(infinicore::context::getStream()); + ret = aclnnFusedInferAttentionScoreV4(workspace, workspace_size, executor, + stream); + aclrtSynchronizeStream(stream); + + if (workspace) { + aclrtFree(workspace); + } + + // Release aclTensor/aclTensorList/aclIntArray resources + aclDestroyTensor(query_acl); + aclDestroyTensorList(key_acl); + aclDestroyTensorList(value_acl); + if (block_table_acl) + aclDestroyTensor(block_table_acl); + aclDestroyTensor(out_acl); + aclDestroyIntArray(actual_seq_q_acl); + aclDestroyIntArray(actual_seq_k_acl); + + if (ret != 0) { + const char *err_msg = aclGetRecentErrMsg(); + throw std::runtime_error( + std::string( + "[mha_varlen/ascend] aclnnFusedInferAttentionScoreV4 failed: ") + + std::to_string(ret) + ", msg: " + (err_msg ? err_msg : "(null)")); + } + + // Copy back if out was not contiguous + if (!p->out->is_contiguous()) { + p->out->copy_from(out_work); + } +} + +void cleanup(void **planned_meta_ptr) { + auto *p = *reinterpret_cast(planned_meta_ptr); + if (p->mask_dev) + aclrtFree(p->mask_dev); + if (p->mask_acl) + aclDestroyTensor(p->mask_acl); + delete p; + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MultiheadAttentionVarlen::plan_dispatcher().registerDevice( + Device::Type::ASCEND, &plan); + MultiheadAttentionVarlen::run_dispatcher().registerDevice( + Device::Type::ASCEND, &run); + MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice( + Device::Type::ASCEND, &cleanup); + return true; +}(); + +} // namespace infinicore::op::mha_varlen_impl::flashattn_ascend +#endif diff --git a/src/infiniop/ops/flash_attention/operator.cc b/src/infiniop/ops/flash_attention/operator.cc index 449c17adc..90ebf942a 100644 --- a/src/infiniop/ops/flash_attention/operator.cc +++ b/src/infiniop/ops/flash_attention/operator.cc @@ -38,6 +38,9 @@ __INFINI_C infiniStatus_t infiniopCreateFlashAttentionDescriptor( #if defined(ENABLE_NVIDIA_API) CREATE(INFINI_DEVICE_NVIDIA, ninetoothed); #endif +#endif +#if defined(ENABLE_ASCEND) + CREATE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/xmake.lua b/xmake.lua index fcbcf423d..49d271688 100644 --- a/xmake.lua +++ b/xmake.lua @@ -251,6 +251,11 @@ if has_config("aten") then and (has_config("nv-gpu") or has_config("metax-gpu") or has_config("qy-gpu") or has_config("hygon-dcu")) then add_defines("ENABLE_FLASH_ATTN") end + -- Ascend flash-attn: enabled when ascend-gpu + aten + flash-attn + if has_config("ascend-npu") and get_config("flash-attn") and get_config("flash-attn") ~= "" then + add_defines("ENABLE_FLASH_ATTN") + add_defines("ENABLE_ASCEND_FLASH_ATTN") + end end -- cuda graph @@ -673,6 +678,16 @@ target("infinicore_cpp_api") local cuda_root = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or get_config("cuda") or "/usr/local/cuda" add_includedirs(cuda_root .. "/include") end + if has_config("ascend-npu") then + local ASCEND_HOME = os.getenv("ASCEND_HOME") or os.getenv("ASCEND_TOOLKIT_HOME") + add_includedirs(ASCEND_HOME .. "/include") + add_includedirs(ASCEND_HOME .. "/include/aclnn") + add_includedirs(ASCEND_HOME .. "/include/aclnnop") + add_linkdirs(ASCEND_HOME .. "/lib64") + add_links("ascendcl", "nnopbase", "opapi", "runtime") + add_linkdirs(ASCEND_HOME .. "/../../driver/lib64/driver") + add_links("ascend_hal") + end if has_config("infiniops") then local infiniops_root = path.absolute(get_config("infiniops-root") or "submodules/InfiniOps", os.projectdir()) if not os.isdir(infiniops_root) then @@ -842,6 +857,33 @@ target("infinicore_cpp_api") "-Wl,-rpath," .. pylib, { force = true } ) + elseif has_config("ascend-npu") then + target:add( + "links", + "torch", + "torch_cpu", + "c10", + { public = true } + ) + -- Detect torch_npu install path + local npu_outdata = os.iorunv("python", {"-c", "import torch_npu, os; print(os.path.dirname(torch_npu.__file__))"}):trim() + local TORCH_NPU_DIR = npu_outdata + target:add( + "linkdirs", + path.join(TORCH_NPU_DIR, "lib"), + { public = true } + ) + target:add( + "links", + "torch_npu", + { public = true } + ) + target:add( + "shflags", + "-Wl,-rpath," .. path.join(TORCH_NPU_DIR, "lib"), + "-Wl,-rpath," .. path.join(TORCH_DIR, "lib"), + { force = true } + ) elseif has_config("hygon-dcu") then target:add( "links", diff --git a/xmake/ascend.lua b/xmake/ascend.lua index 6179cac13..c1b519fa0 100644 --- a/xmake/ascend.lua +++ b/xmake/ascend.lua @@ -5,6 +5,7 @@ local SOC_VERSION = os.getenv("SOC_VERSION") -- Add include dirs add_includedirs(ASCEND_HOME .. "/include") add_includedirs(ASCEND_HOME .. "/include/aclnn") +add_includedirs(ASCEND_HOME .. "/include/aclnnop") add_linkdirs(ASCEND_HOME .. "/lib64") add_links("libascendcl.so") add_links("libnnopbase.so") From f49faedbbaaaa18dd31790315d3ba502751ab216 Mon Sep 17 00:00:00 2001 From: Jingbo-gao Date: Fri, 17 Jul 2026 10:31:27 +0800 Subject: [PATCH 2/3] change format of operator.cc --- src/infiniop/ops/flash_attention/operator.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/infiniop/ops/flash_attention/operator.cc b/src/infiniop/ops/flash_attention/operator.cc index 90ebf942a..9867d8ef2 100644 --- a/src/infiniop/ops/flash_attention/operator.cc +++ b/src/infiniop/ops/flash_attention/operator.cc @@ -40,7 +40,7 @@ __INFINI_C infiniStatus_t infiniopCreateFlashAttentionDescriptor( #endif #endif #if defined(ENABLE_ASCEND) - CREATE(INFINI_DEVICE_ASCEND, ascend); + CREATE(INFINI_DEVICE_ASCEND, ascend); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; From f9e592a7c7d30d97f97090d0b46016644f2cb5ae Mon Sep 17 00:00:00 2001 From: Jingbo-gao Date: Fri, 17 Jul 2026 14:22:54 +0800 Subject: [PATCH 3/3] fix pre-check error --- python/infinicore/__init__.py | 8 ++++---- python/infinicore/device.py | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index cba4aed1e..a74dad116 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -56,6 +56,10 @@ from infinicore.ops.addr import addr from infinicore.ops.all import all from infinicore.ops.argwhere import argwhere +from infinicore.ops.ascend_flash_attn import ( + ascend_flash_attn_decode, + ascend_flash_attn_prefill, +) from infinicore.ops.asin import asin from infinicore.ops.asinh import asinh from infinicore.ops.asum import asum @@ -113,10 +117,6 @@ moore_mate_flash_attn_prefill, ) from infinicore.ops.mrope import mrope -from infinicore.ops.ascend_flash_attn import ( - ascend_flash_attn_decode, - ascend_flash_attn_prefill, -) from infinicore.ops.mul import mul from infinicore.ops.mul_scalar import mul_scalar from infinicore.ops.narrow import narrow diff --git a/python/infinicore/device.py b/python/infinicore/device.py index 3607da183..8858a1235 100644 --- a/python/infinicore/device.py +++ b/python/infinicore/device.py @@ -1,4 +1,3 @@ -import torch from infinicore.lib import _infinicore _infinicore_2_python_dict = {}