diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md index 21bb0f8c3..377402f11 100644 --- a/docs/CONTRIBUTING.md +++ b/docs/CONTRIBUTING.md @@ -26,6 +26,25 @@ uv pip install -e ".[dev]" and editable plugin, and builds the native artifacts. It downloads the Metal toolchain if needed. Restart the server after editing Python files. +### Experimental PyTorch MPS + +From a source checkout, `./install.sh --mps` creates `.venv-vllm-metal-mps`, +installs the compiler/backend revisions pinned in the `mps` extra, and builds +them against the checksum-verified upstream LLVM SDK. The first build takes +several minutes and needs the Xcode setup above. + +```bash +./install.sh --mps +source .venv-vllm-metal-mps/bin/activate +export VLLM_METAL_BACKEND=mps TOOLCHAINS=Metal +vllm serve Qwen/Qwen3-0.6B --dtype float16 +``` + +The installer applies the temporary vLLM CPU-wheel device-detection workaround +only in this MPS environment and verifies an upstream sampling kernel. Use the +default environment for MLX. Bad-word filtering remains disabled pending the +upstream compiler fix. + ## Editing the Metal kernels When changing `.metal` shaders or `paged_ops.cpp`, enable source builds: diff --git a/install.sh b/install.sh index 2528726c3..f5ad19dce 100755 --- a/install.sh +++ b/install.sh @@ -128,6 +128,91 @@ install_vllm() { success "Installed vLLM core" } +install_mps_dependencies() { + section "Building pinned Triton and Apple backend" + # These are build tools, not runtime dependencies of the default MLX path. + uv pip install setuptools wheel 'cmake>=3.20,<4' ninja \ + 'nanobind==2.10.2' 'scikit-build-core>=0.10' requests + + local sources triton_url backend_url backend_repo backend_ref build_dir llvm_dir + sources=$(python - <<'PY' +import tomllib +from packaging.requirements import Requirement + +with open("pyproject.toml", "rb") as f: + extra = tomllib.load(f)["project"]["optional-dependencies"]["mps"] +requirements = {r.name: r for r in map(Requirement, extra)} +for name in ("triton", "triton-apple-backend"): + print(requirements[name].url) +PY + ) + triton_url=$(printf '%s\n' "$sources" | sed -n '1p') + backend_url=$(printf '%s\n' "$sources" | sed -n '2p') + backend_repo="${backend_url#git+}" + backend_repo="${backend_repo%@*}" + backend_ref="${backend_url##*@}" + backend_ref="${backend_ref%%#*}" + build_dir=$(mktemp -d) + register_cleanup_dir "$build_dir" + git init --quiet "$build_dir/backend" + git -C "$build_dir/backend" fetch --quiet --depth 1 "$backend_repo" "$backend_ref" + git -C "$build_dir/backend" checkout --quiet --detach FETCH_HEAD + + # Upstream selects the pinned macOS LLVM SDK and verifies its checksum. + llvm_dir=$(cd "$build_dir" && python backend/ci/download_llvm.py) + ( + export LLVM_SYSPATH="$build_dir/$llvm_dir" + export LLVM_INSTALL_DIR="$LLVM_SYSPATH" + export TRITON_EXT_ENABLED=1 TRITON_BUILD_PROTON=OFF + export MAX_JOBS="${MAX_JOBS:-8}" + export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-8}" + uv pip install --no-build-isolation "$triton_url" + uv pip install --no-build-isolation "$backend_url" + uv pip install --no-build-isolation -e '.[mps]' + ) + # Temporary vLLM 0.30 workaround, confined to this MPS environment. Replace + # the file atomically so uv's shared package cache cannot be modified. + python - <<'PY' +from importlib.util import find_spec +from pathlib import Path +import sys + +path = Path(find_spec("vllm").origin).parent / "triton_utils/importing.py" +if not path.resolve().is_relative_to(Path(sys.prefix).resolve()): + raise RuntimeError("Refusing to patch vLLM outside the MPS environment") +source = path.read_text() +old = 'if "cpu" in version("vllm"):' +new = 'if "cpu" in version("vllm") and current_platform.is_cpu():' +if new not in source: + if source.count(old) != 1: + raise RuntimeError("vLLM Triton detection changed; review the MPS install workaround") + temporary = path.with_suffix(".py.mps") + temporary.write_text(source.replace(old, new)) + temporary.replace(path) +print("Applied the vLLM CPU-wheel Triton workaround in the MPS environment") +PY + VLLM_METAL_BACKEND=mps python - <<'PY' +import torch +import triton +from vllm.triton_utils import HAS_TRITON +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample + +if not HAS_TRITON or not torch.backends.mps.is_available(): + raise RuntimeError("Triton or PyTorch MPS is unavailable") +target = triton.runtime.driver.active.get_current_target() +if target.backend != "mps": + raise RuntimeError(f"Expected the Triton Apple backend, got {target}") +logits = torch.arange(32, dtype=torch.float32, device="mps")[None, :] +mapping = torch.zeros(1, dtype=torch.int32, device="mps") +seed = torch.zeros(1, dtype=torch.int64, device="mps") +sampled = gumbel_sample(logits, mapping, torch.zeros(1, device="mps"), + seed, seed, apply_temperature=False, is_drafting=False) +if sampled.item() != 31: + raise RuntimeError("Upstream Triton sampling verification failed") +print(f"Verified upstream vLLM sampling with Triton {triton.__version__} on {target}") +PY +} + download_and_install_wheel() { local wheel_url="$1" local package_name="$2" @@ -173,9 +258,13 @@ main() { # Override the default dev channel with --stable or VLLM_METAL_CHANNEL. local channel="${VLLM_METAL_CHANNEL:-dev}" + local mps=false for arg in "$@"; do case "$arg" in + --mps) + mps=true + ;; --dev) channel="dev" ;; @@ -184,13 +273,15 @@ main() { ;; -h|--help) cat <<'EOF' -Usage: install.sh [--dev | --stable] +Usage: install.sh [--dev | --stable] [--mps] Options: --dev Install the latest development build cut from main. This is the default and the currently recommended channel. --stable Install the latest tagged stable release. Stable releases are cut by hand and may lag behind the dev channel. + --mps Build the pinned experimental Triton Apple backend from a + source checkout, in a separate .venv-vllm-metal-mps environment. -h, --help Show this help. The channel can also be set with VLLM_METAL_CHANNEL=dev|stable. @@ -245,6 +336,11 @@ EOF exit 1 fi + if $mps && [[ -z "$local_lib" || ! -f "$local_lib" ]]; then + error "Experimental MPS installation requires a source checkout." + exit 1 + fi + if ! ensure_uv; then exit 1 fi @@ -256,6 +352,9 @@ EOF # working directory there instead of the caller's cwd. cd "$script_dir" || exit 1 venv="$script_dir/.venv-vllm-metal" + if $mps; then + venv="${venv}-mps" + fi fi ensure_venv "$venv" @@ -270,8 +369,14 @@ EOF # Source checkouts build native artifacts; release installs use the wheel. uv pip install -e . + if $mps; then + export TOOLCHAINS="${TOOLCHAINS:-Metal}" + fi ensure_metal_toolchain build_native_artifacts + if $mps; then + install_mps_dependencies + fi else local release_data selected release_tag wheel_url vllm_release_tag release_data=$(fetch_release "$repo_owner" "$repo_name" "$channel") @@ -302,6 +407,9 @@ EOF echo "" echo "To use vllm, activate the virtual environment:" echo " source $venv/bin/activate" + if $mps; then + echo " export VLLM_METAL_BACKEND=mps TOOLCHAINS=Metal" + fi echo "" echo "Or add the venv to your PATH:" echo " export PATH=\"$venv/bin:\$PATH\"" diff --git a/pyproject.toml b/pyproject.toml index 0dec08987..6a3fc7e05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,11 @@ dependencies = [ [project.optional-dependencies] mps = [ - # Experimental host launcher for PyTorch's Metal command stream. + # Installed in build order by ./install.sh --mps (requires the LLVM SDK). + "torch==2.13.0", "ninja>=1.11", + "triton @ git+https://github.com/triton-lang/triton.git@4a15f415d8ac4f830ce788c6fca3a5b3b29908e7", + "triton-apple-backend @ git+https://github.com/imperatormk/triton-ext.git@adbf6e843d3d8e6b4328a1895be03213fe3a8ddd#subdirectory=backend/AppleGPU", ] gguf = [ # Local GGUF model loading. Keep aligned with upstream vLLM's lower bound diff --git a/vllm_metal/platform.py b/vllm_metal/platform.py index af905a58a..e044ce7a6 100644 --- a/vllm_metal/platform.py +++ b/vllm_metal/platform.py @@ -124,6 +124,16 @@ class MetalPlatform(Platform): # the DP registration is idempotent across repeated check_and_update_config calls. _dp_ray_hook_registered: bool = False + @classmethod + def num_compute_units(cls, device_id: int = 0) -> int: + if envs.VLLM_METAL_BACKEND != "mps": + return super().num_compute_units(device_id) + from vllm.triton_utils import triton + + return triton.runtime.driver.active.utils.get_device_properties(device_id)[ + "multiprocessorCount" + ] + @classmethod def get_device_name(cls, device_id: int = 0) -> str: """Get the name of the Metal device. @@ -294,29 +304,18 @@ def validate_request(cls, processed_inputs: object, params: object) -> None: if enabled ] if envs.VLLM_METAL_BACKEND == "mps": - unsupported_controls += [ - name - for name, enabled in ( - ("temperature", params.temperature != 0), - ("presence_penalty", params.presence_penalty != 0), - ("frequency_penalty", params.frequency_penalty != 0), - ("repetition_penalty", params.repetition_penalty != 1), - ("allowed_token_ids", bool(params.allowed_token_ids)), - ("bad_words", bool(params.bad_words)), - ) - if enabled - ] - unsupported_controls += [ + unsupported_controls = [ name for name in ( - "logprobs", - "logprob_token_ids", "prompt_logprobs", "structured_outputs", "thinking_token_budget", ) if getattr(params, name, None) is not None ] + # The Apple Triton backend cannot yet compile this upstream kernel. + if params.bad_words or params.bad_words_token_ids: + unsupported_controls.append("bad_words") if unsupported_controls: controls = ", ".join(unsupported_controls) raise VLLMValidationError( diff --git a/vllm_metal/pytorch_backend/runtime.py b/vllm_metal/pytorch_backend/runtime.py index c4c445fa2..3cd1aeea7 100644 --- a/vllm_metal/pytorch_backend/runtime.py +++ b/vllm_metal/pytorch_backend/runtime.py @@ -77,7 +77,6 @@ def staged_write(self): def install(): from vllm.v1.worker.gpu import buffer_utils, model_runner - from vllm.v1.worker.gpu.sample import sampler from vllm_metal.pytorch_backend import input_ops @@ -96,6 +95,4 @@ def install(): buffer_utils.UvaBackedTensor.copy_to_uva = copy_changed_state input_ops.install() - # validate_request limits this baseline to plain greedy generation. - sampler.gumbel_sample = lambda logits, *args, **kwargs: logits.argmax(dim=-1) buffer_utils._metal_installed = True diff --git a/vllm_metal/pytorch_backend/worker.py b/vllm_metal/pytorch_backend/worker.py index 67c1a056a..4014c8137 100644 --- a/vllm_metal/pytorch_backend/worker.py +++ b/vllm_metal/pytorch_backend/worker.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 """Opt-in vLLM worker for the MPS proof of concept.""" -from functools import wraps - import psutil import torch from vllm.utils.torch_utils import set_random_seed @@ -11,31 +9,10 @@ from vllm_metal.v1.worker import MetalWorker, init_worker_distributed_environment -def _patch_mps_mrv2_validation(): - from vllm.config import VllmConfig - - from vllm_metal import envs - - original = VllmConfig._validate_v2_model_runner - if getattr(original, "_metal_mrv2", False): - return - - @wraps(original) - def validate(config): - if envs.VLLM_METAL_BACKEND != "mps": - return original(config) - # Keep upstream feature validation; MPS replaces the Triton operations. - unsupported = config._get_v2_model_runner_unsupported_features() - if unsupported: - raise ValueError(f"MPS MRV2 does not support: {', '.join(unsupported)}") - - validate._metal_mrv2 = True - VllmConfig._validate_v2_model_runner = validate - - def configure_mps(config): from vllm import envs as vllm_envs from vllm.config.compilation import CompilationMode, CUDAGraphMode + from vllm.triton_utils import HAS_TRITON, triton from vllm_metal import envs @@ -44,6 +21,11 @@ def configure_mps(config): raise ValueError("Experimental MPS requires VLLM_USE_V2_MODEL_RUNNER=1") if not vllm_envs.VLLM_USE_HW_AGNOSTIC: raise ValueError("Experimental MPS requires VLLM_USE_HW_AGNOSTIC=1") + if ( + not HAS_TRITON + or triton.runtime.driver.active.get_current_target().backend != "mps" + ): + raise ValueError("MPS sampling requires a working Triton Apple backend") model = config.model_config if ( @@ -81,7 +63,6 @@ def configure_mps(config): f"Experimental MPS does not support {', '.join(unsupported)}" ) model.model_impl = "transformers" - _patch_mps_mrv2_validation() config.cache_config.block_size = 16 config.scheduler_config.async_scheduling = False config.compilation_config.mode = CompilationMode.NONE