[GPU] Enable RMS fusion to match RMS without gamma#36865
Conversation
github-actions
Bot
added
the
category: transformations
| if (enable_without_gamma) { | ||
| // Pattern 2: RMS without gamma, but multiplied with dynamic input | ||
| // RMS(x) * scale where scale is non-constant (e.g., gate, activation, residual) | ||
| // 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; | ||
| } |
There was a problem hiding this comment.
Why do we even need the enable_without_gamma flag and the if/else construct? Can't we just express everything using auto rms_mul = std::make_shared<pattern::op::Or>(OutputVector{mul_with_gamma, mul_with_scale, mul_or_div});?
You don't seem like using the flag later.
There was a problem hiding this comment.
Basically I agree your comment, but RMSFusion is a common pass shared by both GPU and CPU. I'm not sure but it is likely that the CPU RMS kernel only supports the gamma (scale) path from what I can see, RMSNormExecutor requires scale_size and execute() unconditionally accesses inputs[1]. So the opt-in flag lets each backend enable only the patterns its kernel supports.
There was a problem hiding this comment.
I'm not sure I understand you. The enable_without_gamma flag is only used for pattern building. Actually, when you have it set to true, you include the case in the else part. Unless I don't really understand something about the CPU specifics. @v-Golubev
There was a problem hiding this comment.
Same question as for the code below. Do we really need the flag to express alternatives in the pattern? Try using a single OR only.
There was a problem hiding this comment.
I totally agree your comment. Several months ago an unknown CPU CI pytorch layer test failure occurred when this pattern was enabled for CPU, so I temporarily guarded it behind a flag. I'll remove the flag, always enable the div_x pattern, and re-run CI to confirm.
There was a problem hiding this comment.
@CuriousPanCake Thanks for the thorough review - removed the enable_div_x flag and CI passes cleanly. The code is much simpler now!
Add Pattern 3 branch to handle unit normalization (x * rsqrt(mean(x^2) + eps)) without a trailing gamma multiply. Removes need for separate standalone pass. Verified on Gemma4 v_norm: 90 nodes fused (q_norm, k_norm, v_norm). Signed-off-by: hyunback <hyunback.kim@intel.com>
Signed-off-by: hyunback <hyunback.kim@intel.com>
Signed-off-by: hyunback <hyunback.kim@intel.com>
Signed-off-by: hyunback <hyunback.kim@intel.com>
hyunback
force-pushed
the
gemma4_rms_vnorm_fusion
branch
from
99c7cb3 to
22098bc
Compare
2 participants
Problem:
Gemma4 model uses unit normalization (v_norm) that applies RMS normalization without a learnable gamma parameter:
x * 1/Sqrt(ReduceMean(x^2, axes) + eps). This pattern was not matched by the existing RMSFusion pass, which required a trailing Multiply with gamma.Solution:
Extend the existing RMSFusion pass with a third pattern branch (mul_or_div as terminal node in the Or pattern:
(x * rsqrt(mean(x^2) + eps))without a trailing gamma multiply) to match RMS without gamma.Tickets:
AI Assistance: