Nvidia TensorCore

Introduction
TensorCore 历史概述
序号
时间
架构
异步数据搬运机制
编程模型 & 计算指令
数据类型
里程碑意义
1
2017
Volta (V100)
无异步机制
ld./st.
Warp 级 wmma 指令
FP16
开启混合精度计算，第 1 代 TensorCore

2
2018
Turing (T4)
同 Volta 架构
同 Volta 架构

• int8, int4, int1
  低精度推理，推动 AI 推理普及
  3
  2020.05
  Ampere (A100)
  软件管理的异步拷贝
  cp.async.*
  Warp 级 mma.sync，可编程 Thread 级数据排布
• fp32, bf16
  大幅提升 FP32 训练性能，结构化稀疏

4
2022.03
Hopper (H100)

硬件 TMA
cp.async.bulk.tensor.*
thread block 级 wgmma & 硬件 TMA
引入 Transformer 引擎

• fp8
  引入FP8和动态精度切换，为 LLM 训练优化
  5
  2024.03
  Blackwell (B200)
  专用的 Tensor Memory
  跨SM协作的 UTCMMA.2SM & Tensor Memory
  第 2 代 Transformer 引擎
• fp4, fp6

专为万亿参数LLM推理和训练设计

演进的核心脉络

3 个关键优化点：

1. 内存能力。核心是异步数据搬运机制 & 内存层级设计。AsyncCopy -> 硬件 TMA -> Tensor Memory。
2. 数据格式。越来越丰富，特别是 FP8 / FP4 的引入。
3. 编程模型。提升并行层级：Warp (mma) => Thread Block (wgmma) => 跨 SM (UTCMMA.2SM)。

计算指令的演进

待补充：指令 shape & 单指令算力 flops。

AI Compiler，比如 Triton。不需要支持 wmma 指令。用 mma 足以取代 wmma 的功能。

MMA & WMMA - A100

都是 warp level 的执行模型。
H100 之前，并存的 2 个主要指令。

MMA 指令特性:

• 线程级控制：支持手动控制线程级别的数据排布和寄存器分配。
• 灵活性高，性能好。支持自定义数据布局优化。
• 编程复杂度：高，需要显式管理寄存器分配和数据流水。
  WMMA 指令特性
• 封装抽象：封装了完整的 Shared Memory → MMA → Shared Memory 流程。
• 易用性优先：降低编程门槛，适合快速原型开发。
• 灵活性限制：难以规避 SMEM bank conflict，性能优化空间有限，无法自定义数据排布。

尽管 mma 是通过一个 warp 去调度实现的，但还是粒度为线程级别时更容易编写。
以下场景 WMMA 无法实现，但 MMA 支持：

1. 数据排布优化：自定义Shared Memory → MMA的swizzle layout
2. 计算融合：将elementwise操作直接融合到MMA累加寄存器
3. 复杂数据重用：实现更高效的数据局部性优化

指令文档：https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-matrix-functions

wmma 内部分成了如下多条指令
void load_matrix_sync(fragment<…> &a, const T* mptr, unsigned ldm);
void load_matrix_sync(fragment<…> &a, const T* mptr, unsigned ldm, layout_t layout);
void store_matrix_sync(T* mptr, const fragment<…> &a, unsigned ldm, layout_t layout);
void mma_sync(fragment<…> &d, const fragment<…> &a, const fragment<…> &b, const fragment<…> &c, bool satf=false);

*_sync 类的指令，等待所有 warp lanes 都 barrier.arrived。
load_matrix_sync 和 mma_sync，需要 warp 内所有 thread 都调用，否则结果未定义。
store_matrix_sync 不要求 warp 内所有 thread 都调用。（比如，warp reduce 之后，只需要 store 1 element）

WGMMA - H100

H100 引入 wgmma，4 个连续的 warp 可以组成一个 warp group。
只有 TMA + wgmma 才能打满 TensorCore 理论算力。

内存层次优化
传统路径: SMEM → RF → TensorCore (存在 RF 瓶颈)
WGMMA路径: SMEM → TensorCore (直接访问 SMEM，消除 RF 瓶颈)

核心优势：

