Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects (2511.02132v1)

Abstract: 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.

• The paper proposes a NUMA-aware scheduling strategy, Swizzled Head-first Mapping, that realigns attention workloads with GPU memory topology to achieve up to 50% performance gains.
• It demonstrates that optimizing spatial locality leads to L2 cache hit rates of 80–97% across large head counts and sequence lengths on AMD MI300X.
• The approach requires minimal code changes and scales across both MHA and GQA, offering key insights for future AI accelerator designs.

NUMA-Aware Attention Optimization for Disaggregated GPU Architectures

Introduction and Motivation

The proliferation of large-scale transformer models has driven the evolution of GPU architectures toward disaggregated, chiplet-based designs to meet the demands of compute and memory bandwidth scaling. This architectural shift introduces pronounced non-uniform memory access (NUMA) effects, fundamentally altering the memory hierarchy and access patterns within the GPU. The paper "Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects" (2511.02132) systematically analyzes these NUMA effects in the context of attention workloads and proposes a NUMA-aware scheduling strategy—Swizzled Head-first Mapping—that aligns computational work with the physical memory topology of multi-die GPUs, specifically AMD's MI300X.

Figure 1: Evolution of GPU architectures from monolithic to multi-die chiplet designs, highlighting the emergence of NUMA effects in modern AI accelerators.

The work demonstrates that conventional attention kernel scheduling, which assumes uniform memory access, is suboptimal on such architectures. By exploiting the spatial locality inherent in attention mechanisms, the proposed mapping strategy achieves up to 50% higher performance and sustains L2 cache hit rates of 80–97% on MI300X, establishing NUMA-aware scheduling as a critical optimization for next-generation AI accelerators.

NUMA Effects in Modern GPU Architectures

Disaggregated GPU architectures, such as AMD MI300X and NVIDIA Blackwell, partition compute and memory resources across multiple chiplets (dies), each with private L2 caches and memory controllers. This design introduces a spatial hierarchy of memory access costs: local accesses to a die's own memory and cache are low-latency and high-bandwidth, while remote accesses incur significant penalties due to inter-die communication.

Figure 2: Workgroup scheduling in multi-die chiplet architectures can fragment cache locality, leading to redundant memory fetches and reduced cache efficiency.

The default round-robin workgroup scheduling policy, while ensuring load balance, can inadvertently fragment spatial locality, causing workgroups that process adjacent tensor regions to be dispatched to different dies. This results in redundant memory fetches and poor cache utilization, as each die's L2 cache is private and not shared across chiplets.

To address this, the paper leverages swizzling—algorithmic remapping of workgroup IDs—to co-locate spatially adjacent computational work within the same NUMA domain, thereby maximizing cache reuse and minimizing off-chip memory traffic.

Spatial Locality in Attention Mechanisms

The attention mechanism, particularly in the context of multi-head attention (MHA) and grouped query attention (GQA), exhibits strong spatial locality. In FlashAttention2 (FA2), the query matrix $Q$ is partitioned into row blocks, with each workgroup processing a block and accessing the full key ($K$) and value ($V$) tensors. All workgroups within the same attention head (or GQA group) share $K$ and $V$, creating natural opportunities for cache reuse.

Figure 3: FlashAttention2 partitions the query matrix into row blocks, with each workgroup accessing the full $K$ and $V$ tensors, establishing spatial locality within attention heads.

This data-sharing pattern persists in both forward and backward passes, motivating the definition of an Attention Compute Cluster (ACC): the set of workgroups that share $K$ and $V$ tensors. In MHA, each head forms an ACC; in GQA, each group of heads sharing $K$ and $V$ forms an ACC.

Figure 4: ACC organization in transformer attention—each head (MHA) or group (GQA) forms an ACC, with shared $K$ and $V$ tensors.

The key insight is that co-locating all workgroups of an ACC on the same die maximizes L2 cache utilization, as shared data is loaded once and reused, avoiding cache fragmentation and redundant memory accesses.

Mapping Strategies for Attention Kernels

The paper evaluates four mapping strategies for assigning workgroups to chiplets:

