From 297ad2df002d13cb9a7a4993def5116762a00aea Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Tue, 28 Jul 2026 14:41:22 -0700 Subject: [PATCH] Add grid-stride + ROCm cap to split optimizer update kernel Summary: The standalone split optimizer path launches split_{optimizer}_update_kernel with grid = div_round_up(grad_dev_indices.numel(), kMaxThreads / kThreadGroupSize) and block = dim3(kThreadGroupSize, kMaxThreads / kThreadGroupSize, 1), so total threads ~= num_unique_indices * kThreadGroupSize exceeds the HIP 2^32 threads-per-launch limit on ROCm for large index counts. Cap the launch with utils::cuda::cap_grid_dim_x(..., OverflowOnly) and add a ROCm grid-stride loop over run_id to the kernel (single guard converted to the loop bound; no internal early-returns). No-op on CUDA. Reviewed By: henrylhtsang Differential Revision: D113351687 --- .../embedding_optimizer_split_kernel_template.cu | 12 ++++++++++++ .../optimizer/embedding_optimizer_split_template.cu | 5 ++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/fbgemm_gpu/codegen/training/optimizer/embedding_optimizer_split_kernel_template.cu b/fbgemm_gpu/codegen/training/optimizer/embedding_optimizer_split_kernel_template.cu index 1e6b089b16..ca27702012 100644 --- a/fbgemm_gpu/codegen/training/optimizer/embedding_optimizer_split_kernel_template.cu +++ b/fbgemm_gpu/codegen/training/optimizer/embedding_optimizer_split_kernel_template.cu @@ -37,10 +37,19 @@ void split_{{ optimizer }}_update_kernel( at::PhiloxCudaState stochastic_rounding_philox_args, {{ args.split_kernel_args | replace_pta_namespace() | join(",\n ") }} ) { + // On ROCm the launch caps the grid to stay within the HIP 2^32 + // threads-per-launch limit, so we grid-stride to cover the full workload. + // On CUDA the grid is not capped and the loop body runs once per warp. +#ifdef USE_ROCM + for (auto run_id = blockIdx.x * blockDim.y + threadIdx.y; + run_id < grad_dev_indices.size(0); + run_id += blockDim.y * gridDim.x) { +#else const auto run_id = blockIdx.x * blockDim.y + threadIdx.y; if (run_id >= grad_dev_indices.size(0)) { return; } +#endif #ifdef FBGEMM_USE_SUBWARP_SHUFFLE const unsigned int shfl_sync_mask = @@ -99,6 +108,9 @@ void split_{{ optimizer }}_update_kernel( shfl_sync_mask, kMaxVecsPerThread, {{ args.split_kernel_arg_names | join(", ") }}); +#ifdef USE_ROCM + } // for run_id (grid-stride loop, ROCm only) +#endif } {%- for use_subwarp in [True, False] %} diff --git a/fbgemm_gpu/codegen/training/optimizer/embedding_optimizer_split_template.cu b/fbgemm_gpu/codegen/training/optimizer/embedding_optimizer_split_template.cu index 1de913efad..bea8a58fe4 100644 --- a/fbgemm_gpu/codegen/training/optimizer/embedding_optimizer_split_template.cu +++ b/fbgemm_gpu/codegen/training/optimizer/embedding_optimizer_split_template.cu @@ -192,7 +192,10 @@ void split_embedding_{{ optimizer }}_update( kMaxVecsPerThread, kThreadGroupSize, 4>), - div_round_up(grad_dev_indices.numel(), kMaxThreads / kThreadGroupSize), + utils::cuda::cap_grid_dim_x( + div_round_up(grad_dev_indices.numel(), kMaxThreads / kThreadGroupSize), + kMaxThreads, + at::cuda::getCurrentCUDAStream()), dim3(kThreadGroupSize, kMaxThreads / kThreadGroupSize, 1), 0, // Shared memory is not needed because uint8_t is not supported at::cuda::getCurrentCUDAStream(),