.. _environment_variables: Environment Variables ===================== MLX uses environment variables to configure compilation, numerical precision, backend behavior, and distributed execution. Set them before starting the process. Many variables are read when the corresponding subsystem is first initialized and changing them later may have no effect. Boolean variables use ``0`` to disable and a nonzero integer to enable unless otherwise noted. General ------- .. envvar:: MLX_DISABLE_COMPILE Disable compilation globally. This variable is enabled by its presence, so setting it to ``0`` also disables compilation. Calling :func:`mlx.core.enable_compile` overrides it. .. envvar:: MLX_ENABLE_TF32 Allow reduced-precision ``float32`` matrix-multiplication family operations on supported hardware. The default is ``1``. Set it to ``0`` to keep these operations in full ``float32`` precision. See :doc:`precision`. Distributed ----------- The variables required to initialize a distributed process depend on the backend. :doc:`distributed` describes their formats and how ``mlx.launch`` sets them. .. envvar:: MLX_RANK The zero-based rank of the current process. This is used by the Ring, JACCL, and NCCL backends. ``JACCL_RANK`` is accepted as a higher-priority alias by JACCL. .. envvar:: MLX_HOSTFILE The path to the JSON host file used by the Ring backend. .. envvar:: MLX_RING_VERBOSE Enable verbose logging for the Ring backend. This variable is enabled by its presence. .. envvar:: MLX_IBV_DEVICES The path to the JSON device-connectivity file used by JACCL. ``JACCL_IBV_DEVICES`` is accepted as a higher-priority alias. .. envvar:: MLX_JACCL_COORDINATOR The coordinator address in ``IP:port`` form used to establish JACCL connections. ``JACCL_COORDINATOR`` is accepted as a higher-priority alias. .. envvar:: MLX_JACCL_RING Prefer a ring topology for JACCL. This variable is enabled by its presence. ``JACCL_RING`` is accepted as a higher-priority alias. .. envvar:: MLX_WORLD_SIZE The total number of processes in an NCCL group. .. envvar:: MLX_NCCL_TIMEOUT The timeout in milliseconds for establishing NCCL bootstrap connections. The default is ``300000``. .. envvar:: MLX_MPI_LIBNAME Override the MPI dynamic-library name. The default is ``libmpi.dylib`` on macOS and ``libmpi.so`` on other platforms. The NCCL backend also requires ``NCCL_HOST_IP`` and ``NCCL_PORT``. Setting ``NCCL_DEBUG=INFO`` enables additional logging while MLX establishes the bootstrap connection. NCCL itself recognizes additional `NCCL environment variables `_. ``CUDA_VISIBLE_DEVICES`` selects the local CUDA device for each process and is handled by the CUDA runtime. Metal ----- .. envvar:: MLX_METAL_FAST_SYNCH Enable the faster Metal CPU/GPU synchronization path. The default is ``0``. This requires Metal 3.2 or later (macOS 15 or later, or iOS 18 or later). Advanced tuning --------------- These variables tune MLX implementation details. They are primarily intended for development, diagnostics, and performance experiments. The defaults are selected automatically for the current hardware and are appropriate for most users. Their behavior may change as the implementation evolves. .. envvar:: MLX_BFS_MAX_WIDTH Set the breadth-first-search width limit used when constructing an evaluation tape. The default is ``20``. .. envvar:: MLX_MAX_OPS_PER_BUFFER Override the maximum number of operations encoded in one Metal command buffer or CUDA graph. The default depends on the device. .. envvar:: MLX_MAX_MB_PER_BUFFER Override the approximate memory limit, in megabytes, for one Metal command buffer or CUDA graph. The default depends on the device. .. envvar:: MLX_METAL_GPU_ARCH Override the Metal GPU architecture string reported to MLX. This affects architecture-specific kernel and scheduling choices, but does not change the capabilities of the physical GPU. Forcing an architecture that does not match the GPU can select incompatible kernels and produce incorrect results. .. envvar:: MLX_SDPA_BLOCKS Override the number of reduction blocks used by the Metal scaled dot-product attention kernel. Positive values are rounded up to a multiple of ``32``. CUDA ---- The MLX-prefixed variables in this section are advanced CUDA backend controls. .. envvar:: MLX_USE_CUDA_GRAPHS Enable CUDA graph capture and replay. The default is ``1``. .. envvar:: MLX_SAVE_CUDA_GRAPHS_DOT_FILE Use the specified value as the filename prefix when writing captured CUDA graphs to numbered DOT files. An unset or empty value disables the output. .. envvar:: MLX_PTX_CACHE_DIR Override the directory used to cache runtime-compiled PTX. By default MLX uses an ``mlx//ptx`` directory under the system temporary directory. .. envvar:: MLX_CUDA_USE_CUDNN_SDPA Allow the CUDA backend to use cuDNN scaled dot-product attention when the inputs and device are supported. The default is ``1``. .. envvar:: MLX_CUDA_CONV_CACHE_SIZE Set the CUDA convolution cache capacity. The default is ``128``. .. envvar:: MLX_CUDA_FFT_CACHE_SIZE Set the CUDA FFT plan cache capacity. The default is ``128``. .. envvar:: MLX_CUDA_GRAPH_CACHE_SIZE Set the CUDA graph cache capacity. The default is ``400``. .. envvar:: MLX_CUDA_SDPA_CACHE_SIZE Set the CUDA forward scaled dot-product attention cache capacity. The default is ``256``. .. envvar:: MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE Set the CUDA backward scaled dot-product attention cache capacity. The default is ``64``. .. envvar:: MLX_ENABLE_CACHE_THRASHING_CHECK Detect repeated CUDA cache misses and raise an error suggesting a larger cache capacity. The default is ``1``. MLX also uses ``CUDA_HOME`` or ``CUDA_PATH`` to locate CUDA headers for runtime kernel compilation when they cannot be found in the Python environment.