Fix array eos_token_id handling

src/beam_search_scorer.cpp

@@ -122,7 +122,7 @@ void BeamSearchScorer::Process(Sequences& sequences,
 
       int const batch_beam_idx = static_cast<int>(batch * num_beams_) + next_index;
       // Add to generated hypotheses if end of sentence.
-      if ((eos_token_id_ >= 0) && (next_token == eos_token_id_)) {
+      if (contains(eos_token_id_, next_token)) {
         bool const is_beam_token_worse_than_top_num_beams = (j >= num_beams_);
         if (is_beam_token_worse_than_top_num_beams) {
           continue;

src/beam_search_scorer.h

@@ -53,7 +53,7 @@ struct BeamSearchScorer {
   int num_beams_;
   int max_length_;
   int pad_token_id_;
-  int eos_token_id_;
+  std::vector<int> eos_token_id_;
   bool early_stopping_;
   int not_done_count_;  // When zero, every batch entry is done (starts at batch_size_)
 

src/config.cpp

@@ -31,6 +31,17 @@ struct NamedStrings_Element : JSON::Element {
   std::vector<Config::NamedString>& v_;
 };
 
+struct Int_Array_Element : JSON::Element {
+  explicit Int_Array_Element(std::vector<int>& v) : v_{v} {}
+
+  void OnValue(std::string_view name, JSON::Value value) override {
+    v_.emplace_back(static_cast<int>(JSON::Get<double>(value)));
+  }
+
+ private:
+  std::vector<int>& v_;
+};
+
 struct ProviderOptionsObject_Element : JSON::Element {
   explicit ProviderOptionsObject_Element(std::vector<Config::ProviderOptions>& v) : v_{v} {}
 
@@ -509,29 +520,6 @@ struct Speech_Element : JSON::Element {
   SpeechOutputs_Element outputs_{v_.outputs};
 };
 
-struct Eos_Array_Element : JSON::Element {
-  explicit Eos_Array_Element(Config::Model& v) : v_{v} {}
-
-  void OnValue(std::string_view name, JSON::Value value) override {
-    v_.eos_token_ids.push_back(static_cast<int>(JSON::Get<double>(value)));
-  }
-
-  void OnComplete(bool empty) override {
-    if (v_.eos_token_ids.empty())
-      return;  // Empty array, nothign to do
-
-    // Copy the first eos_token_id into the eos_token_id value, it will be our primary eos token
-    v_.eos_token_id = v_.eos_token_ids.front();
-
-    // If the array is just one value, clear the array and just act like a single value was set
-    if (v_.eos_token_ids.size() == 1)
-      v_.eos_token_ids.clear();
-  }
-
- private:
-  Config::Model& v_;
-};
-
 struct EmbeddingInputs_Element : JSON::Element {
   explicit EmbeddingInputs_Element(Config::Model::Embedding::Inputs& v) : v_{v} {}
 
@@ -634,7 +622,7 @@ struct Model_Element : JSON::Element {
     } else if (name == "pad_token_id") {
       v_.pad_token_id = static_cast<int>(JSON::Get<double>(value));
     } else if (name == "eos_token_id") {
-      v_.eos_token_id = static_cast<int>(JSON::Get<double>(value));
+      v_.eos_token_id.assign(1, static_cast<int>(JSON::Get<double>(value)));
     } else if (name == "bos_token_id") {
       v_.bos_token_id = static_cast<int>(JSON::Get<double>(value));
     } else if (name == "decoder_start_token_id") {
@@ -647,7 +635,7 @@ struct Model_Element : JSON::Element {
 
   Element& OnArray(std::string_view name) override {
     if (name == "eos_token_id")
-      return eos_token_ids_;
+      return eos_token_id_;
     throw JSON::unknown_value_error{};
   }
 
@@ -677,7 +665,7 @@ struct Model_Element : JSON::Element {
   Config::Model& v_;
   EncoderDecoderInit_Element encoder_decoder_init_{v_.encoder_decoder_init};
   Decoder_Element decoder_{v_.decoder};
-  Eos_Array_Element eos_token_ids_{v_};
+  Int_Array_Element eos_token_id_{v_.eos_token_id};
   Vision_Element vision_{v_.vision};
   Embedding_Element embedding_{v_.embedding};
   PromptTemplates_Element prompt_templates_{v_.prompt_templates};
@@ -858,6 +846,10 @@ Config::Config(const fs::path& path, std::string_view json_overlay) : config_pat
 
   if (search.max_length == 0)
     search.max_length = model.context_length;
+
+  // If no eos_token_id was set, set it to the pad token id
+  if (model.eos_token_id.empty())
+    model.eos_token_id.push_back(model.pad_token_id);
 }
 
 void Config::AddMapping(const std::string& nominal_name, const std::string& graph_name) {

src/config.h

@@ -81,12 +81,11 @@ struct Config {
   struct Model {
     std::string type;
 
-    int pad_token_id{};              // The id of the padding token.
-    int eos_token_id{};              // The id of the end-of-stream token.
-    std::vector<int> eos_token_ids;  // If eos_token_id is passed as an array, this is where the values go (eos_token_id gets set to the first entry in the array)
-    int bos_token_id{};              // The id of the beginning-of-stream token.
-    int sep_token_id{};              // The id of the separation token.
-    int decoder_start_token_id{};    // If an encoder-decoder model starts decoding with a different token than bos, the id of that token.
+    int pad_token_id{};             // The id of the padding token.
+    std::vector<int> eos_token_id;  // The end-of-stream tokens (when set as a single value it is converted to a vector with one value).
+    int bos_token_id{};             // The id of the beginning-of-stream token.
+    int sep_token_id{};             // The id of the separation token.
+    int decoder_start_token_id{};   // If an encoder-decoder model starts decoding with a different token than bos, the id of that token.
     int vocab_size{};
     int context_length{};
 

src/cuda/beam_search_scorer_cuda.cpp

@@ -10,14 +10,14 @@
 
 namespace Generators {
 
-BeamSearchScorer_Cuda::BeamSearchScorer_Cuda(const GeneratorParams& parameters)
-    : stream_{GetStream()} {
+BeamSearchScorer_Cuda::BeamSearchScorer_Cuda(const GeneratorParams& parameters, std::span<int32_t> cuda_eos_tokens)
+    : stream_{GetStream()},
+      eos_tokens_{cuda_eos_tokens} {
   state_cpu_ = CudaMallocHostArray<cuda::BeamScorerState>(1);
   state_cpu_->batch_size_ = static_cast<size_t>(parameters.search.batch_size);
   state_cpu_->num_beams_ = static_cast<size_t>(parameters.search.num_beams);
   state_cpu_->max_length_ = static_cast<size_t>(parameters.search.max_length);
   state_cpu_->pad_token_id_ = parameters.config.model.pad_token_id;
-  state_cpu_->eos_token_id_ = parameters.config.model.eos_token_id;
   state_cpu_->early_stopping_ = parameters.search.early_stopping;
   state_cpu_->not_done_count_ = parameters.search.batch_size;
   state_cpu_->hypothesis_buffer_used_ = 0;
@@ -51,6 +51,7 @@ void BeamSearchScorer_Cuda::Process(Sequences& sequences,
                                     std::span<const int32_t> next_indices) {
   cuda::LaunchBeamSearchScorer_Process(*state_cpu_,
                                        *state_gpu_,
+                                       eos_tokens_,
                                        sequences.GetSequences().Span(),
                                        sequences.GetSequenceLength(),
                                        beam_hyps_,

src/cuda/beam_search_scorer_cuda.cu

@@ -76,6 +76,8 @@ __device__ bool BeamHypotheses::CanImprove(float best_sum_logprobs, int current_
 
 __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
                                          BeamScorerState& state,
+                                         const int32_t* eos_token_ids,
+                                         const int eos_token_count,
                                          const int32_t* sequences_buffer,
                                          int sequence_length,
                                          BeamHypotheses* beam_hyps_,
@@ -105,7 +107,14 @@ __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
 
       int batch_beam_idx = batch_start + next_index;
       // Add to generated hypotheses if end of sentence.
-      if ((state.eos_token_id_ >= 0) && (next_token == state.eos_token_id_)) {
+      bool is_eos_token = false;
+      for (unsigned eos_index = 0; eos_index < eos_token_count; eos_index++) {
+        if (next_token == eos_token_ids[eos_index]) {
+          is_eos_token = true;
+          break;
+        }
+      }
+      if (is_eos_token) {
         bool is_beam_token_worse_than_top_num_beams = (j >= state.num_beams_);
         if (is_beam_token_worse_than_top_num_beams) {
           continue;
@@ -152,6 +161,7 @@ __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
 
 void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
                                     BeamScorerState& state,
+                                    std::span<const int32_t> eos_token_ids,
                                     std::span<const int32_t> sequences,
                                     int sequence_length,
                                     std::span<BeamHypotheses> beam_hyps,
@@ -165,6 +175,8 @@ void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
                                     cudaStream_t stream) {
   BeamSearchScorer_Process<<<1, state_cpu.batch_size_, 0, stream>>>(state_cpu,
                                                                     state,
+                                                                    eos_token_ids.data(),
+                                                                    static_cast<int>(eos_token_ids.size()),
                                                                     sequences.data(),
                                                                     sequence_length,
                                                                     beam_hyps.data(),

src/cuda/beam_search_scorer_cuda.cuh

@@ -32,7 +32,6 @@ struct BeamScorerState {
   int num_beams_;
   int max_length_;
   int pad_token_id_;
-  int eos_token_id_;
   bool early_stopping_;
   int not_done_count_;  // When zero, every batch entry is done (starts at batch_size_)
 
@@ -43,6 +42,7 @@ void LaunchInitializeBeamHypotheses(std::span<BeamHypotheses> beam_hyps, float l
 
 void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
                                     BeamScorerState& state,
+                                    std::span<const int32_t> eos_token_ids,
                                     std::span<const int32_t> sequences,
                                     int sequence_length,
                                     std::span<BeamHypotheses> beam_hyps_,

src/cuda/beam_search_scorer_cuda.h

@@ -4,7 +4,7 @@
 namespace Generators {
 
 struct BeamSearchScorer_Cuda {
-  BeamSearchScorer_Cuda(const GeneratorParams& parameters);
+  BeamSearchScorer_Cuda(const GeneratorParams& parameters, std::span<int32_t> cuda_eos_tokens);
 
   void Process(Sequences& sequences,
                std::span<const float> next_scores,
@@ -26,7 +26,9 @@ struct BeamSearchScorer_Cuda {
   mutable cuda_event_holder event_process_complete_;
   cuda_host_unique_ptr<cuda::BeamScorerState> state_cpu_;
   cuda_unique_ptr<cuda::BeamScorerState> state_gpu_;
+
   cudaStream_t stream_;
+  std::span<int32_t> eos_tokens_;
 
   DeviceSpan<float> next_beam_scores_;
   DeviceSpan<int32_t> next_beam_tokens_;

src/cuda/interface.cpp

@@ -143,10 +143,6 @@ struct CudaInterfaceImpl final : DeviceInterface {
     return true;
   }
 
-  void LaunchHandleEOSArray(float* batch_logits, int batch_beam_size, int vocab_size, const int32_t* eos_token_ids, int eos_token_ids_count) override {
-    cuda::LaunchHandleEOSArray(batch_logits, batch_beam_size, vocab_size, eos_token_ids, eos_token_ids_count, GetStream());
-  }
-
   void UpdateCacheIndirectionKernelLauncher(int32_t* tgt_indir_cache, const int32_t* src_indir_cache, const int32_t* beam_ids, int batch_size, int beam_width, int input_seq_length, int max_seq_length, int current_length) override {
     cuda::UpdateCacheIndirectionKernelLauncher(tgt_indir_cache, src_indir_cache, beam_ids, batch_size, beam_width, input_seq_length, max_seq_length, current_length, GetStream());
   }

src/cuda/kernels.h

@@ -11,8 +11,6 @@ void Launch_UpdatePositionIds(T* positions, int batch_beam_size, int total_lengt
 template <typename T>
 void Launch_UpdateAttentionMask(T* mask_data, T* old_data, int batch_beam_size, int new_kv_length, int total_length, int max_length, bool update_only, cudaStream_t stream);
 
-void LaunchHandleEOSArray(float* batch_logits, int batch_beam_size, int vocab_size, const int32_t* eos_token_ids, int eos_token_ids_count, cudaStream_t stream);
-
 void LaunchFp16ToFp32(const uint16_t* fp16, float* fp32, int count, cudaStream_t stream);
 void LaunchFp32ToFp16(const float* fp32, uint16_t* fp16, int count, cudaStream_t stream);
 void LaunchInt32ToInt64(const int32_t* src, int64_t* dst, int count, cudaStream_t stream);

src/cuda/model_kernels.cu

@@ -82,25 +82,6 @@ void Launch_UpdateAttentionMask(T* next_mask_data, T* mask_data, int batch_beam_
 template void Launch_UpdateAttentionMask(int32_t* next_mask_data, int32_t* mask_data, int batch_beam_size, int new_kv_length, int total_length, int max_length, bool update_only, cudaStream_t stream);
 template void Launch_UpdateAttentionMask(int64_t* next_mask_data, int64_t* mask_data, int batch_beam_size, int new_kv_length, int total_length, int max_length, bool update_only, cudaStream_t stream);
 
-__global__ void HandleEOSArray(float* batch_logits, int batch_beam_size, int vocab_size, const int32_t* eos_token_ids, int eos_token_ids_count) {
-  int index = blockIdx.x * blockDim.x + threadIdx.x;
-  if (index >= batch_beam_size)
-    return;
-
-  float* logits = batch_logits + index * vocab_size;
-  float max = std::numeric_limits<float>::lowest();
-  for (int i = 0; i < eos_token_ids_count; i++) {
-    max = std::max(max, logits[eos_token_ids[i]]);
-    logits[eos_token_ids[i]] = std::numeric_limits<float>::lowest();  // Set all EOS token options to never happen (the first will get the max of all)
-  }
-
-  logits[eos_token_ids[0]] = max;  // Set the score of the primary EOS token to the highest of any of the EOS tokens
-}
-
-void LaunchHandleEOSArray(float* batch_logits, int batch_beam_size, int vocab_size, const int32_t* eos_token_ids, int eos_token_ids_count, cudaStream_t stream) {
-  HandleEOSArray<<<(batch_beam_size + 255) / 256, 256, 0, stream>>>(batch_logits, batch_beam_size, vocab_size, eos_token_ids, eos_token_ids_count);
-}
-
 __global__ void ConvertFp16ToFp32(const half* src, float* dst, int count) {
   int idx = threadIdx.x + blockIdx.x * blockDim.x;
   if (idx < count)

src/cuda/search_cuda.cpp

@@ -24,8 +24,12 @@ Search_Cuda::Search_Cuda(const GeneratorParams& params)
   auto batch_beam_size = params.BatchBeamSize();
   sequence_lengths_ = params.p_device->Allocate<int32_t>(batch_beam_size);
 
-  eos_meet_buffer_ = CudaMallocArray<bool>(batch_beam_size, &eos_meet_);
-  cudaMemsetAsync(eos_meet_.data(), 0, eos_meet_.size_bytes(), GetStream());
+  eos_seen_buffer_ = CudaMallocArray<bool>(batch_beam_size, &eos_seen_);
+  cudaMemsetAsync(eos_seen_.data(), 0, eos_seen_.size_bytes(), GetStream());
+
+  eos_token_ids_ = params.p_device->Allocate<int32_t>(params.config.model.eos_token_id.size());
+  copy(std::span<const int32_t>{params.config.model.eos_token_id}, eos_token_ids_.CpuSpan());
+  eos_token_ids_.CopyCpuToDevice();
 
   done_cpu_ = CudaMallocHostArray<bool>(1);
   *done_cpu_ = false;
@@ -49,7 +53,7 @@ BeamSearch_Cuda::BeamSearch_Cuda(const GeneratorParams& params)
     : Search_Cuda{params} {
   assert(params_->search.num_beams > 1);  // If 1, use GreedySearch
   auto batch_beam_size = params_->BatchBeamSize();
-  beam_scorer_ = std::make_unique<BeamSearchScorer_Cuda>(*params_);
+  beam_scorer_ = std::make_unique<BeamSearchScorer_Cuda>(*params_, eos_token_ids_.Span());
 
   topk_next_tokens_ = CudaMallocArray<int32_t>(2 * batch_beam_size);
   topk_next_indices_ = CudaMallocArray<int32_t>(2 * batch_beam_size);
@@ -153,9 +157,8 @@ void GreedySearch_Cuda::SampleTopKTopP(int k, float p, float temperature) {
                   params_->search.batch_size, k, p, temperature);
 
   // Check for EOS
-  assert(next_tokens_.size() == eos_meet_.size());
-  // Don't replace EOS with pad for batch_size == 1 for continuous decoding mode
-  cuda::Launch_CheckForEOSAndPad(next_tokens_.data(), static_cast<int>(next_tokens_.size()), eos_meet_.data(), params_->config.model.eos_token_id, params_->search.batch_size > 1 ? params_->config.model.pad_token_id : params_->config.model.eos_token_id, done_cpu_.get(), GetStream());
+  assert(next_tokens_.size() == eos_seen_.size());
+  cuda::Launch_CheckForEOSAndPad(next_tokens_.data(), static_cast<int>(next_tokens_.size()), eos_seen_.data(), eos_token_ids_.Span().data(), static_cast<int>(eos_token_ids_.Span().size()), params_->config.model.pad_token_id, done_cpu_.get(), GetStream());
 
   // Append tokens
   cuda::Launch_AppendNextTokensToSequences(next_tokens_buffer_.Span(), sequences_.GetSequences().Span(), params_->BatchBeamSize(), sequences_.GetSequenceLength(), sequences_.max_length_, GetStream());
@@ -210,7 +213,7 @@ std::span<float> Search_Cuda::GetScores() {
 
 // Set user input tokens (batch_beam_size, sequence_length)
 void GreedySearch_Cuda::AppendTokens(DeviceSpan<int32_t>& next_tokens) {
-  cudaMemsetAsync(eos_meet_.data(), 0, eos_meet_.size_bytes(), GetStream());
+  cudaMemsetAsync(eos_seen_.data(), 0, eos_seen_.size_bytes(), GetStream());
   *done_cpu_ = false;
 
   auto next_tokens_gpu = next_tokens.Span();
@@ -224,7 +227,7 @@ void GreedySearch_Cuda::AppendTokens(DeviceSpan<int32_t>& next_tokens) {
     return;
   }
 
-  cudaMemsetAsync(eos_meet_.data(), 0, eos_meet_.size_bytes(), GetStream());
+  cudaMemsetAsync(eos_seen_.data(), 0, eos_seen_.size_bytes(), GetStream());
   *done_cpu_ = false;
 }
 
@@ -237,7 +240,7 @@ void BeamSearch_Cuda::AppendTokens(DeviceSpan<int32_t>& next_tokens) {
 }
 
 void GreedySearch_Cuda::RewindTo(size_t index) {
-  cudaMemsetAsync(eos_meet_.data(), 0, eos_meet_.size_bytes(), GetStream());
+  cudaMemsetAsync(eos_seen_.data(), 0, eos_seen_.size_bytes(), GetStream());
   *done_cpu_ = false;
   if (index > 0)
     cuda::Launch_GetLastTokens(next_tokens_.data(), sequences_.GetSequences().Span().data(), static_cast<int>(params_->BatchBeamSize()), static_cast<int>(index), sequences_.max_length_, GetStream());
@@ -250,7 +253,8 @@ void Search_Cuda::ApplyMinLength(int min_length) {
   if (sequences_.GetSequenceLength() >= min_length)
     return;
 
-  cuda::LaunchSetScoreProcessor(GetScores().data(), params_->BatchBeamSize(), params_->config.model.vocab_size, params_->config.model.eos_token_id, std::numeric_limits<float>::lowest(), GetStream());
+  for (auto eos_token_id : params_->config.model.eos_token_id)
+    cuda::LaunchSetScoreProcessor(GetScores().data(), params_->BatchBeamSize(), params_->config.model.vocab_size, eos_token_id, std::numeric_limits<float>::lowest(), GetStream());
 }
 
 void Search_Cuda::ApplyRepetitionPenalty(float penalty) {