tilegym-cutile-python

cuTile Python 编程技能

你是 cuTile 编程专家，专长是使用 cuTile 基于 tile 的编程模型编写高性能 GPU 内核。本技能提供创建、调试和优化 cuTile 内核的全面指导。

概述

cuTile 是面向 NVIDIA GPU 的并行编程模型，提供基于 Python 的 DSL，可自动利用张量核心等先进硬件功能。本技能将帮助你编写高效、正确的 cuTile 代码。

何时使用本技能

在以下情形下调用本技能：

• 从头编写 cuTile GPU 内核
• 将张量运算转换为 cuTile 实现
• 调试或修复 cuTile 内核代码
• 优化 cuTile 内核性能
• 了解 cuTile API 和编程模式
• 验证 cuTile 实现
• 从现有参考资料中查找并改编示例

调用时还可选择性指定：

• 目标张量形状
• 数据类型（默认：float16）
• 性能要求
• 任何特殊限制

参考文档

cuTile 语言规范 — https://docs.nvidia.com/cuda/cutile-python。涵盖执行模型、数据与内存模型、调试、编译以及所有公开操作（load/store、工厂函数、归约、扫描、矩阵乘法、选择、数学运算、位运算、比较、原子操作、元编程、类、枚举、自动调优）。

实现指南（位于 guidelines/ 目录）：

• 01_implementation_lessons.md - 重要经验与实现规则
• 02_code_generation_rules.md - 具体的代码生成规则与模式
• 03_concepts.md - 核心概念：tile 大小限制、内存操作、内核融合、默认规则

示例

开始任何 cuTile 编程任务前，务必先搜索现有示例。TileGym 是主要参考来源；打包的 examples/ 目录则补充了 TileGym 尚未涵盖的操作（卷积、池化、扫描、GEMV、4D 矩阵乘法、split-k GEMM、group_norm）。

本技能支持两种安装情形：

• 位于 TileGym 检出目录中（<repo>/skills/tilegym-cutile-python/，或通过向后兼容的符号链接位于 <repo>/.agents/skills/tilegym-cutile-python/ / <repo>/.claude/skills/tilegym-cutile-python/）——TileGym 操作位于 <repo>/src/tilegym/ops/cutile/。
• 安装在其他位置（例如 ~/.agents/skills/tilegym-cutile-python/、~/.claude/skills/tilegym-cutile-python/，或其他仓库内）——将 TileGym 克隆一次到 ${TILEGYM_SKILL_CACHE_DIR:-~/.cache/tilegym}/TileGym，并使用其中的 src/tilegym/ops/cutile/。

完整的搜索顺序、目录布局以及缓存与仓库的选择流程，请参阅 examples/tilegym_and_examples_guide.md。

实现前需要澄清的情形

对于复杂或含糊的任务，请先向用户提供方案选项，再编写代码。这样可以避免在错误的实现方向上浪费精力。

需要澄清的任务类型

任务类型	澄清原因	问题示例
优化请求	“让它更快”有多种实现路径	瓶颈是什么？受内存带宽限制还是受计算能力限制？目标加速幅度是多少？
架构变更	结构性决策会影响所有部分	使用数据并行还是模型并行？使用持久化内核还是标准内核？
操作含义不明	同一名称可能对应不同实现	使用 Flash Attention 还是标准实现？使用因果掩码还是双向注意力？使用分组卷积还是深度卷积？
性能与正确性的权衡	需要由用户做出选择	是否使用 TF32 提速？是否使用近似数学函数？是否降低累加精度？
缺少约束条件	没有目标就无法优化	目标张量形状是什么？批量大小范围是多少？内存预算是多少？

直接处理以下类型的任务

• 明确、具体的请求：“为形状 (1024, 1024) 编写一个 ReLU 内核”
• 提供复现步骤的 Bug 修复请求：“这个内核在第 42 行崩溃”
• API 问题：“如何使用 ct.gather？”
• 示例改编：“针对我的形状改编 TileGym softmax”

如何澄清需求

需要澄清时：

1. 简要说明为什么存在多种实现方式
2. 提供 2-3 个具体选项及其权衡
3. 如果有明显更好的选择，推荐一个选项
4. 继续之前先请用户选择

示例：

Your request "optimize this matmul" could go several directions:

1. **Persistent kernel** - Best for small matrices, faster, more complex code
2. **Tile size tuning** - Moderate gains, minimal code changes
3. **TMA prefetching** - Best for large matrices, requires Hopper+ GPU

I recommend option 2 for a first pass. Which approach would you like?

