CUTLASS Task Scheduling（TS）入门教程：用 Copy 内核理解资源、任务与显式调度

CUTLASS Task Scheduling（TS）入门教程：用 Copy 内核理解资源、任务与显式调度

【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass

本篇是 CUTLASS Python DSL 中Task Scheduling（TS）编程模型的入门指南，基于 Tutorial 01: Copy Basics 及其配套内核展开。TS 是面向 NVIDIA GPU 异步、warp 专用化内核的编程模型：它不隐藏底层的 TMA、Tensor Core 等指令，而是为"哪个 warp 组做哪件事、数据何时就绪、屏障如何推进"提供一种可被静态校验的显式表达方式。读完本文，你将掌握 TS 的四大核心概念（Resource、Task、Explicit Schedule、Dependency Graph），并能在实际内核中写出可被编译器验证的生产者/消费者流水线。

说明：本教程系列为教学目的编写，不保证达到当前最优性能（"These kernels are intended only for TS educational purposes. State-of-the-art performance is not guaranteed."）。

为什么需要 TS：显式化异步调度

现代 NVIDIA GPU 暴露了大量异步操作——TMA、tcgen05 MMA、Tensor Memory 拷贝等。高效利用它们需要编写异步的 warp 专用化程序：每个 warp 或 warp 组"拥有"一项任务（如数据搬运或数学计算），任务之间通过共享内存（SMEM）、张量内存（TMEM）等资源通信，因此必须同步以避免并发数据竞争。在裸机（bare-metal）代码中，调度是隐式的：warp 角色、屏障到达（arrival）、阶段推进（phase advancement）散落在各个 warp 分支和内核主体里，没有独立的调度对象可供检查，也就没有死锁或竞争检查器（见 教程总览）。

一个手写的 warp 专用化内核通常长这样：

if warp_idx == load_warp: compute coords acquire empty stage issue TMA commit full stage else: wait full stage consume SMEM release empty stage

TS 将上述模式结构化拆解：

这种拆分是刻意的：TS 鼓励开发者显式声明生产者-消费者流水线中的角色、数据归属与操作顺序，并在将调度 lower 到选定 warp 分支之前，静态校验调度的正确性与流水线初始化。

Kernel 01：GMEM Grid-Stride Copy——TS 词汇全景

01_copy_grid_stride.py 把一维int16张量从一块全局内存（GMEM）经寄存器拷贝到另一块 GMEM。没有 SMEM、没有流水线、只有一个任务，但它是能练习到全部 TS 核心声明的最小内核。

其数据流可以表示为：

InputGmemResource -- item --> OutputGmemResource

不使用 TS 的 CUTLASS Primitives 裸机版本如下（也定义在同一文件中，作为基线gmem_copy_naive_kernel）：

@cute.kernel def gmem_copy_naive_kernel( num_entries: cutlass.Int32, source_tensor: cute.Tensor, destination_tensor: cute.Tensor, num_warps: cutlass.Constexpr, ): gdimx, _, _ = cute.arch.grid_dim() bx, _, _ = cute.arch.block_idx() tx, _, _ = cute.arch.thread_idx() gid = bx * num_warps * 32 + tx for i in cutlass.range(gid, num_entries, gdimx * num_warps * 32): destination_tensor[i] = source_tensor[i]

Step 1：声明资源（MemoryResource）

对于这个拷贝，存在两个独立的物理资源：源 GMEM 与目标 GMEM，各声明一个资源。两者都不拥有 SMEM 或流水线，只拥有张量句柄，以及从资源读/写一个元素的逻辑：

@dataclass(kw_only=True) class InputGmemResource(MemoryResource): source_tensor: cute.Tensor num_entries: cutlass.Int32 num_warps: int ... @dataclass(kw_only=True) class OutputGmemResource(MemoryResource): destination_tensor: cute.Tensor num_entries: cutlass.Int32 num_warps: int ...

Step 2：定义资源上的工作（producer vs consumer）

TS 不会替你生成数据读写逻辑，内核开发者仍需为每个资源写明"数据如何写入（producer work）"与"如何读出（consumer work）"。工作方法用@producer_work或@consumer_work装饰：

• @consumer_work从资源中读出一个值；
• @producer_work向资源中写入一个值。

本内核中，InputGmemResource.get_item()是 consumer work（读出输入元素），OutputGmemResource.set_item()是 producer work（写入输出元素）：

