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

[apsonawane]

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.

[Copilot]

Pull request overview

This PR hardens SamplingState::Init in the generation/transformers greedy-search implementation by moving buffer element-count computations into SafeInt<size_t> so integer overflow can’t lead to under-allocation and downstream memory errors.

Changes:

• Compute batch_size * vocab_size using SafeInt<size_t> to prevent overflow before buffer allocation.
• Reuse the checked total_count across CPU/CUDA allocations in SamplingState.

[tianleiwu]

The production change in SamplingState::Init is correct and a genuine hardening win. int total_count = batch_size * vocab_size; is replaced with a checked SafeInt<size_t> product before any value can positive-wrap, and the two remaining in-int sub-expressions (batch_size + 1, batch_size * max_iter) are fixed the same way. Since batch_size (BatchBeamSize()) and vocab_size are model-controlled ints, this is a reachable defensive fix.

One concern on the test file (inline): the added tests re-implement the SafeInt expression inside the test itself and never call the production code, so they cannot fail if someone reverts SamplingState::Init to int * int — which contradicts the test's own comment. They only prove that SafeInt<size_t> throws on negative operands, a property of SafeInt rather than of this PR.

Minor (optional): safeint.h already exposes SafeMul<size_t>(batch_size, vocab_size) which does exactly this multiply; not required since the explicit form matches the existing next_token_scores pattern.

Verdict: COMMENT — production fix looks good; please make the regression test actually exercise the production path (or correct the misleading comment).

[tianleiwu]

The revised approach resolves my earlier concern: the size computation is now in a shared production helper (SamplingBufferElementCount in sampling_buffer_element_count.h) that both SamplingState::Init and the new test call, so the regression test genuinely guards against a revert. The core fix is correct — SafeMul<size_t> casts each operand to size_t before multiplying, which catches the original int * int overflow (previously it wrapped in the int domain before the old SafeInt<size_t> cast ever saw it) as well as negative operands. Moving the +1 into the checked domain for d_offset (SafeInt<size_t>(batch_size) + 1) is also a correct hardening. LGTM aside from one minor test nit below.

Note (pre-existing, out of scope): the downstream byte-count math still multiplies in the int domain — e.g. sizeof(int) * batch_size * vocab_size in generation_device_helper.cc and batch_size * vocab_size in generation_cuda_impl.cu. These are now effectively guarded because Init throws before allocation when the product would overflow, so the buffers are either correctly sized or the run aborts. No change needed here, just flagging for follow-up hardening.