ORCA:基于持續批處理的LLM推理性能優化技術詳解
大語言模型(LLMs)推理過程中的批處理優化面臨顯著挑戰,這主要源于其推理過程的迭代特性。核心問題在于批處理中的各個請求完成時間存在差異,這導致資源釋放和新請求整合的復雜性顯著提高,特別是在處理不同完成階段的請求時。當批處理中序列的生成長度差異較大時,GPU資源利用率往往不夠理想。如下圖所示,序列1、3和4中,結束符號(EOS)之后的白色區塊直觀地展示了GPU資源的閑置狀態,凸顯了當前批處理方法的效率損失。
ORCA系統創新性地提出了持續批處理概念,通過引入迭代級調度和選擇性批處理機制,有效解決了大語言模型批處理中的關鍵技術挑戰。
特性1 — 迭代級調度
下圖詳細展示了基于迭代級調度的ORCA系統架構及其工作流程。ORCA系統對外提供標準化接口(支持HTTPS或gRPC協議)用于處理推理請求。系統將接收到的請求存儲于請求池中,由調度器進行統一監控和管理。調度器的核心職責包括:從請求池中篩選請求、調度執行引擎執行模型迭代、接收引擎輸出的標記序列,并據此更新請求池狀態。執行引擎負責進行張量運算,支持跨多GPU和多機器的并行處理能力。與傳統批處理方法不同,ORCA的調度器能夠在每次迭代中動態調整處理請求的組合,提供了更靈活的請求處理機制。當某個請求處理完成后,系統將其從請求池中移除,并通過接口返回響應結果。這種方法突破了傳統批處理必須等待所有請求完成全部迭代的限制。
特性2 — 選擇性批處理
選擇性批處理技術通過在張量運算中實現差異化的批處理策略,顯著提升了模型處理的靈活性。這種創新方法特別適用于處理形狀不規則的輸入張量場景,這類張量通常難以直接合并為統一的大規模張量進行批處理。在非注意力機制的矩陣乘法運算和層歸一化等操作中,選擇性批處理采用將不規則張量重構為無批處理維度的二維結構的方式,實現統一的大規模矩陣計算。這種優化策略的可行性基于這些操作的逐標記處理特性,使得來自不同請求的標記可以等同處理。但是對于需要請求專屬鍵值對的注意力機制操作,系統采用分離式批處理策略,確保每個請求獨立處理。注意力操作完成后,輸出結果會重新合并為統一張量用于后續處理。ORCA通過K/V管理器持久化存儲先前注意力操作產生的鍵值對,使增量請求能夠在多輪迭代中復用這些信息。
分布式架構設計
ORCA系統在Transformer模型并行化方面融合了兩種關鍵技術:層內并行與層間并行。層內并行技術(這也是FasterTransformer采用的方法)主要針對矩陣乘法運算(包括線性變換和注意力機制)及其參數在多GPU間進行分配。而層間并行則致力于在多個GPU之間分配Transformer層。ORCA采用均衡分配策略,確保每個GPU負責處理相同數量的Transformer層。例如,對于一個4層GPT模型,可以先劃分為2個層間并行分區,每個分區再細分為3個層內并行分區,最終形成6個GPU的協同計算架構。
FasterTransformer是NVIDIA開發的開源優化庫,專注于加速Transformer模型,支持BERT、GPT-2、GPT-J和T5等主流架構。目前,FasterTransformer的技術演進已轉移至TensorRT-LLM平臺。
ORCA執行引擎采用分布式執行策略,將工作進程分配到不同物理機器上,每個工作進程負責管理一個或多個控制GPU的CPU線程。在具體實現中,每個工作進程專注于處理模型的特定層間分區,而線程數量則由層內并行度決定。當系統為一批請求調度模型迭代時,引擎主控程序首先向第一個工作進程(Worker1)傳遞必要的計算信息,包括標記序列和控制指令。Worker1隨后將這些信息分發給其管理的GPU控制線程,觸發相應的GPU計算內核。同時,Worker1采用非阻塞方式將控制指令轉發給下一個工作進程(Worker2)。最后一個工作進程在確保其GPU任務全部完成后,負責匯總輸出標記并將結果返回給引擎主控程序。這種架構設計實現了跨設備、跨機器的高效分布式計算。
調度算法設計
ORCA調度器采用基于到達時間的優先級機制,在最大批處理大小約束下動態選擇請求進行處理。系統在首次調度請求時會預分配存儲鍵值對所需的內存空間。調度過程中Select
函數確保嚴格遵循請求的時間優先級進行批次選擇。調度器通過評估最大標記數需求,動態驗證新請求的內存資源可用性。在確認內存充足的情況下,系統保證為請求的完整生命周期預留必要的鍵值對緩存空間。整個調度機制采用并發處理模式,支持請求的動態添加和完成請求的及時移除。
ORCA調度器通過任務流水線化技術優化工作進程執行效率。系統不采用傳統的串行批次處理模式,而是維持并發運行批次數與工作進程數的動態平衡。這種策略確保了工作進程的持續高效運轉。例如在三工作進程配置下,最大批處理大小為二的場景中,調度器可以同時處理多個批次(如AB、CD、EF)。在調度第一個批次(AB)后,系統立即安排后續兩個批次(CD和EF)的處理,無需等待AB批次的完成。當AB批次處理完成后,如果這些請求仍具有時間優先性,則會被重新納入調度流程。這種機制顯著提升了系統的資源利用率和整體吞吐性能。
性能評估
得益于創新的迭代級調度和選擇性批處理機制,ORCA系統在保持相同延遲水平的前提下,實現了相比FasterTransformer高達37倍的吞吐量提升,展現了顯著的性能優勢。