H200// TMA
← H200 速查主页 / 特性卡片墙

H200 · HOPPER GH100 · 特性专题

TMA · Tensor Memory Accelerator

★ Hopper 搬运主线 · 把地址计算与指令发射开销整个卸到硬件

01

是什么

DEFINITION

TMA 是 Hopper 引入的专用硬件拷贝单元——一条指令 + 一个 host 预编码的 tensor-map 描述符，由硬件完成整块多维 tile 在 global ↔ shared 之间的异步搬运（含 multicast 广播）。线程不再逐元素算地址、不再循环发小拷贝。

02

为什么需要（vs A100 cp.async）

RATIONALE · VS A100 CP.ASYNC
维度A100 cp.async / LDGSTSHopper TMA
发起粒度逐线程，每线程算源/目地址1 个线程（elect_one）发 1 条指令搬整块 tile
寻址线程内循环算 stride/offset硬件按 tensor-map 自动算多维地址
维度仅 1D1D–5D（boxDim ≤ 256/维）
完成同步commit + waitmbarrier::complete_tx::bytes（记字节数）
广播无multicast::cluster，一份拷贝落多个 CTA smem
越界处理软件判断硬件用 0 / NaN 填充

结果：地址计算开销归零；搬一个 128×128 tile 只占 1 线程的 1 条指令；multicast 省 L2/global 带宽。TMA + wgmma + mbarrier 三者构成 Hopper 的端到端异步流水，是 FlashAttention-3 / CUTLASS 双缓冲能真正重叠搬运与计算的前提。

03

关键数字

KEY NUMBERS · CONSTRAINTS
项值 / 约束注
tensor map 维度上限≤ 5D（tensorRank ≤ 5）cuTensorMapEncodeTiled
boxDim（每维 tile）≤ 256boxDim[0]×elemsize 必须 16B 对齐
globalStrides16 的倍数，< 2⁴⁰字节单位
地址对齐base 16B（INTERLEAVE_32B 要 32B）CUtensorMap 结构自身 64B 对齐
tensor map 创建位置HOST运行期换基址用 cuTensorMapReplaceAddress
传 kernel 方式__grid_constant__ const CUtensorMap*强制 grid_constant 语义
multicast 前提必须 cluster launch+ 16-bit multicast_mask（每位=1个CTA）
发射线程整 CTA 仅 1 个（elect_one）多线程同发 UB
04

可编译示例

COMPILABLE EXAMPLE · VERBATIM

逐字摘自 learn-cuda/02e_matmul_sm100/common.h 与 CUTLASS copy_sm90_tma.hpp。

① Host 端：编码 2D tensor map

// 描述一个 bf16 [H,W] 矩阵，tile 尺寸 [sH,sW] 落到 shared mem
inline void init_tmap_2d(CUtensorMap* tmap, const nv_bfloat16* ptr,
                 uint64_t gH, uint64_t gW, uint32_t sH, uint32_t sW,
                 CUtensorMapSwizzle swizzle) {
  constexpr uint32_t rank = 2;
  uint64_t globalDim[rank]        = {gW, gH};                     // {宽,高}
  uint64_t globalStrides[rank-1]  = {gW * sizeof(nv_bfloat16)}; // 字节, %16==0
  uint32_t boxDim[rank]           = {sW, sH};                     // tile in smem
  uint32_t elementStrides[rank]   = {1, 1};
  cuTensorMapEncodeTiled(
      tmap, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, rank, (void*)ptr,
      globalDim, globalStrides, boxDim, elementStrides,
      CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle,
      CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
}

② Device 端：发起 2D TMA global→shared + mbarrier 同步

// 单线程发起异步整块搬运；完成后 mbarrier 归零
__device__ __forceinline__ void tma_2d_g2s(
    uint32_t smem_dst, const void* tmap, uint32_t x, uint32_t y, uint32_t mbar) {
  asm volatile(
    "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes "
    "[%0], [%1, {%2, %3}], [%4];"
    :: "r"(smem_dst), "l"(tmap), "r"(x), "r"(y), "r"(mbar) : "memory");
}

// 使用顺序（cluster/单 CTA 通用）：
//   1) mbarrier_init(mbar, 1)
//   2) mbarrier_arrive_expect_tx(mbar, tile_bytes)   ← 预告期望字节数
//   3) if (elect_one) tma_2d_g2s(...)
//   4) mbarrier_wait_parity(mbar, phase)             ← 等硬件搬完归零

③ Multicast 变体（cluster 内广播到多个 CTA）

// 逐字自 CUTLASS SM90_TMA_LOAD_MULTICAST_2D —— 必须在 cluster launch 下
asm volatile(
  "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes"
  ".multicast::cluster.L2::cache_hint [%0], [%1, {%4,%5}], [%2], %3, %6;"
  :
  : "r"(smem_int_ptr), "l"(gmem_desc), "r"(smem_int_mbar),
    "h"(multicast_mask),               // 16-bit：每位 = cluster 内一个 CTA
    "r"(crd0), "r"(crd1), "l"(cache_hint)
  : "memory");
05

注意事项 / 陷阱

PITFALLS · CHECKLIST
CAUTION · 陷阱清单
  • tensor map 必须 HOST 创建；运行期换基地址用 cuTensorMapReplaceAddress，不要在 device 侧重建。
  • 传 kernel 必须 __grid_constant__ const CUtensorMap*——否则地址/语义 UB。
  • base 16B 对齐（INTERLEAVE_32B 要 32B）；swizzle 必须与 smem 访问布局一致，否则数据错位不报错。
  • Multicast 必须 cluster launch，目标用 .shared::cluster；非 cluster 下用 multicast 是 UB。
  • mbarrier 用前必须 mbarrier.init；先 arrive.expect_tx(size) 再发 TMA，顺序反了会少计字节。
  • 整 CTA 只能 1 个线程（elect_one）发射 cp.async.bulk.tensor，多线程同发 UB。
06

数据源

SOURCES