Skip to content

Latest commit

 

History

11 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

kbench — Transformer Attention Micro-benchmark Harness

A C++17 + CUDA harness that benchmarks transformer scaled-dot-product attention across three radically different execution targets behind one problem definition and one report schema:

  1. NVIDIA GPU — three hand-written CUDA kernels (naive, fused, flash) timed with CUDA events and profiled with Nsight Compute / perf.
  2. CPU reference — a golden fp32 implementation, profiled with perf_event_open for LLC cache-miss behavior.
  3. ARM Cortex-M4 — a bare-metal Q15 fixed-point implementation on an STM32F407, reached over UART, reporting DWT cycle counts plus on-die temperature / supply-voltage diagnostics.

Every target's output is diffed against the fp32 reference, and the whole run is emitted as CSV + Markdown. The point is to turn the hardware/software co-design question — where should attention run, and what would a dedicated datapath buy? — into measurements. See docs/codesign.md.


Architecture

flowchart TB
    subgraph host["host harness (host/, C++17)"]
        main["main.cpp<br/>sweeps seq, dim, heads, kernels"]
        ref["reference.cpp<br/>reference_fp32 golden output<br/>reference_q15 fixed-point model"]
        perf["perf_counters.cpp<br/>perf_event_open, LLC misses"]
        rep["report.cpp<br/>CSV + Markdown, one schema for all targets"]
        mcuc["mcu_client.cpp<br/>McuClient over Serial"]
        proto["protocol.cpp<br/>COBS + CRC-16/CCITT framing"]
        serial["serial.cpp<br/>termios, 8N1, 921600 baud"]
        stub["cuda_stub.cpp<br/>linked when no CUDA toolkit is present"]
    end

    subgraph gpu["GPU backend (gpu/, nvcc)"]
        launch["launch.cu<br/>run_attention dispatch, CUDA event timing"]
        knaive["attention_naive.cu<br/>3 pass, S x S scores in DRAM"]
        kfused["attention_fused.cu<br/>one block per row, scores in shared memory"]
        kflash["attention_flash.cu<br/>tiled, online softmax"]
        devinfo["device_info.cu"]
    end

    subgraph mcu["STM32F407 (arm/, bare metal)"]
        fw["main.c command loop"]
        q15["attention_q15.c<br/>Q15 fixed point"]
        dwt["DWT cycle counter"]
        diag["diag.c<br/>ADC die temperature and VDDA"]
    end

    subgraph tools["Analysis"]
        nsight["scripts/run_nsight.sh<br/>Nsight Compute / nsys"]
        perfsh["scripts/run_perf.sh"]
        roof["scripts/roofline.py"]
        codes["scripts/codesign.py"]
    end

    main --> ref
    main --> perf
    main --> rep
    main -->|"run_attention(kernel, cfg)"| launch
    launch --> knaive
    launch --> kfused
    launch --> kflash
    launch --> devinfo
    main -.->|"no CUDA toolkit"| stub
    main -->|"one head, fp32 Q/K/V"| mcuc --> proto --> serial
    serial <-->|"UART frames"| fw
    fw --> q15
    fw --> dwt
    fw --> diag
    ref -->|"max abs error per target"| rep
    rep --> roof
    rep --> codes
    nsight --> launch
    perfsh --> main
Loading

One MCU measurement, end to end — the DWT window deliberately excludes the link:

sequenceDiagram
    participant H as Host harness
    participant P as protocol.cpp
    participant M as STM32F407
    participant A as q15_attention

    H->>P: RUN_ATTN request, seq, dim, causal, Q K V as fp32
    P->>M: CRC-16 appended, COBS encoded, 0x00 delimited
    M->>M: decode, verify CRC, convert to Q15
    M->>A: start DWT cycle counter
    A-->>M: O in Q15, stop counter
    M->>M: sample die temperature and VDDA
    M-->>P: response: status, cycles, O, temp_milli_c, vdda_milli_v
    P-->>H: decoded payload
    H->>H: diff against reference_fp32, cycles / mcu_clock = compute latency
