import unittest import torch from utils import precision from sglang.srt.layers.rotary_embedding import ( MRotaryEmbedding, RotaryEmbedding, ) from sglang.srt.layers.rotary_embedding.rope_variant import ( DeepseekScalingRotaryEmbedding, apply_rotary_pos_emb_native, ) from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.test.test_utils import CustomTestCase torch.manual_seed(1234) class TestROPE(CustomTestCase): def test_mrope(self): torch.manual_seed(100) 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 q_head_dim = 192 qk_nope_head_dim = 128 qk_rope_head_dim = 64 max_pos = 256 k_dim = 576 rotary_dim = 64 is_neox_style = False set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) # Create cos_sin_cache freqs = torch.rand(max_pos, qk_rope_head_dim // 2) cos = freqs.cos() * 0.7 sin = freqs.sin() * 0.7 cos_sin_cache = torch.cat((cos, sin), dim=-1).to(torch.bfloat16) positions = torch.randint(0, max_pos, (seq_len,)) rope = DeepseekScalingRotaryEmbedding( qk_rope_head_dim, rotary_dim, max_pos, 16, # not used since cos_sin_cache is provided is_neox_style, 1.0, torch.bfloat16, device="cpu", ) rope.register_buffer("cos_sin_cache", cos_sin_cache) for dtype in [torch.bfloat16]: enable_autocast = True with torch.no_grad(), torch.amp.autocast("cpu", enabled=enable_autocast): q = torch.randn(seq_len, num_head, q_head_dim, dtype=dtype) q_clone = q.clone() k = torch.randn(seq_len, 1, k_dim, dtype=dtype) k_clone = k.clone() _, q_pe = q.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1) _, q_pe_clone = q_clone.split( [qk_nope_head_dim, qk_rope_head_dim], dim=-1 ) k_pe = k[:, :, k_dim - qk_rope_head_dim :] k_pe_clone = k_clone[:, :, k_dim - qk_rope_head_dim :] # ref kernel q_pe, k_pe = rope.forward_native( query=q_pe, key=k_pe, positions=positions, ) # fused rope kernel q_pe_clone, k_pe_clone = torch.ops.sgl_kernel.rotary_embedding_cpu( positions, q_pe_clone, k_pe_clone, rope.head_size, cos_sin_cache, False, ) atol = rtol = precision[q_pe.dtype] torch.testing.assert_close(q_pe, q_pe_clone, atol=atol, rtol=rtol) torch.testing.assert_close(k_pe, k_pe_clone, atol=atol, rtol=rtol) torch.testing.assert_close(k_pe, k_pe_clone) def test_origin_rope(self): def single_test( head_size: int, rotary_dim: int, max_position_embeddings: int, base: int, dims: int, is_neox_style: bool, dtype: torch.dtype, device: str, batch_size: int, seq_len: int, num_q_heads: int, num_kv_heads: int, ): set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) torch.manual_seed(100) rope_ref = RotaryEmbedding( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, ).to(device) pos_ids = torch.arange(seq_len, device=device).repeat(batch_size) query = torch.randn( batch_size * seq_len, num_q_heads * head_size, dtype=dtype, device=device, ) key = torch.randn( batch_size * seq_len, num_kv_heads * head_size, dtype=dtype, device=device, ) if dims == 4: query = query.view(batch_size, seq_len, num_q_heads, head_size) key = key.view(batch_size, seq_len, num_kv_heads, head_size) query_ref, key_ref = query.clone(), key.clone() query_cpu, key_cpu = query.clone(), key.clone() query_ref_out, key_ref_out = rope_ref.forward_native( pos_ids, query_ref, key_ref ) query_cpu_out, key_cpu_out = torch.ops.sgl_kernel.rotary_embedding_cpu( pos_ids, query_cpu, key_cpu, rope_ref.head_size, rope_ref.cos_sin_cache.to(query.dtype), rope_ref.is_neox_style, ) torch.testing.assert_close( query_ref_out, query_cpu_out, atol=1e-2, rtol=1e-2 ) torch.testing.assert_close(key_ref_out, key_cpu_out, atol=1e-2, rtol=1e-2) test_config = [ (64, 64, 32, 8000, True, torch.bfloat16, "cpu", 32, 32, 1, 1), (256, 128, 4096, 10000, True, torch.bfloat16, "cpu", 2, 512, 32, 8), (512, 128, 311, 10000, True, torch.bfloat16, "cpu", 3, 39, 4, 2), (128, 128, 2048, 10000, False, torch.bfloat16, "cpu", 2, 512, 32, 8), (128, 128, 2048, 10000, False, torch.bfloat16, "cpu", 2, 512, 16, 4), (512, 128, 311, 10000, False, torch.bfloat16, "cpu", 3, 39, 4, 2), ] for ( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads, ) in test_config: for dim in [2, 4]: single_test( head_size, rotary_dim, max_position_embeddings, base, dim, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads, ) def test_apply_rotary_pos_emb(self): num_tokens = 1024 num_heads = 8 head_size = 72 qkv = torch.randn(num_tokens, num_heads * head_size * 3).to(torch.bfloat16) query, key, _ = qkv.split( [num_heads * head_size, num_heads * head_size, num_heads * head_size], dim=-1, ) query = query.view(num_tokens, num_heads, head_size) key = key.view(num_tokens, num_heads, head_size) for sincos_dtype in [torch.float32, torch.bfloat16]: cos = torch.rand(num_tokens, head_size).to(sincos_dtype) sin = torch.rand(num_tokens, head_size).to(sincos_dtype) q_out_ref, k_out_ref = apply_rotary_pos_emb_native(query, key, cos, sin) q_out_sgl, k_out_sgl = torch.ops.sgl_kernel.apply_rotary_pos_emb_cpu( query, key, cos, sin ) torch.testing.assert_close(q_out_ref, q_out_sgl, atol=1e-2, rtol=1e-2) torch.testing.assert_close(k_out_ref, k_out_sgl, atol=1e-2, rtol=1e-2) if __name__ == "__main__": unittest.main()