如何吃透 oMLX specprefillGemma4 assistant drafter 加速本地推理的完整解析【免費(fèi)下載鏈接】omlxLLM inference server with continuous batching SSD caching for Apple Silicon — managed from the macOS menu bar項(xiàng)目地址: https://gitcode.com/GitHub_Trending/om/omlxoMLX 是一款為 Apple Silicon 打造的 LLM 推理服務(wù)器支持連續(xù)批處理continuous batching與 SSD 緩存并可通過 macOS 菜單欄統(tǒng)一管理。它最讓人眼前一亮的性能特性就是本文的主角specprefill 稀疏預(yù)填充與基于Gemma4 assistant drafter的 VLM MTP 投機(jī)解碼。這篇文章帶你從機(jī)制到落地完整看懂 oMLX 是如何在 M 系列芯片上把長(zhǎng)提示詞的首字延遲打下來的。什么是 SpecPrefill兩個(gè)階段的核心思路SpecPrefill 要解決的問題很明確長(zhǎng)提示詞的prefill預(yù)填充階段是整個(gè)請(qǐng)求中最耗時(shí)的一段尤其是系統(tǒng)提示詞之后的對(duì)話部分。oMLX 的做法是把這件事拆成兩個(gè)階段全部代碼集中在 omlx/specprefill/ 目錄下階段負(fù)責(zé)模塊作用① Draft 打分omlx/specprefill/draft.py用輕量草稿模型給全部待處理 token 打分選出最值得保留的子集② 準(zhǔn)入策略omlx/specprefill/policy.py決定哪些請(qǐng)求值得走 SpecPrefill超過閾值才啟用③ 目標(biāo)規(guī)劃omlx/specprefill/planning.py從選中的 token 索引推導(dǎo)出目標(biāo)模型的預(yù)填充計(jì)劃④ Target 稀疏預(yù)填充omlx/specprefill/target.py目標(biāo)模型只預(yù)填充系統(tǒng)提示詞 被選中的稀疏對(duì)話 token準(zhǔn)入策略不是所有請(qǐng)求都值得加速policy.py 中的plan_specprefill_scoring有兩道準(zhǔn)入門檻剩余 token 數(shù)量必須超過閾值再排除已緩存的系統(tǒng)前綴后剩余數(shù)量仍要超過閾值。兩個(gè)條件都滿足請(qǐng)求才會(huì)進(jìn)入 SpecPrefill 流程——這保證了短請(qǐng)求不會(huì)被加速邏輯拖慢加速收益始終大于開銷。Draft 階段草稿模型替目標(biāo)模型預(yù)讀draft.py 中的run_specprefill_draft_scoring會(huì)先嘗試從草稿前綴緩存中恢復(fù)可復(fù)用的 KV 狀態(tài)然后分塊對(duì) token 打分并把打分進(jìn)度實(shí)時(shí)上報(bào)給 prefill 進(jìn)度跟蹤器界面上看到的specprefill_*階段標(biāo)記就來自這里。打分與選塊的實(shí)際計(jì)算由 omlx/patches/specprefill.py 提供。Target 階段只算值得算的 tokentarget.py 中的run_specprefill_target_prefill執(zhí)行真正的稀疏預(yù)填充系統(tǒng)提示詞全量計(jì)算對(duì)話部分只計(jì)算 draft 階段選中的稀疏 token 子集。它還支持一個(gè)細(xì)節(jié)優(yōu)化Issue #2177靜態(tài)系統(tǒng)前綴復(fù)用——第一個(gè)請(qǐng)求把系統(tǒng)提示詞的 KV 狀態(tài)存入分層前綴緩存后續(xù)請(qǐng)求直接恢復(fù)完全跳過系統(tǒng)部分的重算。Gemma4 assistant drafterVLM MTP 運(yùn)行時(shí)機(jī)制如果說 SpecPrefill 加速的是讀入那么 assistant drafter 加速的是輸出。oMLX 通過VLM MTPMulti-Token Prediction投機(jī)解碼讓視覺語(yǔ)言模型在 decode 階段一次產(chǎn)出多個(gè) token。它是什么gemma4_assistant是一個(gè)專為 Gemma 4 視覺語(yǔ)言模型訓(xùn)練的草稿模型model_type為gemma4_assistant屬于外部 drafter 家族之一另一類是 Qwen 3.5/3.6 的qwen3_5_mtp。兩者在 mlx-vlm 中統(tǒng)一解析為draft_kindmtp共享同一套 MTP 輪次循環(huán)。運(yùn)行時(shí)如何工作核心封裝在 omlx/speculative/vlm_mtp.pyPrefill 階段目標(biāo) VLM 以return_hiddenTrue、return_shared_kvTrue完成預(yù)填充產(chǎn)出末位 hidden state、共享 KV 快照與首個(gè) bonus tokenDrafter 掛載引擎加載時(shí)把 assistant drafter 掛在模型上見 omlx/engine/vlm.py 中的_vlm_mtp_drafter字段Decode 輪次每輪 drafter 先打草出若干個(gè)候選 token目標(biāo)模型一次性驗(yàn)證接受多少算多少——驗(yàn)證失敗則回滾 KV正確性零損失調(diào)度接管omlx/scheduler.py 檢測(cè)到 drafter 已掛載后把請(qǐng)求路由進(jìn) MTP 解碼路徑繞開常規(guī) BatchGenerator 的單 token 循環(huán)。該封裝層刻意把 mlx-vlm 的內(nèi)部符號(hào)隔離在一個(gè)文件內(nèi)——上游 API 變化時(shí)只需改這一處這也是 oMLX 補(bǔ)丁體系的典型設(shè)計(jì)思路。如何開啟模型設(shè)置里的三個(gè)開關(guān)在 oMLX 的模型設(shè)置中omlx/model_settings.py 定義了全部配置項(xiàng)vlm_mtp_enabled總開關(guān)啟用 VLM MTP 投機(jī)解碼vlm_mtp_draft_modelassistant drafter 的路徑或模型庫(kù)例如gemma-4-26B-A4B-it-assistantvlm_mtp_draft_block_size每輪起草的 token 數(shù)留空則使用 mlx-vlm 默認(rèn)值。?? 一個(gè)需要注意的互斥規(guī)則MTP 解碼路徑繞過了 logits processors因此它無(wú)法與思考預(yù)算、結(jié)構(gòu)化輸出等處理器類設(shè)置同時(shí)啟用。model_settings.py 中的vlm_mtp_processor_conflicts會(huì)自動(dòng)檢測(cè)沖突并在加載時(shí)關(guān)閉vlm_mtp_enabled并給出日志提示避免配置靜默失效。相關(guān)測(cè)試與延伸閱讀想驗(yàn)證自己的環(huán)境是否跑通可以關(guān)注這些測(cè)試tests/test_specprefill.py 與 tests/test_specprefill_draft.pySpecPrefill 主流程與 draft 打分tests/test_vlm_mtp.py、tests/test_vlm_mtp_thinking_budget.pyassistant drafter 的 MTP 輪次與沖突處理tests/integration/test_specprefill_static_prefix_real_model.py靜態(tài)系統(tǒng)前綴復(fù)用的真實(shí)模型集成測(cè)試。 總結(jié)一句SpecPrefill 用草稿模型幫目標(biāo)模型跳讀長(zhǎng)提示詞Gemma4 assistant drafter 用 MTP 投機(jī)解碼幫目標(biāo)模型搶跑生成——前者省 prefill后者省 decode兩者疊加正是 oMLX 在 Apple Silicon 上實(shí)現(xiàn)快而穩(wěn)推理的底層武器?！久赓M(fèi)下載鏈接】omlxLLM inference server with continuous batching SSD caching for Apple Silicon — managed from the macOS menu bar項(xiàng)目地址: https://gitcode.com/GitHub_Trending/om/omlx創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考