python3Packages.tokenspeed-{mla,triton}: init (#527879)

2026-06-11 09:25:25 +00:00
parent bd2e3b09cd c9e6c11040
commit 42e30dc06e
6 changed files with 391 additions and 0 deletions
@@ -0,0 +1,86 @@
+{
+  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;
+    badPlatforms = [
+      # clang++: error: clang frontend command failed with exit code 139
+      "x86_64-darwin"
+    ];
+  };
+})
@@ -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;
+  };
+})
@@ -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;
+  };
+})
@@ -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;
+  };
+})
@@ -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"
+    ];
+  };
+})
@@ -19857,6 +19857,14 @@ self: super: with self; {

  tokenlib = callPackage ../development/python-modules/tokenlib { };

+  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 { };
+
  tokentrim = callPackage ../development/python-modules/tokentrim { };

  tololib = callPackage ../development/python-modules/tololib { };