design-cutile-dsl-kernel

cuTile Python DSL — 言語固有のガイダンス

前提条件：共有の命名規則、バージョン管理、KernelPlan の構造、合成パターン、クローンワークフロー、開発ログのテンプレート、設計者の出力契約については、/design-kernel を読み込んでください。このスキルでは、cuTile 固有のランタイムパターンと制約のみを扱います。

cuTile Python DSL を使用する場合

次の最適化が以下を通じて表現できる範囲である場合は、cuTile Python DSL（cutile-dsl）を使用してください。

• タイルサイズとテンソルレイアウトの選択
• ct.bid() / ct.num_blocks() を使った CTA の再マッピングと作業の分割
• 占有率とクラスターサイズのヒント（num_ctas、occupancy）
• latency / allow_tma のヒントなど、コンパイラーへの指示
• KernelPipeline / ConcurrentKernels による複数段階の合成

適合性の判定

cuTile で明示的な CTA 内スケジューリングを表現できるとは考えないでください。公開されている実行モデルでは、スレッド、ワープ、ワープグループの ID は公開されておらず、ブロック内の明示的な同期や通信もできません。同様に、num_ctas を使って公開されているクラスターのプログラミングモデルを利用することもできません。公開されているクラスターランク、クラスターのバリア、DSM プリミティブ、クラスターのメモリースコープはありません。

プロファイリングの結果、カーネルが SM あたり 1 つの CTA に固定されている場合、これは設計上の重大な制約になります。この状態で最先端のカーネルが通常採用する解決策は、CTA 内で明示的にスケジューリングすることです。たとえば、ワープグループのシーソースケジュール、プロデューサー／コンシューマーパイプライン、バリアの調整、クラスター協調型のオーバーラップなどがあります。必要な最適化がこのカテゴリに当てはまる場合は、cuTile での反復を中止し、CuTe DSL に切り替えてください（/design-cute-dsl-kernel）。

FlashMLA の開発ログ（docs/kernels/flash-mla.md）が具体例です。公式の FlashMLA 設計は明示的なスケジューリングを使うため、SM あたり 1 つの常駐ブロックでも高い生産性を維持できますが、cuTile の FlashMLA バージョンではその動作を再現できません。

命名とレイアウト

• 公開言語キー：cutile-dsl
• Python パッケージパス：cutile
• カーネルのレイアウト：src/mla_var3/kernel/cutile/<layer>/<design>/<design>[_vN]/
• モジュールファイル：<design>[_vN].py
• メインのランタイムラッパー：CtKernel

CtKernel のランタイムパターン

CtKernel(Kernel) は、grid_fn および args_fn の呼び出し可能オブジェクトで cuTile カーネルをラップします。

from mla_var3.runtime import CtKernel, KernelPlan, Tiling, ConstInt, ConstBool

@ct.kernel
def my_kernel(Tensor, ..., Bm: ConstInt, Bn: ConstInt, EVEN_N: ConstBool):
  bid_x = ct.bid(0)
  # block-level operations: ct.load, ct.mma, ct.store, ct.reshape, ...

@dataclass
class CtMyKernel(KernelPlan):
  b: int = 64; s: int = 1; t: int = 4096
  tiling: MyTiling = field(default_factory=MyTiling)

  def plan(self, *inputs) -> CtKernel:
    Out = torch.empty_like(inputs[0])

    def grid_fn(cfg):
      return (math.ceil(s / cfg.Bm), math.ceil(h / cfg.Bh), b)

    def args_fn(cfg):
      return (inputs[0], Out, cfg.Bm, cfg.Bn, (t % cfg.Bn) == 0)

    return CtKernel(
      input_tensors=inputs,
      output_tensors=(Out,),
      kernel_fn=my_kernel,
      grid_fn=grid_fn,
      args_fn=args_fn,
      tiling=self.tiling,
      autotune_configs=self._autotune_configs(),
      algorithmic_flops_bytes_fn=self._algorithmic_flops_bytes,
    )

