B-08. Hopper TMA:从单线程 Bulk Copy 到吞吐墙

B-07 解决的是 Ampere 路径:cp.async/cuda::pipeline,用多 stage 藏 GMEM 延迟。
当 tile 更大、地址更不规则、算地址 + 发一堆 copy 本身吃指令带宽时,下一站是把搬运卸给专用引擎。
本章进入 Hopper+TMA:单线程 issue、descriptor / tensor map、mbarrier交接——以及什么时候固定开销反而让你变慢


TL;DR(工程结论)

  1. TMA 是专用异步拷贝引擎(GMEM↔SMEM;亦可 cluster DSM):与 B-07 的 SM 内cp.async、B-06 的 Host CE不是一层
  2. 收益来自单线程 issue + 引擎搬运 ∥ compute,不是「换 TMA 指令本身更快」。本机(RTX 5090):立刻 wait 的bulk1d/tensor2d0.86~1.05×;2-stagepipe2低 AI 可达约1.69×
  3. 该上:需要大 tile / 多维搬运、或要把 copy 从计算 warp 卸掉(warp-spec);上之前先用sweep证明pipe2段 >1。
  4. 别上 / 先别上:只换引擎立刻 wait、对齐搞不定、已 compute-bound(本机fma≥128pipe2已回落到 ~1.03×)、B-07 pipeline 已够用。
  5. 同步模型:G2S 用mbarrier+expect_tx/arrive_tx;S2G 常走bulk async-group——两套完成模型不要混用。

1. 问题:B-07 之后还卡在哪

B-07 的四层梯子里,第 ④ 层只留了钩子:

① sync ② pipeline_memcpy ③ cuda::pipeline ④ TMA bulk gmem→reg→smem Ampere 异步拷贝 多 stage 流水线 专用引擎 ▲ 本章

Amperecp.async已经旁路寄存器,但每个线程仍要算自己的地址并发指令。tile 变大、维度变多时:地址/predication 吃指令带宽、挤寄存器;warp-spec 时 copy warp「发令」过贵。

TMA 的回答:一个线程把 descriptor + 坐标交给硬件引擎,其余线程去算(或等 barrier)。

左:每线程各自算地址、发一串小拷贝。右:elect 线程发一条控制;TMA 引擎整块搬 tile。

层级API / 硬件谁算地址典型证据
B-06 Host↔DeviceCE + Stream驱动/CENSYS
B-07 Devicecp.asyncLDGSTS / pipeline每线程intensity 曲线
B-08 Device TMAbulk / tensor map描述符 + 引擎intensity 曲线(本章)

2. 物理模型:引擎、描述符、mbarrier

┌─────────────┐ elect thread ───►│ 提交命令 │ (tensor_map, coords) 或 (ptr, bytes) │ │ 一条 issue │ │ └──────┬──────┘ │ │ 异步 │ ▼ GMEM ════════════════► TMA engine ════════════════► SMEM tile (HBM) │ (可旁路 L1) │ complete_tx(按字节记账) ▼ mbarrier ──wait──► 全体线程开始读 SMEM / 计算

时间线上,两种用法差很多:

上:bulk1d/tensor2d立刻 wait → copy 与 compute 串行(本机 ≈ 0.9~1.0×)。
下:pipe2算当前 tile 时预取下一块 → 才能藏延迟(本机低 AI ≈ 1.7×)。

工程含义:

  • 文献(Luo et al.):完整路径可有约+170 cycle固定开销。
  • 本机:4 KiB tile 立刻 wait 几乎不赚;重叠才赚。
  • 对齐不满足时memcpy_async_tx/cp.async.bulkUB——对齐必须写死。

3. API 分层与同步

心里记四层(由浅到深):

① 1D cp.async.bulk / cuda::device::memcpy_async_tx ← 指针+长度,无 tensor map ② cuTensorMapEncode* + cp.async.bulk.tensor ← 多维 tile + 坐标 ③ cuda::memcpy_async(Hopper+ 满足对齐时也可走 TMA)← 有回退;对照实验慎用 ④ CuTe / CUTLASS PipelineTmaAsync ← 生产 GEMM,扩展阅读

3.1 1D bulk:指针 + 字节数

  • GMEM / SMEM:16B 对齐;size:16 的倍数
  • ptx::elect_sync(或invoke_one)选举 issue 线程,避免if (threadIdx.x==0)被插 peeling loop。
  • memcpy_async_tx总是走 TMA;tx 字节用barrier_arrive_tx/expect_tx显式声明。

