What is the tokamax autotuning cache key made of?

The names of the arguments mapped to their shapes and dtypes, and nothing else. The key builder reads the positional parameter names off the forward implementation signature, zips them with the positional arguments, and concatenates that with the keyword arguments, running everything through an abstraction step that replaces each array with a jax.ShapeDtypeStruct first. The result is an immutabledict, which is hashable, so it can index a plain dictionary. Two consequences follow. Passing different numbers in arrays of the same shape always hits the same cache entry, which is deliberate: tiling depends on shape, not on values. And the key is built from the private _fwd signature with the config parameter removed, not from the public bind signature, so what appears in the key is whatever _fwd declares.

What happens on a tokamax autotuning cache miss?

By default you get a logged warning and an unmeasured heuristics config, and your program runs. The lookup indexes the on-disk cache directly and catches the KeyError, serialises the key back to JSON, and logs a warning naming the op, the device kind and the full key, then returns None. The ladder above it then falls through to the backend heuristics. A global flag, tokamax_autotuning_cache_miss_fallback, changes that: set to autotune, the tuner runs inline on that first call and compiles every candidate config in your process; set to error, you get a ValueError naming the bound arguments instead. The default is heuristics. Note what a miss does not do: it never changes which implementation runs. That was decided before the config was ever resolved.

Why does a tokamax op with no config never read the autotuning cache?

Because of the order of the tests in get_config, which differs from the order its docstring lists. After the check for an explicit config on the op, the next thing the method does is compute the heuristics config and compare it by identity against _NULL_CONFIG, the single module-level instance of the empty config dataclass. If it is that singleton, the method returns immediately, before the cache is consulted at all. An op whose config_cls is NullConfig has nothing to tune, so its base heuristics method hands back that singleton and the lookup is skipped. If you are looking at a shipped cache JSON for such an op and wondering why nothing reads it, that file is measurement data for comparison, not a dispatch input.

How does tokamax tune a kernel offline from a production job?

By writing the whole call into the compiled program and recovering it later. Before dispatching, the op is copied with the resolved config and no VJP function, its arguments are replaced with shape-and-dtype descriptors, and the pair is dumped to JSON by a Pydantic adapter. That string, prefixed with tokamax:, is set as an XLA metadata payload for the duration of the forward call, so it lands as an mhlo.frontend_attributes entry on every operation the kernel emits and survives into the lowered module. Tooling then walks a captured module or an XProf trace, finds the custom calls, pulls the payload off each, balances braces to isolate the JSON, and validates it back into a BoundArguments. Those can be tuned on a separate machine, because the payload carries shapes and dtypes rather than tensors. A nested tokamax call appends to the enclosing payload rather than replacing it, so the outermost op appears first.

Why does tokamax wrap every op in jax.custom_vjp even without a gradient function?

Because of a composition gap between two JAX features, and the file says so in a comment. Tokamax uses custom_batching to capture the vmap axes of a call, since the autotuner needs the batched shapes and an ordinary function body cannot see them. custom_batching does not support jax.vjp, so the op has to be wrapped in a custom_vjp even when no vjp function was supplied; in that case the backward rule is a one-liner that simply invokes the closure the forward pass returned. There is one escape: an op that sets supports_batched_args_capture to False skips the wrapper entirely and calls the plain function. A related workaround sits nearby, a small _AlwaysEqual wrapper that makes two closures from separate traces compare equal, without which the surrounding jit would recompile on every step.

The tokamax Op and Its Autotuner, Line by Line

A fast kernel is never one kernel. It is a family of implementations and a family of tilings, and something has to choose between them on every call. Tokamax, the fused-kernel library Google builds on Pallas, answers that in one 701-line Python file. This walks it whole at a pinned commit: the frozen Op dataclass and the two hooks a backend overrides, the twelve lines that turn a call into a dispatch key made only of argument shapes and dtypes, the five-step ladder that decides whether you get an explicit config, a cached measurement, an unmeasured guess or an error, the measurement loop that fills the cache in, and the payload the file writes into the compiled program so a captured production trace can be tuned on a machine that never runs your model. The signature exhibit is that real key, opened against the library’s own shipped cache: pick a ragged-dot workload and a chip, and read every cached implementation with its winning tiling and its measured median.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The tokamax Op and Its Autotuner, Line by Line — built to be understood, not skimmed.

