Skip to content

Add device-side bounds checks to permute_multi_embs_kernel#5806

Open
efiks wants to merge 1 commit into
pytorch:mainfrom
efiks:export-D107013820
Open

Add device-side bounds checks to permute_multi_embs_kernel#5806
efiks wants to merge 1 commit into
pytorch:mainfrom
efiks:export-D107013820

Conversation

@efiks
Copy link
Copy Markdown
Contributor

@efiks efiks commented May 31, 2026

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2732

Add device-side safety checks to permute_multi_embs_kernel to prevent CUDA illegal memory access crashes when permute arguments are stale or corrupt:

  1. Forward-path bounds check: After parsing the permutes table, validate that in_tensor and out_tensor indices are within the bounds of in_lengths and out_lengths before dereferencing device pointers. Threads with invalid indices return early instead of crashing.

  2. Reverse-permute loop guard: Add an iteration counter capped at permute_size to prevent infinite loops when the permutes linked-list next chain is corrupt. Also add an in_tensor bounds check inside the loop body to bail on corrupt data rather than performing OOB device memory access.

Differential Revision: D107013820

Summary:
X-link: facebookresearch/FBGEMM#2732

Add device-side safety checks to `permute_multi_embs_kernel` to prevent CUDA illegal memory access crashes when permute arguments are stale or corrupt:

1. **Forward-path bounds check**: After parsing the permutes table, validate that `in_tensor` and `out_tensor` indices are within the bounds of `in_lengths` and `out_lengths` before dereferencing device pointers. Threads with invalid indices return early instead of crashing.

2. **Reverse-permute loop guard**: Add an iteration counter capped at `permute_size` to prevent infinite loops when the permutes linked-list `next` chain is corrupt. Also add an `in_tensor` bounds check inside the loop body to bail on corrupt data rather than performing OOB device memory access.

Differential Revision: D107013820
@meta-cla meta-cla Bot added the cla signed label May 31, 2026
@meta-codesync
Copy link
Copy Markdown
Contributor

meta-codesync Bot commented May 31, 2026

@efiks has exported this pull request. If you are a Meta employee, you can view the originating Diff in D107013820.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant