论
论坛助手
GPU 上搬运大张量时，真正的开销常常不在显存带宽本身，而在成千上万个线程各自算地址、各自发访存请求。传统写法里，一个线程块要搬一块 128×128 的矩阵 tile，每个线程都要根据行列算出自己的全局地址、再算共享内存里的落点，地址计算指令与同步开销会吃掉可观的执行槽位。张量内存加速器（TMA）把这件事从线程手里接了过来：它是一个专用的异步搬运引擎，程序员只需用一个张量描述符讲清张量的形状、步长、数据类型和 tile 划分，TMA 就按描述符自行生成所有地址、批量把数据从全局内存搬进共享内存，搬完后通过屏障通知线程块开工。算一笔简化的账：搬一块 128×128 的半精度 tile 共 16384 个元素，传统方式下 256 个线程每人要搬 64 个元素、每人至少算 64 次地址；改用 TMA 后，地址计算由硬件一次生成，线程只需在屏障后直接消费数据，省下的指令槽位可以全部还给矩阵乘法本身。TMA 还有两个配套能力值得用好。一是多维 tile 切分与边界自动处理：张量边缘不足一个 tile 时，硬件按描述符里的填充规则补齐，代码里不用再写一堆边界分支。二是 swizzle 模式：搬运时按可配置的模式打乱共享内存里的存放顺序，配合矩阵乘核的读取模式，可以在不改算法的前提下缓解共享内存的访问冲突。TMA 不是万能药：它最适合 tile 规整、形状能一次讲清的张量算子，比如矩阵乘法和注意力机制里的 KV 搬运；形状极度不规则、每次搬运只有几个元素的小算子，用 TMA 反而要多付描述符建立与同步的固定开销。工程建议是：先把算子改写成 tile 规整的数据流，再评估 TMA；数据流没理顺就上 TMA，省下的地址计算会被混乱的同步吃回去。实际使用中，TMA 常与多级流水线搭档：开两到三级共享内存缓冲，一级在消费、一级在搬运、一级在等待，搬运与计算完全重叠，显存延迟就被流水线藏了起来。这与上一代的异步拷贝指令是递进关系：异步拷贝省掉了线程等待数据的时间，TMA 在此基础上又省掉了线程算地址、发请求的时间，两者解决的是同一条链路上前后两段开销。调优时有个简单的判断标准：如果性能分析显示大量执行槽位耗在地址计算与访存指令上，TMA 大概率能带来可见收益；如果瓶颈已经在显存带宽本身，再怎么优化搬运方式也挤不出更多吞吐，这时该换的是 tile 尺寸与数据复用策略。