feat(trtllm): separate request and response loop

The executor_status_looper spend CPU time polling at the number of
tokens. Because the function is protected by mutex inside, this also
interferes with the Executor.

Because now the TensorRtLlmBackendImpl is interior mutable, we can mark
it as `Send` and share it in multiple threads. Therefore, the loop can
be split into request and response parts, and we can await for tokens
instead of constantly polling.

Author: leejuyuu

backends/trtllm/csrc/backend.cpp

@@ -46,10 +46,6 @@ namespace huggingface::tgi::backends::trtllm {
     backend_t::backend_t(std::filesystem::path &engines_folder, std::filesystem::path &executor_worker_path)
             : workspace(engines_folder, executor_worker_path), executor_(executor_factory_initializer(workspace)) {}
 
-    size_t backend_t::num_tokens_ready() const noexcept {
-        return executor_.getNumResponsesReady();
-    }
-
     std::expected<request_id_t, backend_error_t>
     backend_t::submit(std::span<const token_id_t> token_ids, const generation_params_t g_params,
                       const sampling_params_t s_params) noexcept {

backends/trtllm/csrc/backend.hpp

@@ -175,13 +175,6 @@ namespace huggingface::tgi::backends::trtllm {
         submit(std::span<const token_id_t> token_ids, generation_params_t generation_params,
                sampling_params_t sampling_params) noexcept;
 
-        /**
-         * Query the number of tokens available across all in-flight generations
-         * @return
-         */
-        [[nodiscard("Pulling out the number of tokens")]]
-        size_t num_tokens_ready() const noexcept;
-
         /**
          * Pull out newly generated tokens from the executor
          * @return

backends/trtllm/csrc/ffi.hpp

@@ -120,8 +120,6 @@ namespace huggingface::tgi::backends::trtllm {
                   m_created_time {created_time}
         {}
 
-        size_t num_tokens_ready() const noexcept { return inner_.num_tokens_ready(); }
-
         request_id_t submit(
                 rust::Slice<const uint32_t> tokens,
                 uint32_t max_new_tokens,
@@ -153,27 +151,22 @@ namespace huggingface::tgi::backends::trtllm {
         }
 
         std::unique_ptr<std::vector<generation_step_t>> pull_tokens() const noexcept {
-            if (num_tokens_ready() > 0) [[likely]] {
-                const auto responses = inner_.pull_tokens();
+            const auto responses = inner_.pull_tokens();
 
-                SPDLOG_TRACE("[FFI] Successfully pulled out {:d} responses from executor", responses.size());
+            SPDLOG_TRACE("[FFI] Successfully pulled out {:d} responses from executor", responses.size());
 
-                auto f = [this](const tle::Response &r){
-                    return as_generation_step(r, m_created_time);
-                };
-                // Transform tle::Response to generation_step_t
+            auto f = [this](const tle::Response &r){
+                return as_generation_step(r, m_created_time);
+            };
+            auto steps = std::make_unique<std::vector<generation_step_t>>();
+            // Transform tle::Response to generation_step_t
 #ifdef __cpp_lib_ranges_to_container
-                auto steps = responses | std::views::transform(f) | std::ranges::to<std::vector>();
+            *steps = responses | std::views::transform(f) | std::ranges::to<std::vector>();
 #else
-                auto steps = std::vector<generation_step_t>();
-                steps.reserve(responses.size());
-                std::transform(responses.begin(), responses.end(), std::back_inserter(steps), f);
+            steps->reserve(responses.size());
+            std::transform(responses.begin(), responses.end(), std::back_inserter(steps), f);
 #endif
-                return std::make_unique<std::vector<generation_step_t>>(steps);
-
-            } else {
-                return std::make_unique<std::vector<generation_step_t>>();
-            }
+            return steps;
         }
 
         void cancel(request_id_t request_id) const noexcept {

backends/trtllm/src/lib.rs

@@ -80,8 +80,6 @@ mod ffi {
             executor_worker: &str,
         ) -> Result<UniquePtr<TensorRtLlmBackendImpl>>;
 
-        fn num_tokens_ready(self: &TensorRtLlmBackendImpl) -> usize;
-
         fn submit(
             self: &TensorRtLlmBackendImpl,
             tokens: &[u32],