Skip to content

[GPU] Enable RMS fusion to match RMS without gamma #36865

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hyunback wants to merge 4 commits into
base: master
Choose a base branch
from
Open

[GPU] Enable RMS fusion to match RMS without gamma #36865

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d2478a3
[GPU] Enable RMS fusion to match RMS without gamma
hyunback Jul 14, 2026
8042b47
Fix CI format-style error
hyunback Jul 14, 2026
a6f86f0
Fix CI test failure issue
hyunback Jul 15, 2026
22098bc
Apply code-review comment
hyunback Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/common/transformations/include/transformations/common_optimizations/rms_fusion.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ namespace pass {
class RMSFusion : public ov::pass::MatcherPass {
public:
OPENVINO_MATCHER_PASS_RTTI("RMSFusion");
RMSFusion(bool force_tail_convert = true, bool enable_div_x = false, bool enable_without_gamma = false);
RMSFusion(bool force_tail_convert = true, bool enable_without_gamma = false);
};

} // namespace pass
Expand Down
26 changes: 15 additions & 11 deletions src/common/transformations/src/transformations/common_optimizations/rms_fusion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ namespace op_util = ov::op::util;

namespace ov::pass {

RMSFusion::RMSFusion(bool force_tail_convert, bool enable_div_x, bool enable_without_gamma) {
RMSFusion::RMSFusion(bool force_tail_convert, bool enable_without_gamma) {
// Detect RMS decomposition pattern
// x * 1/Sqrt(ReduceMean(x^2,axes)+eps) * gamma
auto x = pattern::any_input();
Expand Down Expand Up @@ -87,15 +87,9 @@ RMSFusion::RMSFusion(bool force_tail_convert, bool enable_div_x, bool enable_wit
// x * 1/Sqrt(ReduceMean(x^2,axes)+eps)
auto mul1 = pattern::wrap_type<v1::Multiply>({x, div_or_pow});

std::shared_ptr<pattern::op::Or> mul_or_div;
// TODO: Check div_x pattern failed in CPU CI Pytorch layer test.
if (enable_div_x) {
// x / Sqrt(ReduceMean(x^2,axes)+eps)
auto div_x = pattern::wrap_type<v1::Divide>({x, sqrt});
mul_or_div = std::make_shared<pattern::op::Or>(OutputVector{mul1, div_x});
} else {
mul_or_div = std::make_shared<pattern::op::Or>(OutputVector{mul1});
}
// x / Sqrt(ReduceMean(x^2,axes)+eps)
auto div_x = pattern::wrap_type<v1::Divide>({x, sqrt});
auto mul_or_div = std::make_shared<pattern::op::Or>(OutputVector{mul1, div_x});

// Pattern 1: RMS with gamma (learnable parameter)
// x * 1/Sqrt(ReduceMean(x^2,axes)+eps) * gamma (gamma is constant)
Expand All @@ -110,7 +104,9 @@ RMSFusion::RMSFusion(bool force_tail_convert, bool enable_div_x, bool enable_wit
// This allows partial fusion: only fuse up to mul_or_div
auto scale = pattern::any_input(pattern::class_other_than<v0::Constant>());
auto mul_with_scale = pattern::wrap_type<v1::Multiply>({mul_or_div, scale});
rms_mul = std::make_shared<pattern::op::Or>(OutputVector{mul_with_gamma, mul_with_scale});
// Pattern 3: RMS without gamma (unit normalization, e.g. Gemma v_norm)
// x * 1/Sqrt(ReduceMean(x^2,axes)+eps) — no trailing Multiply
rms_mul = std::make_shared<pattern::op::Or>(OutputVector{mul_with_gamma, mul_with_scale, mul_or_div});
} else {
rms_mul = mul_with_gamma;
}
Expand Down Expand Up @@ -139,6 +135,14 @@ RMSFusion::RMSFusion(bool force_tail_convert, bool enable_div_x, bool enable_wit
auto mul_or_div_node = pattern_map.at(mul_or_div).get_node_shared_ptr();
bool elementwise_affine = pattern_map.count(mul_with_gamma);

// Avoid partial fusion (fusing without gamma) when gamma Multiply follows
if (!elementwise_affine && m.get_match_root() == mul_or_div_node) {
for (auto& target : mul_or_div_node->output(0).get_target_inputs()) {
if (ov::is_type<v1::Multiply>(target.get_node()))
return false;
}
}

std::shared_ptr<ov::Node> gamma_node;
if (elementwise_affine) {
gamma_node = pattern_map.at(gamma).get_node_shared_ptr();
Expand Down
37 changes: 31 additions & 6 deletions src/common/transformations/tests/common_optimizations/rms_norm_decomposition_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -273,7 +273,7 @@ TEST_F(TransformationTestsF, RMSNormFusionTest8) {
auto comp = std::make_shared<ov::op::v0::Convert>(mul, ov::element::f16);

model = std::make_shared<ov::Model>(ov::OutputVector{comp}, ov::ParameterVector{input});
manager.register_pass<RMSFusion>(true, true);
manager.register_pass<RMSFusion>(true);
}
{
auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{1, 2, 6});
Expand Down Expand Up @@ -307,7 +307,7 @@ TEST_F(TransformationTestsF, RMSNormFusionTest9) {
auto mul = std::make_shared<ov::op::v1::Multiply>(gamma, div);

model = std::make_shared<ov::Model>(ov::OutputVector{mul}, ov::ParameterVector{input});
manager.register_pass<RMSFusion>(false, true);
manager.register_pass<RMSFusion>(false);
}
{
auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{1, 2, 6});
Expand Down Expand Up @@ -342,7 +342,7 @@ TEST_F(TransformationTestsF, RMSNormFusionTest10) {
auto mul2 = std::make_shared<ov::op::v1::Multiply>(mul1, scale);

model = std::make_shared<ov::Model>(ov::OutputVector{mul2}, ov::ParameterVector{input, scale});
manager.register_pass<RMSFusion>(false, false, true);
manager.register_pass<RMSFusion>(false, true);
}
{
auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{1, 2, 6});
Expand Down Expand Up @@ -374,7 +374,7 @@ TEST_F(TransformationTestsF, RMSNormFusionTest11) {
auto mul2 = std::make_shared<ov::op::v1::Multiply>(mul1, scale);

model = std::make_shared<ov::Model>(ov::OutputVector{mul2}, ov::ParameterVector{input, scale});
manager.register_pass<RMSFusion>(false, false, true);
manager.register_pass<RMSFusion>(false, true);
}
{
auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 6});
Expand Down Expand Up @@ -517,7 +517,7 @@ TEST_F(TransformationTestsF, RMSNormFusionTest15_PowerNegHalf_NoGammaScale) {
auto mul2 = std::make_shared<ov::op::v1::Multiply>(mul1, scale);

model = std::make_shared<ov::Model>(ov::OutputVector{mul2}, ov::ParameterVector{input, scale});
manager.register_pass<RMSFusion>(false, false, true);
manager.register_pass<RMSFusion>(false, true);
}
{
auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 6});
Expand Down Expand Up @@ -655,7 +655,7 @@ TEST_F(TransformationTestsF, RMSNormFusionTest19_MulSquare_NoGamma_DynamicScale)
auto mul2 = std::make_shared<ov::op::v1::Multiply>(mul1, scale);

model = std::make_shared<ov::Model>(ov::OutputVector{mul2}, ov::ParameterVector{input, scale});
manager.register_pass<RMSFusion>(false, false, true);
manager.register_pass<RMSFusion>(false, true);
}
{
auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 2048});
Expand All @@ -666,4 +666,29 @@ TEST_F(TransformationTestsF, RMSNormFusionTest19_MulSquare_NoGamma_DynamicScale)
model_ref = std::make_shared<ov::Model>(ov::OutputVector{mul}, ov::ParameterVector{input, scale});
}
comparator.enable(FunctionsComparator::CmpValues::ACCURACY);
}

