Delta-Matching for reliable FP8 attention training

01How does delta become stale?

Write down softmax backward

For one attention head with queries $Q$, keys $K$, values $V$, and head dimension $d$,

Write $dX$ for the loss gradient with respect to $X$. Backward computes

For a fixed query row $i$, $\delta_i$ does not depend on the key index $j$. In exact arithmetic, $\sum_jP_{ij}=1$. Expanding the row sum and substituting the definition of $\delta_i$ gives

The third line replaces $\delta_i$ with its weighted sum. The next line uses row normalization. The two remaining sums are identical because $j$ and $\ell$ are just different names for the same key index.

This is the zero-row-sum constraint. Why should it hold? Adding the same constant to every logit in a row leaves softmax unchanged, so the gradient must sum to zero along that direction.

Why can we reuse the forward output?

Start from $\delta_i$ and substitute $dP=dO\,V^\top$. Here $k$ indexes value channels,

We swap the finite sums and take $dO_{ik}$ outside the sum over $j$, since it does not depend on $j$. The remaining inner sum is $O_{ik}=\sum_jP_{ij}V_{jk}$, which gives the final dot product.

Attention kernels can therefore compute the correction from the saved output and recompute probabilities tile by tile during backward, without storing the full attention matrix. The shortcut works because both expressions use the same operands.

What changes in FP8?

Let $\widehat P$, $\widehat V$, and $\widehat{dO}$ denote the decoded FP8 operands, including their block scales. Here $P$ is the FP32 softmax output from the FP8 attention path, before probability quantization. It is not the probability matrix from a separate BF16 reference run.

The forward and backward products are

These hats label products of quantized operands; the products themselves use FP32 accumulation. To isolate operand rounding, assume that $P$ is normalized, both passes use the same effective $\widehat V$, and accumulation, reduction, and output-storage rounding introduce no further error.

Backward now uses $\widehat{dP}$, but the conventional correction still uses the output gradient before its FP8 cast,

Expand the stale gradient's row sum in the same way,

Normalization still removes $\sum_jP_{ij}$, but the two remaining terms need not match. To see their operands explicitly, substitute both FP8 products and exchange the finite sums,

Separate the two mismatches

Let $M_i=\langle\widehat{dO}_i,\widehat P_i\widehat V\rangle$. Adding $-M_i+M_i=0$ preserves the residual but lets us group one pair with the same output gradient and another with the same weighted values,

In the first bracket, $\widehat{dO}_i$ is shared, leaving $(P_i-\widehat P_i)\widehat V$. In the second, $\widehat P_i\widehat V$ is shared, leaving $\widehat{dO}_i-dO_i$.

The first term comes from probability quantization. The score gradient uses $P$, while the saved output uses $\widehat P$. The second comes from output-gradient quantization. The backward matmul uses $\widehat{dO}$, while the saved-output correction uses $dO$.

When these errors do not cancel, the score gradient has a nonzero row sum. We call the mismatched correction stale delta. A related saved-output rounding problem, and the direct backward correction, were studied in BF16 training by Qiu and Yao.

02What does stale delta do to training?

Keep $P$ and $\widehat{dP}$ fixed, and define a centered correction and gradient,

The residual above is $\varepsilon_i=\delta_i^{\mathrm{centered}}-\delta_i^{\mathrm{stale}}$. Subtracting the gradients shows where the extra term enters,

The two $P_{ij}\widehat{dP}_{ij}$ terms cancel. Rearranging gives $dS_{ij}^{\mathrm{stale}}=dS_{ij}^{\mathrm{centered}}+\varepsilon_iP_{ij}$.

The extra term is proportional to the entire probability row. In a small gradient-descent step on the logits, this term alone can sharpen the row when $\varepsilon_i<0$ and flatten it when $\varepsilon_i>0$.

Positive and negative residuals can nearly cancel across rows without canceling within individual channels. Channel gains are the learned RMSNorm scales on query and key channels. In our diagnostics, channels favored by the sharpened rows also show consistent upward pressure in Adam's gain updates.

Around step 6,000, the last attention layer's key RMSNorm gain for channel 7 reaches about 5.4 under stale delta, versus 1.2 with Delta-Matching. Later, a few channels account for more than 99% of the sum of squared query-key gain products in affected layers. This concentration can amplify selected logit contributions and push softmax toward saturation.

These diagnostics associate the residual with gain growth, without isolating its contribution to gain updates. We therefore track channel gains alongside loss.

03Does matching the backward operands fix it?

We can fix this mismatch by computing the correction from the operands used by softmax backward,

Substitute this correction into the row sum, just as in the first proof,

After using $\sum_jP_{ij}=1$, the two weighted sums contain exactly the same backward operands and cancel. This equality uses the normalization and rounding assumptions stated above.

The probabilities stay in FP32 for this correction, and $\widehat{dP}$ is the FP32 result of an FP8 matmul. All seven standard attention-core matmuls can use FP8, while softmax and reductions use higher precision.

Could we just match $dO$ in the saved-output shortcut? Replacing $dO$ with $\widehat{dO}$ removes the output-gradient mismatch, but leaves the probability mismatch. We test both choices in the default 1.67B experiment.

Mean ± standard error over two initialization seeds, after 30B Nemotron-CC tokens. The GDN/GQA model interleaves Gated DeltaNet and grouped-query attention. Its reference already uses FP8 FFNs; the FP8 methods also quantize attention projections and the attention core.

In the kernel, Delta-Matching restores zero row sums up to FP32 rounding before $dS$ is cast to FP8. That cast introduces another row-sum residual. In our kernel probes, normalized post-cast row-sum residuals remain about 5.3 to 6.3 times larger with stale delta than with Delta-Matching. Ordinary FP8 rounding remains, and downstream scores vary across benchmarks. For example, default RULER-8K is 47.8% with Delta-Matching versus 52.3% for the reference.

04What changes as the model gets larger?

We pretrain 569M and 5.29B GDN/GQA hybrids for 30B tokens each. Both use attention head dimension 256 and the default training recipe and precision scopes.

Mean ± standard error over two initialization seeds.

At 569M, stale delta leaves a modest loss gap and largely preserves downstream performance. At 5.29B, its RULER-8K score falls from the reference's 53.5% to 16.3%, while Delta-Matching scores 58.6%.

The 569M model shows little excess channel-gain growth. At 5.29B, selected gains grow rapidly, more steeply than in the 1.67B trace above over the same first 6,000 updates. Depth, width, and head geometry also change across these models, so the results establish a scaling trend across the tested configurations.

05How do architecture and optimizer choices affect training?

Architecture and optimizer choices change how strongly stale delta affects training. We compare each configuration with its own BF16/FP32 reference.

Attention variants change when failure appears

We test three variants of the 1.67B GDN/GQA model.

• QK normalization prevents outright collapse here, but leaves a large loss gap. Removing query and key RMSNorm makes stale-delta training diverge.
• NoPE delays degradation. We remove RoPE from GQA while the Gated DeltaNet layers remain order-sensitive. Stale delta tracks the reference for longer, then develops a loss gap.
• A larger head dimension does not fix the mismatch. At head dimension 256, stale delta ends at validation CE 2.0709 versus 1.4135 for its reference.

Delta-Matching stays within 0.0011 validation CE of all three references, with broadly comparable downstream performance. The head-dimension-256 study keeps FFNs and attention projections in BF16; NoPE and no QK normalization use the default precision scopes.

Model architecture can hide the damage

Replacing GQA with multi-head latent attention (MLA), or Gated DeltaNet with Kimi Delta Attention (KDA), does not remove the failure. After 30B tokens, stale delta raises validation CE by 0.5090 in GDN/MLA and 0.4230 in KDA/GQA. Delta-Matching stays within 0.0006 of their references, changing only the softmax-attention backward pass.

Near-reference loss does not always mean healthy attention. After 30B tokens, the 1.67B Mamba-2/GQA hybrid retains broadly comparable downstream performance. A single-initialization probe nevertheless finds severely inflated gains in its last attention layer. Masking that layer's attention output hurts less than masking the reference's, suggesting compensation elsewhere in the network.

Last-attention-layer probes find little gain corruption in the 1.475B pure-GQA model, but severe corruption in a single-seed 2.58B run despite a small loss gap. Masking results again suggest compensation, possibly through its other attention layers. Head geometry, precision scope, and implementation also change between these models, so this comparison does not isolate parameter count.

Muon still needs a matched correction

We also train the default 1.67B hybrid for the full 28,610 updates, about 30B tokens, with Muon. The three methods share one initialization and data order. Muon updates hidden weight matrices; AdamW still updates the other parameters, including normalization gains. The reference uses FP8 FFNs, while the FP8 methods also quantize attention and its projections.

Full-training results from one initialization, without across-seed error bars. RULER-8K and CS9 are percentages; CS9 averages nine individual benchmark accuracies.

Muon improves reference validation CE over the paired AdamW run, from 1.4174 to 1.3988, but stale delta still degrades loss and downstream performance. Delta-Matching matches the Muon reference's validation loss and preserves downstream quality. A probe of the first 4,500 updates also shows excess channel-gain growth under both optimizer recipes. Switching hidden-matrix updates to Muon does not repair the backward mismatch.

06How much does the long-context token budget matter?

We next continue pretraining the 1.67B GDN/GQA model from 8K to 64K context. Each method starts from the same completed BF16/FP32 checkpoint for its initialization seed, with a fresh optimizer. Training runs for 28,000 updates, about 29.4B additional tokens, at a peak learning rate of $8\times10^{-4}$ instead of the base recipe's $2.4\times10^{-3}$.

Stale-delta FP8 and cuDNN/TE FP8 build excess channel gain while their losses initially remain close to the reference. Visible loss gaps appear later. For stale delta, the gap exceeds 0.01 only after roughly 19,000 optimizer steps, or 20B tokens. Delta-Matching maintains near-reference gains and tracks BF16/FP32 training loss.

With tokens per optimizer step fixed, a larger training-token budget also means more updates. A short test of loss alone would miss the later degradation in these runs. The corpus expanded at step 18,000, however, and TE's amplitude history reset while optimizer state was preserved. This experiment does not isolate training duration, unique-data volume, or learning rate as the cause of the onset.

At 64K context, the mean RULER score is 45.6% for Delta-Matching and 46.0% for BF16/FP32, compared with 40.1% for stale delta and 38.5% for cuDNN/TE FP8.