[PyTorch-MPS] Add output logprobs and MLX-LM parity by WindChimeRan · Pull Request #1001 · vllm-project/vllm-metal

WindChimeRan · 2026-10-04T23:32:39Z

Stacked on #880.

Add MPS output-token logprobs using upstream vLLM scoring so the existing mlx-lm parity tool can check PyTorch-MPS serving. Both serving backends use the same reference and shared comparison/reporting code.

Parity

Every serving run below uses PyTorch-MPS. Both tables use matching unquantized BF16 checkpoints and input token IDs: 40 prompts, 16 output tokens, request batch sizes 1 and 2, prefix caching off, and MPS CPU fallback disabled.

Cells report EXACT / TOP_K_MATCH / FAIL. Top-5 requires mutual membership at the first divergence; later continuation tokens are not compared.

PyTorch-MPS vs MLX-LM

Measured on 39d97baa using vLLM's default dtype resolution. All four MPS servers resolved to BF16.

Model	Batch 1: exact / top-5 / fail	Batch 2: exact / top-5 / fail
Qwen3-0.6B	29 / 11 / 0	29 / 11 / 0
Hunyuan-1.8B-Instruct	38 / 2 / 0	37 / 3 / 0
OLMo-2-0425-1B-Instruct	29 / 11 / 0	32 / 8 / 0
LFM2.5-1.2B-Instruct	31 / 9 / 0	34 / 6 / 0

VLLM_METAL_BACKEND=mps PYTORCH_ENABLE_MPS_FALLBACK=0 MLX_ENABLE_TF32=0 \
  python tools/check_parity.py --model /path/to/checkpoint \
  --top-k 5 --max-tokens 16 --batch-size 1 2 \
  --gpu-memory-utilization 0.15 --output-dir parity-mlx-lm

PyTorch-MPS vs Transformers/MPS

I tried with Transformers/MPS and found it's unnecessary, so I delete it

Historical cross-check from 788f0786, preserved from before the Transformers adapter was removed. These are measured results; the current tool uses MLX-LM only.

Model	Batch 1: exact / top-5 / fail	Batch 2: exact / top-5 / fail
Qwen3-0.6B	32 / 8 / 0	30 / 10 / 0
Hunyuan-1.8B-Instruct	38 / 2 / 0	37 / 3 / 0
OLMo-2-0425-1B-Instruct	32 / 8 / 0	34 / 6 / 0
LFM2.5-1.2B-Instruct	31 / 9 / 0	32 / 8 / 0

ericcurtin · 2026-10-07T08:32:19Z

Thanks! Blocked on #858 (stacked). Please also drop the no-op VLLM_METAL_MEMORY_FRACTION in benchmark_mps.py, set PYTORCH_ENABLE_MPS_FALLBACK=0 for the MPS arm, and reuse metric_counts from tools/benchmark/dspark_serving_benchmark.py instead of hand-parsing /metrics. Setting to draft, please reopen when ready for re-review.

WindChimeRan · 2026-10-13T23:00:51Z

Post review changes

Refactored the tool, by integrating with [Tools] Separate speculative parity from serving benchmarks #1065
Added “always set PYTORCH_ENABLE_MPS_FALLBACK=0” rule to AGENTS.md and the MPS docs
Removed the unused memory-fraction setting
benchmark & custom parsing are no longer in the PR.

ricky-chaoju

Tried the MPS parity with Qwen3-0.6B and got 33/7/0 on both batch sizes, logprobs look right too. LGTM

Use upstream vLLM scoring for MPS output-token logprobs. Reuse the existing MLX-LM parity flow with CPU fallback disabled and leave dtype resolution to the model loaders. Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>

