OK那么现在也是在华子准备过夜了，今晚的任务就是学习一下TMA，刚好这边有分配的超大屏😋，那就他奶奶的把好东西端上来吧！

啊哈哈哈——鸡汤来咯！

骗你的，其实是TMA😋

本篇基本参考知乎uu的解析 (opens new window)，并且主要用于自我学习，在这里感谢他的贡献

拷贝的前传

在最初的最初，GPU都是使用线程进行搬运；可以在Global memory（即平常说的HBM：整个GPU的所有线程都可以访问的存储，容量最大但速度也最慢）与Shared memory（GPU 芯片上的 SM 内部 SRAM，几十到几百KB大小）之间搬运数据；比如说计算需要用到某些数据的时候，就需要把数据从Global memory搬运到Shared memory，然后进行计算，计算完成后，再从Shared memory搬运回Global memory，而在拷贝的过程中，事实上还需要经过寄存器这个中间buffer（比如说，从Share memory -> 寄存器 -> Global memory）；那么在基模实习过若干次的朋友都知道，整个计算kernel的延迟不仅取决于计算的速度，还取决于访存的速度，因此优化拷贝的靠小可以优化kernel的性能；

上述所说的拷贝是同步拷贝，顾名思义，在进行拷贝的时候，指令是被阻塞的，无法继续执行其他指令，因此线程只能空转；事实上除了这个缺点，还存在这样几个问题

1. 每次拷贝都需要经过寄存器这个中间buffer，因此会占用寄存器的空间
2. 一次复制涉及许多指令的操作，包括：地址计算，LDG（加载），STS（存储）指令；GPU是L/S架构的，那么其指令的操作数就必须有一个是寄存器；
3. 若干无数据依赖的指令不得不串行执行；显然，这是其同步的特性造成的
4. 可能导致线程空转；比如16个线程去搬运18个数据，有两个线程被分配了搬运两个数据的任务，而剩下的都是只搬运一个数据，这种情况下其他线程需要等待那两个线程完成，不得不空转从而浪费GPU；

针对以上几点，我们可以思考一些对策

1. 首先是减少指令的数量，比如合并一些操作为一条指令——将LDG和STS指令合并为一条
2. 然后显然是减少不必要的串行等待，如果可以做成异步的搬运，那可以极大提高并发程度，提高带宽利用率————TMA
3. 避免大数据量下的循环展开指令（因为每次拷贝都要经过寄存器，而寄存器的大小非常之小，因此每次搬运的数据量也很小，如果整体要搬运很大的数据块的话，需要循环搬运很多次，从而会发射大量的指令）如果说，针对任何数据量（假设是连续的）我们都可以使用一条指令搞定，那极大节省了指令开销——TMA

于是乎，它来了！

Ampere架构——异步拷贝

该特性支持异步从 Global Memory 拷贝到 Shared Memory（不支持反方向的传输，反方向依然只能用STS），可以在拷贝期间执行其他计算操作（除了读取Shared Memory相关的计算），这样实现拷贝和计算重叠，提升线程执行的性能；同时，该特性还减少了寄存器的占用，进一步释放资源，提高算力(如下图)

另外，它也将指令进行合并，将 LDG 和 STS 指令合并为一种新指令，减少指令数了一定的指令数，并且更有利于做计算与拷贝的交叠；

由于异步拷贝肯定也存在需要同步等待的场景，比如预取某数据，但是在计算之前必须确保数据已经准备好的场景；因此异步拷贝的使用方式大致为：

1. 获取当前 block；
2. 在 Shared Memory 上创建 cuda::barrier 对象，并使用一个线程对其初始化；
3. 调用 cuda::memcpy_async，将上一步创建的 cuda::barrier 对象放到 barrier 参数中；
4. 使用 barrier 同步异步拷贝

Hopper架构——TMA

前面的Ampere架构虽然以及极大提高了效率，不过可以想到他主要是在指令结构上进行优化，指令本身的行为还是大差不大的，因此其实还是遗留了几个问题

