Skip to content

Use 64-bit indices in cdist backward kernel to fix large-input OOB write - #6174

Open
DamJanusz wants to merge 1 commit into
mainfrom
djanuszx/cdist_64bit_ind
Open

DamJanusz wants to merge 1 commit into
mainfrom
djanuszx/cdist_64bit_ind

Conversation

@DamJanusz

@DamJanusz DamJanusz commented Nov 12, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #5660

Root cause

cdist_backward_kernel_impl_func held its per-row index y, batch index l, in-batch index k and per-batch buffer stride l_size in 32-bit int. Two shapes overflow:

  • (b=2, r=1024, m=2048): l_size = r1 * r2 * m = 2^31 overflows, so l * l_size wraps and buff_i points outside the buffer.
  • (b=32, r=8192, m=1): count = dist.numel() = 2^31, so y itself truncates.

CUDA had the same bug pytorch/pytorch#128791 and fixed it in pytorch/pytorch#198452.

Changes

After this change the XPU index math matches aten/src/ATen/native/cuda/DistanceKernel.cu (cdist_backward_kernel_cuda_impl)

PR to remove @skipIfXpu on upstream pytorch/pytorch#204593 needs to land after this change propagates to commit pin.

Testing

The regression test lives in PyTorch, not torch-xpu-ops: test/test_torch.py::TestTorchDeviceTypeXPU::test_cdist_backward_large_index_* is currently disabled on XPU with @skipIfXpu(msg=".../issues/5660"). Removing the skip before the pin bump would crash XPU CI with a GPU segfault.

  cd pytorch/test
  pytest -q test_torch.py -k "test_cdist_backward_large_index and xpu"
  # 2 passed, 2083 deselected

@github-actions github-actions Bot added disable_e2e Disable all e2e test jobs for the PR disable_distributed Disable distributed UT test jobs for the PR labels Nov 12, 2026
  Fixes #5660

  `y`, `l`, `k` and `l_size` in cdist_backward_kernel_impl_func were `int`.
  For (b=2, r=1024, m=2048), r_size * m = 2^31 overflows, so `l * l_size`
  wraps and the buffer write lands out of bounds (GPU page fault, process
  abort). For (b=32, r=8192, m=1), count = dist.numel() = 2^31 and `y`
  itself truncates. Widening these to int64_t mirrors the CUDA fix in
  pytorch/pytorch#128791. `init`/`stride` stay `int`; they are bounded by
  the dim-0 global range, matching upstream.

  Test Plan:
  With @skipIfXpu removed locally from test_cdist_backward_large_index:

      cd pytorch/test
      pytest -q test_torch.py -k "test_cdist_backward_large_index and xpu"
      # 2 passed

@dvrogozh dvrogozh 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.

LGTM, @DamJanusz I've slightly modified PR description for the clarity.

@dvrogozh

Copy link
Copy Markdown
Contributor

don't close issue before test is unskipped

@DamJanusz , this comment does not belong to PR description as it will be used as a commit message (consider: what use will reader have from this message reading git log?). Leave a comment in #5660 or whatever issue you've meant as the one not to be closed.

@dvrogozh
dvrogozh requested a review from pbielak November 12, 2026 22:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

disable_distributed Disable distributed UT test jobs for the PR disable_e2e Disable all e2e test jobs for the PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[upstream_ut] test_cdist_backward_large_index GPU segfault: int32 overflow in XPU cdist backward buffer offset

2 participants