apple / apple/coremltools

FlashAttention with MPS

Open
#2,543 1 comment 0 reactions 0 assignees View on GitHub
feature request
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🌱 Describe your Feature Request
To compete with Windows/Linux & Nvidia in genAI
MacOS need containerized easy to use flash attention that can be easily integrated with Transformers in our custom code.

- Running LLM models efficientnly on any GPU requires FlashAttention

## How can this feature be used?
Used by developers to run models on macOS.

## Describe alternatives you've considered
No Flashattention, or Linux/Windows PC with Nvidia.

Use existing swift-based flashattention if you are developing in Swift.
https://github.com/philipturner/metal-flash-attention

## Additional context
Metal-Flash-Attention is cool, but the most of LLM models that you as a developer interactiot is in Python.
Say you have a new model available but you canot run it with an attention on macOS with MPS acceleration because there's no flashattention modules in python.

https://huggingface.co/meta-llama/Llama-4-Maverick-17B-128E-Instruct

Contributor guide

Open the contributing guide

Research direction

No files, tests, or entry points are named. Start by locating the Python integration points for model execution and MPS acceleration in coremltools, then compare the requested macOS FlashAttention behavior with the linked Swift implementation; done means a usable Python integration for Transformers on macOS.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, swift
Domain
machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.