diff --git a/pkgs/development/python-modules/exllamav3/default.nix b/pkgs/development/python-modules/exllamav3/default.nix index c734222c2421..f50965fcc98f 100644 --- a/pkgs/development/python-modules/exllamav3/default.nix +++ b/pkgs/development/python-modules/exllamav3/default.nix @@ -8,6 +8,7 @@ setuptools, flash-attn, + flash-linear-attention, formatron, kbnf, marisa-trie, @@ -28,14 +29,14 @@ let in buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { pname = "exllamav3"; - version = "0.0.38"; + version = "0.0.39"; pyproject = true; src = fetchFromGitHub { owner = "turboderp-org"; repo = "exllamav3"; tag = "v${finalAttrs.version}"; - hash = "sha256-WlHIbnQX1Jd7y5yQzlqXVgBLQ92rnDSWy4z9bEm3WLA="; + hash = "sha256-auAOnsNOr22TTIBR9L81tp9ZCrSLY4RxXWAJ1E39EwM="; }; pythonRelaxDeps = [ @@ -60,6 +61,7 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { dependencies = [ flash-attn + flash-linear-attention formatron kbnf marisa-trie diff --git a/pkgs/development/python-modules/flash-linear-attention/default.nix b/pkgs/development/python-modules/flash-linear-attention/default.nix new file mode 100644 index 000000000000..9cce12efe716 --- /dev/null +++ b/pkgs/development/python-modules/flash-linear-attention/default.nix @@ -0,0 +1,66 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + pythonOlder, + + # build-system + setuptools, + + # dependencies + einops, + torch, + transformers, + + # optional-dependencies + causal-conv1d, + matplotlib, + datasets, + pytest, +}: + +buildPythonPackage (finalAttrs: { + pname = "flash-linear-attention"; + version = "0.5.0"; + pyproject = true; + + disabled = pythonOlder "3.10"; + + src = fetchFromGitHub { + owner = "fla-org"; + repo = "flash-linear-attention"; + tag = "v${finalAttrs.version}"; + hash = "sha256-g66yGHaBwEyjb+of76tKTtV/7as/2xQuqcjbGs4E3rU="; + }; + + build-system = [ setuptools ]; + + dependencies = [ + einops + torch + transformers + ]; + + optional-dependencies = { + # tilelang = [ tilelang ]; + conv1d = [ causal-conv1d ]; + benchmark = [ + matplotlib + datasets + ]; + test = [ pytest ]; + }; + + # Tests require a GPU + doCheck = false; + + pythonImportsCheck = [ "fla" ]; + + meta = { + description = "Triton-based implementations of causal linear attention"; + homepage = "https://github.com/fla-org/flash-linear-attention"; + changelog = "https://github.com/fla-org/flash-linear-attention/releases/tag/v${finalAttrs.version}"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ BatteredBunny ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index f4ff36be4137..ef6865cf518f 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -5709,6 +5709,8 @@ self: super: with self; { flash-attn-4 = callPackage ../development/python-modules/flash-attn-4 { }; + flash-linear-attention = callPackage ../development/python-modules/flash-linear-attention { }; + flash-mla = callPackage ../development/python-modules/flash-mla { }; flashinfer = callPackage ../development/python-modules/flashinfer { };