学习CUTLASS,很多人一上来就冲进GEMM kernel的源码里,结果被层层模板参数和坐标变换砸得晕头转向。我的建议是先看 03_visualize_layout 这个示例,它不写kernel、不碰shared memory同步,只做一件事——把数据布局(Layout)的逻辑坐标到物理内存偏移的映射关系,用矩阵的形式打印出来。就是这个看起来"没什么技术含量"的小工具,把CUTLASS(以及它底层的CuTe模板库)里最核心、也最影响性能的Layout抽象讲透了。
这个示例适合两类人。一类是准备深入读CUTLASS源码、却总在Shape、Stride、Fragment坐标里迷路的读者;另一类是只想用CUTLASS写高效算子、但想搞清楚"为什么同样一个矩阵,RowMajor和ColumnMajor的性能能差出几十倍"的工程师。我自己的体会是,花一个下午把这份示例的每一行输出看懂,之后再去读 mma、cp.async、ldmatrix 这些真正跑在GPU上的代码,会有一种"地图在手里"的踏实感。
下面我把整个示例按我读源码时的思路拆开讲,包括Layout的本质、主流程拆解、几类典型布局的可视化对照,以及如何把它改造成自己的布局调试工具。
1. 为什么先啃"打印布局"这个示例:它是整个CUTLASS的地基
1.1 它在整套示例里的位置
CUTLASS的examples目录从01排到50多,越往后越复杂。01、02通常是矩阵乘法入门,到了03突然插入一个"可视化布局"的示例,很多人会觉得突兀,甚至直接跳过。实际上这个排序是故意的——03之后的每个示例都要用Layout来表达"数据怎么摆、线程怎么取",如果这一步没建立直觉,后面全是空中楼阁。
这个示例的特殊之处在于:它是纯host端代码,不启动任何kernel。你在跑完编译后执行它,看到的只是一堆数字排列成网格状。但这些数字恰恰是GPU高性能计算里最底层的"地质结构"。
1.2 弄懂它之后你能获得什么
读完这个示例,我认为你至少能收获三样东西:
- 看懂Layout表达式。以后在任何CUTLASS代码里看到
Layout<Shape<_4,_2>, Stride<_2,_1>>,你能立刻在脑子里浮现出一个4行2列的矩阵,以及它在内存里的一维偏移序列。 - 理解坐标变换的方向。逻辑坐标
(row, col)和物理地址之间的映射不是"想当然的",它由Stride精确控制。这个映射决定了访存是否合并、是否发生bank conflict。 - 为读kernel源码做铺垫。CUTLASS 3.x里的GEMM kernel大量使用
cute::make_layout、composition、LogicalToPhysical这些操作,本质上都是在操作03示例里打印出来的那种映射关系。
我自己读这套源码的路径是:先跑03,再跑05(基于CuTe的GEMM),最后才回头看07、08那些带tile scheduler的复杂版本。03这一步省不得。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. Layout的内存映射本质:Shape、Stride与offset公式
2.1 从二维坐标到一维地址
GPU显存本质上是一维的,地址从0递增到某个大数。但我们写矩阵运算时习惯用 (row, col) 这种二维坐标描述元素。Layout就是连接这两套坐标系的桥梁。
一个Layout由两部分组成:
- Shape:逻辑空间的形状,描述"这个布局有多少行、多少列"。
- Stride:步长,描述"逻辑坐标每个维度增加1时,内存偏移增加多少"。
计算某个逻辑坐标对应的内存偏移,公式极其简单:
code复制offset(row, col) = row * stride_row + col * stride_col
以4行2列的矩阵为例,RowMajor布局写作 Layout<Shape<_4,_2>, Stride<_2,_1>>,意思是行方向步长为2、列方向步长为1。于是:
(0,0)-> 0(0,1)-> 1(1,0)-> 2(3,1)-> 7
把整个矩阵打印出来,就是:
text复制0 1
2 3
4 5
6 7
这就是"行优先"的本质——同一行的元素在内存里相邻,换行时要跳过2个元素(其中1个是上一行最后一个、1个是下一行第一个)。
2.2 为什么用模板常量而不是普通整数
CUTLASS/CuTe里到处是 _4、_2、_1 这种带下划线的类型,初见很劝退。它们其实是编译期整数常量,类型是 Int<4>。把它写成模板参数的原因只有一个:让编译器在编译期就完成所有offset计算。
举个例子,`Layout<Shape<4,
