Skip to content

fix oob read access in GatherBlockQuantized#28718

Merged
guschmue merged 1 commit into
mainfrom
gs/wgpu-gatherblkq-robustness
Jun 9, 2026
Merged

fix oob read access in GatherBlockQuantized#28718
guschmue merged 1 commit into
mainfrom
gs/wgpu-gatherblkq-robustness

Conversation

@guschmue

Copy link
Copy Markdown
Contributor

GatherBlockQuantized shader could read oob for ill formed indicies.
The correct solution would be to throw and exception but the indicies input sits on gpu and we don't want to take the performance hit by copying it to cpu.
Instead, we only check inside the shader to not read oob.

@guschmue guschmue added the ep:WebGPU ort-web webgpu provider label May 29, 2026
@guschmue
guschmue requested a review from Copilot June 1, 2026 15:50

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Fixes potential out-of-bounds reads in the WebGPU GatherBlockQuantized shader when indices fall outside the gather-axis range. Since the indices live on the GPU, the check is performed inside the shader (rather than on CPU) to avoid a host copy; invalid indices simply produce zero outputs.

Changes:

  • Compute gather_axis_dim once, then bounds-check the (possibly negative-adjusted) index and write 0 + early return on OOB.
  • Add a WebGPU regression test exercising both positive (2) and negative (-3) out-of-range indices for int32 and int64 index types on 2-bit/uint8 data.

Reviewed changes

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

File Description
onnxruntime/contrib_ops/webgpu/quantization/gather_block_quantized.cc Adds shader-side bounds check after negative-index normalization; OOB lanes write 0 and return.
onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc New WebGpu_InvalidIndices_2Bits_Uint8 test covering invalid indices for both int32 and int64 index tensors.

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

@guschmue
guschmue merged commit 17724b3 into main Jun 9, 2026
92 of 93 checks passed
@guschmue
guschmue deleted the gs/wgpu-gatherblkq-robustness branch June 9, 2026 18:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ep:WebGPU ort-web webgpu provider

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants