SystolicAttention in FSA Architecture

• SystolicAttention is a scheduling algorithm and hardware enhancement that integrates matrix multiplication with softmax operations on a single systolic array.
• It achieves high utilization by fusing operations such as rowmax, exponentiation, and matrix multiplication in a fully pipelined, deterministic design.
• Evaluations on a 128×128 array at 1 GHz show 1.8–4.8× higher utilization versus current accelerators, underscoring its practical benefits for large sequence tasks.

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 (Lin et al., 15 Jul 2025).

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)$:

• $K \in \mathbb{R}^{L \times d}$0
• $K \in \mathbb{R}^{L \times d}$1
• $K \in \mathbb{R}^{L \times d}$2 (where $K \in \mathbb{R}^{L \times d}$3)
• $K \in \mathbb{R}^{L \times d}$4
• $K \in \mathbb{R}^{L \times d}$5
• $K \in \mathbb{R}^{L \times d}$6

SystolicAttention organizes these computations in a fully-pipelined schedule on a $K \in \mathbb{R}^{L \times d}$7 array (selecting $K \in \mathbb{R}^{L \times d}$8 and $K \in \mathbb{R}^{L \times d}$9), with key innovations:

• The upward dataflow within the array simultaneously computes $V \in \mathbb{R}^{L \times d}$0 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 $V \in \mathbb{R}^{L \times d}$1 with $V \in \mathbb{R}^{L \times d}$2 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 $V \in \mathbb{R}^{L \times d}$3 multiplication.

The core pseudocode for a single $V \in \mathbb{R}^{L \times d}$4 tile iteration is:

$Q_i \in \mathbb{R}^{B_r \times d}$0 (Lin et al., 15 Jul 2025)

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 ($V \in \mathbb{R}^{L \times d}$5) and fractional ($V \in \mathbb{R}^{L \times d}$6) parts, so $V \in \mathbb{R}^{L \times d}$7, with $V \in \mathbb{R}^{L \times d}$8.
• Row-sum accumulation and output calculation strictly match the original sum order in FlashAttention.
• Final output updating follows the relations:
  • $V \in \mathbb{R}^{L \times d}$9
  • $Q$0

No reductions or FP operation orderings are reordered. The error introduced by the piecewise-linear exp2 is bounded ($Q$1 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 $Q$2 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$3) are loaded via DMA, while two on-chip SRAMs (approx. 192 KiB scratchpad, 64 KiB accumulation for a $Q$4 array) are used for buffer and accumulator storage. The controller utilizes only two cycle counters or simple state machines.

Hardware Component	Array Overhead (%)	Area ($Q$5)
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

(Lin et al., 15 Jul 2025)

4. Performance Evaluation and Utilization

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

• AWS NeuronCore-v2: $Q$7 @ 2.8 GHz yields 91.75 TFLOP/s (matrix multiplication), vector unit at 2.3 TFLOP/s.
• Google TPUv5e: Four $Q$8 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 $Q$9K–16K):

• $T_r$0 higher FLOPs/s utilization than NeuronCore-v2
• $T_r$1 higher than TPUv5e

Concrete example: For $T_r$2, TPUv5e achieves $T_r$320% utilization, Neuron-v2 ~30%, while FSA delivers $T_r$460% utilization.

In summary, the complete attention (QK, softmax, and KV) for an $T_r$5 tile completes in $T_r$6 cycles, and area overhead v. a plain array is $T_r$710% (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 $T_r$8 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 $T_r$9 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 (Lin et al., 15 Jul 2025).