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.