• Naive Block-first: Iterates over all heads for each row block, assigning workgroups in round-robin fashion, splitting ACCs across dies.
• Swizzled Block-first: Applies swizzling to co-locate GQA groups within dies, effective only when the number of groups matches the die count.
• Naive Head-first: Iterates over all blocks of a head before moving to the next, but still stripes heads across dies.
• Swizzled Head-first (proposed): Assigns all blocks of a head (or ACC) to a single die before moving to the next, ensuring each die services one ACC at a time.

  Figure 5: Naive Block-first mapping splits ACCs across dies, reducing cache efficiency.

  

  Figure 6: Swizzled Block-first mapping co-locates GQA groups within dies, effective when group count matches die count.

  

  Figure 7: Naive Head-first mapping stripes heads across dies, still splitting ACCs.

  

  Figure 8: Swizzled Head-first mapping confines each ACC to a single die, maximizing cache locality.

The Swizzled Head-first mapping is implemented with minimal code changes using workgroup ID swizzling, as shown in the provided Triton pseudocode. This approach is robust to varying batch sizes, head counts, and sequence lengths, and generalizes across both MHA and GQA.

Empirical Evaluation

Multi-Head Attention Sensitivity

The Swizzled Head-first mapping consistently outperforms other strategies across a wide range of configurations on MI300X. For high head counts ($H_Q \geq 64$) and long sequences ($N_{CTX} \geq 32K$), block-first approaches degrade to 64–70% of the Swizzled Head-first baseline, with performance gaps reaching 50% at $H_Q = 128$, $N_{CTX} = 128K$.

Figure 9: MHA performance relative to Swizzled Head-first baseline; block-first approaches show widening performance gaps at scale.

L2 cache hit rates corroborate these results: Swizzled Head-first maintains 90–96% hit rates under extreme configurations, while block-first approaches collapse to ~1%, and Naive Head-first drops to 40–60%.

Figure 10: L2 cache hit rates for MHA; Swizzled Head-first sustains high hit rates, while others degrade sharply at scale.

Grouped Query Attention

For GQA (e.g., Llama-3 models), both Swizzled Head-first and Swizzled Block-first perform well when the number of KV heads matches the die count. Naive Block-first degrades significantly at high query head counts and long sequences due to cache splitting.

Figure 11: GQA performance normalized to Swizzled Head-first; block-first mapping degrades at high head counts and long sequences.

Real-World Case Study: DeepSeekV3

In DeepSeekV3's prefill phase (128 heads), Swizzled Head-first achieves the highest performance, with block-first approaches dropping below 65% efficiency at 128K tokens.

Figure 12: DeepSeekV3 prefill performance; block-first approaches degrade substantially at long sequence lengths.

Backward Pass

Swizzled Head-first also provides consistent speedup in the FlashAttention2 backward pass, with increasing gains at longer sequences (up to 1.10× at 128K tokens).

Figure 13: Backward pass speedup of Swizzled Head-first over Naive Block-first, increasing with sequence length.

Implementation Considerations

• Minimal Code Changes: The swizzling logic can be implemented in a few lines of Triton code, making it practical to integrate into existing attention kernels.
• Portability: The approach is robust to changes in hardware scheduling policies and generalizes across MHA and GQA.
• Resource Requirements: The method is most beneficial when the number of ACCs (heads or groups) is much larger than the number of dies, as in large LLMs.
• Scalability: The performance benefits increase with head count, sequence length, and batch size, making the approach essential for scaling to future model and hardware generations.
• Limitations: For GQA, Swizzled Block-first is effective only when the number of groups matches the die count; otherwise, Swizzled Head-first is preferred.

Implications and Future Directions

This work establishes that NUMA-aware kernel design is fundamental for efficient utilization of chiplet-based AI accelerators. As GPU architectures continue to scale via disaggregation, spatially-aware scheduling will become increasingly critical for all memory-bound workloads, not just attention. The demonstrated performance gains suggest that similar strategies should be applied to other key primitives (e.g., GEMM, convolution) and that hardware/software co-design for explicit NUMA management will be a major area of research.

Potential future directions include:

• Extending NUMA-aware mapping to distributed multi-GPU systems.
• Dynamic runtime adaptation of mapping strategies based on workload and hardware topology.
• Hardware support for programmable workgroup scheduling and cache management.

Conclusion

The paper provides a comprehensive analysis of NUMA effects in modern GPU architectures and introduces a practical, high-impact optimization for attention workloads. By aligning computational work with the physical memory topology via Swizzled Head-first Mapping, it achieves substantial performance and cache efficiency gains on AMD MI300X. The results underscore the necessity of hardware-aware algorithm design as chiplet-based architectures become the standard for AI acceleration. This work will inform both future kernel development and hardware design for scalable, efficient AI systems.