| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2024 The vLLM team. |
| 5 | # |
| 6 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 7 | # you may not use this file except in compliance with the License. |
| 8 | # You may obtain a copy of the License at |
| 9 | # |
| 10 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | # |
| 12 | # Unless required by applicable law or agreed to in writing, software |
| 13 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | # See the License for the specific language governing permissions and |
| 16 | # limitations under the License. |
| 17 | """Wrapper around `transformers` models.""" |
| 18 | |
| 19 | from typing import TYPE_CHECKING |
| 20 | |
| 21 | import torch.nn.functional as F |
| 22 | from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| 23 | |
| 24 | from vllm.model_executor.models.transformers.base import VLLM_ATTN_ATTR |
| 25 | from vllm.model_executor.models.transformers.causal import CausalMixin |
| 26 | from vllm.model_executor.models.transformers.legacy import LegacyMixin |
| 27 | from vllm.model_executor.models.transformers.moe import MoEMixin |
| 28 | from vllm.model_executor.models.transformers.multimodal import ( |
| 29 | MultiModalDummyInputsBuilder, |
| 30 | MultiModalMixin, |
| 31 | MultiModalProcessingInfo, |
| 32 | MultiModalProcessor, |
| 33 | ) |
| 34 | from vllm.model_executor.models.transformers.pooling import ( |
| 35 | EmbeddingMixin, |
| 36 | SequenceClassificationMixin, |
| 37 | ) |
| 38 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 39 | |
| 40 | if TYPE_CHECKING: |
| 41 | import torch |
| 42 | |
| 43 | from vllm.model_executor.layers.attention import Attention, MLAAttention |
| 44 | |
| 45 | |
| 46 | def check_sinks( |
| 47 | module: "torch.nn.Module", |
| 48 | self_attn: "Attention | MLAAttention", |
| 49 | s_aux: "torch.Tensor | None", |
| 50 | ): |
| 51 | """Fail loudly if the model applies a sink the attention layer will not. |
| 52 | |
| 53 | Only the attention impl can fold a sink into the softmax denominator, so a sink |
| 54 | that never reached `Attention` is dropped and every softmax is subtly wrong. |
| 55 | """ |
| 56 | if s_aux is None or getattr(self_attn, "has_sink", False): |
| 57 | return |
| 58 | raise ValueError( |
| 59 | f"{type(module).__name__} applies attention sinks, but they were not passed " |
| 60 | f"to {type(self_attn).__name__}, so the output would be wrong. Either the " |
| 61 | "Transformers modeling backend could not find the parameter holding them, or " |
| 62 | "vLLM does not support sinks for this kind of attention. Please open an issue " |
| 63 | "at https://github.com/vllm-project/vllm/issues/new" |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | def vllm_attention_forward( |
| 68 | # Transformers args |
| 69 | module: "torch.nn.Module", |
| 70 | query: "torch.Tensor", |
| 71 | key: "torch.Tensor", |
| 72 | value: "torch.Tensor", |
| 73 | attention_mask: "torch.Tensor", |
| 74 | **kwargs, |
| 75 | ): |
| 76 | self_attn = getattr(module, VLLM_ATTN_ATTR) |
| 77 | check_sinks(module, self_attn, kwargs.get("s_aux")) |
| 78 | hidden = query.shape[-2] |
| 79 | head_dim_qk = query.shape[-1] |
| 80 | head_dim_v = value.shape[-1] |
| 81 | query, key, value = (x.transpose(1, 2) for x in (query, key, value)) |
| 82 | query, key, value = (x.reshape(hidden, -1) for x in (query, key, value)) |
| 83 | # Pad `value` up to the query/key head size when it is smaller (expanded |
| 84 | # MLA). A larger last dim just means `value` isn't split per head, e.g. |
| 85 | # packed grouped/multi-query projections, and needs no padding. |
| 86 | pad_value = head_dim_v < head_dim_qk |
| 87 | if pad_value: |
| 88 | value = F.pad(value.view(-1, head_dim_v), (0, head_dim_qk - head_dim_v)) |
| 89 | value = value.reshape(hidden, -1) |
| 90 | attn_output = self_attn.forward(query, key, value) |
| 91 | if pad_value: |
| 92 | attn_output = attn_output.view(-1, head_dim_qk)[..., :head_dim_v] |
| 93 | attn_output = attn_output.reshape(hidden, -1) |
| 94 | return attn_output, None |
| 95 | |
| 96 | |
| 97 | def vllm_mla_attention_forward( |
| 98 | # Transformers args |
| 99 | module: "torch.nn.Module", |
| 100 | query: "torch.Tensor", |
| 101 | kv_c_normed: "torch.Tensor", |
| 102 | k_pe: "torch.Tensor", |
| 103 | attention_mask: "torch.Tensor", |
| 104 | **kwargs, |
| 105 | ): |
| 106 | self_attn = getattr(module, VLLM_ATTN_ATTR) |
| 107 | check_sinks(module, self_attn, kwargs.get("s_aux")) |
| 108 | # [batch=1, heads, num_tokens, qk_head_dim] -> [num_tokens, heads, qk_head_dim] |
| 109 | query = query.transpose(1, 2).flatten(0, 1) |
| 110 | num_tokens, num_heads = query.shape[:2] |
| 111 | # [batch=1, num_tokens, kv_lora_rank] -> [num_tokens, kv_lora_rank] |
| 112 | kv_c_normed = kv_c_normed.reshape(-1, kv_c_normed.shape[-1]) |
| 113 | # [batch=1, heads=1, num_tokens, qk_rope] -> [num_tokens, 1, qk_rope] |
| 114 | k_pe = k_pe.reshape(-1, 1, k_pe.shape[-1]) |
| 115 | attn_output = self_attn.forward( |
| 116 | query, |
| 117 | kv_c_normed, |
| 118 | k_pe, |
| 119 | output_shape=(num_tokens, num_heads * self_attn.v_head_dim), |
| 120 | ) |
| 121 | return attn_output, None |
| 122 | |
| 123 | |
| 124 | ALL_ATTENTION_FUNCTIONS.register("vllm", vllm_attention_forward) |
| 125 | ALL_ATTENTION_FUNCTIONS.register("vllm_mla", vllm_mla_attention_forward) |
| 126 | |
| 127 | |
| 128 | # Text only models |
| 129 | class TransformersForCausalLM(CausalMixin): ... |
| 130 | |
| 131 | |
| 132 | class TransformersMoEForCausalLM(MoEMixin, CausalMixin): ... |
| 133 | |
| 134 | |
| 135 | # Multimodal models |
| 136 | @MULTIMODAL_REGISTRY.register_processor( |
| 137 | MultiModalProcessor, |
| 138 | info=MultiModalProcessingInfo, |
| 139 | dummy_inputs=MultiModalDummyInputsBuilder, |
| 140 | ) |
| 141 | class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin): ... |
| 142 | |
| 143 | |
| 144 | @MULTIMODAL_REGISTRY.register_processor( |
| 145 | MultiModalProcessor, |
| 146 | info=MultiModalProcessingInfo, |
| 147 | dummy_inputs=MultiModalDummyInputsBuilder, |
| 148 | ) |
| 149 | class TransformersMultiModalMoEForCausalLM(MoEMixin, MultiModalMixin, CausalMixin): ... |
| 150 | |
| 151 | |
| 152 | # Embedding models |
| 153 | class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin): ... |
| 154 | |
| 155 | |
| 156 | class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin): ... |
| 157 | |
| 158 | |
| 159 | @MULTIMODAL_REGISTRY.register_processor( |
| 160 | MultiModalProcessor, |
| 161 | info=MultiModalProcessingInfo, |
| 162 | dummy_inputs=MultiModalDummyInputsBuilder, |
| 163 | ) |
| 164 | class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin): ... |
| 165 | |
| 166 | |
| 167 | # Sequence classification models |
| 168 | class TransformersForSequenceClassification( |
| 169 | SequenceClassificationMixin, LegacyMixin |
| 170 | ): ... |
| 171 | |
| 172 | |
| 173 | class TransformersMoEForSequenceClassification( |
| 174 | SequenceClassificationMixin, MoEMixin |
| 175 | ): ... |
| 176 | |
| 177 | |
| 178 | @MULTIMODAL_REGISTRY.register_processor( |
| 179 | MultiModalProcessor, |
| 180 | info=MultiModalProcessingInfo, |
| 181 | dummy_inputs=MultiModalDummyInputsBuilder, |
| 182 | ) |
| 183 | class TransformersMultiModalForSequenceClassification( |
| 184 | SequenceClassificationMixin, MultiModalMixin |
| 185 | ): ... |
| 186 | |
| 187 | |
| 188 | def __getattr__(name: str): |
| 189 | """Handle imports of non-existent classes with a helpful error message.""" |
| 190 | if name not in globals(): |
| 191 | raise AttributeError( |
| 192 | "The Transformers modeling backend does not currently have a class to " |
| 193 | f"handle the requested model type: {name}. Please open an issue at " |
| 194 | "https://github.com/vllm-project/vllm/issues/new" |
| 195 | ) |
| 196 | return globals()[name] |