dropbox/gemlite

何を解決するか

GemLiteはLLM推論における低ビット行列乗算のパフォーマンスボトルネックに対処します。デフォルトのTorch AOカーネルと比較して、プリフェイルおよびデコードステージの実行を大幅に高速化する高度に最適化されたカーネル群を提供し、量子化モデルのレイテンシを低減し、スループットを向上させます。

仕組み

GemLiteはTritonを使用して、さまざまなワークロードに特化した複数の行列乗算カーネルを実装しています：

• GEMM: 一般的な行列乗算向けの標準的なTensor Coreベースのカーネル。
• GEMM Split-K: バッチサイズ2〜32のバッチ処理LLMデコードを最適化し、K次元を複数のジョブに分割して部分和を計算します。
• GEMV: 小規模バッチ（M=1）向けに設計され、活性化値を1次元のチャンクに分割します。
• GEMV RevSplit-K: 新しいアルゴリズムで、バッチサイズ1のデコードを最適化し、1つのプログラムあたりのワークロードを2倍にすることで、スケールおよびゼロのロードオーバーヘッドを削減します。

幅広い精度をサポートしており、FP16、BF16、FP8、INT8、および低ビット重み（8、4、2、1ビット）に加え、Blackwellなどの最新ハードウェア向けのMXFPおよびNVFP形式も対応しています。

対象ユーザー

GemLiteは、特にNVIDIA GPU上で高いパフォーマンスを維持しつつモデルサイズとメモリ帯域幅の要件を削減するために量子化を活用するLLM推論最適化に取り組む開発者および研究者を対象としています。

主な特徴

• 大幅な高速化: デフォルトのTorch AOカーネルと比較して、プリフェイルが最大7〜8倍、デコードが最大3〜6倍高速化されます。
• 広範な精度サポート: FP16、BF16、FP8、INT8、およびMXFP/NVFP4と互換性があります。
• 統合性: vLLM（HQQまたはTorchAO経由）、SGLangと連携可能で、torch.compileと互換性があります。
• オートチューニングキャッシュ: 最適なカーネル設定を保存・読み込み可能で、起動時の繰り返しオートチューニング時間を排除できます。
• 柔軟なパッキング: HQQスタイルの重みパッキングおよび重みと活性化のチャネルワイズスケーリングをサポートしています。