Kernels

One header declares every kernel, and each backend implements it.

kernels.h dispatching to AVX2, NEON, scalar; CUDA Mamba, flow, and DP heads are device-resident

The production backend is picked at build time:

cmake -B build -DFLOWEDGE_BACKEND=cpu

cpu selects AVX2, NEON, or scalar code for the host. avx512 is an explicit x86 build option. scalar explicitly selects the portable implementation for compatibility builds, differential testing, sanitizers, and CPUs where an AVX2 deployment baseline is unsuitable. cuda requires nvcc and links the ISA kernel surface through src/core/kernels/cuda/. The flow head keeps flow.* weights and ODE scratch on device after load; sample / sampler_advance do not cudaMalloc after that. The diffusion head keeps dp.* weights and U-Net/DDIM scratch on device after load; denoise / sample do not cudaMalloc after that. Dense Diffusion Policy ops (conv1d, conv_transpose1d, group_norm, mish, film, timestep embedding) also live in the CUDA TU as host-span launches. Short horizons use a launch that matches L; transpose is closed-form; L=4 dense conv splits IC across a warp; GroupNorm is one block per group. flowedge_cuda_dp_mix prints a launch census for one native DDIM so isolate xN is the resident mix, not a CPU estimate. Mamba host std::span arguments still copy in and out. See CUDA. Metal here is Tenstorrent TT-Metal, not Apple. Vulkan remains unimplemented.

The CPU backend splits by ISA. kernels_avx2.cc, kernels_neon.cc, and kernels_scalar.cc all implement the same public surface, and CMake picks the best file for the host. Every kernel takes std::span, owns no memory, and never allocates over 64-byte-aligned caller buffers.

The set is matmul for F32 and BF16 weights, conv1d_causal, conv1d_step, dense conv1d / conv_transpose1d (optional packed-im2col workspace through matmul), rmsnorm, rmsnorm_bf16, silu, softplus, gate_silu, apply_rope, grouped_query_attention, cached_causal_attention, layer_norm, gelu, softmax, round_to_bf16, and discretize_and_scan.

The exp and log kernels

exp and log use the Cephes method, because the compiler folds a scalar exp back into a per-element libm call. exp reduces its argument to a small interval and then evaluates a polynomial:

\[ n = \left\lfloor x \log_2 e \right\rceil, \qquad r = x - n \ln 2, \qquad \exp(x) = 2^{n}\, P(r) \]

\(P\) is a degree-5 minimax polynomial on \(r \in [-\tfrac{\ln 2}{2}, \tfrac{\ln 2}{2}]\), and \(2^{n}\) is applied by writing \(n\) into the float exponent. The result is within about 1 ULP. The vector forms are exp8 on AVX2 and exp4 on NEON.

Why this way

Almost everything here is bandwidth-bound. That sets the optimization order.

matmul dominates the runtime, about 84 percent of a layer. A square GEMM would be compute-bound, with \(O(N^3)\) FLOPs over \(O(N^2)\) bytes, but that is not the regime a policy runs in. The prefix is short, so each weight is read only a handful of times and the arithmetic intensity falls to roughly \(2\,\text{rows}\) FLOPs per byte, well under the machine balance. Measurement agrees: in_proj and out_proj have very different shapes and both land near 15 GB/s, which is a DRAM ceiling rather than an FMA one, while the small cache-resident x_proj runs half again faster. So the lever on matmul is weight traffic, not FMA blocking. Storing weights in BF16 halves the bytes, INT8 quarters them, and threading buys more memory parallelism. See ADR 0007.

The scan path is bandwidth-bound too, so the CPU backend fuses discretization and state update into discretize_and_scan. It streams \(\bar{A}\) and \(\bar{B}u\) once with almost no reuse, so intensity is near one FLOP per byte and the ALU waits on memory. Fusing removes an intermediate pass and keeps the model code small.

matmul, the activations, and the fused scan carry explicit intrinsics on SIMD backends. Packed Diffusion Policy convolution reuses matmul. rmsnorm and the Transformer helpers stay scalar unless a profile says otherwise.

The scan stores its state in a \([t][n][c]\) layout so that the inner loop over channels is unit-stride, which turns into contiguous vector loads over full cache lines. A \([t][c][n]\) layout would stride the reduction and waste bandwidth on partial lines. See ADR 0001.

Threading

To saturate memory bandwidth, large matrix multiplications dispatch concurrently across an SPMC lock-free thread pool. Workers spin briefly (_mm_pause / yield) to absorb gaps inside a forward pass, then park with C++ atomic::wait. Enqueue increments a work epoch and wakes a worker, avoiding both a lost-wakeup race and permanently burning CPU between inference requests. Threads are pinned to CPU cores where the platform allows it. The task ring and worker storage are arena-carved, power-of-two sized, and isolated on destructive-interference boundaries.

Pool capacity and per-kernel parallelism are separate. The runtime defaults to at most four workers, while an explicit override can provision up to eight. Each matrix uses a C++23 power-of-two task selector based on rows, input width, output width, and F32 versus BF16 work. Small projections stay on the caller thread; large projections use only the useful 2/4/8 tier. This avoids dispatch-dominated kernels and SMT/virtualization oversubscription.

BF16 Weight Widening

in_proj, out_proj, and Mamba projection weights may be stored as BF16 in the arena. The loader exposes a typed WeightView, and model code dispatches through weight_ops.h, so F32 and BF16 weights share one high-level call. BF16 is widened to F32 inline during the innermost SIMD loop of matmul.

This doubles effective DRAM bandwidth at the negligible cost of zero-extension and shift instructions (_mm256_castsi256_ps on AVX2, or vreinterpretq_f32_u32 on NEON).

Because the model only ever calls this header, a new backend is the same ten functions and a link-time switch.