Skip to content

Hopper TMA 与 cuTensorMap 完全指南:从 descriptor 到 OOB 语义

在现代大模型算子(如 FlashAttention、Paged Attention、GEMM)的高性能实现中,计算单元(Tensor Core / Matrix Core)的浮点吞吐(FLOPS)呈现爆炸式增长,而片外显存(HBM)带宽与片上共享内存(Shared Memory / LDS)带宽的供需矛盾日益尖锐。然而,比物理带宽更早遭遇瓶颈的,是 SM(Streaming Multiprocessor)内部负责访存地址计算与边界检查的整数计算流水线及通用寄存器堆。

在传统 GPU 编程范式中,即便是搬运规则的矩阵分块(Tile),由于张量通常以高维(4D/5D)格式排布于一维显存空间,Warp 内的 32 个线程必须各自执行繁重的整数乘加指令计算显存物理字节偏移,并配合条件分支或谓词寄存器进行边界检查。这造成了严重的通用寄存器压力(VGPR Pressure)和指令发射槽位浪费(Issue Slot Waste)。

NVIDIA 从 Hopper 架构(SM90 / Compute Capability 9.0)开始,引入了革命性的硬件级搬运单元——TMA(Tensor Memory Accelerator)与驱动级描述符 cuTensorMap。TMA 不再只是一个简单的异步 DMA 控制器,而是第一个完全知晓多维张量几何形状(Tensor Geometry)、步长跨度(Strides)、片上 Bank 布局(Swizzle)以及越界填充语义(OOB Fill)的专用硬件协处理器。

本文将以极其详尽的技术深度,从底层驱动描述符定义、硬件地址生成微架构、皇冠级的 OOB 越界语义、PTX/SASS 汇编实测、Ampere cp.async 前身演进,直至跨体系的四方对照,为读者建立无懈可击的技术认知。


方法论:事实归入与严格标注准则

Section titled “方法论:事实归入与严格标注准则”

沿用上一篇的技术纪律,本文对涉及硬件行为、指令语法、执行语义及微架构特性的所有关键陈述,均严格归入以下三类并显式标注:

  1. 【documented】:NVIDIA 官方规范明确记载的事实。必须能追溯至官方一手资料,包括:
  2. 【observed】:作者在真实集群机器(环境:NVIDIA H100 80GB HBM3,CUDA 12.8 / 13.2,架构目标 sm_90 / sm_90a)上通过 ptxas 编译与 nvdisasm 反汇编实际复现验证的汇编与指令行为,正文均附带完整编译命令与 SASS 指令输出。
  3. 【inferred】:厂商未公开芯片底层微架构流水线细节,但根据指令行为、边界测试以及体系结构硬件实现原理所反推的合理微架构推断。严禁把 inferred 伪装成 documented。

第一部分:TensorMap 深度剖析与逐参数拆解

Section titled “第一部分:TensorMap 深度剖析与逐参数拆解”

1.1 CUtensorMap 类型定义:为什么是 128 字节 Opaque 结构?

Section titled “1.1 CUtensorMap 类型定义:为什么是 128 字节 Opaque 结构?”

在 CUDA 驱动 API 头文件 /usr/local/cuda/include/cuda.h 中,CUtensorMap 的类型声明如下:

#define CU_TENSOR_MAP_NUM_QWORDS 16
/**
* Tensor map descriptor. Requires compiler support for aligning to 64 bytes.
*/
typedef struct CUtensorMap_st {
#if defined(__cplusplus) && (__cplusplus >= 201103L)
alignas(64)
#elif __STDC_VERSION__ >= 201112L
_Alignas(64)
#endif
cuuint64_t opaque[CU_TENSOR_MAP_NUM_QWORDS];
} CUtensorMap;

物理事实核实【documented: cuda.h:3472-3485】

Section titled “物理事实核实【documented: cuda.h:3472-3485】”
  • 内存占用:opaque[16] 包含 16 个 64 位四字(Qword),即精确占用 16×8=12816 \times 8 = 128 字节(1024 位)。
  • 对齐要求:强制声明为 64 字节内存对齐(alignas(64))。

为什么设计为 Opaque?为什么只能通过驱动 API 构造?【inferred & documented】

Section titled “为什么设计为 Opaque?为什么只能通过驱动 API 构造?【inferred & documented】”
  1. 硬件解耦与微架构隐藏【inferred】: CUtensorMap 是硬件 TMA 单元的硬件描述符在显存/主机内存中的投影。SM 内的 TMA AGU 硬件直接从常量内存(Constant Memory)、全局显存(Global Memory)或共享内存(Shared Memory)拉取该 128 字节并解码其内部位域(bitfields)。NVIDIA 不对应用层公开其内部 1024 位的具体比特排布(例如哪个比特表示 rank、哪个比特段存 stride),以允许在未来微架构演进中(如 Blackwell sm_100 / Rubin)自由调整硬件字段定义,而无需破坏用户代码的二进制兼容性(ABI Stability)。
  2. 硬件不可变量的强制前置校验【documented】: 驱动 API cuTensorMapEncodeTiled 充当了“描述符微码编译器”的角色。构造一个合法的描述符需要满足数十项严格的几何约束(如全局步长必须是 16 的倍数、内部维度字节必须与 swizzle 模式匹配等)。若允许开发者像海光 MLS 拼接 vec4_uint 那样手动写入位域,任何细微的位计算错误都可能在硬件 TMA 引擎解码时导致难以调试的未定义行为或硬件断言。驱动 API 在构造时进行强类型的静态合法性断言,阻断了无效描述符下发到 GPU 的可能。

1.2 cuTensorMapEncodeTiled 逐参数全景拆解

Section titled “1.2 cuTensorMapEncodeTiled 逐参数全景拆解”

用于构造切片分块张量描述符的标准 CUDA Driver API 原型如下【documented: cuda.h:23910】:

CUresult CUDAAPI cuTensorMapEncodeTiled(
CUtensorMap *tensorMap,
CUtensorMapDataType tensorDataType,
cuuint32_t tensorRank,
void *globalAddress,
const cuuint64_t *globalDim,
const cuuint64_t *globalStrides,
const cuuint32_t *boxDim,
const cuuint32_t *elementStrides,
CUtensorMapInterleave interleave,
CUtensorMapSwizzle swizzle,
CUtensorMapL2promotion l2Promotion,
CUtensorMapFloatOOBfill oobFill
);

