fix(math-cuda): driver-independent cuda builds (cudarc pin + cubin AOT)

.github/workflows/benchmark-gpu.yml

@@ -358,14 +358,13 @@ jobs:
           # the build through the CUDA prover path. NOTE: requires this PR's bench_abba.sh change
           # (the BENCH_FEATURES env) to be on main — i.e. it only takes effect after merge.
           # REBUILD=1: each Vast box is fresh, GPU-specific hardware — always rebuild both
-          # binaries (PTX is compiled for the detected arch); never trust a cached binary.
-          # CUDARC_PIN: pin cudarc to a fixed CUDA version (cuda-12080 = CUDA 12.8, matching the
-          # cuda_max_good>=12.8 offer floor) and drop fallback-latest, so cudarc binds a known
-          # symbol set instead of its newest. With fallback-latest cudarc requested a symbol the
-          # box's driver doesn't export (e.g. cuDevSmResourceSplit) -> runtime panic. This is the
-          # too-new end of the same compatibility window that MIN_DRIVER>=580 guards at the
-          # too-old end (older drivers lack cuCtxGetDevice_v2 and the GPU path falls back to CPU).
-          # nvidia-smi is logged for diagnosing driver issues.
+          # binaries (cubin is compiled for the detected arch); never trust a cached binary.
+          # CUDARC_PIN: compat shim for pre-pin baseline shas. cudarc's CUDA version is now pinned
+          # permanently in crypto/math-cuda/Cargo.toml (cuda-12080), so this no-ops on shas that
+          # carry the pin and only rewrites older baselines (where fallback-latest could request a
+          # symbol the box's driver doesn't export, e.g. cuDevSmResourceSplit -> runtime panic).
+          # MIN_DRIVER>=580 still guards the too-old end (older drivers lack cuCtxGetDevice_v2 and
+          # the GPU path falls back to CPU). nvidia-smi is logged for diagnosing driver issues.
           REMOTE="set -e; cd /workspace/lambda_vm; \
             command -v python3 >/dev/null || { apt-get update -qq && apt-get install -y -qq python3; }; \
             nvidia-smi || true; \

.github/workflows/gpu-tests.yml

@@ -287,7 +287,8 @@ jobs:
             ''|*[!A-Za-z0-9._/-]*) echo "::error::invalid ref: '$REF'"; exit 1 ;;
           esac
           # Check out the ref under test on the box, then run the CUDA test groups.
-          # gpu_test.sh owns the CUDARC_PIN / SYSROOT_DIR defaults — don't duplicate them here.
+          # gpu_test.sh owns the SYSROOT_DIR default — don't duplicate it here. (cudarc's CUDA
+          # version is pinned in crypto/math-cuda/Cargo.toml, so no CUDARC_PIN is needed.)
           REMOTE="set -e; cd /workspace/lambda_vm; \
             git fetch --force origin '$REF'; \
             git checkout -f FETCH_HEAD; \

README.md

@@ -233,10 +233,15 @@ The CUDA test groups run only on a machine with an NVIDIA GPU and `nvcc`:
 - `make test-prover-cuda` — the prover/stark/crypto/ecsm suite with the GPU path enabled
 - `make test-prover-comprehensive-cuda` — the comprehensive all-instructions prove on the GPU path
 
-The kernels are compiled by `nvcc` into PTX that the driver JIT-compiles at load, so the GPU's
-driver must be new enough for the toolkit — an older driver rejects the PTX with
-`CUDA_ERROR_UNSUPPORTED_PTX_VERSION`. These groups run automatically on a rented GPU in the merge
-queue via `.github/workflows/gpu-tests.yml` (which filters offers on `cuda_max_good`).
+The kernels are AOT-compiled by `nvcc` into native cubin (SASS) for the host GPU's real arch
+(detected via `nvidia-smi`, or overridden with `CUDARC_NVCC_ARCH`), not PTX. This sidesteps the
+PTX-ISA JIT version check, so a CUDA toolkit *newer* than the driver still loads and runs — no
+`CUDA_ERROR_UNSUPPORTED_PTX_VERSION` and no need to hand-match the toolkit to the driver. The only
+requirement is that the toolkit knows the GPU's compute capability (a too-old toolkit fails loudly
+at `nvcc` build time). cudarc's host-side driver-API symbol set is likewise pinned to a safe floor
+(`cuda-12080`) in `crypto/math-cuda/Cargo.toml`, so no `CUDARC_CUDA_VERSION` env is needed either.
+These groups run automatically on a rented GPU in the merge queue via
+`.github/workflows/gpu-tests.yml` (which filters offers on `cuda_max_good`).
 
 ## Benchmarking & Profiling
 

crypto/math-cuda/Cargo.toml

@@ -6,12 +6,26 @@ edition = "2024"
 license.workspace = true
 
 [dependencies]
+# cudarc CUDA version is PINNED to `cuda-12080` (CUDA 12.8) — do NOT restore
+# `cuda-version-from-build-system` + `fallback-latest`. Rationale:
+#   * That auto-detect binds the newest symbol set the *build toolkit* knows
+#     (e.g. a CUDA 13.1 toolkit pulls in `cuDevSmResourceSplit`, gated behind
+#     `cuda-13010`/`cuda-13020`). cudarc eagerly resolves those symbols at CUDA
+#     init; a driver that predates them (e.g. 580.x = CUDA 13.0 max) has no such
+#     export, so the `dynamic-loading` resolver `.expect()`s and PANICS.
+#   * This crate's cudarc surface is entirely CUDA-11-era
+#     (CudaContext/CudaFunction/CudaSlice/CudaStream/LaunchConfig/PushKernelArg/
+#     DriverError/Ptx — no green contexts). The 12.8 symbol set is a strict
+#     subset every >=12.8 driver exports, so pinning it resolves cleanly on any
+#     supported driver and a *newer* driver loses nothing we use.
+#   * This replaces the fragile per-script `sed` pin in scripts/gpu_test.sh.
+# To move the floor (e.g. to use a newer driver-API symbol), bump this one
+# feature deliberately — see crypto/math-cuda/build.rs and README "GPU Tests".
 cudarc = { version = "0.19", default-features = false, features = [
     "driver",
     "nvrtc",
     "std",
-    "cuda-version-from-build-system",
-    "fallback-latest",
+    "cuda-12080",
     "dynamic-loading",
 ] }
 math = { path = "../math" }

crypto/math-cuda/build.rs

@@ -15,10 +15,11 @@ fn nvcc_path() -> PathBuf {
 }
 
 /// Query `nvidia-smi` for the local GPU's compute capability (e.g. "12.0"
-/// for Blackwell). Returns a `compute_XX` target on success, falling back
-/// to `compute_89` (Ada) when no GPU is visible or the query fails.
+/// for Blackwell) and return a *real* arch (`sm_XX`) suitable for cubin
+/// (SASS) generation. Falls back to `sm_89` (Ada) when no GPU is visible or
+/// the query fails.
 fn detect_arch() -> String {
-    const FALLBACK: &str = "compute_89";
+    const FALLBACK: &str = "sm_89";
     let output = match Command::new("nvidia-smi")
         .args(["--query-gpu=compute_cap", "--format=csv,noheader"])
         .output()
@@ -40,13 +41,26 @@ fn detect_arch() -> String {
         None => return FALLBACK.to_string(),
     };
     if major.chars().all(|c| c.is_ascii_digit()) && minor.chars().all(|c| c.is_ascii_digit()) {
-        format!("compute_{major}{minor}")
+        format!("sm_{major}{minor}")
     } else {
         FALLBACK.to_string()
     }
 }
 
