Gate on avg-SL threshold via zero-stall prev-iter caching

Author: avbokovoy

fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu

@@ -833,6 +833,19 @@ Tensor {{ embedding_cuda_op }}(
 #endif
     const int used_shared_bytes = used_shared_kb << 10;
 
+#ifdef USE_ROCM
+    {%- if enable_optimized_hip_mixed_D_kernel %}
+    const bool use_hip_kernel_flag = fbgemm_gpu::config::is_feature_enabled(
+        fbgemm_gpu::config::FeatureGateName::TBE_ROCM_HIP_BACKWARD_KERNEL);
+    // Declared here so the dispatch lambdas below can capture it.
+    // Populated from the pinned buffer after transpose_embedding_input().
+    struct NumUniqueEntry { at::Tensor pinned_buf; };
+    static thread_local std::unordered_map<const void*, NumUniqueEntry>
+        s_num_unique_cache;
+    int64_t num_unique_prev = 0;
+    {%- endif %}
+#endif // USE_ROCM
+
     Tensor linear_indices, linear_indices_sorted, infos_sorted,
         sorted_linear_indices_run, sorted_linear_indices_run_lengths,
         sorted_linear_indices_num_runs,
@@ -865,6 +878,31 @@ Tensor {{ embedding_cuda_op }}(
             {%- endif %}
         );
 
+#ifdef USE_ROCM
+    {%- if enable_optimized_hip_mixed_D_kernel %}
+    if (use_hip_kernel_flag) {
+        auto& entry = s_num_unique_cache[hash_size_cumsum.data_ptr()];
+        if (entry.pinned_buf.defined()) {
+            // Read the exact num_unique from the previous backward call.
+            num_unique_prev =
+                static_cast<int64_t>(entry.pinned_buf.data_ptr<int32_t>()[0]);
+        } else {
+            // First call: allocate pinned buffer.
+            entry.pinned_buf = at::empty(
+                {1},
+                at::TensorOptions()
+                    .dtype(at::kInt)
+                    .device(at::kCPU)
+                    .pinned_memory(true));
+        }
+        // Schedule async D2H copy for the next backward call.
+        entry.pinned_buf.copy_(
+            sorted_linear_indices_num_runs.slice(0, 0, 1),
+            /*non_blocking=*/true);
+    }
+    {%- endif %}
+#endif // USE_ROCM
+
     {%- if not dense %}
     Tensor {{ locs_or_addrs_tensor }}_sorted = {{ locs_or_addrs_tensor }};
     Tensor table_unique_indices_offsets;
@@ -1348,15 +1386,13 @@ Tensor {{ embedding_cuda_op }}(
                     // condition 'use_hip_kernel' is True and "is_optimized_hip_kernel_supported_mode" is True. If no optimization is available for current
                     // condition , it will fallback to the default kernel.
                     {%- if vbe %}
-                    // Apply the same avg_SL threshold as the non-VBE mixed_D path.
                     if (use_hip_kernel &&
-                        total_L <= 2 * static_cast<int64_t>(sorted_linear_indices_num_runs[0].item<int32_t>())) {
+                        num_unique_prev > 0 &&
+                        total_L <= 2 * num_unique_prev) {
                     {%- else %}
-                    // Use hip_mixed_d only when avg segment length <= 2
-                    // (total_L / num_unique_rows <= 2), i.e. when the momentum preload
-                    // benefit outweighs the serial inner-loop serialization cost.
                     if (use_hip_kernel && mixed_D &&
-                        total_L <= 2 * static_cast<int64_t>(sorted_linear_indices_num_runs[0].item<int32_t>())) {
+                        num_unique_prev > 0 &&
+                        total_L <= 2 * num_unique_prev) {
                     {%- endif %}
                         backward_warp_per_row_kernel =
                         {{ hip_mixed_d_warp_kernel }}