When and Why to Use TMA: A Workload-Grounded Characterization of Tensor Memory Acceleration on NVIDIA Hopper GPUs
Christin David Bose, Cesar Avalos, Junrui Pan, Timothy Rogers
The Tensor Memory Accelerator (TMA), introduced on NVIDIA Hopper, is used pervasively in production GPU kernels yet its usage is conspicuously selective: expert kernel developers in CUTLASS, FlashAttention-3, and FlashInfer apply TMA to some tensor operands within a kernel but fall back to cp.async, direct global loads, or distributed shared memory (DSM) for others. Vendor documentation describes the API, and recent microbenchmarking work characterizes TMA throughput in isolation, but neither answers the question practitioners actually face: given a memory operand in a kernel, should it use TMA, and if so, with or without cluster multicast? We address this gap with a workload-grounded characterization study on H100. We catalog the TMA, multicast, DSM, and cp.async choices made inside FlashAttention-3 forward, isolate the discriminating features through targeted microbenchmarks (transfer-size crossover, multicast scaling), and produce rules of thumb for the load-mechanism choice on Hopper. We further compare TMA multicast against DSM as competing cluster-broadcast primitives, and isolate the FlashAttention-3 paged-key-value (KV) TMA-vs-cp.async dispatch on H100 across page size, datatype, model shape, and execution phase. Our results yield concrete, predictive guidance on when TMA pays off and when multicast helps or hurts on Hopper, and we extend the analysis to server Blackwell (SM100) via a CUTLASS GEMM study on B200 that shows the qualitative dispatch rules carry across generations while absolute crossovers shift with tensor-core and bandwidth scaling.
尚无 AI 解读，展示原文摘要