5eeef1b908 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/model_executor/models/bert.py

+2 -2

@@ -28,8 +28,8 @@ from vllm.model_executor.pooling_metadata import PoolingMetadata

28 from vllm.sequence import IntermediateTensors

29 from vllm.tasks import PoolingTask

30

31 -from .interfaces import (SupportsCrossEncoding, SupportsQuant,

32 - default_pooling_type)

31 +from .interfaces import SupportsCrossEncoding, SupportsQuant

32 +from .interfaces_base import default_pooling_type

33 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix

34

35

vllm/model_executor/models/bert_with_rope.py

+3 -2

@@ -27,13 +27,14 @@ from vllm.model_executor.layers.rotary_embedding import get_rope

27 from vllm.model_executor.layers.vocab_parallel_embedding import (

28 VocabParallelEmbedding)

29 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

30 -from vllm.model_executor.models.interfaces import (SupportsQuant,

31 - default_pooling_type)

30 from vllm.model_executor.models.utils import WeightsMapper

31 from vllm.model_executor.utils import set_weight_attrs

32 from vllm.platforms import current_platform

33 from vllm.sequence import IntermediateTensors

34

35 +from .interfaces import SupportsQuant

36 +from .interfaces_base import default_pooling_type

37 +

38

39 class BertWithRopeEmbedding(nn.Module):

40

vllm/model_executor/models/gritlm.py

+1 -1

@@ -20,7 +20,7 @@ from vllm.sequence import PoolerOutput

20 from vllm.tasks import PoolingTask

21 from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config

22

23 -from .interfaces import default_pooling_type

23 +from .interfaces_base import default_pooling_type

24

25 logger = init_logger(__name__)

26

vllm/model_executor/models/interfaces.py

+1 -18

@@ -3,7 +3,7 @@

3

4 from collections.abc import Iterable, Mapping, MutableSequence

5 from typing import (TYPE_CHECKING, ClassVar, Literal, Optional, Protocol,

6 - TypeVar, Union, overload, runtime_checkable)

6 + Union, overload, runtime_checkable)

7

8 import numpy as np

9 import torch

