← All projects

Attention Accelerator:
Python → CUDA → RTL

I am building a single-head scaled dot-product attention engine, the core math behind every transformer, from a Python reference all the way down to synthesizable SystemVerilog, rebuilding the exact same operation at six levels of abstraction and diffing each one against the last.

Result so far: scaled dot-product attention built and cross-verified across NumPy, Q8.8 fixed-point Python, naive and tiled CUDA, and SystemVerilog RTL, each level diffed against the one above. The CUDA kernels agree with the float64 reference to about one float32 epsilon. On an RTX 4070 Laptop GPU, profiling with Nsight Compute took my FlashAttention-lite CUDA kernel from 28.3 ms to 2.23 ms (12.7x, 3.86 TFLOP/s) at N=4096, and my Triton version runs at 74 to 131% of cuDNN. On the FPGA, every N=16 configuration now builds for the ZUBoard 1CG and meets 150 MHz, and the fastest computes a 16-token head in 10.3 us, using about a third of the LUTs. On-board latency is next and tracked in the status table below.

Attention dataflow: Q, K, V flow through QK-transpose scaled by root-d, softmax, and PV to produce output O

Attention is three matrix operations. Given queries Q, keys K and values V, you compute a score matrix, turn it into a probability distribution with softmax, and use those probabilities to blend the values:

S = Q · K⊤ / √d        // token-by-token scores
P = softmax(S)          // each row sums to 1
O = P · V               // weighted blend of the value rows

That math is simple enough to check by hand at toy dimensions, which is exactly the point. I run everything at N = d = d_v = 4 with a fixed random seed so the inputs never drift. When a lower level disagrees with the golden model, it is always a real bug in my code, never a change in the data.

Why build it six times

The goal is not another attention library (PyTorch already has one). The goal is to understand the operation deeply enough to implement it on hardware, and to build the verification discipline that hardware demands. Each level strips away an abstraction and forces me to own something the level above hid from me: floating point, then fixed-point arithmetic, then GPU parallelism, then the memory hierarchy, then bit-level RTL and handshaking.

Six abstraction levels from Python float golden model down to FlashAttention-lite, each with status and error metric

How to think about the cost

The lens I carry across every level is the roofline question: is this step limited by compute (doing FLOPs) or by memory (moving bytes)? Many attention configurations, especially decode and any kernel that materializes the intermediate score matrix, are limited by memory traffic more than arithmetic; the balance shifts with shape, datatype, and hardware. That lens explains most of the design work below. The naive GPU kernel re-reads every key and value row once per query, so it moves N times more bytes than necessary; a single-block tiled kernel stages each input tile into on-chip memory once and reuses it, raising arithmetic intensity (FLOPs per byte). At realistic sizes the full tensors do not fit on chip, so tiles are re-loaded across thread blocks and the win is fewer global reads per tile rather than a single global load. On the FPGA the same tension returns as an area-versus-cycles trade: one multiplier reused over D cycles is small and slow, D multipliers in parallel are large and fast.

Key idea

Correctness first, then cost. Each level is proven against the one above (same seed, same inputs) before I reason about whether it is compute-bound or memory-bound. A performance win that moves the answer is not a win.

Architecture & results so far

Every level is validated against the one above it, so I always have an unambiguous oracle. Here is where each stage stands today:

Parallel QK-transpose engine block diagram: four dot4 MAC units compute one score row at once with valid-ready handshake and 1 over root d scaling
LevelWhat it isResultStatus
M1Python float64 golden modelself-checking reference (softmax rows sum to 1)Done
M2Python Q8.8 fixed-point twinmax abs err 0.00631 vs golden (~0.6%)Done
M3CUDA naive kernel (float32)max abs err 1.192093e-07 = the float32 floor (2-23)Done
M4CUDA tiled kernel (shared memory)max abs err 1.192093e-07, bit-identical to M3Done
M5SystemVerilog RTL core (Q8.8)full datapath, 52/52 values exact vs the integer spec at all four stagesDone
M6Host to FPGA link (AXI-Stream + DMA, PYNQ)packaged as an AXI IP with per-phase cycle counters and a completion interrupt; caught a counter race in integrationIn progress
M7FPGA bring-up on ZUBoard 1CG (XCZU1CG)every N=16 build meets 150 MHz (was 65 to 111); a parallel fold cuts compute from 7,328 to 1,520 cycles, bit-exact; on-hardware diff nextIn progress
M8FlashAttention-lite, RTL (online softmax)360 cy on one multiplier, beating the naive core on four (515 cy); N² buffers goneDone
M8FlashAttention-lite, CUDA and Tritonon an RTX 4070 Laptop GPU: Triton at 74 to 131% of cuDNN; CUDA kernel 12.7x faster after following Nsight and moving to tensor cores (3.86 TFLOP/s); found and fixed a shared-memory raceDone
Stacked bar chart of RTL cycles per stage at one, two and four MAC lanes: the matmul segments shrink while the softmax segment stays fixed at 449 cycles

