naklecha/simple-llm

解決する課題

SimpleLLM は、OpenAI の gpt-oss-20b および gpt-oss-120b モデルを単一の NVIDIA H100 GPU で実行するために特別に設計された、軽量で高性能な推論エンジンです。特定のワークロードにおいて vLLM よりも高速なデコード速度を提供しつつ、元のモデルの精度を維持することを目指しています。

仕組み

このエンジンは約 1,351 行の Python で実装され、GPU カーネルに Triton を利用しています。公式の MXFP4 チェックポイントを再量子化なしで直接読み込みます。主な性能最適化は以下の通りです：

• カスタムカーネル： sinks、スライディングウィンドウ、GQA、KV append を処理する単一のデコードアテンションカーネルと、融合された RoPE カーネル。
• メモリ管理： フルアテンション層にはページングされた KV キャッシュ、スライディングウィンドウ層にはリングバッファを使用し、コンテキスト処理を最適化。
• 実行効率： アクティブスロットには CUDA グラフを使用し、貪欲な argmax/次のステップの入力は CUDA グラフ内で直接計算。
• MoE 最適化： テンソルコア上での MXFP4 混合エキスパート（MoE）実装と、エキスパートごとにタイルテーブルを書き込む最適化されたルーティング。

対象読者

単一の H100 で OpenAI の gpt-oss モデルを実行し、特に小〜中規模のバッチサイズや長いコンテキスト長で最大のデコードスループットを求める開発者や研究者。

ハイライト

• 性能： バッチ 1 および 64 で、vLLM 0.31.0 よりも 11〜15% 高速にデコード。
• 効率： 非常にコンパクトなコードベース（1,400 行未満の Python）。
• 精度： OpenAI が公開した GPQA-Diamond および AIME 2025 のスコアと一致。
• ハードウェア最適化： MXFP4 ウェイトを使用し、NVIDIA H100（80 GB）に特化して調整。