Skip to content
Published on

Triton Gluon — コンパイラが隠していたレイアウトを手で書く言語

シェア
Authors

はじめに — Tritonが行き詰まるのは「遅いから」ではない

Tritonの魅力は明らかです。Pythonでタイル単位のカーネルを書けば、コンパイラがレイアウト、メモリ割り当て、データ移動、同期を自動で処理してくれます。多くの場合これは良い取引です — CUDA C++まで降りなくても、使える性能が出るからです。

問題は、その取引が崩れる瞬間です。Gluonチュートリアルの導入部がこの状況を正確に述べています。Tritonコンパイラは広い範囲のカーネルに対して効率的なコードをうまく生成しますが、手でチューニングした低レベルコードに負けることがあります。そしてそれが起きたとき — チュートリアルの表現では「すべての詳細が隠されているため、ユーザーが性能を有意に改善するためにできることはほとんどありません」。

ここが核心です。Tritonの限界は「遅い」ことではなく、行き詰まったときに抜け出す扉がないことです。抽象化はうまく働いているときは恵みで、ずれた瞬間には壁になります。これまで、その壁の前での選択肢は一つしかありませんでした — CUDA C++でゼロから書き直すことです。

Gluonは第三の選択肢を作ろうとする試みです。

Gluonとは何か — 同じコンパイラスタック、逆転した契約

公式ドキュメントの定義は短いものです。GluonはTritonの下位レベルGPUプログラミングモデルであり、「レイアウト、共有メモリ、ワープ特殊化、ターゲット固有の機能を直接露出させ、高度なカーネルが利便性を制御権と引き換えられるようにする」ものです。

重要なのは、Gluonが別個のプロジェクトではないという点です。チュートリアルが明かすとおり、GluonとTritonは同じコンパイラスタックを使います。どちらもタイルベースのSPMDプログラミングモデルを実装し、どちらも同じフロントエンドとJITインフラを共有するPython DSLです。ホスト側のコードもほぼ同じです — カーネルを呼ぶ方法、グリッドを指定する方法、triton.autotuneさえそのまま使えます。

変わるのはデバイスコードの書き方です。チュートリアルのまとめをそのまま移せば、Gluonはデバイスコードの書き方を変え、ホストコードはカーネルのハイパーパラメータが増える程度にしか変わりません。

最も単純なGluonカーネルはTritonと区別がつきません。

import torch
import triton
from triton.experimental import gluon
from triton.experimental.gluon import language as gl


@gluon.jit
def copy_scalar_kernel(in_ptr, out_ptr):
    value = gl.load(in_ptr)
    gl.store(out_ptr, value)

@triton.jitの代わりに@gluon.jittlの代わりにgl。ここまでは名前が変わっただけのように見えます。

そして、そのimportパスに注目してください — triton.experimentalです。本稿の最後で改めて触れます。

最初のカーネルはがっかりする — 666 GB/s

チュートリアルは正直にも失敗から始まります。スカラーを一つずつコピーするmemcpyカーネルを書き、XBLOCKをオートチューンし、GB200上で8 GBをコピーします。チュートリアルのソースにコメントとして残されている結果はこうです。

Time:        24.00 ms
Throughput: 666.24 GB/s

オートチューナーが選んだ最適なXBLOCKは2048で、結果は約666 GB/sでした。チュートリアルはこの数字について「GPUのピーク帯域幅8 TB/sにはるかに届かない」と記しています。おおよそピークの8%程度です。

理由はカーネルを見ればすぐにわかります。

@gluon.jit
def memcpy_kernel(in_ptr, out_ptr, xnumel, XBLOCK: gl.constexpr):
    pid = gl.program_id(0)
    start = pid * XBLOCK
    end = min(start + XBLOCK, xnumel)
    for i in range(start, end):
        value = gl.load(in_ptr + i)
        gl.store(out_ptr + i, value)

各プログラム(GPUのCTA、つまりスレッドブロック)が一度に要素を一つずつコピーします。GPUの並列性をまったく使えていないコードです。複数要素を一度に動かすには、スカラーではなくタイルをロードする必要があり、タイルをロードするには — レイアウトを選ばなければなりません

