From 616e3e737e02033369cfe059fc149a0f67aee55a Mon Sep 17 00:00:00 2001 From: PanZezhong Date: Wed, 15 Jul 2026 08:19:58 +0800 Subject: [PATCH] feat(metax) support mha op --- .../infinicore/adaptor/flash_attention_adaptor.hpp | 11 +++++++---- .../ops/multi_head_attention/mha_flashattn.cc | 12 ++++++++---- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/include/infinicore/adaptor/flash_attention_adaptor.hpp b/include/infinicore/adaptor/flash_attention_adaptor.hpp index c5bf14858..2c195cf0f 100644 --- a/include/infinicore/adaptor/flash_attention_adaptor.hpp +++ b/include/infinicore/adaptor/flash_attention_adaptor.hpp @@ -9,10 +9,13 @@ namespace flash { #endif std::vector -mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) - const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) - const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) - std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) +mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) + const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) + const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) + std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) +#if defined(ENABLE_METAX_API) + std::optional &softmax_lse_, // MetaX flash-attn dense fwd ABI includes an optional preallocated LSE tensor +#endif std::optional &alibi_slopes_, // num_heads or batch_size x num_heads const float p_dropout, const float softmax_scale, diff --git a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc index 13c96b94d..5bd0a9ac7 100644 --- a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc @@ -39,10 +39,9 @@ void *plan(Tensor out, namespace { -// Only support nv for now -#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_NVIDIA_API) // MetaX/hpcc pip `flash_attn_2_cuda` exports `mha_fwd` at global scope (no namespace), // while NVIDIA `flash-attn-nvidia.so` uses `flash::mha_fwd`. +#if defined(ENABLE_FLASH_ATTN) && (defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API)) #if defined(ENABLE_METAX_API) #define INFINICORE_FLASH_OP(name) ::name #else @@ -53,8 +52,7 @@ namespace { } // namespace void run(void *planned_meta) { -// Only support nv for now -#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_NVIDIA_API) +#if defined(ENABLE_FLASH_ATTN) && (defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API)) c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); auto *p = reinterpret_cast(planned_meta); @@ -67,6 +65,9 @@ void run(void *planned_meta) { auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); auto out = std::optional(out_work); +#if defined(ENABLE_METAX_API) + std::optional softmax_lse = std::nullopt; +#endif auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; auto scale = p->scale; auto is_causal = p->is_causal; @@ -81,6 +82,9 @@ void run(void *planned_meta) { k, v, out, +#if defined(ENABLE_METAX_API) + softmax_lse, +#endif alibi_slopes, 0.0, scale,