Skip to content

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

Closed
WindChimeRan wants to merge 2 commits into
stack/mps-upstream-modelsfrom
stack/mps-sampling-dispatch
Closed

WindChimeRan wants to merge 2 commits into
stack/mps-upstream-modelsfrom
stack/mps-sampling-dispatch

Conversation

@WindChimeRan

@WindChimeRan WindChimeRan commented Sep 28, 2026 •

Copy link
Copy Markdown
Collaborator

Stacked on #859. Run vLLM's unchanged sampling kernels through Triton's Apple backend.

  • ./install.sh --mps installs the pinned Torch/Triton/Apple-backend dependencies in .venv-vllm-metal-mps, fetches the checksum-verified LLVM SDK, builds in dependency order, and verifies an upstream sampling kernel.
  • The installer applies the temporary vLLM 0.30 CPU-wheel detection workaround only in that isolated environment. Default MLX installation is unchanged.
  • No local sampling algorithms or copied kernels. Bad-word filtering stays disabled pending the deferred upstream compiler fix.
./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

Validated: clean install, repeat install, dependency consistency, Qwen3 sampling/logprobs and prefix reuse. The compiler backend is pinned to triton-ext #126; this remains an experimental draft.

Qwen3 FP16, T=0.6: 337.8 → 455.4 tokens/s (+34.8%) versus MLX, two-run means with the same Sonnet 1024/128, 100-request, 10/s, C32, 2,340-block, prefix-off settings as #858. MLX native sampling is enabled, including its mixed-step CPU fallback.

@WindChimeRan
WindChimeRan added this pull request to stack #860 September 28, 2026 09:18
@WindChimeRan
WindChimeRan force-pushed the stack/mps-upstream-models branch from 676fdfe to c6b41fa Compare September 28, 2026 09:30
@WindChimeRan
WindChimeRan force-pushed the stack/mps-sampling-dispatch branch from 718581c to 7515140 Compare September 28, 2026 09:30
@WindChimeRan WindChimeRan changed the title [PyTorch-MPS][MRv2] Use upstream sampling kernel dispatch [PyTorch-MPS][MRv2] Add MPS sampling with upstream kernel dispatch Sep 28, 2026
Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>
@WindChimeRan
WindChimeRan force-pushed the stack/mps-sampling-dispatch branch from 7515140 to e0bff59 Compare September 28, 2026 11:20
@WindChimeRan WindChimeRan changed the title [PyTorch-MPS][MRv2] Add MPS sampling with upstream kernel dispatch [PyTorch-MPS][MRv2] Use upstream Triton sampling directly on MPS Sep 28, 2026
Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>
@WindChimeRan

Copy link
Copy Markdown
Collaborator Author

Closing this draft because it depends on the unmerged Triton Apple backend, triton-lang/triton-ext#126.

We’ll follow up after vllm-project/vllm#43048 reaches a vLLM release, using TileLang sampling kernels shared by the MLX and PyTorch MPS paths with small framework-specific adapters.

@WindChimeRan
WindChimeRan removed this pull request from stack #860 September 28, 2026 20:31
@imperatormk

imperatormk commented Sep 28, 2026 •

Copy link
Copy Markdown

@WindChimeRan Thank you for getting this to run although I think it's gonna take some while until the backend gets to be released officially. Fixed the bad words kernel bug either way and now the system runs properly e2e for me (after some UVA related fixes in copy_changed_state, might post it up as an issue but I assume it's inevitable you'll come across that yourself)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants