{"slug": "tokengate-adapting-llms-for-pre-fill-with-learned-skimming-for-11x-speedup", "title": "TokenGate: Adapting LLMs for pre-fill with learned skimming for ~11x speedup", "summary": "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.", "body_md": "**LLMs that can *skim* their context to selectively read the tokens that matter.**\n\n**Skimming beats full reads: up to 11x speedup & 3x lower memory for long-context prefill.**\n\n[Paper](https://arxiv.org/abs/2610.26754) ·\n[Quick start](#-quick-start) ·\n[How it works](#-how-it-works) ·\n[Reproducing results](#-reproducing-results)\n\nCurrent 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.\n\nTokenGate 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.\n\n```\nuv venv && uv pip install -e .     # Python ≥ 3.10, CUDA PyTorch ≥ 2.9; \".[train]\" for training\n```\n\nRecommended: `python kernels/fa2bias/build.py` builds the fast\nFlashAttention-2 key-bias kernel that the dense-Qwen3 defaults use (needs\n`nvcc`). Everything also runs without it via FlexAttention: inference falls\nback automatically with a warning; training asks for an explicit\n`attn_backend=\"flex\".\n\nIf you're already familiar with 🤗 transformers (`from_pretrained` /\n`apply_chat_template` / `generate`), getting started with TokenGate will be a\nbreeze. Two things are new: a `{state}` placeholder marks the prunable\ndocument in your context, and `generate`\ntakes a keep-budget `max_kept_prefill_tokens` — the number of state tokens\nthat are retained for full processing during prefill.\n\nThe 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:\n\n``` python\nimport json\nfrom transformers import AutoTokenizer\nfrom tokengate import load_model, apply_chat_template\n\n# one object: base model + gate + paired adapter; gate layer and base model\n# are read from the checkpoint's own metadata\nmodel = load_model(\"liana313/tokengate-niah-qwen3-1.7b\")\ntok = AutoTokenizer.from_pretrained(model.base_model)\n\n# ~32k tokens of essays with one planted fact about San Francisco\nessays = json.loads(next(open(\"examples/data/niah_32k_samples.jsonl\")))[\"body\"]\n\nprompt = (\"Answer the question using only the document below. \"\n          \"Keep your answer short and direct.\\n\\n\"\n          \"Question: What is the best thing to do in San Francisco?\\n\\n\"\n          \"Document:\\n{state}\\n\\nAnswer:\")\nids = apply_chat_template(tok, prompt, state=essays)\n\nfor k in (None, 2000, 500, 100):          # None = keep every state token\n    out = model.generate([ids], tok, max_kept_prefill_tokens=k)\n    print(k, f\"{out.prefill_s * 1e3:5.0f} ms\", out.answers[0])\n\n# K      prefill   answer\n# all    1471 ms   eat a sandwich and sit in Dolores Park on a sunny day\n# 2000    153 ms   eat a sandwich and sit in Dolores Park on a sunny day\n# 500     157 ms   eat a sandwich and sit in Dolores Park on a sunny day\n# 100     130 ms   eat a sandwich and sit in Dolores Park on a sunny day\n```\n\nOne 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.\n\nTraining 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:\n\n```\npython datasets/niah_kamradt/generate_niah_kamradt.py --output-file niah.jsonl --n-train 12000\n```\n\nThen train:\n\n``` python\nimport json\nfrom tokengate import TokenGateConfig, TokenGateTrainer\n\nniah = [json.loads(l) for l in open(\"niah.jsonl\")]\nrows = lambda split: [{\"article\": f\"{r['body']}\\n\\nQuestion: {r['question']}\",\n                       \"answer\": r[\"answer\"]}\n                      for r in niah if r[\"split\"] == split]\n\nconfig = TokenGateConfig(\n    base_model=\"Qwen/Qwen3-1.7B\",\n    gate_layer=1,                         # layer L whose hidden states the gate scores\n    output_dir=\"e2e_gates/niah-1p7b\",\n    learning_rate=1e-4,\n    epochs=20,\n    k_eval=[2000],                        # in-training eval K sweep\n    # LoRA co-adaptation (the backbone itself stays frozen)\n    lora_rank=16, lora_alpha=32, lora_lr=1e-4,\n    lora_targets=[\"q_proj\", \"v_proj\"],\n    wandb_project=\"token-gate\",           # omit to disable logging\n)\ntrainer = TokenGateTrainer(\n    config,\n    train_dataset=rows(\"train\"),          # [{\"article\": ..., \"answer\": ...}, ...]\n    eval_dataset=rows(\"test\"),\n    task=\"niah_kamradt\",                  # registered prompt + metric (datasets/tasks)\n)\ntrainer.train()                           # → e2e_gates/niah-1p7b/gate_final.pt\n```\n\nEach row is `{\"article\": <the prunable document>, \"answer\": <target>}`\n(`\"state\"` is an alias for `\"article\"`; any HF Dataset with those columns\nworks). The trainer tokenizes the rows, carves a deterministic eval split\nfrom `eval_dataset`, and launches a DDP `torchrun` run on all visible GPUs\n(`num_gpus` overrides), blocking until the final gate and its name-paired\nLoRA land in `output_dir`. Already-prepared data under `datasets/` can be\npassed by name instead — `train_dataset=\"biodex/biodex_reactions\"` — and\nuses the store's own val split.\n\nThe gate trains against randomly sampled keep-budgets, so one trained gate\nserves any K at inference. The Qwen3 defaults train on the key-bias kernel\nfrom [Installation](#-installation) and fail fast without it; set\n`attn_backend=\"flex\"` to train on the no-build kernels (the Qwen3.5 hybrid\ndefaults to flex, and other model families must set it explicitly).\n\nA gate module at an early layer *L* scores every token of the prunable input\n(the **state**) and keeps the top-*K*. The layers above *L* process only\nthose *K* tokens and store only their KV-cache entries — pruned tokens never\nconsume compute or memory above the gate, which is where both the\ntime-to-first-token and the peak-memory savings come from. Everything outside\nthe `{state}` placeholder (instructions, the question, the generated answer)\nis never pruned.\n\nThe top-*K* selection introduces a challenge for end-to-end training since top-*K is non-differentiable.\nTo address this we introduce *Token-gated attention**, a mechanism which re-weights each kept token's\nattention logic using its token importance score, providing a gradient path for e2e training.\nThe gate trains end-to-end with standard next-token prediction, jointly with\nthe model backbone.\n\n`examples/niah_demo.py` downloads the trained gate for Qwen3-1.7B (gate\nlayer 1) from the Hub and runs the classic Kamradt needle-in-a-haystack on\ncommitted ~32k-token Paul Graham essay haystacks: for each keep-budget it\nreports retrieval accuracy, prefill time-to-first-token, and how much of the\nneedle sentence itself survived the pruning:\n\n```\npython examples/niah_demo.py                 # 5 haystacks, K = all/2000/500/100\npython examples/niah_demo.py --k 500 --n 2   # one budget, fewer rows\n```\n\nValidated output (5 haystacks, needle depths 11–100%, one GPU):\n\n```\n kept tokens    acc   TTFT(ms)  speedup  needle kept\n------------------------------------------------------\n         all   1.00     1533.6    1.00x        100%\n        2000   1.00      153.4   10.00x        100%\n         500   1.00      144.3   10.63x        100%\n         100   1.00      137.2   11.18x        100%\n```\n\nEvery needle is retrieved at every budget — even keeping 100 of ~32k tokens, the gate's selection lands squarely on the needle sentence.\n\nA second demo on real data follows the same pattern:\n`examples/biodexr_demo.py` (BioDEX adverse-drug-reaction extraction; mean\nexact-match set F1 0.622 at K=2000 over the committed val articles). Both\npipelines are covered by `python -m tokengate.e2e_gate.tests.test_niah_demo`\nand `...tests.test_biodexr_demo`, and `--gate-attn flex` avoids the\ncustom-kernel build at some speed cost.\n\nFor a full accuracy + latency comparison against standard prefill,\n`benchmarks/run_benchmark.py` measures, on the same samples and base model,\n**accuracy** (the task's own metric) and **TTFT** (prefill-to-first-token,\nbatch size 1) for three arms: standard full prefill of the base model,\nTokenGate with nothing pruned (the accuracy reference), and TokenGate at each\nconfigured K. Each config in `benchmarks/configs/` (niah, repobench, biodexr)\npairs a released checkpoint with the prepared dataset's pinned eval split —\n`--split val` is the 50-row in-training eval partition, `--split test` the\nfull reserved held-out partition:\n\n```\npython benchmarks/run_benchmark.py --config benchmarks/configs/niah.yaml --split test\n```\n\nValidated held-out **test** results (every reserved test row, batch 1, one\nH100; `base` = the unmodified model with full prefill):\n\n| dataset | n | metric | base | TokenGate (no pruning) | TokenGate @K=2000 | TTFT speedup @K=2000 | \n|---|---|---|---|---|---|---|\n| NIAH (16k) | 150 | acc | 0.993 | 1.000 | 1.000 | 6.3× | \n| RepoBench-P | 450 | ES | 0.285 | 0.587 | 0.550 | 3.8× | \n| BioDEX-Reactions | 450 | F1 | 0.056 | 0.284 | 0.292 | 2.1× | \n\nTighter budgets trade little: NIAH stays at 1.000 acc down to K=100 (8.1×),\nRepoBench holds 0.533 ES at K=500 (5.6×). The no-pruning arm doubles as the\nfine-tuning reference — its gap to `base` is what the task-trained\ngate+adapter add. The TTFT gap widens with context length (the paper's\n1.4–11.3× range comes from the longer-context tiers). The 50-row `val`\npartitions reproduce the same ordering (e.g. RepoBench 0.576/0.559,\nBioDEX 0.377/0.368 unpruned/@K=2000).\n\nReleased checkpoints: the benchmark\nset above plus `liana313/tokengate-2wikimqa-qwen3-1.7b`,\n`liana313/tokengate-multinews-qwen3-1.7b`, and the hybrid-attention\n`liana313/tokengate-niah-qwen3.5-2b` — each repo holds the exact gate+LoRA\ncheckpoint the paper's benchmarks used.\n\nThe released checkpoints train from the lane configs under `train-scripts/`\n(one script per dataset, gate layer 1 — each sets every hyperparameter of\nits run and calls the shared driver). Install with the `train` extras (see\n[Installation](#-installation)); the shell drivers read the interpreter from\n`PYTHON` — point it at your venv (`export PYTHON=$PWD/.venv/bin/python`) —\nand log to W&B (`wandb login`, or `WANDB_MODE=disabled`).\n\nThe trainer reads pre-prepped tokenized stores. For BioDEX-Reactions (the\nofficial splits of\n[BioDEX](https://huggingface.co/datasets/BioDEX/BioDEX-Reactions), ~11.5k\ntrain articles):\n\n```\n# one-time download of the official dataset into the HF cache\npython -c \"import datasets; datasets.load_dataset('BioDEX/BioDEX-Reactions')\"\n\n# tokenize into {train,val}.pkl for the base model\npython datasets/biodex/biodex_reactions/prepare.py\n\n# train (the packed memory-mapped token store is built automatically at launch)\nWANDB_PROJECT=<your-project> bash train-scripts/example/run_e2e_gate_biodexr_example.sh\n```\n\nCommon overrides (environment variables; see the script for the full list):\n\n```\nNUM_GPUS=1    bash train-scripts/example/run_e2e_gate_biodexr_example.sh   # single GPU\nMAX_STEPS=50  bash train-scripts/example/run_e2e_gate_biodexr_example.sh   # smoke run\n```\n\nOutputs land under `e2e_gates/<model>/<run-name>/baseline/L<layer>/`:\n`gate_final.pt` + `lora_final.pt` (and SWA variants), periodic resumable\ncheckpoints, `params.json`, and training-curve CSVs/plots. During training\nthe model is evaluated every `EVAL_FREQ` steps on the pinned val split and\nthe reserved held-out test split, at each K in `K_EVAL_VALUES`, alongside\nno-compression and random-selection reference baselines. The released demo\ncheckpoint was trained with exactly this config. (For a synthetic dataset\nwith no external download, see the NIAH lanes:\n`datasets/niah_kamradt/generate_niah_kamradt.py`.)\n\nEach dataset is a directory `datasets/<dataset>/<task>/` with a\n`prepare.py` that emits rows of `{article (the state), answer, id, metadata}`\nand tokenizes them into `tokenized/<model>/{train,val}.pkl`. A training run\npins its dataset through a YAML manifest\n(`tokengate/datasets/manifests/iso_<target>.yaml`) and pins its val/test\nsplit through an eval-split manifest — the explicit source of eval-set sizes.\nScoring lives in `tokengate/datasets/tasks/` (one `Task` class per task: prompt\nformat, metric, judge config).\n\nIf you find TokenGate useful, please cite our [paper](https://arxiv.org/abs/2610.26754)\n\n```\n@article{patel2026tokengate,\n  title={TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning},\n  author={Patel, Liana and Ranjan, Rishabh and Stoica, Ion and Zaharia, Matei and Guestrin, Carlos},\n  journal={arXiv preprint arXiv:2610.26754},\n  year={2026}\n}\n```\n\n", "url": "https://wpnews.pro/news/tokengate-adapting-llms-for-pre-fill-with-learned-skimming-for-11x-speedup", "canonical_source": "https://github.com/token-gate-llms/TokenGate", "published_at": "2026-10-26 07:52:39+00:00", "updated_at": "2026-10-26 08:15:24.757020+00:00", "lang": "en", "topics": ["large-language-models", "ai-research", "machine-learning", "artificial-intelligence"], "entities": ["TokenGate", "arXiv", "Qwen3-1.7B", "FlashAttention-2", "FlexAttention", "Hugging Face transformers", "TRL", "PEFT"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/tokengate-adapting-llms-for-pre-fill-with-learned-skimming-for-11x-speedup", "markdown": "https://wpnews.pro/news/tokengate-adapting-llms-for-pre-fill-with-learned-skimming-for-11x-speedup.md", "text": "https://wpnews.pro/news/tokengate-adapting-llms-for-pre-fill-with-learned-skimming-for-11x-speedup.txt", "jsonld": "https://wpnews.pro/news/tokengate-adapting-llms-for-pre-fill-with-learned-skimming-for-11x-speedup.jsonld"}}