| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This file registers Oink implementations for vLLM IR ops. |
| 4 | |
| 5 | vLLM does not depend on the external Oink repository/package. When an external |
| 6 | plugin registers torch.library.custom_op entrypoints under the `oink::` |
| 7 | namespace (e.g. via vLLM's general_plugins mechanism), these ops will be marked |
| 8 | as supported. To dispatch to those ops, set kernel_config.ir_op_priority.<op> to oink. |
| 9 | Alternatively, `VLLM_USE_OINK_OPS=1` will add this to priority by default. |
| 10 | """ |
| 11 | |
| 12 | import torch |
| 13 | from torch import Tensor |
| 14 | |
| 15 | from vllm import ir |
| 16 | from vllm.platforms import current_platform |
| 17 | |
| 18 | OINK_AVAILABLE = current_platform.has_device_capability(100) and hasattr( |
| 19 | torch.ops, "oink" |
| 20 | ) |
| 21 | |
| 22 | |
| 23 | def has_oink_op(name: str) -> bool: |
| 24 | """Check if a specific oink op is registered.""" |
| 25 | return OINK_AVAILABLE and hasattr(torch.ops.oink, name) |
| 26 | |
| 27 | |
| 28 | def _can_view_as_2d(x: Tensor) -> bool: |
| 29 | """Return True if x.view(-1, x.shape[-1]) is viewable (no copy).""" |
| 30 | if x.dim() < 2: |
| 31 | return False |
| 32 | if x.dim() == 2: |
| 33 | return True |
| 34 | # For a view(-1, N) to be valid, all leading dims must be contiguous with |
| 35 | # respect to each other (size-1 dims are ignored). |
| 36 | for dim in range(x.dim() - 1): |
| 37 | # Strides for size-1 dims are irrelevant and can be arbitrary. |
| 38 | if x.size(dim + 1) != 1 and x.stride(dim) != x.stride(dim + 1) * x.size( |
| 39 | dim + 1 |
| 40 | ): |
| 41 | return False |
| 42 | return True |
| 43 | |
| 44 | |
| 45 | def _is_oink_stride_compatible_2d(x_2d: Tensor) -> bool: |
| 46 | """Return True if x_2d meets Oink's pointer-path stride constraints.""" |
| 47 | if x_2d.dim() != 2: |
| 48 | return False |
| 49 | if x_2d.stride(1) != 1: |
| 50 | return False |
| 51 | # Match Oink's vectorization constraint: stride(0) divisible by 256b. |
| 52 | if x_2d.dtype in (torch.float16, torch.bfloat16): |
| 53 | divby = 16 |
| 54 | elif x_2d.dtype == torch.float32: |
| 55 | divby = 8 |
| 56 | else: |
| 57 | return False |
| 58 | return (x_2d.stride(0) % divby) == 0 |
| 59 | |
| 60 | |
| 61 | oink_rms_supported = ( |
| 62 | lambda x, weight, epsilon, variance_size=None: variance_size is None |
| 63 | and weight is not None |
| 64 | and x.dim() >= 2 |
| 65 | and x.dtype == weight.dtype |
| 66 | and weight.is_contiguous() |
| 67 | and _can_view_as_2d(x) |
| 68 | and _is_oink_stride_compatible_2d(x.view(-1, x.shape[-1])) |
| 69 | ) |
| 70 | """ |
| 71 | Oink rms only supports 2d-like inputs with contiguous weight |
| 72 | and no variance_size override. |
| 73 | """ |
| 74 | |
| 75 | |
| 76 | @ir.ops.rms_norm.register_impl( |
| 77 | "oink", supports_args=oink_rms_supported, supported=has_oink_op("rmsnorm") |
| 78 | ) |
| 79 | def rms_norm( |
| 80 | x: Tensor, |
| 81 | weight: Tensor | None, |
| 82 | epsilon: float, |
| 83 | variance_size: int | None = None, |
| 84 | ) -> Tensor: |
| 85 | assert variance_size is None |
| 86 | x_2d = x.view(-1, x.shape[-1]) |
| 87 | return torch.ops.oink.rmsnorm(x_2d, weight, epsilon).view_as(x) |
| 88 | |
| 89 | |
| 90 | oink_add_rms_supported = ( |
| 91 | lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None |
| 92 | and weight is not None |
| 93 | and x.dim() >= 2 |
| 94 | and x.dtype == weight.dtype |
| 95 | and weight.is_contiguous() |
| 96 | and _can_view_as_2d(x) |
| 97 | and _is_oink_stride_compatible_2d(x.view(-1, x.shape[-1])) |
| 98 | # residual must have 2d-compatible strides and match x shape/dtype |
| 99 | and x.dtype == x_residual.dtype |
| 100 | and x.shape == x_residual.shape |
| 101 | and _can_view_as_2d(x_residual) |
| 102 | and _is_oink_stride_compatible_2d(x_residual.view(-1, x_residual.shape[-1])) |
| 103 | ) |
| 104 | """ |
| 105 | Oink fused_add_rms_norm has the same constraints as rms_norm, |
| 106 | and residual must be 2d-like with compatible strides. |
| 107 | """ |
| 108 | |
| 109 | |
| 110 | @ir.ops.fused_add_rms_norm.register_impl( |
| 111 | "oink", |
| 112 | supports_args=oink_add_rms_supported, |
| 113 | supported=has_oink_op("fused_add_rms_norm"), |
| 114 | inplace=True, |
| 115 | ) |
| 116 | def fused_add_rms_norm( |
| 117 | x: Tensor, |
| 118 | x_residual: Tensor, |
| 119 | weight: Tensor | None, |
| 120 | epsilon: float, |
| 121 | variance_size: int | None = None, |
| 122 | ) -> tuple[Tensor, Tensor]: |
| 123 | assert variance_size is None |
| 124 | x_2d = x.view(-1, x.shape[-1]) |
| 125 | residual_2d = x_residual.view(-1, x_residual.shape[-1]) |
| 126 | torch.ops.oink.fused_add_rms_norm(x_2d, residual_2d, weight, epsilon) |
| 127 | return x, x_residual |