下面逐一拆解这 12 个参数,讲透其类型、取值范围、约束条件以及为什么硬件需要知道这个信息。

  • 类型:CUtensorMap *
  • 约束【documented】:输出指针,目标地址必须严格按 64 字节对齐。
  • 为什么硬件需要:TMA 引擎在常量内存、全局内存或共享内存中加载描述符时,利用 64B 宽度的对齐访存事务(Aligned 64B Transaction)以最低延迟将描述符加载至内部 AGU 缓存。
  • 类型:枚举 CUtensorMapDataType【documented: cuda.h:23660-23680】
    typedef enum CUtensorMapDataType_enum {
    CU_TENSOR_MAP_DATA_TYPE_UINT8 = 0, // 1 字节
    CU_TENSOR_MAP_DATA_TYPE_UINT16, // 2 字节
    CU_TENSOR_MAP_DATA_TYPE_UINT32, // 4 字节
    CU_TENSOR_MAP_DATA_TYPE_INT32, // 4 字节
    CU_TENSOR_MAP_DATA_TYPE_UINT64, // 8 字节
    CU_TENSOR_MAP_DATA_TYPE_INT64, // 8 字节
    CU_TENSOR_MAP_DATA_TYPE_FLOAT16, // 2 字节
    CU_TENSOR_MAP_DATA_TYPE_FLOAT32, // 4 字节
    CU_TENSOR_MAP_DATA_TYPE_FLOAT64, // 8 字节
    CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, // 2 字节
    CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ, // 4 字节, Flush-To-Zero
    CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, // 4 字节 (TensorFloat-32)
    CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ, // 4 字节 (TF32 FTZ)
    CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B, // 4 位紧凑打包, 8 字节对齐
    CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B, // 4 位打包, 16 字节对齐(含空洞)
    CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B // 6 位打包, 16 字节对齐(含空洞)
    } CUtensorMapDataType;
  • 约束【documented】:必须为上述合法枚举值。对于 Packed Sub-byte 类型,对对齐和尺寸有额外特殊限制。
  • 为什么硬件需要:
    • 坐标至字节地址折算:硬件 TMA 单元接收的是逻辑元素坐标(如 x=16x = 16),硬件必须知晓每个元素占用多少字节,才能在内部执行 byte_offset=coord×sizeof(dtype)\text{byte\_offset} = \text{coord} \times \text{sizeof(dtype)};
    • 分块总字节数计算:硬件需要依据元素大小核算 Bounding Box 的总字节数,并作为 Transaction 字节数打入 mbarrier;
    • 数据解包与对齐:对于 sub-byte 类型(如 4 位或 6 位),硬件 TMA 在搬运进共享内存时自动执行解包(unpacking)与填充。
  • 类型:cuuint32_t
  • 约束【documented】:1≤tensorRank≤51 \le \text{tensorRank} \le 5。最大支持 5 维张量。若 interleave != CU_TENSOR_MAP_INTERLEAVE_NONE,则 tensorRank 必须 ≥3\ge 3。
  • 为什么硬件需要:硬件 AGU 内部包含定长级联的多维步长展开管线。硬件需要根据 Rank 激活相应数量的坐标加法与步长乘法器通道。
  • 类型:void *(64 位虚地址)
  • 约束【documented】:基地址必须至少 16 字节对齐。特别地:
    • 若 interleave == CU_TENSOR_MAP_INTERLEAVE_32B,必须 32 字节对齐;
    • 若数据类型为 16U6_ALIGN16B 或 16U4_ALIGN16B,必须 32 字节对齐。
  • 为什么硬件需要:显存系统的突发读取以 16 字节(128 位宽总线节拍)或更粗粒度为基准对齐。硬件 AGU 将展开后的多维相对偏移叠加在 globalAddress 上生成物理显存请求。
  • 类型:const cuuint64_t *(长度为 tensorRank 的数组)
  • 约束【documented】:每个维度的尺寸必须非零且 ≤232\le 2^{32}(每个维度最大 4G4\text{G} 个元素)。
    • 对于 Packed 类型 16U6_ALIGN16B 和 16U4_ALIGN16B,globalDim[0] 必须是 128 的整数倍;
    • 对于 16U4_ALIGN8B,globalDim[0] 必须是 2 的整数倍。
  • 为什么硬件需要:定义逻辑张量的物理边界! 硬件 TMA 单元在收到逻辑坐标后,将遍历窗口的每个位置与 globalDim 进行比较。一旦超出 globalDim,硬件直接判定为 Out-Of-Bounds(OOB),并在硅片电路层级切断对显存总线的请求,实施越界填充。
  • 类型:const cuuint64_t *(长度为 tensorRank−1\text{tensorRank} - 1 的数组,单位:字节 Bytes)
  • 约束【documented】:
    • 每个步长必须是 16 的倍数;
    • 若 interleave == CU_TENSOR_MAP_INTERLEAVE_32B 或 packed 类型,步长必须是 32 的倍数;
    • 步长值必须严格 <240< 2^{40}(支持最大 1 TiB 的单维跨度寻址);
    • 高维步长必须单调递增并包含前序维度空间:
      globalStrides[0] = globalDim[0] * elementSize + padding[0];
      globalStrides[i] = globalStrides[i-1] * (globalDim[i] + padding[i]);
      assert(globalStrides[i] >= globalDim[i]);
  • 为什么硬件需要:连续排布的最低维(Dimension 0)由元素字节大小确定,不需要额外步长。从第 1 维到第 R−1R-1 维,硬件 AGU 通过如下公式在硬件电路中单周期展开多维物理地址: Effective Address=globalAddress+coord0×sizeof(dtype)+∑i=1R−1coordi×globalStrides[i−1]\text{Effective Address} = \text{globalAddress} + \text{coord}_0 \times \text{sizeof(dtype)} + \sum_{i=1}^{R-1} \text{coord}_i \times \text{globalStrides}[i-1]
  • 类型:const cuuint32_t *(长度为 tensorRank 的数组,单位:元素个数)
  • 约束【documented】:
    • 每个维度的尺寸必须非零且 ≤256\le 256;
    • 当 interleave == CU_TENSOR_MAP_INTERLEAVE_NONE 时,最内层维度的总字节数 {boxDim[0]×elementSize}\{\text{boxDim}[0] \times \text{elementSize}\} 必须是 16 字节的倍数;
    • 若数据类型为 16U6_ALIGN16B 或 16U4_ALIGN16B,boxDim[0] 必须固定为 128。
  • 为什么硬件需要:指定了单次 TMA 指令搬运的 Bounding Box 几何尺寸(Tile Shape)。硬件据此规划总数据量、在共享内存中划分目标分块拓扑、并向 mbarrier 注册预期接收的字节计数。
  • 类型:const cuuint32_t *(长度为 tensorRank 的数组)
  • 约束【documented】:
    • 数组每个元素必须非零且 ≤8\le 8;
    • 若 interleave == CU_TENSOR_MAP_INTERLEAVE_NONE,elementStrides[0]\text{elementStrides}[0] 被硬件完全忽略(最低维内部不支持跳跃式子采样,必须连续);
    • 当某个维度 elementStrides[i]>1\text{elementStrides}[i] > 1 时,TMA 实际加载的元素个数为 ⌈boxDim[i]/elementStrides[i]⌉\lceil \text{boxDim}[i] / \text{elementStrides}[i] \rceil。
  • 为什么硬件需要:支持下采样(Sub-sampling / Strided Slicing)。硬件在遍历 Bounding Box 时,按该跨步跃进取样,而无需在加载到寄存器后手动丢弃。
  • 类型:枚举 CUtensorMapInterleave【documented: cuda.h:23685-23695】
    typedef enum CUtensorMapInterleave_enum {
    CU_TENSOR_MAP_INTERLEAVE_NONE = 0,
    CU_TENSOR_MAP_INTERLEAVE_16B,
    CU_TENSOR_MAP_INTERLEAVE_32B
    } CUtensorMapInterleave;
  • 约束【documented】:支持特定交错格式,如卷积与图像处理中的通道紧凑存储格式(如 NC/8HWC8NC/8HWC8 采用 16B 交错,NC/16HWC16NC/16HWC16 采用 32B 交错)。当启用交错时,张量维度必须 ≥3\ge 3。
  • 为什么硬件需要:硬件数据重组单元在向 Shared Memory 写入时,能够直接将交错排布的通道切片解构为规整的连续张量,消除软件手动内存转置的巨额开销。
  • 类型:枚举 CUtensorMapSwizzle【documented: cuda.h:23715-23730】
    typedef enum CUtensorMapSwizzle_enum {
    CU_TENSOR_MAP_SWIZZLE_NONE = 0,
    CU_TENSOR_MAP_SWIZZLE_32B, // 在 32B 跨度内对 16B 数据块进行异或打散
    CU_TENSOR_MAP_SWIZZLE_64B, // 在 64B 跨度内对 16B 数据块进行异或打散
    CU_TENSOR_MAP_SWIZZLE_128B, // 在 128B 跨度内对 16B 数据块进行异或打散
    CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B, // 在 128B 跨度内以 32B 为原子块进行打散
    CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B, // 32B 原子打散并在隔行交换高低 8 字节
    CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B // 在 128B 跨度内以 64B 为原子块进行打散
    } CUtensorMapSwizzle;
  • 约束【documented】:
    • 若 interleave == NONE 且 swizzle != NONE,Bounding Box 最内层维度的字节跨度(boxDim[0]×elementSize\text{boxDim}[0] \times \text{elementSize})必须 ≤\le swizzle 模式跨度(例如 64B Swizzle 要求最内层 ≤64\le 64 字节;128B Swizzle 要求最内层 ≤128\le 128 字节);
    • 若 interleave == CU_TENSOR_MAP_INTERLEAVE_32B,swizzle 必须为 CU_TENSOR_MAP_SWIZZLE_32B。
  • 为什么硬件需要:彻底消除 Shared Memory Bank Conflicts! 共享内存由 32 个 4 字节(32-bit)的 Banks 构成。Tensor Core(WGMMA)以特定的交织行/列模式并发读取数据。TMA 搬运引擎直连片上 Crossbar,在将显存数据写入 Shared Memory 时,直接利用硬件 XOR 网络对目标字节地址的高低位进行异或映射,将原本会落在同一 Bank 的数据天然打散至 32 个不同 Banks 中,实现 Tensor Core 读取时的零冲突(Zero Bank Conflict)。
  • 类型:枚举 CUtensorMapL2promotion【documented: cuda.h:23755-23765】
    typedef enum CUtensorMapL2promotion_enum {
    CU_TENSOR_MAP_L2_PROMOTION_NONE = 0,
    CU_TENSOR_MAP_L2_PROMOTION_L2_64B,
    CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
    CU_TENSOR_MAP_L2_PROMOTION_L2_256B
    } CUtensorMapL2promotion;
  • 约束【documented】:指定 L2 Cache 从外部 DRAM(HBM)拉取缓存行时的填充粒度(Fetch Granularity)。
  • 为什么硬件需要:优化片外显存总线带宽。对于连续稠密张量,提升至 128B 或 256B 填充能大幅提高总线利用率;而对于稀疏或小跨步切片,较小的粒度(64B 或 NONE)可有效防止 L2 缓存污染并节省不必要的带宽损耗。
  • 类型:枚举 CUtensorMapFloatOOBfill【documented: cuda.h:23770-23780】
    typedef enum CUtensorMapFloatOOBfill_enum {
    CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE = 0,
    CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA
    } CUtensorMapFloatOOBfill;
  • 约束【documented】:CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA 仅允许在浮点数据类型(FLOAT16, FLOAT32, FLOAT64, BFLOAT16, TFLOAT32 等)且非 Packed 类型时使用;整型张量必须配置为 NONE。
  • 为什么硬件需要:告诉硬件当 Bounding Box 越界时,向 Shared Memory 写入何种物理数值! 是直接填零还是填入特定的 OOB-NaN。这是整个 TMA 体系中最关键的容错与边界保护机制。

