SpAtten: Co-Design for Sparse Attention

• SpAtten is an end-to-end co-design framework that employs cascade token and head pruning to reduce the O(N²) complexity in Transformer attention.
• It leverages progressive quantization and a custom top-k hardware accelerator to minimize data movement and computation costs.
• The approach achieves substantial speedups and energy savings by dynamically optimizing precision and pruning non-essential tokens and heads.

SpAtten is an end-to-end algorithm-architecture co-design framework targeting efficient sparse attention in Transformer-style self-attention, focusing on layerwise dynamic token and head pruning, together with progressive quantization and a tailored hardware accelerator architecture. The approach addresses the quadratic computational and memory complexity of attention, specifically the $O(N^2)$ scaling with input sequence length $N$ and the low arithmetic intensity, which are bottlenecks for both compute- and memory-bound regimes in natural language processing tasks.

1. Motivation and Design Principles

Self-attention incurs $O(N^2)$ work and DRAM traffic, particularly pronounced during sequence generation (as in GPT-2), leading to a memory-bound workload on conventional CPUs and GPUs. In BERT-style summarization—primarily compute-bound—overhead persists due to intricate data movement (split, transpose, concatenate) that general-purpose hardware inadequately supports. To address the "memory-compute gap," SpAtten targets simultaneous reductions in compute, data movement, and bit-width via:

• Cascade Token Pruning: Layerwise elimination of unimportant tokens, reducing effective sequence length $N' \ll N$.
• Cascade Head Pruning: Dynamic selection of relevant attention heads, decreasing total head count $H \to H'$.
• Progressive Quantization: Two-stage quantization of queries, keys, and values (Q/K/V), initially using only the most significant bits and selectively increasing precision as dictated by softmax distribution sharpness.
• Custom Accelerator: A pipelined hardware design that avoids fetching or computing on pruned data and dynamically adapts to precision requirements, utilizing on-chip buffers and high-throughput top-$k$ engines.

This unified strategy aims to minimize end-to-end memory and compute demands while maintaining model accuracy (Wang et al., 2020).

2. Cascade Token Pruning Mechanism

Token Importance Scoring

For each layer $\ell$, SpAtten computes a cumulative importance score $s_t \in \mathbb{R}^{L_1}$ for K/V tokens ($j=1...L_1$):

$$p_h^{(\ell)}[i,j] = \operatorname{Softmax}_j \left( Q_h^{(\ell)}[i] \cdot K_h^{(\ell)}[j] / \sqrt{D} \right)$$

The importance accumulator is updated:

%%%%10%%%%

This summation provides a global measure across heads and query positions, capturing aggregate token influence.

Pruning Schedule and On-the-Fly Top-k

A global token-keep ratio $R \in (0,1)$ prescribes the number of active tokens $$K_\ell = \lceil R_\ell \times \text{previous\_length} \rceil$$ at layer $\ell$, using an interpolated schedule $R_\ell$. Once a token is pruned at any layer, it is permanently removed. Selection utilizes a custom $O(n)$ quick-select hardware engine emitting $K_\ell$ highest-scoring tokens in a streaming manner without full sorting.

Algorithmic Summary

At each layer:

1. Compute attention probabilities over all heads, queries, and tokens.
2. Update cumulative importance scores $s_t$.
3. Use the accelerator's top-$k$ engine to select surviving tokens.
4. Shrink K/V (and Q if applicable) tensors accordingly.

This enables a staged ("cascade") reduction in sequence width, with early layers keeping more tokens and later layers more aggressively pruning.

3. Layerwise Cascade Head Pruning

Head Importance

After each head $h$’s output $E_h \in \mathbb{R}^{L_0 \times D}$ is computed, a global score $s_h[h]$ is accumulated:

$$s_h[h] \leftarrow s_h[h] + \sum_{i=1}^{L_0} \sum_{d=1}^D |E_h[i,d]|$$

The magnitude signifies contribution to the block output. Heads are pruned according to a scheduled per-layer "head-keep" ratio $R_{h,\ell}$ analogous to the token regimen.

Complexity Reduction

Attention computation per layer is reduced from $O(H L_0 L_1 D)$ to $O(H'_\ell L_0 L'_1 D)$, with further savings in downstream two-layer feed-forward (FFN) blocks due to reduced token counts.

4. Hardware-Aware Top-k Selection

Pipeline and Optimizations

The hardware top-$k$ engine streams score arrays into a quick-select core using random pivots (via LFSR), partitions input into less-than/greater-than FIFOs, identifies the k-th threshold, and compacts the surviving indices. It leverages:

• 16 parallel comparators per cycle.
• Double-buffered FIFOs (64 entries/stage) to hide DRAM latency.
• Multi-stage pipelining for sustained throughput.

This achieves average $O(n)$ latency with minimal passes ($\approx 3-5$), maintaining streaming data order and avoiding global shuffles.

5. Progressive Quantization for Q/K/V

Two-Stage Quantization Process

Q, K, and V vectors in DRAM are stored as concatenated most-significant bit (MSB, $b_M$) and least-significant bit (LSB, $b_L$) planes ($b_M=6...12$, $b_L=4$ typical):

1. First stage: Fetch only $b_M$-bit MSBs; compute attention and softmax in reduced precision.
2. Dynamic confidence test: If rowwise $\max_j p_{ij}^{(1)} < \tau$ (e.g., $\tau = 0.1$, indicating a "flat" distribution and possible quantization error), fetch $b_L$-bit LSBs, reconstruct Q/K, and recompute.
3. Otherwise: Retain coarse computation output for "peaked" distributions.

Error Propagation

Quantization errors $\Delta s_i$ in softmax translate as: $\Delta p_i = p_i(1-p_i)\Delta s_i$ (for $i=j$), and $-\sum_{j \neq i} p_i p_j \Delta s_j$ otherwise, with aggregate error $\le 0.5 |\Delta s|$. This ensures error vanishes for confident ($p_i \to 1$) softmax outputs.

6. Accelerator Architecture and Dataflow

Block Structure

• HBM2 DRAM (16×128-bit@2GHz)
• Q/K/V fetchers, programmable crossbars, address FIFOs
• Double-buffered key/value SRAM, 512×12-bit multiplier arrays, floating-point softmax units
• Progressive quantization control
• Token/head top-$k$ engines, accumulation modules, controller sequencer

Computation Pipeline

Per layer:

1. Load Q.
2. Apply token-top-$k$ on $s_t$; generate addresses of surviving K_j.
3. Fetch MSB-only K_j into key SRAM.
4. Attend and softmax.
5. Verify quantization confidence; fetch LSB if necessary.
6. Update $s_t$.
7. Prob×V for attention output E_h.
8. Update $s_h$ for head importance.
9. After all queries/heads, apply head-top-$k$.
10. Prune Q/K/V chunks.
11. Residual, layer-norm, FFN handled on GPU or shared multipliers.

Architecture maximizes pipeline utilization through double buffering, crossbars, and coarse-to-fine bitwidth conversion modules.

7. Performance Results and Comparisons

Empirical Outcomes

SpAtten evaluated on 30 NLP benchmarks shows:

• 10× average DRAM access reduction (no accuracy loss)
• 2.1× overall computation reduction (1.9× via token pruning, 1.1× via head pruning)
• End-to-end speedup: 162× over TITAN Xp GPU, 347× over Intel Xeon, 1095× over Jetson Nano, 5071× over Raspberry Pi
• Energy savings: 1193× (TITAN), 4059× (Xeon), 406× (Nano), 1910× (Pi) (Wang et al., 2020)

Comparative Summary with MNNFast and A³

Feature	MNNFast	A³	SpAtten
Token Pruning	local only	per-head local	global cascade
Head Pruning	—	—	cascade on-the-fly
Quantization	fixed	fixed	progressive MSB→LSB
DRAM Reduction	none	none	10×
GOP/s	120	221	360
Energy Eff. (GOP/J)	120	269	382
Area Eff. (GOP/s/mm²)	—	106	238

Both MNNFast and A³ must fetch all Q/K/V before pruning, yielding no DRAM savings and only accelerating compute-bound BERT. SpAtten’s cascade approach yields global savings in both self-attention and downstream FFN layers and is applicable to generative models (e.g., GPT-2). Progressive quantization closely approaches HBM bandwidth limits in memory-bound settings.

Roofline Analysis

For BERT (high arithmetic intensity), SpAtten achieves 1.61 TFLOPS against a 2 TFLOPS hardware roof, compared to ~0.45 TFLOPS on an NVIDIA GTX1080Ti. For GPT-2 (low intensity), SpAtten approaches the 512 GB/s HBM bandwidth ceiling, while the GPU is limited by memory efficiency.

---

SpAtten exemplifies a synergy of algorithmic sparsification and dynamic quantization, tightly integrated with a streaming, pipelined accelerator design. The result is significant reductions in DRAM traffic, computational load, and energy consumption for Transformer attention, establishing a new efficiency benchmark in hardware/software co-design for deep sequence models (Wang et al., 2020).