Kernels

In short: every optimized kernel in rml sits next to a slow, obviously correct reference and a test that holds the two together (exact for integer arithmetic, a stated tolerance for floating point). Five kernel families exist at the pin: Gemm (Goto/BLIS packing with AVX-512, AVX2 and portable micro-kernels chosen at run time, plus matrix-vector and small-batch streaming paths), VecDot (weights in their GGML block format against int8 activations, with sdot on Apple Silicon), Attend (tiled online softmax over an f32, f16 or int8 KV cache), the decoder fusions in rml/llm/kernels.h, and StridedCopy. All numbers so far come from the 4-vCPU Linux x86_64 development container, not the Apple Silicon target; the NEON paths are compiled and tested only in macOS CI, and the Accelerate comparison, the smmla prefill path, Metal and the backend dispatch module are planned.

The rule comes from the project's guidance: every optimized kernel has a scalar reference and a test comparing the two (AGENTS.md L62–63), and each section of docs/kernels.md states what the kernel computes, its reference, its tiling and register budget, its test, and measured speed (kernels.md L1–7). The concepts behind them are on /ml-models: matrix multiplication, quantization, attention and profiling.

Which kernel runs where

GEMM

C = alpha · A B + beta · C on f32, over MatrixViews with a row and a column stride, so a transpose is a view (kernels.md L9–17). GemmRef is the textbook triple loop accumulating in double, slow on purpose so that tolerances on faster kernels reflect only their own rounding. Gemm is the Goto/BLIS scheme with a micro-kernel picked at run time from the CPU's features (L19–30):

The x86 kernels are compiled with __attribute__((target(...))) so a baseline build still runs on any x86-64 (gemm.cc L133, L163), and the portable tile is 8 × 12 on __aarch64__, 6 × 8 elsewhere (L97–104). The design points (kernels.md L32–80):

• Blocking. k in 256-deep slices; B packed into NR-column panels and each 96-row block of A into MR-row panels, so the micro-kernel reads two contiguous streams whatever the strides. A packed A block (96 KB) sits in L2, a B panel (at most 32 KB) in L1.
• Matrix-vector (one row of A, the decode shape) skips packing: each row of B is added, scaled, into a chunk of C with a vectorized axpy, split across threads.
• Up to 32 rows also skip packing: register tiles of 8 × 48 (AVX-512) or 4 × 24 (AVX2), fully unrolled (GCC otherwise spilled all 24 accumulators every step, halving speed), with B rows prefetched 8 steps ahead by hand.
• Threads split (row block, column chunk) pairs; with fewer than four pairs per thread the driver widens chunks to balance the slowest thread.
• Constant weights with more than 32 rows on the other side are packed once when a graph compiles (PackedB / GemmPackedB).

gemm_test compares every kernel the CPU supports with GemmRef over shapes from 1 × 1 × 1 to 200 × 401 × 257, every transpose, non-trivial alpha and beta, a padded ldc and 1–8 threads, within 4e-6 · (k + 1), also under ASan, UBSan and TSan; the NEON tile only in CI on macos-14 (L82–97).

kernels.md L99–110. The AVX-512 kernel is 6–7× the SSE-only portable kernel it replaced on that machine. The streaming path's numbers (80 GFLOP/s for 8 × 256 × 1024 on one core) and the effect of balanced chunks and packed weights on the zoo's graphs are on Linear algebra and Lowering and fast CPU.

Quantized dot products

