HelionがvLLMに統合：PyTorchネイティブDSLがCUDAカーネルを代替し、Qwen3モデルのFP8推論を高速化

PyTorchチームはRed Hatとの共同ブログポストで、Qwen3モデルのFP8推論に向けた9つのHelionカーネルのvLLMへの統合を説明した。著者はRed HatのSean ChenとMeta PlatformsのYanan Cao（PyTorchチーム）だ。

Helionとは何か、なぜ開発者は注目すべきか？

Helionは「タイル付きPyTorch」とも表現できるドメイン固有言語（DSL、domain-specific language）であり、直接的なCUDAコードを書かずにタイルプログラミング（tile-programming）の概念でPyTorchの構文を拡張して高性能なGPUカーネルを書くことができる。目標はハードウェアの可搬性だ。同じカーネルがNVIDIA H100、B200、そして将来のアーキテクチャで手動移植なしに動作するべきだ。

HelionはAOT（事前コンパイル）オートチューニングの高度なインフラを提供する——各ターゲットGPUアーキテクチャと入力次元に対して最適なタイルサイズとメモリレイアウトを自動的に探索する。カーネル開発は速い。著者たちは9つのカーネルの大部分が1日の作業時間内に実装・検証されたと述べている。

vLLMに統合されたカーネル

9つのカーネルはFP8推論のフォワードパス（前向き計算）のほぼ全体をカバーする。

1. dynamic_per_token_scaled_fp8_quant
2. rms_norm_dynamic_per_token_quant
3. silu_and_mul_dynamic_per_token_quant — 新しいフュージョンカーネル
4. fused_qk_norm_rope
5. per_token_group_fp8_quant
6. rms_norm_per_block_quant
7. silu_and_mul_per_block_quant
8. scaled_mm
9. scaled_mm_blockwise

この統合によりper-token量子化のフォワードパスでのカーネル呼び出し回数が11回以上から10回のフュージョン演算に削減される。

個別カーネルレベルのパフォーマンス

非GEMMカーネルは従来の実装（torch.ops._C）と比較してH100で一貫した大幅な高速化を示す。

• rms_norm_per_block_quant: +67%（torch.ops._Cに対して2.055×）
• silu_and_mul_per_block_quant: +73%（torch.ops._Cに対して2.269×）
• rms_norm_dynamic_per_token_quant: torch.ops._Cに対して1.80×
• fused_qk_norm_rope: torch.compileに対して1.38×

GEMMカーネル（scaled_mm）は混合した結果を示す。H100で1.08×だがB200で0.739×。著者たちはBlackwellアーキテクチャに対するTritoncodegenの制限が原因であることをオープンに述べており、その問題を解決するはずのCuteDSLバックエンドが開発中だ。

モデル全体でのスループット向上

テストはH100とB200でQwen3モデル（1.7B、8B、32B）に対して実施された。

• Per-token量子化（H100）：Qwen3-1.7BでH100上の全体スループット約1.05×
• 投機的デコーディング、Qwen3-8B、H100：全体スループット最大1.09×（バッチサイズ16でのTTFT高速化1.15×）
• Per-group量子化（H100とB200）：すべてのワークロードで約1.05×

Qwen3-8Bで3つの投機的トークンを使った投機的デコーディングではTTFTの1.33×高速化（time to first token）が計測された。

実際の制限事項

著者たちはトレードオフについて透明だ。scaled_mmのような大規模カーネルのオートチューニングには、3つのQwen3バリアントの全入力形状（168の異なる次元）をカバーするのに約1日かかる。Helionランタイムディスパッチはカーネル呼び出しごとに数十マイクロ秒のレイテンシーオーバーヘッドをもたらすため、最適なパフォーマンスにはCUDAグラフキャプチャが必須だ。

HelionのPer-group量子化カーネルはB200でDeepGEMMが必要とするUE8M0フォーマットとまだ互換性がなく、もう一つの未解決事項となっている。

今後の展望

プロジェクトはGitHub issue vllm-project/vllm#32962として公開されており、9つすべてのカーネルの実装が含まれている。目標はHelionサポートを追加のアーキテクチャに拡張し、CuteDSLバックエンドでBlackwellのGEMMパフォーマンスを解決することだ。この問題でのRed HatとMeta PyTorchチームの協力は、PyTorchエコシステムが高性能な推論カーネルのハードウェア可搬性という問題に真剣に取り組んでいるシグナルだ。