跳到主要內容
AI News HubLIVE
站內改寫2 分鐘閱讀

Cohere 的 North Mini Code 解碼大核心服務引擎

文章摘要

Cohere 釋出了圍繞解碼大核心(decode megakernel)構建的 North Mini Code 服務引擎。在單張 H100 上使用 BF16,端到端吞吐比 vLLM 快 1.25–1.41 倍;批大小為 1 時達到 62% 的記憶體頻寬理論峰值,且支援生產級特性與工具呼叫,程式碼已開源。

Cohere 的 North Mini Code 解碼大核心服務引擎
回報錯誤

更正管道尚未開通,可先複製下方文章資訊留存。

查看更正說明
直接讀正文

Cohere 今日釋出了一個面向 North Mini Code 模型、以解碼大核心(decode megakernel)為核心的服務引擎。官方給出的測試結果顯示:在單張 H100 上採用 BF16 精度,其端到端吞吐比 vLLM 快 1.25~1.41 倍;並且這一套引擎的程式碼已開源在 GitHub 上,方便開發者直接檢視或複用。

目前大多數 LLM 服務框架仍把每一次前向計算拆成一連串獨立核心:先啟動 QKV 投影,等待;再啟動注意力,等待;接著啟動 MoE,再等待。單獨看每一個核心都足夠高效,問題出在核心之間的等待上。在小批次解碼時,GPU 有相當大比例的時鐘週期並不是在計算,而是在等上一個核心結束、等下一個網格被排程。自迴歸解碼,尤其是小批次場景,從根本上說是記憶體受限的:每一步都要從 HBM 讀取大量引數,但真正完成的計算量卻不多。

North Mini Code 是一個 30B 模型,每個 token 啟用 3.3B 引數;在 BF16 下,每解碼一步大約要從 HBM 搬運 6.6 GB 權重,外加 8K 上下文時約 0.5 GB 的 KV cache。H100 的 HBM 頻寬是 3.35 TB/s,因此理論上的光速(Speed of Light)約為 470 token/s;而 vLLM 只能跑到 185 token/s,僅為光速的 39%。

大核心(megakernel)正是為縮小這個差距而設計的:不再一次啟動幾十個小核心,而是將整個 forward pass 放進同一個常駐核心中。Cohere 的做法是,在每個 SM 上只常駐一個執行緒塊,由它從一個由 host 準備好的任務列表中讀取待執行任務;任務間的依賴關係不再用核心邊界表示,而是透過全域性記憶體中的計數器來顯式表達。於是排程的粒度從“一個運算元”縮小到“一個運算元的一個 tile”,同步的粒度也從整個 GPU 縮小到真正依賴的那些生產者。

速度提升主要來自四個方面。第一是啟動和同步開銷顯著減少;第二是減輕 wave quantization:當 GEMM 的 tile 數量不是 SM 數的整數倍時,傳統核心會產生閒置波次,而大核心可以把空閒 SM 立即用後續任務“回填”。North Mini Code 採用並行 transformer 層,注意力與 MoE 可平行計算,因此這種回填還能做得更激進。第三是消除虛假依賴:核心邊界意味著全網格同步,所有 SM 都要等最慢的那個;而細粒度計數屏障允許某一路 O-proj 在對應注意力輸出完成後立刻開始,不必等所有注意力組都完成。第四是對不依賴當前啟用的權重做預取,例如 router 和 QKV 的權重可以在上一層的 O-proj 尚未結束時就開始搬運,利用本來會閒置的頻寬。

效果方面,在 batch size 為 1 時,大核心達到 292 token/s,約為光速的 62%,是 vLLM 的 1.58 倍;這一優勢在不同 batch size 以及最長 256K 上下文範圍內都成立,且沒有可測量的精度損失。更重要的是,Cohere 強調這是一個完整的服務系統:支援 continuous batching、paged attention、變長序列,並透過 OpenAI 相容介面提供 tool calling,因此可以直接把 OpenCode 指向這個端點來進行編碼。與此同時,實現本身並不複雜:整個核心就是單個 CUDA 檔案,既不需要編譯器,也沒有引入新的程式設計正規化或抽象,甚至附帶了把已有核心改造成大核心的“配方”。

在技術淵源上,Cohere 明確借鑑了 Hazy Research 的“Look Ma, No Bubbles!”開創性工作,沿用了一些關鍵技術:SM 上的任務直譯器模式、基於計數器的同步、以及任務邊界處的重疊。不同之處在於,Cohere 的實現更依賴 tensor core 指令,即使 batch size 為 1 也使用 wgmma;另外,他們沒有采用共享記憶體分頁來做權重預取,因為這會引入複雜簿記和開銷。作為替代,每個 opcode 有一套 warp-specialized 流水線,共享記憶體佈局在編譯期靜態確定;同時利用同型別相鄰 GEMM 任務之間的流水線階段重疊,以及在 GEMM 內部的權重預取策略,讓等待輸入的任務仍然在搬運權重位元組。這些做法讓大核心在一個可部署服務中真正落地。

展開要點與分析

文章情報

工程師進階

要點

  • 將整個解碼前向過程融合為單個常駐大核心,減少核心啟動與全網格同步開銷。
  • 批大小 1 時吞吐 292 tok/s,是 vLLM 的 1.58 倍;優勢在不同批大小和最長 256K 上下文下保持,且無明顯精度損失。
  • 支援連續批處理、分頁注意力、變長序列以及 OpenAI 相容介面和工具呼叫。
  • 實現為單一 CUDA 檔案,無需編譯器或新程式設計抽象,並提供移植現有核心的指南。

要點與分析由自動化流程生成,可能有誤,請結合原始來源核實。