VecDot(type, w, x, n) is one weight row in its stored format against activations quantized to int8 on the fly (kernels.md L112–136). Q8_0 and Q4_0 pair with Q8_0 activations (32 int8 and an f16 scale per block); Q4_K, Q5_K and Q6_K with Q8_K (256 int8, an f32 scale, and the sum of every 16, so a K-quant's minimum term costs one multiply). Within a block every product is exact int32; only block scales are float. Each block is unpacked to an aligned int8 array and dotted 16 or 32 lanes at a time: sdot with ARMv8.4 dotprod (every Apple Silicon core; dot.cc L11), a widening loop that SSE2 lowers to pmaddwd elsewhere. Because both functions share the float step, VecDotRef and VecDot are compared with EXPECT_FLOAT_EQ.

kernels.md L138–147: SSE2 has no int8 multiply, so these are a floor; the target is the sdot path, verified only on macos-14 CI. llm::MatMul uses them whenever the caller passes scratch for the quantized input, which DecoderLM sizes once, so decoding stays allocation-free; rounding activations to int8 moves toy-model logits by about 1e-2, as in llama.cpp (L149–155).

Attention

Attend is one pass of online softmax over the KV cache: a task is one kv head and one query position, so the query heads sharing a kv head read each key and value row once; keys are tiled 32 at a time, the running output stays in L1, and the output is rescaled once per tile. The test holds it within 2e-5 of AttendRef and threaded output bit-identical to serial (kernels.md L157–168). One decode step of a Llama-3-8B layer (32 heads, 8 kv heads, 128 dims) on one x86 core, --config=opt:

kernels.md L170–178; the int8 column was measured in a later session, in which f32 took 1.8 / 23 ms and f16 1.7 / 25 ms. An f16 row is widened by HalfToFloat (dtype.h L63), branch-free and bit-exact against Float16::ToFloat on all 65536 halves; without F16C on SSE2 that costs more than the halved traffic saves, while on Apple Silicon fcvtl makes it nearly free. int8 widens with one multiply and is under half of f16's size, so it is the fastest of the three at 4096 tokens there (L180–194). The kernel itself is drawn on The decoder.

Decoder fusions

Each fusion in rml/llm/kernels.h removes a pass over an activation buffer, or the buffer (kernels.md L196–212):

kernels.md L214–228. The model is random-weight and Llama-shaped (hidden 2048, 2 layers, FFN 5632, vocab 8192), and the times are one decode step. The difference is within run-to-run noise: at decode the weight reads dominate and the buffers removed are a few KB. What the fusions buy is fewer pool dispatches per dense layer (4 instead of 7) and a smaller workspace, which rml expects to matter more at Apple Silicon's per-token speeds. kernels_test holds every fused kernel bit-identical to the unfused sequence for F32, F16, Q8_0 and Q4_0, serial and pooled.

Strided copies

Contiguous, Clone and CopyFrom all go through StridedCopy, which graph ops call whenever a view must be materialized (a reshape of a broadcast, a concat, a conv input). It drops size-1 axes and merges axes both operands step through as one, copies the innermost two axes as a block (memcpy for unit strides, fill_n for a broadcast source, a strided loop otherwise), and splits copies of 64K elements or more over the default pool (kernels.md L230–247; tensor.cc L161). The old loop handled one innermost row per callback on one thread, about 4 ns an element on a nearest-2× upsample's 2-element rows.

kernels.md L249–256.

A negative result: folding the recurrence bias

An RNN or LSTM step at batch 64 is one 64 × 256 × N GEMM followed by bias and input-projection adds and a fused gate epilogue. Three attempts folded the adds (and then the activation) into the GEMM, or ran all column slices in one dispatch; none beat the unfused path at 4 threads on torch_rnn_b64 or torch_lstm_b64. In BM_RnnSteps a step took 103 µs fused against 86 µs unfused on one thread, and 54–62 µs against 58–62 µs on four. The per-step GEMM already runs at 245–290 GFLOP/s on 4 threads, about the machine's peak, so the code was not kept; the LSTM's remaining gap to PyTorch (8.8 ms against 13–16 ms) is the GEMM itself (kernels.md L258–273). The same story from the graph side is on Lowering and fast CPU.

The graph kernels

docs/kernels.md has no section for rml/graph/kernels.h, the graph executor's fast f32 kernels for elementwise math, broadcasting binary ops, softmax, layer norm, row reductions, im2col and pooling. Its header states the same contract: each fast kernel has a scalar reference (libm, double accumulation) that the tests compare it with; loops are compiled once per instruction set (AVX-512, AVX2, baseline) with target_clones and chosen at load time; the transcendental functions are branch-free polynomials so the loops vectorize (kernels.h L4–12). They are described with the fusion layer that drives them on Lowering and fast CPU.

Planned

Done in Phase 5 (progress.md L61–71): the packed f32 GEMM, the quantized GEMV with on-the-fly int8 activations (its sdot path verified only on macOS CI), fused tiled attention, the f16 KV cache, the decoder fusions, and zero allocations per decoded token, verified by decoder_alloc_test. The Mac comparison against llama.cpp, which rml's headline goal (decode within about 1.5× of llama.cpp's CPU speed; AGENTS.md L13–15) depends on, has not been run.

Tests

How these run in CI, including the macos-14 jobs that are the only place the NEON and sdot paths execute, is on Build and testing.