第二部分:核心追问:为什么不能只传 pointer + byte offset?

Section titled “第二部分:核心追问:为什么不能只传 pointer + byte offset?”

在讨论具体指令之前,我们必须站在计算机体系结构师(Computer Architect)的角度直面这个根本问题:既然传统的 GPU 内存加载只需要 void* ptr 和一个 32 位/64 位偏移,为什么 TMA 偏偏要大费周章引入一个包含 12 项配置的 128 字节描述符?

答案在于抹平计算指令集与高层张量代数之间的“语义鸿沟”(Semantic Gap)。

传统指令集 (Volta / Turing / Ampere)
┌────────────────────────────────────────────────────────┐
│ 软件 (Software Warp) │
│ - 32 个线程分别计算多维张量坐标 │
│ - 32 个线程分别计算 64 位物理字节地址 │
│ - 32 个线程各自执行 bounds 检查与谓词分流 │
│ - 32 个线程分别发射 16B load 指令 (每 Warp 32 条指令) │
└─────────────────────────┬──────────────────────────────┘
│ 传入: 裸指针 + 字节偏移
▼
┌────────────────────────────────────────────────────────┐
│ 硬件显存控制器 (Memory Controller) │
│ - 仅知晓 1D 线性虚拟地址 (Byte Address) │
│ - 不知晓张量维度、边界、步长、切片形状 │
└────────────────────────────────────────────────────────┘
Hopper TMA 体系 (SM90)
┌────────────────────────────────────────────────────────┐
│ 软件 (Software Warp) │
│ - 仅由 1 个线程提交逻辑多维坐标 {coord0, coord1, ...} │
│ - 单条 TMA 指令启动整个分块搬运 (每 Warp 1 条指令) │
└─────────────────────────┬──────────────────────────────┘
│ 传入: CUtensorMap (128B) + 逻辑坐标
▼
┌────────────────────────────────────────────────────────┐
│ TMA 硬件异步引擎 (TMA Hardware Engine) │
│ 1. globalDim 比较器 ──► 硅片级瞬时完成 OOB 判定 │
│ 2. globalStrides 定点单元 ──► 硬件展开多维物理地址 │
│ 3. boxDim 发生器 ──► 生成成批 Burst 访存事务 │
│ 4. Swizzle XOR 网络 ──► 写入即为 Tensor Core 零冲突布局│
│ 5. oobFill 状态机 ──► 越界区域自动硬件补零 / 补 NaN │
└────────────────────────────────────────────────────────┘

从上图可以看出,如果只向硬件传递 pointer + byte offset:

  1. 丢失边界定义能力:硬件只看到一个线性地址数值,根本无法识别该地址是否越过了第 1 维的边界、还是第 2 维的边界。硬件无法定义“什么是越界”,边界检查被迫全盘堆积在软件层;
  2. 失去硬件展开多维地址的可能:缺乏 globalStrides,硬件无法自动将 (z,y,x)(z, y, x) 映射到物理行与物理面,SM 必须耗费成百上千个 ALU 周期与寄存器去手工算地址;
  3. 失去片上排布的预知性:缺乏 boxDim 与 swizzle,搬运硬件不知道目标分块在 Shared Memory 中的 2D/3D 几何形状,无法在写入时直接通过 Crossbar 完成打散,软件必须在搬入后再做一次昂贵的共享内存内部转置;
  4. 失去硬件阻断与安全置零的依据:缺乏 oobFill,硬件无法在切断越界物理访存的同时,向目标内存安全注入零或 NaN,软件不得不保留复杂的掩码逻辑。

在上一篇《MLS 三个访存 Bug 深度剖析》第三部分中,我们详细分析了海光 MLS 的底层描述符。海光 MLS 基于 AMD GCN/CDNA 架构的标量描述符 SRD(vec4_uint,128 位)。

在 AMD 标准规范中,SRD 的第 2/3 字原本设计有 NUM_RECORDS 字段用于硬件缓冲区的边界检查。然而:

海光 MLS 为了在极度受限的 128 位寄存器中塞入矩阵的 1D 行步长,直接将原本用于边界保护的 NUM_RECORDS 字段挪用来存放 row_stride!【见上一篇 §3.2 与 intrinsic_mls_ds.h】

这种架构层面的妥协导致了毁灭性的后果:MLS 硬件彻底丧失了对张量总尺寸与物理边界的感知能力! 硬件只能依赖一个粗粒度的 nm_filter 行掩码;当未对齐尾块的基址落入未映射虚页时,MLS 依然拿着裸地址向 MMU 索取翻译,从而直接引发了致命的硬件 VM Fault 崩溃(MLS Bug 2)。

而 NVIDIA Hopper 的 CUtensorMap 拥有整整 128 字节(1024 位) 的巨大物理容量,将 globalDim(每个维度独立 32 位)、globalStrides(每个维度独立 40 位)、boxDim、swizzle 与 oobFill 完整容纳。NVIDIA 根本不需要在“步长展开”与“边界检查”之间做生死二选一,从而在硬件体系结构的最底层彻底抹平了这一代际缺陷。


第三部分:OOB 语义完全解析(皇冠深度详解)

Section titled “第三部分:OOB 语义完全解析(皇冠深度详解)”

在所有关于 TMA 的微架构设计中,Out-Of-Bounds(OOB)越界处理机制是其最精妙、最具开创性的皇冠部分。读者在面试或工程研发中,必须对以下细节建立极度严密的逻辑区分。

3.1 彻底切开三类本质不同的异常概念

Section titled “3.1 彻底切开三类本质不同的异常概念”

在高性能编程中,“越界”常被混为一谈。在 TMA 体系中,必须在概念上严密切开以下三者:

┌─────────────────────────┬────────────────────────────────────────────────────────┐
│ 概念分类 │ 触发条件与硬件行为 │
├─────────────────────────┼────────────────────────────────────────────────────────┤
│ 1. logical OOB │ 描述符完全合法,但请求的 Tile 坐标超出了 globalDim 范围。│
│ (逻辑越界, 完全合法) │ 【documented】这是正常执行路径!硬件绝不报错,硬件自动截断 │
│ │ 物理显存事务,向 Shared Memory 填零或 OOB-NaN。 │
├─────────────────────────┼────────────────────────────────────────────────────────┤
│ 2. invalid descriptor │ 描述符构造错误(如地址未 16B 对齐、Rank>5、boxDim>256 等)。│
│ (非法描述符, 驱动报错│ 【documented】在 Host 端驱动直接返回 CUDA_ERROR_INVALID_ │
│ 或硬件未定义行为) │ VALUE;若在 Device 篡改破坏描述符,触发 Illegal Instruction。│
├─────────────────────────┼────────────────────────────────────────────────────────┤
│ 3. invalid pointer / │ globalAddress 指向非法、已释放或未向 GPU MMU 映射物理页 │
│ unmapped page (缺页) │ 的虚拟地址。即使逻辑坐标合法,访问该基址所在虚页也会触发 │
│ │ 硬件级 GPU Page Fault / VM Fault 崩溃。 │
└─────────────────────────┴────────────────────────────────────────────────────────┘

与 MLS Bug 2 的根本对照: 海光 MLS Bug 2 的本质错误在于:把第 1 类(逻辑越界)错误地转化成了第 3 类(未映射页缺页)。因为软件计算的基址直接跨入了不存在的虚拟页,而硬件又没有张量边界感知; 而 Hopper TMA 则将第 1 类(逻辑越界)通过 globalDim 的硬件比较器收敛为完全合法、安全、有硬件规范保障的片上纯内存操作,物理显存总线与 MMU 甚至根本不会察觉到这次越界!


