millerson-overlay.nix/packages/kyojin/package.nix

{
  lib,
  stdenv,
  pkgs,
  python3,
  fetchFromGitHub,
  makeWrapper,
  rocmPackages,
  # Override only the serving toolchain; torch and the extension keep their ROCm.
  jitRocmSdk ? null,
}:

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.7.1";

  src = fetchFromGitHub {
    owner = "Yamz-Labs";
    repo = "kyojin";
    rev = "v${version}";
    hash = "sha256-84lRbN7H5eNG/RG8Ygvo9cByGOVOp7moFi0GS/YF+mU=";
  };

  # 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
      ninja
      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 loads exllamav3/exllamav3_ext/build_config.py by path (and
    # pyproject.toml ships the .hip sources, qsa_proof markers, and dense-GEMM
    # tuning seed), so no source patches are needed.

    # 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 ]);
  jitSdk = if jitRocmSdk == null then rocmToolkit else jitRocmSdk;
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`.
  #
  # Upstream `tools/strix_halo/env.sh` is sourced before serving, and on a
  # distro box that brings a system `cc` (the AMD Triton backend compiles a
  # CPython glue module at runtime, and again per kernel launcher) plus a hipcc
  # for the JIT HIP kernels (qsa prefill, gdn, kda, ple). Nix gives the wrapper
  # neither, so export them here instead: Triton takes the compiler from $CC,
  # exllamav3 finds hipcc under $EXL3_ROCM_SDK, and hipclang only finds the
  # amdgcn bitcode through $HIP_DEVICE_LIB_PATH (hipcc's own DEVICE_LIB_PATH is
  # not enough for a --genco call). The ROCm paths stay out of LD_LIBRARY_PATH
  # (upstream Trap 1/5: a second libhsa segfaults torch).
  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} \
        --set CC ${stdenv.cc}/bin/cc \
        --set EXL3_ROCM_SDK ${jitSdk} \
        --set HIP_DEVICE_LIB_PATH ${rocmPackages."rocm-device-libs"}/amdgcn/bitcode \
        --set EXL3_EXPANDABLE_SEGMENTS 0 \
        --prefix PATH : ${lib.makeBinPath [ stdenv.cc ]} \
        --add-flags "${src}/tools/$model/serve.py"
    done

    runHook postInstall
  '';

  passthru = {
    inherit python rocmArch;
    pythonPackage = library;
    jitRocmSdk = jitSdk;
    tests.jit =
      pkgs.runCommand "kyojin-jit-check"
        {
          EXL3_ROCM_SDK = jitSdk;
          # CLR/COMGR also reads this variable. Newer bitcode crashes torch
          # matmul; a custom hipcc must select its own bitcode in its subprocess.
          HIP_DEVICE_LIB_PATH = "${rocmPackages."rocm-device-libs"}/amdgcn/bitcode";
          EXL3_EXPANDABLE_SEGMENTS = "0";
          KYOJIN_REQUIRE_SDK = lib.boolToString (jitRocmSdk != null);
        }
        ''
          export HOME="$TMPDIR"
          ${python}/bin/python ${./check-jit.py}
          touch "$out"
        '';
    category = "AI Inference";
    updateScript = [
      "nix-update"
      "--flake"
      ".#kyojin"
    ];
  };

  meta = library.meta // {
    mainProgram = "kyojin-serve-qwen";
  };
}