AMD's NUNA tech cuts GPU collectives 1.8x, boosting ML inference speed

As GPUs scale up to meet compute and memory demands, intra-socket wire transfer delays are growing. Previous research focused on compute and memory locality within a socket, but the spatial impact on inter-GPU communication was largely unexplored. AMD researchers (Conor James Green, William Won, Tuan Ta, Bradford M. Beckmann) introduce NUNA (non-uniform network access) to characterize this emerging dimension. They target latency-sensitive collective communication, which is critical for machine learning inference. The paper proposes two complementary optimizations: NUNA-aware placement (NAP), which places threadblocks and data near I/O to reduce inter-GPU traffic, and NUNA-aware routing (NAR), which chooses optimized, spatially-aware inter-GPU paths in large scale-up network topologies.

Across 15 figures and 2 tables, the results are compelling. NAP alone achieves up to 1.5x collective speedups over a locality-unaware baseline, while combining NAP with NAR yields up to 1.8x faster collectives. On real ML inference workloads, this translates to a 7% mean (28% max) speedup in time-per-output-token. The paper highlights that as GPU systems adopt multi-die scale-up architectures, software must become NUNA-aware to avoid leaving significant performance on the table. This work lays a foundation for future runtime schedulers, network topology design, and compiler optimizations that treat inter-GPU communication topology as a first-class citizen.