Fix heap overflow

[Akshay Sonawane (apsonawane)]

This pull request adds input validation to the ExpandBuffer helper functions for both CPU and CUDA implementations, ensuring that the input sequence length does not exceed the specified maximum sequence length. It also introduces corresponding unit tests to verify this validation logic.

Input validation improvements:

• Added a check in both GenerationCpuDeviceHelper::ExpandBuffer (generation_device_helper.cc) and GenerationCudaDeviceHelper::ExpandBuffer (generation_device_helper.cc) to return an error if the input sequence length exceeds the maximum allowed sequence length. [1] [2]

Testing enhancements:

• Added a new test ExpandBufferRejectsSequenceLengthExceedingMaximum to beam_search_test.cc for CPU, which verifies that the function properly rejects inputs with sequence length greater than the maximum.
• Added a similar test ExpandBufferRejectsSequenceLengthExceedingMaximum to generation_cuda_impl_test.cc for CUDA, ensuring the same validation logic is enforced for CUDA.
• Included necessary headers in generation_cuda_impl_test.cc to support the new test.

[Copilot]

🟢 Approval recommended

The validation precedes allocation and copying, with equivalent regression coverage for both implementations.

Pull request overview

Adds CPU and CUDA validation to prevent oversized KV-cache copies from overflowing expanded buffers.

Changes:

• Rejects sequence lengths exceeding the configured maximum.
• Adds CPU and CUDA regression tests.

File summaries

File	Description
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc	Adds CPU bounds validation.
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc	Adds CUDA bounds validation.
onnxruntime/test/contrib_ops/beam_search_test.cc	Tests CPU rejection behavior.
onnxruntime/test/providers/cuda/test_cases/generation_cuda_impl_test.cc	Tests CUDA rejection behavior.

Review details

• Files reviewed: 4/4 changed files
• Comments generated: 0
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.