feat: add logit_scale to PoolerConfig for affine score calibration (#39435)

Signed-off-by: Jesus Federico <jefp@amazon.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

vllm/model_executor/layers/pooler/seqwise/heads.py

@@ -104,6 +104,7 @@ class ClassifierPoolerHead(SequencePoolerHead):
         self,
         classifier: ClassifierFn | None = None,
         logit_bias: float | None = None,
+        logit_scale: float | None = None,
         head_dtype: torch.dtype | str | None = None,
         activation: ActivationFn | None = None,
     ) -> None:
@@ -111,6 +112,7 @@ class ClassifierPoolerHead(SequencePoolerHead):
 
         self.classifier = classifier
         self.logit_bias = logit_bias
+        self.logit_scale = logit_scale
         self.head_dtype = head_dtype
         self.activation = activation
 
@@ -140,6 +142,8 @@ class ClassifierPoolerHead(SequencePoolerHead):
         # logits shape: [batchsize, num_labels]
         if self.logit_bias is not None:
             logits -= self.logit_bias
+        if self.logit_scale is not None:
+            logits *= self.logit_scale
 
         if self.activation is not None:
             flags = [p.use_activation for p in pooling_params]

vllm/model_executor/layers/pooler/seqwise/poolers.py

@@ -119,6 +119,7 @@ def pooler_for_classify(
         head_dtype=model_config.head_dtype,
         classifier=classifier,
         logit_bias=model_config.pooler_config.logit_bias,
+        logit_scale=model_config.pooler_config.logit_scale,
         activation=resolve_classifier_act_fn(
             model_config, static_num_labels=True, act_fn=act_fn
         ),

vllm/model_executor/layers/pooler/tokwise/heads.py

@@ -93,6 +93,7 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
         self,
         classifier: ClassifierFn | None = None,
         logit_bias: float | None = None,
+        logit_scale: float | None = None,
         head_dtype: torch.dtype | str | None = None,
         activation: ActivationFn | None = None,
     ) -> None:
@@ -100,6 +101,7 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
 
         self.classifier = classifier
         self.logit_bias = logit_bias
+        self.logit_scale = logit_scale
         self.head_dtype = head_dtype
         self.activation = activation
 
@@ -127,6 +129,8 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
 
         if self.logit_bias is not None:
             logits -= self.logit_bias
+        if self.logit_scale is not None:
+            logits *= self.logit_scale
 
         if self.activation is not None and pooling_param.use_activation:
             logits = self.activation(logits)

vllm/model_executor/layers/pooler/tokwise/poolers.py

@@ -128,6 +128,7 @@ def pooler_for_token_classify(
         head_dtype=model_config.head_dtype,
         classifier=classifier,
         logit_bias=model_config.pooler_config.logit_bias,
+        logit_scale=model_config.pooler_config.logit_scale,
         activation=resolve_classifier_act_fn(
             model_config, static_num_labels=False, act_fn=act_fn
         ),