针对多芯片显卡的内存延迟差异，提出新型调度策略提升注意力计算效率。
Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects
- 按GPU NUMA域布局注意力头，利用片内缓存复用
- 在MI300X上性能提升最高达50%，缓存命中率80%-97%
- 适合大规模AI训练与推理场景下的显卡优化
随着拆分式AI显卡的发展，大规模注意力计算面临非统一内存访问（NUMA）的关键瓶颈。当多芯片设计成为扩展计算能力的标准时，不同计算区域的内存延迟和带宽差异显著，破坏了传统显卡内核调度策略所依赖的均匀内存访问假设。我们揭示了NUMA效应如何扭曲多头注意力（MHA）中的局部性，并提出一种空间感知的调度策略——混洗头优先映射（Swizzled Head-first Mapping），将注意力头对齐至GPU NUMA域，以利用片内缓存复用。在AMD MI300X架构上，该方法相比现有基于常规调度的先进注意力算法，性能最高提升50%，并保持80%-97%的稳定L2缓存命中率。结果表明，NUMA感知调度已成为在下一代拆分式显卡上实现全效性能的基础，为可扩展的AI训练与推理提供了可行路径。
原文摘要 · Abstract (English)
The rise of disaggregated AI GPUs has exposed a critical bottleneck in large-scale attention workloads: non-uniform memory access (NUMA). As multi-chiplet designs become the norm for scaling compute capabilities, memory latency and bandwidth vary sharply across compute regions, undermining the performance of traditional GPU kernel scheduling strategies that assume uniform memory access. We identify how these NUMA effects distort locality in multi-head attention (MHA) and present Swizzled Head-first Mapping, a spatially-aware scheduling strategy that aligns attention heads with GPU NUMA domains to exploit intra-chiplet cache reuse. On AMD's MI300X architecture, our method achieves up to 50% higher performance over state-of-the-art attention algorithms using conventional scheduling techniques and sustains consistently high L2 cache hit rates of 80-97%. These results demonstrate that NUMA-aware scheduling is now fundamental to achieving full efficiency on next-generation disaggregated GPUs, offering a path forward for scalable AI training and inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。