Add position_ids bounds validation to WebGPU/JS RotaryEmbedding kernels

[titaiwangms]

This PR adds position_ids bounds checking to WebGPU and JS RotaryEmbedding implementations, completing the security fix started in PR #27597 (commit 056bab3) which covered CPU and CUDA.

Problem

The com.microsoft::RotaryEmbedding kernel uses position_ids as row indices into cos_cache/sin_cache without bounds validation. While PR #27597 fixed CPU and CUDA paths, WebGPU and JS implementations were still missing bounds checks, which could produce silently wrong results (WebGPU hardware clamps OOB reads).

Changes

• contrib_ops/webgpu/bert/rotary_embedding.cc: Host-side validation (ORT_MAKE_STATUS) + shader-side defense-in-depth (pass-through on OOB)
• core/providers/webgpu/llm/rotary_embedding.cc: Host-side validation with format-0 awareness
• js/web/lib/wasm/jsep/webgpu/ops/rotary-embedding.ts: TypeScript validation using getBigInt64Array
• 7 new C++ OOB test cases across contrib and ONNX domains targeting WebGPU EP

Security

Addresses the same vulnerability as #27597 (OOB read via position_ids, CVSS 7.5-9.1) for WebGPU/JS execution providers.

Testing

• 7 new unit tests (3 contrib + 4 ONNX domain) with GTEST_SKIP when WebGPU EP unavailable
• JS/TS error tests not feasible with current JSONC test format (documented)
• Build environment lacks C++20/emsdk for full compilation verification; validated structurally

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

This PR completes the RotaryEmbedding position_ids bounds-checking security fix for the WebGPU EP and the JS WebGPU implementation, aligning them with the earlier CPU/CUDA hardening work.

Changes:

• Add host-side position_ids bounds validation to WebGPU RotaryEmbedding kernels (ONNX + contrib/MS domain), plus shader-side pass-through on OOB as defense-in-depth.
• Add TypeScript-side position_ids bounds validation for the JS WebGPU RotaryEmbedding op.
• Add new WebGPU-focused unit tests covering negative and out-of-range position_ids cases.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/core/providers/webgpu/llm/rotary_embedding.cc	Adds host-side position_ids bounds validation and forces CPU input for validation.
onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc	Adds host-side bounds validation, marks position_ids as CPU input, and adds shader-side OOB pass-through.
js/web/lib/wasm/jsep/webgpu/ops/rotary-embedding.ts	Adds JS-side bounds checks for position_ids using BigInt64 reads.
onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc	Adds ONNX-domain WebGPU EP failure tests for OOB/negative position_ids.
onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc	Adds contrib/MS-domain WebGPU EP failure tests for OOB/negative position_ids.

---

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

[tianleiwu]

Review Summary

The validation logic and shader-side defense-in-depth checks are well-designed and mirror the CPU/CUDA approach correctly. The test coverage is thorough. However, both kernel registrations have a critical issue that will cause a runtime crash.

Critical: InputMemoryType(OrtMemTypeCPUInput) + AddInputs() incompatibility

Both kernel registrations add .InputMemoryType(OrtMemTypeCPUInput, ...) to keep position_ids on CPU for host-side validation. However, ComputeInternal still passes position_ids to AddInputs() for the shader. The WebGPU program execution path in webgpu_context.cc (around L464) does:

bind_buffers.push_back(reinterpret_cast<WGPUBuffer>(const_cast<void*>(inputs[i].tensor->DataRaw())));

This assumes all program inputs have GPU buffer handles in DataRaw(). A CPU-resident tensor's DataRaw() returns a heap pointer, producing an invalid WGPUBuffer handle that will crash or corrupt at dispatch.

No existing WebGPU kernel combines InputMemoryType(OrtMemTypeCPUInput) with AddInputs() on the same tensor. Existing kernels (Reduce, Pad, Resize, GQA, etc.) read CPU-pinned inputs for host-side configuration only and never submit them as shader bindings.

Recommended fix

The simplest approach is to remove both InputMemoryType directives and the host-side C++ validation, relying solely on the shader-side defense-in-depth checks (which are already correct in this PR). This matches CUDA's approach: device-side bounds checking with pass-through on OOB. The TS validation in rotary-embedding.ts can remain since it runs on CPU during validateInputs() before shader dispatch.