考虑一个 1D/2D 张量,在最内层维度上:

  • 张量全局维度:globalDim[0] = 1000;
  • TMA 分块形状:boxDim[0] = 128;
  • 某个 CTA 请求的 Tile 起始坐标:coord = 896。
  • 该 Tile 覆盖的逻辑区间为 [896,1024)[896, 1024)。

此时发生边界切分:

  • 合法数据区间:[896,1000)[896, 1000),共计 1000−896=1041000 - 896 = 104 个元素;
  • 越界数据区间:[1000,1024)[1000, 1024),共计 1024−1000=241024 - 1000 = 24 个元素。
张量逻辑空间 (globalDim = 1000)
0 896 1000 1024
├──────────────────────────────────────────────┼──────────────┼───────────┤
│ 合法显存区域 │ 合法切片 │ 逻辑越界 │
│ │ (104 elem) │ (24 elem) │
└──────────────────────────────────────────────┴──────────────┴───────────┘
▲ ▲
│◄────── boxDim=128 ──────►│
│
TMA 发起请求起始坐标
目标 Shared Memory 分块 (128 elem 空间)
┌──────────────────────────────────────────────┬──────────────────────────┐
│ 从物理显存 Burst 读取的真实数据 (104 elem) │ 硬件 TMA 自动填充区域 │
│ [896, 1000) │ (24 elem: 全 0 或 OOB-NaN)│
└──────────────────────────────────────────────┴──────────────────────────┘
  1. 是否对越界区域发出 global memory transaction?
    • 【documented: PTX ISA 5.5.3.3】:PTX 规范明确保证:落在张量物理边界外的元素不会从显存中加载,而是由硬件自动填充为 0 或 OOB-NaN。
    • 【inferred: 微架构请求抑制】:在硅片微架构层面,TMA AGU 将 Bounding Box 拆解为面向 L2 Cache / DRAM 的底层事务(通常以 32B/64B/128B Cache Line 为粒度)。对于完全位于 1000 元素之后的 Cache Lines,TMA AGU 的硬件比较器直接抑制(Suppress)生成访存事务;对于跨越 1000 边界的单条 Cache Line,硬件依据元素掩码进行 Byte-level Masking,确保绝不触碰非法边界外的物理存储单元。
  2. 目标 Shared Memory 被写入什么?
    • 若 oobFill == CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE:Shared Memory 中对应 [1000,1024)[1000, 1024) 的 24 个元素位置被精确硬件写入字节 0x00(即浮点数 +0.0f);
    • 若 oobFill == CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA:Shared Memory 中对应 [1000,1024)[1000, 1024) 的 24 个元素位置被精确写入特定的 OOB-NaN 位模式。
  3. mbarrier 事务计数如何满足?
    • 无论是否发生 Partial OOB,Shared Memory 接收到的总写入数据量依然是完整的 128 个元素大小。TMA 硬件在完成真实数据搬运与越界填充后,向关联的 mbarrier 发出 complete_tx 信号,其核减的字节数精确等于完整的分块大小,上层流水线无需对 tail 分块做任何特殊的同步逻辑处理。

考虑更为极端但实际工程中极其常见的场景:

  • 张量全局维度:globalDim = 1000;
  • 某个尾部 CTA 计算出的起始坐标:coord = 2048(远大于 1000,例如网格向上取整分块后产生的完全多余的 CTA)。
  • 此时该 Tile 覆盖的逻辑区间为 [2048,2176)[2048, 2176),整块数据完全处于张量范围之外。
  1. 这是合法的调用吗?
    • 【documented: PTX ISA 5.5.3.3】:完全合法! 驱动 API 与 PTX 指令集对 Fully OOB 拥有完备且良定义的语义,软件无需用 if (tile_coord < total_tiles) 去保护 TMA 指令的发射。
  2. 硬件微架构行为是什么?
    • 【inferred】:TMA AGU 在接收到坐标后,首级比较器判定 coord >= globalDim,整个 Bounding Box 完全命中 OOB 条件。
    • TMA AGU 完全不向 L2 Cache 或 HBM 显存控制器发射任何读事务(Zero Memory Transactions)!
    • TMA 引擎直接驱动片上 Shared Memory 写入管线,将整个目标分块(128×elem_size128 \times \text{elem\_size} 字节)全部刷入 0 或 OOB-NaN;
    • 写入完成后,硬件立即向 mbarrier 冲刷事务计数信号,通知下游消费者 Warp。
  3. 与 MLS Bug 2 的本质区别:
    • MLS Bug 2 的悲剧:海光 FlashAttention 中,尾块计算出的基地址直接指向了不存在的物理页,由于硬件缺乏张量维度定义,AGU 强行发起虚地址翻译,导致 GPU 发生不可恢复的致命缺页(Page Fault);
    • TMA 的降维打击:TMA 中传入的基地址永远是张量合法的首地址 globalAddress(该首地址所在的虚页必然是有效映射的),软件传入的只是整数坐标 coord = 2048。TMA 硬件在纯数值比较阶段就拦截了所有向外层总线的请求。基址是合法的,坐标越界被内部消化为片上内存填零,从物理电路上彻底抹除了触发 MMU Page Fault 的可能性!

3.4 两种 Float OOB Fill 模式的本质差异

Section titled “3.4 两种 Float OOB Fill 模式的本质差异”

CUDA 驱动中定义的两种枚举模式如下【documented: cuda.h:23770-23780】:

  1. CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  2. CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA
  • CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE:
    • 命名中虽带有 NONE,但其实际语义为:“无特殊 NaN 请求,越界元素直接物理置零(Zero-Fill)”【documented】。
    • 对于 IEEE 754 浮点数,写入的位模式为全部比特置 0(即 +0.0,对于 FP32 为 0x00000000,对于 FP16 为 0x0000)。
  • CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA:
    • 名字逐段拆解:
      • NAN:越界元素在 Shared Memory 中被填充为一个专用的、带特定 payload 的 OOB-NaN 浮点常量;
      • REQUEST_ZERO_FMA:该特殊常量携带了一个对下游算术管线的“协议请求”——要求后续支持该语义的 FMA 指令(如 fma.oob)在执行乘加运算时,将该乘法操作数强制当作 0.00.0 处理!

核心追问:为什么 OOB-NaN 比简单的 Zero-Fill 更“高级、丰富”?

Section titled “核心追问:为什么 OOB-NaN 比简单的 Zero-Fill 更“高级、丰富”?”

很多开发者会产生直觉困惑:既然最终都要当零算,为什么不直接全填零,非要多此一举搞一个 OOB-NaN?

  1. 信息量的保留与“有效性”标记:
    • 在科学计算与神经网络中,数值 0.00.0 本身是一个完全合法的数值(例如稀疏矩阵的真实零、ReLU 激活后的截断零、特征归一化后的中心零)。
    • 如果硬件在数据搬运阶段简单粗暴地将所有越界位置填零,下游算术代码将彻底失去区分“这是一个真实的零”还是“这是一个越界填充的虚假占位符”的能力。
    • 例如在 FlashAttention 的 Softmax 归约中,若需计算每行的 max⁡(x)\max(x),真实数据可能全为负数(如 [−10,−5,−2][-10, -5, -2])。若尾部填充被粗暴置为 0.00.0,则计算出的最大值会被错误地抬高为 0.00.0!而如果是 OOB-NaN,下游逻辑可以在归约时直接跳过 NaN,或者通过专用比较指令识别出无效槽位。
  2. 数据通路上的硬件级断言保护:
    • OOB-NaN 在未经 fma.oob 保护的普通算术流水线中具有 IEEE 754 强传染性。如果算子开发者的后续逻辑(如 Epilogue 阶段)意外使用了非 OOB 指令消费了越界区域,NaN 会立即扩散爆发,从而在单元测试阶段直接暴露边界处理缺陷,而不会像普通 0.00.0 那样引发极难排查的“静默精度损毁”(Silent Data Corruption)。
  • CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA 仅对标准浮点类型有效(FLOAT16, BFLOAT16, FLOAT32, FLOAT64, TFLOAT32 等);
  • 对于整型数据类型(UINT8, INT32 等)或 Packed Sub-byte 类型,驱动 API 严格禁止配置此项,只能使用 CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE(整型越界填 0)。

3.5 fma.oob 与 OOB-NaN 的配套关系及实测验证

Section titled “3.5 fma.oob 与 OOB-NaN 的配套关系及实测验证”

PTX 语法规范【documented: PTX ISA 9.7.4.8】

Section titled “PTX 语法规范【documented: PTX ISA 9.7.4.8】”

PTX ISA 8.1 针对 SM90(Hopper)引入了带 .oob 修饰符的 FMA 指令:

