提示与讲解

CuTe 专题：用“布局代数”管住 kernel 里所有的下标

CuTe 是 CUTLASS 3 和 CuTe DSL 的核心：它把 GPU kernel 里五花八门的“谁拿哪个数、数放在哪”统一写成同一种对象 Layout，再用一小套代数运算把它们组合起来。这一页先讲它要解决的问题和思路，再用四个可以动手玩的演示把 Layout、composition、logical_divide 和线程划分一步步拆开。演示背后跑的就是参考解 g2_cute_layout.py 的算法（移植到网页里，并和 Python 版逐个对拍过）。

怎么读这一页

这一页讲“为什么”和“怎么想”；G2 讲解讲参考解每个函数的细节，两页可以对照着看。所有演示都可以点：按钮切换例子，▶| 一步步走，鼠标悬停看细节。参考解：drillsg2_cute_layout.py。

① 问题：kernel 里的“下标地狱”

翻一下你在这套题里写过的 kernel，真正难写、最容易错的不是乘法，而是下标：

题目	这一行在算什么	代码
C3 转置	线程 → tile 里的位置，tile 在 shared memory 里的地址	`tile[threadIdx.x][threadIdx.y + j]`，`tile[32][33]`
C4 GEMM	协作加载时第 e 个元素在哪	`m = e / BK, k = e % BK; As[k][m] = A[gm * K + gk]`
C4 GEMM	每个线程负责哪个 4×4 小块	`tx = tid % (BN/TN); ty = tid / (BN/TN); As[k][ty*TM + i]`
T4 Triton GEMM	program 负责 C 的哪个 tile	`pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)`
L2 paged decode	第 p 页的 KV 在哪	`phys * stride_cb + offs_t * stride_ct + kvh * stride_ch + offs_d`
G1	一个 lane 持有哪几个元素	`reptile + warpspttpw + lanespt + s`

每一行都是手写的整数算术。它们有三个共同的问题：

手写下标的问题悬停节点查看连接

一个高性能 GEMM 里，这四层映射要同时正确：数据在显存里怎么放、在 shared memory 里怎么放、每个线程搬哪些、每个线程的寄存器怎么喂给 Tensor Core。换一代 GPU（Ampere 的 mma.sync → Hopper 的 wgmma → Blackwell 的 tcgen05），最后两层就全变了。

类比 · 车同轨，书同文

战国时每个国家的车轴宽度、文字、度量衡都不一样，货物从一国运到另一国，每过一道边境都要换车、换秤、重新登记。手写下标就是这样：全局内存一套编号，shared memory 一套，线程一套，Tensor Core 又一套，每一层之间都靠人手写的公式“换车”。CuTe 做的事相当于统一度量衡：所有层都用同一种东西描述（Layout），层与层之间的转换也变成标准运算，不再是每次手写的一次性公式。

② 思路：一切映射都是 Layout，再加一小套代数

CuTe 的核心观察是：上面那些下标公式，本质上都是“坐标 → 编号”的函数。不管是“矩阵的 (i, j) 在内存的第几个位置”，还是“第 t 个线程的第 v 个值对应矩阵的哪个位置”，都可以用同一种对象表示：

Layout = (Shape, Stride)      一个函数：坐标 → 整数
         offset = Σ coord_i × stride_i

然后提供一小套运算，把这些函数组合、切分、化简：

CuTe 的层次悬停节点查看连接

越往右越高层，但每一层都建立在 Layout 上。参考解实现的就是最左边三层（IntTuple、Layout、布局代数），这也是理解整个 CuTe 的门槛。

类比 · 图形学里的变换矩阵

早期的图形程序里，平移、旋转、缩放、投影各写各的公式，组合起来很麻烦。后来大家发现它们都可以写成 4×4 矩阵，“先旋转再平移”就是两个矩阵相乘，求逆、复合都有现成的运算。CuTe 对下标做了同样的事：所有映射都是 Layout，“先切 tile 再分给线程”就是 Layout 的组合（composition / divide），你不再需要手推每一层的整数公式。

③ Layout：从坐标到偏移的函数

(4,2):(1,4) 的意思是：shape 是 4 × 2，stride 是 (1, 4)。坐标 (i, j) 映射到偏移 i·1 + j·4。一个整数下标先按列字典序（第一个模式变化最快）拆成坐标，再算偏移。