ここがGluonとTritonの分かれ道です。Tritonならコンパイラが自動で選んでいたものを、Gluonではあなたが書きます。

レイアウト — Gluonがあなたに渡すもの

レイアウトは、テンソルの要素がスレッドブロック内のスレッドにどう分配されるかを規定します。分配はGPUの階層構造に従います — スレッドブロック、ワープ、レーン、そして各レーンの個別レジスタです。

最も一般的な種類がBlockedLayoutです。

gl.BlockedLayout(
    size_per_thread=[2, 4],
    threads_per_warp=[16, 2],
    warps_per_cta=[2, 2],
    order=[1, 0],
)

読み方はこうです。

  • size_per_thread=[2, 4] — 各スレッドが連続する2x4のサブタイルを所有し、それはそのスレッドのレジスタに収まります。
  • threads_per_warp=[16, 2] — スレッドタイルをこう積み上げてワープタイルを作ります。積はワープあたりのスレッド数と一致しなければなりません(NVIDIAハードウェアでは32)。
  • warps_per_cta=[2, 2] — ワープタイルをこう積み上げてブロック全体を作ります。
  • order=[1, 0] — どちらの次元からタイル化するか。行優先(row-major)です。

3つを要素ごとに掛け合わせるとブロック形状が出ます — ここでは[64, 16]です。

orderがなぜ重要かは、一つのスレッドのタイルを描いてみるとすぐにわかります。order=[1, 0]なら、レジスタは内側の次元に沿って増えます。

[[T:0, T:1, T:2, T:3],
 [T:4, T:5, T:6, T:7]]

order=[0, 1]なら、同じ2x4タイルはこう変わります。

[[T:0, T:2, T:4, T:6],
 [T:1, T:3, T:5, T:7]]

同じ要素、同じスレッド、違う配置。そして後で見るように、この違い一つが帯域幅を何倍にも分けます。

レイアウトを明示することには副次的な利点もあります — レジスタ予算が計算可能になるのです。チュートリアルの例に沿うと、上のレイアウト(ブロック[64, 16])で128x128のf32テンソルを扱うと、ブロックは[2, 8]にタイル化され、スレッドあたりブロックごとに8個、テンソル全体では8 * 16 = 128個のレジスタを使います。Triton上ではコンパイラの裏に隠れていた数字です。

注意すべき落とし穴も一緒に出てきます。テンソルがブロックより小さいとブロードキャストされます。32x8のf32テンソルは要素が256個しかありませんが、ブロックのタイル化自体は変わらないため、各プログラムでは64 * 16 = 1024個の物理レジスタに格納されます。小さいテンソルがレジスタを4倍食う計算です — 明示的なレイアウトがなければ気づきにくい種類のコストです。

数字で見るレイアウト

ここから先は、Tritonリポジトリのチュートリアルソース(python/tutorials/gluon/02-layouts.py)にコメントとして記録されている測定値です。上流の開発者がGB200上で実際に測ってソースに残した値であり、私が再現したものではありません。

1D memcpyをきちんとしたタイルロードに書き直し、XBLOCK=2048num_warps=4を固定した上で、size_per_thread(以下のR)だけを変えて計測しています。

gl.BlockedLayout(
    size_per_thread=[R],
    threads_per_warp=[32],
    warps_per_cta=[4],
    order=[0],
)

記録された結果です。

R=1   6.574 TB/s
R=2   6.476 TB/s
R=4   6.474 TB/s
R=8   6.502 TB/s
R=16  6.214 TB/s

読み取れることが2つあります。まず、スカラーループの0.666 TB/sから6.574 TB/sへと変わりました — 記録された2つの値を単純に割ると約10倍です(これは私の計算で、上流がそう表現しているわけではありません)。タイルロードとレイアウトがもたらすもののスケールはこの程度です。

次に、Rを大きくしても良くなりません。R=1が最速で、R=16が最も遅いのです。チュートリアルはここでSASSを開き、LDG.E/STG.E命令を直接数える方向に進みます。これはGluonの性格をよく示しています — レイアウトを露出するということは「より良いデフォルト」を与えることではなく、選択肢と、その選択を検証する責任を渡すことを意味します。幅を広げれば速くなるという直感は、ここでは外れます。

レイアウトを間違えると性能が崩れる

