设备硬件

Tensor MemoryTensor Memory

Tensor Memory（张量内存）是某些 GPU（如 B200）的流式多处理器 (SM) 中的一种专用内存，用于存储 Tensor Core 的输入和输出。

张量内存的访问受到严格限制: 数据必须在线程束组 warpgroup 中的四个线程束 (warp) 中协同移动，并且仅支持特定模式的内存操作，例如在张量内存与寄存器 (register) 之间传输数据、将共享内存 (shared memory) 写入张量内存，或向 Tensor Core 发送矩阵乘加（MMA）指令以使用张量内存存储特定操作数。这便是 "统一计算"设备架构的复杂之处！

具体来说，对于计算 D += A @ B 的 tcgen05.mma 并行线程执行 (Parallel Thread eXecution) 指令若要使用张量内存，"累加器"矩阵 D 必须位于张量内存中，左侧矩阵 A 可以位于张量内存或共享内存 (shared memory) 中，而右侧矩阵 B 必须位于共享内存 (shared memory) 中，而非张量内存。这种设计虽复杂但并非随意——在矩阵乘法运算期间，累加器的访问频率远高于分块矩阵（tiles），因此更能从专用硬件中获益（例如 Tensor Core 与张量内存之间更短、更简单的布线）。请注意，所有矩阵都不在寄存器 (register) 中。

注意：张量内存与 Tensor Memory Accelerator 没有直接关系，后者是将数据加载到 L1 数据缓存中。大致流程为：仅当执行张量核心 Tensor Core 操作时，数据才会从 L1 缓存移至张量内存，后续需显式移出以进行后续处理，例如神经网络中矩阵乘法之后的非线性激活。

有关张量内存及其在矩阵乘法中的使用模式详情，请参见 GTC 2025 的《使用 CUTLASS 编程 Blackwell Tensor Core》演讲。

英文原文(Modal GPU Glossary)

Tensor Memory is a specialized memory in the Streaming Multiprocessor (SM) of certain GPUs, like the B200, for storing the inputs and outputs of Tensor Cores.

Tensor Memory access is highly restricted. Data must be moved collectively by four warps in a warpgroup, and they can move memory only in specific patterns between Tensor Memory and registers, write shared memory to Tensor Memory, or issue matrix-multiply-accumulate (MMA) instructions to Tensor Cores that use Tensor Memory for specific operands. So much for a "compute-unified" device architecture!

Specifically, for a tcgen05.mma Parallel Thread eXecution instruction computing D += A @ B to use Tensor Memory, the "accumulator" matrix D _must_ be in Tensor Memory, the left-hand matrix A _may_ be in Tensor Memory or shared memory, and the right-hand matrix B _must_ be in shared memory, not Tensor Memory. This is complex, but not arbitrary -- accumulators are accessed more frequently during matmuls than are the tiles, so they benefit more from specialized hardware, e.g. from shorter, simpler wiring between the Tensor Cores and the Tensor Memory. Note that none of the matrices are in the registers.

Beware: Tensor Memory is not directly related to the Tensor Memory Accelerator, which instead loads into the L1 data cache. Roughly speaking, data is moved from that cache into Tensor Memory only as a result of a Tensor Core operation and then is explicitly moved out for post-processing, e.g. the non-linearity after a matrix multiplication in a neural network.

For details on tensor memory and patterns for its use in matrix multiplications, see the _Programming Blackwell Tensor Cores with CUTLASS_ talk from GTC 2025.

相关词条