[CPU] Add mrope kernel for Qwen3-vl (#12531)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
blzheng
2026-03-18 22:12:48 -07:00
committed by GitHub
co-authored by Ma Mingfei
parent 4052b53227
commit dc6aa26ce9
5 changed files with 428 additions and 2 deletions
@@ -306,6 +306,17 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
at::Tensor& cos_sin_cache,
bool is_neox);
// mrope
std::tuple<at::Tensor, at::Tensor> multimodal_rotary_embedding_cpu(
at::Tensor& positions,
at::Tensor& query,
at::Tensor& key,
int64_t head_size,
at::Tensor& cos_sin_cache,
const std::optional<std::vector<int64_t>>& mrope_section,
bool mrope_interleaved,
bool is_neox);
// CPU and memory binding
std::string init_cpu_threads_env(const std::string& cpu_ids);
@@ -518,6 +529,11 @@ 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);
// multimodal rope
m.def(
"multimodal_rotary_embedding_cpu(Tensor positions, Tensor query, Tensor key, int head_size, Tensor "
"cos_sin_cache, int[]? mrope_section, bool mrope_interleaved, bool is_neox) -> (Tensor, Tensor)");
m.impl("multimodal_rotary_embedding_cpu", torch::kCPU, &multimodal_rotary_embedding_cpu);
// CPU and memory binding
m.def("init_cpu_threads_env(str cpu_ids) -> str");