Skip to main content

Cuda Tensor Core

要保持张量核心持续运行并不容易。

研究人员发现GPU硬件具有一些特性,对于保持矩阵乘法的运行非常重要:

  • WGMMA指令虽然是必要的,但使用起来颇为麻烦。

  • 共享内存的速度并不如预期的快,使用时还需格外注意。

  • 生成地址的成本较高。

  • 保持高占用率对于提升性能是有益的,寄存器至关重要

早期GPU中的张量核心指令如wmma.mma.sync和mma.sync,要求SM一个子单元内的32个线程的一个warp同步传输数据块至张量核心并等待结果。

wgmma.mma_async指令则不同。它允许128个连续线程跨SM所有子单元协作同步,并从共享内存及寄存器(可选)异步启动矩阵乘法。这使得这些warp在等待矩阵乘法结果时可以处理其他任务。

然而,这些指令的内存布局极其复杂。未重排的共享内存布局合并性差,需要额外的L2带宽。重排的内存布局记录不准确,研究人员花费了大量时间才弄明白。最终发现,这些布局只适用于特定矩阵形状,并与wgmma.mma_async指令的其他部分不兼容,例如硬件仅在未重排的布局下转置子矩阵。

此外,未重排的wgmma布局内存合并性差且有bank conflicts。尽管TMA和L2缓存在如flash attention这类内核上能较好地掩盖这些问题,但要充分利用硬件,必须精心控制内存请求的合并和避免bank conflicts。

TMA允许用户在全局和共享内存中指定多维张量布局,命令其异步提取张量的一部分,并在完成后触发一个屏障。这大大节省了地址生成的开销,并简化了pipelines的构建。

 

 VoltaTuringAmpereHopper
Base Size*64个F16(4x4x4)64个F16(4x4x4)256个F16(8x4x8)512个F16(8x4x16)
tensorCore个数672(8/SM)576(8/SM)512(4/SM)576(4/SM)
SM内算力*512个F16512个F161024个F162048个F16
总算力 FOPS(F16)125T130.5T312T1000T
boost clock(MHz)153014551410≈A100x1.3
FP64--
TF32--
BF16--
FP16
FP8---
INT8/UINT8-
INT4/UINT4--
INT1--
VoltaTuringAmpereHopper
SM/TPC2222
process blocks/SM4444
FP64/PB8-816
FP32/PB16161632
INT32/PB16161616
tensorCore/PB2211
LSU/PB8488
register file64KB(16384*32bit)64KB(16384*32bit)64KB(16384*32bit)64KB(16384*32bit)
L0 ICache/PB1111
L1/SHM/SM128KB96KB192KB256KB
warp scheduler/PB1(32thread/clk)1(32thread/clk)1(32thread/clk)1(32thread/clk)
dispatch uint/PB1(32thread/clk)1(32thread/clk)1(32thread/clk)1(32thread/clk)

image.png

tensorCore的启动是依靠SM内的thread(warp)共同完成的,每个thread能提供的数据个数是相同的,因此统计这个算力能有助于分析SM内部的thread和tensorCore之间的调度方式