[CPU] Add mrope kernel for Qwen3-vl (#12531)
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
@@ -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");
|
||||
|
||||
Reference in New Issue
Block a user