chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,234 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for vllm.model_executor.layers.pooler.activations."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm.model_executor.layers.pooler.activations import (
|
||||
LambdaPoolerActivation,
|
||||
PoolerClassify,
|
||||
PoolerIdentity,
|
||||
PoolerMultiLabelClassify,
|
||||
PoolerNormalize,
|
||||
get_act_fn,
|
||||
resolve_classifier_act_fn,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PoolerIdentity
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestPoolerIdentity:
|
||||
def test_returns_input_unchanged(self):
|
||||
pooler = PoolerIdentity()
|
||||
x = torch.randn(4, 128)
|
||||
out = pooler(x)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_forward_list(self):
|
||||
pooler = PoolerIdentity()
|
||||
tensors = [torch.randn(128), torch.randn(256)]
|
||||
out = pooler(tensors)
|
||||
assert len(out) == 2
|
||||
for orig, result in zip(tensors, out):
|
||||
assert torch.equal(orig, result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PoolerNormalize
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestPoolerNormalize:
|
||||
def test_output_has_unit_norm(self):
|
||||
pooler = PoolerNormalize()
|
||||
x = torch.randn(4, 128)
|
||||
out = pooler(x)
|
||||
norms = torch.linalg.norm(out, dim=-1)
|
||||
assert torch.allclose(norms, torch.ones(4), atol=1e-5)
|
||||
|
||||
def test_single_vector(self):
|
||||
pooler = PoolerNormalize()
|
||||
x = torch.randn(1, 64)
|
||||
out = pooler(x)
|
||||
norm = torch.linalg.norm(out, dim=-1)
|
||||
assert torch.allclose(norm, torch.ones(1), atol=1e-5)
|
||||
|
||||
def test_forward_list(self):
|
||||
pooler = PoolerNormalize()
|
||||
tensors = [torch.randn(1, 64), torch.randn(1, 128)]
|
||||
out = pooler(tensors)
|
||||
for t in out:
|
||||
norm = torch.linalg.norm(t, dim=-1)
|
||||
assert torch.allclose(norm, torch.ones(1), atol=1e-5)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PoolerMultiLabelClassify
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestPoolerMultiLabelClassify:
|
||||
def test_output_in_zero_one(self):
|
||||
pooler = PoolerMultiLabelClassify()
|
||||
x = torch.randn(4, 10)
|
||||
out = pooler(x)
|
||||
assert (out >= 0).all() and (out <= 1).all()
|
||||
|
||||
def test_large_positive_maps_near_one(self):
|
||||
pooler = PoolerMultiLabelClassify()
|
||||
x = torch.full((1, 3), 100.0)
|
||||
out = pooler(x)
|
||||
assert torch.allclose(out, torch.ones(1, 3), atol=1e-4)
|
||||
|
||||
def test_large_negative_maps_near_zero(self):
|
||||
pooler = PoolerMultiLabelClassify()
|
||||
x = torch.full((1, 3), -100.0)
|
||||
out = pooler(x)
|
||||
assert torch.allclose(out, torch.zeros(1, 3), atol=1e-4)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PoolerClassify
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestPoolerClassify:
|
||||
def test_infers_from_shape_when_num_labels_none(self):
|
||||
pooler = PoolerClassify(num_labels=None)
|
||||
assert pooler.num_labels is None
|
||||
x = torch.randn(2, 5)
|
||||
out = pooler(x)
|
||||
sums = out.sum(dim=-1)
|
||||
assert torch.allclose(sums, torch.ones(2), atol=1e-5)
|
||||
|
||||
def test_sigmoid_when_num_labels_lt_2(self):
|
||||
pooler = PoolerClassify(num_labels=1)
|
||||
x = torch.zeros(1, 1)
|
||||
out = pooler(x)
|
||||
assert torch.allclose(out, torch.tensor([[0.5]]), atol=1e-5)
|
||||
|
||||
def test_num_labels_zero_uses_sigmoid(self):
|
||||
pooler = PoolerClassify(num_labels=0)
|
||||
assert pooler.num_labels == 0
|
||||
x = torch.zeros(1, 3)
|
||||
out = pooler(x)
|
||||
assert torch.allclose(out, torch.full((1, 3), 0.5), atol=1e-5)
|
||||
|
||||
def test_num_labels_ge_2_uses_softmax(self):
|
||||
pooler = PoolerClassify(num_labels=4)
|
||||
assert pooler.num_labels == 4
|
||||
x = torch.randn(2, 4)
|
||||
out = pooler(x)
|
||||
sums = out.sum(dim=-1)
|
||||
assert torch.allclose(sums, torch.ones(2), atol=1e-5)
|
||||
|
||||
def test_default_num_labels_is_none(self):
|
||||
pooler = PoolerClassify()
|
||||
assert pooler.num_labels is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LambdaPoolerActivation
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestLambdaPoolerActivation:
|
||||
def test_applies_custom_fn(self):
|
||||
pooler = LambdaPoolerActivation(nn.ReLU())
|
||||
x = torch.tensor([[-1.0, 2.0, -3.0]])
|
||||
out = pooler(x)
|
||||
expected = torch.tensor([[0.0, 2.0, 0.0]])
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
def test_forward_list(self):
|
||||
pooler = LambdaPoolerActivation(nn.ReLU())
|
||||
tensors = [torch.tensor([-1.0, 2.0]), torch.tensor([3.0, -4.0])]
|
||||
out = pooler(tensors)
|
||||
assert torch.equal(out[0], torch.tensor([0.0, 2.0]))
|
||||
assert torch.equal(out[1], torch.tensor([3.0, 0.0]))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_act_fn factory
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestGetActFn:
|
||||
@staticmethod
|
||||
def _make_config(**kwargs):
|
||||
return SimpleNamespace(**kwargs)
|
||||
|
||||
def test_regression(self):
|
||||
cfg = self._make_config(problem_type="regression")
|
||||
result = get_act_fn(cfg)
|
||||
assert isinstance(result, PoolerIdentity)
|
||||
|
||||
def test_single_label_classification(self):
|
||||
cfg = self._make_config(
|
||||
problem_type="single_label_classification", num_labels=3
|
||||
)
|
||||
result = get_act_fn(cfg)
|
||||
assert isinstance(result, PoolerClassify)
|
||||
assert result.num_labels == 3
|
||||
|
||||
def test_multi_label_classification(self):
|
||||
cfg = self._make_config(problem_type="multi_label_classification")
|
||||
result = get_act_fn(cfg)
|
||||
assert isinstance(result, PoolerMultiLabelClassify)
|
||||
|
||||
def test_sentence_transformers_activation(self):
|
||||
cfg = self._make_config(
|
||||
problem_type="",
|
||||
sentence_transformers={
|
||||
"activation_fn": "torch.nn.modules.activation.Sigmoid"
|
||||
},
|
||||
)
|
||||
result = get_act_fn(cfg)
|
||||
assert isinstance(result, PoolerClassify)
|
||||
|
||||
def test_sbert_activation(self):
|
||||
cfg = self._make_config(
|
||||
problem_type="",
|
||||
sbert_ce_default_activation_function=(
|
||||
"torch.nn.modules.activation.Sigmoid"
|
||||
),
|
||||
)
|
||||
result = get_act_fn(cfg)
|
||||
assert isinstance(result, PoolerClassify)
|
||||
|
||||
def test_default_fallback(self):
|
||||
cfg = self._make_config(problem_type="")
|
||||
result = get_act_fn(cfg)
|
||||
assert isinstance(result, PoolerClassify)
|
||||
|
||||
def test_sentence_transformers_takes_priority(self):
|
||||
cfg = self._make_config(
|
||||
problem_type="",
|
||||
sentence_transformers={"activation_fn": "torch.nn.modules.linear.Identity"},
|
||||
sbert_ce_default_activation_function=(
|
||||
"torch.nn.modules.activation.Sigmoid"
|
||||
),
|
||||
)
|
||||
result = get_act_fn(cfg)
|
||||
assert isinstance(result, PoolerIdentity)
|
||||
|
||||
def test_rejects_non_torch_activation(self):
|
||||
cfg = self._make_config(
|
||||
problem_type="",
|
||||
sentence_transformers={"activation_fn": "os.system"},
|
||||
)
|
||||
with pytest.raises(ValueError, match="restricted"):
|
||||
get_act_fn(cfg)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# resolve_classifier_act_fn
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestResolveClassifierActFn:
|
||||
def test_delegates_to_get_act_fn_when_none(self):
|
||||
model_config = SimpleNamespace(
|
||||
hf_config=SimpleNamespace(num_labels=3, problem_type="")
|
||||
)
|
||||
result = resolve_classifier_act_fn(model_config, act_fn=None)
|
||||
assert isinstance(result, PoolerClassify)
|
||||
assert result.num_labels == 3
|
||||
|
||||
def test_passes_through_provided_act_fn(self):
|
||||
custom = PoolerIdentity()
|
||||
result = resolve_classifier_act_fn(None, act_fn=custom)
|
||||
assert result is custom
|
||||
@@ -0,0 +1,481 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for sequence and token pooler head classes."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm.model_executor.layers.pooler.activations import PoolerNormalize
|
||||
from vllm.model_executor.layers.pooler.seqwise.heads import (
|
||||
ClassifierPoolerHead,
|
||||
EmbeddingPoolerHead,
|
||||
)
|
||||
from vllm.model_executor.layers.pooler.tokwise.heads import (
|
||||
TokenClassifierPoolerHead,
|
||||
TokenEmbeddingPoolerHead,
|
||||
)
|
||||
from vllm.pooling_params import PoolingParams
|
||||
from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
|
||||
|
||||
_HIDDEN = 16
|
||||
_BATCH = 3
|
||||
|
||||
|
||||
def _make_params(
|
||||
n: int,
|
||||
*,
|
||||
task: str = "embed",
|
||||
dimensions: int | None = None,
|
||||
use_activation: bool | None = None,
|
||||
) -> list[PoolingParams]:
|
||||
return [
|
||||
PoolingParams(task=task, dimensions=dimensions, use_activation=use_activation)
|
||||
for _ in range(n)
|
||||
]
|
||||
|
||||
|
||||
def _make_metadata(pooling_params: list[PoolingParams]) -> PoolingMetadata:
|
||||
n = len(pooling_params)
|
||||
return PoolingMetadata(
|
||||
prompt_lens=torch.ones(n, dtype=torch.long),
|
||||
prompt_token_ids=None,
|
||||
prompt_token_ids_cpu=None,
|
||||
pooling_params=pooling_params,
|
||||
pooling_states=[PoolingStates() for _ in range(n)],
|
||||
)
|
||||
|
||||
|
||||
def _linear(in_f: int, out_f: int) -> nn.Linear:
|
||||
torch.manual_seed(42)
|
||||
return nn.Linear(in_f, out_f, bias=False)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# EmbeddingPoolerHead
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestEmbeddingPoolerHead:
|
||||
def test_supported_tasks(self):
|
||||
head = EmbeddingPoolerHead()
|
||||
assert head.get_supported_tasks() == {"embed"}
|
||||
|
||||
def test_passthrough(self):
|
||||
head = EmbeddingPoolerHead()
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH))
|
||||
out = head(x, meta)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_head_dtype(self):
|
||||
head = EmbeddingPoolerHead(head_dtype=torch.float16)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH))
|
||||
out = head(x, meta)
|
||||
assert out.dtype == torch.float16
|
||||
|
||||
def test_projector(self):
|
||||
proj = _linear(_HIDDEN, 8)
|
||||
head = EmbeddingPoolerHead(projector=proj)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH))
|
||||
out = head(x, meta)
|
||||
assert out.shape == (_BATCH, 8)
|
||||
assert torch.allclose(out, proj(x))
|
||||
|
||||
def test_matryoshka_uniform(self):
|
||||
head = EmbeddingPoolerHead()
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
params = _make_params(_BATCH, dimensions=4)
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert out.shape == (_BATCH, 4)
|
||||
assert torch.equal(out, x[..., :4])
|
||||
|
||||
def test_matryoshka_mixed(self):
|
||||
head = EmbeddingPoolerHead()
|
||||
x = torch.randn(2, _HIDDEN)
|
||||
params = [
|
||||
PoolingParams(task="embed", dimensions=4),
|
||||
PoolingParams(task="embed", dimensions=8),
|
||||
]
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert isinstance(out, list)
|
||||
assert len(out) == 2
|
||||
assert out[0].shape[-1] == 4
|
||||
assert out[1].shape[-1] == 8
|
||||
|
||||
def test_matryoshka_mixed_with_none(self):
|
||||
head = EmbeddingPoolerHead()
|
||||
x = torch.randn(2, _HIDDEN)
|
||||
params = [
|
||||
PoolingParams(task="embed", dimensions=4),
|
||||
PoolingParams(task="embed", dimensions=None),
|
||||
]
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert isinstance(out, list)
|
||||
assert out[0].shape[-1] == 4
|
||||
assert torch.equal(out[1], x[1])
|
||||
|
||||
def test_activation_uniform_true(self):
|
||||
head = EmbeddingPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
params = _make_params(_BATCH, use_activation=True)
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
norms = torch.linalg.norm(out, dim=-1)
|
||||
assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
|
||||
|
||||
def test_activation_uniform_false(self):
|
||||
head = EmbeddingPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
params = _make_params(_BATCH, use_activation=False)
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_activation_mixed_flags(self):
|
||||
head = EmbeddingPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(2, _HIDDEN)
|
||||
params = [
|
||||
PoolingParams(task="embed", use_activation=True),
|
||||
PoolingParams(task="embed", use_activation=False),
|
||||
]
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert isinstance(out, list)
|
||||
norm_0 = torch.linalg.norm(out[0], dim=-1)
|
||||
assert torch.allclose(norm_0, torch.ones(1), atol=1e-5)
|
||||
assert torch.equal(out[1], x[1])
|
||||
|
||||
def test_list_input_gets_stacked(self):
|
||||
head = EmbeddingPoolerHead()
|
||||
tensors = [torch.randn(_HIDDEN) for _ in range(_BATCH)]
|
||||
meta = _make_metadata(_make_params(_BATCH))
|
||||
out = head(tensors, meta)
|
||||
assert out.shape == (_BATCH, _HIDDEN)
|
||||
expected = torch.stack(tensors)
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
def test_projector_then_matryoshka(self):
|
||||
proj = _linear(_HIDDEN, 8)
|
||||
head = EmbeddingPoolerHead(projector=proj)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
params = _make_params(_BATCH, dimensions=4)
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert out.shape == (_BATCH, 4)
|
||||
assert torch.equal(out, proj(x)[..., :4])
|
||||
|
||||
def test_matryoshka_then_activation(self):
|
||||
head = EmbeddingPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
params = _make_params(_BATCH, dimensions=4, use_activation=True)
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert out.shape == (_BATCH, 4)
|
||||
norms = torch.linalg.norm(out, dim=-1)
|
||||
assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
|
||||
|
||||
def test_empty_batch(self):
|
||||
head = EmbeddingPoolerHead()
|
||||
x = torch.randn(0, _HIDDEN)
|
||||
meta = _make_metadata([])
|
||||
out = head(x, meta)
|
||||
assert out.shape == (0, _HIDDEN)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ClassifierPoolerHead
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestClassifierPoolerHead:
|
||||
def test_supported_tasks(self):
|
||||
head = ClassifierPoolerHead()
|
||||
assert head.get_supported_tasks() == {"classify"}
|
||||
|
||||
def test_passthrough(self):
|
||||
head = ClassifierPoolerHead()
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH, task="classify"))
|
||||
out = head(x, meta)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_head_dtype(self):
|
||||
head = ClassifierPoolerHead(head_dtype=torch.float16)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH, task="classify"))
|
||||
out = head(x, meta)
|
||||
assert out.dtype == torch.float16
|
||||
|
||||
def test_classifier(self):
|
||||
clf = _linear(_HIDDEN, 3)
|
||||
head = ClassifierPoolerHead(classifier=clf)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH, task="classify"))
|
||||
out = head(x, meta)
|
||||
assert out.shape == (_BATCH, 3)
|
||||
assert torch.allclose(out, clf(x))
|
||||
|
||||
def test_logit_mean(self):
|
||||
head = ClassifierPoolerHead(logit_mean=2.0)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH, task="classify"))
|
||||
out = head(x, meta)
|
||||
assert torch.allclose(out, x - 2.0)
|
||||
|
||||
def test_logit_sigma(self):
|
||||
head = ClassifierPoolerHead(logit_sigma=0.5)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH, task="classify"))
|
||||
out = head(x, meta)
|
||||
assert torch.allclose(out, x / 0.5)
|
||||
|
||||
def test_platt_scaling_combined(self):
|
||||
head = ClassifierPoolerHead(logit_mean=1.0, logit_sigma=2.0)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH, task="classify"))
|
||||
out = head(x, meta)
|
||||
assert torch.allclose(out, (x - 1.0) / 2.0)
|
||||
|
||||
def test_activation_uniform_true(self):
|
||||
head = ClassifierPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
params = _make_params(_BATCH, task="classify", use_activation=True)
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
norms = torch.linalg.norm(out, dim=-1)
|
||||
assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
|
||||
|
||||
def test_activation_uniform_false(self):
|
||||
head = ClassifierPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
params = _make_params(_BATCH, task="classify", use_activation=False)
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_activation_mixed_flags(self):
|
||||
head = ClassifierPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(2, _HIDDEN)
|
||||
params = [
|
||||
PoolingParams(task="classify", use_activation=True),
|
||||
PoolingParams(task="classify", use_activation=False),
|
||||
]
|
||||
meta = _make_metadata(params)
|
||||
out = head(x, meta)
|
||||
assert isinstance(out, list)
|
||||
norm_0 = torch.linalg.norm(out[0], dim=-1)
|
||||
assert torch.allclose(norm_0, torch.ones(1), atol=1e-5)
|
||||
assert torch.equal(out[1], x[1])
|
||||
|
||||
def test_list_input_gets_stacked(self):
|
||||
head = ClassifierPoolerHead()
|
||||
tensors = [torch.randn(_HIDDEN) for _ in range(_BATCH)]
|
||||
meta = _make_metadata(_make_params(_BATCH, task="classify"))
|
||||
out = head(tensors, meta)
|
||||
assert out.shape == (_BATCH, _HIDDEN)
|
||||
expected = torch.stack(tensors)
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
def test_classifier_then_platt_scaling(self):
|
||||
clf = _linear(_HIDDEN, 3)
|
||||
head = ClassifierPoolerHead(classifier=clf, logit_mean=1.0, logit_sigma=2.0)
|
||||
x = torch.randn(_BATCH, _HIDDEN)
|
||||
meta = _make_metadata(_make_params(_BATCH, task="classify"))
|
||||
out = head(x, meta)
|
||||
expected = (clf(x) - 1.0) / 2.0
|
||||
assert torch.allclose(out, expected)
|
||||
|
||||
def test_empty_batch(self):
|
||||
head = ClassifierPoolerHead()
|
||||
x = torch.randn(0, _HIDDEN)
|
||||
meta = _make_metadata([])
|
||||
out = head(x, meta)
|
||||
assert out.shape == (0, _HIDDEN)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TokenEmbeddingPoolerHead
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestTokenEmbeddingPoolerHead:
|
||||
def test_supported_tasks(self):
|
||||
head = TokenEmbeddingPoolerHead()
|
||||
assert head.get_supported_tasks() == {"token_embed"}
|
||||
|
||||
def test_passthrough(self):
|
||||
head = TokenEmbeddingPoolerHead()
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_embed")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_none_chunked_prefill(self):
|
||||
head = TokenEmbeddingPoolerHead()
|
||||
param = PoolingParams(task="token_embed")
|
||||
out = head.forward_chunk(None, param)
|
||||
assert out is None
|
||||
|
||||
def test_head_dtype(self):
|
||||
head = TokenEmbeddingPoolerHead(head_dtype=torch.float16)
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_embed")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert out.dtype == torch.float16
|
||||
|
||||
def test_projector(self):
|
||||
proj = _linear(_HIDDEN, 8)
|
||||
head = TokenEmbeddingPoolerHead(projector=proj)
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_embed")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert out.shape == (5, 8)
|
||||
assert torch.allclose(out, proj(x))
|
||||
|
||||
def test_matryoshka_truncation(self):
|
||||
head = TokenEmbeddingPoolerHead()
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_embed", dimensions=4)
|
||||
out = head.forward_chunk(x, param)
|
||||
assert out.shape == (5, 4)
|
||||
assert torch.equal(out, x[..., :4])
|
||||
|
||||
def test_activation_true(self):
|
||||
head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_embed", use_activation=True)
|
||||
out = head.forward_chunk(x, param)
|
||||
norms = torch.linalg.norm(out, dim=-1)
|
||||
assert torch.allclose(norms, torch.ones(5), atol=1e-5)
|
||||
|
||||
def test_activation_false(self):
|
||||
head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_embed", use_activation=False)
|
||||
out = head.forward_chunk(x, param)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_projector_then_matryoshka(self):
|
||||
proj = _linear(_HIDDEN, 8)
|
||||
head = TokenEmbeddingPoolerHead(projector=proj)
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_embed", dimensions=4)
|
||||
out = head.forward_chunk(x, param)
|
||||
assert out.shape == (5, 4)
|
||||
assert torch.equal(out, proj(x)[..., :4])
|
||||
|
||||
def test_matryoshka_then_activation(self):
|
||||
head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_embed", dimensions=4, use_activation=True)
|
||||
out = head.forward_chunk(x, param)
|
||||
assert out.shape == (5, 4)
|
||||
norms = torch.linalg.norm(out, dim=-1)
|
||||
assert torch.allclose(norms, torch.ones(5), atol=1e-5)
|
||||
|
||||
def test_forward_mixed_batch_chunked_prefill(self):
|
||||
head = TokenEmbeddingPoolerHead()
|
||||
pooled_data = [torch.randn(5, _HIDDEN), None, torch.randn(3, _HIDDEN)]
|
||||
params = _make_params(3, task="token_embed")
|
||||
meta = _make_metadata(params)
|
||||
out = head(pooled_data, meta)
|
||||
assert len(out) == 3
|
||||
assert torch.equal(out[0], pooled_data[0])
|
||||
assert out[1] is None
|
||||
assert torch.equal(out[2], pooled_data[2])
|
||||
|
||||
def test_forward_empty_batch(self):
|
||||
head = TokenEmbeddingPoolerHead()
|
||||
meta = _make_metadata([])
|
||||
out = head([], meta)
|
||||
assert out == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TokenClassifierPoolerHead
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestTokenClassifierPoolerHead:
|
||||
def test_supported_tasks(self):
|
||||
head = TokenClassifierPoolerHead()
|
||||
assert head.get_supported_tasks() == {"token_classify"}
|
||||
|
||||
def test_passthrough(self):
|
||||
head = TokenClassifierPoolerHead()
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_classify")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_none_chunked_prefill(self):
|
||||
head = TokenClassifierPoolerHead()
|
||||
param = PoolingParams(task="token_classify")
|
||||
out = head.forward_chunk(None, param)
|
||||
assert out is None
|
||||
|
||||
def test_head_dtype(self):
|
||||
head = TokenClassifierPoolerHead(head_dtype=torch.float16)
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_classify")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert out.dtype == torch.float16
|
||||
|
||||
def test_classifier(self):
|
||||
clf = _linear(_HIDDEN, 3)
|
||||
head = TokenClassifierPoolerHead(classifier=clf)
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_classify")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert out.shape == (5, 3)
|
||||
assert torch.allclose(out, clf(x))
|
||||
|
||||
def test_logit_mean(self):
|
||||
head = TokenClassifierPoolerHead(logit_mean=2.0)
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_classify")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert torch.allclose(out, x - 2.0)
|
||||
|
||||
def test_logit_sigma(self):
|
||||
head = TokenClassifierPoolerHead(logit_sigma=0.5)
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_classify")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert torch.allclose(out, x / 0.5)
|
||||
|
||||
def test_platt_scaling_combined(self):
|
||||
head = TokenClassifierPoolerHead(logit_mean=1.0, logit_sigma=2.0)
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_classify")
|
||||
out = head.forward_chunk(x, param)
|
||||
assert torch.allclose(out, (x - 1.0) / 2.0)
|
||||
|
||||
def test_activation_true(self):
|
||||
head = TokenClassifierPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_classify", use_activation=True)
|
||||
out = head.forward_chunk(x, param)
|
||||
norms = torch.linalg.norm(out, dim=-1)
|
||||
assert torch.allclose(norms, torch.ones(5), atol=1e-5)
|
||||
|
||||
def test_activation_false(self):
|
||||
head = TokenClassifierPoolerHead(activation=PoolerNormalize())
|
||||
x = torch.randn(5, _HIDDEN)
|
||||
param = PoolingParams(task="token_classify", use_activation=False)
|
||||
out = head.forward_chunk(x, param)
|
||||
assert torch.equal(out, x)
|
||||
|
||||
def test_forward_mixed_batch_chunked_prefill(self):
|
||||
head = TokenClassifierPoolerHead()
|
||||
pooled_data = [torch.randn(5, _HIDDEN), None, torch.randn(3, _HIDDEN)]
|
||||
params = _make_params(3, task="token_classify")
|
||||
meta = _make_metadata(params)
|
||||
out = head(pooled_data, meta)
|
||||
assert len(out) == 3
|
||||
assert torch.equal(out[0], pooled_data[0])
|
||||
assert out[1] is None
|
||||
assert torch.equal(out[2], pooled_data[2])
|
||||
|
||||
def test_forward_empty_batch(self):
|
||||
head = TokenClassifierPoolerHead()
|
||||
meta = _make_metadata([])
|
||||
out = head([], meta)
|
||||
assert out == []
|
||||
@@ -0,0 +1,499 @@
|
||||
# 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 unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.pooler.seqwise.methods import (
|
||||
CLSPool,
|
||||
LastPool,
|
||||
MeanPool,
|
||||
get_seq_pooling_method,
|
||||
)
|
||||
from vllm.model_executor.layers.pooler.tokwise.methods import (
|
||||
AllPool,
|
||||
StepPool,
|
||||
get_tok_pooling_method,
|
||||
)
|
||||
from vllm.pooling_params import PoolingParams
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
def _make_metadata(
|
||||
prompt_lens: list[int],
|
||||
*,
|
||||
tasks: list[str] | 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,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeVllmConfig:
|
||||
scheduler_config: _FakeSchedulerConfig
|
||||
|
||||
|
||||
class TestAllPool:
|
||||
@staticmethod
|
||||
def _make_all_pool(*, chunked: bool = False) -> AllPool:
|
||||
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 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
|
||||
|
||||
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([chunk1, chunk2], dim=0)
|
||||
assert torch.equal(out2[0], expected)
|
||||
|
||||
def test_chunked_prefill_single_shot_matches_non_chunked(self):
|
||||
pooler = self._make_all_pool(chunked=True)
|
||||
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_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")
|
||||
@@ -0,0 +1,91 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.is_cuda():
|
||||
pytest.skip(
|
||||
"ROCm skinny GEMM tests are not supported on CUDA.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.model_executor.layers import utils
|
||||
|
||||
|
||||
def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
|
||||
x = torch.randn(1, 64, dtype=torch.float16)
|
||||
weight = torch.randn(128, 64, dtype=torch.float16)
|
||||
|
||||
monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
|
||||
monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
|
||||
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
|
||||
|
||||
wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
|
||||
llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
|
||||
|
||||
out = utils.rocm_unquantized_gemm_impl(x, weight, None)
|
||||
ref = torch.nn.functional.linear(x, weight, None)
|
||||
|
||||
wvsplitk_mock.assert_called_once()
|
||||
llmm1_mock.assert_not_called()
|
||||
assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch):
|
||||
# wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must
|
||||
# fall back to torch.nn.functional.linear.
|
||||
x = torch.randn(6, 64, dtype=torch.float16)
|
||||
weight = torch.randn(128, 64, dtype=torch.float16)
|
||||
|
||||
monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
|
||||
monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
|
||||
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
|
||||
|
||||
wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
|
||||
llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
|
||||
|
||||
out = utils.rocm_unquantized_gemm_impl(x, weight, None)
|
||||
ref = torch.nn.functional.linear(x, weight, None)
|
||||
|
||||
wvsplitk_mock.assert_not_called()
|
||||
llmm1_mock.assert_not_called()
|
||||
assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch):
|
||||
x = torch.randn(16, 1024, dtype=torch.float16)
|
||||
weight = torch.randn(256, 1024, dtype=torch.float16)
|
||||
|
||||
monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
|
||||
monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
|
||||
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
|
||||
|
||||
wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "wvSplitKrc", wvsplitkrc_mock)
|
||||
wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
|
||||
|
||||
out = utils.rocm_unquantized_gemm_impl(x, weight, None)
|
||||
ref = torch.nn.functional.linear(x, weight, None)
|
||||
|
||||
wvsplitkrc_mock.assert_called_once()
|
||||
wvsplitk_mock.assert_not_called()
|
||||
assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
|
||||
Reference in New Issue
Block a user