From 803aca42a986990dab49891db3738377278f0e86 Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Mon, 27 Jul 2026 10:28:42 -0700 Subject: [PATCH] Add grid-stride + ROCm cap to grad_mean kernel Summary: The MEAN-pooling TBE backward launches grad_mean{,_vbe}_kernel with grid = div_round_up(total_B, kMaxThreads / grad_mean_warp_size) and block = dim3(grad_mean_warp_size, kMaxThreads / grad_mean_warp_size) (total_B = B * T), so total threads ~= total_B * grad_mean_warp_size exceeds the HIP 2^32 threads-per-launch limit on ROCm for large total_B. Cap the launch with utils::cuda::cap_grid_dim_x(..., OverflowOnly) and add a ROCm grid-stride loop over b_t to the kernel (no internal early-returns; single guard converted to the loop bound). No-op on CUDA. Reviewed By: cthi Differential Revision: D113351689 --- .../embedding_backward_split_grad_template.cu | 14 +++++++++++++- .../backward/embedding_backward_split_template.cu | 5 ++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_grad_template.cu b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_grad_template.cu index 5b2aefe8d7..a2653de0ca 100644 --- a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_grad_template.cu +++ b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_grad_template.cu @@ -158,14 +158,23 @@ __global__ __launch_bounds__(kMaxThreads) void grad_mean{{ vdesc }}_kernel( {% endif %} ) { int32_t T = D_offsets.size(0) - 1; - auto b_t = blockIdx.x * blockDim.y + threadIdx.y; [[maybe_unused]] int32_t b; int32_t t; const auto total_B = offsets.size(0) - 1; + // 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 b_t = blockIdx.x * blockDim.y + threadIdx.y; + b_t < total_B; + b_t += blockDim.y * gridDim.x) { +#else + auto b_t = blockIdx.x * blockDim.y + threadIdx.y; if (b_t >= total_B) { return; } +#endif {% if vbe %} const auto info = reinterpret_cast(&b_t_map[b_t])[0]; @@ -205,6 +214,9 @@ __global__ __launch_bounds__(kMaxThreads) void grad_mean{{ vdesc }}_kernel( grad_out_vec.store(&shifted_grad_output_mean[d * 4]); } } +#ifdef USE_ROCM + } // for b_t (grid-stride loop, ROCm only) +#endif } //////////////////////////////////////////////////////////////////////////////// diff --git a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu index e905484b49..4ec27e7f12 100644 --- a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu +++ b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu @@ -952,7 +952,10 @@ Tensor {{ embedding_cuda_op }}( const int grad_mean_warp_size = at::cuda::warp_size(); FBGEMM_LAUNCH_KERNEL( (grad_mean{{ vdesc }}_kernel), - div_round_up(total_B, kMaxThreads / grad_mean_warp_size), + utils::cuda::cap_grid_dim_x( + div_round_up(total_B, kMaxThreads / grad_mean_warp_size), + kMaxThreads, + at::cuda::getCurrentCUDAStream()), dim3(grad_mean_warp_size, kMaxThreads / grad_mean_warp_size), 0, at::cuda::getCurrentCUDAStream(),