Reduce Triton TBE unweighted backward reduction width (meta-pytorch#4488)

Summary:

We could drop that, small tuning on buffer_size.

Reviewed By: axeisghost

Differential Revision: D114403703

Author: stashuk-olek

torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py

@@ -541,7 +541,7 @@ def triton_tbe_backward_short_run_unweighted(
 ) -> None:
     """Backward kernel for short runs only. Each program handles one short run."""
     col_offsets = tl.arange(0, BLOCK_SIZE)
-    buffer_size: tl.constexpr = 16
+    buffer_size: tl.constexpr = 8
     buffer_offsets = tl.arange(0, buffer_size)
 
     if USE_CLC:
@@ -1037,7 +1037,7 @@ def triton_tbe_backward_long_run_grad_accum_unweighted(
     and atomically adds the partial result into a temp gradient buffer.
     """
     col_offsets = tl.arange(0, BLOCK_SIZE)
-    buffer_size: tl.constexpr = 16
+    buffer_size: tl.constexpr = 8
     buffer_offsets = tl.arange(0, buffer_size)
 
     pid = tl.program_id(0)

torchrec/distributed/triton_tbe/triton_tbe_backward_long_run_fused.py

@@ -62,7 +62,7 @@ def triton_tbe_backward_long_run_fused_weighted(
     applies the optimizer update — eliminating a separate apply kernel launch.
     """
     col_offsets = tl.arange(0, BLOCK_SIZE)
-    buffer_size: tl.constexpr = 16
+    buffer_size: tl.constexpr = 8
     buffer_offsets = tl.arange(0, buffer_size)
 
     clc_phase_producer = 1