fma.rn.oob{.relu}.type d, a, b, c;
  • 支持的数据类型【documented】:.f16, .f16x2, .bf16, .bf16x2。
  • 硬件语义【documented】:当指令带有 .oob 修饰符时,若输入操作数中检测到特定的 OOB-NaN 常量,硬件在算术电路层级强行将该输入操作数的乘法乘积项截断置为 +0.0。
  • SIMD 独立性【documented】:对于 .f16x2 和 .bf16x2,硬件在两个独立的 16 位 SIMD 切片上各自执行 OOB-NaN 检测与置零,互不干扰。

集群实测验证与 SASS 反汇编【observed】

Section titled “集群实测验证与 SASS 反汇编【observed】”

为了彻底验证编译器与底层硬件指令的真实形态,我们在 H100 集群(CUDA 12.8,-arch=sm_90)上进行了全类型拓展实测。

测试 1:半精度向量 fma.rn.oob.f16x2
Section titled “测试 1:半精度向量 fma.rn.oob.f16x2”

编写 PTX 测试文件:

.version 8.1
.target sm_90
.address_size 64
.visible .entry test_fma_f16x2(.param .u64 p_out, .param .u64 p_in) {
.reg .u64 po, pi;
.reg .b32 a, b, c, d;
ld.param.u64 po, [p_out]; ld.param.u64 pi, [p_in];
ld.global.b32 a, [pi]; ld.global.b32 b, [pi+4]; ld.global.b32 c, [pi+8];
fma.rn.oob.f16x2 d, a, b, c;
st.global.b32 [po], d;
ret;
}

执行编译与反汇编:

Terminal window
ptxas -arch=sm_90 test_fma_f16x2.ptx -o test_fma_f16x2.cubin
nvdisasm test_fma_f16x2.cubin

真实 SASS 输出【observed】:

/*0070*/ HFMA2.OOB R7, R0, R7, R6 ;

编写对应 PTX 并执行编译反汇编,真实 SASS 输出【observed】:

/*0070*/ HFMA2.OOB R0, R0.H0_H0, R7.H0_H0, R6.H0_H0 ;

微架构观察:标量 f16 的 OOB FMA 依然复用了 32 位的半精度向量单元 HFMA2.OOB,通过子寄存器广播标志 .H0_H0 实现标量运算!

测试 3:BFloat16 变体 fma.rn.oob.bf16 与 fma.rn.oob.bf16x2
Section titled “测试 3:BFloat16 变体 fma.rn.oob.bf16 与 fma.rn.oob.bf16x2”

编写对应 PTX 并执行编译反汇编,真实 SASS 输出【observed】:

// bf16 标量 SASS
/*0070*/ HFMA2.BF16_V2.OOB R0, R0.H0_H0, R7.H0_H0, R6.H0_H0 ;
// bf16x2 向量 SASS
/*0070*/ HFMA2.BF16_V2.OOB R7, R0, R7, R6 ;
测试 4:拓展测试:是否存在 fma.rn.oob.f32 或 fma.rn.oob.f64?
Section titled “测试 4:拓展测试:是否存在 fma.rn.oob.f32 或 fma.rn.oob.f64?”

我们将类型修改为单精度 .f32 进行编译:

fma.rn.oob.f32 d, a, b, c;

执行编译命令:

Terminal window
ptxas -arch=sm_90 test_fma_f32.ptx -o test_fma_f32.cubin

真实输出报错【observed】:

ptxas test_fma_f32.ptx, line 16; error : Illegal modifier '.oob' for instruction 'fma'
ptxas fatal : Ptx assembly aborted due to errors

同样,测试 .f64 时触发完全相同的语法错误。 实测定论【observed & documented】: fma.oob 在硬件与指令集层面严格仅支持半精度(.f16, .f16x2, .bf16, .bf16x2),在 32 位单精度(FFMA)与 64 位双精度(DFMA)硬件执行单元中不存在 .OOB 电路修饰符!

严格厘清四者的关系(绝不造假)

Section titled “严格厘清四者的关系(绝不造假)”

在技术面试与架构审查中,必须极度精准地厘清以下四者的边界:

┌─────────────────────────┬────────────────────────────────────────────────────────┐
│ 实体对象 │ 硬件/文档规范行为 │
├─────────────────────────┼────────────────────────────────────────────────────────┤
│ 1. 普通 IEEE 754 NaN │ 标准 quiet/signaling NaN。普通 FMA 执行 0.0 * NaN 结果 │
│ │ 依然为 NaN!【documented: IEEE 754 标准】 │
├─────────────────────────┼────────────────────────────────────────────────────────┤
│ 2. TMA 专用 OOB-NaN │ TMA 越界填充时写入的保留 Payload NaN 常量。 │
│ │ 【documented: PTX ISA 5.5.3.3】 │
├─────────────────────────┼────────────────────────────────────────────────────────┤
│ 3. fma.oob (HFMA2.OOB) │ 专属 CUDA Core SIMD 指令。输入命中 OOB-NaN 时,将结果 │
│ │ 强制截断为 +0.0。【documented: PTX ISA 9.7.4.8】 │
│ │ 【注意】对普通 NaN 不起置零作用,普通 NaN 依然扩散! │
├─────────────────────────┼────────────────────────────────────────────────────────┤
│ 4. MMA / WGMMA │ 【documented / 严正声明】NVIDIA 官方 PTX/CUDA 文档中, │
│ (Tensor Core) │ 从未声明 WGMMA 具备自动吞掉 OOB-NaN 的能力! │
│ │ WGMMA 指令集没有任何 .oob 修饰符。 │
│ │ 【工程决策】在使用 Tensor Core (WGMMA) 的矩阵乘中, │
│ │ 防范 NaN 污染必须在 TMA 阶段直接配置 FILL_NONE (填 0)!│
└─────────────────────────┴────────────────────────────────────────────────────────┘

[!CAUTION] 绝对认知红线: 某些网络博客常声称“Hopper Tensor Core / WGMMA 会自动识别并吞掉所有 OOB-NaN”。这是严重缺乏文档依据与硬件常识的错误断言! 官方 PTX ISA 中,拥有 .oob 修饰符的只有 CUDA Core 的标量/向量半精度 FMA 指令(fma.rn.oob.*)。在 Tensor Core 的异步矩阵乘流水线(wgmma.mma_async)中,硬件要求输入操作数必须是合法的数值;如果向 Tensor Core 馈入 NaN,即便在 TMA 描述符中设置了 NAN_REQUEST_ZERO_FMA,累加器也会被迅速毒化。 在 GEMM 和 FlashAttention 的 QK/PV 矩阵乘阶段,TMA 描述符必须无条件使用 CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE(填 0)! 那么 NAN_REQUEST_ZERO_FMA 是为谁设计的?它是为算子后续在通用 SM 核心上执行的向量化非矩阵乘流水线(如 FlashAttention 的 Softmax 归约、偏置融合、LayerNorm 等由 CUDA Core 执行的 HFMA2.OOB 计算)量身定制的!


3.6 TMA 的 OOB 语义如何消灭海光 MLS 的三大 Bug?

Section titled “3.6 TMA 的 OOB 语义如何消灭海光 MLS 的三大 Bug?”

回到学习计划的核心主线,我们逐一审查 TMA 对上一篇《MLS 三个访存 Bug 深度剖析》中三大 Bug 的终结机制:

1. 对照 MLS Bug 1(SRD 低 32 位累加导致跨 4 GiB 回绕)

Section titled “1. 对照 MLS Bug 1(SRD 低 32 位累加导致跨 4 GiB 回绕)”
  • 海光 MLS 缺陷根源:软件必须在循环中用 32 位加法拼接更新标量描述符 srsrc[0:1],由于疏漏未向高 32 位传递进位,导致显存跨 4 GiB 边界时低 32 位回绕,发生毁灭性的非法显存覆盖;
  • TMA 的消除机制:
    • 张量物理基地址 globalAddress 作为一个完整的 64 位指针,在 Host 端构造 CUtensorMap 时一次性写入描述符;
    • Kernel 端在执行分块搬运时,完全不需要、也不可能去修改 64 位基地址!
    • Kernel 端只向指令提交 32 位的逻辑张量坐标 {coord0, coord1, ...};
    • 物理地址展开(Address Generation)完全由 TMA 硬件 AGU 在内部以全位宽硬件加法器执行;
    • 定论:这类“64 位指针进位丢失导致的显存回绕”在 TMA 体系下被从物理根源上彻底消灭。
    • 边界区分:软件若算错坐标(例如 tile 索引算错),只会导致加载到错误的逻辑 tile 或触发 OOB 填零,但绝不会破坏 64 位物理地址的虚页映射结构。

2. 对照 MLS Bug 2(全掩码尾块触发硬件缺页 VM Fault)

