SIGN IN SIGN UP

Fix CUDA illegal memory access in embedding_bag per_sample_weights backward (#193408)

Fixes #192445.

## Root cause

`_embedding_bag_per_sample_weights_backward_kernel` in
`aten/src/ATen/native/cuda/EmbeddingBag.cu` indexed the embedding `weight` by
`embedding_idx = indices[sample_idx]` with no bounds check, unlike the forward
`EmbeddingBag` kernels in the same file which validate the index against
`numRows`. When an index is out of range, the kernel reads past the `weight`
allocation, causing an illegal global memory access (undefined behavior /
potential silent corruption). This can be triggered when `indices` is mutated in
place after the forward pass (e.g. via a DLPack alias), so the bad value is only
seen at backward time.

A secondary defect was the lossy `const int embedding_idx = (int)indices[...]`
cast: for int64 indices this truncates to 32 bits, so a large in-range int64
index could alias the wrong row or defeat a naive check.

## Fix

Thread `weight.size(0)` into the kernel as `num_embeddings`, read the index with
its native `index_t` type (no truncation), and add a `CUDA_KERNEL_ASSERT` bounds
check against `[0, num_embeddings)` before the weight read, using the same
message convention as the forward kernels. This converts the illegal access into
a deterministic device-side assert.

## Impact

The change is confined to the CUDA per-sample-weights backward path. In-range
indices behave exactly as before (the assert is a no-op on the valid range); the
only behavioral change is that a previously-illegal out-of-range read now fails
fast with a device-side assert instead of reading out of bounds.

## Validation

New device-generic test `test_embedding_bag_per_sample_weights_mutated_indices`
reproduces the issue out-of-process (a device-side assert corrupts the CUDA
context) by mutating `indices` via a DLPack alias after the forward pass, then
running backward, and asserts the device-side assert markers appear in stderr.
It is gated `@onlyOn(["cuda"])` for both int32 and int64.

```
python test/nn/test_embedding.py -k test_embedding_bag_per_sample_weights_mutated_indices -v
```

Validation limitation: the CUDA test and `lintrunner -a` could not be run in the
authoring environment (macOS, no CUDA/GPU, PyTorch not built, lintrunner/spin
unavailable). These need to be run on a CUDA box; CI is expected to exercise them
here.

---

This pull request was prepared with the assistance of an AI coding assistant. The
author has reviewed and understands the code change and takes responsibility for
it.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/193408
Approved by: https://github.com/eqy
Y
Yunqian Fan committed
550fa843275302768fc1fe758f308b2dccf543b2
Parent: 9bb4e0c
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/24/2026, 9:43:18 AM