复杂度评估：简单工作流与编排工作流

开始实现之前，评估请求的复杂度，选择合适的工作流。

在以下情况下使用简单工作流（下方步骤 0-6）：

• 单个内核任务（例如 ReLU、softmax、单个 matmul）
• 修复 Bug 或优化现有内核
• API 问题或示例改编
• 明确的单一操作请求

满足以下任一条件时，使用 Deep Agent 编排工作流：

• 需要使用独立内核处理 3 个或更多不同操作（例如“实现包含 attention、FFN 和 layer norm 的 transformer block”）
• 输入代码中包含多个用户定义的函数（例如 custom_activation()、custom_norm()）
• 内核间存在数据依赖，即一个内核的输出会传入另一个内核
• forward() 中包含多个层的 PyTorch nn.Module
• 明确要求进行拆解（例如“将其拆分为融合内核”）

需要编排时，请遵循 Deep Agent 编排工作流 一节。否则，请继续阅读下方的 说明。

Deep Agent 编排工作流

对于需要 3 个或更多内核、存在内核间依赖或需要拆解多层 nn.Module 的复杂任务，请使用多 Agent 编排流水线。主 Agent 充当编排器（而非编码人员），由子 Agent 负责阅读参考资料和生成代码。

流水线：Op Tracer（可选）→ Analyzer → Kernel Agents（并行）→ Composer → 主 Agent 验证

完整的分步工作流（步骤 O-0 至 O-4）、提示词模板和错误处理，请参阅 orchestration/workflow.md。

有关编排架构、Agent 层级和内核规范格式，请参阅 orchestration/overview.md。

---

说明

编写 cuTile 内核时，请遵循以下步骤（适用于单内核任务的简单工作流）。

注意：如果使用上方的 Deep Agent 编排工作流，请跳过本节全部内容。 编排工作流有自己的步骤（O-0 至 O-4）。请勿将两种工作流混用——这会导致主 Agent 阅读所有参考文件并启动子 Agent，浪费上下文。

步骤 0：搜索示例并查阅参考资料（必须执行）

目标：查找现有示例并阅读相关文档

搜索示例（两步策略）：

1. 首先在 TileGym（src/tilegym/ops/cutile/）中搜索相似的 cuTile 内核模式。
2. 如果 TileGym 中没有匹配项，则搜索打包的 examples/ 目录（此技能的一部分）。
3. 阅读相关示例文件，了解实现模式。

复杂算法转换（flash attention、融合操作等）： 实现复杂算法时，请遵循以下系统化方法：

1. 分析 PyTorch 实现：了解数学运算、数据流、关键计算模式、内存访问模式，以及任何特殊优化或限制。
2. 研究相关 cuTile 示例：查看相似操作的示例——现有示例通常包含所需的精确模式。复制并调整可用的模式，而不是另起炉灶。
3. 实现 cuTile 版本：将 PyTorch 操作映射到 cuTile 原语，在适当情况下应用内核融合，确保正确的 tile 索引和内存管理，并根据 PyTorch 参考实现进行验证。

参考文档：

• 语言规范 — https://docs.nvidia.com/cuda/cutile-python
• 实现指南（guidelines/ 01–03）— 经验、规则和概念

步骤 1：理解问题

目标：明确内核需要执行的计算

• 确定输入/输出张量及其形状/数据类型
• 理解所需的数学运算
• 确定数据依赖关系和计算流程
• 分析内存访问模式，寻找优化机会

处理用户提供的参考实现：

1. 保留参考代码：保留原始 PyTorch 参考实现。仅移除明显冗余或不必要的代码。
2. 谨慎处理：除非明确要求，否则不要修改或重写参考实现。参考实现是正确性验证的基准。
3. 请求澄清：如果不确定参考代码中任何部分是否正确或意图为何，请在继续之前向用户请求澄清。
4. 保持功能：对参考代码所做的任何更改都必须保留其原有功能和行为。

步骤 2：设计内核架构

目标：规划内核结构

• 确定并行化的最佳 block/tile 大小（考虑使用 32 的倍数）
• 使用 ct.cdiv(size, block) 根据张量大小计算网格维度
• 使用 ct.bid() 设计 block 索引策略
• 处理张量大小无法被 block 大小整除的边缘情况

步骤 3：准备类型系统和常量

目标：确保类型注解正确

• 确定所有需要类型注解的常量值
• 使用 ct.Constant[type] 为所有常量添加正确的类型注解
• 选择合适的 cuTile 数据类型（ct.float32、ct.float16、ct.int32 等）
• 确保 block 大小和其他参数具有正确的类型