Questions

What is the tokamax autotuning cache key made of?: The names of the arguments mapped to their shapes and dtypes, and nothing else. The key builder reads the positional parameter names off the forward implementation signature, zips them with the positional arguments, and concatenates that with the keyword arguments, running everything through an abstraction step that replaces each array with a jax.ShapeDtypeStruct first. The result is an immutabledict, which is hashable, so it can index a plain dictionary. Two consequences follow. Passing different numbers in arrays of the same shape always hits the same cache entry, which is deliberate: tiling depends on shape, not on values. And the key is built from the private _fwd signature with the config parameter removed, not from the public bind signature, so what appears in the key is whatever _fwd declares.
What happens on a tokamax autotuning cache miss?: By default you get a logged warning and an unmeasured heuristics config, and your program runs. The lookup indexes the on-disk cache directly and catches the KeyError, serialises the key back to JSON, and logs a warning naming the op, the device kind and the full key, then returns None. The ladder above it then falls through to the backend heuristics. A global flag, tokamax_autotuning_cache_miss_fallback, changes that: set to autotune, the tuner runs inline on that first call and compiles every candidate config in your process; set to error, you get a ValueError naming the bound arguments instead. The default is heuristics. Note what a miss does not do: it never changes which implementation runs. That was decided before the config was ever resolved.
Why does a tokamax op with no config never read the autotuning cache?: Because of the order of the tests in get_config, which differs from the order its docstring lists. After the check for an explicit config on the op, the next thing the method does is compute the heuristics config and compare it by identity against _NULL_CONFIG, the single module-level instance of the empty config dataclass. If it is that singleton, the method returns immediately, before the cache is consulted at all. An op whose config_cls is NullConfig has nothing to tune, so its base heuristics method hands back that singleton and the lookup is skipped. If you are looking at a shipped cache JSON for such an op and wondering why nothing reads it, that file is measurement data for comparison, not a dispatch input.
How does tokamax tune a kernel offline from a production job?: By writing the whole call into the compiled program and recovering it later. Before dispatching, the op is copied with the resolved config and no VJP function, its arguments are replaced with shape-and-dtype descriptors, and the pair is dumped to JSON by a Pydantic adapter. That string, prefixed with tokamax:, is set as an XLA metadata payload for the duration of the forward call, so it lands as an mhlo.frontend_attributes entry on every operation the kernel emits and survives into the lowered module. Tooling then walks a captured module or an XProf trace, finds the custom calls, pulls the payload off each, balances braces to isolate the JSON, and validates it back into a BoundArguments. Those can be tuned on a separate machine, because the payload carries shapes and dtypes rather than tensors. A nested tokamax call appends to the enclosing payload rather than replacing it, so the outermost op appears first.
Why does tokamax wrap every op in jax.custom_vjp even without a gradient function?: Because of a composition gap between two JAX features, and the file says so in a comment. Tokamax uses custom_batching to capture the vmap axes of a call, since the autotuner needs the batched shapes and an ordinary function body cannot see them. custom_batching does not support jax.vjp, so the op has to be wrapped in a custom_vjp even when no vjp function was supplied; in that case the backward rule is a one-liner that simply invokes the closure the forward pass returned. There is one escape: an op that sets supports_batched_args_capture to False skips the wrapper entirely and calls the plain function. A related workaround sits nearby, a small _AlwaysEqual wrapper that makes two closures from separate traces compare equal, without which the surrounding jit would recompile on every step.

The tokamax Op and Its Autotuner, Line by Line

Questions

Related explainers