SuperAPI AI 工具集

design-cute-dsl-kernel

编程与开发

查看项目

更新于 10-16 09:50

技能介绍

用途:指导在 cuTile 控制面不足时改用 CuTe Python DSL(cute-dsl)编写 GPU 内核,并配套共享的命名、版本与工作流规范(需同时加载 /design-kernel),API 参考与执行模型见 /cute-dsl-ref。

适用时机:当优化需要 cuTile 未暴露的能力时采用,包括线程/warp/warpgroup 身份控制、CTA 内同步(barrier)、warpgroup 调度与生产者/消费者角色、TMA 多级异步流水线、cluster 编程与分布式共享内存、寄存器级数据搬运、自定义 epilogue。若优化仍可通过 tile 尺寸、CTA 重映射、占用与 num_ctas 提示、延迟提示或 allow_tma 表达,则应留在 cuTile。典型切换触发条件为:1 CTA/SM、可调度 warp 少,且修复需要显式 warpgroup 或 barrier 调度。

命名与布局:语言键为 cute-dsl,Python 包路径为 cute_python;内核位于 src/mla_var3/kernel/cute_python/<layer>/<design>/<design>[_vN]/,模块为 <design>[_vN].py,封装类为 CuteKernel,运行时封装位于 src/mla_var3/runtime/cute_kernel.py。

使用方式:采用主机/设备两级模式,@cute.jit 主机函数负责建立 TMA 描述符、计算 grid 并启动,@cute.kernel 设备函数承载 GPU 代码;通过 KernelPlan.plan() 返回 CuteKernel,传入 kernel_fn、launch_fn 闭包、输入输出张量、tiling、自动调优配置与 flops/bytes 函数。自动调优逐配置用 do_bench_cudagraph 基准测试并取最快者,失败配置跳过;CuTe DSL JIT 自动缓存编译产物。tiling 常含 num_warpgroups、num_pipeline_stages、num_tma_buffers、cluster_m/cluster_n 等字段,validate 需校验维度整除与寄存器、SMEM 资源限制。

限制与陷阱:不得臆造 CuTe DSL API,需对照参考文档;@cute.kernel 函数名必须与模块文件名一致;TMA 描述符仅在主机侧设置;每线程寄存器上限 255;SM100 共享内存上限 228 KB/SM、SM120 为 96 KB,不可默认 Blackwell 等同数据中心限制;流水级数增加会提高 SMEM 占用;barrier 同步错误会静默产生错误结果,需用 --check 测试;cluster 编程要求 grid 维度兼容。运行、编译或剖析前先用 nvidia-smi --query-gpu=name --format=csv,noheader 核对设备 SM 版本,不匹配则跳过,并在开发日志记录原因。