748 lines
27 KiB
Python
748 lines
27 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Unit tests for sequence and token pooling methods and their factories."""
|
|
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from transformers import PreTrainedConfig
|
|
|
|
from vllm.config import PoolerConfig, set_current_vllm_config
|
|
from vllm.model_executor.layers.pooler import PoolingParamsUpdate
|
|
from vllm.model_executor.layers.pooler.activations import PoolerNormalize
|
|
from vllm.model_executor.layers.pooler.seqwise.methods import (
|
|
CLSPool,
|
|
LastPool,
|
|
MeanPool,
|
|
get_seq_pooling_method,
|
|
)
|
|
from vllm.model_executor.layers.pooler.special import BOSEOSFilter, DispatchPooler
|
|
from vllm.model_executor.layers.pooler.tokwise.heads import TokenEmbeddingPoolerHead
|
|
from vllm.model_executor.layers.pooler.tokwise.methods import (
|
|
AllPool,
|
|
StepPool,
|
|
get_tok_pooling_method,
|
|
)
|
|
from vllm.model_executor.layers.pooler.tokwise.poolers import TokenPooler
|
|
from vllm.pooling_params import LateChunkingParams, PoolingParams
|
|
from vllm.tasks import PoolingTask
|
|
from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates
|
|
|
|
_CPU = torch.device("cpu")
|
|
|
|
|
|
def _make_pooling_cursor(
|
|
prompt_lens: list[int],
|
|
*,
|
|
num_scheduled_tokens: list[int] | None = None,
|
|
seq_lens: list[int] | None = None,
|
|
device: torch.device = _CPU,
|
|
) -> PoolingCursor:
|
|
"""Build a PoolingCursor from a list of per-sequence prompt lengths."""
|
|
prompt_lens_cpu = torch.tensor(prompt_lens, dtype=torch.long)
|
|
if num_scheduled_tokens is None:
|
|
num_scheduled_tokens_cpu = prompt_lens_cpu.clone()
|
|
else:
|
|
num_scheduled_tokens_cpu = torch.tensor(num_scheduled_tokens, dtype=torch.long)
|
|
if seq_lens is None:
|
|
seq_lens_cpu = prompt_lens_cpu.clone()
|
|
else:
|
|
seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.long)
|
|
|
|
cumsum = torch.zeros(len(prompt_lens) + 1, dtype=torch.long, device=device)
|
|
torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
|
|
|
|
return PoolingCursor(
|
|
first_token_indices_gpu=cumsum[: len(prompt_lens)].to(device),
|
|
last_token_indices_gpu=(cumsum[1:] - 1).to(device),
|
|
prompt_lens_cpu=prompt_lens_cpu,
|
|
seq_lens_cpu=seq_lens_cpu,
|
|
num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
|
|
partial_prefill=not torch.equal(prompt_lens_cpu, num_scheduled_tokens_cpu),
|
|
finished_mask=torch.eq(prompt_lens_cpu, seq_lens_cpu).tolist(),
|
|
)
|
|
|
|
|
|
def _make_metadata(
|
|
prompt_lens: list[int],
|
|
*,
|
|
tasks: list[PoolingTask] | None = None,
|
|
token_ids: list[list[int]] | None = None,
|
|
pooling_params: list[PoolingParams] | None = None,
|
|
num_scheduled_tokens: list[int] | None = None,
|
|
seq_lens: list[int] | None = None,
|
|
device: torch.device = _CPU,
|
|
) -> PoolingMetadata:
|
|
"""Build a minimal PoolingMetadata for testing pooling methods."""
|
|
n_seqs = len(prompt_lens)
|
|
if tasks is None:
|
|
tasks = ["embed"] * n_seqs
|
|
if pooling_params is None:
|
|
pooling_params = [PoolingParams(task=t) for t in tasks]
|
|
|
|
prompt_lens_tensor = torch.tensor(prompt_lens, dtype=torch.long)
|
|
|
|
prompt_token_ids_cpu = None
|
|
prompt_token_ids = None
|
|
if token_ids is not None:
|
|
max_len = max(len(t) for t in token_ids)
|
|
padded = [t + [0] * (max_len - len(t)) for t in token_ids]
|
|
prompt_token_ids_cpu = torch.tensor(padded, dtype=torch.long)
|
|
prompt_token_ids = prompt_token_ids_cpu.to(device)
|
|
|
|
cursor = _make_pooling_cursor(
|
|
prompt_lens,
|
|
num_scheduled_tokens=num_scheduled_tokens,
|
|
seq_lens=seq_lens,
|
|
device=device,
|
|
)
|
|
|
|
pooling_states = [PoolingStates() for _ in range(n_seqs)]
|
|
|
|
return PoolingMetadata(
|
|
prompt_lens=prompt_lens_tensor,
|
|
prompt_token_ids=prompt_token_ids,
|
|
prompt_token_ids_cpu=prompt_token_ids_cpu,
|
|
pooling_params=pooling_params,
|
|
pooling_states=pooling_states,
|
|
pooling_cursor=cursor,
|
|
)
|
|
|
|
|
|
def test_build_pooling_cursor_caches_cpu_status() -> None:
|
|
metadata = PoolingMetadata(
|
|
prompt_lens=torch.tensor([3, 4, 2], dtype=torch.long),
|
|
prompt_token_ids=None,
|
|
prompt_token_ids_cpu=None,
|
|
pooling_params=[PoolingParams(task="embed") for _ in range(3)],
|
|
pooling_states=[PoolingStates() for _ in range(3)],
|
|
)
|
|
metadata.build_pooling_cursor(
|
|
np.array([3, 2, 2], dtype=np.int32),
|
|
seq_lens_cpu=torch.tensor([3, 2, 1], dtype=torch.long),
|
|
device=_CPU,
|
|
)
|
|
|
|
cursor = metadata.get_pooling_cursor()
|
|
|
|
assert cursor.is_partial_prefill() is True
|
|
assert cursor.get_finished_mask() == [True, False, False]
|
|
assert cursor[:1].is_partial_prefill() is False
|
|
assert cursor[1:2].is_partial_prefill() is True
|
|
assert cursor[1:].get_finished_mask() == [False, False]
|
|
|
|
|
|
def test_pooling_params_update_combines_token_id_requirements() -> None:
|
|
update = PoolingParamsUpdate() | PoolingParamsUpdate(requires_token_ids=True)
|
|
params = PoolingParams(task="embed")
|
|
|
|
update.apply(params)
|
|
|
|
assert params.requires_token_ids is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLSPool
|
|
# ---------------------------------------------------------------------------
|
|
class TestCLSPool:
|
|
def test_extracts_first_token(self):
|
|
hidden = torch.tensor(
|
|
[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
|
|
)
|
|
metadata = _make_metadata([2, 3])
|
|
pooler = CLSPool()
|
|
out = pooler(hidden, metadata)
|
|
expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
|
|
assert torch.equal(out, expected)
|
|
|
|
def test_rejects_partial_prefill(self):
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
|
|
metadata = _make_metadata([3], num_scheduled_tokens=[2])
|
|
pooler = CLSPool()
|
|
with pytest.raises(RuntimeError, match="partial prefill"):
|
|
pooler(hidden, metadata)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LastPool
|
|
# ---------------------------------------------------------------------------
|
|
class TestLastPool:
|
|
def test_extracts_last_token(self):
|
|
hidden = torch.tensor(
|
|
[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
|
|
)
|
|
metadata = _make_metadata([2, 3])
|
|
pooler = LastPool()
|
|
out = pooler(hidden, metadata)
|
|
expected = torch.tensor([[3.0, 4.0], [9.0, 10.0]])
|
|
assert torch.equal(out, expected)
|
|
|
|
def test_partial_prefill_extracts_last_scheduled(self):
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
metadata = _make_metadata([4], num_scheduled_tokens=[2])
|
|
pooler = LastPool()
|
|
out = pooler(hidden, metadata)
|
|
expected = torch.tensor([[3.0, 4.0]])
|
|
assert torch.equal(out, expected)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MeanPool
|
|
# ---------------------------------------------------------------------------
|
|
class TestMeanPool:
|
|
def test_computes_mean(self):
|
|
hidden = torch.tensor(
|
|
[[1.0, 2.0], [3.0, 4.0], [10.0, 20.0]], dtype=torch.float32
|
|
)
|
|
metadata = _make_metadata([2, 1])
|
|
pooler = MeanPool()
|
|
out = pooler(hidden, metadata)
|
|
expected = torch.tensor([[2.0, 3.0], [10.0, 20.0]], dtype=torch.float32)
|
|
assert torch.allclose(out, expected, atol=1e-5)
|
|
|
|
def test_single_token_is_identity(self):
|
|
hidden = torch.tensor([[5.0, 10.0]], dtype=torch.float32)
|
|
metadata = _make_metadata([1])
|
|
pooler = MeanPool()
|
|
out = pooler(hidden, metadata)
|
|
assert torch.allclose(out, hidden, atol=1e-5)
|
|
|
|
def test_uniform_values_return_same(self):
|
|
hidden = torch.full((4, 3), 7.0, dtype=torch.float32)
|
|
metadata = _make_metadata([4])
|
|
pooler = MeanPool()
|
|
out = pooler(hidden, metadata)
|
|
expected = torch.full((1, 3), 7.0, dtype=torch.float32)
|
|
assert torch.allclose(out, expected, atol=1e-5)
|
|
|
|
def test_multiple_sequences(self):
|
|
hidden = torch.tensor(
|
|
[
|
|
[0.0, 0.0],
|
|
[2.0, 4.0],
|
|
[4.0, 8.0],
|
|
[10.0, 10.0],
|
|
],
|
|
dtype=torch.float32,
|
|
)
|
|
metadata = _make_metadata([3, 1])
|
|
pooler = MeanPool()
|
|
out = pooler(hidden, metadata)
|
|
expected = torch.tensor([[2.0, 4.0], [10.0, 10.0]], dtype=torch.float32)
|
|
assert torch.allclose(out, expected, atol=1e-5)
|
|
|
|
def test_empty_batch(self):
|
|
hidden = torch.empty((0, 8), dtype=torch.float32)
|
|
metadata = _make_metadata([])
|
|
pooler = MeanPool()
|
|
out = pooler(hidden, metadata)
|
|
assert out.shape == (0, 8)
|
|
|
|
def test_rejects_partial_prefill(self):
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
|
|
metadata = _make_metadata([3], num_scheduled_tokens=[2])
|
|
pooler = MeanPool()
|
|
with pytest.raises(RuntimeError, match="partial prefill"):
|
|
pooler(hidden, metadata)
|
|
|
|
def test_chunked_accumulation(self):
|
|
hidden = torch.arange(20, dtype=torch.float32).reshape(5, 4)
|
|
metadata = _make_metadata([3, 2])
|
|
pooler = MeanPool()
|
|
with patch(
|
|
"vllm.model_executor.layers.pooler.seqwise.methods"
|
|
"._MEAN_POOL_ACCUMULATION_CHUNK_BYTES",
|
|
16,
|
|
):
|
|
out = pooler(hidden, metadata)
|
|
expected_seq0 = hidden[:3].float().mean(dim=0, keepdim=True)
|
|
expected_seq1 = hidden[3:].float().mean(dim=0, keepdim=True)
|
|
expected = torch.cat([expected_seq0, expected_seq1], dim=0)
|
|
assert torch.allclose(out, expected, atol=1e-5)
|
|
|
|
def test_upcasts_to_float32(self):
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float16)
|
|
metadata = _make_metadata([2])
|
|
pooler = MeanPool()
|
|
out = pooler(hidden, metadata)
|
|
assert out.dtype == torch.float32
|
|
expected = torch.tensor([[2.0, 3.0]], dtype=torch.float32)
|
|
assert torch.allclose(out, expected, atol=1e-2)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# get_seq_pooling_method factory
|
|
# ---------------------------------------------------------------------------
|
|
class TestGetSeqPoolingMethod:
|
|
def test_cls(self):
|
|
assert isinstance(get_seq_pooling_method("CLS"), CLSPool)
|
|
|
|
def test_last(self):
|
|
assert isinstance(get_seq_pooling_method("LAST"), LastPool)
|
|
|
|
def test_mean(self):
|
|
assert isinstance(get_seq_pooling_method("MEAN"), MeanPool)
|
|
|
|
def test_unknown_raises(self):
|
|
with pytest.raises(NotImplementedError, match="UNKNOWN"):
|
|
get_seq_pooling_method("UNKNOWN")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AllPool
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class _FakeSchedulerConfig:
|
|
enable_chunked_prefill: bool = False
|
|
async_scheduling: bool = False
|
|
|
|
|
|
@dataclass
|
|
class _FakeVllmConfig:
|
|
scheduler_config: _FakeSchedulerConfig
|
|
|
|
|
|
@pytest.mark.parametrize("tok_pooling_type", ["ALL", "STEP"])
|
|
def test_dispatch_seq_cls_honors_token_pooling_type(tok_pooling_type):
|
|
"""Filter STEP rows before classification; preserve all rows for resolved ALL."""
|
|
pooler_config = PoolerConfig(
|
|
seq_pooling_type="LAST", tok_pooling_type=tok_pooling_type
|
|
)
|
|
config = SimpleNamespace(
|
|
scheduler_config=_FakeSchedulerConfig(),
|
|
model_config=SimpleNamespace(
|
|
pooler_config=pooler_config,
|
|
head_dtype=None,
|
|
hf_config=PreTrainedConfig(num_labels=3),
|
|
),
|
|
)
|
|
classifier = torch.nn.Linear(4, 3)
|
|
with set_current_vllm_config(config):
|
|
pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=classifier)
|
|
|
|
token_ids = [[2, 99, 3, 99], [4, 99], [5, 6]]
|
|
params = [
|
|
PoolingParams(
|
|
task="token_classify",
|
|
use_activation=False,
|
|
step_tag_id=99 if tok_pooling_type == "STEP" else None,
|
|
)
|
|
for _ in token_ids
|
|
]
|
|
metadata = _make_metadata([4, 2, 2], token_ids=token_ids, pooling_params=params)
|
|
hidden = torch.arange(32, dtype=torch.float32).reshape(8, 4)
|
|
outputs = pooler(hidden, metadata)
|
|
|
|
expected_rows = (
|
|
[[1, 3], [5], []]
|
|
if tok_pooling_type == "STEP"
|
|
else [[0, 1, 2, 3], [4, 5], [6, 7]]
|
|
)
|
|
assert len(outputs) == len(expected_rows)
|
|
for output, rows in zip(outputs, expected_rows):
|
|
torch.testing.assert_close(output, classifier(hidden[rows]))
|
|
assert pooler.get_pooling_updates("token_classify").requires_token_ids == (
|
|
tok_pooling_type == "STEP"
|
|
)
|
|
|
|
|
|
class TestAllPool:
|
|
@staticmethod
|
|
def _make_all_pool(
|
|
*, chunked: bool = False, async_scheduling: bool = False
|
|
) -> AllPool:
|
|
fake_config = _FakeVllmConfig(
|
|
scheduler_config=_FakeSchedulerConfig(
|
|
enable_chunked_prefill=chunked,
|
|
async_scheduling=async_scheduling,
|
|
),
|
|
)
|
|
with patch(
|
|
"vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
|
|
return_value=fake_config,
|
|
):
|
|
return AllPool()
|
|
|
|
def test_splits_by_sequence(self):
|
|
pooler = self._make_all_pool()
|
|
hidden = torch.tensor(
|
|
[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
|
|
)
|
|
metadata = _make_metadata([2, 3])
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 2
|
|
assert torch.equal(out[0], hidden[:2])
|
|
assert torch.equal(out[1], hidden[2:])
|
|
|
|
def test_single_sequence(self):
|
|
pooler = self._make_all_pool()
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
|
|
metadata = _make_metadata([3])
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
assert torch.equal(out[0], hidden)
|
|
|
|
def test_chunked_prefill_returns_none_for_unfinished(self):
|
|
pooler = self._make_all_pool(chunked=True)
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
metadata = _make_metadata(
|
|
[4],
|
|
num_scheduled_tokens=[2],
|
|
seq_lens=[2],
|
|
)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
assert out[0] is None
|
|
|
|
def test_chunked_prefill_returns_concat_when_finished(self):
|
|
pooler = self._make_all_pool(chunked=True)
|
|
|
|
chunk1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
metadata1 = _make_metadata(
|
|
[4],
|
|
num_scheduled_tokens=[2],
|
|
seq_lens=[2],
|
|
)
|
|
out1 = pooler(chunk1, metadata1)
|
|
assert out1[0] is None
|
|
expected_chunk1 = chunk1.clone()
|
|
chunk1.zero_()
|
|
|
|
chunk2 = torch.tensor([[5.0, 6.0], [7.0, 8.0]])
|
|
metadata2 = _make_metadata(
|
|
[4],
|
|
num_scheduled_tokens=[2],
|
|
seq_lens=[4],
|
|
)
|
|
metadata2.pooling_states = metadata1.pooling_states
|
|
out2 = pooler(chunk2, metadata2)
|
|
assert out2[0] is not None
|
|
expected = torch.cat([expected_chunk1, chunk2], dim=0)
|
|
assert torch.equal(out2[0], expected)
|
|
|
|
@pytest.mark.parametrize("async_scheduling", [False, True])
|
|
def test_chunked_prefill_single_shot_matches_non_chunked(
|
|
self, async_scheduling: bool
|
|
):
|
|
pooler = self._make_all_pool(chunked=True, async_scheduling=async_scheduling)
|
|
hidden = torch.tensor(
|
|
[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
|
|
)
|
|
expected = [hidden[:2].clone(), hidden[2:].clone()]
|
|
metadata = _make_metadata([2, 3])
|
|
out = pooler(hidden, metadata)
|
|
if async_scheduling:
|
|
# Finished outputs must own their storage under async scheduling.
|
|
hidden.zero_()
|
|
assert len(out) == 2
|
|
assert torch.equal(out[0], expected[0])
|
|
assert torch.equal(out[1], expected[1])
|
|
|
|
def test_non_chunked_owns_output_under_async_scheduling(self):
|
|
pooler = self._make_all_pool(async_scheduling=True)
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
|
|
expected = [hidden[:2].clone(), hidden[2:].clone()]
|
|
metadata = _make_metadata([2, 1])
|
|
out = pooler(hidden, metadata)
|
|
hidden.zero_()
|
|
assert len(out) == 2
|
|
assert torch.equal(out[0], expected[0])
|
|
assert torch.equal(out[1], expected[1])
|
|
|
|
def test_chunked_prefill_mixed_finished_unfinished(self):
|
|
pooler = self._make_all_pool(chunked=True)
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
|
|
metadata = _make_metadata(
|
|
[2, 4],
|
|
num_scheduled_tokens=[2, 1],
|
|
seq_lens=[2, 1],
|
|
)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 2
|
|
assert torch.equal(out[0], hidden[:2])
|
|
assert out[1] is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# StepPool
|
|
# ---------------------------------------------------------------------------
|
|
class TestStepPool:
|
|
@staticmethod
|
|
def _make_step_pool(*, chunked: bool = False) -> StepPool:
|
|
fake_config = _FakeVllmConfig(
|
|
scheduler_config=_FakeSchedulerConfig(
|
|
enable_chunked_prefill=chunked,
|
|
),
|
|
)
|
|
with patch(
|
|
"vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
|
|
return_value=fake_config,
|
|
):
|
|
return StepPool()
|
|
|
|
def test_filters_by_step_tag_id(self):
|
|
pooler = self._make_step_pool()
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
|
|
token_ids = [[10, 99, 10, 20]]
|
|
params = [PoolingParams(task="token_classify", step_tag_id=10)]
|
|
metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
|
|
assert torch.equal(out[0], expected)
|
|
|
|
def test_filters_by_returned_token_ids(self):
|
|
pooler = self._make_step_pool()
|
|
hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
|
|
token_ids = [[10, 20]]
|
|
params = [PoolingParams(task="token_classify", returned_token_ids=[0, 2])]
|
|
metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
expected = torch.tensor([[1.0, 3.0], [4.0, 6.0]])
|
|
assert torch.equal(out[0], expected)
|
|
|
|
def test_no_filtering_without_params(self):
|
|
pooler = self._make_step_pool()
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
token_ids = [[10, 20]]
|
|
params = [PoolingParams(task="token_classify")]
|
|
metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
assert torch.equal(out[0], hidden)
|
|
|
|
def test_combined_step_tag_and_returned_token_ids(self):
|
|
pooler = self._make_step_pool()
|
|
hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
|
|
token_ids = [[99, 10, 99]]
|
|
params = [
|
|
PoolingParams(
|
|
task="token_classify",
|
|
step_tag_id=10,
|
|
returned_token_ids=[0, 2],
|
|
)
|
|
]
|
|
metadata = _make_metadata([3], token_ids=token_ids, pooling_params=params)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
expected = torch.tensor([[4.0, 6.0]])
|
|
assert torch.equal(out[0], expected)
|
|
|
|
def test_step_tag_id_no_match_returns_empty(self):
|
|
pooler = self._make_step_pool()
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
token_ids = [[10, 20]]
|
|
params = [PoolingParams(task="token_classify", step_tag_id=999)]
|
|
metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
assert out[0].shape == (0, 2)
|
|
|
|
def test_chunked_prefill_propagates_none_for_unfinished(self):
|
|
pooler = self._make_step_pool(chunked=True)
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
token_ids = [[10, 20, 30, 40]]
|
|
params = [PoolingParams(task="token_classify", step_tag_id=10)]
|
|
metadata = _make_metadata(
|
|
[4],
|
|
token_ids=token_ids,
|
|
pooling_params=params,
|
|
num_scheduled_tokens=[2],
|
|
seq_lens=[2],
|
|
)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
assert out[0] is None
|
|
|
|
def test_chunked_prefill_filters_when_finished(self):
|
|
pooler = self._make_step_pool(chunked=True)
|
|
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
|
|
token_ids = [[10, 99, 10, 20]]
|
|
params = [PoolingParams(task="token_classify", step_tag_id=10)]
|
|
metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params)
|
|
out = pooler(hidden, metadata)
|
|
assert len(out) == 1
|
|
expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
|
|
assert torch.equal(out[0], expected)
|
|
|
|
def test_requires_token_ids_update(self):
|
|
pooler = self._make_step_pool()
|
|
update = pooler.get_pooling_updates("token_classify")
|
|
assert update.requires_token_ids is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# get_tok_pooling_method factory
|
|
# ---------------------------------------------------------------------------
|
|
class TestGetTokPoolingMethod:
|
|
def test_all(self):
|
|
fake_config = _FakeVllmConfig(
|
|
scheduler_config=_FakeSchedulerConfig(
|
|
enable_chunked_prefill=False,
|
|
),
|
|
)
|
|
with patch(
|
|
"vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
|
|
return_value=fake_config,
|
|
):
|
|
assert isinstance(get_tok_pooling_method("ALL"), AllPool)
|
|
|
|
def test_step(self):
|
|
fake_config = _FakeVllmConfig(
|
|
scheduler_config=_FakeSchedulerConfig(
|
|
enable_chunked_prefill=False,
|
|
),
|
|
)
|
|
with patch(
|
|
"vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
|
|
return_value=fake_config,
|
|
):
|
|
assert isinstance(get_tok_pooling_method("STEP"), StepPool)
|
|
|
|
def test_unknown_raises(self):
|
|
with pytest.raises(NotImplementedError, match="UNKNOWN"):
|
|
get_tok_pooling_method("UNKNOWN")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("chunk_size", "dtype", "head_dtype"),
|
|
[
|
|
(1, torch.float32, None),
|
|
(2, torch.float16, None),
|
|
(2, torch.bfloat16, torch.float32),
|
|
(7, torch.float32, None),
|
|
(10, torch.float32, None),
|
|
(2**63 - 1, torch.float32, None),
|
|
],
|
|
)
|
|
def test_late_chunk_pool_matches_fp32_means_before_projection_and_normalization(
|
|
chunk_size, dtype, head_dtype
|
|
):
|
|
# Unequal norms expose the incorrect normalize-then-average implementation.
|
|
hidden = (torch.arange(28).reshape(7, 4) - 10).to(dtype)
|
|
projector = torch.nn.Linear(4, 3, bias=True).to(head_dtype or dtype)
|
|
head = TokenEmbeddingPoolerHead(head_dtype, projector, PoolerNormalize())
|
|
pooler = TokenPooler(TestAllPool._make_all_pool(), head)
|
|
params = PoolingParams(
|
|
task="token_embed",
|
|
late_chunking_params=LateChunkingParams(chunk_size=chunk_size),
|
|
use_activation=True,
|
|
)
|
|
metadata = _make_metadata([7], pooling_params=[params])
|
|
actual = pooler(hidden, metadata)[0]
|
|
means = torch.stack(
|
|
[hidden[a : a + chunk_size].float().mean(0) for a in range(0, 7, chunk_size)]
|
|
).to(head_dtype or dtype)
|
|
expected = torch.nn.functional.normalize(projector(means), p=2, dim=-1)
|
|
torch.testing.assert_close(actual, expected)
|
|
assert actual.dtype == (head_dtype or dtype)
|
|
|
|
|
|
def test_late_chunk_pool_reduces_before_nonlinear_projector():
|
|
hidden = torch.tensor([[2.0, 0.0], [0.0, 1.0]])
|
|
projector = torch.nn.Tanh()
|
|
head = TokenEmbeddingPoolerHead(projector=projector, activation=PoolerNormalize())
|
|
params = PoolingParams(
|
|
task="token_embed",
|
|
late_chunking_params=LateChunkingParams(chunk_size=2),
|
|
use_activation=True,
|
|
)
|
|
actual = head.forward_chunk(hidden, params)
|
|
expected = torch.nn.functional.normalize(projector(hidden.mean(0, keepdim=True)))
|
|
wrong_order = torch.nn.functional.normalize(projector(hidden).mean(0, keepdim=True))
|
|
torch.testing.assert_close(actual, expected)
|
|
assert not torch.allclose(actual, wrong_order)
|
|
|
|
|
|
def test_late_chunk_pool_isolates_mixed_requests_and_owns_finished_storage():
|
|
hidden = torch.arange(36, dtype=torch.float32).reshape(9, 4)
|
|
params = [
|
|
PoolingParams(
|
|
task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=2)
|
|
),
|
|
PoolingParams(task="token_embed"),
|
|
PoolingParams(
|
|
task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=1)
|
|
),
|
|
]
|
|
pooler = TokenPooler(
|
|
TestAllPool._make_all_pool(async_scheduling=True), TokenEmbeddingPoolerHead()
|
|
)
|
|
outputs = pooler(hidden, _make_metadata([3, 4, 2], pooling_params=params))
|
|
expected = [
|
|
torch.stack([hidden[:2].mean(0), hidden[2]]),
|
|
hidden[3:7].clone(),
|
|
hidden[7:].clone(),
|
|
]
|
|
hidden.fill_(-999)
|
|
for actual, reference in zip(outputs, expected):
|
|
torch.testing.assert_close(actual, reference)
|
|
|
|
|
|
@pytest.mark.parametrize("chunk_size", [1, 2, 8])
|
|
@pytest.mark.parametrize("embedding_size", [1, 4])
|
|
def test_bos_eos_filter_preserves_chunk_vectors_in_mixed_batch(
|
|
chunk_size, embedding_size
|
|
):
|
|
hidden = torch.arange(10 * embedding_size, dtype=torch.float32).reshape(
|
|
10, embedding_size
|
|
)
|
|
pooler = BOSEOSFilter(
|
|
TokenPooler(TestAllPool._make_all_pool(), TokenEmbeddingPoolerHead()),
|
|
bos_token_id=101,
|
|
eos_token_id=102,
|
|
)
|
|
metadata = _make_metadata(
|
|
[5, 5],
|
|
token_ids=[[101, 1, 2, 3, 102]] * 2,
|
|
pooling_params=[
|
|
PoolingParams(
|
|
task="token_embed", late_chunking_params=LateChunkingParams(chunk_size)
|
|
),
|
|
PoolingParams(task="token_embed"),
|
|
],
|
|
)
|
|
chunked, ordinary = pooler(hidden, metadata)
|
|
expected = torch.stack([part.mean(0) for part in hidden[:5].split(chunk_size)])
|
|
torch.testing.assert_close(chunked, expected)
|
|
torch.testing.assert_close(ordinary, hidden[6:9].squeeze(-1))
|
|
|
|
|
|
def test_bos_eos_filter_waits_for_complete_chunk_vectors():
|
|
pooler = BOSEOSFilter(
|
|
TokenPooler(
|
|
TestAllPool._make_all_pool(chunked=True), TokenEmbeddingPoolerHead()
|
|
),
|
|
bos_token_id=101,
|
|
)
|
|
params = [
|
|
PoolingParams(task="token_embed", late_chunking_params=LateChunkingParams(3))
|
|
]
|
|
first = _make_metadata(
|
|
[4],
|
|
token_ids=[[101, 1, 2, 102]],
|
|
pooling_params=params,
|
|
num_scheduled_tokens=[2],
|
|
seq_lens=[2],
|
|
)
|
|
hidden = torch.arange(8, dtype=torch.float32).reshape(4, 2)
|
|
assert pooler(hidden[:2], first) == [None]
|
|
last = _make_metadata(
|
|
[4],
|
|
token_ids=[[101, 1, 2, 102]],
|
|
pooling_params=params,
|
|
num_scheduled_tokens=[2],
|
|
seq_lens=[4],
|
|
)
|
|
last.pooling_states = first.pooling_states
|
|
actual = pooler(hidden[2:], last)[0]
|
|
torch.testing.assert_close(actual, torch.stack([hidden[:3].mean(0), hidden[3]]))
|
|
assert not last.pooling_states[0].hidden_states_cache
|