[None][perf] optimize encoder-decoder PyTorch performance

cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp

@@ -32,6 +32,7 @@
 #include "tensorrt_llm/runtime/torchView.h"
 
 #include <ATen/ATen.h>
+#include <algorithm>
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/chrono.h>
 #include <nanobind/stl/optional.h>
@@ -40,6 +41,7 @@
 #include <nanobind/stl/tuple.h>
 #include <nanobind/stl/unique_ptr.h>
 #include <nanobind/stl/vector.h>
+#include <numeric>
 #include <torch/extension.h>
 #include <tuple>
 
@@ -577,6 +579,153 @@ void initBindings(nb::module_& m)
         "Add new tokens to multiple LLM requests. The tokens vector should contain tokens for beam beam_idx of all "
         "requests in order.");
 
+    m.def(
+        "prepare_encoder_decoder_inputs",
+        [](std::vector<std::shared_ptr<tb::LlmRequest>> const& contextRequests,
+            std::vector<std::shared_ptr<tb::LlmRequest>> const& generationRequests, at::Tensor const& inputIds,
+            at::Tensor const& positionIds, at::Tensor const& sequenceLengths, at::Tensor const& promptLengths,
+            at::Tensor const& cachedTokenLengths, at::Tensor const& kvLengths, at::Tensor const& encoderKvLengths,
+            at::Tensor const& previousBatchIndices, SizeType32 positionIdOffset)
+        {
+            auto checkIntBuffer = [](at::Tensor const& tensor, char const* name)
+            {
+                TLLM_CHECK_WITH_INFO(tensor.device().is_cpu(), "%s must be a CPU tensor", name);
+                TLLM_CHECK_WITH_INFO(tensor.scalar_type() == at::kInt, "%s must have torch.int32 dtype", name);
+                TLLM_CHECK_WITH_INFO(tensor.is_contiguous(), "%s must be contiguous", name);
+            };
+            checkIntBuffer(inputIds, "input_ids");
+            checkIntBuffer(positionIds, "position_ids");
+            checkIntBuffer(sequenceLengths, "sequence_lengths");
+            checkIntBuffer(promptLengths, "prompt_lengths");
+            checkIntBuffer(cachedTokenLengths, "cached_token_lengths");
+            checkIntBuffer(kvLengths, "kv_lengths");
+            checkIntBuffer(encoderKvLengths, "encoder_kv_lengths");
+            checkIntBuffer(previousBatchIndices, "previous_batch_indices");
+
+            auto const numSequences = contextRequests.size() + generationRequests.size();
+            TLLM_CHECK_WITH_INFO(sequenceLengths.numel() >= static_cast<int64_t>(numSequences),
+                "sequence_lengths capacity is smaller than the batch");
+            TLLM_CHECK_WITH_INFO(promptLengths.numel() >= static_cast<int64_t>(numSequences),
+                "prompt_lengths capacity is smaller than the batch");
+            TLLM_CHECK_WITH_INFO(cachedTokenLengths.numel() >= static_cast<int64_t>(numSequences),
+                "cached_token_lengths capacity is smaller than the batch");
+            TLLM_CHECK_WITH_INFO(kvLengths.numel() >= static_cast<int64_t>(numSequences),
+                "kv_lengths capacity is smaller than the batch");
+            TLLM_CHECK_WITH_INFO(encoderKvLengths.numel() >= static_cast<int64_t>(numSequences),
+                "encoder_kv_lengths capacity is smaller than the batch");
+            TLLM_CHECK_WITH_INFO(previousBatchIndices.numel() >= static_cast<int64_t>(generationRequests.size()),
+                "previous_batch_indices capacity is smaller than the generation batch");
+
+            auto* inputIdsPtr = inputIds.data_ptr<SizeType32>();
+            auto* positionIdsPtr = positionIds.data_ptr<SizeType32>();
+            auto* sequenceLengthsPtr = sequenceLengths.data_ptr<SizeType32>();
+            auto* promptLengthsPtr = promptLengths.data_ptr<SizeType32>();
+            auto* cachedTokenLengthsPtr = cachedTokenLengths.data_ptr<SizeType32>();
+            auto* kvLengthsPtr = kvLengths.data_ptr<SizeType32>();
+            auto* encoderKvLengthsPtr = encoderKvLengths.data_ptr<SizeType32>();
+            auto* previousBatchIndicesPtr = previousBatchIndices.data_ptr<SizeType32>();
+
+            std::vector<tb::LlmRequest::RequestIdType> requestIds;
+            std::vector<SizeType32> encoderSequenceLengths;
+            std::vector<SizeType32> encoderCachedTokenLengths;
+            requestIds.reserve(numSequences);
+            encoderSequenceLengths.reserve(numSequences);
+            encoderCachedTokenLengths.reserve(numSequences);
+
+            SizeType32 numTokens{0};
+            SizeType32 numContextTokens{0};
+            SizeType32 numPreviousBatchRequests{0};
+            SizeType32 cachedKvTokens{0};
+            SizeType32 contextKvTokens{0};
+            SizeType32 generationKvTokens{0};
+            SizeType32 maxKvLength{0};
+            SizeType32 contextEncoderKvTokens{0};
+            SizeType32 generationEncoderKvTokens{0};
+            SizeType32 maxEncoderKvLength{0};
+            for (auto const& request : contextRequests)
+            {
+                auto const sequenceIdx = requestIds.size();
+                auto const begin = request->getContextCurrentPosition();
+                auto const chunkSize = request->getContextChunkSize();
+                auto const& tokens = request->getTokens(0);
+                TLLM_CHECK_WITH_INFO(begin + chunkSize <= static_cast<SizeType32>(tokens.size()),
+                    "Context chunk exceeds the request token count");
+                TLLM_CHECK_WITH_INFO(inputIds.numel() >= static_cast<int64_t>(numTokens + chunkSize),
+                    "input_ids capacity is smaller than the packed context");
+                TLLM_CHECK_WITH_INFO(positionIds.numel() >= static_cast<int64_t>(numTokens + chunkSize),
+                    "position_ids capacity is smaller than the packed context");
+
+                std::copy_n(tokens.data() + begin, chunkSize, inputIdsPtr + numTokens);
+                std::iota(positionIdsPtr + numTokens, positionIdsPtr + numTokens + chunkSize, begin + positionIdOffset);
+                sequenceLengthsPtr[sequenceIdx] = chunkSize;
+                promptLengthsPtr[sequenceIdx] = chunkSize;
+                cachedTokenLengthsPtr[sequenceIdx] = begin;
+                auto const kvLength = begin + chunkSize;
+                kvLengthsPtr[sequenceIdx] = kvLength;
+                auto const encoderKvLength = request->getEncoderOutputLen();
+                encoderKvLengthsPtr[sequenceIdx] = encoderKvLength;
+                cachedKvTokens += begin;
+                contextKvTokens += kvLength;
+                maxKvLength = std::max(maxKvLength, kvLength);
+                contextEncoderKvTokens += encoderKvLength;
+                maxEncoderKvLength = std::max(maxEncoderKvLength, encoderKvLength);
+                numTokens += chunkSize;
+                numContextTokens += chunkSize;
+
+                requestIds.push_back(request->mRequestId);
+                encoderSequenceLengths.push_back(encoderKvLength);
+                encoderCachedTokenLengths.push_back(0);
+            }
+
+            bool sawDummyRequest{false};
+            for (auto const& request : generationRequests)
+            {
+                auto const sequenceIdx = requestIds.size();
+                auto const isDummy = request->isDummyRequest();
+                sawDummyRequest = sawDummyRequest || isDummy;
+                TLLM_CHECK_WITH_INFO(
+                    isDummy || !sawDummyRequest, "CUDA graph dummy requests must follow real generation requests");
+                TLLM_CHECK_WITH_INFO(
+                    positionIds.numel() > numTokens, "position_ids capacity is smaller than the packed batch");
+
+                auto const pastSeenTokens = request->getMaxBeamNumTokens() - (isDummy ? 1 : 0);
+                positionIdsPtr[numTokens] = pastSeenTokens + positionIdOffset;
+                sequenceLengthsPtr[sequenceIdx] = 1;
+                promptLengthsPtr[sequenceIdx] = request->mPromptLen;
+                cachedTokenLengthsPtr[sequenceIdx] = pastSeenTokens;
+                auto const kvLength = pastSeenTokens + 1;
+                kvLengthsPtr[sequenceIdx] = kvLength;
+                auto const encoderKvLength = request->getEncoderOutputLen();
+                encoderKvLengthsPtr[sequenceIdx] = encoderKvLength;
+                cachedKvTokens += pastSeenTokens;
+                generationKvTokens += kvLength;
+                maxKvLength = std::max(maxKvLength, kvLength);
+                generationEncoderKvTokens += encoderKvLength;
+                maxEncoderKvLength = std::max(maxEncoderKvLength, encoderKvLength);
+                ++numTokens;
+
+                if (!isDummy)
+                {
+                    TLLM_CHECK_WITH_INFO(
+                        request->mSeqSlot.has_value(), "A real generation request must have a sequence slot");
+                    previousBatchIndicesPtr[numPreviousBatchRequests++] = request->mSeqSlot.value();
+                }
+
+                requestIds.push_back(request->mRequestId);
+                encoderSequenceLengths.push_back(0);
+                encoderCachedTokenLengths.push_back(encoderKvLength);
+            }
+
+            return std::make_tuple(requestIds, encoderSequenceLengths, encoderCachedTokenLengths, numTokens,
+                numContextTokens, numPreviousBatchRequests, cachedKvTokens, contextKvTokens, generationKvTokens,
+                maxKvLength, contextEncoderKvTokens, generationEncoderKvTokens, maxEncoderKvLength);
+        },
+        nb::arg("context_requests"), nb::arg("generation_requests"), nb::arg("input_ids"), nb::arg("position_ids"),
+        nb::arg("sequence_lengths"), nb::arg("prompt_lengths"), nb::arg("cached_token_lengths"), nb::arg("kv_lengths"),
+        nb::arg("encoder_kv_lengths"), nb::arg("previous_batch_indices"), nb::arg("position_id_offset") = 0,
+        nb::call_guard<nb::gil_scoped_release>(),
+        "Prepare the persistent CPU input buffers for a simple encoder-decoder batch.");
+
     m.def(
         "make_decoding_batch_input",
         [](tb::DecoderInputBuffers& decoderInputBuffers, runtime::decoder::DecoderState& decoderState,