From 78937945928f08ced8289136838bb644929ab3b0 Mon Sep 17 00:00:00 2001 From: Sizhe Zhao Date: Thu, 4 Jun 2026 17:16:20 +0800 Subject: [PATCH 1/6] tokenspeed-triton-llvm: init at 23.0.0-unstable-2026-04-08 --- .../to/tokenspeed-triton-llvm/package.nix | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 pkgs/by-name/to/tokenspeed-triton-llvm/package.nix diff --git a/pkgs/by-name/to/tokenspeed-triton-llvm/package.nix b/pkgs/by-name/to/tokenspeed-triton-llvm/package.nix new file mode 100644 index 000000000000..568542fccfe3 --- /dev/null +++ b/pkgs/by-name/to/tokenspeed-triton-llvm/package.nix @@ -0,0 +1,82 @@ +{ + lib, + stdenv, + fetchFromGitHub, + + cmake, + python3, + + llvmTargetsToBuild ? [ "NATIVE" ], # "NATIVE" resolves into x86 or aarch64 depending on stdenv + llvmProjectsToBuild ? [ + # Required for building triton>=3.7.0 + # https://github.com/triton-lang/triton/blob/a34f373ba47899831f3de3c83cd8f4877bf23d69/third_party/nvidia/CMakeLists.txt#L14 + "clang" + + # Required for building triton>=3.5.0 + # https://github.com/triton-lang/triton/blob/c3c476f357f1e9768ea4e45aa5c17528449ab9ef/third_party/amd/CMakeLists.txt#L6 + "lld" + + "llvm" + "mlir" + ], +}: +let + llvmNativeTarget = + if stdenv.hostPlatform.isx86_64 then + "X86" + else if stdenv.hostPlatform.isAarch64 then + "AArch64" + else + throw "Currently unsupported LLVM platform '${stdenv.hostPlatform.config}'"; + + inferNativeTarget = t: if t == "NATIVE" then llvmNativeTarget else t; + llvmTargetsToBuild' = [ + "AMDGPU" + "NVPTX" + ] + ++ map inferNativeTarget llvmTargetsToBuild; +in +stdenv.mkDerivation (finalAttrs: { + pname = "tokenspeed-triton-llvm"; + version = "23.0.0-unstable-2026-04-08"; # See cmake/Modules/LLVMVersion.cmake + __structuredAttrs = true; + strictDeps = true; + + # See https://github.com/lightseekorg/triton/blob/v3.7.10.post20260531/cmake/llvm-info.json + src = fetchFromGitHub { + owner = "llvm"; + repo = "llvm-project"; + rev = "87717bf9f81f7b29466c5d9a30a3453bdfc93941"; + hash = "sha256-8+Q19pOgovZgpN0it5TDrrQfXZFGiIRoP0Ha5dLQJp0="; + }; + + nativeBuildInputs = [ + cmake + python3 + ]; + + preConfigure = '' + cd llvm + ''; + + cmakeFlags = [ + (lib.cmakeFeature "LLVM_TARGETS_TO_BUILD" (lib.concatStringsSep ";" llvmTargetsToBuild')) + (lib.cmakeFeature "LLVM_ENABLE_PROJECTS" (lib.concatStringsSep ";" llvmProjectsToBuild)) + (lib.cmakeBool "LLVM_INSTALL_UTILS" true) + ]; + + requiredSystemFeatures = [ "big-parallel" ]; + + meta = { + description = "Collection of modular and reusable compiler and toolchain technologies"; + homepage = "https://github.com/llvm/llvm-project"; + license = + with lib.licenses; + AND [ + ncsa + (WITH asl20 llvm-exception) + ]; + maintainers = with lib.maintainers; [ prince213 ]; + platforms = with lib.platforms; aarch64 ++ x86; + }; +}) From 013498f6d1365d98a592ba63193e75a71e96c384 Mon Sep 17 00:00:00 2001 From: Sizhe Zhao Date: Thu, 4 Jun 2026 16:13:40 +0800 Subject: [PATCH 2/6] python3Packages.tokenspeed-triton: init at 3.7.10.post20260531 --- .../tokenspeed-triton/default.nix | 92 +++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 94 insertions(+) create mode 100644 pkgs/development/python-modules/tokenspeed-triton/default.nix diff --git a/pkgs/development/python-modules/tokenspeed-triton/default.nix b/pkgs/development/python-modules/tokenspeed-triton/default.nix new file mode 100644 index 000000000000..f63b1973cee8 --- /dev/null +++ b/pkgs/development/python-modules/tokenspeed-triton/default.nix @@ -0,0 +1,92 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + pythonAtLeast, + + # build-system + cmake, + nanobind, + ninja, + setuptools, + + # nativeBuildInputs + writableTmpDirAsHomeHook, + + # buildInputs + tokenspeed-triton-llvm, + nlohmann_json, + zlib, +}: +buildPythonPackage (finalAttrs: { + pname = "tokenspeed-triton"; + version = "3.7.10.post20260531"; + pyproject = true; + __structuredAttrs = true; + + src = fetchFromGitHub { + owner = "lightseekorg"; + repo = "triton"; + tag = "v${finalAttrs.version}"; + hash = "sha256-xsV63z2NtB5BM0rF0J+cnMH2RYzoWkpsSXHQI2nIEdQ="; + }; + + postPatch = '' + substituteInPlace setup.py \ + --replace-fail "TRITON_VERSION = " 'TRITON_VERSION = "${finalAttrs.version}" # ' + + sed -i '/def is_git_repo()/a\\ return False' setup.py + + substituteInPlace pyproject.toml \ + --replace-fail "cmake>=3.20,<4.0" "cmake>=3.20" \ + --replace-fail "nanobind==2.10.2" "nanobind>=2.10.2" + ''; + + build-system = [ + cmake + nanobind + ninja + setuptools + ]; + + dontUseCmakeConfigure = true; + + nativeBuildInputs = [ + writableTmpDirAsHomeHook + ]; + + # https://github.com/lightseekorg/triton/blob/v3.7.10.post20260531/.github/workflows/wheels.yml#L109-L117 + env = { + TRITON_OFFLINE_BUILD = true; + TRITON_BUILD_RELEASE = true; + TRITON_BUILD_PROTON = false; + TRITON_STABLE_ABI = pythonAtLeast "3.12"; + LLVM_SYSPATH = tokenspeed-triton-llvm; + JSON_SYSPATH = nlohmann_json; + NIX_CFLAGS_COMPILE = "-Wno-stringop-overflow"; + }; + + buildInputs = [ + zlib + ]; + + pythonImportsCheck = [ + "tokenspeed_triton" + ]; + + # tests import triton instead of tokenspeed_triton + doCheck = false; + + meta = { + description = "Language and compiler for custom Deep Learning operations"; + homepage = "https://github.com/lightseekorg/triton"; + downloadPage = "https://pypi.org/project/tokenspeed-triton/#files"; + changelog = "https://github.com/lightseekorg/triton/releases/tag/${finalAttrs.src.tag}"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ prince213 ]; + platforms = [ + "aarch64-linux" + "x86_64-linux" + ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index edeaf390c28d..413bf4899c55 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19800,6 +19800,8 @@ self: super: with self; { tokenlib = callPackage ../development/python-modules/tokenlib { }; + tokenspeed-triton = callPackage ../development/python-modules/tokenspeed-triton { }; + tokentrim = callPackage ../development/python-modules/tokentrim { }; tololib = callPackage ../development/python-modules/tololib { }; From 4d411b1251455e9b90672c0f3f9cb3b68f5432ba Mon Sep 17 00:00:00 2001 From: Sizhe Zhao Date: Thu, 4 Jun 2026 14:38:22 +0800 Subject: [PATCH 3/6] python3Packages.tokenspeed-triton-bin: init at 3.7.10.post20260531 --- .../python-modules/tokenspeed-triton/bin.nix | 73 +++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 75 insertions(+) create mode 100644 pkgs/development/python-modules/tokenspeed-triton/bin.nix diff --git a/pkgs/development/python-modules/tokenspeed-triton/bin.nix b/pkgs/development/python-modules/tokenspeed-triton/bin.nix new file mode 100644 index 000000000000..6b7b834a98de --- /dev/null +++ b/pkgs/development/python-modules/tokenspeed-triton/bin.nix @@ -0,0 +1,73 @@ +{ + lib, + buildPythonPackage, + fetchPypi, + python, + pythonAtLeast, + stdenv, + + # nativeBuildInputs + autoPatchelfHook, + pypaInstallHook, + wheelUnpackHook, + + # buildInputs + zlib, +}: +let + inherit (stdenv.hostPlatform) system; + pythonVersion = if pythonAtLeast "3.12" then "3.12" else python.pythonVersion; + abiTag = if pythonAtLeast "3.12" then "abi3" else pythonTag; + pythonTag = "cp${lib.replaceStrings [ "." ] [ "" ] pythonVersion}"; + + hashes = { + aarch64-linux = { + cp310 = "sha256-rRex0Y8vvOsLsJO+gxeZNkRF+iT/79IWQ5ShvwZJm3Y="; + cp311 = "sha256-Ds3Y1+o63Z19ckG4NcgMEthCZSvWsg3NhNsYojXhYIQ="; + abi3 = "sha256-Fs0KP8HP/rRYp+A+hohxT0n98LWhCL/KmZ9GWXw/qrs="; + }; + x86_64-linux = { + cp310 = "sha256-GGrclq8qFWW1Z27JIvYSZdf+u5pK6KK/5ioTigYMJ7Y="; + cp311 = "sha256-Ur4JbdxSJcbyNF4NCQrStWRmqFu+uGbMEaIOD+Pd4b8="; + abi3 = "sha256-uQrEHn8VkzeXVF/xqegDqdi+tMqbpw9tQang/CZIT1w="; + }; + }; +in +buildPythonPackage (finalAttrs: { + pname = "tokenspeed-triton"; + version = "3.7.10.post20260531"; + pyproject = false; + __structuredAttrs = true; + + src = fetchPypi { + format = "wheel"; + pname = "tokenspeed_triton"; + inherit (finalAttrs) version; + dist = pythonTag; + python = pythonTag; + abi = abiTag; + platform = "manylinux_2_27_${stdenv.hostPlatform.uname.processor}.manylinux_2_28_${stdenv.hostPlatform.uname.processor}"; + hash = hashes.${system}.${abiTag} or (throw "Unsupported system: ${system}"); + }; + + nativeBuildInputs = [ + autoPatchelfHook + pypaInstallHook + wheelUnpackHook + ]; + + buildInputs = [ zlib ]; + + pythonImportsCheck = [ "tokenspeed_triton" ]; + + meta = { + description = "Language and compiler for custom Deep Learning operations"; + homepage = "https://github.com/lightseekorg/triton"; + downloadPage = "https://pypi.org/project/tokenspeed-triton/#files"; + changelog = "https://github.com/lightseekorg/triton/releases/tag/v${finalAttrs.version}"; + license = lib.licenses.mit; + sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ]; + maintainers = with lib.maintainers; [ prince213 ]; + platforms = lib.attrNames hashes; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 413bf4899c55..3fa73bbd6cea 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19802,6 +19802,8 @@ self: super: with self; { tokenspeed-triton = callPackage ../development/python-modules/tokenspeed-triton { }; + tokenspeed-triton-bin = callPackage ../development/python-modules/tokenspeed-triton/bin.nix { }; + tokentrim = callPackage ../development/python-modules/tokentrim { }; tololib = callPackage ../development/python-modules/tololib { }; From cf6aef351e3d98f41e554f0f4e0737cfc82aadae Mon Sep 17 00:00:00 2001 From: Sizhe Zhao Date: Thu, 4 Jun 2026 21:15:14 +0800 Subject: [PATCH 4/6] python3Packages.tokenspeed-mla: init at 0.1.5 --- .../python-modules/tokenspeed-mla/default.nix | 53 +++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 55 insertions(+) create mode 100644 pkgs/development/python-modules/tokenspeed-mla/default.nix diff --git a/pkgs/development/python-modules/tokenspeed-mla/default.nix b/pkgs/development/python-modules/tokenspeed-mla/default.nix new file mode 100644 index 000000000000..4813cd099604 --- /dev/null +++ b/pkgs/development/python-modules/tokenspeed-mla/default.nix @@ -0,0 +1,53 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + setuptools, + + # dependencies + apache-tvm-ffi, + nvidia-cutlass-dsl, + tokenspeed-triton, + torch, +}: +buildPythonPackage (finalAttrs: { + pname = "tokenspeed-mla"; + version = "0.1.5"; + pyproject = true; + __structuredAttrs = true; + + src = fetchFromGitHub { + owner = "lightseekorg"; + repo = "tokenspeed"; + rev = "a39b3854dfd9b08a410028dbe5260eda08ef6b63"; + hash = "sha256-rl+cpZabmK24nMcam5Ud4GqnpLA3TqpVRznlX6lz6Xs="; + }; + + sourceRoot = "${finalAttrs.src.name}/tokenspeed-mla"; + + build-system = [ + setuptools + ]; + + dependencies = [ + apache-tvm-ffi + nvidia-cutlass-dsl + tokenspeed-triton + torch + ]; + + pythonImportsCheck = [ "tokenspeed_mla" ]; + + # no tests + doCheck = false; + + meta = { + description = "Speed-of-light TokenSpeed MLA kernels for Blackwell SM100 and SM103"; + homepage = "https://github.com/lightseekorg/tokenspeed/tree/main/tokenspeed-mla"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ prince213 ]; + broken = !torch.cudaSupport; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 3fa73bbd6cea..6cb712605dff 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19800,6 +19800,8 @@ self: super: with self; { tokenlib = callPackage ../development/python-modules/tokenlib { }; + tokenspeed-mla = callPackage ../development/python-modules/tokenspeed-mla { }; + tokenspeed-triton = callPackage ../development/python-modules/tokenspeed-triton { }; tokenspeed-triton-bin = callPackage ../development/python-modules/tokenspeed-triton/bin.nix { }; From 5387185762a64e5e7f001bccb53baaf782e7dd33 Mon Sep 17 00:00:00 2001 From: Sizhe Zhao Date: Thu, 4 Jun 2026 11:35:47 +0800 Subject: [PATCH 5/6] python3Packages.tokenspeed-mla-bin: init at 0.1.5 --- .../python-modules/tokenspeed-mla/bin.nix | 79 +++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 81 insertions(+) create mode 100644 pkgs/development/python-modules/tokenspeed-mla/bin.nix 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 { }; From c9e6c110409d7adf8fa319b03a2d56e57ccda0d5 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Sat, 6 Jun 2026 19:30:59 +0000 Subject: [PATCH 6/6] tokenspeed-triton-llvm: mark as broken on x86_64-darwin --- pkgs/by-name/to/tokenspeed-triton-llvm/package.nix | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pkgs/by-name/to/tokenspeed-triton-llvm/package.nix b/pkgs/by-name/to/tokenspeed-triton-llvm/package.nix index 568542fccfe3..9c039a344ef0 100644 --- a/pkgs/by-name/to/tokenspeed-triton-llvm/package.nix +++ b/pkgs/by-name/to/tokenspeed-triton-llvm/package.nix @@ -78,5 +78,9 @@ stdenv.mkDerivation (finalAttrs: { ]; maintainers = with lib.maintainers; [ prince213 ]; platforms = with lib.platforms; aarch64 ++ x86; + badPlatforms = [ + # clang++: error: clang frontend command failed with exit code 139 + "x86_64-darwin" + ]; }; })