saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

update

Browse source
This commit is contained in:
Tianlei Wu 2025-02-05 02:03:56 -08:00
parent 8d6c1556a5
commit 259bef34c7
1 changed files with 11 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu
View file

@ -429,9 +429,18 @@ __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
#endif
if (beam_hyp.beams_used_ == state.num_beams_) {
bool is_done = state.early_stopping_;
if (is_done) {
#ifdef DEBUG_GENERATION
printf("\n --- state.early_stopping_ (batch %d)\n", batch);
#endif
}
if (!is_done) {
float best_sum_logprobs = *std::max_element(next_scores + batch_start, next_scores + batch_start + top_k);
is_done = !beam_hyp.CanImprove(*std::max_element(next_scores + batch_start, next_scores + batch_start + top_k), sequence_length);
is_done = !beam_hyp.CanImprove(best_sum_logprobs, sequence_length);
#ifdef DEBUG_GENERATION
printf("\n --- CanImprove=%d with best_sum_logprobs=%f, sequence_length=%d\n", is_done ? 0 : 1, best_sum_logprobs, sequence_length);
#endif
}
if (is_done) {
@ -440,6 +449,7 @@ __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
state_cpu.not_done_count_ = 0; // Update the CPU side
#ifdef DEBUG_GENERATION
printf("\n --- BeamSearchScorer_Process updated cpu state for batch %d\n", batch);
printf("\n --- is_done (state.early_stopping_=%d state_cpu.not_done_count_=%d) for batch %d\n", state.early_stopping_ ? 1 : 0, state_cpu.not_done_count_, batch);
#endif
}
}