---
title: SystolicAttention in FSA Architecture
url: https://www.emergentmind.com/topics/systolicattention
type: topic
---

# SystolicAttention in FSA Architecture

SystolicAttention is a scheduling algorithm and hardware architecture enhancement within the FSA (Fused Systolic Array) system, introduced to map the FlashAttention algorithm entirely onto a single systolic array. It addresses the fundamental mismatch between typical systolic arrays—optimized for large, consecutive matrix multiplication—and the interleaved matrix–softmax–matrix workflow of FlashAttention. SystolicAttention enables all operations required for scaled dot-product attention (SDPA), including the non-matrix components of softmax (row-wise max, exponentials, row-wise sums), to execute efficiently and with high utilization inside a single array, without external vector unit involvement. This structural integration preserves the exact floating-point (FP) operation order and numerical stability of FlashAttention, while substantially improving hardware utilization and efficiency compared to commercial accelerators [2507.11331].

## 1. Formal Definition and Scheduling Algorithm

Let $Q \in \mathbb{R}^{L \times d}$, $K \in \mathbb{R}^{L \times d}$, $V \in \mathbb{R}^{L \times d}$. These are partitioned into tiles: $Q$ into $T_r$ blocks $Q_i \in \mathbb{R}^{B_r \times d}$, and $K,V$ into $T_c$ blocks $K_j, V_j \in \mathbb{R}^{B_c \times d}$. The FlashAttention inner loop then computes for each $(i,j)$:

- $S = Q_i K_j^T \in \mathbb{R}^{B_r \times B_c}$
- $m = \operatorname{rowmax}(S) \in \mathbb{R}^{B_r}$
- $N = (S - m \cdot 1^T) / \sigma$ (where $\sigma = \sqrt{d}$)
- $P = \exp(N)$
- $\ell = \operatorname{rowsum}(P)$
- $O_{\text{partial}} = P V_j$

SystolicAttention organizes these computations in a fully-pipelined schedule on a $N_\text{ROWS} \times N_\text{COLS}$ array (selecting $B_r = N_\text{COLS}$ and $B_c = N_\text{ROWS} = d$), with key innovations:

- The *upward dataflow* within the array simultaneously computes $Q_i K_j^T$ and determines rowmax in flight.
- An *injected element-wise exponential* (via a piecewise linear exp2 implemented by reusing the multiply-accumulate in each processing element, or PE) computes exponentials directly in place.
- The *downward dataflow* initiates as soon as elements are ready (starting from the top-left PE), multiplying elements of $P$ with $V_j$ for accumulation.

All loops are fully pipelined, with the microprogram of each PE being entirely counter-based and deterministic. This approach achieves overlap such that while later rows perform rowmax, earlier rows already launch the $P V_j$ multiplication.

The core pseudocode for a single $(i,j)$ tile iteration is:

```python
for c in 0..B_c-1:                        
    for r in 0..B_r-1:                    
        PE(r,c).acc += Q_i[r,*] · K_j[c,*]     # upward mat-mul partial sum
        PE(r).cmp_max(acc)                    # update rowmax m[r]
# after QK, S and m are established
for c in 0..B_c-1:
    for r in 0..B_r-1:
        N = PE(r,c).acc - m[r]
        x_i, x_f = split(N/σ)
        P = (2^x_i) * PWL_interp(x_f)
        rowsum[r] += P
        PE(r,c).acc = P
# immediately begin second mat-mul via downward path
for k in 0..d-1:
    for r in 0..B_c-1:
        PE(r,k).acc2 += PE(r,*).P · V_j[*,k]   # downward mat-mul
```
[2507.11331]

## 2. Numerical Stability and Preservation of Floating-Point Operation Order

A defining requirement of FlashAttention is the arrangement of operations to avoid FP overflow/underflow, particularly by subtracting the running maximum from each row before the exponentiation. SystolicAttention ensures:

- Identical ordering of partial dot-product accumulation, subtractions, and softmax steps per row/column.
- The exponential is approximated per PE using exp2, split into integer ($x_i$) and fractional ($x_f$) parts, so $\exp(x) \approx 2^{x_i} \cdot 2^{x_f}$, with $2^{x_f} \in (0.5,1]$.
- Row-sum accumulation and output calculation strictly match the original sum order in FlashAttention.
- Final output updating follows the relations:
  - $\text{new}_\ell[r] = \text{old}_\ell[r] \cdot b[r] + \sum_c P[r, c]$
  - $\text{new}_O[r, *] = b[r] \cdot \text{old}_O[r, *] + \sum_c P[r, c] \cdot V_j[c, *]$

No reductions or FP operation orderings are reordered. The error introduced by the piecewise-linear exp2 is bounded ($\sim 10^{-3} - 10^{-2}$ relative per element), and end-to-end relative error in fp16 is maintained at similar levels.

## 3. FSA Microarchitectural Enhancements for SystolicAttention

The FSA architecture augments a conventional input-stationary $N \times N$ systolic array to enable SystolicAttention by incorporating:

- An *upward data path* in each PE to support both upward and downward data propagation.
- A *top row* of compare units (CMP) to perform on-the-fly rowmax computation for each row, cycle by cycle, removing the need for offloading this reduction.
- A *split unit* in each PE, which decomposes an FP value into its integer and fractional parts for exp2 computation, using per-segment piecewise linear interpolation.
  
With these enhancements, all softmax computations (rowmax, exponentials, summations) are internalized within the array, eliminating external vector unit dependencies for the attention mechanism. Tile data ($Q, K, V$) are loaded via DMA, while two on-chip SRAMs (approx. 192 KiB scratchpad, 64 KiB accumulation for a $128 \times 128$ array) are used for buffer and accumulator storage. The controller utilizes only two cycle counters or simple state machines.

| Hardware Component    | Array Overhead (%) | Area ($\mu m^2$)  |
|----------------------|-------------------|-------------------|
| Base array           | 89.7              | 17.74 M           |
| Upward path          | 4.8               | 0.96 M            |
| Split units          | 4.8               | 0.94 M            |
| CMP row              | 0.7               | 0.13 M            |
| **Total overhead**   | 10.3              | 2.03 M            |

[2507.11331]

## 4. Performance Evaluation and Utilization

FSA's performance was evaluated with a single $128 \times 128$ array at 1 GHz in 16 nm technology. For comparison:

- **AWS NeuronCore-v2**: $128 \times 128$ @ 2.8 GHz yields 91.75 TFLOP/s (matrix multiplication), vector unit at 2.3 TFLOP/s.
- **Google TPUv5e**: Four $128 \times 128$ arrays @ 1.5 GHz, total 196.6 TFLOP/s, with dedicated vector lanes.

SystolicAttention, via FSA, achieves mean utilization rates as follows (for sequence lengths $L=2$K–16K):

- $1.77\times$ higher FLOPs/s utilization than NeuronCore-v2
- $4.83\times$ higher than TPUv5e

Concrete example: For $L=8K$, TPUv5e achieves $\sim$20% utilization, Neuron-v2 ~30%, while FSA delivers $\sim$60% utilization.

In summary, the complete attention (QK, softmax, and KV) for an $N \times N$ tile completes in $5N + 10$ cycles, and area overhead v. a plain array is $\sim$10% (excluding SRAM & DMA).

## 5. Trade-Offs, Limitations, and Open Directions

- The batch decoding phase in LLMs (Q length = 1) is highly memory-bound, and FSA's array tiling (e.g., 128 × 128) leads to padding and inefficiency for such cases.
- Piecewise-linear exp2 introduces $\approx 10^{-3} - 10^{-2}$ relative error per element; though the total MRE remains similar for attention, some domains may require greater precision.
- FSA removes the need for a vector unit in attention, but non-attention activations (e.g., GELU, layernorm) still necessitate a modest SFU/vector pipe.
- Extensions to larger array sizes or dynamic tile sizing necessitate nontrivial system-scale buffer and DMA scheduling design.

A plausible implication is that while SystolicAttention markedly improves array utilization for large sequence lengths typical in pretraining and fine-tuning, decoder deployments requiring highly variable or small batch sizes may require supplementary architectural strategies.

## 6. Significance and Future Research

SystolicAttention demonstrates that efficient, fully-on-chip execution of FlashAttention is attainable by small, targeted register-transfer-level enhancements and a deterministic counter-based scheduling algorithm. By enabling fusion of matrix and softmax phases, it increases sustained array utilization by $1.8-4.8\times$ compared to current-generation accelerators, within modest silicon area budgets. The approach leaves the sequence and order of floating-point operations untouched from the canonical FlashAttention computation, preserving agreed-upon numerical stability guarantees.

Open challenges for future research include further optimizing FSA for variable traffic patterns (for example, in decoder or inference workloads), reducing approximation error for applications that require higher precision, and generalizing the hardware-software co-design strategy to other irregular, interleaved operator workloads in emerging deep learning models [2507.11331].

Source: https://www.emergentmind.com/topics/systolicattention