learn-cute-dsl

CUTLASS の例から CuTe DSL を学ぶ

CUTLASS Blackwell の例を読み、必要に応じてプロジェクトのランタイムインフラストラクチャに取り込み、プロファイリングを行い、再利用可能なパターンをナレッジベースにまとめることで、CuTe Python DSL のスキルを体系的に身につけるワークフローです。

例のソース

多数の例が third_party/cutlass/examples/python/CuTeDSL にあります。ドキュメントは docs/cute-dsl/（CuTe Python DSL）と docs/cutlass-cpp/cute/（CuTe C++ の概念）にあります。

カテゴリ	例	サブディレクトリ
MLA デコード	mla_decode_fp16.py、mla_decode_fp8.py、mla_helpers.py	mla/
アテンション	fmha.py、fmha_bwd.py	（トップレベル）
混合アテンション	mixed_input_fmha_decode.py、mixed_input_fmha_prefill_*.py	mixed_input_fmha/
GEMM（基本）	dense_gemm.py	（トップレベル）
GEMM（パイプライン化）	dense_gemm_software_pipeline.py	（トップレベル）
GEMM（永続型）	dense_gemm_persistent.py、dense_gemm_persistent_prefetch.py	（トップレベル）
GEMM（動的永続型）	dense_gemm_persistent_dynamic.py	（トップレベル）
GEMM チュートリアル	fp16_gemm_0.py から fp16_gemm_6.py	tutorial_gemm/
ブロックスケーリング GEMM	dense_blockscaled_gemm_persistent.py、*_prefetch.py、*_amax.py	（トップレベル）
混合入力 GEMM	mixed_input_gemm.py、grouped_mixed_input_gemm.py	mixed_input_gemm/
エピローグ	custom_epilogue_dense_gemm.py、activation_custom_epilogue_dense_gemm.py	epilogue/
リダクション	reduce.py、rmsnorm.py	（トップレベル）
依存関係のある起動	programmatic_dependent_launch.py	（トップレベル）
Mamba2	mamba2_ssd.py	mamba2_ssd/
ブロック単位 GEMM	blockwise_gemm.py、contiguous_grouped_gemm.py	blockwise_gemm/

その他の例は、ユーザーから直接提供してもらうこともできます。

ワークフロー

1. SELECT    -> Pick examples based on learning goals
2. READ      -> Read and annotate, identify CuTe DSL patterns
3. IMPORT    -> (Optional) Wrap in CuteKernel + KernelPlan for profiling
4. PROFILE   -> Profile imported kernels, compare with existing cuTile versions
5. EXTRACT   -> Distill reusable CuTe DSL patterns and anti-patterns
6. DOCUMENT  -> Update knowledge base (docs/knowledge/languages/cute-dsl/)

---

ステップ 1：選択

学びたい内容に基づいて例を選びます。

学習目標	最初に読む例
CuTe DSL の基本（グリッド、TMA、MMA）	tutorial_gemm/fp16_gemm_0.py から fp16_gemm_6.py
ソフトウェアパイプライン化	dense_gemm_software_pipeline.py
永続カーネルの構造	dense_gemm_persistent.py
Warpgroup の特化	dense_gemm_persistent_prefetch.py
アテンション向け TMA パイプライン化	fmha.py
MLA デコード（直接参照）	mla/mla_decode_fp16.py + mla/mla_helpers.py
FP8 を使用した MLA デコード	mla/mla_decode_fp8.py
クラスタプログラミング	dense_gemm_persistent_dynamic.py
カスタムエピローグ	epilogue/custom_epilogue_dense_gemm.py
依存関係のあるカーネル起動	programmatic_dependent_launch.py

読む前に、具体的な疑問を書き出します。例：

• 「MLA デコードカーネルは KV キャッシュ用の TMA ディスクリプターをどのように設定するか？」
• 「永続 GEMM では、Warpgroup にプロデューサー／コンシューマーの役割をどのように割り当てるか？」
• 「FMHA は K ループのパイプラインにどのようなバリアパターンを使用するか？」

---

ステップ 2：読む