The result that taught me the most is the one in that chart. I parameterized the MAC count in the RTL and swept it, expecting the usual area-for-speed trade. Four times the multipliers bought 1.28x. The matmuls really did get 3.4x faster, but the serial divider inside softmax costs the same 449 cycles regardless, so adding lanes only grew its share of the runtime, from 68% to 87%. The fix was not more hardware. Rewriting softmax to stream (M8) got 1.83x at identical area and removed the N² score buffers entirely, which at N=128 and head dim 64 is the difference between needing roughly 64 KB of BRAM and needing 66 words. I wrote that up in the dev blog.

The two CUDA numbers are the ones I am most proud of. 1.192093e-07 is about one float32 epsilon near unity (2-23 is the ULP spacing near 1.0). Because I compute the golden model in float64 and cast the exported inputs to float32, a difference this small means close numerical agreement for this test case, not a bug in the kernel. Getting there twice, once naive and once tiled, is strong evidence the kernels are correct, not just close.

M8 on a real GPU

The FlashAttention-lite kernels were verified at toy sizes long before I had a GPU to time them on. Once I ran them on an RTX 4070 Laptop GPU at N=4096 and head dim 128, the first CUDA version took 28.3 ms. My traffic model predicted it was memory-bound, and it was, but Nsight Compute showed the bottleneck was not where I expected: L1 sat at 98% while L2 was only at 42%, because the K loads were uncoalesced and each 32-byte sector delivered about 7 useful bytes. The model got the traffic right and the bottleneck wrong.

Horizontal bar chart of kernel time at N=4096, D=128: original 28.3 ms, coalesced K 19.4 ms, query-tiled 7.59 ms, tensor cores 2.23 ms, and cuDNN for reference at 0.29 ms.

Each step followed the profiler. Coalescing the K loads got it to 19.4 ms. Tiling queries so each block reuses every K and V row across many queries got 7.59 ms. Moving the two matmuls onto fp16 tensor cores with WMMA got 2.23 ms, a 12.7x speedup at 3.86 TFLOP/s. Along the way compute-sanitizer found a shared-memory race in the block reduction (21 hazards to 0) that the small tests had never exposed.

Log-log line chart of time against sequence length at D=128. The original fp32 kernel rises to 28.3 ms at N=4096, the query-tiled kernel to 7.59 ms, the tensor-core kernel to 2.23 ms, and PyTorch SDPA to 0.29 ms.

For honesty, the reference point is PyTorch's own attention (SDPA), which uses cuDNN here and takes 0.29 ms. My hand-written kernel reaches 13 to 30% of it depending on N; the Triton version is within range of cuDNN outright. The roofline shows why the gap remains. I measured the ceilings myself: 11.8 TFLOP/s FP32, 39.2 TFLOP/s on fp16 tensor cores, and 1.30 TB/s of L2 read bandwidth. That last number corrected my own chart: I first plotted a copy-benchmark figure and had the original kernel at 72% of its roof, when it was really at 53%.

Log-log roofline with an L2 read bandwidth slope of 1.30 TB/s and flat ceilings at 11.8 TFLOP/s for FP32 and 39.2 TFLOP/s for fp16 tensor cores. The original kernel sits at 0.44 FLOP per byte and 303 GFLOP/s, the coalesced kernel just above it, the query-tiled kernel at 1.13 TFLOP/s and the tensor-core kernel at 3.86 TFLOP/s.

What is left is known: sharing K and V tiles through shared memory with asynchronous copies (cp.async), and keeping the accumulator in registers with mma.sync instead of round-tripping it through shared memory. The full story, including the numbers I got wrong first, is in the dev blog.

M7 on the FPGA

