CuTe Layout 到 Tensor Core：从 Shape/Stride 到 Core Matrix

一句话结论

CuTe 最值得建立的 mental model 是：

外围矩阵用 Shape + Stride 描述“逻辑坐标 → 存储位置”；进入 Tensor Core 后，仍然使用 Layout，但此时描述的是“thread/value → MMA 逻辑矩阵坐标”。

也就是说，CuTe 并没有为 Global Memory、Shared Memory、Register Fragment、Tensor Core 分别发明四套坐标系统，而是统一成一个核心抽象：

[ \boxed{\text{Layout} = \text{coordinate mapping}} ]

区别只是不同层级上，Layout 的输入和输出分别代表什么。

这也是理解 CuTe API 最重要的一条主线。

1. 最外围：矩阵首先就是一个 Shape + Stride Layout

CuTe 官方把 Layout 定义成一对 IntTuple：

[ \text{Layout}=(\text{Shape},\text{Stride}). ]

Shape 定义坐标空间，Stride 定义这些坐标如何映射到 index。

例如一个逻辑矩阵

[ A\in\mathbb{R}^{M\times K} ]

如果 K 维连续，那么可以写成：

auto shape_A  = make_shape(M, K);
auto stride_A = make_stride(K, Int<1>{});
auto layout_A = make_layout(shape_A, stride_A);

它表达的是：

[ (m,k)\mapsto mK+k. ]

所以最外围的矩阵，本质上就是：

[ \boxed{(M,K):(K,1)} ]

这样的 Shape/Stride 映射。

如果换成 M 维连续，则是：

[ (M,K):(1,\mathrm{ldA}). ]

CuTe 因此更喜欢直接谈某个 mode 是否 stride-1，而不是先纠结传统 BLAS 的 row-major / column-major / transpose 标签。

2. Tensor = Pointer + Layout

Layout 本身只有“坐标映射”，还没有数据。

真正的数据对象是 Tensor：

Tensor mA = make_tensor(
    make_gmem_ptr(A),
    make_layout(
        make_shape(M, K),
        make_stride(ldA, Int<1>{})
    )
);

可以把它理解成：

[ \boxed{\text{Tensor}=\text{Engine/Pointer}+\text{Layout}} ]

Pointer 回答：

数据在哪里？

Layout 回答：

给我一个逻辑坐标，我应该访问哪个位置？

因此：

mA(m, k)

从概念上相当于：

A[layout_A(m, k)]

这一步已经把“矩阵是什么”和“矩阵怎么存”分开了。

3. 为什么 size 和 cosize 不一样？

因为 Layout 本质是函数。

假设：

[ L:S\rightarrow I ]

那么：

size(L)：domain 中有多少个逻辑元素；
cosize(L)：为了覆盖所有可能 index，底层至少需要多少个存储位置。

对于普通 contiguous matrix，两者往往相同。

但一旦出现 padding、swizzle、hierarchical layout，它们就可能不同。

这也是为什么 Shared Memory buffer 常写成：

__shared__ TA smemA[cosize_v<ASmemLayout>];

而不是简单地写：

BLK_M * BLK_K

因为真正决定“地址能走到哪里”的是 Layout。

4. 从整矩阵到 CTA Tile：local_tile 没有换掉坐标体系

假设 GEMM 使用：

auto bM = Int<128>{};
auto bN = Int<128>{};
auto bK = Int< 32>{};

auto cta_tiler = make_shape(bM, bN, bK);

完整矩阵还是：

[ A:(M,K),\qquad B:(N,K),\qquad C:(M,N). ]

然后用：

auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _);

Tensor gA = local_tile(
    mA, cta_tiler, cta_coord,
    Step<_1, X, _1>{}
);

Tensor gB = local_tile(
    mB, cta_tiler, cta_coord,
    Step<X, _1, _1>{}
);

Tensor gC = local_tile(
    mC, cta_tiler, cta_coord,
    Step<_1, _1, X>{}
);

