TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning

TokenGate. Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning. By Liana Patel and four others, from Stanford University and UC Berkeley. Published on the nineteenth of October. Its code repository is worked out: about one hundred source files, with a readme and a licence.

Expert human readers handle large inputs, a code repository or a long document, by skimming. They read closely only what matters. Language models do the opposite. Every input token is processed through every model layer, whether it matters or not, and for long-context inputs that uniform treatment carries a significant processing cost. Prefill, the stage where the whole input is taken in before the first token comes out, is where the cost lands, in time and in memory.

TokenGate gives an existing trained model something like the skimming habit. It is a small, drop-in component that extends the model so that token pruning over long-context inputs becomes learnable end to end. Tokens can be dropped during prefill rather than carried through every layer, and the goal is a prefill that is both fast and memory-efficient, without building a new model from scratch.

The obstacle to learning this end to end is the top-K operator that does the pruning. It selects the highest-scoring tokens, and a hard selection like that is not differentiable, so a learning signal cannot pass through it. The answer is token-gated attention, a mechanism inspired by Mixture-of-Experts, which re-weights attention logits using token importance scores. Adapting a model this way is simple and lightweight, comparable to standard parameter-efficient fine-tuning, and it natively extends both standard transformers such as Qwen3 and hybrid-attention models such as Qwen3.5.

Across diverse tasks, TokenGate cuts time to first token by between one point four and eleven point three times. Peak memory falls by between one point five and three point four times. Both come with no accuracy degradation for prefill in transformers. Against the best-performing prior methods for token pruning, the margin is up to six point four times faster prefill, and two times lower peak memory.

The skimming reader from the opening carries the whole design. Skimming is a learned judgment about what to skip, and TokenGate turns that judgment into something a trained model can acquire, through a gate soft enough to learn through. Because it attaches to models that already exist, adopting it costs an adaptation comparable to ordinary fine-tuning rather than a new model.