[CPU]Support Rope for GPT-OSS

Author: zhangYiIntel

src/common/transformations/include/ov_ops/rotary_positional_embeddings.hpp

@@ -33,6 +33,7 @@ class TRANSFORMATIONS_API RoPE : public Op {
         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 cos_sin_ndims = 0;       // last dimension of con/sin table
         size_t head_cnt = 0;
         size_t head_size = 0;
         int gather_position_arg_id =

src/common/transformations/include/transformations/common_optimizations/fuse_rotary_positional_embeddings.hpp

@@ -21,6 +21,7 @@ class TRANSFORMATIONS_API RoPEFusionIOSlicing;
 class TRANSFORMATIONS_API RoPEFusionPreprocess;
 class TRANSFORMATIONS_API RoPEFusionCosSinPreprocess;
 class TRANSFORMATIONS_API RoPEShareCosSin;
+class TRANSFORMATIONS_API RoPEFusionGPTOSS;
 
 }  // namespace pass
 }  // namespace ov
@@ -91,6 +92,12 @@ class ov::pass::RoPEShareCosSin : public ov::pass::MatcherPass {
     std::vector<std::shared_ptr<ov::Node>> m_shared_inputs{2, nullptr};
 };
 
+class ov::pass::RoPEFusionGPTOSS : public ov::pass::MatcherPass {
+public:
+    OPENVINO_MATCHER_PASS_RTTI("RoPEFusionGPTOSS");
+    RoPEFusionGPTOSS();
+};
+
 /**
  * @ingroup ov_transformation_common_api
  * @brief Fuses special sub-graph into an internal Rotary Positional Embedding operation

src/common/transformations/src/transformations/common_optimizations/fuse_rotary_positional_embeddings.cpp

@@ -61,6 +61,7 @@ bool ov::pass::RoPEFusion::run_on_model(const std::shared_ptr<ov::Model>& model)
     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>();
+    symbolic_ctx_manager->register_pass<ov::pass::RoPEFusionGPTOSS>();
     // optional heads & tails are fused in separate matcher pass,
     // after RoPENode has been created.
     symbolic_ctx_manager->register_pass<ov::pass::RoPEFusionCosSinPreprocess>();
@@ -1111,3 +1112,74 @@ ov::pass::RoPEShareCosSin::RoPEShareCosSin() {
     auto m = std::make_shared<pattern::Matcher>(result, matcher_name);
     this->register_matcher(m, callback);
 }
+
+ov::pass::RoPEFusionGPTOSS::RoPEFusionGPTOSS() {
+    using namespace ov::op::util;
+    MATCHER_SCOPE(RoPEFusionGPTOSS);
+
+    // gpt-oss style
+    // first_half, second_half = torch.chunk(x, 2, dim=-1)
+    // first_ = first_half * cos - second_half * sin
+    // second_ = second_half * cos + first_half * sin
+    // return torch.cat((first_, second_), dim=-1)
+    auto x = pattern::any_input(pattern::rank_equals(4));
+    auto t_cos = pattern::any_input(pattern::shape_matches("[?, 1, ?, half_ndims]"));
+    auto t_sin = pattern::any_input(pattern::shape_matches("[?, 1, ?, half_ndims]"));
+    auto varsplit = pattern::wrap_type<v1::VariadicSplit>({x, -1, {"half_ndims", "?"}});
+    varsplit->set_output_size(2);
+
+    auto first_half_mul_cos = pattern::wrap_type<v1::Multiply>({varsplit->output(0), t_cos}, {{"auto_broadcast", "numpy"}});
+    auto second_half_mul_sin = pattern::wrap_type<v1::Multiply>({varsplit->output(1), t_sin}, {{"auto_broadcast", "numpy"}});
+    auto neg = pattern::wrap_type<v1::Multiply>({second_half_mul_sin, -1.0f}, {{"auto_broadcast", "numpy"}});
+    auto sub_Subtract = pattern::wrap_type<v1::Add>({first_half_mul_cos, neg}, {{"auto_broadcast", "numpy"}});
+
+    auto second_half_mul_cos = pattern::wrap_type<v1::Multiply>({varsplit->output(1), t_cos}, {{"auto_broadcast", "numpy"}});
+    auto first_half_mul_sin = pattern::wrap_type<v1::Multiply>({varsplit->output(0), t_sin}, {{"auto_broadcast", "numpy"}});
+    auto add_Add = pattern::wrap_type<v1::Add>({second_half_mul_cos, first_half_mul_sin}, {{"auto_broadcast", "numpy"}});
+    auto concat_result = pattern::wrap_type<opset1::Concat>({sub_Subtract, add_Add}, {{"axis", -1}});
+
+    auto result = concat_result;
+
+    matcher_pass_callback callback = [OV_CAPTURE_CPY_AND_THIS](ov::pass::pattern::Matcher& m) {
+        const auto& pattern_map = m.get_pattern_value_map();
+        auto root = m.get_match_root();
+        const auto& x_val = pattern_map.at(x);
+        const auto& v_cos = pattern_map.at(t_cos);
+
+        auto symbols = m.get_symbols();
+        const auto& half_ndims = symbols["half_ndims"];
+        if (!half_ndims.is_integer()) {
+            return false;
+        }
+
+        op::internal::RoPE::Config config;
+        OutputVector new_args;
+        config.rotary_ndims = 2ul * static_cast<size_t>(half_ndims.i());
+        config.cos_sin_ndims = static_cast<size_t>(half_ndims.i());
+
+        new_args.push_back(x_val);
+        new_args.push_back(v_cos);
+        new_args.push_back(pattern_map.at(t_sin));
+        auto old_node = root;
+        auto new_node = std::make_shared<internal::RoPE>(new_args, config);
+        new_node->set_friendly_name(old_node->get_friendly_name());
+        ov::copy_runtime_info({
+                            pattern_map.at(neg).get_node_shared_ptr(),
+                            pattern_map.at(sub_Subtract).get_node_shared_ptr(),
+                            pattern_map.at(first_half_mul_cos).get_node_shared_ptr(),
+                            pattern_map.at(first_half_mul_sin).get_node_shared_ptr(),
+                            pattern_map.at(second_half_mul_cos).get_node_shared_ptr(),
+                            pattern_map.at(second_half_mul_sin).get_node_shared_ptr(),
+                            pattern_map.at(add_Add).get_node_shared_ptr(),
+                            pattern_map.at(result).get_node_shared_ptr()},
+                            new_node);         
+        ov::replace_node(old_node, new_node);
+
+        // this new node may match following additional matchers
+        register_new_node(new_node);
+        return true;
+    };
+
+    auto m = std::make_shared<ov::pass::pattern::Matcher>(result, matcher_name);
+    this->register_matcher(m, callback);
+}

src/plugins/intel_cpu/src/nodes/kernels/x64/rope_kernel.cpp

@@ -90,9 +90,12 @@ void jit_rotary_kernel<isa>::rotary_half(size_t step) {
     store(reg_dst, vmm_dst0, m_jcp.dst_prc, step);
 
     // cos[i + halfRotaryNdims]
-    load(vmm_cos, reg_cos, ov::element::f32, step, false, half_rotary_ndims * sizeof(float));
+    // if con/sin table is not the same size, it is reused for both halves.
+    if (!m_jcp.cos_sin_ndims)
+        load(vmm_cos, reg_cos, ov::element::f32, step, false, half_rotary_ndims * sizeof(float));
     // sin[i + halfRotaryNdims]
-    load(vmm_sin, reg_sin, ov::element::f32, step, false, half_rotary_ndims * sizeof(float));
+    if (!m_jcp.cos_sin_ndims)
+        load(vmm_sin, reg_sin, ov::element::f32, step, false, half_rotary_ndims * sizeof(float));
     // cos[i + half_rotary_dims] * src1
     uni_vmulps(vmm_dst0, vmm_cos, vmm_src1);
     // cos[i + half_rotary_dims] * src1 + sin[i + half_rotary_dims] * src0

src/plugins/intel_cpu/src/nodes/kernels/x64/rope_kernel.hpp

@@ -27,6 +27,7 @@ struct jit_rotary_compile_params {
     ov::element::Type src_prc;
     ov::element::Type dst_prc;
     size_t rotary_ndims = 0UL;
+    size_t cos_sin_ndims = 0UL;
     bool interleave = false;
     bool mix_cos_sin = false;
 };

src/plugins/intel_cpu/src/nodes/rope.cpp

@@ -111,6 +111,7 @@ struct RoPE::RoPEExecutorRotateHalf : public RoPE::Executor {
         jcp.dst_prc = precision_of<T>::value;
         jcp.rotary_ndims = config.rotary_ndims;
         jcp.interleave = false;
+        jcp.cos_sin_ndims = config.cos_sin_ndims;
         m_rotaryKernel = createJitKernel(jcp);
     }
 
@@ -157,7 +158,8 @@ struct RoPE::RoPEExecutorRotateHalf : public RoPE::Executor {
         auto head_cnt = t_src.size(1);
         auto seq_len = t_src.size(2);
         auto feature_size = t_src.size(3);
-
+        auto half_rotary_dims = rotary_dims / 2;
+        const size_t cos_sin_offset = (m_config.cos_sin_ndims == half_rotary_dims) ? 0 : half_rotary_dims;
         parallel_for3d(batch_size, head_cnt, seq_len, [&](size_t b, size_t h, size_t p) {
             auto cos_pos = p;
             if (gather) {
@@ -181,7 +183,7 @@ struct RoPE::RoPEExecutorRotateHalf : public RoPE::Executor {
                     auto src0 = src[i];
                     auto src1 = src[i + half_rotary_dims];
                     dst[i] = cos[i] * src0 - sin[i] * src1;
-                    dst[i + half_rotary_dims] = cos[i + half_rotary_dims] * src1 + sin[i + half_rotary_dims] * src0;
+                    dst[i + half_rotary_dims] = cos[i + cos_sin_offset] * src1 + sin[i + cos_sin_offset] * src0;
                 }
             }
             if (!can_inplace) {

src/plugins/intel_cpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp

@@ -79,5 +79,12 @@ INSTANTIATE_TEST_SUITE_P(smoke_RoPETestQwenVL,
                             ::testing::ValuesIn(vit_param)),
                          RoPETestQwenVL::getTestCaseName);
 
+INSTANTIATE_TEST_SUITE_P(smoke_RoPETestGPTOSS,
+                         RoPETestGPTOSS,
+                         ::testing::Combine(
+                            ::testing::Values(ov::element::f32),
+                            ::testing::Values(ov::test::utils::DEVICE_CPU)),
+                         RoPETestGPTOSS::getTestCaseName);
+
 }  // namespace test
 }  // namespace ov

src/tests/functional/plugin/shared/include/shared_test_classes/subgraph/rotary_pos_emb.hpp

@@ -198,5 +198,19 @@ class RoPETestChatGLMHF : public SubgraphBaseTest, public testing::WithParamInte
     static std::string getTestCaseName(const testing::TestParamInfo<rope_params>& obj);
 };
 
+class RoPETestGPTOSS : public SubgraphBaseTest, public testing::WithParamInterface<rope_params> {
+private:
+    std::shared_ptr<ov::Model> buildROPE_GPTOSS(int num_head,
+                                              int rotary_dims,
+                                              ov::element::Type element_type);
+
+protected:
+    void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
+    void SetUp() override;
+
+public:
+    static std::string getTestCaseName(const testing::TestParamInfo<rope_params>& obj);
+};
+
 }  // namespace test
 }  // namespace ov

src/tests/functional/plugin/shared/include/subgraph_tests/rotary_pos_emb.hpp

@@ -115,5 +115,12 @@ TEST_P(RoPETestChatGLMHF, CompareWithRefs) {
     CheckNumberOfNodesWithType(function, {"RoPE"}, 1);
 };
 
+TEST_P(RoPETestGPTOSS, CompareWithRefs) {
+    SKIP_IF_CURRENT_TEST_IS_DISABLED();
+    run();
+    auto function = compiledModel.get_runtime_model();
+    CheckNumberOfNodesWithType(function, {"RoPE"}, 1);
+};
+
 }  // namespace test
 }  // namespace ov

src/tests/functional/plugin/shared/src/subgraph/rotary_pos_emb.cpp

@@ -1457,5 +1457,83 @@ std::string RoPETestChatGLMHF::getTestCaseName(const testing::TestParamInfo<rope
     return result.str();
 }
 
+std::shared_ptr<ov::Model> RoPETestGPTOSS::buildROPE_GPTOSS(int num_head,
+                                                            int rotary_dims,
+                                                            ov::element::Type element_type) {
+    auto int32_max = std::numeric_limits<std::int32_t>::max();
+    auto input = std::make_shared<ov::opset1::Parameter>(element_type, PartialShape{-1, -1, num_head, rotary_dims});
+    input->set_friendly_name("input");
+    auto permute_Transpose = makeOP<opset1::Transpose>({input, {0, 2, 1, 3}}, {});
+    auto cos = std::make_shared<ov::opset1::Parameter>(element_type, PartialShape{-1, 1, -1, rotary_dims/2});
+    cos->set_friendly_name("cos");
+    auto sin = std::make_shared<ov::opset1::Parameter>(element_type, PartialShape{-1, 1, -1, rotary_dims/2});
+    sin->set_friendly_name("sin");
+    auto variadicSplit = makeOP<opset1::VariadicSplit>({permute_Transpose, -1, {rotary_dims/2, -1}});
+    auto first_half_mul_cos = makeOP<opset1::Multiply>({variadicSplit->output(0), cos}, {{"auto_broadcast", "numpy"}});
+    
+    auto second_half_mul_sin = makeOP<opset1::Multiply>({variadicSplit->output(1), sin}, {{"auto_broadcast", "numpy"}});
+    auto neg = makeOP<opset1::Multiply>({second_half_mul_sin, -1.000000f}, {{"auto_broadcast", "numpy"}});   //  tensor_array<f32[?,64,?,32]> Multiply_9680(__module.model.layers.0.self_attn/aten::mul/Multiply_1, Constant_9679)
+    auto sub_Subtract = makeOP<opset1::Add>({first_half_mul_cos, neg}, {{"auto_broadcast", "numpy"}});   //  tensor_array<f32[?,64,?,32]> __module.model.layers.0.self_attn/aten::sub/Subtract(__module.model.layers.0.self_attn/aten::mul/Multiply, Multiply_9680)
+    
+    
+    auto second_half_mul_cos = makeOP<opset1::Multiply>({variadicSplit->output(1), cos}, {{"auto_broadcast", "numpy"}});
+    auto first_half_mul_sin = makeOP<opset1::Multiply>({variadicSplit->output(0), sin}, {{"auto_broadcast", "numpy"}});   
+    auto add_Add = makeOP<opset1::Add>({second_half_mul_cos, first_half_mul_sin}, {{"auto_broadcast", "numpy"}});   
+    auto cat_Concat = makeOP<opset1::Concat>({sub_Subtract, add_Add}, {{"axis", -1}});   
+    return std::make_shared<ov::Model>(cat_Concat, ov::ParameterVector{input, cos, sin});
+}
+
+void RoPETestGPTOSS::generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) {
+    const auto& funcInputs = function->inputs();
+
+    auto& input_shape = targetInputStaticShapes[0];
+    auto& cos_shape = targetInputStaticShapes[1];
+    auto& sin_shape = targetInputStaticShapes[2];
+
+    ov::test::utils::InputGenerateData in_data;
+    in_data.start_from = -1;
+    in_data.range = 2;
+    in_data.resolution = 32768;
+    auto cos_data = in_data;
+    cos_data.seed = 10;
+
+    auto sin_data = in_data;
+    sin_data.seed = 20;
+    ov::Tensor t_input = utils::create_and_fill_tensor(funcInputs[0].get_element_type(), input_shape, in_data);
+    ov::Tensor t_cos_cache =
+        utils::create_and_fill_tensor(funcInputs[1].get_element_type(), cos_shape, cos_data);
+    ov::Tensor t_sin_cache =
+        utils::create_and_fill_tensor(funcInputs[2].get_element_type(), sin_shape, sin_data);
+
+    inputs.clear();
+    inputs.insert({funcInputs[0].get_node_shared_ptr(), t_input});
+    inputs.insert({funcInputs[1].get_node_shared_ptr(), t_cos_cache});
+    inputs.insert({funcInputs[2].get_node_shared_ptr(), t_sin_cache});
+}
+
+std::string RoPETestGPTOSS::getTestCaseName(const testing::TestParamInfo<rope_params>& obj) {
+    const auto& [element_type, targetDevice] = obj.param;
+    std::ostringstream result;
+    result << "type=" << element_type << "_"
+           << "targetDevice=" << targetDevice;
+    return result.str();
+}
+
+void RoPETestGPTOSS::SetUp() {
+    const auto& [element_type, _targetDevice] = this->GetParam();
+    targetDevice = _targetDevice;
+
+    const int batch = 5;
+    const int seq_length = 7;
+    const int num_head = 8;
+    const int rotary_dims = 64;
+
+    InputShape input = {{batch, seq_length, num_head, rotary_dims}, {{batch, seq_length, num_head, rotary_dims}}};
+    InputShape cos = {{batch, 1, seq_length, rotary_dims/2}, {{batch, 1, seq_length, rotary_dims/2}}};
+    InputShape sin = {{batch, 1, seq_length, rotary_dims/2}, {{batch, 1, seq_length, rotary_dims/2}}};
+    init_input_shapes({input, cos, sin});
+    function = buildROPE_GPTOSS(num_head, rotary_dims, element_type);
+}
+
 }  // namespace test
 }  // namespace ov