Fix out-of-bounds error

Merged
AP
5mo ago
+11-1

This pull request introduces important safety checks to prevent out-of-bounds access in the logits processing code for transformers. The main updates ensure that token IDs are validated against the vocabulary size before being used, which improves robustness and prevents potential crashes.

Safety and robustness improvements:

  • Added bounds checking for token IDs in the RepetitionPenaltyLogitsProcessor<T>::Process method to ensure only valid IDs are used when accessing beam_token_scores.
  • Added bounds checking for token IDs in the NoRepeatNGramLogitsProcessor<T>::Process method to prevent out-of-bounds writes to beam_token_scores.
  • Updated the NextTokenScores::SetScore method to return early if the provided token_id is out of bounds, replacing the previous assert with a safe check.