I re-implemented this one for macOS over the last couple of months: https://github.com/mmastrac/diffgemma
I like the model a lot and it's fairly good at reasoning. You can also really bend it to your needs. It's designed for machines with more compute than memory bandwidth but IMO does really well on metal.
I've got it up to ~15tok/s on M3-class machines, but I wager there's a bunch of perf on M5 that I just don't have hardware access to unlock.
I tried to implement MTP using the other Gemma MTP heads but I failed to move that perf needle. There's some interesting research to be done about pre-seeding the diffusion canvas from draft models. DiffusionGemma with the right drafter can hit 20-30 tok/s on my machine, but I've been unable to combine the two together to make it faster than what it's been running at so far.
I've also re-implemented DiffusionGemma from scratch! But with Rust+CUDA rather than macOS, I'm sure our club is a fairly small one :)
About the drafting/MTP, how would that work? As I understand, MTP/drafting is mostly useful for autoregressive models, not for a diffusion model, because each canvas (in the case of DiffusionGemma at least) is fast enough, what you care about is the "convergence" time before it figures out that the accuracy is good enough to move on to another canvas. This is my understanding at least. Would this drafting for a diffusion model like this mean you'd generate multiple canvas at the same time? That'd also seem strange, as then you'll cut performance for generating a complete canvas... I'm genuinely curious about how it could be applied to DiffusionGemma though, so please do expand :)
You can save steps when diffusing from a pre-drafted canvas instead of seeding with random noise! If you have a drafter that is good enough, steps drop by ~50% or more, and tok/s can definitely increase.
I'd be curious to see your CUDA approach - happy to combine efforts as well.
I have M5, but diffusion is a bad fit for metal: https://eamag.me/2026/why-parallel-diffusion-llms-are-slow-o...
That's mostly correct, but I believe the INT8 support on M5 changes the equation pretty drastically. I also believe that diffusion models generally output _better_ results if they have shaping for code (they often will miss delimiter matching inside of a canvas).