今年的兩篇最佳論文一作均為華人。
近日,國際系統領域頂會 MLSys 2025 公佈了最佳論文獎。
今年的大獎頒發給了來自華盛頓大學、英偉達、Perplexity AI、卡耐基梅隆大學的 FlashInfer,以及瑞典查爾摩斯理工大學的《The Hidden Bloat in Machine Learning Systems》。
對此,英偉達第一時間發出祝賀,並表示「FlashInfer」的「LLM 推理核心能力」已經被整合到 vLLM 專案、SGLang 以及自定義推理引擎中。

FlashInfer 最初是華盛頓大學 Paul G. Allen 電腦科學院、卡耐基梅隆大學及陳天奇的創業公司 OctoAI 共同發起的合作研究專案,旨在建立一個靈活的大語言模型(LLM)推理核心庫,提供 LLM GPU 核心的高效能實現,如 FlashAttention、SparseAttention、PageAttention、Sampling 等。
英偉達表示,首個適用於 DeepSeek MLA 的 Blackwell 核心也出自 FlashInfer。
FlashInfer 獨立於引擎,高度優化,並且易於擴充套件以支援新技術,例如鍵值快取複用演算法。現在,它已成為一個蓬勃發展的開源專案,可生產部署,並得到了來自整個 AI 系統社羣研發團隊的貢獻。

-
論文:FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving -
論文連結:https://arxiv.org/abs/2501.01005 -
專案主頁:https://flashinfer.ai/ -
GitHub 倉庫:https://github.com/flashinfer-ai/flashinfer
Transformer 結構以注意力機制(Attention Mechanism)為核心,是大型語言模型(Large Language Models,LLMs)的基礎。隨著模型規模的不斷擴大,高效的 GPU 注意力計算模組(attention kernel)對於實現高吞吐和低延遲(即更快、更高效)的推理至關重要。面對多樣化的 LLM 應用場景,亟需靈活且高效能的注意力解決方案。
研究團隊提出了「FlashInfer」:一款可定製且高效的注意力推理引擎,專為 LLM 部署優化設計。其主要特點如下:
優化記憶體訪問並減少冗餘:FlashInfer 通過採用塊稀疏格式(block-sparse format)與可組合格式(composable formats)解決鍵值快取(KV-cache)儲存的異構性問題。
可定製的注意力計算模板:支援基於即時編譯(Just-In-Time,JIT)的靈活配置,以應對不同應用需求。
高效的任務排程機制:FlashInfer 引入了基於負載均衡的排程演算法,既能適應使用者請求的動態變化,又與要求靜態配置的 CUDAGraph 保持相容性。
經過核心級(kernel-level)及端到端(end-to-end)的完整評估,FlashInfer 在多種推理場景下顯著提升了計算效能:與當前最先進的 LLM 部署方案相比,FlashInfer 在 token 間延遲方面提高顯著,相較通用編譯器後端,能減少 29% 至 69% 的 inter-token 延遲;在長上下文推理任務中延遲降低 28% 至 30%;在並行生成場景下,推理速度提升達 13% 至 17%。
系統設計

近年來,為提升記憶體效率,KV-Cache 儲存機制(如 PageAttention、RadixAttention)開始採用非連續記憶體佈局,以塊或 token 為最小儲存單元。
FlashInfer 證明,這些不同的非連續 KV-Cache 數據結構都可以統一抽象建模為塊稀疏矩陣格式(如圖 2 所示)。

