EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Optimize rope in sgl kernel (#4267)

Browse Source
This commit is contained in:
Lianmin Zheng
2025-03-10 10:07:45 -07:00
committed by GitHub
parent 3d56585a97
commit cf0ccd406e
2 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sgl-kernel/csrc/elementwise/rope.cu
View File

@@ -65,7 +65,7 @@ void apply_rope_pos_ids_cos_sin_cache(
static_cast<c_type*>(q_rope.data_ptr()), static_cast<c_type*>(q_rope.data_ptr()),
static_cast<c_type*>(k_rope.data_ptr()), static_cast<c_type*>(k_rope.data_ptr()),
static_cast<float*>(cos_sin_cache.data_ptr()), static_cast<float*>(cos_sin_cache.data_ptr()),
static_cast<int32_t*>(pos_ids.data_ptr()), static_cast<int64_t*>(pos_ids.data_ptr()),
nnz, nnz,
num_qo_heads, num_qo_heads,
num_kv_heads, num_kv_heads,

3
sgl-kernel/python/sgl_kernel/elementwise.py
View File

@@ -139,14 +139,13 @@ def apply_rope_with_cos_sin_cache_inplace(
if cos_sin_cache.dtype != torch.float32: if cos_sin_cache.dtype != torch.float32:
raise ValueError("cos_sin_cache should be float32") raise ValueError("cos_sin_cache should be float32")
positions = positions.int()
torch.ops.sgl_kernel.apply_rope_pos_ids_cos_sin_cache( torch.ops.sgl_kernel.apply_rope_pos_ids_cos_sin_cache(
q=query.view(query.shape[0], -1, head_size), q=query.view(query.shape[0], -1, head_size),
k=key.view(key.shape[0], -1, head_size), k=key.view(key.shape[0], -1, head_size),
q_rope=query.view(query.shape[0], -1, head_size), q_rope=query.view(query.shape[0], -1, head_size),
k_rope=key.view(key.shape[0], -1, head_size), k_rope=key.view(key.shape[0], -1, head_size),
cos_sin_cache=cos_sin_cache, cos_sin_cache=cos_sin_cache,
pos_ids=positions, pos_ids=positions.long(),
interleave=(not is_neox), interleave=(not is_neox),
cuda_stream=get_cuda_stream(), cuda_stream=get_cuda_stream(),
) )