diff --git a/pkgs/development/python-modules/triton-kernels/default.nix b/pkgs/development/python-modules/triton-kernels/default.nix new file mode 100644 index 000000000000..24e34f23f2d2 --- /dev/null +++ b/pkgs/development/python-modules/triton-kernels/default.nix @@ -0,0 +1,113 @@ +{ + lib, + buildPythonPackage, + + # build-system + setuptools, + + # dependencies + numpy, + torch, + triton, + + # tests + pytestCheckHook, + writableTmpDirAsHomeHook, + + # passthru + triton-kernels, +}: + +buildPythonPackage (finalAttrs: { + pname = "triton-kernels"; + # triton_kernels lives in Triton's repository and is released with it, while its own version has been "1.0.0" + # since its introduction; use Triton's version to tell the releases apart. + inherit (triton) version; + pyproject = true; + __structuredAttrs = true; + + inherit (triton) src; + + sourceRoot = "${finalAttrs.src.name}/python/triton_kernels"; + + # Make the wheel's version match `version`, which the metadata check requires. + postPatch = '' + substituteInPlace pyproject.toml \ + --replace-fail \ + 'version = "1.0.0"' \ + 'version = "${finalAttrs.version}"' + ''; + + build-system = [ + setuptools + ]; + + # torch and triton are imported throughout but not declared, as the package expects to be + # installed alongside the Triton it comes from. + dependencies = [ + numpy + torch + triton + ]; + + # Declared, but only the tests use it. + pythonRemoveDeps = [ "pytest" ]; + + pythonImportsCheck = [ + "triton_kernels" + "triton_kernels.matmul" + "triton_kernels.reduce" + "triton_kernels.swiglu" + "triton_kernels.tensor" + "triton_kernels.topk" + ]; + + nativeCheckInputs = [ + pytestCheckHook + # Triton writes its kernel cache under $HOME + writableTmpDirAsHomeHook + ]; + + disabledTestPaths = [ + # Spawns one process per GPU and uses torch's symmetric memory, which needs several GPUs + "tests/test_distributed.py" + + # Need Triton's proton profiler (triton.profiler), which triton is built without (TRITON_BUILD_PROTON = "OFF") + "tests/test_roofline.py" + "tests/test_topk.py" + ]; + + # The tests launch kernels; passthru.gpuCheck runs them. + doCheck = false; + + passthru = { + # The version follows Triton's; update Triton instead. + updateScript = null; + + gpuCheck = triton-kernels.overridePythonAttrs { + requiredSystemFeatures = [ "cuda" ]; + doCheck = true; + + disabledTests = [ + # triton.runtime.errors.OutOfResources: out of resource: shared memory + # These batched MXFP4 matmuls need more shared memory per block than consumer GPUs have (e.g. 101376 bytes on + # an RTX 4090, compute capability 8.9). + "1024-1024-1024-batched-mxfloat4_e2m1-bfloat16-bfloat16-10-9" + "1024-1024-1024-batched-mxfloat4_e2m1-float16-bfloat16-10-9" + ]; + }; + }; + + meta = { + description = "Triton kernels for matmul, mixture-of-experts top-k, reductions, SwiGLU and MXFP conversion"; + homepage = "https://github.com/triton-lang/triton/tree/main/python/triton_kernels"; + changelog = "https://github.com/triton-lang/triton/releases/tag/${finalAttrs.src.tag}"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ + GaetanLepage + connorbaker + ]; + teams = [ lib.teams.cuda ]; + inherit (triton.meta) platforms; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 6f687b1bcbf4..d84e3f631310 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -21334,6 +21334,8 @@ self: super: with self; { triton-cuda = self.triton.override { cudaSupport = true; }; + triton-kernels = callPackage ../development/python-modules/triton-kernels { }; + triton-no-cuda = self.triton.override { cudaSupport = false; }; tritonclient = callPackage ../development/python-modules/tritonclient { };