ZAYA1-8B-Diffusion-Preview Released by AMD

View organization page for Zyphra

3,813 followers

Today we're releasing ZAYA1-8B-Diffusion-Preview, our first diffusion language model, and the first diffusion language model trained on AMD . Almost every LLM in production today is autoregressive, generating text one token at a time. This is fundamentally bottlenecked by memory bandwidth. Each new token requires reloading the model's working memory (the KV-cache), leaving the GPU sitting idle or underutilized. As inference workloads scale and reasoning models generate longer outputs, this bottleneck has become one of the largest cost and latency drivers in deploying LLMs. Diffusion language models take a different approach. Rather than generating one token at a time, they draft entire blocks of tokens in parallel and then verify them in parallel. Because the model produces many tokens in a single forward pass (using the same KV-cache), the GPU stays compute-bound and operates closer to peak utilization. We built ZAYA for diffusion from the start, co-designing the architecture around AMD hardware. Our FLOP-efficient CCA attention variant lets us diffuse more tokens at once before becoming compute-bound, dramatically increasing inference speed. Training a diffusion LLM from scratch is challenging, and few proven recipes exist. With ZAYA1-8B-Diffusion-Preview, we take a different approach: building on the TiDAR recipe, we convert our existing autoregressive ZAYA1-8B model into a discrete diffusion model through additional mid-training. This lets us reuse our existing pretraining stack while gaining the inference benefits of diffusion. The result is the first MoE diffusion model converted from an autoregressive LLM. The preview model: - Diffuses blocks of 16 tokens simultaneously, achieving up to 4.6x decoding speedup with a lossless sampler and up to 7.7x with our new logit-mixing sampler - Shows minimal performance degradation versus the autoregressive base, with gains on several evaluation benchmarks - Was trained and converted on AMD, with our compute-efficient CCA attention variant and the larger VRAM capacity of AMD hardware making diffusion training practical. Beyond raw inference speed, diffusion enables more expressive generation than left-to-right autoregression and substantially reduces the cost of on-policy rollouts. This translates to more reinforcement learning and test-time compute scaling for the same compute budget, areas where diffusion's advantages compound over time. ZAYA1-8B-Diffusion-Preview is an early step. We plan to post-train and release a full diffusion model soon. Read more: https://lnkd.in/gsrQCFc3

  • chart, bar chart

Great! I've implemented MLX (Apple Silicon) inference for ZAYA here; https://github.com/ml-explore/mlx-lm/pull/1261 Do you think we can use this model for DFlash ? MLX DFlash implementation is already implemented. This could boost performance but I guess a strong quantization is required to not run into issues with memory bandwith. I get 30 token/sec generated on MacBook Air M4 with my implementation (the standard ZAYA-8B-Base Q4)

To view or add a comment, sign in

Explore content categories