这里最重要的不是记住 API，而是理解：

local_tile 并没有真的“制造一个新矩阵”，它是在原 Tensor 上应用 tiler，再 slice 出当前 CTA 对应的坐标区域。

例如 gA 最终得到：

[ (\mathrm{BLK_M},\mathrm{BLK_K},k_\text{tiles}) ]

第三维表示 CTA 还要沿 K 方向遍历多少 tile。

所以这一层本质是：

[ \text{Global Matrix} \xrightarrow{\text{tiling}} \text{CTA-local logical view}. ]

5. 到 Shared Memory：逻辑 tile 相同，但物理 Layout 可以完全不同

例如 CTA 需要一个：

[ 128\times32 ]

的 A tile。

Shared Memory 可以定义自己的 layout：

auto sA_layout =
    make_layout(make_shape(bM, bK), LayoutRight{});

__shared__ TA smemA[cosize_v<decltype(sA_layout)>];

Tensor sA =
    make_tensor(make_smem_ptr(smemA), sA_layout);

此时 gA 和 sA 可以代表同一块逻辑数据：

[ A_\text{tile}(m,k), ]

但二者 Layout 完全不同。

gA 的 Layout 继承自 Global Memory 的 leading dimension。

sA 则可以：

K-major；
M-major；
padding；
swizzle；
为 ldmatrix / TMA / WGMMA 特殊组织。

因此 CuTe 的一个核心思想是：

[ \boxed{\text{logical coordinates stay consistent, physical layout may change}} ]

copy 做的正是：

[ \text{same logical coordinate} \quad gA(m,k)\rightarrow sA(m,k). ]

6. 到这里为止，Layout 还是“矩阵坐标 → memory index”

前面的 Layout 都可以理解成：

[ (m,k)\rightarrow \text{memory offset} ]

或者：

[ (n,k)\rightarrow \text{memory offset}. ]

这就是最直观的 Shape/Stride 语义。

但进入 Tensor Core 之后，有一个非常重要的变化：

Layout 仍然是 Layout，但它不一定再描述内存地址。

它开始描述：

[ (\text{thread},\text{value}) \rightarrow (\text{logical matrix coordinate}). ]

这一步就是从“普通 Tensor Layout”进入“MMA Layout”的关键。

7. MMA_Atom：把一条 Tensor Core 指令包装成 Layout 问题

CuTe 对一条 MMA 指令大致分成三层：

[ \text{Operation} \rightarrow \text{MMA_Traits} \rightarrow \text{MMA_Atom}. ]

Operation

Operation 最底层基本就是 PTX MMA 指令 wrapper。

MMA_Traits

MMA_Traits 描述这条指令的所有元信息。

典型字段包括：

Shape_MNK
ThrID
ALayout
BLayout
CLayout

例如官方 Volta 示例：

using Shape_MNK = Shape<_8,_8,_4>;
using ThrID     = SM70_QuadPair;
using ALayout   = SM70_8x4_Col;
using BLayout   = SM70_8x4_Col;
using CLayout   = SM70_8x8_32b;

这几个 Layout 已经不是“大矩阵在 HBM 里的 stride”。

它们回答的是：

这一条硬件 MMA 指令里，哪个 logical thread 的第几个 value，对应 A/B/C 中哪个逻辑元素？

8. ALayout / BLayout / CLayout 真正表达的东西

以 CLayout 为例：

[ (\text{logical thread id},\text{logical value id}) \rightarrow (m,n). ]

CuTe 内部通常会把 ((m,n)) 编码成一个整数 index，所以形式上仍然是普通 Layout。

但语义已经变成：

[ (T,V)\rightarrow C(m,n). ]

例如某个 Tensor Core 指令可能规定：

T0 的 V0 对应 C(0,0)；
T0 的 V1 对应 C(0,1)；
T1 的 V0 对应 C(1,0)；
某些更远的 value 又跳到另外一块矩阵区域。

