Fix HIP 2^32 grid-overflow in batch_index_select_dim0 forward kernels (#6053)

[cx-yin]

Summary:

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

The batch_index_select_dim0 forward kernel crashes on AMD/ROCm for large workloads. MAST job fire-changbin-train-UFM-10EBF_full_length-4d036 failed with:
[batch_index_select_dim0_codegen_forward_kernel] [grid dim 10773753 x 1 x 1] [block dim 64 x 8 x 1]:
Total number of threads 5516161536 is greater than the HIP limit of 2^32

Mirror the D94944619 cap-plus-grid-stride pattern for the index-select forward:

1. Launch sites (embedding_forward_split_template.cu): wrap both the small- and main-kernel grids in utils::cuda::cap_grid_dim_x(..., BlockCapPolicy::OverflowOnly), gated on is_index_select.
2. Kernels (embedding_forward_split_kernel_template.cu, ..._nobag_small_template.cu): add an #ifdef USE_ROCM grid-stride loop for the is_index_select path; convert t >= T -> break (t is monotonic in b_t) and L_start >= total_L -> continue.

Also complete the grid-x cap for permute_multi_embs_kernel (permute_multi_embedding_ops.cu), which shares this ROCm CI test suite. The prior cap (PR #6040) passed only kMaxThreads as the per-grid-x thread count, ignoring the non-grid-strided batch_id axis (gridDim.y * gridDim.z). For large permute_size with batch_size > 1 the launch still exceeded 2^32 (test_permute_multi_embedding_large_grid: grid 262145 x 32 x 1, block 64 x 16 x 1 -> 8.59e9 threads). Fold grid_dim_y * grid_dim_z into the cap's per-grid-x thread count, matching the sparse_index_add.cu / jagged_softmax precedent. This unblocks the ROCm CI test suite.

Reviewed By: q10

Differential Revision: D113306267

[meta-codesync]

@cx-yin has exported this pull request. If you are a Meta employee, you can view the originating Diff in D113306267.