TensorFlow
XLA
CUDA
Configuration
GPU

How to let TensorFlow XLA know the CUDA path

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

When TensorFlow XLA cannot find CUDA tools or GPU support files, the fix is usually not "install CUDA again." The more specific issue is that XLA does not know where to find the CUDA toolkit data directory, especially ptxas and related GPU compilation assets.

The most common way to point XLA at CUDA is to set XLA_FLAGS with --xla_gpu_cuda_data_dir=/path/to/cuda. On many Linux systems that path is something like /usr/local/cuda.

What XLA Is Looking For

XLA generates optimized GPU code paths, and for some workflows it needs access to the local CUDA toolkit layout. Depending on the environment, the problem may involve:

  • CUDA toolkit binaries
  • 'ptxas'
  • libdevice files
  • mismatched driver or toolkit installation paths

That is why GPU visibility alone is not always enough. TensorFlow may see the GPU, while XLA still fails during compilation.

The Usual Environment Variable Fix

On Linux or macOS shells:

bash
export XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/local/cuda
python train.py

If your CUDA toolkit is installed elsewhere, use that path instead:

bash
export XLA_FLAGS=--xla_gpu_cuda_data_dir=/opt/cuda-12.2
python train.py

On Windows PowerShell:

powershell
$env:XLA_FLAGS="--xla_gpu_cuda_data_dir=C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v12.2"
python train.py

The key is that XLA needs the toolkit directory, not just the driver.

Also Set the Normal CUDA Variables

It is still worth making the standard environment consistent:

bash
export CUDA_HOME=/usr/local/cuda
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

These variables help the broader runtime environment, while XLA_FLAGS is the XLA-specific hint.

A Small Verification Script

The script below does not prove XLA is perfectly configured, but it does confirm that TensorFlow sees GPU devices and can execute a simple operation:

python
1import tensorflow as tf
2
3print("GPUs:", tf.config.list_physical_devices("GPU"))
4
5@tf.function(jit_compile=True)
6def add_one(x):
7    return x + 1
8
9result = add_one(tf.constant([1.0, 2.0, 3.0]))
10print(result)

If this still fails with XLA CUDA-path complaints, the next step is usually checking the toolkit location and version compatibility rather than changing model code.

Version Compatibility Still Matters

Pointing XLA to a directory does not fix version mismatches. You still need:

  • a TensorFlow build compatible with the installed CUDA stack
  • a compatible NVIDIA driver
  • a toolkit layout that actually contains the expected files

If those are inconsistent, XLA may fail even with the right path.

Practical Troubleshooting Order

Use this order:

  1. confirm TensorFlow sees the GPU
  2. confirm the CUDA toolkit path exists
  3. set XLA_FLAGS=--xla_gpu_cuda_data_dir=...
  4. verify normal PATH and library environment
  5. check version compatibility if errors remain

That sequence prevents random configuration changes.

When You May Not Need a Manual Path

Some packaged TensorFlow environments already bundle or expect the required CUDA components in known locations. In those cases, manually setting XLA_FLAGS is unnecessary. You only need the override when XLA cannot resolve the toolkit location on its own or when multiple CUDA installations make the default resolution ambiguous.

Common Pitfalls

  • Setting only CUDA_HOME and assuming XLA will infer everything automatically.
  • Pointing --xla_gpu_cuda_data_dir at the wrong directory level.
  • Confusing the driver installation with the CUDA toolkit installation.
  • Ignoring TensorFlow-CUDA version compatibility and blaming the path alone.
  • Testing GPU visibility but never testing an actual XLA-compiled function.

Summary

  • XLA usually learns the CUDA toolkit location through XLA_FLAGS=--xla_gpu_cuda_data_dir=....
  • Standard CUDA variables such as CUDA_HOME, PATH, and library paths should also be consistent.
  • The relevant path is the CUDA toolkit directory, not just the GPU driver.
  • A visible GPU does not guarantee XLA can compile GPU code.
  • If the path is correct and failures continue, check version compatibility next.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.