External Cannbot Ops Lab Tilelang Skills Tilelang Op Developer

TileLang-Ascend 算子代码生成

基于设计文档（design.md）和已有示例，生成可运行的算子实现与测试。

1. 从 design.md 中提取的信息（只取这些）

design.md 可能很长，只提取以下字段，忽略其余内容：

明确忽略的内容（这些容易误导）：

• 模式选型的分析推理过程
• 内存预算的计算过程和多轮优化迭代
• 风险点与注意事项（过于笼统）
• 交付清单（仅是文件列表）
• 任何标注为"待确认"的内容

2. 参考来源（优先级高于 design.md 伪代码）

当 design.md 伪代码与 examples/ 中同类实现有冲突时，以 examples/ 为准。

2.1 API 用法和模式选择

• API 用法：查阅 tilelang-api-best-practices SKILL.md 及其 references 目录
• 编程模式和 pass_configs：查阅 tilelang-programming-model-guide SKILL.md 及其 references 目录

2.2 同类算子示例

生成代码前，必须查阅 examples/ 中的同类算子：

查阅示例时关注：

1. Kernel 结构：T.Kernel 参数、cid/vid 用法
2. Buffer 分配方式：shape 和 dtype
3. pass_configs 配置：该类算子实际使用哪些开关
4. 数据搬运：T.copy 的索引写法

3. 代码生成流程

步骤 1：读取设计文档

读取 design.md，按 §1 的表格提取字段。

步骤 2：查找参考示例

在 examples/ 中找到最相似的算子实现，完整阅读其代码。

步骤 3：生成实现代码

基于 design.md 的 API 映射 + 参考示例的代码风格，生成 example_{op}.py。

文件结构：

import tilelang
from tilelang import DataType, language as T
import torch

tilelang.cache.clear_cache()

# ========== 算子实现 ==========
@tilelang.jit(out_idx=[...], pass_configs={...})
def op_name(M, N, block_M, block_N, dtype="float"):
    # 分块计算
    m_num = T.ceildiv(M, block_M)
    n_num = T.ceildiv(N, block_N)
    VEC_NUM = 2

    @T.prim_func
    def main(Input: T.Tensor((M, N), dtype), Output: T.Tensor((M, N), dtype)):
        with T.Kernel(..., is_npu=True) as (cid, vid):
            # buffer 分配
            # 数据搬入
            # 计算
            # 数据搬出
            pass

    return main

# ========== 测试 ==========
if __name__ == "__main__":
    torch.manual_seed(0)
    test_configs = [...]  # 来自 design.md §8

    for config in test_configs:
        # 1. 创建 kernel
        # 2. 生成输入数据
        # 3. 执行 kernel
        # 4. golden 对比
        # 5. 精度检查
        pass

    print("All tests passed!")

步骤 4：运行验证

python examples/{op}/example_{op}.py

如果报错，按以下顺序排查：

1. 编译错误 → 检查 buffer 大小、API 参数、对齐
2. 运行错误 → 检查索引越界、同步缺失
3. 精度错误 → 检查计算公式、数据类型、容差设置

4. 关键编码规范

Buffer 分配

# VEC_NUM = 2，每个 vector 核处理 block_M // VEC_NUM 行
a_ub = T.alloc_ub([block_M // VEC_NUM, block_N], dtype)

数据搬运索引

# 标准索引模式
row_start = bx * block_M + vid * block_M // VEC_NUM
T.copy(A[row_start, by * block_N], a_ub)
T.copy(a_ub, B[row_start, by * block_N])

同步

# Expert 模式：手��同步
with T.Scope("V"):
    T.copy(A[...], a_ub)
    T.barrier_all()
    T.tile.exp(a_ub, a_ub)
    T.barrier_all()
    T.copy(a_ub, B[...])

# Developer 模式 + 自动同步：无需手动 barrier
pass_configs = {
    tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True,
    tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True,
}

广播

# 归约结果 [M, 1] 广播到 [M, N]
max_ub = T.alloc_ub([block_M // VEC_NUM, 1], dtype)
max_2d_ub = T.alloc_ub([block_M // VEC_NUM, block_N], dtype)
T.tile.broadcast(max_2d_ub, max_ub)

测试模板

# golden 对比
ref_output = torch.nn.functional.softmax(input_data, dim=-1)  # 或手写 golden
torch.testing.assert_close(output.cpu(), ref_output.cpu(), rtol=rtol, atol=atol)

5. Checklist

生成代码后逐项检查：