@consumer_work(returns=item) @cute.jit def get_item(self, stage_info: StageInfo) -> cutlass.Int16: gid = stage_info.loop_offset val = cutlass.Int16(0) if gid < self.num_entries: val = self.source_tensor[gid] return val

@producer_work @cute.jit def set_item(self, stage_info: StageInfo, data: cutlass.Int16) -> None: gid = stage_info.loop_offset if gid < self.num_entries: self.destination_tensor[gid] = data

每个工作方法有两条硬性要求：

1. 装饰器顺序：外层必须是@producer_work/@consumer_work，内层必须是@cute.jit，顺序不能颠倒。TS 装饰器把方法注册到资源上，@cute.jit编译其函数体。
2. 固定签名形状：第一个参数是self，第二个参数恒为stage_info: StageInfo（由 TS 注入，携带当前循环与流水线上下文，例如stage_info.loop_offset）。工作需要的额外输入以后置关键字或位置参数声明（如set_item的data: cutlass.Int16）；需要向下游发射的值通过装饰器上的returns=...加上匹配的 Pythonreturn声明。

从 StageInfo 源码 可以看到它携带的完整上下文：loop_offset（当前循环迭代下标）、loop_start/loop_end/loop_step（循环边界与步长）、stage_idx（流水线阶段下标，无流水线时为None）、label（编译期工作标签）、barrier（当前流水线阶段的 mbarrier 指针）、work_tile（瓦片调度坐标）、context（统一上下文，含smem_base等）。

Step 3：值如何在资源间流动（TaskLocalVariable）

get_item()产生的值必须被set_item()消费。该值声明为拥有它的资源（即产生它的输入资源）上的TaskLocalVariable：

@dataclass(kw_only=True) class InputGmemResource(MemoryResource): item: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() def __post_init__(self) -> None: self.item = TaskLocalVariable( dtype=cutlass.Int16, default=cutlass.Int16(0), docs="Input element loaded for the current grid-stride iteration.", )

从源码看，TaskLocalVariable是"资源拥有的逻辑变量的公共身份"，在生成代码中物化为 task-local 寄存器状态，default/default_factory提供任何 producer 写入之前及 SSA 汇合处的安全初值（见 resources.py）。

一条规则：不要在 producer/consumer work 内部直接读写self.item。该字段是调度构建器用来追踪数据流边的元数据，不是运行时活载荷；载荷通过方法的return值与调度中的命名绑定传播（下一步）。

Step 4：定义调度（@schedule + domain_loop）

调度是每个异步任务显式规定的操作序列，用@schedule装饰的函数按顺序调用资源工作方法，即 grid-stride 循环的 TS 版本：

@schedule def schedule_fn( input_gmem: InputGmemResource, output_gmem: OutputGmemResource, ) -> None: threads_per_block = num_warps * 32 start = bx * threads_per_block + tx step = gdimx * threads_per_block with domain_loop(start, num_entries, step, unroll=unroll): res = input_gmem.get_item() output_gmem.set_item(data=res)

这里发生两件事，二者都是捕获规则而非普通 Python 执行：

• domain_loop(...)是捕获的循环控制流。捕获期间它不会真正跑 Python 循环，而是记录循环边界，并把当前坐标通过stage_info.loop_offset暴露给每个工作方法。从 schedule_builder.py 的源码可见，它接受 1～3 个类似range()的边界（domain_loop(end)/(start, end)/(start, end, step)，省略的 start 默认 0、step 默认 1），step不能为 0，且任意边界都可以是可调用对象——这使循环边界可以是动态的（动态域在教程 03 中展开）。
• res = input_gmem.get_item()在捕获期间并不真正加载int16。TS 记录的是：get_item()产生一个 task-local 值，且稍后set_item(data=res)在同一任务中消费同一个值。

更好的读法是：把调度看作 TS 将要生成循环的描述。domain_loop会成为真正的 grid-stride 循环（使用你传入的start、num_entries、step、unroll），每个工作方法调用被内联到其出现位置。本内核捕获的调度大致 lower 为：

threads_per_block = num_warps * 32 start = bx * threads_per_block + tx step = gdimx * threads_per_block for gid in cutlass.range(start, num_entries, step, unroll=unroll): # inlined input_gmem.get_item(), with stage_info.loop_offset == gid res = cutlass.Int16(0) if gid < num_entries: res = source_tensor[gid] # inlined output_gmem.set_item(data=res) if gid < num_entries: destination_tensor[gid] = res

