fix(reranker): avoid re-computing the first batch in predict()'s batch-size probe

Existing Issue:
KaLMReranker.predict() runs a full forward pass on the first batch twice: once as a throwaway OOM-size probe, and once again for the real computation. The probe's result is computed but never used. This roughly doubles reranking latency for any call where the total number of documents fits within a single batch (the common case, since batch_size defaults to 32).


Proposed Fix:
Keep the probe's result and reuse it as the first batch's scores, starting the real loop after the first batch instead of from the beginning.

kalm_reranker.py

@@ -282,10 +282,12 @@ class KaLMReranker:
         )
         sorted_pairs = [validated_pairs[index] for index in length_sorted_indices]
 
+
         tested_batch_size = effective_batch_size
+        first_batch_scores: List[float] = []
         while tested_batch_size > 1:
             try:
-                self._predict_batch(
+                first_batch_scores = self._predict_batch(
                     sorted_pairs[: min(len(sorted_pairs), tested_batch_size)],
                     effective_instruction,
                 )
@@ -295,9 +297,9 @@ class KaLMReranker:
                     torch.cuda.empty_cache()
                 tested_batch_size = max(1, tested_batch_size * 3 // 4)
 
-        sorted_scores: List[float] = []
+        sorted_scores: List[float] = list(first_batch_scores)
         try:
-            for start in range(0, len(sorted_pairs), tested_batch_size):
+            for start in range(tested_batch_size, len(sorted_pairs), tested_batch_size):
                 sorted_scores.extend(
                     self._predict_batch(
                         sorted_pairs[start : start + tested_batch_size],