flake.nix

{
  description = "Flake for Hopper Flash Attention kernel";

  inputs = {
    kernel-builder.url = "github:huggingface/kernel-builder";
  };

  outputs =
    {
      self,
      kernel-builder,
    }:
    kernel-builder.lib.genFlakeOutputs {
      path = ./.;
      rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
      # Building with CDUA later than 12.4 fails with:
      #
      # error: 'ptxas' died due to signal 11 (Invalid memory reference)
      #
      # So, build for 12.4 only and copy to all the other build variants
      # by hand (which works fine thanks to backward compat).
      torchVersions = [
        {
          torchVersion = "2.6";
          cudaVersion = "12.4";
          cxx11Abi = false;
          systems = [ "x86_64-linux" ];
          upstreamVariant = true;
        }
        {
          torchVersion = "2.6";
          cudaVersion = "12.4";
          cxx11Abi = true;
          systems = [ "x86_64-linux" ];
          upstreamVariant = true;
        }
        {
          torchVersion = "2.7";
          cudaVersion = "12.4";
          cxx11Abi = true;
          systems = [
            "x86_64-linux"
            "aarch64-linux"
          ];
          upstreamVariant = true;
        }
      ];
    };
}