Loading

The GPU and MCU paths are independently optional. With no CUDA toolkit the build links a stub and skips GPU rows; with no --serial port the MCU comparison is skipped. Same binary, same report schema, runs on a GPU CI box or an engineer's desk with the dev board.


Repository layout

kernel_bench/
├── CMakeLists.txt              # host harness + optional CUDA backend
├── gpu/
│   ├── include/kbench/
│   │   ├── attention.hpp        # shared problem def + CUDA-free dispatch API
│   │   └── device_softmax.cuh   # warp/block reductions, online softmax
│   ├── kernels/
│   │   ├── attention_naive.cu   # 3-pass, S×S scores in DRAM (baseline)
│   │   ├── attention_fused.cu   # 1 block/row, scores in shared memory
│   │   └── attention_flash.cu   # tiled, online softmax, O(Bc·d) shared mem
│   ├── src/{launch,device_info}.cu
│   └── bench/gpu_microbench.cu  # standalone GPU sweep (run under Nsight)
├── host/
│   ├── include/kbench/          # serial, protocol, reference, perf, report, mcu
│   └── src/                     # implementations + main.cpp orchestrator
├── arm/                         # STM32F407 firmware (arm-none-eabi, no SDK)
│   ├── Makefile  stm32f407.ld
│   ├── include/                 # protocol/board/uart/diag/attention_q15
│   └── src/                     # main, startup (PLL+vectors), uart, diag, q15
├── scripts/
│   ├── run_perf.sh              # perf stat cache-miss profile of CPU ref
│   ├── run_nsight.sh            # ncu/nsys GPU kernel profiles
│   ├── roofline.py              # roofline placement from results.csv
│   └── codesign.py              # energy + notional-ASIC trade-off model
├── tests/                       # dependency-free COBS/CRC + reference tests
├── docs/                        # codesign.md, protocol.md
└── reports/                     # CSV/MD output + sample_results.md

Implementation details

The three GPU kernels

All compute O = softmax(Q Kᵀ / √d) V over a flattened (batch·heads) axis, laid out [head, seq, dim] row-major. They share reductions from device_softmax.cuh.

  • naive — the memory-bound baseline. Three global passes: a GEMM writing the full S×S score matrix to DRAM, a row-wise softmax over it, then a second GEMM for P·V. Streams O(S²) scores through DRAM, so it pins ~peak DRAM bandwidth and never gets faster per-FLOP as S grows.
  • fused — one block per query row keeps the S-length score vector in shared memory, collapsing the three passes into one and eliminating the score round-trip. Shared-memory footprint is O(S), so it caps out at moderate S — that ceiling is the comparison point against flash.
  • flash — FlashAttention-style. Streams keys/values in tiles of kBc and maintains an online softmax (running max m, denominator l) so the score matrix is never materialized; shared memory is O(kBc·d) regardless of S. Per-key dot products use a block-wide reduction across the dim threads, so it is correct for any block size, not just a single warp.

Timing uses CUDA events around the kernel only — no H2D/D2H in the window — with configurable warmup, and reports mean / p50 / p99 / min plus achieved GFLOP/s and effective GB/s (gpu/src/launch.cu).

CPU reference + cache profiling

reference_fp32 is the golden output every other target is diffed against (max-abs and relative-L2). reference_q15 is a bit-accurate model of the MCU's fixed-point path — Q15 inputs, Q15.16 score accumulation, a 64-entry exp LUT — so the board's quantization error is predictable before hardware is attached (tests/test_reference.cpp pins it under 5% rel-L2).

PerfGroup (host/src/perf_counters.cpp) opens a perf_event_open group — LLC loads, LLC load-misses, instructions, cycles on one leader fd so they're scheduled together — around the reference run, yielding miss-rate and IPC. It degrades gracefully (counters disabled) where perf_event_paranoid or missing CAP_PERFMON block it.