在下面的演示里切换不同的 layout，用 ▶| 一步步看“下标 → 坐标 → 偏移”。也可以在输入框里自己写一个 layout（比如 (3,4):(4,1)）。

演示 1：Layout 浏览器

自己写：

内存（偏移 0 … cosize−1）里放的是哪个 (行, 列)

几个值得点开看的例子：

列主序 vs 行主序：同样是 4×2，只是 stride 不同，“第 5 个元素”在内存里的位置就不同。
列主序 + padding (4,2):(1,5)：每列多空一格，内存条里出现虚线的空位。这就是 C3 的 tile[32][33]。
广播 (4,2):(0,1)：第一个模式的 stride 是 0，4 行全都映射到同一个位置。这就是 G1 里 SliceLayout 的广播，也是 CuTe 表达“同一个值给好几个线程”的方式。
分块 ((2,2),(2,2)):((1,4),(2,8))：下一节专门讲。

类比 · 储物柜

Layout 像商场储物柜的编号规则。你只知道“第 i 排第 j 列”（坐标），柜子上贴的却是一个号码（偏移）。stride 就是规则：往下一排号码加几，往右一列号码加几。同一排柜子，可以按“先往下数”编号（列主序），也可以按“先往右数”编号（行主序），还可以每列之间空一个柜子（padding），甚至让一整排共用一个柜子（广播）。柜子本身没动，变的只是编号规则。

④ 为什么要嵌套：一个维度里藏着两层结构

CuTe 的 shape 可以嵌套，比如 ((2,2),(2,2))。它表示一个逻辑上 4×4 的矩阵，但每一维内部又分成了两层：“块内第几个”和“第几个块”。

在演示 1 里点 ((2,2),(2,2)):((1,4),(2,8))：

行坐标 r（0..3）被拆成 (r % 2, r / 2)，偏移贡献是 (r%2)·1 + (r/2)·4
列坐标 c（0..3）被拆成 (c % 2, c / 2)，偏移贡献是 (c%2)·2 + (c/2)·8

结果是：每个 2×2 小块在内存里是连续的 4 个数（0–3、4–7……），小块之间再按列主序排。网格里看起来还是普通的 4×4，内存条上每 4 个颜色相近的格子就是一个小块。

这正是 tiling 的表达方式：逻辑坐标不变（还是 4×4），物理排列按块组织。kernel 代码照样用 (r, c) 访问，块怎么排全藏在 layout 里。换一种分块方式，只需要换 layout，访问代码一行都不用改。

类比 · 楼—单元—层—户

“5 栋 2 单元 3 楼 1 号”是一个四层嵌套的地址，但你也可以只说“第 37 户”。CuTe 的嵌套 shape 就是这样：一个逻辑维度（第几户）里藏着好几层物理结构（楼、单元、层）。代码可以只用“第几户”（扁平坐标），也可以用完整的四层地址（层次坐标），两者通过列字典序自动互相转换。

⑤ 布局代数：每个运算解决一个具体问题

运算	要解决的问题	一句话
`coalesce`	层次很深的 layout 不好读、也不好优化	同一个函数，化简成最少的模式
`composition` A∘B	“先按 B 选元素，再按 A 找地址”	用一个 layout 去选另一个 layout 的元素
`complement`	选了一部分，剩下的怎么描述	找出 A 在 [0, M) 里没覆盖到的“空隙”
`logical_divide`	怎么把一个大张量切成 tile	tile 的形状由 B 给，tile 怎么排由 complement 给

composition：先选，再查地址

R = A ∘ B 的意思是 R(i) = A(B(i))：B 决定“要 A 的第几个元素”，A 决定“这个元素在内存哪里”。演示里左边是 B 选出的下标，右边是它们在 A 的内存里的位置。

演示 2：composition，R(i) = A(B(i))

A（数据在内存里怎么放）：

B（要 A 的哪些元素）：

i → B(i)（B 选出 A 的第几个）

A 的坐标网格（格子里是偏移），被选中的描边

值得试的组合：

A = 4×4 行主序，B = (4,4):(4,1)：用“行主序的读法”去读一个行主序矩阵，结果 R = (4,4):(1,4)，相当于转置。C3 里你手写的坐标对调，在 CuTe 里就是一次 composition。
A = 4×2 行主序，B = 4:2：B 的步子跨过了 A 的第一维边界，结果被拆成两个模式 (2,2):(4,1)（G2 讲解第 ④ 节的例子）。
B = (2,2):(1,4)：从 4×4 里选出左上角 2×2 的子块。取 tile 本质上就是 composition。

类比 · 两次换乘

composition 像坐地铁换乘：B 是第一段线路，把你从“起点站 i”送到“换乘站 B(i)”；A 是第二段线路，把你从换乘站送到“终点 A(B(i))”。CuTe 的本事是直接给你一张合并后的直达线路图 R，以后不用每次都换乘。

complement 和 logical_divide：把一个张量切成 tile

切 tile 需要回答两个问题：一个 tile 长什么样（tile 内部的元素怎么选），以及 tile 之间怎么排。第一个由你给的 tiler B 决定；第二个就是 B 的 complement：B 没覆盖到的“空隙”，正好是 tile 的重复方式。

logical_divide(A, B) = A ∘ ( B , complement(B, size(A)) )
                             └ tile 内 ┘ └──── tile 之间 ────┘

演示 3：logical_divide，把 A 切成 tile

A：

tiler B：

结果：每一行是一个 tile（模式 1），每一列是 tile 内第几个（模式 0），格子里是偏移

A 的内存：颜色 = 属于第几个 tile

tiler 4:1：最普通的切法，4 个连续的一块，complement 是 4:4（每块起点相差 4）。
tiler 4:2：每个 tile 是“隔一个取一个”的 4 个元素，complement 是 (2,2):(1,8)，所以 tile 0 和 tile 1 交错，tile 2、3 在后半段。
A = 16:2：A 本身是跨步存放的，切出来的偏移也跟着翻倍，tile 的结构不变。这就是“tile 的形状”和“数据怎么存”被分开描述的好处。

二维的情况（GEMM 里把 M×N 的矩阵切成 BM×BN 的 tile）是对每一维分别做一次 divide，CuTe 叫它 zipped_divide / local_tile。

类比 · 切蛋糕

tiler 是你手里的模具，决定每一块蛋糕的形状（4 个连续的，或者隔一个取一个的 4 个）；complement 是“模具一次次往哪里放”，决定这些块在整个蛋糕上怎么排开，保证不重叠、不遗漏。logical_divide 就是“用这个模具把整个蛋糕切完”，结果是一张“第几块的第几口在哪”的表。蛋糕本身放在什么盘子里（A 的 stride），模具不用关心。

⑥ 用起来：Tensor、local_tile 和线程划分

有了 Layout，CuTe 的 Tensor = 指针 + Layout。在 Tensor 上做 divide，得到的还是 Tensor，所以“切 tile”和“分给线程”都变成了对 Tensor 的操作。CUTLASS 的教程 cutlasssgemm_1.cu 就是 C4 那个 GEMM 的 CuTe 版本：

同一个 GEMM：C4 手写的下标 vs CuTe悬停节点查看连接

左边每一行都是一条手写的整数公式；右边每一行都是一次 layout 运算。CuTe 版本里没有任何 / 和 %：tile 尺寸和线程排布都写在 layout 里（host 端几行 make_layout），换配置只改 layout，kernel 主体不变。

对应的源码行：

构造全局张量：cutlassmake_tensor(make_gmem_ptr(A), shape, stride)
取本 block 的 tile：cutlasslocal_tile(mA, cta_tiler, cta_coord, …)
每个线程搬哪些：cutlasslocal_partition(gA, tA, threadIdx.x)
每个线程算哪些：cutlasslocal_partition(sA, tC, threadIdx.x, Step…)
host 端的配置全是 layout：cutlassbM、bN、bK，sA、sB，tA = 32×8 个线程，tC = 16×16 个线程

线程划分：同一个 tile，换一种线程 layout

local_partition(tile, thr_layout, tid) 的意思是：把 tile 按线程 layout 的形状切开，第 tid 个线程拿走属于它的那些元素。下面的演示把一个 8×8 的 tile 分给几个线程，对比两种排法。kernel 代码完全一样，变的只是一个 layout。

演示 4：把 8×8 的 tile 分给线程

线程 layout：

分法：

每格写的是哪个线程拿它（悬停一个格子，高亮这个线程的全部元素）

