utils.py 8.87 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Utility methods for model layers."""
4

5
from collections.abc import Callable
6
7
8

import torch

9
10
from vllm import _custom_ops as ops
from vllm import envs
11
from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.logger import init_logger
13
from vllm.platforms import CpuArchEnum, current_platform
14
from vllm.utils.platform_utils import get_cu_count
15
from vllm.utils.torch_utils import direct_register_custom_op
16

17
18
logger = init_logger(__name__)

19
20
21
22
23
24
25
26
27
28
29
30
31
32
MOE_LAYER_ROUTER_GATE_SUFFIXES = {
    "gate",
    "router",
    "router_gate",
    "shared_expert_gate",
    "expert_gate",
}


def is_layer_moe_router_gate(prefix: str) -> bool:
    if not prefix:
        return False
    return prefix.rsplit(".", 1)[-1] in MOE_LAYER_ROUTER_GATE_SUFFIXES

33

34
35
36
37
38
39
40
41
42
43
44
45
46
def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
    # Shuffle weight along the last dimension so that
    # we folded the weights to adjance location
    # Example:
    # input:
    #       [[1, 2, 3, 4, 5, 6],
    #        [7, 8, 9, 10, 11, 12]]
    # output:
    #       [[1, 4, 2, 5, 3, 6],
    #        [7, 10, 8, 11, 9, 12]]
    # This will be used together with triton swiglu kernel
    shape = w.shape
    N = shape[-1]
47
48
    first = w[..., : N // 2]
    second = w[..., N // 2 :]
49
50
51
52
53
54

    stacked = torch.stack((first, second), dim=-1)
    w_shuffled = stacked.reshape(shape)
    return w_shuffled


55
56
57
58
def get_token_bin_counts_and_mask(
    tokens: torch.Tensor,
    vocab_size: int,
    num_seqs: int,
59
) -> tuple[torch.Tensor, torch.Tensor]:
60
61
    # Compute the bin counts for the tokens.
    # vocab_size + 1 for padding.
62
63
64
    bin_counts = torch.zeros(
        (num_seqs, vocab_size + 1), dtype=torch.long, device=tokens.device
    )
65
66
67
68
69
70
71
    bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens))
    bin_counts = bin_counts[:, :vocab_size]
    mask = bin_counts > 0

    return bin_counts, mask


72
73
74
75
76
77
78
79
def apply_penalties(
    logits: torch.Tensor,
    prompt_tokens_tensor: torch.Tensor,
    output_tokens_tensor: torch.Tensor,
    presence_penalties: torch.Tensor,
    frequency_penalties: torch.Tensor,
    repetition_penalties: torch.Tensor,
) -> torch.Tensor:
80
81
82
    """
    Applies penalties in place to the logits tensor
    logits : The input logits tensor of shape [num_seqs, vocab_size]
83
84
85
86
    prompt_tokens_tensor: A tensor containing the prompt tokens. The prompts
        are padded to the maximum prompt length within the batch using
        `vocab_size` as the padding value. The value `vocab_size` is used
        for padding because it does not correspond to any valid token ID
