HelionをvLLMに統合、量子化GEMMを自動調整する推論バックエンド

概要

大規模言語モデルの推論では、量子化されたlinear layerの行列積がスループットを大きく左右します。vLLMはFP8、INT8、INT4、NVFP4などの形式に対して、CUTLASS、DeepGEMM、FlashInferといった最適化ライブラリを利用しています。しかし、最適なGEMM戦略は、モデル、バッチ形状、量子化方式、GPUによって変わります。汎用的なカーネルだけでは、特定のデプロイ環境で取りこぼしが生じる可能性があります。

PyTorchのブログでは、この問題への別のアプローチとして、HelionをvLLMのlinear backendに組み込んだ取り組みが紹介されています。HelionはPyTorchネイティブでハードウェアに依存しにくいカーネルDSLです。開発者が多くの専用カーネルを個別に管理する代わりに、アルゴリズムやタイル、メモリ配置などの選択肢をパラメーターとして公開し、AOT autotunerに探索させます。

主なポイント

• 複数のGEMM方式を一つの実装で扱う。 Standard GEMMに加え、K方向を分割して並列度を高めるSplit-K、小さなMでの利用効率改善を狙うSwap-ABを同じ実装に組み込めます。
• ディスパッチ規則を自動化する。 従来は入力形状を見てカーネルを切り替えるヒューリスティックを手作業で設計する必要がありました。Helionでは方式の選択自体をチューニング対象にでき、各形状に対して方式と詳細設定をまとめて探索します。
• Hopper向けの量子化GEMMに焦点を置く。 現段階ではHelionのTriton backendを使い、FP8 dynamic、W8A8 INT8、Block-FP8を対象としています。Blackwell向けのCuteDSL backendでも初期のGEMM結果は競争力を示していますが、拡張にはバックエンドの成熟が必要です。
• 混合ディスパッチで比較。 Hopperでは形状別チューニングとhybrid dispatchを組み合わせ、評価したモデルでvLLM標準のCUTLASSおよびDeepGEMMを上回りました。一部ワークロードではエンドツーエンドのスループットが10%を超えて改善しています。ただし、これは評価されたハードウェア、モデル、形状に基づく結果です。

導入時のトレードオフ

自動チューニングはコストをなくすのではなく、体系化します。細かなAOT探索にはなお数時間かかる場合があります。また、vLLMの起動時にCUDA Graph captureがHelionのJITコンパイルを誘発し、コールドスタートを遅くする可能性があります。コンパイル済み成果物をキャッシュすれば、ウォームスタートでは影響を抑えられます。CUDA Graphの範囲外では、CPU側のディスパッチとカーネル起動がGPU上の改善を相殺することもあります。

さらに、モデルごとの調整済み設定を上流で多数配布すると、設定ファイルの保守やCIでの検証が難しくなります。性能、使いやすさ、保守性の三者は簡単には同時に最大化できません。高い専用性能を求めるほど、調整時間または設定管理の負担が増えるためです。

意義

この取り組みは、推論基盤におけるカーネル開発の役割分担を変える可能性があります。開発者は一つの高水準実装で複数のアルゴリズムを表現し、チューナーがハードウェアとワークロードに応じた探索を担います。そのため、GPUカーネルの専門知識が限られる運用チームでも、自社モデルに合わせた最適化を試しやすくなります。一方、形状が頻繁に変わるサービスや起動時間を重視する用途では、従来の事前最適化バックエンドの方が扱いやすい場合があります。Helionの実用性は、GPU性能の向上がコンパイル、ディスパッチ、保守のコストを上回るかで決まります。

PyTorch Blog