Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866912686726447104 |
|---|---|
| author | Choudhary, Mansi Sangaiah, Karthik Singh, Sonali Osama, Muhammad Wills, Lisa Wu Dasika, Ganesh |
| author_facet | Choudhary, Mansi Sangaiah, Karthik Singh, Sonali Osama, Muhammad Wills, Lisa Wu Dasika, Ganesh |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2511_02132 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects Choudhary, Mansi Sangaiah, Karthik Singh, Sonali Osama, Muhammad Wills, Lisa Wu Dasika, Ganesh Hardware Architecture Distributed, Parallel, and Cluster Computing Machine Learning Performance 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. |
| title | Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects |
| topic | Hardware Architecture Distributed, Parallel, and Cluster Computing Machine Learning Performance |
| url | https://arxiv.org/abs/2511.02132 |