[XPU] Add xpu forward in Gemma3RMSNorm & Add test and benchmark for Gemma3RMSNorm (#36278)

python/sglang/srt/layers/layernorm.py

@@ -1317,6 +1317,19 @@ class Gemma3RMSNorm(BaseFusedOp):
             return gemma_rmsnorm(x, self.weight.data, self.eps)
         return self.forward_native(x)
 
+    def forward_xpu(self, x, residual: Optional[torch.Tensor] = None):
+        if residual is not None and x.dim() == 2:
+            # The decoder residual is token-major and contiguous. The fused
+            # kernel updates both tensors in place: x becomes the normalized
+            # output and residual becomes x + residual for the next layer.
+            gemma_fused_add_rmsnorm(x, residual, self.weight.data, self.eps)
+            return x, residual
+        # The XPU kernel flattens leading dims internally, so 2D/3D/4D inputs
+        # can all go through it directly without a Python-side reshape.
+        elif residual is None and x.dim() in (2, 3, 4):
+            return gemma_rmsnorm(x, self.weight.data, self.eps)
+        return self.forward_native(x, residual)
+
     def forward_musa(self, x, residual: Optional[torch.Tensor] = None):
         # sgl_kernel's gemma norm ops are built for MUSA; follow the CUDA path.
         return self.forward_cuda(x, residual)