Comment by mmastrac
6 hours ago
I have a project that's almost ready to run DiffusionGemma as well. The two project might potentially work well together. I'm getting ~20tok/s on a 36GB M3 and there's strong possibility we might be able to crib faster kernels from each other.
Feel free to reach out.
(currently at https://github.com/mmastrac/diffgemma but not in a releasable state yet)
Cool project! I looked into it recently and thought that running diffusion models locally doesn't really make sense: https://eamag.me/2026/why-parallel-diffusion-llms-are-slow-o...
What are your thoughts on this?
TBH, I think there's some truth to that. I spent _ages_ tuning the kernels to match the tested FLOP count of my M3's processor. I only have an M3 though and wasn't able to push int8 very far on it, but I think there's a chance that M5-class machines and higher might have more capability in this regard.
What I also learned is that MLX/vLLM is probably within ~20% or so of the absolute max perf on Mac. I found some improvements over what they were doing, but we're at the point where it's challenging to optimize without per-stepping kernels.
I found a few improvements over stock DiffusionGemma along the way, like using top-k attention, which drastically improves perf on my mac without sacrificing any of the benchmarks I was able to throw at it.
FWIW some of the issues with Gemma being slow on Mac are specific choices they've made in the architecture that make it challenging to make use various optimizations that have popped up recently. I think a Kimi K3-style network hybrid with the diffusion bits of DiffusionGemma could have some serious sway.
I think that diffusion still has an edge locally, but with some architecture tweaks and CPU improvements it would actually be a winner (ie: training the network for smaller token batch sizes or flexibility in attention heads, a less expensive attention mechanism, and others).
It is super cool! Diffusion Gemma was released around the middle of my project, and I seriously considered switching to it. But I decided to finish the project as it was.
I believe it would be a perfect match!
Feel free to use any parts of my project or drop me a message. There’s my LinkedIn link at the end of the readme. Or I will drop you a message later!
Awesome. I may need to finally bite the bullet and upgrade my macOS to test out the MPP approach you've taken.
I've got a number of tiled-load kernels, and a top-k attention kernel that you might find interesting.