Elementwise and softmax baselines
Scalar add, mul, relu, and stable softmax baselines give every future fused kernel a correctness and speed reference.
Four scalar kernels
ops/ gets one file per operation, all on flat float arrays:
void ew_add(const float *a, const float *b, float *out, size_t n);
void ew_mul(const float *a, const float *b, float *out, size_t n);
void ew_relu(const float *in, float *out, size_t n);
void softmax_rows(const float *in, float *out, size_t rows, size_t cols);
Piecewise loops, -O2, no intrinsics. These are the oracle for every fused kernel later.
Softmax, stably
Two passes per row: find the row max, then compute exp(x - max) and its sum, then divide. Subtracting the max changes nothing mathematically, because softmax is shift-invariant, and it keeps every exponent at or below zero so exp cannot overflow.
Generate test data with a wide range, including values near 100. In f32, exp(100) exceeds the largest finite float and becomes inf, so the naive form produces inf/inf, which is NaN. The stable form agrees with the naive one on easy input, and only the hard input separates them, which is why the wide range belongs in the baseline rather than in a comment. The mirror case is a row of large negative values: the naive form underflows every term to zero and divides zero by zero, while the max-subtracted form returns a uniform row.
Bandwidth is the number that matters
An add reads 8 bytes and writes 4 per element for one flop. Reporting only GFLOP/s hides the limit, so print GB/s beside it; a later fused kernel’s win then has a unit that matches its bottleneck. ReLU needs one decision recorded in the header: x > 0 ? x : 0 sends NaN to zero, matching LLVM’s maxnum, not maximum.
Things that break
NaN in the output is the signature of a missing max subtraction, and it often appears only on rows your generator happened to make large. Assert that no output element is NaN before recording a time, so the failure is loud instead of a plausible-looking table row.
Sources
- The log-sum-exp trickblog
- FlashAttention: exact attention without materialising the score matrixpaper
- LLVM Language Reference: minnum, maxnum and NaN semanticsdocs
A version records the documentation the claim was checked against, not a house pin. Confirm flags with --help on the release you installed.
Check yourself
Why subtract the row maximum before calling `exp`?
A row holds only large-magnitude values near -1000. What does a stable softmax return?
A fused add-then-relu kernel is checked against what?