Bound cross_qk layer and head indices in CUDA Whisper beam search

[Akshay Sonawane (apsonawane)]

CopyCrossQKSingleDecodeStepKernel took the layer and head indices from the cross_qk_layer_head model input and used them to index qk_layer_pointers and to offset into the layer buffer without any range check, even though num_layers was already passed to the kernel for that purpose. Out-of-range pairs now zero their output slice instead of dereferencing a wild pointer.

Also require the cross_qk_layer_head input to have shape [layer_head_pair_count, 2], since the pair count is taken from dim 0 and the kernel reads two values per pair.

[Copilot]

Pull request overview

Hardens Whisper beam search cross-attention Q*K collection on CUDA by treating cross_qk_layer_head as untrusted input, preventing out-of-range indexing into layer/head buffers and enforcing the expected pair tensor shape.

Changes:

• Add layer/head bounds checks in CopyCrossQKSingleDecodeStepKernel, zero-filling output for invalid pairs.
• Enforce cross_qk_layer_head input shape to be [layer_head_pair_count, 2] before using its dim 0 as the pair count.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu	Adds GPU-side layer/head range checks and safe zero-fill behavior for invalid pairs.
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h	Validates cross_qk_layer_head tensor rank/shape before consuming it.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Ti-Tai Wang (titaiwangms)]

Major review finding

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h:313

The new pair_count <= num_layers * num_heads restriction rejects schema-valid lists containing repeated (layer, head) pairs. The WhisperBeamSearch schema defines this input only as a list shaped [number of pairs to keep, 2], and defines the output dimension from cross_qk_layer_head.shape[0]; it does not require uniqueness.

For example, [[0, 0], [0, 0]] is a valid request for two output slices even when the model has one layer and one head, but this check rejects it.

Please retain the actual launcher bound (< 65536) and per-pair layer/head validation without inferring a uniqueness-based count limit that the operator contract does not impose.

[Tianlei Wu (tianleiwu)]

Requesting changes for one additional CUDA edge case. The separate repeated-pair contract regression is already tracked in #31998 (comment), so I have not duplicated that feedback here. The new device bounds checks otherwise correctly prevent out-of-range source indexing and deterministically zero-fill invalid pair slices. Please add focused coverage for an empty pair list and invalid layer/head values.