@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9101dc756c

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 07:45 UTC

[Model] Avoid hardcoding pooling type (#32119)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jan 12, 2026 at 13:28 UTC 9101dc756c14b6467c192e7bfc67036f40d76320

6 files changed +47 -22

vllm/model_executor/models/bert.py

+10 -4

@@ -25,11 +25,11 @@ from vllm.model_executor.layers.pooler import (

25 PoolingParamsUpdate,

26 )

27 from vllm.model_executor.layers.pooler.seqwise import (

28 - CLSPool,

28 SequencePooler,

29 SequencePoolerHeadOutput,

30 SequencePoolerOutput,

31 SequencePoolingMethodOutput,

32 + get_seq_pooling_method,

33 )

34 from vllm.model_executor.layers.pooler.tokwise import (

35 pooler_for_token_classify,

@@ -94,9 +94,9 @@ class BertEmbedding(nn.Module):

96 class BertPooler(SequencePooler):

97 - def __init__(self, config: BertConfig):

97 + def __init__(self, config: BertConfig, pooler_config: PoolerConfig):

98 super().__init__(

99 - pooling=CLSPool(),

99 + pooling=get_seq_pooling_method(pooler_config.seq_pooling_type),

100 head=self.head,

101 )

102

@@ -450,7 +450,11 @@ class BertPoolingModel(BertModel):

450 )

451

452 config = vllm_config.model_config.hf_config

453 - self.pooler = BertPooler(config)

453 +

454 + pooler_config = vllm_config.model_config.pooler_config

455 + assert pooler_config is not None

456 +

457 + self.pooler = BertPooler(config, pooler_config)

458

459 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

460 other_weights, loaded_stacked_params = self._load_weights(weights)

@@ -711,6 +715,8 @@ class BertSpladeSparseEmbeddingModel(BertEmbeddingModel):

715 layer_norm_eps=getattr(cfg, "layer_norm_eps", 1e-12),

716 )

717

718 + # None of vLLM's built-in sequence pooling types are

719 + # applicable so it is overwritten by SPLADESparsePooler

720 pooling_mode = getattr(self, "_splade_pooling", "max")

721

722 cls_id = getattr(cfg, "cls_token_id", None)

vllm/model_executor/models/bert_with_rope.py

+9 -1

@@ -453,6 +453,7 @@ class BertWithRope(nn.Module, SupportsQuant):

453 add_pooling_layer: bool = False,

454 ):

455 super().__init__()

456 +

457 self.vllm_config = vllm_config

458 self.add_pooling_layer = add_pooling_layer

459 self.config = vllm_config.model_config.hf_config

@@ -463,7 +464,14 @@ class BertWithRope(nn.Module, SupportsQuant):

464 rotary_kwargs=self.config.rotary_kwargs,

465 prefix=f"{prefix}.encoder",

466 )

466 - self.pooler = BertPooler(self.config) if add_pooling_layer else None

467 +

468 + if add_pooling_layer:

469 + pooler_config = vllm_config.model_config.pooler_config

470 + assert pooler_config is not None

471 +

472 + self.pooler = BertPooler(self.config, pooler_config)

473 + else:

474 + self.pooler = None

475

476 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

477 return self.embeddings(input_ids)

vllm/model_executor/models/gritlm.py

+9 -4

@@ -5,7 +5,7 @@ from collections.abc import Set

5 import numpy as np

6 import torch

8 -from vllm.config import ModelConfig, VllmConfig

8 +from vllm.config import ModelConfig, PoolerConfig, VllmConfig

9 from vllm.logger import init_logger

10 from vllm.model_executor.layers.pooler import (

11 DispatchPooler,

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.pooler.seqwise import (

17 SequencePoolerHeadOutput,

18 SequencePoolingMethod,

19 SequencePoolingMethodOutput,

20 + get_seq_pooling_method,

21 )

22 from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed

23 from vllm.model_executor.models.llama import LlamaForCausalLM

@@ -177,9 +178,13 @@ class GritLMMeanPool(SequencePoolingMethod):

178

179

180 class GritLMPooler(SequencePooler):

180 - def __init__(self, model_config: ModelConfig):

181 + def __init__(self, model_config: ModelConfig, pooler_config: PoolerConfig):

182 super().__init__(

182 - pooling=GritLMMeanPool(model_config),

183 + pooling=(

184 + GritLMMeanPool(model_config)

185 + if pooler_config.seq_pooling_type == "MEAN"

186 + else get_seq_pooling_method(pooler_config.seq_pooling_type)

187 + ),

188 head=self.head,

189 )

190

@@ -235,6 +240,6 @@ class GritLM(LlamaForCausalLM):

240 self.pooler = DispatchPooler(

241 {

242 "token_embed": pooler_for_token_embed(pooler_config),

238 - "embed": GritLMPooler(vllm_config.model_config),

243 + "embed": GritLMPooler(vllm_config.model_config, pooler_config),

244 }

245 )

vllm/model_executor/models/modernbert.py

+12 -4

@@ -8,7 +8,7 @@ from transformers import ModernBertConfig

8 from transformers.activations import ACT2FN

10 from vllm.compilation.decorators import support_torch_compile

11 -from vllm.config import VllmConfig

11 +from vllm.config import PoolerConfig, VllmConfig

12 from vllm.distributed import get_tensor_model_parallel_world_size

13 from vllm.model_executor.layers.attention.encoder_only_attention import (

14 EncoderOnlyAttention,

@@ -282,9 +282,14 @@ class ModernBertModel(nn.Module):

282

283

284 class ModernBertPooler(SequencePooler):

285 - def __init__(self, config: ModernBertConfig):

285 + def __init__(self, config: ModernBertConfig, pooler_config: PoolerConfig):

286 + hf_pooling_type = config.classifier_pooling.upper()

287 + # vllm_pooling_type = pooler_config.seq_pooling_type

288 + # Currently we don't have a way to see if the user set the pooling type

289 + # explicitly or not, so we always use the HF pooling type for now.

290 +

291 super().__init__(

287 - pooling=get_seq_pooling_method(config.classifier_pooling.upper()),

292 + pooling=get_seq_pooling_method(hf_pooling_type),

293 head=self.head,

294 )

295

@@ -314,7 +319,9 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):

319

320 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

321 super().__init__()

322 +

323 config = vllm_config.model_config.hf_config

324 +

325 self.config = config

326 self.model = ModernBertModel(

327 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "modernbert")

@@ -324,11 +331,12 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):

331 config.num_labels,

332 dtype=vllm_config.model_config.head_dtype,

333 )

327 - self.pooling = ModernBertPooler(config)

334

335 pooler_config = vllm_config.model_config.pooler_config

336 assert pooler_config is not None

337

338 + self.pooling = ModernBertPooler(config, pooler_config)

339 +

340 self.pooler = DispatchPooler.for_seq_cls(

341 pooler_config,

342 pooling=self.pooling,

vllm/model_executor/models/roberta.py

+1 -3

@@ -9,7 +9,6 @@ from transformers import RobertaConfig

10 from vllm.config import ModelConfig, VllmConfig

11 from vllm.model_executor.layers.pooler import DispatchPooler

12 -from vllm.model_executor.layers.pooler.seqwise import CLSPool

12 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

13 from vllm.model_executor.models.bert import (

14 TOKEN_TYPE_SHIFT,

@@ -86,7 +85,7 @@ class RobertaClassificationHead(nn.Module):

85 )

87 def forward(self, x: torch.Tensor) -> torch.Tensor:

89 - # CLSPool has already been applied in `pooling`

88 + # Token extraction has already been applied in `pooler.pooling`

89 x = self.dense(x)

90 x = torch.tanh(x)

91 x = self.out_proj(x)

@@ -194,7 +193,6 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):

193

194 self.pooler = DispatchPooler.for_seq_cls(

195 pooler_config,

197 - pooling=CLSPool(),

196 classifier=self.classifier,

197 )

198

vllm/model_executor/models/transformers/pooling.py

+6 -6

@@ -23,7 +23,6 @@ from transformers import AutoModelForSequenceClassification

24 from vllm.config.utils import getattr_iter

25 from vllm.model_executor.layers.pooler import DispatchPooler

26 -from vllm.model_executor.layers.pooler.seqwise import CLSPool

26 from vllm.model_executor.models.interfaces import SupportsCrossEncoding

27 from vllm.model_executor.models.interfaces_base import VllmModelForPooling

@@ -32,7 +31,7 @@ if TYPE_CHECKING:

33 class EmbeddingMixin(VllmModelForPooling):

35 - default_pooling_type = "CLS"

34 + default_seq_pooling_type = "CLS"

36 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

37 # Skip VllmModelForPooling.__init__ and call the next class in MRO

@@ -47,7 +46,7 @@ class EmbeddingMixin(VllmModelForPooling):

48 class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):

50 - default_pooling_type = "CLS"

49 + default_seq_pooling_type = "CLS"

51 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

52 # Skip VllmModelForPooling.__init__ and call the next class in MRO

@@ -85,8 +84,10 @@ class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):

84 self.init_parameters(self.classifier, dtype=self.model_config.head_dtype)

86 class ClassifierWithReshape(self.classifier.__class__):

88 - """CLSPool has already been applied in `pooling`.

89 - Add dim to match expected input shape of `classifier.forward`."""

87 + """

88 + Token extraction has already been applied in `pooler.pooling`.

89 + Add dim to match expected input shape of `classifier.forward`.

90 + """

92 def forward(self, *args, **kwargs):

93 if len(args) > 0:

@@ -97,6 +98,5 @@ class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):

99 self.pooler = DispatchPooler.for_seq_cls(

100 pooler_config,

100 - pooling=CLSPool(),

101 classifier=self.classifier,

102 )