Skip to content

Add grid-stride processing to jagged unique-index kernels - #6344

Open
q10 wants to merge 1 commit into
pytorch:mainfrom
q10:export-D116334991
Open

q10 wants to merge 1 commit into
pytorch:mainfrom
q10:export-D116334991

Conversation

@q10

@q10 q10 commented Sep 21, 2026

Copy link
Copy Markdown
Contributor

Summary:
NOTE: No linked task. Please associate a task with this diff.

Add grid-stride loops to the linearization and group-length kernels so large launches remain correct when grid.x is capped. On ROCm, apply the occupancy/safety cap on every launch; CUDA retains overflow-only behavior.

Preserve empty groups and validate that sample counts fit 32-bit accessors, batch size defines at least one table, and HIP table launches stay within the inclusive thread limit. Add a ROCm regression sized one block beyond the cap so both kernels execute a second grid-stride iteration.

Reviewed By: jianyuh

Differential Revision: D116334991

Summary:
NOTE: No linked task. Please associate a task with this diff.

Add grid-stride loops to the linearization and group-length kernels so large launches remain correct when grid.x is capped. On ROCm, apply the occupancy/safety cap on every launch; CUDA retains overflow-only behavior.

Preserve empty groups and validate that sample counts fit 32-bit accessors, batch size defines at least one table, and HIP table launches stay within the inclusive thread limit. Add a ROCm regression sized one block beyond the cap so both kernels execute a second grid-stride iteration.

Reviewed By: jianyuh

Differential Revision: D116334991
@meta-cla meta-cla Bot added the cla signed label Sep 21, 2026
@meta-codesync

meta-codesync Bot commented Sep 21, 2026

Copy link
Copy Markdown
Contributor

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

This branch has not been deployed

No deployments
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