August 2026 Summaries

Training speculative-decoding draft heads for long-context language models faces two major bottlenecks: the memory required to materialize vocabulary logits for language-model losses and the inefficient sparse attention patterns used by EAGLE-3 draft heads. Streaming Cross Entropy addresses logits memory by calculating token-level cross-entropy loss in chunks rather than constructing logits for every batch token simultaneously, preserving the original objective when targets, masks, and head-specific shifts are aligned before chunking. This can reduce the loss-path memory requirement from hundreds of gigabytes to an amount determined by chunk size, although excessively small chunks reduce throughput through added iteration and matrix-multiplication overhead. Block-sparse FlashAttention addresses EAGLE-3’s banded diagonal attention by computing only block-level regions allowed by causal, document-packing, and draft-head constraints, avoiding dense quadratic attention processing and enabling efficient FlashAttention-style kernels. Benchmarks cited for GPT-OSS-20B-like configurations show substantially lower memory use and higher throughput than a JAX baseline, while combining both techniques reportedly reduced GPU memory utilization from 44% to 25% and raised throughput from 6,900 to 11,200 tokens per second. Together, the approaches support training workload-specific speculative draft models with larger vocabularies, longer contexts, and multiple prediction heads in Nebius Token Factory.