SuperAPI AI 工具集

learn-cute-dsl

编程与开发

查看项目

更新于 10-16 07:57

技能介绍

用途:建立 CuTe Python DSL 的实践能力,围绕 CUTLASS Blackwell 示例内核进行阅读、导入、剖析与模式提炼。示例位于 third_party/cutlass/examples/python/CuTeDSL,文档在 docs/cute-dsl/ 与 docs/cutlass-cpp/cute/。

主要能力:提供六步工作流——SELECT(按学习目标选例,如 tutorial_gemm/fp16_gemm_0.py 至 fp16_gemm_6.py、dense_gemm_software_pipeline.py、mla/mla_decode_fp16.py、fmha.py 等)、READ(标注主机侧 @cute.jit 的 TMA 描述符、网格与簇维度、共享内存预算、启动参数,以及设备侧 @cute.kernel 的线程角色、流水线、TMA/MMA 操作与同步)、IMPORT(可选,将示例包装为 KernelPlan 与 CuteKernel 以便剖析)、PROFILE(先锁频,再用 annotation 或 ncu 模式)、EXTRACT(提炼模式与反模式)、DOCUMENT(写入知识库并更新索引)。

使用方式:导入时创建内核包与 KernelPlan 子类,实现 prepare_inputs、reference_fn、_autotune_configs、_algorithmic_flops_bytes、plan,并添加 main.py;正确性检查命令为 source .venv/bin/activate && python -m mla_var3.kernel cutlass_mla_decode --prof_type=disabled --check;剖析前用 nvidia-smi --query-gpu=name --format=csv,noheader 核对架构。

适用场景:学习 CuTe DSL 基础、软件流水线、持久化内核、warpgroup 特化、TMA 流水线、MLA 解码、自定义 epilogue 等。

限制:许多示例针对 SM100(B200/B300),使用 tcgen05 的示例无法在 SM120(RTX 5090/5080)运行;架构不匹配时应跳过剖析,仅基于源码提炼并记录不匹配。部分示例依赖复杂辅助模块,导入可能不顺畅。其余细节暂未提供。