Validate min_tokens_to_keep in the Sampling op to prevent an out-of-bounds access

[David Fan (jiafatom)]

Summary

The Sampling contrib operator (com.microsoft, default CPU EP) reads the integer graph attribute min_tokens_to_keep without validation and uses it in an unsigned subtraction against vocab_size to bound a heap-buffer loop. When min_tokens_to_keep > vocab_size, the expression (size_t)vocab_size - (size_t)min_tokens_to_keep underflows to a value near 2^64, so cumulate_and_filter() walks the cumulative_probs buffer (size batch_size * vocab_size floats) past its end, causing an out-of-bounds read-modify-write and a crash. This happens at the first Run() with no special flags.

Root cause

onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h, cumulate_and_filter():

for (size_t j = 1; j < (size_t)vocab_size - (size_t)min_tokens_to_keep; j++)
  cumulative_probs[j + offset] += cumulative_probs[j + offset - 1];

vocab_size is validated ([1, decoder-logits-width] in GreedySearchProcessLogits), but min_tokens_to_keep is copied verbatim from the model (sampling_parameters.cc:18; schema default 0 at contrib_defs.cc:1369).

Fix

Validate min_tokens_to_keep is in [0, vocab_size) at the top of Sample() — before any buffer access — so an out-of-range model fails cleanly with a Status error instead of accessing memory out of bounds. Applied to both the CPU (sampling_cpu_helper.h) and CUDA (sampling_cuda_helper.h, where the same attribute feeds LaunchFilterLogitsKernel) paths. vocab_size is already validated positive before Sample() runs.

Testing

• Added SamplingTest.InvalidMinTokensToKeep_CPU: loads a model with min_tokens_to_keep=1000000 and asserts Run() fails with a clear error rather than crashing.
• New testdata tiny_gpt2_sampling_invalid_min_tokens.onnx (clone of tiny_gpt2_sampling.onnx with the attribute overwritten).
• Verified the guard admits exactly the safe range (0..vocab_size-1) and rejects values that underflow the loop bound.

[Copilot]

Pull request overview

This PR hardens the com.microsoft contrib Sampling operator by validating the min_tokens_to_keep graph attribute before it can be used in CPU/CUDA filtering logic, preventing arithmetic underflow/invalid indexing that could lead to heap out-of-bounds access during Run().

Changes:

• Add runtime validation in CPU and CUDA Sample() helpers to require min_tokens_to_keep be in [0, vocab_size).
• Add a CPU regression test that runs an invalid model (min_tokens_to_keep set far above vocab_size) and asserts Run() fails with a clear error instead of crashing.

Reviewed changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/test/contrib_ops/sampling_test.cc	Adds a CPU regression test to ensure invalid min_tokens_to_keep fails cleanly at runtime.
onnxruntime/contrib_ops/cuda/transformers/sampling_cuda_helper.h	Adds a guard rejecting out-of-range min_tokens_to_keep before launching CUDA filtering logic.
onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h	Adds a guard rejecting out-of-range min_tokens_to_keep before CPU-side cumulative filtering.