| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | from torch import Tensor |
| 5 | |
| 6 | from vllm import ir |
| 7 | from vllm.platforms import current_platform |
| 8 | |
| 9 | current_platform.import_kernels() |
| 10 | |
| 11 | CUDA_ALIKE = current_platform.is_cuda_alike() |
| 12 | """Most kernels in this file are supported on all CUDA-alike platforms.""" |
| 13 | IS_ROCM = current_platform.is_rocm() |
| 14 | """ROCm needs shape normalization before calling some vLLM C kernels.""" |
| 15 | GPGPU_DEVICE = CUDA_ALIKE or current_platform.is_xpu() |
| 16 | |
| 17 | rms_no_var_size = lambda x, weight, epsilon, variance_size=None: ( |
| 18 | variance_size is None and (weight is None or weight.dtype == x.dtype) |
| 19 | ) |
| 20 | """vLLM kernel requires no variance_size override and matching input/weight dtype.""" |
| 21 | |
| 22 | |
| 23 | @ir.ops.rms_norm.register_impl( |
| 24 | "vllm_c", supports_args=rms_no_var_size, supported=GPGPU_DEVICE |
| 25 | ) |
| 26 | def rms_norm( |
| 27 | x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None |
| 28 | ) -> Tensor: |
| 29 | assert variance_size is None |
| 30 | # ROCm's vLLM C RMSNorm kernel operates on contiguous 2D tensors. |
| 31 | # Higher-rank callers still normalize over the last dimension, so flatten |
| 32 | # all leading dims. reshape handles strided views from q/k/v splits. |
| 33 | if IS_ROCM and (x.dim() > 2 or not x.is_contiguous()): |
| 34 | original_shape = x.shape |
| 35 | x = x.reshape(-1, original_shape[-1]) |
| 36 | # empty_like preserves the strides of transposed inputs, but the |
| 37 | # libtorch-stable kernel requires a contiguous output tensor. |
| 38 | output = torch.empty(x.shape, device=x.device, dtype=x.dtype) |
| 39 | torch.ops._C.rms_norm(output, x, weight, epsilon) |
| 40 | return output.reshape(original_shape) |
| 41 | |
| 42 | output = torch.empty(x.shape, device=x.device, dtype=x.dtype) |
| 43 | torch.ops._C.rms_norm(output, x, weight, epsilon) |
| 44 | return output |
| 45 | |
| 46 | |
| 47 | rms_add_no_var_size = lambda x, x_residual, weight, epsilon, variance_size=None: ( |
| 48 | variance_size is None and (weight is None or weight.dtype == x.dtype) |
| 49 | ) |
| 50 | """vLLM Kernel does not support variance_size parameter and requires |
| 51 | matching input/weight dtype.""" |
| 52 | |
| 53 | |
| 54 | @ir.ops.fused_add_rms_norm.register_impl( |
| 55 | "vllm_c", |
| 56 | supports_args=rms_add_no_var_size, |
| 57 | supported=GPGPU_DEVICE, |
| 58 | inplace=True, |
| 59 | ) |
| 60 | def fused_add_rms_norm( |
| 61 | x: Tensor, |
| 62 | x_residual: Tensor, |
| 63 | weight: Tensor | None, |
| 64 | epsilon: float, |
| 65 | variance_size: int | None = None, |
| 66 | ) -> tuple[Tensor, Tensor]: |
| 67 | assert variance_size is None |
| 68 | if IS_ROCM and (not x.is_contiguous() or not x_residual.is_contiguous()): |
| 69 | output, residual = ir.ops.fused_add_rms_norm.impls["native"].impl_fn( |
| 70 | x, x_residual, weight, epsilon |
| 71 | ) |
| 72 | x.copy_(output) |
| 73 | x_residual.copy_(residual) |
| 74 | return x, x_residual |
| 75 | |
| 76 | # ROCm's vLLM C RMSNorm kernel operates on contiguous 2D tensors. |
| 77 | # Higher-rank callers still normalize over the last dimension, so flatten |
| 78 | # all leading dims. |
| 79 | if IS_ROCM and x.dim() > 2: |
| 80 | original_shape = x.shape |
| 81 | x = x.view(-1, original_shape[-1]) |
| 82 | x_residual = x_residual.view(-1, original_shape[-1]) |
| 83 | torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon) |
| 84 | return x.view(original_shape), x_residual.view(original_shape) |
| 85 | |
| 86 | torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon) |
| 87 | return x, x_residual |