Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215# NOTE: Since 0.6.x, upstream no longer pre-compiles kernels here (the AOT half now lives in the# separate `flashinfer-jit-cache` distribution). Every kernel is therefore compiled at runtime via# PyTorch's JIT, which requires the CUDA toolkit (via nvcc) to be available.## This means that if you plan to use flashinfer, you will need to set the environment variable# `CUDA_HOME` to `cudatoolkit`.{ lib, buildPythonPackage, fetchFromGitHub, symlinkJoin,
# build-system apache-tvm-ffi, packaging, setuptools,
# nativeBuildInputs ninja, cudaPackages, writableTmpDirAsHomeHook,
# dependencies click, cuda-bindings, cuda-core, cuda-tile, einops, nccl4py, numpy, nvidia-cudnn-frontend, nvidia-cutlass-dsl, nvidia-ml-py, requests, tabulate, torch, tqdm,
# tests mpi4py, numba, nvidia-cutlass, pytestCheckHook, pytest-timeout, responses,
cudaSupport ? torch.cudaSupport,}:
buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { pname = "flashinfer-python"; version = "0.7.2.post1"; pyproject = true; __structuredAttrs = true;
src = fetchFromGitHub { owner = "flashinfer-ai"; repo = "flashinfer"; tag = "v${finalAttrs.version}"; fetchSubmodules = true; hash = "sha256-a63PMJ/OE9214w7Alj+RTcMZhJ3aqqQeHFTNmfWaH0s="; };
build-system = [ apache-tvm-ffi packaging setuptools ];
nativeBuildInputs = [ ninja (lib.getBin cudaPackages.cuda_nvcc) # flashinfer creates its JIT workspace in $HOME at import time writableTmpDirAsHomeHook ];
buildInputs = with cudaPackages; [ cccl cuda_cudart libcublas libcurand ];
env = { FLASHINFER_CUDA_ARCH_LIST = lib.concatStringsSep " " torch.cudaCapabilities; # Prevent the build backend from trying to pip-install cuda-tile/NIXL wheels FLASHINFER_BUILD_NO_PIP = "1"; # Skip the (best-effort) NIXL-EP submodule build BUILD_NVEP = "0"; };
pythonRemoveDeps = [ # `cuda-python` is a meta-package pulling in the sub-packages of github:NVIDIA/cuda-python. # Only the ones flashinfer actually imports are substituted below. "cuda-python"
# Unpackaged. Only provides the optional `nccl.ep` backend of `flashinfer.moe_ep`, which is # probed at runtime. "nccl-extensions" ]; dependencies = [ apache-tvm-ffi click cuda-bindings cuda-core cuda-tile einops nccl4py ninja numpy nvidia-cudnn-frontend nvidia-cutlass-dsl nvidia-ml-py packaging requests tabulate torch tqdm ];
pythonImportsCheck = [ "flashinfer" ];
nativeCheckInputs = [ # nvshmem4py mpi4py numba nvidia-cutlass # pycute pytestCheckHook pytest-timeout responses ];
preCheck = # Make sure that the installed package (with its `data/` directory) is tested '' rm -rf flashinfer '' # The JIT compiler expects a regular CUDA toolkit layout (`bin/nvcc`, `include`, `lib64`) + '' export CUDA_HOME=${finalAttrs.passthru.cudaHome} '' # Only compile kernels for the GPU available at test time + '' unset FLASHINFER_CUDA_ARCH_LIST '';
# JIT-compiled modules reference template instantiations that are never called (nor compiled). # This only works with lazy binding, which `bindnow` (`-z now`) breaks: "undefined symbol". hardeningDisable = [ "bindnow" ];
# The full test suite has >200k (parametrized) tests, each JIT-compiling kernels. # Run the core subset that upstream runs in `scripts/task_jit_run_tests_part1.sh`. enabledTestPaths = [ "tests/attention/test_batch_decode_kernels.py" "tests/attention/test_logits_cap.py" "tests/attention/test_sliding_window.py" "tests/attention/test_tensor_cores_decode.py" ];
disabledTests = [ # Flaky: sporadic wrong results / illegal memory access, which then breaks all subsequent tests "test_batch_decode_tensor_cores_with_fast_plan" ];
# Tests require access to a GPU doCheck = false; passthru = { cudaHome = symlinkJoin { name = "flashinfer-cuda-home"; paths = lib.concatMap (p: map (o: p.${o}) (lib.remove "static" p.outputs)) ( lib.filter (p: p != null) ( with cudaPackages; [ cccl cuda_crt cuda_cudart cuda_nvcc libcublas libcurand ] ) ); postBuild = '' ln -s lib $out/lib64 ''; };
gpuCheck = finalAttrs.finalPackage.overrideAttrs { requiredSystemFeatures = [ "cuda" ]; doInstallCheck = true; }; };
meta = { broken = !cudaSupport; homepage = "https://flashinfer.ai/"; downloadPage = "https://github.com/flashinfer-ai/flashinfer"; changelog = "https://github.com/flashinfer-ai/flashinfer/releases/tag/${finalAttrs.src.tag}"; description = "Library and kernel generator for Large Language Models"; longDescription = '' FlashInfer is a library and kernel generator for Large Language Models that provides high-performance implementation of LLM GPU kernels such as FlashAttention, PageAttention and LoRA. FlashInfer focus on LLM serving and inference, and delivers state-of-the-art performance across diverse scenarios. ''; license = lib.licenses.asl20; mainProgram = "flashinfer"; teams = [ lib.teams.cuda ]; maintainers = with lib.maintainers; [ GaetanLepage ]; };})