2Dに進むと話がさらに鋭くなります。

8 GBの32K x 64KテンソルをXBLOCK=1YBLOCK=2048で、1Dで勝ったR=1と同じように振る舞うよう設計したレイアウトでコピーします。

layout = gl.BlockedLayout([1, 1], [1, 32], [1, 4], [1, 0])

結果は6.260 TB/s — 1Dより5%遅いです。チュートリアルは2Dの算術がより複雑であることなどを理由に挙げ、この程度は納得できるとしています。

では、入力テンソルを転置するとどうなるでしょうか。同じカーネル、同じレイアウトです。

0.774 TB/s

チュートリアルの表現では、性能が「暴落(craters)」します。理由は単純です — 内側の次元がもう連続していないため、結合(coalescing)が失われます。32個のレーンが隣接するアドレスを読んでこそ一つのトランザクションにまとまるのですが、転置されたテンソルではレーンごとに遠く離れたアドレスを指してしまいます。

直し方も単純です。ブロックサイズを変え、レイアウトを転置します。

layout = gl.BlockedLayout([1, 1], [32, 1], [4, 1], [0, 1])
# XBLOCK=2048, YBLOCK=1
6.590 TB/s

0.774から6.590へ。カーネルのアルゴリズムは一文字も変わっておらず、変わったのはレイアウト記述子だけです。

ここでチュートリアルが付け加える観察が面白いところです。直したバージョン(6.590 TB/s)は1D memcpy(6.574 TB/s)よりもわずかに速いのです。転置あり・なし両方で各プログラムがメモリにアクセスするパターンは同一なのに差が出ており、チュートリアルはこれをプログラムがGPU上のどこにスケジュールされるか、つまりデータの局所性のせいだと説明します — TLBの仮想アドレス変換キャッシング、H100でL2キャッシュが互いに通信するパーティションに分かれていること、といった要因です。この層では「同じアクセスパターン」ですら同じ性能を意味しません。

ここから得られる実務上の教訓 — 最適なレイアウトはカーネルではなく、グローバルメモリ上に置かれたテンソルのストライドに依存します。だからチュートリアルは最終的にこんなヘルパーに収束します。

def get_layout_for_gmem_access(tensor, num_warps):
    if len(tensor.shape) == 1:
        return gl.BlockedLayout([1], [32], [num_warps], [0])

    assert len(tensor.shape) == 2, "only 1D and 2D tensors are supported"
    assert 1 in tensor.stride(), "expected at least 1 contiguous dimension"
    if tensor.stride(1) == 1:
        return gl.BlockedLayout([1, 1], [1, 32], [1, num_warps], [1, 0])
    else:
        return gl.BlockedLayout([1, 1], [32, 1], [num_warps, 1], [0, 1])

レイアウトを定数として埋め込むのではなく、実行時のストライドから導いています。Gluonカーネルが柔軟でありながら速くあろうとすると、こういうコードが必要になります。

PyTorchに勝つ場面、そしてその代償

Gluonの2D memcpyが価値を発揮する場面もチュートリアルは取り上げています。8 GBのテンソルで、一行おきに飛ばしたビュー(非連続テンソル)を連続テンソルへコピーする作業です。PyTorchのx.contiguous()と同じ仕事です。

2D memcpy:                6.258 TB/s
torch.Tensor.contiguous:  2.946 TB/s
2D memcpy (transposed):   6.398 TB/s

チュートリアルの表現は「すでにPyTorch実装より2倍以上速い」です。これはGluonが実際に価値を発揮する種類の事例です — 入力や出力が変則的なレイアウトを持つ場合です。ただし、1D memcpyは入力と出力がどちらも連続メモリブロックのビューであるときにしか使えない、という条件がつきます。

では、入力と出力のレイアウトが互いに逆だとどうなるでしょうか。チュートリアルはこの場合も測っています。入力はdim 1に沿って、出力はdim 0に沿って連続している32K x 32Kテンソルです。

2D memcpy (order=[1, 0]): 0.978 TB/s
2D memcpy (order=[0, 1]): 1.674 TB/s

どちらのレイアウトを選んでもひどい結果です。片方を満たせばもう片方が崩れるからです。解決策はロードとストアに異なるレイアウトを使い、その間で変換することです。

@gluon.jit
def memcpy_2d_inout_kernel(in_ptr, out_ptr, ..., layout_in: gl.constexpr, layout_out: gl.constexpr, ...):
    ...
    value = gl.load(in_ptr + in_offsets, mask=mask_in)

    # Use `gl.convert_layout` to perform layout conversions.
    value = gl.convert_layout(value, layout_out)

    gl.store(out_ptr + out_offsets, value, mask=mask_out)
2D memcpy (in/out layouts): 4.814 TB/s

0.978から4.814へ上がりましたが — 6.5 TB/s台には届きません。そしてチュートリアルはその理由を隠しません。convert_layoutは共有メモリを使う必要があり、共有メモリはGPU上で貴重な資源です。レイアウト変換に共有メモリを使うと占有率(occupancy)と最大パイプライン深度が下がり、性能に悪影響を与えます。チュートリアルの結論は「この場合は変換コストが避けられず、非効率なグローバルメモリアクセスのコストよりはるかに安く済む」というものです。

つまり4.814 TB/sは勝利ではなく、最も悪くない選択です。こうやってコストを目に見えるようにするのがGluonの性格です。

リリースタイムライン — どこまで来たか

Gluonがいつどこまで来たかは、リリースノートで追えます。以下の日付はGitHubのリリースAPIで確認した値です。

