From 2472d0049a4b5c46a53e6eb1934228af2a08750d Mon Sep 17 00:00:00 2001
From: jioffe502 <jioffe@nvidia.com>
Date: Thu, 6 Aug 2026 23:34:33 +0000
Subject: [PATCH] perf(rerank): batch vLLM pair scoring across queries

vLLM already memory-schedules aligned N-to-N score inputs. Submit all query-document pairs in one call so candidates from multiple queries share the engine queue instead of paying one synchronous score round trip per query.
---
 .../models/local/nemotron_rerank_vl_v2.py     | 41 ++++++++-----------
 .../tests/test_nemotron_rerank_vl_v2.py       | 35 +++++++++++++++-
 2 files changed, 50 insertions(+), 26 deletions(-)

diff --git a/nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py b/nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py
index 6eb7d0c1a9..76fa8363eb 100644
--- a/nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py
+++ b/nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py
@@ -297,27 +297,20 @@ def score_pairs(
         if not pairs:
             return []
 
-        # Group consecutive pairs by query so each (query, [docs...]) goes to
-        # vLLM in a single batched score() call instead of one round-trip per pair.
-        scores: List[float] = [0.0] * len(pairs)
-        i = 0
-        while i < len(pairs):
-            q = pairs[i][0]
-            j = i
-            doc_inputs: list[Any] = []
-            while j < len(pairs) and pairs[j][0] == q:
-                _, d = pairs[j]
-                img = images_b64[j] if (images_b64 is not None and j < len(images_b64)) else None
-                d = self._truncate_doc_text(q, d, has_image=bool(img))
-                doc_inputs.append(self._build_document(d, img))
-                j += 1
-            outputs = self._llm.score(
-                q,
-                doc_inputs,
-                chat_template=SCORE_TEMPLATE,
-            )
-            for k, out in enumerate(outputs):
-                scores[i + k] = out.outputs.score
-            i = j
-
-        return scores
+        query_inputs: list[str] = []
+        doc_inputs: list[Any] = []
+        for index, (query, document) in enumerate(pairs):
+            image = images_b64[index] if images_b64 is not None and index < len(images_b64) else None
+            document = self._truncate_doc_text(query, document, has_image=bool(image))
+            query_inputs.append(query)
+            doc_inputs.append(self._build_document(document, image))
+
+        # vLLM schedules this aligned N-to-N list within its memory limit. A single
+        # call lets it combine candidates from multiple queries in one engine queue.
+        outputs = self._llm.score(
+            query_inputs,
+            doc_inputs,
+            use_tqdm=False,
+            chat_template=SCORE_TEMPLATE,
+        )
+        return [output.outputs.score for output in outputs]
diff --git a/nemo_retriever/tests/test_nemotron_rerank_vl_v2.py b/nemo_retriever/tests/test_nemotron_rerank_vl_v2.py
index 24f47831b9..b27bef1f9f 100644
--- a/nemo_retriever/tests/test_nemotron_rerank_vl_v2.py
+++ b/nemo_retriever/tests/test_nemotron_rerank_vl_v2.py
@@ -162,17 +162,48 @@ def test_score_with_images(self, reranker):
         assert docs[1] == "Paris is..."
 
     def test_score_pairs_text_only(self, reranker):
+        from nemo_retriever.models.local.nemotron_rerank_vl_v2 import SCORE_TEMPLATE
+
         out1 = MagicMock()
         out1.outputs.score = 1.0
         out2 = MagicMock()
         out2.outputs.score = 2.0
 
-        reranker._llm.score.side_effect = [[out1], [out2]]
+        reranker._llm.score.return_value = [out1, out2]
 
         scores = reranker.score_pairs([("q1", "d1"), ("q2", "d2")])
 
         assert scores == [1.0, 2.0]
-        assert reranker._llm.score.call_count == 2
+        reranker._llm.score.assert_called_once_with(
+            ["q1", "q2"],
+            ["d1", "d2"],
+            use_tqdm=False,
+            chat_template=SCORE_TEMPLATE,
+        )
+
+    def test_score_pairs_batches_multimodal_inputs(self, reranker):
+        from nemo_retriever.models.local.nemotron_rerank_vl_v2 import SCORE_TEMPLATE
+
+        outputs = [MagicMock(), MagicMock(), MagicMock()]
+        for index, output in enumerate(outputs):
+            output.outputs.score = float(index)
+        reranker._llm.score.return_value = outputs
+
+        scores = reranker.score_pairs(
+            [("q1", "d1"), ("q1", "d2"), ("q2", "d3")],
+            images_b64=["image-1", None, "image-3"],
+        )
+
+        assert scores == [0.0, 1.0, 2.0]
+        query_inputs, doc_inputs = reranker._llm.score.call_args.args
+        assert query_inputs == ["q1", "q1", "q2"]
+        assert doc_inputs[0]["content"][0]["image_url"]["url"].endswith("image-1")
+        assert doc_inputs[1] == "d2"
+        assert doc_inputs[2]["content"][0]["image_url"]["url"].endswith("image-3")
+        assert reranker._llm.score.call_args.kwargs == {
+            "use_tqdm": False,
+            "chat_template": SCORE_TEMPLATE,
+        }
 
     def test_score_chat_template_passed(self, reranker):
         from nemo_retriever.models.local.nemotron_rerank_vl_v2 import SCORE_TEMPLATE