Skip to content

fix(metax): support flash-attn 2.6.3 forward ABI #1555

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zxyy-bys wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(metax): support flash-attn 2.6.3 forward ABI #1555

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
446a85b
fix(metax): support flash-attn 2.6.3 forward ABI
zxyy-bys Sep 8, 2026
010ccd3
feat(metax): select flash-attn ABI from wheel symbols
zxyy-bys Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
189 changes: 122 additions & 67 deletions include/infinicore/adaptor/flash_attention_adaptor.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,46 @@
#pragma once
#include "aten_adaptor.hpp"

#if defined(ENABLE_METAX_API)

// MetaX flash-attn wheels declare their API in terms of c10::optional. On torch >= 2.1 that
// is an alias of std::optional (identical mangling); spelling the declarations the same way
// as the wheel guarantees identical mangling even on older torch stacks. When the c10 header
// no longer exists, torch has fully migrated to std::optional -- and so have wheels built
// against it.
#if __has_include(<c10/util/Optional.h>)
#include <c10/util/Optional.h>
#define INFINICORE_FA_OPTIONAL c10::optional
#else
#define INFINICORE_FA_OPTIONAL std::optional
#endif

// MetaX flash-attn ships two incompatible forward ABIs:
// INFINICORE_METAX_FA_ABI 253 -- flash_attn 2.5.3 (MACA/HPCC 2.x):
// mha_fwd / mha_varlen_fwd / mha_fwd_kvcache take 13 / 18 / 18 arguments
// INFINICORE_METAX_FA_ABI 263 -- flash_attn 2.6.3+metax (MACA/HPCC 3.x):
// the same functions take 16 / 23 / 21 arguments (softcap, leftpad_k, varlen
// block_table, s_aux and return_max_logit are appended)
// xmake (xmake/metax.lua) injects INFINICORE_METAX_FA_ABI after inspecting the dynamic
// symbols of the actual `flash_attn_2_cuda` wheel that will be linked. When it is not
// injected (e.g. the wheel is absent at configure time), fall back to the HPCC/MACA
// toolkit major-version probe: >= 3 ships the 2.6.3 ABI, otherwise the legacy 2.5.3 ABI.
#if !defined(INFINICORE_METAX_FA_ABI)
#if defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
#define INFINICORE_METAX_FA_ABI 263
#else
#define INFINICORE_METAX_FA_ABI 253
#endif
#endif
#define INFINICORE_METAX_FA263 (INFINICORE_METAX_FA_ABI >= 263)

#else // !ENABLE_METAX_API

#define INFINICORE_FA_OPTIONAL std::optional
#define INFINICORE_METAX_FA263 1

#endif // ENABLE_METAX_API

// 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.
Expand All @@ -14,40 +54,46 @@
namespace flash {
#endif
std::vector<at::Tensor>
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<at::Tensor> &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)
INFINICORE_FA_OPTIONAL<at::Tensor> &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
INFINICORE_FA_OPTIONAL<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
#if defined(ENABLE_METAX_API)
std::optional<at::Tensor> &softmax_lse_, // MetaX flash-attn dense fwd ABI includes an optional preallocated LSE tensor
INFINICORE_FA_OPTIONAL<at::Tensor> &attn_mask_,
#endif
std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
const float p_dropout,
const float softmax_scale,
bool is_causal,
int window_size_left,
int window_size_right,
#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263
const float softcap,
#endif
const bool return_softmax,
std::optional<at::Generator> gen_
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
// MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
INFINICORE_FA_OPTIONAL<at::Generator> gen_
#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263
// MetaX `flash_attn_2_cuda` 2.6.3+ (MACA/HPCC 3.x) appends these arguments vs the
// 2.5.3 wheel and upstream Dao-AILab flash-attn.
,
std::optional<at::Tensor> &flash_attn_mars_ext_
INFINICORE_FA_OPTIONAL<at::Tensor> &s_aux_,
bool return_max_logit_
#endif
);