バージョン日付Gluon関連
v3.4.02025-07-30「Gluon Framework Comprehensive Enhancement」— static_assert、TensorDescriptorカーネル引数、非同期TMA、テンソルメモリ、同期バリア、split/join/reshape・リダクション
v3.5.02025-10-21Hopper WGMMA + async wait (#7300, #7313)、Gluonへのlibdevice露出 (#7890)、公開APIドキュメント文字列 (#7323)
v3.6.02026-01-212CTAモードの初期対応 (#8644, #8653)、num_ctasの実装 (#8602)、AMD gfx1250向けasync_copy (#8622)、実験的なTriton → Gluon変換器 (#8417)
v3.7.02026-05-072CTAのエンドツーエンド対応、AMDワープパイプライン (#8586, #8975, #8980)、gfx1250向け4・8ワープStream-K Gluonカーネル (#9370)、TDM L2プリフェッチ (#9086, #9148)、きめ細かなクラスタバリア (#9206)、sm103のtcgen05.ld.red (#9151)、local scatter/gather (#8480)
v3.7.12026-06-18パッチリリース — リリースノートによると「新機能やAPI変更なし」、リグレッション2件を修正

読み取れる流れがあります。2025年半ばに骨格が入り、2026年に入ってマルチCTAと最新ハードウェア機能(Blackwellのtcgen05、クラスタバリア)へと拡張されているところです。そしてv3.7.1が純粋なバグ修正だったという事実は — 最新の安定リリース時点でGluonの表面が1か月ほど静かだったということでもあります。

3.6.0のTriton → Gluon変換器は別途触れる価値があります。マイグレーションを楽にしてくれる道具のように聞こえますが、該当PR(#8417、2025-10-11マージ)の説明文は期待を正確に抑え込んでいます。「これはプロダクション用ではなく、Tritonカーネルを素朴な(naive)Gluon版に変換できるようにするだけだ。もちろん、そのGluon版はかなり遅くなるだろう。」

正直な文章です。そしてGluonの本質を要約しています — Gluonへ移すこと自体は何の性能も与えません。性能は移した後、あなたがレイアウトを正しく選んだときに出ます。

NVIDIA専用ではない — AMD gfx950とMI355

上の表で目を引くのは、AMDの項目が繰り返し出てくることです。gfx1250向けasync_copy、AMDワープパイプライン、gfx1250向けStream-Kカーネル、AMDクラスタバリア。Gluonは NVIDIA専用の脱出口として設計されたわけではありません。

AMDもこれを推し進めています。2026年5月22日、ROCmブログにLixun Zhang、Jason Furmanek、Peng Sun、Emad Barsoumが書いた「From Naive to Near-Peak: Building High-Performance GEMM Kernels with Gluon」が公開されました。gfx950(CDNA4)ベースのMI350/MI355上で、GluonによるGEMMカーネルをv0からv9まで段階的に最適化していくチュートリアルです。

以下の数字はすべてAMDのベンダー自己測定です。 ブログが明かした測定条件はMI355一枚、ROCm 7.0、gfx950-tutorial-v0.1タグでビルドしたTritonです。

カーネル形状結果
FP16 GEMM v0 (naive)4096x4096x8192520 TFLOPS、MFMA効率25%
FP16 GEMM v9 (最適化)4096x4096x81921489 TFLOPS、MFMA効率98.75%
BF8 GEMM4096x4096x163843257 TFLOPS、MFMA効率99.72%
MXFP4 GEMM4096x4096x327685255 TFLOPS、MFMA効率92.41%

ここで誤解しやすい点をはっきりさせておきます。520 → 1489 TFLOPSの約3倍は、GluonがTritonに3倍勝つという意味ではありません。 v0もv9も全部Gluonカーネルです — 同じリポジトリ(ROCm/gfx950-gluon-tutorials)のv0_naiveからv9_beyond_hotloopまで、教育目的で設計された最適化の旅の両端にすぎません。素朴なGluonカーネルが遅いのは当然で(先の変換器PRが言っていたとおりです)、3倍は「レイアウトとスケジューリングをきちんとやればここまで来る」という話であって、言語間の比較ではありません。

そしてブログ自身が線を引いています — これはhipBLASLtのようなプロダクションライブラリの代替ではないと明記しています。私が確認した本文には、hipBLASLtとの直接比較の数値はありません。つまり「Gluonカーネルがベンダーライブラリに勝った」という主張は、この資料からは出てきません。

方法論の面でこのブログが正直な部分もあります。主な指標としてTFLOPSではなくMFMA効率を使っており、その理由としてMFMA効率がクロックと無関係で、実行間の再現性がraw TFLOPSより優れている点を挙げています。ブーストクロックによって揺れ動くTFLOPSだけを掲げるベンダー資料よりも一枚上手な態度です。(もちろん、MFMA効率99%がそのままワークロードが速いことを意味するわけではありません — ホットループの中で行列エンジンが遊んでいる時間がほとんどないという意味であり、それはカーネルが正しいアルゴリズムを使っているときにのみ意味を持ちます。)

レイアウトの底 — 線形レイアウト

もう一段深く入ると、線形レイアウト(linear layout)があります。チュートリアルが説明するとおり、Gluonには正準(canonical)なレイアウト表現がありません — 複数のレイアウトが同じ要素マッピングを表しうるのです。たとえば以下の2つは同等です。

gl.BlockedLayout([1], [32], [4], [0])
gl.SliceLayout(1, gl.BlockedLayout([1, 1], [32, 1], [4, 1], [1, 0]))

同等だと分かっていれば、gl.convert_layout(x, layout, assert_trivial=True)で、その変換が実際にタダなのか(レジスタの並べ替え以上のことをしていないか)をコンパイラに検証させられます。明示的な言語らしいAPIです。

そしてすべてのGluonレイアウトは線形レイアウトとして表現可能です。チュートリアルはこれを「最も表現力があり強力な表現」と呼びながら、同時に「比較的まれで理解しにくいことがある」とも付け加えています。基盤となる理論は論文として出ています — Keren Zhou、Mario Lezcano、Jeff Niu、Phil Tillet、Thomas Raouxらによる「Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using F2」(arXiv:2505.23819、2025年5月28日提出、2026年3月6日改訂)。レイアウトをGF(2)上の二進行列としてモデル化し、既存システムの二次的な複雑さを避けるというアプローチです。

実務者にとっての要点はこうです — ほとんどの場合BlockedLayoutSliceLayoutで十分であり、線形レイアウトは必要になるまで知らなくてもかまいません。ただし、コストゼロのsplit/join/reshape/permuteが必要になる瞬間が来たら、それがある場所はここです。

正直なトレードオフ — いつ使うべきでないか

ここからは正直になる部分です。

まだexperimentalです。 importパスはtriton.experimental.gluonで、リポジトリ上でもpython/triton/experimental/gluonに置かれています。公式のGluon概要ドキュメントには実験的な状態や安定性について別途の警告文はありませんが、パスそのものがAPIの安定性を約束していないと語っています。3.4.0から3.7.0までリリースのたびにAPIが追加・変更され続けてきた経緯も、この解釈と一致します。プロダクションカーネルをここに載せるなら、Tritonのバージョンを上げるたびに壊れる覚悟が必要です。

要求される知識の量が違います。 チュートリアルの導入部は率直です — Gluonカーネルを書くには、GPUハードウェアとGPUプログラミングのさまざまな側面についてより深い理解が必要です。ワープあたりのスレッドが32であること、結合がなぜ内側の次元に依存するのか、共有メモリのバンク衝突とは何か、レジスタ圧力が占有率にどう作用するかを知らなければ、BlockedLayoutの4つの引数を埋めることはできません。Tritonがこれを隠してくれていたのは怠慢ではなく、サービスでした。

Gluonへ移すこと自体はタダの性能ではありません。 変換器PRが言うとおり — 素朴に移したGluonはかなり遅くなります。そして先に見たRのスイープが示すように、直感的な選択(幅を広げること)がむしろ悪化させることもあります。自動化されたコンパイラがやっていた仕事を人間がやることにしたなら、人間がその仕事をより上手くやっているという証拠を、毎回ベンチマークで作らなければなりません。

チュートリアル全体がmemcpyです。 本稿で引用したGB200の数字は、上流の開発者が測定してソースに記録した値であり、私が再現したものではありません。そして対象はメモリ帯域幅に律速されたmemcpyです — あなたのカーネルが演算律速なら、レイアウトのスイープで10倍は出ません。

まとめると、判断基準はこうなります。

使って価値がある場合

  • すでにTritonで書いており、プロファイリングでコンパイラが出したコードがハードウェアの限界にはるかに届かないことを測定で確認した。
  • ボトルネックがレイアウト・データ移動・同期にあり、Tritonのレベルでは触れる手段がない。
  • 最新のハードウェア機能(Blackwellのtcgen05、Hopper WGMMA、2CTA、TMA、CDNA4のMXFP4)を直接叩く必要がある。
  • カーネルが長く生き、頻繁に動くため、チューニングの時間とバージョンアップのたびの維持コストを回収できる。

過剰な場合

  • まだTritonでプロファイリングしていない。(これが大半です。)
  • torch.compileが出すコードで十分。
  • カーネルが演算律速で、レイアウトがボトルネックではない。
  • ベンダーライブラリ(cuBLAS、hipBLASLt、cuDNN)がすでにその演算をカバーしている — AMDのブログ自身も、自分たちのチュートリアルはhipBLASLtの代替ではないと明言しています。
  • チームにGPUマイクロアーキテクチャを知る人がいない。

おわりに

Gluonは「より良いTriton」ではありません。別の契約です。

Tritonの契約は「詳細は私が処理するから、あなたはアルゴリズムだけを書け」です。多くの場合これは良い取引であり、これからも大半のカーネルはTritonにとどまるべきです。Gluonの契約は「詳細はあなたが処理しろ、その代わり行き詰まったときに壁ではなく扉がある」です。

このスペクトルが重要です。かつてはTritonとCUDA C++の間に断崖がありました — コンパイラが出してくれなければ、言語とエコシステムをまるごと乗り換える必要がありました。Gluonはその断崖に階段を架けます。同じコンパイラスタック、同じフロントエンド、同じJIT、同じオートチューナーにとどまったまま、必要な部分だけ下へ降ります。そしてその階段がNVIDIAとAMDの両方に架けられつつあること — gfx950チュートリアルをAMD自身が書き、リリースノートのたびにAMDの項目がつくこと — の方がもしかすると重要なニュースかもしれません。

ただし順序は守ってください。Tritonで書き、プロファイリングし、コンパイラが実際に負けていることを確認し、それからGluonを取り出してください。0.774 TB/sと6.590 TB/sを分けたのがアルゴリズムではなくレイアウト記述子一行だったという事実は諸刃の剣です — うまく選べば10倍を得ますが、選び間違えれば10倍を失います。Tritonはあなたがその選択をしなくて済むようにしてくれました。Gluonはその選択をあなたに返します。それが贈り物か請求書かは、あなたがプロファイラを開いたかどうかにかかっています。

参考資料