Section titled “2. 对照 MLS Bug 2(全掩码尾块触发硬件缺页 VM Fault)”
  • 海光 MLS 缺陷根源:尾块虽然将行掩码设为全过滤(nm_filter = 16),但描述符基地址被软件算到了超出物理显存分配的未映射虚拟页。硬件 AGU 缺乏全局张量几何感知,依然向 MMU 索要地址翻译,触发硬件 Page Fault;
  • TMA 的消除机制:
    • TMA 描述符内建了精确的张量尺寸 globalDim;
    • 当尾块完全处于张量范围之外时(Fully OOB),TMA AGU 在首级数值比较器上直接判定请求越界;
    • TMA 硬件在微架构上直接切断访存总线请求,绝不向 L2/MMU 发射任何物理显存事务!
    • TMA 内部管线直接向 Shared Memory 写入 0 并打卡 mbarrier;
    • 定论:这类“尾块地址触碰未映射页导致 VM Fault”在 TMA 体系下被彻底消除。

3. 对照 MLS Bug 3(预取越界读入脏显存引发 NaN 毒化累加器)

Section titled “3. 对照 MLS Bug 3(预取越界读入脏显存引发 NaN 毒化累加器)”
  • 海光 MLS 缺陷根源:双 Tile 预取循环在边界处错位,读入未初始化显存中的垃圾数据,恰好命中浮点 NaN 位模式,在后续 MMA 计算中因 IEEE 754 0.0×NaN=NaN0.0 \times \text{NaN} = \text{NaN} 导致累加器被不可逆毒化;
  • TMA 的消除机制:
    • 第一道防线(入口物理置零):TMA 在数据搬运的第一步,就由硬件强行向 Shared Memory 越界槽位注入全零(FILL_NONE),从源头上切断了未初始化垃圾数据进入片上的物理通路;
    • 第二道防线(算术级硬件截断):若配置为 OOB_FILL_NAN,在后续 CUDA Core 的半精度计算中,fma.oob(HFMA2.OOB)在检测到 OOB-NaN 时在电路层强制将乘积项置为 +0.0f,双重保险彻底封死累加器毒化通路。

第四部分:前身:Ampere cp.async(不可逾越的技术阶梯)

Section titled “第四部分:前身:Ampere cp.async(不可逾越的技术阶梯)”

在很多初学者的认知中,常误以为“只有到了 Hopper 架构才具备异步安全处理尾块的能力”。这是严重的认知断层!

事实上,NVIDIA 在 SM80(Ampere 架构)就引入了第一代硬件级异步搬运机制 cp.async。要深刻理解 TMA 究竟把哪些责任移交给了硬件,必须先彻底吃透 cp.async。

4.1 SM80 cp.async 核心机制与指令细分

Section titled “4.1 SM80 cp.async 核心机制与指令细分”

1. cp.async.ca vs cp.async.cg【documented: PTX ISA 9.7.9.26.3.1】

Section titled “1. cp.async.ca vs cp.async.cg【documented: PTX ISA 9.7.9.26.3.1】”
  • cp.async.ca(Cache at All levels):
    • 数据在从全局显存搬往共享内存的同时,在 L1 数据缓存和 L2 缓存中均建立拷贝;
    • 支持 4 字节、8 字节、16 字节 拷贝粒度。
  • cp.async.cg(Cache at Global level):
    • 绕过(Bypass)SM 本地的 L1 缓存,仅在 L2 缓存保留副本,数据直通 Shared Memory;
    • 严格只支持 16 字节拷贝;
    • 架构价值:在深度学习 GEMM 中,矩阵块通常只在 Shared Memory 中复用一次,若经过 L1 会迅速将 SM 有限的 L1 缓存行(Line)冲刷殆尽。cp.async.cg 完美避免了 L1 缓存污染。

传统的 @p cp.async.ca.shared.global [dst], [src], 16;: 当谓词 @p 为 false 时,指令整条被屏蔽(NOP)。这能防止非法地址越界读,但它无法解决目标共享内存的置零问题——共享内存原本是什么脏数据,就依然保留什么脏数据。

3. src-size < cp-size 与硬件自动 Zero-Fill【documented】

Section titled “3. src-size < cp-size 与硬件自动 Zero-Fill【documented】”

为了解决尾块置零,Ampere 硬件在指令中引入了动态的 src-size 参数:

cp.async.ca.shared.global [dst], [src], cp-size, src-size;
  • cp-size:编译期静态常量(4, 8, 16 字节),代表写入 Shared Memory 的总字节数;
  • src-size:运行期 32 位寄存器动态值,代表实际从全局显存读取的有效字节数(必须 ≤cp-size\le \text{cp-size})。
  • 硬件执行语义【documented】: 硬件只从全局显存读取 src-size 个字节写入目标内存,而目标内存中剩余的 (cp-size - src-size) 个字节由硬件自动填充为 0!

在 PTX ISA 7.5(SM80)中,NVIDIA 进一步引入了 ignore-src 谓词操作数:

cp.async.cg.shared.global [dst], [src], 16, p;
  • 当谓词寄存器 p 为 True 时,硬件完全忽略全局显存地址 [src](不发起任何显存读事务,即使 [src] 是非法未映射地址也不会触发异常),同时直接向共享内存 [dst] 写入 16 字节的全零!

Ampere cp.async SASS 反汇编实测【observed】

Section titled “Ampere cp.async SASS 反汇编实测【observed】”

我们在集群上编写包含 src-size 和 ignore-src 的 PTX 代码并反汇编:

Terminal window
ptxas -arch=sm_80 test_cp_async.ptx -o test_cp_async.cubin
nvdisasm test_cp_async.cubin

底层真实 SASS 输出【observed】:

/*00f0*/ LDGSTS.E.128.ZFILL [R7], [R4.64], P1 ;
/*0100*/ LDGSTS.E.BYPASS.128 [R7+0x10], [R2.64+0x10], !P0 ;
/*0110*/ LDGDEPBAR ;
/*0120*/ DEPBAR.LE SB0, 0x0 ;

指令解密【observed】:

  1. LDGSTS:Load Global Store Shared(硬件直通搬运助记符);
  2. .ZFILL:对应 src-size,硬件执行自动零填充;
  3. .BYPASS:对应 .cg,绕过 L1 Cache;后面的操作数 !P0 即为 ignore-src 谓词!
  4. LDGDEPBAR:对应 cp.async.commit_group;
  5. DEPBAR.LE:对应 cp.async.wait_group,基于计数的异步依赖屏障。

4.2 三代微架构演进:软硬件职责的全面下沉

Section titled “4.2 三代微架构演进:软硬件职责的全面下沉”
Volta / Turing (SM70 / SM75)
┌───────────────────────────────────────────────────────────────┐
│ 软件全包模式 │
│ 1. 软件线程各自计算 64 位全局地址与 smem 偏移 │
│ 2. 软件执行 branch / predicate 分流 │
│ 3. 硬件发射 LDG 加载数据至通用寄存器 (VGPR 中转) │
│ 4. 软件执行分支写入: 合法写 STS, 越界写 STS 0 (消耗巨额指令) │
└───────────────────────────────┬───────────────────────────────┘
│ 硬件进化: 引入异步直通与硬件补零
▼
Ampere (SM80: cp.async)
┌───────────────────────────────────────────────────────────────┐
│ 异步直通 + 线程级硬件补零 │
│ 1. 硬件支持 Global ──► Shared 直通 (绕过 VGPR 中转) │
│ 2. 硬件引入 LDGSTS.ZFILL 与 ignore-src (单线程粒度硬件补零) │
│ 3. 遗留缺陷: 每个线程依然必须独立算 64 位地址与谓词, │
│ 每个 Warp 依然必须发射 32 条 LDGSTS 指令! │
└───────────────────────────────┬───────────────────────────────┘
│ 硬件进化: 引入多维张量感知与集群组播
▼
Hopper (SM90: TMA & cuTensorMap)
┌───────────────────────────────────────────────────────────────┐
│ 描述符驱动 + 张量级硬件协处理 │
│ 1. 软件只需 1 个线程提交逻辑张量坐标 (UR 寄存器) │
│ 2. 硬件 AGU 自主展开多维地址与跨步跳转 │
│ 3. 硬件比较器统一裁定 Partial / Fully OOB 越界截断与填零 │
│ 4. 硬件 Crossbar 写入时自动完成 Bank 冲突 Swizzle │
│ 5. 硬件网络支持跨 CTA 集群组播 (Multicast) │
└───────────────────────────────────────────────────────────────┘

架构核心结论: TMA 的革命性提升,绝不仅仅是“做了一个更快的 memcpy”,而是把原本散落在 thread-level software 里的 tensor address computation、bounds checking、zero-fill 以及 shared-memory bank swizzle,完整、彻底地下沉到了专用的硬件搬运引擎!


第五部分:TMA 为什么是一个好抽象?(五个维度透视)

