Skip to content

feat(moore): call MATE flash attention natively - #987

Merged
voltjia merged 7 commits into
masterfrom
feat/moore-flash-attn-native
Oct 14, 2026
Merged

voltjia merged 7 commits into
masterfrom
feat/moore-flash-attn-native

Conversation

@voltjia

@voltjia voltjia commented Sep 15, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Replace the per-call Python/MATE bridge from feat(moore): support flash-attn via MooreThreads/mate v0.2.5 #983 with a TVM-FFI native execution path.
  • Use MATE only once per operator to resolve/compile the selected FMHA and combine modules, then cache their loaded DSOs and __tvm_ffi_* entry points.
  • Write FlashAttention output and LSE directly into InfiniOps-owned tensors, removing provider allocations and D2D result copies.
  • Route the Moore MATE glue through the existing system-C++ Torch object path so the vendor MUSA compiler does not parse ATen headers.

Validation

  • Moore minimal linked build.
  • Existing FlashAttention varlen/KV-cache correctness suites.
  • Repeat-call check confirming Python is bootstrap-only.
  • Non-default stream and split-KV combine coverage.

This is a draft while Moore hardware validation is in progress.

@voltjia
voltjia marked this pull request as ready for review September 16, 2026 01:32
@voltjia
voltjia requested review from a team and spike-zhu September 16, 2026 01:32
@voltjia

voltjia commented Sep 24, 2026

Copy link
Copy Markdown
Collaborator Author

Fix-up for the Moore build/runtime failures in \19c6490d:\n\n- Load TVM-FFI's DLPack header before ATen's vendored DLPack header, so the TVM-FFI ABI sees \DLManagedTensorVersioned\ and the current \DLTensor\ layout.\n- Keep ATen/MATE runtime state in each .cc\ via a forward-declared \MateState; generated MUSA dispatch units therefore only consume a lightweight provider declaration.\n- Use PyTorch's official \�t::Tensor\ pybind caster, mark the MATE loader lambdas mutable, and convert external softmax-LSE views before handing them to MATE.\n\nValidated on the new Moore S5000 host using Torch 2.7.1, MATE 0.2.5, and \�pache_tvm_ffi 0.1.9.post3+musa.1. I first built current InfiniRT master with \WITH_CPU=ON -DWITH_MOORE=ON, then built this PR with \WITH_LINKED=ON\ and the two FlashAttention operators selected. Compilation and linking passed.\n\nExisting tests, no new smoke tests: \pytest tests/test_flash_attn_varlen_func.py tests/test_flash_attn_with_kvcache.py --devices moore -k 'musa-16 and (q_lens0 or dense or paged)'\ -> 13 passed, 5 skipped, 94 deselected. The skipped cases are existing capability-specific skips.

@spike-zhu

Copy link
Copy Markdown
Contributor

Approved. This PR limits Python/MATE to first-call bootstrap, caches the selected TVM-FFI DSO entry points, and keeps subsequent FlashAttention calls on the native path while preserving zero-copy output and LSE handling. The latest state isolation and DLPack/ATen boundary fixes also make the implementation much cleaner.

spike-zhu
spike-zhu previously approved these changes Oct 10, 2026
@voltjia

voltjia commented Oct 12, 2026

Copy link
Copy Markdown
Collaborator Author

Rebased onto master \4c014ca5\ and resolved the overlap in \ ests/test_flash_attn_varlen_func.py\ by retaining both master's Iluvatar regression cases and this branch's non-ALiBi Moore paged case. Head is now \�33050f6; GitHub reports the PR mergeable.\n\nRe-validated the linked Moore build on the S5000 host after the rebase. Compilation and linking passed with the two FlashAttention linked operators selected. Existing test subset: 13 passed, 9 skipped, 120 deselected. The additional skips compared with the pre-rebase run come from newly selected master cases that are outside this provider's coverage; there were no failures.

@voltjia
voltjia merged commit 5653c00 into master Oct 14, 2026
14 of 20 checks passed
@voltjia
voltjia deleted the feat/moore-flash-attn-native branch October 14, 2026 08:26
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.

2 participants