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

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 5, 07:15 UTC

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>

Jesus Federico committed Apr 10, 2026 at 13:21 UTC b87575d24b752ee2afffaec00bfb244570d5c95a

6 files changed +44 -2

docs/models/pooling_models/classify.md

+27 -2

@@ -267,9 +267,34 @@ You can modify the `problem_type` via problem_type in the Hugging Face config. T

267

268 Implement alignment with transformers [ForSequenceClassificationLoss](https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92).

269

270 -### Logit bias

270 +### Affine Score Calibration

271

272 -You can modify the `logit_bias` (aka `sigmoid_normalize`) through the logit_bias parameter in `vllm.config.PoolerConfig`.

272 +Affine Score Calibration, also known as [Platt Scaling](https://en.wikipedia.org/wiki/Platt_scaling) (Platt, 1999), is the most widely used method for calibrating classifier outputs into well-calibrated probabilities.

273 +

274 +The calibration follows the transformation:

275 +

276 +`activation(logit_scale * (logit - logit_bias))`

277 +

278 +| Parameter | Default | Description |

279 +| --------- | ------- | ----------- |

280 +| `logit_bias` | `None` | Bias subtracted from logits before activation |

281 +| `logit_scale` | `None` | Scale factor applied to logits after bias subtraction |

282 +

283 +Note: `logit_bias` is **subtracted** from the logits (not added), consistent with the `sigmoid_normalize` convention where `sigmoid(x - bias)` centers the sigmoid around the bias value.

284 +

285 +The computation order is as follows:

286 +

287 +```python

288 +logits -= logit_bias # subtract bias (center scores)

289 +logits *= logit_scale # scale logits

290 +logits = activation(logits) # e.g. sigmoid

291 +```

292 +

293 +Example configuration:

294 +

295 +```bash

296 +--pooler-config '{"use_activation": true, "logit_bias": 4.5, "logit_scale": 1.0}'

297 +```

298

299 ## Removed Features

300

vllm/config/pooler.py

@@ -83,6 +83,13 @@ class PoolerConfig:

83 If provided, apply classification logit biases. Defaults to None.

84 """

86 + logit_scale: float | None = None

87 + """

88 + If provided, scale the classification logits by this factor before

89 + activation. Combined with logit_bias, enables affine score calibration:

90 + activation(logit_scale * (score - logit_bias)). Defaults to None.

91 + """

92 +

93 ## for reward models

94 step_tag_id: int | None = None

95 """

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

@@ -104,6 +104,7 @@ class ClassifierPoolerHead(SequencePoolerHead):

104 self,

105 classifier: ClassifierFn | None = None,

106 logit_bias: float | None = None,

107 + logit_scale: float | None = None,

108 head_dtype: torch.dtype | str | None = None,

109 activation: ActivationFn | None = None,

110 ) -> None:

@@ -111,6 +112,7 @@ class ClassifierPoolerHead(SequencePoolerHead):

112

113 self.classifier = classifier

114 self.logit_bias = logit_bias

115 + self.logit_scale = logit_scale

116 self.head_dtype = head_dtype

117 self.activation = activation

118

@@ -140,6 +142,8 @@ class ClassifierPoolerHead(SequencePoolerHead):

142 # logits shape: [batchsize, num_labels]

143 if self.logit_bias is not None:

144 logits -= self.logit_bias

145 + if self.logit_scale is not None:

146 + logits *= self.logit_scale

147

148 if self.activation is not None:

149 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(

119 head_dtype=model_config.head_dtype,

120 classifier=classifier,

121 logit_bias=model_config.pooler_config.logit_bias,

122 + logit_scale=model_config.pooler_config.logit_scale,

123 activation=resolve_classifier_act_fn(

124 model_config, static_num_labels=True, act_fn=act_fn

125 ),

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

@@ -93,6 +93,7 @@ class TokenClassifierPoolerHead(TokenPoolerHead):

93 self,

94 classifier: ClassifierFn | None = None,

95 logit_bias: float | None = None,

96 + logit_scale: float | None = None,

97 head_dtype: torch.dtype | str | None = None,

98 activation: ActivationFn | None = None,

99 ) -> None:

@@ -100,6 +101,7 @@ class TokenClassifierPoolerHead(TokenPoolerHead):

101

102 self.classifier = classifier

103 self.logit_bias = logit_bias

104 + self.logit_scale = logit_scale

105 self.head_dtype = head_dtype

106 self.activation = activation

107

@@ -127,6 +129,8 @@ class TokenClassifierPoolerHead(TokenPoolerHead):

129

130 if self.logit_bias is not None:

131 logits -= self.logit_bias

132 + if self.logit_scale is not None:

133 + logits *= self.logit_scale

134

135 if self.activation is not None and pooling_param.use_activation:

136 logits = self.activation(logits)

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

@@ -128,6 +128,7 @@ def pooler_for_token_classify(

128 head_dtype=model_config.head_dtype,

129 classifier=classifier,

130 logit_bias=model_config.pooler_config.logit_bias,

131 + logit_scale=model_config.pooler_config.logit_scale,

132 activation=resolve_classifier_act_fn(

133 model_config, static_num_labels=False, act_fn=act_fn

134 ),