Section titled “第五部分:TMA 为什么是一个好抽象?(五个维度透视)”

5.1 Address Generation Offload(算术与寄存器卸载)

Section titled “5.1 Address Generation Offload(算术与寄存器卸载)”

在 Ampere 及之前,若要搬运一个 128×64128 \times 64 的 FP16 Tile:

  • Warp 内 32 个线程必须各自根据 threadIdx.x、行索引、列索引、Batch 步长执行大量的 IMAD(整数乘加)、LEA(地址加载)、ISETP(谓词比较)指令;
  • 每个线程需要霸占 4~8 个 32 位通用寄存器(VGPR)专门保存地址与谓词;
  • 每个 Warp 必须发射 32 条独立的内存指令,严重占满 SM 的指令发射槽位(Issue Slots)。

而在 Hopper TMA 下:

  • Warp 内通过 elect.sync 仅选派 1 个 Leader 线程 发射单条 TMA 指令;
  • 传入的坐标全为标量整数,在底层 SASS 中被分配至标量统一寄存器(Uniform Registers, UR)【observed】;
  • SM 整数算术流水线开销直降 90% 以上,通用寄存器(VGPR)被完全释放给 Tensor Core 累加器,SM 的调度器可以 100% 专注于高密度的矩阵乘计算。

5.2 Boundary Handling(不变量固化的全局边界)

Section titled “5.2 Boundary Handling(不变量固化的全局边界)”

海光 MLS 的惨痛教训证明:以裸指针加粗粒度行过滤去处理张量边界,是典型的“头痛医头”。张量维度的全局尺寸(M,N,KM, N, K)是张量在数学上的不可变元(Invariant)。 TMA 将 globalDim 显式写入描述符,使搬运引擎具备了全局空间拓扑意识。无论软件如何切分 Tile、如何跨步迭代,边界检查在硬件内部是一致、确定、不可穿透的。

5.3 Shared-Memory Layout Awareness(无冲突 Swizzle 原生集成)

Section titled “5.3 Shared-Memory Layout Awareness(无冲突 Swizzle 原生集成)”

Shared Memory 划分为 32 个 4 字节宽的 Banks。当 Tensor Core 执行 wgmma 读取矩阵操作数时,每条指令要求特定的数据排布。如果简单按行平铺,多个线程必定并发访问同一 Bank 导致严重的 Bank Conflict,吞吐骤降数倍。

在以往,算子工程师必须在 Shared Memory 写入时进行复杂的软件异或索引映射(XOR Swizzling),耗费昂贵指令。 TMA 将 swizzle 模式编码进描述符,搬运引擎直连 Shared Memory 交叉开关(Crossbar),在数据从显存总线写入片上内存的瞬间,硬件线路自动执行地址位异或重排。写入完成即是 Tensor Core 最佳布局,实现了真正的“零开销重排”。

5.4 Async Synchronization(基于 mbarrier 的事务计数流水线)

Section titled “5.4 Async Synchronization(基于 mbarrier 的事务计数流水线)”

传统异步加载使用 cp.async.commit_group 配合 __syncthreads(),整个 CTA 的所有线程必须周期性停顿等待数据就绪,造成巨大的流水线空泡(Pipeline Bubble)。

TMA 配合异步屏障 mbarrier,引入了革命性的 Transaction Counting(事务字节计数):

  1. 预告期望字节数:Leader 线程调用 mbarrier.arrive.expect_tx [mbar], bytes;,告知屏障即将到来的传输总量;
  2. 硬件自动打卡:TMA 硬件在后台搬运数据,每写入一部分并最终完成时,硬件自动向 mbarrier 发送 complete_tx 信号核减字节;
  3. 计算线程无锁轮询:消费者线程通过 mbarrier.try_wait 检查屏障状态。这彻底奠定了 Warp Specialization(生产者 Warp 专职搬运,消费者 Warp 专职计算) 的纯异步、无锁双缓冲流水线基石。

TMA + mbarrier 集群编译实测【observed】

Section titled “TMA + mbarrier 集群编译实测【observed】”

我们在集群上编译包含 mbarrier 与 2D TMA 的完整流水线,反汇编验证硬件指令:

/*0010*/ ELECT P0, URZ, PT ; // 选派 Leader 线程
/*0080*/ SYNCS.EXCH.64 URZ, [UR6], UR4 ; // mbarrier.init
/*0090*/ @P0 SYNCS.ARRIVE.TRANS64 RZ, [UR6], R0 ; // mbarrier.arrive.expect_tx (TRANS64 事务计数!)
/*00a0*/ @!P0 EXIT ;
/*0100*/ UTMALDG.2D [UR4], [UR8] ; // 硬件统一 TMA 2D 加载!

架构观察【observed】: TMA 加载指令在 SASS 中呈现为 UTMALDG.2D(Uniform Tensor Memory Access Load Global 2D),所有坐标与描述符操作数全部使用标量统一寄存器 UR4, UR8,彻底解放通用矢量寄存器!

5.5 Multicast(集群组播打破 HBM 带宽墙)

Section titled “5.5 Multicast(集群组播打破 HBM 带宽墙)”

在大型矩阵乘或多头注意力中,同一个 K/V Tile 或 Weight Tile 经常被属于同一个 Thread Block Cluster 的多个 CTA 同时需要。在传统架构下,每个 CTA 都必须各自独立向 HBM 发射读请求,重复消耗数倍的宝贵显存带宽。

TMA 原生支持 Cluster Multicast(集群组播):

  • 语法:cp.async.bulk.tensor....multicast::cluster [sdata], [desc, {coords}], [mbar], ctaMask;
  • SASS 指令【observed】:
    UTMALDG.2D.MULTICAST [UR8], [UR6], UR4 ;
  • 微架构行为【documented】:Cluster 内的一个 CTA 发射 TMA 加载指令,数据从全局显存读入后,通过片上专用的 Cluster Network 同时广播写入 ctaMask 指定的多个不同 CTA 的 Shared Memory 中!HBM 读带宽消耗直接降低到原来的 1N\frac{1}{N},彻底击碎片外带宽墙。

下表将通用软件加载、海光 MLS、Ampere cp.async 以及 Hopper TMA 放置在统一的技术标尺下进行全面对比:

评估维度 (Dimension)传统软件加载 (Volta/Turing)海光 MLS (Hygon DCU)Ampere cp.async (SM80)Hopper TMA (SM90)
地址抽象 (Address Abstraction)1D 线性虚拟地址(64 位指针)1D 虚拟基地址 + 1D 标量行步长1D 线性虚拟地址(64 位指针)多维张量坐标空间 (c0,c1,…,cR−1)(c_0, c_1, \dots, c_{R-1})
描述符抽象 (Descriptor Abstraction)无描述符(仅寄存器指针)128 位寄存器 SRD (vec4_uint),步长挤占了边界字段无硬件描述符(仅静态指令参数与运行时谓词)128 字节 Opaque 描述符 (CUtensorMap),64 字节对齐
地址展开者 (Who Computes Addr)软件 Warp 内 32 个线程各自通过 ALU 展开软件算好 64 位基地址,硬件内按 1D 步长展开软件 Warp 内 32 个线程各自通过 ALU 展开TMA 硬件专用 AGU 依据多维步长自动展开
张量阶数感知 (Tensor Rank Awareness)完全不知晓(Rank = Unknown)完全不知晓(Rank = Unknown)完全不知晓(Rank = Unknown)原生硬件感知 1≤Rank≤51 \le \text{Rank} \le 5
全局边界感知 (Bounds Awareness)无(完全依赖软件分支/谓词)无全局边界(无 extent,无法识别跨行越界)无(仅单线程粒度字节长度)硬件原生感知 globalDim(各维度独立边界比较器)
局部尾块处理 (Partial Tail Handling)软件执行谓词分流,显式写零仅粗粒度 nm_filter 行数过滤,不处理列尾块硬件支持 src-size < cp-size 自动写零硬件自动按元素截断物理请求,精确填充目标区域
全无效尾块 (Fully Invalid Tile)软件必须用分支绕过,否则非法访存必然发生 VM Fault 崩溃(MLS Bug 2)软件使用 ignore-src 绕过物理读取并置零硬件完全合法短路(无访存事务,自动填零/NaN)
硬件自动补零 (Zero Fill)不支持(软件循环显式写零)不支持(无物理置零机制,引发 MLS Bug 3)原生支持(.ZFILL 与 ignore-src)原生支持(CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)
OOB-NaN 语义支持无无无原生支持(NAN_REQUEST_ZERO_FMA + fma.oob)
共享内存布局感知 (Smem Layout)搬入后由软件手动执行 XOR Swizzle支持硬件转置与 4-interleave 打散不感知(仍需软件维护无冲突布局)原生集成 Swizzle(32B/64B/128B XOR 硬件重排)
通用寄存器压力 (VGPR Pressure)极大(每线程 4~8 个寄存器存地址)较小(基地址与步长保存在标量 SGPR)极大(每线程占用寄存器计算物理地址)极小(单线程发射,标量坐标全入 Uniform Register UR)
每线程地址指令数 (Addr Insns)极多(大量 IMAD, LEA, ISETP)极少(仅 Leader 线程或标量单元拼接 SRD)极多(每线程均发射 64 位地址算术)趋近于零(仅 Leader 线程发射单条多维 TMA 指令)
异步搬运机制 (Async Behavior)同步阻塞(必须经由 VGPR 中转)硬件直通 LDS(绕过 VGPR,标量控制)硬件直通 Shared Memory(绕过 VGPR,Warp 控制)独立硬件异步 DMA 协处理器(独立流水线执行)
同步原语 (Synchronization)全 CTA 显式 __syncthreads()标量等待或硬件计数等待基于阶段计数的 cp.async.wait_group基于事务字节计数的硬件 mbarrier(无锁流水线)
集群组播 (Multicast)不支持不支持不支持原生硬件组播(UTMALDG.2D.MULTICAST 跨 CTA 广播)
程序员主要职责 (Responsibility)计算 64 位物理地址、管理谓词、手动防越界拼接 SRD、计算行过滤、对全越界块手动重定向计算 64 位物理地址、计算 src-size、配置谓词在 Host 端构造一次描述符,Kernel 端仅提交逻辑坐标
典型失效模式 (Failure Modes)越界读脏显存、分支分流导致性能断崖低 32 位进位丢失(Bug 1)、尾块 VM Fault(Bug 2)物理地址溢出、src-size 大于 cp-size 产生未定义行为Host 端描述符参数构造非法(对齐/尺寸越限)

