LlamaIndex 團隊在部署其最新文檔解析模型 OCRia-1 時發現，推理延遲是客户最關心的問題之一。經過一系列優化，他們成功將單張 B200 GPU 上的 token 生成速度從約 285 tokens/s 提升到約 2800 tokens/s，實現了近 10 倍的改進。這一成果不僅加快了內部基準測試的迭代速度，也讓最終用户能夠更快獲得解析結果。
OCRia 的推理過程主要分為兩個階段：預填充（prefill）和解碼（decode）。預填充是計算密集型操作，而解碼是內存帶寬密集型操作。對於 B200 和該團隊的工作負載，解碼階段嚴重受限於內存帶寬——GPU 花費大量時間在移動模型權重和狀態上，而非實際計算。因此，優化方向聚焦於減少每步解碼移動的數據量、在單次目標模型前向傳播中生成更多 token，以及消除推理引擎中的各種開銷。
首先，團隊採用了工作負載感知的量化策略。他們並未簡單地對所有參數進行統一量化，而是根據各層對量化的敏感度使用混合精度：敏感層保留較高精度，其餘層則更激進地壓縮。通過使用真實輸入分佈進行激活感知量化校準，並評估不同量化組合的精度與速度，最終在視覺編碼器上保持不量化以保護視覺準確度，同時證明量化在目標任務上幾乎無損。他們混合使用了 BF16、FP8 和 NVFP4 量化權重。
其次，投機解碼是提升吞吐的關鍵。傳統自迴歸模型每步只生成一個 token，且每一步都需要讀取全部模型權重，導致在較小批量下計算資源大量閒置。團隊引入了一個小型草稿模型 DFlash，它基於目標模型的隱藏狀態特徵，通過一次前向傳播並行預測多個未來 token。目標模型隨後在一次前向傳播中並行驗證這些 token，並保證接受結果與目標模型分佈一致，從而實現數學上無損的加速。最初的 DFlash 接受率為 9 個 token，升級到 DFlash2 後（增加了局部卷積和候選選擇器），接受率超過 12 個 token。此時，在批量大小為 8 到 12 時，每用户吞吐已從約 200 tokens/s 提升至 1400 tokens/s 以上。
為進一步挖掘性能，團隊深入優化了投機驗證環節。OCRia 架構使用了門控 Delta 網絡（GDN），它不像全注意力層那樣保留 KV 緩存，而是用固定大小的循環狀態總結歷史。在投機解碼中，當草稿 token 被拒絕時，傳統方法需要恢復完整狀態，造成大量內存搬運。團隊適配了 ReplaySSM 技術：保存循環狀態的檢查點以及近期輸入緩衝區，僅在週期刷新時更新完整檢查點；當拒絕發生時，只需調整緩衝區指針來保留已接受前綴，避免全狀態恢復。隨後他們優化了該 kernel，將每層耗時從上游 Triton 實現的 72–80 微秒降低到自定義 CUDA 實現的 23–24 微秒，提速約 3.3 倍。
此外，團隊還進行了算子融合以減少不必要的內存傳輸。例如，將 MLP 中的 SiLU 激活與乘法及其 FP8 量化融合，將門控 RMSNorm 與 FP8 量化融合，並將 NVFP4 激活量化集成到產生激活的 kernel 中。僅 SiLU 與量化融合一項就帶來了 2.3%–2.6% 的吞吐提升。同時，他們還為 GDN 層實現了更快的因果卷積。
CPU 端的準備工作也被證明是延遲的重要來源。推理引擎在每個解碼步驟前都需要構建描述批次的元數據（序列長度、token 偏移、緩存位置等）。由於 vLLM 按 KV 緩存組組織層，而目標模型和草稿模型共有四種不同的緩存佈局（24 個 GDN 層、8 個目標註意力層、5 個草稿滑動窗口層和 1 個草稿全注意力層），分組規則導致生成了 38 個組，每組都觸發一次元數據構建，儘管大量信息是共享的。團隊改為對兼容的構建器類和緩存規格只構建一次共享元數據，然後為其他組提供帶有自身塊表和槽映射的淺拷貝。對於 GDN 元數據，他們將與組無關的計算移到 NumPy 上，並利用草稿請求位於批次前幾行的特點簡化掩碼，最終將每步約 290 次小型 GPU 啓動減少到僅 1 次。
最後，團隊還實施了另外 17 項優化，包括針對模型架構定製的 Fused Triton CUDA kernel，以及選擇特定的現有 kernel 和 cuDNN 組件。每項微優化單獨只能帶來 2%–5% 的吞吐提升，但疊加起來效果顯著。這些優化共同構成了 OCRia 高效推理棧的基礎，並且部分改進（如 ReplaySSM 適配）仍以草稿 PR 形式貢獻給 vLLM 社區。