mcpField
C
cute-dsl-ref

cute-dsl-ref

github.com/pepperu96开发082 次安装32 周前更新
代码仓库

关于 cute-dsl-ref

cute-dsl-ref 是面向使用 NVIDIA CuTe Python DSL 编写 GPU 内核的开发者的参考技能，涵盖主机与设备执行模型、API、常见模式及架构约束，包括 SM100 与 SM120 的差异。

主要功能

  1. 01说明主机与设备内核的职责
  2. 02汇总布局、张量和 MMA API
  3. 03介绍 TMA、共享内存与同步
  4. 04说明 JIT 编译与缓存
  5. 05列出架构约束和常见模式

适用场景

  1. 01构建 CuTe DSL GPU 内核
  2. 02选择布局和 MMA 操作
  3. 03实现 TMA 流水线与屏障
  4. 04运行前核对 GPU 架构约束

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/pepperu96/hyper-mla --skill cute-dsl-ref

安装后在 Claude Code 中输入/cute-dsl-ref确认已加载。

附带文件references/architecture-ops.mdreferences/docs-index.md

详细介绍

CuTe Python DSL 参考

执行模型

CuTe Python DSL 是 NVIDIA CuTe 布局代数的 Python 接口。与 cuTile 的块级抽象不同，CuTe DSL 可显式控制线程、warp 和 warpgroup，并提供 TMA 流水线、屏障编排和共享内存管理。

两级主机/设备模式

每个 CuTe DSL 内核都包含两个函数：

  1. @cute.jit 主机函数 — 在 CPU 上运行，设置 TMA 描述符、计算网格、分配共享内存并启动内核
  2. @cute.kernel 设备函数 — 在 GPU 上运行，包含实际计算
import cutlass
import cutlass.cute as cute

@cute.kernel
def my_kernel(tiled_mma: cute.TiledMma, ...):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()
    # ... GPU code

@cute.jit
def host_fn(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
    # Setup TMA descriptors, compute grid, allocate SMEM
    my_kernel(...).launch(grid=grid_shape, block=block_shape, smem=smem_bytes)
编译流水线

三阶段 JIT 编译：

  1. 预暂存阶段：Python AST 重写
  2. 元阶段：Python 解释器执行元程序（解析编译时常量）
  3. 对象阶段：编译器后端生成 PTX → SASS

JIT 产物会自动缓存。使用 CUTE_DSL_CACHE_DIR 和 CUTE_DSL_DISABLE_FILE_CACHING 控制缓存。

与 cuTile 的主要区别
方面cuTileCuTe DSL
抽象级别块级（无线程标识）线程、warp、warpgroup 级
TMA隐式（allow_tma=True）在主机端显式设置描述符
共享内存编译器管理显式分配和布局
同步不提供屏障、到达/等待、命名屏障
流水线不提供显式多阶段异步流水线
MMA对 tile 使用 ct.mma(A, B, C)对分区片段使用 cute.gemm(tiled_mma, d, a, b, c)
编译cuTile 字节码 → MLIRCuTe DSL JIT → PTX → SASS

核心 API 表

装饰器和编译
API用途关键参数
@cute.kernel设备内核装饰器函数名必须与模块文件名一致
@cute.jit主机 JIT 函数装饰器—
kernel(...).launch(grid, block, smem, cluster)启动内核grid=(x,y,z)、block=(threads,1,1)、smem=bytes、cluster=(cx,cy,cz)
cute.compile(fn, *args, options="...")提前编译"opt-level=3"、"keep-ptx" 等
线程/块/网格索引
API用途
cute.arch.thread_idx()返回 (tidx, tidy, tidz)
cute.arch.block_idx()返回 (bidx, bidy, bidz)
cute.arch.block_dim()块维度
cute.arch.grid_dim()网格维度
cute.arch.warp_idx()块内的 warp 索引
cute.arch.block_idx_in_cluster()集群内的块索引（Hopper+）
布局和张量构造
API用途说明
cute.make_layout(shape, stride)创建布局核心抽象：将坐标映射为索引
cute.make_ordered_layout(shape, order)按步幅顺序创建布局例如，列优先布局使用 order=(1,0)
cute.make_identity_layout(shape)创建恒等布局用于谓词处理
cute.make_tensor(ptr, layout)创建张量视图将指针与布局配对
cute.make_fragment(layout, dtype)分配寄存器张量用于累加器
cute.make_fragment_like(src, dtype)创建与源匹配的寄存器张量—
cute.make_ptr(dtype, addr, memspace)创建带地址空间标记的指针cute.AddressSpace.smem/gmem
cute.from_dlpack(tensor)转换 PyTorch/JAX 张量返回带有 .mark_layout_dynamic() 的 cute.Tensor
布局代数
API用途
cute.size(t, mode=None)总大小或指定模式的大小
cute.shape(t)获取形状元组
cute.stride(t)获取步幅元组
cute.rank(t)模式数量
cute.local_tile(tensor, tiler, coord, proj)提取块的分块
cute.logical_divide(tensor, divisor)划分为分块和剩余部分
cute.composition(layout1, layout2)组合布局
cute.complement(layout)求布局的补布局
cute.coalesce(layout)简化布局
cute.flatten(t)展平张量/布局
cute.group_modes(layout, start, end)将多个模式组合在一起
cute.tile_to_shape(tile, shape)将分块扩展到目标形状
cute.ceil_div(a, b)向上取整除法
MMA 操作
API用途备注
cute.make_tiled_mma(op)根据操作创建分块 MMA例如，tcgen05.MmaF16BF16Op(...)
tiled_mma.get_slice(thread_idx)获取线程的分区返回 ThrMma
thr_mma.partition_A(tensor)对 A 操作数进行分区线程级视图
thr_mma.partition_B(tensor)对 B 操作数进行分区线程级视图
thr_mma.partition_C(tensor)对 C 操作数进行分区线程级视图
tiled_mma.partition_shape_C(shape)获取 C 分区的形状用于片段分配
tiled_mma.make_fragment_A(tensor)创建 A 片段寄存器分配
tiled_mma.make_fragment_B(tensor)创建 B 片段寄存器分配
tiled_mma.make_fragment_C(shape)创建 C 片段寄存器分配
cute.gemm(tiled_mma, d, a, b, c)执行 MMA：d = a @ b + c分派到硬件 MMA
复制和 TMA 操作
API用途备注
cute.make_tiled_copy(atom)创建分块复制用于批量数据移动
cute.copy(atom, src, dst, **kw)执行复制tma_bar_ptr=、mcast_mask=
cute.basic_copy(src, dst)逐元素复制无需 atom
cute.prefetch(atom, src)预取 TMA 描述符—
cute.nvgpu.make_tiled_tma_atom_A(op, tensor, smem_layout, tile, mma)为 A 创建 TMA atom仅主机端
cute.nvgpu.make_tiled_tma_atom_B(op, tensor, smem_layout, tile, mma)为 B 创建 TMA atom仅主机端
共享内存
API用途
cutlass.utils.SmemAllocator()创建 SMEM 分配器
smem.allocate_tensor(dtype, layout, align, swizzle)在 SMEM 中分配张量
smem.allocate(struct_type)在 SMEM 中分配结构体
张量内存（仅限 SM100）
API用途
cutlass.utils.TmemAllocator(...)创建 TMEM 分配器
tmem.allocate(num_cols)分配列
tmem.wait_for_alloc()等待分配完成
tmem.retrieve_ptr(dtype)获取指针
tmem.free(ptr)释放内存
同步
API用途
cute.arch.mbar.init(barrier_ptr, count)初始化屏障
cute.arch.mbar.arrive(barrier_ptr)到达屏障
cute.arch.mbar.wait(barrier_ptr, phase)等待屏障
cute.arch.mbar.arrive_and_expect_tx(barrier_ptr, bytes)到达屏障并指定预期的 TX 字节数
cute.arch.mbar.try_wait(barrier_ptr, phase)非阻塞等待
Pipeline classes (see references)多阶段异步流水线
数学函数（对 TensorSSA 逐元素运算）
API用途API用途
cute.exp(x)e^xcute.exp2(x)2^x
cute.log(x)ln(x)cute.log2(x)log2(x)
cute.sqrt(x)sqrtcute.rsqrt(x)1/sqrt(x)
cute.sin(x)sincute.cos(x)cos
cute.tanh(x)tanhcute.erf(x)erf

所有数学函数都支持 fastmath=True，以使用近似硬件内建函数。

Tensor Operations
API用途
cute.full(shape, val, dtype)填充指定值
cute.zeros_like(tensor)创建与张量形状相同的零张量
cute.where(cond, x, y)条件选择
cute.any_(tensor)逻辑任一
cute.all_(tensor)逻辑全部
调试
API用途
cute.printf(fmt, ...)设备端 printf
cute.print_tensor(tensor)打印张量内容
print(...)编译时打印（仅限元阶段）

关键约束

  1. @cute.kernel 函数名必须与模块文件名一致——与 cuTile 的规则相同

  2. TMA 描述符仅限主机端使用——在 @cute.jit 中创建，并传递给 @cute.kernel

  3. 寄存器预算：每线程最多 255 个——根据 docs/devices/ 规格进行验证

  4. SMEM 限制因设备和 SM 版本而异——SM100（B200/B300）：228 KB/SM，227 KB/块（选择启用）；SM120（RTX 5090）：96 KB

  5. 流水线阶段会占用 SMEM——每个阶段都需要自己的缓冲区；请验证总量

  6. 屏障同步错误不会显式报错——会导致结果错误，而非程序崩溃；务必使用 --check 进行测试

  7. 集群维度必须与网格兼容——集群形状必须能整除网格

  8. print() 仅用于编译时打印——设备端输出请使用 cute.printf()

  9. 循环中不支持提前退出——请改用谓词控制

  10. 32 位布局代数——形状/步长限制为 32 位整数

  11. 架构支持：Ampere（SM80）及以上；SM100（B200/B300）支持包括 tcgen05 和 TMEM 在内的全部功能

  12. 特定架构的内核：部分内核面向特定 SM 版本（例如，tcgen05 MMA 操作要求 SM100，而非笼统的“Blackwell”）。Blackwell 这一营销名称涵盖多个 SM 版本，它们具有不同的指令集：

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

    运行、编译或分析性能前，务必检查设备：nvidia-smi --query-gpu=name --format=csv,noheader。然后与内核的目标 SM 版本进行匹配——仅凭 GPU 名称并不足够；你必须了解它对应的 SM 版本。如果 SM 版本不匹配，请跳过运行/性能分析步骤，不要尝试执行，以免失败或产生误导性结果。

控制流

构造行为备注
for i in range(N)如果 N 是静态值，则展开循环标准 Python range
for i in cutlass.range(N)运行时循环生成 MLIR 循环
for i in cutlass.range_constexpr(N)编译时展开循环N 必须是静态值
if cutlass.const_expr(cond)编译时分支在编译时消除
if cond运行时分支两个分支都必须通过类型检查

Blackwell 数据中心 MMA 操作（tcgen05）——仅限 SM100

需要 SM100。 tcgen05 操作仅适用于 B200/B300（SM100）。消费级 Blackwell GPU（如 RTX 5090（SM120/sm_120a））不支持这些操作。

from cutlass.cute.nvgpu import tcgen05

# Create MMA operation for SM100 (B200/B300) tensor cores
op = tcgen05.MmaF16BF16Op(
    dtype=cutlass.Float16,           # Input type
    acc_dtype=cutlass.Float32,       # Accumulator type
    shape=(128, 128, 64),            # M x N x K tile shape
    cta_group=tcgen05.CtaGroup.ONE,  # ONE or TWO CTAs
)

tiled_mma = cute.make_tiled_mma(op)
操作输入类型形状
MmaF16BF16OpFP16/BF16 → FP32各种 M×N×K
MmaF8F6F4OpFP8/FP6/FP4 → FP32窄精度
分块缩放变体使用缩放因子参见示例

常见模式

使用 TMA 流水线的 GEMM（简化版）
@cute.kernel
def gemm_kernel(tiled_mma, tma_a, tma_b, smem_a, smem_b, gmem_c, ...):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()

    # Get thread's MMA partition
    thr_mma = tiled_mma.get_slice(tidx)

    # Allocate accumulator in registers
    acc = tiled_mma.partition_shape_C(tile_shape)
    tCrC = cute.make_fragment(acc, cutlass.Float32)

    # K-loop with TMA pipeline
    for k_tile in range(num_k_tiles):
        # TMA copy: global → shared (async)
        cute.copy(tma_a, gA_tile, sA_tile, tma_bar_ptr=barrier)
        cute.copy(tma_b, gB_tile, sB_tile, tma_bar_ptr=barrier)

        # Wait for TMA to complete
        cute.arch.mbar.wait(barrier, phase)

        # Partition shared memory for this thread
        tCsA = thr_mma.partition_A(sA)
        tCsB = thr_mma.partition_B(sB)

        # MMA: accumulate into registers
        cute.gemm(tiled_mma, tCrC, tCsA, tCsB, tCrC)

    # Epilogue: write back to global memory
    cute.basic_copy(tCrC, gC_partition)
Warpgroup 专项化（生产者/消费者）
warp_idx = cute.arch.warp_idx()
is_producer = warp_idx < num_producer_warps

if is_producer:
    # Issue TMA copies for upcoming pipeline stages
    for stage in range(num_stages):
        cute.copy(tma_atom, src, dst, tma_bar_ptr=barriers[stage])
        cute.arch.mbar.arrive_and_expect_tx(barriers[stage], bytes_per_stage)
else:
    # Consume data from shared memory, execute MMA
    for stage in range(num_stages):
        cute.arch.mbar.wait(barriers[stage], phase)
        cute.gemm(tiled_mma, acc, sA_stage, sB_stage, acc)
持久化内核循环
@cute.kernel
def persistent_kernel(tiled_mma, num_tiles, ...):
    bidx, _, _ = cute.arch.block_idx()
    num_blocks = cute.arch.grid_dim()[0]

    tile_id = bidx
    while tile_id < num_tiles:
        # Process tile
        # ... TMA copy, MMA, epilogue ...
        tile_id += num_blocks

详细参考

评论

支持 Markdown