design-cutile-dsl-kernel

cuTile Python DSL — 语言专属指南

前置要求：加载 /design-kernel，了解通用命名、版本管理、KernelPlan 结构、组合模式、克隆工作流、开发日志模板和设计者输出约定。本技能仅介绍 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 内部的显式调度。其公开执行模型不提供线程、warp 或 warpgroup 标识，也不允许在一个块内显式同步或通信。同样，num_ctas 并未提供公开的集群编程模型：没有公开的集群 rank、集群屏障、DSM 原语或集群内存作用域。

当性能分析显示内核被限制为每个 SM 一个 CTA 时，这将成为硬性设计约束。在这种情况下，先进内核通常采用的补救方式是在 CTA 内部显式调度：warpgroup 交替调度、生产者/消费者流水线、屏障编排或集群协作式重叠。如果所需优化属于这一类，请停止在 cuTile 中迭代，改用 CuTe DSL（/design-cute-dsl-kernel）。

FlashMLA 开发日志（docs/kernels/flash-mla.md）是一个具体示例：官方 FlashMLA 设计使用显式调度，因此即使每个 SM 只有一个常驻块，仍能有效运行；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 限制）
• 共享内存：默认每个 block 为 48KB，可选择启用 99KB，每个 SM 总计 100KB
• 为提高占用率而缩小 tile 时，请确认计算效率不会下降
• 不要假设 num_ctas 可以启用显式的集群协作算法；在 cuTile 中，它是一个提示，而非集群编程接口
• 如果 NCU 显示每个 SM 只有一个 CTA、可运行的 warp 数量较少，而问题需要通过显式的 warpgroup 或屏障调度来解决，那么仅靠进一步调整 cuTile 参数不太可能缩小差距

知识链接

• 共享优化知识：docs/knowledge/optimizations/
• 共享反模式：docs/knowledge/anti-patterns/
• cuTile 专属优化补充内容：docs/knowledge/languages/cutile-dsl/
• cuTile 专属优化目录：加载 /optimization-catalog-cutile-dsl

参考资料

• 运行时 API 详情：请参阅 /design-kernel 技能的 runtime-api 参考，其中包含完整的 CtKernel API 文档
• Kernel 模板：请参阅 /design-kernel 技能的 kernel-templates 参考，获取可复制的模板代码
• FlashMLA 调度限制：请参阅 docs/kernels/flash-mla.md，了解 cuTile FlashMLA 与官方显式调度实现的对比诊断