Designing Spatial Architectures for Sparse Attention: STAR Accelerator via Cross-Stage Tiling

TL;DR

This paper introduces STAR, a novel accelerator and algorithm co-design for sparse attention in large language models, achieving significant speedup and energy efficiency improvements through cross-stage coordination and optimized memory access.

Contribution

The paper presents STAR, a cross-stage optimized accelerator and algorithm for sparse attention, enabling efficient Transformer inference under large token parallelism.

Findings

Up to 9.2× speedup over A100

71.2× energy efficiency improvement

20.1× throughput increase in spatial architecture

Abstract

Large language models (LLMs) rely on self-attention for contextual understanding, demanding high-throughput inference and large-scale token parallelism (LTPP). Existing dynamic sparsity accelerators falter under LTPP scenarios due to stage-isolated optimizations. Revisiting the end-to-end sparsity acceleration flow, we identify an overlooked opportunity: cross-stage coordination can substantially reduce redundant computation and memory access. We propose STAR, a cross-stage compute- and memory-efficient algorithm-hardware co-design tailored for Transformer inference under LTPP. STAR introduces a leading-zero-based sparsity prediction using log-domain add-only operations to minimize prediction overhead. It further employs distributed sorting and a sorted updating FlashAttention mechanism, guided by a coordinated tiling strategy that enables fine-grained stage interaction for improved memory efficiency and latency. These optimizations are supported by a dedicated STAR accelerator architecture, achieving up to 9.2$\times$ speedup and 71.2$\times$ energy efficiency over A100, and surpassing SOTA accelerators by up to 16.1$\times$ energy and 27.1$\times$ area efficiency gains. Further, we deploy STAR onto a multi-core spatial architecture, optimizing dataflow and execution orchestration for ultra-long sequence processing. Architectural evaluation shows that, compared to the baseline design, Spatial-STAR achieves a 20.1$\times$ throughput improvement.