斯坦福大學與加州大學伯克利分校的研究人員釋出了 TokenGate，這是一個面向長上下文預填充（prefill）的模組。TokenGate 透過可學習的 token 剪枝，讓現有的大語言模型能夠對輸入進行“快速瀏覽”：只有真正重要的 token 會獲得完整的逐層處理。研究團隊報告稱，在不損失準確率的情況下，首個 token 生成時間（TTFT）可提升 1.4 倍至 11.3 倍，預填充階段的峰值記憶體下降 1.5 倍至 3.4 倍。程式碼與檢查點已在 GitHub 上以 Apache-2.0 許可證釋出。
TL;DR（要點速覽）
MLP 門控：在 Qwen3-1.7B 或 Qwen3.5-2B 上加入一個小型 MLP 門控模組和 LoRA（r=16），測試上下文長度最高達 128k token。
執行環境：論文中的計時實驗使用 1 塊 A100、批大小為 1；訓練使用 4 塊 A100。
效能表現：在 5 項任務上相對 6 個基線方法取得帕累托最優。
最佳情況：NIAH 16k 任務上 TTFT 快 11.3 倍、記憶體降低 3.42 倍；在 128k 上下文中 TTFT 增益可達 13.0 倍。
最差情況：在部分任務上加速比降至 1.4 倍。
TokenGate 是什麼？
TokenGate 是一個可直接嵌入（drop-in）的元件，為預訓練大語言模型增加可訓練的 token 剪枝能力。標準 Transformer 會把每一個 token 送入每一層，而注意力開銷隨序列長度呈二次增長。免訓練的剪枝方法成本低，但在高剪枝率下會損失準確率；而修改模型架構通常需要從頭預訓練。
TokenGate 的可學習 token 剪枝如何工作？
前向傳播被拆分為三個階段：
標準前向：門控層 L 以下的層照常處理全部 N 個 token。
門控模組：一個 MLP 根據第 L 層注意力後的殘差為每個 token 打分，top-K 選擇器保留得分最高的 token。
剪枝後前向：L 層以上的層只處理這 K 個 token，且不會為被剪掉的 token 建立 KV 條目。
只有提示詞中的文件部分可被剪枝，問題本身永遠不會被剪枝。
什麼是 token 門控注意力？
Top-K 選擇不可微，這會阻斷端到端訓練。TokenGate 借用了混合專家（MoE）路由器的技巧：每個被保留 token 的得分被轉化為門控層以上注意力 logits 的偏置，從而為評分器提供梯度回傳路徑。一個可學習的逐頭增益和可學習溫度進一步對此進行細化，溫度用於校正每個查詢所看到的鍵變少的問題。訓練使用普通的下一 token 預測損失，每一步對全上下文損失和隨機預算損失取平均，因此一個訓練好的門控可在多種預算下通用。
TokenGate 在基準測試上表現如何？
研究團隊在 RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX 以及大海撈針（NIAH）任務上測試了 Qwen3-1.7B，基線包括 FastKV、SlimInfer、DASH、GemFilter、FlexPrefill 和 MInference，均在相同的 LoRA 骨幹上進行對比。
相對完整預填充：TTFT 快 1.4 倍至 11.3 倍，記憶體降低 1.5 倍至 3.4 倍。
相對次優基線：TTFT 快 1.2 倍至 6.4 倍，記憶體最多降低 2.1 倍。
2WikiMultihopQA：在達到完整上下文準確率 95% 的情況下，預填充快 3.0 倍、記憶體減少 2.7 倍。
更長上下文：在 16k 上訓練的門控在 32k 上達到 12.8 倍，在 64k 和 128k 上達到 13.0 倍；記憶體節省維持在約 3.43 倍。
在相同層上，可學習選擇器優於 FastKV 的啟發式方法；即便隨機 K 選擇在 TokenGate 骨幹上得分也更高，研究團隊將此歸功於骨幹協同適配。
它能在混合注意力模型上工作嗎？
可以。Qwen3.5-2B 混合了 Gated DeltaNet 層與全注意力層。在 NIAH 32k 上，第 3 層使用 TokenGate 達到與完整預填充相當的準確率，執行速度比完整預填充快 3.8 倍，比 FastKV 快 2.1 倍。在第 1 或第 2 層進行門控可達到 7.8 倍加速，準確率差距在 6% 以內。FastKV 只能剪枝注意力層。
TokenGate 與 FastKV、SlimInfer、GemFilter、DASH 的對比
特性對比：
Token 選擇：TokenGate 使用可學習 MLP 門控；FastKV 基於最後查詢的注意力；SlimInfer 採用塊注意力、多階段；GemFilter 在過濾層使用最後查詢注意力；DASH 基於隱狀態穩定性。
是否需要訓練：TokenGate 需要（門控 + LoRA r=16）；其餘四個均不需要。
TokenGate 的 Qwen3-1.7B 執行中的剪枝層：113（預設）；FastKV 為 8、17、25；GemFilter 為 12；DASH 為 11。
是否在 Gated DeltaNet 層剪枝：TokenGate 可以；FastKV 不可以（僅注意力層）；SlimInfer、GemFilter、DASH 未披露。
自報最佳預填充增益（各自設定下）：TokenGate 為 11.3 倍 TTFT、3.4 倍記憶體；FastKV 為 1.82 倍預填充；SlimInfer 為 2.53 倍 TTFT（LLaMA3.1-8B，RTX 4090）；GemFilter 相比 SOTA 快 2.4 倍、GPU 記憶體少 30%；DASH 未披露。
程式碼：TokenGate、FastKV、SlimInfer、GemFilter 均已在 GitHub 上，DASH 僅公佈。
注意：各項自報增益使用不同實驗設定，不能直接橫向比較。剪枝層資訊來自 TokenGate 論文附錄 C.2。
如何部署 TokenGate？
GitHub 倉庫需要 Python 3.10+ 和 CUDA PyTorch 2.9+。可選的 FlashAttention-2 key-bias 核心可加速稠密 Qwen3，開銷為普通呼叫的 1.0 倍至 1.1 倍；不使用時推理回退到 FlexAttention。
使用 {state} 佔位符標記可剪枝的文件部分。generate 接受 max_kept_prefill_tokens 預算引數。檢查點覆蓋 Qwen3-1.7B 上的全部 5 項任務，另有一個 Qwen3.5-2B 的 NIAH 門控。示例：liana313/tokengate-niah-qwen3-1.7b。
訓練方式與 TRL/PEFT 配方一致，每一步比普通 LoRA 慢 1.3 倍；5 項任務中有 4 項收斂更快，平均訓練時間為 1.8 個 A100 小時，而普通 LoRA 為 2.2 個 A100 小時。
關鍵要點
TokenGate 透過小型門控和 LoRA 為現有大語言模型加入可學習 token 剪枝。
預填充速度提升 1.4 倍至 11.3 倍，峰值記憶體降低 1.5 倍至 3.4 倍。
token 門控注意力以 MoE 路由器的方式使 top-K 選擇可訓練。
在 TTFT 上比次優剪枝基線最高快 6.4 倍。
它可在混合架構的 Qwen3.5-2B 上工作，而 FastKV 無法剪枝 DeltaNet 層。
可查閱論文、GitHub 倉庫和 Hugging Face 論文頁面。