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