CUDA一強に風穴を開けるTritonの正体

TECH

深層学習の世界で「速いコードを書く」といえば、長らくCUDAが絶対的な存在だった。NVIDIAのGPUを直接叩ける唯一無二の手段として、研究者もエンジニアも避けて通れない道になっていた。だがその反面、CUDAを使いこなすにはかなりの低レベル知識が必要で、演算ライブラリの範囲を超えた新しいアイデアを実装するには、大きな労力がかかっていた。

そこで注目を集めているのが「Triton」だ。これはGPUカーネルをもっと簡単に、しかも高速に書けるように設計されたオープンソースの言語兼コンパイラである。見た目はCやPythonに近いが、その裏側ではLLVMを利用してGPU向けのコードを最適化し、cuBLASやcuDNNといったベンダー製ライブラリと肩を並べる性能を叩き出すことが報告されている。

Tritonの面白さは「タイル」と呼ばれる考え方にある。巨大な行列やテンソルを細かいブロックに分け、その単位でデータ移動や計算を制御する。これにより、研究者が思いついた新しい演算も、既存のライブラリに頼らず効率的に記述できるようになる。たとえば特殊な畳み込みやスパースな演算など、既存ライブラリではカバーできないニッチな処理を自分で最適化できるわけだ。

さらに動きは広がりつつある。PyTorchに統合されたことで、研究者はお馴染みの環境から自然にTritonを呼び出せるようになった。AMDのGPUでも動作させる取り組みが進み、NVIDIA一強の構図に変化の兆しが見えてきた。加えて、最近ではAIを使ってTritonカーネルを自動生成・最適化する研究まで登場しており、手作業でチューニングする時代を終わらせる可能性すらある。

もっとも、TritonがすぐにCUDAを駆逐するわけではない。依然としてNVIDIA GPUが主戦場であり、ハードウェアレベルの最適化では手書きCUDAに軍配が上がる場面もあるだろう。それでも、「CUDA以外にも道がある」と示したこと自体が大きい。これまで閉じられていた門に、新しい入口が開いたのだ。

CUDAの牙城に挑むのか、それともむしろNVIDIA自身が新しい武器として取り込むのか。いずれにせよ、TritonがGPU計算の世界に風穴を開けつつあるのは間違いない。