MLX 环境变量完全指南编译、数值精度、分布式与后端行为配置【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx导读本文以 MLX 官方文档 docs/src/usage/environment_variables.rst 为骨架系统梳理 MLXApple silicon 上的数组框架提供的全部环境变量它们分别控制编译开关、float32 矩阵乘精度、Ring/JACCL/NCCL/MPI 分布式执行、Metal 同步路径以及 CUDA 后端的高级缓存与调度行为。读完本文你将掌握每个变量的语义、默认值、生效时机并能在启动进程前正确设置它们从而完成精度控制、分布式调试、性能实验等实战任务。一、总览与通用约定MLX 通过环境变量配置编译、数值精度、后端行为与分布式执行。它们必须在进程启动之前设置很多变量只在对应子系统首次初始化时被读取进程运行中再修改通常不会生效见 mlx/utils.h 中大量以static局部变量缓存读值结果的设计。布尔类变量遵循统一约定0表示关闭任意非零整数表示开启除非另有说明。一个值得注意的特例是存在即启用型变量MLX_DISABLE_COMPILE、MLX_RING_VERBOSE、MLX_JACCL_RING这类变量只要被设置哪怕设为0即为启用。原因在于源码直接以std::getenv(...)的返回值是否为空来判断而不是解析数值——例如 mlx/compile.cpp 中就是if (std::getenv(MLX_DISABLE_COMPILE))。二、通用开关编译与数值精度2.1 MLX_DISABLE_COMPILE全局禁用编译MLX_DISABLE_COMPILE用于全局禁用编译compilation。它由存在即启用因此设置MLX_DISABLE_COMPILE0同样会禁用编译——想恢复编译不能靠赋0而应彻底不设置该变量。调用 mlx.core.enable_compile 可以覆盖优先于该环境变量。在 docs/src/usage/compile.rst 中MLX_DISABLE_COMPILE也被用作运行时的兜底开关当mlx.compile编译出的图因某些原因不可用时可通过该变量强制回退到即时执行模式。典型用法MLX_DISABLE_COMPILE1 python my_script.py2.2 MLX_ENABLE_TF32float32 矩阵乘的精度权衡MLX_ENABLE_TF32控制在受支持的硬件上是否允许矩阵乘法matmul系列运算使用降低精度的 float32TF32。默认值为1开启设置为0则这些运算保持完整的float32精度。其读取逻辑位于 mlx/utils.hinline bool enable_tf32() { static bool enable_tf32_ get_var(MLX_ENABLE_TF32, 1); return enable_tf32_; }在精度敏感场景如数值验证、单元测试中官方测试就显式关闭了 TF32——见 python/tests/run.py 的os.environ[MLX_ENABLE_TF32] 0。实战关闭方式摘自 docs/src/usage/precision.rstMLX_ENABLE_TF320 python my_script.py三、分布式执行相关变量分布式变量的设置取决于所选后端。官方文档 docs/src/usage/distributed.rst 详细描述了各后端的变量格式且mlx.launch会在启动子进程时自动注入这些变量实现见 python/mlx/_distributed_utils/launch.py。下面按后端归类说明。3.1 Ring 后端Apple silicon 上的多机/多卡Ring 后端需要两个变量MLX_RANK当前进程的零基 rank0-basedRing、JACCL、NCCL 三个后端都会用到。MLX_HOSTFILERing 后端使用的JSON 主机文件路径描述参与节点与拓扑。MLX_RING_VERBOSE为 Ring 后端启用详细日志存在即启用。Ring 后端的初始化代码mlx/distributed/ring/ring.cpp会在缺少 rank 或 hostfile 时直接报错const char* hostfile std::getenv(MLX_HOSTFILE); const char* rank_str std::getenv(MLX_RANK); const char* ring_verbose std::getenv(MLX_RING_VERBOSE); // 缺少时抛出[ring] You need to provide via environment variables both // a rank (MLX_RANK) and a hostfile (MLX_HOSTFILE) ...mlx.launch对应注入逻辑见 python/mlx/_distributed_utils/launch.pyfiles {MLX_HOSTFILE: hostfile} # ... if args.verbose: env.append(MLX_RING_VERBOSE1)3.2 JACCL 后端InfiniBand 集群JACCL 后端的变量都支持JACCL_*优先别名即设置了JACCL_*时优先使用否则回退到MLX_*。源码见 mlx/distributed/jaccl/lib/jaccl/jaccl.cppconst char* dev_file getenv(JACCL_IBV_DEVICES, MLX_IBV_DEVICES); getenv(JACCL_COORDINATOR, MLX_JACCL_COORDINATOR); const char* rank_str getenv(JACCL_RANK, MLX_RANK); const char* ring getenv(JACCL_RING, MLX_JACCL_RING);MLX_RANK/JACCL_RANK进程零基 rankJACCL 场景下 JACCL_RANK 优先级更高。MLX_IBV_DEVICES/JACCL_IBV_DEVICES描述设备连通性的JSON 文件路径。MLX_JACCL_COORDINATOR/JACCL_COORDINATOR协调者coordinator地址格式为IP:port用于建立 JACCL 连接。MLX_JACCL_RING/JACCL_RING存在时优先使用ring 拓扑存在即启用。mlx.launch的注入逻辑见 python/mlx/_distributed_utils/launch.pyenv.append(fMLX_JACCL_COORDINATOR{coordinator}:{args.starting_port}) if args.ring: env.append(MLX_JACCL_RING1) files {MLX_IBV_DEVICES: json.dumps([h.rdma for h in hosts])}3.3 NCCL 后端NVIDIA GPU / CUDANCCL 后端在MLX_RANK之外还要求MLX_WORLD_SIZENCCL 组内的进程总数。MLX_NCCL_TIMEOUT建立 NCCL bootstrap 连接的超时时间毫秒默认300000即 5 分钟。NCCL_HOST_IP与NCCL_PORTNCCL bootstrap 连接所需由mlx.launch注入见 python/mlx/_distributed_utils/launch.py。NCCL_DEBUGINFO在 MLX 建立 bootstrap 连接期间输出额外日志便于排查。CUDA_VISIBLE_DEVICES由 CUDA runtime 处理用于为每个进程选择本地 CUDA 设备——mlx.launch中正是用它做设备绑定见 python/mlx/_distributed_utils/launch.py 的CUDA_VISIBLE_DEVICES{rank % args.repeat_hosts}。NCCL 后端的环境变量解析位于 mlx/distributed/nccl/nccl.cppstd::string host detail::get_env_var_or_throw(NCCL_HOST_IP, strict); std::string port detail::get_env_var_or_throw(NCCL_PORT, strict); std::string rank_str detail::get_env_var_or_throw(MLX_RANK, strict); std::string world_size detail::get_env_var_or_throw(MLX_WORLD_SIZE, strict);3.4 MPI 后端MLX_MPI_LIBNAME覆盖 MPI动态库名称。默认值在 macOS 上为libmpi.dylib其他平台为libmpi.soMPI 后端通过dlopen动态加载该库见 mlx/distributed/mpi/mpi.cppstatic const char* get_libmpi_name() { const char* libname std::getenv(MLX_MPI_LIBNAME); if (libname ! nullptr) { return libname; } #if defined(__APPLE__) return libmpi.dylib; #else return libmpi.so; #endif }四、Metal 后端MLX_METAL_FAST_SYNCHMLX_METAL_FAST_SYNCH用于启用更快的 Metal CPU/GPU 同步路径。默认值为0关闭。该特性要求Metal 3.2 或更高版本macOS 15或 iOS 18。其默认读值与注释同样出现在 mlx/utils.h 与 mlx/fence.h 中。适用场景对 CPU/GPU 同步开销敏感的高频小算子调度若运行环境不满足 Metal 3.2 的系统要求则应保持默认关闭。五、高级调优变量Advanced tuning以下变量用于调节 MLX 的实现细节主要面向开发、诊断与性能实验。默认值已针对当前硬件自动选择对绝大多数用户是合适的这些变量的行为可能随实现演进而变化。5.1 MLX_BFS_MAX_WIDTH求值磁带构造宽度设置构造求值磁带evaluation tape时广度优先搜索BFS的宽度上限默认20。实现在 mlx/utils.hinline int bfs_max_width() { static int bfs_max_width_ get_var(MLX_BFS_MAX_WIDTH, 20); return bfs_max_width_; }该值影响 MLX 懒求值lazy evaluation图展开时的并发分支宽度通常在分析极端宽/深的计算图时才需要调整。5.2 MLX_MAX_OPS_PER_BUFFER 与 MLX_MAX_MB_PER_BUFFERMLX_MAX_OPS_PER_BUFFER覆盖单个 Metal command buffer 或 CUDA graph 中编码的最大操作数默认值取决于设备。MLX_MAX_MB_PER_BUFFER覆盖单个 Metal command buffer 或 CUDA graph 的近似内存上限MB默认值取决于设备。两者都通过get_var(name, default_value)读取默认值由调用方设备相关传入见 mlx/utils.h。CUDA 后端在 mlx/backend/cuda/device.cpp 的注释中明确说明这两个变量可用于调节 CUDA graph 的批处理规模。5.3 MLX_METAL_GPU_ARCH覆盖 GPU 架构字符串覆盖 MLX 内部报告的Metal GPU 架构字符串影响架构相关的 kernel 与调度选择但不会改变物理 GPU 的实际能力。强制一个与真实 GPU 不匹配的架构可能选中不兼容的 kernel 并产生错误结果因此仅在开发调试时使用。读取位置见 mlx/utils.hinline const std::string metal_gpu_arch() { static std::string gpu_arch_ get_var(MLX_METAL_GPU_ARCH, ); return gpu_arch_; }5.4 MLX_SDPA_BLOCKS缩放点积注意力归约块数覆盖 Metal 缩放点积注意力scaled dot-product attention, SDPAkernel 使用的归约块reduction blocks数量。正值会被向上取整为 32 的倍数。实现见 mlx/backend/metal/scaled_dot_product_attention.cppif (int blocks_env env::get_var(MLX_SDPA_BLOCKS, 0); blocks_env 0) { // 按 32 向上取整后使用 }该参数用于在长序列 attention 场景下调节 kernel 的并行粒度。5.5 源码中定义的其他诊断变量延伸阅读从 mlx/utils.h 可见同一套env::get_var机制还支持两个官方文档未单列的诊断变量可作为排查 Metal 显存驻留residency行为的补充MLX_RESIDENCY_SET_MAX_PCT每个 residency set 占设备推荐最大 working-set 的百分比默认5取值0或100时所有内容放入单一 set。MLX_RESIDENCY_DEBUG在每个 residency set 创建时输出日志默认0。六、CUDA 后端高级控制本节变量均为 CUDA 后端的进阶控制项。6.1 CUDA graph 捕获与回放MLX_USE_CUDA_GRAPHS启用 CUDA graph 捕获与回放默认1。读取位置见 mlx/backend/cuda/device.cppstatic bool use_graphs env::get_var(MLX_USE_CUDA_GRAPHS, true);。MLX_SAVE_CUDA_GRAPHS_DOT_FILE将捕获的 CUDA graph 以编号的 DOT 文件写出其值作为文件名前缀未设置或为空时禁用输出见 mlx/backend/cuda/device.cpp。DOT 文件可用 Graphviz 等工具可视化是分析 kernel 调度结构的利器。6.2 运行时编译缓存MLX_PTX_CACHE_DIR覆盖运行时编译 PTX 的缓存目录。默认目录为系统临时目录下的mlx/version/ptx。实现见 mlx/backend/cuda/jit_module.cppif (auto c std::getenv(MLX_PTX_CACHE_DIR); c)。在多用户共享机器上可借此将缓存定向到可写路径避免重复编译开销。6.3 各功能缓存容量Cache Capacity下表汇总 CUDA 后端各类缓存的容量环境变量及默认值实现位置分别对应 mlx/backend/cuda/conv.cpp、mlx/backend/cuda/fft.cu、mlx/backend/cuda/device.cpp、mlx/backend/cuda/scaled_dot_product_attention.cpp环境变量控制对象默认值MLX_CUDA_CONV_CACHE_SIZE卷积convolution缓存容量128MLX_CUDA_FFT_CACHE_SIZEFFT plan 缓存容量128MLX_CUDA_GRAPH_CACHE_SIZECUDA graph 缓存容量400MLX_CUDA_SDPA_CACHE_SIZESDPA 前向缓存容量256MLX_CUDA_SDPA_BACKWARD_CACHE_SIZESDPA 反向缓存容量646.4 其他 CUDA 控制项MLX_CUDA_USE_CUDNN_SDPA允许 CUDA 后端在输入与设备受支持时使用cuDNN 缩放点积注意力默认1见 mlx/backend/cuda/scaled_dot_product_attention.cpp。需要排查 SDPA 数值差异时可临时置0回退到自有实现。MLX_ENABLE_CACHE_THRASHING_CHECK检测重复的 CUDA 缓存未命中cache miss命中时抛出错误并建议增大缓存容量默认1见 mlx/backend/cuda/lru_cache.h。在特定访问模式下若遇到缓存抖动提示可按建议调大对应缓存容量。CUDA_HOME/CUDA_PATH当 Python 环境中找不到 CUDA 头文件时MLX 使用这两个变量定位 CUDA 头文件以支持运行时 kernel 编译见 mlx/backend/cuda/jit_module.cpp查找顺序为CUDA_HOME优先其次CUDA_PATH。七、实战建议与注意事项设置时机所有变量必须在进程启动前设置。由于读取结果多被static变量缓存见 mlx/utils.h进程内通过os.environ修改后通常不会生效。布尔语义绝大多数变量用0/非零区分开关但MLX_DISABLE_COMPILE、MLX_RING_VERBOSE、MLX_JACCL_RING是存在即启用赋0也无法关闭需彻底不设置。分布式调试优先使用mlx.launch启动分布式任务它会自动注入MLX_RANK、MLX_WORLD_SIZE、MLX_HOSTFILE、MLX_JACCL_*、NCCL_HOST_IP/NCCL_PORT等变量见 python/mlx/_distributed_utils/launch.py手动设置时务必保证各进程变量一致。精度控制数值敏感场景如单元测试、精度对比建议设置MLX_ENABLE_TF320与 python/tests/run.py 的做法保持一致。性能实验CUDA 缓存的容量类变量conv/FFT/graph/SDPA与MLX_MAX_OPS_PER_BUFFER、MLX_MAX_MB_PER_BUFFER适合做吞吐实验MLX_METAL_GPU_ARCH、MLX_SDPA_BLOCKS等属于高风险项仅限开发诊断勿在生产环境随意改动。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考