fixup! [nvbugs/5274894] fix: Sort requests for functional correctness and performance (adapted from NVIDIA#4608)

Signed-off-by: Robin Kobus <[email protected]>

Author: Funatiq

cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp

@@ -309,7 +309,7 @@ std::tuple<RequestVector, RequestVector> MicroBatchScheduler::operator()(Request
         }
     }
 
-    utils::sortRequests(contextRequests, generationRequests);
+    utils::sortRequests(contextRequests, generationRequests, !allContextRequestsFit);
 
     TLLM_LOG_DEBUG(
         "batchSize (num ctx/enc requests + num gen requests): %u", contextRequests.size() + generationRequests.size());

cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp

@@ -39,25 +39,31 @@ TensorPtr collectRequestIds(RequestVector const& contextRequests, RequestVector
     return requestIds;
 }
 
-void sortRequests(RequestVector& contextRequests, RequestVector& generationRequests)
+void sortRequests(RequestVector& contextRequests, RequestVector& generationRequests, bool chunksPresent)
 {
     TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__);
 
-    // Move context requests that reached the last context chunk to the end of the vector.
-    // This order is required for moveFinishedContextRequestsToGeneration.
-    auto firstFinished = std::partition(contextRequests.begin(), contextRequests.end(),
-        [](auto const& llmReq) { return !llmReq->isLastContextChunk(); });
-
     auto sortByLoraId = [](RequestVector::iterator begin, RequestVector::iterator end)
     {
         std::sort(
             begin, end, [](auto const& lhs, auto const& rhs) { return lhs->getLoraTaskId() < rhs->getLoraTaskId(); });
     };
 
-    // Sort context requests by lora task id, but keep finished requests separate.
-    sortByLoraId(contextRequests.begin(), firstFinished);
-    sortByLoraId(firstFinished, contextRequests.end());
-    // Sort generation requests by lora task id.
+    if (chunksPresent)
+    {
+        // Move context requests that reached the last context chunk to the end of the vector.
+        // This order is required for moveFinishedContextRequestsToGeneration.
+        auto firstFinished = std::partition(contextRequests.begin(), contextRequests.end(),
+            [](auto const& llmReq) { return !llmReq->isLastContextChunk(); });
+
+        // Sort context requests by lora task id, but keep finished requests separate.
+        sortByLoraId(contextRequests.begin(), firstFinished);
+        sortByLoraId(firstFinished, contextRequests.end());
+    }
+    else
+    {
+        sortByLoraId(contextRequests.begin(), contextRequests.end());
+    }
     sortByLoraId(generationRequests.begin(), generationRequests.end());
 
     TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__);

cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h

@@ -40,7 +40,8 @@ TensorPtr collectRequestIds(RequestVector const& contextRequests, RequestVector
 //!          Sort requests by lora task id for performance.
 //! @param contextRequests The context requests.
 //! @param generationRequests The generation requests.
-void sortRequests(RequestVector& contextRequests, RequestVector& generationRequests);
+//! @param chunksPresent Whether context chunks are present.
+void sortRequests(RequestVector& contextRequests, RequestVector& generationRequests, bool chunksPresent);
 
 //! @brief Move finished context requests to generation requests.
 //! @details This function assumes that the context requests are sorted so that requests with isLastContextChunk() are