Optimize group_index_select_or_add_2d_kernel on ROCm by adding a separate codepath for small embedding dimensions#5233

Open

aryaman-gupta wants to merge 6 commits intopytorch:mainfrom

ROCm:aryaman/group-index-subwarp

Commits on Dec 12, 2025

adds optimized path for small dimension sizes to group_index_select_or_add_2d_kernel
aryaman-gupta
committed

Commits on Dec 16, 2025

sparse_group_index.cu: edits some comments
aryaman-gupta
committed
adds USE_ROCM guards to subwarp optimizations for group_index_select_or_add_2d_kernel
aryaman-gupta
committed

Commits on Dec 18, 2025