[Apple Silicon] Add custom Metal RoPE kernel with fused KV cache store

[adityavaid]

Motivation

This PR adds the first AOT-compiled Metal kernel to SGLang's Apple Silicon backend: a fused NeoX RoPE + KV pool scatter that replaces two mx.fast.rope calls and the per-request KV-pool scatter in batched decode with a single C++ entry that emits one Metal command buffer.

On Apple Silicon, each kernel dispatch incurs ~1-4ms of GPU idle time due to Metal command buffer overhead (profiled in #22114). Reducing dispatch count directly improves decode latency. This kernel consolidates the Q and K rotation into one dispatch, cutting the RoPE step from ~224us (2 × 112us) to ~127us — a 1.8x speedup for the RoPE operation.
Part of the Apple Device Support roadmap (#19137)

NOTE : this builds on #23449's scaffolding

Kernel logic

A single .metallib produced at build time by Apple's xcrun metal +
xcrun metallib (same pipeline MLX itself uses), exporting three
specialised kernels per dtype (f16, bf16, f32). All three are dispatched
into one Metal command buffer per batched-decode forward, per layer:

Kernel	Work
rope_q	NeoX RoPE on Q
rope_k_pool	NeoX RoPE on K and scatter rotated K to k_pool[slot]
v_to_pool	Scatter V to v_pool[slot]

Loaded once at import via register_library(path). Per-shape pipelines
are specialised through Metal function constants so a single .metallib
serves every model configuration.

Optimisations over a naïve fused-RoPE kernel

( Inspired by MLX's own rope.metal )
Baseline : single 1-D-grid kernel that recovers (token, head, dim) via div/mod and branches on is_q.

Layered on top of that:

1. AOT .metallib instead of runtime JIT — Apple's full compiler-optimisation pass at build time.
   2.. 3D thread-grid (pos.x = dim, pos.y = token, pos.z = head) — eliminates the div/mod hot path.
2. Smart threadgroup sizing — fills 4–8 SIMD groups per threadgroup instead of the naïve (32, 1, 1) that uses 1.
3. Function-constant specialisation of shape parameters (HEAD_DIM, NUM_QO_HEADS, NUM_KV_HEADS, INV_DIM_LOG2_BASE, HEADS_PER_THREAD) — compiler folds constants per pipeline.
4. Inline trig (metal::exp2 + metal::fast::cos/sin) — no precomputed cos/sin cache.
5. Heads-per-thread amortisation infra (HEADS_PER_THREAD fn-const, SGLANG_RPF_N env override). Default N=1 — kernel is launch-overhead-bound at our shapes; N>1 left for future hardware.
6. mlx::core::Primitive integration — required for kernel writes to be visible to MLX's lazy graph and downstream ops.
7. Fused KV-pool write — rotated K + V land in kv_pool[slot] in the same command buffer, eliminating the separate pool[slots] = … scatters that would otherwise be needed after RoPE.
8. Zero-copy pool donation — copy_shared_buffer returns the input pool buffer as the output; the user-side handle is rebound, GPU memory is the same allocation.

Accuracy Tests

( Adding Unit Tests )

E2E server test with Qwen/Qwen3-0.6B:

Check	Result
Startup log: Custom RoPE kernel enabled	Present
Single-request inference	Correct text output
KV pool initialized	42165 slots × 28 layers

Benchmarking and Profiling

Kernel microbenchmark on Apple M4 Pro (48GB), Qwen3-0.6B dimensions (32 tokens, 16 Q heads, 2 KV heads, head_dim=64)

Focused local profile timing:

Case	MLX baseline	AOT fused	Speedup
LLaMA-like bf16 BS=512 (nq=32, nk=8, hd=128)	88.13 us	64.23 us	1.37x
Qwen-like bf16 BS=512 (nq=16, nk=2, hd=64)	28.22 us	39.96 us	0.71x

Large-batch bf16 sweep:

Config	MLX median	AOT fused median	Speedup
Qwen3 bf16 BS=128	14.56 us	17.65 us	0.82x
Qwen3 bf16 BS=256	16.88 us	24.78 us	0.68x
Qwen3 bf16 BS=512	21.40 us	33.22 us	0.64x
Qwen3 bf16 BS=1024	32.85 us	38.40 us	0.86x
LLaMA bf16 BS=128	29.21 us	33.46 us	0.87x
LLaMA bf16 BS=256	45.73 us	41.07 us	1.11x
LLaMA bf16 BS=512	79.57 us	57.74 us	1.38x
LLaMA bf16 BS=1024	156.89 us	115.72 us	1.36x

Takeaway: this is not a universal speedup. The fused AOT path helps on larger
LLaMA-like shapes, especially larger bf16 batches, but regresses on Qwen-like
small-KV-head shapes. The PR keeps the implementation guarded and easy to A/B.

How to Run Tests

cd /path/to/sglang
source .venv-mlx-dev/bin/activate

# Full server E2E (downloads Qwen3-0.6B on first run)
SGLANG_USE_MLX=1 python3 -m sglang.launch_server \
  --model-path Qwen/Qwen3-0.6B --port 43440 --mem-fraction-static 0.5
# Look for "Custom RoPE kernel enabled" in logs

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.

• /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci

1. After green CI and required approvals, ask Merge Oncalls to merge.

---

CI States

Latest PR Test (Base): ❌ Run #26619584932
Latest PR Test (Extra): ❌ Run #26619584836

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive MLX-based KV cache system and model runner for Apple Silicon, enabling features like radix-cache prefix sharing and custom Metal RoPE kernels. The implementation includes a flat KV pool (MlxKVPool), contiguous request-local buffers (ContiguousKVCache), and attention patching to support batched decoding. However, several critical issues were identified: the custom Metal RoPE kernels do not correctly handle models with partial RoPE (where rope_dim < head_dim), leading to corrupted KV caches. Additionally, the write_token method in ContiguousKVCache lacks necessary allocation and bounds checks, which will cause runtime errors during the first token write. There is also an opportunity to optimize the batched decode logic by replacing a Python loop with a more efficient batched concatenation of pre-allocated buffers.

[yeahdongcn]

Nice work! Since this is the first Metal kernel for SGLang, it's important to establish the right approach for integrating it. Future kernels will likely follow this pattern, so getting the foundation right is critical.

You may take a look at how https://github.com/vllm-project/vllm-metal/tree/880af64e95f75832649d55d8260ad823244fc8b0/vllm_metal/metal/kernels_v2 works.

[alexnails]

I will have more comments, just doing a quick review

[yeahdongcn]

@Jonahcb is also working on custom Metal kernels (paged attention). Maybe we can discuss the best way to integrate them.

[alexnails]

do you have profile?

[adityavaid]

@alexnails @yeahdongcn Updated the PR .

[yeahdongcn]

Should we still care about these?

[adityavaid]

Keeping it now to facilitate fallback and validation

[yeahdongcn]

I think _import_sgl_kernel_metal already provides sufficient protection at the call site, so we can remove these checks.

[yeahdongcn]

Based on the performance data, I didn't see a clear overall win (large batch size seems to be meaningless for mac) from the AOT kernels (though this is our first attempt at adding custom AOT kernels to the MLX backend). So I think we should introduce something like SGLANG_MLX_USE_CUSTOM_ROPE and keep the default value as false.

[adityavaid]

Ack, Changing

[yeahdongcn]

Please consider introducing an explicit abstraction before adding more fields directly to the runner/context, for example:

@dataclass
class MlxAOTKernelContext:
    rope: Optional[MlxAOTRoPEContext] = None

@dataclass
class MlxAOTRoPEContext:
    config: dict
    base: float
    kv_pool: MlxKVPool
    new_token_slots: Optional[mx.array]

Then the model runner can have one helper such as _build_aot_kernel_context(...), and BatchedDecodeContext only needs a single aot field. That makes it clear which data belongs to optional AOT kernels and keeps the regular MLX decode path from accumulating kernel-specific kwargs.

[adityavaid]

done, Claude is amazing.

[yeahdongcn]

Just pushed a new commit to make the MLX AOT kernel selection into a backend-level registry (we will have more AOT kernels in the future). @adityavaid feel free to drop or update.

[yeahdongcn]

/tag-and-rerun-ci

[adityavaid]

/rerun-failed-ci

dismiss stale review