3.2 2D tensor map:主机编码 +__grid_constant__

  1. Host:cuTensorMapEncodeTiled(可用cudaGetDriverEntryPointByVersion)。
  2. Kernel:const __grid_constant__ CUtensorMap
  3. Device:cp.async.bulk.tensor+ 坐标;目标 SMEM 128B 对齐
  4. OOB 区域 G2S 零填充——硬件行为,不是 bug。

3.3 G2S vs S2G 完成模型

本章 MVP只做 G2S。S2G 常见于 GEMM epilogue,完成模型不同:

方向完成跟踪
GMEM→SMEM(本章)mbarrier+ tx 字节
SMEM→GMEM(扩展)cp.async.bulk.commit_group+wait_group_read

S2G 写回前:fence_proxy_async__syncthreads(),否则引擎可能读到旧 SMEM。

3.4 和 B-07 pipeline 的关系

本章 mode对应 B-07时间线(见图)本机结果
bulk1d/tensor2dasync1立刻 wait,串行≈ 0.9~1.0×
pipe2pipe2prefetch ∥ compute低 AI 最高约 1.69×

生产级还会拆 producer warp(只发 TMA)/ consumer warpgroup(WGMMA)——本章不实现,见 §7。


4. 决策表:何时上 TMA、何时回退 B-07

