(jax-501-compilation-cache)=
# Persistent compilation cache

<!--* freshness: { reviewed: '2026-09-22' } *-->

JAX has an optional disk cache for compiled programs. If enabled, JAX
stores copies of compiled programs on disk, which can save recompilation time
when running the same or similar tasks repeatedly.

Note: if the compilation cache is not on a local filesystem,
[etils](https://pypi.org/project/etils/) needs to be installed.

```sh
pip install etils
```

```{warning}
The compilation cache is considered trusted. Do not share a compilation cache
with users you do not trust. For example, if you put the compilation cache in
a directory to which others may write, those users can trigger your JAX
process to run arbitrary code. Sharing a compilation cache is equivalent to
allowing anyone who can write to the cache directory to run code on your
machine.
```

## Usage

### Quick start

```python
import jax
import jax.numpy as jnp

jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
jax.config.update("jax_persistent_cache_enable_xla_caches", "xla_gpu_per_fusion_autotune_cache_dir")

@jax.jit
def f(x):
  return x + 1

x = jnp.zeros((2, 2))
f(x)
```

### Setting cache directory

The compilation cache is enabled when the cache location
(`jax_compilation_cache_dir`) is set. This should be done before the first
compilation. Set the location in one of these ways:

(1) Using environment variable

In shell, before running the script:

```sh
export JAX_COMPILATION_CACHE_DIR="/tmp/jax_cache"
```

Or on the top of the Python script:

```python
import os
os.environ["JAX_COMPILATION_CACHE_DIR"] = "/tmp/jax_cache"
```

(2) Using `jax.config.update()`

```python
jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
```

(3) Using {func}`~jax.experimental.compilation_cache.compilation_cache.set_cache_dir`

```python
from jax.experimental.compilation_cache import compilation_cache as cc
cc.set_cache_dir("/tmp/jax_cache")
```

### Caching thresholds

* `jax_persistent_cache_min_compile_time_secs`: A computation will only be
   written to the persistent cache if the compilation time is longer than
   the specified value. The default is 1.0 second.

* `jax_persistent_cache_min_entry_size_bytes`: The minimum size (in bytes)
   of an entry that will be cached in the persistent compilation cache:

   *  `-1`: disable the size restriction and prevent overrides.

   *  Leave at default (`0`) to allow for overrides. The override will
      typically ensure that the minimum size is optimal for the file system
      being used for the cache.

   *  `> 0`: the actual minimum size desired; no overrides.

Note that both criteria need to be satisfied for a function to be cached.

### Additional caching

XLA supports additional caching mechanisms, which can be enabled alongside
JAX's persistent compilation cache to further improve recompilation time.

* `jax_persistent_cache_enable_xla_caches`: Possible values:

   * `all`: enable all XLA caching features

   * `none`: don't enable any extra XLA caching features

   * `xla_gpu_kernel_cache_file`: only enable the kernel cache

   * `xla_gpu_per_fusion_autotune_cache_dir`: (default value) only enable the
      autotuning cache


### Google Cloud

When running on Google Cloud, the compilation cache can be placed on a Google
Cloud Storage (GCS) bucket. We recommend the following configuration:

*  Create the bucket in the same region as where the workload will run.

*  Create the bucket in the same project as the workload’s VM(s). Ensure that
   permissions are set so that the VM(s) can write to the bucket.

*  There is no need for replication for smaller workloads. Larger workloads
   could benefit from replication.

*  Use “Standard” for the default storage class for the bucket.

*  Set the soft delete policy to its shortest: 7 days.

*  Set the object lifecycle to the expected duration of the workload run.
   For example, if the workload is expected to run for 10 days, set the object
   lifecycle to 10 days. That should cover restarts that occur during the entire
   run. Use `age` for the lifecycle condition and `Delete` for the action. See
   [Object Lifecycle Management](https://cloud.google.com/storage/docs/lifecycle)
   for details. If the object lifecycle is not set, the cache will continue to
   grow since there is no eviction mechanism implemented.

*  All encryption policies are supported.

It is **recommended** to use
[Google Cloud Storage Fuse](https://cloud.google.com/storage/docs/cloud-storage-fuse)
to mount the GCS bucket as a local directory. This is because when running JAX
in a multi-node setup, multiple nodes might try to write to the cache
simultaneously, leading to GCS rate-limit errors. GCSFuse handles this by
ensuring that only one process can write to a file at a time, preventing these
errors.

To set up GCSFuse, follow instructions for
[GCE](https://cloud.google.com/storage/docs/cloud-storage-fuse/mount-bucket) or
[GKE](https://cloud.google.com/kubernetes-engine/docs/how-to/cloud-storage-fuse-csi-driver-setup).
For better performance, enable file caching
([GCE](https://cloud.google.com/storage/docs/cloud-storage-fuse/file-caching) and
[GKE](https://cloud.google.com/kubernetes-engine/docs/how-to/cloud-storage-fuse-csi-driver-perf#enable-and-use-file-caching)).

Once GCSFuse is configured, set the JAX cache directory to the GCSFuse mount
point:

```python
# Example assuming the GCS bucket is mounted at /gcs/my-bucket
jax.config.update("jax_compilation_cache_dir", "/gcs/my-bucket/jax-cache")
```

**Direct GCS access:**

If you choose not to use GCSFuse, you can point the cache directly to a GCS
bucket.

Assuming that `gs://jax-cache` is the GCS bucket, set cache location as
follows:

```python
jax.config.update("jax_compilation_cache_dir", "gs://jax-cache")
```

## How it works

The cache key is a hash of the following, which together identify a
compiled program:

*  The computation performed by the function, as captured by the
   non-optimized HLO of the JAX function being hashed

*  The jaxlib version, and the version of the backend (platform)

*  Relevant XLA compilation flags, and the compile options

*  The device configuration, captured by the number of devices and their
   topology. Currently for GPUs, the topology only contains a string
   representation of the GPU name

*  Compression algorithm used to compress the compiled executable

*  A string produced by `jax._src.cache_key.custom_hook()`, which returns an
   empty string by default. This hook is private, not a supported extension
   point, and may change without notice.

## Caching on multiple nodes

The first time a program is run (when the persistent cache is cold or empty),
all processes compile, but only the process with rank 0 in the global
communication group writes to the persistent cache. In subsequent runs, all
processes attempt to read from the persistent cache, so it's important for the
persistent cache to be in a shared file system (e.g., NFS) or remote storage
(e.g., GCS). If the persistent cache is local to rank 0, then all processes
except rank 0 will compile again in subsequent runs, because they miss the
cache.

### Pre-compiling multi-node programs on single node

JAX can populate the compilation cache with compiled programs for multiple nodes
on a single node. Preparing the cache on a single node helps to decrease the costly
compilation time on a cluster. To compile and run multi-node programs on a single
node, users can create fake remote devices using
the `jax_mock_gpu_topology` configuration option.

For instance, the snippet below instructs JAX to mock a cluster with four
nodes, each node running eight processes with each process attached to one GPU.

```python
jax.config.update("jax_mock_gpu_topology", "4x8x1")
```

After populating the cache with this config, users can run the program
without recompilation on four nodes, eight processes per node,
one GPU per process.

Important notes:

* The process running the mocked program must have the same number of GPUs
  and the same GPU model as the nodes that would use the cache. For instance,
  a mocked topology `8x4x2` must run in a process with two GPUs.

* When running programs with mocked topology, the results of communications
  with other nodes are undefined, so the outputs of JAX programs running
  in mocked environments will likely be incorrect.


## Logging cache activity

When debugging, it can be helpful to examine exactly what the persistent
compilation cache is doing. Here are a few suggestions for how to begin.

You can enable logging in the relevant source files by placing

```python
import os
os.environ["JAX_DEBUG_LOG_MODULES"] = "jax._src.compiler,jax._src.lru_cache"
```

at the top of the script. Alternatively, you can change JAX's global logging
level with

```python
import os
os.environ["JAX_LOGGING_LEVEL"] = "DEBUG"
# or locally with
jax.config.update("jax_logging_level", "DEBUG")
```

### Examining cache misses

To understand why there are cache misses, JAX includes a configuration flag
that logs cache misses, with explanations for misses in the tracing cache.
With the persistent cache enabled, it also logs each persistent compilation
cache miss, and why a compiled program wasn't written to the cache (for
example, because it compiled faster than the minimum compile time). Enable it
with:

```python
jax.config.update("jax_explain_cache_misses", True)
```

## Pitfalls

* Currently the persistent cache doesn't work with a function that uses
  primitives that implement their own `custom_partitioning`.
  - The HLO of the function contains a pointer to the `custom_partitioning`
    callback, which leads to different cache keys for the same computation
    across runs.
  - In this situation, caching still proceeds, but a different key is
    produced every time, making the cache ineffective.

### Working around `custom_partitioning`

As mentioned, the compilation cache doesn't work with a function that is
composed of primitives that implement `custom_partitioning`. However, you can
use `shard_map` to circumvent `custom_partitioning` for those primitives, and
make the compilation cache work as expected.

Let's pretend we have a function `F` that implements a layernorm followed by a matrix multiplication using a primitive `LayerNorm` that implements `custom_partitioning`:

```python
import jax

def F(x1, x2, gamma, beta):
   ln_out = LayerNorm(x1, gamma, beta)
   return ln_out @ x2
```
If we were to compile this function without `shard_map`, the cache key for
`layernorm_matmul_without_shard_map` would be different every time we ran
the same code:

```python
layernorm_matmul_without_shard_map = jax.jit(F, in_shardings=(...), out_shardings=(...))(x1, x2, gamma, beta)
```

However, if we wrap the layernorm primitive in `shard_map` and define a
function `G` that performs the same computation, the cache key for
`layernorm_matmul_with_shard_map` will be the same every time, even though
`LayerNorm` implements `custom_partitioning`:

```python
import jax

def G(x1, x2, gamma, beta, mesh, ispecs, ospecs):
   ln_out = jax.shard_map(LayerNorm, mesh=mesh, in_specs=ispecs, out_specs=ospecs, check_vma=False)(x1, gamma, beta)
   return ln_out @ x2

ispecs = jax.sharding.PartitionSpec(...)
ospecs = jax.sharding.PartitionSpec(...)
mesh = jax.sharding.Mesh(...)
layernorm_matmul_with_shard_map = jax.jit(G, static_argnames=['mesh', 'ispecs', 'ospecs'])(x1, x2, gamma, beta, mesh, ispecs, ospecs)
```

Note that for this workaround, the primitive that implements
`custom_partitioning` must itself be wrapped in `shard_map`. It is
insufficient to wrap the outer function `F` in `shard_map`.