mcpField
D
design-cute-dsl-kernel

design-cute-dsl-kernel

github.com/pepperu96开发084 次安装32 周前更新
代码仓库

关于 design-cute-dsl-kernel

面向 GPU 内核开发者的 CuTe Python DSL 设计技能，用于决定何时从 cuTile 转向 CuTe DSL，并指导内核结构、自动调优与架构兼容性检查。

主要功能

  1. 01评估转用 CuTe DSL 的条件
  2. 02指导主机与设备端内核结构
  3. 03提供 CuteKernel 自动调优模式
  4. 04验证 GPU 架构与资源限制
  5. 05提示同步和内存使用陷阱

适用场景

  1. 01为需要 warpgroup 控制的优化选型
  2. 02设计显式 TMA 流水线内核
  3. 03检查 Blackwell GPU 指令集兼容性

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/pepperu96/hyper-mla --skill design-cute-dsl-kernel

安装后在 Claude Code 中输入/design-cute-dsl-kernel确认已加载。

详细介绍

CuTe Python DSL 内核设计

还要始终加载 /design-kernel，以获取通用的命名、版本管理和工作流规范。 还要加载 /cute-dsl-ref，以获取 API 参考、执行模型和架构操作相关内容。

何时使用 CuTe Python DSL

当 cuTile 的公共控制能力已不再满足需求，但仍适合采用 Python 编写内核的工作流时，请使用 CuTe Python DSL（cute-dsl）。

适用性判断

当下一步优化需要以下任何一种 cuTile 未提供的控制能力时，CuTe DSL 就是合适的选择：

  • 线程/warp/warpgroup 标识——显式控制各线程的工作内容
  • CTA 内同步——屏障、命名屏障、到达/等待模式
  • Warpgroup 调度——生产者/消费者 warpgroup 角色、持久化 warpgroup 循环
  • TMA 流水线控制——带屏障同步的显式多阶段异步复制流水线
  • 集群编程——跨 CTA 共享内存访问、分布式共享内存
  • 寄存器级数据移动——显式的寄存器间 shuffle、warp 级原语
  • 自定义 epilogue——具备细粒度控制的融合后处理
何时继续使用 cuTile

如果仍可通过 tile 大小、CTA 重映射、occupancy/num_ctas 提示、延迟提示或 allow_tma 标志来实现优化，请继续使用 cuTile。CuTe DSL 会增加复杂性，因此只有在这份复杂性不可或缺时才使用它。

硬性设计约束

当性能分析显示每个 SM 仅有 1 个 CTA、符合条件的 warp 数较少，且解决方案需要显式的 warpgroup 或屏障调度时，cuTile 无法弥补差距。这是切换到 CuTe DSL 的典型触发条件。

命名与布局

  • 公共语言键：cute-dsl
  • Python 包路径：cute_python
  • 内核布局：src/mla_var3/kernel/cute_python/<layer>/<design>/<design>[_vN]/
  • 模块：<design>[_vN].py
  • 包装器：CuteKernel

CuteKernel 运行时模式

运行时包装器位于 src/mla_var3/runtime/cute_kernel.py。

CuTe DSL 内核采用两级主机/设备模式：

  1. @cute.jit 主机函数负责设置 TMA 描述符、计算网格并启动内核
  2. @cute.kernel 设备函数包含 GPU 代码

CuteKernel 数据类封装了此模式：

from mla_var3.runtime.cute_kernel import CuteKernel

# In KernelPlan.plan():
def plan(self, *inputs):
    # Build a closure that captures inputs and calls the host function
    def launch_fn(tiling):
        # Convert tensors, set up TMA descriptors, compute grid
        host_fn(tiling, *converted_inputs)

    return CuteKernel(
        kernel_fn=device_kernel,       # The @cute.kernel function (for naming)
        launch_fn=launch_fn,           # Closure: (tiling) -> launches kernel
        input_tensors=list(inputs),
        output_tensors=[output],
        tiling=self.tiling,
        autotune_configs=self._autotune_configs(),
        algorithmic_flops_bytes_fn=self._algorithmic_flops_bytes,
    )
与 CtKernel 的主要区别
AspectCtKernel (cuTile)CuteKernel (CuTe DSL)
启动机制grid_fn + args_fn + cuTile 编译器封装 @cute.jit 主机函数的 launch_fn 闭包
编译cuTile 字节码 + MLIRCuTe DSL JIT（自动缓存）
自动调优来自 cuda.tile_experimental 的 autotune_launch()针对每种配置调用 triton.testing.do_bench_cudagraph
网格设置grid_fn(cfg) -> (x, y, z)在 launch_fn / 主机函数内部
TMA 描述符隐式（由 cuTile 处理）在主机函数中显式设置
自动调优

