[refactor] refactoring for sequence (#140)

this refaoctoring includes:
1> moved incremental decoding logic into a central place.
2> pass in capacity for sequence and move all GFlags into a central
place.
3> added unittests for speculative decoding for sequence.

Author: guocuimi

src/engine/batch.cpp

@@ -285,7 +285,7 @@ void Batch::process_sample_output(const SampleOutput& sample_output) {
       // add the next token to sequence
       const int32_t next_token_id =
           static_cast<int32_t>(next_tokens[output_idx++].item<int64_t>());
-      seq->append_new_token_id(next_token_id);
+      seq->append_token(next_token_id);
     }
     CHECK_EQ(output_idx, num_seqs);
   }
@@ -307,7 +307,7 @@ void Batch::process_validate_output(const torch::Tensor& accepted_ids) {
         ids.data_ptr<int64_t>(), static_cast<size_t>(ids.numel())};
 
     // validate the draft tokens with accepted tokens
-    seq->validate_token_ids(accepted_token_ids);
+    seq->validate_tokens(accepted_token_ids);
   }
   CHECK_EQ(output_idx, num_seqs);
 }

src/engine/batch_test.cpp

@@ -34,40 +34,37 @@ TEST(BatchTest, Basic) {
   auto block_0 = allocator.allocate();
   EXPECT_EQ(block_0.id(), 0);
 
-  SamplingParameter sampling_param;
-  sampling_param.frequency_penalty = 0.1;
-  StoppingCriteria stopping_criteria;
-  stopping_criteria.max_tokens = 20;
+  Sequence::Options options;
+  options.sampling_param.frequency_penalty = 0.1;
+  options.stopping_criteria.max_tokens = 20;
+  const size_t capacity = 100;
 
   // prepare sequences
   // sequence in prefill phase
-  Sequence seq1(/*token_ids=*/{1, 3, 5, 7, 5, 4, 3, 2, 1},
-                sampling_param,
-                stopping_criteria,
-                /*echo=*/false,
-                /*on_stream=*/nullptr);
+  Sequence seq1(/*prompt=*/"",
+                /*token_ids=*/{1, 3, 5, 7, 5, 4, 3, 2, 1},
+                capacity,
+                options);
   seq1.append_blocks(allocator.allocate(3));  // [1, 2, 3]
 
   // seq in decode phase
-  Sequence seq2(/*token_ids=*/{2, 4, 6, 8, 6, 4, 2},
-                sampling_param,
-                stopping_criteria,
-                /*echo=*/false,
-                /*on_stream=*/nullptr);
+  Sequence seq2(/*prompt=*/"",
+                /*token_ids=*/{2, 4, 6, 8, 6, 4, 2},
+                capacity,
+                options);
   seq2.append_blocks(allocator.allocate(4));  // [4, 5, 6, 7]
   seq2.commit_kv_cache(/*size=*/7);
-  seq2.append_new_token_id(100);
+  seq2.append_token(100);
 
   // seq in decode phase
   Sequence seq3(
+      /*prompt=*/"",
       /*token_ids=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19},
-      sampling_param,
-      stopping_criteria,
-      /*echo=*/false,
-      /*on_stream=*/nullptr);
+      capacity,
+      options);
   seq3.append_blocks(allocator.allocate(5));  // [8, 9, 10, 11, 12]
   seq3.commit_kv_cache(/*size=*/15);
-  seq3.append_new_token_id(200);
+  seq3.append_token(200);
 
   // define outputs
   Batch batch({&seq1, &seq2, &seq3});

src/handlers/chat_handler.cpp

@@ -1,5 +1,6 @@
 #include "chat_handler.h"
 
+#include <absl/strings/escaping.h>
 #include <glog/logging.h>
 #include <grpcpp/grpcpp.h>
 #include <torch/torch.h>
@@ -17,9 +18,9 @@
 #include "scheduler/scheduler.h"
 #include "utils.h"
 
