feat(cuda): selectable compilation target via -K cuda_arch, traced by Lake

Device code is compiled *for* a GPU architecture, but the native kernels were compiled
with no -arch/-gencode at all, leaving nvcc on its built-in default. That default is a
property of the toolkit, not of the machine: CUDA 13.0 emits sm_75 SASS plus compute_75
PTX. On any other GPU the kernels then run through forward PTX JIT rather than native
SASS, or on an older architecture do not load at all, and nothing about the build says so.

Add -K cuda_arch=<spec>, with TORCHLEAN_CUDA_ARCH as a fallback so a container image or CI
job can select a target without rewriting its build command. A bare spec becomes
-arch=<spec>, which covers sm_86, compute_86, native, all, and all-major; a spec starting
with `-` is passed to nvcc verbatim, which is how a multi-architecture binary is requested
and which keeps this package free of any policy about which architecture carries the PTX.

Selecting a target is only meaningful if changing it rebuilds, so also split
buildNativeBackendLib's compiler arguments the way buildO intends: include paths stay in
weakArgs, where a moved checkout does not invalidate every object, and the flags that
change what the compiler emits move to traceArgs, which buildO hashes. Previously every
argument sat in weakArgs, so a changed optimization level or compilation target left the
existing objects looking current. nvcc's own NVCC_APPEND_FLAGS has the same hole and Lake
cannot close it — the option and the environment variable added here both participate in
the trace, and the documentation says which to prefer.

scripts/checks/cuda_arch_target.sh asserts the whole surface — flags reaching nvcc,
recompilation on a changed target, no recompilation on an unchanged one, the environment
fallback, its precedence, and verbatim pass-through — against a recording stand-in for
nvcc, so it needs neither a CUDA toolkit nor a GPU and runs in CI.

One consequence worth expecting: moving the flags into the trace changes the trace of
every object built by buildNativeBackendLib, so the first build after this lands
recompiles those four objects once — the four .cu kernels under CUDA, or the four C
stubs on a CPU-only machine.

Author: NicolasRouquette

.github/workflows/ci.yml

@@ -47,6 +47,10 @@ jobs:
       - name: Run curated test suite
         run: ~/.elan/bin/lake test
 
+      # Runs against a recording stand-in for nvcc, so this needs no CUDA toolkit and no GPU.
+      - name: Check the CUDA compilation target
+        run: LAKE=~/.elan/bin/lake scripts/checks/cuda_arch_target.sh
+
   slow_proofs:
     name: CI (slow proofs)
     runs-on: ubuntu-latest

home_page/cuda/index.md

