electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

we got it to BUILD

Browse Source
This commit is contained in:
Evan
2026-03-20 15:38:16 +00:00
parent 581a1fcd79
commit 532a8f0b07
42 changed files with 456 additions and 235 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/exo_core/pyproject.toml
+2 -6
View File
@@ -3,13 +3,9 @@ name = "exo-core"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
authors = [
{ name = "Evan", email = "[email protected]" }
]
authors = [{ name = "Evan", email = "[email protected]" }]
requires-python = ">=3.13"
dependencies = [
"pydantic",
]
dependencies = ["pydantic"]
[build-system]
requires = ["uv_build>=0.9.24,<0.10.0"]
python/exo_core/src/exo_core/engine.py
+9 -4
View File
@@ -2,12 +2,19 @@ from abc import ABC, abstractmethod
from collections.abc import Callable, Iterable
from typing import Self
class TaskId(str): ...
class Cancelled: ...
class Finished: ...
CANCEL_ALL_TASKS = TaskId("CANCEL_TALL_TASKS")
class Engine[TaskType, ResponseType](ABC):
_cancelled_tasks: set[TaskId]
@@ -32,13 +39,12 @@ class Engine[TaskType, ResponseType](ABC):
@abstractmethod
def step(
self,
) -> Iterable[
tuple[TaskId, ResponseType | Cancelled | Finished]
]: ...
) -> Iterable[tuple[TaskId, ResponseType | Cancelled | Finished]]: ...
@abstractmethod
def close(self) -> None: ...
class EngineBuilder[SetupType, TaskType, ResponseType](ABC):
@classmethod
@abstractmethod
@@ -62,4 +68,3 @@ class EngineBuilder[SetupType, TaskType, ResponseType](ABC):
@abstractmethod
def close(self) -> None: ...
python/exo_core/src/exo_core/model_cards.py
+3 -3
View File
@@ -5,6 +5,9 @@ import aiofiles
import aiofiles.os as aios
import tomlkit
from anyio import Path, open_file
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
from exo.utils.pydantic_ext import CamelCaseModel
from huggingface_hub import model_info
from loguru import logger
from pydantic import (
@@ -23,9 +26,6 @@ from exo.shared.constants import (
EXO_ENABLE_IMAGE_MODELS,
RESOURCES_DIR,
)
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
from exo.utils.pydantic_ext import CamelCaseModel
# kinda ugly...
# TODO: load search path from config.toml
python/exo_core/src/exo_core/types/instances.py
+2 -2
View File
@@ -1,10 +1,10 @@
from enum import Enum
from exo.shared.types.common import Host, Id, NodeId
from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
from pydantic import model_validator
from exo_core.model_cards import ModelTask
from exo.shared.types.common import Host, Id, NodeId
from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
from exo_core.pydantic_ext import CamelCaseModel, TaggedModel
python/exo_core/src/exo_core/types/runner_response.py
+2 -1
View File
@@ -1,6 +1,8 @@
from collections.abc import Generator
from typing import Any, Literal
from exo.utils.pydantic_ext import TaggedModel
from exo.api.types import (
FinishReason,
GenerationStats,
@@ -9,7 +11,6 @@ from exo.api.types import (
TopLogprobItem,
Usage,
)
from exo.utils.pydantic_ext import TaggedModel
class BaseRunnerResponse(TaggedModel):
python/exo_core/src/exo_core/types/runners.py
+1 -1
View File
@@ -3,9 +3,9 @@ from collections.abc import Mapping
from pydantic import model_validator
from exo_core.model_cards import ModelId
from exo_core.pydantic import CamelCaseModel, TaggedModel
from exo_core.types.common import Id, NodeId
from exo_core.types.shards import ShardMetadata
from exo_core.pydantic import CamelCaseModel, TaggedModel
class RunnerId(Id):
python/exo_core/src/exo_core/types/tasks.py
+2 -2
View File
@@ -6,12 +6,12 @@ from exo.api.types import (
ImageEditsTaskParams,
ImageGenerationTaskParams,
)
from exo_core.pydantic import TaggedModel
from exo_core.types.common import CommandId, Id
from exo_core.types.text_generation import TextGenerationTaskParams
from exo_core.types.instances import BoundInstance, InstanceId
from exo_core.types.runners import RunnerId
from exo_core.types.shards import ShardMetadata
from exo_core.pydantic import TaggedModel
from exo_core.types.text_generation import TextGenerationTaskParams
class TaskId(Id):
python/nanobind_cmake.patch
+25
View File
@@ -0,0 +1,25 @@
diff --git a/cpp/nanobind/CMakeLists.txt b/cpp/nanobind/CMakeLists.txt
index ebfd3da..6ef06e3 100644
--- a/cpp/nanobind/CMakeLists.txt
+++ b/cpp/nanobind/CMakeLists.txt
@@ -18,8 +18,18 @@ target_sources(python_methods PRIVATE python_methods.cc)
target_link_libraries(python_methods PUBLIC xgrammar)
# Any code that uses nanobind directly lives here
-nanobind_add_module(xgrammar_bindings LTO nanobind.cc)
-target_link_libraries(xgrammar_bindings PRIVATE python_methods)
+nanobind_build_library(nanobind)
+add_library(xgrammar_bindings MODULE nanobind.cc)
+target_link_libraries(xgrammar_bindings PRIVATE
+ python_methods
+ nanobind
+)
+nanobind_opt_size(xgrammar_bindings)
+nanobind_lto(xgrammar_bindings)
+nanobind_set_visibility(xgrammar_bindings)
+nanobind_extension(xgrammar_bindings)
+nanobind_compile_options(xgrammar_bindings)
+nanobind_link_options(xgrammar_bindings)
if(DEFINED SKBUILD_PROJECT_NAME)
# Building wheel through scikit-build-core
python/parts.nix
+224 -31
View File
@@ -38,38 +38,231 @@
python = pkgs.python313;
# Overlay to provide build systems and custom packages
buildSystemsOverlay = final: prev: {
# mlx-lm is a git dependency that needs setuptools
mlx-lm = prev.mlx-lm.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
buildSystemsOverlay = final: prev:
let
addSetupTools = (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
torchLibs = [
final.nvidia-cuda-runtime
final.nvidia-cuda-nvrtc
final.nvidia-cuda-cupti
final.nvidia-nvjitlink
final.nvidia-cudnn-cu13
final.nvidia-cusparse
final.nvidia-cusparselt-cu13
final.nvidia-cufile
final.nvidia-nvshmem-cu13
final.nvidia-nccl-cu13
final.nvidia-cublas
final.nvidia-cufft
final.nvidia-curand
final.nvidia-cusolver
];
});
# rouge-score and sacrebleu don't declare setuptools as a build dependency
rouge-score = prev.rouge-score.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
cutlass = pkgs.fetchFromGitHub {
name = "cutlass-source";
owner = "NVIDIA";
repo = "cutlass";
tag = "v4.2.1";
hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
};
triton-kernels = pkgs.fetchFromGitHub {
owner = "triton-lang";
repo = "triton";
tag = "v3.5.0";
hash = "sha256-F6T0n37Lbs+B7UHNYzoIQHjNNv3TcMtoXjNrT8ZUlxY=";
};
cutlass-flashmla = pkgs.fetchFromGitHub {
owner = "NVIDIA";
repo = "cutlass";
rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
};
flashmla = pkgs.stdenv.mkDerivation {
pname = "flashmla";
version = "1.0.0";
src = pkgs.fetchFromGitHub {
name = "FlashMLA-source";
owner = "vllm-project";
repo = "FlashMLA";
rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
};
dontConfigure = true;
buildPhase = ''
rm -rf csrc/cutlass
ln -sf ${cutlass-flashmla} csrc/cutlass
'';
installPhase = ''
cp -rva . $out
'';
};
qutlass = pkgs.fetchFromGitHub {
name = "qutlass-source";
owner = "IST-DASLab";
repo = "qutlass";
rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
};
vllm-flash-attn = pkgs.stdenv.mkDerivation {
pname = "vllm-flash-attn";
version = "2.7.2.post1";
src = pkgs.fetchFromGitHub {
name = "flash-attention-source";
owner = "vllm-project";
repo = "flash-attention";
rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
};
patches = [
(pkgs.fetchpatch {
url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
})
(pkgs.fetchpatch {
url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
})
];
dontConfigure = true;
buildPhase = ''
rm -rf csrc/cutlass
ln -sf ${cutlass} csrc/cutlass
'';
installPhase = ''
cp -rva . $out
'';
};
mergedCudaLibraries = with pkgs.cudaPackages_13; [
cuda_cudart # cuda_runtime.h, -lcudart
cuda_cccl
libcurand # curand_kernel.h
libcusparse # cusparse.h
libcusolver # cusolverDn.h
cuda_nvtx
cuda_nvrtc
# cusparselt # cusparseLt.h
libcublas
];
});
sacrebleu = prev.sacrebleu.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
sqlitedict = prev.sqlitedict.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
word2number = prev.word2number.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
} // lib.optionalAttrs isDarwin {
# Use our pure Nix-built MLX with Metal support (macOS only)
mlx = self'.packages.mlx;
};
cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" {} ''
mkdir -p $out/include
ln -s ${pkgs.cudaPackages_13.cuda_cccl}/include $out/include/cccl
'';
cudaToolkitRoot = pkgs.symlinkJoin {
name = "cuda-merged-exo";
paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (mergedCudaLibraries ++ [ pkgs.cudaPackages_13.cuda_nvcc cuda_cccl_compat ]);
};
in
{
# mlx-lm is a git dependency that needs setuptools
mlx-lm = prev.mlx-lm.overrideAttrs addSetupTools;
# rouge-score and sacrebleu don't declare setuptools as a build dependency
rouge-score = prev.rouge-score.overrideAttrs addSetupTools;
sacrebleu = prev.sacrebleu.overrideAttrs addSetupTools;
sqlitedict = prev.sqlitedict.overrideAttrs addSetupTools;
word2number = prev.word2number.overrideAttrs addSetupTools;
fastsafetensors = prev.fastsafetensors.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools final.pybind11 ]; });
torch = prev.torch.overrideAttrs (old: {
propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ torchLibs ++ [ final.typing-extensions final.numpy ];
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
});
torchaudio = prev.torchaudio.overrideAttrs (old:
{
buildInputs = (old.buildInputs or [ ]) ++ [
final.torch
];
preFixup = (old.preFixup or "") + ''
addAutoPatchelfSearchPath "${final.torch}"
'';
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
});
torchvision = prev.torchvision.overrideAttrs (old:
{
buildInputs = (old.buildInputs or [ ]) ++ [
final.torch
];
preFixup = (old.preFixup or "") + ''
addAutoPatchelfSearchPath "${final.torch}"
'';
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
});
xgrammar = prev.xgrammar.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools final.scikit-build-core final.packaging final.pathspec pkgs.cmake final.nanobind ];
prePatch = ''
cat cpp/nanobind/CMakeLists.txt
'';
patches = (old.patches or [ ]) ++ [ ./nanobind_cmake.patch ];
});
vllm = prev.vllm.overrideAttrs (old: {
patches = (old.patches or [ ]) ++ [ ./vllm_uv2nix_cmake.patch ];
nativeBuildInputs = with pkgs.cudaPackages_13; (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
final.setuptools-scm
final.scikit-build-core
pkgs.cmake
cuda_nvcc
final.jinja2
final.wheel
final.markupsafe
pkgs.ninja
pkgs.autoAddDriverRunpath
];
buildInputs = with pkgs.cudaPackages_13; [
libcufile
cudnn
nccl
] ++ mergedCudaLibraries;
propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ torchLibs ++ [ final.torch ];
CUDA_HOME = "${cudaToolkitRoot}";
VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
VLLM_TARGET_DEVICE = "cuda";
TORCH_CUDA_ARCH_LIST = "12.0;12.1";
TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
CAFFE2_USE_CUDNN = "ON";
CAFFE2_USE_CUFILE = "ON";
CUTLASS_ENABLE_CUBLAS = "ON";
CUTLASS_NVCC_ARCHS_ENABLED = "12.1;12.1";
UV2NIX_CMAKE_FLAGS_JSON = builtins.toJSON [
"-DFETCHCONTENT_SOURCE_DIR_CUTLASS=${lib.getDev cutlass}"
"-DFLASH_MLA_SRC_DIR=${lib.getDev flashmla}"
"-DVLLM_FLASH_ATTN_SRC_DIR=${lib.getDev vllm-flash-attn}"
"-DQUTLASS_SRC_DIR=${lib.getDev qutlass}"
"-DTORCH_CUDA_ARCH_LIST=12.0;12.1"
"-DCUTLASS_NVCC_ARCHS_ENABLED=${pkgs.cudaPackages_13.flags.cmakeCudaArchitecturesString}"
"-DCUDA_HOME=${cudaToolkitRoot}"
"-DCAFFE2_USE_CUDNN=ON"
"-DCAFFE2_USE_CUFILE=ON"
"-DCUTLASS_ENABLE_CUBLAS=ON"
];
});
} // lib.optionalAttrs isDarwin {
# Use our pure Nix-built MLX with Metal support (macOS only)
mlx = self'.packages.mlx;
};
# Additional overlay for Linux-specific fixes (type checking env).
# Native wheels have shared lib dependencies we don't need at type-check time.
@@ -128,7 +321,7 @@
venvIgnoreCollisions = venvCollisionPaths;
};
exoCudaVenv = (pythonSet.mkVirtualEnv "exo-env" {
exo = lib.optionals isLinux [ "cuda" ];
exo = [ "cuda" ];
exo-pyo3-bindings = [ ];
}).overrideAttrs {
venvIgnoreCollisions = venvCollisionPaths;
python/vllm_engine/pyproject.toml
+3 -4
View File
@@ -12,14 +12,13 @@ dependencies = [
]
[tool.uv]
environments = [
"sys_platform == 'linux' and platform_machine == 'aarch64'"
]
environments = ["sys_platform == 'linux' and platform_machine == 'aarch64'"]
[tool.uv.sources]
vllm = { git = "https://github.com/hmellor/vllm.git", rev = "b99bedc737166ae5ca98cb9e3534b96e0c8c69aa" }
torch = [{ index = "pytorch-cu130", marker = "platform_machine == 'aarch64'" },
torch = [
{ index = "pytorch-cu130", marker = "platform_machine == 'aarch64'" },
{ index = "pytorch-cpu", marker = "platform_machine == 'x86_64'" },
]
python/vllm_engine/src/vllm_runner/__init__.py → python/vllm_engine/src/vllm_engine/__init__.py
View File
python/vllm_engine/src/vllm_runner/growable_cache.py → python/vllm_engine/src/vllm_engine/growable_cache.py
+1 -2
View File
@@ -339,9 +339,8 @@ def _patch_get_computed_blocks() -> None:
):
return original(self, request)
from vllm.utils.math_utils import cdiv # type: ignore[reportMissingImports]
from exo.worker.engines.vllm.vllm_generator import _build_layer_groups
from vllm.utils.math_utils import cdiv # type: ignore[reportMissingImports]
num_groups = len(self.kv_cache_config.kv_cache_groups)
null_block = self.block_pool.null_block
python/vllm_engine/src/vllm_runner/kv_cache.py → python/vllm_engine/src/vllm_engine/kv_cache.py
View File
python/vllm_engine/src/vllm_runner/prompt_format.py → python/vllm_engine/src/vllm_engine/prompt_format.py
+2 -2
View File
@@ -1,9 +1,9 @@
from exo.shared.types.common import ModelId
from exo.shared.types.text_generation import TextGenerationTaskParams
from mlx_lm.tokenizer_utils import TokenizerWrapper
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.llm_engine import LLMEngine
from exo.shared.types.common import ModelId
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.worker.engines.mlx.utils_mlx import (
apply_chat_template,
get_eos_token_ids_for_model,
python/vllm_engine/src/vllm_runner/vllm_generator.py → python/vllm_engine/src/vllm_engine/vllm_generator.py
+7 -7
View File
@@ -8,11 +8,6 @@ from collections.abc import Callable, Generator
from dataclasses import dataclass, field
import torch
from vllm.engine.arg_utils import EngineArgs
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.llm_engine import LLMEngine
from vllm.v1.kv_cache_interface import KVCacheConfig
from exo.shared.types.api import (
CompletionTokensDetails,
GenerationStats,
@@ -24,8 +19,6 @@ from exo.shared.types.memory import Memory
from exo.shared.types.tasks import TaskId
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.runner_response import GenerationResponse
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.mlx.utils_mlx import get_eos_token_ids_for_model
from exo.worker.engines.vllm.growable_cache import (
get_model_runner,
patch_vllm,
@@ -36,6 +29,13 @@ from exo.worker.engines.vllm.prompt_format import (
format_vllm_prompt,
make_vllm_sampling_params,
)
from vllm.engine.arg_utils import EngineArgs
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.llm_engine import LLMEngine
from vllm.v1.kv_cache_interface import KVCacheConfig
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.mlx.utils_mlx import get_eos_token_ids_for_model
from exo.worker.runner.bootstrap import logger
from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
python/vllm_uv2nix_cmake.patch
+12
View File
@@ -0,0 +1,12 @@
diff --git a/setup.py b/setup.py
index 68861fe4b..fd4738089 100644
--- a/setup.py
+++ b/setup.py
@@ -150,6 +150,7 @@ class cmake_build_ext(build_ext):
cmake_args = [
"-DCMAKE_BUILD_TYPE={}".format(cfg),
"-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
+ *json.loads(os.environ.get("UV2NIX_CMAKE_FLAGS_JSON", "[]"))
]
verbose = envs.VERBOSE