Gemmini Online Attention Module
and Attention Fusion

Research problem

Gemmini accelerates transformer matmul on a 16×16 systolic array. After QK^T, int32 scores sit in the accumulator; PV needs int8 weights in the scratchpad. The official path materializes the full attention matrix in DRAM between those two on-chip memories — pure data-movement overhead for a value consumed immediately.

A natural first fix is to fuse the store locally (skip DRAM). That works at short sequences but collapses at longer ones because Gemmini's built-in softmax still requires a full-row scan across every J-tile. The bottleneck is not missing softmax hardware — it is the wrong softmax algorithm mapped onto tiled accumulator memory.

Attention fusion — keeping weights on-chip

The baseline Gemmini attention path computes QK^T into the accumulator, softmaxes via the Normalizer, mvouts the N×N attention matrix to DRAM, then mvins it back for PV — even though PV only needs those weights in the scratchpad. That is pure data-movement waste.

Three pipelines compared

Performance

Verilator simulation, IBertGemminiRocketConfig, BERT-base attention sublayer (QKV proj + 12-head attention + W_o + LayerNorm). Same simulator for all configs.

Full sublayer breakdown

Attention core only (where fusion + online softmax matter)

Key takeaways

1. Attention fusion removes pointless DRAM traffic. Routing softmax output to scratchpad (GEMMINI_LOCAL_STORE_FLAG, A=NULL for PV) eliminates the N×N matrix round-trip.
2. Fusion alone does not scale. Exp 1 proves the data-movement analysis was right (611K→358K attn @ seq=128) but full-row softmax causes ~6.4× slowdown at seq=256.
3. Online softmax hardware completes the design. OnlineAttention makes fusion work at seq=512 with attn core 245K vs 611K baseline @ seq=128.

OnlineAttention hardware (completes the fusion)

OnlineAttention.scala — a new Gemmini module dispatched via RoCC funct=23. It reads accumulator scores directly (raw ACC bypass), maintains per-row online softmax state (max, sum, rescale; up to 256 rows), runs 16 parallel iexp units per DIM chunk, and writes int8 weights to scratchpad.

End-to-end pipeline (Exp 5): QK^T → online softmax in hardware → PV with A=NULL. Three RoCC commands per head; zero attention-matrix DRAM traffic.

Why default Gemmini softmax fails at scale (DIM=16)

One attention row spans SEQ_LEN columns in accumulator SRAM, tiled as J-tiles of 16 (DIM). The built-in Normalizer on mvout scans every J-tile in three passes — MAX, SUM_EXP, normalize — before any weight is valid.

OnlineAttention: streaming K-blocks

Instead of three full-row passes, the module keeps running state per row and processes one K-block (16 columns) at a time — the same tiling Gemmini already uses, but with FlashAttention-style online updates.

Design journey

An iterative experiment path — useful to discuss in research interviews:

Further analysis

The sections above are the headline story. Below is the supporting analysis from the full report — useful if you want the “why” behind the numbers without reading 16 pages of PDF.

Prior art and framing

FlashAttention streams K/V blocks and keeps online softmax statistics (m, ℓ) so the full N×N matrix never sits in GPU HBM. Our setting is different: a RoCC-attached systolic accelerator with separate scratchpad and accumulator, a built-in Normalizer on the mvout path, and i-BERT fixed-point iexp. The idea is the same — fuse softmax with where scores already live — but the mechanism is Gemmini-specific RTL, not a CUDA kernel.

Why optimize attention if QKV is ~45% of cycles?

At seq=128 the baseline spends ~960K cycles on QKV projection vs ~611K on the attention core. QKV must load weights and activations from DRAM at least once. The attention matrix P is different: it is produced on-chip, immediately consumed by PV, yet the official path still writes and reads it through DRAM (~393 KiB/layer at seq=128; ~6.3 MiB at seq=512). That traffic is avoidable — the main removable cost.

Why Exp 1 collapses: the O(N²) control storm

Fusion removes DRAM, but Exp 1 still drives built-in full-row softmax through thousands of small mvout-style operations. With Q_block=16, each Q-block pays roughly 48 × (N/16) store passes across the full J dimension; multiplied by N/16 Q-blocks per head, total mvout-class work scales as O(N²).

OnlineAttention replaces per-slice mvout with batched ACC reads and one OP_FUSED_BATCH per Q-block — that is the mechanism change Exp 1 could not get by redirecting softmax output alone.

Q_block sweep — granularity dominates after fusion

Once attention-matrix DRAM traffic is zero, performance is often set by how many times the scalar core launches matmul and custom-op sequences. Exp 5 swept Q_block ∈ {16, 64, 128, 256} on the attention core (K_block=N):

At seq=128 and 256, larger Q_block wins: RoCC dispatch savings outweigh accumulator spill. At seq=512, Q=256 becomes worse than Q=128 (+3.2%) — a U-curve where ACC overflow DMA exceeds the benefit of fewer Q-blocks. ACC_ROWS=2048 caps how large Q_block can go; doubling accumulator depth would likely shift the optimum rightward.

Bottleneck shift across the project

What Exp 5 fixed over Exp 4

• OP_FUSED_BATCH — one RoCC command for UPDATE→WEIGHTS per Q-block (vs two)
• Zero printf inside timed regions (~10²K cycles noise at seq=128)
• Q_block chosen from sweep evidence, not fixed at 16
• Baseline, Exp 1, and Exp 5 all on the same Verilator binary

Correctness

Cycle counts come from timed Verilator regions only. Numerical checks use hardware-matched reference softmax outside the timed bracket — outputs were validated during bring-up (Exp 4 micro-tests through full BERT-scale integration); the reported speedups are for the measured pipeline, not unchecked RTL.

Full detail: Final Report (PDF) · experiment timeline · architecture notes

Artifacts & reproduction

Repository layout

• hardware/.../OnlineAttention.scala — Chisel module
• benchmarks/exp{1,4,5}/ — C workloads + cycle counters
• patches/ — unified diff vs upstream Gemmini

export CHIPYARD=/path/to/chipyard
cd scripts
VERILATOR_THREADS=16 ./exp5/run_exp5_full_seq128.sh

Ivan Lok, Data-Movement-Aware Optimization of Transformer Attention on Gemmini, ELEC 5140 Advanced Computer Architecture, HKUST, May 2026.