# 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.

Website: https://arxiv.org/abs/2606.10820