diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index 0a36550de..237528fd1 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -19,10 +19,11 @@ from sglang.srt.layers.rotary_embedding.yarn import ( yarn_linear_ramp_mask, ) from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import is_cuda, is_npu +from sglang.srt.utils import cpu_has_amx_support, is_cuda, is_npu _is_cuda = is_cuda() _is_npu = is_npu() +_is_cpu_amx_available = cpu_has_amx_support() if _is_cuda: from sglang.jit_kernel.rope import apply_rope_with_cos_sin_cache_inplace @@ -164,6 +165,26 @@ class MRotaryEmbedding(RotaryEmbedding): key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) return query, key + def forward_cpu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + fused_set_kv_buffer_arg=None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if _is_cpu_amx_available: + return torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu( + positions, + query, + key, + self.head_size, + self.cos_sin_cache, + self.mrope_section if self.mrope_section else None, + self.mrope_interleaved, + self.is_neox_style, + ) + return self.forward_native(positions, query, key, fused_set_kv_buffer_arg) + def forward_cuda( self, positions: torch.Tensor, diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index f71e83ec4..b98f531cd 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -188,6 +188,19 @@ def register_fake_ops(): else: return torch.empty_like(query), torch.empty_like(key) + @torch.library.register_fake("sgl_kernel::multimodal_rotary_embedding_cpu") + def _( + positions, + query, + key, + head_size, + cos_sin_cache, + mrope_section, + mrope_interleaved, + is_neox, + ): + return query, key + @torch.library.register_fake("sgl_kernel::qkv_proj_with_rope_fused_weight") def _( hidden_states, diff --git a/sgl-kernel/csrc/cpu/rope.cpp b/sgl-kernel/csrc/cpu/rope.cpp index a6ff0cdf2..7efc816b4 100644 --- a/sgl-kernel/csrc/cpu/rope.cpp +++ b/sgl-kernel/csrc/cpu/rope.cpp @@ -169,6 +169,100 @@ void rotary_embedding_neox_4D_kernel_impl( } } +template +inline scalar_t* get_cache_ptr( + int64_t j, + scalar_t* cache_t_ptr, + scalar_t* cache_h_ptr, + scalar_t* cache_w_ptr, + int64_t mrope_section_t, + int64_t mrope_section_h, + int64_t mrope_section_w, + bool mrope_interleaved) { + if (mrope_interleaved) { + if (j % 3 == 1 && j <= mrope_section_h * 3) return cache_h_ptr; + if (j % 3 == 2 && j <= mrope_section_w * 3) return cache_w_ptr; + return cache_t_ptr; + } + if (j < mrope_section_t) return cache_t_ptr; + if (j < mrope_section_t + mrope_section_h) return cache_h_ptr; + return cache_w_ptr; +} + +template +void multimodal_rotary_embedding_neox_2D_kernel_impl( + int64_t* __restrict__ positions, + scalar_t* __restrict__ query, + scalar_t* __restrict__ key, + scalar_t* __restrict__ cos_sin_cache, + int64_t rotary_dim, + int64_t query_stride_s, + int64_t key_stride_s, + int64_t num_heads, + int64_t num_kv_heads, + int64_t head_size, + int64_t num_tokens, + int64_t mrope_section_t, + int64_t mrope_section_h, + int64_t mrope_section_w, + int64_t positions_stride0, + bool mrope_interleaved) { + int64_t embed_dim = rotary_dim / 2; + auto compute_loop = + [&](int64_t token_head, scalar_t* cache_t_ptr, scalar_t* cache_h_ptr, scalar_t* cache_w_ptr, scalar_t* qk) { + for (int64_t j = 0; j < embed_dim; ++j) { + int64_t x_index = j; + int64_t y_index = embed_dim + j; + + int64_t out_x = token_head + x_index; + int64_t out_y = token_head + y_index; + + scalar_t* cache_ptr = get_cache_ptr( + j, + cache_t_ptr, + cache_h_ptr, + cache_w_ptr, + mrope_section_t, + mrope_section_h, + mrope_section_w, + mrope_interleaved); + float _cos = cache_ptr[x_index]; + float _sin = cache_ptr[y_index]; + + float _q_x = qk[out_x]; + float _q_y = qk[out_y]; + + qk[out_x] = _q_x * _cos - _q_y * _sin; + qk[out_y] = _q_y * _cos + _q_x * _sin; + } + }; + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { + int64_t token_idx = {0}; + data_index_init(begin, token_idx, num_tokens); + for (int i = begin; i < end; ++i) { + int64_t pos_t = positions[token_idx]; + int64_t pos_h = positions[positions_stride0 + token_idx]; + int64_t pos_w = positions[positions_stride0 * 2 + token_idx]; + scalar_t* cache_t_ptr = cos_sin_cache + pos_t * rotary_dim; + scalar_t* cache_h_ptr = cos_sin_cache + pos_h * rotary_dim; + scalar_t* cache_w_ptr = cos_sin_cache + pos_w * rotary_dim; + + for (int64_t i = 0; i < num_heads; ++i) { + int64_t head_idx = i; + int64_t token_head = token_idx * query_stride_s + head_idx * head_size; + compute_loop(token_head, cache_t_ptr, cache_h_ptr, cache_w_ptr, query); + } + + for (int64_t i = 0; i < num_kv_heads; ++i) { + int64_t head_idx = i; + int64_t token_head = token_idx * key_stride_s + head_idx * head_size; + compute_loop(token_head, cache_t_ptr, cache_h_ptr, cache_w_ptr, key); + } + data_index_step(token_idx, num_tokens); + } + }); +} + template void rotary_embedding_4D_kernel_impl( int64_t* __restrict__ positions, @@ -248,6 +342,87 @@ void rotary_embedding_4D_kernel_impl( }); } +template +void multimodal_rotary_embedding_2D_kernel_impl( + int64_t* __restrict__ positions, + scalar_t* __restrict__ query, + scalar_t* __restrict__ key, + scalar_t* __restrict__ cos_sin_cache, + int64_t rotary_dim, + int64_t query_stride_s, + int64_t key_stride_s, + int64_t num_heads, + int64_t num_kv_heads, + int64_t head_size, + int64_t num_tokens, + int64_t mrope_section_t, + int64_t mrope_section_h, + int64_t mrope_section_w, + int64_t positions_stride0, + bool mrope_interleaved) { + int64_t embed_dim = rotary_dim / 2; + auto compute_loop = [&](scalar_t* cache_t_ptr, scalar_t* cache_h_ptr, scalar_t* cache_w_ptr, scalar_t* head_query) { + for (int64_t j = 0; j < embed_dim; j += 1) { + int64_t rot_offset = j; + int64_t x_index = 2 * rot_offset; + int64_t y_index = 2 * rot_offset + 1; + + scalar_t* cache_ptr = get_cache_ptr( + j, + cache_t_ptr, + cache_h_ptr, + cache_w_ptr, + mrope_section_t, + mrope_section_h, + mrope_section_w, + mrope_interleaved); + float cos = cache_ptr[rot_offset]; + float sin = cache_ptr[rot_offset + embed_dim]; + + float x = head_query[x_index]; + float y = head_query[y_index]; + + head_query[x_index] = x * cos - y * sin; + head_query[y_index] = y * cos + x * sin; + } + }; + at::parallel_for(0, num_tokens * num_heads, GRAIN_SIZE / rotary_dim, [&](int64_t begin, int64_t end) { + int64_t token_idx = {0}, i = {0}; + data_index_init(begin, token_idx, num_tokens, i, num_heads); + for ([[maybe_unused]] auto z : c10::irange(begin, end)) { + int64_t pos_t = positions[token_idx]; + int64_t pos_h = positions[positions_stride0 + token_idx]; + int64_t pos_w = positions[positions_stride0 * 2 + token_idx]; + scalar_t* cache_t_ptr = cos_sin_cache + pos_t * rotary_dim; + scalar_t* cache_h_ptr = cos_sin_cache + pos_h * rotary_dim; + scalar_t* cache_w_ptr = cos_sin_cache + pos_w * rotary_dim; + int64_t head_idx = i; + int64_t token_head = token_idx * query_stride_s + head_idx * head_size; + scalar_t* head_query = token_head + query; + compute_loop(cache_t_ptr, cache_h_ptr, cache_w_ptr, head_query); + data_index_step(token_idx, num_tokens, i, num_heads); + } + }); + + at::parallel_for(0, num_tokens * num_kv_heads, GRAIN_SIZE / rotary_dim, [&](int64_t begin, int64_t end) { + int64_t token_idx{0}, i = {0}; + data_index_init(begin, token_idx, num_tokens, i, num_kv_heads); + for ([[maybe_unused]] auto z : c10::irange(begin, end)) { + int64_t pos_t = positions[token_idx]; + int64_t pos_h = positions[positions_stride0 + token_idx]; + int64_t pos_w = positions[positions_stride0 * 2 + token_idx]; + scalar_t* cache_t_ptr = cos_sin_cache + pos_t * rotary_dim; + scalar_t* cache_h_ptr = cos_sin_cache + pos_h * rotary_dim; + scalar_t* cache_w_ptr = cos_sin_cache + pos_w * rotary_dim; + int64_t head_idx = i; + int64_t token_head = token_idx * key_stride_s + head_idx * head_size; + scalar_t* head_key = key + token_head; + compute_loop(cache_t_ptr, cache_h_ptr, cache_w_ptr, head_key); + data_index_step(token_idx, num_tokens, i, num_kv_heads); + } + }); +} + } // namespace std::tuple rotary_embedding_cpu( @@ -385,3 +560,130 @@ std::tuple rotary_embedding_cpu( }); return std::make_tuple(query_out, key_out); } + +// positions: [num_tokens] (text only) or [3, num_tokens] (T/H/W positions with multimodal inputs) +// query: [num_tokens, num_heads * head_size] +// key: [num_tokens, num_kv_heads * head_size] +// cos_sin_cache: [max_position_embeddings, rotary_dim] +// mrope_section: [t, h, w] +std::tuple multimodal_rotary_embedding_cpu( + at::Tensor& positions, + at::Tensor& query, + at::Tensor& key, + int64_t head_size, + at::Tensor& cos_sin_cache, + const std::optional>& mrope_section, + bool mrope_interleaved, + bool is_neox) { + RECORD_FUNCTION("sgl-kernel::multimodal_rotary_embedding_cpu", std::vector({query, key})); + TORCH_CHECK(positions.dim() == 1 || positions.dim() == 2, "positions must be a 1D or 2D tensor"); + CHECK_DIM(2, query); + CHECK_DIM(2, key); + CHECK_DIM(2, cos_sin_cache); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(query); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(key); + int64_t rotary_dim = cos_sin_cache.size(1); + int64_t num_tokens = positions.size(-1); + CHECK_EQ(key.size(0), num_tokens); + CHECK_EQ(query.size(0), num_tokens); + const auto input_dtype = query.scalar_type(); + TORCH_CHECK(positions.scalar_type() == at::kLong, "expect positions to be int64, got ", positions.scalar_type()); + TORCH_CHECK(input_dtype == key.scalar_type(), "query and key must have the same data type"); + TORCH_CHECK(input_dtype == cos_sin_cache.scalar_type(), "query and cos_sin_cache must have the same data type"); + + int64_t num_heads = query.size(-1) / head_size; + int64_t num_kv_heads = key.size(-1) / head_size; + int64_t key_stride_s = key.stride(0); + int64_t query_stride_s = query.stride(0); + + if (positions.dim() == 2) { + TORCH_CHECK(mrope_section.has_value(), "mrope_section must be provided when positions is 2D"); + auto mrope_section_val = mrope_section.value(); + CHECK_EQ(mrope_section_val.size(), 3); + CHECK_EQ(positions.size(0), 3); + int64_t mrope_section_t = mrope_section_val[0]; + int64_t mrope_section_h = mrope_section_val[1]; + int64_t mrope_section_w = mrope_section_val[2]; + int64_t positions_stride0 = positions.stride(0); + AT_DISPATCH_REDUCED_FLOATING_TYPES(input_dtype, "rotary_embedding_cpu", [&] { + if (is_neox) { + multimodal_rotary_embedding_neox_2D_kernel_impl( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + rotary_dim, + query_stride_s, + key_stride_s, + num_heads, + num_kv_heads, + head_size, + num_tokens, + mrope_section_t, + mrope_section_h, + mrope_section_w, + positions_stride0, + mrope_interleaved); + } else { + multimodal_rotary_embedding_2D_kernel_impl( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + rotary_dim, + query_stride_s, + key_stride_s, + num_heads, + num_kv_heads, + head_size, + num_tokens, + mrope_section_t, + mrope_section_h, + mrope_section_w, + positions_stride0, + mrope_interleaved); + } + }); + } else { // positions.dim() == 1 + AT_DISPATCH_REDUCED_FLOATING_TYPES(input_dtype, "rotary_embedding_cpu", [&] { + if (is_neox) { + rotary_embedding_neox_4D_kernel_impl( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + rotary_dim, + 0, + query_stride_s, + head_size, + 0, + key_stride_s, + head_size, + num_heads, + num_kv_heads, + head_size, + 1, + num_tokens); + } else { + rotary_embedding_4D_kernel_impl( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + rotary_dim, + 0, + query_stride_s, + head_size, + 0, + key_stride_s, + head_size, + num_heads, + num_kv_heads, + head_size, + 1, + num_tokens); + } + }); + } + return std::make_tuple(query, key); +} diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 6b29da857..cc487c3af 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -306,6 +306,17 @@ std::tuple rotary_embedding_cpu( at::Tensor& cos_sin_cache, bool is_neox); +// mrope +std::tuple multimodal_rotary_embedding_cpu( + at::Tensor& positions, + at::Tensor& query, + at::Tensor& key, + int64_t head_size, + at::Tensor& cos_sin_cache, + const std::optional>& mrope_section, + bool mrope_interleaved, + bool is_neox); + // CPU and memory binding std::string init_cpu_threads_env(const std::string& cpu_ids); @@ -518,6 +529,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "rotary_embedding_cpu(Tensor positions, Tensor query, Tensor key, int head_size, Tensor cos_sin_cache, " "bool is_neox) -> (Tensor, Tensor)"); m.impl("rotary_embedding_cpu", torch::kCPU, &rotary_embedding_cpu); + // multimodal rope + m.def( + "multimodal_rotary_embedding_cpu(Tensor positions, Tensor query, Tensor key, int head_size, Tensor " + "cos_sin_cache, int[]? mrope_section, bool mrope_interleaved, bool is_neox) -> (Tensor, Tensor)"); + m.impl("multimodal_rotary_embedding_cpu", torch::kCPU, &multimodal_rotary_embedding_cpu); // CPU and memory binding m.def("init_cpu_threads_env(str cpu_ids) -> str"); diff --git a/test/srt/cpu/test_rope.py b/test/srt/cpu/test_rope.py index 5eeac5337..97475a607 100644 --- a/test/srt/cpu/test_rope.py +++ b/test/srt/cpu/test_rope.py @@ -3,7 +3,10 @@ import unittest import torch from utils import precision -from sglang.srt.layers.rotary_embedding.base import RotaryEmbedding +from sglang.srt.layers.rotary_embedding import ( + MRotaryEmbedding, + RotaryEmbedding, +) from sglang.srt.layers.rotary_embedding.rope_variant import ( DeepseekScalingRotaryEmbedding, ) @@ -14,6 +17,77 @@ torch.manual_seed(1234) class TestROPE(CustomTestCase): + def test_mrope(self): + head_size = 128 + seq_len = 512 + num_heads = 16 + num_kv_heads = 1 + rotary_dim = 128 + max_pos = 262144 + base = 5000000 + is_neox_style = True + dtype = torch.bfloat16 + mrope_section = [24, 20, 20] + mrope_interleaved = True + positions_mrope = torch.randint(0, max_pos, (3, seq_len)) + positions_text = torch.randint(0, max_pos, (seq_len,)) + set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) + + test_config = [ + # (dtype, is_neox_stype, mrope_interleaved, positions, mrope_section) + (torch.bfloat16, False, True, positions_mrope, mrope_section), + (torch.bfloat16, False, False, positions_mrope, mrope_section), + (torch.bfloat16, False, False, positions_text, None), + (torch.bfloat16, True, True, positions_mrope, mrope_section), + (torch.bfloat16, True, False, positions_mrope, mrope_section), + (torch.bfloat16, True, False, positions_text, None), + ] + for ( + dtype, + is_neox_style, + mrope_interleaved, + positions, + mrope_section, + ) in test_config: + rope = MRotaryEmbedding( + head_size, + rotary_dim, + max_pos, + base, + is_neox_style, + dtype, + mrope_section, + mrope_interleaved, + ) + enable_autocast = True + + with torch.no_grad(), torch.amp.autocast("cpu", enabled=enable_autocast): + q = torch.randn(seq_len, num_heads * head_size, dtype=dtype) + q_clone = q.clone() + k = torch.randn(seq_len, num_kv_heads * head_size, dtype=dtype) + k_clone = k.clone() + + # ref kernel + q_ref, k_ref = rope.forward_native( + query=q, + key=k, + positions=positions, + ) + # fused rope kernel + q_sgl, k_sgl = torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu( + positions, + q_clone, + k_clone, + rope.head_size, + rope.cos_sin_cache, + rope.mrope_section, + rope.mrope_interleaved, + is_neox_style, + ) + atol = rtol = precision[q_ref.dtype] + torch.testing.assert_close(q_ref, q_sgl, atol=atol, rtol=rtol) + torch.testing.assert_close(k_ref, k_sgl, atol=atol, rtol=rtol) + def test_deepseek_v2_rope(self): num_head = 16 seq_len = 1024