Cohere、デコード処理全体を単一のCUDA megakernelに統合、H100のエンドツーエンドスループットを25～41％向上

Cohereは、North Mini Code向けの推論サーバーを公開した。その中核は、個々のGEMMやattention kernelをさらに調整することではなく、単一の常駐型CUDA megakernelでデコードのforward pass全体を実行する点にある。従来のエンジンはRMSNorm、QKV、attention、MoEなどのkernelを順番に起動し、各処理の境界でGPU全体の同期完了を待つ必要がある。低バッチかつ自己回帰デコードのようにメモリ帯域幅とレイテンシに制約されるワークロードでは、起動間の空白時間や1 waveに満たない処理量によってSMがアイドル状態になる。

新エンジンは、H100の各SMに1つのthread blockを常駐させる。ホスト側であらかじめ細粒度のタスクテーブルを構築し、GEMM tileとattention splitを利用可能なSMに割り当てる。データ依存関係はグローバルメモリ上のカウンターで表現されるため、処理を受け取る側は自身の入力が準備でき次第、grid全体を待たずに直ちに実行できる。また、North Mini CodeのattentionとMoEは同じ正規化済み入力から並行して処理でき、空いているSMは別経路に残った終盤の処理を引き受けられる。システムはcontinuous batching、paged KV cache、prefix caching、preemption、tool calling、OpenAI-compatible APIも引き続きサポートする。

単一H100、BF16、batch size 1という条件で行われた公式テストでは、megakernelは毎秒292 tokenに達し、推定メモリ帯域幅上限の約62％を実現した。これはvLLM 0.24の1.58倍の速度に相当する。実際のprompt、prefill、dynamic batchingを加えた場合、5種類のワークロードにおける平均デコードスループットの向上幅は1.25～1.41倍に縮小した。この結果は、量子化に依存しなくても、サービス層のスケジューリングと同期構造を改善することで大幅な性能向上が得られる可能性を示している。

制約も明確だ。現時点でサポートされているのはH100 SM90a、CUDA 13、BF16、batch size 1～8のみで、スケジューリングもNorth Mini Codeに特化している。prefillには引き続きPyTorch kernelが使われ、その間はdecodeが一時停止するほか、samplingのオプションも完全にはそろっていない。今後は、ほかのMoEモデルにも汎用化できるか、prefillとdecodeの混在に対応できるか、さらに異なるリクエスト分布やより高い並列度でも優位性を維持できるかを見極める必要がある。既存データは主に公開元が提供したものであり、独立した再現検証はまだ不足している。