ZAYA1-8B-Diffusion-Preview: Converting an AR MoE to Diffusion on AMD

TL;DR

Zyphra converts the autoregressive ZAYA1-8B-base MoE checkpoint into a discrete block-diffusion language model with no systematic eval degradation. The conversion follows the TiDAR recipe — 600B tokens of diffusion mid-training at 32k context, 500B tokens of native context extension to 128k, then a diffusion SFT phase — and produces a model that diffuses blocks of 16 tokens simultaneously. The reported acceptance/speedup is 4.6× with a lossless speculative sampler and 7.7× with a “mixed-logit” sampler that averages the diffusion and AR logits at verification time. Notable as the first MoE diffusion LM converted from an AR base, and the first diffusion LM trained end-to-end on AMD. The architectural justification is that ZAYA1’s CCGQA attention (4:1 query:key head ratio) plus Zyphra’s CCA variant makes attention memory-bandwidth-bound at standard decode, so converting decoding into prefill-shaped block diffusion gets large speedups before the operation re-saturates compute.

Key claims

ZAYA1-8B-Diffusion-Preview converts an AR base into a discrete block-diffusion model with no systematic eval degradation under the lossless sampler; some benchmarks (e.g. LCB-v6) improve, attributed partly to improved mid-training data since the original ZAYA1-8B midtrain [§Introduction, §Theoretical speedup figure caption].
Block size 16 tokens per diffusion draft; tokens are “accepted” by a sampling criterion similar to speculative decoding, with both speculator and verifier living in the same forward pass [§Results].
Lossless speculative sampler achieves 4.6× theoretical speedup over AR decoding; mixed-logit sampler achieves 7.7× at some quality cost [§Results, speedup figure caption].
Lossless sampler uses the standard speculative-decoding criterion min(1, p(x)/q(x)) with p from the AR mode and q from the diffusion mode; on rejection the next token is drawn from the residual p(x) − q(x) [§Results].
Mixed-logit sampler first averages the AR and diffusion logits, then uses the mixed distribution for verification — pulls verification logits closer to the diffusion proposal, improving acceptance at the cost of some quality [§Results].
The architecture is a single-step speculative diffusion model with order-constrained generation: the diffusion model can only generate tokens in a contiguous subsequence starting from the prefix, which sharply improves training stability versus unconstrained mask-diffusion objectives or set block decoding [§Architecture].
Conversion builds directly on the TiDAR recipe (arXiv 2511.08923) [§Autoregression to Diffusion, §Architecture figure caption].
ZAYA1’s CCGQA attention (4:1 query-to-key head ratio) was chosen partly to keep decoding memory-bandwidth-bound; block diffusion accesses the same cache but scales arithmetic intensity with block size, allowing a single forward pass to verify one block while speculating ~3 block-sized proposals on MI300X (bf16) or ~5 on MI355X [§Architecture].
CCGQA’s 2× compression also made the additional training-time FLOPs of TiDAR affordable in midtraining; AMD’s larger per-GPU VRAM enabled the larger working set [§Architecture].
First MoE diffusion language model converted from an AR LLM, and first diffusion-language model trained on AMD [§Introduction, §Architecture].

Method

Starting from the public ZAYA1-8B-base checkpoint, Zyphra applies 600B tokens of diffusion-conversion mid-training at 32k context, then 500B tokens of native context-extension to 128k, then a diffusion SFT phase. The training recipe is TiDAR-style block diffusion with order-constrained masking: each block of 16 tokens is unmasked simultaneously in a single forward pass, but blocks are processed left-to-right and only contiguous subsequences starting from the prefix can be generated. Order constraint is the key stability lever versus unconstrained mask-diffusion training.

At inference, each forward pass speculates on the next 16-token block while simultaneously verifying the previously accepted tokens via the same model. Accepted tokens are appended to the sequence; the next forward pass speculates against the new prefix. Two samplers are described:

Lossless speculative sampler — accepts each token with probability min(1, p_AR(x)/q_diff(x)); rejected positions are resampled from the residual p − q. This is the standard speculative-decoding rule applied within a single model.
Mixed-logit sampler — averages p_AR and q_diff first, then uses the averaged distribution for verification. Acceptance is higher because the verification target now sits closer to the proposal, but the output distribution drifts from pure-AR.

Block-level arithmetic intensity is the architectural lever: with a 4:1 query-to-key head ratio under CCGQA, AR decoding is memory-bandwidth-bound, so converting decoding into block prefill (which scales arithmetic intensity with block size) lifts utilization. On MI300X bf16 they fit ~3 block-sized proposals per forward; on MI355X, ~5.

Results

4.6× theoretical speedup over AR with the lossless sampler; 7.7× with the mixed-logit sampler [§Results, speedup figure caption].
“Minimal evaluation degradation compared to the base autoregressive checkpoint” under the lossless sampler; LCB-v6 improves [§Introduction, speedup figure caption].
Mixed-logit sampler shows “some degradation” but substantial additional speedup — explicitly framed as a runtime quality/performance trade-off [§Results, speedup figure caption].
The post reports the speedup exceeds alternative speculative methods such as multi-token prediction (MTP), EAGLE3, and is comparable to dFlash with the additional benefit that TiDAR-style diffusion uses a single forward pass [§Results].
All eval numbers are pass@-style on the base mid-train checkpoint (no diffusion-RL post-training yet) [§Theoretical speedup figure caption].

(No quantitative benchmark table is included in the blog body — the figure with per-task numbers is described in the caption but not transcribed.)

Why it’s interesting

This is a direct sibling to Introspective Diffusion Language Models (I-DLM) and WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference (WeDLM) — all three convert pretrained AR models into block-diffusion LMs and recover wall-clock speedup by exploiting the fact that block-parallel decoding is more arithmetic-intensive than per-token AR decode. The differences are worth tracking. I-DLM trains a strict-causal “introspective” objective on a concatenated [masked | clean] input and uses min(1, p/q) within a single forward pass; WeDLM uses topological reordering to keep masked positions cache-compatible under unmodified causal attention; ZAYA1-Diffusion-Preview follows TiDAR’s order-constrained block diffusion and the same min(1, p/q) rule but is the first MoE base and the first AMD-native diffusion LM. The architectural co-design claim — that ZAYA1’s CCGQA was chosen partly because it leaves decoding memory-bandwidth-bound, leaving headroom for block diffusion to fill — is a stronger version of the “end-to-end architecture/inference codesign” framing in Interaction Models: A Scalable Approach to Human-AI Collaboration. Pairs naturally with the AR companion ZAYA1-8B Technical Report (same MoE, RL-trained for reasoning) and with Speculative Decoding: Performance or Illusion?‘s diagnosis that verification compute dominates production wall-clock in standard SD — block-diffusion-within-one-model collapses draft and verify into a single forward pass, the very thing SSD (Speculative Speculative Decoding) tries to parallelize across hardware.