[PyTorch-MPS][MRv2] Use upstream Triton sampling directly on MPS

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:

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\""

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

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(

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

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