Vendor candle-kernels with a CUDA 12.2+ guard; CUDA build green again (#86)

.github/workflows/build-cuda-wheel.yml

@@ -49,6 +49,10 @@ on:
       - "torchnative/rust/torch_c/Cargo.toml"
       - "torchnative/rust/torch_c/build.rs"
       - "torchnative/rust/torch_c/src/device.rs"
+      - "torchnative/rust/torch_c/Cargo.lock"
+      - "torchnative/rust/vendor/candle-kernels/**"
+      - "torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch"
+      - "torchnative/rust/vendor/candle-core/**"
 
 permissions:
   contents: read

scripts/vendor/vendor_candle.sh

@@ -3,8 +3,9 @@
 # the committed copies are exactly those forks.
 #
 #   sh scripts/vendor/vendor_candle.sh           write  torchnative/rust/vendor/candle-core/
-#                                              and torchnative/rust/vendor/candle-metal-kernels/
-#   sh scripts/vendor/vendor_candle.sh --check   rebuild both in a temp dir and diff
+#                                              torchnative/rust/vendor/candle-metal-kernels/
+#                                              and torchnative/rust/vendor/candle-kernels/
+#   sh scripts/vendor/vendor_candle.sh --check   rebuild all three in a temp dir and diff
 #
 # TWO crates, because `DType::I8` is two halves that are useless apart. The
 # `candle-core` half can build an int8 buffer on Metal; the
@@ -29,6 +30,12 @@
 #   candle-metal-kernels 0.11.0, pinned by the sha256 the lock already
 #   recorded, plus torchnative/rust/vendor/int8-candle-metal-kernels-0.11.0.patch.
 #
+#   candle-kernels 0.11.0 (the CUDA kernels), pinned by the sha256 the lock
+#   recorded, plus torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch.
+#   This third fork has nothing to do with I8: the published compatibility.cuh
+#   redefines __hmax_nan/__hmin_nan that CUDA 12.2+ headers already define, so
+#   no CUDA 12.6 build compiles (issue #86). The patch guards that one block.
+#
 # The *published* crate rather than a clone of candle's repository, on purpose:
 # `cargo package` normalises `Cargo.toml`, so the published crate stands alone.
 # The repository's `candle-core/Cargo.toml` inherits `version.workspace = true`
@@ -60,14 +67,16 @@ esac
 # and `TORCHNATIVE_CANDLE_CRATE` stay scoped to candle-core, which is what
 # test_int8.py's wrong-crate refusal test drives; the second crate has its own
 # pair so a test can drive either without disturbing the other.
-crates="candle-core candle-metal-kernels"
+crates="candle-core candle-metal-kernels candle-kernels"
 
 sha256_for() {
     case "$1" in
     candle-core)
         echo 5ecb245093b0f791b89d3420c3df9c6d49c60ab63ba54db896bf8a3baf486706 ;;
     candle-metal-kernels)
         echo 242e83c6acf639bb273c929d73c67a882bb4dd08a140f121096e19ba2f213d3e ;;
+    candle-kernels)
+        echo 67450168a281bbb195a14cc85cf164c7a12023a54a03409a3324f476346e0789 ;;
     *) echo "sha256_for: unknown crate $1" >&2; exit 2 ;;
     esac
 }
@@ -76,6 +85,7 @@ patch_for() {
     case "$1" in
     candle-core)          echo "$repo/torchnative/rust/vendor/int8-candle-$version-cpu.patch" ;;
     candle-metal-kernels) echo "$repo/torchnative/rust/vendor/int8-candle-metal-kernels-$version.patch" ;;
+    candle-kernels)       echo "$repo/torchnative/rust/vendor/cuda12-candle-kernels-$version.patch" ;;
     *) echo "patch_for: unknown crate $1" >&2; exit 2 ;;
     esac
 }
@@ -84,13 +94,15 @@ dest_for() {
     case "$1" in
     candle-core)          echo "${TORCHNATIVE_CANDLE_DIR:-$repo/torchnative/rust/vendor/candle-core}" ;;
     candle-metal-kernels) echo "${TORCHNATIVE_CANDLE_KERNELS_DIR:-$repo/torchnative/rust/vendor/candle-metal-kernels}" ;;
+    candle-kernels)       echo "${TORCHNATIVE_CANDLE_CUDA_KERNELS_DIR:-$repo/torchnative/rust/vendor/candle-kernels}" ;;
     esac
 }
 
 override_for() {
     case "$1" in
     candle-core)          echo "${TORCHNATIVE_CANDLE_CRATE:-}" ;;
     candle-metal-kernels) echo "${TORCHNATIVE_CANDLE_KERNELS_CRATE:-}" ;;
+    candle-kernels)       echo "${TORCHNATIVE_CANDLE_CUDA_KERNELS_CRATE:-}" ;;
     esac
 }
 

tests/numerics/test_int8.py

@@ -41,6 +41,7 @@
 _SCRIPT = os.path.join(_REPO_ROOT, "scripts", "vendor", "vendor_candle.sh")
 _FORK = os.path.join(_REPO_ROOT, "torchnative", "rust", "vendor", "candle-core")
 _KERNELS_FORK = os.path.join(_REPO_ROOT, "torchnative", "rust", "vendor", "candle-metal-kernels")
+_CUDA_KERNELS_FORK = os.path.join(_REPO_ROOT, "torchnative", "rust", "vendor", "candle-kernels")
 
 
 # --------------------------------------------------------------------------
@@ -152,8 +153,8 @@ def test_the_committed_metal_kernels_fork_is_the_pinned_crate_plus_the_patch():
     assert proc.returncode == 0, "exit %d\n%s\n%s" % (proc.returncode, proc.stdout, proc.stderr)
     assert ("sha256 242e83c6acf639bb273c929d73c67a882bb4dd08a140f121096e19ba2f213d3e verified"
             in proc.stdout), proc.stdout
