| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import functools |
| 4 | |
| 5 | import torch |
| 6 | from torch import Tensor |
| 7 | from torch.library import Library |
| 8 | |
| 9 | from vllm import ir |
| 10 | from vllm.platforms import current_platform |
| 11 | from vllm.utils.torch_utils import direct_register_custom_op |
| 12 | |
| 13 | current_platform.import_kernels() |
| 14 | |
| 15 | |
| 16 | def is_aiter_found() -> bool: |
| 17 | from importlib.util import find_spec |
| 18 | |
| 19 | return find_spec("aiter") is not None |
| 20 | |
| 21 | |
| 22 | aiter_lib = Library("vllm_aiter", "FRAGMENT") |
| 23 | """ |
| 24 | This library holds torch custom ops for wrapped AITER ops. |
| 25 | Many AITER ops want to remain invisible to torch.compile even after lowering. |
| 26 | They are thus wrapped into torch custom ops inside the IR op implementations. |
| 27 | """ |
| 28 | |
| 29 | direct_register_aiter_op = functools.partial( |
| 30 | direct_register_custom_op, target_lib=aiter_lib |
| 31 | ) |
| 32 | """Syntactic sugar for registering AITER custom ops.""" |
| 33 | |
| 34 | AITER_SUPPORTED = is_aiter_found() |
| 35 | """Most kernels in this file are supported if AITER is installed.""" |
| 36 | |
| 37 | |
| 38 | def flatten_to_2d_is_free(x: Tensor) -> bool: |
| 39 | """Whether ``x.reshape(-1, x.shape[-1])`` is a view with unit-stride rows. |
| 40 | |
| 41 | The AITER norm kernels only take dense 2D inputs, so the wrappers below |
| 42 | flatten the leading dims with ``Tensor.reshape``, which silently falls back |
| 43 | to ``contiguous()`` when the flattened shape is not expressible with the |
| 44 | existing strides, adding a whole-tensor device copy that costs more than |
| 45 | the kernel saves. A strided last dim flattens without a copy but not into a |
| 46 | dense buffer, so it is rejected too. |
| 47 | """ |
| 48 | if x.dim() <= 2: |
| 49 | return True |
| 50 | if x.stride(-1) != 1: |
| 51 | return False |
| 52 | expected_stride = x.size(-1) |
| 53 | for i in range(x.dim() - 2, -1, -1): |
| 54 | if x.size(i) != 1 and x.stride(i) != expected_stride: |
| 55 | return False |
| 56 | expected_stride *= x.size(i) |
| 57 | return True |
| 58 | |
| 59 | |
| 60 | rms_no_var_16bit_only = ( |
| 61 | lambda x, weight, epsilon, variance_size=None: variance_size is None |
| 62 | and x.dtype in (torch.float16, torch.bfloat16) |
| 63 | and (weight is None or weight.dtype == x.dtype) |
| 64 | and flatten_to_2d_is_free(x) |
| 65 | ) |
| 66 | """AITER rms_norm only supports float16 and bfloat16 acts, no var_size override, |
| 67 | requires weight dtype to match x dtype, and requires flattening the activation |
| 68 | to 2D to be free.""" |
| 69 | |
| 70 | |
| 71 | @ir.ops.rms_norm.register_impl( |
| 72 | "aiter", supports_args=rms_no_var_16bit_only, supported=AITER_SUPPORTED |
| 73 | ) |
| 74 | def rms_norm( |
| 75 | x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None |
| 76 | ) -> Tensor: |
| 77 | assert variance_size is None |
| 78 | assert x.dtype in (torch.float16, torch.bfloat16) |
| 79 | if weight is None: |
| 80 | weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) |
| 81 | return torch.ops.vllm_aiter.rms_norm(x, weight, epsilon) |
| 82 | |
| 83 | |
| 84 | def _rms_norm_impl(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor: |
| 85 | from aiter import rms_norm |
| 86 | |
| 87 | if x.dim() > 2: |
| 88 | x_original_shape = x.shape |
| 89 | x = x.reshape(-1, x_original_shape[-1]) |
| 90 | x = rms_norm(x, weight, variance_epsilon) |
| 91 | return x.reshape(x_original_shape) |
| 92 | |
| 93 | return rms_norm(x, weight, variance_epsilon) |
| 94 | |
| 95 | |
| 96 | def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor: |
| 97 | return torch.empty_like(x) |
| 98 | |
| 99 | |
| 100 | direct_register_aiter_op( |
| 101 | op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake |
| 102 | ) |
| 103 | rms_add_no_var_16bit_only = ( |
| 104 | lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None |
| 105 | and x.dtype in (torch.float16, torch.bfloat16) |
| 106 | and (weight is None or weight.dtype == x.dtype) |
| 107 | and flatten_to_2d_is_free(x) |
| 108 | and flatten_to_2d_is_free(x_residual) |
| 109 | ) |
| 110 | """ |
| 111 | AITER fused_add_rms_norm only supports 16-bit activations and no var_size override. |
| 112 | Requires weight dtype to match x dtype, and flattening both the activation and the |
| 113 | residual to 2D to be free. |
| 114 | """ |
| 115 | |
| 116 | |
| 117 | @ir.ops.fused_add_rms_norm.register_impl( |
| 118 | "aiter", supports_args=rms_add_no_var_16bit_only, supported=AITER_SUPPORTED |
| 119 | ) |
| 120 | def fused_add_rms_norm( |
| 121 | x: Tensor, |
| 122 | x_residual: Tensor, |
| 123 | weight: Tensor | None, |
| 124 | epsilon: float, |
| 125 | variance_size: int | None = None, |
| 126 | ) -> tuple[Tensor, Tensor]: |
| 127 | assert variance_size is None |
| 128 | assert x.dtype in (torch.float16, torch.bfloat16) |
| 129 | if weight is None: |
| 130 | weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) |
| 131 | return torch.ops.vllm_aiter.fused_add_rms_norm(x, x_residual, weight, epsilon) |
| 132 | |
| 133 | |
| 134 | def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( |
| 135 | x: torch.Tensor, |
| 136 | residual: torch.Tensor, |
| 137 | weight: torch.Tensor, |
| 138 | variance_epsilon: float, |
| 139 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 140 | from aiter import rmsnorm2d_fwd_with_add |
| 141 | |
| 142 | x_shape = x.shape |
| 143 | residual_shape = residual.shape |
| 144 | hidden_size = x.shape[-1] |
| 145 | x = x.reshape(-1, hidden_size) |
| 146 | residual = residual.reshape(-1, hidden_size) |
| 147 | |
| 148 | # TODO can out = x and residual_out = residual to save memory? |
| 149 | # Need to check if the kernel supports in-place residual output |
| 150 | # (if yes set mutates_args and inplace) |
| 151 | residual_out = torch.empty_like(residual) |
| 152 | out = torch.empty_like(x) |
| 153 | rmsnorm2d_fwd_with_add( |
| 154 | out, # output |
| 155 | x, # input |
| 156 | residual, # residual input |
| 157 | residual_out, # residual output |
| 158 | weight, |
| 159 | variance_epsilon, |
| 160 | ) |
| 161 | return out.reshape(x_shape), residual_out.reshape(residual_shape) |
| 162 | |
| 163 | |
| 164 | def _rocm_aiter_rmsnorm2d_fwd_with_add_fake( |
| 165 | x: torch.Tensor, |
| 166 | residual: torch.Tensor, |
| 167 | weight: torch.Tensor, |
| 168 | variance_epsilon: float, |
| 169 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 170 | residual_out = torch.empty_like(residual) |
| 171 | out = torch.empty_like(x) |
| 172 | return out, residual_out |
| 173 | |
| 174 | |
| 175 | direct_register_aiter_op( |
| 176 | op_name="fused_add_rms_norm", |
| 177 | op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl, |
| 178 | fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake, |
| 179 | ) |