Skip to content

feat(linked): add MetaX FlashAttention providers - #934

Draft
voltjia wants to merge 1 commit into
masterfrom
feat/metax-linked-flash-attn
Draft

feat(linked): add MetaX FlashAttention providers#934
voltjia wants to merge 1 commit into
masterfrom
feat/metax-linked-flash-attn

Conversation

@voltjia

@voltjia voltjia commented Aug 18, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add MetaX linked providers for the canonical flash_attn_varlen_func and flash_attn_with_kvcache operators using the installed flash-attn shared library.
  • Register both providers in linked implementation slot 16 so builds can select them through the existing ops.json mechanism.

Motivation

InfiniCore's basic Llama/Qwen prefill and decode paths should call the canonical InfiniOps FlashAttention interfaces on MetaX. The installed MetaX flash-attn package exports an ABI with one additional optional extension tensor, so it needs a platform-specific linked adapter while preserving the canonical public operator contract.

Type of Change

  • feat - new feature / new operator / new platform
  • fix - bug fix
  • perf - performance improvement (no behavioral change)
  • refactor - code restructuring without behavior change
  • test - adding or fixing tests only
  • docs - documentation only
  • build / ci - build system or CI configuration
  • chore - tooling, formatting, or other non-code changes
  • Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

MetaX, clean build using an ops.json selection:
-- Operator selection: /tmp/codex-ops-json-selection-20260819/InfiniOps/build/ops.json
-- Resolved 2 linked operator source(s).
-- Codegen torch op allowlist: argmax;index_select
[100%] Built target infiniops

The generated linked manifest contains only:

src/linked/torch/metax/ops/flash_attn_varlen_func/flash_attn.cc
src/linked/torch/metax/ops/flash_attn_with_kvcache/flash_attn.cc

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA No N/A - not affected
Iluvatar No N/A - not affected
MetaX Yes full build passed Qwen3-0.6B prefill and decode passed through the dependent InfiniCore branch
Cambricon No N/A - not affected
Moore No N/A - not affected
Ascend No N/A - not affected

Additional checks:

  • python3 -m pytest tests/test_ops_config.py tests/test_generate_wrappers.py tests/test_resolve_linked_ops.py (84 passed)
  • clang-format 21.1.8 --dry-run --Werror on all new C++ provider files
  • git diff --check origin/master...HEAD

Benchmark / Performance Impact

N/A. This PR adds a MetaX adapter for an existing external implementation.

Notes for Reviewers

  • Validated against flash-attn 2.6.3+metax3.2.1.3torch2.4. The MetaX ABI adds a trailing std::optional<at::Tensor>&; the adapter owns that extension argument and delegates through the canonical InfiniOps interface.
  • The dependent build selects slot 16 for these two operators through ops.json, using the operator implementation selection added by feat(build): select operator implementations from ops.json #931. This PR adds no separate linked-provider selector or environment variable.

@voltjia
voltjia force-pushed the feat/metax-linked-flash-attn branch from 8c5389a to fa3f8ea Compare August 19, 2026 03:39
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