Chord、Kimi K2.x向けにINT4 MoEカーネルを再設計、推論フェーズに応じてGPUスケジューリングを選択

Novita AIとvLLMチームは、Kimi K2.xシリーズ向けのW4A16 MoE CUDA演算子「Chord」をオープンソース化した。入力はBF16のまま維持し、重みをINT4に量子化したうえで、重み32個ごとに1グループとするBF16 scaleを用いて逆量子化する。汎用的なGEMMの代替ではなく、疎モデルにおけるgate、up、down projectionの実際のルーティング形状に合わせ、prefill、decode、テンソル並列、エキスパート並列それぞれのスケジューリングを個別に調整している。

Chordは2系統のカーネルを提供する。`indexed`パスはHummingのルーティングインターフェースを踏襲し、H200ではEP8 prefill、decode、TP8の混合ワークロードを、B200／B300ではEP8 decodeをサポートする。`grouped_contiguous`と`grouped_masked`はDeepGEMMのデータレイアウトを基にしており、それぞれprefillとdecodeを処理する。スケジューラは、バッチ全体のtoken数だけを見るのではなく、各expertが実際に受け取ったtoken数に応じて、blockサイズ、stream-K、CTA常駐数、パイプライン深度を選択する。また、カーネルはWGMMA、TMA、warp specialization、キャッシュ済みのcubin descriptorを使用し、小規模なdecodeのたびに設定を再探索する処理を回避している。

作者らは、同一のGPU、形状、ルーティング分布を用いて公開版Hummingと比較した。その結果、H200のEP8 prefillは約1.11～1.20倍、H200のdecodeは約1.16～1.24倍高速だった。B300のdecodeでは1.81～2.15倍と報告されている。ただし、B300のベースラインには、まだチューニングされていないHummingのデフォルト戦略が用いられている。また、測定にはルーティング、通信、activationが含まれていないため、vLLMサービス全体が同じ倍率で高速化されることを意味しない。パッケージはHumming互換のエントリーポイントを介して一部のvLLMバージョンに組み込めるものの、groupedパスの正式な統合は現在も進行中である。デプロイ前には、WNA16 group-scaleのサポート、量子化形式、自身のtoken-per-expert分布も検証する必要がある。