Skip to content

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

Open
voltjia wants to merge 6 commits into
masterfrom
feat/moore-flash-attn-native
Open

voltjia wants to merge 6 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
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