AKO

Adaptive Kernel Optimizer

A tool that checks that an optimized version of a computation is still correct before using it, then keeps the fastest, case by case.A tool that checks that an optimized computation is still correct before using it, then keeps the fastest, by shape.A for swapping in an : each candidate is checked before use, then chosen by shape.

Measured on a single laptop (RTX 3070 graphics card). Nothing is claimed for other machines.Measured on an . No claim about other .Measured on an . No claim about other .

See the code on GitHub Go to the results

Diagram of the flow, not a measurement. The truncated candidate stands for a kernel that is fast but wrong, rejected before it is measured.

The problem

An optimized computation is not faster everywhere, and it can be wrong without saying so.An optimized computation is not faster everywhere, and it can be wrong without saying so.A is not faster everywhere, and it can be wrong without saying so.

The “fast” version sometimes loses to . It can also win by skipping part of the data. You have to check before you choose.Across 72 configurations tested, the optimized version loses to in 22 cases, by up to 55%. A version can also win by ignoring data, and a fixed error threshold rejects correct versions. AKO brings a correctness contract.On the RTX 3070, the loses to on 22 of the 72 × configurations, by up to 55% (: 8 shapes out of 36 in , 14 out of 36 in ). An whose BLOCK_SIZE caps at 8192 truncates wider rows and still wins the . A fixed max_abs_error < 1e-2 threshold rejects a correct bf16 kernel that is only 1.00 off. Swapping a kernel comes with no numerical contract: AKO provides one.

Approach

Validate before measuring, measure before choosing, keep the reason for the choice.

The first time a computation comes up, AKO tries every version, drops the ones that get it wrong, times the rest and keeps the fastest.On the first call, AKO runs every candidate version, rejects those whose error exceeds a tolerance matched to the numbers' precision, measures the rest and records the decision with its reasons. The standard version always stays a candidate.On the first call for a given operation, , and , AKO runs every registered candidate (, F.rms_norm, , ). Those whose error exceeds the budget are rejected; a candidate that raises an exception is dropped. This budget, expressed in , depends on the dtype's precision and the result's magnitude; it is never hard-coded. The survivors are measured (), and the decision is written to a readable cache (~/.ako/kernel_cache.json), with its metrics.

Three operations covered, each tested automatically.Operations covered: , fused residual + RMSNorm, (neox and interleaved). The always stays a candidate. tested without a via TRITON_INTERPRET=1; LLaMA-style checked against transformers (exact match in ). 107 tests without a GPU, 18 GPU tests.

empty cache
validation
benchmark
cached decision

Unfused referencebenchmarked
F.rms_norm (PyTorch)benchmarked
torch.compilebenchmarked
Tritonchosen, cached
Truncated Tritonover budget

Illustration of the flow, not a measurement.Illustration of the flow, not a measurement.Illustration of the flow for a given , not a measurement. The winner changes with the shape and the .

Results

Against 's best option: 1.03 to 1.17 times faster, median. That is the real gain. Against an unoptimized base version: 7.3 times, a number that mostly measures how slow that version is.Against the best : 1.03× to 1.17× median. That is the real gain. Against the : 7.3× on , a number that mostly measures its memory traffic.Medians on an , 2.6, 12.4, WSL2. The big number mostly measures the memory traffic of the ; the useful column is the best (almost always ).

Triton kernel, median gain by operation and dtype
Operation	dtype	vs unfused reference	vs best baseline	Shapes lost
RMSNorm	fp16	7.30×	1.07×	8 / 36
RMSNorm	bf16	6.95×	1.03×	14 / 36
RoPE	fp16	5.31×	1.17×	1 / 8
RoPE	bf16	5.40×	1.07×	3 / 8

The project's value: check each version before using it, and choose case by case, because none wins everywhere.The project's value: on the configurations measured, choosing per avoids up to a 55% penalty ( forced everywhere) and up to 314% on ( forced everywhere).Table below, then penalties and sensitivity.

Bar chart: median speedup of Triton kernels for RMSNorm and RoPE in fp16 and bf16. Against the unfused reference: 5.4 to 7.4 times. Against the best alternative: 1.03 to 1.18 times. — The choice of baseline decides the headline number: up to 7.4× against an unfused version, 1.03× to 1.18× against the best alternative.

Two horizontal bar charts, RMSNorm (72 configurations) and RoPE (16). Always torch.compile: worst case 2.50 times and 4.14 times slower than the best choice. Always Triton: 1.55 times and 1.32 times. Medians stay close to 1. — No fixed choice is safe: the median stays close to 1×, but the cost shows up in the worst cases (up to 4.14× slower).

Scatter plot: gain of the Triton kernel over the best alternative by tensor size (1 to 500 MB), in fp16 and bf16. Points range from 0.65 to 2.5: the same kernel wins or loses depending on size. — The same kernel wins or loses depending on tensor size: below the line at 1, the Triton kernel is slower.

Real gain over what PyTorch already provides: 1.0× to 1.2×.
Penalty against the best per-shape choice (median / geometric mean / worst case): Triton forced 1.00× / 1.04× / 1.55×; torch.compile forced 1.04× / 1.09× / 2.50×; F.rms_norm forced 6.65× / 5.66× / 8.81×; per-shape selection 1.00× throughout. On RoPE, forcing torch.compile costs up to 314%.
Sensitivity: the same kernel wins 2.5× at 4 MiB and loses 35% at 8 MiB, same dtype.
Example cached decision: rmsnorm_triton p50 0.319 ms versus rmsnorm_compiled 0.341 ms (1.07×), max error 0.00195 for a budget of 0.0234.

Limits

When generating text, compiled stays ahead (0.69×), cause under investigation. Estimated impact on the whole model: about 1%. Only one computer tested.In , 0.69× against the , cause under investigation. The operations covered account for about 24% of a 's time: overall impact of about 1%. Only one measured; attention and matrix multiplication not covered.

In decode, 0.69× against the compiled stack. Regression unexplained to date; the sweeps include no decode shapes.
Operations covered ≈ 24% of a decoder's time: overall impact estimated at about 1%.
torch.compile slower than eager on the stack (10.75 ms versus 10.39 ms in prefill, recompilations): the 1.24× in prefill is optimistic.
INT8/INT4 quantized matmul (~50% of the time) and IO-aware attention (~25%) not started. No GEMM.
Only one GPU measured: no claim of hardware portability.
Not an inference server, not a kernel library (3 operations), NVIDIA only.
Not on PyPI, no semantic versioning.

See the code on GitHub (github.com/Farx1/AKO)Next project: QCompress