在此基礎上,FlashInfer 進一步引入組合式稀疏格式(Composable Sparse Formats)來提升記憶體效率。與單一固定塊大小的格式不同,組合式格式允許在同一稀疏矩陣中靈活採用多種塊稀疏形式,從而顯著提高記憶體利用率。
單一塊格式的侷限在於其固定大小(特別是行塊大小 Br)導致的碎片化風險和共享記憶體互訪限制。組合式格式則可以根據資料特徵(如共享字首形成邏輯上的稠密子矩陣)選用合適的塊形式(例如為稠密子矩陣選用較大的 Br)。
如圖 3 所示,這種方式無需資料移動,僅通過索引即可實現基於共享記憶體的高速訪問,進一步優化記憶體效率。
FlashInfer 的 KV-Cache 具體實現採用塊稀疏行格式(Block-Sparse Row, BSR),其塊大小 (Br,Bc) 可根據應用動態配置,核心支援對任意配置的高效處理。
2、計算抽象
在計算抽象層面,FlashInfer 基於 CUDA/CUTLASS 模板實現了對 FlashAttention 演算法的高效支援。該實現相容 NVIDIA Turing 到 Hopper 全系架構(即 sm75 至 sm90a),能夠原生處理稠密和塊稀疏的矩陣運算,與上層統一的 KV-Cache 儲存格式緊密結合。
具體實現上,FlashInfer 針對 Ada(sm89)及以下的架構採用了 FlashAttention2(FA2)演算法,而對於最新的 Hopper 架構則使用了優化後的 FlashAttention3(FA3)演算法,以充分利用各代硬體特性。
3、計算核心優化與硬體適配
爲了最大化 GPU 計算單元(如 Tensor Core)的效率,高效的資料載入是關鍵。FlashInfer 設計了專門機制,將 KV 快取資料從全域性記憶體快速載入共享記憶體,尤其針對其支援的塊稀疏及非連續格式。載入時,根據塊是稀疏或稠密採用不同地址計算(利用 BSR 索引 vs 仿射變換)。
KV 快取最後一個維度(即注意力頭維度 d)保持連續,這使得高效的合併訪問(coalesced access)成為可能,顯著提升訪存效率。數據傳輸利用 GPU 的非同步複製及硬體加速器(如 Hopper TMA,如圖 4 所示,需注意其適用性),並自動回退通用路徑。
載入到共享記憶體後,FlashInfer 基於 FlashAttention 演算法的計算核心進一步優化。為適應大語言模型執行時多變的計算強度和不同 GPU 架構資源(如 Ada 共享記憶體限制),核心支援多種不同的 tile 尺寸組合。
這解決了傳統固定尺寸在短序列解碼等場景效率不高的問題(例如在 A100 上預填充表現最佳的 (128,64))。
FlashInfer 提供一系列 tile 範圍,並結合啓發式策略自適應選擇最優尺寸:例如,基於平均查詢長度選擇行塊大小 Br,或通過資源消耗模型最大化 SM 利用率。這種方法確保 FlashInfer 在不同負載和硬件平臺均能高效能執行。
4、注意力變體靈活性、使用者介面與動態排程
當前大語言模型使用了大量注意力變體,為每個變體定製高效能 CUDA 核心難以持續。借鑑 Flex Attention,FlashInfer 構建了可自定義的 CUDA 模板和 JIT 編譯器。
使用者只需通過定義一系列函式子模組(functor)來指定變體邏輯(如 Query/Key/Value/Output/Logits 的變換與掩碼),JIT 編譯器即可生成優化的變體專用核心。
這種設計使得 FlashInfer 能靈活支援自定義掩碼、滑動視窗、非 Softmax 變體(如 FlashSigmoid,如圖 5 所示其對映方式)以及將 RoPE、歸一化、投影等前處理操作融合進 Attention 核心,實現高效計算。使用者甚至可以直接傳入 CUDA 程式碼。
生成的 CUDA 程式碼由 JIT 編譯器(如 PyTorch 的 JIT)編譯為自定義運算元,並通過 DLPack 等介面支援框架無關部署。
FlashInfer 提供面向使用者的 API,通過初始化 wrapper、傳入變體描述和任務資訊來完成核心的 JIT 編譯和快取。對於組合格式,會建立對應 wrapper。根據 KV 快取配置(如平均查詢長度),FlashInfer 在後臺將不同核心編譯並捕獲到 CUDAGraphs 中。執行時,服務框架動態選擇最合適的 CUDAGraph 執行。
FlashInfer 採用 Plan/Run 模式處理不規則負載:
Plan 函式在 CPU 上執行,處理序列長度等資料,生成負載均衡的排程方案(可快取複用);
Run 函式在 GPU 上執行,接受輸入和排程方案資料,執行注意力計算(可被 CUDAGraph 捕獲)。這種 Plan/Run 分離設計借鑑了不規則平行計算中的 Inspector-Executor (IE) 模式。
作者介紹
FlashInfer 的第一作者葉子豪(Zihao Ye)是華盛頓大學(University of Washington)的博士生,導師為 Luis Ceze。
葉子豪與陳天奇在機器學習編譯器方面有密切合作,也曾在英偉達實習,合作者為 Vinod Grover,他還曾獲得 2024-2025 年度英偉達研究生獎學金。
加入華盛頓大學之前,葉子豪在 AWS 工作了兩年,期間跟隨王敏捷、張崢研究機器學習系統。他的本科畢業於上海交通大學 ACM 班。
葉子豪目前的主要研究方向是機器學習編譯器和稀疏計算。
The Hidden Bloat in Machine Learning Systems
共同獲得今年最佳論文獎項的還有來自瑞典查爾摩斯理工大學的一篇論文。
論文:The Hidden Bloat in Machine Learning Systems
連結:https://arxiv.org/abs/2503.14226
軟體臃腫(software bloat)是指軟體在執行時未使用的程式碼和功能。對於機器學習系統而言,臃腫是其技術債務的主要原因,會導致效能下降和資源浪費。
在該工作中,作者提出了 Negativa-ML,可通過分析 ML 框架的共享庫來識別和消除其中的臃腫。作者提出的方法包括一些新技術,用於檢測和定位裝置程式碼中不必要的程式碼——這是現有研究(主要關注主機程式碼)所忽略的關鍵領域。
作者使用四種流行的 ML 框架,在 300 個共享庫的十個工作負載中對 Negativa-ML 進行了評估。結果表明,ML 框架在裝置和主機程式碼端都存在嚴重的臃腫問題。
平均而言,Negativa-ML 可將這些框架中的裝置程式碼大小減少高達 75%,主機程式碼大小減少高達 72%,最終檔案總大小減少高達 55%。
裝置程式碼是 ML 框架中臃腫的主要來源。通過減少膨脹,新方法分別可將峰值主機記憶體使用量、峰值 GPU 記憶體使用量和執行時間減少了多達 74.6%、69.6% 和 44.6%。
Negativa-ML 概覽,該研究中提出的元件以黃色矩形突出顯示。
作者介紹
本文一作 Huaifeng Zhang 是查爾姆斯理工大學的博士生,主要研究方向是 MLsys、雲端計算、分散式系統和區塊鏈。
參考連結:
https://mlsys.org/virtual/2025/awards_detail