std::vector<at::Tensor>
mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
std::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &cu_seqlens_q, // b+1
const at::Tensor &cu_seqlens_k, // b+1
std::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
std::optional<const at::Tensor> &leftpad_k_, // batch_size
std::optional<at::Tensor> &block_table_, // batch_size x max_num_blocks_per_seq
std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
INFINICORE_FA_OPTIONAL<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &cu_seqlens_q, // b+1
const at::Tensor &cu_seqlens_k, // b+1
INFINICORE_FA_OPTIONAL<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263
INFINICORE_FA_OPTIONAL<const at::Tensor> &leftpad_k_, // batch_size
INFINICORE_FA_OPTIONAL<at::Tensor> &block_table_, // batch_size x max_num_blocks_per_seq
#endif
INFINICORE_FA_OPTIONAL<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
int max_seqlen_q,
const int max_seqlen_k,
const float p_dropout,
Expand All @@ -56,50 +102,54 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_hea
bool is_causal,
int window_size_left,
int window_size_right,
#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263
const float softcap,
#endif
const bool return_softmax,
std::optional<at::Generator> gen_
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
// MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
INFINICORE_FA_OPTIONAL<at::Generator> gen_
#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263
// MetaX `flash_attn_2_cuda` 2.6.3+ (MACA/HPCC 3.x) appends these arguments vs the
// 2.5.3 wheel and upstream Dao-AILab flash-attn.
,
std::optional<at::Tensor> &flash_attn_mars_ext_
INFINICORE_FA_OPTIONAL<at::Tensor> &s_aux_,
bool return_max_logit_
#endif
);

std::vector<at::Tensor>
mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
const at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size
const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size
const at::Tensor &out, // batch_size x seqlen_q x num_heads x head_size
const at::Tensor &softmax_lse, // b x h x seqlen_q
std::optional<at::Tensor> &dq_, // batch_size x seqlen_q x num_heads x head_size
std::optional<at::Tensor> &dk_, // batch_size x seqlen_k x num_heads_k x head_size
std::optional<at::Tensor> &dv_, // batch_size x seqlen_k x num_heads_k x head_size
std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
const float p_dropout, // probability to drop
mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
const at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size
const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size
const at::Tensor &out, // batch_size x seqlen_q x num_heads x head_size
const at::Tensor &softmax_lse, // b x h x seqlen_q
INFINICORE_FA_OPTIONAL<at::Tensor> &dq_, // batch_size x seqlen_q x num_heads x head_size
INFINICORE_FA_OPTIONAL<at::Tensor> &dk_, // batch_size x seqlen_k x num_heads_k x head_size
INFINICORE_FA_OPTIONAL<at::Tensor> &dv_, // batch_size x seqlen_k x num_heads_k x head_size
INFINICORE_FA_OPTIONAL<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
const float p_dropout, // probability to drop
const float softmax_scale,
const bool is_causal,
int window_size_left,
int window_size_right,
const float softcap,
const bool deterministic,
std::optional<at::Generator> gen_,
std::optional<at::Tensor> &rng_state);
INFINICORE_FA_OPTIONAL<at::Generator> gen_,
INFINICORE_FA_OPTIONAL<at::Tensor> &rng_state);

std::vector<at::Tensor>
mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size
const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &out, // total_q x num_heads x head_size
const at::Tensor &softmax_lse, // h x total_q, softmax logsumexp
std::optional<at::Tensor> &dq_, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
std::optional<at::Tensor> &dk_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
std::optional<at::Tensor> &dv_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &cu_seqlens_q, // b+1
const at::Tensor &cu_seqlens_k, // b+1
std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size
const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &out, // total_q x num_heads x head_size
const at::Tensor &softmax_lse, // h x total_q, softmax logsumexp
INFINICORE_FA_OPTIONAL<at::Tensor> &dq_, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
INFINICORE_FA_OPTIONAL<at::Tensor> &dk_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
INFINICORE_FA_OPTIONAL<at::Tensor> &dv_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &cu_seqlens_q, // b+1
const at::Tensor &cu_seqlens_k, // b+1
INFINICORE_FA_OPTIONAL<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
const int max_seqlen_q,
const int max_seqlen_k, // max sequence length to choose the kernel
const float p_dropout, // probability to drop
Expand All @@ -110,34 +160,39 @@ mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads,
int window_size_right,
const float softcap,
const bool deterministic,
std::optional<at::Generator> gen_,
std::optional<at::Tensor> &rng_state);
INFINICORE_FA_OPTIONAL<at::Generator> gen_,
INFINICORE_FA_OPTIONAL<at::Tensor> &rng_state);