CuteKernel 会遍历 autotune_configs，使用 do_bench_cudagraph 对每种配置进行基准测试，并选择最快的配置。失败的配置会被捕获并跳过。

编译

CuTe DSL JIT 会自动缓存产物。compile() 方法会创建输出目录，但不会提取显式产物（可按内核扩展）。

分块数据类指南

CuTe DSL 分块通常包含 cuTile 分块没有的字段：

@dataclass
class Tiling:
    # Standard tile dimensions
    tile_m: int
    tile_n: int
    tile_k: int

    # CuTe DSL-specific fields
    num_warpgroups: int          # Warpgroups per CTA (typically 2-4)
    num_pipeline_stages: int     # Async copy pipeline depth
    num_tma_buffers: int         # TMA double/triple buffering
    cluster_m: int = 1           # Cluster shape (M dimension)
    cluster_n: int = 1           # Cluster shape (N dimension)

    def validate(self, pd) -> bool:
        # Validate against problem dimensions and device limits
        ...

具体字段取决于内核设计。validate 方法应检查分块维度能否整除问题维度，并确保资源用量（寄存器、SMEM）不超过设备限制。

架构兼容性检查

内核可能会针对需要特定 SM 版本的架构专用指令。营销名称“Blackwell”涵盖多个 SM 版本，且它们的指令集不同：

  • SM100（B200、B300）——数据中心 GPU，支持 tcgen05 MMA、TMEM 和完整的集群功能
  • SM120（GeForce RTX 5090、RTX 5080）——消费级 Blackwell，使用 sm_120a，不支持 tcgen05 操作

使用 tcgen05 操作的内核需要 SM100，即使 RTX 5090（SM120）和它们都属于“Blackwell”，该内核也无法在 RTX 5090 上运行。

运行、编译或分析 CuTe DSL 内核之前，务必确认可用设备：

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

然后根据 GPU 名称确认其 SM 版本。如果内核的目标 SM 版本不匹配，请跳过运行/分析步骤。不要尝试执行——这要么会在编译时失败，要么会产生误导性结果。在内核开发日志中记录架构要求，并注明已跳过。

常见陷阱

  • 切勿臆造 CuTe DSL API——请查阅 /cute-dsl-ref 技能、docs/cute-dsl/ 文档或 CUTLASS 示例内核进行核实
  • @cute.kernel 函数名必须与模块文件名一致（与 cuTile 规则相同）
  • TMA 描述符只能在主机端设置——没有适当的 @cute.jit 主机端设置时，不要尝试在 @cute.kernel 内执行 TMA 操作
  • 寄存器预算：每个线程最多 255 个；请通过 docs/devices/ 和 src/mla_var3/conf/devices.json，根据活动设备的 SM 限制进行验证
  • 共享内存因 SM 版本而异：SM100（B200/B300）每个 SM 最多 228 KB，每个区块可选择启用 227 KB；SM120（RTX 5090）为 96 KB——不要假设“Blackwell”就意味着数据中心级 SMEM 限制
  • 流水线阶段数会影响 SMEM 用量（每个阶段都需要自己的缓冲区）——增加阶段数之前，请验证 SMEM 总量
  • 屏障同步错误不会显式报错，会导致结果错误而非崩溃——务必使用 --check 进行测试
  • 集群编程要求启动时使用兼容集群的网格维度

参考资源

  • API 参考：加载 /cute-dsl-ref，查看核心 API 表、执行模型和架构操作
  • 官方文档：docs/cute-dsl/（CuTe Python DSL）和 docs/cutlass-cpp/cute/（CuTe C++ 概念）
  • 示例内核和学习路径：加载 /learn-cute-dsl，查看分类整理的 CUTLASS 示例内核索引

知识链接

  • 共享优化知识：docs/knowledge/optimizations/
  • 共享反模式：docs/knowledge/anti-patterns/
  • CuTe Python DSL 覆盖文档：docs/knowledge/languages/cute-dsl/

开发日志条目

使用 docs/kernels/<kernel>.md，并以 Python 包路径格式记录实现位置：

src/mla_var3/kernel/cute_python/mla/<kernel>/<kernel>[_vN]/

评论

支持 Markdown