From a6f359d1033945247179bc28216c204152e99fa7 Mon Sep 17 00:00:00 2001 From: Tejas Dharamsi Date: Sun, 10 May 2026 21:48:37 -0700 Subject: [PATCH] [Rerank] Use heapq.nlargest for top_n to avoid full sort (#24871) --- .../sglang/srt/entrypoints/openai/serving_rerank.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/serving_rerank.py b/python/sglang/srt/entrypoints/openai/serving_rerank.py index 57f5af4ca..f5d6e4743 100644 --- a/python/sglang/srt/entrypoints/openai/serving_rerank.py +++ b/python/sglang/srt/entrypoints/openai/serving_rerank.py @@ -1,3 +1,4 @@ +import heapq import logging from typing import Any, Dict, List, Optional, Union @@ -593,11 +594,11 @@ class OpenAIServingRerank(OpenAIServingBase): ) ) - # Sort by score in descending order (highest relevance first) + # When top_n is set, nlargest avoids fully sorting the candidate list + # (O(N log top_n) vs O(N log N)) — meaningful for large rerank batches. + # Validator (V1RerankReqInput.validate_top_n) guarantees top_n >= 1. + if request.top_n is not None: + return heapq.nlargest(request.top_n, responses, key=lambda x: x.score) + responses.sort(key=lambda x: x.score, reverse=True) - - # Apply top_n limit if specified - if request.top_n is not None and request.top_n > 0: - responses = responses[: request.top_n] - return responses