[transformation][CPU][GPU] support 3D rope (#31949)

tiger100256-hu · zaixing-wang · web-flow · commit f412d66bfc4f · 2025-10-14T05:51:29.000Z
### Details: - *support 3D rope fusing in Qwen2.5-vl-3b and Qwen2-vl-7b* ### Tickets: - *CVS-171588, CVS-171442* --------- Signed-off-by: HU Yuan2 <yuan2.hu@intel.com> Co-authored-by: zaixing.wang <zaixing.wang@intel.com>
diff --git a/src/common/transformations/include/ov_ops/rotary_positional_embeddings.hpp b/src/common/transformations/include/ov_ops/rotary_positional_embeddings.hpp
@@ -32,6 +32,7 @@ class TRANSFORMATIONS_API RoPE : public Op {
                                         // each head. change input order to [batch, head_cnt, 4608] to support 2d rope
         bool is_qwen = false;           // Qwen is special which overrides other setting
         bool use_rope_cache = false;    // use precomputed RoPE cache for trigonometric values (cosine and sine)
+        bool support_3d_rope = false;   // use same logic as RoPEFusionGPTNEOX(4), used by gpu plugin
         size_t head_cnt = 0;
         size_t head_size = 0;
         int gather_position_arg_id =
diff --git a/src/common/transformations/include/transformations/common_optimizations/fuse_rotary_positional_embeddings.hpp b/src/common/transformations/include/transformations/common_optimizations/fuse_rotary_positional_embeddings.hpp
@@ -28,7 +28,7 @@ class TRANSFORMATIONS_API RoPEShareCosSin;
 class ov::pass::RoPEFusionGPTNEOX : public ov::pass::MatcherPass {
 public:
     OPENVINO_MATCHER_PASS_RTTI("RoPEFusionGPTNEOX");
-    RoPEFusionGPTNEOX();
+    RoPEFusionGPTNEOX(int rank);
 };
 
 class ov::pass::RoPEFusionFlux : public ov::pass::MatcherPass {
diff --git a/src/common/transformations/src/ov_ops/rotary_positional_embeddings.cpp b/src/common/transformations/src/ov_ops/rotary_positional_embeddings.cpp
@@ -95,6 +95,7 @@ bool RoPE::visit_attributes(ov::AttributeVisitor& visitor) {
     visitor.on_attribute("rotary_ndims", m_config.rotary_ndims);
     visitor.on_attribute("is_chatglm", m_config.is_chatglm);
     visitor.on_attribute("support_2d_rope", m_config.support_2d_rope);
+    visitor.on_attribute("support_3d_rope", m_config.support_3d_rope);
     visitor.on_attribute("is_qwen", m_config.is_qwen);
     visitor.on_attribute("use_rope_cache", m_config.use_rope_cache);
     visitor.on_attribute("head_cnt", m_config.head_cnt);
diff --git a/src/common/transformations/src/transformations/common_optimizations/fuse_rotary_positional_embeddings.cpp b/src/common/transformations/src/transformations/common_optimizations/fuse_rotary_positional_embeddings.cpp
@@ -58,7 +58,8 @@ bool ov::pass::RoPEFusion::run_on_model(const std::shared_ptr<ov::Model>& model)
     auto symbolic_ctx_manager = symbolic_optimizations.get_manager();
 
     symbolic_ctx_manager->register_pass<ov::pass::RoPEFusionFlux>();
-    symbolic_ctx_manager->register_pass<ov::pass::RoPEFusionGPTNEOX>();
+    symbolic_ctx_manager->register_pass<ov::pass::RoPEFusionGPTNEOX>(4);
+    symbolic_ctx_manager->register_pass<ov::pass::RoPEFusionGPTNEOX>(3);
     symbolic_ctx_manager->register_pass<ov::pass::RoPEFusionGPTJ>();
     // optional heads & tails are fused in separate matcher pass,
     // after RoPENode has been created.
@@ -164,7 +165,7 @@ ov::pass::RoPEFusionFlux::RoPEFusionFlux() {
     this->register_matcher(m, callback);
 }
 
-ov::pass::RoPEFusionGPTNEOX::RoPEFusionGPTNEOX() {
+ov::pass::RoPEFusionGPTNEOX::RoPEFusionGPTNEOX(int rank) {
     using namespace ov::op::util;
     MATCHER_SCOPE(RoPEFusionGPTNEOX);
 
@@ -177,19 +178,19 @@ ov::pass::RoPEFusionGPTNEOX::RoPEFusionGPTNEOX() {
     // branch.
     // so here we use a WA, only match the path of rotate_hal(x)*sin and check the x*cos path
     // in the callback
-    auto x = pattern::any_input(pattern::rank_equals(4));
-    auto x_or_cos1 = pattern::any_input(pattern::rank_equals(4));
-    auto x_or_cos2 = pattern::any_input(pattern::rank_equals(4));
-    auto t_sin = pattern::any_input(pattern::rank_equals(4));
+    auto x = pattern::any_input(pattern::rank_equals(rank));
+    auto x_or_cos1 = pattern::any_input(pattern::rank_equals(rank));
+    auto x_or_cos2 = pattern::any_input(pattern::rank_equals(rank));
+    auto t_sin = pattern::any_input(pattern::rank_equals(rank));
 
-    auto varsplit = pattern::wrap_type<v1::VariadicSplit>({x, 3, {"half_ndims", "?"}});
+    auto varsplit = pattern::wrap_type<v1::VariadicSplit>({x, rank - 1, {"half_ndims", "?"}});
     varsplit->set_output_size(2);
 
     auto int32_max = std::numeric_limits<std::int32_t>::max();
 
-    auto x2 = NewGenSlice(x, "half_ndims", int32_max, 1, 3);
+    auto x2 = NewGenSlice(x, "half_ndims", int32_max, 1, rank - 1);
     auto x2neg = pattern::wrap_type<v1::Multiply>({x2 | varsplit->output(1), -1.0f}, {{"auto_broadcast", "numpy"}});
-    auto x1 = NewGenSlice(x, 0, "half_ndims", 1, 3);
+    auto x1 = NewGenSlice(x, 0, "half_ndims", 1, rank - 1);
     auto x_rotate_half = pattern::wrap_type<v0::Concat>({x2neg, x1 | varsplit->output(0)}, {{"axis", -1}});
 
     auto mul_cos = pattern::wrap_type<v1::Multiply>({x_or_cos1, x_or_cos2}, {{"auto_broadcast", "numpy"}});
@@ -220,6 +221,9 @@ ov::pass::RoPEFusionGPTNEOX::RoPEFusionGPTNEOX() {
 
         op::internal::RoPE::Config config;
         OutputVector new_args;
+        if (rank == 3) {
+            config.support_3d_rope = true;
+        }
         config.rotary_ndims = 2ul * static_cast<size_t>(half_ndims.i());
 
         new_args.push_back(pattern_map.at(x));
diff --git a/src/common/transformations/tests/common_optimizations/fuse_rotary_positional_embeddings.cpp b/src/common/transformations/tests/common_optimizations/fuse_rotary_positional_embeddings.cpp
diff --git a/src/plugins/intel_cpu/src/nodes/rope.cpp b/src/plugins/intel_cpu/src/nodes/rope.cpp
@@ -129,6 +129,11 @@ struct RoPE::RoPEExecutorRotateHalf : public RoPE::Executor {
             t_src = t_src.slice(3, m_config.slice_start, m_config.slice_stop);
             can_inplace = false;
         }
+        if (t_src.m_rank == 3) {
+            t_src = t_src.reshape({1, t_src.size(0), t_src.size(1), t_src.size(2)});
+            t_dst = t_dst.reshape({1, t_dst.size(0), t_dst.size(1), t_dst.size(2)});
+        }
+
         if (m_config.input_trans0213) {
             t_src = t_src.permute({0, 2, 1, 3});
             can_inplace = false;
@@ -139,9 +144,13 @@ struct RoPE::RoPEExecutorRotateHalf : public RoPE::Executor {
 
         if (t_cos.m_rank == 2) {
             t_cos = t_cos.reshape({1, 1, t_cos.size(0), t_cos.size(1)});
+        } else if (t_cos.m_rank == 3) {
+            t_cos = t_cos.reshape({1, t_cos.size(0), t_cos.size(1), t_cos.size(2)});
         }
         if (t_sin.m_rank == 2) {
             t_sin = t_sin.reshape({1, 1, t_sin.size(0), t_sin.size(1)});
+        } else if (t_sin.m_rank == 3) {
+            t_sin = t_sin.reshape({1, t_sin.size(0), t_sin.size(1), t_sin.size(2)});
         }
 
         auto batch_size = t_src.size(0);
diff --git a/src/plugins/intel_cpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp b/src/plugins/intel_cpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp
@@ -70,5 +70,14 @@ INSTANTIATE_TEST_SUITE_P(smoke_RoPETestChatGLM,
                             ::testing::Values(ov::test::utils::DEVICE_CPU)),
                          RoPETestChatGLM2DRoPEStridedSlice::getTestCaseName);
 
+const std::vector<std::string> vit_param = {"VariadicSplit", "Slice", "StridedSlice"};
+INSTANTIATE_TEST_SUITE_P(smoke_RoPETestQwenVL,
+                         RoPETestQwenVL,
+                         ::testing::Combine(
+                            ::testing::Values(ov::element::f32),
+                            ::testing::Values(ov::test::utils::DEVICE_CPU),
+                            ::testing::ValuesIn(vit_param)),
+                         RoPETestQwenVL::getTestCaseName);
+
 }  // namespace test
 }  // namespace ov
diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/rope_opt.cl b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/rope_opt.cl
@@ -344,7 +344,6 @@ KERNEL(rope_opt)
         input_idx += SLICED_FROM_START;
     #endif
 #endif
-
 uint cos_sin_p = p;
 #ifdef ENABLE_GATHER
     uint gather_b = b < INPUT3_BATCH_NUM ? b : 0;
@@ -378,7 +377,7 @@ uint cos_sin_p = p;
     uint cos_sin_h = 0;
     cos_sin_p = cos_sin_p < INPUT1_BATCH_NUM ? cos_sin_p : 0;
 
-    #ifndef SIN_COS_HAVE_DYNAMIC_PADDINGS
+#ifndef SIN_COS_HAVE_DYNAMIC_PADDINGS
     uint cos_sin_idx = INPUT1_GET_INDEX(cos_sin_p, 0, 0, 0);
 
     uint cos_idx = cos_sin_idx;
@@ -387,8 +386,20 @@ uint cos_sin_p = p;
     uint cos_idx = INPUT1_GET_INDEX(cos_sin_p, 0, 0, 0);
     uint sin_idx = INPUT2_GET_INDEX(cos_sin_p, 0, 0, 0);
 #endif
+#elif INPUT1_DIMS == 3 && INPUT2_DIMS == 3
+    uint cos_sin_b = b < INPUT1_BATCH_NUM ? b : 0;
+    uint cos_sin_h = h < INPUT1_FEATURE_NUM ? h : 0;
+#ifndef SIN_COS_HAVE_DYNAMIC_PADDINGS
+    uint cos_sin_idx = INPUT1_GET_INDEX(cos_sin_b, cos_sin_h, 0, 0);
+
+    uint cos_idx = cos_sin_idx;
+    uint sin_idx = cos_sin_idx;
 #else
-#   error "rope_opt.cl - 4 or 2 of INPUT1_DIMS/INPUT2_DIMS is supported only"
+    uint cos_idx = INPUT1_GET_INDEX(cos_sin_b, cos_sin_h, 0, 0);
+    uint sin_idx = INPUT2_GET_INDEX(cos_sin_b, cos_sin_h, 0, 0);
+#endif
+#else
+#   error "rope_opt.cl - 2, 3 or 4 of INPUT1_DIMS/INPUT2_DIMS is supported only"
 #endif
 
     uint output_idx = OUTPUT_GET_INDEX(b, h, p, 0);
@@ -415,6 +426,7 @@ uint cos_sin_p = p;
     *(OUTPUT_VEC_TYPE*)(output + output_idx + r) = out1;
     *(OUTPUT_VEC_TYPE*)(output + output_idx + HALF_ROTARY_NDIMS + r) = out2;
 #endif
+
 }
 #endif
 
diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/rope_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/rope_opt.cpp
@@ -152,8 +152,10 @@ class RopeGenerator : public KernelGenerator {
                     auto b = extract_channel(ChannelName::BATCH, out_l);
                     auto f = extract_channel(ChannelName::FEATURE, out_l);
                     auto y = extract_channel(ChannelName::Y, out_l);
-
                     wgs.global = {b, f, y * cfg.rotary_ndims / 2ul / vec_size};
+                    if (cfg.support_3d_rope) {
+                        wgs.global = {b, f, cfg.rotary_ndims / 2ul / vec_size};
+                    }
                 }
 
                 wgs.local = ov::intel_gpu::get_optimal_lws(wgs.global, params.get_device_info(), in_l.format, out_l.format, dims_by_gws);
diff --git a/src/plugins/intel_gpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp b/src/plugins/intel_gpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp
@@ -14,6 +14,33 @@ INSTANTIATE_TEST_SUITE_P(smoke_RoPETestFlux,
                             ::testing::Values(ov::test::utils::DEVICE_GPU)),
                          RoPETestFlux::getTestCaseName);
 
+class GPURoPETestQwenVL : public RoPETestQwenVL {
+protected:
+    void SetUp() override {
+        RoPETestQwenVL::SetUp();
+        const auto& [element_type, _targetDevice, split_op_type] = this->GetParam();
+        if (element_type == ov::element::f16) {
+            abs_threshold = 0.015f;
+        }
+    }
+};
+
+TEST_P(GPURoPETestQwenVL, CompareWithRefs) {
+    SKIP_IF_CURRENT_TEST_IS_DISABLED();
+    run();
+    auto function = compiledModel.get_runtime_model();
+    CheckNumberOfNodesWithType(function, {"RoPE"}, 1);
+};
+
+const std::vector<std::string> vit_param = {"VariadicSplit", "Slice", "StridedSlice"};
+INSTANTIATE_TEST_SUITE_P(smoke_RoPEQwenVL,
+                         GPURoPETestQwenVL,
+                         ::testing::Combine(
+                            ::testing::Values(ov::element::f16, ov::element::f32),
+                            ::testing::Values(ov::test::utils::DEVICE_GPU),
+                            ::testing::ValuesIn(vit_param)),
+                         GPURoPETestQwenVL::getTestCaseName);
+
 INSTANTIATE_TEST_SUITE_P(smoke_RoPETestChatGLM,
                          RoPETestChatGLMStridedSlice,
                          ::testing::Combine(
diff --git a/src/tests/functional/plugin/shared/include/shared_test_classes/subgraph/rotary_pos_emb.hpp b/src/tests/functional/plugin/shared/include/shared_test_classes/subgraph/rotary_pos_emb.hpp
@@ -11,6 +11,7 @@ namespace test {
 
 using rope_params = std::tuple<ov::element::Type, std::string>;
 using rope_params_2 = std::tuple<bool, ov::element::Type, std::string>;
+using rope_params_qwenvit = std::tuple<ov::element::Type, std::string, std::string>;
 
 class RoPETestFlux : public SubgraphBaseTest, public testing::WithParamInterface<rope_params> {
 private:
@@ -19,6 +20,7 @@ class RoPETestFlux : public SubgraphBaseTest, public testing::WithParamInterface
                                                int num_head,
                                                int ndims,
                                                ov::element::Type element_type);
+
 protected:
     void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
     void SetUp() override;
@@ -27,13 +29,29 @@ class RoPETestFlux : public SubgraphBaseTest, public testing::WithParamInterface
     static std::string getTestCaseName(const testing::TestParamInfo<rope_params>& obj);
 };
 
+class RoPETestQwenVL : public SubgraphBaseTest, public testing::WithParamInterface<rope_params_qwenvit> {
+private:
+    std::shared_ptr<ov::Model> buildROPE_QwenVL(ov::element::Type element_type,
+                                                ov::PartialShape input_shape,
+                                                ov::PartialShape cos_shape,
+                                                ov::PartialShape sin_shape,
+                                                std::string split_op_type);
+
+protected:
+    void SetUp() override;
+
+public:
+    static std::string getTestCaseName(const testing::TestParamInfo<rope_params_qwenvit>& obj);
+};
+
 class RoPETestLlama2StridedSlice : public SubgraphBaseTest, public testing::WithParamInterface<rope_params> {
 private:
     std::shared_ptr<ov::Model> buildROPE_Llama2(int batch,
                                                 int seq_length,
                                                 int max_position_embeddings,
                                                 int num_head,
                                                 int ndims);
+
 protected:
     ov::Tensor create_i32_tensor(const ov::Shape& shape, int start, int step = 1);
     ov::OutputVector makeCosSinCache(int max_position_embeddings, int rotary_ndims);
@@ -46,7 +64,11 @@ class RoPETestLlama2StridedSlice : public SubgraphBaseTest, public testing::With
 
 class RoPETestChatGLMStridedSlice : public SubgraphBaseTest, public testing::WithParamInterface<rope_params> {
 private:
-    std::shared_ptr<ov::Model> buildROPE_ChatGLM(int batch, int head_cnt, int rotary_dims, ov::element::Type element_type);
+    std::shared_ptr<ov::Model> buildROPE_ChatGLM(int batch,
+                                                 int head_cnt,
+                                                 int rotary_dims,
+                                                 ov::element::Type element_type);
+
 protected:
     ov::Tensor create_i32_tensor(const ov::Shape& shape, int start, int step = 1);
     void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
@@ -59,6 +81,7 @@ class RoPETestChatGLMStridedSlice : public SubgraphBaseTest, public testing::Wit
 class RoPETestQwen7bStridedSlice : public SubgraphBaseTest, public testing::WithParamInterface<rope_params_2> {
 private:
     std::shared_ptr<ov::Model> buildROPE_QWen7b(bool specialReshape, ov::element::Type element_type);
+
 protected:
     void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
     void SetUp() override;
@@ -74,6 +97,7 @@ class RoPETestGPTJStridedSlice : public SubgraphBaseTest, public testing::WithPa
                                               int rotary_dims,
                                               bool hasShapeOf,
                                               ov::element::Type element_type);
+
 protected:
     void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
     void SetUp() override;
@@ -91,6 +115,7 @@ class RoPETestRotateHalfWithoutTranspose : public SubgraphBaseTest, public testi
                                                                     int max_position_embeddings,
                                                                     int num_head,
                                                                     int ndims);
+
 protected:
     void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
     void SetUp() override;
@@ -106,38 +131,49 @@ class RoPETestLlama2Slice : public RoPETestLlama2StridedSlice {
                                                 int max_position_embeddings,
                                                 int num_head,
                                                 int ndims);
+
 protected:
     void SetUp() override;
 };
 
 class RoPETestChatGLMSlice : public RoPETestChatGLMStridedSlice {
 private:
-    std::shared_ptr<ov::Model> buildROPE_ChatGLM(int batch, int head_cnt, int rotary_dims, ov::element::Type element_type);
+    std::shared_ptr<ov::Model> buildROPE_ChatGLM(int batch,
+                                                 int head_cnt,
+                                                 int rotary_dims,
+                                                 ov::element::Type element_type);
+
 protected:
     void SetUp() override;
 };
 
 class RoPETestQwen7bSlice : public RoPETestQwen7bStridedSlice {
 private:
     std::shared_ptr<ov::Model> buildROPE_Qwen7b(bool specialReshape, ov::element::Type element_type);
+
 protected:
     void SetUp() override;
 };
 
 class RoPETestGPTJSlice : public RoPETestGPTJStridedSlice {
 private:
     std::shared_ptr<ov::Model> buildROPE_GPTJ(int num_head,
-                                                int hidden_dims,
-                                                int rotary_dims,
-                                                bool hasShapeOf,
-                                                ov::element::Type element_type);
+                                              int hidden_dims,
+                                              int rotary_dims,
+                                              bool hasShapeOf,
+                                              ov::element::Type element_type);
+
 protected:
     void SetUp() override;
 };
 
 class RoPETestChatGLM2DRoPEStridedSlice : public SubgraphBaseTest, public testing::WithParamInterface<rope_params> {
 private:
-    std::shared_ptr<ov::Model> buildROPE_ChatGLM(int batch, int head_cnt, int rotary_dims, ov::element::Type element_type);
+    std::shared_ptr<ov::Model> buildROPE_ChatGLM(int batch,
+                                                 int head_cnt,
+                                                 int rotary_dims,
+                                                 ov::element::Type element_type);
+
 protected:
     ov::Tensor create_i32_tensor(const ov::Shape& shape, int start, int step = 1);
     void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
@@ -153,6 +189,7 @@ class RoPETestChatGLMHF : public SubgraphBaseTest, public testing::WithParamInte
                                                    int num_heads,
                                                    int rotary_ndims,
                                                    ov::element::Type element_type);
+
 protected:
     void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
     void SetUp() override;
diff --git a/src/tests/functional/plugin/shared/include/subgraph_tests/rotary_pos_emb.hpp b/src/tests/functional/plugin/shared/include/subgraph_tests/rotary_pos_emb.hpp
@@ -31,6 +31,13 @@ TEST_P(RoPETestFlux, CompareWithRefs) {
     CheckNumberOfNodesWithType(function, {"RoPE"}, 1);
 };
 
+TEST_P(RoPETestQwenVL, CompareWithRefs) {
+    SKIP_IF_CURRENT_TEST_IS_DISABLED();
+    run();
+    auto function = compiledModel.get_runtime_model();
+    CheckNumberOfNodesWithType(function, {"RoPE"}, 1);
+};
+
 TEST_P(RoPETestLlama2StridedSlice, CompareWithRefs) {
     SKIP_IF_CURRENT_TEST_IS_DISABLED();
     run();
diff --git a/src/tests/functional/plugin/shared/src/subgraph/rotary_pos_emb.cpp b/src/tests/functional/plugin/shared/src/subgraph/rotary_pos_emb.cpp
