[PyTorch-MPS] Fuse Q/K RoPE with a bundled Metal kernel by WindChimeRan · Pull Request #1052 · vllm-project/vllm-metal

WindChimeRan · 2026-10-07T22:25:38Z

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.

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 b32095fe, 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

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

Checkpoints and environment

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`

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%

Benchmark setup and individual runs

PR4 3f0c9afd vs PR5 b32095fe, 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:

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:

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

Register the upstream RoPE layer through register_oot and retain its cosine/sine cache. Rotate packed Q/K rows in one dispatch, preserving FP16/BF16 intermediate rounding and projection strides. Compile the shader with the existing release artifacts and load it through torch.mps.load_metallib. No runtime compiler dependency or new C++ launcher is needed. Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>

## PR series - PR1 [vllm-omni #8943](vllm-project/vllm-omni#8943): external operator registration and compatibility fixes (merged) - PR1.5 [vllm-omni #9285](vllm-project/vllm-omni#9285): publication-ordering fix (merged) - PR2 [vllm-omni #8946](vllm-project/vllm-omni#8946): shared-memory storage hooks (merged) - PR3 [vllm-omni #8949](vllm-project/vllm-omni#8949): eager predictor KV reuse (optimization, merged) - [vllm-metal #1109](#1109): Qwen3-TTS serving on Apple Silicon - **[vllm-metal #1117](#1117): Torch nightly for earlier GPU submission (optimization)** - [vllm-metal #1052](#1052): fused Q/K RoPE (optimization, merged) ## Core motivation * Fresh MPS/MLX top-5 parity passed on Qwen3-0.6B and LFM2.5-1.2B; details below. * Earlier measurements show MPS speedups on the text models below and Qwen3-TTS in #1109. ## Summary Pin PyTorch `2.16.0.dev20261013` for [upstream earlier MPS submission](pytorch/pytorch#200181). The tested nightly contains it; PyTorch 2.15.0 does not. Install it after vLLM and build matching native extensions. PyTorch now submits queued GPU work earlier, roughly every 16 encoded operations, subject to an in-flight limit, so GPU execution can overlap CPU encoding. Each iteration still encodes its operations. user installation is the same ## Benchmark Earlier throughput gains versus Torch 2.13 on M5 Pro (BF16, Sonnet 1024/128, 100 requests per setting). These measurements predate the current rebase and were not rerun: | Model | Concurrency 1 | Concurrency 8 | |---|---:|---:| | Qwen3-0.6B | +30% | +33% | | LFM2.5-1.2B | +11% | +21% | <details> <summary>Validation</summary> Runtime validation at `c1d27dca`, based on main `487e7316`: vLLM 0.32.0, Torch `2.16.0.dev20261013`, MLX 0.32.4; `PYTORCH_ENABLE_MPS_FALLBACK=0`. - `pytest tests/test_platform.py tests/test_metal_build_cache.py -q`: **228 passed, 10 skipped**. `VLLM_METAL_BACKEND=mps pytest tests/test_mps_backend.py -q`: **30 passed**. - Fresh `tools/check_parity.py` checks: the first eight shared corpus prompts, 16 output tokens, matching checkpoints/input IDs, top-k 5, batches 1/8. **15 EXACT, 17 TOP_K_MATCH, 0 FAIL** across both models. Qwen3: 4/4 at each batch; LFM2.5: 3/5 then 4/4 (EXACT/TOP_K_MATCH). Top-k matches are checked at the first divergence. - Native artifacts rebuilt with nightly. The built wheel includes the Torch requirements and loads its prebuilt MPS launcher without compiler tooling. All three pinned Python 3.12/Apple Silicon wheels remain downloadable. - ShellCheck, Bash/Ruby syntax, Ruff lint/format and whitespace checks passed. - Homebrew fix at `93ae0bc3`: `brew style Formula/vllm-metal.rb` passed. A cleanup/relocation probe preserved both required `RECORD` files; TVM-FFI library discovery and xgrammar native parsing passed. - The earlier benchmark used one serving run per setting and completed all 800 requests; it measures the complete nightly upgrade. That benchmark and a clean Homebrew installation were not rerun. </details> ## Related PRs mps "cuda graph like" record & replay pytorch/pytorch#189019 ## Retirement Plan * Replace the nightly with a stable Torch release containing the optimization. * Remove the override once the supported vLLM release declares a compatible Torch version. --------- Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>

## PR series All merged! upstream vllm-omni plumbing - PR1 [vllm-omni #8943](vllm-project/vllm-omni#8943): external operator registration and compatibility fixes (merged) - PR1.5 [vllm-omni #9285](vllm-project/vllm-omni#9285): publication-ordering fix (merged) - PR2 [vllm-omni #8946](vllm-project/vllm-omni#8946): shared-memory storage hooks (merged) - PR3 [vllm-omni #8949](vllm-project/vllm-omni#8949): eager predictor KV reuse (optimization, merged) vllm-metal PRs. - **[vllm-metal #1109](#1109): Qwen3-TTS serving on Apple Silicon** - [vllm-metal #1117](#1117): Torch nightly for earlier GPU submission (optimization) - [vllm-metal #1052](#1052): fused Q/K RoPE (optimization, merged) ## Purpose Adds optional Qwen3-TTS speech serving on Apple Silicon through vLLM-Omni, with streaming audio and concurrent requests. This preview uses greedy decoding. See the [usage guide](https://github.com/vllm-project/vllm-metal/blob/ran/omni-qwen3-tts/docs/qwen3-tts-mps.md). Rebased onto Metal main `0433ac61`. MPS defaults live in Metal's packaged deployment profile, selected with `--deploy-config` as shown in the guide. Predictor KV reuse is automatic upstream. All Omni prerequisites are merged. The `omni` extra, also included in `all`, now pins [upstream `c4ac0642`](vllm-project/vllm-omni@c4ac064). Validation on the rebased head: **206 passed, 10 skipped**, including speech streaming at C=1/4/8. Fresh Transformers/MPS trace parity for Base, CustomVoice and VoiceDesign at B=1/3: **322 EXACT, 2 TOP_K_MATCH, 0 FAIL**. Both top-5 matches are tied highest logits in the reference. Tested with the pinned Omni revision and Torch 2.13. ## Framework comparison Pre-rebase Qwen3-TTS 0.6B Base measurements on 100 [Seed-TTS samples](https://huggingface.co/datasets/zhaochenyang20/seed-tts-eval), split evenly between English and Chinese: ![Historical Qwen3-TTS throughput: pre-rebase PR and earlier RoPE/nightly experiments](https://raw.githubusercontent.com/WindChimeRan/vllm-metal/d708d12ca24cd98cfa00f67b5ff37752bd68eb6c/research/qwen3-tts-milestone-20261014.png) All results below predate this rebase. “Proposed PRs (fresh)” in the figure refers to `f2887704`; the RoPE and nightly panels are earlier experiments. These are not measurements of the rebased head. | Version | C | Total time (s) | Audio s/s | First audio (ms) | | --- | ---: | ---: | ---: | ---: | | mlx-audio (earlier) | 1 | 149.65 | 3.491 | 169.3 | | vllm-omni-mlx (earlier) | 1 | 147.45 | 3.543 | 66.5 | | PR before rebase (`f2887704`) | 1 | 271.46 | 1.884 | 407.9 | | PR before rebase (`f2887704`) | 4 | 104.62 | 4.894 | 539.7 | | PR before rebase (`f2887704`) | 8 | 72.18 | 7.090 | 753.9 | | + RoPE (earlier) | 1 | 241.63 | 2.117 | 407.9 | | + RoPE (earlier) | 4 | 96.63 | 5.325 | 532.7 | | + RoPE (earlier) | 8 | 67.82 | 7.621 | 767.8 | | + RoPE + nightly (earlier) | 1 | 143.28 | 3.603 | 196.4 | | + RoPE + nightly (earlier) | 4 | 59.90 | 8.596 | 290.1 | | + RoPE + nightly (earlier) | 8 | 43.66 | 11.848 | 431.4 | <details> <summary>Measurement setup</summary> - M5 Pro, 20 GPU cores, 64 GiB, macOS 27.0.1. Same unquantized checkpoint, ordered texts, reference audio and eight excluded warmups. Real HTTP streaming, greedy decoding, repetition penalty 1.5, automatic language selection, and a 1024-token cap. One 100-request pass per setting. C denotes client concurrency, not a fixed decode batch size. - Pre-rebase PR runs: Metal `f2887704`, installed Omni preview `4bb27684`, PyTorch 2.13.0, vLLM 0.31.0. Installed runtime sources were checked against the pinned preview. Stage-0 active-request peaks reached 1/4/8. No thermal warnings or swap-outs during timed windows; desktop activity was not isolated. - Earlier runs: Metal `85ef619`, Omni `d3e17ab`; RoPE used PyTorch 2.13.0, and the nightly arm used 2.16.0.dev20261013. The nightly includes [pytorch/pytorch#200181](pytorch/pytorch#200181) and other upstream changes. - Earlier baselines: mlx-audio 0.5.8; vllm-omni-mlx `210cc4e`; MLX 0.32.1. Audio s/s is generated audio duration divided by wall time; first audio is mean client latency. Framework-native precision, prompting, repetition history and chunking differ. Single passes do not estimate run-to-run variation. </details> <details> <summary>Previously reported validation (before this rebase)</summary> - All 300 pre-rebase timed requests completed below the output cap. C1 PCM matches the prior validated run for all 100 utterances. C4/C8 match 21/100 and 2/100; concurrent generation is not claimed to be bit-identical. ASR was not rerun for those concurrent outputs. For model additions, follow the [contributor pattern](https://github.com/vllm-project/vllm-metal/blob/ran/omni-qwen3-tts/docs/pytorch-mps.md#adding-omni-models). </details> ## Earlier KV reuse measurement Base 0.6B, BF16, M5 Pro (64 GiB), PyTorch 2.13.0, vLLM 0.32.0. Same 50 EN + 50 ZH Seed-TTS requests in each setting; only predictor KV reuse changes. | HTTP concurrency | KV off: audio s/s | KV on: audio s/s | Speedup | | --- | ---: | ---: | ---: | | C=1 | 1.795 | 1.917 | **1.07×** | | C=8 | 5.688 | 7.132 | **1.25×** | C is client concurrency, not a fixed decode batch. The fixed-batch predictor measurements are in #8949. <details> <summary>Validation and benchmark setup</summary> Metal `8073051b`, Omni preview `c2b3c229`. Native artifacts rebuilt from this checkout; `PYTORCH_ENABLE_MPS_FALLBACK=0`. - **167 passed, 10 skipped:** platform 135/10; MPS backend 28; real CustomVoice HTTP 4 (streaming at C=1/4/8 and rejection of random sampling). - Base serving benchmark: **400/400 completed, zero failures**, using upstream `vllm-omni bench serve --omni --backend openai-audio-speech`. One pass per setting; 50 requests per locale, with four warmups and the tool's initial probe per locale. The server used the built Metal adapter and packaged profile. - C=1 wall time: 290.19 → 271.79 s; generated audio: 520.96 → 521.04 s. C=8: 91.16 → 73.12 s; audio: 518.48 → 521.52 s. Throughput is total audio / summed benchmark time. BF16 outputs need not be bit-identical; no ASR rerun. - Actual worker KV off/on dispatch confirmed separately. AC power; no observed thermal warnings or new swap-outs. Desktop not isolated; single passes do not estimate variance. - Ruff lint/format and whitespace checks passed. Built wheel includes the MPS profile and updated preview pin; `all` includes build, Omni and development extras. ```bash PYTORCH_ENABLE_MPS_FALLBACK=0 pytest tests/test_platform.py -q VLLM_METAL_BACKEND=mps PYTORCH_ENABLE_MPS_FALLBACK=0 pytest tests/test_mps_backend.py -q QWEN3_TTS_TEST_MODEL=<local-CustomVoice-checkpoint> PYTORCH_ENABLE_MPS_FALLBACK=0 \ pytest tests/test_mps_omni_serving.py -q ``` </details> --------- Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>

WindChimeRan added this pull request to stack #871 October 7, 2026 22:25

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

WindChimeRan force-pushed the ran/mps-rope branch from ae566b9 to 90a8958 Compare October 8, 2026 05:10

WindChimeRan force-pushed the ran/mps-rope branch from 90a8958 to 809fb71 Compare October 8, 2026 05:38

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

WindChimeRan mentioned this pull request Oct 13, 2026

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

Merged

WindChimeRan force-pushed the ran/mps-rope branch from 809fb71 to b32095f Compare October 14, 2026 01:12

WindChimeRan marked this pull request as ready for review October 14, 2026 01:15

WindChimeRan force-pushed the ran/mps-rope branch 2 times, most recently from 8bba8ef to 47354f3 Compare October 14, 2026 05:02

WindChimeRan mentioned this pull request Oct 14, 2026

[PyTorch-MPS][Omni] Serve Qwen3-TTS through upstream Omni #1109

Merged

Base automatically changed from ran/mps-backend-benchmark to main October 14, 2026 20:18

WindChimeRan force-pushed the ran/mps-rope branch from 47354f3 to a1ef0aa Compare October 14, 2026 20:18

ricky-chaoju approved these changes Oct 15, 2026

View reviewed changes

WindChimeRan merged commit f2e26b3 into main Oct 15, 2026
3 checks passed

Navigation Menu

Repository navigation

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

[PyTorch-MPS] Fuse Q/K RoPE with a bundled Metal kernel#1052
WindChimeRan merged 1 commit into
mainfrom
ran/mps-rope

WindChimeRan commented Oct 7, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

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

Uh oh!

Parity

Performance

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

WindChimeRan commented Oct 7, 2026 •

edited

Loading