真实的 kernel 往往两种都用：搬数据用交错分法（合并访存），算的时候用分块分法（寄存器里复用）。sgemm_1 里搬运用 tA（cutlasslocal_partition(gA, tA, …)），计算用另一个线程 layout tC（cutlasslocal_partition(sA, tC, …)）。两者是两个不同的 layout，同一份数据在两次划分之间通过 shared memory 交接，这正是 C4 里协作加载和寄存器 tile 的关系。

类比 · 工位分配表

一批货（tile）进仓库，要分给几个工人（线程）。“交错分法”像流水线：一排货依次滑过，第 0 个工人拿第 0 件，第 1 个拿第 1 件，大家同时伸手，一次拿走一整排（合并访存）。“分块分法”像包干到人：每人认领一个货架，自己的货都在一起，加工时不用走动（寄存器复用）。CuTe 里换分法只是换一张“工位分配表”（线程 layout），工人干活的步骤（kernel 代码）一个字都不用改。

⑦ 硬件也用 Layout 说话

CuTe 最强的地方在最后一层：GPU 的矩阵指令和拷贝指令，本身就有固定的数据排列要求，CuTe 把这些要求也写成 layout。

指令的要求也是 layout悬停节点查看连接

一条 mma.sync m16n8k16 指令要求 32 个线程的寄存器按一种特定的方式持有 A、B、C 的元素。CuTe 把这个要求写成一个“(线程号, 值号) → (行, 列)”的 layout（TV layout）。换到 Hopper 的 wgmma，只是换一个 Atom，上层用同样的 divide / partition 去对齐它。

TV layout：G1 里你写的“lane 持有哪几个元素”的计算器，就是一个 TV layout。CuTe 用它描述每条 MMA 指令的寄存器排列。
Swizzle：C3 讲过的异或技巧也是一个 layout（cutlassSwizzle），和 shared memory 的 layout 组合起来就消除了 bank conflict，不需要 padding。
文档：cutlass0t_mma_atom.md · cutlass0x_gemm_tutorial.md · cutlass0z_tma_tensors.md

这也是为什么 FlashAttention-3/4、FlashInfer 的 Hopper / Blackwell 后端都建立在 CuTe 上：新一代硬件来了，只需要补上新指令的 Atom（它的 layout），上层的 tiling 和 partition 代码大部分可以复用。CuTe DSL 则把同一套概念搬到了 Python 里。

⑧ 参考解 ↔ CuTe

参考解 drillsg2_cute_layout.py	CuTe C++	pycute（官方 Python 参考）
drillsLayout、call、size、cosize	cutlasslayout.hpp	cutlassclass Layout
drills_idx2crd（列字典序）	`idx2crd`	—
drillscoalesce	cutlasscoalesce	cutlasscoalesce
drillscomposition	cutlasscomposition	cutlasscomposition
drillscomplement	cutlasscomplement	cutlasscomplement
drillslogical_divide	cutlasslogical_divide	cutlasslogical_divide

这一页的四个演示用的就是参考解这六个函数的 JavaScript 移植版，和 Python 版在 55 个例子上逐个对拍过，结果完全一致。

自测 1：(4,4):(1,4) 和 (4,4):(4,1) 合成（前者 ∘ 后者）是什么？意味着什么？

A = 列主序，B 按“行主序”的顺序选下标：R(i) = A(B(i))。B 把下标 i 当作行主序的 (i/4, i%4) 去取 A 的元素，结果是 (4,4):(4,1)。在演示 2 里选 A = (4,4):(1,4)、B = (4,4):(4,1) 就能看到。和第 ⑤ 节的“转置”例子正好对称。

自测 2：logical_divide(32:1, 8:1) 是什么？complement 是什么？

complement(8:1, 32) = 4:8（4 个 tile，起点相差 8）。logical_divide = (8,4):(1,8)：4 个连续的 tile，每个 8 个元素。在演示 3 里选 A = 32:1、B = 8:1 验证。

源码与资料索引

本仓库：drillssolutions/g2_cute_layout.py · drillstests/test_g2_cute_layout.py · drillscontent/g2.md · G2 讲解 · G1 讲解

CuTe 文档：cutlass00_quickstart · cutlass01_layout · cutlass02_layout_algebra · cutlass03_tensor · cutlass0x_gemm_tutorial

代码：cutlasssgemm_1.cu · cutlasspycute/layout.py · cutlasslayout.hpp · cutlassCuTe DSL notebook