std::vector<at::Tensor>
mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
const at::Tensor &kcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
const at::Tensor &vcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
std::optional<const at::Tensor> &k_, // batch_size x seqlen_knew x num_heads_k x head_size
std::optional<const at::Tensor> &v_, // batch_size x seqlen_knew x num_heads_k x head_size
std::optional<const at::Tensor> &seqlens_k_, // batch_size
std::optional<const at::Tensor> &rotary_cos_, // seqlen_ro x (rotary_dim / 2)
std::optional<const at::Tensor> &rotary_sin_, // seqlen_ro x (rotary_dim / 2)
std::optional<const at::Tensor> &cache_batch_idx_, // indices to index into the KV cache
std::optional<const at::Tensor> &leftpad_k_, // batch_size
std::optional<at::Tensor> &block_table_, // batch_size x max_num_blocks_per_seq
std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
std::optional<at::Tensor> &out_, // batch_size x seqlen_q x num_heads x head_size
mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
const at::Tensor &kcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
const at::Tensor &vcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
INFINICORE_FA_OPTIONAL<const at::Tensor> &k_, // batch_size x seqlen_knew x num_heads_k x head_size
INFINICORE_FA_OPTIONAL<const at::Tensor> &v_, // batch_size x seqlen_knew x num_heads_k x head_size
INFINICORE_FA_OPTIONAL<const at::Tensor> &seqlens_k_, // batch_size
INFINICORE_FA_OPTIONAL<const at::Tensor> &rotary_cos_, // seqlen_ro x (rotary_dim / 2)
INFINICORE_FA_OPTIONAL<const at::Tensor> &rotary_sin_, // seqlen_ro x (rotary_dim / 2)
INFINICORE_FA_OPTIONAL<const at::Tensor> &cache_batch_idx_, // indices to index into the KV cache
#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263
INFINICORE_FA_OPTIONAL<const at::Tensor> &leftpad_k_, // batch_size
#endif
INFINICORE_FA_OPTIONAL<at::Tensor> &block_table_, // batch_size x max_num_blocks_per_seq
INFINICORE_FA_OPTIONAL<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
INFINICORE_FA_OPTIONAL<at::Tensor> &out_, // batch_size x seqlen_q x num_heads x head_size
const float softmax_scale,
bool is_causal,
int window_size_left,
int window_size_right,
#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263
const float softcap,
#endif
bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2
int num_splits
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
// MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263
// MetaX `flash_attn_2_cuda` 2.6.3+ (MACA/HPCC 3.x) appends this argument vs the
// 2.5.3 wheel and upstream Dao-AILab flash-attn.
,
std::optional<at::Tensor> &flash_attn_mars_ext_
INFINICORE_FA_OPTIONAL<at::Tensor> &s_aux_
#endif
);

Expand Down
14 changes: 10 additions & 4 deletions src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,9 @@ void run(void *planned_meta) {
std::optional<const at::Tensor> rotary_cos = std::nullopt;
std::optional<const at::Tensor> rotary_sin = std::nullopt;
std::optional<const at::Tensor> cache_batch_idx = std::nullopt;
#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263
std::optional<const at::Tensor> leftpad_k = std::nullopt;
#endif

const bool use_dynamic_out = q.dim() == 4 && k_cache.dim() == 4
&& q.size(1) == 1 && q.size(2) > k_cache.size(2)
Expand All @@ -233,8 +235,8 @@ void run(void *planned_meta) {
auto out = use_dynamic_out ? std::optional<at::Tensor>(std::nullopt)
: std::optional<at::Tensor>(out_tensor);

#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
std::optional<at::Tensor> flash_attn_mars_ext = std::nullopt;
#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263
std::optional<at::Tensor> s_aux = std::nullopt;
#endif

auto result = INFINICORE_FLASH_OP(mha_fwd_kvcache)(
Expand All @@ -247,20 +249,24 @@ void run(void *planned_meta) {
rotary_cos,
rotary_sin,
cache_batch_idx,
#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263
leftpad_k,
#endif
block_table,
alibi_slopes,
out,
p->scale,
true,
-1,
-1,
#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263
0.0f,
#endif
false,
0
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263
,
flash_attn_mars_ext
s_aux
#endif
);

Expand Down
Loading
Loading