Conversation
3ec7f15 to
a01e579
Compare
a01e579 to
446a85b
Compare
MetaX flash-attn ships two incompatible forward ABIs: flash_attn 2.5.3 wheels (MACA/HPCC 2.x, mha_fwd/mha_varlen_fwd/mha_fwd_kvcache take 13/18/18 args) and flash_attn 2.6.3+metax wheels (MACA/HPCC 3.x, the same functions take 16/23/21 args with softcap, leftpad_k, varlen block_table, s_aux and return_max_logit appended). Select the ABI at configure time from the wheel that will actually be linked: xmake/metax.lua inspects its demangled dynamic symbols in on_load (same approach as Cambricon in xmake/bang.lua) and injects INFINICORE_METAX_FA_ABI=253|263 as a public define, so it also reaches infinicore-test and removes the previous ODR hazard between the two targets. The HPCC/MACA Version.txt probe moves from xmake.lua into the same hook as a fallback when the wheel cannot be inspected, and a new --metax-fa-abi=auto|253|263 option allows overriding the result. The shared header derives a single INFINICORE_METAX_FA263 gate macro from the injected ABI (falling back to the HPCC major version when it is absent), replacing the scattered INFINICORE_HPCC_VERSION_MAJOR conditions in the declarations and the three flashattn call sites. Optional parameters are respelled through INFINICORE_FA_OPTIONAL (c10::optional under MetaX when <c10/util/Optional.h> exists, std::optional otherwise) so declarations mangle exactly like the wheel on any torch version; non-MetaX platforms expand to std::optional and are unaffected (verified by symbol comparison against origin/main). Also refuse paged KV explicitly on the 2.5.3 varlen path (no block_table parameter) instead of silently dropping it, and scope the kvcache leftpad_k local to the gate that uses it. Co-Authored-By: Claude Code <noreply@anthropic.com>
|
We encountered a related Flash Attention ABI mismatch while validating InfiniLM #571 on a MetaX server with MACA 3.1. The installed The mismatch was observed for all three forward entry points:
Therefore, this appears to overlap with this PR more closely than #1558, since #1558 currently detects the ABI only for One concern is that selecting a single ABI profile for the whole wheel, or falling back to the MACA toolkit version, may not cover every deployed wheel. In our environment, the MACA major version alone did not reliably identify the exported signatures. Would it make sense to combine the approaches of this PR and #1558:
We can provide the demangled symbols and validate a generalized implementation on our MACA 3.1 environment if that would be useful. Related: #1558 |
2 participants
Problem
The MetaX MACA 3.x FlashAttention 2.6.3 extension exports a forward ABI that differs from the legacy declaration used by InfiniCore. The deployed flash_attn_2_cuda.fwd expects out, alibi_slopes, attn_mask, softcap, s_aux, and return_max_logit in its argument sequence, while the old InfiniCore path passed a preallocated softmax-LSE tensor and omitted the newer arguments. This shifts the positional arguments at the extension boundary and can cause FlashAttention runtime failures during end-to-end inference, including multimodal request processing.
Older MetaX FlashAttention 2.20.2 has a different legacy layout as well: dense forward includes attn_mask but no softcap or FA2 tail; varlen and KV-cache omit newer leftpad, softcap, and auxiliary arguments. Without version-specific guards, those arguments are also shifted or passed to the wrong ABI.
Summary
Scope
Validation
The MetaX 3.x/flash-attn 2.6.3 path was exercised end to end. The 2.20.2 layout was verified against the public source branch; no legacy MetaX GPU was available for runtime execution.