| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Custom normalization layers.""" |
| 4 | |
| 5 | import torch |
| 6 | import torch.nn as nn |
| 7 | import torch.nn.functional as F |
| 8 | |
| 9 | # Import kernels |
| 10 | import vllm.kernels # noqa: F401 |
| 11 | from vllm import envs, ir |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.model_executor.custom_op import CustomOp |
| 14 | from vllm.model_executor.determinism.batch_invariant import rms_norm_batch_invariant |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | def poly_norm( |
| 20 | x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, variance_epsilon: float |
| 21 | ) -> torch.Tensor: |
| 22 | from vllm import _custom_ops as ops |
| 23 | |
| 24 | out = torch.empty_like(x) |
| 25 | ops.poly_norm( # type: ignore[attr-defined] |
| 26 | out, |
| 27 | x, |
| 28 | weight, |
| 29 | bias, |
| 30 | variance_epsilon, |
| 31 | ) |
| 32 | return out |
| 33 | |
| 34 | |
| 35 | # --8<-- [start:rms_norm] |
| 36 | @CustomOp.register("rms_norm") |
| 37 | class RMSNorm(CustomOp): |
| 38 | """Root mean square normalization. |
| 39 | |
| 40 | Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. |
| 41 | Refer to https://arxiv.org/abs/1910.07467 |
| 42 | """ |
| 43 | |
| 44 | # --8<-- [end:rms_norm] |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | hidden_size: int, |
| 49 | eps: float = 1e-6, |
| 50 | var_hidden_size: int | None = None, |
| 51 | has_weight: bool = True, |
| 52 | dtype: torch.dtype | None = None, |
| 53 | ) -> None: |
| 54 | super().__init__() |
| 55 | |
| 56 | self.hidden_size = hidden_size |
| 57 | self.variance_epsilon = eps |
| 58 | self.variance_size_override = ( |
| 59 | None if var_hidden_size == hidden_size else var_hidden_size |
| 60 | ) |
| 61 | weight_dtype = dtype or torch.get_default_dtype() |
| 62 | self.has_weight = has_weight |
| 63 | self.weight = torch.ones(hidden_size, dtype=weight_dtype) |
| 64 | if self.has_weight: |
| 65 | self.weight = nn.Parameter(self.weight) |
| 66 | |
| 67 | # When has_weight=False, pass weight=None so implementations that |
| 68 | # support a weightless path can skip the per-channel multiply. |
| 69 | # Implementations that require weight (e.g. oink) fall back via IR |
| 70 | # op priority when weight=None is unsupported. |
| 71 | self.pass_weight = self.has_weight |
| 72 | self.pass_weight_add = self.has_weight |
| 73 | |
| 74 | def forward_native( |
| 75 | self, |
| 76 | x: torch.Tensor, |
| 77 | residual: torch.Tensor | None = None, |
| 78 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 79 | """PyTorch-native implementation equivalent to forward().""" |
| 80 | if residual is None: |
| 81 | return ir.ops.rms_norm( |
| 82 | x, |
| 83 | self.weight.data if self.pass_weight else None, |
| 84 | self.variance_epsilon, |
| 85 | self.variance_size_override, |
| 86 | ) |
| 87 | else: |
| 88 | return ir.ops.fused_add_rms_norm.maybe_inplace( |
| 89 | x, |
| 90 | residual, |
| 91 | self.weight.data if self.pass_weight_add else None, |
| 92 | self.variance_epsilon, |
| 93 | self.variance_size_override, |
| 94 | ) |
| 95 | |
| 96 | def forward_cuda( |
| 97 | self, |
| 98 | x: torch.Tensor, |
| 99 | residual: torch.Tensor | None = None, |
| 100 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 101 | if envs.VLLM_BATCH_INVARIANT: |
| 102 | assert self.variance_size_override is None, ( |
| 103 | "Batch invariance is not supported for variance_size_override" |
| 104 | ) |
| 105 | pass_weight = ( |
| 106 | self.pass_weight_add if residual is not None else self.pass_weight |
| 107 | ) |
| 108 | return rms_norm_batch_invariant( |
| 109 | x, |
| 110 | self.weight.data if pass_weight else None, |
| 111 | self.variance_epsilon, |
| 112 | residual=residual, |
| 113 | ) |
| 114 | |
| 115 | return self.forward_native(x, residual) |
| 116 | |
| 117 | def forward_xpu( |
| 118 | self, |
| 119 | x: torch.Tensor, |
| 120 | residual: torch.Tensor | None = None, |
| 121 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 122 | if envs.VLLM_BATCH_INVARIANT and residual is not None: |
| 123 | assert self.variance_size_override is None, ( |
| 124 | "Batch invariance is not supported for variance_size_override" |
| 125 | ) |
| 126 | weight = self.weight.data if self.pass_weight_add else None |
| 127 | return ir.ops.fused_add_rms_norm.impls["native"].impl_fn( |
| 128 | x, residual, weight, self.variance_epsilon |
| 129 | ) |
| 130 | return self.forward_cuda(x, residual) |
| 131 | |
| 132 | def extra_repr(self) -> str: |
| 133 | s = f"hidden_size={self.weight.data.size(0)}" |
| 134 | s += f", eps={self.variance_epsilon}" |
| 135 | return s |
| 136 | |
| 137 | |
| 138 | # --8<-- [start:gemma_rms_norm] |
| 139 | @CustomOp.register("gemma_rms_norm") |
| 140 | class GemmaRMSNorm(CustomOp): |
| 141 | """RMS normalization for Gemma. |
| 142 | |
| 143 | Two differences from the above RMSNorm: |
| 144 | 1. x * (1 + w) instead of x * w. |
| 145 | 2. (x * w).to(orig_dtype) instead of x.to(orig_dtype) * w. |
| 146 | """ |
| 147 | |
| 148 | # --8<-- [end:gemma_rms_norm] |
| 149 | |
| 150 | def __init__( |
| 151 | self, |
| 152 | hidden_size: int, |
| 153 | eps: float = 1e-6, |
| 154 | ) -> None: |
| 155 | super().__init__() |
| 156 | self.weight = nn.Parameter(torch.zeros(hidden_size)) |
| 157 | self.variance_epsilon = eps |
| 158 | |
| 159 | def forward_native( |
| 160 | self, |
| 161 | x: torch.Tensor, |
| 162 | residual: torch.Tensor | None = None, |
| 163 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 164 | """PyTorch-native implementation equivalent to forward().""" |
| 165 | weight = self.weight.float() + 1.0 |
| 166 | if residual is None: |
| 167 | return ir.ops.rms_norm(x, weight, self.variance_epsilon) |
| 168 | return ir.ops.fused_add_rms_norm(x, residual, weight, self.variance_epsilon) |
| 169 | |
| 170 | def forward_cuda( |
| 171 | self, |
| 172 | x: torch.Tensor, |
| 173 | residual: torch.Tensor | None = None, |
| 174 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 175 | return self.forward_native(x, residual) |
| 176 | |
| 177 | def forward_xpu( |
| 178 | self, |
| 179 | x: torch.Tensor, |
| 180 | residual: torch.Tensor | None = None, |
| 181 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 182 | if envs.VLLM_BATCH_INVARIANT: |
| 183 | if residual is None: |
| 184 | return self.forward_cuda(x, residual) |
| 185 | return ir.ops.fused_add_rms_norm.impls["native"].impl_fn( |
| 186 | x, |
| 187 | residual, |
| 188 | self.weight.float() + 1.0, |
| 189 | self.variance_epsilon, |
| 190 | ) |
| 191 | import vllm._xpu_ops # noqa: F401 registers torch.ops.vllm.xpu_gemma_rms_norm |
| 192 | |
| 193 | # Fall back to the native path if the fused gemma kernels are not |
| 194 | # available in the installed vllm-xpu-kernels package. |
| 195 | if not hasattr(torch.ops._C, "gemma_rms_norm"): |
| 196 | return self.forward_native(x, residual) |
| 197 | |
| 198 | # Pass the raw (bf16/fp16) weight; the +1 offset and the fp32 multiply |
| 199 | # are folded into the kernel (matches forward_native numerics). |
| 200 | if residual is not None: |
| 201 | torch.ops.vllm.xpu_fused_add_gemma_rms_norm( |
| 202 | x, residual, self.weight.data, self.variance_epsilon |
| 203 | ) |
| 204 | return x, residual |
| 205 | # empty_like preserves x's strides, but the kernel requires a |
| 206 | # contiguous out (unlike x, which it can handle non-contiguous). |
| 207 | out = torch.empty(x.shape, device=x.device, dtype=x.dtype) |
| 208 | torch.ops.vllm.xpu_gemma_rms_norm( |
| 209 | out, x, self.weight.data, self.variance_epsilon |
| 210 | ) |
| 211 | return out |
| 212 | |
| 213 | |
| 214 | @CustomOp.register("layer_norm") |
| 215 | class LayerNorm(CustomOp): |
| 216 | """Standard (mean-centered) LayerNorm. |
| 217 | |
| 218 | Drop-in for a bare `nn.LayerNorm` that dispatches to a fused XPU kernel |
| 219 | when one is available, and to the native implementation otherwise. |
| 220 | |
| 221 | Normalization runs in the wider of the input and parameter dtypes and the |
| 222 | result is cast back to the input dtype, so `dtype=torch.float32` gives an |
| 223 | fp32 reduction for a lower-precision input. |
| 224 | """ |
| 225 | |
| 226 | def __init__( |
| 227 | self, |
| 228 | hidden_size: int, |
| 229 | eps: float = 1e-5, |
| 230 | elementwise_affine: bool = True, |
| 231 | bias: bool = True, |
| 232 | dtype: torch.dtype | None = None, |
| 233 | ) -> None: |
| 234 | super().__init__() |
| 235 | self.normalized_shape = (hidden_size,) |
| 236 | self.eps = eps |
| 237 | weight_dtype = dtype or torch.get_default_dtype() |
| 238 | self.weight: nn.Parameter | None = None |
| 239 | self.bias: nn.Parameter | None = None |
| 240 | if elementwise_affine: |
| 241 | self.weight = nn.Parameter(torch.ones(hidden_size, dtype=weight_dtype)) |
| 242 | if bias: |
| 243 | self.bias = nn.Parameter(torch.zeros(hidden_size, dtype=weight_dtype)) |
| 244 | |
| 245 | def _compute_dtype(self, x: torch.Tensor) -> torch.dtype: |
| 246 | if self.weight is None: |
| 247 | return x.dtype |
| 248 | return torch.promote_types(x.dtype, self.weight.dtype) |
| 249 | |
| 250 | def forward_native(self, x: torch.Tensor) -> torch.Tensor: |
| 251 | compute_dtype = self._compute_dtype(x) |
| 252 | if compute_dtype != x.dtype: |
| 253 | return F.layer_norm( |
| 254 | x.to(compute_dtype), |
| 255 | self.normalized_shape, |
| 256 | self.weight, |
| 257 | self.bias, |
| 258 | self.eps, |
| 259 | ).type_as(x) |
| 260 | return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) |
| 261 | |
| 262 | def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: |
| 263 | return self.forward_native(x) |
| 264 | |
| 265 | def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: |
| 266 | import vllm._xpu_ops # noqa: F401 registers torch.ops.vllm.xpu_layer_norm |
| 267 | |
| 268 | if ( |
| 269 | self.weight is None |
| 270 | or self.bias is None |
| 271 | or not hasattr(torch.ops._C, "layer_norm") |
| 272 | # The kernel reduces in the input dtype, so a widened reduction |
| 273 | # takes the native path. |
| 274 | or self._compute_dtype(x) != x.dtype |
| 275 | ): |
| 276 | return self.forward_native(x) |
| 277 | # empty_like preserves x's strides, but the kernel requires a |
| 278 | # contiguous out (unlike x, which it can handle non-contiguous). |
| 279 | out = torch.empty(x.shape, device=x.device, dtype=x.dtype) |
| 280 | torch.ops.vllm.xpu_layer_norm(out, x, self.weight, self.bias, self.eps) |
| 281 | return out |
| 282 | |
| 283 | def extra_repr(self) -> str: |
| 284 | return f"hidden_size={self.normalized_shape[0]}, eps={self.eps}" |
| 285 | |
| 286 | |
| 287 | # --8<-- [start:rms_norm_gated] |
| 288 | @CustomOp.register("rms_norm_gated") |
| 289 | class RMSNormGated(CustomOp): |
| 290 | """RMS Normalization with optional gating. |
| 291 | |
| 292 | This is a native PyTorch implementation that supports: |
| 293 | - Standard RMS normalization |
| 294 | - Group RMS normalization |
| 295 | - Optional gating with SiLU activation |
| 296 | """ |
| 297 | |
| 298 | # --8<-- [end:rms_norm_gated] |
| 299 | |
| 300 | def __init__( |
| 301 | self, |
| 302 | hidden_size: int, |
| 303 | eps: float = 1e-5, |
| 304 | group_size: int | None = None, |
| 305 | norm_before_gate: bool = False, |
| 306 | device: torch.device | None = None, |
| 307 | dtype: torch.dtype | None = None, |
| 308 | activation: str = "swish", |
| 309 | ): |
| 310 | """Initialize RMSNormGated. |
| 311 | |
| 312 | Args: |
| 313 | hidden_size: Size of the hidden dimension |
| 314 | eps: Epsilon for numerical stability |
| 315 | group_size: If not None, do GroupNorm with each group |
| 316 | having group_size elements. |
| 317 | group_size=None is equivalent to group_size=hidden_size |
| 318 | (i.e. there's only 1 group). |
| 319 | norm_before_gate: If True and z is provided: out = norm(x) * silu(z) |
| 320 | If False and z is provided: out = norm(x * silu(z)) |
| 321 | device: Device to create parameters on |
| 322 | dtype: Data type for parameters |
| 323 | activation: Activation function name for gating |
| 324 | |
| 325 | """ |
| 326 | factory_kwargs = {"device": device, "dtype": dtype} |
| 327 | super().__init__() |
| 328 | self.eps = eps |
| 329 | self.activation = activation |
| 330 | self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) |
| 331 | self.register_parameter("bias", None) |
| 332 | self.group_size = group_size |
| 333 | self.norm_before_gate = norm_before_gate |
| 334 | self.reset_parameters() |
| 335 | |
| 336 | def reset_parameters(self): |
| 337 | torch.nn.init.ones_(self.weight) |
| 338 | |
| 339 | @staticmethod |
| 340 | def forward_static( |
| 341 | x: torch.Tensor, |
| 342 | z: torch.Tensor | None, |
| 343 | weight: torch.Tensor, |
| 344 | epsilon: float, |
| 345 | orig_dtype: torch.dtype, |
| 346 | group_size: int | None = None, |
| 347 | norm_before_gate: bool = True, |
| 348 | activation: str = "swish", |
| 349 | ) -> torch.Tensor: |
| 350 | """Pure-PyTorch RMS normalization with optional gating. |
| 351 | |
| 352 | This static method contains the full native logic so that both |
| 353 | ``forward_native`` and ``MatcherRMSNormGated`` (used by the |
| 354 | compilation pattern matcher) can share the same implementation. |
| 355 | |
| 356 | If *z* is not None and *norm_before_gate* is True: |
| 357 | ``out = rms_norm(x) * act(z)`` |
| 358 | If *z* is not None and *norm_before_gate* is False: |
| 359 | ``out = rms_norm(x * act(z))`` |
| 360 | """ |
| 361 | x = x.float() |
| 362 | weight = weight.float() |
| 363 | if z is not None: |
| 364 | z = z.float() |
| 365 | |
| 366 | assert activation in ["silu", "sigmoid", "swish"] |
| 367 | act_fn = F.sigmoid if activation == "sigmoid" else F.silu |
| 368 | |
| 369 | if z is not None and not norm_before_gate: |
| 370 | x = x * act_fn(z) |
| 371 | |
| 372 | if group_size is None: |
| 373 | variance = x.pow(2).mean(dim=-1, keepdim=True) |
| 374 | x_normed = x * torch.rsqrt(variance + epsilon) |
| 375 | out = x_normed * weight |
| 376 | else: |
| 377 | from einops import rearrange |
| 378 | |
| 379 | x_group = rearrange(x, "... (g d) -> ... g d", d=group_size) |
| 380 | variance = x_group.pow(2).mean(dim=-1, keepdim=True) |
| 381 | x_normed = x_group * torch.rsqrt(variance + epsilon) |
| 382 | out = rearrange(x_normed, "... g d -> ... (g d)") * weight |
| 383 | |
| 384 | if z is not None and norm_before_gate: |
| 385 | out = out * act_fn(z) |
| 386 | |
| 387 | return out.to(orig_dtype) |
| 388 | |
| 389 | def forward_native( |
| 390 | self, x: torch.Tensor, z: torch.Tensor | None = None |
| 391 | ) -> torch.Tensor: |
| 392 | """PyTorch-native implementation equivalent to forward().""" |
| 393 | return self.forward_static( |
| 394 | x, |
| 395 | z, |
| 396 | self.weight, |
| 397 | self.eps, |
| 398 | x.dtype, |
| 399 | group_size=self.group_size, |
| 400 | norm_before_gate=self.norm_before_gate, |
| 401 | activation=self.activation, |
| 402 | ) |
| 403 | |
| 404 | def forward_cuda( |
| 405 | self, x: torch.Tensor, z: torch.Tensor | None = None |
| 406 | ) -> torch.Tensor: |
| 407 | from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( |
| 408 | rmsnorm_fn, |
| 409 | ) |
| 410 | |
| 411 | return rmsnorm_fn( |
| 412 | x, |
| 413 | self.weight, |
| 414 | self.bias, |
| 415 | z=z, |
| 416 | eps=self.eps, |
| 417 | group_size=self.group_size, |
| 418 | norm_before_gate=self.norm_before_gate, |
| 419 | activation=self.activation, |
| 420 | ) |
| 421 | |
| 422 | def forward_xpu( |
| 423 | self, x: torch.Tensor, z: torch.Tensor | None = None |
| 424 | ) -> torch.Tensor: |
| 425 | return self.forward_cuda(x, z) |