From 25e607cfdbf5336f10c9682d87d2ed7450a491a6 Mon Sep 17 00:00:00 2001
From: morluto <76467478+morluto@users.noreply.github.com>
Date: Wed, 5 Aug 2026 04:31:08 +0800
Subject: [PATCH] build: use the validated NVCC compiler

Signed-off-by: morluto <76467478+morluto@users.noreply.github.com>
---
 README.md |  6 ++++++
 setup.py  | 14 ++++++++++----
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 5f203a1..68d6bf3 100644
--- a/README.md
+++ b/README.md
@@ -53,6 +53,12 @@ To build for SM100 instead, set the `MOK_ARCH` environment variable:
 MOK_ARCH=SM100 pip install . --no-build-isolation
 ```
 
+To use a specific CUDA compiler, set `MOK_NVCC` (or `NVCC`). The selected compiler is used for both the CUDA version check and the build:
+
+```bash
+MOK_NVCC=/path/to/cuda-13.0/bin/nvcc pip install . --no-build-isolation
+```
+
 To verify the installation:
 
 ```bash
diff --git a/setup.py b/setup.py
index 82cca15..efe9a8f 100644
--- a/setup.py
+++ b/setup.py
@@ -65,8 +65,12 @@ def check_pytorch() -> None:
         warnings.warn(f"MoK should run with PyTorch CUDA {torch.version.cuda}, but it is built and tested for CUDA 13.0.", RuntimeWarning, stacklevel=2)
 
 
-def check_nvcc() -> None:
-    nvcc_name = os.environ.get("NVCC", "nvcc")
+def get_nvcc_name() -> str:
+    return os.environ.get("MOK_NVCC", os.environ.get("NVCC", "nvcc"))
+
+
+def check_nvcc() -> str:
+    nvcc_name = get_nvcc_name()
     nvcc = shutil.which(nvcc_name)
     if nvcc is None:
         raise RuntimeError(
@@ -93,6 +97,8 @@ def check_nvcc() -> None:
             "Install the matching CUDA toolkit or ensure PATH points to its bin directory."
         )
 
+    return nvcc
+
 
 def check_thunderkittens() -> None:
     thunderkittens_header = THUNDERKITTENS_ROOT / "include" / "kittens.cuh"
@@ -119,7 +125,7 @@ def build_extension(self, ext) -> None:
         check_make()
         check_python()
         check_pytorch()
-        check_nvcc()
+        nvcc = check_nvcc()
         check_thunderkittens()
 
         pytorch_includes = " ".join(f"-I{path}" for path in include_paths())
@@ -130,7 +136,7 @@ def build_extension(self, ext) -> None:
         subprocess.check_call([
             "make",
             f"SRC={REPO_ROOT / 'csrc' / 'bindings.cu'}",
-            f"NVCC={os.environ.get('MOK_NVCC', 'nvcc')}",
+            f"NVCC={nvcc}",
             f"ARCH={os.environ.get('MOK_ARCH', 'SM103')}",
             f"PYTHON={sys.executable}",
             f"THUNDERKITTENS_ROOT={THUNDERKITTENS_ROOT}",