今日已更新 305 条资讯 | 累计 19865 条内容
关于我们

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

Shixin Zhang 2026年07月15日 11:41 0 次阅读 来源:Dev.to

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, expecta

本文内容来源于互联网,版权归原作者所有
查看原文