葉子豪、陳天奇等人開源專案FlashInfer入選，MLSys2025最佳論文獎公佈

機器之心報道

編輯：澤南、+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 模式處理不規則負載：