Ollamaのプレリリース版、MLXカスタムカーネルをCUDAに移植　Nemotronの限定テストでプロンプト処理スループットが倍増

Ollamaは10月14日、v0.40.5-rc0のプレリリース版を公開し、前日にマージされたMLXのCUDAカーネル移植を取り込んだ。この更新は、NVIDIA GPUでの実際の実行経路を対象としている。従来、Metal実装しかなかった一部のカスタム演算はCUDA上で計算グラフ演算にフォールバックしていたが、新版では対応するカーネルが追加され、対象モデルが専用実装を利用できるようになった。リリースノート

移植対象にはmamba2_scan、depthwise_conv_silu、gated_deltaと、関連する状態およびバイアント変種が含まれる。作者はスレッド境界のチェックとスカラー読み取り方法も調整し、一部のキャッシュを各warpのレジスタに置くことで、計算グラフの正規化計算を繰り返す代わりに、共有メモリと同期バリアを不要にした。こうした詳細は、バックエンド間の移植にもGPUの実行モデルやメモリモデルに合わせた調整が必要であることを示している。技術提案

作者はLinuxとRTX 6000 Adaで、メインブランチと変更ブランチを比較した。NVFP4のnemotron3:33bでは、プロンプト処理速度が毎秒323.9 tokensから679.0に、生成速度が112.3から149.6に向上した。nemotron-3.5-lightning:30b-a3bでは、プロンプト処理が313.3から695.6に、生成が113.2から129.7に向上した。プロンプト処理の伸びはそれぞれ約110%と122%、生成速度の伸びは約33%と15%だった。これは作者による限定的なテストであり、すべてのモデルや本番サービスで同じ加速が得られるとは限らない。テストデータ

MLXの公式ドキュメントには、ソース文字列からCUDAカスタムカーネルを実行時にコンパイルするインターフェースがすでにあり、入力レイアウトや動的共有メモリを指定できる。今回の変更の技術的な意義は、この既存インターフェースをOllamaの具体的なモデル演算に適用し、バックエンド間の性能差を埋めた点にある。MLXドキュメント

デプロイ担当者は今後、正式版への収録状況を追い、自社のプロンプト長、量子化設定、GPUを使って、初回トークンのレイテンシと継続生成速度を再測定するとよい。数値の一貫性も検証が必要だ。提案によると、融合カーネルと計算グラフではリダクションの順序が異なり、bf16の結果に最下位ビット1つ分の差が生じる場合がある。また、レビュアーはCUDA実行をローカルで検証していないと明記している。現時点の証拠は特定の構成における性能改善を裏付けているが、デバイスをまたいだ性能を保証するには不十分だ。検証に関する説明