fix: use executor API

janhq · Aug 1, 2024 · ce07bdb · ce07bdb
1 parent 163b3c0
commit ce07bdb
Show file tree

Hide file tree

Showing 4 changed files with 292 additions and 229 deletions.
diff --git a/cpp/tensorrt_llm/cortex.tensorrt-llm/examples/server/server.cc b/cpp/tensorrt_llm/cortex.tensorrt-llm/examples/server/server.cc
@@ -52,7 +52,7 @@ int main(int argc, char** argv) {
         };
     resp.set_chunked_content_provider("text/event-stream",
                                       chunked_content_provider,
-                                      [](bool) { LOG_INFO << "Done"; });
+                                      [](bool) { });
   };
 
   const auto handle_load_model = [&](const httplib::Request& req,

diff --git a/cpp/tensorrt_llm/cortex.tensorrt-llm/src/models/load_model_request.h b/cpp/tensorrt_llm/cortex.tensorrt-llm/src/models/load_model_request.h
@@ -7,6 +7,7 @@ namespace tensorrtllm::model {
 struct LoadModelRequest {
     int ctx_len = 2048;
     int n_parallel = 1;
+    int batch_size = 1;
     std::string model_path;
     std::string user_prompt = "";
     std::string ai_prompt = "";
@@ -18,6 +19,7 @@ inline LoadModelRequest fromJson(std::shared_ptr<Json::Value> json_body) {
   if (json_body) {
     request.ctx_len       = json_body->get("ctx_len", 2048).asInt();
     request.n_parallel    = json_body->get("n_parallel", 1).asInt();
+    request.batch_size    = json_body->get("batch_size", 1).asInt();
     request.model_path   = json_body->get("model_path", "").asString();
     request.user_prompt   = json_body->get("user_prompt", "").asString();
     request.ai_prompt     = json_body->get("ai_prompt", "").asString();