chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[CPU] Add 4D input support for ROPE in sgl-kernel (#9337)

Browse Source
This commit is contained in:
blzheng
2025-12-16 17:27:39 +08:00
committed by GitHub
parent 272c5fe43e
commit e27635a02d
2 changed files with 102 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
test/srt/cpu/test_rope.py
View File

@@ -88,6 +88,7 @@ class TestROPE(CustomTestCase):
rotary_dim: int,
max_position_embeddings: int,
base: int,
dims: int,
is_neox_style: bool,
dtype: torch.dtype,
device: str,
@@ -119,7 +120,9 @@ class TestROPE(CustomTestCase):
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()
@@ -161,19 +164,21 @@ class TestROPE(CustomTestCase):
num_q_heads,
num_kv_heads,
) in test_config:
single_test(
head_size,
rotary_dim,
max_position_embeddings,
base,
is_neox_style,
dtype,
device,
batch_size,
seq_len,
num_q_heads,
num_kv_heads,
)
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,
)
if __name__ == "__main__":