diff --git a/README.md b/README.md index 8511f50..3bfec0f 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,7 @@ A custom Nix overlay and flake providing additional packages not found in upstre | `shardr` | Decentralized LLM repository: content-addressed model store, BitTorrent-based sync, OpenAI-compatible serving via llama-server | AI Inference | | `skillsmcp` | MCP server that exposes Agent Skills to AI agents via the Model Context Protocol | MCP Servers | | `kubernetes-mcp-server` | Model Context Protocol (MCP) server for Kubernetes and OpenShift | MCP Servers | +| `kyojin` | ExLlamaV3-based inference engine with speculative decoding for 100 GB-class MoE packs on one 128 GB AMD Strix Halo (gfx1151, ROCm 7) | AI Inference | | `llama-cpp` | llama.cpp pinned to build b11439, built for Strix Halo (gfx1151) with ROCm + Vulkan backends | AI Inference | | `loop` | Corporate messenger for your team | Communication | | `radar` | Modern Kubernetes visibility — topology, event timeline, service traffic, resource browsing, Helm management, and GitOps support | Kubernetes | @@ -199,6 +200,7 @@ nix-overlay/ │ ├── haivemind/ # Multi-model AI consensus runner │ ├── hipengine/ # ROCm-native LLM inference engine for AMD GPUs │ ├── kubernetes-mcp-server/ # MCP server for Kubernetes and OpenShift +│ ├── kyojin/ # ExLlamaV3-based inference engine for Strix Halo (gfx1151) │ ├── loop/ # Corporate messenger for your team │ ├── mcp-gateway/ # MCP protocol gateway │ ├── omniroute/ # Unified AI router with 160+ providers diff --git a/packages/kyojin/default.nix b/packages/kyojin/default.nix new file mode 100644 index 0000000..a543bca --- /dev/null +++ b/packages/kyojin/default.nix @@ -0,0 +1,12 @@ +{ + pkgs, + inputs, + ... +}: +# The HIP extension has to be compiled against a ROCm build of PyTorch, and the +# only ROCm torch in this flake with a cached binary is the one from the pinned +# nixpkgs-torch211 input (2.11 + ROCm 7.2, python 3.13). Current nixos-unstable +# ships torch 2.13, whose ROCm build fails to configure in nixpkgs and is not +# on cache.nixos.org. +inputs.nixpkgs-torch211.legacyPackages.${pkgs.stdenv.hostPlatform.system}.callPackage ./package.nix + { } diff --git a/packages/kyojin/package.nix b/packages/kyojin/package.nix new file mode 100644 index 0000000..f7ee3d3 --- /dev/null +++ b/packages/kyojin/package.nix @@ -0,0 +1,204 @@ +{ + lib, + stdenv, + pkgs, + python3, + fetchFromGitHub, + makeWrapper, + rocmPackages, +}: + +let + py = python3.pkgs; + + # Kyojin is the Yamz fork of ExLlamaV3: a Python package plus a C++ extension + # that torch's cpp_extension cross-compiles to HIP code objects. Upstream + # only supports AMD Strix Halo, so the kernels are built for gfx1151 alone. + rocmArch = "gfx1151"; + + version = "1.3"; + + src = fetchFromGitHub { + owner = "Yamz-Labs"; + repo = "kyojin"; + rev = "v${version}"; + hash = "sha256-/DPoqW/iaLXRAH8XF5RdE+K5YS4jZT6Sic6QeCDQMlU="; + }; + + # setup.py passes EXL3_ROCM_DEV_INCLUDE as -I to every translation unit. The + # AMD wheels ship one devel tree holding hipsparse/ and thrust/; nixpkgs + # splits them per package, so join the include dirs. torch's public HIP + # headers also pull hipsolver/, and it no longer vendors pybind11. + rocmDevInclude = pkgs.symlinkJoin { + name = "kyojin-rocm-dev-include"; + paths = map lib.getDev ( + [ py.pybind11 ] + ++ [ + rocmPackages.clr + rocmPackages.hipblas + rocmPackages.hipblaslt + rocmPackages.hipcub + rocmPackages.hipfft + rocmPackages.hiprand + rocmPackages.hipsparse + rocmPackages.hipsolver + rocmPackages.rocblas + rocmPackages.rocprim + rocmPackages.rocrand + rocmPackages.rocsolver + rocmPackages.rocsparse + rocmPackages.rocthrust + ] + ); + }; + + # torch's HIP path in cpp_extension expects one ROCM_HOME holding bin/hipcc, + # the headers and the amdgcn bitcode; nixpkgs ships them separately. + rocmToolkit = pkgs.symlinkJoin { + name = "kyojin-rocm-toolkit"; + paths = [ + rocmPackages."rocm-core" + rocmPackages."rocm-runtime" + rocmPackages."rocm-device-libs" + rocmPackages."rocm-comgr" + rocmPackages.clr + (lib.getBin rocmPackages.hipcc) + rocmPackages.rocblas + rocmPackages.hipblas + rocmPackages.hipblas-common + rocmPackages.hipblaslt + rocmPackages.hipsparse + rocmPackages.hipsolver + rocmPackages.hiprand + rocmPackages.rocrand + rocmPackages.hipfft + rocmPackages.rocsparse + rocmPackages.rocsolver + rocmPackages.rocthrust + rocmPackages.rocprim + rocmPackages.hipcub + ]; + # The setup hooks of the joined packages would rewrite build flags; the + # compiler only needs the files. + postBuild = '' + rm -rf $out/nix-support + ''; + }; + + # The library: exllamav3 plus the compiled exllamav3_ext module. + library = py.buildPythonPackage { + inherit src version; + pname = "kyojin"; + + format = "pyproject"; + build-system = with py; [ + setuptools + wheel + ]; + + nativeBuildInputs = [ + py.ninja + rocmPackages."rocm-runtime" + ]; + + buildInputs = [ (lib.getBin rocmPackages.hipcc) ]; + + propagatedBuildInputs = with py; [ + aiohttp + huggingface-hub + jinja2 + llguidance + marisa-trie + numpy + pillow + pydantic + pyyaml + rich + safetensors + tokenizers + torchWithRocm + typing-extensions + ]; + + env = { + PYTORCH_ROCM_ARCH = rocmArch; + ROCM_PATH = "${rocmToolkit}"; + ROCM_HOME = "${rocmToolkit}"; + HIP_PATH = "${rocmToolkit}"; + HIPCC = "${rocmToolkit}/bin/hipcc"; + HIP_DEVICE_LIB_PATH = "${rocmPackages."rocm-device-libs"}/amdgcn/bitcode"; + # Upstream build scripts look these up to find the ROCm devel tree. + EXL3_ROCM_SDK = "${rocmToolkit}"; + EXL3_ROCM_DEV_INCLUDE = "${rocmDevInclude}/include"; + # Default defines of upstream build.sh (tools/strix_halo/rebuild.sh). + EXL3_HIP_DEFINES = "EXL3_HIP_STG_PAD"; + }; + + # setup.py pulls the extension source list and the arch helper with + # `from exllamav3...import ...`, which runs the package __init__ and + # therefore exllamav3/ext.py. That module JIT-compiles the extension when no + # precompiled exllamav3_ext is importable, i.e. during every wheel build + # (and it wants a writable $HOME for it). Register an empty `exllamav3` + # package first so only the two leaf modules are imported, and load + # build_config without the package wrapper. + postPatch = '' + sed -i '1i import sys as _sys, types as _types; _exl3_pkg = _types.ModuleType("exllamav3"); _exl3_pkg.__path__ = ["exllamav3"]; _sys.modules.setdefault("exllamav3", _exl3_pkg)' setup.py + sed -i 's|^from exllamav3\.exllamav3_ext\.build_config import get_sources as _get_sources$|import sys as _exl3_sys; _exl3_sys.path.insert(0, "exllamav3/exllamav3_ext"); from build_config import get_sources as _get_sources; _exl3_sys.path.pop(0)|' setup.py + grep -q 'from build_config import' setup.py + ''; + + # tests/ needs a real gfx1151 GPU and the published model packs. + doCheck = false; + + meta = with lib; { + description = "Inference engine for 100 GB-class EXL3 MoE models on AMD Strix Halo (gfx1151, ROCm 7)"; + homepage = "https://github.com/Yamz-Labs/kyojin"; + changelog = "https://github.com/Yamz-Labs/kyojin/releases"; + license = licenses.mit; + sourceProvenance = with sourceTypes; [ fromSource ]; + # HIP code objects are gfx1151-only and ROCm is Linux-only. + platforms = platforms.linux; + }; + }; + + python = py.python.withPackages (_: [ library ]); +in +stdenv.mkDerivation { + inherit src version; + pname = "kyojin"; + + # Only the wrappers are installed here: the serve scripts are plain scripts, + # they import their siblings by path, and the module itself is in `library`. + dontConfigure = true; + dontBuild = true; + + nativeBuildInputs = [ makeWrapper ]; + + installPhase = '' + runHook preInstall + + mkdir -p $out/bin + for model in glm mimo qwen; do + makeWrapper ${python}/bin/python $out/bin/kyojin-serve-$model \ + --set PYTORCH_ROCM_ARCH ${rocmArch} \ + --add-flags "${src}/tools/$model/serve.py" + done + + runHook postInstall + ''; + + passthru = { + inherit python rocmArch; + pythonPackage = library; + category = "AI Inference"; + updateScript = [ + "nix-update" + "--flake" + ".#kyojin" + ]; + }; + + meta = library.meta // { + mainProgram = "kyojin-serve-qwen"; + }; +}