• 消除瓶颈：绕过寄存器文件，直接 SMEM → TensorCore
• 能效提升：减少数据移动，降低指令开销和功耗
• 编程简化：减少显式数据搬移指令
• 规模扩展：支持更大矩阵运算，warp group 协同计算

From https://llvm.org/devmtg/2025-04/slides/technical_talk/ozen_blackwell.pdf
Page 21

PTX 中有 4 条 WGMMA 相关的指令:

• wgmma.mma_async。异步的计算指令，执行时间不确定（因为要读 shared memory）
• wgmma.fence。确保 fence 前所有 MMA 完成后再发射后续 MMA。(比如多条 mma 在 K 方向累加)
• wgmma.commit_group。将之前发射的异步 MMA 打包成一个 group 提交。便于跟踪一组相关计算任务。
• wgmma.wait_group。等待指定数量 N 的未完成计算组。N=0 等待所有组完成。N 是常数。

完整执行流程
// 1. 数据准备
Load matrices A, B, D to registers or shared memory;

// 2. 内存同步（fence operations）
wgmma.fence; // 确保 warp group 内 register/shared-memory 数据可见
fence.proxy.async; // 使通用代理操作对 async proxy 可见

// 3. 异步计算
wgmma.mma_async; // 在 async proxy 中 Issue 矩阵乘加

// 4. 组提交
wgmma.commit_group; // 打包提交所有未决的 MMA 操作。Create a wgmma-group and commit all the prior outstanding wgmma.mma_async operations into the group

// 5. 等待完成
wgmma.wait_group; // 等待指定组完成

// 6. 结果就绪
// 所有 wgmma.mma_async 操作已完成

关键约束条件

• A 矩阵：Register 或 Shared Memory
• B 矩阵：必须从 Shared Memory 直接读取
• 矩阵形状（fp16 示例）：
  • M = 64 (固定)
  • K = 16 (固定)
  • N = 8~256 (8的倍数，编译时常数，设置在指令 encoding 中。不能运行时修改)

SASS 层是 HGMMA 指令。

推测：一条 wgmma 执行时，在一个SM内部，矩阵A是被切分为4组，每一组是16*16，分别执行在对应4个tensor core上；而矩阵B是需要被4个tensor core共享的，所以它只能放在 shared memory 上。
如果采用 mma 指令，则需 4 条 MMA.m16n8k16 指令。所以，wgmma 可以减少指令数。
// 一条Tensor core PTX指令wgmma(64n8k16)
wgmma.mma_async.sync.aligned.m64n8k16.f16.f16.f16

// 对应三条SASS指令
WARPGROUP.ARRIVE
HGMMA.64x8x16.F16
WARPGROUP.DEPBAR.LE

UTCMMA.2SM - B200

B200 引入。

Blackwell GEMM https://mp.weixin.qq.com/s/HgXWsNa5xkKQ_gmNgmRvbQ

TensorCore 架构设计 - 原理科普

每 cycle 完成的计算：D = A * B + C

常见的，matmul 算子的输入 & 输出，都是 FP16 类型。
底层的 PTX 指令，一般是 fp16 输入 + fp32 输出，以避免累加时精度丢失。

向量乘 FMA：1 行 * 1 列

FMA 指令，实现 D = A * B + C

绿色是 register。

1. 横的是 16 bit-register，存 input or 中间数据
2. 竖的是 32 bit-register，存 累积结果 or 中间高精度数据

2 个核心加速点：

1. 所有的乘法，可以并行触发，并行计算。单周期完成。
2. 所有的加法，树形归约 Reduction Tree。N 个 加法，只需要 $$log_{2}^{N}$$ 步完成。看硬件的频率设定，1 个 cycle 可能完成多步。

矩阵乘 MMA - M 行 * N 列

1. 1 个蓝色方块，就是上面的 一行 * 一列的 电路。每个蓝色方块对应一个输出位置。输出 shape 是 4*4 矩阵。
2. A 和 B 矩阵的 register，也按二维矩阵设计。每个 cycle，广播 A 的一行，B 的一列到对应的蓝色方块上。
3. 所有 16 个蓝色方块，可在同一 cycle 启动计算，完成各自的 FMA 计算。

