[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
co-authored by
Xiaodong Ye
Xiaodong Ye
parent
7dff4118b9
commit
b2eed9e16d
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
)
|
||||
@@ -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, :]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -1,21 +1,17 @@
|
||||
# sgl-kernel Metal kernels
|
||||
|
||||
Custom Apple Metal kernels for the MLX backend on Apple Silicon. Shader sources (`*.metal`) and C++ host / nanobind sources (`*.cpp`) in this directory are compiled by [`sgl-kernel/setup_metal.py`](../../setup_metal.py) into the `sgl_kernel._metal` extension and the `sgl_metal_kernels.metallib` archive, and exposed through Python wrappers in [`python/sgl_kernel/metal.py`](../../python/sgl_kernel/metal.py).
|
||||
Custom Apple Metal kernels for the MLX backend on Apple Silicon. Shader sources (`*.metal`) and C++ host / nanobind sources (`*.cpp`) in this directory are compiled by [`sgl-kernel/setup_metal.py`](../../setup_metal.py) into the native Metal extension and the `sgl_metal_kernels.metallib` archive, then exposed through public Python wrappers in [`python/sgl_kernel/metal.py`](../../python/sgl_kernel/metal.py).
|
||||
|
||||
## Kernels
|
||||
|
||||
| Kernel | Description | Tested on |
|
||||
| --- | --- | --- |
|
||||
| _none yet_ | — | — |
|
||||
| `rope_pool_fused` | Fused NeoX RoPE for Q/K plus K/V scatter into the MLX KV pool. | Apple Silicon / MLX |
|
||||
|
||||
## Adding a new Metal kernel
|
||||
|
||||
1. Add the shader under `csrc/metal/<kernel>.metal`.
|
||||
2. Add the C++ host / nanobind binding under `csrc/metal/<kernel>.cpp`, exporting the entry point on the `sgl_kernel._metal` module.
|
||||
2. Add the C++ host / nanobind binding under `csrc/metal/<kernel>.cpp`, exporting the native entry point for the wrapper in [`python/sgl_kernel/metal.py`](../../python/sgl_kernel/metal.py).
|
||||
3. Append both files to `metal_shader_sources` and `cxx_sources` in [`sgl-kernel/setup_metal.py`](../../setup_metal.py).
|
||||
4. Add a Python wrapper in [`python/sgl_kernel/metal.py`](../../python/sgl_kernel/metal.py) that validates input shapes/dtypes and calls `mx.eval` on its operands before invoking the AOT C++ entry point.
|
||||
4. Add a Python wrapper in [`python/sgl_kernel/metal.py`](../../python/sgl_kernel/metal.py) that validates input shapes/dtypes and invokes the native AOT entry point without forcing MLX evaluation.
|
||||
5. Add a test under [`sgl-kernel/tests/`](../../tests) and update the **Kernels** table above with a short description and the hardware / OS / MLX version the kernel was validated on.
|
||||
|
||||
## Note on `placeholder.metal` / `placeholder.cpp`
|
||||
|
||||
`placeholder.metal` and `placeholder.cpp` are intentionally empty. They exist only so that `setup_metal.py` has at least one shader source and one C++ source to compile, allowing the `sgl_kernel._metal` extension and the `sgl_metal_kernels.metallib` archive to build successfully before any real Metal kernels have been added. Both files (and their entries in `metal_shader_sources` / `cxx_sources` in `setup_metal.py`) MUST be removed once the first real kernel lands.
|
||||
|
||||
@@ -0,0 +1,325 @@
|
||||
// Combined optimal: real AOT .metallib + Primitive integration + optimized
|
||||
// 3-kernel + 3D-grid dispatch + fused KV pool write.
|
||||
|
||||
#include <nanobind/nanobind.h>
|
||||
#include <nanobind/stl/string.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
#include "mlx/allocator.h"
|
||||
#include "mlx/array.h"
|
||||
#include "mlx/backend/metal/device.h"
|
||||
#include "mlx/mlx.h"
|
||||
#include "mlx/primitives.h"
|
||||
#include "mlx/stream.h"
|
||||
|
||||
namespace nb = nanobind;
|
||||
using namespace mlx::core;
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr const char* kLibraryName = "sgl_metal_kernels";
|
||||
|
||||
MTL::Library* g_library = nullptr;
|
||||
|
||||
const char* dtype_suffix(Dtype dt) {
|
||||
switch (dt) {
|
||||
case float16:
|
||||
return "f16";
|
||||
case bfloat16:
|
||||
return "bf16";
|
||||
case float32:
|
||||
return "f32";
|
||||
default:
|
||||
throw std::runtime_error("rope_pool_fused: unsupported dtype");
|
||||
}
|
||||
}
|
||||
|
||||
void register_library_impl(const std::string& path) {
|
||||
if (path.empty()) {
|
||||
throw std::runtime_error("register_library requires a non-empty path");
|
||||
}
|
||||
auto& d = metal::device(Device::gpu);
|
||||
g_library = d.get_library(kLibraryName, path);
|
||||
if (g_library == nullptr) {
|
||||
throw std::runtime_error("failed to load .metallib from: " + path);
|
||||
}
|
||||
}
|
||||
|
||||
MTL::Size pick_tg(uint32_t gx, uint32_t gy, uint32_t gz) {
|
||||
constexpr uint32_t kMaxThreads = 256;
|
||||
uint32_t tx = std::min<uint32_t>(gx, 32u);
|
||||
uint32_t ty = std::min<uint32_t>(gy, kMaxThreads / std::max<uint32_t>(tx, 1u));
|
||||
uint32_t tz = std::min<uint32_t>(gz, kMaxThreads / std::max<uint32_t>(tx * ty, 1u));
|
||||
while (ty > 1 && (gy % ty) != 0)
|
||||
--ty;
|
||||
while (tz > 1 && (gz % tz) != 0)
|
||||
--tz;
|
||||
return MTL::Size::Make(tx, std::max<uint32_t>(ty, 1u), std::max<uint32_t>(tz, 1u));
|
||||
}
|
||||
|
||||
uint32_t pick_heads_per_thread(uint32_t nh) {
|
||||
if (nh == 0) return 1;
|
||||
if (const char* e = std::getenv("SGLANG_RPF_N")) {
|
||||
uint32_t v = static_cast<uint32_t>(std::atoi(e));
|
||||
if (v >= 1 && nh % v == 0) return v;
|
||||
}
|
||||
return 1u;
|
||||
}
|
||||
|
||||
class RopePoolFused : public Primitive {
|
||||
public:
|
||||
RopePoolFused(Stream stream, int head_dim, int num_qo_heads, int num_kv_heads, float rope_base)
|
||||
: Primitive(stream),
|
||||
head_dim_(head_dim),
|
||||
num_qo_heads_(num_qo_heads),
|
||||
num_kv_heads_(num_kv_heads),
|
||||
rope_base_(rope_base) {}
|
||||
|
||||
void eval_cpu(const std::vector<array>&, std::vector<array>&) override {
|
||||
throw std::runtime_error("rope_pool_fused: CPU eval not supported");
|
||||
}
|
||||
|
||||
void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs) override {
|
||||
if (g_library == nullptr) {
|
||||
throw std::runtime_error("rope_pool_fused: register_library() not called yet");
|
||||
}
|
||||
auto& q = inputs[0];
|
||||
auto& k = inputs[1];
|
||||
auto& v = inputs[2];
|
||||
auto& positions = inputs[3];
|
||||
auto& slots = inputs[4];
|
||||
auto& k_pool_in = inputs[5];
|
||||
auto& v_pool_in = inputs[6];
|
||||
|
||||
auto& q_out = outputs[0];
|
||||
auto& k_out = outputs[1];
|
||||
auto& k_pool_out = outputs[2];
|
||||
auto& v_pool_out = outputs[3];
|
||||
|
||||
q_out.set_data(allocator::malloc(q_out.nbytes()));
|
||||
k_out.set_data(allocator::malloc(k_out.nbytes()));
|
||||
// Donate input pool buffers to outputs - zero-copy in-place semantics.
|
||||
k_pool_out.copy_shared_buffer(k_pool_in);
|
||||
v_pool_out.copy_shared_buffer(v_pool_in);
|
||||
|
||||
auto& d = metal::device(stream().device);
|
||||
|
||||
const uint32_t hd = static_cast<uint32_t>(head_dim_);
|
||||
const uint32_t nq = static_cast<uint32_t>(num_qo_heads_);
|
||||
const uint32_t nk = static_cast<uint32_t>(num_kv_heads_);
|
||||
const uint32_t half_dim = hd / 2;
|
||||
const uint32_t num_tokens = static_cast<uint32_t>(q.shape(0));
|
||||
const uint32_t hpt_q = pick_heads_per_thread(nq);
|
||||
const uint32_t hpt_k = pick_heads_per_thread(nk);
|
||||
const uint32_t hpt_v = pick_heads_per_thread(nk);
|
||||
|
||||
const float inv_dim_log2_base = std::log2(rope_base_) / static_cast<float>(head_dim_);
|
||||
|
||||
auto build_consts = [&](const uint32_t& hpt) {
|
||||
return metal::MTLFCList{
|
||||
{&hd, MTL::DataType::DataTypeUInt, 0},
|
||||
{&nq, MTL::DataType::DataTypeUInt, 1},
|
||||
{&nk, MTL::DataType::DataTypeUInt, 2},
|
||||
{&inv_dim_log2_base, MTL::DataType::DataTypeFloat, 3},
|
||||
{&hpt, MTL::DataType::DataTypeUInt, 4},
|
||||
};
|
||||
};
|
||||
|
||||
auto build_hash = [&](const std::string& kname, uint32_t hpt) {
|
||||
return kname + "_hd" + std::to_string(head_dim_) + "_q" + std::to_string(num_qo_heads_) + "_k" +
|
||||
std::to_string(num_kv_heads_) + "_n" + std::to_string(hpt) + "_b" +
|
||||
std::to_string(static_cast<int>(rope_base_));
|
||||
};
|
||||
|
||||
const std::string rect_kname = std::string("rope_pool_rect_") + dtype_suffix(q.dtype());
|
||||
const std::string q_kname = std::string("rope_q_") + dtype_suffix(q.dtype());
|
||||
const std::string k_kname = std::string("rope_k_pool_") + dtype_suffix(k.dtype());
|
||||
const std::string v_kname = std::string("v_to_pool_") + dtype_suffix(v.dtype());
|
||||
|
||||
// Single rectangular dispatch uses one HEADS_PER_THREAD value for both
|
||||
// Q and KV heads. Keep it valid for both head counts.
|
||||
uint32_t hpt = std::min(hpt_q, hpt_k);
|
||||
if (nq % hpt != 0 || nk % hpt != 0) hpt = 1;
|
||||
|
||||
auto& enc = metal::get_command_encoder(stream());
|
||||
|
||||
const bool use_rect_dispatch = q.dtype() == bfloat16 && hd >= 128 && nk >= 8 && num_tokens >= 256;
|
||||
if (use_rect_dispatch) {
|
||||
auto rect_consts = build_consts(hpt);
|
||||
auto* rect_pipe = d.get_kernel(rect_kname, g_library, build_hash(rect_kname, hpt), rect_consts);
|
||||
if (!rect_pipe) {
|
||||
throw std::runtime_error("rope_pool_fused: failed to resolve rectangular kernel");
|
||||
}
|
||||
const uint32_t max_heads = std::max(nq, nk);
|
||||
const uint32_t gz = (max_heads + hpt - 1) / hpt;
|
||||
enc.set_compute_pipeline_state(rect_pipe);
|
||||
enc.set_input_array(q, 0);
|
||||
enc.set_input_array(k, 1);
|
||||
enc.set_input_array(v, 2);
|
||||
enc.set_output_array(q_out, 3);
|
||||
enc.set_output_array(k_out, 4);
|
||||
enc.set_output_array(k_pool_out, 5);
|
||||
enc.set_output_array(v_pool_out, 6);
|
||||
enc.set_input_array(positions, 7);
|
||||
enc.set_input_array(slots, 8);
|
||||
enc.dispatch_threads(MTL::Size::Make(hd, num_tokens, gz), pick_tg(hd, num_tokens, gz));
|
||||
} else {
|
||||
auto q_consts = build_consts(hpt_q);
|
||||
auto k_consts = build_consts(hpt_k);
|
||||
auto v_consts = build_consts(hpt_v);
|
||||
|
||||
auto* q_pipe = d.get_kernel(q_kname, g_library, build_hash(q_kname, hpt_q), q_consts);
|
||||
auto* k_pipe = d.get_kernel(k_kname, g_library, build_hash(k_kname, hpt_k), k_consts);
|
||||
auto* v_pipe = d.get_kernel(v_kname, g_library, build_hash(v_kname, hpt_v), v_consts);
|
||||
if (!q_pipe || !k_pipe || !v_pipe) {
|
||||
throw std::runtime_error("rope_pool_fused: failed to resolve kernels");
|
||||
}
|
||||
|
||||
// Kernel 1: Q rope
|
||||
{
|
||||
enc.set_compute_pipeline_state(q_pipe);
|
||||
enc.set_input_array(q, 0);
|
||||
enc.set_output_array(q_out, 1);
|
||||
enc.set_input_array(positions, 2);
|
||||
const uint32_t gz = (nq + hpt_q - 1) / hpt_q;
|
||||
enc.dispatch_threads(MTL::Size::Make(half_dim, num_tokens, gz), pick_tg(half_dim, num_tokens, gz));
|
||||
}
|
||||
|
||||
// Kernel 2: K rope + pool write
|
||||
{
|
||||
enc.set_compute_pipeline_state(k_pipe);
|
||||
enc.set_input_array(k, 0);
|
||||
enc.set_output_array(k_out, 1);
|
||||
enc.set_output_array(k_pool_out, 2);
|
||||
enc.set_input_array(positions, 3);
|
||||
enc.set_input_array(slots, 4);
|
||||
const uint32_t gz = (nk + hpt_k - 1) / hpt_k;
|
||||
enc.dispatch_threads(MTL::Size::Make(half_dim, num_tokens, gz), pick_tg(half_dim, num_tokens, gz));
|
||||
}
|
||||
|
||||
// Kernel 3: V copy to pool
|
||||
{
|
||||
enc.set_compute_pipeline_state(v_pipe);
|
||||
enc.set_input_array(v, 0);
|
||||
enc.set_output_array(v_pool_out, 1);
|
||||
enc.set_input_array(slots, 2);
|
||||
enc.dispatch_threads(MTL::Size::Make(hd, num_tokens, nk), pick_tg(hd, num_tokens, nk));
|
||||
}
|
||||
}
|
||||
// No commit / synchronize - MLX's lazy graph batches into one buffer.
|
||||
}
|
||||
|
||||
const char* name() const override {
|
||||
return "RopePoolFused";
|
||||
}
|
||||
|
||||
bool is_equivalent(const Primitive& other) const override {
|
||||
auto* o = dynamic_cast<const RopePoolFused*>(&other);
|
||||
return o != nullptr && o->head_dim_ == head_dim_ && o->num_qo_heads_ == num_qo_heads_ &&
|
||||
o->num_kv_heads_ == num_kv_heads_ && o->rope_base_ == rope_base_;
|
||||
}
|
||||
|
||||
std::vector<Shape> output_shapes(const std::vector<array>& inputs) override {
|
||||
return {inputs[0].shape(), inputs[1].shape(), inputs[5].shape(), inputs[6].shape()};
|
||||
}
|
||||
|
||||
private:
|
||||
int head_dim_;
|
||||
int num_qo_heads_;
|
||||
int num_kv_heads_;
|
||||
float rope_base_;
|
||||
};
|
||||
|
||||
// Python entry: returns 4 arrays (q_rot, k_rot, k_pool_new, v_pool_new).
|
||||
nb::tuple rope_pool_fused_py(
|
||||
nb::handle q_h,
|
||||
nb::handle k_h,
|
||||
nb::handle v_h,
|
||||
nb::handle positions_h,
|
||||
nb::handle slots_h,
|
||||
nb::handle k_pool_h,
|
||||
nb::handle v_pool_h,
|
||||
int head_dim,
|
||||
int num_qo_heads,
|
||||
int num_kv_heads,
|
||||
float rope_base) {
|
||||
auto& q = *nb::inst_ptr<array>(q_h);
|
||||
auto& k = *nb::inst_ptr<array>(k_h);
|
||||
auto& v = *nb::inst_ptr<array>(v_h);
|
||||
auto& positions = *nb::inst_ptr<array>(positions_h);
|
||||
auto& slots = *nb::inst_ptr<array>(slots_h);
|
||||
auto& k_pool = *nb::inst_ptr<array>(k_pool_h);
|
||||
auto& v_pool = *nb::inst_ptr<array>(v_pool_h);
|
||||
|
||||
if (q.ndim() != 3 || k.ndim() != 3 || v.ndim() != 3) throw std::runtime_error("rope_pool_fused: q/k/v must be 3-D");
|
||||
if (positions.ndim() != 1 || slots.ndim() != 1)
|
||||
throw std::runtime_error("rope_pool_fused: positions/slots must be 1-D");
|
||||
if (k_pool.ndim() != 3 || v_pool.ndim() != 3) throw std::runtime_error("rope_pool_fused: pools must be 3-D");
|
||||
if (positions.dtype() != int32 || slots.dtype() != int32)
|
||||
throw std::runtime_error("rope_pool_fused: positions/slots must be int32");
|
||||
if (q.dtype() != k.dtype() || q.dtype() != v.dtype() || q.dtype() != k_pool.dtype() || q.dtype() != v_pool.dtype())
|
||||
throw std::runtime_error("rope_pool_fused: all float arrays must share dtype");
|
||||
if ((head_dim & 1) != 0) throw std::runtime_error("rope_pool_fused: head_dim must be even");
|
||||
// Shape cross-checks (catch any drift between Python pre-flight state
|
||||
// and actual tensors at dispatch time).
|
||||
const int num_tokens = q.shape(0);
|
||||
if (k.shape(0) != num_tokens || v.shape(0) != num_tokens || positions.shape(0) != num_tokens ||
|
||||
slots.shape(0) != num_tokens)
|
||||
throw std::runtime_error("rope_pool_fused: q/k/v/positions/slots must agree on token dim");
|
||||
if (q.shape(1) != num_qo_heads || k.shape(1) != num_kv_heads || v.shape(1) != num_kv_heads)
|
||||
throw std::runtime_error("rope_pool_fused: head-count mismatch with num_qo_heads/num_kv_heads");
|
||||
if (q.shape(2) != head_dim || k.shape(2) != head_dim || v.shape(2) != head_dim)
|
||||
throw std::runtime_error("rope_pool_fused: head_dim mismatch with q/k/v last dim");
|
||||
if (k_pool.shape(1) != num_kv_heads || v_pool.shape(1) != num_kv_heads || k_pool.shape(2) != head_dim ||
|
||||
v_pool.shape(2) != head_dim)
|
||||
throw std::runtime_error("rope_pool_fused: pool layout must be [pool_size, num_kv_heads, head_dim]");
|
||||
|
||||
auto stream = default_stream(Device::gpu);
|
||||
auto primitive = std::make_shared<RopePoolFused>(stream, head_dim, num_qo_heads, num_kv_heads, rope_base);
|
||||
|
||||
auto outs = array::make_arrays(
|
||||
{q.shape(), k.shape(), k_pool.shape(), v_pool.shape()},
|
||||
{q.dtype(), k.dtype(), k_pool.dtype(), v_pool.dtype()},
|
||||
primitive,
|
||||
{q, k, v, positions, slots, k_pool, v_pool});
|
||||
|
||||
// Cross-module nb cast doesn't work cleanly - explicitly construct.
|
||||
nb::module_ mx_core = nb::module_::import_("mlx.core");
|
||||
nb::object py_array_type = mx_core.attr("array");
|
||||
|
||||
nb::list result;
|
||||
for (auto& a : outs) {
|
||||
nb::object py_obj = py_array_type(0);
|
||||
auto* dst = nb::inst_ptr<array>(py_obj);
|
||||
new (dst) array(std::move(a));
|
||||
nb::inst_mark_ready(py_obj);
|
||||
result.append(py_obj);
|
||||
}
|
||||
return nb::tuple(result);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
NB_MODULE(_metal, m) {
|
||||
m.def("register_library", ®ister_library_impl, nb::arg("path"));
|
||||
|
||||
m.def(
|
||||
"rope_pool_fused",
|
||||
&rope_pool_fused_py,
|
||||
nb::arg("q"),
|
||||
nb::arg("k"),
|
||||
nb::arg("v"),
|
||||
nb::arg("positions"),
|
||||
nb::arg("slots"),
|
||||
nb::arg("k_pool"),
|
||||
nb::arg("v_pool"),
|
||||
nb::arg("head_dim"),
|
||||
nb::arg("num_qo_heads"),
|
||||
nb::arg("num_kv_heads"),
|
||||
nb::arg("rope_base"));
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
// SGLang Apple Silicon Metal kernel: NeoX RoPE fused with KV pool scatter.
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
constant uint HEAD_DIM [[function_constant(0)]];
|
||||
constant uint NUM_QO_HEADS [[function_constant(1)]];
|
||||
constant uint NUM_KV_HEADS [[function_constant(2)]];
|
||||
constant float INV_DIM_LOG2_BASE [[function_constant(3)]];
|
||||
// Heads-per-thread amortization (MLX uses 8). Each thread computes cos/sin
|
||||
// once for its (token,dim) and reuses it across N heads. Saves N-1 trig calls.
|
||||
constant uint HEADS_PER_THREAD [[function_constant(4)]];
|
||||
|
||||
constant uint HALF_DIM = HEAD_DIM / 2;
|
||||
|
||||
// ----------------------------------------------------------------------
|
||||
// Kernel 1: Q rope (no branch, no pool write) - heads-per-thread amortized
|
||||
// grid: (HALF_DIM, num_tokens, NUM_QO_HEADS / HEADS_PER_THREAD)
|
||||
// Each thread processes HEADS_PER_THREAD consecutive Q heads, sharing one
|
||||
// cos/sin computation across them.
|
||||
// ----------------------------------------------------------------------
|
||||
template <typename T>
|
||||
inline void rope_q_impl(
|
||||
const device T* q_in,
|
||||
device T* q_out,
|
||||
const device int32_t* positions,
|
||||
uint3 pos
|
||||
) {
|
||||
const uint dim_idx = pos.x;
|
||||
const uint token_id = pos.y;
|
||||
const uint head_block = pos.z;
|
||||
const uint head_start = head_block * HEADS_PER_THREAD;
|
||||
|
||||
// Trig is independent of head_id, compute once and reuse.
|
||||
const float pos_f = float(positions[token_id]);
|
||||
const float theta = pos_f * metal::exp2(
|
||||
-float(2u * dim_idx) * INV_DIM_LOG2_BASE);
|
||||
const float c = metal::fast::cos(theta);
|
||||
const float s = metal::fast::sin(theta);
|
||||
|
||||
// Apply to HEADS_PER_THREAD heads. Compiler unrolls when N is a fn-const.
|
||||
for (uint h = 0; h < HEADS_PER_THREAD; ++h) {
|
||||
const uint head_id = head_start + h;
|
||||
// Boundary: when num_qo_heads is not a multiple of N, skip extras.
|
||||
if (head_id >= NUM_QO_HEADS) break;
|
||||
|
||||
const uint base = (token_id * NUM_QO_HEADS + head_id) * HEAD_DIM;
|
||||
const uint i1 = base + dim_idx;
|
||||
const uint i2 = base + HALF_DIM + dim_idx;
|
||||
|
||||
const float x1 = float(q_in[i1]);
|
||||
const float x2 = float(q_in[i2]);
|
||||
q_out[i1] = static_cast<T>(x1 * c - x2 * s);
|
||||
q_out[i2] = static_cast<T>(x1 * s + x2 * c);
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------
|
||||
// Kernel 2: K rope + write rotated K to pool slots (no Q branch)
|
||||
// grid: (HALF_DIM, num_tokens, NUM_KV_HEADS)
|
||||
// - Same as Kernel 1 but reads from k_in, writes both k_out and k_pool[slot]
|
||||
// - slots[token_id] < 0 means "skip pool write"
|
||||
// ----------------------------------------------------------------------
|
||||
template <typename T>
|
||||
inline void rope_k_pool_impl(
|
||||
const device T* k_in,
|
||||
device T* k_out,
|
||||
device T* k_pool,
|
||||
const device int32_t* positions,
|
||||
const device int32_t* slots,
|
||||
uint3 pos
|
||||
) {
|
||||
const uint dim_idx = pos.x;
|
||||
const uint token_id = pos.y;
|
||||
const uint head_block = pos.z;
|
||||
const uint head_start = head_block * HEADS_PER_THREAD;
|
||||
|
||||
const float pos_f = float(positions[token_id]);
|
||||
const float theta = pos_f * metal::exp2(
|
||||
-float(2u * dim_idx) * INV_DIM_LOG2_BASE);
|
||||
const float c = metal::fast::cos(theta);
|
||||
const float s = metal::fast::sin(theta);
|
||||
|
||||
// Hoist slot lookup; same for all heads of this token.
|
||||
const int32_t slot = slots[token_id];
|
||||
const bool write_pool = slot >= 0;
|
||||
|
||||
for (uint h = 0; h < HEADS_PER_THREAD; ++h) {
|
||||
const uint head_id = head_start + h;
|
||||
if (head_id >= NUM_KV_HEADS) break;
|
||||
|
||||
const uint base = (token_id * NUM_KV_HEADS + head_id) * HEAD_DIM;
|
||||
const uint i1 = base + dim_idx;
|
||||
const uint i2 = base + HALF_DIM + dim_idx;
|
||||
|
||||
const float x1 = float(k_in[i1]);
|
||||
const float x2 = float(k_in[i2]);
|
||||
const T r1 = static_cast<T>(x1 * c - x2 * s);
|
||||
const T r2 = static_cast<T>(x1 * s + x2 * c);
|
||||
k_out[i1] = r1;
|
||||
k_out[i2] = r2;
|
||||
|
||||
if (write_pool) {
|
||||
const uint pool_base =
|
||||
((uint)slot * NUM_KV_HEADS + head_id) * HEAD_DIM;
|
||||
k_pool[pool_base + dim_idx] = r1;
|
||||
k_pool[pool_base + HALF_DIM + dim_idx] = r2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------
|
||||
// Kernel 3: V copy to pool slots
|
||||
// grid: (HEAD_DIM, num_tokens, NUM_KV_HEADS)
|
||||
// - Pure memcpy from v_in[token, head, dim] to v_pool[slot, head, dim]
|
||||
// - No trig, no rotation
|
||||
// ----------------------------------------------------------------------
|
||||
template <typename T>
|
||||
inline void v_to_pool_impl(
|
||||
const device T* v_in,
|
||||
device T* v_pool,
|
||||
const device int32_t* slots,
|
||||
uint3 pos
|
||||
) {
|
||||
const uint dim_idx = pos.x;
|
||||
const uint token_id = pos.y;
|
||||
const uint head_id = pos.z;
|
||||
|
||||
const int32_t slot = slots[token_id];
|
||||
if (slot < 0) return;
|
||||
|
||||
const uint src = (token_id * NUM_KV_HEADS + head_id) * HEAD_DIM + dim_idx;
|
||||
const uint dst = ((uint)slot * NUM_KV_HEADS + head_id) * HEAD_DIM + dim_idx;
|
||||
v_pool[dst] = v_in[src];
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------
|
||||
// Experimental single-dispatch rectangular kernel.
|
||||
// grid: (HEAD_DIM, num_tokens, max(NUM_QO_HEADS, NUM_KV_HEADS) / HPT)
|
||||
// - dim < HALF_DIM lanes rotate Q for Q heads
|
||||
// - dim < HALF_DIM lanes rotate K and write K pool for KV heads
|
||||
// - all dim lanes copy V for KV heads
|
||||
// This avoids packed div/mod region decoding but wastes lanes when Q and KV
|
||||
// head counts differ.
|
||||
// ----------------------------------------------------------------------
|
||||
template <typename T>
|
||||
inline void rope_pool_fused_rect_impl(
|
||||
const device T* q_in,
|
||||
const device T* k_in,
|
||||
const device T* v_in,
|
||||
device T* q_out,
|
||||
device T* k_out,
|
||||
device T* k_pool,
|
||||
device T* v_pool,
|
||||
const device int32_t* positions,
|
||||
const device int32_t* slots,
|
||||
uint3 pos
|
||||
) {
|
||||
const uint dim_idx = pos.x;
|
||||
const uint token_id = pos.y;
|
||||
const uint head_start = pos.z * HEADS_PER_THREAD;
|
||||
|
||||
const bool rope_lane = dim_idx < HALF_DIM;
|
||||
float c = 0.0f;
|
||||
float s = 0.0f;
|
||||
if (rope_lane) {
|
||||
const float pos_f = float(positions[token_id]);
|
||||
const float theta = pos_f * metal::exp2(
|
||||
-float(2u * dim_idx) * INV_DIM_LOG2_BASE);
|
||||
c = metal::fast::cos(theta);
|
||||
s = metal::fast::sin(theta);
|
||||
}
|
||||
|
||||
const int32_t slot = slots[token_id];
|
||||
const bool write_pool = slot >= 0;
|
||||
|
||||
for (uint h = 0; h < HEADS_PER_THREAD; ++h) {
|
||||
const uint head_id = head_start + h;
|
||||
|
||||
if (rope_lane && head_id < NUM_QO_HEADS) {
|
||||
const uint q_base = (token_id * NUM_QO_HEADS + head_id) * HEAD_DIM;
|
||||
const uint q_i1 = q_base + dim_idx;
|
||||
const uint q_i2 = q_base + HALF_DIM + dim_idx;
|
||||
const float x1 = float(q_in[q_i1]);
|
||||
const float x2 = float(q_in[q_i2]);
|
||||
q_out[q_i1] = static_cast<T>(x1 * c - x2 * s);
|
||||
q_out[q_i2] = static_cast<T>(x1 * s + x2 * c);
|
||||
}
|
||||
|
||||
if (head_id >= NUM_KV_HEADS) continue;
|
||||
|
||||
const uint kv_base = (token_id * NUM_KV_HEADS + head_id) * HEAD_DIM;
|
||||
const uint pool_base =
|
||||
write_pool ? ((uint)slot * NUM_KV_HEADS + head_id) * HEAD_DIM : 0u;
|
||||
|
||||
if (write_pool) {
|
||||
v_pool[pool_base + dim_idx] = v_in[kv_base + dim_idx];
|
||||
}
|
||||
|
||||
if (!rope_lane) continue;
|
||||
|
||||
const uint k_i1 = kv_base + dim_idx;
|
||||
const uint k_i2 = kv_base + HALF_DIM + dim_idx;
|
||||
const float x1 = float(k_in[k_i1]);
|
||||
const float x2 = float(k_in[k_i2]);
|
||||
const T r1 = static_cast<T>(x1 * c - x2 * s);
|
||||
const T r2 = static_cast<T>(x1 * s + x2 * c);
|
||||
k_out[k_i1] = r1;
|
||||
k_out[k_i2] = r2;
|
||||
|
||||
if (write_pool) {
|
||||
k_pool[pool_base + dim_idx] = r1;
|
||||
k_pool[pool_base + HALF_DIM + dim_idx] = r2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------
|
||||
// dtype-specialized entry points
|
||||
// ----------------------------------------------------------------------
|
||||
#define INSTANTIATE(NAME, T) \
|
||||
[[host_name("rope_pool_rect_" #NAME)]] [[kernel]] void rope_pool_rect_##NAME(\
|
||||
const device T* q_in [[buffer(0)]], \
|
||||
const device T* k_in [[buffer(1)]], \
|
||||
const device T* v_in [[buffer(2)]], \
|
||||
device T* q_out [[buffer(3)]], \
|
||||
device T* k_out [[buffer(4)]], \
|
||||
device T* k_pool [[buffer(5)]], \
|
||||
device T* v_pool [[buffer(6)]], \
|
||||
const device int32_t* positions [[buffer(7)]], \
|
||||
const device int32_t* slots [[buffer(8)]], \
|
||||
uint3 pos [[thread_position_in_grid]]) { \
|
||||
rope_pool_fused_rect_impl<T>(q_in, k_in, v_in, q_out, k_out, k_pool, \
|
||||
v_pool, positions, slots, pos); \
|
||||
} \
|
||||
[[host_name("rope_q_" #NAME)]] [[kernel]] void rope_q_##NAME( \
|
||||
const device T* q_in [[buffer(0)]], \
|
||||
device T* q_out [[buffer(1)]], \
|
||||
const device int32_t* positions [[buffer(2)]], \
|
||||
uint3 pos [[thread_position_in_grid]]) { \
|
||||
rope_q_impl<T>(q_in, q_out, positions, pos); \
|
||||
} \
|
||||
[[host_name("rope_k_pool_" #NAME)]] [[kernel]] void rope_k_pool_##NAME( \
|
||||
const device T* k_in [[buffer(0)]], \
|
||||
device T* k_out [[buffer(1)]], \
|
||||
device T* k_pool [[buffer(2)]], \
|
||||
const device int32_t* positions [[buffer(3)]], \
|
||||
const device int32_t* slots [[buffer(4)]], \
|
||||
uint3 pos [[thread_position_in_grid]]) { \
|
||||
rope_k_pool_impl<T>(k_in, k_out, k_pool, positions, slots, pos); \
|
||||
} \
|
||||
[[host_name("v_to_pool_" #NAME)]] [[kernel]] void v_to_pool_##NAME( \
|
||||
const device T* v_in [[buffer(0)]], \
|
||||
device T* v_pool [[buffer(1)]], \
|
||||
const device int32_t* slots [[buffer(2)]], \
|
||||
uint3 pos [[thread_position_in_grid]]) { \
|
||||
v_to_pool_impl<T>(v_in, v_pool, slots, pos); \
|
||||
}
|
||||
|
||||
INSTANTIATE(f16, half)
|
||||
INSTANTIATE(bf16, bfloat)
|
||||
INSTANTIATE(f32, float)
|
||||
@@ -16,7 +16,8 @@ try:
|
||||
_metallib_path = Path(_metal.__file__).resolve().parent / _METALLIB_NAME
|
||||
if not _metallib_path.is_file():
|
||||
raise ImportError(
|
||||
f"{_METALLIB_NAME} not found next to sgl_kernel._metal at {_metallib_path}"
|
||||
f"{_METALLIB_NAME} not found next to the native Metal extension "
|
||||
f"at {_metallib_path}"
|
||||
)
|
||||
_metal.register_library(str(_metallib_path))
|
||||
except ImportError as _exc: # pragma: no cover - import guarded at call time
|
||||
@@ -25,6 +26,87 @@ except ImportError as _exc: # pragma: no cover - import guarded at call time
|
||||
else:
|
||||
_IMPORT_ERROR = None
|
||||
|
||||
# Python wrappers for the compiled `_metal.*` entry points go below. Each
|
||||
# wrapper validates input shapes/dtypes and calls `mx.eval` on its operands
|
||||
# before invoking the AOT C++ entry point.
|
||||
# Python wrappers for the compiled `_metal.*` entry points go below. Wrappers
|
||||
# validate input shapes/dtypes and then invoke AOT C++ entry points. They do
|
||||
# not force `mx.eval`, so MLX can keep these calls inside its lazy graph.
|
||||
|
||||
|
||||
def rope_pool_fused(
|
||||
q: "mx.array",
|
||||
k: "mx.array",
|
||||
v: "mx.array",
|
||||
positions: "mx.array",
|
||||
slots: "mx.array",
|
||||
k_pool: "mx.array",
|
||||
v_pool: "mx.array",
|
||||
*,
|
||||
head_dim: int,
|
||||
num_qo_heads: int,
|
||||
num_kv_heads: int,
|
||||
rope_base: float,
|
||||
) -> tuple["mx.array", "mx.array", "mx.array", "mx.array"]:
|
||||
"""Apply NeoX RoPE to Q/K and scatter K/V into the MLX KV pool.
|
||||
|
||||
Args:
|
||||
q: Query tensor with shape `[num_tokens, num_qo_heads, head_dim]`.
|
||||
k: Key tensor with shape `[num_tokens, num_kv_heads, head_dim]`.
|
||||
v: Value tensor with shape `[num_tokens, num_kv_heads, head_dim]`.
|
||||
positions: int32 positions with shape `[num_tokens]`.
|
||||
slots: int32 KV-pool slots with shape `[num_tokens]`; values `< 0`
|
||||
skip the pool write for that token.
|
||||
k_pool: Existing K pool with shape `[pool_size, num_kv_heads, head_dim]`.
|
||||
v_pool: Existing V pool with shape `[pool_size, num_kv_heads, head_dim]`.
|
||||
|
||||
Returns:
|
||||
`(q_rot, k_rot, k_pool_new, v_pool_new)`.
|
||||
"""
|
||||
if q.ndim != 3 or k.ndim != 3 or v.ndim != 3:
|
||||
raise ValueError("rope_pool_fused expects q/k/v to be 3-D")
|
||||
if positions.ndim != 1 or slots.ndim != 1:
|
||||
raise ValueError("rope_pool_fused expects positions/slots to be 1-D")
|
||||
if k_pool.ndim != 3 or v_pool.ndim != 3:
|
||||
raise ValueError("rope_pool_fused expects pool tensors to be 3-D")
|
||||
q_shape = tuple(q.shape)
|
||||
k_shape = tuple(k.shape)
|
||||
v_shape = tuple(v.shape)
|
||||
positions_shape = tuple(positions.shape)
|
||||
slots_shape = tuple(slots.shape)
|
||||
k_pool_shape = tuple(k_pool.shape)
|
||||
v_pool_shape = tuple(v_pool.shape)
|
||||
|
||||
if q_shape != (q_shape[0], num_qo_heads, head_dim):
|
||||
raise ValueError(
|
||||
"q shape must be [num_tokens, num_qo_heads, head_dim], " f"got {q.shape}"
|
||||
)
|
||||
if k_shape != (q_shape[0], num_kv_heads, head_dim):
|
||||
raise ValueError(
|
||||
"k shape must be [num_tokens, num_kv_heads, head_dim], " f"got {k.shape}"
|
||||
)
|
||||
if v_shape != k_shape:
|
||||
raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}")
|
||||
if positions_shape != (q_shape[0],) or slots_shape != (q_shape[0],):
|
||||
raise ValueError("positions/slots must have one entry per token")
|
||||
if k_pool_shape[1:] != (num_kv_heads, head_dim):
|
||||
raise ValueError(f"k_pool has incompatible shape {k_pool.shape}")
|
||||
if v_pool_shape != k_pool_shape:
|
||||
raise ValueError(
|
||||
f"v_pool shape must match k_pool shape, got {v_pool.shape} vs {k_pool.shape}"
|
||||
)
|
||||
if q.dtype != k.dtype or q.dtype != v.dtype:
|
||||
raise ValueError("q/k/v dtypes must match")
|
||||
if k_pool.dtype != q.dtype or v_pool.dtype != q.dtype:
|
||||
raise ValueError("pool dtypes must match q/k/v dtype")
|
||||
|
||||
return _metal.rope_pool_fused(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
positions,
|
||||
slots,
|
||||
k_pool,
|
||||
v_pool,
|
||||
head_dim,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
float(rope_base),
|
||||
)
|
||||
|
||||
@@ -95,10 +95,10 @@ metallib_name = "sgl_metal_kernels.metallib"
|
||||
# Metal shader sources (compiled with `xcrun metal`) and C++ host sources
|
||||
# (compiled with `c++`). Add new kernels by appending to these lists.
|
||||
metal_shader_sources = [
|
||||
"csrc/metal/placeholder.metal",
|
||||
"csrc/metal/rope_pool_fused.metal",
|
||||
]
|
||||
cxx_sources = [
|
||||
"csrc/metal/placeholder.cpp",
|
||||
"csrc/metal/rope_pool_fused.cpp",
|
||||
]
|
||||
|
||||
# Header search paths shared by both the Metal shader compiler and the C++
|
||||
|
||||
Reference in New Issue
Block a user