learn-cute-dsl

从 CUTLASS 示例学习 CuTe DSL

通过阅读 CUTLASS Blackwell 示例，逐步掌握 CuTe Python DSL；可选择将示例导入项目的运行时基础设施并对其进行性能分析，然后提炼可复用模式并整理到知识库中。

示例来源

大量示例位于 third_party/cutlass/examples/python/CuTeDSL。文档位于 docs/cute-dsl/（CuTe Python DSL）和 docs/cutlass-cpp/cute/（CuTe C++ 概念）。

类别	示例	子目录
MLA 解码	mla_decode_fp16.py、mla_decode_fp8.py、mla_helpers.py	mla/
Attention	fmha.py、fmha_bwd.py	（顶层）
混合 Attention	mixed_input_fmha_decode.py、mixed_input_fmha_prefill_*.py	mixed_input_fmha/
GEMM（基础）	dense_gemm.py	（顶层）
GEMM（流水线）	dense_gemm_software_pipeline.py	（顶层）
GEMM（持久化）	dense_gemm_persistent.py、dense_gemm_persistent_prefetch.py	（顶层）
GEMM（动态持久化）	dense_gemm_persistent_dynamic.py	（顶层）
GEMM 教程	fp16_gemm_0.py 至 fp16_gemm_6.py	tutorial_gemm/
块缩放 GEMM	dense_blockscaled_gemm_persistent.py、*_prefetch.py、*_amax.py	（顶层）
混合输入 GEMM	mixed_input_gemm.py、grouped_mixed_input_gemm.py	mixed_input_gemm/
Epilogue	custom_epilogue_dense_gemm.py、activation_custom_epilogue_dense_gemm.py	epilogue/
归约	reduce.py、rmsnorm.py	（顶层）
依赖启动	programmatic_dependent_launch.py	（顶层）
Mamba2	mamba2_ssd.py	mamba2_ssd/
分块 GEMM	blockwise_gemm.py、contiguous_grouped_gemm.py	blockwise_gemm/

用户也可以直接提供其他示例。

工作流

1. SELECT    -> Pick examples based on learning goals
2. READ      -> Read and annotate, identify CuTe DSL patterns
3. IMPORT    -> (Optional) Wrap in CuteKernel + KernelPlan for profiling
4. PROFILE   -> Profile imported kernels, compare with existing cuTile versions
5. EXTRACT   -> Distill reusable CuTe DSL patterns and anti-patterns
6. DOCUMENT  -> Update knowledge base (docs/knowledge/languages/cute-dsl/)

---

步骤 1：选择

根据学习目标挑选示例。

学习目标	从以下内容开始
CuTe DSL 基础知识（grid、TMA、MMA）	tutorial_gemm/fp16_gemm_0.py 至 fp16_gemm_6.py
软件流水线	dense_gemm_software_pipeline.py
持久化内核结构	dense_gemm_persistent.py
Warpgroup 专项化	dense_gemm_persistent_prefetch.py
Attention 的 TMA 流水线	fmha.py
MLA 解码（直接参考）	mla/mla_decode_fp16.py + mla/mla_helpers.py
使用 FP8 的 MLA 解码	mla/mla_decode_fp8.py
Cluster 编程	dense_gemm_persistent_dynamic.py
自定义 epilogue	epilogue/custom_epilogue_dense_gemm.py
依赖内核启动	programmatic_dependent_launch.py

阅读之前，先写下具体问题。例如：

• “MLA 解码内核如何为 KV 缓存设置 TMA 描述符？”
• “在持久化 GEMM 中，warpgroups 如何分配生产者／消费者角色？”
• “FMHA 在 K 循环流水线中使用了什么屏障模式？”

---

步骤 2：阅读

对于每个选定的示例，阅读完整源代码，并标注以下 CuTe DSL 特有的方面：

主机端（@cute.jit）

1. TMA 描述符创建——哪些张量会获取 TMA 描述符，tile 形状是什么？
2. Grid 和 cluster 维度——如何计算 grid？是否使用了 clustering？
3. 共享内存大小——如何根据流水线阶段和 tile 形状计算 SMEM 预算？
4. 启动参数——max_register_count、cluster 形状、SMEM carveout

设备端（@cute.kernel）

1. 线程／warp／warpgroup 标识——如何分配角色？
2. 流水线结构——阶段数量、屏障初始化、生产者／消费者循环
3. TMA 操作——cute.copy(tma_desc, ...)、到达模式
4. MMA 操作——累加器设置、MMA 循环主体、epilogue
5. 同步——cute.barrier_arrive()、cute.barrier_wait()、cute.cluster_barrier_*
6. 共享内存布局——分配、swizzling、在 warpgroups 间进行分区

模式检查清单

阅读时，标记以下模式的实例：