Stacked on #1001 (PR5). Qwen3's eager MPS RoPE path performs separate cache lookup, rotation, and output assembly operations. This PR applies packed Q/K rotation in one Metal dispatch, using vLLM's existing model definitions and cosine/sine caches. The kernel handles FP16/BF16 full-head NeoX RoPE. Other rotary variants retain vLLM's native implementation. The shader is bundled in the wheel and loaded through `torch.mps.load_metallib`; no new dependency or C++ launcher is added. ```bash VLLM_METAL_BACKEND=mps PYTORCH_ENABLE_MPS_FALLBACK=0 \ vllm serve Qwen/Qwen3-0.6B --dtype bfloat16 ``` ## Parity PyTorch-MPS vs native MLX-LM on [b32095f](b32095f), using matching BF16 checkpoints and input token IDs: 40 prompts, 16 output tokens, request batches 1 and 2, prefix caching off, and optional CPU fallback disabled. Cells report **EXACT / TOP_K_MATCH / FAIL**. Top-5 requires mutual membership at the first divergence; later continuation tokens are not compared. All 320 MPS output token sequences are identical to PR4. | Model | Batch 1: exact / top-5 / fail | Batch 2: exact / top-5 / fail | |---|---:|---:| | Qwen3-0.6B | 29 / 11 / 0 | 29 / 11 / 0 | | Hunyuan-1.8B-Instruct | 38 / 2 / 0 | 37 / 3 / 0 | | OLMo-2-0425-1B-Instruct | 29 / 11 / 0 | 32 / 8 / 0 | | LFM2.5-1.2B-Instruct | 31 / 9 / 0 | 34 / 6 / 0 | ```bash VLLM_METAL_BACKEND=mps PYTORCH_ENABLE_MPS_FALLBACK=0 MLX_ENABLE_TF32=0 \ python tools/check_parity.py --model /path/to/checkpoint \ --top-k 5 --max-tokens 16 --batch-size 1 2 \ --gpu-memory-utilization 0.15 --output-dir rope-parity ``` <details> <summary>Checkpoints and environment</summary> Apple M5 Pro, 64 GiB; vLLM 0.31.0, PyTorch 2.13.0, MLX 0.32.1, MLX-LM 0.32.0. Runs use prebuilt native artifacts. The built wheel also completed Qwen3 generation with compiler imports blocked and TileLang absent; 448 fused RoPE calls were observed. | Checkpoint | Revision | |---|---| | Qwen/Qwen3-0.6B | `c1899de289a04d12100db370d81485cdf75e47ca` | | tencent/Hunyuan-1.8B-Instruct | `de940610ab5d45b4fedbd37e08444d0dc502837c` | | allenai/OLMo-2-0425-1B-Instruct | `48d788eca847d4d7548f375ad03d3c9312f6139e` | | LiquidAI/LFM2.5-1.2B-Instruct | `df58c174f05ff733f83f8cae10ea9298224c8006` | </details> ## Performance Qwen3-0.6B BF16, Sonnet 1024/128, 100 requests at 10/s, concurrency 32. Means of two fresh-server runs per path, in PR4/PR5/PR5/PR4 order. | Metric | PR4 MPS | PR5 MPS | Change | |---|---:|---:|---:| | Output tok/s ↑ | 461.49 | 503.93 | +9.2% | | Mean TTFT (ms) ↓ | 371.38 | 279.79 | -24.7% | | Mean TPOT (ms) ↓ | 57.77 | 52.70 | -8.8% | <details> <summary>Benchmark setup and individual runs</summary> PR4 [3f0c9af](3f0c9af) vs PR5 [b32095f](b32095f), on the same M5 Pro and software environment as the parity runs. Both paths use BF16, async scheduling, max model length/batched tokens 2048, block size 16, 2,340 KV blocks, prefix caching off, seed 0, greedy decoding, four warmups, and no logprobs. Each run completed 100 timed requests with 101,230 input tokens, 12,800 output tokens, identical input lengths, and zero preemptions. AC power; no recorded thermal/performance warnings. Desktop activity was not isolated. About 7.7 GiB of swap was already allocated; 7.5–33.8 MiB of swap-ins and no swap-outs were recorded per run. The individual results show the observed variation. | Run | Output tok/s | TTFT (ms) | TPOT (ms) | |---|---:|---:|---:| | 1-pr4 | 437.30 | 374.59 | 60.81 | | 2-pr5 | 501.00 | 276.59 | 53.02 | | 3-pr5 | 506.86 | 282.98 | 52.38 | | 4-pr4 | 485.68 | 368.17 | 54.74 | For each checkout, start a fresh server: ```bash VLLM_METAL_BACKEND=mps PYTORCH_ENABLE_MPS_FALLBACK=0 \ VLLM_USE_V2_MODEL_RUNNER=1 VLLM_USE_HW_AGNOSTIC=1 \ VLLM_ENABLE_V1_MULTIPROCESSING=1 VLLM_METAL_BUILD_FROM_SOURCE=0 MLX_ENABLE_TF32=0 \ vllm serve /path/to/Qwen3-0.6B --host 127.0.0.1 --port 8000 \ --dtype bfloat16 --model-impl auto --block-size 16 \ --max-model-len 2048 --max-num-seqs 32 --max-num-batched-tokens 2048 \ --gpu-memory-utilization 0.3 --kv-cache-memory-bytes 4294967296 \ --num-gpu-blocks-override 2340 --async-scheduling \ --no-enable-prefix-caching --generation-config vllm --seed 0 ``` After the server is ready, run the client from another terminal. Use a separate result directory for each run: ```bash vllm bench serve --backend vllm --base-url http://127.0.0.1:8000 \ --endpoint /v1/completions --model /path/to/Qwen3-0.6B \ --dataset-name sonnet --dataset-path sonnet.txt \ --sonnet-input-len 1024 --sonnet-output-len 128 \ --num-prompts 100 --request-rate 10 --max-concurrency 32 \ --temperature 0 --ignore-eos --seed 0 --num-warmups 4 \ --ready-check-timeout-sec 0 --save-result --save-detailed \ --result-dir /path/to/run-results --result-filename result.json ``` </details> Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>

