Tensor memory accelerator usage rules improve GPU kernel performance
When and Why to Use TMA: A Workload-Grounded Characterization of Tensor Memory Acceleration on NVIDIA Hopper GPUs
Distributed, Parallel, and Cluster ComputingHardware Architecture
Summary
Modern NVIDIA GPUs have a feature called the Tensor Memory Accelerator (TMA) that can speed up data loading in complex computations. But experts don’t always use TMA for every part of their work, switching to other methods depending on the situation. The paper studies when TMA and related techniques are best used, based on real-world workloads. This helps programmers decide which memory method to pick to get faster results on NVIDIA’s Hopper and Blackwell GPUs.
What this means in practice
- •For gpu kernel developers: Optimize memory load strategies in GPU kernels for better performance on NVIDIA Hopper and Blackwell architectures using the paper’s concrete rules.
- •For ai infrastructure teams: Improve efficiency of large-scale transformer model inference by choosing the right memory acceleration technique in GPU workloads informed by this workload-grounded study.
Authors
Christin David Bose, Cesar Avalos, Junrui Pan, Timothy Rogers
Abstract
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.