From 419e4d1822f7a8bd747f2c9e093a2abfca87ddee Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sun, 13 Apr 2025 15:06:30 +0200
Subject: [PATCH] nixos/hardware/nvidia: eager load nvidia-uvm for open driver

Fixes contingent CUDA issues with open driver:
https://github.com/NixOS/nixpkgs/issues/334180

Co-authored-by: Someone Serge <else@someonex.net>
---
 nixos/modules/hardware/video/nvidia.nix | 28 ++++++++++++++-----------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/nixos/modules/hardware/video/nvidia.nix b/nixos/modules/hardware/video/nvidia.nix
index 068e23aa3347..48f8ded73c64 100644
--- a/nixos/modules/hardware/video/nvidia.nix
+++ b/nixos/modules/hardware/video/nvidia.nix
@@ -343,14 +343,23 @@ in
             ];
 
             # Don't add `nvidia-uvm` to `kernelModules`, because we want
-            # `nvidia-uvm` be loaded only after `udev` rules for `nvidia` kernel
-            # module are applied.
+            # `nvidia-uvm` be loaded only after the GPU device is available, i.e. after `udev` rules
+            # for `nvidia` kernel module are applied.
+            # This matters on Azure GPU instances: https://github.com/NixOS/nixpkgs/pull/267335
             #
             # Instead, we use `softdep` to lazily load `nvidia-uvm` kernel module
             # after `nvidia` kernel module is loaded and `udev` rules are applied.
             extraModprobeConfig = ''
               softdep nvidia post: nvidia-uvm
             '';
+
+            # Exception is the open-source kernel module failing to load nvidia-uvm using softdep
+            # for unknown reasons.
+            # It affects CUDA: https://github.com/NixOS/nixpkgs/issues/334180
+            # Previously nvidia-uvm was explicitly loaded only when xserver was enabled:
+            # https://github.com/NixOS/nixpkgs/pull/334340/commits/4548c392862115359e50860bcf658cfa8715bde9
+            # We are now loading the module eagerly for all users of the open driver (including headless).
+            kernelModules = lib.optionals useOpenModules [ "nvidia_uvm" ];
           };
           systemd.tmpfiles.rules = lib.mkIf config.virtualisation.docker.enableNvidia [
             "L+ /run/nvidia-docker/bin - - - - ${nvidia_x11.bin}/origBin"
@@ -639,16 +648,11 @@ in
           boot = {
             extraModulePackages = if useOpenModules then [ nvidia_x11.open ] else [ nvidia_x11.bin ];
             # nvidia-uvm is required by CUDA applications.
-            kernelModules =
-              lib.optionals config.services.xserver.enable [
-                "nvidia"
-                "nvidia_modeset"
-                "nvidia_drm"
-              ]
-              # With the open driver, nvidia-uvm does not automatically load as
-              # a softdep of the nvidia module, so we explicitly load it for now.
-              # See https://github.com/NixOS/nixpkgs/issues/334180
-              ++ lib.optionals (config.services.xserver.enable && useOpenModules) [ "nvidia_uvm" ];
+            kernelModules = lib.optionals config.services.xserver.enable [
+              "nvidia"
+              "nvidia_modeset"
+              "nvidia_drm"
+            ];
 
             # If requested enable modesetting via kernel parameters.
             kernelParams =