K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling
Official code for K-Forcing, which distills an autoregressive language model into a push-forward model that decodes several tokens jointly in one forward pass instead of one token at a time, for faster batch inference.
K-Forcing distills an autoregressive language model into a push-forward language model that maps k independent noise variables to k future tokens jointly in a single forward pass, modeling their joint distribution rather than generating one token per step. The output length is always fixed and the original autoregressive backbone is reused as-is, which the authors position as better suited to batch serving than either standard autoregressive decoding or speculative decoding, whose variable output length breaks regular batching.
The repository includes a script for autoregressive, push-forward and masked-diffusion-style baseline inference from text prefixes with KV-caching, and released checkpoints for autoregressive and push-forward models trained on OpenWebText and the One Billion Word benchmark.
Setup uses a Python virtual environment with uv, and a separate flash-attention wheel download for the environments the authors tested on.