From f8679ea325aafa728c5fd6253b59e67210d48f77 Mon Sep 17 00:00:00 2001 From: Jingbo-gao Date: Thu, 16 Jul 2026 21:23:52 +0800 Subject: [PATCH 1/7] 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/7] 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/7] 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 = {} From 9f81b521c8d4c0a8915c0d3a5c6e871e47a55413 Mon Sep 17 00:00:00 2001 From: Jingbo-gao Date: Mon, 20 Jul 2026 14:58:56 +0800 Subject: [PATCH 4/7] fix format of mha_varlen_flashattn_ascend.cc --- .../mha_varlen_flashattn_ascend.cc | 574 +++++++++--------- 1 file changed, 286 insertions(+), 288 deletions(-) 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 index 8bd51839b..0ecad4ddc 100644 --- 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 @@ -3,7 +3,6 @@ #include "infinicore/context/context.hpp" #include "infinicore/ops/mha_varlen.hpp" - #include #include #include @@ -16,328 +15,327 @@ 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"); - } + 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()); + 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; + 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; - } + if (p->mask_acl != nullptr) { + return p->mask_acl; + } - int64_t mask_size = 2048 * 2048; + 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; - } + 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; + + 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}; + 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) { + 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); + 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)); + + 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); } - } - - 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; + 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; + 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 From d70ad6c9e0af86441d9b64c39be66a36d40021f0 Mon Sep 17 00:00:00 2001 From: Jingbo-gao Date: Mon, 20 Jul 2026 15:08:06 +0800 Subject: [PATCH 5/7] fix format of mha_kvcache_flashattn_ascend.cc --- .../mha_kvcache_flashattn_ascend.cc | 512 +++++++++--------- 1 file changed, 252 insertions(+), 260 deletions(-) diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc index da9273cf0..4c396a837 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc @@ -3,7 +3,6 @@ #include "infinicore/context/context.hpp" #include "infinicore/ops/mha_kvcache.hpp" - #include #include #include @@ -16,219 +15,256 @@ 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"); - } + 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()); + 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; + 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}; + 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() + 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) { + 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); @@ -236,77 +272,33 @@ void run(void *planned_meta) { 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)); + + 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); } - } - - 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; + 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; + 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 From dcd51a9827fd3543fd7e83e4f23faf9567da62ad Mon Sep 17 00:00:00 2001 From: Jingbo-gao Date: Mon, 20 Jul 2026 15:51:09 +0800 Subject: [PATCH 6/7] fix format of mha_kvcache_flashattn_ascend.cc --- src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc index 4c396a837..cb98e9ccd 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_ascend.cc @@ -54,7 +54,7 @@ void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, graph::GraphTensor(seqlens_k), graph::GraphTensor(block_table), alibi_slopes ? std::optional( - graph::GraphTensor(*alibi_slopes)) + graph::GraphTensor(*alibi_slopes)) : std::nullopt, scale}; } From e97fb4fbfafd5f34ce0f726773c4b614a41fea01 Mon Sep 17 00:00:00 2001 From: Jingbo-gao Date: Mon, 20 Jul 2026 15:58:07 +0800 Subject: [PATCH 7/7] fix format of ascend_flash_attn.py --- python/infinicore/ops/ascend_flash_attn.py | 24 ++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/python/infinicore/ops/ascend_flash_attn.py b/python/infinicore/ops/ascend_flash_attn.py index 01379c866..ec960039d 100644 --- a/python/infinicore/ops/ascend_flash_attn.py +++ b/python/infinicore/ops/ascend_flash_attn.py @@ -159,12 +159,20 @@ def ascend_flash_attn_prefill( 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, + 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, + 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 @@ -189,7 +197,9 @@ def ascend_flash_attn_prefill( 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) + 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, @@ -204,5 +214,7 @@ def ascend_flash_attn_prefill( causal=True, ) torch.npu.synchronize() - assert list(out_prefill.shape) == list(q_prefill.shape), f"Expected {q_prefill.shape}, got {out_prefill.shape}" + 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")