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

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

Browse Source 
Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Co-authored-by: Xiaodong Ye <yeahdongcn@gmail.com>
This commit is contained in:
Aditya Sharma
2026-05-29 15:09:33 +08:00
committed by GitHub
co-authored by Xiaodong Ye
parent 7dff4118b9
commit b2eed9e16d
12 changed files with 1066 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/metal/README.md
+4 -8
View File
@@ -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.
sgl-kernel/csrc/metal/placeholder.cpp
View File
sgl-kernel/csrc/metal/placeholder.metal
View File
sgl-kernel/csrc/metal/rope_pool_fused.cpp
+325
View File
@@ -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", &register_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"));
}
sgl-kernel/csrc/metal/rope_pool_fused.metal
+262
View File
@@ -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)
sgl-kernel/python/sgl_kernel/metal.py
+86 -4
View File
@@ -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),
)
sgl-kernel/setup_metal.py
+2 -2
View File
@@ -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++