| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Copyright (c) 2024, Tri Dao. |
| 4 | # Adapted from https://github.com/state-spaces/mamba/blob/60dadf2e0ee730ac337035d5533de10bc26e4847/mamba_ssm/ops/triton/layernorm_gated.py |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.triton_utils import tl, triton |
| 9 | |
| 10 | |
| 11 | @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) |
| 12 | @triton.heuristics({"HAS_Z": lambda args: args["Z"] is not None}) |
| 13 | @triton.jit |
| 14 | def _layer_norm_fwd_1pass_kernel( |
| 15 | X, # pointer to the input |
| 16 | Y, # pointer to the output |
| 17 | W, # pointer to the weights |
| 18 | B, # pointer to the biases |
| 19 | Z, # pointer to the other branch |
| 20 | Mean, # pointer to the mean |
| 21 | Rstd, # pointer to the 1/std |
| 22 | stride_x_row: tl.int64, |
| 23 | stride_y_row: tl.int64, |
| 24 | stride_z_row: tl.int64, |
| 25 | M: tl.int64, # number of rows in X |
| 26 | N: tl.int64, # number of columns in X |
| 27 | eps, # epsilon to avoid division by zero |
| 28 | BLOCK_N: tl.constexpr, |
| 29 | HAS_BIAS: tl.constexpr, |
| 30 | HAS_Z: tl.constexpr, |
| 31 | NORM_BEFORE_GATE: tl.constexpr, |
| 32 | IS_RMS_NORM: tl.constexpr, |
| 33 | ): |
| 34 | # Map the program id to the row of X and Y it should compute. |
| 35 | row = tl.program_id(0) |
| 36 | group = tl.program_id(1) |
| 37 | X += row * stride_x_row + group * N |
| 38 | Y += row * stride_y_row + group * N |
| 39 | if HAS_Z: |
| 40 | Z += row * stride_z_row + group * N |
| 41 | if not IS_RMS_NORM: |
| 42 | Mean += group * M |
| 43 | Rstd += group * M |
| 44 | W += group * N |
| 45 | if HAS_BIAS: |
| 46 | B += group * N |
| 47 | # Compute mean and variance |
| 48 | cols = tl.arange(0, BLOCK_N) |
| 49 | x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) |
| 50 | if HAS_Z and not NORM_BEFORE_GATE: |
| 51 | z = tl.load(Z + cols, mask=cols < N).to(tl.float32) |
| 52 | x *= z * tl.sigmoid(z) |
| 53 | if not IS_RMS_NORM: |
| 54 | mean = tl.sum(x, axis=0) / N |
| 55 | tl.store(Mean + row, mean) |
| 56 | xbar = tl.where(cols < N, x - mean, 0.0) |
| 57 | var = tl.sum(xbar * xbar, axis=0) / N |
| 58 | else: |
| 59 | xbar = tl.where(cols < N, x, 0.0) |
| 60 | var = tl.sum(xbar * xbar, axis=0) / N |
| 61 | rstd = 1 / tl.sqrt(var + eps) |
| 62 | tl.store(Rstd + row, rstd) |
| 63 | # Normalize and apply linear transformation |
| 64 | mask = cols < N |
| 65 | w = tl.load(W + cols, mask=mask).to(tl.float32) |
| 66 | if HAS_BIAS: |
| 67 | b = tl.load(B + cols, mask=mask).to(tl.float32) |
| 68 | x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd |
| 69 | y = x_hat * w + b if HAS_BIAS else x_hat * w |
| 70 | if HAS_Z and NORM_BEFORE_GATE: |
| 71 | z = tl.load(Z + cols, mask=mask).to(tl.float32) |
| 72 | y *= z * tl.sigmoid(z) |
| 73 | # Write output |
| 74 | tl.store(Y + cols, y, mask=mask) |
| 75 | |
| 76 | |
| 77 | def _layer_norm_fwd( |
| 78 | x, |
| 79 | weight, |
| 80 | bias, |
| 81 | eps, |
| 82 | z=None, |
| 83 | out=None, |
| 84 | group_size=None, |
| 85 | norm_before_gate=True, |
| 86 | is_rms_norm=False, |
| 87 | ): |
| 88 | M, N = x.shape |
| 89 | if group_size is None: |
| 90 | group_size = N |
| 91 | assert N % group_size == 0 |
| 92 | ngroups = N // group_size |
| 93 | assert x.stride(-1) == 1 |
| 94 | if z is not None: |
| 95 | assert z.stride(-1) == 1 |
| 96 | assert z.shape == (M, N) |
| 97 | assert weight.shape == (N,) |
| 98 | assert weight.stride(-1) == 1 |
| 99 | if bias is not None: |
| 100 | assert bias.stride(-1) == 1 |
| 101 | assert bias.shape == (N,) |
| 102 | # allocate output |
| 103 | if out is not None: |
| 104 | assert out.shape == x.shape |
| 105 | else: |
| 106 | out = torch.empty_like(x) |
| 107 | assert out.stride(-1) == 1 |
| 108 | mean = ( |
| 109 | torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) |
| 110 | if not is_rms_norm |
| 111 | else None |
| 112 | ) |
| 113 | rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) |
| 114 | # Less than 64KB per feature: enqueue fused kernel |
| 115 | MAX_FUSED_SIZE = 65536 // x.element_size() |
| 116 | BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) |
| 117 | if group_size > BLOCK_N: |
| 118 | raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") |
| 119 | # heuristics for number of warps |
| 120 | num_warps = min(max(BLOCK_N // 256, 1), 8) |
| 121 | grid = (M, ngroups) |
| 122 | with torch.accelerator.device_index(x.device.index): |
| 123 | _layer_norm_fwd_1pass_kernel[grid]( |
| 124 | x, |
| 125 | out, |
| 126 | weight, |
| 127 | bias, |
| 128 | z, |
| 129 | mean, |
| 130 | rstd, |
| 131 | x.stride(0), |
| 132 | out.stride(0), |
| 133 | z.stride(0) if z is not None else 0, |
| 134 | M, |
| 135 | group_size, |
| 136 | eps, |
| 137 | BLOCK_N=BLOCK_N, |
| 138 | NORM_BEFORE_GATE=norm_before_gate, |
| 139 | IS_RMS_NORM=is_rms_norm, |
| 140 | num_warps=num_warps, |
| 141 | ) |
| 142 | return out, mean, rstd |
| 143 | |
| 144 | |
| 145 | def rms_norm_gated( |
| 146 | x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True |
| 147 | ): |
| 148 | x_shape_og = x.shape |
| 149 | # reshape input data into 2D tensor |
| 150 | x = x.reshape(-1, x.shape[-1]) |
| 151 | if x.stride(-1) != 1: |
| 152 | x = x.contiguous() |
| 153 | if z is not None: |
| 154 | assert z.shape == x_shape_og |
| 155 | z = z.reshape(-1, z.shape[-1]) |
| 156 | if z.stride(-1) != 1: |
| 157 | z = z.contiguous() |
| 158 | weight = weight.contiguous() |
| 159 | if bias is not None: |
| 160 | bias = bias.contiguous() |
| 161 | y, _, _ = _layer_norm_fwd( |
| 162 | x, |
| 163 | weight, |
| 164 | bias, |
| 165 | eps, |
| 166 | z=z, |
| 167 | group_size=group_size, |
| 168 | norm_before_gate=norm_before_gate, |
| 169 | is_rms_norm=True, |
| 170 | ) |
| 171 | |
| 172 | return y.reshape(x_shape_og) |