283 lines
10 KiB
Python
283 lines
10 KiB
Python
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()
|