这种 mapping 如果手写 CUDA，就会表现为大量：

lane_id；
lane_id / 4；
lane_id % 4；
register index；
xor / permutation；
特定架构硬编码。

CuTe 的做法是：

[ \boxed{\text{把 ISA 的 fragment ownership 直接写成 Layout}} ]

于是线程拓扑也进入了 Layout algebra。

9. Hopper 文档里的 Core Matrix，正好对应这个 mental model

Hopper GMMA/WGMMA 的官方 CuTe 文档明确从一个 core matrix 开始描述 accumulator mapping。

以 FP16 accumulator 为例，可以先考虑一个较小的 core matrix。

然后完整 GMMA tile 通过：

沿 M 方向堆叠 core matrix；
再沿 N 方向复制这些 core matrix；

最终得到完整的：

[ M\times N ]

accumulator tile。

因此可以把它理解成：

[ \boxed{ \text{Core Matrix} \rightarrow \text{MMA Atom internal layout} \rightarrow \text{full GMMA tile} } ]

这里的 Core Matrix 不是一个新的 CuTe 类型。

它更像是：

Tensor Core / GMMA 硬件规定的最小规则矩阵拓扑，而 CuTe 用 hierarchical Layout 把它编码出来。

这正是为什么 Hopper 的 CLayout 会长成层次化 Shape/Stride：

Layout<
  Shape <
    Shape < ... >,
    Shape < ... >
  >,
  Stride<
    Stride< ... >,
    Stride< ... >
  >
>

层次不是为了炫技，而是在直接表达：

[ \text{core matrix} \times \text{M repetition} \times \text{N repetition} \times \text{warp / warpgroup ownership}. ]

10. 一个特别重要的区分：外层 Layout 和 MMA Layout 不是同一种“物理含义”

因此最好不要简单说：

CuTe Layout 就是矩阵的内存 stride。

更准确的是：

CuTe Layout 是一个通用 coordinate mapping。

在外围矩阵上：

[ (M,K):(d_M,d_K) ]

表示：

[ (m,k)\rightarrow \text{memory index}. ]

而在 MMA_Traits 中：

[ (T,V):(\cdots) ]

表达：

[ (thread,value)\rightarrow \text{encoded }(m,k)/(n,k)/(m,n). ]

所以同一个 Layout vocabulary 被复用在两个完全不同的层级。

这恰恰是 CuTe 最漂亮的地方。

11. TiledMMA：把一个 MMA Atom 在更大的 Tile 上复制、交织

一个 MMA_Atom 只描述一条基本硬件 MMA。

如果想构造更大的 MMA tile，就使用：

make_tiled_mma(...)

例如官方示例：

TiledMMA mma = make_tiled_mma(
    SM70_8x8x4_F32F16F16F32_NT{},
    Layout<
        Shape <_2,_2>,
        Stride<_2,_1>
    >{}
);

原本一个 Atom 是：

[ 8\times8\times4. ]

通过一个：

[ (2,2):(2,1) ]

的 Atom Layout 复制之后，可以形成更大的：

[ 16\times16\times4 ]

TiledMMA。

所以 TiledMMA 可以理解为：

[ \boxed{ \text{Atom hardware topology} + \text{Atom replication layout} + \text{M/N/K tiler} } ]

它不是简单地把矩阵切块，而是在扩展“谁负责哪些 fragment”的硬件执行拓扑。

12. get_slice：从整个 TiledMMA 中取出“我这个执行者”的视角

当 TiledMMA 确定之后：

ThrMMA thr_mma = mma.get_slice(threadIdx.x);

可以理解为：

给定当前 thread id，告诉我在整个 TiledMMA 中我负责哪部分 A/B/C。

对于 Warp MMA，一般是 thread slice。

对于 Hopper WGMMA，参与者是 128-thread warpgroup，因此更准确的执行者视角会围绕 warpgroup 构造。

但抽象不变：

