flake.nix

{
  description = "Flake for megablocks_moe kernel";

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

  outputs =
    {
      self,
      kernel-builder,
    }:
    kernel-builder.lib.genFlakeOutputs {
      path = ./.;
      rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;

      pythonCheckInputs = pkgs: with pkgs; [ 
        tqdm
        py-cpuinfo
        importlib-metadata
        torchmetrics
      ];

      torchVersions = defaultVersions: defaultVersions ++ [
        {
          torchVersion = "2.7";
          rocmVersion = "6.4.2";
          cxx11Abi = true;
          systems = [ "x86_64-linux" ];
          bundleBuild = true;
        }
      ];
    };
}