[Apple Silicon] Add custom Metal RoPE kernel with fused KV cache store (#22868)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Co-authored-by: Xiaodong Ye <yeahdongcn@gmail.com>
This commit is contained in:
Aditya Sharma
2026-05-29 15:09:33 +08:00
committed by GitHub
co-authored by Xiaodong Ye
parent 7dff4118b9
commit b2eed9e16d
12 changed files with 1066 additions and 40 deletions
@@ -59,6 +59,7 @@ SGLANG_USE_MLX=1 python -m sglang.launch_server \
1. `SGLANG_USE_MLX=1` - Enables the use of MLX as the SGLang runtime backend (if disabled, SGLang will fall back to `torch.mps`, which has less support)
2. `--disable-cuda-graph` - Disables usage of CUDA graph, which is not relevant for Apple Metal.
3. `--disable-overlap-schedule` - Disables overlap scheduling (enabled/not present by default) achieved using MLX's `async_eval()`
4. `SGLANG_MLX_USE_CUSTOM_ROPE=1` - Enables the optional custom Metal RoPE kernel. It is disabled by default, so the MLX backend uses the standard RoPE path unless you opt in for A/B testing.
## Quantization
@@ -84,7 +85,6 @@ The MLX backend supports two quantization paths on Apple Silicon:
```
The MLX backend silently ignores `--quantization mlx_q4` when the model is already quantized in its HF config (path 1), so the same flag is safe to pass either way.
## Benchmarking with Requests
`sglang.benchmark_one_batch` calls the synchronous prefill/decode methods directly without going through the scheduler and the overlap code path.