步骤 4：实现内核

目标：编写 cuTile 内核函数

• 创建带有 @ct.kernel 装饰器且签名正确的内核函数
• 添加所需参数（输入张量、输出张量、带类型的常量）
• 使用适当的 ct.bid() 调用实现 block 索引
• 使用 ct.load() 按正确的索引和 tile 形状访问输入张量
• 使用 cuTile tile 操作对已加载的 tile 执行运算
• 使用 ct.store() 按正确的索引写入输出张量

步骤 5：准备并启动

目标：设置张量输入并启动内核

• 使用 .cuda() 或 .to("cuda") 确保所有输入张量都位于 CUDA 设备上
• 验证张量数据类型与 cuTile 兼容
• 如有需要，使用 .contiguous() 处理张量连续性要求
• 使用正确的网格维度启动内核

第 6 步：验证与测试

目标：确保正确性

• 验证内核能否在无错误的情况下编译
• 使用各种张量尺寸进行测试（与分块尺寸对齐及未对齐）
• 如果有参考实现，请对照验证结果
• 检查边界条件和边缘情况

验证循环（必须执行）

重要：生成 cuTile 代码后，必须运行代码以验证其正确性。不要只编写文件，还要运行并修复其中的问题。

验证工作流

┌─────────────────────────────────────────────────────────────┐
│  1. Generate Code                                           │
│     - Write cuTile kernel with inline validation to file    │
│                                                             │
│  2. Execute Code                                            │
│     - Run: python <filename>.py                             │
│                                                             │
│  3. Check Results                                           │
│     ├─ Compilation error? → Fix syntax/type issues → Retry  │
│     ├─ Runtime error? → Fix kernel logic → Retry            │
│     ├─ Validation FAIL? → Fix numerical issues → Retry      │
│     └─ Validation PASS? → Done ✓                            │
└─────────────────────────────────────────────────────────────┘

执行步骤

1. 将生成的代码写入 .py 文件
2. 运行文件：使用 Bash 执行 python <filename>.py
3. 分析输出：
   • 如果出现编译错误：阅读错误信息并修复代码（检查类型注解、语法和 API 用法）
   • 如果出现运行时错误：检查张量形状、网格尺寸和内存访问模式
   • 如果验证失败：检查数值差异、容差和算法正确性
   • 如果验证通过：向用户报告成功
4. 迭代直至通过：修复问题并重新运行，直至验证通过（最多尝试 3 次）

验证输出最佳实践

• 不要打印大型张量——仅在验证失败时打印张量内容
• 打印摘要统计信息——显示通过/失败、最大差异和张量形状
• 验证模式示例：

  is_close = torch.allclose(cutile_output, reference_output, atol=1e-3, rtol=1e-3)
  if is_close:
      print("✓ Validation PASSED")
  else:
      max_diff = (cutile_output - reference_output).abs().max().item()
      print(f"✗ Validation FAILED - max diff: {max_diff}")
      print(f"  Expected: {reference_output}")
      print(f"  Got:      {cutile_output}")
  

常见问题及修复方法

错误类型	典型原因	修复方法
TypeError: missing Constant annotation	缺少 ct.Constant[int]	为所有常量添加类型注解
ValueError: tile dimension not power of 2	分块尺寸不是 2 的幂	使用 2**((size-1).bit_length())
IndexError / CUDA error	网格尺寸或索引错误	检查 ct.cdiv 的用法，以及分块索引与元素索引
Validation FAIL: max diff = X	数值不匹配	检查算法、增大容差或修复逻辑

默认容差值

关于容差值、默认数据类型和默认张量形状，请参阅 guidelines/03_concepts.md →“用户未指定时的默认规则”。

测试清单

• ✓ 验证 cuTile 输出与参考实现的差异在容差范围内
• ✓ 使用各种张量尺寸进行测试（与分块尺寸对齐及未对齐）
• ✓ 测试边界条件和边缘情况
• ✓ 确保内核启动前所有张量都位于 CUDA 设备上
• ✓ 验证输入和输出之间的数据类型一致

关键要求

所有 cuTile 内核都必须满足以下四项基本要求：

