Skip to content

fix(metax): detect Flash Attention varlen ABI from linked library - #1558

Open
big-hip wants to merge 3 commits into
InfiniTensor:InfiniLM-v0.2.9cfrom
big-hip:fix/metax-flash-attention-abi
Open

big-hip wants to merge 3 commits into
InfiniTensor:InfiniLM-v0.2.9cfrom
big-hip:fix/metax-flash-attention-abi

Conversation

@big-hip

@big-hip big-hip commented Sep 16, 2026

Copy link
Copy Markdown

Problem and change

On MetaX Flash Attention wheels that append return_max_logit to mha_varlen_fwd, InfiniCore builds a caller with the wrong C++ signature and fails to import with an undefined symbol. This change inspects the extension that will actually be linked and selects the matching declaration and call, passing false only when the argument exists.

  • Recognizes the base signature, the optional-tensor extension, and the optional tensor plus trailing bool. Missing, unknown or ambiguous signatures fail during the build with an actionable diagnostic.
  • Uses the same extension resolver for ABI detection and linking, with an explicit FLASH_ATTN_2_CUDA_SO override. Other entry points retain their existing version handling; NVIDIA and other device adapters are unchanged.
  • Adds detection fixtures and concise configuration documentation. Experiment evidence is linked from this PR rather than included in the source diff. Target branch: InfiniLM-v0.2.9c.

Technical benefit and measured motivation

Direct before/after: affected-wheel import fails -> import, Prefill, Decode and model inference succeed. This restores access to the existing MetaX Flash Attention backend. There is no latency ratio against an import failure.

The following archived experiment quantifies the practical value of that backend. Both arms use the same ABI-compatible prototype library, changing only attn_backend from paged-attn to flash-attn. These are backend-selection benefits, not kernel acceleration implemented by this ABI patch, and not a performance rerun of the final generalized resolver.

Model Native -> vendor TTFT Reduction Native -> vendor ITL Reduction
Qwen3-0.6B 472.52 -> 78.80 ms 83.32% 17.13 -> 6.88 ms 59.85%
Qwen3-4B 1,212.21 -> 212.09 ms 82.50% 26.00 -> 12.87 ms 50.50%

Conditions: C500 50% compute / 32,000 MiB slice, six CPU cores; MACA3.5.3.20, driver3.8.30, torch2.8.0+metax3.5.3.9, flash-attn2.6.3+metax3.5.3.9torch2.8. BF16 model/KV, TP1/PP1, dense, eager, page256/pool32, chunking and prefix reuse off. 2048 input / 16 output tokens, greedy with EOS ignored for timing; two measurements after warmup per model/backend. TTFT is their mean; ITL is the mean of request medians. All eight measured outputs match HF eager and paired-backend controls. Full-library hashes match within each comparison.

Finite single-request output rates including Prefill were 21.93 -> 87.92 and 9.99 -> 39.46 token/s; these are not sustained serving throughput. Only two samples on a shared slice and two same-family models were measured. Gains come from the vendor attention implementation; InfiniLM still owns scheduling and KV-cache management.

Detailed conditions, attribution and reproduction · per-request measurements and source/binary hashes.

Validation

The result JSON, screenshot and detailed experiment report are preserved at archive revision 9ecc781f; the final diff contains only the ABI change, build resolver, fixture test and short configuration guide. Hardware measurements were not repeated for this cleanup.

  • xmake lua tests/xmake/test_metax_flash_abi.lua: three accepted signatures and three negative cases passed; also checked against the installed C500 extension using FLASH_ATTN_2_CUDA_SO.
  • Final adapter build/import/execute passed: six Prefill + six Decode FP16/BF16 cases against explicit FP32 attention, including GQA, HD64/128, permuted physical pages, historical KV and causal masking.
  • The resulting library passed archived Qwen3-0.6B/4B cache/chunk/graph checks. The companion InfiniLM #573 now scopes chunking to eager Prefill + existing Decode graphs; experimental Prefill-graph cases remain historical evidence outside that submission. This ABI fix does not implement graph compilation.
  • Fork Build and test: Four Windows/Linux matrix jobs passed, including CPU tests. The existing matrix labels jobs debug/release but does not forward matrix.type to the install command, so this is not evidence of distinct debug and release builds. Validation branch has the same implementation as f02ec327, with only an added workflow trigger; it is not a MetaX hardware CI result. Ruff passed.
  • Repository formatting and git diff --check passed. The scope cleanup only removes result attachments and trims the guide; the ABI implementation and fixtures are unchanged. The portable fixture was rerun successfully.

Screenshot of the audited saved-results report (not a new GPU run):

Audited MetaX compatibility and backend results

Limits and remaining review gates

Older signatures have fixture coverage only; execution against older installed wheels remains untested. The MetaX hardware build reused unchanged objects; the successful clean CPU CI does not establish a clean accelerator build. Backend timings preceded the generalized signature resolver and are presented as motivation; no new-kernel speedup is claimed. No Ascend/Moore or unrelated device-support claim is added.

The upstream workflow reports action_required with zero jobs started and needs repository-side approval. Fork CI success does not replace upstream required checks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant