How vca-epfl/fsa works
FSA is a FlashAttention accelerator implemented as a single systolic array, with a Python API that schedules attention work and runs it on Verilator or an XDMA-connected FPGA. The principal workflow constructs tiled attention instructions, transfers tensors and instructions to an execution engine, runs DMA and matrix operations in hardware, and retrieves the output for optional comparison with Torch or the Python reference. The README documents both TileLink and direct AXI4 integration. Hardware links reflect the documented top-level architecture and sampled module responsibilities; unsampled subsystem wiring is not expanded.
Main entry points: main.py Attention runner · Decoder.scala Instruction decoder · MatrixEngineController.scala Matrix controller · SystolicArray.scala Systolic array
Components 24
Python workflow
- Attention runner
python/main.py - Attention kernel
python/fsa/kernel.py - Execution engines
python/fsa/engine.py - Tensor model
python/fsa/tensor.py - Reference attention
python/fa_ref.py
Instruction frontend
- Instruction decoder
src/main/scala/fsa/frontend/Decoder.scala - Instruction formats
src/main/scala/fsa/isa/ISA.scala - Execution plans
src/main/scala/fsa/ExecutionPlan.scala - Matrix controller
src/main/scala/fsa/MatrixEngineController.scala
Attention compute
- Systolic array
src/main/scala/fsa/sa/SystolicArray.scala - Processing elements
src/main/scala/fsa/sa/PE.scala - Compare units
src/main/scala/fsa/sa/CMP.scala - Accumulator
src/main/scala/fsa/Accumulator.scala - FP arithmetic
src/main/scala/fsa/arithmetic/FPArithmeticImpl.scala
Data movement
- DMA controller
src/main/scala/fsa/dma/DMA.scala - Load/store queues
src/main/scala/fsa/dma/LSQ.scala - Banked SRAM
src/main/scala/fsa/BankedSRAM.scala
Simulation and integration
- FSA top
src/main/scala/fsa/AXI4FSA.scala - Verilator simulator
python/fsa/engine.py - FPGA engine
python/fsa/engine.py
Other components
- Caller
- XDMA devices
- Backing memory
- PyTorch reference
Connections 27
- Caller → Attention runner: runs
- Attention runner → Attention kernel: builds kernel
- Attention runner → Tensor model: converts tensors
- Attention runner → Execution engines: executes kernel
- Attention runner → Reference attention: compares results
- Attention runner → PyTorch reference: compares results
- Attention kernel → Execution engines: provides instructions
- Execution engines → Verilator simulator: selects simulator
- Execution engines → FPGA engine: selects FPGA
- Verilator simulator → FSA top: runs RTL
- FPGA engine → XDMA devices: transfers data
- FPGA engine → FSA top: controls device
- FSA top → Instruction decoder: dispatches instructions
- Instruction decoder → Instruction formats: decodes formats
- Instruction decoder → Matrix controller: dispatches matrix ops
- Instruction decoder → DMA controller: dispatches DMA ops
- Matrix controller → Execution plans: generates controls
- Matrix controller → Systolic array: drives compute
- Systolic array → Processing elements: composes mesh
- Systolic array → Compare units: uses compare units
- Systolic array → FP arithmetic: uses arithmetic
- Matrix controller → Accumulator: controls accumulation
- DMA controller → Load/store queues: issues requests
- Load/store queues → Banked SRAM: reads and writes
- Load/store queues → Backing memory: transfers AXI data
- FSA top → Backing memory: exposes AXI channels
- XDMA devices → Backing memory: accesses device memory