Fix signed-int overflow in SamplingState::Init to prevent heap-buffer-overflow

Merged
AP
4mo ago
+124-12

This pull request improves the safety of buffer size calculations in the SamplingState initialization logic by ensuring that all multiplications involving batch_size and vocab_size are safely performed using SafeInt<size_t>. This prevents potential integer overflow bugs that could lead to under-allocated buffers and memory errors.

Buffer allocation safety improvements:

  • All buffer size calculations that multiply batch_size and vocab_size now use SafeInt<size_t> to ensure checked arithmetic, preventing silent integer overflows that could cause heap-buffer-overflow issues. This includes allocations for both CPU and CUDA buffers in SamplingState. [1] [2]
  • The calculation for the buffer size of h_sampled_all now also safely casts max_iter to size_t before multiplication, further protecting against overflow.

These changes make the code more robust and secure, especially when handling large or model-controlled input sizes.