Universal Metal Flash Attention — Metal向けユニバーサルFFIアダプタ

AI/ML
Universal Metal Flash Attention — Metal向けユニバーサルFFIアダプタ

Universal Metal Flash Attention — Metal向けユニバーサルFFIアダプタ

Metal向けの高速なFlash Attentionライブラリ（Philip Turnerの実装）を、CのForeign Function Interface（FFI）を通して汎用的に利用できるようにするアダプタ実装です。Swiftで実装されたコアをC ABIでラップし、Rust/Python/Julia/Go/Objective-Cなど任意の言語からシームレスに呼び出せるほか、PyTorch向けのゼロコピーバインディングも提供します。性能面でもネイティブに近い実行速度を報告しており、Metalベースの注意機構をマルチ言語環境で活用したい開発者に有用です。（約300字）

https://github.com/bghira/universal-metal-flash-attention

概要

Universal Metal Flash Attentionは、Philip Turner氏のMetal実装のFlash Attentionを「言語非依存」に使えるようにするCのFFIアダプタです。コア実装はSwiftをベースにしつつ、C ABIを介してRust、Python、Julia、Go、Objective-Cなど任意のC FFI対応言語から呼び出せるAPI（Flash Attention 3スタイル）を提供します。特にPyTorch向けにゼロコピーの統合を実現しており、既存のPyTorchワークフローへ性能劣化なく組み込める点が特徴です。READMEでは各FFIのベンチマーク（Rust: 1135 GINSTRS/s、Python: PyTorch SDPA比1.87x高速、Objective-C: 1148 GINS…）が報告されています。

GitHub

リポジトリの統計情報

  • スター数: 2
  • フォーク数: 0
  • ウォッチャー数: 2
  • コミット数: 11
  • ファイル数: 9
  • メインの言語: Swift

主な特徴

  • Metal上のFlash Attentionを汎用C FFIとして公開し、さまざまな言語から利用可能にする。
  • Flash Attention 3スタイルAPIをサポートし、既存の実装との親和性を確保。
  • PyTorch向けにゼロコピーバインディングを提供し、PyTorchのSDPA実装より高速な結果を報告。
  • Swiftベースで実装され、Swift Package ManagerやMakefileによるビルドが可能。

技術的なポイント

本プロジェクトの中心は「Metalで高速に動作するFlash Attentionの実体（Philip Turnerの実装）を、Cの呼び出し規約でラップしてどの言語からも使えるようにする」ことです。Swiftで書かれたコアをC ABIへ露出することで、RustやGo、Julia、Objective-Cなどの言語が直接ネイティブ呼び出しできるインターフェースを用意しています。READMEに示されたベンチマークは、ラッパーを挟んでもネイティブ相当の性能が出ている点を強調しており、特にRust向けは1135 GINSTRS/s、Objective-Cでも1148 GINSTRS/sといった数値が報告されています。Pythonバインディングは「ゼロコピー」を謳っており、PyTorchのテンソルとメモリを共有することでデータ移動のオーバーヘッドを避け、PyTorchのSDPA実装よりも約1.87倍高速という実測結果が示されています。プロジェクト構成からはSwift Package Manager（Package.swift）を使ったビルドやサブモジュール（.gitmodules）の管理、Makefileによるビルド手順の用意が確認でき、クロス言語での組み込みを念頭に置いた設計になっています。これにより、Metal対応のMac/Apple Silicon環境で高速注意機構を他言語プロジェクトへ比較的容易に導入できます。

プロジェクトの構成

主要なファイルとディレクトリ：

  • .gitignore: file
  • .gitmodules: file
  • Makefile: file
  • Package.swift: file
  • README.md: file

…他 4 ファイル

まとめ

Metalベースの高性能Flash Attentionを多言語から使えるようにする実用的なFFIアダプタです（50字程度）。

リポジトリ情報：

READMEの抜粋：

Universal Metal Flash Attention

A universal C Foreign Function Interface (FFI) for the Metal Flash Attention library, providing Flash Attention 3-style API for seamless integration with Rust, Python, Julia, Go, and any language supporting C FFI.

🎉 All FFI Implementations Working & Production Ready

  • ✅ Rust FFI: 1135 GINSTRS/s (matches native Swift performance)
  • ✅ Python FFI: 1.87x faster than PyTorch SDPA with zero-copy PyTorch integration
  • ✅ Objective-C FFI: 1148 GINS…