-DEFINE_bool(disable_default_chat_template,
-            false,
-            "Disable default chat template");
+DEFINE_bool(enable_jinja_chat_template, false, "Enable Jinja chat template");
+
+DECLARE_int32(num_speculative_tokens);
 
 namespace llm {
 
@@ -87,10 +88,9 @@ bool send_delta_to_client(ChatCallData* call_data,
                           Request* request,
                           uint32_t index,
                           bool first_message,
-                          const std::string& delta,
-                          FinishReason reason) {
+                          const SequenceDeltaOutput& output) {
   // send delta to client
-  if (!delta.empty()) {
+  if (!output.delta.empty()) {
     ChatResponse response;
     response.set_object("chat.completion.chunk");
     response.set_id(request->id);
@@ -104,22 +104,22 @@ bool send_delta_to_client(ChatCallData* call_data,
     if (first_message) {
       message->set_role("assistant");
     }
-    message->set_content(delta);
+    message->set_content(output.delta);
     if (!call_data->write(std::move(response))) {
       return false;
     }
   }
 
   // send finish reason as a separate message
-  if (reason != FinishReason::NONE) {
+  if (output.finish_reason != FinishReason::NONE) {
     ChatResponse response;
     response.set_object("chat.completion");
     response.set_id(request->id);
     response.set_created(request->created_time);
     // response.set_model(request->model);
     auto* choice = response.add_choices();
     choice->set_index(index);
-    choice->set_finish_reason(finish_reason_to_string(reason));
+    choice->set_finish_reason(finish_reason_to_string(output.finish_reason));
     if (!call_data->write(std::move(response))) {
       return false;
     }
@@ -129,7 +129,7 @@ bool send_delta_to_client(ChatCallData* call_data,
 
 bool send_result_to_client(ChatCallData* call_data,
                            Request* request,
-                           const std::vector<SequenceResult>& seq_results,
+                           const std::vector<SequenceOutput>& seq_results,
                            const Status& /*status*/,
                            const Statistics& stats) {
   ChatResponse response;
@@ -145,7 +145,7 @@ bool send_result_to_client(ChatCallData* call_data,
     choice->set_index(i);
     auto* message = choice->mutable_message();
     message->set_role("assistant");
-    message->set_content(seq_result.output_text);
+    message->set_content(seq_result.text);
     if (seq_result.finish_reason != FinishReason::NONE) {
       choice->set_finish_reason(
           finish_reason_to_string(seq_result.finish_reason));
@@ -206,9 +206,21 @@ std::unique_ptr<Request> grpc_request_to_request(ChatCallData* call_data,
     return nullptr;
   }
 
+  uint32_t max_tokens = 0;
+  if (grpc_request.has_max_tokens()) {
+    max_tokens = grpc_request.max_tokens();
+  } else {
+    const uint32_t kDefaultMaxTokens = 16;
+    max_tokens = kDefaultMaxTokens;
+  }
+
   const uint32_t num_seqs = grpc_request.has_n() ? grpc_request.n() : 1;
+  // allocate enough capacity for prompt tokens, max tokens, and speculative
+  // tokens
+  const size_t capacity = prompt_tokens.size() + max_tokens +
+                          FLAGS_num_speculative_tokens + /*bouns_token*/ 1;
   auto request = std::make_unique<Request>(
-      generate_request_id(), "", num_seqs, prompt_tokens);
+      generate_request_id(), "", prompt_tokens, capacity, num_seqs);
 
   // construct sampling parameters
   auto& sampling_param = request->sampling_param;
@@ -232,16 +244,9 @@ std::unique_ptr<Request> grpc_request_to_request(ChatCallData* call_data,
 
   // construct stopping criteria
   auto& stopping_criteria = request->stopping_criteria;
-  auto max_tokens =
-      static_cast<uint32_t>(max_context_len - prompt_tokens.size());
-  if (grpc_request.has_max_tokens()) {
-    max_tokens = std::min(max_tokens, grpc_request.max_tokens());
-  } else {
-    const uint32_t kDefaultMaxTokens = 128;
-    max_tokens = std::min(max_tokens, kDefaultMaxTokens);
-  }
   stopping_criteria.max_tokens = max_tokens;
-  stopping_criteria.max_context_length = model_args.max_position_embeddings();
+  stopping_criteria.max_context_len =
+      max_context_len - FLAGS_num_speculative_tokens;
   // stopping_criteria.ignore_eos_token = false;
   stopping_criteria.eos_token_id = model_args.eos_token_id();
 
@@ -280,14 +285,14 @@ std::unique_ptr<Request> grpc_request_to_request(ChatCallData* call_data,
   // set callbacks
   if (request->stream) {
     // set callback for stream delta
-    request->on_stream_delta = [call_data, request = request.get()](
-                                   size_t index,
-                                   bool first_message,
-                                   const std::string& delta,
-                                   FinishReason reason) -> bool {
-      return send_delta_to_client(
-          call_data, request, index, first_message, delta, reason);
-    };
+    request->on_stream_delta =
+        [call_data, request = request.get(), first_message = true](
+            size_t index, const SequenceDeltaOutput& output) mutable {
+          const auto ret = send_delta_to_client(
+              call_data, request, index, first_message, output);
+          first_message = false;
+          return ret;
+        };
 
     // set callback for stream request
     request->on_stream_finish = [call_data](const Status& /*status*/) -> bool {
@@ -296,7 +301,7 @@ std::unique_ptr<Request> grpc_request_to_request(ChatCallData* call_data,
   } else {
     // set callback for non-stream request
     request->on_finish = [call_data, request = request.get()](
-                             const std::vector<SequenceResult>& seq_results,
+                             const std::vector<SequenceOutput>& seq_results,
                              const Status& status,
                              const Statistics& stats) -> bool {
       return send_result_to_client(
@@ -323,15 +328,15 @@ ChatHandler::ChatHandler(Scheduler* scheduler, const Engine* engine)
   // construct chat template
   auto factory = ModelRegistry::get_default_chat_template_factory(
       model_args_.model_type());
-  if (!FLAGS_disable_default_chat_template && factory) {
-    LOG(INFO) << "Use default chat template for model type: "
+  if (!FLAGS_enable_jinja_chat_template && factory) {
+    LOG(INFO) << "Using default chat template for model type: "
               << model_args_.model_type();
     chat_template_ = factory();
   } else {
     const auto& tokenizer_args = engine->tokenizer_args();
     if (!tokenizer_args.chat_template().empty()) {
-      LOG(INFO) << "Use chat template from tokenizer args for model type: "
-                << model_args_.model_type();
+      LOG(INFO) << "Using jinja chat template: "
+                << absl::CEscape(tokenizer_args.chat_template());
       chat_template_ = std::make_unique<JinjaChatTemplate>(
           tokenizer_args.chat_template(), /*add_generation_prompt=*/true);
     }

src/handlers/completion_handler.cpp

@@ -14,6 +14,8 @@
 #include "scheduler/scheduler.h"
 #include "utils.h"
 
+DECLARE_int32(num_speculative_tokens);
+
 namespace llm {
 
 namespace {
@@ -102,30 +104,29 @@ bool verify_request_arguments(CompletionCallData* call_data) {
 bool send_delta_to_client(CompletionCallData* call_data,
                           Request* request,
                           uint32_t index,
-                          const std::string& delta,
-                          FinishReason reason) {
-  if (!delta.empty()) {
+                          const SequenceDeltaOutput& output) {
+  if (!output.delta.empty()) {
     CompletionResponse response;
     response.set_object("text_completion");
     response.set_id(request->id);
     response.set_created(request->created_time);
     // response.set_model(request->model);
     auto* choice = response.add_choices();
     choice->set_index(index);
-    choice->set_text(delta);
+    choice->set_text(output.delta);
     if (!call_data->write(std::move(response))) {
       return false;
     }
   }
 
-  if (reason != FinishReason::NONE) {
+  if (output.finish_reason != FinishReason::NONE) {
     CompletionResponse response;
     response.set_object("text_completion");
     response.set_id(request->id);
     response.set_created(request->created_time);
     // response.set_model(request->model);
     auto* choice = response.add_choices();
-    choice->set_finish_reason(finish_reason_to_string(reason));
+    choice->set_finish_reason(finish_reason_to_string(output.finish_reason));
     if (!call_data->write(std::move(response))) {
       return false;
     }
@@ -135,7 +136,7 @@ bool send_delta_to_client(CompletionCallData* call_data,
 
 bool send_result_to_client(CompletionCallData* call_data,
                            Request* request,
-                           const std::vector<SequenceResult>& seq_results,
+                           const std::vector<SequenceOutput>& outputs,
                            const Status& /*status*/,
                            const Statistics& stats) {
   CompletionResponse response;
@@ -145,15 +146,14 @@ bool send_result_to_client(CompletionCallData* call_data,
   // response.set_model(request->model);
 
   // add choices into response
-  for (uint32_t i = 0; i < seq_results.size(); ++i) {
-    const auto& seq_result = seq_results[i];
+  for (uint32_t i = 0; i < outputs.size(); ++i) {
+    const auto& output = outputs[i];
     auto* choice = response.add_choices();
     choice->set_index(i);
-    choice->set_text(seq_result.output_text);
+    choice->set_text(output.text);
     // choice->set_logprobs(0);
-    if (seq_result.finish_reason != FinishReason::NONE) {
-      choice->set_finish_reason(
-          finish_reason_to_string(seq_result.finish_reason));
+    if (output.finish_reason != FinishReason::NONE) {
+      choice->set_finish_reason(finish_reason_to_string(output.finish_reason));
     }
   }
 
@@ -191,9 +191,24 @@ std::unique_ptr<Request> grpc_request_to_request(CompletionCallData* call_data,
     return nullptr;
   }
 
+  uint32_t max_tokens = 0;
+  if (grpc_request.has_max_tokens()) {
+    max_tokens = grpc_request.max_tokens();
+  } else {
+    const uint32_t kDefaultMaxTokens = 16;
+    max_tokens = kDefaultMaxTokens;
+  }
+  // allocate enough capacity for prompt tokens, max tokens, and speculative
+  // tokens
+  const size_t capacity = prompt_tokens.size() + max_tokens +
+                          FLAGS_num_speculative_tokens + /*bouns_token*/ 1;
+
   const uint32_t num_seqs = grpc_request.has_n() ? grpc_request.n() : 1;
-  auto request = std::make_unique<Request>(
-      generate_request_id(), grpc_request.prompt(), num_seqs, prompt_tokens);
+  auto request = std::make_unique<Request>(generate_request_id(),
+                                           grpc_request.prompt(),
+                                           prompt_tokens,
+                                           capacity,
+                                           num_seqs);
 
   // construct sampling parameters
   auto& sampling_param = request->sampling_param;
@@ -217,16 +232,9 @@ std::unique_ptr<Request> grpc_request_to_request(CompletionCallData* call_data,
 
   // construct stopping criteria
   auto& stopping_criteria = request->stopping_criteria;
-  auto max_tokens =
-      static_cast<uint32_t>(max_context_len - prompt_tokens.size());
-  if (grpc_request.has_max_tokens()) {
-    max_tokens = std::min(max_tokens, grpc_request.max_tokens());
-  } else {
-    const uint32_t kDefaultMaxTokens = 128;
-    max_tokens = std::min(max_tokens, kDefaultMaxTokens);
-  }
   stopping_criteria.max_tokens = max_tokens;
-  stopping_criteria.max_context_length = model_args.max_position_embeddings();
+  stopping_criteria.max_context_len =
+      max_context_len - FLAGS_num_speculative_tokens;
   // stopping_criteria.ignore_eos_token = false;
   stopping_criteria.eos_token_id = model_args.eos_token_id();
 
@@ -263,10 +271,8 @@ std::unique_ptr<Request> grpc_request_to_request(CompletionCallData* call_data,
   if (request->stream) {
     request->on_stream_delta = [call_data, request = request.get()](
                                    size_t index,
-                                   bool /*first_message*/,
-                                   const std::string& delta,
-                                   FinishReason reason) -> bool {
-      return send_delta_to_client(call_data, request, index, delta, reason);
+                                   const SequenceDeltaOutput& output) -> bool {
+      return send_delta_to_client(call_data, request, index, output);
     };
 
     // add on_stream_finish callback
@@ -276,7 +282,7 @@ std::unique_ptr<Request> grpc_request_to_request(CompletionCallData* call_data,
   } else {
     // add on_finish callback
     request->on_finish = [call_data, request = request.get()](
-                             const std::vector<SequenceResult>& seq_results,
+                             const std::vector<SequenceOutput>& seq_results,
                              const Status& status,
                              const Statistics& stats) -> bool {
       return send_result_to_client(