// Power(-0.5) direct path without gamma, no trailing Multiply (unit normalization)
TEST_F(TransformationTestsF, RMSNormFusionTest20_PowerNegHalf_NoGamma) {
{
auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{1, 8, 256});
auto power_const = ov::op::v0::Constant::create(ov::element::f32, ov::Shape{}, {2.f});
auto power = std::make_shared<ov::op::v1::Power>(input, power_const);
auto mean_axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1});
auto mean = std::make_shared<ov::op::v1::ReduceMean>(power, mean_axes, true);
auto eps = ov::op::v0::Constant::create(ov::element::f32, ov::Shape{}, {1e-6f});
auto add_eps = std::make_shared<ov::op::v1::Add>(mean, eps);
auto neg_half = ov::op::v0::Constant::create(ov::element::f32, ov::Shape{}, {-0.5f});
auto rsqrt = std::make_shared<ov::op::v1::Power>(add_eps, neg_half);
auto mul = std::make_shared<ov::op::v1::Multiply>(input, rsqrt);

model = std::make_shared<ov::Model>(ov::OutputVector{mul}, ov::ParameterVector{input});
manager.register_pass<RMSFusion>(false, true);
}
{
auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{1, 8, 256});
auto rms = std::make_shared<ov::op::internal::RMS>(input, 1e-6f, ov::element::f32);
model_ref = std::make_shared<ov::Model>(ov::OutputVector{rms}, ov::ParameterVector{input});
}
comparator.enable(FunctionsComparator::CmpValues::ACCURACY);
}
2 changes: 1 addition & 1 deletion src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -747,7 +747,7 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
const int32_t vec_size = 8;
return static_cast<int32_t>((gamma_shape.back() / vec_size)) > static_cast<int32_t>(device_info.max_work_group_size);
});
manager.register_pass<ov::pass::RMSFusion>(false, true, true);
manager.register_pass<ov::pass::RMSFusion>(false, true);
manager.register_pass<DisableFP16CompForGemma3RMSPattern>();
manager.register_pass<DisableFP16ComForGPTOSSROPEPattern>();
manager.register_pass<DisableFP16CompCumSumSinGen>();
Expand Down
Loading