Use Llama RMSNorm for Gemma

[WoosukKwon]

Gemma's RMSNorm is only slightly different from Llama's RMSNorm. Thus, we can use the existing custom op for it. This optimization leads to ~10% latency reduction.

[Yard1]

Nice!

[Yard1]

Before we merge, let's make sure it doesn't change the outputs (maybe we could add a test like we have for other models, using transformers as a reference).

[WoosukKwon]

For a note, using the custom op brings a slight numerical difference in handling the residual connection.

While the original implementation uses the current dtype (f16 of bf16) in hidden_states + residual, the fused RMSNorm op upcasts both to FP32 before addition:

vllm/vllm/model_executor/layers/layernorm.py

Line 35 in 8fbd84b

x = x + residual.to(torch.float32)