| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang |
| 4 | # |
| 5 | # This file contains code copied from the flash-linear-attention project. |
| 6 | # The original source code was licensed under the MIT license and included |
| 7 | # the following copyright notice: |
| 8 | # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang |
| 9 | |
| 10 | import os |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from vllm.triton_utils import tl, triton |
| 15 | |
| 16 | BT_LIST = [8, 16, 32, 64, 128] |
| 17 | |
| 18 | USE_DEFAULT_FLA_NORM = int(os.getenv("USE_DEFAULT_FLA_NORM", "0")) |
| 19 | |
| 20 | |
| 21 | @triton.autotune( |
| 22 | configs=[ |
| 23 | triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16, 32] |
| 24 | ], |
| 25 | key=["D"], |
| 26 | ) |
| 27 | @triton.jit |
| 28 | def l2norm_fwd_kernel1( |
| 29 | x, |
| 30 | y, |
| 31 | D, |
| 32 | BD: tl.constexpr, |
| 33 | eps, |
| 34 | ): |
| 35 | i_t = tl.program_id(0) |
| 36 | x += i_t * D |
| 37 | y += i_t * D |
| 38 | # Compute mean and variance |
| 39 | cols = tl.arange(0, BD) |
| 40 | mask = cols < D |
| 41 | b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) |
| 42 | b_var = tl.sum(b_x * b_x, axis=0) |
| 43 | b_rstd = 1 / tl.sqrt(b_var + eps) |
| 44 | # tl.store(Rstd + i_t, rstd) |
| 45 | # Normalize and apply linear transformation |
| 46 | b_y = b_x * b_rstd |
| 47 | tl.store(y + cols, b_y, mask=mask) |
| 48 | |
| 49 | |
| 50 | @triton.autotune( |
| 51 | configs=[ |
| 52 | triton.Config({"BT": BT}, num_warps=num_warps) |
| 53 | for num_warps in [1, 2, 4, 8, 16] |
| 54 | for BT in BT_LIST |
| 55 | ], |
| 56 | key=["D"], |
| 57 | ) |
| 58 | @triton.jit(do_not_specialize=["NB"]) |
| 59 | def l2norm_fwd_kernel( |
| 60 | x, |
| 61 | y, |
| 62 | eps, |
| 63 | NB, |
| 64 | T, |
| 65 | D: tl.constexpr, |
| 66 | BT: tl.constexpr, |
| 67 | BD: tl.constexpr, |
| 68 | ): |
| 69 | i_t = tl.program_id(0) |
| 70 | p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) |
| 71 | b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) |
| 72 | b_var = tl.sum(b_x * b_x, axis=1) |
| 73 | b_y = b_x / tl.sqrt(b_var + eps)[:, None] |
| 74 | p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) |
| 75 | tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) |
| 76 | |
| 77 | |
| 78 | @triton.jit |
| 79 | def l2norm_fwd_kernel2( |
| 80 | X, Y, eps, M, N: tl.constexpr, BD: tl.constexpr, MBLOCK: tl.constexpr |
| 81 | ): |
| 82 | xoffset = tl.program_id(0) * MBLOCK |
| 83 | row_idx = xoffset + tl.arange(0, MBLOCK)[:, None] |
| 84 | xmask = row_idx < M |
| 85 | rindex = tl.arange(0, BD)[None, :] |
| 86 | cmask = rindex < N |
| 87 | mask = xmask & cmask |
| 88 | xs = tl.load(X + (rindex + N * row_idx), mask, other=0.0).to(tl.float32) |
| 89 | square = tl.broadcast_to(xs * xs, [MBLOCK, BD]) |
| 90 | square_sum = tl.sum(tl.where(xmask, square, 0), 1)[:, None] |
| 91 | rsqrt = tl.rsqrt(square_sum + eps) |
| 92 | tl.store(Y + (rindex + N * row_idx), xs * rsqrt, mask) |
| 93 | |
| 94 | |
| 95 | def l2norm_fwd( |
| 96 | x: torch.Tensor, eps: float = 1e-6, output_dtype: torch.dtype | None = None |
| 97 | ): |
| 98 | x_shape_og = x.shape |
| 99 | x = x.view(-1, x.shape[-1]) |
| 100 | # allocate output |
| 101 | if output_dtype is None: |
| 102 | y = torch.empty_like(x) |
| 103 | else: |
| 104 | y = torch.empty_like(x, dtype=output_dtype) |
| 105 | assert y.stride(-1) == 1 |
| 106 | T, D = x.shape[0], x.shape[-1] |
| 107 | # rstd = torch.empty((T,), dtype=torch.float32, device=x.device) |
| 108 | # Less than 64KB per feature: enqueue fused kernel |
| 109 | MAX_FUSED_SIZE = 65536 // x.element_size() |
| 110 | BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) |
| 111 | if D > BD: |
| 112 | raise RuntimeError("This layer doesn't support feature dim >= 64KB.") |
| 113 | |
| 114 | if not USE_DEFAULT_FLA_NORM: |
| 115 | MBLOCK = 32 |
| 116 | # M, N = x.shape |
| 117 | l2norm_fwd_kernel2[(triton.cdiv(T, MBLOCK),)]( |
| 118 | x, |
| 119 | y, |
| 120 | eps, |
| 121 | T, |
| 122 | D, |
| 123 | BD, |
| 124 | MBLOCK, |
| 125 | ) |
| 126 | else: |
| 127 | if D <= 512: |
| 128 | NB = triton.cdiv(T, 2048) |
| 129 | |
| 130 | def grid(meta): |
| 131 | return (triton.cdiv(T, meta["BT"]),) |
| 132 | |
| 133 | l2norm_fwd_kernel[grid]( |
| 134 | x, |
| 135 | y, |
| 136 | eps, |
| 137 | NB=NB, |
| 138 | T=T, |
| 139 | D=D, |
| 140 | BD=BD, |
| 141 | ) |
| 142 | else: |
| 143 | l2norm_fwd_kernel1[(T,)]( |
| 144 | x, |
| 145 | y, |
| 146 | eps=eps, |
| 147 | D=D, |
| 148 | BD=BD, |
| 149 | ) |
| 150 | |
| 151 | return y.view(x_shape_og) |