• TMA 多阶段流水线（有多少个阶段，每个阶段是否有屏障）
• Warpgroup 专门化（哪些 warpgroup 负责生产，哪些负责消费）
• 持久化内核主循环（工作分块调度、退出条件）
• 在线 softmax 或流式累积
• 寄存器级累加器管理
• 共享内存 swizzle 模式
• 集群级通信（分布式共享内存）
• Epilogue 模式（寄存器内、经由 SMEM 或融合）
• 边界处理（针对不完整分块的谓词判断、掩码）
• 流水线排空（如何刷新最后几个流水线阶段）

交叉参考

检查知识库中是否已有相关模式：

• docs/knowledge/optimizations/——共享模式（例如，pipeline-driven-low-occupancy）
• docs/knowledge/languages/cute-dsl/——现有的 CuTe DSL 覆盖文档

记录示例是确认、完善还是与现有知识相矛盾。

---

第 3 步：导入（可选）

将 CUTLASS 示例导入项目，使其可通过标准 CLI 和性能分析基础设施运行。

适合导入的情况

• 你想在自己的硬件上分析该示例的性能
• 你想以该示例为起点，编写自定义 CuTe DSL 内核
• 该示例涉及的工作负载与现有 cuTile 内核相近（例如，MLA 解码）

不适合导入的情况

• 你只需要理解该模式，不需要进行基准测试
• 该示例依赖复杂，或使用了难以适配的辅助基础设施

导入流程

1. 创建内核包：

   src/mla_var3/kernel/cute_python/<layer>/<design>/<design>/
   

   CUTLASS MLA 解码 FP16 示例：

   src/mla_var3/kernel/cute_python/mla/cutlass_mla_decode/cutlass_mla_decode/
   

2. 创建 KernelPlan 子类（cutlass_mla_decode.py）：

   • prepare_inputs()——分配与示例预期形状匹配的张量（Q、KV 缓存、输出）
   • reference_fn()——用于正确性验证的 PyTorch 参考实现
   • _autotune_configs()——从示例中提取的分块参数
   • _algorithmic_flops_bytes()——屋顶线分析
   • plan()——返回一个封装示例主机端/设备端函数的 CuteKernel

3. 创建 __main__.py：

   from .cutlass_mla_decode import CutlassMlaDecodeKernel
   
   if __name__ == "__main__":
       kernel_plan = CutlassMlaDecodeKernel()
       kernel_plan.benchmark_kernel_argparse()
   

4. 在每个包层级创建 __init__.py

5. 处理辅助模块：如果示例使用了辅助模块（例如，mla_helpers.py），可以：

   • 将其复制到内核包中并调整导入语句
   • 从 CUTLASS 路径导入（不够稳妥，但便于快速探索）

6. 测试正确性：

   source .venv/bin/activate && python -m mla_var3.kernel cutlass_mla_decode --prof_type=disabled --check
   

---

第 4 步：性能分析

架构兼容性检查

开始性能分析前，确认可用设备与内核的目标架构匹配：

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

许多 CUTLASS 示例（尤其是 mla/ 中的示例以及任何使用 tcgen05 MMA 操作的示例）专为 SM100（B200、B300 数据中心 GPU）设计。请注意，“Blackwell”是一个营销名称，涵盖多个 SM 版本，而这些版本的指令集各不相同：

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

即使 RTX 5090（SM120）和 SM100 GPU 的产品都以“Blackwell”命名，使用 tcgen05 操作的内核也无法在 RTX 5090 上运行。请根据 GPU 的实际 SM 版本，而非营销代际名称来判断是否匹配。

如果内核的目标 SM 版本与可用设备不匹配，请完全跳过“性能分析”步骤——执行要么会在编译时失败，要么会产生误导性结果。此时，请直接进入第 5 步（提取），仅根据源代码阅读结果总结见解，并在任何文档中注明架构不匹配的情况。

性能分析流程

1. 锁定 GPU 时钟：bash scripts/lock-gpu-clock.sh

2. 先使用注释模式：

   source .venv/bin/activate && python -m mla_var3.kernel <kernel> \
       --b=32 --s=1 --t=4096 --prof_type=annotation
   

3. 如果注释结果揭示了有意思的模式，则使用 NCU 深入分析：

   source .venv/bin/activate && python -m mla_var3.kernel <kernel> \
       --b=32 --s=1 --t=4096 --prof_type=ncu
   

重点关注领域

既然目标是学习（而非单纯进行基准测试），就应重点了解内核如何实现高性能：

• 流水线效率：TMA 复制是否与计算充分重叠？检查 nsys 时间线中是否存在空闲间隙。
• Warpgroup 利用率：所有 warpgroup 是否都在忙碌？检查 NCU 中的调度器统计信息。
• 寄存器压力与占用率的权衡：每个线程使用多少寄存器？低占用率是否为有意选择？
• SMEM 使用情况：使用了多少可用的 SMEM 预算？能容纳多少个流水线阶段？
• MMA 吞吐量：达到了峰值 TC 利用率的多少？

与 cuTile 版本比较

如果导入的内核涵盖了 cuTile 中已有的工作负载（例如 MLA 解码），请进行比较：

指标	cuTile 版本	CUTLASS CuTe DSL	差异	见解
耗时（us）
TC%
DRAM%
占用率
每线程寄存器数
每个线程块的 SMEM（KB）

这些差异揭示了显式控制（warpgroup 调度、TMA 流水线）相较于 cuTile 的线程块级抽象所带来的收益。

---

步骤 5：提取

将观察结果提炼为可复用的 CuTe DSL 模式。

分类

类型	存放位置
通用算法／硬件模式	docs/knowledge/optimizations/<name>.md
通用故障模式	docs/knowledge/anti-patterns/<name>.md
CuTe DSL API 模式或实现细节	docs/knowledge/languages/cute-dsl/optimizations/<name>.md
CuTe DSL 陷阱或易错点	docs/knowledge/languages/cute-dsl/anti-patterns/<name>.md

CUTLASS 示例中的大多数模式都将是CuTe DSL 特有的（例如 API 模式、TMA 设置惯用写法、屏障编排）。通用模式（例如流水线驱动调度等算法见解）可能已从 cuTile 工作中总结出来。

提取标准

符合以下条件时，记录该模式：

• 可迁移到其他 CuTe DSL 内核（而不只是当前示例）
• 仅凭 CuTe DSL 文档难以显而易见
• 具有性能影响，且可通过性能分析观察到，或从架构上清晰可见
• 细化现有通用模式，并补充 CuTe DSL 特有的实现细节

模式模板

# [Pattern Name]

## When to Apply
- [CuTe DSL kernel type, workload shape, bottleneck condition]

## Mechanism
[How the pattern works in CuTe DSL -- reference specific APIs]

## Affected Metrics
- [Metric 1]
- [Metric 2]

## Implementation
```python
# CuTe DSL code snippet, generalized from the CUTLASS example

来源

[CUTLASS 示例名称和路径]

性能证据

来源：[示例名称、设备、配置]

配置	指标	数值	上下文

泛化

[与设备无关的要点]

易错点

• [CuTe DSL 特有的故障模式]

交互

• [该模式如何与其他 CuTe DSL 模式交互]

---

## Step 6: DOCUMENT

1. Write pattern files to `docs/knowledge/languages/cute-dsl/optimizations/` or `anti-patterns/`
2. Update the index in `/optimization-catalog-cute-dsl`
3. If a pattern is shared (not DSL-specific), write to `docs/knowledge/optimizations/` and update `/optimization-catalog`
4. If you imported and profiled a kernel, write a devlog entry in `docs/kernels/`:

```markdown
# [Kernel Name] (Imported from CUTLASS)

## Overview
[What this kernel does and why it was imported for study]

## Source
Original: `third_party/cutlass/examples/python/CuTeDSL/<path>`
Imported: `src/mla_var3/kernel/cute_python/mla/<design>/`

## Performance
[Profiling results table]

## Patterns Extracted
- [Pattern] -> `docs/knowledge/languages/cute-dsl/optimizations/<name>.md`

## Insights
[What was learned that informs future CuTe DSL kernel design]

---

推荐学习路径

路径 A：CuTe DSL 基础

在编写任何 MLA 内核之前，先建立基本熟练度：

1. tutorial_gemm/fp16_gemm_0.py  -> Minimal kernel: grid, TMA load, MMA, store
2. tutorial_gemm/fp16_gemm_3.py  -> Software pipelining basics
3. tutorial_gemm/fp16_gemm_6.py  -> Full persistent GEMM with warpgroup specialization
4. dense_gemm_software_pipeline.py -> Production-quality pipelining
5. dense_gemm_persistent.py      -> Persistent scheduling pattern

路径 B：MLA 专属模式

直接学习 CUTLASS 如何实现 MLA：

1. mla/mla_helpers.py            -> Shared utilities, tensor layouts, TMA setup
2. mla/mla_decode_fp16.py        -> FP16 MLA decode with full pipeline
3. mla/mla_decode_fp8.py         -> FP8 variant (quantization handling)
4. fmha.py                       -> Attention forward (different design choices)

路径 C：高级模式

学习特定优化技术：

1. dense_gemm_persistent_prefetch.py  -> Prefetching in persistent kernels
2. dense_gemm_persistent_dynamic.py   -> Dynamic tile scheduling + clusters
3. epilogue/custom_epilogue_dense_gemm.py -> Fused epilogue patterns
4. programmatic_dependent_launch.py   -> Dependent kernel launch