選んだ各例について、ソース全体を読み、次の CuTe DSL 固有の側面に注釈を付けます。

ホスト側（@cute.jit）

1. TMA ディスクリプターの作成 -- どのテンソルに TMA ディスクリプターを設定するか、タイルの形状はどうなっているか。
2. グリッドとクラスタの次元 -- グリッドをどのように計算するか、クラスタリングを使用するか。
3. 共有メモリのサイズ -- パイプラインのステージ数とタイル形状から SMEM の予算をどのように計算するか。
4. 起動パラメーター -- max_register_count、クラスタ形状、SMEM カーブアウト

デバイス側（@cute.kernel）

1. スレッド／ワープ／Warpgroup の識別 -- 役割をどのように割り当てるか。
2. パイプラインの構造 -- ステージ数、バリアの初期化、プロデューサー／コンシューマーのループ
3. TMA 操作 -- cute.copy(tma_desc, ...)、到着パターン
4. MMA 操作 -- アキュムレーターの設定、MMA ループ本体、エピローグ
5. 同期 -- cute.barrier_arrive()、cute.barrier_wait()、cute.cluster_barrier_*
6. 共有メモリのレイアウト -- 割り当て、スウィズリング、Warpgroup 間の分割

パターンのチェックリスト

読みながら、以下に該当する箇所に印を付けます。

• TMA マルチステージパイプライン（ステージ数、各ステージのバリア）
• ワープグループの特化（どのワープグループが生成し、どのワープグループが消費するか）
• 永続カーネルのメインループ（ワークタイルのスケジューリング、終了条件）
• オンラインソフトマックスまたはストリーミング累積
• レジスタレベルのアキュムレーター管理
• 共有メモリのスウィズルパターン
• クラスターレベルの通信（分散共有メモリ）
• エピローグのパターン（レジスター内、SMEM 経由、または融合）
• 境界処理（部分タイルの述語評価、マスキング）
• パイプラインの排出（最後のパイプラインステージのフラッシュ方法）

相互参照

ナレッジベースに既存のパターンがあるか確認します。

• docs/knowledge/optimizations/ — 共有パターン（例：pipeline-driven-low-occupancy）
• docs/knowledge/languages/cute-dsl/ — 既存の CuTe DSL オーバーレイ

この例が既存の知識を裏付けているか、精緻化しているか、または反証しているかを記録します。

---

ステップ 3：インポート（任意）

CUTLASS の例をプロジェクトにインポートし、標準 CLI とプロファイリング基盤で実行できるようにします。

インポートする場合

• お使いのハードウェアで例をプロファイリングする場合
• カスタム CuTe DSL カーネルの出発点が必要な場合
• 既存の cuTile カーネルと同等のワークロード（例：MLA デコード）を例が扱う場合

インポートしない場合

• ベンチマークではなく、パターンの理解だけが目的の場合
• 例に複雑な依存関係や、簡単に対応付けられないヘルパー基盤がある場合

インポート手順

1. カーネルパッケージを作成します。

   src/mla_var3/kernel/cute_python/<layer>/<design>/<design>/
   

   CUTLASS MLA デコード FP16 の例：

   src/mla_var3/kernel/cute_python/mla/cutlass_mla_decode/cutlass_mla_decode/
   

2. KernelPlan サブクラスを作成します（cutlass_mla_decode.py）。

   • prepare_inputs() — 例で想定されている形状に合わせてテンソル（Q、KV キャッシュ、出力）を割り当てます
   • reference_fn() — 正確性確認用の PyTorch リファレンス
   • _autotune_configs() — 例に基づくタイル分割パラメーター
   • _algorithmic_flops_bytes() — ルーフライン分析
   • plan() — 例のホスト関数／デバイス関数をラップする CuteKernel を返します

3. __main__.py を作成します。

   from .cutlass_mla_decode import CutlassMlaDecodeKernel
   
   if __name__ == "__main__":
       kernel_plan = CutlassMlaDecodeKernel()
       kernel_plan.benchmark_kernel_argparse()
   