信号建议
需要大 tile / 多维 / 卸掉 copy 指令压力试 TMA:先看pipe2sweep,再考虑嵌算子
bulk1d/tensor2d≤1,但pipe2>1(本机典型)保留 prefetch;不要只换 API 立刻 wait
已 compute-bound(本机约fma≥128停;TMA 协议变纯开销
B-07sweep已够、且无多维/指令带宽墙先留在 B-07
对齐搞不定回退 sync / B-07;勿硬上memcpy_async_tx(UB)
cluster 多播同一份 K/V概念上 TMA multicast(B-01);本章不做
冲 GEMM/Attention SOLModule D / CUTLASS / FA3;本 micro-bench 不复现库性能

5. 实验怎么设计

配套代码:examples/02_memory_optim/08_tma_intro.cu
主命令(一条就够):

./bin/02_memory_optim_08_tma_intro--modesweep
mode回答什么
sync协作 sync load 整 tile(公平基线)
bulk1d1D TMA + 立刻 wait
tensor2d2D tensor-map + 立刻 wait
pipe22-stage TMA prefetch
sweepfma-iters→ 加速比 vs intensity(主证据

Tile 固定1024 floats(32×32,4 KiB)n须整除 tile。sync整 tile 协作 load(不是 B-07 的 per-thread 单元素),才能和 TMA bulk 公平对照。

sweep输出 CSV:

fma_iters,sync_ms,bulk1d_ms,tensor2d_ms,pipe2_ms,speedup_*,...

证据优先级:① 裸跑sweepdocs/results/B-08_*.csv;② 可选 NCU / SASS。
不要ncu附着时读程序自己打印的 ms。需要sm_90+;Blackwell 建议-DCMAKE_CUDA_ARCHITECTURES=120。增删.cu后请重新cmake

5.1 预期形状(本机已验证)

speedup (sync / mode) ▲ │ pipe2 ╭── 低 AI:明显 >1 │ ╱ │──────╱──────── 1.0 │ ╱ bulk1d ≈ 1;tensor2d 常 <1 │ ╱ 高 AI:全体 → 1 └──────────────────────► fma-iters

5.2 RTX 5090 实测(裸跑)

  • GPU:RTX 5090,sm_120sharedMemPerBlock=48 KB
  • 载荷:n=4194304tiles/block=64block=256;2D 视图2048×2048
  • 口径:CUDA eventmedian
  • 明细:docs/results/B-08_tma.md

Intensity sweep

fma_iterssync_msbulk1dtensor2dpipe2sp_bulk1dsp_tensor2dsp_pipe2
10.03680.03500.04040.02181.052×0.913×1.688×
20.03620.03730.04210.02380.970×0.859×1.521×
40.04090.04290.04750.02880.952×0.860×1.419×
80.05050.05090.05660.04020.992×0.893×1.258×
160.06950.07100.07540.05880.979×0.922×1.183×
320.10740.10950.11470.09640.981×0.937×1.114×
640.18360.18520.19030.17320.991×0.964×1.060×
1280.33630.33760.34320.32540.996×0.980×1.034×
2560.64180.64330.64760.62970.998×0.991×1.019×

数据源:docs/results/B-08_sweep.csv;重画:python scripts/plot_b08_tma.py

固定 mode(fma_iters=8,取自 sweep 同行)

modemedian (ms)相对 sync一句话
sync0.05051.00×协作 sync load 基线
bulk1d0.05090.99×换引擎立刻 wait ≈ 打平
tensor2d0.05660.89×2D 立刻 wait 更慢
pipe20.04021.26×prefetch 才赚

数据源:docs/results/B-08_modes.csv

怎么读

  1. 藏延迟靠 overlappipe2fma=11.69×fma=81.26×
  2. 换引擎 ≠ 加速bulk1d/tensor2d立刻 wait 全程约0.86~1.05×
  3. 高 AI 回落fma=256pipe2≈1.02×——算力淹没协议开销,与 B-07 同形。

6. 工程边界

6.1 对齐与 UB

路径关键对齐
1D bulkGMEM/SMEM 16B;size ×16
2D tensorGMEM 16B;stride ×16;SMEM 128B
memcpy_async可不满足时回退;不能当「一定是 TMA」的证据

6.2 落地后仍要管 bank / swizzle

TMA 只负责「把字节放到 shared」。描述符 swizzle(32B/64B/128B)是给 WGMMA 友好布局用的,不是bank conflict 免死金牌(见 B-02)。

6.3 Cluster multicast / DSM

B-01 提过 TMA multicast。本章 MVP不做;需要时进 Cluster / Module D。

6.4 Blackwell 仍保留 cp.async

B-07 在消费级新卡上仍然有用;TMA 是另一条梯子,不是「有 Blackwell 就必须重写一切」。


7. 扩展阅读(不抢 Module D)

  • Colfax, Mastering the Hopper TMA
  • Colfax / CUTLASS:warp-specialized +PipelineTmaAsync
  • Shah et al.,FlashAttention-3(NeurIPS’24 / arXiv:2407.08608)
  • ACTA(GPGPU’25):自动选 tile/queue
  • Yadav et al.,Cypress(PLDI’25 / arXiv:2504.07004)
  • PyTorch, Hopper TMA for FP8 GEMMs(descriptor 开销反例)
  • Luo et al., arXiv:2501.12084(TMA 延迟/吞吐 microbench)

8. 工程 SOP 与常见误区

建议流程

  1. 确认 sm_90+
  2. 直接跑./bin/02_memory_optim_08_tma_intro --mode sweep
  3. speedup_pipe2:低 AI >1、高 AI →1 → 标题成立
  4. 若只有bulk1d/tensor2d≤1、没有 overlap 计划 →不要为 TMA 而 TMA
  5. 有收益再嵌真实算子;需要布局 → B-09

判停

  • pipe2低 AI 明显 >1(本机约 1.3~1.7×)→ 值得在 latency-bound 路径保留 prefetch
  • 全程只有立刻 wait ≤1 → 回退 B-07,或加大 tile / 上真 overlap 后再测
  • tensor2d异常慢 → 先查 encode / 坐标 / 128B 对齐

高频误区

  1. 把 HostcudaMemcpyAsync、B-07memcpy_async、TMA bulk 当成同一件事
  2. 写了 TMA 却立刻 wait,又期待「自动变快」(本机已打脸)
  3. threadIdx.x==0issue 却不 elect
  4. G2S / S2G 完成模型混用
  5. 不对齐仍调memcpy_async_tx(UB)
  6. 每次 launch 重建/重传巨大 descriptor
  7. 以为上了 TMA 就不用管 bank / swizzle / 布局

9. 小结与下一章

三句话:

  1. 先证明你需要 overlap(或指令带宽墙),再上 TMA——本机pipe2极低 AI 约1.69×
  2. 赚的是引擎搬运 ∥ compute,不是指令名字(立刻 wait 的bulk1d≈1.0×tensor2d≈0.89×);
  3. sweep判停——高 AI 回落到 ~1 是正常物理结果。

下一章B-09回到数据布局:AoS/SoA/Transpose。TMA 再强,也救不了从第一天就选错的布局。


10. 参考文献

官方文档

  1. CUDA Programming Guide — Asynchronous Data Copies
  2. NVIDIA Hopper Architecture In-Depth
  3. CCCL — Tensor Memory Accelerator (TMA)
  4. NVIDIA Hopper Tuning Guide

工程教程

  1. Colfax, Mastering the Hopper TMA
  2. MLC.ai, Pipelining GEMM with TMA
  3. PyTorch, Deep Dive on the Hopper TMA Unit for FP8 GEMMs

实证 / 前沿

  1. Luo et al., IPDPS’24 / arXiv:2402.13499
  2. Luo et al., arXiv:2501.12084
  3. Shah et al., FlashAttention-3(NeurIPS’24 / arXiv:2407.08608)
  4. ACTA(GPGPU’25 / DOI:10.1145/3725798.3725802)
  5. Yadav et al., Cypress(PLDI’25 / arXiv:2504.07004)