[ \text{TiledMMA} \xrightarrow{\text{get_slice}} \text{ThrMMA / participant-local view}. ]

13. partition_A/B/C：把普通矩阵 Tile 对齐到 Tensor Core 的 ownership

有了 ThrMMA，就可以：

Tensor tCsA = thr_mma.partition_A(sA);
Tensor tCsB = thr_mma.partition_B(sB);
Tensor tCgC = thr_mma.partition_C(gC);

典型 shape 会类似：

tCsA: (MMA, MMA_M, MMA_K)
tCsB: (MMA, MMA_N, MMA_K)
tCgC: (MMA, MMA_M, MMA_N)

这里最值得理解的是：

partition_A/B/C 不是随便给 thread 分数据，而是根据 TiledMMA 内部已经定义好的 ALayout/BLayout/CLayout 去分。

也就是说：

[ \text{Matrix Tile} \xrightarrow{\text{MMA thread-value layout}} \text{Tensor Core compatible partition}. ]

所以它完成的是“普通矩阵逻辑坐标”到“Tensor Core fragment ownership”的桥接。

14. 从 Shared Memory 到 Register Fragment

对传统 register-source MMA，可以进一步得到每个线程真正持有的 register fragment。

最简单的 accumulator 可以写成：

Tensor tCrC = make_tensor_like(tCgC);

它在逻辑 shape 上和 partition_C 一致，但 storage engine 变成线程私有 register。

于是整个过程变成：

[ C_\text{CTA tile} \rightarrow C_\text{MMA partition} \rightarrow C_\text{register fragment}. ]

现代 CuTe DSL 还会显式使用：

tiled_mma.make_fragment_A(...)
tiled_mma.make_fragment_B(...)
tiled_mma.make_fragment_C(...)

在 Hopper/Blackwell 上，fragment 甚至可能不只是普通 RMEM：

A/B 可以来自 Shared Memory descriptor；
C accumulator 可以位于 Register；
Blackwell 某些路径则会进入 TMEM。

但逻辑 shape 仍然由 TiledMMA 决定。

15. 最后 cute::gemm 为什么可以看起来这么简单？

最后计算可能只有一句：

gemm(tCsA, tCsB, tCrC);

或者显式携带 tiled MMA：

gemm(tiled_mma, ...);

之所以 API 能这么简单，不是因为底层简单，而是因为前面已经通过 Layout 把所有约束编码进 Tensor 了：

A 哪些元素属于哪个 thread；
B 哪些元素属于哪个 thread；
C 哪些 accumulator 属于哪个 thread；
每个 thread 中 value 的顺序；
Atom 如何重复；
M/N/K 如何 tile。

所以真正复杂的东西不在 gemm 调用本身，而在：

[ \boxed{\text{Layout construction + partitioning}} ]

这也是 CuTe 的设计哲学。

16. 把整条链串起来

一个典型 GEMM 可以画成：

Global Matrix
  │
  │ make_layout(shape, stride)
  ▼
Global Tensor
  │
  │ local_tile
  ▼
CTA Tile
  │
  │ copy / TiledCopy
  ▼
Shared Memory Tensor
  │
  │ TiledMMA.get_slice
  │ partition_A/B/C
  ▼
MMA-partitioned Tensor
  │
  │ fragment / register ownership
  ▼
Tensor Core Core-Matrix / Fragment Topology
  │
  │ cute::gemm
  ▼
Accumulator

对应两个最重要的映射：

外围

[ (m,k)\rightarrow \text{memory offset} ]

Tensor Core 内部

[ (thread,value)\rightarrow(m,k)/(n,k)/(m,n) ]

二者都叫 Layout。

17. 为什么这套抽象对 FlashAttention 特别重要？

FlashAttention 里经常出现：

[ QK^T \rightarrow P \rightarrow PV. ]

第一发 MMA 的 accumulator 可能成为 softmax 的输入。

经过 softmax 后，P 又要成为第二发 MMA 的 operand。

于是问题变成：

