---
title: 'SpAtten: Co-Design for Sparse Attention'
url: https://www.emergentmind.com/topics/spatten-algorithm-architecture-co-design
type: topic
---

# SpAtten: Co-Design for Sparse Attention

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 [2012.09852].

## 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:

$$
s_t[j] \leftarrow s_t[j] + \sum_{h=1}^H \sum_{i=1}^{L_0} p_h^{(\ell)}[i,j]
$$

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) [2012.09852]

### 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 [2012.09852].

Source: https://www.emergentmind.com/topics/spatten-algorithm-architecture-co-design