fix spec decoding

flexflow · Jan 12, 2024 · 197e308 · 197e308
1 parent 9c85a4f
commit 197e308
Show file tree

Hide file tree

Showing 6 changed files with 6 additions and 6 deletions.
diff --git a/deps/legion b/deps/legion
diff --git a/inference/models/falcon.cc b/inference/models/falcon.cc
@@ -39,7 +39,7 @@ void FALCON::create_falcon_model(FFModel &ff,
   Tensor input;
   {
     // assert(falcon_config.max_num_tokens <= BatchConfig::MAX_NUM_TOKENS);
-    int const token_dims[] = {mode == TREE_VERIFY_MODE
+    int const token_dims[] = {(mode == TREE_VERIFY_MODE || mode == BEAM_SEARCH_MODE)
                                   ? BatchConfig::max_verify_tokens_per_batch()
                                   : BatchConfig::max_tokens_per_batch(),
                               1};

diff --git a/inference/models/llama.cc b/inference/models/llama.cc
@@ -41,7 +41,7 @@ void LLAMA::create_llama_model(FFModel &ff,
 
   Tensor input;
   {
-    int const token_dims[] = {mode == TREE_VERIFY_MODE
+    int const token_dims[] = {(mode == TREE_VERIFY_MODE || mode == BEAM_SEARCH_MODE)
                                   ? BatchConfig::max_verify_tokens_per_batch()
                                   : BatchConfig::max_tokens_per_batch(),
                               1};

diff --git a/inference/models/mpt.cc b/inference/models/mpt.cc
@@ -40,7 +40,7 @@ void MPT::create_mpt_model(FFModel &ff,
   //------------------------------ build the model --------------------------
   Tensor input;
   {
-    int const token_dims[] = {mode == TREE_VERIFY_MODE
+    int const token_dims[] = {(mode == TREE_VERIFY_MODE || mode == BEAM_SEARCH_MODE)
                                   ? BatchConfig::max_verify_tokens_per_batch()
                                   : BatchConfig::max_tokens_per_batch(),
                               1};

diff --git a/inference/models/opt.cc b/inference/models/opt.cc
@@ -42,7 +42,7 @@ void OPT::create_opt_model(FFModel &ff,
   Tensor position_input;
   ff.set_position_offset(2);
   {
-    int const token_dims[] = {mode == TREE_VERIFY_MODE
+    int const token_dims[] = {(mode == TREE_VERIFY_MODE || mode == BEAM_SEARCH_MODE)
                                   ? BatchConfig::max_verify_tokens_per_batch()
                                   : BatchConfig::max_tokens_per_batch(),
                               1};

diff --git a/inference/models/starcoder.cc b/inference/models/starcoder.cc
@@ -48,7 +48,7 @@ void STARCODER::create_starcoder_model(
   ff.set_position_offset(0);
   {
     // assert(startcoder_config.max_num_tokens <= BatchConfig::MAX_NUM_TOKENS);
-    int const token_dims[] = {mode == TREE_VERIFY_MODE
+    int const token_dims[] = {(mode == TREE_VERIFY_MODE || mode == BEAM_SEARCH_MODE)
                                   ? BatchConfig::max_verify_tokens_per_batch()
                                   : BatchConfig::max_tokens_per_batch(),
                               1};