上一发 MMA 产生的 register ownership，能不能低成本重解释成下一发 MMA 所需要的 operand ownership？

如果你把 fragment 当成“乱七八糟的一堆 lane/register”，这个问题会非常难处理。

但用 CuTe 的视角，它变成：

[ \text{Layout}_1 \rightarrow \text{layout algebra / retile} \rightarrow \text{Layout}_2. ]

很多时候并不是“搬动整个矩阵”，而是在重新建立逻辑坐标与 thread-value 的关系。

这就是 CuTe 对 FA / GEMM kernel 很有价值的地方。

18. 面试里最推荐的说法

可以用下面这段：

CuTe 的 Layout 不是单纯的 memory layout，而是一个统一的 coordinate mapping abstraction。对于 Global/Shared Memory Tensor，它通常用 Shape+Stride 描述逻辑矩阵坐标到物理 offset 的映射；经过 local_tile 后，我们得到 CTA-local view。进入 Tensor Core 层以后，MMA_Traits 仍然用 Layout，但 ALayout/BLayout/CLayout 的语义变成 thread-value 到 MMA 逻辑矩阵坐标的映射。MMA_Atom 编码一条硬件 MMA 的最小 fragment topology，Hopper 文档里甚至直接从 core matrix 出发构造完整 GMMA CLayout；TiledMMA 再把 Atom 沿 M/N/K 复制和交织，最后通过 get_slice 和 partition_A/B/C 把普通矩阵 tile 投影成每个 thread/warpgroup 应该持有的 fragment。

一句话压缩就是：

[ \boxed{ \text{CuTe Layout 统一了 Memory Geometry 与 Tensor Core Execution Geometry} } ]

19. 最容易混淆的三个点

1. Layout 不等于 memory layout

Memory layout 只是 Layout 的一种用途。

MMA thread-value mapping 同样是 Layout。

2. partition_A/B/C 不等于普通 thread tiling

它不是单纯：

[ \text{tile}/\text{num threads}. ]

它受到硬件 MMA Atom 的 fragment ownership 约束。

3. Core Matrix 不是 CuTe 的一个类

Core Matrix 是理解 Hopper GMMA accumulator topology 的硬件/几何概念。

CuTe 真正提供的是：

Layout；
MMA_Traits；
MMA_Atom；
TiledMMA；
ThrMMA；
partition / fragment API。

Core Matrix 的规则最终被编码进这些 Layout 中。

20. 最终 mental model

看到下面这些 API：

make_layout(...)
make_tensor(...)
local_tile(...)
make_tiled_mma(...)
get_slice(...)
partition_A(...)
partition_B(...)
partition_C(...)
gemm(...)

不要把它们记成一堆独立接口。

应该读成：

[ \boxed{ \text{定义坐标空间} \rightarrow \text{切出局部坐标空间} \rightarrow \text{把坐标空间映射到线程/值} \rightarrow \text{执行硬件 MMA} } ]

所以 CuTe 真正做的事情不是“帮你写 GEMM”。

它是在提供一套可以一直从：

[ \text{HBM Matrix} ]

推导到：

[ \text{Tensor Core Core Matrix} ]

的统一坐标代数。

Reference

NVIDIA CUTLASS Documentation, CuTe Layouts
https://docs.nvidia.com/cutlass/latest/media/docs/cpp/cute/01_layout.html
NVIDIA CUTLASS Documentation, CuTe dense matrix-matrix multiply tutorial
https://docs.nvidia.com/cutlass/latest/media/docs/cpp/cute/0x_gemm_tutorial.html
NVIDIA CUTLASS Documentation, CuTe’s support for Matrix Multiply-Accumulate instructions
https://docs.nvidia.com/cutlass/latest/media/docs/cpp/cute/0t_mma_atom.html
NVIDIA CUTLASS Documentation, CUTLASS 3.x GEMM API
https://docs.nvidia.com/cutlass/latest/media/docs/cpp/gemm_api_3x.html