minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[CPU] add kernel apply_rotary_pos_emb_cpu for Qwen3-VL and Qwen3-Omni (#13121)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
blzheng
2026-03-29 23:43:46 -07:00
committed by GitHub
co-authored by Ma Mingfei
parent 6da8f5f69e
commit ed01e1d5d6
4 changed files with 303 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+5
View File
@@ -318,6 +318,8 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
int64_t head_size,
at::Tensor& cos_sin_cache,
bool is_neox);
std::tuple<at::Tensor, at::Tensor>
apply_rotary_pos_emb_cpu(at::Tensor& query, at::Tensor& key, at::Tensor& cos, at::Tensor& sin);
// mrope
std::tuple<at::Tensor, at::Tensor> multimodal_rotary_embedding_cpu(
@@ -572,6 +574,9 @@ 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);
m.def("apply_rotary_pos_emb_cpu(Tensor query, Tensor key, Tensor cos, Tensor sin) -> (Tensor, Tensor)");
m.impl("apply_rotary_pos_emb_cpu", torch::kCPU, &apply_rotary_pos_emb_cpu);
// multimodal rope
m.def(
"multimodal_rotary_embedding_cpu(Tensor positions, Tensor query, Tensor key, int head_size, Tensor "