Token都要過FFN:RedKnot Token選擇性稀疏FFN與自適應(yīng)專家Top-K執(zhí)行策略詳解)
并非每個(gè)Token都要過FFNRedKnot Token選擇性稀疏FFN與自適應(yīng)專家Top-K執(zhí)行策略詳解【免費(fèi)下載鏈接】RedKnotEfficient Long-Context LLM Serving with Head-Aware KV Reuse and SegPagedAttention項(xiàng)目地址: https://gitcode.com/gh_mirrors/re/RedKnotRedKnot 是一個(gè)面向長上下文大模型推理的加速方案基于 SGLang核心思想是按需計(jì)算通過Head-Aware KV 復(fù)用、Token 選擇性稀疏 FFN 和自適應(yīng)專家 Top-K三層策略讓不重要的 Token 跳過 FFN 前饋計(jì)算、讓每個(gè) Token 只激活真正需要的專家。在 HotpotQA 長上下文場景下它減少了約71.9% 的 Prefill 計(jì)算量端到端首 Token 延遲TTFT提升2.2×3.1×且答案質(zhì)量不降反升。 為什么 FFN 是長上下文的隱形瓶頸在長上下文推理中大家習(xí)慣盯著注意力Attention做優(yōu)化但 RedKnot 的剖析文檔給出了一個(gè)反直覺的事實(shí)在 2K8K Token 的 Prefill 階段FFN前饋網(wǎng)絡(luò)占到了總耗時(shí)的 57%62%比注意力還大。原因在于注意力的 KV 可以跨請求復(fù)用但 FFN 的矩陣乘法目前是對每個(gè) Token 都要算一遍的硬性成本。RedKnot 的關(guān)鍵洞察是——并非每個(gè) Token 都需要過 FFN。它的推理流程論文 Algorithm 1可以簡化為四步S ← SelectImportantTokens(A)根據(jù)恢復(fù)的注意力信號選出重要 Token集合 SZ[S] ← FFN(Y[S])只有 S 中的 Token 真正執(zhí)行 FFNZ[~S] ← 0其余 Token 走殘差恒等路徑FFN 輸出置零等于直通X_next ← Y Z殘差合并正常進(jìn)入下一層這一機(jī)制的完整實(shí)現(xiàn)位于 sparse_ffn.py與 Head 分類配置、SegPagedAttention 共同組成 RedKnot 的 P0 核心組件。 Token 選擇性稀疏 FFN怎么判斷 Token重要判斷 Token 重要性的信號來自上一層恢復(fù)的注意力輸出每個(gè) Token 取其注意力貢獻(xiàn)向量的L2 范數(shù)作為重要性分?jǐn)?shù)直覺很好理解一個(gè) Token 如果注意力輸出接近零說明它的隱狀態(tài)沒什么變化走殘差直通幾乎無損選 Token 的策略不是固定數(shù)量而是Top-P 式累計(jì)質(zhì)量閾值mass threshold按重要性降序排列累計(jì)到覆蓋總質(zhì)量mass_thresh為止就停止。再疊加兩道保險(xiǎn)保護(hù)機(jī)制作用recent_n如 128最近 N 個(gè) Token無條件保留保證局部上下文新鮮min_keep≥1每序列至少保留 1 個(gè) Token防止整層退化為全直通、凍結(jié)殘差流淺層/質(zhì)量分?jǐn)?shù)為全零的退化行會(huì)自動(dòng)回退到密集 FFN保證數(shù)值安全當(dāng)mass_thresh1.0時(shí)輸出與密集 FFN逐位一致可以先驗(yàn)證等價(jià)性再開啟稀疏。三層調(diào)度淺層密集、中層保守、深層激進(jìn)SparseFFNSchedule把模型按深度分成三檔稀疏力度逐級放大層級判定條件策略淺層layer_idx dense_until全密集 FFN保護(hù)早期殘差流中層dense_until ≤ idx deep_layer_start稀疏使用保守閾值mass_thresh深層idx ≥ deep_layer_start稀疏使用更激進(jìn)的mass_thresh_deep官方文檔中還提供了四檔開箱即用的 FFN 預(yù)設(shè)來自 FFN Presets預(yù)設(shè)dense_untilmass_threshrecent_n定位speed30.364速度優(yōu)先balanced200.5128均衡quality320.7128質(zhì)量優(yōu)先dense801.0128等價(jià)密集基線想自己掃參數(shù)找甜點(diǎn)倉庫提供了單模型加載、多配置掃描的腳本 sweep_sparse_ffn.py一次加載 Qwen3-32B遍歷(dense_until, mass_thresh)網(wǎng)格直接輸出每個(gè)配置的 F1/EM、TTFT、加速比和深層 Token 激活率。 自適應(yīng)專家 Top-K每個(gè) Token 激活多少個(gè)專家對 MoE 模型稀疏還可以發(fā)生在專家路由這一層。RedKnot 的 Pro-0813 配置策略見 scale_policy.py給出了兩套機(jī)制1. 漸進(jìn)式 Top-K 調(diào)度Progressive Top-K Schedule按層深度逐步收縮每 Token 激活的專家數(shù)0-15:6, 16-39:5, 40-60:4淺層0-15 層每個(gè) Token 激活 6 個(gè)專家中層16-39降到 5 個(gè)深層40-60進(jìn)一步降到 4 個(gè)——淺層承擔(dān)更多特征提取工作深層則更惜算。2. 基于路由概率質(zhì)量的自適應(yīng) Top-K不固定 K 值而是按每個(gè) Token 的路由概率累計(jì)質(zhì)量默認(rèn)閾值0.50在候選桶(3, 4, 5, 6)中動(dòng)態(tài)選擇置信度高的 Token 激活少而精的專家概率分散的 Token 多激活幾個(gè)專家兜底。相關(guān)硬件側(cè)參數(shù)也隨模型規(guī)模同步調(diào)整Pro-0813 版本有384 個(gè)路由專家Flash 版為 256 個(gè)在線索引 Top-K 從 512 提升到1024原生每 Token 專家數(shù)為 6。此外Token 級稀疏還配有上下界約束min_full_ratio0.10、max_full_ratio0.50配合固定預(yù)算塊選擇以 128 Token 為塊做 Top-K讓路由后的 MoE 行形成少量連續(xù)區(qū)間而非上千個(gè)零散索引——這直接決定了線上 gather/scatter 路徑的效率且全程在 GPU 上完成不引入主機(jī)同步。 實(shí)測效果計(jì)算量 -71.9%TTFT 提速 2.2×3.1×RedKnot 團(tuán)隊(duì)記錄了完整的優(yōu)化賬本TTFT_OPTIMIZATION.md基線是標(biāo)準(zhǔn) FlashAttention-2 密集 PrefillH200 INT4 Qwen3-32B上下文長度標(biāo)準(zhǔn) FA-2 TTFTRedKnot TTFT加速比16K3.4 s1.3–1.5 s2.2–2.6×40K10.4 s3.3 s3.1×算法級計(jì)算節(jié)省 71.9%3.56× FLOPs 減少注意力 ?83.8%FFN ?79.8%激進(jìn)的 Sparse FFNmass0.2下深層層平均只有約 13% 的 Token真正執(zhí)行 FFN加速比隨長度增長——基線 Prefill 是 O(L2)而 RedKnot 的局部 Head KV 可復(fù)用質(zhì)量不降HotpotQA 的 SQuAD F1/EM 上持平或反超密集基線早期 F1 下降已被定位為解碼路徑 Bug 并修復(fù)一個(gè)誠實(shí)的細(xì)節(jié)FLOPs 減少 3.56×墻鐘只快 2.42×40K。差距來自全局 Head 重新 Prefill 時(shí)不可跳過的投影 GEMM——這是真實(shí)計(jì)算而非開銷是墻鐘加速的理論下限。 如何快速上手最簡單的體驗(yàn)方式是跑 RAG 端到端 Demorag_redknot_demo.py加上--sparse-ffn相關(guān)參數(shù)即可對比稀疏 FFN 前后python examples/redknot/rag_redknot_demo.py \ --model-path Qwen/Qwen3-32B \ --dataset hotpotqa/hotpot_qa --dataset-config distractor --split validation \ --n-samples 1 --n-segments 6 --tokens-per-segment 5000 \ --sparse-ffn --ffn-dense-until 20 --ffn-mass-thresh 0.5 --ffn-recent-n 128 \ --output /tmp/redknot_rag_sparseffn.json輸出 JSON 中會(huì)包含深層層選中 Token 比例、FFN FLOPs 節(jié)省與加速比。若將--ffn-mass-thresh設(shè)為1.0FFN 退化為逐位密集的基線適合先做等價(jià)性校驗(yàn)。也可以在正式服務(wù)中通過獨(dú)立后端啟用 Head 分類的 SegPaged 解碼路徑python -m sglang.launch_server \ --attention-backend redknot \ --redknot-head-config-path /path/to/head_config.json \ --redknot-segpaged-decode --redknot-page-size 64 總結(jié)一張表看懂 RedKnot 的按需計(jì)算機(jī)制解決什么浪費(fèi)關(guān)鍵參數(shù)源碼位置Token 選擇性稀疏 FFN深層 FFN 對不重要 Token 的無效計(jì)算dense_until / mass_thresh / recent_nsparse_ffn.py漸進(jìn)式 自適應(yīng)專家 Top-KMoE 專家激活數(shù)量的僵化層間 K 值調(diào)度、路由質(zhì)量桶scale_policy.pyHead 分類 KV 復(fù)用 SegPaged注意力對局部 Head 的全量重算Head 配置 JSONglobal/localexamples/redknot/README.md長上下文推理提速的下一步不在于算得更快而在于算得更聰明——RedKnot 用稀疏 FFN 與自適應(yīng) Top-K 證明只要選對 Token、選對專家省下的算力可以直接兌現(xiàn)為用戶可感知的響應(yīng)速度?!久赓M(fèi)下載鏈接】RedKnotEfficient Long-Context LLM Serving with Head-Aware KV Reuse and SegPagedAttention項(xiàng)目地址: https://gitcode.com/gh_mirrors/re/RedKnot創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考