Optimize Triton TBE forward for small long-bag tables (meta-pytorch#4486)

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

Differential Revision: D114273027

Author: stashuk-olek

torchrec/distributed/batched_embedding_kernel.py

@@ -3894,6 +3894,7 @@ def __init__(
         )
         output_dtype = output_dtype_sparse.as_dtype()
         stochastic_rounding = fused_params.get("stochastic_rounding", True)
+        bag_size_hints: Optional[List[int]] = fused_params.get("bag_size_hints")
 
         # Create Triton TBE module with feature_table_map for correct batch size handling
         self._emb_module: TritonTableBatchedEmbeddingBags = (
@@ -3907,6 +3908,7 @@ def __init__(
                 eps=eps,
                 optimizer=optimizer,
                 device=device,
+                bag_size_hints=bag_size_hints,
             )
         )
 

torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py

@@ -185,7 +185,10 @@ def table_batched_embedding_bag_forward_weighted_kernel(
 
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < embedding_dim
-    bag_output = tl.zeros((BLOCK_SIZE,), dtype=tl.float64)
+    accumulator_dtype: tl.constexpr = (
+        tl.float64 if weight_ptr.dtype.element_ty == tl.float32 else tl.float32
+    )
+    bag_output = tl.zeros((BLOCK_SIZE,), dtype=accumulator_dtype)
 
     # without type hint the unrolling performance will downgrade
     step: tl.constexpr = 4
@@ -270,6 +273,8 @@ def table_batched_embedding_bag_forward_unweighted_kernel(
     T: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     vbe: tl.constexpr = False,
+    FEATURE_START: tl.constexpr = 0,
+    FEATURE_END: tl.constexpr = -1,
     FUSED_BOUNDS_CHECK: tl.constexpr = False,
 ) -> None:
 
@@ -281,7 +286,8 @@ def table_batched_embedding_bag_forward_unweighted_kernel(
     else:
         output_row_base = output_ptr + b * total_embedding_dim
 
-    for t in range(T):
+    feature_end: tl.constexpr = T if FEATURE_END < 0 else FEATURE_END
+    for t in range(FEATURE_START, feature_end):
         if vbe:
             # VBE: check if this batch index is within feature t's batch size
             B_start = tl.load(B_offsets_ptr + t).to(tl.int64)
@@ -308,7 +314,10 @@ def table_batched_embedding_bag_forward_unweighted_kernel(
             end = tl.load(offsets_ptr + b_t + 1)
 
             mask = col_offsets < embedding_dim
-            bag_output = tl.zeros((BLOCK_SIZE,), dtype=tl.float64)
+            accumulator_dtype: tl.constexpr = (
+                tl.float64 if weight_ptr.dtype.element_ty == tl.float32 else tl.float32
+            )
+            bag_output = tl.zeros((BLOCK_SIZE,), dtype=accumulator_dtype)
 
             step: tl.constexpr = 4
             ns = (end - start) // step
@@ -401,6 +410,100 @@ def table_batched_embedding_bag_forward_unweighted_kernel(
         tl.atomic_add(bounds_check_warning_ptr, warning_count.to(tl.int64))
 
 
+@triton.jit
+def table_batched_embedding_bag_forward_small_table_kernel(
+    output_ptr,
+    indices_ptr,
+    offsets_ptr,
+    weight_ptr,
+    table_offsets_ptr,
+    embedding_dims_ptr,
+    embedding_offsets_ptr,
+    feature_table_map_ptr,
+    bounds_check_warning_ptr,
+    total_embedding_dim: tl.constexpr,
+    B,
+    FEATURE: tl.constexpr,
+    NUM_ROWS: tl.constexpr,
+    ROW_BINS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    FUSED_BOUNDS_CHECK: tl.constexpr = False,
+) -> None:
+    bags_per_program: tl.constexpr = 16
+    histogram_chunk_size: tl.constexpr = 256
+
+    bag_slots = tl.arange(0, bags_per_program)
+    bags = tl.program_id(0).to(tl.int64) * bags_per_program + bag_slots
+    bag_mask = bags < B
+    starts = tl.load(offsets_ptr + FEATURE * B + bags, mask=bag_mask, other=0)
+    ends = tl.load(offsets_ptr + FEATURE * B + bags + 1, mask=bag_mask, other=0)
+    lengths = ends - starts
+
+    positions = tl.arange(0, histogram_chunk_size)
+    input_mask = bag_mask[:, None] & (positions[None, :] < lengths[:, None])
+    row_indices, invalid_indices = _load_checked_index(
+        indices_ptr,
+        starts[:, None] + positions[None, :],
+        NUM_ROWS,
+        input_mask,
+        FUSED_BOUNDS_CHECK,
+    )
+    row_indices = row_indices.to(tl.int32)
+    warning_count = tl.sum(invalid_indices.to(tl.int32))
+    encoded_indices = row_indices + bag_slots[:, None] * ROW_BINS
+    counts = tl.histogram(
+        encoded_indices.reshape((bags_per_program * histogram_chunk_size,)),
+        bags_per_program * ROW_BINS,
+        mask=input_mask.reshape((bags_per_program * histogram_chunk_size,)),
+    ).reshape((bags_per_program, ROW_BINS))
+
+    table_idx = tl.load(feature_table_map_ptr + FEATURE)
+    table_offset = tl.load(table_offsets_ptr + table_idx)
+    embedding_dim = tl.load(embedding_dims_ptr + FEATURE)
+    embedding_offset = tl.load(embedding_offsets_ptr + FEATURE)
+    rows = tl.arange(0, ROW_BINS)
+    columns = tl.arange(0, BLOCK_SIZE)
+    table = tl.load(
+        weight_ptr + table_offset + rows[:, None] * embedding_dim + columns[None, :],
+        mask=(rows[:, None] < NUM_ROWS) & (columns[None, :] < embedding_dim),
+        other=0,
+    )
+    bag_output = tl.dot(counts.to(tl.float16), table)
+
+    tail_lengths = tl.maximum(lengths - histogram_chunk_size, 0)
+    for tail in range(0, tl.max(tail_lengths)):
+        active = bag_mask & (tail < tail_lengths)
+        row_idx, invalid = _load_checked_index(
+            indices_ptr,
+            starts + histogram_chunk_size + tail,
+            NUM_ROWS,
+            active,
+            FUSED_BOUNDS_CHECK,
+        )
+        if FUSED_BOUNDS_CHECK:
+            warning_count += tl.sum(invalid.to(tl.int32))
+        row = tl.load(
+            weight_ptr
+            + table_offset
+            + row_idx[:, None] * embedding_dim
+            + columns[None, :],
+            mask=active[:, None] & (columns[None, :] < embedding_dim),
+            other=0,
+        )
+        bag_output += row.to(tl.float32)
+
+    tl.store(
+        output_ptr
+        + bags[:, None] * total_embedding_dim
+        + embedding_offset
+        + columns[None, :],
+        bag_output,
+        mask=bag_mask[:, None] & (columns[None, :] < embedding_dim),
+    )
+    if FUSED_BOUNDS_CHECK and warning_count > 0:
+        tl.atomic_add(bounds_check_warning_ptr, warning_count.to(tl.int64))
+
+
 @triton.jit
 def triton_tbe_backward_short_run_unweighted(
     dout_ptr,
@@ -1153,6 +1256,10 @@ def forward(
         precomputed_total_B: int = 0,
         precomputed_max_B: int = 0,
         hoist_transpose_to_forward: bool = True,
+        histogram_feature: int = -1,
+        histogram_num_rows: int = 0,
+        histogram_row_bins: int = 0,
+        histogram_block_size: int = 0,
         bounds_check_warning: Optional[torch.Tensor] = None,
         fused_bounds_check: bool = False,
     ) -> torch.Tensor:
@@ -1368,8 +1475,16 @@ def forward(
                 num_warps=num_warps,
             )
         else:
-            if is_amd():
-                _amd_fwd_unweighted_kernel[(B,)](
+            use_small_table_kernel = (
+                histogram_feature >= 0
+                and not vbe
+                and not is_amd()
+                and weight.dtype == torch.float16
+            )
+            if use_small_table_kernel:
+                table_batched_embedding_bag_forward_small_table_kernel[
+                    (triton.cdiv(B, 16),)
+                ](
                     output,
                     indices,
                     offsets,
@@ -1378,17 +1493,19 @@ def forward(
                     embedding_dims,
                     embedding_offsets,
                     feature_table_map,
-                    row_output_offsets_ptr,
-                    B_offsets_ptr,
+                    bounds_check_warning_ptr,
                     total_embedding_dim,
                     B,
-                    T,
-                    BLOCK_SIZE=block_size,
-                    vbe=vbe,
-                    num_warps=num_warps,
+                    FEATURE=histogram_feature,
+                    NUM_ROWS=histogram_num_rows,
+                    ROW_BINS=histogram_row_bins,
+                    BLOCK_SIZE=histogram_block_size,
+                    FUSED_BOUNDS_CHECK=fused_bounds_check,
+                    num_warps=1,
                 )
-            else:
-                table_batched_embedding_bag_forward_unweighted_kernel[(B,)](
+
+            if is_amd():
+                _amd_fwd_unweighted_kernel[(B,)](
                     output,
                     indices,
                     offsets,
@@ -1397,18 +1514,50 @@ def forward(
                     embedding_dims,
                     embedding_offsets,
                     feature_table_map,
-                    rows_cumsum,
-                    bounds_check_warning_ptr,
                     row_output_offsets_ptr,
                     B_offsets_ptr,
                     total_embedding_dim,
                     B,
                     T,
                     BLOCK_SIZE=block_size,
                     vbe=vbe,
-                    FUSED_BOUNDS_CHECK=fused_bounds_check,
                     num_warps=num_warps,
                 )
+            else:
+                feature_ranges = (
+                    [
+                        (0, histogram_feature),
+                        (histogram_feature + 1, T),
+                    ]
+                    if use_small_table_kernel
+                    else [(0, T)]
+                )
+                for feature_start, feature_end in feature_ranges:
+                    if feature_start >= feature_end:
+                        continue
+                    table_batched_embedding_bag_forward_unweighted_kernel[(B,)](
+                        output,
+                        indices,
+                        offsets,
+                        weight,
+                        table_offsets,
+                        embedding_dims,
+                        embedding_offsets,
+                        feature_table_map,
+                        rows_cumsum,
+                        bounds_check_warning_ptr,
+                        row_output_offsets_ptr,
+                        B_offsets_ptr,
+                        total_embedding_dim,
+                        B,
+                        T,
+                        BLOCK_SIZE=block_size,
+                        vbe=vbe,
+                        FEATURE_START=feature_start,
+                        FEATURE_END=feature_end,
+                        FUSED_BOUNDS_CHECK=fused_bounds_check,
+                        num_warps=num_warps,
+                    )
 
         # Record a CUDA event to mark forward kernel completion.
         # This is needed for synchronization before NCCL collectives.
@@ -1906,6 +2055,10 @@ def backward(ctx, dout) -> Tuple[None, ...]:
             None,  # hoist_transpose_to_forward
             None,
             None,
+            None,
+            None,
+            None,
+            None,
         )
 
 
@@ -1927,6 +2080,7 @@ def __init__(
         optimizer: OptimType = OptimType.EXACT_SGD,
         device: Optional[torch.device] = None,
         hoist_transpose_to_forward: bool = False,
+        bag_size_hints: Optional[List[int]] = None,
         fused_bounds_check: bool = False,
     ) -> None:
         super().__init__()
@@ -2005,6 +2159,35 @@ def __init__(
         self.output_dtype = (
             output_dtype if output_dtype is not None else weights_precision
         )
+        if bag_size_hints is not None and len(bag_size_hints) != self.T:
+            raise ValueError(
+                f"bag_size_hints must have {self.T} entries, "
+                f"got {len(bag_size_hints)}"
+            )
+
+        self._histogram_feature = -1
+        self._histogram_num_rows = 0
+        self._histogram_row_bins = 0
+        self._histogram_block_size = 0
+        if (
+            bag_size_hints is not None
+            and weights_precision == torch.float16
+            and self.output_dtype == torch.float32
+        ):
+            candidates = []
+            for feature, table in enumerate(feature_table_map):
+                num_rows = hash_sizes[table]
+                dim = feature_dims[feature]
+                bag_size = bag_size_hints[feature]
+                if num_rows <= 64 and 64 <= dim <= 128 and bag_size >= 64:
+                    candidates.append((bag_size * dim, feature, num_rows, dim))
+            if candidates:
+                _, feature, num_rows, dim = max(candidates)
+                self._histogram_feature = feature
+                self._histogram_num_rows = num_rows
+                self._histogram_row_bins = max(32, triton.next_power_of_2(num_rows))
+                self._histogram_block_size = triton.next_power_of_2(dim)
+
         self.stochastic_rounding = stochastic_rounding
         self.learning_rate = learning_rate
         self.eps = eps
@@ -2269,6 +2452,10 @@ def forward(
             total_B,
             max_B,
             self.hoist_transpose_to_forward,
+            self._histogram_feature,
+            self._histogram_num_rows,
+            self._histogram_row_bins,
+            self._histogram_block_size,
             self.bounds_check_warning,
             use_fused_bounds_check,
         )