| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from unittest.mock import MagicMock, patch |
| 5 | |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | |
| 11 | from vllm.model_executor.models.bert import ( |
| 12 | BertMLMHead, |
| 13 | SPLADESparsePooler, |
| 14 | ) |
| 15 | from vllm.pooling_params import PoolingParams |
| 16 | from vllm.v1.pool.late_interaction_runner import LateInteractionRunner |
| 17 | from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates |
| 18 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 19 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner |
| 20 | from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner |
| 21 | from vllm.v1.worker.gpu.states import RequestState |
| 22 | |
| 23 | # --------------------------------------------------------------------- |
| 24 | # Functional test: SPLADE formula correctness (no HF download needed) |
| 25 | # --------------------------------------------------------------------- |
| 26 | |
| 27 | |
| 28 | @pytest.mark.parametrize("B,T,H,V", [(2, 3, 5, 7)]) |
| 29 | @torch.inference_mode |
| 30 | def test_splade_pooler_matches_reference_formula(B, T, H, V): |
| 31 | """Ensure SPLADESparsePooler forward() matches the mathematical formula: |
| 32 | log1p(relu(logits)) -> max over sequence length (after masking).""" |
| 33 | torch.manual_seed(0) |
| 34 | |
| 35 | # Prepare [B] sequences of shape [T, H] |
| 36 | hs_list = [torch.randn(T, H) for _ in range(B)] |
| 37 | hs_tenser = torch.cat(hs_list) |
| 38 | |
| 39 | # Simulate PoolingMetadata (only required fields) |
| 40 | prompt_lens = [T, T - 1] |
| 41 | prompt_lens_tenser = torch.tensor(prompt_lens, dtype=torch.int32) |
| 42 | token_ids = torch.tensor( |
| 43 | [ |
| 44 | [101, 5, 102], # Batch 0: [CLS], token, [SEP] |
| 45 | [101, 6, 6], # Batch 1: [CLS], token, token (last token ignored) |
| 46 | ], |
| 47 | dtype=torch.long, |
| 48 | ) |
| 49 | meta = PoolingMetadata( |
| 50 | prompt_lens=prompt_lens_tenser, |
| 51 | prompt_token_ids=token_ids, |
| 52 | prompt_token_ids_cpu=token_ids, |
| 53 | pooling_params=[PoolingParams(task="embed")] * B, |
| 54 | pooling_states=[PoolingStates() for _ in range(B)], |
| 55 | ) |
| 56 | |
| 57 | # MLM head (prefer BertMLMHead, fallback to Linear if unavailable) |
| 58 | try: |
| 59 | mlm_head = BertMLMHead(hidden_size=H, vocab_size=V, layer_norm_eps=1e-12) |
| 60 | except Exception: |
| 61 | mlm_head = nn.Linear(H, V, bias=True) |
| 62 | |
| 63 | # Forward pass through SPLADE pooler |
| 64 | pooler = SPLADESparsePooler(mlm_head=mlm_head, pooling="max", remove_cls_sep=True) |
| 65 | pooled = pooler(hidden_states=hs_tenser, pooling_metadata=meta) # list of [V] |
| 66 | |
| 67 | # Basic output checks |
| 68 | assert isinstance(pooled, torch.Tensor) and len(pooled) == B |
| 69 | for vec in pooled: |
| 70 | assert vec.shape == (V,) |
| 71 | assert torch.isfinite(vec).all() |
| 72 | assert (vec >= 0).all(), "SPLADE outputs must be non-negative." |
| 73 | |
| 74 | # Reference implementation for comparison |
| 75 | def ref_one(hs: torch.Tensor, L: int, tid_row: torch.Tensor) -> torch.Tensor: |
| 76 | keep = torch.ones(L, dtype=torch.bool) |
| 77 | if L > 0 and tid_row[0].item() == 101: # remove CLS |
| 78 | keep[0] = False |
| 79 | if L > 0 and tid_row[L - 1].item() == 102: # remove SEP |
| 80 | keep[L - 1] = False |
| 81 | |
| 82 | valid = hs[:L][keep[:L]] |
| 83 | if valid.numel() == 0: |
| 84 | return torch.zeros(V, dtype=torch.float32) |
| 85 | |
| 86 | logits = mlm_head(valid) # [L', V] |
| 87 | scores = torch.log1p(torch.relu(logits)) # [L', V] |
| 88 | return scores.max(dim=0).values.to(torch.float32) |
| 89 | |
| 90 | torch.testing.assert_close( |
| 91 | pooled[0], |
| 92 | ref_one(hs_list[0], prompt_lens[0], token_ids[0]), |
| 93 | rtol=1e-4, |
| 94 | atol=1e-4, |
| 95 | ) |
| 96 | torch.testing.assert_close( |
| 97 | pooled[1], |
| 98 | ref_one(hs_list[1], prompt_lens[1], token_ids[1]), |
| 99 | rtol=1e-4, |
| 100 | atol=1e-4, |
| 101 | ) |
| 102 | |
| 103 | |
| 104 | def test_pooling_runner_gathers_required_token_ids() -> None: |
| 105 | runner = PoolingRunner.__new__(PoolingRunner) |
| 106 | pooling_params = PoolingParams(task="embed", requires_token_ids=True) |
| 107 | runner.pooling_params = {1: pooling_params, 3: pooling_params} |
| 108 | runner.pooling_states = {1: PoolingStates(), 3: PoolingStates()} |
| 109 | runner.prompt_token_ids = { |
| 110 | 1: torch.tensor([101, 102]), |
| 111 | 3: torch.tensor([101, 11, 102]), |
| 112 | } |
| 113 | |
| 114 | input_batch = MagicMock(spec=InputBatch) |
| 115 | input_batch.idx_mapping_np = np.array([3, 1], dtype=np.int32) |
| 116 | input_batch.num_reqs = 2 |
| 117 | req_states = MagicMock(spec=RequestState) |
| 118 | req_states.prompt_len = MagicMock(np=np.array([0, 2, 0, 3], dtype=np.int32)) |
| 119 | metadata = runner._get_pooling_metadata( |
| 120 | input_batch, req_states, torch.device("cpu") |
| 121 | ) |
| 122 | |
| 123 | expected = torch.tensor([[101, 11, 102], [101, 102, 0]]) |
| 124 | assert metadata.prompt_token_ids_cpu is not None |
| 125 | assert metadata.prompt_token_ids is None |
| 126 | assert metadata.prompt_token_ids_cpu.is_pinned() is False |
| 127 | torch.testing.assert_close( |
| 128 | metadata.prompt_lens, torch.tensor([3, 2], dtype=torch.int32) |
| 129 | ) |
| 130 | torch.testing.assert_close(metadata.prompt_token_ids_cpu, expected) |
| 131 | |
| 132 | |
| 133 | def test_pooling_runner_stores_only_required_token_ids() -> None: |
| 134 | runner = PoolingRunner.__new__(PoolingRunner) |
| 135 | runner.model = MagicMock() |
| 136 | runner.supported_tasks = frozenset({"embed"}) |
| 137 | runner.pooling_params = {} |
| 138 | runner.pooling_states = {} |
| 139 | runner.prompt_token_ids = {} |
| 140 | runner.late_interaction_runner = MagicMock() |
| 141 | |
| 142 | runner.add_request("req-1", 1, PoolingParams(task="embed"), [101, 102]) |
| 143 | runner.add_request( |
| 144 | "req-2", |
| 145 | 2, |
| 146 | PoolingParams(task="embed", requires_token_ids=True), |
| 147 | [101, 11, 102], |
| 148 | ) |
| 149 | |
| 150 | assert 1 not in runner.prompt_token_ids |
| 151 | torch.testing.assert_close(runner.prompt_token_ids[2], torch.tensor([101, 11, 102])) |
| 152 | |
| 153 | |
| 154 | def test_pooling_runner_releases_aborted_late_interaction_doc() -> None: |
| 155 | runner = PoolingRunner.__new__(PoolingRunner) |
| 156 | runner.late_interaction_runner = LateInteractionRunner() |
| 157 | |
| 158 | query_key = "query-abort" |
| 159 | late_interaction_runner = runner.late_interaction_runner |
| 160 | late_interaction_runner._query_cache[query_key] = torch.ones(2, 4) |
| 161 | late_interaction_runner._query_uses[query_key] = 1 |
| 162 | late_interaction_runner._doc_query_keys["doc-req"] = query_key |
| 163 | |
| 164 | runner.on_requests_finished({"doc-req"}) |
| 165 | |
| 166 | assert not late_interaction_runner._query_cache |
| 167 | assert not late_interaction_runner._query_uses |
| 168 | assert not late_interaction_runner._doc_query_keys |
| 169 | |
| 170 | |
| 171 | def test_encoder_cache_reset_clears_late_interaction_state() -> None: |
| 172 | # Cached query embeddings are only invalidated by weight reloads, which |
| 173 | # reach the runner via reset_encoder_cache. Resetting the multi-modal |
| 174 | # cache is unrelated and must leave them intact. |
| 175 | runner = GPUModelRunner.__new__(GPUModelRunner) |
| 176 | runner.encoder_cache = MagicMock() |
| 177 | runner.pooling_runner = MagicMock() |
| 178 | |
| 179 | runner.reset_mm_cache() |
| 180 | |
| 181 | runner.encoder_cache.reset_mm_cache.assert_called_once_with() |
| 182 | runner.pooling_runner.clear.assert_not_called() |
| 183 | |
| 184 | runner.reset_encoder_cache() |
| 185 | |
| 186 | runner.encoder_cache.reset_encoder_cache.assert_called_once_with() |
| 187 | runner.pooling_runner.clear.assert_called_once_with() |
| 188 | |
| 189 | |
| 190 | def test_pooling_runner_rejects_unsupported_selected_task() -> None: |
| 191 | model = MagicMock() |
| 192 | model.pooler.get_supported_tasks.return_value = {"embed", "plugin"} |
| 193 | vllm_config = MagicMock() |
| 194 | vllm_config.scheduler_config.max_num_seqs = 2 |
| 195 | vllm_config.model_config.attn_type = "encoder_only" |
| 196 | vllm_config.model_config.get_pooling_task.return_value = "plugin" |
| 197 | |
| 198 | with ( |
| 199 | patch.object(PoolingRunner, "get_supported_tasks", return_value=["embed"]), |
| 200 | pytest.raises(ValueError, match="selects 'plugin'"), |
| 201 | ): |
| 202 | PoolingRunner(model, vllm_config) |
| 203 | |
| 204 | |
| 205 | def test_pooling_runner_supports_encoder_token_classification() -> None: |
| 206 | model = MagicMock() |
| 207 | model.pooler.get_supported_tasks.return_value = {"token_classify"} |
| 208 | vllm_config = MagicMock() |
| 209 | vllm_config.scheduler_config.max_num_seqs = 2 |
| 210 | vllm_config.model_config.attn_type = "encoder_only" |
| 211 | vllm_config.model_config.get_pooling_task.return_value = "token_classify" |
| 212 | |
| 213 | runner = PoolingRunner(model, vllm_config) |
| 214 | |
| 215 | assert runner.supported_tasks == {"token_classify"} |
| 216 | |
| 217 | |
| 218 | def test_pooling_runner_supports_encoder_token_embedding() -> None: |
| 219 | model = MagicMock() |
| 220 | model.pooler.get_supported_tasks.return_value = {"token_embed"} |
| 221 | vllm_config = MagicMock() |
| 222 | vllm_config.scheduler_config.max_num_seqs = 2 |
| 223 | vllm_config.model_config.attn_type = "encoder_only" |
| 224 | vllm_config.model_config.get_pooling_task.return_value = "token_embed" |
| 225 | |
| 226 | runner = PoolingRunner(model, vllm_config) |
| 227 | |
| 228 | assert runner.supported_tasks == {"token_embed"} |
| 229 | |
| 230 | |
| 231 | def test_pooling_runner_supports_decoder_token_classification() -> None: |
| 232 | model = MagicMock() |
| 233 | model.pooler.get_supported_tasks.return_value = {"token_classify"} |
| 234 | vllm_config = MagicMock() |
| 235 | vllm_config.scheduler_config.max_num_seqs = 2 |
| 236 | vllm_config.model_config.attn_type = "decoder" |
| 237 | vllm_config.model_config.get_pooling_task.return_value = "token_classify" |
| 238 | |
| 239 | runner = PoolingRunner(model, vllm_config) |
| 240 | |
| 241 | assert runner.supported_tasks == {"token_classify"} |
| 242 | |
| 243 | |
| 244 | def test_pooling_runner_keeps_decoder_token_classification() -> None: |
| 245 | model = MagicMock() |
| 246 | model.pooler.get_supported_tasks.return_value = {"embed", "token_classify"} |
| 247 | vllm_config = MagicMock() |
| 248 | vllm_config.scheduler_config.max_num_seqs = 2 |
| 249 | vllm_config.model_config.attn_type = "decoder" |
| 250 | vllm_config.model_config.get_pooling_task.return_value = "embed" |
| 251 | |
| 252 | runner = PoolingRunner(model, vllm_config) |
| 253 | |
| 254 | assert runner.supported_tasks == {"embed", "token_classify"} |
| 255 | |
| 256 | |
| 257 | def test_pooling_runner_keeps_decoder_token_embedding() -> None: |
| 258 | model = MagicMock() |
| 259 | model.pooler.get_supported_tasks.return_value = {"embed", "token_embed"} |
| 260 | vllm_config = MagicMock() |
| 261 | vllm_config.scheduler_config.max_num_seqs = 2 |
| 262 | vllm_config.model_config.attn_type = "decoder" |
| 263 | vllm_config.model_config.get_pooling_task.return_value = "embed" |
| 264 | |
| 265 | runner = PoolingRunner(model, vllm_config) |
| 266 | |
| 267 | assert runner.supported_tasks == {"embed", "token_embed"} |
| 268 | |
| 269 | |
| 270 | def test_pooling_runner_supports_embed_token_classification() -> None: |
| 271 | model = MagicMock() |
| 272 | model.pooler.get_supported_tasks.return_value = {"embed", "embed&token_classify"} |
| 273 | vllm_config = MagicMock() |
| 274 | vllm_config.scheduler_config.max_num_seqs = 2 |
| 275 | |
| 276 | vllm_config.model_config.attn_type = "encoder_only" |
| 277 | vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify" |
| 278 | assert PoolingRunner(model, vllm_config).supported_tasks == { |
| 279 | "embed", |
| 280 | "embed&token_classify", |
| 281 | } |