From e4d4b7fa70589f205bf26889967cb8ea3ae04a95 Mon Sep 17 00:00:00 2001 From: Elliot Berman Date: Fri, 5 Jun 2026 13:18:34 -0700 Subject: [PATCH] cuda-modules: Allow SBSA Orin to use CUDA 13+ NVIDIA announced an SBSA software stack for the Jetson Orin (capability 8.7) using CUDA 13.2 [1]. Update the redist system selection logic and assertions to support Orin SBSA. [1]: https://forums.developer.nvidia.com/t/jetpack-7-2-jetson-linux-r39-2-on-jetson-agx-orin-developer-kit-getting-started-and-feedback-thread/372156 --- .../cuda-modules/_cuda/lib/default.nix | 1 + .../cuda-modules/_cuda/lib/redist.nix | 34 ++++++++++++++++--- .../cuda-modules/backendStdenv/default.nix | 24 +++++++------ 3 files changed, 43 insertions(+), 16 deletions(-) diff --git a/pkgs/development/cuda-modules/_cuda/lib/default.nix b/pkgs/development/cuda-modules/_cuda/lib/default.nix index 376abcd27f7bd..53ee8a1247a3e 100644 --- a/pkgs/development/cuda-modules/_cuda/lib/default.nix +++ b/pkgs/development/cuda-modules/_cuda/lib/default.nix @@ -29,6 +29,7 @@ # See ./redist.nix for documentation. inherit (import ./redist.nix { inherit _cuda lib; }) + _getJetsonMinSbsaCapability _redistSystemIsSupported getNixSystems getRedistSystem diff --git a/pkgs/development/cuda-modules/_cuda/lib/redist.nix b/pkgs/development/cuda-modules/_cuda/lib/redist.nix index 3138097c16948..bee8f9fc16b98 100644 --- a/pkgs/development/cuda-modules/_cuda/lib/redist.nix +++ b/pkgs/development/cuda-modules/_cuda/lib/redist.nix @@ -108,13 +108,35 @@ else [ ]; + /** + The lowest Jetson CUDA capability which uses the `linux-sbsa` redist (rather than `linux-aarch64`) for a given + CUDA version. Jetson capabilities below this threshold use `linux-aarch64`. + + The threshold depends on the CUDA version because NVIDIA moved Jetson onto the SBSA software stack incrementally: + - CUDA 12 (JetPack 6): only Thor (10.1) is SBSA; Orin (8.7) is `linux-aarch64`. + - CUDA 13 (JetPack 7.2): Orin (8.7) and Thor (11.0) are both SBSA. + + # Type + + ``` + _getJetsonMinSbsaCapability :: (cudaMajorMinorVersion :: String) -> CudaCapability + ``` + + # Inputs + + `cudaMajorMinorVersion` + + : The major and minor version of CUDA (e.g. "12.6") + */ + _getJetsonMinSbsaCapability = + cudaMajorMinorVersion: if lib.versionAtLeast cudaMajorMinorVersion "13.0" then "8.7" else "10.1"; + /** Maps a Nix system to a NVIDIA redistributable system. NOTE: Certain Nix systems can map to multiple NVIDIA redistributable systems. In particular, ARM systems can map to either `linux-sbsa` (for server-grade ARM chips) or `linux-aarch64` (for Jetson devices). Complicating matters - further, as of CUDA 13.0, Jetson Thor devices use `linux-sbsa` instead of `linux-aarch64`. (It is unknown whether - NVIDIA plans to make the Orin series use `linux-sbsa` as well for the CUDA 13.0 release.) + further, as of CUDA 13.0, Jetson Thor and Orin devices use `linux-sbsa` instead of `linux-aarch64`. NOTE: This function *will* be called by unsupported systems because `cudaPackages` is evaluated on all systems. As such, we need to handle unsupported systems gracefully. @@ -189,11 +211,13 @@ if system == "x86_64-linux" then "linux-x86_64" else if system == "aarch64-linux" then - # If all the Jetson devices are at least 10.1 (Thor, CUDA 12.9; CUDA 13.0 and later use 11.0 for Thor), then - # we've got SBSA. + # If all the Jetson devices are SBSA-compatible, then we've got SBSA. if + let + sbsaJetsonCapability = _cuda.lib._getJetsonMinSbsaCapability cudaMajorMinorVersion; + in lib.all ( - cap: _cuda.db.cudaCapabilityToInfo.${cap}.isJetson -> lib.versionAtLeast cap "10.1" + cap: _cuda.db.cudaCapabilityToInfo.${cap}.isJetson -> lib.versionAtLeast cap sbsaJetsonCapability ) cudaCapabilities then "linux-sbsa" diff --git a/pkgs/development/cuda-modules/backendStdenv/default.nix b/pkgs/development/cuda-modules/backendStdenv/default.nix index 762e0901f88b6..b0358fa442e3a 100644 --- a/pkgs/development/cuda-modules/backendStdenv/default.nix +++ b/pkgs/development/cuda-modules/backendStdenv/default.nix @@ -24,6 +24,7 @@ let inherit (_cuda.lib) _cudaCapabilityIsDefault _cudaCapabilityIsSupported + _getJetsonMinSbsaCapability _mkFailedAssertionsString getRedistSystem mkVersionedName @@ -121,11 +122,12 @@ let assertions = let - # Jetson devices (pre-Thor) cannot be targeted by the same binaries which target non-Jetson devices. While + # Jetson devices (pre-Orin) cannot be targeted by the same binaries which target non-Jetson devices. While # NVIDIA provides both `linux-aarch64` and `linux-sbsa` packages, which both target `aarch64`, # they are built with different settings and cannot be mixed. - preThorJetsonCudaCapabilities = filter (flip versionOlder "10.1") passthruExtra.requestedJetsonCudaCapabilities; - postThorJetsonCudaCapabilities = filter (flip versionAtLeast "10.1") passthruExtra.requestedJetsonCudaCapabilities; + sbsaJetsonCapability = _getJetsonMinSbsaCapability cudaMajorMinorVersion; + preSbsaJetsonCudaCapabilities = filter (flip versionOlder sbsaJetsonCapability) passthruExtra.requestedJetsonCudaCapabilities; + postSbsaJetsonCudaCapabilities = filter (flip versionAtLeast sbsaJetsonCapability) passthruExtra.requestedJetsonCudaCapabilities; # Remove all known capabilities from the user's list to find unrecognized capabilities. unrecognizedCudaCapabilities = subtractLists allSortedCudaCapabilities passthruExtra.cudaCapabilities; @@ -166,25 +168,25 @@ let } { message = - "Requested pre-Thor (10.1) Jetson CUDA capabilities (${toJSON preThorJetsonCudaCapabilities}) cannot be " - + "specified with other capabilities (${toJSON (subtractLists preThorJetsonCudaCapabilities passthruExtra.cudaCapabilities)})"; + "Requested pre-SBSA (${sbsaJetsonCapability}) Jetson CUDA capabilities (${toJSON preSbsaJetsonCudaCapabilities}) cannot be " + + "specified with other capabilities (${toJSON (subtractLists preSbsaJetsonCudaCapabilities passthruExtra.cudaCapabilities)})"; assertion = # If there are preThorJetsonCudaCapabilities, they must be the only requested capabilities. - preThorJetsonCudaCapabilities != [ ] - -> preThorJetsonCudaCapabilities == passthruExtra.cudaCapabilities; + preSbsaJetsonCudaCapabilities != [ ] + -> preSbsaJetsonCudaCapabilities == passthruExtra.cudaCapabilities; } { message = - "Requested pre-Thor (10.1) Jetson CUDA capabilities (${toJSON preThorJetsonCudaCapabilities}) require " + "Requested pre-SBSA (${sbsaJetsonCapability}) Jetson CUDA capabilities (${toJSON preSbsaJetsonCudaCapabilities}) require " + "computed NVIDIA hostRedistSystem (${passthruExtra.hostRedistSystem}) to be linux-aarch64"; assertion = - preThorJetsonCudaCapabilities != [ ] -> passthruExtra.hostRedistSystem == "linux-aarch64"; + preSbsaJetsonCudaCapabilities != [ ] -> passthruExtra.hostRedistSystem == "linux-aarch64"; } { message = - "Requested post-Thor (10.1) Jetson CUDA capabilities (${toJSON postThorJetsonCudaCapabilities}) require " + "Requested post-SBSA (${sbsaJetsonCapability}) Jetson CUDA capabilities (${toJSON postSbsaJetsonCudaCapabilities}) require " + "computed NVIDIA hostRedistSystem (${passthruExtra.hostRedistSystem}) to be linux-sbsa"; - assertion = postThorJetsonCudaCapabilities != [ ] -> passthruExtra.hostRedistSystem == "linux-sbsa"; + assertion = postSbsaJetsonCudaCapabilities != [ ] -> passthruExtra.hostRedistSystem == "linux-sbsa"; } ];