Skip to content

Kernel Library

Jump to bottom
Richard Huang edited this page Aug 3, 2026 · 1 revision

Kernel Library

Kernels are the actual implementations that run on the hardware — the code that performs a matrix multiply, an attention block, a normalization. The Graph Compiler decides what to compute and in what shape; kernels determine how fast that computation actually goes.

A kernel library is where hardware capability becomes realized performance. Peak throughput on a datasheet means nothing until a kernel achieves it.

Objectives

  • Approach peak — get as close to hardware limits as the operation allows
  • Cover the model — every operation needs an implementation
  • Stay correct — including at numerical edge cases
  • Remain maintainable — hand-tuned assembly that nobody can modify is a liability
  • Adapt to shapes — a kernel fast at one size is often slow at another

Where It Fits

flowchart LR
    A[Graph Compiler] --> B[Kernel Library]
    B --> C[Runtime]
    D[Memory Management] --> B
    style B fill:#2d6a9f,color:#fff
Loading

The Kernel Taxonomy

Class Examples Bound by Notes
Matrix multiply Linear projections, FFN Compute (large) / memory (small) The dominant cost
Attention Scaled dot-product attention Memory Deserves a specialized fused kernel
Elementwise Activations, residual add Memory Should almost always be fused, never standalone
Reduction Normalization, softmax Memory Fusible into neighbors
Data movement Transpose, gather, concat Memory Pure overhead; eliminate where possible
Quantize / dequantize Format conversion Memory Must be fused or it defeats Quantization

The pattern: almost everything except large matrix multiplies is memory bound, and the correct response to a memory-bound operation is usually to stop treating it as a separate kernel at all.

Matrix Multiplication

The core operation. Two regimes matter, and they need different kernels:

Regime Shape Character
GEMM (matrix-matrix) Large × large Prefill; compute bound; classic tiling and reuse
GEMV (matrix-vector) Large × 1 Decode; memory bound; every weight byte read once, used once

GEMV is the decode bottleneck and it cannot be fixed by faster arithmetic. Reading an entire weight matrix to multiply it by a single vector has arithmetic intensity near 1 — utterly memory bound. The only real levers are: read fewer bytes (Quantization), read them once for several vectors (batching, see Batching and Scheduling), or skip reading some entirely (sparsity).

Tiling for GEMM

The standard hierarchy mirrors the memory hierarchy:

Global memory  →  on-chip buffer  →  register / accumulator
    (big tile)        (small tile)        (element)

Each level's tile size is chosen so its working set fits in that level's capacity, maximizing reuse before data must be re-fetched. Double buffering at each level overlaps the next load with the current compute.

Accumulation precision

Multiply in low precision, accumulate in higher precision. Accumulating INT8 products into an INT8 accumulator overflows almost immediately; accumulating into INT32, or FP16 products into FP32, is standard. Getting this wrong produces plausible-looking but wrong results — one of the nastier bug classes in this area.

Attention Kernels

Naive attention computes the full score matrix, applies softmax, then multiplies by values. The score matrix is O(n²) in sequence length and must be written to memory and read back.

The fused approach processes attention in tiles that never materialize the full score matrix: for each block of keys and values, compute partial scores, update a running softmax, and accumulate into the output. Memory traffic drops from O(n²) to O(n), which is the difference between attention being the bottleneck and attention being reasonable.

The mechanism that makes this work is online softmax — computing a numerically stable softmax incrementally without seeing all inputs at once, by tracking a running maximum and rescaling accumulated results as new maxima appear.

Decode attention is a special case worth its own kernel: the query is a single token, so it's matrix-vector against the entire KV cache. This is bandwidth bound on cache reads, which makes cache layout and cache quantization the levers that matter.

Numerical Stability

Operation Hazard Standard fix
Softmax exp of large values overflows Subtract the row maximum first
Accumulation Precision loss over long sums Higher-precision accumulator; pairwise summation
Normalization Division by near-zero variance Epsilon in the denominator
Low-bit matmul Overflow in the accumulator Size the accumulator for the worst case

These are not optional refinements. A kernel that is fast and occasionally produces NaN is not a working kernel.

Implementation Strategies

Strategy Strengths Weaknesses
Hand-written Maximum performance Expensive to write, hard to maintain, one shape
Template / parameterized Reusable across shapes Needs good parameter selection
Autotuned Finds good parameters automatically Tuning time; needs a search space
Compiler-generated Scales to many operators Usually below hand-tuned peak

Mature stacks use all four: generated kernels for coverage, templates for the common cases, autotuning to pick parameters, and hand-written code for the two or three kernels that dominate runtime. Optimizing anything else by hand is usually wasted effort.

Metrics

Metric Why it matters
Achieved throughput vs. peak The headline efficiency number
Achieved bandwidth vs. peak The right metric for memory-bound kernels
Roofline position Tells you which limit you're actually against
Occupancy / utilization Whether the hardware is kept busy
Numerical error vs. reference Correctness under tolerance
Shape coverage Fraction of required shapes with a fast path

For memory-bound kernels, comparing against peak compute is meaningless. Compare against peak bandwidth — a decode GEMV at 90% of memory bandwidth is an excellent kernel even at 5% of peak FLOPs.

Common Problems

Problem Root cause Fix
Far below peak on GEMM Poor tiling; no reuse Retune tile sizes to capacity
Compute units idle No overlap of load and compute Double buffer; software pipeline
Fast at one shape, slow at another Single tuned configuration Multiple variants selected by shape
NaN in output Softmax overflow or accumulator overflow Max subtraction; wider accumulator
Quantized kernel no faster Dequantization not fused Fuse unpacking into the inner loop
Correct but slow attention Score matrix materialized Use a fused tiled attention kernel
Unaligned access penalty Layout mismatch Pad and align; fix layout in the compiler

Best Practices

  • Profile before optimizing. Two or three kernels usually dominate; the rest don't matter.
  • Know which roof you're under. Optimizing arithmetic in a memory-bound kernel accomplishes nothing.
  • Accumulate wide. Always.
  • Fuse the small stuff. Standalone elementwise kernels are almost always a mistake.
  • Test edge shapes — sequence length 1, non-multiple-of-tile sizes, empty batches. These are where kernels break.
  • Keep a naive reference for every kernel. Correctness testing needs ground truth.
  • Autotune per shape bucket rather than picking one configuration and hoping.
  • Write the fast path last. Get it correct, measure, then optimize what the measurement identifies.

See also: Graph Compiler · Quantization · Memory Management · Performance Optimization · Inference Software

Clone this wiki locally