Large-Scale TensorCircuit Contractions on GPUs: Disabling XLA GPU Autotuning

When running large-scale tensor-network contractions with TensorCircuit-NG and the JAX GPU backend, the following runtime configuration is worth testing:

XLA_PYTHON_CLIENT_PREALLOCATE=false XLA_FLAGS=--xla_gpu_autotune_level=0 python your_script.py

Its main benefit is not speed, but lower persistent GPU memory usage from XLA GPU autotuning, which makes memory behavior during compilation and on the first visible GPU more predictable. In the TensorCircuit contraction workloads we tested, disabling autotuning also slightly improved steady-state runtime, but the memory savings were the more important result.

Key takeaway

XLA GPU autotuning evaluates alternative algorithms or workspace configurations for certain GPU kernels and custom calls, then selects an implementation. This can be valuable for convolutions, large GEMMs, and deep-learning workloads with fixed shapes. For large TensorCircuit contractions, however, the contraction path is already determined by OMECO or cotengra, leaving relatively little optimization freedom for autotuning while still potentially incurring substantial persistent memory overhead during compilation and tuning.

For TensorCircuit contractions, run this A/B test by default:

# Baseline
XLA_PYTHON_CLIENT_PREALLOCATE=false python your_script.py

# Test configuration
XLA_PYTHON_CLIENT_PREALLOCATE=false XLA_FLAGS=--xla_gpu_autotune_level=0 python your_script.py

Both environment variables must be set before the Python process starts and before JAX is imported. This recommendation primarily concerns GPUs; CPU backends do not exhibit the same GPU-kernel autotuning behavior.

Representative results

All results below use a fixed contraction path so that path-search randomness does not affect the comparison.

Workload	Autotuning	Post-compile memory	Peak memory	Steady-state runtime
100 qubits × 24 layers, amplitude	Default	8.7 GiB	8.7 GiB	0.37 s
100 qubits × 24 layers, amplitude	autotune=0	0.5 GiB	4.6 GiB	0.40 s
28 qubits × 12 layers, expectation gradient	Default	4.7 GiB	21.4 GiB	0.18 s
28 qubits × 12 layers, expectation gradient	autotune=0	0.5 GiB	17.2 GiB	0.15 s
28 qubits × 13 layers, expectation gradient	Default	10.9 GiB	43.8 GiB	0.46 s
28 qubits × 13 layers, expectation gradient	autotune=0	0.5 GiB	33.5 GiB	0.42 s
28 qubits × 14 layers, expectation gradient	Default	17.0 GiB	64.5 GiB	0.58 s
28 qubits × 14 layers, expectation gradient	autotune=0	0.5 GiB	64.5 GiB	0.58 s

The first three workload groups show the most common benefit: disabling autotuning reduces compilation-stage memory and the final peak. The 28 × 14 case appears different: post-compile memory falls from 17.0 GiB to 0.5 GiB, but the final peak is unchanged. This behavior is related to the JAX GPU allocator.

Why the final peak is unchanged for the 28 × 14 case

Even with XLA_PYTHON_CLIENT_PREALLOCATE=false, the default JAX GPU allocator tends to retain and reuse GPU memory that it has already allocated. The memory reported by nvidia-smi is therefore the amount held by the process, not the total size of tensors that are currently live.

In the 28 × 14 example, default autotuning did consume approximately 17 GiB of additional memory during compilation. During the first real backward contraction, however, the runtime buffers themselves also required a large allocation. The default allocator could reuse blocks allocated earlier, so the final peak was not simply the sum of runtime memory and autotuning memory.

Using XLA_PYTHON_CLIENT_ALLOCATOR=platform as a diagnostic exposes this difference:

Workload	Allocator	Autotuning	Post-compile memory	Peak memory
28 qubits × 14 layers, expectation gradient	Default	Default	17.0 GiB	64.5 GiB
28 qubits × 14 layers, expectation gradient	Default	autotune=0	0.5 GiB	64.5 GiB
28 qubits × 14 layers, expectation gradient	platform	Default	17.0 GiB	82.8 GiB
28 qubits × 14 layers, expectation gradient	platform	autotune=0	0.5 GiB	66.3 GiB

The platform allocator retains less reusable GPU memory, which makes it useful for diagnosing where allocations originate. Because it also reuses fewer large blocks allocated earlier, however, the runtime may request additional memory and produce a higher peak. Prefer the default allocator for normal execution; use platform mainly to confirm whether autotuning introduces extra memory usage.

Interpreting peak memory

Here, peak memory is the maximum per-process GPU memory observed by polling nvidia-smi, with snapshots recorded at stages such as after_compile and after_first_run. XLA's compiled.memory_analysis().peak_memory_in_bytes measures something different: it more closely reflects the computation graph's buffer assignment and is usually lower than the actual process memory reported by nvidia-smi.

Why the first GPU may use more memory in multi-GPU processes

In a JAX/XLA process that can see multiple GPUs, backend initialization, the default device, compilation services, autotuning caches, executable caches, or allocator state may be placed preferentially on the first GPU in the visible-device list. Consequently, the first GPU can hold an extra block of memory. Disabling xla_gpu_autotune_level often reduces this first-device memory tax.

For multiple independent single-GPU jobs, use CUDA_VISIBLE_DEVICES so that each process sees only its assigned GPU:

CUDA_VISIBLE_DEVICES=1 XLA_PYTHON_CLIENT_PREALLOCATE=false XLA_FLAGS=--xla_gpu_autotune_level=0 python your_script.py

For genuine multi-GPU parallel or distributed jobs, do not hide GPUs that must participate in the computation merely to avoid the first-device memory tax. Preserve the correct device set, then measure how disabling autotuning changes memory usage on each GPU.

Practical recommendations

For large TensorCircuit-NG contractions, use the following procedure:

1. Fix the contraction path instead of running a new stochastic path search for every benchmark.
2. Set XLA_PYTHON_CLIENT_PREALLOCATE=false to prevent JAX from reserving a large block of GPU memory at startup.
3. A/B test default autotuning against XLA_FLAGS=--xla_gpu_autotune_level=0.
4. Record compile time, first-run time, steady-state runtime, post-compile memory, and post-first-run memory separately.
5. Even if disabling autotuning does not improve steady-state runtime, prefer autotune=0, especially near the OOM limit or when scheduling multiple GPUs.

In short, disabling XLA GPU autotuning does not change the contraction path; it reduces an additional source of memory usage in the GPU compilation and execution layers. For large TensorCircuit contractions, this is usually a low-risk configuration well worth testing.