TokenGate v0.1.0

Initial public release, accompanying the paper TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning.

Highlights

• load_model / generate inference API with max_kept_prefill_tokens — one gate serves every keep-budget K; the README quickstart runs as-is on the shipped sample haystack.
• TRL/PEFT-style training interface (TokenGateConfig + TokenGateTrainer) plus the paper's per-dataset training lanes.
• Benchmark harness (accuracy + TTFT vs standard prefill) with per-dataset configs and frozen eval splits.
• Six released checkpoints on Hugging Face — the exact gates the paper's benchmarks used: NIAH, BioDEX, 2WikiMQA, MultiNews, RepoBench (Qwen3-1.7B), and NIAH on Qwen3.5-2B.
• Optional FlashAttention-2 key-bias kernel (python kernels/fa2bias/build.py); inference falls back to FlexAttention automatically without it.

Licensed under Apache-2.0.