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 1, 22:30 UTC
[Misc] Add unit tests for pooler head classes (#44471)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jun 4, 2026 at 13:59 UTC
8d9536a775a87eae5e0a1904603aa4725562765e
1 file changed
+481
tests/model_executor/layers/test_pooler_heads.py
new
+481
@@ -0,0 +1,481 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for sequence and token pooler head classes."""
4
+
5
+import torch
6
+import torch.nn as nn
7
+
8
+from vllm.model_executor.layers.pooler.activations import PoolerNormalize
9
+from vllm.model_executor.layers.pooler.seqwise.heads import (
10
+ ClassifierPoolerHead,
11
+ EmbeddingPoolerHead,
12
+)
13
+from vllm.model_executor.layers.pooler.tokwise.heads import (
14
+ TokenClassifierPoolerHead,
15
+ TokenEmbeddingPoolerHead,
16
+)
17
+from vllm.pooling_params import PoolingParams
18
+from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
19
+
20
+_HIDDEN = 16
21
+_BATCH = 3
22
+
23
+
24
+def _make_params(
25
+ n: int,
26
+ *,
27
+ task: str = "embed",
28
+ dimensions: int | None = None,
29
+ use_activation: bool | None = None,
30
+) -> list[PoolingParams]:
31
+ return [
32
+ PoolingParams(task=task, dimensions=dimensions, use_activation=use_activation)
33
+ for _ in range(n)
34
+ ]
35
+
36
+
37
+def _make_metadata(pooling_params: list[PoolingParams]) -> PoolingMetadata:
38
+ n = len(pooling_params)
39
+ return PoolingMetadata(
40
+ prompt_lens=torch.ones(n, dtype=torch.long),
41
+ prompt_token_ids=None,
42
+ prompt_token_ids_cpu=None,
43
+ pooling_params=pooling_params,
44
+ pooling_states=[PoolingStates() for _ in range(n)],
45
+ )
46
+
47
+
48
+def _linear(in_f: int, out_f: int) -> nn.Linear:
49
+ torch.manual_seed(42)
50
+ return nn.Linear(in_f, out_f, bias=False)
51
+
52
+
53
+# ---------------------------------------------------------------------------
54
+# EmbeddingPoolerHead
55
+# ---------------------------------------------------------------------------
56
+class TestEmbeddingPoolerHead:
57
+ def test_supported_tasks(self):
58
+ head = EmbeddingPoolerHead()
59
+ assert head.get_supported_tasks() == {"embed"}
60
+
61
+ def test_passthrough(self):
62
+ head = EmbeddingPoolerHead()
63
+ x = torch.randn(_BATCH, _HIDDEN)
64
+ meta = _make_metadata(_make_params(_BATCH))
65
+ out = head(x, meta)
66
+ assert torch.equal(out, x)
67
+
68
+ def test_head_dtype(self):
69
+ head = EmbeddingPoolerHead(head_dtype=torch.float16)
70
+ x = torch.randn(_BATCH, _HIDDEN)
71
+ meta = _make_metadata(_make_params(_BATCH))
72
+ out = head(x, meta)
73
+ assert out.dtype == torch.float16
74
+
75
+ def test_projector(self):
76
+ proj = _linear(_HIDDEN, 8)
77
+ head = EmbeddingPoolerHead(projector=proj)
78
+ x = torch.randn(_BATCH, _HIDDEN)
79
+ meta = _make_metadata(_make_params(_BATCH))
80
+ out = head(x, meta)
81
+ assert out.shape == (_BATCH, 8)
82
+ assert torch.allclose(out, proj(x))
83
+
84
+ def test_matryoshka_uniform(self):
85
+ head = EmbeddingPoolerHead()
86
+ x = torch.randn(_BATCH, _HIDDEN)
87
+ params = _make_params(_BATCH, dimensions=4)
88
+ meta = _make_metadata(params)
89
+ out = head(x, meta)
90
+ assert out.shape == (_BATCH, 4)
91
+ assert torch.equal(out, x[..., :4])
92
+
93
+ def test_matryoshka_mixed(self):
94
+ head = EmbeddingPoolerHead()
95
+ x = torch.randn(2, _HIDDEN)
96
+ params = [
97
+ PoolingParams(task="embed", dimensions=4),
98
+ PoolingParams(task="embed", dimensions=8),
99
+ ]
100
+ meta = _make_metadata(params)
101
+ out = head(x, meta)
102
+ assert isinstance(out, list)
103
+ assert len(out) == 2
104
+ assert out[0].shape[-1] == 4
105
+ assert out[1].shape[-1] == 8
106
+
107
+ def test_matryoshka_mixed_with_none(self):
108
+ head = EmbeddingPoolerHead()
109
+ x = torch.randn(2, _HIDDEN)
110
+ params = [
111
+ PoolingParams(task="embed", dimensions=4),
112
+ PoolingParams(task="embed", dimensions=None),
113
+ ]
114
+ meta = _make_metadata(params)
115
+ out = head(x, meta)
116
+ assert isinstance(out, list)
117
+ assert out[0].shape[-1] == 4
118
+ assert torch.equal(out[1], x[1])
119
+
120
+ def test_activation_uniform_true(self):
121
+ head = EmbeddingPoolerHead(activation=PoolerNormalize())
122
+ x = torch.randn(_BATCH, _HIDDEN)
123
+ params = _make_params(_BATCH, use_activation=True)
124
+ meta = _make_metadata(params)
125
+ out = head(x, meta)
126
+ norms = torch.linalg.norm(out, dim=-1)
127
+ assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
128
+
129
+ def test_activation_uniform_false(self):
130
+ head = EmbeddingPoolerHead(activation=PoolerNormalize())
131
+ x = torch.randn(_BATCH, _HIDDEN)
132
+ params = _make_params(_BATCH, use_activation=False)
133
+ meta = _make_metadata(params)
134
+ out = head(x, meta)
135
+ assert torch.equal(out, x)
136
+
137
+ def test_activation_mixed_flags(self):
138
+ head = EmbeddingPoolerHead(activation=PoolerNormalize())
139
+ x = torch.randn(2, _HIDDEN)
140
+ params = [
141
+ PoolingParams(task="embed", use_activation=True),
142
+ PoolingParams(task="embed", use_activation=False),
143
+ ]
144
+ meta = _make_metadata(params)
145
+ out = head(x, meta)
146
+ assert isinstance(out, list)
147
+ norm_0 = torch.linalg.norm(out[0], dim=-1)
148
+ assert torch.allclose(norm_0, torch.ones(1), atol=1e-5)
149
+ assert torch.equal(out[1], x[1])
150
+
151
+ def test_list_input_gets_stacked(self):
152
+ head = EmbeddingPoolerHead()
153
+ tensors = [torch.randn(_HIDDEN) for _ in range(_BATCH)]
154
+ meta = _make_metadata(_make_params(_BATCH))
155
+ out = head(tensors, meta)
156
+ assert out.shape == (_BATCH, _HIDDEN)
157
+ expected = torch.stack(tensors)
158
+ assert torch.equal(out, expected)
159
+
160
+ def test_projector_then_matryoshka(self):
161
+ proj = _linear(_HIDDEN, 8)
162
+ head = EmbeddingPoolerHead(projector=proj)
163
+ x = torch.randn(_BATCH, _HIDDEN)
164
+ params = _make_params(_BATCH, dimensions=4)
165
+ meta = _make_metadata(params)
166
+ out = head(x, meta)
167
+ assert out.shape == (_BATCH, 4)
168
+ assert torch.equal(out, proj(x)[..., :4])
169
+
170
+ def test_matryoshka_then_activation(self):
171
+ head = EmbeddingPoolerHead(activation=PoolerNormalize())
172
+ x = torch.randn(_BATCH, _HIDDEN)
173
+ params = _make_params(_BATCH, dimensions=4, use_activation=True)
174
+ meta = _make_metadata(params)
175
+ out = head(x, meta)
176
+ assert out.shape == (_BATCH, 4)
177
+ norms = torch.linalg.norm(out, dim=-1)
178
+ assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
179
+
180
+ def test_empty_batch(self):
181
+ head = EmbeddingPoolerHead()
182
+ x = torch.randn(0, _HIDDEN)
183
+ meta = _make_metadata([])
184
+ out = head(x, meta)
185
+ assert out.shape == (0, _HIDDEN)
186
+
187
+
188
+# ---------------------------------------------------------------------------
189
+# ClassifierPoolerHead
190
+# ---------------------------------------------------------------------------
191
+class TestClassifierPoolerHead:
192
+ def test_supported_tasks(self):
193
+ head = ClassifierPoolerHead()
194
+ assert head.get_supported_tasks() == {"classify"}
195
+
196
+ def test_passthrough(self):
197
+ head = ClassifierPoolerHead()
198
+ x = torch.randn(_BATCH, _HIDDEN)
199
+ meta = _make_metadata(_make_params(_BATCH, task="classify"))
200
+ out = head(x, meta)
201
+ assert torch.equal(out, x)
202
+
203
+ def test_head_dtype(self):
204
+ head = ClassifierPoolerHead(head_dtype=torch.float16)
205
+ x = torch.randn(_BATCH, _HIDDEN)
206
+ meta = _make_metadata(_make_params(_BATCH, task="classify"))
207
+ out = head(x, meta)
208
+ assert out.dtype == torch.float16
209
+
210
+ def test_classifier(self):
211
+ clf = _linear(_HIDDEN, 3)
212
+ head = ClassifierPoolerHead(classifier=clf)
213
+ x = torch.randn(_BATCH, _HIDDEN)
214
+ meta = _make_metadata(_make_params(_BATCH, task="classify"))
215
+ out = head(x, meta)
216
+ assert out.shape == (_BATCH, 3)
217
+ assert torch.allclose(out, clf(x))
218
+
219
+ def test_logit_mean(self):
220
+ head = ClassifierPoolerHead(logit_mean=2.0)
221
+ x = torch.randn(_BATCH, _HIDDEN)
222
+ meta = _make_metadata(_make_params(_BATCH, task="classify"))
223
+ out = head(x, meta)
224
+ assert torch.allclose(out, x - 2.0)
225
+
226
+ def test_logit_sigma(self):
227
+ head = ClassifierPoolerHead(logit_sigma=0.5)
228
+ x = torch.randn(_BATCH, _HIDDEN)
229
+ meta = _make_metadata(_make_params(_BATCH, task="classify"))
230
+ out = head(x, meta)
231
+ assert torch.allclose(out, x / 0.5)
232
+
233
+ def test_platt_scaling_combined(self):
234
+ head = ClassifierPoolerHead(logit_mean=1.0, logit_sigma=2.0)
235
+ x = torch.randn(_BATCH, _HIDDEN)
236
+ meta = _make_metadata(_make_params(_BATCH, task="classify"))
237
+ out = head(x, meta)
238
+ assert torch.allclose(out, (x - 1.0) / 2.0)
239
+
240
+ def test_activation_uniform_true(self):
241
+ head = ClassifierPoolerHead(activation=PoolerNormalize())
242
+ x = torch.randn(_BATCH, _HIDDEN)
243
+ params = _make_params(_BATCH, task="classify", use_activation=True)
244
+ meta = _make_metadata(params)
245
+ out = head(x, meta)
246
+ norms = torch.linalg.norm(out, dim=-1)
247
+ assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
248
+
249
+ def test_activation_uniform_false(self):
250
+ head = ClassifierPoolerHead(activation=PoolerNormalize())
251
+ x = torch.randn(_BATCH, _HIDDEN)
252
+ params = _make_params(_BATCH, task="classify", use_activation=False)
253
+ meta = _make_metadata(params)
254
+ out = head(x, meta)
255
+ assert torch.equal(out, x)
256
+
257
+ def test_activation_mixed_flags(self):
258
+ head = ClassifierPoolerHead(activation=PoolerNormalize())
259
+ x = torch.randn(2, _HIDDEN)
260
+ params = [
261
+ PoolingParams(task="classify", use_activation=True),
262
+ PoolingParams(task="classify", use_activation=False),
263
+ ]
264
+ meta = _make_metadata(params)
265
+ out = head(x, meta)
266
+ assert isinstance(out, list)
267
+ norm_0 = torch.linalg.norm(out[0], dim=-1)
268
+ assert torch.allclose(norm_0, torch.ones(1), atol=1e-5)
269
+ assert torch.equal(out[1], x[1])
270
+
271
+ def test_list_input_gets_stacked(self):
272
+ head = ClassifierPoolerHead()
273
+ tensors = [torch.randn(_HIDDEN) for _ in range(_BATCH)]
274
+ meta = _make_metadata(_make_params(_BATCH, task="classify"))
275
+ out = head(tensors, meta)
276
+ assert out.shape == (_BATCH, _HIDDEN)
277
+ expected = torch.stack(tensors)
278
+ assert torch.equal(out, expected)
279
+
280
+ def test_classifier_then_platt_scaling(self):
281
+ clf = _linear(_HIDDEN, 3)
282
+ head = ClassifierPoolerHead(classifier=clf, logit_mean=1.0, logit_sigma=2.0)
283
+ x = torch.randn(_BATCH, _HIDDEN)
284
+ meta = _make_metadata(_make_params(_BATCH, task="classify"))
285
+ out = head(x, meta)
286
+ expected = (clf(x) - 1.0) / 2.0
287
+ assert torch.allclose(out, expected)
288
+
289
+ def test_empty_batch(self):
290
+ head = ClassifierPoolerHead()
291
+ x = torch.randn(0, _HIDDEN)
292
+ meta = _make_metadata([])
293
+ out = head(x, meta)
294
+ assert out.shape == (0, _HIDDEN)
295
+
296
+
297
+# ---------------------------------------------------------------------------
298
+# TokenEmbeddingPoolerHead
299
+# ---------------------------------------------------------------------------
300
+class TestTokenEmbeddingPoolerHead:
301
+ def test_supported_tasks(self):
302
+ head = TokenEmbeddingPoolerHead()
303
+ assert head.get_supported_tasks() == {"token_embed"}
304
+
305
+ def test_passthrough(self):
306
+ head = TokenEmbeddingPoolerHead()
307
+ x = torch.randn(5, _HIDDEN)
308
+ param = PoolingParams(task="token_embed")
309
+ out = head.forward_chunk(x, param)
310
+ assert torch.equal(out, x)
311
+
312
+ def test_none_chunked_prefill(self):
313
+ head = TokenEmbeddingPoolerHead()
314
+ param = PoolingParams(task="token_embed")
315
+ out = head.forward_chunk(None, param)
316
+ assert out is None
317
+
318
+ def test_head_dtype(self):
319
+ head = TokenEmbeddingPoolerHead(head_dtype=torch.float16)
320
+ x = torch.randn(5, _HIDDEN)
321
+ param = PoolingParams(task="token_embed")
322
+ out = head.forward_chunk(x, param)
323
+ assert out.dtype == torch.float16
324
+
325
+ def test_projector(self):
326
+ proj = _linear(_HIDDEN, 8)
327
+ head = TokenEmbeddingPoolerHead(projector=proj)
328
+ x = torch.randn(5, _HIDDEN)
329
+ param = PoolingParams(task="token_embed")
330
+ out = head.forward_chunk(x, param)
331
+ assert out.shape == (5, 8)
332
+ assert torch.allclose(out, proj(x))
333
+
334
+ def test_matryoshka_truncation(self):
335
+ head = TokenEmbeddingPoolerHead()
336
+ x = torch.randn(5, _HIDDEN)
337
+ param = PoolingParams(task="token_embed", dimensions=4)
338
+ out = head.forward_chunk(x, param)
339
+ assert out.shape == (5, 4)
340
+ assert torch.equal(out, x[..., :4])
341
+
342
+ def test_activation_true(self):
343
+ head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
344
+ x = torch.randn(5, _HIDDEN)
345
+ param = PoolingParams(task="token_embed", use_activation=True)
346
+ out = head.forward_chunk(x, param)
347
+ norms = torch.linalg.norm(out, dim=-1)
348
+ assert torch.allclose(norms, torch.ones(5), atol=1e-5)
349
+
350
+ def test_activation_false(self):
351
+ head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
352
+ x = torch.randn(5, _HIDDEN)
353
+ param = PoolingParams(task="token_embed", use_activation=False)
354
+ out = head.forward_chunk(x, param)
355
+ assert torch.equal(out, x)
356
+
357
+ def test_projector_then_matryoshka(self):
358
+ proj = _linear(_HIDDEN, 8)
359
+ head = TokenEmbeddingPoolerHead(projector=proj)
360
+ x = torch.randn(5, _HIDDEN)
361
+ param = PoolingParams(task="token_embed", dimensions=4)
362
+ out = head.forward_chunk(x, param)
363
+ assert out.shape == (5, 4)
364
+ assert torch.equal(out, proj(x)[..., :4])
365
+
366
+ def test_matryoshka_then_activation(self):
367
+ head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
368
+ x = torch.randn(5, _HIDDEN)
369
+ param = PoolingParams(task="token_embed", dimensions=4, use_activation=True)
370
+ out = head.forward_chunk(x, param)
371
+ assert out.shape == (5, 4)
372
+ norms = torch.linalg.norm(out, dim=-1)
373
+ assert torch.allclose(norms, torch.ones(5), atol=1e-5)
374
+
375
+ def test_forward_mixed_batch_chunked_prefill(self):
376
+ head = TokenEmbeddingPoolerHead()
377
+ pooled_data = [torch.randn(5, _HIDDEN), None, torch.randn(3, _HIDDEN)]
378
+ params = _make_params(3, task="token_embed")
379
+ meta = _make_metadata(params)
380
+ out = head(pooled_data, meta)
381
+ assert len(out) == 3
382
+ assert torch.equal(out[0], pooled_data[0])
383
+ assert out[1] is None
384
+ assert torch.equal(out[2], pooled_data[2])
385
+
386
+ def test_forward_empty_batch(self):
387
+ head = TokenEmbeddingPoolerHead()
388
+ meta = _make_metadata([])
389
+ out = head([], meta)
390
+ assert out == []
391
+
392
+
393
+# ---------------------------------------------------------------------------
394
+# TokenClassifierPoolerHead
395
+# ---------------------------------------------------------------------------
396
+class TestTokenClassifierPoolerHead:
397
+ def test_supported_tasks(self):
398
+ head = TokenClassifierPoolerHead()
399
+ assert head.get_supported_tasks() == {"token_classify"}
400
+
401
+ def test_passthrough(self):
402
+ head = TokenClassifierPoolerHead()
403
+ x = torch.randn(5, _HIDDEN)
404
+ param = PoolingParams(task="token_classify")
405
+ out = head.forward_chunk(x, param)
406
+ assert torch.equal(out, x)
407
+
408
+ def test_none_chunked_prefill(self):
409
+ head = TokenClassifierPoolerHead()
410
+ param = PoolingParams(task="token_classify")
411
+ out = head.forward_chunk(None, param)
412
+ assert out is None
413
+
414
+ def test_head_dtype(self):
415
+ head = TokenClassifierPoolerHead(head_dtype=torch.float16)
416
+ x = torch.randn(5, _HIDDEN)
417
+ param = PoolingParams(task="token_classify")
418
+ out = head.forward_chunk(x, param)
419
+ assert out.dtype == torch.float16
420
+
421
+ def test_classifier(self):
422
+ clf = _linear(_HIDDEN, 3)
423
+ head = TokenClassifierPoolerHead(classifier=clf)
424
+ x = torch.randn(5, _HIDDEN)
425
+ param = PoolingParams(task="token_classify")
426
+ out = head.forward_chunk(x, param)
427
+ assert out.shape == (5, 3)
428
+ assert torch.allclose(out, clf(x))
429
+
430
+ def test_logit_mean(self):
431
+ head = TokenClassifierPoolerHead(logit_mean=2.0)
432
+ x = torch.randn(5, _HIDDEN)
433
+ param = PoolingParams(task="token_classify")
434
+ out = head.forward_chunk(x, param)
435
+ assert torch.allclose(out, x - 2.0)
436
+
437
+ def test_logit_sigma(self):
438
+ head = TokenClassifierPoolerHead(logit_sigma=0.5)
439
+ x = torch.randn(5, _HIDDEN)
440
+ param = PoolingParams(task="token_classify")
441
+ out = head.forward_chunk(x, param)
442
+ assert torch.allclose(out, x / 0.5)
443
+
444
+ def test_platt_scaling_combined(self):
445
+ head = TokenClassifierPoolerHead(logit_mean=1.0, logit_sigma=2.0)
446
+ x = torch.randn(5, _HIDDEN)
447
+ param = PoolingParams(task="token_classify")
448
+ out = head.forward_chunk(x, param)
449
+ assert torch.allclose(out, (x - 1.0) / 2.0)
450
+
451
+ def test_activation_true(self):
452
+ head = TokenClassifierPoolerHead(activation=PoolerNormalize())
453
+ x = torch.randn(5, _HIDDEN)
454
+ param = PoolingParams(task="token_classify", use_activation=True)
455
+ out = head.forward_chunk(x, param)
456
+ norms = torch.linalg.norm(out, dim=-1)
457
+ assert torch.allclose(norms, torch.ones(5), atol=1e-5)
458
+
459
+ def test_activation_false(self):
460
+ head = TokenClassifierPoolerHead(activation=PoolerNormalize())
461
+ x = torch.randn(5, _HIDDEN)
462
+ param = PoolingParams(task="token_classify", use_activation=False)
463
+ out = head.forward_chunk(x, param)
464
+ assert torch.equal(out, x)
465
+
466
+ def test_forward_mixed_batch_chunked_prefill(self):
467
+ head = TokenClassifierPoolerHead()
468
+ pooled_data = [torch.randn(5, _HIDDEN), None, torch.randn(3, _HIDDEN)]
469
+ params = _make_params(3, task="token_classify")
470
+ meta = _make_metadata(params)
471
+ out = head(pooled_data, meta)
472
+ assert len(out) == 3
473
+ assert torch.equal(out[0], pooled_data[0])
474
+ assert out[1] is None
475
+ assert torch.equal(out[2], pooled_data[2])
476
+
477
+ def test_forward_empty_batch(self):
478
+ head = TokenClassifierPoolerHead()
479
+ meta = _make_metadata([])
480
+ out = head([], meta)
481
+ assert out == []