4. 各パッケージレベルに __init__.py を作成します

5. ヘルパーモジュールを処理します：例でヘルパー（例：mla_helpers.py）を使用している場合、次のいずれかを行います。

   • カーネルパッケージにコピーして、インポートを調整する
   • CUTLASS のパスからインポートする（壊れやすいものの、探索にはより高速）

6. 正確性をテストします。

   source .venv/bin/activate && python -m mla_var3.kernel cutlass_mla_decode --prof_type=disabled --check
   

---

ステップ 4：プロファイル

アーキテクチャの互換性チェック

プロファイリングの前に、利用可能なデバイスがカーネルの対象アーキテクチャと一致することを確認します。

nvidia-smi --query-gpu=name --format=csv,noheader

CUTLASS の例の多く（特に mla/ 内の例や、tcgen05 MMA 演算を使用する例）は、特に SM100（B200、B300 データセンター GPU）を対象としています。「Blackwell」は複数の SM バージョンにまたがるマーケティング上の名称であり、命令セットが異なることに注意してください。

• SM100（B200、B300）— データセンター GPU。tcgen05 MMA、TMEM、完全なクラスター機能をサポートします
• SM120（GeForce RTX 5090、RTX 5080）— コンシューマー向け Blackwell。sm_120a を使用し、tcgen05 演算はサポートしません

tcgen05 演算を使用するカーネルは、どちらも「Blackwell」として販売されていても、RTX 5090（SM120）では実行できません。マーケティング上の世代ではなく、GPU 名と実際の SM バージョンを照合してください。

カーネルの対象 SM バージョンが利用可能なデバイスと一致しない場合は、PROFILE のステップを完全にスキップします。コンパイル時に失敗するか、誤解を招く結果が出るためです。その場合は、ソースコードを読んで得た知見のみをもとに、ステップ 5（EXTRACT）に進み、ドキュメントにはアーキテクチャの不一致を記載してください。

プロファイリングの手順

1. GPU クロックを固定します：bash scripts/lock-gpu-clock.sh

2. まずアノテーションモードを実行します。

   source .venv/bin/activate && python -m mla_var3.kernel <kernel> \
       --b=32 --s=1 --t=4096 --prof_type=annotation
   

3. アノテーションで興味深いパターンが見つかった場合は、NCU で詳細に調査します。

   source .venv/bin/activate && python -m mla_var3.kernel <kernel> \
       --b=32 --s=1 --t=4096 --prof_type=ncu
   

注目する領域

目的は単なるベンチマークではなく学習であるため、カーネルがどのように性能を実現しているかを理解することに重点を置きます。

• パイプライン効率：TMA コピーは計算と完全にオーバーラップしていますか？nsys のタイムラインでアイドル状態の空白を確認します。
• Warpgroup の稼働率：すべての warpgroup が稼働していますか？NCU でスケジューラーの統計を確認します。
• レジスター圧力と占有率のトレードオフ：スレッドあたりのレジスター数はいくつですか？占有率が低いのは意図的ですか？
• SMEM の使用状況：利用可能な SMEM 予算のうち、どの程度を使用していますか？パイプラインのステージはいくつ収まりますか？
• MMA のスループット：ピーク TC 利用率に対して、どの程度を達成していますか？

cuTile のバージョンと比較する

インポートしたカーネルが cuTile に存在するワークロード（例：MLA デコード）に対応している場合は、比較します。

指標	cuTile バージョン	CUTLASS CuTe DSL	差分	考察
実行時間（us）
TC%
DRAM%
占有率
スレッドあたりのレジスター数
ブロックあたりの SMEM（KB）

違いから、明示的な制御（warpgroup スケジューリング、TMA パイプライン処理）によって、cuTile のブロックレベルの抽象化に何が加わるかがわかります。

---

ステップ 5：抽出

観察結果を、再利用可能な CuTe DSL パターンとしてまとめます。

分類

