ChordがKimi K2.x向けINT4 MoEカーネルを公開

概要

MoEモデルの推論性能は、行列演算の速度だけでは決まりません。各エキスパートに何個のトークンが送られるか、ルーティング後にどれだけパディングが発生するか、そしてGPU上で十分なwarpを常駐させられるかによって、最適なカーネルのスケジュールは変わります。Novita AIが公開したChordは、こうした差を考慮し、Kimi K2.xのINT4サービング形状を対象にしたCUDAカーネルです。

ChordはW4A16方式を採用します。BF16のアクティベーション、INT4の重み、group-32のスケール係数を組み合わせる構成です。現在の実装には二つの独立したカーネル系統があります。indexedパスは公開版Hummingをベースにしており、vLLMのsorted_ids、expert_ids、num_tokens_paddedなどのルーティング情報を利用します。もう一つのgrouped_contiguous（prefill）とgrouped_masked（decode）は、DeepGEMMから着想を得たSM90向けの別実装で、異なる重みパッキング形式を使います。

主な測定結果

• 公開Hummingの対応パスとの比較では、H200 EP8 prefillで約1.11倍から1.20倍、H200 TP8の単一インスタンスサービングで約1.17倍から1.33倍でした。
• H200 EP8 decodeでは約1.16倍から1.24倍で、downステージでは1.31倍に達しました。
• B300 EP8 decodeでは1.81倍から2.15倍が報告されています。ただし、公開HummingにはSM100/SM103向けの調整表がないため、B300の比較対象は未調整のデフォルト設定です。
• groupedパスでは、H200 EP8 prefillが約1.00倍から1.31倍、decodeが約1.16倍から1.35倍でした。EP16とEP32の測定表も公開されています。

Chordの重要な考え方は、すべてのリクエストに同じMoEカーネルを使わないことです。prefillとdecodeでは、エキスパートごとのルーティングトークン数が大きく異なる可能性があります。そこでindexedパスは、その値をblock-Mやoccupancyの選択に反映します。H200のprefillでは、一定以上のトークン数に対してエキスパートのパディング行数に合わせてタイルを調整します。また、一部の中規模タイルではレジスタ使用量を制限し、常駐warpを増やして、データ収集や逆量子化の待ち時間を隠します。重みのロード、INT4の逆量子化、WGMMA演算を重ねるパイプラインも採用されています。

導入時の注意点

対応するvLLMのリビジョンでは、ChordをインストールしてHummingバックエンドを明示的に選択できます。indexedパスは既存の統合を利用できますが、groupedパスとvLLM Hummingバックエンドの統合はまだ開発中です。Kimi K2.xで使われるuint4、group-32、BF16スケール、および圧縮INT4重み形式に対応し、未対応の量子化方式はロード時に失敗する設計です。

なお、公開値はカーネルレベルのベンチマークであり、サービス全体が同じ割合で高速化することを保証するものではありません。実際の効果は、エキスパート並列度、バッチ構成、ルーティング分布、GPU世代、重み形式、vLLMのバージョンに左右されます。Chordが示すのは、MoE推論バックエンドを固定設定から、実際のルーティング形状に応じて調整する方向です。

出典：vLLM Blog