调度体看起来像普通命令式代码，但必须遵守捕获规则：值名连接 producer 调用与 consumer 调用，调用顺序定义 TS 将强制执行的次序。@schedule装饰器的源码实现也印证了这一点——被装饰函数收到的是资源的ResourceProxy包装，方法调用会记录调度条目与路由边，domain_loop标记结构边界（见 schedule_builder.py）。

Step 5：定义依赖图（resource_dependency_graph）

除了捕获值，TS 还需要资源级的排序信息。依赖图声明输出资源依赖于输入资源：

resource_dependency_graph = { output_gmem_resource: [input_gmem_resource], }

依赖图不是通过调度传值res的替代品，含义也未必相同：图记录的是资源之间的依赖——输出资源在输入资源被消费之前不能被生产；而捕获值说明的是哪个标量值从get_item()流向set_item()。有时开发者出于种种原因不让consumer_work向producer_work返回值，但依赖图仍然必须记录两个资源之间的关系。

Step 6：定义任务（Task）

任务是 warp 专用化的单元，它把一个捕获的调度绑定到一段 warp 范围，并声明自己消费、生产哪些资源：

task = Task( name="CopyTask", src_resources=[input_gmem_resource], dst_resources=[output_gmem_resource], warp_idx=0, num_warps=num_warps, schedule=schedule_fn(input_gmem_resource, output_gmem_resource), )

资源放在哪一侧仍遵循"从资源视角看生产者/消费者"的规则：

• 任务调用某资源的consumerwork（从资源读出）→ 该资源进src_resources；
• 任务调用某资源的producerwork（向资源写入）→ 该资源进dst_resources。

本任务从输入 GMEM 读出、向输出 GMEM 写入，所以 input 是 source、output 是 destination。schedule是调用schedule_fn(...)返回的捕获调度对象本身。

Step 7：TaskManager 与固定生命周期

TaskManager校验骨架并为每个 warp 执行选定的调度：

task_manager = TaskManager( tasks=[task], resource_dependency_graph=resource_dependency_graph, ) task_manager.setup_resources_and_tasks() task_manager.run()

生命周期是固定的，且必须按上文的步骤顺序出现在内核中：

1. 创建资源（resources）；
2. 创建依赖图（dependency graph）；
3. 捕获调度（schedules）；
4. 创建任务（tasks）；
5. 创建TaskManager（此处 TS 校验并接线骨架）；
6. 调用setup_resources_and_tasks()（此处 TS 初始化共享内存屏障）；
7. 调用屏障初始化 fence 并同步线程（prims.fence_mbarrier_init()/prims.barrier_cta_sync(0)，见02_copy_tma.py）；
8. 调用run()（为每个 warp 选择合适任务，执行其附带显式调度）。

从 task_manager.py 的模块文档可知一个关键事实：整个 TS 框架（TaskManager、Task、MemoryResource、acquire/release/commit）在 DSL 编译期间被完全 trace 掉，生成的 PTX 是单一整体循环，与手写裸机内核等价，生成的 GPU 代码中没有运行时任务分发、没有框架开销。若 TS 内核比裸机内核慢，原因在生成代码本身（指令序列、寄存器使用、循环结构），而非框架抽象。

生命周期同样出现在之后每个教程中。

Kernel 02：TMA Copy——第一个真正的生产者/消费者流水线

02_copy_tma.py 通过 SMEM 暂存 128 列的 tile，把二维 FP16 矩阵从 GMEM 搬到 GMEM。它复用了内核 01 的全部概念，并引入第一条真正的生产者/消费者流水线。可以把它理解为"内核 01，但值改经共享缓冲区而非寄存器流动"。

资源布局为何变化

内核 01 中值经寄存器直接从一块 GMEM 到另一块 GMEM；这里硬件路径不同：TMA 从 GMEM 拷贝进入共享内存（SMEM），store warp 再从 SMEM 读出并写回 GMEM。SMEM 是被两组不同 warp 触碰的缓冲区，所以它拥有自己的资源：

InputGmemResource -> SmemResource -> OutputGmemResource coordinates staged tile GMEM output

关键新点：SmemResource同时是 destination 和 source——

• TMA 写入 SMEM →SmemResource是生产者目标（tma_load是producer_work）；
• 从 SMEM 读到寄存器 →SmemResource是消费者来源（read_smem是consumer_work）。