種類	保存先
共通のアルゴリズム／ハードウェアパターン	docs/knowledge/optimizations/<name>.md
共通の失敗パターン	docs/knowledge/anti-patterns/<name>.md
CuTe DSL API パターンまたは実装の詳細	docs/knowledge/languages/cute-dsl/optimizations/<name>.md
CuTe DSL の落とし穴や注意点	docs/knowledge/languages/cute-dsl/anti-patterns/<name>.md

CUTLASS の例から得られるパターンの大半は、CuTe DSL 固有（API パターン、TMA のセットアップ方法、バリアの連携方法）です。共通パターン（パイプライン駆動のスケジューリングなどのアルゴリズムに関する知見）は、cuTile の作業ですでに存在している場合があります。

抽出基準

次の条件に当てはまる場合は、パターンとして記録します。

• 他の CuTe DSL カーネルにも応用できる（この例だけに限られない）
• CuTe DSL のドキュメントだけでは明らかでない
• プロファイリングで確認できる、またはアーキテクチャ上明確な性能への影響がある
• 既存の共通パターンを CuTe DSL 固有の実装詳細で補足する

パターンのテンプレート

# [Pattern Name]

## When to Apply
- [CuTe DSL kernel type, workload shape, bottleneck condition]

## Mechanism
[How the pattern works in CuTe DSL -- reference specific APIs]

## Affected Metrics
- [Metric 1]
- [Metric 2]

## Implementation
```python
# CuTe DSL code snippet, generalized from the CUTLASS example

ソース

[CUTLASS の例の名前とパス]

性能の根拠

ソース：[例の名前、デバイス、構成]

構成	指標	値	コンテキスト

一般化

[デバイスに依存しない要点]

注意点

• [CuTe DSL 固有の失敗パターン]

相互作用

• [他の CuTe DSL パターンとの相互作用]

---

## Step 6: DOCUMENT

1. Write pattern files to `docs/knowledge/languages/cute-dsl/optimizations/` or `anti-patterns/`
2. Update the index in `/optimization-catalog-cute-dsl`
3. If a pattern is shared (not DSL-specific), write to `docs/knowledge/optimizations/` and update `/optimization-catalog`
4. If you imported and profiled a kernel, write a devlog entry in `docs/kernels/`:

```markdown
# [Kernel Name] (Imported from CUTLASS)

## Overview
[What this kernel does and why it was imported for study]

## Source
Original: `third_party/cutlass/examples/python/CuTeDSL/<path>`
Imported: `src/mla_var3/kernel/cute_python/mla/<design>/`

## Performance
[Profiling results table]

## Patterns Extracted
- [Pattern] -> `docs/knowledge/languages/cute-dsl/optimizations/<name>.md`

## Insights
[What was learned that informs future CuTe DSL kernel design]

---

推奨する学習経路

経路 A：CuTe DSL の基礎

MLA カーネルを書く前に、基本的な知識を身につけるための経路です。

1. tutorial_gemm/fp16_gemm_0.py  -> Minimal kernel: grid, TMA load, MMA, store
2. tutorial_gemm/fp16_gemm_3.py  -> Software pipelining basics
3. tutorial_gemm/fp16_gemm_6.py  -> Full persistent GEMM with warpgroup specialization
4. dense_gemm_software_pipeline.py -> Production-quality pipelining
5. dense_gemm_persistent.py      -> Persistent scheduling pattern

経路 B：MLA 固有のパターン

CUTLASS における MLA の実装方法を直接学ぶための経路です。

1. mla/mla_helpers.py            -> Shared utilities, tensor layouts, TMA setup
2. mla/mla_decode_fp16.py        -> FP16 MLA decode with full pipeline
3. mla/mla_decode_fp8.py         -> FP8 variant (quantization handling)
4. fmha.py                       -> Attention forward (different design choices)

経路 C：高度なパターン

特定の最適化手法を学ぶための経路です。

1. dense_gemm_persistent_prefetch.py  -> Prefetching in persistent kernels
2. dense_gemm_persistent_dynamic.py   -> Dynamic tile scheduling + clusters
3. epilogue/custom_epilogue_dense_gemm.py -> Fused epilogue patterns
4. programmatic_dependent_launch.py   -> Dependent kernel launch