Enable subwarp only for unweighted (#2051)

Summary:
Pull Request resolved: #2051

Disable the subwarp optimization for weighted TBE

- The cause of the NE issue has to be investigated.
- This does not affect the overall E2E performance much since the
  majority of TBEs in the model is unweighted

Reviewed By: jasonjk-park

Differential Revision: D49159119

fbshipit-source-id: 805639d94d4ce7b3be8f275db4dfd0ecc95a539a

Author: sryap

fbgemm_gpu/codegen/embedding_forward_split_kernel_v2_template.cu

@@ -940,6 +940,7 @@ __global__ void split_embedding_codegen_forward_{{ wdesc }}_v2_kernel(
 
     // Tail warp
     // STEP_MASK computation assumes STEP = 4
+    {% if not weighted %}
     if (load_D - load_d < kWarpSize) {
       const auto tail_warp_size = load_D % kWarpSize;
       if (tail_warp_size <= 8) {
@@ -955,6 +956,9 @@ __global__ void split_embedding_codegen_forward_{{ wdesc }}_v2_kernel(
     else {
       INVOKE_PROCESS_ALL_INDICES(large_Ls, 32, 0xf)
     }
+    {% else %}
+    INVOKE_PROCESS_ALL_INDICES(large_Ls, 32, 0xf)
+    {% endif %}
 
 #undef INVOKE_PROCESS_ALL_INDICES_HELPER
 #undef INVOKE_PROCESS_ALL_INDICES