teenygrad kernels / Making It Fast
Numerics
A kernel that is fast and slightly wrong is worse than one that is slow and right, because the wrongness is invisible until it is expensive.
This chapter is about choosing dtypes and accumulators deliberately.
The type hierarchy
The dtype bounds you put on a kernel are a real constraint, checked at compile time:
pub trait Dtype: Copy + Clone {}
pub trait Num: Dtype { const BITS: u8; }
pub trait Float: Num { const ZERO: Self; const ONE: Self; /* from_f64 */ }
pub trait Int: Num {}
pub trait Bool: Dtype + Copy {}
Pick the tightest one that admits your operations:
| Bound | Admits | Use when |
|---|---|---|
Float |
f32, f64 |
You call exp, log, sqrt, sigmoid |
Int |
i8…i64, u8…u64 |
Bitwise operations, integer atomics |
Num |
both | Arithmetic only |
Dtype |
everything | Pure data movement |
Float also gives you D::from_f64(...), which is how you get a constant of
the right type into a kernel — T::full(&[BLOCK], D::from_f64(0.5)). There is
no way to write a float literal of a generic float type without it.
Choosing the tightest bound is not pedantry. D: Num on a kernel that calls
T::exp is a compile error, and that error is the one Python Triton would have
given you on a GPU, at run time, in production.
What is actually available
DtypeRepr — the runtime, type-erased tag — has thirteen variants including
F16 and BF16. The macro’s dtype-set inference does not:
| Bound | Dtypes you get |
|---|---|
Float |
f32, f64 |
Int |
i8, i16, i32, i64, u8, u16, u32, u64 |
Num |
all of the above |
Bool |
bool |
Dtype |
all of the above |
f16 and bf16 are absent, deliberately: the macro’s own comment says they are
marker-only and cannot be monomorphized, because there is no concrete Rust
implementation to instantiate a kernel against.
This is the largest gap in the book. Half precision is a substantial part of
why people write GPU kernels — it halves memory traffic and doubles Tensor Core
throughput — and today you cannot write a #[kernel] that dispatches to it. It
is item 4 in
KNOWN-GAPS.md.
What follows is therefore mostly about f32, which is what you can use.
The formats, for when this is fixed
Worth knowing, because the trade-offs are what motivate everything below.
| Format | Bits | Exponent | Mantissa | Character |
|---|---|---|---|---|
f64 |
64 | 11 | 52 | Rarely worth it on a GPU |
f32 |
32 | 8 | 23 | The default |
| TF32 | 19 stored in 32 | 8 | 10 | f32’s range, less precision, Tensor Cores |
bf16 |
16 | 8 | 7 | f32’s range, much less precision |
f16 |
16 | 5 | 10 | More precision than bf16, far less range |
The exponent column is the important one. bf16 has the same range as f32, so
a value that fits in f32 fits in bf16 — you lose precision, not magnitude.
f16 has five exponent bits and overflows above about 65,504, which is why
training in f16 needs loss scaling and training in bf16 mostly does not.
What f64 actually costs
examples/numerics.rs runs an elementwise add over 16M elements, once as f32
and once as f64:
cargo run --release -p teeny-triton --features cuda --example numerics
On an RTX 5070 (sm_120), CUDA 13.3, driver 610.43.02:
| dtype | time | bandwidth | moved |
|---|---|---|---|
f32 |
339.5 µs | 593.1 GB/s | 201 MB |
f64 |
676.7 µs | 595.0 GB/s | 403 MB |
f64 costs 1.99× the time — and the same bandwidth.
That equality is the whole story. This kernel is memory-bound, so f64 costs
exactly what its extra bytes cost and not one thing more. Both runs saturate the
card at ~594 GB/s, the same ceiling Chapters 16 and 17 hit with different
kernels.
Which means the folk wisdom “f64 is catastrophically slow on a consumer GPU”
is, for this kernel, wrong. It is catastrophically slow when arithmetic is the
bottleneck — consumer cards run f64 arithmetic at a small fraction of their
f32 rate — and this kernel barely does any. Move a byte, add once, move a
byte.
So the honest rule: on a memory-bound kernel f64 costs 2×; on a
compute-bound one it costs far more. Chapter 1’s question — which kind is
this? — decides which number applies.
Accumulators
Here is the rule that matters most:
Accumulate in
f32, whatever you multiply in.
T::dot has two type parameters precisely for this:
fn dot<D: Num, O: Num>(a: Tensor<D>, b: Tensor<D>, acc: Option<Tensor<O>>, ...) -> Tensor<O>;
D is the inputs, O is the accumulator. Mixing them is not a compromise; it is
the correct arrangement, and the Tensor Cores are built for it.
The reason is that error accumulates with the number of additions. Multiplying
two f16 values gives a result f16 can represent fine. Adding a thousand of
them in f16 does not — each addition rounds, and after K additions the error
has grown roughly with sqrt(K) at best. In f32, with sixteen more mantissa
bits, the same sequence stays accurate.
The same applies to hand-written reductions. Summing a long f16 row into an
f16 total loses precision that summing into f32 does not.
Measured
The same example sums 16M f32 values four ways. The reference is the identical
data summed in f64, so the only thing varying is how the sum is accumulated:
| how | result | relative error |
|---|---|---|
exact (f64 reference) |
2181038.0593 | — |
GPU block reduction, f64 accumulator |
2181037.9700 | 4.09e-8 |
GPU block reduction, f32 accumulator |
2181055.7500 | 8.11e-6 |
sequential f32 loop on the CPU |
2158069.0000 | 1.05e-2 |
Read the last row first. A plain f32 loop is wrong in the fifth
significant figure — a 1% error, from nothing but adding numbers up in order.
Once the running total reaches ~2 million, adding 0.1 to it barely moves it, and
16 million such additions lose most of what they should have contributed.
The GPU’s block reduction is over a thousand times more accurate than that loop, and it is not because the GPU is careful. It is because a tree adds numbers of similar magnitude to each other: 16M values pair down through 24 levels, and no partial sum ever dwarfs what is being added to it. The parallel algorithm is more accurate than the obvious sequential one, which is the opposite of what most people expect.
Widening only the final accumulator — the host-side sum of the per-block
partials — from f32 to f64 gains another 200×. That is one cast, on
16384 values, and it is free next to the kernel.
So the practical shape of an accurate reduction is: reduce in blocks on the device, accumulate the partials in something wider. You get the tree for free and the wide accumulator for almost nothing.
And it is not reproducible
The same data through the same kernel at two different block sizes:
BLOCK=256 → 2181037.969997
BLOCK=1024 → 2181037.969986
A different block size is a different tree shape, and floating-point addition is not associative, so the answers differ — here by 5.2e-12 relative. Tiny, real, and enough to break an exact-equality assertion. This is the concrete version of the warning below.
InputPrecision
For f32 × f32, T::dot lets you choose what the hardware actually does:
| Value | Behaviour |
|---|---|
TF32 |
Round inputs to 19 bits, use Tensor Cores. Fastest. Default on capable hardware |
TF32x3 |
Three TF32 products combined to recover most of f32’s precision |
IEEE |
True f32 arithmetic. Slowest, exact |
TF32 keeps f32’s range and drops thirteen mantissa bits. For neural network
training that is almost always fine, which is why it is the default.
IEEEturns off the Tensor Cores. It is not a slightly slower, slightly more accurate mode — Triton only routes adotto the matrix hardware for the reduced-precision modes, soIEEEfalls back to software fused-multiply-add. The fused conv kernel in this tree was written withIEEEto match cuDNN’s accumulation and silently lost its Tensor Cores; it now usesTF32.
So “it felt safer” is an expensive instinct here. Choose IEEE when you need
exact f32 more than you need the hardware, knowing that is the trade.
TF32x3 is the middle option — three TF32 products recovering most of f32’s
precision while staying on the Tensor Cores — and is worth knowing about because
most people do not know it exists.
Casting
fn cast<Src: Dtype, Dst: Dtype>(x: Tensor<Src>, rounding: Option<FpDowncastRounding>, bitcast: bool) -> Tensor<Dst>;
Two arguments to be careful with.
FpDowncastRounding applies when narrowing: Rtne rounds to nearest with ties
to even, Rtz truncates toward zero. Rtne is what you want; Rtz biases every
value toward zero, and a systematic bias through a training loop compounds in a
way that random rounding error does not.
bitcast: true reinterprets the bits rather than converting the value. It is
occasionally what you want, and it is never what you want by accident.
Where precision goes
Catastrophic cancellation. Subtracting two nearly equal large numbers leaves
a result dominated by their rounding error. sqrt(a² - b²) is the classic; use
sqrt((a-b)*(a+b)).
Summing many values. Error grows with the count. A long reduction in f32
is fine; the same in f16 is not.
Exponentials. exp(x) overflows f32 above about 88, which is why softmax
subtracts the maximum first — Chapter 10.
Division by something near zero. Normalisation layers add an epsilon for this reason.
Determinism
Two things make GPU results non-reproducible, and both are expected rather than broken:
Floating-point addition is not associative. A reduction combines in an unspecified order, so the last bits can differ between a GPU and a CPU reference, or between two block sizes — measured above at 5.2e-12 relative for a block size change alone.
Atomics arrive in an unspecified order. Chapter 14. So a backward pass using them is not bit-reproducible run to run.
Consequences for your tests: compare with a tolerance, never exact equality. The
tests in this tree use 1e-5 for a forward pass and 1e-6 for a backward one.
And when a test does fail, check whether the difference is at the tolerance
boundary — a mismatch in the last two bits is arithmetic; a mismatch in the first
digit is a bug.
A checklist
- Tightest bound that admits your operations.
Floatif you useexp. - Accumulate wider than you store.
f32forf16inputs;f64for the host-side sum off32partials, which costs nothing and bought 200× above. - Reduce in a tree, not a loop. You get this for free from
T::sum, and it was 1000× more accurate than the sequential version. - Leave
InputPrecisionatTF32unless you can say why not —IEEEturns the Tensor Cores off. Rtnewhen narrowing, notRtz.- Subtract the max before exponentiating.
- Compare with a tolerance in tests, and know which tolerance and why.
f64costs 2× on a memory-bound kernel and far more on a compute-bound one. Know which you have before ruling it out.- Half precision is not available yet. Design for it; do not depend on it.
End of Part 4
You can choose a block size, reason about layout, measure honestly, and pick dtypes deliberately.
Every number in this part was measured on one card, an RTX 5070, and each
chapter names it. Three independent kernels — the block-size sweep, the
coalescing comparison, and the f32/f64 add — all plateau within 1% of
594 GB/s, which is the most useful single fact in these four chapters: it is
this card’s ceiling, and it is what a new kernel should be judged against.
Your card’s number will differ. The method will not.