Training speculative decoders: removing the logits and attention bottlenecks

Training speculative decoding draft heads at scale introduces two bottlenecks that are easy to miss: the memory cost of the language-modeling loss, and the attention pattern used by EAGLE-style draft heads.

This post walks through two engineering techniques we use to make that training practical. Streaming Cross Entropy avoids materializing full vocabulary logits for every token at once. Block-sparse FlashAttention lets EAGLE-3 draft-head training use efficient attention kernels instead of falling back to dense or less efficient implementations.

This is the engineering layer behind the broader Custom Speculator Training direction in Nebius Token Factory: helping teams train workload-specific drafters for speculative decoding, then evaluate them against real production traffic.

How this is different from generic long-context fine-tuningHow this is different from generic long-context fine-tuning

Other long-context training explainers often focus on extending context length for standard fine-tuning. The problem here is narrower and more specific: training speculative decoding draft heads.

Draft-head training adds extra token-prediction objectives. Each head may predict a different future token, which makes target alignment, masking, loss computation, and attention efficiency more delicate than in ordinary language-model training. Streaming Cross Entropy addresses the logits-memory bottleneck; block-sparse FlashAttention addresses the draft-head attention pattern.

The logits bottleneck in long-context draft-head trainingThe logits bottleneck in long-context draft-head training

Training long-context LLMs with large vocabularies introduces a memory bottleneck that is easy to underestimate: the final language-modeling loss.

In a standard language-modeling objective, the last transformer layer produces hidden states, which are projected by the LM head into vocabulary logits. The expensive part is not only the projection itself, but the full path:

hidden states -> LM head -> logits -> Cross Entropy -> loss

We will refer to this as the LM loss path. A naive implementation may materialize the full logits tensor before applying Cross Entropy.

The hidden states have shape: [batch_size, sequence_length, hidden_dim];

The logits tensor has shape: [batch_size, sequence_length, vocab_size].

Hidden dimension is significantly smaller than vocabulary size. Logits tensor scales with the number of tokens and the vocabulary size. For example, Llama 3.1 8B with the following setup:

batch_size = 16
sequence_length = 32k
vocab_size = 128256 # 128k (llama 3.1 8b)
logits dtype = fp32

The logits tensor alone requires: 16 × 32768 × 128256 × 4 bytes ≈ 250.5 GiB. While hidden states require only: 16 × 32768 × 4096 × 2 bytes ≈ 4 GiB. If both logits and its gradient are materialized, the peak memory associated with this path is roughly doubled: ~501GiB.

Furthermore, it is crucial for both small and large models:

• Llama-3.2-1B has hidden_dim=2048 and vocab_size=128k;

• Llama 3.1 405B has hidden_dim=16384 and vocab_size=128k.

In this case 1B model will have the 8-times smaller hidden-states but same-size logits tensor as 405B model. This is clearly impractical on a single GPU or even across an 8GPU node. Without changing the loss implementation, one has to reduce the local token count, i.e. use smaller microbatches, use higher degrees of context parallelism, or reduce precision while it is usually undesirable.

The LM loss path is not always the largest memory consumer in LLM training. Depending on the model architecture and parallelism strategy, attention activations, MLP activations, optimizer state, communication buffers, or MoE buffers may dominate. However, with large vocabularies, long contexts or multiple speculative draft heads, the LM loss path can become a major activation-memory bottleneck. Streaming Cross Entropy addresses this bottleneck directly.

Streaming Cross EntropyStreaming Cross Entropy

Streaming Cross Entropy avoids materializing logits for all tokens at once.

The key observation is that typical LM objective, cross entropy, is just an independent sum of per-token CE losses. So we can process tokens in smaller chunks independently, sum their losses, and divide by the total valid tokens at the end, instead of computing logits for all tokens at once.

Conceptually, the dense version does this:

logits = lm_head(hidden_states)  # [batch_size, sequence_length, vocab_size]
loss = cross_entropy(logits, targets)

Streaming Cross Entropy does this instead:

total_loss = 0.0
total_weight = 0

for chunk in split_into_chunks(
    hidden_states,
    targets,
    loss_mask,
    chunk_size,
):
    chunk_logits = lm_head(chunk.hidden_states)

    chunk_loss_sum, chunk_weight = cross_entropy_sum(
        chunk_logits,
        chunk.targets,
        chunk.mask,
    )

    total_loss += chunk_loss_sum
    total_weight += chunk_weight

loss = total_loss / total_weight

The important detail is that chunking happens over tokens.

A token chunk does not have to correspond strictly to the sequence axis. It may be a slice of the sequence dimension, a slice of the batch dimension, or a flattened batch-sequence token set. What matters is that every valid token contributes exactly once with the same target, mask, and scaling as in the dense loss.

Correctness: alignment before chunkingCorrectness: alignment before chunking

Streaming Cross Entropy must be lossless. Every token that contributes to the dense loss must contribute exactly once to the streaming loss, with the same target, mask, and scaling factor. Otherwise, the optimization objective changes.

The safe order of operations is:

1. Align hidden states, targets, and masks for the objective;

2. Apply loss masks;

3. Split the aligned token set into chunks;

4. Accumulate summed loss and valid-token count;

5. Normalize using the original loss denominator.

This is especially important for packed sequences and speculative decoding objectives. For standard LM training, targets are usually shifted by one token. For speculative decoding or draft-head training, each head may predict a different future token, so the target shift depends on the head. If chunking is applied before target alignment, chunk boundaries may cut through prediction-target pairs. This can lead to incorrect targets, dropped boundary tokens, or wrong loss masks.

Here we see the draft-heads-shifting masks most of the boundary tokens as each head should rely on a previous token, but the previous token is stored in a previous chunk. Therefore, target shifting and mask alignment must happen before chunking.

Hidden states, targets, and masks are first transformed into the exact token-level objective, and only then split into chunks. This keeps Streaming Cross Entropy equivalent to dense Cross Entropy, except for small floating-point differences caused by reduction order.

Backward passBackward pass

The forward pass computes the loss chunk by chunk. The backward pass follows the same idea. For each token chunk, we compute the local logits, Cross Entropy, and gradients. The input gradients are produced for the corresponding hidden-state chunk. The LM-head weight gradients are accumulated across chunks.

The production implementation wraps this logic in a custom backward rule, so the caller sees the same interface as a regular loss function while the implementation avoids storing the full logits tensor.

Conceptually:

for each token chunk:
    compute logits for the chunk
    compute chunk loss
    compute gradients for chunk hidden states
    compute gradient for LM-head
    accumulate gradients for LM-head weights

As a result, the peak memory is bounded by: chunk_size x vocab_size x dtype_size x 2.

For example, with:

chunk_size = 1024
vocab_size = 128k
dtype = fp32

the logits plus gradient buffer require: 1024 x 150_000 x 4 x 2 ~= 0.97GiB. Compare this with the dense example above, where the same path required hundreds of GiB.

This is conceptually similar to gradient accumulation, but the split happens inside the loss computation rather than across optimizer steps. We are not changing the training batch. We are changing how many token logits are materialized at one time.

Gradient accumulation reference:

$L_b = \sum_{1}^{n} L_{mb_i} * \frac{|tokens_{mb_i}|}{|tokens_{batch}|}$, where mb — minibatch, n — number of minibatches in batch.

Streaming cross entropy loss:

$$L_{minibatch} = \sum_{chunks} L_{c_i} * \frac{|tokens_{c_i}|}{|tokens_{minibatch}|}$$

The final loss must use the same normalization as the original dense objective. For token-level supervised fine-tuning, this usually means normalizing by the number of valid, non-masked tokens. It should not be normalized by the number of chunks.

Streaming Cross Entropy is orthogonal to the surrounding parallelism strategy. It can be used with FSDP, data parallelism, context parallelism, tensor parallelism, or no model parallelism at all. The technique only changes how many token logits are materialized locally at a time; it does not require changing the global training topology.

Choosing the chunk sizeChoosing the chunk size

The chunk_size defines the maximum number of tokens for which logits are materialized at once. For example: chunk_size = 1024, means the largest logits tensor has shape: [1024, vocab_size] regardless of how those tokens are distributed across batch and sequence dimensions.

Lowering chunk_size reduces logits memory, but only up to a point. Once the LM loss path is no longer the peak memory consumer, making chunks smaller does not reduce the overall training peak. It only adds overhead: more loop iterations, more kernel launches, and less efficient matrix multiplications.

Therefore, the goal is not to make chunk_size as small as possible. The goal is to choose the largest chunk size that keeps the LM loss path below the next major memory peak in the training step. A useful heuristic is: streaming_lm_peak <= transformer_block_peak.

Once this condition holds, the LM loss path is no longer the dominant memory peak. Further reducing chunk_size mostly hurts performance without improving the actual peak memory of the compiled graph.