The build scripts for the ZUBoard 1CG ran for the first time in Vivado 2024.1, and needed four fixes before they made a bitstream. That first bitstream missed 150 MHz, so the PL clock is now 100 MHz. I then swept the number of dot-product lanes at N=16. Area barely moves: 2 to 16 lanes costs 2,500 LUTs and 14 DSPs on a part with 37,440 and 216. The clock is what moves.

Line chart of fmax against BLK for N=D=16 on the ZUBoard: 111 MHz at BLK 2, 102 at BLK 4, 96 at BLK 8 and 65 at BLK 16, against a 100 MHz target. Critical-path logic levels grow from 25 to 50.

Every build has the same critical path: from a lane's DSP, through the max across lanes, into the exponential lookup and the running sum. The max is a chain of comparisons, so the path grows from 25 logic levels at 2 lanes to 50 at 16. A comparison tree with a register after it is the next RTL change, along with a fold step that does DV multiplies at once, since today extra lanes only speed up the dot products. Both are now done, below.

Line chart of fmax against BLK at N=D=16. Before pipelining: 111, 102, 96 and 65 MHz at BLK 2, 4, 8 and 16. After: 156, 150, 155 and 155 MHz, all at or above the 150 MHz target.

Following the timing report one path at a time (the lane max, then the output scaling multiply, then the max tree again at 16 lanes) brought every configuration to 150 MHz or more, for two extra cycles per block of keys. Then a parallel fold, one multiplier per output column, made extra lanes actually pay: N=16 compute went from 7,328 cycles to 1,520, or 10.3 us, for 45 more DSPs, and the output is bit-identical. Next is running it on the board and diffing the output against the fixed-point model. The details are in the dev blog.

The fixed-point path to hardware

FPGAs do not do floating point for free, so before writing a line of HDL I built a Q8.8 fixed-point twin of the model in Python: real numbers scaled by 256 and stored as signed 16-bit integers. The one rule the whole thing lives on is that multiplying two Q8.8 numbers doubles the fraction bits to Q16.16, so after every integer matmul I divide the accumulator back down by SCALE²:

Qq, Kq = quantize(Q), quantize(K)      # real -> Q8.8 integers (saturated)
S_raw   = Qq @ Kq.T                    # integer matmul -> Q16.16
S_fixed = (S_raw / (SCALE * SCALE)) / np.sqrt(d)   # /SCALE^2, then /sqrt(d)

That let me quote a concrete 0.6% worst-case error budget up front, and it surfaced a classic hardware bug where it is cheap to see: quantizing a normalized softmax row breaks its normalization (row sums drift to 255-257 instead of a clean 256). That informs whether the RTL needs a renormalization stage or more fraction bits.

The RTL now under way (M5) is a Q8.8, 16-bit signed datapath with a valid/ready handshake on every module, the same handshake semantics AXI4-Stream uses, so wrapping these blocks as AXI4-Stream later is straightforward. The first module is a time-multiplexed dot-product MAC (dot4): it accumulates the raw Q16.16 products in a wide accumulator, then shifts back to Q8.8 with an arithmetic shift (to keep the sign bit) and saturates exactly once at the end.

assign shifted = acc >>> FRAC;         // Q16.16 -> Q8.8, keep the sign bit
always_comb begin
    if      (shifted > sMAX) s = sMAX;  // clamp high
    else if (shifted < sMIN) s = sMIN;  // clamp low
    else    s = shifted[DW-1:0];        // in range
end

It is verified in simulation across basic, signed, and saturating cases. The QK⊤ engine that reuses this one dot4 over all token pairs is in progress. Full step-by-step write-ups for every level are on the blog.

Toolchain

Python / NumPy for the golden and fixed-point models; CUDA C++ compiled with nvcc and Triton for the GPU kernels, run on an RTX 4070 Laptop GPU and profiled with Nsight Compute and compute-sanitizer, with a byte-exact .npy data bridge back to the golden model; SystemVerilog targeting a Xilinx/AMD FPGA through Vivado, with XSim and Icarus Verilog for simulation (checked in CI) and XDC constraints. The FPGA build runs in Vivado 2024.1 for the ZUBoard 1CG; next is running it under PYNQ and closing timing above 100 MHz.

Read the full build log

I keep a running dev-log with every milestone, every bug, and the lessons from each one: the honest version, mistakes included.