diff --git a/pkgs/development/python-modules/flash-mla/default.nix b/pkgs/development/python-modules/flash-mla/default.nix new file mode 100644 index 000000000000..e301ce5b1454 --- /dev/null +++ b/pkgs/development/python-modules/flash-mla/default.nix @@ -0,0 +1,96 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + replaceVars, + cudaPackages, + + # build-system + setuptools, + torch, + + # buildInputs + pybind11, + + # passthru + nix-update-script, + + config, + cudaSupport ? config.cudaSupport, +}: + +let + inherit (lib) + getBin + optionalAttrs + optionals + ; +in +buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { + pname = "flash-mla"; + version = "0-unstable-2026-03-31"; + pyproject = true; + + src = fetchFromGitHub { + owner = "deepseek-ai"; + repo = "FlashMLA"; + rev = "71c737929f2567bd0a094ae140f8f60f390b1232"; + # Using the cutlass git subodules is necessary to get cutlass/util/command_line.h which is not + # shipped in cudaPackages.cutlass + fetchSubmodules = true; + hash = "sha256-d8Hh+1QFwD6cl9fE8/XSYdWiJJjY9bSRk5h4N2sEV2U="; + }; + + patches = [ + (replaceVars ./inject-git-rev.patch { + git_rev = "+${finalAttrs.src.rev}"; + }) + ]; + + postPatch = '' + substituteInPlace setup.py \ + --replace-fail \ + 'subprocess.run(["git", "submodule", "update", "--init", "csrc/cutlass"])' \ + "" + ''; + + env = optionalAttrs cudaSupport { + CUDA_HOME = (getBin cudaPackages.cuda_nvcc).outPath; + }; + + build-system = [ + setuptools + torch + ]; + + buildInputs = [ + pybind11 + ] + ++ optionals cudaSupport ( + with cudaPackages; + [ + cuda_cudart # cuda_runtime.h + libcublas # cublas_v2.h + libcurand # curand_kernel.h + libcusolver # cusolverDn.h + libcusparse # cusparse.h + ] + ); + + pythonImportsCheck = [ "flash_mla" ]; + + # Tests are not meant to run with pytest + doCheck = false; + + passthru.updateScript = nix-update-script { + extraArgs = [ "--version=branch" ]; + }; + + meta = { + description = "Efficient Multi-head Latent Attention Kernels"; + homepage = "https://github.com/deepseek-ai/FlashMLA"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ GaetanLepage ]; + broken = !cudaSupport; + }; +}) diff --git a/pkgs/development/python-modules/flash-mla/inject-git-rev.patch b/pkgs/development/python-modules/flash-mla/inject-git-rev.patch new file mode 100644 index 000000000000..4b3899f2fdb8 --- /dev/null +++ b/pkgs/development/python-modules/flash-mla/inject-git-rev.patch @@ -0,0 +1,19 @@ +diff --git a/setup.py b/setup.py +index 513b435..2178018 100644 +--- a/setup.py ++++ b/setup.py +@@ -133,13 +133,7 @@ ext_modules.append( + ) + ) + +-try: +- cmd = ['git', 'rev-parse', '--short', 'HEAD'] +- rev = '+' + subprocess.check_output(cmd).decode('ascii').rstrip() +-except Exception as _: +- now = datetime.now() +- date_time_str = now.strftime("%Y-%m-%d-%H-%M-%S") +- rev = '+' + date_time_str ++rev = "@git_rev@" + + + setup( diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index edda4a85e58f..a922352d6f57 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -5613,6 +5613,8 @@ self: super: with self; { flash-attn = callPackage ../development/python-modules/flash-attn { }; + flash-mla = callPackage ../development/python-modules/flash-mla { }; + flashinfer = callPackage ../development/python-modules/flashinfer { }; flashtext = callPackage ../development/python-modules/flashtext { };