FlashAttention FPGA Kernel

I built this project to ask a hardware-specific question: can the memory-saving idea behind FlashAttention become a clear, synthesizable dataflow for an AWS EC2 F2 FPGA? The result is a SystemC attention kernel for Catapult high-level synthesis (HLS), accompanied by software references and an early network-anomaly-detection pipeline that gives the kernel an application context.

I ended up with a tiled attention design whose floating-point reference matches ordinary dense attention to machine precision at the current maximum shape, N=512 and d=64. The SystemC core implements the same dataflow with fixed-point arithmetic, 32 x 32 tiles, 16 row-parallel processing elements, and a piecewise-linear exponential. This is still a research prototype: the standalone kernel simulation uses deterministic internal tensors, its main testbench checks control flow and a nonzero checksum rather than every output value, and the included FPGA performance and resource numbers are analytical estimates rather than post-synthesis or AWS F2 measurements. The anomaly-detection scaffold also uses deterministic fixture weights; it is not a trained detector.

How it works, in one paragraph

The testbench writes a sequence length through a 128-bit ready/valid control interface and then starts FlashAttnCore, which reads query (Q), key (K), and value (V) vectors either from its deterministic test banks or from circular buffers attached by an integration test. The core keeps 32 query rows on chip and walks over K/V in 32-row tiles; 16 parallel processing elements compute 64-wide dot products, while each query row carries only a running maximum, softmax denominator, and unnormalized output vector. Those values are rescaled after every score block using the numerically stable online-softmax recurrence, so the design never allocates a full N x N score or probability matrix. The hardware path uses fixed-point types and a piecewise-linear exponential, then normalizes and stores each output row and exposes a checksum plus a counter that currently represents processed tile pairs. The separate anomaly-detection scaffold shows the intended larger flow: screen packet features, project them into Q/K/V, retain a sliding window, run attention, mean-pool the rows, classify the window, and raise an alert.

How to run it

The main hardware entry point is src/FlashAttnKernel/testbench.cpp, invoked from the repository root with make systemc_sim. It requires the course CAD environment used by setup.csh: Siemens Catapult 2024.2_1, SystemC, MatchLib Connections, Algorithmic C datatypes, and the configured Synopsys tools under /cad. In that environment:

tcsh
source setup.csh
make systemc_sim

Run the module-level SystemC tests in the same shell with:

make -C testbench REPO_TOP="$PWD" run_all

To synthesize the kernel and run the configured SCVerify flow:

make hls

The portable reference is the quickest way to verify the algorithm without the proprietary hardware tools. It requires Python 3 and NumPy:

python3 scripts/reference_attention.py --seq 512

For a repeatable CPU comparison:

cc -O2 -o /tmp/cpu_baseline scripts/cpu_baseline.c -lm
/tmp/cpu_baseline 512 100

The analytical FPGA models can be regenerated with bash scripts/run_estimates.sh; they write ignored artifacts under reports/. The AWS helper targets are not self-contained: using them requires an AWS F2 instance, the aws-fpga HDK/CL environment, AWS credentials and S3 configuration, generated RTL, an FPGA image, and XDMA device nodes. The current synthesizable specification supports at most 512 tokens, while the reference host header still declares 1024; the HLS Tcl also still selects an xcvu9p library instead of the F2 xcvu47p. Both must be reconciled before treating the AWS path as deployable.

How to navigate the repo

Attention kernel and interface

• src/FlashAttnKernel/FlashAttnKernel.h contains the tiled attention datapath, online-softmax update, fixed-point exponential approximation, and the control wrapper presented to the testbench and HLS flow.
• src/FlashAttnKernel/testbench.cpp is the main simulation driver; it programs N=512, launches the kernel, and checks configuration readback, completion, checksum, and the tile counter.
• src/include/FlashAttnSpec.h is the source of truth for tensor dimensions, tile sizes, processing-element count, fixed-point formats, and control/status messages.
• src/include/AxiSpec.h defines the 128-bit AXI-to-ready/valid interface used for configuration and status traffic.
• src/include/helper.h provides the byte-packing utility used to construct testbench register addresses.

Application scaffold

• src/include/NetworkAnomalyDetector.h assembles the intended packet-to-alert hierarchy and its host-visible register map; it is integration scaffolding rather than the current HLS top.
• src/include/Tier1Screening.h applies a small rule-based packet filter so suspicious traffic can trigger attention early.
• src/include/TriggerController.h coordinates packet intake, window freezing, attention launch, classification, and alert reporting as a finite-state machine.
• src/include/EmbeddingProjection.h maps each 32-feature packet into 64-element Q, K, and V vectors with reusable matrix-vector hardware.
• src/include/CircularWindowBuffer.h stores a 512-entry sliding Q/K/V window in a transposed layout intended for on-chip memory.
• src/include/ClassificationHead.h mean-pools attention rows and passes them through a two-layer multilayer perceptron to produce an anomaly score.

Verification and evaluation

• scripts/reference_attention.py compares dense attention with the tiled online-softmax algorithm using the same shapes and deterministic tensors as the kernel.
• scripts/cpu_baseline.c is the portable scalar performance baseline, while scripts/cpu_baseline_optimized.c adds cache-local tiles and manual loop unrolling.
• testbench/ contains focused SystemC tests for the circular buffer, projections, classifier, rule screen, trigger controller, and window-to-attention handoff.
• testbench/test_data/gen_test_vectors.py regenerates the fixed-point packet and weight fixtures stored in testbench/test_data/.
• scripts/estimate_latency.py and scripts/estimate_resources.py expose the assumptions behind the analytical FPGA estimates; scripts/run_estimates.sh runs both and collects their reports.

Build and deployment

• Makefile is the top-level command surface for SystemC simulation, HLS, RTL copying, and AWS build handoff.
• setup.csh loads the repository paths and proprietary CAD tool environment expected by the build files.
• src/FlashAttnKernel/Makefile, hls/FlashAttnKernel/Makefile, and scripts/hls/ connect the source and testbench to Catapult, MatchLib, VCS, and SCVerify.
• design_top/software/runtime/flash_attn_api.h documents the planned host-to-FPGA programming sequence, and flash_attn_host.c demonstrates the corresponding XDMA calls.
• design_top/ and scripts/aws/ hold the AWS F2 handoff and image-management helpers that depend on an external aws-fpga environment.