内核 01 中"从资源视角看生产者/消费者"的规则依然适用，但同一个资源现在取决于你看哪个任务而扮演两个角色。

为什么现在有两个任务

内核 01 里单个任务做完了一切；这里 load 侧与 store 侧跑在不同 warp、位于 SMEM 流水线的两侧。由于 staged 共享内存，只要还有空闲 SMEM 阶段，TMA 加载任务就可以跑在 store 任务之前，所以工作被拆成两个异步任务：

• LoadTask：acquire 空 SMEM 阶段 →tma_load()→ commit 满阶段；
• StoreTask：wait 满 SMEM 阶段 →read_smem()→ release 空阶段。

每个任务有自己的捕获调度。load 调度跑在 warp 4 上，计算 tile 坐标并向 SMEM 发出 TMA：

@schedule def load_schedule(input_gmem: MemoryResource, smem: MemoryResource) -> None: smem.init_load_state() with domain_loop(0, num_rows, box_dim[1]): gmem_idx = input_gmem.compute_coords() smem.try_acquire() smem.acquire() smem.tma_load(gmem_idx=gmem_idx) smem.commit()

gmem_idx与内核 01 中的res是同一种捕获数据流边：compute_coords()（输入资源上的 consumer work）发射它，然后 smem 资源被 acquire——必须有空闲槽位供数据写入——tma_load(gmem_idx=...)（SMEM 资源上的 producer work）消费该坐标并执行 TMA，随后 commit 缓冲区，表示数据已写入 smem 资源。

store 调度跑在 warp 0-3 上，等待满 SMEM 阶段、读取每个线程的元素、释放阶段并写入输出 GMEM：

@schedule def store_schedule(smem: MemoryResource, output_gmem: MemoryResource) -> None: smem.init_read_state() with domain_loop(0, num_rows, box_dim[1]): smem.try_wait() smem.wait() smem_val = smem.read_smem() output_gmem.store(smem_val=smem_val) smem.release()

保留字：TS 管理的流水线操作 vs 开发者工作方法

注意两类调用的区别。compute_coords()、tma_load()、read_smem()、store()是内核开发者自定义的 producer/consumer work，定义在资源上；而try_acquire、acquire、commit、try_wait、wait、release不是开发者定义的工作方法，它们是每个流水线资源都暴露的、由 TS 管理的固定流水线操作，名称保留，在任何异步生产者/消费者流水线中含义都一致：

• 生产者侧：try_acquire→acquire（认领空阶段）→ … →commit（发布填满的阶段）；
• 消费者侧：try_wait→wait（等待满阶段）→ … →release（把阶段还给生产者）。

因为这些名字为流水线协议保留，开发者不能把 producer/consumer work 命名为acquire、commit、wait、release等；工作方法（TMA 加载、SMEM 读、GMEM 写）必须放在这些括号内部。

流水线调用显式地留在捕获调度中而非藏进 TMA 工作方法，有几点好处：

• 调度由用户显式定义，TS 可以校验其正确性；
• 性能优化更容易——开发者可以尝试调整 wait/acquire/commit/release 的位置以更好地重叠阶段或隐藏延迟，而不用担心破坏代码；如果新调度无效，TS 会在生成内核之前失败；
• 资源的生产/消费逻辑与执行顺序分离，可读性与可维护性更好。

依赖图现在有两条边：

resource_dependency_graph = { smem_resource: [input_gmem_resource], output_gmem_resource: [smem_resource], }

调度正确性注记：release 与 DMA 消费者

@schedule def store_schedule(smem: MemoryResource, output_gmem: MemoryResource) -> None: smem.init_read_state() with domain_loop(0, num_rows, box_dim[1]): smem.try_wait() smem.wait() smem_val = smem.read_smem() smem.release() output_gmem.store(smem_val=smem_val)

即在output_gmem.store之前调用smem.release()也是合法调度，TS 校验器允许。提前 release 可能带来更好性能——SMEM 资源更早可用，load 任务的数据加载可与 store 任务写 GMEM 并行。

可以这样做，是因为从共享内存消费的数据已读入寄存器，资源可以在 consumer work 之后释放。但这并非总是可行：例如当依赖资源的 producer work 是 DMA 操作时——tensor-core MMA、发起 TMA 等——流水线不能在相应 producer work 之前释放。关于 TS 因这个原因拒绝某个调度的具体例子，见 Tutorial 02: 调度正确性注记：release 顺序与 DMA 消费者。

PipelineConfig

被两组 warp 触碰的共享缓冲区需要流水线守护。通过给SmemResource一个PipelineConfig声明流水线类型：

pipeline_config = PipelineConfig.create_tma_async_pipeline_cfg( num_stages=num_stages, num_bytes=box_dim[0] * box_dim[1] * cutlass.Float16.width // 8, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup( pipeline.Agent.Thread, num_warps_epilogue * 32, ), )

之所以用create_tma_async_pipeline_cfg工厂，是因为流水线两端的行为，名字同时反映了两者：tma是生产者端——TMA 异步写入 SMEM，并通过带事务字节计数（transaction-byte count）的 mbarrier 发出完成信号；该工厂恰好构建此协议：生产者的commit以期望字节数武装 mbarrier，消费者的wait阻塞直到硬件报告这么多字节落地。async后缀是消费者端——消费者是异步线程，其release经 mbarrier 排序。不同的生产者/消费者（例如直接写 SMEM 的 warp、cp.async拷贝、以 tensor core 作为数据消费者等）会使用不同的工厂，因为"阶段已满/阶段空闲"的信令方式不同（后续教程展开）。从 resources.py 源码可见，该工厂构造的PipelineConfig的pipeline_type为PipelineType.TmaAsync，另有create_tma_umma_pipeline_cfg（TMA 生产者 + UMMA 消费者）等其他工厂。

各参数的含义：

• num_stages——流水线轮转的 in-flight SMEM 缓冲区个数。阶段越多，生产者越能跑在消费者前面。这里取1以保持示例简单；后续内核用更多阶段时调度形状完全相同。
• num_bytes——每阶段传输的字节数，由 tile 形状与元素位宽计算（box_dim[0] * box_dim[1] * cutlass.Float16.width // 8）。这是消费者wait监视的事务字节计数，因此必须与实际tma_load传输量一致。
• producer_group——填充阶段的协作组。生产者是 load 任务，只有一个被选举线程发出 TMA 并信令屏障，所以是单个Agent.Thread。
• consumer_group——排空阶段的协作组。消费者是运行在 4 个 epilogue warp 上的 store 任务，所以是num_warps_epilogue * 32个线程；所有这些线程都必须到达release，阶段才会还给生产者。

TS 会把这些组大小与调用 producer/consumer work 的任务 warp 数交叉核对，因此 pipeline config 与实际调度不匹配会在 setup 阶段被抓住。

SMEM 分配

TS 管理共享内存分配并跟踪每个资源使用了多少 SMEM。它用SmemAllocation定义共享内存区域，并通过get_smem_requirements()说明所需共享内存：

self._alloc_smem = SmemAllocation( "smem_data", dtype=cutlass.Float16, count=num_stages * tile_size, alignment=128, ) def get_smem_requirements(self): return [self._alloc_smem]

内核把该资源交给SmemAllocator：

allocator = SmemAllocator() allocator.add_resource(smem_resource) allocator.compute_layout()

分配器把 SMEM 数据载荷与流水线屏障存储打包进一个共享内存块（这也是02_copy_tma.py中 pipeline config 故意不传barrier_ptr的原因——由分配器统一放置屏障存储，并在setup_resources_and_tasks()期间补丁到 pipeline config 上）。资源在辅助工作方法中物化其带类型的cutlass.Array视图：

@producer_work(work_attrs=WorkAttr.AUXILIARY) def init_load_state(...): ... @consumer_work(work_attrs=WorkAttr.AUXILIARY) def init_read_state(...): ...

辅助工作（auxiliary work）是 producer/consumer work 的特殊类型：它们仍被捕获在调度中，但其在调度中的顺序不参与依赖图与 wait/release/acquire/commit 的校验。辅助工作只用于不触碰实际资源数据的辅助代码，例如把指针设置代码从循环中提升出来。不要把读写资源载荷的工作标记为 auxiliary——TMA 加载、SMEM 读、MMA、GMEM 写等应保持为普通 producer/consumer work，以便 TS 校验其顺序。

try_* 与阻塞调用

最小生产者括号是：

smem.acquire() smem.tma_load(...) smem.commit()

最小消费者括号是：

smem.wait() smem_val = smem.read_smem() smem.release()

教程在阻塞调用前使用try_acquire()与try_wait()：

smem.try_acquire() smem.acquire()

try_*调用提前发起异步屏障查询：若成功，随后的acquire()/wait()立即穿过；否则阻塞直到屏障翻转。当早查询与阻塞调用之间有独立工作可放，或要把不同资源的多个try_*调用分组并发执行时，使用这种拆分形式。本示例中这些可选调用主要起示范作用，以匹配更真实 GEMM 示例中的调度形状。

Warp 专用化与 Padding

TMA copy 使用 8 个 warp：

CUDA 寄存器重分配按四个连续 warp 组成的 warp 组进行。TS 因此检查同一组内每个 warp 都声明了num_registers值。Warp 4-7 是一个组；其中只有 warp 4 做实际加载工作，所以PaddingTask给 warp 5-7 相同的寄存器预算与一个空的捕获调度。若num_registers为None，则不需要 padding 任务。这一校验逻辑在 task_manager.py 的_verify_warp_group_registers中实现：同一 warp 组内不同num_registers会抛ValueError；一旦某组任一任务设置了num_registers，该组全部 4 个 warp 都必须被设置了num_registers的任务覆盖。

与原始内核对比

同一文件包含手写 CUTLASS Python Primitives 版本tma_copy_raw_kernel：它手工创建 SMEM 存储、mbarrier 存储、TMA 流水线参与者，以及if warp_idx == 4分支。TS 版本计算结果相同，但把分支结构与屏障协议表达为：

• 带PipelineConfig的SmemResource；
• LoadTask、StoreTask、PaddingTask；
• load_schedule与store_schedule；
• resource_dependency_graph={smem: [input_gmem], output_gmem: [smem]}；
• TaskManager.setup_resources_and_tasks()与TaskManager.run()。

这正是 TS 存在的原因：底层工作仍然显式，但调度现在是被检查的对象，而不是散落的控制流。

要点回顾

• 资源是所有权单元：它可能拥有物理存储、流水线、task-local 值，或只有坐标逻辑。
• 任务是 warp 专用化的单元：它拥有一段 warp 范围与一个捕获调度。
• 捕获调度是顺序契约：它记录资源方法调用、流水线括号、循环结构与捕获值流动。
• 依赖图是资源顺序契约：它与捕获的标量值相互独立。
• 生产者/消费者命名从资源视角出发：producer work 写入资源，consumer work 从资源读出。
• 流水线生产者用acquire然后commit；流水线消费者用wait然后release。
• TaskManager.setup_resources_and_tasks()是 TS 在run()之前校验并接线资源/任务骨架之处。

另外值得记住的是 TaskManager 提供的校验能力：构造时会打印调度表并运行校验（print_and_verify），用轮转方式模拟所有任务跨 head/loop/tail 阶段的并发执行，ConsumerWait与ProducerAcquire被视为阻塞阶段，可在内核构建期检测简单死锁、未配对的 commit/release（见 task_manager.py 的print_schedule_list）。默认还启用穷尽死锁/竞争检查（exhaustive_deadlock_race_check=True），可通过skip_validation等参数控制校验强度。这也印证了：TS 的同步与执行顺序问题大多在编译期暴露，而不是变成运行时挂起或数据竞争。

如何运行

两个示例位于 examples/python/CuTeDSL/experimental/task_scheduling/blackwell/tutorial/01_copy_basics_ts/，需在有 GPU 的 CUDA 环境中运行：

python 01_copy_grid_stride.py python 02_copy_tma.py --rows_cols 256,512 python 02_copy_tma.py --rows_cols 256,512 --run-raw-kernel

命令行细节（见02_copy_tma.py的__main__）：--rows_cols接收逗号分隔的两个整数（默认256,512），--run-raw-kernel切换为运行手写原始内核做对比。形状约束：行数必须为正，列数必须是tile_size=128的倍数，否则run_tma_copy_kernel_prim抛ValueError。两个内核运行后都会用torch.testing.assert_close校验目标张量与源张量一致并打印PASS。

建议的进阶路线（见 教程总览）：先完整读懂本教程（Tutorial 01），再按顺序阅读 Tutorial 02: Simple FP16/BF16 GEMM（引入GmemAb -> SmemAb -> TmemC -> GmemD更深的资源链与 TMEM 所有权）、Tutorial 03: persistent scheduling 与动态域，再进入教程 04-07 探索集群、PDL、Split-K、PipelineGroup 等硬件特性。生产级 TS 示例可参考blackwell/kernel/dense_gemm_ts/batched_gemm/与blackwell/kernel/attention_ts/fmha/。

【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass