| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score |
| 8 | from vllm.pooling_params import LateInteractionParams, PoolingParams |
| 9 | from vllm.v1.pool.late_interaction import ( |
| 10 | LATE_INTERACTION_MODE_CACHE_QUERY, |
| 11 | build_late_interaction_doc_params, |
| 12 | build_late_interaction_query_params, |
| 13 | compute_maxsim_score_batched, |
| 14 | ) |
| 15 | from vllm.v1.pool.late_interaction_runner import LateInteractionRunner |
| 16 | |
| 17 | |
| 18 | def _make_pooling_params( |
| 19 | late_interaction_params: LateInteractionParams, |
| 20 | ) -> PoolingParams: |
| 21 | return PoolingParams( |
| 22 | task="token_embed", |
| 23 | late_interaction_params=late_interaction_params, |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | def test_postprocess_scores_and_releases_query_cache(): |
| 28 | runner = LateInteractionRunner() |
| 29 | query_key = "query-0" |
| 30 | query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32) |
| 31 | doc_emb = torch.tensor([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]], dtype=torch.float32) |
| 32 | |
| 33 | query_params = _make_pooling_params( |
| 34 | build_late_interaction_query_params(query_key=query_key, query_uses=1) |
| 35 | ) |
| 36 | query_output = runner.postprocess_pooler_output( |
| 37 | raw_pooler_output=[query_emb], |
| 38 | pooling_params=[query_params], |
| 39 | req_ids=["query-req"], |
| 40 | finished_mask=[True], |
| 41 | ) |
| 42 | assert isinstance(query_output, list) |
| 43 | assert query_output[0] is not None |
| 44 | assert query_output[0].shape == torch.Size([]) |
| 45 | |
| 46 | doc_params = _make_pooling_params( |
| 47 | build_late_interaction_doc_params(query_key=query_key) |
| 48 | ) |
| 49 | doc_output = runner.postprocess_pooler_output( |
| 50 | raw_pooler_output=[doc_emb], |
| 51 | pooling_params=[doc_params], |
| 52 | req_ids=["doc-req"], |
| 53 | finished_mask=[True], |
| 54 | ) |
| 55 | assert isinstance(doc_output, list) |
| 56 | assert doc_output[0] is not None |
| 57 | assert torch.allclose(doc_output[0], compute_maxsim_score(query_emb, doc_emb)) |
| 58 | |
| 59 | with pytest.raises(ValueError, match="query cache miss"): |
| 60 | runner.postprocess_pooler_output( |
| 61 | raw_pooler_output=[doc_emb], |
| 62 | pooling_params=[doc_params], |
| 63 | req_ids=["doc-req-2"], |
| 64 | finished_mask=[True], |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | def test_postprocess_scores_docs_in_batch(): |
| 69 | runner = LateInteractionRunner() |
| 70 | query_key = "query-batch" |
| 71 | query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32) |
| 72 | doc_emb_1 = torch.tensor([[1.0, 0.0], [0.5, 0.5]], dtype=torch.float32) |
| 73 | doc_emb_2 = torch.tensor([[0.0, 1.0], [0.3, 0.7], [1.0, 0.0]], dtype=torch.float32) |
| 74 | |
| 75 | query_params = _make_pooling_params( |
| 76 | build_late_interaction_query_params(query_key=query_key, query_uses=2) |
| 77 | ) |
| 78 | runner.postprocess_pooler_output( |
| 79 | raw_pooler_output=[query_emb], |
| 80 | pooling_params=[query_params], |
| 81 | req_ids=["query-req"], |
| 82 | finished_mask=[True], |
| 83 | ) |
| 84 | |
| 85 | doc_params = _make_pooling_params( |
| 86 | build_late_interaction_doc_params(query_key=query_key) |
| 87 | ) |
| 88 | doc_output = runner.postprocess_pooler_output( |
| 89 | raw_pooler_output=[doc_emb_1, doc_emb_2], |
| 90 | pooling_params=[doc_params, doc_params], |
| 91 | req_ids=["doc-req-1", "doc-req-2"], |
| 92 | finished_mask=[True, True], |
| 93 | ) |
| 94 | assert isinstance(doc_output, list) |
| 95 | assert doc_output[0] is not None |
| 96 | assert doc_output[1] is not None |
| 97 | assert torch.allclose(doc_output[0], compute_maxsim_score(query_emb, doc_emb_1)) |
| 98 | assert torch.allclose(doc_output[1], compute_maxsim_score(query_emb, doc_emb_2)) |
| 99 | |
| 100 | with pytest.raises(ValueError, match="query cache miss"): |
| 101 | runner.postprocess_pooler_output( |
| 102 | raw_pooler_output=[doc_emb_1], |
| 103 | pooling_params=[doc_params], |
| 104 | req_ids=["doc-req-3"], |
| 105 | finished_mask=[True], |
| 106 | ) |
| 107 | |
| 108 | |
| 109 | def test_finished_request_releases_unscored_doc_use(): |
| 110 | runner = LateInteractionRunner() |
| 111 | query_key = "query-cancel" |
| 112 | query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32) |
| 113 | doc_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32) |
| 114 | |
| 115 | query_params = _make_pooling_params( |
| 116 | build_late_interaction_query_params(query_key=query_key, query_uses=1) |
| 117 | ) |
| 118 | runner.postprocess_pooler_output( |
| 119 | raw_pooler_output=[query_emb], |
| 120 | pooling_params=[query_params], |
| 121 | req_ids=["query-req"], |
| 122 | finished_mask=[True], |
| 123 | ) |
| 124 | |
| 125 | doc_params = _make_pooling_params( |
| 126 | build_late_interaction_doc_params(query_key=query_key) |
| 127 | ) |
| 128 | runner.register_request("doc-req", doc_params) |
| 129 | runner.on_requests_finished({"doc-req"}) |
| 130 | |
| 131 | with pytest.raises(ValueError, match="query cache miss"): |
| 132 | runner.postprocess_pooler_output( |
| 133 | raw_pooler_output=[doc_emb], |
| 134 | pooling_params=[doc_params], |
| 135 | req_ids=["doc-req-retry"], |
| 136 | finished_mask=[True], |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | @pytest.mark.skip_global_cleanup |
| 141 | def test_release_queries_is_targeted_and_idempotent(): |
| 142 | runner = LateInteractionRunner() |
| 143 | query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32) |
| 144 | |
| 145 | for query_key in ("query-release", "query-keep"): |
| 146 | query_params = _make_pooling_params( |
| 147 | build_late_interaction_query_params(query_key=query_key, query_uses=1) |
| 148 | ) |
| 149 | runner.postprocess_pooler_output( |
| 150 | raw_pooler_output=[query_emb], |
| 151 | pooling_params=[query_params], |
| 152 | req_ids=[f"{query_key}-req"], |
| 153 | finished_mask=[True], |
| 154 | ) |
| 155 | |
| 156 | released_doc_params = _make_pooling_params( |
| 157 | build_late_interaction_doc_params(query_key="query-release") |
| 158 | ) |
| 159 | kept_doc_params = _make_pooling_params( |
| 160 | build_late_interaction_doc_params(query_key="query-keep") |
| 161 | ) |
| 162 | runner.register_request("released-doc", released_doc_params) |
| 163 | runner.register_request("kept-doc", kept_doc_params) |
| 164 | |
| 165 | runner.release_queries(["query-release"]) |
| 166 | runner.release_queries(["query-release"]) |
| 167 | |
| 168 | with pytest.raises(ValueError, match="query cache miss"): |
| 169 | runner.postprocess_pooler_output( |
| 170 | raw_pooler_output=[query_emb], |
| 171 | pooling_params=[released_doc_params], |
| 172 | req_ids=["released-doc"], |
| 173 | finished_mask=[True], |
| 174 | ) |
| 175 | |
| 176 | kept_output = runner.postprocess_pooler_output( |
| 177 | raw_pooler_output=[query_emb], |
| 178 | pooling_params=[kept_doc_params], |
| 179 | req_ids=["kept-doc"], |
| 180 | finished_mask=[True], |
| 181 | ) |
| 182 | assert isinstance(kept_output, list) |
| 183 | assert kept_output[0] is not None |
| 184 | |
| 185 | |
| 186 | def test_invalid_query_uses_raises(): |
| 187 | runner = LateInteractionRunner() |
| 188 | bad_meta = LateInteractionParams( |
| 189 | mode=LATE_INTERACTION_MODE_CACHE_QUERY, |
| 190 | query_key="query-bad", |
| 191 | ) |
| 192 | bad_meta.query_uses = "bad-int" |
| 193 | bad_query_params = _make_pooling_params(bad_meta) |
| 194 | |
| 195 | with pytest.raises(ValueError, match="must be an integer value"): |
| 196 | runner.postprocess_pooler_output( |
| 197 | raw_pooler_output=[torch.ones((2, 2), dtype=torch.float32)], |
| 198 | pooling_params=[bad_query_params], |
| 199 | req_ids=["query-req"], |
| 200 | finished_mask=[True], |
| 201 | ) |
| 202 | |
| 203 | |
| 204 | # --------------------------------------------------------------------------- |
| 205 | # Fused Triton scoring path (PR #40337). |
| 206 | # --------------------------------------------------------------------------- |
| 207 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") |
| 208 | def test_score_flash_matches_fp64_oracle(): |
| 209 | """Shared-query scoring through the fused kernel must match a per-pair |
| 210 | fp64 MaxSim oracle on ragged doc lengths, including length-1 docs.""" |
| 211 | torch.manual_seed(0) |
| 212 | d = 128 |
| 213 | q = torch.randn(32, d, device="cuda", dtype=torch.float16) |
| 214 | doc_lengths = [180, 37, 512, 1, 300, 64, 1030, 256] |
| 215 | docs = [ |
| 216 | torch.randn(ld, d, device="cuda", dtype=torch.float16) for ld in doc_lengths |
| 217 | ] |
| 218 | |
| 219 | runner = LateInteractionRunner() |
| 220 | assert runner._flash_enabled |
| 221 | scores = runner._score([q] * len(docs), docs) |
| 222 | assert runner._flash_enabled, "kernel path must not have fallen back" |
| 223 | |
| 224 | for i, doc in enumerate(docs): |
| 225 | ref = (q.double() @ doc.double().T).max(dim=1).values.sum().float() |
| 226 | torch.testing.assert_close( |
| 227 | scores[i].to(torch.float32), ref, atol=5e-2, rtol=1e-3 |
| 228 | ) |
| 229 | |
| 230 | |
| 231 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") |
| 232 | def test_score_multi_query_uses_reference_path(): |
| 233 | """Distinct queries per pair route to the reference scorer and agree |
| 234 | with it exactly (no kernel dispatch for the N:N pattern).""" |
| 235 | torch.manual_seed(1) |
| 236 | d = 128 |
| 237 | queries = [torch.randn(16, d, device="cuda") for _ in range(3)] |
| 238 | docs = [torch.randn(ld, d, device="cuda") for ld in (50, 3, 77)] |
| 239 | |
| 240 | runner = LateInteractionRunner() |
| 241 | scores = runner._score(queries, docs) |
| 242 | ref = compute_maxsim_score_batched(queries, docs) |
| 243 | for s, r in zip(scores, ref): |
| 244 | torch.testing.assert_close(s, r) |
| 245 | |
| 246 | |
| 247 | def test_score_flash_failure_disables_and_falls_back(monkeypatch): |
| 248 | """A kernel failure must serve the batch via the reference scorer and |
| 249 | disable the kernel path for the rest of the process.""" |
| 250 | runner = LateInteractionRunner() |
| 251 | runner._flash_enabled = True |
| 252 | |
| 253 | def _boom(query, docs): |
| 254 | raise RuntimeError("synthetic kernel failure") |
| 255 | |
| 256 | monkeypatch.setattr(runner, "_score_flash_shared_query", _boom) |
| 257 | q = torch.randn(8, 32) |
| 258 | docs = [torch.randn(5, 32), torch.randn(7, 32)] |
| 259 | # CPU tensors already skip the kernel; force the dispatch condition. |
| 260 | monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) |
| 261 | scores = runner._score([q, q], docs) |
| 262 | assert not runner._flash_enabled |
| 263 | ref = compute_maxsim_score_batched([q, q], docs) |
| 264 | for s, r in zip(scores, ref): |
| 265 | torch.testing.assert_close(s, r) |
| 266 | |
| 267 | |
| 268 | def test_disabled_flag_uses_reference_path(): |
| 269 | runner = LateInteractionRunner(enable_flash=False) |
| 270 | assert not runner._flash_enabled |
| 271 | |
| 272 | |
| 273 | # --------------------------------------------------------------------------- |
| 274 | # Autotune config pruning. |
| 275 | # --------------------------------------------------------------------------- |
| 276 | def _prune(monkeypatch, budget: int, named_args: dict): |
| 277 | from vllm.v1.pool.flash_maxsim import _common |
| 278 | |
| 279 | monkeypatch.setattr(_common, "_smem_budget", lambda: budget) |
| 280 | return _common._prune_configs(_common._get_configs(), named_args) |
| 281 | |
| 282 | |
| 283 | def _est(cfg, d_pad: int) -> int: |
| 284 | bq, bd = cfg.kwargs["BLOCK_Q"], cfg.kwargs["BLOCK_D"] |
| 285 | return (bq * d_pad + bd * d_pad) * 2 + bq * bd * 4 |
| 286 | |
| 287 | |
| 288 | def test_prune_configs_uses_padded_dim(monkeypatch): |
| 289 | """d=513 pads to 1024: survivors must fit the budget at d_pad, not at |
| 290 | the un-padded d (which would admit ~2x-oversized tiles).""" |
| 291 | budget = 166_912 # A100 opt-in |
| 292 | survivors = _prune(monkeypatch, budget, {"Lq": 1024, "d": 513, "d_pad": 1024}) |
| 293 | assert survivors |
| 294 | assert all(_est(c, 1024) <= budget for c in survivors) |
| 295 | |
| 296 | |
| 297 | def test_prune_configs_no_rejected_fallback(monkeypatch): |
| 298 | """When nothing fits, the fallback is the single smallest-footprint |
| 299 | config — never a slice of rejected configs.""" |
| 300 | from vllm.v1.pool.flash_maxsim import _common |
| 301 | |
| 302 | survivors = _prune(monkeypatch, 10_000, {"Lq": 1024, "d_pad": 1024}) |
| 303 | assert len(survivors) == 1 |
| 304 | smallest = min(_common._get_configs(), key=lambda c: _est(c, 1024)) |
| 305 | assert survivors[0].kwargs == smallest.kwargs |