If host-side rejection is desired, the kernel would need to explicitly copy position_ids from GPU to CPU (via DataTransferManager) for validation, then pass the original GPU tensor to AddInputs().

What's good

• Shader-side OOB checks: raw_pos < 0 catches negative i32 from truncated int64; position_id >= max_position catches OOB after u32 conversion + offset. Pass-through behavior matches CUDA.
• FusedQKRotaryEmbeddingProgram bounds check on computed position_id correctly handles both Q and K outputs, including the conditional K-head guard.
• TS validation using getBigInt64Array() avoids int64 truncation issues.
• 7 new tests cover format-0, format-1, negative, and batch OOB comprehensively.

[tianleiwu]

Round-2 review (head 9744bfb).

Summary

The high-severity concern from the prior round — InputMemoryType(OrtMemTypeCPUInput) + AddInputs(position_ids) would crash because the same tensor cannot be both CPU-resident and bound as a WebGPU storage buffer — has been resolved. The author removed the host-side C++ value scanning and InputMemoryType directives and now relies on shader-side defense-in-depth (OOB → pass-through input) for both contrib and ONNX-domain WebGPU kernels, plus host-side validation on the JS path via getBigInt64Array() in validateInputs. Tests were converted to kExpectSuccess with output == input. The pass-through semantics now match the CUDA kernel for cross-EP consistency.

Resolving my two prior threads:

• contrib_ops/webgpu/bert/rotary_embedding.cc (shader as primary protection) — addressed.
• test/contrib_ops/rotary_embedding_op_test.cc (tests must verify pass-through, not failure) — addressed.

Remaining items (all nits, non-blocking)

1. Shader comment understates what raw_pos < 0 catches — onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc

The WebGPU shader storage representation truncates int64 to i32 (see ProgramVariableDataType::Int64 → value type i32 in core/providers/webgpu/shader_variable.cc). Any positive int64 value greater than INT32_MAX will alias to a negative i32 and be caught by the raw_pos < 0 branch. That is safe (still pass-through), but the inline comment only mentions "negative position_ids". Consider expanding the comment to note that very large positive int64 values also fall into this branch by virtue of i32 truncation, so a future maintainer doesn't think the check is incomplete.

2. JSEP getBigInt64Array alignment fallback returns a read-only copy — js/web/lib/wasm/jsep/init.ts and js/web/lib/wasm/jsep/tensor-view.ts

The inline comment in init.ts correctly notes that the unaligned path returns a read-only copy (mutations don't propagate to the WASM heap). All current call sites (pad.ts, split.ts, reduce.ts, slice.ts, tile.ts, expand.ts, cumsum.ts, resize.ts, gather-block-quantized.ts, the new rotary-embedding.ts) are read-only iterations, so this is safe today. Consider promoting that read-only caveat into the TensorView.getBigInt64Array() JSDoc in tensor-view.ts so future consumers don't accidentally write through the returned view expecting WASM-heap propagation.

3. Format-0 detection in TS is contrib-op-specific — js/web/lib/wasm/jsep/webgpu/ops/rotary-embedding.ts

positionIdsElementCount === 1 triggers the format-0 base-offset branch. The ONNX-domain RotaryEmbedding schema requires rank-2 position_ids (no base-offset mode). If this validateInputs is shared with the ONNX-domain op (now or later), the format-0 branch should be guarded by op identity. A short comment confirming the JSEP path scope would prevent silent acceptance of malformed shapes against the ONNX-domain op.

4. Pre-existing related gap (informational) — separately noted by the Copilot reviewer at L159 of contrib_ops/webgpu/bert/rotary_embedding.cc: when position_ids is provided, CPU/CUDA reject sequence_length > max_sequence_length with NOT_IMPLEMENTED, while the WebGPU contrib path (after this PR) silently produces pass-through for the OOB tokens. This was the case before this PR too. Worth tracking as a follow-up if cross-EP behavior parity is desired here, but it does not need to block this security fix.

Verdict

Looks good. The shader bounds checks are correctly designed (both for the standalone RotaryEmbeddingProgram and the FusedQKRotaryEmbeddingProgram Q+K pass-through path), tests are appropriate, and the JS validation is functionally correct.