Maintained fork with PyTorch bindings and macOS 15 support

[imperatormk]

Hey! I've been using metal-flash-attention for a PyTorch training project on Apple Silicon and made some additions:

Fork: imperatormk/metal-flash-attention

Changes:

• ✅ Causal masking support (for autoregressive models)
• ✅ macOS 15 (Sequoia) compatibility fixes
• ✅ Tested with PyTorch MPS backend

PyTorch Bindings: imperatormk/mps-flash-attention

Drop-in Flash Attention for PyTorch on Apple Silicon:

• O(N) memory instead of O(N²)
• 2-5x faster than PyTorch SDPA
• Full backward pass for training
• Pre-compiled kernels for zero cold start

from mps_flash_attn import flash_attention

out = flash_attention(q, k, v, is_causal=True)

If you're interested in merging any of these changes upstream, happy to open a PR!

Thanks for the original library - it's been super helpful for training on M1/M3 Macs.

[bghira]

see previous discussion in #30 and #31, but the level of confidence in your claims of vibe-coded stuff is just kinda unwelcome and causing problems in a more broad sense across Github.

[bghira]

you also can't put pre-compiled kernels into a codegen based repo. there's a large number of issues caused by removing all of the dynamic nature of this repo.

[philipturner]

I designed this library to not need any maintenance after the summer internship was over. I will not accept further commits except typo fixes. However, I keep the repo open for issue threads like this.

Liu made his v2.5 fork after I did a knowledge transfer to teach him the skills to design a high-performance codegen based library. It's in C++ and in his CCV/NNC repo.

With macOS 26, this entire library no longer works. I'm refusing to upgrade any of my devices to macOS 26. Apple intentionally banned SIMD-group async copies to do planned obsolescence of M1 and M2 GPUs.

[imperatormk]

You're right for SIMD but there is a way to do it by compiling with -std=macos-metal2.4 which still allows __asm directives; see all the changes I made: https://github.com/imperatormk/metal-flash-attention to make https://github.com/mpsops/mps-flash-attention work. I'd address precompiled kernels vs codegen approach after the initial concept is stable

[bghira]

that's just not supported in a way that we can rely on. it's the same as using non-public APIs. you can't get into the app store that way. if the initial concept isn't stable, why even open the PR?

[bghira]

also, there's built in flash attention in pytorch - it's just not enabled by default. it's often better-performing and faster than MFA because it's using highly optimised MPS kernels, which are really hard to reach performance parity with - especially when vibe-coding the kernels.

[philipturner]

You're right for SIMD but there is a way to do it by compiling with -std=macos-metal2.4 which still allows __asm directives; see all the changes I made: https://github.com/imperatormk/metal-flash-attention to make https://github.com/mpsops/mps-flash-attention work. I'd address precompiled kernels vs codegen approach after the initial concept is stable

@imperatormk does this option work on macOS 26 and with runtime JIT compilation? If so, we should ping my previous employer at Draw Things.

I think the saturation of software improvements, means we can use Metal 2.4 indefinitely into the future. Just like how Windows 10 stays relevant today and will never become irrelevant, because it's all just artificial planned obsolescence.

it's often better-performing and faster than MFA because it's using highly optimised MPS kernels

On M3 and M4 chips. But performance is probably crap on M1 and M2.

especially when vibe-coding the kernels.

No need to vibe code when a correctly coded code base can autogenerate kernels, with clearly legible explanations of the parameters to tune, and how to test theories about the best possible kernel. Perhaps the AI reasoning model could help you follow the guidelines of how I designed this repo, and continue following them. I haven't taken the time to read over the code in this PR or issue.

There are instances in EUV and X-ray lithography where correctly applied computations can maximize performance of an analog process, which relies on computational lithography. But it requires first prioritizing the performance objective, and the tools used are a 2nd priority. It can be as simple as back-of-the-hand Maxwell's Laws, or as complex as parallel computers or AI.

Even in the real, analog world, simple things require correctly selecting among 10 different controlling variables, with ~1.65 thousand combinations to explicitly test in silico, using nothing but simple physical equations and models. That's how MFA was designed, for all architectures (M1, M2, M3, M4), except those simulations can become IRL tests on a computer, because software is so easy to compile and execute.

[bghira]

these LLMs will go in circles to make things performant, at the cost of correctness. they can create quite zippy kernels, but by disrupting the results. they'll create scalar code that loops over elements to compute eg. multi-head attention because it's very out of distribution for them to "understand" how vector loads work and how to split work units so that GPU occupancy / utilisation is high without simply wasting compute. you can watch the model waste thousands of dollars in tokens trying to "solve the problem" and make its tests produce the results it's looking for. but prepare for it to "forget" how to run the test after the context compacts and decide to start making shortcuts that you'd already told it not to do.

oh right, another problem when i tried this was that the codegen syntax gets entirely ruined by the model and then swiftformat doesn't understand the indents, makes the problem worse, and the code no longer passes checks. it may take 3 to 4 hours to get it to resolve that.

we're not there yet

[bghira]

one of the notable hallucinations on the linked github repo is that it claims to have FP8 support, despite the hardware not having such a thing available.

[imperatormk]

@bghira I think we are somewhat there, this was partially my fault of not being able to guard the development properly but let me reinvestigate.

@philipturner Yes unless I'm missing something it works well; see https://github.com/imperatormk/metal24-asm-repro with results included. And I confirm the vanilla flash attn is significantly slower on my M1 vs this implementation (still always room for improvement though)

[bghira]

you can't do that and still pass Apple store policy checks (and your whole developer license can be revoked for trying)

[bghira]

oh, i see. this primitive is only in Xcode metal. it's not in the system library. you have to have XCode installed for this to be accessible, and it's not accessible on mobile.

[philipturner]

Can we make the existing repo, my main repo, run as-is on an iOS device with A14 or A15 chip, which runs iOS 26? Using just the -std= flag?

Do not dismiss the idea so readily based on developer license revocation and crap. Apple is forcing obsolescence on us, it's their fault not ours.

[philipturner]

We're using chips made with a freaking molten tin laser that shoots 13.5 nm light, developed over 20+ years after an artificial Moore's Law through the early 2010's which was actually just computational lithography boosting the resolution of otherwise 193 nm light.

Even the old chips are worth exploiting, even if new ones on the exact same lithographic technology are "less obsolete". None of the new chips are any substantial progress for humanity in the long run.

[imperatormk]

Unfortunately no can do on iOS with this approach, I was preparing all that just so this can be used for inference with say diffusion models; but I think there is an alternative - will let you know, curious about the results too!

Edit: so I was able to get all these 1D/2D ops working on iOS. Now adapting the whole flash attn is a bigger project but doable I'd say.