Repository: NixOS/nixpkgs
Pull request #408555: flashinfer: init at 0.2.5
Author: login: breakds
Labels: name: 6.topic: python
description: Python is a high-level, general-purpose programming language.
name: 10.rebuild-linux: 1-10
description: This PR causes between 1 and 10 packages to rebuild on Linux.
name: 10.rebuild-darwin: 1-10
description: This PR causes between 1 and 10 packages to rebuild on Darwin.
name: 6.topic: cuda
description: Parallel computing platform and API
Created: 2025-05-18T23:39:01Z | Merged: 2025-05-21T15:25:25Z

<!--
^ Please summarise the changes you have done and explain why they are necessary here ^

For package updates please link to a changelog or describe changes, this helps your fellow maintainers discover breaking updates.
For new packages please briefly describe the package or provide a link to its homepage.
-->

FlashInfer is an open-source library and kernel generator for large language model (LLM) inference, offering high-performance GPU kernels such as FlashAttention, SparseAttention, and PageAttention, with both just-in-time (JIT) and ahead-of-time (AOT) compilation modes, and is integrated into major serving frameworks including [SGLang](https://github.com/sgl-project/sglang) and [vLLM](https://github.com/vllm-project/vllm).

I am packaging it as part of the effort to eventually package `sglang`.

NOTE: I’m new to packaging `CUDA`-enabled Python libraries in nixpkgs and I’d greatly appreciate any guidance on best practices of this. Thanks!

## Things done

I have verified locally that the built package can be used to run a simple `flashinfer` program, given that I have specified `CUDA_HOME` in my devshell.

<!-- Please check what applies. Note that these are not hard requirements but merely serve as information for reviewers. -->

- Built on platform(s)
  - [x] x86_64-linux
  - [ ] aarch64-linux
  - [ ] x86_64-darwin
  - [ ] aarch64-darwin
- For non-Linux: Is sandboxing enabled in `nix.conf`? (See [Nix manual](https://nixos.org/manual/nix/stable/command-ref/conf-file.html))
  - [ ] `sandbox = relaxed`
  - [ ] `sandbox = true`
- [ ] Tested, as applicable:
  - [NixOS test(s)](https://nixos.org/manual/nixos/unstable/index.html#sec-nixos-tests) (look inside [nixos/tests](https://github.com/NixOS/nixpkgs/blob/master/nixos/tests))
  - and/or [package tests](https://github.com/NixOS/nixpkgs/blob/master/pkgs/README.md#package-tests)
  - or, for functions and "core" functionality, tests in [lib/tests](https://github.com/NixOS/nixpkgs/blob/master/lib/tests) or [pkgs/test](https://github.com/NixOS/nixpkgs/blob/master/pkgs/test)
  - made sure NixOS tests are [linked](https://github.com/NixOS/nixpkgs/blob/master/pkgs/README.md#linking-nixos-module-tests-to-a-package) to the relevant packages
- [ ] Tested compilation of all packages that depend on this change using `nix-shell -p nixpkgs-review --run "nixpkgs-review rev HEAD"`. Note: all changes have to be committed, also see [nixpkgs-review usage](https://github.com/Mic92/nixpkgs-review#usage)
- [ ] Tested basic functionality of all binary files (usually in `./result/bin/`)
- [Nixpkgs 25.11 Release Notes](https://github.com/NixOS/nixpkgs/blob/master/doc/release-notes/rl-2511.section.md) (or backporting [24.11](https://github.com/NixOS/nixpkgs/blob/master/nixos/doc/release-notes/rl-2411.section.md) and [25.05](https://github.com/NixOS/nixpkgs/blob/master/doc/manual/release-notes/rl-2505.section.md) Nixpkgs Release notes)
  - [ ] (Package updates) Added a release notes entry if the change is major or breaking
- [NixOS 25.11 Release Notes](https://github.com/NixOS/nixpkgs/blob/master/nixos/doc/manual/release-notes/rl-2511.section.md) (or backporting [24.11](https://github.com/NixOS/nixpkgs/blob/master/nixos/doc/manual/release-notes/rl-2411.section.md) and [25.05](https://github.com/NixOS/nixpkgs/blob/master/nixos/doc/manual/release-notes/rl-2505.section.md) NixOS Release notes)
  - [ ] (Module updates) Added a release notes entry if the change is significant
  - [ ] (Module addition) Added a release notes entry if adding a new NixOS module
- [x] Fits [CONTRIBUTING.md](https://github.com/NixOS/nixpkgs/blob/master/CONTRIBUTING.md).

<!--
To help with the large amounts of pull requests, we would appreciate your
reviews of other pull requests, especially simple package updates. Just leave a
comment describing what you have tested in the relevant package/service.
Reviewing helps to reduce the average time-to-merge for everyone.
Thanks a lot if you do!

List of open PRs: https://github.com/NixOS/nixpkgs/pulls
Reviewing guidelines: https://github.com/NixOS/nixpkgs/blob/master/pkgs/README.md#reviewing-contributions
-->

---

Add a :+1: [reaction] to [pull requests you find important].

[reaction]: https://github.blog/2016-03-10-add-reactions-to-pull-requests-issues-and-comments/
[pull requests you find important]: https://github.com/NixOS/nixpkgs/pulls?q=is%3Aopen+sort%3Areactions-%2B1-desc



Diff:
diff --git a/pkgs/development/python-modules/flashinfer/default.nix b/pkgs/development/python-modules/flashinfer/default.nix
new file mode 100644
index 0000000000000..1600ab2c38a41
--- /dev/null
+++ b/pkgs/development/python-modules/flashinfer/default.nix
@@ -0,0 +1,103 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  config,
+  buildPythonPackage,
+  fetchFromGitHub,
+  setuptools,
+  cudaPackages,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  src_cutlass = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cutlass";
+    # Using the revision obtained in submodule inside flashinfer's `3rdparty`.
+    rev = "df8a550d3917b0e97f416b2ed8c2d786f7f686a3";
+    hash = "sha256-d4czDoEv0Focf1bJHOVGX4BDS/h5O7RPoM/RrujhgFQ=";
+  };
+
+in
+buildPythonPackage {
+  inherit pname version;
+
+  src = fetchFromGitHub {
+    owner = "flashinfer-ai";
+    repo = "flashinfer";
+    tag = "v${version}";
+    hash = "sha256-YrYfatkI9DQkFEEGiF8CK/bTafaNga4Ufyt+882C0bQ=";
+  };
+
+  build-system = [ setuptools ];
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    (lib.getBin cudaPackages.cuda_nvcc)
+  ];
+  dontUseCmakeConfigure = true;
+
+  buildInputs = [
+    cudaPackages.cuda_cudart
+    cudaPackages.libcublas
+    cudaPackages.cuda_cccl
+    cudaPackages.libcurand
+  ];
+
+  postPatch = ''
+    rmdir 3rdparty/cutlass
+    ln -s ${src_cutlass} 3rdparty/cutlass
+  '';
+
+  # FlashInfer offers two installation modes:
+  #
+  # JIT mode: CUDA kernels are compiled at runtime using PyTorch’s JIT, with
+  # compiled kernels cached for future use. JIT mode allows fast installation,
+  # as no CUDA kernels are pre-compiled, making it ideal for development and
+  # testing. JIT version is also available as a sdist in PyPI.
+  #
+  # AOT mode: Core CUDA kernels are pre-compiled and included in the library,
+  # reducing runtime compilation overhead. If a required kernel is not
+  # pre-compiled, it will be compiled at runtime using JIT. AOT mode is
+  # recommended for production environments.
+  #
+  # Here we use opt for the AOT version.
+  preConfigure = ''
+    export FLASHINFER_ENABLE_AOT=1
+    export TORCH_NVCC_FLAGS="--maxrregcount=64"
+  '';
+
+  TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;
+
+  dependencies = [
+    numpy
+    torch
+  ];
+
+  meta = with lib; {
+    broken = !torch.cudaSupport || !config.cudaSupport;
+    homepage = "https://flashinfer.ai/";
+    description = "Library and kernel generator for Large Language Models";
+    longDescription = ''
+      FlashInfer is a library and kernel generator for Large Language Models
+      that provides high-performance implementation of LLM GPU kernels such as
+      FlashAttention, PageAttention and LoRA. FlashInfer focus on LLM serving
+      and inference, and delivers state-of-the-art performance across diverse
+      scenarios.
+    '';
+    license = licenses.asl20;
+    maintainers = with maintainers; [ breakds ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 144581c5e046e..5ac6fe4de7a42 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -5015,6 +5015,8 @@ self: super: with self; {
 
   flasgger = callPackage ../development/python-modules/flasgger { };
 
+  flashinfer = callPackage ../development/python-modules/flashinfer { };
+
   flashtext = callPackage ../development/python-modules/flashtext { };
 
   flask = callPackage ../development/python-modules/flask { };


Reviews:
user: login: ConnorBaker
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T18:02:45Z
user: login: breakds
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T21:37:14Z
user: login: breakds
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T21:47:29Z
user: login: ConnorBaker
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T15:24:07Z

Review comments:
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  inherit (torch) cudaPackages;
+  inherit (cudaPackages) cudaMajorMinorVersion;
+
+  cudaMajorMinorVersionString = lib.replaceStrings [ "." ] [ "" ] cudaMajorMinorVersion;
+
+  src_cutlass = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cutlass";
+    # Using the revision obtained in submodule inside flashinfer's `3rdparty`.
+    rev = "df8a550d3917b0e97f416b2ed8c2d786f7f686a3";
+    hash = "sha256-d4czDoEv0Focf1bJHOVGX4BDS/h5O7RPoM/RrujhgFQ=";
+  };
+
+in
+buildPythonPackage {
+  inherit pname version;
+
+  src = fetchFromGitHub {
+    owner = "flashinfer-ai";
+    repo = "flashinfer";
+    tag = "v${version}";
+    hash = "sha256-YrYfatkI9DQkFEEGiF8CK/bTafaNga4Ufyt+882C0bQ=";
+  };
+
+  build-system = [ setuptools ];
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    cudaPackages.cudatoolkit
+  ];
+  dontUseCmakeConfigure = true;
+
+  postPatch = ''
+    rmdir 3rdparty/cutlass
+    ln -s ${src_cutlass} 3rdparty/cutlass
+  '';
+
+  # FlashInfer offers two installation modes:
+  #
+  # JIT mode: CUDA kernels are compiled at runtime using PyTorch’s JIT, with
+  # compiled kernels cached for future use. JIT mode allows fast installation,
+  # as no CUDA kernels are pre-compiled, making it ideal for development and
+  # testing. JIT version is also available as a sdist in PyPI.
+  #
+  # AOT mode: Core CUDA kernels are pre-compiled and included in the library,
+  # reducing runtime compilation overhead. If a required kernel is not
+  # pre-compiled, it will be compiled at runtime using JIT. AOT mode is
+  # recommended for production environments.
+  #
+  # Here we use opt for the AOT version.
+  preConfigure = ''
+    export FLASHINFER_ENABLE_AOT=1
+    export TORCH_NVCC_FLAGS="--maxrregcount=64"
+  '';
+
+  CUDA_HOME = "${cudaPackages.cudatoolkit}";
+  TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
+
+  dependencies = [
+    numpy
+    torch
+  ];
+
+  meta = with lib; {
+    homepage = "https://flashinfer.ai/";
+    description = ''
+      ;
+            FlashInfer is a library and kernel generator for Large Language Models
+            that provides high-performance implementation of LLM GPU kernels such as
+            FlashAttention, PageAttention and LoRA. FlashInfer focus on LLM serving
+            and inference, and delivers state-of-the-art performance across diverse
+            scenarios.
+    '';
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: ConnorBaker
body: See https://ryantm.github.io/nixpkgs/stdenv/meta/#var-meta-description for descriptions of the various fields and what they should contain.

```suggestion
    description = "Library and kernel generator for Large Language Models";
    longDescription = ''
      FlashInfer is a library and kernel generator for Large Language Models
      that provides high-performance implementation of LLM GPU kernels such as
      FlashAttention, PageAttention and LoRA. FlashInfer focus on LLM serving
      and inference, and delivers state-of-the-art performance across diverse
      scenarios.
    '';
```
created_at: 2025-05-20T17:55:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: ConnorBaker
body: Don't use asserts -- instead, use `meta.broken`.
created_at: 2025-05-20T17:55:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  inherit (torch) cudaPackages;
+  inherit (cudaPackages) cudaMajorMinorVersion;
+
+  cudaMajorMinorVersionString = lib.replaceStrings [ "." ] [ "" ] cudaMajorMinorVersion;
+
+  src_cutlass = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cutlass";
+    # Using the revision obtained in submodule inside flashinfer's `3rdparty`.
+    rev = "df8a550d3917b0e97f416b2ed8c2d786f7f686a3";
+    hash = "sha256-d4czDoEv0Focf1bJHOVGX4BDS/h5O7RPoM/RrujhgFQ=";
+  };
+
+in
+buildPythonPackage {
+  inherit pname version;
+
+  src = fetchFromGitHub {
+    owner = "flashinfer-ai";
+    repo = "flashinfer";
+    tag = "v${version}";
+    hash = "sha256-YrYfatkI9DQkFEEGiF8CK/bTafaNga4Ufyt+882C0bQ=";
+  };
+
+  build-system = [ setuptools ];
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    cudaPackages.cudatoolkit
+  ];
+  dontUseCmakeConfigure = true;
+
+  postPatch = ''
+    rmdir 3rdparty/cutlass
+    ln -s ${src_cutlass} 3rdparty/cutlass
+  '';
+
+  # FlashInfer offers two installation modes:
+  #
+  # JIT mode: CUDA kernels are compiled at runtime using PyTorch’s JIT, with
+  # compiled kernels cached for future use. JIT mode allows fast installation,
+  # as no CUDA kernels are pre-compiled, making it ideal for development and
+  # testing. JIT version is also available as a sdist in PyPI.
+  #
+  # AOT mode: Core CUDA kernels are pre-compiled and included in the library,
+  # reducing runtime compilation overhead. If a required kernel is not
+  # pre-compiled, it will be compiled at runtime using JIT. AOT mode is
+  # recommended for production environments.
+  #
+  # Here we use opt for the AOT version.
+  preConfigure = ''
+    export FLASHINFER_ENABLE_AOT=1
+    export TORCH_NVCC_FLAGS="--maxrregcount=64"
+  '';
+
+  CUDA_HOME = "${cudaPackages.cudatoolkit}";
+  TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
+
+  dependencies = [
+    numpy
+    torch
+  ];
+
+  meta = with lib; {
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: ConnorBaker
body: You should also take `config` as an argument.

```suggestion
  meta = with lib; {
    broken = !torch.cudaSupport || !config.cudaSupport;
```
created_at: 2025-05-20T17:56:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  inherit (torch) cudaPackages;
+  inherit (cudaPackages) cudaMajorMinorVersion;
+
+  cudaMajorMinorVersionString = lib.replaceStrings [ "." ] [ "" ] cudaMajorMinorVersion;
+
+  src_cutlass = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cutlass";
+    # Using the revision obtained in submodule inside flashinfer's `3rdparty`.
+    rev = "df8a550d3917b0e97f416b2ed8c2d786f7f686a3";
+    hash = "sha256-d4czDoEv0Focf1bJHOVGX4BDS/h5O7RPoM/RrujhgFQ=";
+  };
+
+in
+buildPythonPackage {
+  inherit pname version;
+
+  src = fetchFromGitHub {
+    owner = "flashinfer-ai";
+    repo = "flashinfer";
+    tag = "v${version}";
+    hash = "sha256-YrYfatkI9DQkFEEGiF8CK/bTafaNga4Ufyt+882C0bQ=";
+  };
+
+  build-system = [ setuptools ];
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    cudaPackages.cudatoolkit
+  ];
+  dontUseCmakeConfigure = true;
+
+  postPatch = ''
+    rmdir 3rdparty/cutlass
+    ln -s ${src_cutlass} 3rdparty/cutlass
+  '';
+
+  # FlashInfer offers two installation modes:
+  #
+  # JIT mode: CUDA kernels are compiled at runtime using PyTorch’s JIT, with
+  # compiled kernels cached for future use. JIT mode allows fast installation,
+  # as no CUDA kernels are pre-compiled, making it ideal for development and
+  # testing. JIT version is also available as a sdist in PyPI.
+  #
+  # AOT mode: Core CUDA kernels are pre-compiled and included in the library,
+  # reducing runtime compilation overhead. If a required kernel is not
+  # pre-compiled, it will be compiled at runtime using JIT. AOT mode is
+  # recommended for production environments.
+  #
+  # Here we use opt for the AOT version.
+  preConfigure = ''
+    export FLASHINFER_ENABLE_AOT=1
+    export TORCH_NVCC_FLAGS="--maxrregcount=64"
+  '';
+
+  CUDA_HOME = "${cudaPackages.cudatoolkit}";
+  TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: ConnorBaker
body: The result is already a string

```suggestion
  TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;
```
created_at: 2025-05-20T17:56:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  inherit (torch) cudaPackages;
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: ConnorBaker
body: I'd recommend taking `cudaPackages` as a `callPackage`-provided argument instead; I think the pattern of packages providing references to their inputs is a bad one.
created_at: 2025-05-20T17:58:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  inherit (torch) cudaPackages;
+  inherit (cudaPackages) cudaMajorMinorVersion;
+
+  cudaMajorMinorVersionString = lib.replaceStrings [ "." ] [ "" ] cudaMajorMinorVersion;
+
+  src_cutlass = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cutlass";
+    # Using the revision obtained in submodule inside flashinfer's `3rdparty`.
+    rev = "df8a550d3917b0e97f416b2ed8c2d786f7f686a3";
+    hash = "sha256-d4czDoEv0Focf1bJHOVGX4BDS/h5O7RPoM/RrujhgFQ=";
+  };
+
+in
+buildPythonPackage {
+  inherit pname version;
+
+  src = fetchFromGitHub {
+    owner = "flashinfer-ai";
+    repo = "flashinfer";
+    tag = "v${version}";
+    hash = "sha256-YrYfatkI9DQkFEEGiF8CK/bTafaNga4Ufyt+882C0bQ=";
+  };
+
+  build-system = [ setuptools ];
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    cudaPackages.cudatoolkit
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: ConnorBaker
body: Please don't use `cudaPackages.cudatoolkit`, it's a giant `symlinkJoin` of a bunch of CUDA libraries. It's much better to specify the libraries the package actually needs.

```suggestion
    cudaPackages.cuda_nvcc
```

You'll likely need `cudaPackages.cuda_cudart` in `buildInputs`, as well.
created_at: 2025-05-20T18:01:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  inherit (torch) cudaPackages;
+  inherit (cudaPackages) cudaMajorMinorVersion;
+
+  cudaMajorMinorVersionString = lib.replaceStrings [ "." ] [ "" ] cudaMajorMinorVersion;
+
+  src_cutlass = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cutlass";
+    # Using the revision obtained in submodule inside flashinfer's `3rdparty`.
+    rev = "df8a550d3917b0e97f416b2ed8c2d786f7f686a3";
+    hash = "sha256-d4czDoEv0Focf1bJHOVGX4BDS/h5O7RPoM/RrujhgFQ=";
+  };
+
+in
+buildPythonPackage {
+  inherit pname version;
+
+  src = fetchFromGitHub {
+    owner = "flashinfer-ai";
+    repo = "flashinfer";
+    tag = "v${version}";
+    hash = "sha256-YrYfatkI9DQkFEEGiF8CK/bTafaNga4Ufyt+882C0bQ=";
+  };
+
+  build-system = [ setuptools ];
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    cudaPackages.cudatoolkit
+  ];
+  dontUseCmakeConfigure = true;
+
+  postPatch = ''
+    rmdir 3rdparty/cutlass
+    ln -s ${src_cutlass} 3rdparty/cutlass
+  '';
+
+  # FlashInfer offers two installation modes:
+  #
+  # JIT mode: CUDA kernels are compiled at runtime using PyTorch’s JIT, with
+  # compiled kernels cached for future use. JIT mode allows fast installation,
+  # as no CUDA kernels are pre-compiled, making it ideal for development and
+  # testing. JIT version is also available as a sdist in PyPI.
+  #
+  # AOT mode: Core CUDA kernels are pre-compiled and included in the library,
+  # reducing runtime compilation overhead. If a required kernel is not
+  # pre-compiled, it will be compiled at runtime using JIT. AOT mode is
+  # recommended for production environments.
+  #
+  # Here we use opt for the AOT version.
+  preConfigure = ''
+    export FLASHINFER_ENABLE_AOT=1
+    export TORCH_NVCC_FLAGS="--maxrregcount=64"
+  '';
+
+  CUDA_HOME = "${cudaPackages.cudatoolkit}";
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: ConnorBaker
body: Shouldn't need this, or at the least, should be able to set it to

```nix
CUDA_HOME = lib.getBin cudaPackages.cuda_nvcc
```

Typically packages just need CUDA_HOME to contain `bin/nvcc`.
created_at: 2025-05-20T18:02:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  inherit (torch) cudaPackages;
+  inherit (cudaPackages) cudaMajorMinorVersion;
+
+  cudaMajorMinorVersionString = lib.replaceStrings [ "." ] [ "" ] cudaMajorMinorVersion;
+
+  src_cutlass = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cutlass";
+    # Using the revision obtained in submodule inside flashinfer's `3rdparty`.
+    rev = "df8a550d3917b0e97f416b2ed8c2d786f7f686a3";
+    hash = "sha256-d4czDoEv0Focf1bJHOVGX4BDS/h5O7RPoM/RrujhgFQ=";
+  };
+
+in
+buildPythonPackage {
+  inherit pname version;
+
+  src = fetchFromGitHub {
+    owner = "flashinfer-ai";
+    repo = "flashinfer";
+    tag = "v${version}";
+    hash = "sha256-YrYfatkI9DQkFEEGiF8CK/bTafaNga4Ufyt+882C0bQ=";
+  };
+
+  build-system = [ setuptools ];
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    cudaPackages.cudatoolkit
+  ];
+  dontUseCmakeConfigure = true;
+
+  postPatch = ''
+    rmdir 3rdparty/cutlass
+    ln -s ${src_cutlass} 3rdparty/cutlass
+  '';
+
+  # FlashInfer offers two installation modes:
+  #
+  # JIT mode: CUDA kernels are compiled at runtime using PyTorch’s JIT, with
+  # compiled kernels cached for future use. JIT mode allows fast installation,
+  # as no CUDA kernels are pre-compiled, making it ideal for development and
+  # testing. JIT version is also available as a sdist in PyPI.
+  #
+  # AOT mode: Core CUDA kernels are pre-compiled and included in the library,
+  # reducing runtime compilation overhead. If a required kernel is not
+  # pre-compiled, it will be compiled at runtime using JIT. AOT mode is
+  # recommended for production environments.
+  #
+  # Here we use opt for the AOT version.
+  preConfigure = ''
+    export FLASHINFER_ENABLE_AOT=1
+    export TORCH_NVCC_FLAGS="--maxrregcount=64"
+  '';
+
+  CUDA_HOME = "${cudaPackages.cudatoolkit}";
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: breakds
body: I added `CUDA_HOME` because I find this in their `setup.py`: https://github.com/flashinfer-ai/flashinfer/blob/main/setup.py#L151-L154

Now reading it again, I think as long as `nvcc` is in the `PATH`, everything should work. I'll try just adding `cuda_nvcc` to `nativeBuildInputs`.
created_at: 2025-05-20T21:37:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,102 @@
+# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a
+# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always
+# requires the CUDA toolkit (via nvcc) to be available.
+#
+# This means that if you plan to use flashinfer, you will need to set the
+# environment varaible `CUDA_HOME` to `cudatoolkit`.
+{
+  lib,
+  buildPythonPackage,
+  symlinkJoin,
+  fetchFromGitHub,
+  setuptools,
+  cmake,
+  ninja,
+  numpy,
+  torch,
+}:
+
+assert torch.cudaSupport;
+
+let
+  pname = "flashinfer";
+  version = "0.2.5";
+
+  inherit (torch) cudaPackages;
path: pkgs/development/python-modules/flashinfer/default.nix
user: login: breakds
body: Sounds good. Just to confess that I did it this way because

1. I saw other packages doing this (not an execuse though)
2. I already had `torch` built locally and I think in this way I can safely avoid triggering the rebuild of the CUDA family packages, which usually takes a lot of time. I find it strange that sometimes CUDA packages gets rebuild even if I think it shouldn't ...
created_at: 2025-05-20T21:47:29Z
author_association: CONTRIBUTOR

PR comments:
user: login: breakds
created_at: 2025-05-18T23:39:27Z
body: `nixpkgs-review` is running. Will update once it is finished.
author_association: CONTRIBUTOR
user: login: breakds
created_at: 2025-05-18T23:41:28Z
body: > `nixpkgs-review` is running. Will update once it is finished.

Actually, since nothing is depending on this new package, I am not sure whether `nixpkgs-review` is relevant.
author_association: CONTRIBUTOR
user: login: breakds
created_at: 2025-05-19T00:05:47Z
body: > > `nixpkgs-review` is running. Will update once it is finished.
> 
> Actually, since nothing is depending on this new package, I am not sure whether `nixpkgs-review` is relevant.

```console
--------- Report for 'x86_64-linux' ---------
4 packages marked as broken and skipped:
python312Packages.bentoml python312Packages.bentoml.dist python313Packages.bentoml python313Packages.bentoml.dist

2 packages failed to build:
xpraWithNvenc xpraWithNvenc.dist
```

I think the 2 failures are not related to this PR.
author_association: CONTRIBUTOR
user: login: breakds
created_at: 2025-05-20T18:07:00Z
body: @ConnorBaker Thanks a lot for the detailed review 🙇🏼 . Working on addressing them and re-testing. 
author_association: CONTRIBUTOR
user: login: ConnorBaker
created_at: 2025-05-20T18:07:39Z
body: Overall you're on the right track!

A note on building CUDA-enabled packages with Nixpkgs: CUDA support isn't enabled by default. In order to configure Nixpkgs to build packages with CUDA, make sure that you've got something like

```nix
{
  allowUnfree = true;
  cudaSupport = true;
  cudaCapabilities = [ "8.9" ];
}
```

in `~/.config/nixpkgs/config.nix` and that you're building with `--impure` (so Nix can find and use that file to provide arguments when importing Nixpkgs). You should change `cudaCapabilities` to whichever GPUs you want to target -- just be aware that the more of them you specify, the longer builds take (and the larger the resulting binaries are); you can find a mapping here: https://developer.nvidia.com/cuda-gpus. Above, 8.9 corresponds to a the Ada architecture.

Likewise for `nixpkgs-review` to build CUDA packages, you'll need to use pass extra arguments like `--extra-nixpkgs-config "$(<~/.config/nixpkgs/config.nix)"` (assuming you use Bash).
author_association: CONTRIBUTOR
user: login: breakds
created_at: 2025-05-20T21:43:28Z
body: > A note on building CUDA-enabled packages with Nixpkgs: CUDA support isn't enabled by default. In order to configure Nixpkgs to build packages with CUDA, make sure that you've got something like

Thanks - TIL!

I actually tested this package out in my package [here](https://github.com/nixvital/ml-pkgs/blob/main/flake.nix#L34), where I had 

```nix
_module.args.pkgs = import inputs.nixpkgs {
          inherit system;
          config = {
            allowUnfree = true;
            cudaSupport = true;
            cudaForwardCompat = true;
            # NOTE: This is mainly to rule out `sm_90a` from showing up in
            # `torch.cuda.get_arch_list`, which will trigger
            # https://github.com/pytorch/pytorch/issues/144037
            cudaCapabilities = [ "7.5" "8.6" "8.9" ];
          };
          .....
        };
``` 

I was not aware of `~/.config/nixpkgs/config.nix` as I am used to put it behind a flake where I had control on the `pkgs` instantiation.
author_association: CONTRIBUTOR
user: login: breakds
created_at: 2025-05-20T23:31:06Z
body: I think I have all the comments addressed. PTAL. Thanks for the patient and detailed review, @ConnorBaker.
author_association: CONTRIBUTOR
user: login: breakds
created_at: 2025-05-20T23:31:58Z
body: I have tested locally with cudaCapability `8.9` in my `~/.config/nixpkgs/config.nix`, and `nix-build -A python3Packages.flashinfer`.
author_association: CONTRIBUTOR
user: login: breakds
created_at: 2025-05-21T15:26:24Z
body: Thanks a lot for the review and teaching!
author_association: CONTRIBUTOR