@@ -47,6 +47,38 @@ Run the CUDA sanitizer suite when changing native kernels:
 scripts/checks/cuda_sanitize_tests.sh --all-tools
 ```
 
+## Compilation Target
+
+`-K cuda=true` says to compile the native kernels; it does not say which GPU to compile them
+*for*. That second choice belongs to `nvcc`, which without instruction applies a built-in default
+that changes with the toolkit version — CUDA 13.0, for instance, emits `sm_75` machine code plus
+`compute_75` PTX. Such a binary runs at full speed on that architecture, reaches a newer GPU only
+through forward PTX just-in-time compilation, and does not load on an older one at all. Name the
+target when the deployment GPU is known:
+
+```bash
+lake -R -K cuda=true -K cuda_arch=sm_86 build     # an A10G or an RTX A4500
+lake -R -K cuda=true -K cuda_arch=native build    # whatever GPU this machine has
+```
+
+A value starting with `-` is passed to `nvcc` verbatim, which is how one binary is compiled for
+several architectures at once:
+
+```bash
+lake -R -K cuda=true \
+  -K cuda_arch="-gencode arch=compute_75,code=sm_75 \
+                -gencode arch=compute_90,code=[sm_90,compute_90]" build
+```
+
+`TORCHLEAN_CUDA_ARCH` carries the same value when the Lake option is absent, so a container image
+or CI job can select a target without rewriting its build command. Both spellings participate in
+Lake's build trace: changing the target recompiles the kernels. `nvcc`'s own `NVCC_APPEND_FLAGS`
+does not — Lake cannot see it, judges the existing objects current, and links kernels compiled for
+the previous target, which is visible only as unexplained throughput. Prefer the option, and delete
+`.lake/build/torchlean_*.o` if a target was ever set that way.
+
+`scripts/checks/cuda_arch_target.sh` asserts this behavior, and needs neither a toolkit nor a GPU.
+
 ## What CUDA Covers
 
 The CUDA path is used for supported Float32 tensor operations: elementwise arithmetic, reductions,

home_page/installation/index.md

@@ -101,6 +101,10 @@ The two CUDA choices happen at different times. `-K cuda=true` tells Lake to com
 native CUDA implementation. `--device cuda` asks the executable to use it. A CPU-linked executable
 rejects `--device cuda` instead of silently moving the run back to the CPU.
 
+`nvcc` compiles the kernels for its own default architecture unless told otherwise. Add
+`-K cuda_arch=sm_86` (or `native`, or a verbatim `-gencode` list) to compile for the GPU you deploy
+on; the [CUDA guide]({{ '/cuda/' | relative_url }}) explains what the default costs you.
+
 Use `-R` whenever you switch between CPU and CUDA configurations; it forces Lake to recompute the
 build description. The CUDA regression suite is:
 

lakefile.lean

@@ -31,6 +31,52 @@ private def cudaHome : String :=
   | some p => cleanCudaHome p
   | none => "/usr/local/cuda"
 
+/-- CUDA compilation target for the native kernels, from `-K cuda_arch=...`.
+
+Device code is compiled *for* an architecture. With no target, `nvcc` applies its own built-in
+default, which changes with the toolkit version (CUDA 13.0 emits `sm_75` SASS plus `compute_75`
+PTX). A binary then runs natively on that architecture only, reaching newer GPUs through forward
+PTX JIT and older ones not at all, so the default is a portability choice rather than a
+performance one.
+
+Accepted values:
+* a bare architecture, such as `sm_86`, `compute_86`, `native`, `all`, or `all-major`, passed as
+  `-arch=<value>` — `nvcc` shorthand for that architecture's SASS *and* its PTX;
+* a value starting with `-`, such as `-gencode arch=compute_86,code=[sm_86,compute_86]`, split on
+  spaces and passed verbatim, which is how a multi-architecture binary is requested.
+
+The `TORCHLEAN_CUDA_ARCH` environment variable supplies the same value when the Lake option is
+absent, so a container or CI job can select the target without rewriting its `lake` invocation.
+Prefer either spelling over `nvcc`'s own `NVCC_APPEND_FLAGS`: both reach `buildO`'s traced
+arguments below, whereas a variable Lake never reads leaves the existing objects looking current,
+and the build then links kernels compiled for the previous architecture. -/
+private def cudaArchConfig : Option String :=
+  match get_config? cuda_arch with
+  | some v =>
+      let t := v.trimAscii.toString
+      if t.isEmpty then none else some t
+  | none => none
+
+/-- Resolve the CUDA compilation target: the `cuda_arch` Lake option, else `TORCHLEAN_CUDA_ARCH`,
+else none (leaving `nvcc` on its built-in default). -/
+private def resolveCudaArch : SpawnM (Option String) := do
+  match cudaArchConfig with
+  | some spec => return some spec
+  | none =>
+      let env ← IO.getEnv "TORCHLEAN_CUDA_ARCH"
+      return env.bind fun v =>
+        let t := v.trimAscii.toString
+        if t.isEmpty then none else some t
+
+/-- `nvcc` flags for a resolved CUDA compilation target; empty when no target was requested. -/
+private def cudaArchArgs : Option String → Array String
+  | none => #[]
+  | some spec =>
+      if spec.startsWith "-" then
+        ((spec.splitOn " ").filter (!·.isEmpty)).toArray
+      else
+        #[s!"-arch={spec}"]
+
 /-- Optional explicit LibTorch root from `-K libtorch_home=...`. -/
 private def libtorchHomeConfig : Option String :=
   match get_config? libtorch_home with
@@ -240,22 +286,26 @@ private def buildNativeBackendLib (pkg : Package) (spec : NativeBackendLib) := d
   let headerDeps ← nativeHeaderDeps pkg
   let includeArgs := nativeIncludeArgs pkg
   let libFile := pkg.buildDir / nameToStaticLib spec.stem
+  -- Include paths stay in `weakArgs`, where a moved checkout does not invalidate every object.
+  -- The flags that change what the compiler emits belong in `traceArgs`: `buildO` hashes those,
+  -- so switching optimization level or CUDA compilation target rebuilds instead of silently
+  -- reusing objects built for the previous one.
   if cudaEnabled then
     let srcJob ← inputFile (pkg.dir / spec.cudaSrc) false
     let oFile := pkg.buildDir / s!"{spec.stem}.o"
+    let archArgs := cudaArchArgs (← resolveCudaArch)
     let oJob ← buildO oFile srcJob
-      (#[
-        "-I", lean.includeDir.toString,
-        "-I", s!"{cudaHome}/include",
-        "-c", "--std=c++17", "-O2", "-Xcompiler", "-fPIC"
-      ] ++ includeArgs) #[] "nvcc" (pure headerDeps.getTrace)
+      (weakArgs := #["-I", lean.includeDir.toString, "-I", s!"{cudaHome}/include"] ++ includeArgs)
+      (traceArgs := #["-c", "--std=c++17", "-O2", "-Xcompiler", "-fPIC"] ++ archArgs)
+      (compiler := "nvcc") (extraDepTrace := pure headerDeps.getTrace)
     buildStaticLib libFile #[oJob]
   else
     let srcJob ← inputFile (pkg.dir / spec.stubSrc) false
     let oFile := pkg.buildDir / s!"{spec.stem}_stub.o"
     let oJob ← buildO oFile srcJob
-      (#["-I", lean.includeDir.toString] ++ includeArgs ++ #["-O2", "-fPIC"])
-      #[] "cc" (pure headerDeps.getTrace)
+      (weakArgs := #["-I", lean.includeDir.toString] ++ includeArgs)
+      (traceArgs := #["-O2", "-fPIC"])
+      (compiler := "cc") (extraDepTrace := pure headerDeps.getTrace)
     buildStaticLib libFile #[oJob]
 
 /-- Native backend for `torchlean_dgemm_cuda`: CUDA+cuBLAS when `-K cuda=true`, else C stub. -/

scripts/README.md

@@ -25,6 +25,7 @@ These scripts are used by the build, documentation, and local verification paths
 - `checks/example_regression.sh`
 - `checks/cuda_sanitize_tests.sh`
 - `checks/cuda_profile_tests.sh`
+- `checks/cuda_arch_target.sh`
 - `checks/repo_lint.py`
 - `checks/TorchLeanLint.lean`
 - `checks/dependency_audit.py`
@@ -104,6 +105,10 @@ Generated locally:
 - `checks/cuda_sanitize_tests.sh`: CUDA sanitizer runner for the CUDA runtime test suite.
 - `checks/cuda_profile_tests.sh`: optional Nsight Systems / Nsight Compute wrapper for CUDA
   performance reports.
+- `checks/cuda_arch_target.sh`: checks that the `cuda_arch` Lake option and the
+  `TORCHLEAN_CUDA_ARCH` environment fallback reach `nvcc` and that changing the target
+  recompiles the kernels. A recording stand-in for `nvcc` supplies the evidence, so neither a
+  CUDA toolkit nor a GPU is needed.
 - `checks/repo_lint.py`: repository lint used by `lake lint`. It checks source hygiene, public API
   boundaries, import-only aggregators, fixed-rank names in public tensor/model APIs,
   trusted-axiom quarantine, public-example spellings, DocGen/Verso math markup, and module

scripts/checks/cuda_arch_target.sh

@@ -0,0 +1,150 @@
+#!/usr/bin/env bash
+# Check that the CUDA compilation target reaches nvcc and invalidates stale objects.
+set -euo pipefail
+
+usage() {
+  cat <<'EOF'
+Usage: scripts/checks/cuda_arch_target.sh [options]
+
+Verify the `cuda_arch` Lake option end to end without a CUDA toolkit or a GPU.
+
+A recording stand-in for nvcc is placed first on PATH; it logs its argument vector and emits an
+empty object file. One extern library is then built repeatedly under different targets, and the
+log is asserted against what each build should have compiled:
+
+  no target                      no -arch flag; nvcc keeps its built-in default
+  -K cuda_arch=sm_86             -arch=sm_86
+  -K cuda_arch=sm_86 (repeated)  no recompilation
+  -K cuda_arch=sm_89             recompiled, because the target is a traced argument
+  TORCHLEAN_CUDA_ARCH=sm_90      -arch=sm_90, the environment fallback
+  both, option and environment   the option wins
+  -K cuda_arch=-gencode ...      passed through verbatim, for multi-architecture binaries
+
+The stand-in's objects are removed afterwards, so a later real CUDA build cannot mistake them
+for its own. Any genuine object for that library goes with them, which costs one recompilation
+on the next real CUDA build. The stored build configuration is returned to its default, since
+every build here enables CUDA.
+
+Options:
+  --target LIB          Extern library to build. Default: torchlean_dgemm_cuda.
+  -h, --help            Show this help message.
+
+Environment:
+  LAKE                  Lake executable to use (default: lake).
+
+Examples:
+  scripts/checks/cuda_arch_target.sh
+  LAKE=~/.elan/bin/lake scripts/checks/cuda_arch_target.sh
+EOF
+}
+
+LAKE="${LAKE:-lake}"
+lib="torchlean_dgemm_cuda"
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --target) [[ $# -ge 2 ]] || { echo "--target needs a value" >&2; exit 2; }; lib="$2"; shift 2 ;;
+    --target=*) lib="${1#--target=}"; shift ;;
+    -h|--help) usage; exit 0 ;;
+    *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;;
+  esac
+done
+
+repo="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+cd "$repo"
+
+work="$(mktemp -d)"
+build="$repo/.lake/build"
+
+# The stand-in writes into the normal build directory, so its objects have to go whether the
+# check passes or fails: Lake would otherwise record them as up to date and a real CUDA build
+# would archive an empty object instead of recompiling.
+cleanup() {
+  rm -f "$build/$lib.o" "$build/$lib.o.trace" "$build/lib$lib.a" "$build/lib$lib.a.trace"
+  rm -rf "$work"
+  # Lake remembers the last build configuration, and every build below sets `cuda=true`. Put the
+  # default back, or the next plain `lake build` in this checkout reaches for a CUDA toolkit that
+  # the developer running this check need not have.
+  "$LAKE" -R check-build >/dev/null 2>&1 || true
+}
+trap cleanup EXIT
+
+mkdir -p "$work/bin"
+cat > "$work/bin/nvcc" <<'STUB'
+#!/usr/bin/env bash
+# Recording stand-in for nvcc: log the argument vector, emit an empty object at -o.
+printf '%s\n' "$*" >> "${NVCC_LOG:?}"
+out=""; prev=""
+for a in "$@"; do [ "$prev" = "-o" ] && out="$a"; prev="$a"; done
+[ -n "$out" ] || { echo "recording nvcc: no -o in argument vector" >&2; exit 1; }
+printf 'int torchlean_recording_nvcc_probe;\n' | cc -x c -c -o "$out" -
+STUB
+chmod +x "$work/bin/nvcc"
+
+export PATH="$work/bin:$PATH"
+export NVCC_LOG="$work/nvcc.log"
+: > "$NVCC_LOG"
+
+# Start from no object at all, so the first build below is guaranteed to compile.
+rm -f "$build/$lib.o" "$build/$lib.o.trace" "$build/lib$lib.a" "$build/lib$lib.a.trace"
+
+failures=0
+
+# Build once and report how many compilations that build triggered.
+compilations_for() {
+  local before after
+  before="$(wc -l < "$NVCC_LOG")"
+  "$@" >/dev/null 2>&1 || { echo "   build failed: $*" >&2; return 1; }
+  after="$(wc -l < "$NVCC_LOG")"
+  echo "$((after - before))"
+}
+
+check() {
+  local name="$1" expect_compilations="$2" expect_flags="$3"; shift 3
+  local n last
+  n="$(compilations_for "$@")" || { failures=$((failures + 1)); return; }
+  last="$(tail -n 1 "$NVCC_LOG")"
+  if [[ "$n" != "$expect_compilations" ]]; then
+    echo "FAIL $name: expected $expect_compilations compilation(s), saw $n" >&2
+    failures=$((failures + 1))
+    return
+  fi
+  if [[ -n "$expect_flags" && "$last" != *"$expect_flags"* ]]; then
+    echo "FAIL $name: expected argument vector to contain '$expect_flags'" >&2
+    echo "   saw: $last" >&2
+    failures=$((failures + 1))
+    return
+  fi
+  if [[ -z "$expect_flags" && "$n" != "0" && "$last" == *"-arch"* ]]; then
+    echo "FAIL $name: expected no -arch flag" >&2
+    echo "   saw: $last" >&2
+    failures=$((failures + 1))
+    return
+  fi
+  echo "ok   $name"
+}
+
+echo "Checking the CUDA compilation target through a recording nvcc ($lib)"
+
+check "no target leaves nvcc on its default" 1 "" \
+  "$LAKE" -R -K cuda=true build "$lib"
+check "-K cuda_arch=sm_86 compiles for sm_86" 1 "-arch=sm_86" \
+  "$LAKE" -R -K cuda=true -K cuda_arch=sm_86 build "$lib"
+check "an unchanged target does not recompile" 0 "" \
+  "$LAKE" -R -K cuda=true -K cuda_arch=sm_86 build "$lib"
+check "a changed target recompiles" 1 "-arch=sm_89" \
+  "$LAKE" -R -K cuda=true -K cuda_arch=sm_89 build "$lib"
+check "TORCHLEAN_CUDA_ARCH is the fallback" 1 "-arch=sm_90" \
+  env TORCHLEAN_CUDA_ARCH=sm_90 "$LAKE" -R -K cuda=true build "$lib"
+check "the option outranks the environment" 1 "-arch=sm_75" \
+  env TORCHLEAN_CUDA_ARCH=sm_90 "$LAKE" -R -K cuda=true -K cuda_arch=sm_75 build "$lib"
+check "an explicit flag list passes through" 1 "-gencode arch=compute_120,code=[sm_120,compute_120]" \
+  "$LAKE" -R -K cuda=true \
+  -K cuda_arch="-gencode arch=compute_75,code=sm_75 -gencode arch=compute_120,code=[sm_120,compute_120]" \
+  build "$lib"
+
+if [[ "$failures" -ne 0 ]]; then
+  echo "$failures check(s) failed" >&2
+  exit 1
+fi
+echo "All CUDA compilation target checks passed."