多數 LLM serving stack 把每個 forward pass 拆成一連串 kernel：先 launch QKV，等；再 launch attention，等；接著 launch MoE，再等。單看每個 kernel 都沒問題，真正的成本在於中間的等待。decode 階段 batch size 小的時候，GPU 有很大比例的時間在等，而不是在算。
Cohere 在 2026 年 9 月 8 日發表了針對 North Mini Code 的 serving engine，核心是一個 decode megakernel：BF16、單張 H100，端到端比 vLLM 快 1.25 到 1.41 倍。對產品開發者來說，重點不是「又一個加速技巧」，而是它把「kernel 邊界」這個隱形成本攤開來，並示範了怎麼用一個 CUDA 檔做到真實 serving 需要的功能。
為什麼 decode 慢：不是算力不夠，是頻寬沒用滿
Autoregressive decoding 在低 batch size 時是 memory-bound 的。North Mini Code 是 30B 模型，每個 token 只有 3.3B 參數活躍，BF16 下每個 decode step 要從 HBM 搬 6.6 GB 權重，加上 8K context 約 0.5 GB 的 KV cache。H100 的 HBM 頻寬是 3.35 TB/s，理論速度上限（Speed-of-Light）約 470 tok/s。但 vLLM 實際只跑到 185 tok/s，等於只用了 39% 的頻寬。
Cohere 的 megakernel 在 batch size 1 達到 292 tok/s，也就是 62% 的 SoL，比 vLLM 快 1.58 倍。這個差距在不同 batch size 和最高 256K context 下都維持，而且沒有可測量的精度損失。
Megakernel 的運作方式：一個 persistent kernel 取代一百個小 kernel
GPU 上有 100 到 150 個 SM（streaming multiprocessor），每個 SM 跑同一個 kernel 程式，但處理不同資料。Megakernel 的做法是：每個 SM 只 launch 一個 threadblock，整個 decode step 都保持 resident。工作不是由 driver 分派，而是每個 block 從 global memory 讀一份 task list——每個 task 代表一個 tile 的某個 operation。
資料相依性不再靠 kernel 邊界表達，而是用 global memory 裡的 counter：task 完成時 increment 一個 counter，需要輸入時 spin 等待對應的 counter。這樣一來，排程單位從「整個 operation」縮小到「一個 tile」，同步單位從「整顆 GPU」縮小到「特定 producer」。
三個實際的加速來源
Cohere 列出三個比「減少 launch overhead」更重要的因素，依影響排序：
1. 消除 wave quantization
假設一個 kernel 有 200 個 tile，GPU 有 132 個 SM。第一波 132 個 tile 平行跑，剩下 68 個 tile 跑第二波，此時有 64 個 SM 閒置。kernel 越小，這種「進位損失」越嚴重。GEMM tile 形狀受限於矩陣維度和 kernel 設計，tile 總數很少剛好是 SM 數的倍數。Megakernel 沒有邊界要對齊：哪個 SM 有空，ready 的 tile 就開始跑。
North Mini Code 特別受惠，因為它用 parallel transformer layer：attention 和 MoE feed-forward 從同一個 normalized input 計算，最後才用 fused residual add + RMSNorm 合併。這表示 attention 和 MoE 不需要等對方的輸出，megakernel 可以更積極地把 ready 的工作「回填」到閒置的 SM。
2. 移除 false dependency
SM 不會同時完成工作，即使 workload 相同。kernel 邊界等於全 GPU barrier，最慢的 SM 決定所有人的進度。例如 attention 拆成 4 個 KV group，某個 group 先算完，那個 SM 還是得等其他三個，即使它下一步需要的資料已經在 memory 裡。Megakernel 用 fine-grained barrier 移除這種假相依：某個 KV group 的 O-proj 可以在該 group 的 attention output 落地後立刻開始；MoE 的 down projection 也可以在對應 expert 的 up projection 完成後就啟動，不用等所有 expert。
3. 權重 prefetch
權重是 immutable 的，不依賴這一步的 activation。所以 task 可以在 activation dependency 滿足之前，就開始把權重 tile 從 HBM 搬進 shared memory。Cohere 在 router 和 QKV projection 上用得最積極：在前一層 O-proj 的尾聲就開始 prefetch 權重，搶在 RMSNorm 執行之前，把原本閒置的頻寬用掉。
從 Hazy Research 的基礎出發，但做了務實的取捨
Cohere 的設計大量借鑑 Hazy Research 的「Look Ma, No Bubbles!」那篇開創性文章。那篇把 Llama-3.2-1B 的 forward pass 融進單一 kernel，在 batch size 1 達到 H100 頻寬的 78%，而 vLLM 和 SGLang 只有大約一半。Cohere 沿用三個關鍵想法：GPU 上的 task interpreter pattern、counter-based synchronization、跨 task boundary 的 overlap。
但 Cohere 做了幾個不同的決定。他們的 GEMM 大量使用 tensor core instruction（wgmma），即使在 batch size 1 也比 CUDA core 快，而且減少 register pressure。更重要的是，他們放棄了 shared memory paging 來做 weight prefetch——原本以為可以讓 memory load 在 previous task 釋放 buffer 前就開始，但實際 bookkeeping 太複雜、容易出 bug，overhead 反而超過好處。
取而代之的是兩個更便宜的重疊來源：同類型連續 GEMM task 之間，讓 pipeline 的 stage phase 跨 tile 延續，而不是每個 tile 邊界都 drain 再 refill；以及在 GEMM pipeline 內部，讓 producer warp 在等待 cross-SM activation 之前就先發出 weight tile load。
對產品開發者的意義
這篇的價值不只是「Cohere 跑得比 vLLM 快」。它示範了一種不同的思考方式：decode 階段的瓶頸不是 flops，而是 memory bandwidth 和 latency。如果你在評估 serving 方案，與其只看 benchmark 數字，不如問：這個系統在 batch size 1 到 8 的區間，實際用到了多少 HBM 頻寬？kernel 邊界造成的等待佔了多少比例？
Cohere 也強調 megakernel 沒有想像中難寫。他們的實作是單一 CUDA 檔，沒有 compiler、沒有新的程式設計典範，就是把既有的 tiled GEMM 和 paged attention 重新組織成單一 calling convention。這對想要自己動手優化 decode 的團隊來說，是一個務實的參考點。
如果你正在處理類似的效能問題，這篇的完整程式碼在 GitHub 上，可以直接看他們怎麼處理 task scheduler、barrier 和 weight prefetch。另外，如果你對「用更少的資源跑更大的模型」這個方向有興趣，之前寫過的細模型的大優勢：企業如何用小語言模型省錢又高效也討論了類似的取捨邏輯。
Megakernel 不是萬靈丹。它最適合 decode 階段、低 batch size、memory-bound 的場景。如果你的 workload 是 prefill-heavy 或 batch size 很大，傳統的 kernel-per-operation 可能還是比較合適。但當你的服務瓶頸在 decode latency，而且 batch size 多半很小時，把 kernel 邊界拆掉是一個值得認真評估的方向。
參考來源
本文由 AI 協助自上述來源整理，經人工審核後發布。