阵列的核心加速点：单周期完成数据分发。
硬连线的广播网络，无冲突的数据访问。A 行数据→水平广播，B 列数据→垂直广播。

     A 矩阵寄存器 (FP16)                  B 矩阵寄存器 (FP16)
          ↓  ↓  ↓  ↓                        → → → →
       +-----------------+              +-----------------+

Row 0 → | A00 → PE00/01/02/03 | | B00 → PE00/10/20/30 | ← Col 0
Row 1 → | A10 → PE10/11/12/13 | | B11 → PE01/11/21/31 | ← Col 1
Row 2 → | A20 → PE20/21/22/23 | | B22 → PE02/12/22/32 | ← Col 2
Row 3 → | A30 → PE30/31/32/33 | | B33 → PE03/13/23/33 | ← Col 3
+—————–+ +—————–+

Nvidia TensorCore 分析

Volta - V100 架构

调度流程

Volta 架构，1 个 SM 有 4 个 sub-core。
每个 sub-core 有：

• 2 个 TensorCore
• 1 组 CudaCore。（图中的 FP 和 INT 单元）
• 1 个 warp scheduler，跑 1 个 warp。（图中 subcore 图的第 2 行）

sm 内的执行流程：

1. Warp Scheduler 向 Tensor Core 发送矩阵乘法 GEMM 运算指令
2. Tensor Core 接收来自寄存器文件的输入矩阵（A、B、C），
3. 执行多次 4x4x4 矩阵乘法操作，直至完成整个矩阵乘法
4. 将结果矩阵写回寄存器文件(Register File)

subcore 内部的架构，如下图
顶部依然是 L1 Cache，紧随其后的是 L0 Cache，也就是 Register File。
Register File 负责将数据传输到 Warp Scheduler，然后是 Warp Scheduler 做指令调度。

1. CUDACore 的计算，Warp Scheduler 通过 Math Dispatch Unit 分发指令。
2. TensorCore 的计算，Warp Scheduler 直接触发。

subcore 的 2 个 TensorCore 并行计算，结果写回 Register File。
Register File 通过 MIO 的 Data Pipeline 与 Shared Memory 通讯。
[图片]

Tensor Core 微架构

参考文章 https://arxiv.org/pdf/1811.08309
逆向工程了 V100 tensor core 的设计。

V100 有 4 个 SubCore。

每个 SubCore 有 2 个 TensorCore。
每个 Tensor Core 处理 8 个 thread，每个 thread 拥有自己的寄存器。
需要 2 个 cycle 才能用 2 个 tensor core 处理完 32 thread 的一组 mma？
每个 Tensor Core / thread 每次执行 4x4x4 的 FP16 矩阵乘。不支持其它 data type。

因此，在 8 个时钟周期内，可以执行 1024 次 MAC 操作。

每个 subcore，可以通过软件的方式将其扩展到 16x16x16 的规模。
通过局部数据的搬运来提升计算效率，但这并不意味着我们能够轻松地处理所有嵌入的向量或大矩阵。

要探索Tensor Core的微架构，本质上就是要回答两个问题：
（1）1616 的矩阵乘法所需的数据 A,B,C 是如何分配给 1 个 warp 内的 32 个线程的？
（2）Tensor Core具体是按照什么顺序来分步执行完这个1616的矩阵乘法的，耗费了多少个周期？
分别可以通过一段简单的CUDA程序与SASS反汇编来获得。参考文章 https://arxiv.org/pdf/1811.08309
（octet 是 8 位字节）

Ampere - A100 架构

Tensor Core 进行了优化，1 个 Warp 提供了 32 个线程。
Volta 每个Tensor Core 只有 8 个线程。
这样的设计可以减少线程之间的数据搬运次数，从而进一步提高计算效率。

每个 Tensor Core 可以处理 32 条线程，
因此在 8 个时钟周期内，可以寄存 2048 次MAC操作，每个时钟周期处理其中一块的数据。

Tensor core 的指令是 wmma，32 线程间的数据是共享的，不同于 FFMA 指令。
// 一条 Tensor core PTX 指令 wmma(m16n16k16)
wmma.mma.sync.aligned.row.col.m16n16k16.f16.f16 \