CtKernel の主なフィールド：

フィールド	型	目的
kernel_fn	@ct.kernel 関数	cuTile カーネル関数
grid_fn	Callable[[Tiling], tuple]	タイリング設定を 3D グリッドの次元にマッピングします
args_fn	Callable[[Tiling], tuple]	タイリング設定をカーネル引数にマッピングします
input_tensors	tuple[torch.Tensor]	自動チューニングのキャッシュキーに使用します
output_tensors	tuple[torch.Tensor]	__call__ が返す値です
tiling	Tiling	現在のタイリング設定
autotune_configs	list[Tiling]	自動チューニングの探索空間
algorithmic_flops_bytes_fn	Callable	ルーフライン分析に使用します

自動チューニング

CtKernel の自動チューニングでは、cuda.tile_experimental の autotune_launch() を使用します。これにより、占有率／num_ctas のヒントが自動的に処理されます。_apply_hints() メソッドは、チューニング後にヒントを再適用します。

コンパイル

CtKernel.compile() は compile_tile() を使用して .bytecode を出力し、必要に応じて cuda-tile-translate で .mlir に変換します。

Tiling データクラス

cuTile のタイリングには通常、タイルの次元とオプションのコンパイラーヒントが含まれます。

@dataclass
class MyTiling(Tiling):
  Bm: int = 16       # query tile size
  Bn: int = 64       # KV tile size
  Bh: int = 8        # heads per block
  num_ctas: int = None   # CGA size (None = auto)
  occupancy: int = None  # occupancy hint (None = auto)

  def validate(self, pd: "CtMyKernel") -> bool:
    return self.Bm <= pd.s and self.Bn <= pd.t and self.Bh <= pd.h

cuTile の定数型

from mla_var3.runtime import ConstInt, ConstBool, ConstFloat, INV_LOG_2

• ConstInt = ct.Constant[int] — コンパイル時整数
• ConstBool = ct.Constant[bool] — コンパイル時ブール値
• ConstFloat = ct.Constant[float] — コンパイル時浮動小数点数

よくある落とし穴

• cuTile API を決してでっち上げないでください。必ず /cutile-dsl-ref または公式ドキュメントで確認してください
• @ct.kernel の関数名はモジュールのファイル名と一致させる必要があります
• レジスターの上限はスレッドあたり最大 255 です。占有率を高めるには 128 未満を目指してください（docs/devices/ と src/mla_var3/conf/devices.json を参照し、使用中のデバイスの SM 上限に照らして確認してください）
• 共有メモリー：ブロックあたりのデフォルトは 48KB、オプトイン時は 99KB、SM あたり合計 100KB です
• 占有率を高めるためにタイルサイズを縮小する場合は、計算効率が低下しないことを確認してください
• num_ctas によって明示的なクラスター協調アルゴリズムが利用可能になるとは限りません。cuTile ではこれはヒントであり、クラスターをプログラミングするためのインターフェースではありません
• NCU で SM あたり 1 CTA、適格ワープ数の少なさ、そして不足している対策として明示的なワープグループまたはバリアのスケジューリングが示されている場合、cuTile だけでさらにチューニングしても差を埋めることは難しいでしょう

ナレッジリンク

• 共有最適化ナレッジ：docs/knowledge/optimizations/
• 共有アンチパターン：docs/knowledge/anti-patterns/
• cuTile 固有の最適化オーバーレイ：docs/knowledge/languages/cutile-dsl/
• cuTile 固有の最適化カタログ：/optimization-catalog-cutile-dsl を読み込んでください

参考資料

• ランタイム API の詳細：完全な CtKernel API ドキュメントについては、/design-kernel スキルの runtime-api リファレンスを参照してください
• カーネルテンプレート：コピーして使えるテンプレートコードについては、/design-kernel スキルの kernel-templates リファレンスを参照してください
• FlashMLA のスケジューリング上の制限：cuTile FlashMLA と公式の明示的スケジューリング実装を比較した診断については、docs/kernels/flash-mla.md を参照してください