-    assert proc.stdout.count("byte for byte") == 2, (
-        "--check reported %d trees, not 2. Both forks have to be covered, or "
+    assert proc.stdout.count("byte for byte") == 3, (
+        "--check reported %d trees, not 3. Every fork has to be covered, or "
         "one of them can drift silently.\n%s"
         % (proc.stdout.count("byte for byte"), proc.stdout))
 
@@ -196,6 +197,64 @@ def test_the_metal_kernels_script_refuses_a_crate_that_is_not_the_pinned_one():
         shutil.rmtree(tmp)
 
 
+# --------------------------------------------------------------------------
+# the THIRD fork -- `candle-kernels` (CUDA), issue #86
+#
+# Not an I8 half: the published `compatibility.cuh` redefines `__hmax_nan` and
+# `__hmin_nan`, which CUDA 12.2+ headers define for every architecture, so no
+# CUDA 12.6 build compiled. Same discipline as the two forks above. This Mac
+# has no nvcc, so these tests can only prove the fork is wired in and carries
+# the guard; build-cuda-wheel.yml is the compiler that proves it builds.
+# --------------------------------------------------------------------------
+
+
+def test_the_cuda_kernels_patch_path_is_relative_and_is_the_tree_the_script_writes():
+    with open(os.path.join(_CRATE_DIR, "Cargo.toml"), "rb") as f:
+        manifest = tomllib.load(f)
+    path = manifest["patch"]["crates-io"]["candle-kernels"]["path"]
+    assert not os.path.isabs(path), path
+    resolved = os.path.normpath(os.path.join(_CRATE_DIR, path))
+    assert resolved == _CUDA_KERNELS_FORK, (resolved, _CUDA_KERNELS_FORK)
+    with open(os.path.join(resolved, "Cargo.toml"), "rb") as f:
+        package = tomllib.load(f)["package"]
+    assert (package["name"], package["version"]) == ("candle-kernels", "0.11.0"), package
+
+
+def test_the_lock_resolves_cuda_kernels_from_the_fork_not_the_registry():
+    with open(os.path.join(_CRATE_DIR, "Cargo.lock"), "rb") as f:
+        lock = tomllib.load(f)
+    entries = [p for p in lock["package"] if p["name"] == "candle-kernels"]
+    assert len(entries) == 1, entries
+    assert entries[0]["version"] == "0.11.0", entries[0]
+    assert "source" not in entries[0], (
+        "Cargo.lock resolves candle-kernels from %s -- the CUDA build uses the "
+        "unpatched compatibility.cuh and fails on CUDA 12.2+" % entries[0].get("source"))
+
+
+def test_the_cuda_kernels_fork_guards_hmax_nan_on_the_toolkit_version():
+    with open(os.path.join(_CUDA_KERNELS_FORK, "src", "compatibility.cuh"), encoding="utf-8") as f:
+        text = f.read()
+    guard = ("#if __CUDA_ARCH__ < 800 && (__CUDACC_VER_MAJOR__ < 12 || "
+             "(__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 2))\n"
+             "__device__ __forceinline__ __half __hmax_nan(")
+    assert guard in text, "compatibility.cuh defines __hmax_nan without the CUDA 12.2 guard"
+    assert "#if __CUDA_ARCH__ < 800\n" not in text, "the unguarded block is still present"
+
+
+def test_the_cuda_kernels_check_refuses_a_tree_that_drifted_from_the_patch():
+    tmp = tempfile.mkdtemp(prefix="cuda-k-drift-")
+    try:
+        copy = os.path.join(tmp, "candle-kernels")
+        shutil.copytree(_CUDA_KERNELS_FORK, copy)
+        with open(os.path.join(copy, "src", "compatibility.cuh"), "a", encoding="utf-8") as f:
+            f.write("\n// an edit made to the fork and not to the patch\n")
+        proc = _run_script("--check", env_extra={"TORCHNATIVE_CANDLE_CUDA_KERNELS_DIR": copy})
+        assert proc.returncode != 0, proc.stdout
+        assert "compatibility.cuh" in proc.stderr, proc.stderr
+    finally:
+        shutil.rmtree(tmp)
+
+
 def test_the_fork_script_refuses_a_crate_that_is_not_the_pinned_one():
     tmp = tempfile.mkdtemp(prefix="int8-crate-")
     try:

torchnative/rust/torch_c/Cargo.toml

@@ -301,3 +301,14 @@ candle-core = { path = "../vendor/candle-core" }
 # entry for a crate absent from a target's graph costs that target nothing,
 # which is why Android, Linux and wasm builds are unaffected.
 candle-metal-kernels = { path = "../vendor/candle-metal-kernels" }
+
+# The CUDA kernels, forked for a build fix and not for I8 (issue #86).
+# candle-kernels 0.11.0's `src/compatibility.cuh` defines `__hmax_nan` and
+# `__hmin_nan` whenever `__CUDA_ARCH__ < 800`, and CUDA 12.2+ `cuda_fp16.hpp`
+# defines both for every architecture, so every `.cu` file failed with "has
+# already been defined" on the CUDA 12.6 toolkit `build-cuda-wheel.yml` uses.
+# `torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch` adds a toolkit
+# version guard to that one block; vendored, pinned and `--check`ed exactly as
+# the two forks above. Reached only through `candle-core`'s `cuda` feature, so
+# a target without it pays nothing for the entry.
+candle-kernels = { path = "../vendor/candle-kernels" }

torchnative/rust/vendor/candle-core/src/cuda_backend/device.rs

@@ -487,6 +487,13 @@ impl BackendDevice for CudaDevice {
                 let data = self.alloc_zeros::<F8E4M3>(elem_count)?;
                 CudaStorageSlice::F8E4M3(data)
             }
+            DType::I8 => {
+                return Err(CudaError::UnsupportedDtype {
+                    dtype: DType::I8,
+                    op: "zeros",
+                }
+                .into())
+            }
             DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => {
                 return Err(
                     CudaError::InternalError("Dummy types not supported in CUDA backend").into(),
@@ -506,6 +513,7 @@ impl BackendDevice for CudaDevice {
             // TODO: Add support for F16 and BF16 though this is likely to require some upstream
             // cudarc changes.
             DType::U8
+            | DType::I8
             | DType::U32
             | DType::I16
             | DType::I32
@@ -561,6 +569,7 @@ impl BackendDevice for CudaDevice {
         };
         let slice = match dtype {
             DType::U8
+            | DType::I8
             | DType::U32
             | DType::I16
             | DType::I32
@@ -641,6 +650,13 @@ impl BackendDevice for CudaDevice {
                 let data = self.alloc::<F8E4M3>(elem_count)?;
                 CudaStorageSlice::F8E4M3(data)
             }
+            DType::I8 => {
+                return Err(CudaError::UnsupportedDtype {
+                    dtype: DType::I8,
+                    op: "alloc_uninit",
+                }
+                .into())
+            }
             DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => {
                 return Err(
                     CudaError::InternalError("Dummy types not supported in CUDA backend").into(),
@@ -695,7 +711,8 @@ impl BackendDevice for CudaDevice {
                 let data = self.clone_htod(storage)?;
                 CudaStorageSlice::F8E4M3(data)
             }
-            CpuStorageRef::F4(_)
+            CpuStorageRef::I8(_)
+            | CpuStorageRef::F4(_)
             | CpuStorageRef::F6E2M3(_)
             | CpuStorageRef::F6E3M2(_)
             | CpuStorageRef::F8E8M0(_) => {
@@ -754,7 +771,8 @@ impl BackendDevice for CudaDevice {
                 let data = self.clone_htod(storage)?;
                 CudaStorageSlice::F8E4M3(data)
             }
-            CpuStorage::F4(_)
+            CpuStorage::I8(_)
+            | CpuStorage::F4(_)
             | CpuStorage::F6E2M3(_)
             | CpuStorage::F6E3M2(_)
             | CpuStorage::F8E8M0(_) => {
@@ -813,7 +831,8 @@ impl BackendDevice for CudaDevice {
                 let data = self.clone_htod(&storage)?;
                 CudaStorageSlice::F8E4M3(data)
             }
-            CpuStorage::F4(_)
+            CpuStorage::I8(_)
+            | CpuStorage::F4(_)
             | CpuStorage::F6E2M3(_)
             | CpuStorage::F6E3M2(_)
             | CpuStorage::F8E8M0(_) => {

torchnative/rust/vendor/candle-core/src/cuda_backend/mod.rs

@@ -47,6 +47,7 @@ impl crate::scalar::Scalar {
         use crate::scalar::Scalar;
         match self {
             Scalar::U8(v) => builder.arg(v),
+            Scalar::I8(v) => builder.arg(v),
             Scalar::U32(v) => builder.arg(v),
             Scalar::I16(v) => builder.arg(v),
             Scalar::I32(v) => builder.arg(v),
@@ -1367,6 +1368,13 @@ impl CudaStorage {
                 let result = dst_stream.clone_dtod(cuda_slice).w()?;
                 CudaStorageSlice::F8E8M0(result)
             }
+            DType::I8 => {
+                return Err(CudaError::UnsupportedDtype {
+                    dtype: DType::I8,
+                    op: "transfer_to_device",
+                }
+                .into())
+            }
         };
 
         Ok(Self {
@@ -1669,6 +1677,13 @@ impl BackendStorage for CudaStorage {
             DType::I16 | DType::I32 => {
                 return Err(CudaError::InternalError("i16,i32 dtypes are not supported").into())
             }
+            DType::I8 => {
+                return Err(CudaError::UnsupportedDtype {
+                    dtype: DType::I8,
+                    op: "to_dtype",
+                }
+                .into())
+            }
             DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => {
                 return Err(
                     CudaError::InternalError("Dummy types not supported in CUDA backend").into(),

torchnative/rust/vendor/candle-kernels/.cargo_vcs_info.json

@@ -0,0 +1,6 @@
+{
+  "git": {
+    "sha1": "31f35b147389700ed2a178ee66a91c3cc25cc80d"
+  },
+  "path_in_vcs": "candle-kernels"
+}