From dbe6d451ca5a596a85ab518b6408c47533013054 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 14 Jun 2026 01:08:44 +0900 Subject: [PATCH 1/3] feat(cuda): resolve bundled CCCL via /proc/self/exe and MLXCEL_CCCL_DIR The CUDA NVRTC JIT (gather/indexing kernels) needs the CCCL headers at runtime. MLX resolved their bundled location from current_binary_dir(), which uses dladdr; a relative ./mlxcel launch yields a relative dli_fname, so the /../include/cccl lookup missed and the first kernel JIT aborted with "cannot open source file cuda/std/tuple" on any host without the build-machine MLX_CCCL_DIR path. Patch the bundled MLX jit_module.cpp so include_path_args(): - honors an MLXCEL_CCCL_DIR env override first (for embedders / flat layouts), and - resolves the executable directory from /proc/self/exe (canonical absolute) on Linux, falling back to current_binary_dir(), so resolution no longer depends on how the process was launched. Closes #265. Follows up on the CCCL bundling added in #262. --- docs/installation.md | 15 +- .../patches/mlx/backend/cuda/jit_module.cpp | 498 ++++++++++++++++++ 2 files changed, 506 insertions(+), 7 deletions(-) create mode 100644 src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp diff --git a/docs/installation.md b/docs/installation.md index 3a16db662..e0cd4758a 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -129,12 +129,11 @@ headers available on the deployment host, not only the runtime libraries: archives (both aarch64 and x86_64). Each unpacks to `bin/` + `include/cccl/`, the layout MLX's JIT looks for relative to the executable (`/../include/cccl`). Keep `mlxcel`/`mlxcel-server` under `bin/` and - the `include/cccl/` directory beside it; do not flatten them. Launch the - binary by its **absolute path** (`/path/to/bin/mlxcel`, as a service manager - or a parent process spawning a subprocess does). MLX resolves the bundled - header location from the executable's own path, and a relative launch - (`./mlxcel` from inside `bin/`) can defeat that resolution and fail the first - kernel JIT. A future `MLXCEL_CCCL_DIR` override will remove this constraint. + the `include/cccl/` directory beside it; do not flatten them. The runtime + resolves the bundled headers from the executable's canonical path + (`/proc/self/exe`), so any launch style works, including a relative + `./mlxcel`. Set `MLXCEL_CCCL_DIR` to point the JIT at the CCCL headers + explicitly, e.g. when embedding mlxcel and keeping a flat binary layout. - **CUDA toolkit headers** (`cuda_runtime.h` and friends) come from the host. Install the CUDA toolkit and set `CUDA_HOME` (or `CUDA_PATH`) if it is not at `/usr/local/cuda`. Without them the first NVRTC compile fails with @@ -150,8 +149,10 @@ sessions. | Variable | Description | Default | |----------|-------------|---------| -| `CUDA_HOME` | CUDA toolkit root used by the build script | `/usr/local/cuda` when present | +| `CUDA_HOME` | CUDA toolkit root, build-time and for runtime NVRTC headers | `/usr/local/cuda` when present | | `MLX_CUDA_ARCHITECTURES` | CUDA SM target list, build-time | auto-detect via `nvidia-smi`, then `90a` fallback | +| `MLXCEL_CCCL_DIR` | Override for the bundled CCCL (libcu++) header dir used by the CUDA NVRTC JIT | bundled `/../include/cccl`, then build-time fallback | +| `MLX_PTX_CACHE_DIR` | On-disk cache for JIT-compiled CUDA kernels | system temp dir | | `MLXCEL_DEVICE` | Runtime device hint (`gpu` or `cpu`) | auto | | `MLXCEL_WIRED_LIMIT` | Apple Silicon wired-memory ceiling, e.g. `64GB`; `0`/`none` disables it | `max` | | `LLAMA_ARG_*` | Environment-backed server options accepted by clap | unset | diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp b/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp new file mode 100644 index 000000000..5e7e90928 --- /dev/null +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp @@ -0,0 +1,498 @@ +// Copyright © 2025 Apple Inc. +// Modified by mlxcel: CCCL include resolution honors the MLXCEL_CCCL_DIR +// override and resolves the executable directory from /proc/self/exe, so the +// bundled headers are found regardless of how the process is launched (a +// relative ./mlxcel from inside bin/ otherwise yields a relative dli_fname from +// dladdr and the lookup misses). + +#include "mlx/backend/cuda/jit_module.h" +#include "mlx/backend/cuda/device.h" +#include "mlx/version.h" + +#include "cuda_jit_sources.h" + +#include +#include +#include +#include + +#include +#include + +namespace mlx::core::cu { + +namespace { + +#define CHECK_NVRTC_ERROR(cmd) check_nvrtc_error(#cmd, (cmd)) + +void check_nvrtc_error(const char* name, nvrtcResult err) { + if (err != NVRTC_SUCCESS) { + throw std::runtime_error( + fmt::format("{} failed: {}", name, nvrtcGetErrorString(err))); + } +} + +// Return the default path to CUDA toolkit. +const std::filesystem::path& default_cuda_toolkit_path() { +#if defined(_WIN32) + static auto cached_path = []() -> std::filesystem::path { + std::filesystem::path root( + LR"(C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA)"); + for (auto& file : std::filesystem::directory_iterator(root)) { + if (std::filesystem::exists(file.path() / "include" / "cuda.h")) { + return file.path(); + } + } + return {}; + }(); +#else + static std::filesystem::path cached_path = "/usr/local/cuda"; +#endif + return cached_path; +} + +// Return the --include-path args used for invoking NVRTC. +const std::vector& include_path_args() { + static std::vector cached_args = []() { + std::vector args; + // Resolve the directory of the running executable. MLX's current_binary_dir + // uses dladdr, whose dli_fname can be relative when the program is launched + // by a relative argv[0] (e.g. ./mlxcel from inside bin/), which breaks the + // bundled-header lookup below. Prefer the canonical absolute path from + // /proc/self/exe on Linux so resolution does not depend on the launch path. + std::filesystem::path root_dir; +#if defined(__linux__) + { + std::error_code exe_ec; + auto exe_path = std::filesystem::read_symlink("/proc/self/exe", exe_ec); + root_dir = exe_ec ? current_binary_dir() : exe_path.parent_path(); + } +#else + root_dir = current_binary_dir(); +#endif +#if !defined(_WIN32) + root_dir = root_dir.parent_path(); +#endif + // CCCL headers: an explicit MLXCEL_CCCL_DIR override wins (it lets an + // embedder such as Backend.AI:GO point at the headers regardless of the + // on-disk layout), then the bundled layout next to the executable + // (bin/ + include/cccl/), then MLX's compiled-in fallback. + std::filesystem::path path; + if (const char* cccl_override = std::getenv("MLXCEL_CCCL_DIR"); + cccl_override != nullptr && *cccl_override != '\0' && + std::filesystem::exists(cccl_override)) { + path = cccl_override; + } else { + path = root_dir / "include" / "cccl"; + } +#if defined(MLX_CCCL_DIR) + if (!std::filesystem::exists(path)) { + path = MLX_CCCL_DIR; + } +#endif + if (std::filesystem::exists(path)) { + args.push_back(fmt::format("--include-path={}", path.string())); + } + // Add path to CUDA runtime headers, try local-installed python package + // first and then system-installed headers. + path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include"; + if (!std::filesystem::exists(path)) { + const char* home = std::getenv("CUDA_HOME"); + if (!home) { + home = std::getenv("CUDA_PATH"); + } + path = home ? std::filesystem::path(home) : default_cuda_toolkit_path(); + if (!path.empty()) { + path = path / "include"; + } + if (path.empty() || !std::filesystem::exists(path)) { + throw std::runtime_error( + "Can not find locations of CUDA headers, please set environment " + "variable CUDA_HOME or CUDA_PATH."); + } + } + args.push_back(fmt::format("--include-path={}", path.string())); + return args; + }(); + return cached_args; +} + +// Get the cache directory for storing compiled results. +const std::filesystem::path& ptx_cache_dir() { + static std::filesystem::path cache = []() -> std::filesystem::path { + std::filesystem::path cache; + if (auto c = std::getenv("MLX_PTX_CACHE_DIR"); c) { + cache = c; + } else { + cache = + std::filesystem::temp_directory_path() / "mlx" / version() / "ptx"; + } + +#if defined(_WIN32) + // Add "\\?\" prefix to support long file path. + const wchar_t* long_path_prefix = L"\\\\?\\"; + if (cache.is_relative()) { + cache = std::filesystem::absolute(cache); + } + if (!cache.native().starts_with(long_path_prefix)) { + cache = long_path_prefix + cache.native(); + } +#endif + + if (!std::filesystem::exists(cache)) { + std::error_code error; + if (!std::filesystem::create_directories(cache, error)) { + return std::filesystem::path(); + } + } + return cache; + }(); + return cache; +} + +std::filesystem::path get_ptx_path( + const std::filesystem::path& cache_dir, + const std::string& module_name) { + constexpr int max_file_name_length = 245; + if (module_name.size() <= max_file_name_length) { + return cache_dir / (module_name + ".ptx"); + } + + auto ptx_path = cache_dir; + int offset = 0; + while (module_name.size() - offset > max_file_name_length) { + ptx_path /= module_name.substr(offset, max_file_name_length); + offset += max_file_name_length; + } + ptx_path /= module_name.substr(offset) + ".ptx"; + + return ptx_path; +} + +// Try to read the cached |ptx| and |ptx_kernels| from |cache_dir|. +bool read_cached_ptx( + const std::filesystem::path& cache_dir, + const std::string& module_name, + std::string& ptx, + std::vector>& ptx_kernels) { + if (cache_dir.empty()) { + return false; + } + + auto ptx_path = get_ptx_path(cache_dir, module_name); + std::error_code error; + auto ptx_size = std::filesystem::file_size(ptx_path, error); + if (error) { + return false; + } + std::ifstream ptx_file(ptx_path, std::ios::binary); + if (!ptx_file.good()) { + return false; + } + ptx.resize(ptx_size); + ptx_file.read(ptx.data(), ptx_size); + + std::ifstream txt_file(ptx_path.replace_extension(".txt"), std::ios::binary); + std::string line; + while (std::getline(txt_file, line)) { + auto tab = line.find('\t'); + if (tab != std::string::npos) { + ptx_kernels.emplace_back(line.substr(0, tab), line.substr(tab + 1)); + } + } + return true; +} + +// Write the |ptx| and |ptx_kernels| to |cache_dir| with |name|. +void write_cached_ptx( + const std::filesystem::path& cache_dir, + const std::string& module_name, + const std::string& ptx, + const std::vector>& ptx_kernels, + const std::string& source_code) { + if (cache_dir.empty()) { + return; + } + + auto ptx_path = get_ptx_path(cache_dir, module_name); + + // Ensure that the directory exists + auto parent = ptx_path.parent_path(); + if (parent != cache_dir) { + std::filesystem::create_directories(parent); + } + + // Write the compiled code and mangled names + std::ofstream ptx_file(ptx_path, std::ios::binary); + if (!ptx.empty()) { + ptx_file.write(&ptx.front(), ptx.size()); + } + std::ofstream txt_file(ptx_path.replace_extension(".txt"), std::ios::binary); + for (const auto& [name, mangled] : ptx_kernels) { + txt_file << name << "\t" << mangled << std::endl; + } + + // Write the generated code + std::ofstream source_file(ptx_path.replace_extension(".cu")); + source_file << source_code; +} + +// Return if |device|'s version is not newer than |major|.|minor| version. +inline bool version_lower_equal(Device& device, int major, int minor) { + if (device.compute_capability_major() < major) { + return true; + } else if (device.compute_capability_major() == major) { + return device.compute_capability_minor() <= minor; + } else { + return false; + } +} + +// Return whether NVRTC supports compiling to |device|'s SASS code. +bool compiler_supports_device_sass(Device& device) { + int nvrtc_major, nvrtc_minor; + CHECK_NVRTC_ERROR(nvrtcVersion(&nvrtc_major, &nvrtc_minor)); + if (nvrtc_major < 9) { + return false; + } else if (nvrtc_major == 9) { + return version_lower_equal(device, 7, 2); + } else if (nvrtc_major == 10) { + return version_lower_equal(device, 7, 5); + } else if (nvrtc_major == 11 && nvrtc_minor == 0) { + return version_lower_equal(device, 8, 0); + } else if (nvrtc_major == 11 && nvrtc_minor < 8) { + return version_lower_equal(device, 8, 6); + } else { + return true; + } +} + +#define INCLUDE_PREFIX "mlx/backend/cuda/device/" + +constexpr const char* g_include_names[] = { + INCLUDE_PREFIX "atomic_ops.cuh", + INCLUDE_PREFIX "binary_ops.cuh", + INCLUDE_PREFIX "cast_op.cuh", + INCLUDE_PREFIX "config.h", + INCLUDE_PREFIX "complex.cuh", + INCLUDE_PREFIX "fp16_math.cuh", + INCLUDE_PREFIX "hadamard.cuh", + INCLUDE_PREFIX "indexing.cuh", + INCLUDE_PREFIX "scatter_ops.cuh", + INCLUDE_PREFIX "unary_ops.cuh", + INCLUDE_PREFIX "ternary_ops.cuh", + INCLUDE_PREFIX "utils.cuh", +}; + +#undef INCLUDE_PREFIX + +constexpr const char* g_headers[] = { + jit_source_atomic_ops, + jit_source_binary_ops, + jit_source_cast_op, + jit_source_config, + jit_source_complex, + jit_source_fp16_math, + jit_source_hadamard, + jit_source_indexing, + jit_source_scatter_ops, + jit_source_unary_ops, + jit_source_ternary_ops, + jit_source_utils, +}; + +void compile( + Device& device, + const std::string& module_name, + const std::string& source, + const std::vector& kernel_names, + std::string& ptx, + std::vector>& ptx_kernels) { + // Create the program + nvrtcProgram prog; + CHECK_NVRTC_ERROR(nvrtcCreateProgram( + &prog, + source.c_str(), + (module_name + ".cu").c_str(), + std::size(g_headers), + g_headers, + g_include_names)); + std::unique_ptr prog_freer( + &prog, + [](nvrtcProgram* p) { CHECK_NVRTC_ERROR(nvrtcDestroyProgram(p)); }); + for (const auto& name : kernel_names) { + CHECK_NVRTC_ERROR(nvrtcAddNameExpression(prog, name.c_str())); + } + + // Compile program. + std::vector args; + bool use_sass = compiler_supports_device_sass(device); + auto cc = device.compute_capability_major(); + std::string arch_tag = (cc >= 9) ? "a" : ""; + std::string compute = fmt::format( + "--gpu-architecture={}_{}{}{}", + use_sass ? "sm" : "compute", + cc, + device.compute_capability_minor(), + arch_tag); + args.push_back(compute.c_str()); + for (const auto& include : include_path_args()) { + args.push_back(include.c_str()); + } + nvrtcResult compile_result = + nvrtcCompileProgram(prog, args.size(), args.data()); + if (compile_result != NVRTC_SUCCESS) { + size_t log_size; + CHECK_NVRTC_ERROR(nvrtcGetProgramLogSize(prog, &log_size)); + std::vector log(log_size + 1, 0); + CHECK_NVRTC_ERROR(nvrtcGetProgramLog(prog, log.data())); + throw std::runtime_error( + fmt::format("Failed to compile kernel: {}.", log.data())); + } + + // Get mangled names of kernel names. + for (const auto& name : kernel_names) { + const char* mangled; + CHECK_NVRTC_ERROR(nvrtcGetLoweredName(prog, name.c_str(), &mangled)); + ptx_kernels.emplace_back(name, mangled); + } + + // Get ptx data. + size_t ptx_size; + if (use_sass) { + CHECK_NVRTC_ERROR(nvrtcGetCUBINSize(prog, &ptx_size)); + } else { + CHECK_NVRTC_ERROR(nvrtcGetPTXSize(prog, &ptx_size)); + } + ptx.resize(ptx_size); + if (use_sass) { + CHECK_NVRTC_ERROR(nvrtcGetCUBIN(prog, ptx.data())); + } else { + CHECK_NVRTC_ERROR(nvrtcGetPTX(prog, ptx.data())); + } +} + +void load_module( + const std::string& module_name, + const std::string& ptx, + const std::vector>& ptx_kernels, + CUmodule& module_, + std::unordered_map>& + kernels) { + // Load module. + char jit_log[4089] = {}; + CUjit_option options[] = { + CU_JIT_ERROR_LOG_BUFFER, CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES}; + void* values[] = {jit_log, reinterpret_cast(std::size(jit_log) - 1)}; + CUresult jit_result = cuModuleLoadDataEx( + &module_, ptx.data(), std::size(options), options, values); + if (jit_result != CUDA_SUCCESS) { + throw std::runtime_error( + fmt::format( + "Failed to load compiled {} kernel: {}.", module_name, jit_log)); + } + + // Load kernels. + for (const auto& [name, mangled] : ptx_kernels) { + CUfunction kernel; + CHECK_CUDA_ERROR(cuModuleGetFunction(&kernel, module_, mangled.c_str())); + kernels[name] = std::make_tuple(kernel, false, 0); + } +} + +} // namespace + +JitModule::JitModule( + Device& device, + const std::string& module_name, + const KernelBuilder& builder, + bool use_disk_cache) { + // Will hold the actual device executable source code and kernel names + std::string ptx; + std::vector> ptx_kernels; + + // Try to load them from the file cache + if (!read_cached_ptx(ptx_cache_dir(), module_name, ptx, ptx_kernels)) { + auto [precompiled, source_code, kernel_names] = builder(); + + // Get the PTX or cubin + if (precompiled) { + ptx = std::move(source_code); + for (auto& name : kernel_names) { + ptx_kernels.emplace_back(name, name); + } + } else { + compile(device, module_name, source_code, kernel_names, ptx, ptx_kernels); + } + + // If requested save them in the file cache for the next launch + if (use_disk_cache) { + write_cached_ptx( + ptx_cache_dir(), module_name, ptx, ptx_kernels, source_code); + } + } + + // Load the module + load_module(module_name, ptx, ptx_kernels, module_, kernels_); +} + +JitModule::~JitModule() { + CHECK_CUDA_ERROR(cuModuleUnload(module_)); +} + +std::pair JitModule::get_kernel_and_dims( + const std::string& kernel_name, + std::function configure_kernel) { + auto it = kernels_.find(kernel_name); + if (it == kernels_.end()) { + throw std::runtime_error( + fmt::format("There is no kernel named {}.", kernel_name)); + } + + // If it is the first time we run this kernel then configure it. Do it only + // once! + auto kernel = std::get<0>(it->second); + if (!std::get<1>(it->second)) { + if (configure_kernel) { + configure_kernel(kernel); + } + std::get<1>(it->second) = true; + std::get<2>(it->second) = max_occupancy_block_dim(kernel); + } + + return {kernel, std::get<2>(it->second)}; +} + +CUfunction JitModule::get_kernel( + const std::string& kernel_name, + std::function configure_kernel) { + return get_kernel_and_dims(kernel_name, std::move(configure_kernel)).first; +} + +JitModule& get_jit_module( + const mlx::core::Device& device, + const std::string& name, + const KernelBuilder& builder, + bool use_disk_cache) { + // The cache are leak intentionally as user code may still be running JIT + // compiled code after main thread teardown. + static auto* cache = new std::unordered_map; + static auto* mtx = new std::shared_mutex; + + { + std::shared_lock rlock(*mtx); + if (auto it = cache->find(name); it != cache->end()) { + return it->second; + } + } + + std::unique_lock wlock(*mtx); + auto it = cache->find(name); + if (it == cache->end()) { + auto& d = cu::device(device); + it = cache->try_emplace(name, d, name, builder, use_disk_cache).first; + } + return it->second; +} + +} // namespace mlx::core::cu From 8be585e274ceda8bb6d3cd9bda68bba563ca62e6 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 14 Jun 2026 01:32:03 +0900 Subject: [PATCH 2/3] feat(cuda): notify once when CUDA kernels are JIT-compiled on first run The first CUDA generate on a cold PTX cache JIT-compiles gather/indexing kernels via NVRTC, which makes the first response noticeably slower with no indication. Emit a one-time stderr notice from the bundled MLX jit_module.cpp overlay at the actual NVRTC compile site, so both `mlxcel run` and `mlxcel-server` surface it on model load / first generate. Suppress with MLXCEL_QUIET_JIT. Document the knob. --- docs/installation.md | 1 + .../patches/mlx/backend/cuda/jit_module.cpp | 28 +++++++++++++++---- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/docs/installation.md b/docs/installation.md index e0cd4758a..2070da499 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -153,6 +153,7 @@ sessions. | `MLX_CUDA_ARCHITECTURES` | CUDA SM target list, build-time | auto-detect via `nvidia-smi`, then `90a` fallback | | `MLXCEL_CCCL_DIR` | Override for the bundled CCCL (libcu++) header dir used by the CUDA NVRTC JIT | bundled `/../include/cccl`, then build-time fallback | | `MLX_PTX_CACHE_DIR` | On-disk cache for JIT-compiled CUDA kernels | system temp dir | +| `MLXCEL_QUIET_JIT` | Suppress the one-time "compiling CUDA kernels" notice on a cold first run | unset (notice shown) | | `MLXCEL_DEVICE` | Runtime device hint (`gpu` or `cpu`) | auto | | `MLXCEL_WIRED_LIMIT` | Apple Silicon wired-memory ceiling, e.g. `64GB`; `0`/`none` disables it | `max` | | `LLAMA_ARG_*` | Environment-backed server options accepted by clap | unset | diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp b/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp index 5e7e90928..bfc79dfab 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp @@ -1,9 +1,12 @@ // Copyright © 2025 Apple Inc. -// Modified by mlxcel: CCCL include resolution honors the MLXCEL_CCCL_DIR -// override and resolves the executable directory from /proc/self/exe, so the -// bundled headers are found regardless of how the process is launched (a -// relative ./mlxcel from inside bin/ otherwise yields a relative dli_fname from -// dladdr and the lookup misses). +// Modified by mlxcel: +// - CCCL include resolution honors the MLXCEL_CCCL_DIR override and resolves +// the executable directory from /proc/self/exe, so the bundled headers are +// found regardless of how the process is launched (a relative ./mlxcel from +// inside bin/ otherwise yields a relative dli_fname from dladdr and the +// lookup misses). +// - A one-time notice when CUDA kernels are JIT-compiled on a cold cache, so +// the slower first response is not a mystery (suppress with MLXCEL_QUIET_JIT). #include "mlx/backend/cuda/jit_module.h" #include "mlx/backend/cuda/device.h" @@ -11,6 +14,7 @@ #include "cuda_jit_sources.h" +#include #include #include #include @@ -422,6 +426,20 @@ JitModule::JitModule( ptx_kernels.emplace_back(name, name); } } else { + // The first NVRTC compile on a cold cache makes the first response + // noticeably slower (gather/indexing kernels are JIT-built here). Emit a + // one-time notice so the latency is not a mystery; a function-local static + // initializes exactly once and is thread-safe. Suppress with + // MLXCEL_QUIET_JIT. + [[maybe_unused]] static const bool mlxcel_jit_notice = []() { + if (std::getenv("MLXCEL_QUIET_JIT") == nullptr) { + std::fputs( + "Compiling CUDA kernels (first run on this host; cached for later runs)...\n", + stderr); + std::fflush(stderr); + } + return true; + }(); compile(device, module_name, source_code, kernel_names, ptx, ptx_kernels); } From b0b78e78daa56978c35b4e4148caae690a4d1383 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 14 Jun 2026 01:49:19 +0900 Subject: [PATCH 3/3] feat(cuda): default the kernel JIT cache to a persistent MLX-pin-scoped dir MLX's default PTX cache lives in the system temp dir (e.g. /tmp/mlx/0.32.0/ptx), cleared on reboot, so the first-run NVRTC compile is paid again every boot. Default MLX_PTX_CACHE_DIR, when unset, to ${MLXCEL_CACHE_DIR:-~/.cache/mlxcel}/cuda-ptx/, so it is paid once per machine. Scope by the pinned MLX commit because read_cached_ptx is keyed only by kernel name and never validates the source, so entries must not survive an MLX bump. mlxcel-core build.rs exposes the pin via MLXCEL_MLX_COMMIT; both mlxcel and mlxcel-server call ensure_persistent_ptx_cache() at startup. The env-2021 crate keeps set_var safe. Users and embedders can still override MLX_PTX_CACHE_DIR. --- src/bin/mlx_server.rs | 4 ++++ src/lib/mlxcel-core/build.rs | 4 ++++ src/lib/mlxcel-core/src/lib.rs | 29 +++++++++++++++++++++++++++++ src/main.rs | 4 ++++ 4 files changed, 41 insertions(+) diff --git a/src/bin/mlx_server.rs b/src/bin/mlx_server.rs index 37402809f..918d84707 100644 --- a/src/bin/mlx_server.rs +++ b/src/bin/mlx_server.rs @@ -1080,6 +1080,10 @@ struct ServerArgs { async fn main() -> anyhow::Result<()> { let cli = Cli::parse(); + // Default the CUDA kernel JIT cache to a persistent, MLX-pin-scoped dir so + // the first-run kernel compilation is paid once per machine, not every boot. + mlxcel_core::ensure_persistent_ptx_cache(); + match cli.command { // Subcommand-driven dispatch. Currently only `download` // exists; future operational subcommands (e.g. cache inspection) can diff --git a/src/lib/mlxcel-core/build.rs b/src/lib/mlxcel-core/build.rs index 77dcd8ace..3c18d9022 100644 --- a/src/lib/mlxcel-core/build.rs +++ b/src/lib/mlxcel-core/build.rs @@ -21,6 +21,10 @@ use std::{env, path::PathBuf}; fn main() { let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); + // Expose the pinned MLX commit to the crate so the runtime can scope the + // persistent CUDA PTX cache directory by it (see ensure_persistent_ptx_cache). + println!("cargo:rustc-env=MLXCEL_MLX_COMMIT={MLX_EXPECTED_COMMIT}"); + // Build MLX using cmake let mlx_dst = build_mlx(); mark_mlx_cache_valid(&out_dir); diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 38d5d119c..77918691b 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -2315,6 +2315,35 @@ pub use lang_analyzer::{ // tokenizer language-analysis disk cache already uses. pub use lang_analyzer::cache_root; +/// Default the CUDA NVRTC PTX cache to a persistent, MLX-pin-scoped directory +/// under the mlxcel cache root, unless `MLX_PTX_CACHE_DIR` is already set. +/// +/// MLX's own default places the JIT cache in the system temp dir +/// (`$TMPDIR/mlx//ptx`), which is cleared on reboot, so the first-run +/// kernel compilation is paid again every boot. A persistent location pays it +/// once per machine. The directory is scoped by the pinned MLX commit because +/// the cache is keyed only by kernel name and is not validated against the +/// kernel source, so entries must not survive an MLX upgrade. No-op on non-CUDA +/// builds, when `MLX_PTX_CACHE_DIR` is already set, and when the cache root +/// cannot be resolved. Call once at startup before the first inference. +pub fn ensure_persistent_ptx_cache() { + if !cfg!(feature = "cuda") { + return; + } + if std::env::var_os("MLX_PTX_CACHE_DIR").is_some() { + return; + } + let Some(root) = cache_root() else { + return; + }; + let commit = env!("MLXCEL_MLX_COMMIT"); + let scope = &commit[..commit.len().min(12)]; + let dir = root.join("cuda-ptx").join(scope); + if std::fs::create_dir_all(&dir).is_ok() { + std::env::set_var("MLX_PTX_CACHE_DIR", &dir); + } +} + fn use_single_query_maskless_path() -> bool { static ENABLED: std::sync::OnceLock = std::sync::OnceLock::new(); *ENABLED.get_or_init(|| { diff --git a/src/main.rs b/src/main.rs index 00e9711b8..e14fff913 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1635,6 +1635,10 @@ impl Default for DiffusionServeOptions { fn main() -> anyhow::Result<()> { let cli = Cli::parse(); + // Default the CUDA kernel JIT cache to a persistent, MLX-pin-scoped dir so + // the first-run kernel compilation is paid once per machine, not every boot. + mlxcel_core::ensure_persistent_ptx_cache(); + match cli.command { Commands::Run(args) => commands::run_run(args), Commands::Generate(args) => commands::run_generate(args),