minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+1
View File
@@ -376,6 +376,7 @@ class Envs:
# MPS (Apple Silicon)
SGLANG_USE_MLX = EnvBool(False)
SGLANG_MLX_USE_CUSTOM_ROPE = EnvBool(False)
# NPU
SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False)
python/sglang/srt/hardware_backend/mlx/aot.py
+233
View File
@@ -0,0 +1,233 @@
"""AOT kernel selection and decode-context helpers for the MLX backend."""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Optional
import mlx.core as mx
from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.hardware_backend.mlx.kv_cache.contiguous_cache import (
ContiguousKVCache,
)
def _load_metal_rope_pool_fused():
try:
from sgl_kernel import metal
except ImportError as exc:
raise ImportError(
"sgl_kernel.metal is not importable. Install sgl-kernel in the "
"active environment before enabling SGLANG_MLX_USE_CUSTOM_ROPE."
) from exc
import_error = getattr(metal, "_IMPORT_ERROR", None)
if getattr(metal, "_metal", None) is None or import_error is not None:
reason = f" Reason: {import_error}." if import_error is not None else ""
raise ImportError(
"sgl_kernel.metal is importable, but the native Metal extension "
f"or metallib is not available.{reason} Install the Metal kernels "
"with `uv run sgl-kernel/setup_metal.py install` from the SGLang "
"repo root in the active environment."
) from import_error
return metal.rope_pool_fused
@dataclass
class MlxAOTRoPEKernel:
base: float = 0.0
config: dict[str, Any] = field(default_factory=dict)
rope_pool_fused: Optional[Any] = None
@property
def enabled(self) -> bool:
return (
self.base > 0.0 and bool(self.config) and self.rope_pool_fused is not None
)
@dataclass
class MlxAOTKernelBuildInputs:
sample_attn: Any
n_kv_heads: int
head_dim: int
@dataclass(frozen=True)
class MlxAOTKernelSpec:
name: str
kernel_attr: str
is_enabled: Callable[[], bool]
build: Callable[[MlxAOTKernelBuildInputs], Any]
@dataclass
class MlxAOTKernelSet:
rope: MlxAOTRoPEKernel = field(default_factory=MlxAOTRoPEKernel)
selected_kernel_names: tuple[str, ...] = ()
class MlxAOTKernelRegistry:
"""Registry for optional MLX AOT kernels.
Each spec owns one kernel field on ``MlxAOTKernelSet``. The registry is the
only place that checks kernel enablement policy and model support.
"""
def __init__(self, specs: tuple[MlxAOTKernelSpec, ...]):
self._specs = specs
@property
def registered_kernel_names(self) -> tuple[str, ...]:
return tuple(spec.name for spec in self._specs)
def build_kernel_set(
self,
*,
sample_attn: Any,
n_kv_heads: int,
head_dim: int,
) -> MlxAOTKernelSet:
inputs = MlxAOTKernelBuildInputs(
sample_attn=sample_attn,
n_kv_heads=n_kv_heads,
head_dim=head_dim,
)
kernel_set = MlxAOTKernelSet()
selected_kernel_names = []
for spec in self._specs:
if not spec.is_enabled():
continue
kernel = spec.build(inputs)
if getattr(kernel, "enabled", False):
if not hasattr(kernel_set, spec.kernel_attr):
raise ValueError(
f"AOT kernel {spec.name} targets unknown kernel-set "
f"attribute {spec.kernel_attr}"
)
setattr(kernel_set, spec.kernel_attr, kernel)
selected_kernel_names.append(spec.name)
kernel_set.selected_kernel_names = tuple(selected_kernel_names)
if kernel_set.selected_kernel_names:
logger.info(
"MLX AOT kernels selected: %s",
", ".join(kernel_set.selected_kernel_names),
)
return kernel_set
def _build_rope_kernel(inputs: MlxAOTKernelBuildInputs) -> MlxAOTRoPEKernel:
sample_attn = getattr(inputs.sample_attn, "_inner", inputs.sample_attn)
rope = getattr(sample_attn, "rope", None)
if rope is None or getattr(rope, "traditional", False):
return MlxAOTRoPEKernel()
rope_dim = int(getattr(rope, "dims", 0))
if rope_dim == 0:
return MlxAOTRoPEKernel()
if rope_dim != inputs.head_dim:
# AOT kernel currently requires rope_dim == head_dim.
return MlxAOTRoPEKernel()
base = float(getattr(rope, "base", 10000.0))
config = {
"head_dim": int(inputs.head_dim),
"rope_dim": rope_dim,
"num_qo_heads": int(sample_attn.n_heads),
"num_kv_heads": int(inputs.n_kv_heads),
}
try:
rope_pool_fused = _load_metal_rope_pool_fused()
except Exception as exc: # noqa: BLE001
logger.info(
"AOT Metal RoPE kernel not available (%s) - falling back to "
"mx.fast.rope.",
exc,
)
return MlxAOTRoPEKernel()
logger.info(
f"AOT Metal RoPE kernel ENABLED: head_dim={inputs.head_dim}, "
f"n_heads={config['num_qo_heads']}, n_kv={config['num_kv_heads']}, "
f"base={base}"
)
return MlxAOTRoPEKernel(
base=base,
config=config,
rope_pool_fused=rope_pool_fused,
)
MLX_AOT_KERNEL_REGISTRY = MlxAOTKernelRegistry(
specs=(
MlxAOTKernelSpec(
name="metal_rope_pool_fused",
kernel_attr="rope",
is_enabled=lambda: envs.SGLANG_MLX_USE_CUSTOM_ROPE.get(),
build=_build_rope_kernel,
),
)
)
@dataclass
class MlxAOTRoPEContext:
kernel: MlxAOTRoPEKernel
kv_pool: Any
new_token_slots: Optional[mx.array] = None
@dataclass
class MlxAOTKernelContext:
rope: Optional[MlxAOTRoPEContext] = None
@classmethod
def from_decode(
cls,
*,
aot_kernels: MlxAOTKernelSet,
kv_pool: Any | None,
req_ids: list[str],
req_pool_idx: dict[str, int],
req_to_token_pool: Any | None,
layer_caches: list[list[ContiguousKVCache]],
) -> "MlxAOTKernelContext":
"""Build optional AOT context for one batched decode step."""
if not aot_kernels.rope.enabled or kv_pool is None:
return cls()
new_token_slots = None
if req_to_token_pool is not None:
try:
slot_ids = []
for req_idx, req_id in enumerate(req_ids):
pool_idx = req_pool_idx.get(req_id)
if pool_idx is None:
raise KeyError(req_id)
slot = int(
req_to_token_pool.req_to_token[
pool_idx, layer_caches[0][req_idx].offset
].item()
)
slot_ids.append(slot)
new_token_slots = mx.array(slot_ids, dtype=mx.int32)
except Exception as exc: # noqa: BLE001
logger.warning(
"AOT RoPE: failed to resolve new-token slots (%s); "
"falling back to RoPE-only for this decode step",
exc,
)
return cls(
rope=MlxAOTRoPEContext(
kernel=aot_kernels.rope,
kv_pool=kv_pool,
new_token_slots=new_token_slots,
)
)
python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py
+112 -4
View File
@@ -9,6 +9,11 @@ from typing import Any, Optional
import mlx.core as mx
import mlx.nn as nn
from sglang.srt.hardware_backend.mlx.aot import (
MlxAOTKernelContext,
MlxAOTKernelSet,
MlxAOTRoPEContext,
)
from sglang.srt.hardware_backend.mlx.kv_cache.contiguous_cache import ContiguousKVCache
_thread_local = threading.local()
@@ -24,6 +29,11 @@ class BatchedDecodeContext:
# layer_caches[layer_idx][req_idx] = ContiguousKVCache
layer_caches: list[list[ContiguousKVCache]]
# Optional AOT kernel state. Keep kernel-specific fields out of the regular
# MLX decode path so future AOT kernels can be added without growing this
# context one field at a time.
aot: MlxAOTKernelContext = field(default_factory=MlxAOTKernelContext)
# Derived tensors/metadata, shared across all layers in one forward pass.
offsets: mx.array = field(init=False)
max_len: int = field(init=False)
@@ -42,6 +52,38 @@ class BatchedDecodeContext:
self.pad_sizes = [max_seq_len - s for s in seq_lens]
self.positions = mx.arange(self.max_len) if self.needs_padding else None
@classmethod
def from_decode(
cls,
*,
caches: list[list[ContiguousKVCache]],
num_layers: int,
req_ids: list[str],
aot_kernels: MlxAOTKernelSet,
kv_pool: Any | None,
req_pool_idx: dict[str, int],
req_to_token_pool: Any | None,
) -> "BatchedDecodeContext":
batch_size = len(req_ids)
seq_lens = [caches[i][0].offset for i in range(batch_size)]
layer_caches = [
[caches[i][layer_idx] for i in range(batch_size)]
for layer_idx in range(num_layers)
]
return cls(
batch_size=batch_size,
seq_lens=seq_lens,
layer_caches=layer_caches,
aot=MlxAOTKernelContext.from_decode(
aot_kernels=aot_kernels,
kv_pool=kv_pool,
req_ids=req_ids,
req_pool_idx=req_pool_idx,
req_to_token_pool=req_to_token_pool,
layer_caches=layer_caches,
),
)
def set_context(ctx: Optional[BatchedDecodeContext]) -> None:
_thread_local.batched_ctx = ctx
@@ -96,13 +138,26 @@ class MLXAttentionWrapper(nn.Module):
keys = keys.transpose(0, 2, 1, 3)
values = values.transpose(0, 2, 1, 3)
# Vectorized RoPE with per-batch offsets
# Vectorized RoPE with per-batch offsets (cached on the context).
offsets = ctx.offsets
queries = inner.rope(queries, offset=offsets)
keys = inner.rope(keys, offset=offsets)
if ctx.aot.rope is not None:
# AOT path: real .metallib RoPE + fused KV pool scatter.
queries, keys = self._rope_custom_aot(
queries,
keys,
values,
offsets,
layer_idx,
ctx.aot.rope,
)
else:
# Fallback: MLX's built-in mx.fast.rope (used when the AOT kernel
# isn't built or the model uses an unsupported RoPE variant).
queries = inner.rope(queries, offset=offsets)
keys = inner.rope(keys, offset=offsets)
layer_caches = ctx.layer_caches[layer_idx]
max_len = ctx.max_len
pad_sizes = ctx.pad_sizes
# TODO: replace per-request loop with native batched/ragged
@@ -148,3 +203,56 @@ class MLXAttentionWrapper(nn.Module):
output = output.transpose(0, 2, 1, 3).reshape(B, 1, -1)
return inner.o_proj(output)
@staticmethod
def _rope_custom_aot(
queries: mx.array,
keys: mx.array,
values: mx.array,
positions: mx.array,
layer_idx: int,
rope_ctx: MlxAOTRoPEContext,
) -> tuple[mx.array, mx.array]:
"""AOT path: rotate Q/K and scatter K/V into the shared pool.
The kernel call does RoPE on Q/K and scatters
rotated K + (untouched) V into ``kv_pool`` at ``new_token_slots``
for ``layer_idx``.
If ``new_token_slots`` is None, slot=-1 sentinel is used (no pool
write, RoPE-only mode). Returns rotated (queries, keys) in the
original 4-D attention layout. ``values`` is unchanged by RoPE.
"""
# (B, n_heads, 1, head_dim) -> (B, n_heads, head_dim) for kernel
q_flat = queries[:, :, 0, :]
k_flat = keys[:, :, 0, :]
v_flat = values[:, :, 0, :]
B = q_flat.shape[0]
if rope_ctx.new_token_slots is None:
slots = mx.full((B,), -1, dtype=mx.int32)
else:
slots = rope_ctx.new_token_slots.astype(mx.int32)
k_pool = rope_ctx.kv_pool.k_buffer[layer_idx]
v_pool = rope_ctx.kv_pool.v_buffer[layer_idx]
q_rot, k_rot, k_pool_new, v_pool_new = rope_ctx.kernel.rope_pool_fused(
q_flat,
k_flat,
v_flat,
positions,
slots,
k_pool,
v_pool,
head_dim=rope_ctx.kernel.config["head_dim"],
num_qo_heads=rope_ctx.kernel.config["num_qo_heads"],
num_kv_heads=rope_ctx.kernel.config["num_kv_heads"],
rope_base=rope_ctx.kernel.base,
)
# Rebind pool buffers (zero-copy donation result).
rope_ctx.kv_pool.k_buffer[layer_idx] = k_pool_new
rope_ctx.kv_pool.v_buffer[layer_idx] = v_pool_new
# (B, n_heads, head_dim) -> (B, n_heads, 1, head_dim) for SDPA path
return q_rot[:, :, None, :], k_rot[:, :, None, :]
python/sglang/srt/hardware_backend/mlx/model_runner.py
+40 -21
View File
@@ -24,6 +24,10 @@ from mlx.utils import tree_flatten
from mlx_lm import load as mlx_lm_load
from mlx_lm.utils import quantize_model as mlx_lm_quantize_model
from sglang.srt.hardware_backend.mlx.aot import (
MLX_AOT_KERNEL_REGISTRY,
MlxAOTKernelSet,
)
from sglang.srt.hardware_backend.mlx.kv_cache import (
BatchedDecodeContext,
ContiguousKVCache,
@@ -148,6 +152,7 @@ class MlxModelRunner:
self._req_synced_offset: dict[str, int] = {}
self._pool_size = self._compute_pool_size(pool_size)
self._aot_kernels = self._build_aot_kernels()
@staticmethod
def _extract_logits(model_output):
@@ -313,6 +318,19 @@ class MlxModelRunner:
def pool_size(self) -> int:
return self._pool_size
def _build_aot_kernels(self) -> MlxAOTKernelSet:
"""Build model-level set of optional registered AOT kernels."""
layer_list, attn_attr = find_attention_layers(self.model)
if not layer_list:
return MlxAOTKernelSet()
sample_attn = getattr(layer_list[0], attn_attr)
n_kv_heads, head_dim, _ = self._get_attn_config()
return MLX_AOT_KERNEL_REGISTRY.build_kernel_set(
sample_attn=sample_attn,
n_kv_heads=int(n_kv_heads),
head_dim=int(head_dim),
)
def init_kv_pool(self, req_to_token_pool: ReqToTokenPool) -> None:
"""Create MlxKVPool (+1 for padding slot 0) and wire scheduler pools."""
self._req_to_token_pool = req_to_token_pool
@@ -378,7 +396,7 @@ class MlxModelRunner:
end = cache_start + len(slot_ids)
slot_ids_mx = mx.array(slot_ids, dtype=mx.int32)
# TODO: Standardize ContiguousKVCache size to avoid transpose
# Transpose cache (1, n_kv_heads, S, head_dim) → pool (S, n_kv_heads, head_dim)
# Transpose cache (1, n_kv_heads, S, head_dim) to pool (S, n_kv_heads, head_dim)
k_all = mx.stack(
[
cache[i].keys[0, :, cache_start:end, :].transpose(1, 0, 2)
@@ -488,7 +506,7 @@ class MlxModelRunner:
if new_token_count > 0:
extend_tokens = new_token_ids
else:
# Full cache hit — rerun last token to get next-token logits
# Full cache hit - rerun last token to get next-token logits
extend_tokens = full_token_ids[-1:]
for c in cache:
c.offset = max(c.offset - 1, 0)
@@ -500,7 +518,7 @@ class MlxModelRunner:
last_logits = logits[:, -1, :]
lazy_token = mx.argmax(last_logits, axis=-1)
# Convert PoolBackedCache → ContiguousKVCache for decode.
# Convert PoolBackedCache to ContiguousKVCache for decode.
# This appends a lazy slice-assign onto the forward graph; the
# arrays get materialised when the caller evaluates lazy_token.
if prefix_len > 0:
@@ -610,16 +628,16 @@ class MlxModelRunner:
caches=caches,
)
seq_lens = [caches[i][0].offset for i in range(batch_size)]
layer_caches = [
[caches[i][layer_idx] for i in range(batch_size)]
for layer_idx in range(num_layers)
]
ctx = BatchedDecodeContext(
batch_size=batch_size,
seq_lens=seq_lens,
layer_caches=layer_caches,
ctx = BatchedDecodeContext.from_decode(
caches=caches,
num_layers=num_layers,
req_ids=req_ids,
aot_kernels=self._aot_kernels,
kv_pool=self._kv_pool,
req_pool_idx=self._req_pool_idx,
req_to_token_pool=self._req_to_token_pool,
)
seq_lens = ctx.seq_lens
set_context(ctx)
try:
max_offset = max(seq_lens)
@@ -668,7 +686,7 @@ class MlxModelRunner:
# to accommodate dynamic growing like ContiguousKVCache.update_and_fetch.
# After prev's graph ran, each ContiguousKVCache.offset was
# bumped by one per layer — attention wrapper's `write_token`
# bumped by one per layer - attention wrapper's `write_token`
# mutates the Python offset synchronously at graph-build time.
# So layer-0 offsets reflect the position the NEW token will
# be written at in step N+1 (and equivalently the RoPE offset).
@@ -686,15 +704,16 @@ class MlxModelRunner:
caches=caches,
)
layer_caches = [
[caches[i][layer_idx] for i in range(batch_size)]
for layer_idx in range(num_layers)
]
ctx = BatchedDecodeContext(
batch_size=batch_size,
seq_lens=seq_lens,
layer_caches=layer_caches,
ctx = BatchedDecodeContext.from_decode(
caches=caches,
num_layers=num_layers,
req_ids=prev.req_ids,
aot_kernels=self._aot_kernels,
kv_pool=self._kv_pool,
req_pool_idx=self._req_pool_idx,
req_to_token_pool=self._req_to_token_pool,
)
seq_lens = ctx.seq_lens
set_context(ctx)
try:
max_offset = max(seq_lens)