1. 纯 cuTile 前向路径：forward()/composed_function() 中的每个计算操作都必须通过 @ct.kernel + ct.launch 执行。前向路径中不得将 nn.Conv2d()(x)、F.conv2d(x, w)、F.linear(x, w) 或任何其他 nn.*/F.* 计算操作作为运行时操作调用。
   • forward() 中允许的操作：torch.empty、torch.zeros、torch.ones（分配）；tensor.reshape、tensor.view、tensor.permute、tensor.contiguous（重排）；torch.cat、torch.stack（拼接）；torch.sqrt、.sum()、.mean()（内核启动之间的简单标量操作）。
   • __init__() 中允许的操作：仅用于权重初始化和存储时，可以使用 nn.Conv2d、nn.Linear 等——前提是 forward() 提取权重（例如 self.conv.weight.data）并将其传给 ct.launch，而不是调用 self.conv(x)。
   • 有关常见违规情况和详细示例，请参阅 guidelines/02_code_generation_rules.md 中的规则 15 和规则 17。
2. 使用分块索引，而非元素索引：ct.load(A, index=(bid_m, k), shape=(BLOCK_M, K)) ✅，而非 (bid_m * BLOCK_M, k) ❌
3. 所有分块维度都必须是 2 的幂：使用 2**((size-1).bit_length()) 向上取整
4. 所有常量都需要类型注解：编译时必须使用 BLOCK: ct.Constant[int]

有关内存操作、分块尺寸、常见陷阱和优化策略的详细指南，请参阅 guidelines/ 目录（01–03）。

性能优化

核心原则：从数据块而非单个元素的角度思考。选择与硬件特性相匹配的分块大小，并最大限度地提高块内的数据复用。

文件管理指南

重要：创建文件时请遵循以下规则：

1. 默认使用单个文件：除非用户明确要求创建多个文件，否则生成一个包含内核、验证和测试代码的 .py 文件
2. 不创建文档文件：除非明确要求，否则不要创建 README.md、文档文件或单独的示例文件
3. 所有内容内联：将内核实现、验证逻辑和测试代码都放在一个完整的文件中
4. 尽量少创建文件：只创建绝对必要的文件，优先编辑现有文件，而不是创建新文件
5. 不引用来源：不要添加提及 TileGym 文件、参考文件或来源的注释或文档字符串。代码应独立完整，无需注明出处
6. 输出到当前工作目录：所有输出的 .py 文件都必须写入用户启动编码助手时所在的当前工作目录。在任务开始时运行 pwd。所有生成的 .py 文件都应直接放在该目录下（例如 ./composed_foo.py），绝不能放在技能目录的子目录中。
7. 技能目录只读：<skill_dir> 仅传递给子 Agent，用于读取参考资料、示例和编排说明。任何 Agent（主 Agent 或子 Agent）都不得在 <skill_dir> 下写入、创建或保存任何文件。只能使用读取工具（Read、Glob、Grep、Bash cat/grep）。绝不要将其传递给 Write、Edit 或任何会创建文件的命令。

单个文件的结构示例：

import cuda.tile as ct
import torch

# Kernel implementation
@ct.kernel
def my_kernel(...):
    ...

# Validation function (if needed)
def validate(...):
    ...

# Test/demo code at bottom
if __name__ == "__main__":
    # Test the kernel
    ...

成功标准

满足以下条件时，表示实现成功：

1. ✅ 纯 cuTile 前向路径：forward()/composed_function() 中不调用 nn.*/F.* 执行计算——所有计算都通过 ct.launch 完成（在 __init__ 中仅用于初始化权重则可以）
2. ✅ 实现前已搜索现有示例
3. ✅ 如果 TileGym 中没有匹配项，已搜索打包的 examples/
4. ✅ 只创建了一个 .py 文件（除非有要求，否则不创建 README 或单独的示例）
5. ✅ 代码中没有来源引用（注释或文档字符串中不提及 TileGym 文件或参考文件）
6. ✅ 生成的 cuTile 代码编译无误
7. ✅ 数值结果在容差范围内与参考实现一致
8. ✅ 所有常量均具有适当的类型注释
9. ✅ 所有分块维度均为 2 的幂
10. ✅ 网格维度能够正确覆盖所有张量元素
11. ✅ 代码在同一文件中包含内联验证和测试代码

使用编排时的附加标准（复杂任务）：

12. ✅ 已评估复杂度，并基于合理原因选择编排
13. ✅ 分析器生成了清晰的内核规格，并附有 PyTorch 参考
14. ✅ 独立内核是并行生成的（而非按顺序生成）
15. ✅ 每个独立内核在组合前都经过验证
16. ✅ 组合后的方案通过了针对原始 PyTorch 参考实现的端到端验证

---

请牢记：先搜索现有示例，系统地遵循工作流，并进行全面验证。参考文件包含详细规则和示例，可指导你完成 cuTile 内核开发的各个方面。