ARM Cortex-M4 firmware

Bare-metal, no vendor SDK — just the registers the firmware touches (arm/include/board.h). startup.c brings up the 168 MHz PLL, sets the vector table, initializes .data/.bss, and enables the FPU. q15_attention (arm/src/attention_q15.c) mirrors the flash kernel's streaming structure — one query row at a time, online-rescaled softmax — but scalar and fixed-point, because the M4 has an FPU but no hardware exp() and prefers 16-bit data through its cache and SRAM budget. The DWT cycle counter brackets only the compute. diag.c samples the internal temperature sensor and VREFINT via ADC1 for the hardware diagnostic fields. Built with arm-none-eabi-gcc (arm/Makefile).

Wire protocol

COBS-framed, CRC-16/CCITT-checked request/response over raw 8N1 UART, with identical encoders on both ends (host/src/protocol.cpp, arm/src/protocol.c) and a contract test (tests/test_protocol.cpp). Full spec in docs/protocol.md.

Analysis

scripts/roofline.py places each kernel on the device roofline (compute vs memory ceiling, ridge point). scripts/codesign.py converts latency into energy-per-inference and adds a first-order systolic-array ASIC operating point (lanes × clock × pJ/MAC), making the GPU-vs-MCU-vs-ASIC trade-off explicit.


Building

Host harness (+ GPU kernels if CUDA present)

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build -j
ctest --test-dir build --output-on-failure      # COBS/CRC + reference tests

CUDA architectures default to 80;86;89 (A100/Ada); override with -DKBENCH_CUDA_ARCHS=90. With no CUDA toolkit, CMake warns and links the GPU stub — the host tool still builds and runs the CPU + MCU paths.

Firmware

cd arm
make            # -> build/kbench-fw.elf + .bin  (arm-none-eabi-gcc)
make flash      # st-flash write build/kbench-fw.bin 0x08000000

Running

# GPU + CPU sweep, write reports/results.{csv,md}
./build/kbench --seq 128,256,512,1024,2048 --heads 16 --dim 64

# Add the Cortex-M board over UART (single-head, S<=256 cases)
./build/kbench --seq 128,256 --serial /dev/ttyUSB0 --baud 921600

# Profiles + analysis
scripts/run_perf.sh   --seq 1024            # CPU cache-miss breakdown
scripts/run_nsight.sh 512,1024,2048         # GPU L2/DRAM counters + timeline
scripts/roofline.py reports/results.csv --plot reports/roofline.png
scripts/codesign.py reports/results.csv

Expected results

Qualitative story the harness is built to expose (full sample with numbers in reports/sample_results.md):

naive fused flash
DRAM traffic vs S O(S²) O(S·d) O(S·d)
Bound memory (all S) shared-mem capped compute-leaning at large S
Relative latency @ S=2048 1.0× ~2.2× faster ~4.7× faster
  • naive holds near-peak DRAM bandwidth at every S and never improves per-FLOP — the textbook memory-bound signature.
  • flash's effective GB/s falls as S grows (it stops touching DRAM for scores) while its GFLOP/s climbs toward the compute roof. That crossover is the result.
  • CPU reference LLC miss-rate and IPC degrade together as the working set exceeds cache — the CPU-side mirror of the naive kernel's DRAM bound.
  • Cortex-M is ~4 orders of magnitude slower per head than the GPU but at ~0.1 W vs ~400 W; Q15 error ≈ 1e-2, dominated by the exp LUT and matching the reference_q15 model. Die-temperature drift across a sweep shows up in the diagnostic columns.

All GPU/CPU kernels agree with the fp32 reference to ≲1e-5 relative L2 (fp32 accumulation order differences only).


License

Apache-2.0.

About

C++17+CUDA transformer attention micro-benchmark: naive/fused/flash GPU kernels vs a bare-metal ARM Cortex-M Q15 path over UART, with perf/Nsight cache-miss profiling and ASIC co-design analysis

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages