diff --git a/pkgs/development/python-modules/tokenspeed-mla/bin.nix b/pkgs/development/python-modules/tokenspeed-mla/bin.nix new file mode 100644 index 000000000000..832d2d5eff19 --- /dev/null +++ b/pkgs/development/python-modules/tokenspeed-mla/bin.nix @@ -0,0 +1,79 @@ +{ + lib, + buildPythonPackage, + fetchPypi, + python, + stdenv, + + # nativeBuildInputs + autoPatchelfHook, + pypaInstallHook, + pythonRuntimeDepsCheckHook, + wheelUnpackHook, + + # dependencies + apache-tvm-ffi, + nvidia-cutlass-dsl, + nvidia-cutlass-dsl-libs-base, + tokenspeed-triton, + torch, +}: +let + inherit (stdenv.hostPlatform) system; + + hashes = { + aarch64-linux = "sha256-rD4rFrvtQXuICjBOkFfmzj3DaGRC4RHGX5EV4x/3T34="; + x86_64-linux = "sha256-1gW6k7nZT4Luj1LOLWaAs91AHBmPi0XCL11loqaCHqQ="; + }; +in +buildPythonPackage (finalAttrs: { + pname = "tokenspeed-mla"; + version = "0.1.5"; + pyproject = false; + __structuredAttrs = true; + + src = fetchPypi { + format = "wheel"; + pname = "tokenspeed_mla"; + inherit (finalAttrs) version; + dist = "py3"; + python = "py3"; + abi = "none"; + platform = "manylinux_2_28_${stdenv.hostPlatform.uname.processor}"; + hash = hashes.${system} or (throw "Unsupported system: ${system}"); + }; + + nativeBuildInputs = [ + autoPatchelfHook + pypaInstallHook + pythonRuntimeDepsCheckHook + wheelUnpackHook + ]; + + dependencies = [ + apache-tvm-ffi + nvidia-cutlass-dsl + tokenspeed-triton + torch + ]; + + preFixup = '' + # libtvm_ffi.so + addAutoPatchelfSearchPath "${apache-tvm-ffi}/${python.sitePackages}/tvm_ffi/lib" + # libcute_dsl_runtime.so + addAutoPatchelfSearchPath "${nvidia-cutlass-dsl-libs-base}/${python.sitePackages}/nvidia_cutlass_dsl/lib" + ''; + + pythonImportsCheck = [ "tokenspeed_mla" ]; + + meta = { + description = "Speed-of-light TokenSpeed MLA kernels for Blackwell SM100 and SM103"; + homepage = "https://github.com/lightseekorg/tokenspeed/tree/main/tokenspeed-mla"; + downloadPage = "https://pypi.org/project/tokenspeed-mla/#files"; + license = lib.licenses.mit; + sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ]; + maintainers = with lib.maintainers; [ prince213 ]; + platforms = lib.attrNames hashes; + broken = !torch.cudaSupport; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 6cb712605dff..7d0681312a9a 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19802,6 +19802,8 @@ self: super: with self; { tokenspeed-mla = callPackage ../development/python-modules/tokenspeed-mla { }; + tokenspeed-mla-bin = callPackage ../development/python-modules/tokenspeed-mla/bin.nix { }; + tokenspeed-triton = callPackage ../development/python-modules/tokenspeed-triton { }; tokenspeed-triton-bin = callPackage ../development/python-modules/tokenspeed-triton/bin.nix { };