打破記憶體牆：在標準 GPU 上實現 3,000 tokens/s 的推理速度

目前 LLM 推理的狀態通常以總吞吐量來衡量——即一台伺服器在數百名使用者下可以生成多少總 token。雖然這對雲端供應商來說是一個關鍵指標，但對於下一代 AI：自主代理（autonomous agents）而言，這並非正確的指標。代理工作流本質上是序列化的；代理必須在一個循環中進行檢查、規劃、編輯和測試，其中每一步都取決於前一步。在這種情況下，唯一重要的指標是單次請求的解碼速度。

Kog AI 發布了 Kog Inference Engine (KIE) 的技術預覽版，證明了標準資料中心 GPU 在處理 2B 模型時，單次請求可以達到 3,000 tokens/s (在 8× AMD MI300X 上) 和 2,100 tokens/s (在 8× NVIDIA H200 上) 的速度。透過將模型架構、運行時（runtime）和底層 GPU kernel 視為一個共同設計的系統，Kog 正試圖縮小通用 GPU 與專用推理晶片之間的差距。

瓶頸：是記憶體頻寬，而非運算能力

對於單次請求解碼（batch size 1），主要的限制因素並非每秒浮點運算次數 (FLOPS)，而是記憶體頻寬。在自回歸解碼（autoregressive decoding）中，GPU 必須在生成每一個 token 時，將每個活動的模型權重從高頻寬記憶體 (HBM) 移動到運算處理器。

這會造成低算術強度（low arithmetic intensity）的情境，使得系統幾乎完全受限於權重的串流速度。雖然像 H200 這樣的現代 GPU 擁有巨大的理論頻寬，但大多數推理堆疊（inference stacks）都無法有效利用它。在 8× H200 節點上，FP16 模式下 2B 模型理論上的「光速」上限大約是 7,700 tokens/s。大多數框架的操作速度遠低於此，這表明瓶頸不在於硬體，而在於軟體堆疊。

遺失的微秒數

在 3,000 tokens/s 的目標速度下，每個 token 的總時間預算僅為333 微秒。在標準的推理堆疊（例如 vLLM 或 TensorRT-LLM）中，這項預算會被通常在較高 batch size 時可以忽略不計、但對於單次請求延遲而言卻是災難性的開銷所消耗：

  • Kernel Launch Overhead: 啟動和清理一個 kernel 可以花費約 4.5µs。如果一個 Transformer 層需要十個 kernel，二十五層就會在進行任何實際運算之前，產生超過 1,100µs 的開銷。
  • CPU Scheduling: 主機端邏輯和 GPU-CPU 通訊會在採樣（sampling）和執行中引入延遲。
  • Synchronization: GPU 全域同步和 GPU 間的集體通訊（例如 tensor parallelism 中的 AllReduce）會為每次操作增加數微秒。
  • Topology Ignorance: GPU chiplet 架構內的非均勻記憶體存取 (NUMA) 效應可能會引入延遲偏差。

Kog 的解決方案：共同設計與 Monokernel

為了奪回這些遺失的微秒數，Kog 在關鍵路徑中避開了 PyTorch 或 Triton 等高階框架，轉而選擇手工打造的 GPU 程式碼 (CUDA/PTX 和 HIP/CDNA ISA)。他們的方法依賴於三個核心創新：

1. The Monokernel Runtime

Kog 不再執行一系列的 kernel，而是使用一個持續性 monokernel。這是一個單一的 GPU 常駐程式，負責處理整個解碼路徑——MatMul、attention、normalization 和 sampling——而無需將控制權交還給 CPU。這完全消除了 kernel launch overhead 和主機端調度。

2. KCCL (Kog Collective Communication Layer)

標準的通訊函式庫是針對吞吐量優化的，而非微秒級的延遲。KCCL 是一個在組合語言層級進行調優的自定義層，旨在將 AllReduce 延遲控制在 3µs 以下，確保 GPU 間的通訊不會阻礙記憶體串流過程。

3. Laneformer Architecture

Kog 引入了 Laneformer 架構，利用了延遲張量並行 (Delayed Tensor Parallelism, DTP)。DTP 改變了解碼過程的依賴結構，允許跨裝置通訊與運算重疊，而不是阻塞關鍵路徑。

擴展至前沿 MoE 模型

雖然 3,000 tokens/s 的基準測試使用的是 2B 模型，但 Kog 在主張，同樣的原理也適用於大型混合專家模型 (MoE)。因為 MoE 模型在每個 token 時僅會啟動其總參數的一小部分，頻寬瓶頸取決於「活動」參數，而非總參數。

Kog 估計，透過將他們的 monokernel 和 KCCL 優化應用於前沿模型（例如 DeepSeek 或 Qwen），他們可以在標準資料中心 GPU 上將輸出速度推升至 1,000–5,000 tokens/s 的範圍，特別是隨著下一代硬體（如 NVIDIA Rubin 或 AMD MI450）增加 HBM 頻寬。

批判性觀點與權衡

儘管基準測試表現優異，社群對於此技術的實際應用提出了幾個重要觀點：

  • 模型實用性 vs. 速度: 一些使用者指出，如果模型缺乏前沿模型的智能，高 token 吞吐量是沒有意義的。正如一位評論者所說，「3.4k tok/s 的純粹胡言亂語」並非一個可行的產品。
  • Hardware Definitions: 對於何謂「標準 GPU」存在爭議，批評者指出 8× H200 節點是頂尖的資料中心硬體，而非易於取得的「標準」GPU。
  • Scaling Concerns: 懷疑論者質疑，在 2B 模型上看到的線性增益是否能轉化為 30B+ 模型，因為記憶體壓力與通訊開銷會以不同的方式擴展。

最終，Kog 的工作凸顯了 AI 基礎設施領域的一個根本性轉變：隨著我們從聊天機器人轉向自主代理，產業必須從追求最大化總吞吐量，轉向追求最小化單個、序列化思維鏈的延遲。

SUMMARY: Kog AI 推出了一款共同設計的推理引擎，透過消除軟體瓶頸，在資料中心 GPU 上實現了巨大的單次請求解碼速度。

TITLE: 打破記憶體牆：在標準 GPU 上實現 3,000 tokens/s 的推理速度

Sources