87
88
89
90
91
92
93
        in the vocabulary.
    output_tokens_tensor: The output tokens tensor.
    presence_penalties: The presence penalties of shape (num_seqs, )
    frequency_penalties: The frequency penalties of shape (num_seqs, )
    repetition_penalties: The repetition penalties of shape (num_seqs, )
    """
    num_seqs, vocab_size = logits.shape
94
95
96
    _, prompt_mask = get_token_bin_counts_and_mask(
        prompt_tokens_tensor, vocab_size, num_seqs
    )
97
    output_bin_counts, output_mask = get_token_bin_counts_and_mask(
98
99
        output_tokens_tensor, vocab_size, num_seqs
    )
100

101
102
    # Apply repetition penalties as a custom op
    from vllm._custom_ops import apply_repetition_penalties
103
104

    apply_repetition_penalties(logits, prompt_mask, output_mask, repetition_penalties)
105

106
107
    # We follow the definition in OpenAI API.
    # Refer to https://platform.openai.com/docs/api-reference/parameter-details
108
109
    logits -= frequency_penalties.unsqueeze(dim=1) * output_bin_counts
    logits -= presence_penalties.unsqueeze(dim=1) * output_mask
110
    return logits
111
112


113
114
115
116
def default_unquantized_gemm(
    layer: torch.nn.Module,
    x: torch.Tensor,
    weight: torch.Tensor,
117
    bias: torch.Tensor | None = None,
118
):
119
120
121
    return torch.nn.functional.linear(x, weight, bias)


122
123
def use_aiter_triton_gemm(n, m, k, dtype):
    if (
124
        not rocm_aiter_ops.is_triton_gemm_enabled()
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
        # MI300's - fp8nuz=True
        or current_platform.is_fp8_fnuz()
        or dtype not in [torch.float16, torch.bfloat16]
    ):
        return False

    # use hipblaslt for the larger GEMMs
    if n > 2048 and m > 512:
        return False
    return (
        (m == 5120 and k == 2880)
        or (m == 2880 and k == 4096)
        or (m == 128 and k == 2880)
        or (m == 640 and k == 2880)
        or (m == 2880 and k == 512)
    )


143
def rocm_unquantized_gemm_impl(
144
    x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
145
) -> torch.Tensor:
146
    from vllm.platforms.rocm import on_gfx9, on_gfx950
147

148
149
    n = x.numel() / x.size(-1)
    m = weight.shape[0]
150
    k = weight.shape[1]
151

152
153
    import math

154
155
156
157
158
    if use_aiter_triton_gemm(n, m, k, x.dtype):
        from aiter.ops.triton.gemm_a16w16 import gemm_a16w16

        return gemm_a16w16(x, weight, bias)

159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
    use_skinny_reduce_counting = (
        envs.VLLM_ROCM_USE_SKINNY_GEMM
        and on_gfx950()
        and x.dtype in [torch.float16, torch.bfloat16]
        and (
            n >= 16
            and n <= 128
            and k > 512
            and math.ceil(k / 512) * math.ceil(m / 16) < get_cu_count()
        )
        # k == 2880 and (m == 640 or m == 128))
    )
    if use_skinny_reduce_counting:
        cu_count = get_cu_count()
        x_view = x.reshape(-1, x.size(-1))
        out = ops.wvSplitKrc(weight, x_view, cu_count, bias)
        return out.reshape(*x.shape[:-1], weight.shape[0])

177
178
179
180
181
182
    use_skinny = (
        envs.VLLM_ROCM_USE_SKINNY_GEMM
        and on_gfx9()
        and x.dtype in [torch.float16, torch.bfloat16]
        and k % 8 == 0
    )
183
184
185
186

    if use_skinny is not True:
        return torch.nn.functional.linear(x, weight, bias)

187
    x_view = x.reshape(-1, x.size(-1))
188
    if m > 8 and 0 < n <= 4:
189
        cu_count = get_cu_count()
190
        out = ops.wvSplitK(weight, x_view, cu_count, bias)
191
        return out.reshape(*x.shape[:-1], weight.shape[0])
192
    elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
193
        out = ops.LLMM1(weight, x_view, 4)
194
        return out.reshape(*x.shape[:-1], weight.shape[0])
195
196
197
    return torch.nn.functional.linear(x, weight, bias)


198
def rocm_unquantized_gemm_fake(
199
    x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
200
) -> torch.Tensor:
201
202
203
    return x.new_empty((*x.shape[:-1], weight.shape[0]))


204
205
206
207
def rocm_unquantized_gemm(
    layer: torch.nn.Module,
    x: torch.Tensor,
    weight: torch.Tensor,
208
    bias: torch.Tensor | None = None,
209
) -> torch.Tensor:
210
    return torch.ops.vllm.rocm_unquantized_gemm(x, weight, bias)
211
212
213


direct_register_custom_op(
214
    op_name="rocm_unquantized_gemm",
215
    op_func=rocm_unquantized_gemm_impl,
216
    fake_impl=rocm_unquantized_gemm_fake,
217
218
219
)


220
def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype) -> bool:
221
222
223
224
225
226
    return (
        torch._C._cpu._is_amx_tile_supported()
        and (dtype in (torch.bfloat16, torch.int8))
        and k % 32 == 0
        and n % 16 == 0
    )
227
228


229
230
231
232
def dispatch_cpu_unquantized_gemm(
    layer: torch.nn.Module,
    remove_weight: bool,
) -> None:
233
234
235
236
237
    # skip for missing layers
    if layer.weight.is_meta:
        layer.cpu_linear = torch.nn.functional.linear
        return

238
239
    N, K = layer.weight.size()
    dtype = layer.weight.dtype
240

241
242
243
244
245
246
    if envs.VLLM_CPU_SGL_KERNEL and check_cpu_sgl_kernel(N, K, dtype):
        packed_weight = torch.ops._C.convert_weight_packed(layer.weight)
        if getattr(layer, "bias", None) is not None:
            bias_f32 = layer.bias.to(torch.float32)
        else:
            bias_f32 = None
247
248
249
        layer.cpu_linear = lambda x, weight, bias: torch.ops._C.weight_packed_linear(
            x, packed_weight, bias_f32 if bias is not None else None, True
        )
250
        if remove_weight:
251
            layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
252
        return
253
254
255
256
    elif (
        ops._supports_onednn
        and current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC
    ):
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
        try:
            origin_weight = layer.weight
            handler = ops.create_onednn_mm(origin_weight.t(), 32)
            layer.cpu_linear = lambda x, weight, bias: ops.onednn_mm(handler, x, bias)
            if remove_weight:
                layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
            return
        except RuntimeError as e:
            logger.warning_once(
                "Failed to create oneDNN linear, fallback to torch linear."
                f" Exception: {e}"
            )

    # fallback case
    layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear(
        x, weight, bias
    )
274
275


276
277
278
279
def cpu_unquantized_gemm(
    layer: torch.nn.Module,
    x: torch.Tensor,
    weight: torch.Tensor,
280
    bias: torch.Tensor | None = None,
281
):
282
    return layer.cpu_linear(x, weight, bias)
283
284


285
286
287
def dispatch_unquantized_gemm() -> Callable[..., torch.Tensor]:
    if current_platform.is_rocm():
        return rocm_unquantized_gemm
288
289
290
291
    elif current_platform.is_cpu():
        return cpu_unquantized_gemm
    else:
        return default_unquantized_gemm