Something went wrong. Try again.
lol
Something went wrong. Try again.
3.0 kB · 117 lines
Nix
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118# NOTE: At runtime, FlashInfer will fall back to PyTorch’s JIT compilation if a# requested kernel wasn’t pre-compiled in AOT mode, and JIT compilation always# 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, config, buildPythonPackage, fetchFromGitHub,
# build-system apache-tvm-ffi, setuptools,
# nativeBuildInputs cmake, ninja, cudaPackages,
# dependencies click, einops, numpy, nvidia-ml-py, tabulate, torch, tqdm,}:
buildPythonPackage (finalAttrs: { pname = "flashinfer"; version = "0.6.4"; pyproject = true;
src = fetchFromGitHub { owner = "flashinfer-ai"; repo = "flashinfer"; tag = "v${finalAttrs.version}"; fetchSubmodules = true; hash = "sha256-Hq3oTeEJHRvXwThI8vc06E3Ot/FnPP0sZUfze3ISa2o="; };
build-system = [ apache-tvm-ffi setuptools ];
nativeBuildInputs = [ cmake ninja (lib.getBin cudaPackages.cuda_nvcc) ];
dontUseCmakeConfigure = true;
buildInputs = with cudaPackages; [ cuda_cccl cuda_cudart libcublas libcurand ];
# FlashInfer offers two installation modes: # # JIT mode: CUDA kernels are compiled at runtime using PyTorch’s JIT, with # compiled kernels cached for future use. JIT mode allows fast installation, # as no CUDA kernels are pre-compiled, making it ideal for development and # testing. JIT version is also available as a sdist in PyPI. # # AOT mode: Core CUDA kernels are pre-compiled and included in the library, # reducing runtime compilation overhead. If a required kernel is not # pre-compiled, it will be compiled at runtime using JIT. AOT mode is # recommended for production environments. # # Here we use opt for the AOT version. preConfigure = '' export FLASHINFER_ENABLE_AOT=1 export TORCH_NVCC_FLAGS="--maxrregcount=64" export MAX_JOBS="$NIX_BUILD_CORES" '';
env.FLASHINFER_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;
pythonRemoveDeps = [ "nvidia-cudnn-frontend" "nvidia-cutlass-dsl" ]; dependencies = [ click einops numpy nvidia-ml-py tabulate torch tqdm ];
meta = { broken = !torch.cudaSupport || !config.cudaSupport; homepage = "https://flashinfer.ai/"; 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; maintainers = with lib.maintainers; [ breakds daniel-fahey ]; };})