Autotuning and Heuristics

Kernel performance may be significantly impacted by implementation-specific hyperparameters (e.g., tile sizes, layouts, etc.), known as Op Configs in Tokamax. The optimal values for these configs must be empirically determined for every input of interest, via autotuning. Tokamax provides a framework both to perform autotuning, and to simplify the management of this process and its outputs.

High-level Overview

Every tokamax.Op implementation defines its set of tunable hyperparameters, and the range of values each of those configs can assume, in tokamax.Op._get_autotuning_configs (example). The autotuner runs an exhaustive search across this space to identify the configuration with the best performance (lowest execution time). The result may be utilized in different ways: directly in your program as a context manager, serialized and stored in your private cache for later use, or as part of the global library-wide cache.

The optimal config for an op depends on the specific set of inputs used to invoke the op, including shapes, dtypes, etc. The autotuning API accepts this list of inputs in a variety of ways, and in all cases, outputs an AutotuningResult object that contains the best configurations for each op in the input.

Autotuner Input Formats

Callable Function

Given a callable function of interest with one or more Tokamax Ops, the autotuner uses StableHLO metadata to extract information about every Op and its respective inputs, and autotunes each one.

# Autotune all Tokamax kernels in f, which can be a JAX function with multiple Tokamax ops and non-Tokamax ops as well. Assume f takes a dictionary of args as input.
autotune_result: tokamax.AutotuningResult = tokamax.autotune(f, **args)
# Best possible result
with autotune_result:
    out_autotuned = f(**args)

Tokamax VJP ops may be implemented as their own Op class and have their own search space and unique _get_autotuning_configs. Since the VJP ops do not have a direct API that can be called for autotuning, you can take your forward function and call jax.grad to obtain the pullback function with Tokamax VJP ops, and then pass the pullback function into the autotuner.

# This pullback function contains Tokamax VJP ops.
f_grad = jax.grad(f)
autotune_result = tokamax.autotune(f_grad)

with autotune_result:
  out = f_grad()

Sequence of BoundArguments

When a Tokamax Op is called with a set of input arguments, the op will first "bind" to the arguments. This process canonicalizes and validates the inputs, creating a BoundArgument object consisting of the op with its input arguments. The autotuner can take in a sequence of BoundArguments as input, and autotunes each one.

ragged_dot_ba = tokamax.PallasMosaicTpuRaggedDot.bind(x, y, group)
attention_ba = tokamax.PallasMosaicTpuAttention.bind(q, k, v)

autotune_result: tokamax.AutotuningResult = tokamax.autotune([ragged_dot_ba, attention_ba])

with autotune_result:
  ...

Serialized List of BoundArguments

The script xplane_to_bound_args.py takes an xplane proto, generated by XProf during profiling runs, as input and outputs a JSON of BoundArguments. The autotuner can read in this JSON file and autotune all the Tokamax Ops that appear in the profile.

Suppose from profiling you have an xplane my_model.xplane.pb, you can extract and autotune all the Tokamax ops used like this:

 python tokamax/_src/tools/xplane_to_bound_args.py \
   --xplane_file=my_model.xplane.pb \
   --output_file=/tmp/bound_args.json

import tokamax

bound_args = tokamax.autotuning.bound_args_from_json("/tmp/bound_args.json")
autotune_result = tokamax.autotune(bound_args)

Autotuner Output Usage

The autotuning process always returns an AutotuningResult object that contains the best configurations for each op in the input. You may utilize that output in different ways to suit your needs.

Context Manager

You may directly utilize the autotuning result as a context manager in your code.

# Autotune all Tokamax kernels in f, which can be a JAX function with multiple Tokamax ops and non-Tokamax ops as well. Assume f takes a dictionary of args as input.
autotune_result: tokamax.AutotuningResult = tokamax.autotune(f, **args)
# Best possible result
with autotune_result:
    out_autotuned = f(**args)

Serialized AutotuningResult

You may wish to serialize and reuse autotune_result for one of two reasons (1) autotuning can be expensive in terms of compute time, depending on the size of the search space (2) kernel execution times can be noisy, resulting in different “best” choices in different runs, and may result in numerical non-determinism.

autotune_result: tokamax.AutotuningResult = tokamax.autotune(f, **args)
autotune_result_json: str = autotune_result.dumps()
autotune_result = tokamax.AutotuningResult.loads(autotune_result_json)
with autotune_result:
    out_autotuned = f(**args)

Multiple AutotuningResult objects can also be merged together into a single object, by reading the results in as a JSON string using load.

autotune_result_1 = tokamax.AutotuningResult.load(path_to_file_1)
autotune_result_2 = tokamax.AutotuningResult.load(path_to_file_2)
merged_results = autotune_result_1 | autotune_result_2

with merged_results:
  ...