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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

Spatial-Aware Orchestration of LLM Attention on Waferscale Chips

  • Taiquan Wei,
  • Huizheng Wang,
  • Zichuan Wang,
  • Shouyi Yin,
  • Yang Hu

摘要

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.