update

flexflow · Oct 7, 2024 · d2caba8 · d2caba8
1 parent 4a3d1bd
commit d2caba8
Show file tree

Hide file tree

Showing 5 changed files with 259 additions and 40 deletions.
diff --git a/include/flexflow/model.h b/include/flexflow/model.h
@@ -277,6 +277,7 @@ enum TaskIDs {
   RM_PREPARE_NEXT_BATCH_INIT_TASK_ID,
   RM_PREPARE_NEXT_BATCH_BEAM_TASK_ID,
   RM_PREPARE_NEXT_BATCH_VERIFY_TASK_ID,
+  RM_PREPARE_NEXT_BATCH_SUFFIX_DECODE_TASK_ID,
   RM_BACKGROUND_SERVING_TASK_ID,
   // Custom tasks
   CUSTOM_GPU_TASK_ID_FIRST,

diff --git a/include/flexflow/request_manager.h b/include/flexflow/request_manager.h
@@ -19,6 +19,7 @@
 #include "flexflow/inference.h"
 #include "flexflow/model.h"
 #include "flexflow/utils/file_loader.h"
+#include "suffix_decoding.h"
 #include <future>
 #include <mutex>
 #include <tokenizers_cpp.h>
@@ -164,6 +165,7 @@ class RequestManager {
 
   void serve_incr_decoding(FFModel *model);
   void serve_spec_infer(FFModel *model);
+  void serve_suffix_decoding(FFModel *model);
   GenerationResult get_generation_result(RequestGuid const &guid);
   RequestGuid register_new_request(Request const &request_);
   RequestGuid register_new_peft_request(Request const &request_);
@@ -210,6 +212,15 @@ class RequestManager {
       Legion::Context ctx,
       Legion::Runtime *runtime);
 
+  TreeVerifyBatchConfig
+      prepare_next_batch_suffix_decode(TreeVerifyBatchConfig const &old_bc,
+                                       InferenceResult const &result);
+  TreeVerifyBatchConfigFuture prepare_next_batch_suffix_decode(
+      TreeVerifyBatchConfigFuture const &old_bc,
+      InferenceResultFuture const &result,
+      Legion::Context ctx,
+      Legion::Runtime *runtime);
+
   void store_beam_metadata(BeamSearchBatchConfig const &old_bc,
                            BeamInferenceResult const &result);
   void update_beam_metadata(BeamSearchBatchConfig &new_bc,
@@ -280,6 +291,12 @@ class RequestManager {
       Legion::Context ctx,
       Legion::Runtime *runtime);
 
+  static TreeVerifyBatchConfig prepare_next_batch_suffix_decode_task(
+      Legion::Task const *task,
+      std::vector<Legion::PhysicalRegion> const &regions,
+      Legion::Context ctx,
+      Legion::Runtime *runtime);
+
 private:
   // configuration parameters
   int max_requests_per_batch;
@@ -295,6 +312,8 @@ class RequestManager {
   // tree width in each speculative step, if not specified 1
   std::vector<int> spec_infer_tree_width;
 
+  SuffixTree *suffix_tree;
+
   // private fields
   std::unique_ptr<Tokenizer> tokenizer_;
   bool verbose;

diff --git a/src/mapper/mapper.cc b/src/mapper/mapper.cc
@@ -284,6 +284,7 @@ void FFMapper::select_task_options(const MapperContext ctx,
       (task.task_id == RM_PREPARE_NEXT_BATCH_INIT_TASK_ID) ||
       (task.task_id == RM_PREPARE_NEXT_BATCH_BEAM_TASK_ID) ||
       (task.task_id == RM_PREPARE_NEXT_BATCH_VERIFY_TASK_ID) ||
+      (task.task_id == RM_PREPARE_NEXT_BATCH_SUFFIX_DECODE_TASK_ID) ||
       (task.task_id == RM_BACKGROUND_SERVING_TASK_ID)) {
     output.initial_proc = all_cpus[0];
     return;

diff --git a/src/runtime/model.cc b/src/runtime/model.cc
@@ -4822,6 +4822,27 @@ void register_flexflow_internal_tasks(Runtime *runtime,
           RequestManager::prepare_next_batch_verify_task>(registrar);
     }
   }
+  // RequestManager prepare_next_batch_suffix_decode
+  {
+    TaskVariantRegistrar registrar(
+        RM_PREPARE_NEXT_BATCH_SUFFIX_DECODE_TASK_ID,
+        "RequestManager Prepare Next Batch (Suffix Decode)");
+    registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
+    registrar.set_leaf();
+    if (pre_register) {
+      Runtime::preregister_task_variant<
+          TreeVerifyBatchConfig,
+          RequestManager::prepare_next_batch_suffix_decode_task>(
+          registrar, "RequestManager Prepare Next Batch (Suffix Decode) Task");
+    } else {
+      if (enable_control_replication) {
+        registrar.global_registration = false;
+      }
+      runtime->register_task_variant<
+          TreeVerifyBatchConfig,
+          RequestManager::prepare_next_batch_suffix_decode_task>(registrar);
+    }
+  }
   // RequestManager background serving task
   {
     TaskVariantRegistrar registrar(RM_BACKGROUND_SERVING_TASK_ID,