python3Packages.flash-mla: init at 1.0.0-unstable-2026-03-31

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;
+  };
+})

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(

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 { };