@@ -641,23 +641,6 @@ def supports_cross_encoding(

641 return is_pooling_model(model) and _supports_cross_encoding(model)

642

643

644 -_T = TypeVar("_T", bound=type[torch.nn.Module])

645 -

646 -

647 -def default_pooling_type(pooling_type: str):

648 - """Set default_pooling_type decorator. """

649 -

650 - def func(model: _T) -> _T:

651 - model.default_pooling_type = pooling_type # type: ignore

652 - return model

653 -

654 - return func

655 -

656 -

657 -def get_default_pooling_type(model: Union[type[object], object]) -> str:

658 - return getattr(model, "default_pooling_type", "LAST")

659 -

660 -

644 class SupportsQuant:

645 """The interface required for all models that support quantization."""

646

vllm/model_executor/models/interfaces_base.py

+28

@@ -144,6 +144,17 @@ class VllmModelForPooling(VllmModel[T_co], Protocol[T_co]):

144 MRO of your model class.

145 """

146

147 + default_pooling_type: ClassVar[str] = "LAST"

148 + """

149 + Indicates the

150 + [vllm.model_executor.layers.pooler.PoolerConfig.pooling_type][]

151 + to use by default.

152 +

153 + You can use the

154 + [vllm.model_executor.models.interfaces_base.default_pooling_type][]

155 + decorator to conveniently set this field.

156 + """

157 +

158 pooler: Pooler

159 """The pooler is only called on TP rank 0."""

160

@@ -165,3 +176,20 @@ def is_pooling_model(

176 return False

177

178 return getattr(model, "is_pooling_model", False)

179 +

180 +

181 +_T = TypeVar("_T", bound=type[nn.Module])

182 +

183 +

184 +def default_pooling_type(pooling_type: str):

185 + """Decorator to set `VllmModelForPooling.default_pooling_type`."""

186 +

187 + def func(model: _T) -> _T:

188 + model.default_pooling_type = pooling_type # type: ignore

189 + return model

190 +

191 + return func

192 +

193 +

194 +def get_default_pooling_type(model: Union[type[object], object]) -> str:

195 + return getattr(model, "default_pooling_type", "LAST")

vllm/model_executor/models/internlm2.py

+2 -1

@@ -31,7 +31,8 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader

31 from vllm.model_executor.sampling_metadata import SamplingMetadata

32 from vllm.sequence import IntermediateTensors

33

34 -from .interfaces import SupportsLoRA, SupportsPP, default_pooling_type

34 +from .interfaces import SupportsLoRA, SupportsPP

35 +from .interfaces_base import default_pooling_type

36 from .utils import (is_pp_missing_parameter,

37 make_empty_intermediate_tensors_factory, make_layers,

38 maybe_prefix)

vllm/model_executor/models/modernbert.py

+2 -1

@@ -26,7 +26,8 @@ from vllm.model_executor.pooling_metadata import PoolingMetadata

26 from vllm.sequence import IntermediateTensors

27 from vllm.tasks import PoolingTask

28

29 -from .interfaces import SupportsCrossEncoding, default_pooling_type

29 +from .interfaces import SupportsCrossEncoding

30 +from .interfaces_base import default_pooling_type

31 from .utils import WeightsMapper, maybe_prefix

32

33

vllm/model_executor/models/prithvi_geospatial_mae.py

+4 -3

@@ -27,9 +27,6 @@ from transformers import BatchFeature

27 from vllm.config import VllmConfig

28 from vllm.model_executor.layers.pooler import DispatchPooler, Pooler

29 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

30 -from vllm.model_executor.models.interfaces import (

31 - IsAttentionFree, MultiModalEmbeddings, SupportsMultiModalWithRawInput,

32 - default_pooling_type)

30 from vllm.model_executor.models.utils import AutoWeightsLoader

31 from vllm.multimodal import MULTIMODAL_REGISTRY

32 from vllm.multimodal.inputs import (ImageItem, ModalityData,

@@ -43,6 +40,10 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,

40 from vllm.multimodal.profiling import BaseDummyInputsBuilder

41 from vllm.sequence import IntermediateTensors

42

43 +from .interfaces import (IsAttentionFree, MultiModalEmbeddings,

44 + SupportsMultiModalWithRawInput)

45 +from .interfaces_base import default_pooling_type

46 +

47

48 def _prithvi_field_config(hf_inputs: Mapping[str, torch.Tensor]):

49 # This model receives in input a multi-dimensional tensor representing

vllm/model_executor/models/qwen2_rm.py

+2 -1

@@ -18,7 +18,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

18 from vllm.model_executor.layers.pooler import DispatchPooler, Pooler

19 from vllm.sequence import IntermediateTensors

20

21 -from .interfaces import SupportsLoRA, SupportsPP, default_pooling_type

21 +from .interfaces import SupportsLoRA, SupportsPP

22 +from .interfaces_base import default_pooling_type

23 from .qwen2 import Qwen2Model

24 from .utils import AutoWeightsLoader, maybe_prefix

25

vllm/model_executor/models/registry.py

+4 -3

@@ -25,11 +25,12 @@ from vllm.logger import init_logger

25 from vllm.transformers_utils.dynamic_module import (

26 try_get_class_from_dynamic_module)

27

28 -from .interfaces import (get_default_pooling_type, has_inner_state, has_noops,

29 - is_attention_free, is_hybrid, supports_cross_encoding,

28 +from .interfaces import (has_inner_state, has_noops, is_attention_free,

29 + is_hybrid, supports_cross_encoding,

30 supports_multimodal, supports_multimodal_raw_input,

31 supports_pp, supports_transcription, supports_v0_only)

32 -from .interfaces_base import is_pooling_model, is_text_generation_model

32 +from .interfaces_base import (get_default_pooling_type, is_pooling_model,

33 + is_text_generation_model)

34

35 logger = init_logger(__name__)

36

vllm/model_executor/models/roberta.py

+2 -1

@@ -22,7 +22,8 @@ from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,

22 from vllm.sequence import IntermediateTensors

23

24 from .bert_with_rope import BertWithRope, JinaRobertaModel

25 -from .interfaces import SupportsCrossEncoding, default_pooling_type

25 +from .interfaces import SupportsCrossEncoding

26 +from .interfaces_base import default_pooling_type

27

28

29 class RobertaEmbedding(nn.Module):

[Model] Explicit `default_pooling_type` interface (#23736)