Spatial-Aware Orchestration of LLM Attention on Waferscale Chips
摘要
Transformer-based LLMs have driven AI progress but face mounting computational challenges as their scale and context lengths expand. Wafer-scale chips (WSCs) offer compelling alternatives to conventional GPUs with substantially higher transistor density and inter-die bandwidth. However, their rigid 2D mesh topology undermines GPU-optimized ring-attention patterns, while causal attention masks create workload imbalances that conventional token reordering fails to address in wafer-scale settings. We analyze communication overhead in LLM attention blocks on WSCs and develop a spatial-aware cost model tailored to wafer-scale topologies. Our spatial-aware orchestration method optimizes communication patterns and strategically places tensors to leverage high-bandwidth wafer-scale interconnects, reducing latency and balancing workloads. This approach yields 1.5 \(\times \) average performance improvement across diverse LLM architectures compared to SOTA training systems.