zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-15 04:51:13 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

[CUDA] Put version in ptx cache dir path (#2352)

Browse Source
This commit is contained in:
Cheng 2025-07-10 23:24:21 +09:00 committed by GitHub
parent 8fb3e7a26c
commit afb9817599
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mlx/backend/cuda/jit_module.cpp
View File

@ -2,6 +2,7 @@
#include "mlx/backend/cuda/jit_module.h" #include "mlx/backend/cuda/jit_module.h"
#include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/device.h"
#include "mlx/version.h"
#include "cuda_jit_sources.h" #include "cuda_jit_sources.h"
@ -53,10 +54,11 @@ const std::string& cuda_home() {
const std::filesystem::path& ptx_cache_dir() { const std::filesystem::path& ptx_cache_dir() {
static std::filesystem::path cache = []() -> std::filesystem::path { static std::filesystem::path cache = []() -> std::filesystem::path {
std::filesystem::path cache; std::filesystem::path cache;
if (auto c = std::getenv("MLX_PTX_CACHE"); c) { if (auto c = std::getenv("MLX_PTX_CACHE_DIR"); c) {
cache = c; cache = c;
} else { } else {
cache = std::filesystem::temp_directory_path() / "mlx" / "ptx"; cache =
std::filesystem::temp_directory_path() / "mlx" / version() / "ptx";
} }
if (!std::filesystem::exists(cache)) { if (!std::filesystem::exists(cache)) {
std::error_code error; std::error_code error;