cute-dsl-ref

CuTe Python DSL リファレンス

実行モデル

CuTe Python DSL は、NVIDIA の CuTe レイアウト代数を Python から利用するためのインターフェースです。cuTile のブロックレベルの抽象化とは異なり、CuTe DSL では スレッド／ワープ／ワープグループの明示的な制御、TMA パイプライン、バリアの制御、共有メモリの管理が可能です。

2 レベルのホスト／デバイスパターン

CuTe DSL の各カーネルには、次の 2 つの関数があります。

1. @cute.jit ホスト関数 — CPU 上で実行され、TMA ディスクリプターを設定し、グリッドを計算し、共有メモリを確保して、カーネルを起動します
2. @cute.kernel デバイス関数 — GPU 上で実行され、実際の計算を行います

import cutlass
import cutlass.cute as cute

@cute.kernel
def my_kernel(tiled_mma: cute.TiledMma, ...):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()
    # ... GPU code

@cute.jit
def host_fn(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
    # Setup TMA descriptors, compute grid, allocate SMEM
    my_kernel(...).launch(grid=grid_shape, block=block_shape, smem=smem_bytes)

コンパイルパイプライン

3 段階の JIT コンパイルです。

1. プリステージング：Python AST の書き換え
2. メタステージ：Python インタープリターがメタプログラムを実行します（コンパイル時定数を解決します）
3. オブジェクトステージ：コンパイラーバックエンドが PTX → SASS を生成します

JIT アーティファクトは自動的にキャッシュされます。CUTE_DSL_CACHE_DIR と CUTE_DSL_DISABLE_FILE_CACHING で制御します。

cuTile との主な違い

項目	cuTile	CuTe DSL
抽象化レベル	ブロックレベル（スレッド ID なし）	スレッド／ワープ／ワープグループレベル
TMA	暗黙的（allow_tma=True）	ホストで明示的にディスクリプターを設定
共有メモリ	コンパイラーが管理	明示的に確保し、レイアウトを指定
同期	公開されていません	バリア、到着／待機、名前付きバリア
パイプライン	公開されていません	明示的な多段非同期パイプライン
MMA	タイルに対する ct.mma(A, B, C)	分割フラグメントに対する cute.gemm(tiled_mma, d, a, b, c)
コンパイル	cuTile バイトコード → MLIR	CuTe DSL JIT → PTX → SASS

---

コア API 一覧

デコレーターとコンパイル

API	用途	主なパラメーター
@cute.kernel	デバイスカーネルのデコレーター	関数名はモジュールのファイル名と一致する必要があります
@cute.jit	ホスト JIT 関数のデコレーター	—
kernel(...).launch(grid, block, smem, cluster)	カーネルを起動	grid=(x,y,z)、block=(threads,1,1)、smem=bytes、cluster=(cx,cy,cz)
cute.compile(fn, *args, options="...")	事前コンパイル	"opt-level=3"、"keep-ptx" など

スレッド／ブロック／グリッドのインデックス指定

API	用途
cute.arch.thread_idx()	(tidx, tidy, tidz) を返します
cute.arch.block_idx()	(bidx, bidy, bidz) を返します
cute.arch.block_dim()	ブロックの次元
cute.arch.grid_dim()	グリッドの次元
cute.arch.warp_idx()	ブロック内のワープインデックス
cute.arch.block_idx_in_cluster()	クラスター内のブロックインデックス（Hopper 以降）

レイアウトとテンソルの構築

API	用途	備考
cute.make_layout(shape, stride)	レイアウトを作成	中核となる抽象化：座標をインデックスに対応付けます
cute.make_ordered_layout(shape, order)	ストライドの順序を指定して作成	例：列優先の場合は order=(1,0)
cute.make_identity_layout(shape)	恒等レイアウト	述語評価に使用
cute.make_tensor(ptr, layout)	テンソルビューを作成	ポインターとレイアウトを組み合わせます
cute.make_fragment(layout, dtype)	レジスタテンソルを確保	アキュムレーター用
cute.make_fragment_like(src, dtype)	ソースに対応するレジスタテンソル	—
cute.make_ptr(dtype, addr, memspace)	タグ付きポインター	cute.AddressSpace.smem/gmem
cute.from_dlpack(tensor)	PyTorch／JAX テンソルを変換	.mark_layout_dynamic() を持つ cute.Tensor を返します

レイアウト演算

API	目的
cute.size(t, mode=None)	全体のサイズまたはモードのサイズ
cute.shape(t)	形状タプルを取得
cute.stride(t)	ストライドタプルを取得
cute.rank(t)	モード数
cute.local_tile(tensor, tiler, coord, proj)	ブロックのタイルを抽出
cute.logical_divide(tensor, divisor)	タイルと残りに分割
cute.composition(layout1, layout2)	レイアウトを合成
cute.complement(layout)	レイアウトの補レイアウト
cute.coalesce(layout)	レイアウトを簡略化
cute.flatten(t)	テンソル／レイアウトをフラット化
cute.group_modes(layout, start, end)	モードをまとめてグループ化
cute.tile_to_shape(tile, shape)	目標の形状に合わせてタイル化
cute.ceil_div(a, b)	切り上げ除算

MMA 演算

API	目的	備考
cute.make_tiled_mma(op)	演算からタイル化 MMA を作成	例：tcgen05.MmaF16BF16Op(...)
tiled_mma.get_slice(thread_idx)	スレッドのパーティションを取得	ThrMma を返します
thr_mma.partition_A(tensor)	A オペランドをパーティション化	スレッドレベルのビュー
thr_mma.partition_B(tensor)	B オペランドをパーティション化	スレッドレベルのビュー
thr_mma.partition_C(tensor)	C オペランドをパーティション化	スレッドレベルのビュー
tiled_mma.partition_shape_C(shape)	C パーティションの形状	フラグメントの割り当て用
tiled_mma.make_fragment_A(tensor)	A フラグメントを作成	レジスター割り当て
tiled_mma.make_fragment_B(tensor)	B フラグメントを作成	レジスター割り当て
tiled_mma.make_fragment_C(shape)	C フラグメントを作成	レジスター割り当て
cute.gemm(tiled_mma, d, a, b, c)	MMA を実行：d = a @ b + c	ハードウェア MMA にディスパッチ

コピーおよび TMA 演算

API	目的	備考
cute.make_tiled_copy(atom)	タイル化コピーを作成	一括データ転送用
cute.copy(atom, src, dst, **kw)	コピーを実行	tma_bar_ptr=, mcast_mask=
cute.basic_copy(src, dst)	要素単位のコピー	アトムは不要
cute.prefetch(atom, src)	TMA ディスクリプターをプリフェッチ	—
cute.nvgpu.make_tiled_tma_atom_A(op, tensor, smem_layout, tile, mma)	A 用の TMA アトムを作成	ホスト側のみ
cute.nvgpu.make_tiled_tma_atom_B(op, tensor, smem_layout, tile, mma)	B 用の TMA アトムを作成	ホスト側のみ

共有メモリ

API	目的
cutlass.utils.SmemAllocator()	SMEM アロケーターを作成
smem.allocate_tensor(dtype, layout, align, swizzle)	SMEM にテンソルを割り当て
smem.allocate(struct_type)	SMEM に構造体を割り当て

テンソルメモリ（SM100 のみ）

API	目的
cutlass.utils.TmemAllocator(...)	TMEM アロケーターを作成
tmem.allocate(num_cols)	列を割り当て
tmem.wait_for_alloc()	割り当てを待機
tmem.retrieve_ptr(dtype)	ポインターを取得
tmem.free(ptr)	メモリを解放

同期

API	目的
cute.arch.mbar.init(barrier_ptr, count)	バリアを初期化
cute.arch.mbar.arrive(barrier_ptr)	バリアに到達
cute.arch.mbar.wait(barrier_ptr, phase)	バリアを待機
cute.arch.mbar.arrive_and_expect_tx(barrier_ptr, bytes)	予期される TX バイト数を指定してバリアに到達
cute.arch.mbar.try_wait(barrier_ptr, phase)	ノンブロッキングで待機
パイプラインクラス（リファレンスを参照）	複数ステージの非同期パイプライン

数学関数（TensorSSA 上の要素単位演算）

API	目的	API	目的
cute.exp(x)	e^x	cute.exp2(x)	2^x
cute.log(x)	ln(x)	cute.log2(x)	log2(x)
cute.sqrt(x)	sqrt	cute.rsqrt(x)	1/sqrt(x)
cute.sin(x)	sin	cute.cos(x)	cos
cute.tanh(x)	tanh	cute.erf(x)	erf

すべての数学関数で、近似ハードウェア組み込み関数を使用する fastmath=True を指定できます。

テンソル操作

API	目的
cute.full(shape, val, dtype)	値で埋めます
cute.zeros_like(tensor)	同じ形状のゼロテンソルを作成します
cute.where(cond, x, y)	条件に応じて選択します
cute.any_(tensor)	論理和を判定します
cute.all_(tensor)	論理積を判定します

デバッグ

API	目的
cute.printf(fmt, ...)	デバイス側の printf
cute.print_tensor(tensor)	テンソルの内容を出力します
print(...)	コンパイル時の出力（メタステージのみ）

---

主な制約

1. @cute.kernel 関数の名前はモジュールのファイル名と一致する必要があります — cuTile と同じ規則です

2. TMA ディスクリプターはホスト側でのみ使用できます — @cute.jit で作成し、@cute.kernel に渡します

3. レジスターの上限はスレッドあたり 255 個です — docs/devices/ の仕様で確認してください

4. SMEM の上限はデバイスと SM のバージョンによって異なります — SM100（B200/B300）：228 KB/SM、オプトイン時は 227 KB/block、SM120（RTX 5090）：96 KB

5. パイプラインのステージは SMEM を消費します — 各ステージに専用のバッファが必要です。合計を確認してください

6. バリア同期のエラーは通知されません — クラッシュせずに誤った結果を引き起こします。必ず --check でテストしてください

7. クラスターの次元はグリッドと互換性が必要です — クラスターの形状でグリッドを割り切れる必要があります

8. print() はコンパイル時専用です — デバイス側への出力には cute.printf() を使用してください

9. ループ内で早期終了する break は使用できません — 代わりに述語を使用してください

10. レイアウト演算は 32 ビットです — 形状とストライドは 32 ビット整数に制限されます

11. アーキテクチャのサポート：Ampere（SM80）以降に対応します。tcgen05 や TMEM を含むすべての機能を使用するには、SM100（B200/B300）が必要です

12. アーキテクチャ固有のカーネル：特定の SM バージョンを対象とするカーネルがあります（例：tcgen05 MMA 演算には「Blackwell」だけでなく SM100 が必要です）。「Blackwell」というマーケティング名は、命令セットが異なる複数の SM バージョンを指します：

    • SM100（B200、B300）— データセンター向け GPU。tcgen05 MMA、TMEM、およびすべてのクラスター機能に対応します
    • SM120（GeForce RTX 5090、RTX 5080）— コンシューマー向け Blackwell。sm_120a を使用し、tcgen05 演算には対応していません

    実行、コンパイル、プロファイリングの前に、必ずデバイスを確認してください：nvidia-smi --query-gpu=name --format=csv,noheader。次に、カーネルの対象 SM バージョンと照合してください。GPU 名だけでは不十分で、対応する SM バージョンを把握する必要があります。SM バージョンが一致しない場合は、失敗したり誤解を招く結果になったりする実行を試みず、実行／プロファイリングの手順をスキップしてください。

制御フロー

構文	動作	備考
for i in range(N)	N が静的な場合は展開されます	標準の Python range
for i in cutlass.range(N)	実行時ループ	MLIR ループを生成します
for i in cutlass.range_constexpr(N)	コンパイル時に展開されます	N は静的である必要があります
if cutlass.const_expr(cond)	コンパイル時の分岐	コンパイル時に除去されます
if cond	実行時の分岐	両方の分岐で型チェックが必要です

---

Blackwell データセンター向け MMA 演算（tcgen05）— SM100 専用

SM100 が必要です。 tcgen05 演算は B200/B300（SM100）でのみ使用できます。RTX 5090（SM120/sm_120a）などのコンシューマー向け Blackwell GPU では使用できません。

from cutlass.cute.nvgpu import tcgen05

# Create MMA operation for SM100 (B200/B300) tensor cores
op = tcgen05.MmaF16BF16Op(
    dtype=cutlass.Float16,           # Input type
    acc_dtype=cutlass.Float32,       # Accumulator type
    shape=(128, 128, 64),            # M x N x K tile shape
    cta_group=tcgen05.CtaGroup.ONE,  # ONE or TWO CTAs
)

tiled_mma = cute.make_tiled_mma(op)

演算	入力型	形状
MmaF16BF16Op	FP16/BF16 → FP32	さまざまな M×N×K
MmaF8F6F4Op	FP8/FP6/FP4 → FP32	低精度
ブロックスケーリング版	スケール係数を使用	例を参照

---

一般的なパターン

TMA パイプラインを使用する GEMM（簡略版）

@cute.kernel
def gemm_kernel(tiled_mma, tma_a, tma_b, smem_a, smem_b, gmem_c, ...):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()

    # Get thread's MMA partition
    thr_mma = tiled_mma.get_slice(tidx)

    # Allocate accumulator in registers
    acc = tiled_mma.partition_shape_C(tile_shape)
    tCrC = cute.make_fragment(acc, cutlass.Float32)

    # K-loop with TMA pipeline
    for k_tile in range(num_k_tiles):
        # TMA copy: global → shared (async)
        cute.copy(tma_a, gA_tile, sA_tile, tma_bar_ptr=barrier)
        cute.copy(tma_b, gB_tile, sB_tile, tma_bar_ptr=barrier)

        # Wait for TMA to complete
        cute.arch.mbar.wait(barrier, phase)

        # Partition shared memory for this thread
        tCsA = thr_mma.partition_A(sA)
        tCsB = thr_mma.partition_B(sB)

        # MMA: accumulate into registers
        cute.gemm(tiled_mma, tCrC, tCsA, tCsB, tCrC)

    # Epilogue: write back to global memory
    cute.basic_copy(tCrC, gC_partition)

Warpgroup の特殊化（プロデューサー／コンシューマー）

warp_idx = cute.arch.warp_idx()
is_producer = warp_idx < num_producer_warps

if is_producer:
    # Issue TMA copies for upcoming pipeline stages
    for stage in range(num_stages):
        cute.copy(tma_atom, src, dst, tma_bar_ptr=barriers[stage])
        cute.arch.mbar.arrive_and_expect_tx(barriers[stage], bytes_per_stage)
else:
    # Consume data from shared memory, execute MMA
    for stage in range(num_stages):
        cute.arch.mbar.wait(barriers[stage], phase)
        cute.gemm(tiled_mma, acc, sA_stage, sB_stage, acc)

永続カーネルのループ

@cute.kernel
def persistent_kernel(tiled_mma, num_tiles, ...):
    bidx, _, _ = cute.arch.block_idx()
    num_blocks = cute.arch.grid_dim()[0]

    tile_id = bidx
    while tile_id < num_tiles:
        # Process tile
        # ... TMA copy, MMA, epilogue ...
        tile_id += num_blocks

---

詳細なリファレンス

• アーキテクチャ操作（スレッドのインデックス付け、TMA、MMA アトム、バリア、SMEM/TMEM）：references/architecture-ops.md を参照してください
• 公式ドキュメントの索引：references/docs-index.md を参照してください
• カーネルの例：カテゴリ別のサンプルソース索引については、/learn-cute-dsl スキルを参照してください