第七部分:“面试不露怯”硬核速查清单

Section titled “第七部分:“面试不露怯”硬核速查清单”

本清单专为应对最苛刻的架构与底层系统追问设计,每个问答 2~4 句,直击物理本质:

Q1:CUtensorMap 结构体多大?为什么定义成 opaque 数组?

Section titled “Q1:CUtensorMap 结构体多大?为什么定义成 opaque 数组?”

A:物理大小精确为 128 字节(1024 位),内存必须 64 字节对齐。设计为 opaque 是为了对应用层隐藏硬件内部的位域排布,使得微架构演进(如 Blackwell)时能自由重构描述符硬件格式而不破坏 ABI 兼容性,同时强制开发者通过驱动 API 进行严格的参数合法性前置校验。

Q2:谁来构造描述符?Host 端还是 Device 端?能在 Device 端修改它吗?

Section titled “Q2:谁来构造描述符?Host 端还是 Device 端?能在 Device 端修改它吗?”

A:绝大多数场景由 Host 端驱动 API(cuTensorMapEncodeTiled)在 Kernel 启动前一次性构造,并将其指针或副本传入 GPU 常量内存或全局显存。Device 端一般只读;若必须在 Device 端动态调整指针,CUDA 驱动提供了专用的内建指令更新基地址,但不允许随意篡改几何尺寸位域。

Q3:Kernel 端的执行线程到底向 TMA 指令提交什么?

Section titled “Q3:Kernel 端的执行线程到底向 TMA 指令提交什么?”

A:执行线程完全不需要提交 64 位物理内存地址,只需要提交目标张量的多维逻辑整数坐标(如 {coord_x, coord_y, coord_z}),以及目标共享内存的偏移与 mbarrier 地址。所有的多维乘加展开与物理总线地址生成全部由 TMA 硬件 AGU 在后台自动完成。

Q4:Partial OOB 和 Fully OOB 在硬件层面分别怎么处理?

Section titled “Q4:Partial OOB 和 Fully OOB 在硬件层面分别怎么处理?”

A:Partial OOB:TMA AGU 仅向合法显存区域发射 Burst 读事务,对越界区域抑制总线事务,并在目标 Shared Memory 中由硬件写入 0 或 OOB-NaN。Fully OOB:硬件判定整个分块完全越界,完全不向 L2/HBM 发起任何显存读请求(Zero Bus Transactions),内部管线直接向 Shared Memory 刷入 0 或 OOB-NaN,并向 mbarrier 核减事务计数,完全杜绝了非法缺页异常。

Q5:两种 Float OOB Fill 模式的完整枚举名是什么?有何本质差别?

Section titled “Q5:两种 Float OOB Fill 模式的完整枚举名是什么?有何本质差别?”

A:完整枚举名为 CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE 和 CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA。前者在越界时直接物理写入数值零(+0.0f);后者向 Shared Memory 写入特殊的 OOB-NaN 常量,用于保留“该槽位属于越界占位符”的信息,并要求下游的 fma.oob 指令在计算乘积时将其截断为零。

Q6:fma.oob 是什么?它和 OOB-NaN 是怎么配套的?

Section titled “Q6:fma.oob 是什么?它和 OOB-NaN 是怎么配套的?”

A:fma.oob 是 PTX 8.1 引入、专属于 CUDA Core 标量/向量半精度单元(SM90)的指令(SASS 呈现为 HFMA2.OOB)。当它的输入操作数检测到 TMA 填充的 OOB-NaN 时,计算硬件在电路层强行将该乘法项的积截断置为 +0.0f,从而在算术执行级彻底消除了 0.0×NaN=NaN0.0 \times \text{NaN} = \text{NaN} 的累加器毒化隐患。

Q7:Tensor Core(WGMMA)会自动吞掉 OOB-NaN 吗?

Section titled “Q7:Tensor Core(WGMMA)会自动吞掉 OOB-NaN 吗?”

A:不会!官方文档没有任何依据表明 WGMMA 支持 OOB-NaN 自动置零! WGMMA 指令没有任何 .oob 修饰符。在矩阵乘(GEMM)阶段,如果 TMA 使用了 NAN_REQUEST_ZERO_FMA,NaN 依然会毒化 Tensor Core 累加器。因此在 GEMM / Attention 矩阵乘流水线中,TMA 描述符必须无条件配置为 CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE(填 0)!

Q8:TMA 与 Ampere cp.async 的本质区别是什么?

Section titled “Q8:TMA 与 Ampere cp.async 的本质区别是什么?”

A:cp.async 本质上还是一个“线程级指令”——Warp 内 32 个线程必须各自用 ALU 算 64 位物理地址,各自发射指令。而 TMA 是一个“张量级异步协处理器”——Warp 只需派 1 个线程提交逻辑坐标,TMA 硬件在硅片层面全包了多维地址展开、边界判定、自动补零、Bank Swizzle 重排以及集群组播。

Q9:为什么 Swizzle(片上打散)必须由 TMA 搬运引擎来管?

Section titled “Q9:为什么 Swizzle(片上打散)必须由 TMA 搬运引擎来管?”

A:Shared Memory 存在 32 个 Banks,Tensor Core 读取矩阵时有严格的交织模式,极易发生严重 Bank 冲突。TMA 搬运引擎直连片上 Crossbar,在显存数据写入 Shared Memory 的同时,由硬件异或网络一步到位完成重排,使软件彻底省去了落入内存后的二次转置开销。

Q10:TMA Multicast 是如何工作的?它解决了什么痛点?

Section titled “Q10:TMA Multicast 是如何工作的?它解决了什么痛点?”

A:在一个 Thread Block Cluster 中,同一个数据 Tile(如 Attention 的 K/V)常被多个 CTA 共享。使用 TMA Multicast 时,单 CTA 发起 TMA 读取,数据到达片上后通过内部互联网络直接广播写入 ctaMask 指定的多个 CTA 的 Shared Memory 中,将片外 HBM 读带宽消耗直接降低到 1N\frac{1}{N}。


从 Volta/Turing 的纯软件手工算地址,到 Ampere cp.async 的线程级直通与硬件填零,再到 Hopper TMA 与 cuTensorMap 的多维张量认知与硬件级协处理,GPU 访存架构完成了一次从“一维线性裸内存搬运”向“高维张量代数原生加速”的伟大飞跃。

通过将张量几何(Rank / Extent / Strides)、片上内存拓扑(Swizzle)、边界语义(OOB Fill)与异步事务计数(mbarrier)固化进 128 字节的描述符中,NVIDIA 不仅从物理电路层彻底消灭了我们在上一篇海光 MLS 中所见的所有经典访存缺陷,更释放了 SM 通用核心的全部算力潜能。

在下一篇技术专栏中,我们将把视野正式切入基于 TMA、WGMMA、Warp Specialization 与 Cluster Multicast 的终极巅峰之作——FlashAttention-3 算法与系统级全景实现剖析,敬请期待。