diff --git a/pkgs/development/python-modules/flashinfer/default.nix b/pkgs/development/python-modules/flashinfer/default.nix
index 5882d0ee8164..2e1119edc390 100644
--- a/pkgs/development/python-modules/flashinfer/default.nix
+++ b/pkgs/development/python-modules/flashinfer/default.nix
@@ -87,7 +87,7 @@ buildPythonPackage rec {
     export MAX_JOBS="$NIX_BUILD_CORES"
   '';
 
-  FLASHINFER_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;
+  env.FLASHINFER_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;
 
   pythonRemoveDeps = [
     "nvidia-cudnn-frontend"