{ lib, pkgs, fetchFromGitHub, }: let # FreeToken is an NVIDIA-GPU inference engine: setup.py builds two C++ # extensions that link the CUDA runtime, and the runtime needs a # CUDA-enabled torch. CUDA libraries carry the unfree "CUDA EULA" license, # so this package re-imports the flake's nixpkgs with allowUnfree enabled — # scoped to freetoken only, leaving the rest of the overlay unchanged. nixpkgsUnfree = import pkgs.path { inherit (pkgs.stdenv.hostPlatform) system; config.allowUnfree = true; }; py = nixpkgsUnfree.python3Packages; # torch>=2.11,<2.12: torch-bin 2.11 is this nixpkgs' CUDA build (it links # the cuda12.9-* libraries). The extensions must link the same libcudart # instance torch-bin links, so CUDA_HOME below is built from the matching # cudaPackages.cuda_cudart. cudart = nixpkgsUnfree.cudaPackages.cuda_cudart; # cuda_runtime_api.h includes crt/host_defines.h & co., which ship with # nvcc's header set, not with cudart. nvccHeaders = nixpkgsUnfree.cudaPackages.cuda_nvcc; # setup.py requires CUDA_HOME containing CUDA headers and libcudart. # Single-output cudart provides ${cudart}/include and ${cudart}/lib; the # crt/ headers come from nvcc; lib64/ is added for torch's cpp_extension, # lib/ is picked up by setup.py's cuda-home/lib fallback. cudaHome = nixpkgsUnfree.runCommand "freetoken-cuda-home" { } '' mkdir -p $out/include $out/lib64 ln -s ${cudart}/include/* $out/include/ ln -s ${nvccHeaders}/include/crt $out/include/crt ln -s ${cudart}/lib $out/lib for f in ${cudart}/lib/*.so*; do ln -s "$f" "$out/lib64/$(basename "$f")" done ''; # flashlib==0.3.0, pinned by freetoken, is not in nixpkgs. torch and triton # come from the prebuilt wheel builds (torch-bin / triton-bin), the same # instances torch-bin propagates — the scope defaults (source-built torch # and triton) would put two tritons/torches into the closure. flashlib = py.callPackage ./flashlib.nix { torch = py.torch-bin; triton = py.triton-bin; }; # freetoken pins apache-tvm-ffi==0.1.13.post3; nixpkgs has 0.1.10. # 0.1.13 needs cython>=3.2.8 (nixpkgs: 3.2.4), so build it with a newer # cython, scoped to this package only. apache-tvm-ffi = (py.apache-tvm-ffi.override { cython = py.cython.overridePythonAttrs (old: rec { version = "3.3.0"; src = fetchFromGitHub { owner = "cython"; repo = "cython"; tag = version; hash = "sha256-gIEqq8DAJF197gH1cMuq5JxI4rV/VHS70vg8hywXDqw="; }; }); }).overridePythonAttrs (old: { version = "0.1.13.post3"; src = fetchFromGitHub { owner = "apache"; repo = "tvm-ffi"; tag = "v0.1.13-post3"; fetchSubmodules = true; hash = "sha256-AN7AqBl62T8DqnRt7KRvGjqo/c0SJ66QZrCQQ5yicHw="; }; # 0.1.13 runs its suite with pytest-xdist (pyproject addopts = [ "-n" # "auto" ]) and touches GPUs; the nixpkgs check inputs have neither # xdist nor a GPU. The imports check still runs (and passes). dontUsePytestCheck = true; }); in py.buildPythonApplication (finalAttrs: { pname = "freetoken"; version = "0.1.2"; pyproject = true; src = fetchFromGitHub { owner = "FlashML-org"; repo = "FreeToken"; tag = "v${finalAttrs.version}"; hash = "sha256-0MhuubuTjNvtQZxisC2cg1dJeR+A6wZ901H5FRv+l+c="; }; # setup.py imports torch.utils.cpp_extension at build time and compiles two # C++ extensions (freetoken.kernel._pinned_tensor and _cpu_moe) that link # libcudart from CUDA_HOME. ninja is required by BuildExtension. build-system = with py; [ ninja setuptools torch-bin wheel ]; nativeBuildInputs = [ nixpkgsUnfree.autoPatchelfHook ]; # cudart: RPATH target for the built extensions (autoPatchelf finds # libcudart.so.12 here); libtorch comes from torch-bin via dependencies. buildInputs = [ cudart ]; env.CUDA_HOME = cudaHome; dependencies = with py; [ apache-tvm-ffi einops fastapi flashlib gguf huggingface-hub # freetoken's floor is modelscope>=1.37 (nixpkgs has 1.36.2); the only # API used is modelscope.snapshot_download (server/args.py), which is # stable across both. modelscope msgpack numpy openai partial-json-parser prompt-toolkit pydantic pyzmq safetensors torch-bin # triton==3.6.0 pin, satisfied with the wheel build (triton-bin) — the # same instance torch-bin propagates, keeping one triton in the closure. triton-bin tqdm transformers uvicorn ]; # Tests need an NVIDIA GPU and real model checkpoints (pytest markers # slow/needs_weights) — impossible in the build sandbox. doCheck = false; pythonImportsCheck = [ "freetoken" ]; # Wheel metadata constraints that nixpkgs packages do not match: # - gguf is versioned by upstream git rev in nixpkgs ("8951") rather than # PyPI semver; the reader API freetoken uses is stable. # - modelscope floor is 1.37 (nixpkgs: 1.36.2); only the stable # snapshot_download API is used. pythonRelaxDeps = [ "gguf" "modelscope" ]; passthru = { category = "AI Inference"; updateScript = [ "nix-update" "--flake" ".#freetoken" ]; }; meta = { description = "Local MoE-offload LLM inference runtime with OpenAI- and Anthropic-compatible APIs"; longDescription = '' FreeToken is an edge-native MoE serving engine for running frontier-scale open-weight models on consumer NVIDIA GPUs (RTX 30/40/50 series). It treats GPUs, CPUs and host memory as one elastic inference platform, with semantic-aware KV/expert caching and runtime VRAM re-allocation between caches. It serves OpenAI- and Anthropic-compatible APIs. Nix packaging notes: - Built against this nixpkgs' CUDA 12.9 stack (torch-bin 2.11 + cudaPackages.cuda_cudart) instead of upstream's cu130 wheels; CUDA 12.9 supports the same GPU range, including RTX 50 (sm_120). - The optional native accel extras (flashinfer, sglang-kernel) and the freetoken-kernel-cache companion wheel are not packaged; the runtime falls back to its pure-Triton kernels where those would be used. - Requires an NVIDIA GPU at runtime. ''; homepage = "https://github.com/FlashML-org/FreeToken"; changelog = "https://github.com/FlashML-org/FreeToken/releases/tag/v${finalAttrs.version}"; # Upstream code is Apache-2.0, but this derivation links CUDA libraries # under the unfree CUDA EULA (same as nixpkgs' torch-bin). license = lib.licenses.unfree; platforms = [ "x86_64-linux" ]; mainProgram = "ft"; }; })