diff --git a/docs_new/docs/hardware-platforms/apple_metal.mdx b/docs_new/docs/hardware-platforms/apple_metal.mdx index 9b45fd2db..bac5f4fc2 100644 --- a/docs_new/docs/hardware-platforms/apple_metal.mdx +++ b/docs_new/docs/hardware-platforms/apple_metal.mdx @@ -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. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1fd02be23..744d3aa28 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/hardware_backend/mlx/aot.py b/python/sglang/srt/hardware_backend/mlx/aot.py new file mode 100644 index 000000000..880514a77 --- /dev/null +++ b/python/sglang/srt/hardware_backend/mlx/aot.py @@ -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, + ) + ) diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py index d9524ab11..dc77bdc53 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py @@ -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, :] diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index 1fb5bb33b..61f4c086c 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner.py @@ -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) diff --git a/sgl-kernel/csrc/metal/README.md b/sgl-kernel/csrc/metal/README.md index 602f480c1..4ca8a6b53 100644 --- a/sgl-kernel/csrc/metal/README.md +++ b/sgl-kernel/csrc/metal/README.md @@ -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/.metal`. -2. Add the C++ host / nanobind binding under `csrc/metal/.cpp`, exporting the entry point on the `sgl_kernel._metal` module. +2. Add the C++ host / nanobind binding under `csrc/metal/.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. diff --git a/sgl-kernel/csrc/metal/placeholder.cpp b/sgl-kernel/csrc/metal/placeholder.cpp deleted file mode 100644 index e69de29bb..000000000 diff --git a/sgl-kernel/csrc/metal/placeholder.metal b/sgl-kernel/csrc/metal/placeholder.metal deleted file mode 100644 index e69de29bb..000000000 diff --git a/sgl-kernel/csrc/metal/rope_pool_fused.cpp b/sgl-kernel/csrc/metal/rope_pool_fused.cpp new file mode 100644 index 000000000..afa93fa58 --- /dev/null +++ b/sgl-kernel/csrc/metal/rope_pool_fused.cpp @@ -0,0 +1,325 @@ +// Combined optimal: real AOT .metallib + Primitive integration + optimized +// 3-kernel + 3D-grid dispatch + fused KV pool write. + +#include +#include + +#include +#include +#include +#include +#include + +#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(gx, 32u); + uint32_t ty = std::min(gy, kMaxThreads / std::max(tx, 1u)); + uint32_t tz = std::min(gz, kMaxThreads / std::max(tx * ty, 1u)); + while (ty > 1 && (gy % ty) != 0) + --ty; + while (tz > 1 && (gz % tz) != 0) + --tz; + return MTL::Size::Make(tx, std::max(ty, 1u), std::max(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(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&, std::vector&) override { + throw std::runtime_error("rope_pool_fused: CPU eval not supported"); + } + + void eval_gpu(const std::vector& inputs, std::vector& 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(head_dim_); + const uint32_t nq = static_cast(num_qo_heads_); + const uint32_t nk = static_cast(num_kv_heads_); + const uint32_t half_dim = hd / 2; + const uint32_t num_tokens = static_cast(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(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(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(&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 output_shapes(const std::vector& 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(q_h); + auto& k = *nb::inst_ptr(k_h); + auto& v = *nb::inst_ptr(v_h); + auto& positions = *nb::inst_ptr(positions_h); + auto& slots = *nb::inst_ptr(slots_h); + auto& k_pool = *nb::inst_ptr(k_pool_h); + auto& v_pool = *nb::inst_ptr(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(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(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")); +} diff --git a/sgl-kernel/csrc/metal/rope_pool_fused.metal b/sgl-kernel/csrc/metal/rope_pool_fused.metal new file mode 100644 index 000000000..1fd10e8bf --- /dev/null +++ b/sgl-kernel/csrc/metal/rope_pool_fused.metal @@ -0,0 +1,262 @@ +// SGLang Apple Silicon Metal kernel: NeoX RoPE fused with KV pool scatter. + +#include +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 +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(x1 * c - x2 * s); + q_out[i2] = static_cast(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 +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(x1 * c - x2 * s); + const T r2 = static_cast(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 +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 +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(x1 * c - x2 * s); + q_out[q_i2] = static_cast(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(x1 * c - x2 * s); + const T r2 = static_cast(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(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(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(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(v_in, v_pool, slots, pos); \ + } + +INSTANTIATE(f16, half) +INSTANTIATE(bf16, bfloat) +INSTANTIATE(f32, float) diff --git a/sgl-kernel/python/sgl_kernel/metal.py b/sgl-kernel/python/sgl_kernel/metal.py index 4bc136f2e..e369ee5cb 100644 --- a/sgl-kernel/python/sgl_kernel/metal.py +++ b/sgl-kernel/python/sgl_kernel/metal.py @@ -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), + ) diff --git a/sgl-kernel/setup_metal.py b/sgl-kernel/setup_metal.py index 132631ff6..a08af03da 100644 --- a/sgl-kernel/setup_metal.py +++ b/sgl-kernel/setup_metal.py @@ -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++