Bound sequence_token_count in CUDA RemovePadding

[Akshay Sonawane (apsonawane)]

The getTokenOffset kernel used each sequence_token_count element directly as a loop bound while writing into token_offset, which holds exactly batch_size * sequence_length entries. A value above sequence_length, or a negative one in the padding loop, walked past the end of that buffer and also produced an inconsistent total token count used to size the output.

Clamp each per-row count to [0, sequence_length] in the kernel and verify the sequence_token_count shape is (batch_size) before launching. Adds tests for out-of-range and negative counts.

[Copilot]

Pull request overview

This PR hardens the CUDA RemovePadding contrib op against malformed sequence_token_count values by clamping per-sequence token counts inside the getTokenOffset kernel and validating the sequence_token_count input shape before kernel launch, preventing out-of-bounds writes and inconsistent output sizing.

Changes:

• Clamp each sequence_token_count[i] to [0, sequence_length] inside the CUDA getTokenOffset kernel to keep writes within the token_offset buffer.
• Add runtime validation that sequence_token_count has shape (batch_size) in the CUDA kernel wrapper before launching work.
• Add unit tests covering out-of-range (too large) and negative sequence_token_count values.

Reviewed changes

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

File	Description
onnxruntime/test/contrib_ops/remove_padding_op_test.cc	Adds test coverage for oversized and negative sequence_token_count inputs to ensure safe, clamped behavior.
onnxruntime/contrib_ops/cuda/bert/remove_padding.cc	Validates sequence_token_count rank/shape matches (batch_size) prior to launching CUDA kernels.
onnxruntime/contrib_ops/cuda/bert/bert_padding.cu	Clamps per-row token counts in getTokenOffset to prevent OOB writes and keep derived token counts consistent.

---

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