1. 对于大数据量下，依然要循环发送很多指令，每个指令会搬运一块比较小的显存块，并且每个指令也还要计算显存的起始地址（其实思考一下可以发现，这个计算甚至是多余的，因为如果我们需要搬运一块连续的内存，只需要知道内存起点和长度即可，为什么要反复计算这么多次分段的起始地址呢？）
2. 进行地址计算时实际上也占用了一定的寄存器，因此潜在的算力其实还可以释放

于是，我们迎来了TMA！

TMA支持这几个功能

1. 大块的异步拷贝——这次只需要一次描述就可以进行一整块搬运，其中地址计算直接由硬件/descriptor 处理
2. 多维拷贝，减少复杂地址计算；Hopper TMA 的做法是：把多维 tensor 的布局信息提前编码成一个 descriptor，也就是常说的 CUtensorMap；然后 kernel 里发起 TMA 拷贝时，不再为每个小片段手动算地址，而是把这个descriptor提交给硬件，硬件根据 descriptor 自动完成地址生成。其中descriptor 包含类似信息：

• tensor 维度
• 每一维大小
• 每一维 stride
• 元素大小
• global memory 基地址
• swizzle / interleave 等布局信息

3. 极大解放线程：只需要一个 warp 或少量线程发起 TMA 请求；而TMA 硬件异步完成大块搬运，其他 warp 可以继续计算
4. 支持 Shared Memory → Global Memory

TMA的工作流程：

1. 初始化 barrier
2. 发起 TMA copy，并绑定 barrier
3. TMA 完成后更新 barrier
4. consumer 等待 barrier
5. consumer 使用 shared memory 数据

该同步的工作流程的优势在于

1. 数据到达通知更精细；
2. 更适合 warp-group 分工；
3. 更容易和 WGMMA pipeline 配合；
4. 计算和搬运重叠更自然；
5. 减少全 CTA 粗粒度同步的需求。

TMA 的其他特性（单卡 / Cluster 内）

上面说的多是「本 SM 的 shared memory ↔ 本 GPU 的 global memory」。Hopper 上 TMA 还有几类很实用的扩展，和 Thread Block Cluster、Distributed Shared Memory（DSMEM） 绑在一起：

1. Cluster 内 Multicast（组播加载）
   同一个 Thread Block Cluster 里的多个 CTA 常常需要同一块 global memory 数据（典型如 GEMM 里共享的 A 矩阵条带）。TMA 可以 只从 global memory 读一次，硬件把数据 同时写入 cluster 内多个 CTA 的 shared memory，省 global memory 带宽。CUTLASS / CuTe 里常叫 TMA multicast load。

2. Cluster 内 DSMEM 搬运
   除 global ↔ shared 外，TMA 还支持在 同一 cluster 内不同 SM 的 shared memory 区域之间 做 bulk 搬运。配合 DSMEM，cluster 里的 block 可以互相读写的 shared memory 范围更大，适合「单 block 放不下、又不想立刻打回 global」的中间态。

3. 写回 global 时的归约（Reduction）
   从 shared memory 写回 global memory 时，TMA 可在硬件里附带 逐元素归约：常见有 add / min / max，以及部分类型的 bitwise and / or。适合多 producer 往同一块 global 输出做累加类写回，而不必每个线程自己 atomic。

4. Descriptor 里的布局信息
   除维度、stride、基地址外，descriptor（CUtensorMap / tensor map）还可编码 swizzle、interleave 等布局，让硬件按 Tensor Core / WGMMA 喜欢的排布把数据搬进 shared memory，减少后续重排。

5. 与 mbarrier / pipeline 的深度绑定
   TMA 是 mbarrier（异步 barrier） 驱动的：发起拷贝时绑定 barrier，硬件完成后 自动 arrive / 更新 phase；consumer warp 在 barrier 上 wait 即可。这也是 CUTLASS 3.x pipeline、WGMMA mainloop 能「TMA 预取 + 计算」流水起来的基础。

区分两个「multicast」：

• Cluster multicast：单 GPU 内、cluster 内多 CTA 共享一次 global 读；
• NVSwitch multicast（下文）：多 GPU 间通过 multicast 地址做广播——名字相近，层次完全不同。

Symmetric Memory 与跨 rank TMA