-fn compile_ptx(src: &str, out_name: &str, have_nvcc: bool) {
+/// Normalize a user-supplied `CUDARC_NVCC_ARCH` override to a *real* arch
+/// (`sm_XX`). cubin (SASS) generation rejects the *virtual* `compute_XX`
+/// form, but we accept it (and a bare `XX`) for backwards compatibility.
+fn to_real_arch(arch: &str) -> String {
+    if let Some(n) = arch.strip_prefix("compute_") {
+        format!("sm_{n}")
+    } else if arch.starts_with("sm_") {
+        arch.to_string()
+    } else {
+        format!("sm_{arch}")
+    }
+}
+
+fn compile_kernel(src: &str, out_name: &str, have_nvcc: bool) {
     let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
     let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
     let src_path = manifest_dir.join("kernels").join(src);
@@ -57,30 +71,39 @@ fn compile_ptx(src: &str, out_name: &str, have_nvcc: bool) {
     println!("cargo:rerun-if-env-changed=CUDA_PATH");
     println!("cargo:rerun-if-env-changed=CUDARC_NVCC_ARCH");
 
-    // When nvcc is missing from PATH, emit an empty PTX stub so the crate
-    // still compiles. include_str! in src/device.rs needs the file to exist
-    // at build time. Any runtime kernel call panics in cudarc when loading
-    // the empty module. We can't run GPU code without nvcc on the build
-    // host anyway.
+    // When nvcc is missing from PATH, emit an empty cubin stub so the crate
+    // still compiles. include_bytes! in src/device.rs needs the file to exist
+    // at build time. Any runtime kernel call fails to load the empty module and
+    // the caller falls back to CPU. We can't run GPU code without nvcc on the
+    // build host anyway.
     if !have_nvcc {
-        fs::write(&out_path, "").expect("failed to write empty PTX stub");
+        fs::write(&out_path, "").expect("failed to write empty cubin stub");
         return;
     }
 
-    // Emit PTX for a virtual architecture; the CUDA driver JIT-compiles it for the
-    // actual GPU at load time. Override with CUDARC_NVCC_ARCH to pin a specific
-    // compute capability. If unset, try `nvidia-smi` to match the host GPU
-    // (avoids JIT failures like nvcc-13.0 PTX rejected on Blackwell drivers);
-    // fall back to compute_89 (Ada) when detection fails.
+    // AOT-compile each kernel to a native cubin (SASS) for the host GPU's real
+    // arch, NOT to PTX. This sidesteps the driver's PTX-ISA JIT version check:
+    // a toolkit's PTX ISA is fixed by its CUDA version (e.g. CUDA 13.1 emits PTX
+    // .version 9.1), and a driver older than that toolkit rejects the module at
+    // load with CUDA_ERROR_UNSUPPORTED_PTX_VERSION -> every kernel silently
+    // falls back to CPU. A cubin carries pre-compiled SASS for a real arch, so
+    // the driver loads it directly as long as it supports that GPU (which the
+    // driver installed for that GPU always does) — regardless of the toolkit's
+    // CUDA version. See README "GPU Tests".
     //
-    // NOTE: this `-arch` only sets the *virtual arch*, not the PTX ISA version, which is
-    // fixed by this nvcc's CUDA toolkit. The runtime driver must support that toolkit's CUDA
-    // version or it rejects the PTX with CUDA_ERROR_UNSUPPORTED_PTX_VERSION — i.e. the box's
-    // driver CUDA must be >= the build toolkit's CUDA. See README "GPU Tests".
-    let arch = env::var("CUDARC_NVCC_ARCH").unwrap_or_else(|_| detect_arch());
+    // Trade-off: a cubin is arch-specific (an `sm_120` cubin runs only on
+    // `sm_120`). We build+run on the same GPU box in every flow and detect the
+    // arch from that box's `nvidia-smi`, so this is exactly right. Override with
+    // CUDARC_NVCC_ARCH (compute_XX / sm_XX / bare XX all accepted) to
+    // cross-compile for a different arch; fall back to sm_89 (Ada) when
+    // detection fails. If the toolkit is too old to know the GPU's arch, nvcc
+    // fails loudly here at build time (better than a silent runtime fallback).
+    let arch = env::var("CUDARC_NVCC_ARCH")
+        .map(|a| to_real_arch(&a))
+        .unwrap_or_else(|_| detect_arch());
 
     let status = Command::new(nvcc_path())
-        .args(["--ptx", "-O3", "-std=c++17", "-arch", &arch, "-o"])
+        .args(["--cubin", "-O3", "-std=c++17", "-arch", &arch, "-o"])
         .arg(&out_path)
         .arg(&src_path)
         .status()
@@ -107,19 +130,19 @@ fn main() {
         .unwrap_or(false);
     if !have_nvcc {
         println!(
-            "cargo:warning=math-cuda: nvcc not found at {} — emitting empty PTX stubs. \
-             Runtime GPU calls will panic. Install CUDA and rebuild for a working backend.",
+            "cargo:warning=math-cuda: nvcc not found at {} — emitting empty cubin stubs. \
+             Runtime GPU calls fall back to CPU. Install CUDA and rebuild for a working backend.",
             nvcc_path().display()
         );
     }
 
-    compile_ptx("arith.cu", "arith.ptx", have_nvcc);
-    compile_ptx("ntt.cu", "ntt.ptx", have_nvcc);
-    compile_ptx("keccak.cu", "keccak.ptx", have_nvcc);
-    compile_ptx("barycentric.cu", "barycentric.ptx", have_nvcc);
-    compile_ptx("deep.cu", "deep.ptx", have_nvcc);
-    compile_ptx("fri.cu", "fri.ptx", have_nvcc);
-    compile_ptx("inverse.cu", "inverse.ptx", have_nvcc);
-    compile_ptx("logup.cu", "logup.ptx", have_nvcc);
-    compile_ptx("constraint_interp.cu", "constraint_interp.ptx", have_nvcc);
+    compile_kernel("arith.cu", "arith.cubin", have_nvcc);
+    compile_kernel("ntt.cu", "ntt.cubin", have_nvcc);
+    compile_kernel("keccak.cu", "keccak.cubin", have_nvcc);
+    compile_kernel("barycentric.cu", "barycentric.cubin", have_nvcc);
+    compile_kernel("deep.cu", "deep.cubin", have_nvcc);
+    compile_kernel("fri.cu", "fri.cubin", have_nvcc);
+    compile_kernel("inverse.cu", "inverse.cubin", have_nvcc);
+    compile_kernel("logup.cu", "logup.cubin", have_nvcc);
+    compile_kernel("constraint_interp.cu", "constraint_interp.cubin", have_nvcc);
 }

crypto/math-cuda/src/device.rs

@@ -90,16 +90,21 @@ impl Drop for PinnedStaging {
     }
 }
 
-const ARITH_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/arith.ptx"));
-const NTT_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/ntt.ptx"));
-const KECCAK_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/keccak.ptx"));
-const BARY_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/barycentric.ptx"));
-const DEEP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/deep.ptx"));
-const FRI_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/fri.ptx"));
-const INVERSE_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/inverse.ptx"));
-const LOGUP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/logup.ptx"));
-const CONSTRAINT_INTERP_PTX: &str =
-    include_str!(concat!(env!("OUT_DIR"), "/constraint_interp.ptx"));
+// Kernels are AOT-compiled to native cubin (SASS) by build.rs, embedded here,
+// and loaded via `Ptx::from_binary` (cubin bytes -> cuModuleLoadData). This
+// avoids the PTX-ISA/driver-version JIT check — see build.rs `compile_kernel`.
+// An empty slice (nvcc-less stub build) fails to load at runtime and the caller
+// falls back to CPU.
+const ARITH_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/arith.cubin"));
+const NTT_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ntt.cubin"));
+const KECCAK_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/keccak.cubin"));
+const BARY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/barycentric.cubin"));
+const DEEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/deep.cubin"));
+const FRI_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/fri.cubin"));
+const INVERSE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/inverse.cubin"));
+const LOGUP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/logup.cubin"));
+const CONSTRAINT_INTERP_CUBIN: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/constraint_interp.cubin"));
 
 /// Number of CUDA streams in the pool. Larger pools let many rayon-parallel
 /// callers overlap on the GPU without serializing on stream ownership. The
@@ -125,7 +130,7 @@ pub struct Backend {
     /// [`detect_vram_budget_bytes`].
     vram_budget_bytes: u64,
 
-    // arith.ptx
+    // arith.cubin
     pub vector_add_u64: CudaFunction,
     pub gl_add: CudaFunction,
     pub gl_sub: CudaFunction,
@@ -135,7 +140,7 @@ pub struct Backend {
     pub ext3_add: CudaFunction,
     pub ext3_sub: CudaFunction,
 
-    // ntt.ptx
+    // ntt.cubin
     pub bit_reverse_permute: CudaFunction,
     pub ntt_dit_level: CudaFunction,
     pub ntt_dit_8_levels: CudaFunction,
@@ -152,7 +157,7 @@ pub struct Backend {
     pub pointwise_mul_row_major: CudaFunction,
     pub matrix_transpose_strided: CudaFunction,
 
-    // keccak.ptx
+    // keccak.cubin
     pub keccak256_leaves_base_row_major_row_pair: CudaFunction,
     pub keccak256_leaves_base_batched: CudaFunction,
     pub keccak256_leaves_base_row_pair_batched: CudaFunction,
@@ -162,22 +167,22 @@ pub struct Backend {
     pub keccak_merkle_level: CudaFunction,
     pub merkle_gather_paths: CudaFunction,
 
-    // barycentric.ptx
+    // barycentric.cubin
     pub barycentric_base_batched: CudaFunction,
     pub barycentric_ext3_batched: CudaFunction,
     pub barycentric_base_batched_strided: CudaFunction,
     pub barycentric_ext3_batched_strided: CudaFunction,
     pub gather_rows_base: CudaFunction,
     pub gather_rows_ext3: CudaFunction,
 
-    // deep.ptx
+    // deep.cubin
     pub deep_composition_ext3_row: CudaFunction,
 
-    // fri.ptx
+    // fri.cubin
     pub fri_fold_ext3: CudaFunction,
     pub fri_update_twiddles: CudaFunction,
 
-    // inverse.ptx
+    // inverse.cubin
     pub compute_denoms_ext3: CudaFunction,
     pub block_inclusive_scan_fwd_ext3: CudaFunction,
     pub apply_block_offsets_fwd_ext3: CudaFunction,
@@ -192,7 +197,7 @@ pub struct Backend {
     pub logup_finalize_accum_ext3: CudaFunction,
     pub logup_assemble_aux_ext3: CudaFunction,
 
-    // constraint_interp.ptx
+    // constraint_interp.cubin
     pub constraint_interp_kernel: CudaFunction,
     pub constraint_composition_kernel: CudaFunction,
 
@@ -291,15 +296,16 @@ impl Backend {
         // we keep the current behaviour.
         retain_default_mempool(&ctx);
 
-        let arith = ctx.load_module(Ptx::from_src(ARITH_PTX))?;
-        let ntt = ctx.load_module(Ptx::from_src(NTT_PTX))?;
-        let keccak = ctx.load_module(Ptx::from_src(KECCAK_PTX))?;
-        let bary = ctx.load_module(Ptx::from_src(BARY_PTX))?;
-        let deep = ctx.load_module(Ptx::from_src(DEEP_PTX))?;
-        let fri = ctx.load_module(Ptx::from_src(FRI_PTX))?;
-        let inverse = ctx.load_module(Ptx::from_src(INVERSE_PTX))?;
-        let logup = ctx.load_module(Ptx::from_src(LOGUP_PTX))?;
-        let constraint_interp = ctx.load_module(Ptx::from_src(CONSTRAINT_INTERP_PTX))?;
+        let arith = ctx.load_module(Ptx::from_binary(ARITH_CUBIN.to_vec()))?;
+        let ntt = ctx.load_module(Ptx::from_binary(NTT_CUBIN.to_vec()))?;
+        let keccak = ctx.load_module(Ptx::from_binary(KECCAK_CUBIN.to_vec()))?;
+        let bary = ctx.load_module(Ptx::from_binary(BARY_CUBIN.to_vec()))?;
+        let deep = ctx.load_module(Ptx::from_binary(DEEP_CUBIN.to_vec()))?;
+        let fri = ctx.load_module(Ptx::from_binary(FRI_CUBIN.to_vec()))?;
+        let inverse = ctx.load_module(Ptx::from_binary(INVERSE_CUBIN.to_vec()))?;
+        let logup = ctx.load_module(Ptx::from_binary(LOGUP_CUBIN.to_vec()))?;
+        let constraint_interp =
+            ctx.load_module(Ptx::from_binary(CONSTRAINT_INTERP_CUBIN.to_vec()))?;
 
         let mut streams = Vec::with_capacity(STREAM_POOL_SIZE);
         for _ in 0..STREAM_POOL_SIZE {
@@ -504,7 +510,26 @@ pub fn backend() -> Result<&'static Backend> {
     if let Some(b) = BACKEND.get() {
         return Ok(b);
     }
-    let b = Backend::init()?;
+    let b = match Backend::init() {
+        Ok(b) => b,
+        Err(e) => {
+            // Backend init failing means every GPU entry point silently falls
+            // back to CPU. That is expected on a GPU-less host, but it also
+            // fires when the AOT cubins won't load — most often a build-host vs
+            // run-host GPU-arch mismatch (cubins are compiled for the detected
+            // `sm_XX`) or an empty nvcc-less stub. Warn once so the fallback is
+            // never silent: rebuild on the run host, or set `CUDARC_NVCC_ARCH`.
+            static WARNED: std::sync::Once = std::sync::Once::new();
+            WARNED.call_once(|| {
+                eprintln!(
+                    "math-cuda: GPU backend unavailable ({e}) — running on CPU. \
+                     If a GPU is present this is likely a kernel-cubin arch mismatch; \
+                     rebuild on the run host or set CUDARC_NVCC_ARCH to its sm_XX."
+                );
+            });
+            return Err(e);
+        }
+    };
     let _ = BACKEND.set(b);
     Ok(BACKEND.get().expect("backend just initialised"))
 }