機器之心報道

編輯：澤南、+0

今年的兩篇最佳論文一作均為華人。

近日，國際系統領域頂會 MLSys 2025 公布了最佳論文獎。

今年的大獎頒發給了來自華盛頓大學、英偉達、Perplexity AI、卡耐基梅隆大學的 FlashInfer，以及瑞典查爾摩斯理工大學的《The Hidden Bloat in Machine Learning Systems》。

對此，英偉達第一時間發出祝賀，并表示「FlashInfer」的「LLM 推理內核能力」已經被集成到 vLLM 項目、SGLang 以及自定義推理引擎中。

FlashInfer

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%。

系統設計

FlashInfer 系統設計概覽：在編譯時提供注意變量規范、任務信息和鍵值緩存布局細節，用于 JIT 編譯；在運行時輸入序列長度信息，用于動態調度。

1、優化的 KV-Cache 存儲結構

近年來，為提升內存效率，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