AMD MI300X GPU向けカスタムカーネル開発によるHugging Face Llama 3.1 405Bモデルの推論パフォーマンス最適化
Creating custom kernels for the AMD MI300
Hugging FaceはAMDと協力し、AMD MI300X GPU上でLlama 3.1 405Bモデルの推論パフォーマンスを向上させるためのオープンソース最適化カーネルを開発した。特に、FP8形式での推論に焦点を当て、Fused residual connection, RMS norm and FP8 conversion kernel、Fused SwiGLU activation and FP8 conversion kernel、Skinny GEMM kernelの3つのカスタムカーネルを組み合わせることで、VLLMフレームワーク上での大幅な高速化を実現した。MI300Xのアーキテクチャ(スレッド、ワープ、スレッドブロック、コンピュートユニット、XCD、キャッシュ階層)を解説し、GPUの並列処理能力とメモリ階層の特性を踏まえたカーネル最適化手法(メモリロードの効率化、キャッシュ利用、共有メモリの活用)について詳述している。特にRMS normカーネルでは、PyTorchと比較して10倍の高速化を達成し、さらに最適化を進めることでVLLMの実装をも上回る性能を示した。SwiGLUカーネルでは、計算効率の向上に焦点を当てている。これらのカーネルは、hf-rocm-kernelsリポジトリで公開されており、開発者はこれを基盤として独自のカーネル開発も可能である。
- Hugging Face が AMD と協力し、MI300X GPU 向けに Llama 3.1 405B モデルの推論を最適化する 3 つのカスタムカーネルを開発し、オープンソースとして公開した
AMDのMI300X GPU向けにHugging Faceが独自のカスタムカーネルを開発し、Llama 3.1 405Bという大規模モデルの推論を大幅に高速化した点は、NVIDIA一強のGPU市場における代替選択肢の実用性を示す重要な進展。特にFP8フォーマットでの最適化や、PyTorch比10倍の性能向上は、AI推論ワークロードにおけるAMD GPUの競争力を具体的に証明しており、半導体業界の競争構造変化やクラウドプロバイダーのマルチベンダー戦略に影響を与えうる。また、オープンソースで公開されているため、コミュニティベースでのさらなる最適化が期待され、AMDエコシステムの拡大につながる可能性がある。