TokenGate: Adapting LLMs for pre-fill with learned skimming for ~11x speedup A research paper posted to arXiv (2610.26754) introduces TokenGate, a drop-in neural network component that lets existing trained transformers and hybrid models learn to selectively process only the context tokens that matter during prefill, delivering up to 11.3x faster prefill and 3.4x lower memory with no accuracy degradation. In the released needle-in-a-haystack checkpoint (liana313/tokengate-niah-qwen3-1.7b), a gate trained once serves every keep-budget: answering a question planted in 32,191 tokens of essays with only 144 tokens (0.4%) retained in the KV cache cut prefill from 1471 ms to 130 ms at roughly 11x lower time-to-first-token. TokenGate is installed as a Python package (Python ≥ 3.10, CUDA PyTorch ≥ 2.9) and trained like standard PEFT, with a {state} placeholder marking the prunable document and a max_kept_prefill_tokens parameter controlling the keep budget. LLMs that can skim their context to selectively read the tokens that matter. Skimming beats full reads: up to 11x speedup & 3x lower memory for long-context prefill. Paper https://arxiv.org/abs/2610.26754 · Quick start -quick-start · How it works -how-it-works · Reproducing results -reproducing-results Current language models read all tokens of their input context, processing each one through all model layers. TokenGated Language Models can skim , selectively processing only context that matters through learned token selection for up to 11.3x faster prefill with 3.4x lower memory at no accuracy degradation . TokenGate extends existing models to enable conditional computing, akin to MoE layers, but for a different scaling dimension: long-context inputs. TokenGate is a drop-in neural network component that extends existing trained models e.g., trained transformers or hybrid architectures to enable learnable token selection with light-weight adaptation. You can adapt your model with TokenGate as a replacement to standard fine-tuning or PEFT . While standard fine-tuning allows models to improve task-specific accuracy, training models with TokenGate allows models to improve both accuracy and prefill efficiency for Pareto-optimal inference performance. Training with TokenGate is as simple and efficient as standard PEFT training. uv venv && uv pip install -e . Python ≥ 3.10, CUDA PyTorch ≥ 2.9; ". train " for training Recommended: python kernels/fa2bias/build.py builds the fast FlashAttention-2 key-bias kernel that the dense-Qwen3 defaults use needs nvcc . Everything also runs without it via FlexAttention: inference falls back automatically with a warning; training asks for an explicit attn backend="flex". If you're already familiar with 🤗 transformers from pretrained / apply chat template / generate , getting started with TokenGate will be a breeze. Two things are new: a {state} placeholder marks the prunable document in your context, and generate takes a keep-budget max kept prefill tokens — the number of state tokens that are retained for full processing during prefill. The example below loads the released needle-in-a-haystack checkpoint and answers a question about one fact planted in ~32,000 tokens of essays. The sample haystack ships with the repo, so it runs as-is: python import json from transformers import AutoTokenizer from tokengate import load model, apply chat template one object: base model + gate + paired adapter; gate layer and base model are read from the checkpoint's own metadata model = load model "liana313/tokengate-niah-qwen3-1.7b" tok = AutoTokenizer.from pretrained model.base model ~32k tokens of essays with one planted fact about San Francisco essays = json.loads next open "examples/data/niah 32k samples.jsonl" "body" prompt = "Answer the question using only the document below. " "Keep your answer short and direct.\n\n" "Question: What is the best thing to do in San Francisco?\n\n" "Document:\n{state}\n\nAnswer:" ids = apply chat template tok, prompt, state=essays for k in None, 2000, 500, 100 : None = keep every state token out = model.generate ids , tok, max kept prefill tokens=k print k, f"{out.prefill s 1e3:5.0f} ms", out.answers 0 K prefill answer all 1471 ms eat a sandwich and sit in Dolores Park on a sunny day 2000 153 ms eat a sandwich and sit in Dolores Park on a sunny day 500 157 ms eat a sandwich and sit in Dolores Park on a sunny day 100 130 ms eat a sandwich and sit in Dolores Park on a sunny day One gate serves every K: the planted needle is still answered verbatim with 144 of 32,191 prompt tokens 0.4% in the KV cache, at ~11× lower time-to-first-token. Training looks like a TRL/PEFT recipe: a config object plus a trainer that takes your dataset. The example below trains the same needle-in-a-haystack gate the inference example uses. First generate the synthetic haystacks: python datasets/niah kamradt/generate niah kamradt.py --output-file niah.jsonl --n-train 12000 Then train: python import json from tokengate import TokenGateConfig, TokenGateTrainer niah = json.loads l for l in open "niah.jsonl" rows = lambda split: {"article": f"{r 'body' }\n\nQuestion: {r 'question' }", "answer": r "answer" } for r in niah if r "split" == split config = TokenGateConfig base model="Qwen/Qwen3-1.7B", gate layer=1, layer L whose hidden states the gate scores output dir="e2e gates/niah-1p7b", learning rate=1e-4, epochs=20, k eval= 2000 , in-training eval K sweep LoRA co-adaptation the backbone itself stays frozen lora rank=16, lora alpha=32, lora lr=1e-4, lora targets= "q proj", "v proj" , wandb project="token-gate", omit to disable logging trainer = TokenGateTrainer config, train dataset=rows "train" , {"article": ..., "answer": ...}, ... eval dataset=rows "test" , task="niah kamradt", registered prompt + metric datasets/tasks trainer.train → e2e gates/niah-1p7b/gate final.pt Each row is {"article":