Validate BeamSearch vocab_size against logits width

[tianleiwu]

Summary

This PR hardens generation-path vocab-size handling by rejecting mismatches between the user-supplied vocab_size and the decoder logits width, instead of relying on unsafe span/copy logic.

Key Changes

• Added session-time validation in the beam-search parameter setup path to reject explicit vocab_size values that do not match the decoder subgraph logits width.
• Added runtime validation in CPU and CUDA generation helpers before logits span/copy logic uses vocab_size.
• Added regression tests covering mismatch rejection and the default-size inference path.

Why

A mismatched vocab_size could drive generation helpers to read or copy beyond the real logits width. These guards stop that class of out-of-bounds / DoS risk at both setup time and runtime.

Testing

• PATH=/home/tlwu/onnxruntime/.venv/bin:$PATH cmake --build . --target onnxruntime_provider_test -j2
• ./onnxruntime_provider_test --gtest_filter='BeamSearchParametersTest.SetSubgraphParametersRejectsMismatchedVocabSize:BeamSearchTest.*:GreedySearchTest.*'

[Copilot]

Pull request overview

This PR aims to harden the generation (BeamSearch/GreedySearch) path by validating that the configured vocab_size is compatible with the decoder subgraph logits last-dimension, preventing out-of-bounds reads/copies driven by a mismatched size.

Changes:

• Added validation in BeamSearchParameters::SetSubgraphParameters to reject mismatched vocab_size vs inferred decoder logits width.
• Added runtime validation in CPU/CUDA generation helpers to ensure vocab_size matches logits width before span/copy logic.
• Added unit tests for mismatch rejection and default vocab-size inference.

Reviewed changes

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

File	Description
onnxruntime/test/contrib_ops/beam_search_test.cc	Adds parameter-level regression tests for vocab-size mismatch handling and default inference.
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc	Adds runtime validation for CUDA BeamSearch/GreedySearch logits processing to guard against vocab/logits width mismatches.
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc	Adds runtime validation for CPU BeamSearch/GreedySearch logits processing to guard against vocab/logits width mismatches.
onnxruntime/contrib_ops/cpu/transformers/beam_search_parameters.cc	Adds setup-time validation to reject incompatible vocab_size at subgraph-parameter binding time.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.