单卡 TMA 解决的是 SM 内部 / 单 GPU HBM 的搬运效率。多卡训练、推理里，还希望 在 kernel 里直接拉 peer GPU 的 HBM，并把 NVLink 通信和计算融在一起——这就需要 Symmetric Memory（对称内存）。

对称内存是什么

现代 NVLink 互联下，peer GPU 的 HBM 可以通过虚拟内存管理（VMM）映射进本 GPU 的虚拟地址空间。各 rank 各自分配一块 同尺寸 的 buffer，经 collective rendezvous（交换 handle、完成 map）后，每个 rank 都拿到：

• 本 rank 本地 buffer 的指针；
• 各 peer rank buffer 在本 GPU VA 空间里的指针（读 peer 就像读「远端的 global memory」）；
• （若启用 NVSwitch multicast）multicast 地址：一次写可广播到参与组内的所有 GPU。

PyTorch 侧可用 torch.distributed._symmetric_memory：symm_mem.empty() 分配 → rendezvous() 建联 → handle 提供 buffer_ptrs、multicast_ptr 等；CUDA 侧则对应 cuMemCreate / cuMemMap、multicast object（cuMulticastCreate 等）那一套。

跨 rank 时 TMA 怎么用

关键观察：TMA 的 tensor map 描述的是「global 侧源/目的」的布局，并不限定这块 global memory 一定在本卡 HBM——只要地址在本 GPU 的 VA 里合法可访问，peer 映射进来的地址也可以作为 TMA 的 global 端。

因此典型流程是：

1. Host / 启动前：为本地张量、各 peer 张量（及可选 multicast 区域）分别建好 TMA descriptor（或 kernel 参数里传 peer 指针 + 同一套 layout 元数据）；
2. Kernel 内：用 少量线程（甚至单线程） 发起 cp.async.bulk.tensor / cuda::memcpy_async 一类 TMA 操作：
   • P2P load：global = peer_rank_k 的映射地址 → shared（把别的卡上的激活/权重拉进本 SM）；
   • P2P store / multicast store：shared → peer 或 multicast 地址（向对端或全体广播一块 tile）；
3. 同步：仍用 mbarrier 等待 TMA 完成；跨 rank 时还要叠加 signal pad / symm_mem 提供的跨 GPU 同步原语（或 NVSHMEM 的 put/get + quiet），保证对端 buffer 已就绪或写已完成。

这和传统 NCCL all-gather + 单独 compute kernel 相比，优势是 通信粒度可细到 tile 级，且 TMA 异步、单线程发起，便于做 「算当前 tile 的同时，TMA 预取下一 tile / 上一 tile 的 reduce-scatter」 这类融合。Hazy Research 的 PGL（Parallel Global Layout）就是把 peer / multicast 指针和 TMA descriptor 封进同一套 layout，让 tma::load_async(tile, pgl[peer_idx], …) 和单卡写法几乎一致。

使用时的注意点

1. 带宽与缓存：peer / multicast 访问走 NVLink，不会像本卡 HBM 那样享受本地 L2 的完整缓存行为；带宽也远低于 HBM。跨 rank TMA 适合 大块、连续、对齐 的 tensor tile，细粒度、跨 warp 不规整访问会很亏。
2. 谁来做通信：实测往往 8–16 个 SM 专门发 TMA / NVLink 就能接近饱和带宽，其余 SM 专心算力活——即 SM 专业化（communication SM vs compute SM），比「全 SM 又算又通信」更稳。
3. 归约 vs 广播：TMA + copy engine 对 multicast 写（广播） 支持较好；NVSwitch 上的 in-fabric reduction 有时需配合 multimem.red / multimem.ld_reduce 等寄存器级指令，而不是单靠 TMA store。
4. 跨节点：单节点 NVLink domain 内用 symmetric memory 最直接；跨机 通常还要 NVSHMEM / rocSHMEM 等（PyTorch symm_mem 也有 NVSHMEM Triton 插件），语义上仍是「把远端内存纳入寻址 + kernel 内发起」，但底层走 RDMA 而非 NVLink。

一句话：Symmetric Memory 解决「多卡地址怎么在同一个 kernel 里可见」；TMA 解决「看见了之后怎么高效、异步、大块地搬」——两者叠加，才是跨 rank 融合算子里「一条 kernel 里既通信又计算」的常见打法。