WindChimeRan added this pull request to stack #871 October 4, 2026 23:33

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from a1b3e96 to 3d0bd58 Compare October 5, 2026 00:16

WindChimeRan force-pushed the ran/mps-backend-benchmark branch 2 times, most recently from 506f4ef to db92aba Compare October 5, 2026 03:57

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from db92aba to 813c02a Compare October 5, 2026 04:37

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from 813c02a to 536381f Compare October 5, 2026 05:02

This was referenced Oct 5, 2026

[PyTorch-MPS][MRv2] Add opt-in PyTorch MPS backend using MRv2 #858

Merged

Roadmap #1005

Open

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from 536381f to 58876db Compare October 5, 2026 06:46

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from 58876db to 363849a Compare October 5, 2026 06:49

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from 363849a to ec4db35 Compare October 5, 2026 18:04

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from ec4db35 to 2cacb77 Compare October 5, 2026 18:25

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from 2cacb77 to 33a040a Compare October 5, 2026 18:36

WindChimeRan force-pushed the ran/mps-backend-benchmark branch 2 times, most recently from bc64072 to fca6c6d Compare October 5, 2026 20:29

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from fca6c6d to 689b165 Compare October 6, 2026 19:11

WindChimeRan changed the title ~~[PyTorch-MPS] Add a matched MLX/MPS serving benchmark~~ [PyTorch-MPS] Add GPU parity and matched serving benchmarks Oct 6, 2026

WindChimeRan marked this pull request as ready for review October 6, 2026 21:10

ericcurtin marked this pull request as draft October 7, 2026 08:32

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from e3f1d61 to 5c25823 Compare October 7, 2026 20:46

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from 5c25823 to 3c96c4e Compare October 7, 2026 21:25

WindChimeRan mentioned this pull request Oct 7, 2026

[PyTorch-MPS] Fuse Q/K RoPE with a bundled Metal kernel #1052

Merged

WindChimeRan marked this pull request as ready for review October 8, 2026 04:01

WindChimeRan force-pushed the ran/mps-backend-benchmark branch 2 times, most recently from 7481f46 to cc924fe Compare October 8, 2026 05:37

WindChimeRan marked this pull request as draft October 13, 2026 06:40

WindChimeRan force-pushed the ran/mps-backend-benchmark branch 2 times, most recently from 78c38b3 to 6e95be3 Compare October 13, 2026 22:03

WindChimeRan changed the title ~~[PyTorch-MPS] Add GPU parity and matched serving benchmarks~~ [PyTorch-MPS] Add output logprobs and MLX-LM parity Oct 13, 2026

WindChimeRan force-pushed the ran/mps-backend-benchmark branch 2 times, most recently from 39d97ba to 3f0c9af Compare October 13, 2026 22:53

WindChimeRan marked this pull request as ready for review October 13, 2026 23:01

WindChimeRan requested a review from ricky-chaoju October 13, 2026 23:02

ricky-chaoju approved these changes Oct 14, 2026

View reviewed changes

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from 3f0c9af to 1e85406 Compare October 14, 2026 04:45

Enable MPS output logprobs and MLX-LM parity

c5a51cf

Use upstream vLLM scoring for MPS output-token logprobs. Reuse the existing MLX-LM parity flow with CPU fallback disabled and leave dtype resolution to the model loaders. Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>

WindChimeRan force-pushed the ran/mps-backend-benchmark branch from 1e85406 to c5a51cf Compare October 14, 2026 05:02

WindChimeRan merged commit 6a56199 into main Oct 14, 2026
3 checks passed

Navigation Menu

Repository navigation

[PyTorch-MPS] Add output logprobs and MLX-LM parity - #1001

[PyTorch-MPS] Add output logprobs and MLX-LM parity#1001
WindChimeRan merged 1 commit into
stack/mps-lfmfrom
ran/mps-backend-benchmark

WindChimeRan commented Oct 4, 2026 •

edited

Loading

Uh oh!

ericcurtin commented Oct 7, 2026

Uh oh!

WindChimeRan commented Oct 13, 2026 •

edited

Loading

Uh oh!

ricky-chaoju left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

WindChimeRan commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Parity

PyTorch-MPS vs MLX-LM

PyTorch-MPS vs Transformers/MPS

Uh oh!

ericcurtin commented Oct 7, 2026

Uh oh!

WindChimeRan commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricky-chaoju left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

WindChimeRan commented Oct 4, 2026 •

edited

Loading

WindChimeRan commented Oct 13, 2026 •

edited

Loading