From 1c949843bcccb1015a24db79f80a94c81fd32fa7 Mon Sep 17 00:00:00 2001 From: Oleksandr Stashuk Date: Fri, 7 Aug 2026 23:13:46 -0700 Subject: [PATCH 1/8] Select CUDA v2 bounds checking for Triton TBE (#4484) Summary: Use BoundsCheckMode.V2 for Triton TBE standalone index validation. This preserves the existing warning and offset-repair behavior while reducing validation overhead before Triton forward. Reviewed By: axeisghost, TroyGarden Differential Revision: D114279989 --- .../triton_table_batched_embeddings.py | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py index 8c90cdfcd..99f608d3b 100644 --- a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py +++ b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py @@ -1891,10 +1891,16 @@ def __init__( rows_per_feature, dtype=torch.int64, device=device ) self.bounds_check_warning = torch.tensor([0], device=device, dtype=torch.int64) - # Use WARNING mode by default. We don't support environment variable override - # because TritonTBE only uses bounds check v1 kernel, while the env var - # (FBGEMM_TBE_BOUNDS_CHECK_MODE) can also set v2 modes (V2_IGNORE, V2_WARNING, V2_FATAL). - self.bounds_check_mode: BoundsCheckMode = BoundsCheckMode.WARNING + self.bounds_check_mode: BoundsCheckMode = BoundsCheckMode.V2_WARNING + + def _bounds_check_config(self) -> Tuple[BoundsCheckMode, int]: + is_v2 = self.bounds_check_mode.name.startswith("V2_") + mode = ( + BoundsCheckMode[self.bounds_check_mode.name[3:]] + if is_v2 + else self.bounds_check_mode + ) + return mode, 1 + int(is_v2) def prepare_inputs( self, @@ -1929,14 +1935,16 @@ def prepare_inputs( ): per_sample_weights = per_sample_weights.float() - if self.bounds_check_mode != BoundsCheckMode.NONE: + bounds_check_mode, bounds_check_version = self._bounds_check_config() + if bounds_check_mode != BoundsCheckMode.NONE: torch.ops.fbgemm.bounds_check_indices( self.rows_per_table, indices, offsets, - self.bounds_check_mode, + bounds_check_mode, self.bounds_check_warning, per_sample_weights, + bounds_check_version=bounds_check_version, ) return indices, offsets, per_sample_weights @@ -2023,12 +2031,13 @@ def forward( ) # Bounds check (VBE-aware) - if self.bounds_check_mode != BoundsCheckMode.NONE: + bounds_check_mode, bounds_check_version = self._bounds_check_config() + if bounds_check_mode != BoundsCheckMode.NONE: torch.ops.fbgemm.bounds_check_indices( self.rows_per_table, indices, offsets, - self.bounds_check_mode, + bounds_check_mode, self.bounds_check_warning, per_sample_weights, B_offsets=vbe_metadata.B_offsets if vbe_metadata is not None else None, @@ -2036,6 +2045,7 @@ def forward( b_t_map=b_t_map, info_B_num_bits=info_B_num_bits if info_B_num_bits > 0 else -1, info_B_mask=info_B_mask if info_B_mask > 0 else -1, + bounds_check_version=bounds_check_version, ) return TritonTBE.apply( From fe73e5d1ec7301a9e731eae53abe364cd200aefd Mon Sep 17 00:00:00 2001 From: Oleksandr Stashuk Date: Fri, 7 Aug 2026 23:13:46 -0700 Subject: [PATCH 2/8] Fuse Triton TBE forward bounds checking (#4485) Summary: Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass. Reviewed By: TroyGarden, axeisghost Differential Revision: D114284802 --- .../triton_table_batched_embeddings.py | 253 +++++++++++++++--- 1 file changed, 221 insertions(+), 32 deletions(-) diff --git a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py index 99f608d3b..7d3e8d17f 100644 --- a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py +++ b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py @@ -19,6 +19,7 @@ import torch import triton # @manual import triton.language as tl # @manual +from fbgemm_gpu.config import FeatureGateName from fbgemm_gpu.split_embedding_configs import EmbOptimType as OptimType, SparseType from fbgemm_gpu.split_table_batched_embeddings_ops_common import ( BoundsCheckMode, @@ -77,6 +78,67 @@ def lengths_to_offsets(lengths: List[int], keep_last: bool = False) -> List[int] return offsets +@triton.jit +def _bounds_check_offsets_kernel( + offsets_ptr, + warning_ptr, + num_indices, + total_B, + BLOCK_SIZE: tl.constexpr, +) -> None: + positions = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + active = positions < total_B + starts = tl.load(offsets_ptr + positions, mask=active, other=0) + raw_ends = tl.load(offsets_ptr + positions + 1, mask=active, other=0) + ends = tl.where(positions == total_B - 1, num_indices, raw_ends) + invalid = active & ((starts < 0) | (starts > ends) | (ends > num_indices)) + if tl.program_id(0) == 0: + last_offset = tl.load(offsets_ptr + total_B) + if last_offset != num_indices: + tl.atomic_add(warning_ptr, 1) + + warning_count = tl.sum(invalid.to(tl.int64)) + if warning_count > 0: + tl.atomic_add(warning_ptr, warning_count) + + +@triton.jit +def _repair_offsets_kernel( + offsets_ptr, + warning_ptr, + num_indices, + total_B, +) -> None: + if tl.load(warning_ptr) > 0: + current = tl.maximum(0, tl.minimum(tl.load(offsets_ptr), num_indices)) + tl.store(offsets_ptr, current) + position = 0 + while position < total_B: + raw_end = tl.load(offsets_ptr + position + 1) + raw_end = tl.where(position == total_B - 1, num_indices, raw_end) + current = tl.maximum(current, tl.minimum(raw_end, num_indices)) + tl.store(offsets_ptr + position + 1, current) + position += 1 + + +@triton.jit +def _load_checked_index( + indices_ptr, + position, + num_rows, + mask, + FUSED_BOUNDS_CHECK: tl.constexpr, +): + row_idx = tl.load(indices_ptr + position, mask=mask, other=0) + if FUSED_BOUNDS_CHECK: + invalid = mask & (row_idx != -1) & ((row_idx < 0) | (row_idx >= num_rows)) + tl.store(indices_ptr + position, 0, mask=invalid) + row_idx = tl.where(invalid, 0, row_idx) + else: + invalid = row_idx < row_idx + return row_idx, invalid + + @triton.jit def table_batched_embedding_bag_forward_weighted_kernel( output_ptr, @@ -196,6 +258,8 @@ def table_batched_embedding_bag_forward_unweighted_kernel( embedding_dims_ptr, embedding_offsets_ptr, feature_table_map_ptr, + rows_cumsum_ptr, + bounds_check_warning_ptr, # VBE-specific pointers (only used when vbe=True) # pyre-fixme[2]: Parameter must be annotated. row_output_offsets_ptr, @@ -206,10 +270,12 @@ def table_batched_embedding_bag_forward_unweighted_kernel( T: tl.constexpr, BLOCK_SIZE: tl.constexpr, vbe: tl.constexpr = False, + FUSED_BOUNDS_CHECK: tl.constexpr = False, ) -> None: b = tl.program_id(0).to(tl.int64) col_offsets = tl.arange(0, BLOCK_SIZE) + warning_count = 0 if vbe: output_row_base = output_ptr # unused, VBE uses row_output_offsets else: @@ -233,6 +299,10 @@ def table_batched_embedding_bag_forward_unweighted_kernel( table_offset = tl.load(table_offsets_ptr + table_idx) # embedding_dim and embedding_offset are indexed by feature embedding_dim = tl.load(embedding_dims_ptr + t) + if FUSED_BOUNDS_CHECK: + num_rows = tl.load(rows_cumsum_ptr + table_idx + 1) - tl.load( + rows_cumsum_ptr + table_idx + ) start = tl.load(offsets_ptr + b_t) end = tl.load(offsets_ptr + b_t + 1) @@ -245,10 +315,41 @@ def table_batched_embedding_bag_forward_unweighted_kernel( endn = start + step * ns for idx in range(start, endn, step): - row_idx_0 = tl.load(indices_ptr + idx + 0) - row_idx_1 = tl.load(indices_ptr + idx + 1) - row_idx_2 = tl.load(indices_ptr + idx + 2) - row_idx_3 = tl.load(indices_ptr + idx + 3) + row_idx_0, invalid_0 = _load_checked_index( + indices_ptr, + idx + 0, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_1, invalid_1 = _load_checked_index( + indices_ptr, + idx + 1, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_2, invalid_2 = _load_checked_index( + indices_ptr, + idx + 2, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_3, invalid_3 = _load_checked_index( + indices_ptr, + idx + 3, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + if FUSED_BOUNDS_CHECK: + warning_count += ( + invalid_0.to(tl.int32) + + invalid_1.to(tl.int32) + + invalid_2.to(tl.int32) + + invalid_3.to(tl.int32) + ) row_start_ptr_0 = weight_ptr + table_offset + row_idx_0 * embedding_dim row_start_ptr_1 = weight_ptr + table_offset + row_idx_1 * embedding_dim @@ -273,7 +374,15 @@ def table_batched_embedding_bag_forward_unweighted_kernel( ) for idx in range(endn, end): - row_idx = tl.load(indices_ptr + idx) + row_idx, invalid = _load_checked_index( + indices_ptr, + idx, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + if FUSED_BOUNDS_CHECK: + warning_count += invalid.to(tl.int32) row_start_ptr = weight_ptr + table_offset + row_idx * embedding_dim row_ptrs = row_start_ptr + col_offsets row = tl.load(row_ptrs, mask=mask, other=0) @@ -288,6 +397,9 @@ def table_batched_embedding_bag_forward_unweighted_kernel( bag_output_original = bag_output.to(tl.float32) tl.store(output_row_ptrs, bag_output_original, mask=mask) + 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( @@ -1041,6 +1153,8 @@ def forward( precomputed_total_B: int = 0, precomputed_max_B: int = 0, hoist_transpose_to_forward: bool = True, + bounds_check_warning: Optional[torch.Tensor] = None, + fused_bounds_check: bool = False, ) -> torch.Tensor: # VBE support: use pre-computed metadata if available, otherwise compute vbe = batch_size_per_feature_per_rank is not None @@ -1120,6 +1234,15 @@ def forward( (B, total_embedding_dim), device=weight.device, dtype=output_dtype ) + weighted = per_sample_weights is not None and per_sample_weights.numel() > 0 + if fused_bounds_check and ( + bounds_check_warning is None + or hoist_transpose_to_forward + or weighted + or vbe + ): + raise ValueError("Invalid fused bounds-check configuration") + # For VBE backward, save row_output_offsets, B_offsets, and b_t_map if vbe: assert vbe_metadata.B_offsets is not None @@ -1204,8 +1327,9 @@ def forward( num_warps = 1 - weighted = per_sample_weights is not None and per_sample_weights.numel() > 0 - + bounds_check_warning_ptr = ( + bounds_check_warning if bounds_check_warning is not None else indices + ) # Prepare VBE pointers (use dummy tensor if not VBE) row_output_offsets_ptr = ( row_output_offsets @@ -1244,29 +1368,47 @@ def forward( num_warps=num_warps, ) else: - fwd_kernel = ( - _amd_fwd_unweighted_kernel - if is_amd() - else table_batched_embedding_bag_forward_unweighted_kernel - ) - fwd_kernel[(B,)]( - output, - indices, - offsets, - weight, - table_offsets, - embedding_dims, - embedding_offsets, - feature_table_map, - row_output_offsets_ptr, - B_offsets_ptr, - total_embedding_dim, - B, - T, - BLOCK_SIZE=block_size, - vbe=vbe, - num_warps=num_warps, - ) + if is_amd(): + _amd_fwd_unweighted_kernel[(B,)]( + output, + indices, + offsets, + weight, + table_offsets, + embedding_dims, + embedding_offsets, + feature_table_map, + row_output_offsets_ptr, + B_offsets_ptr, + total_embedding_dim, + B, + T, + BLOCK_SIZE=block_size, + vbe=vbe, + num_warps=num_warps, + ) + else: + 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, + 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. @@ -1762,6 +1904,8 @@ def backward(ctx, dout) -> Tuple[None, ...]: None, # precomputed_total_B None, # precomputed_max_B None, # hoist_transpose_to_forward + None, + None, ) @@ -1783,6 +1927,7 @@ def __init__( optimizer: OptimType = OptimType.EXACT_SGD, device: Optional[torch.device] = None, hoist_transpose_to_forward: bool = False, + fused_bounds_check: bool = False, ) -> None: super().__init__() logging.info("TritonTableBatchedEmbeddingBags init args: %s", locals()) @@ -1865,6 +2010,7 @@ def __init__( self.eps = eps self.optimizer = optimizer self.hoist_transpose_to_forward = hoist_transpose_to_forward + self.fused_bounds_check = fused_bounds_check # Initialize optimizer state rows = [spec[0] for spec in embedding_specs] @@ -1892,6 +2038,9 @@ def __init__( ) self.bounds_check_warning = torch.tensor([0], device=device, dtype=torch.int64) self.bounds_check_mode: BoundsCheckMode = BoundsCheckMode.V2_WARNING + self._disable_offsets_adjustment = ( + FeatureGateName.DISABLE_OFFSETS_ADJUSTMENT.is_enabled() + ) def _bounds_check_config(self) -> Tuple[BoundsCheckMode, int]: is_v2 = self.bounds_check_mode.name.startswith("V2_") @@ -2030,9 +2179,25 @@ def forward( offsets, batch_size_per_feature_per_rank ) - # Bounds check (VBE-aware) bounds_check_mode, bounds_check_version = self._bounds_check_config() - if bounds_check_mode != BoundsCheckMode.NONE: + use_fused_bounds_check = ( + self.fused_bounds_check + and bounds_check_mode == BoundsCheckMode.WARNING + and batch_size_per_feature_per_rank is None + and (per_sample_weights is None or per_sample_weights.numel() == 0) + and not self.hoist_transpose_to_forward + and not is_amd() + ) + + if use_fused_bounds_check: + if indices.dim() != 1 or offsets.dim() != 1: + raise RuntimeError("indices and offsets must be one-dimensional") + if offsets.numel() == 0 or (offsets.numel() - 1) % self.T != 0: + raise RuntimeError("offsets size must equal B * T + 1") + if indices.device != offsets.device or indices.device != self.weight.device: + raise RuntimeError("TBE inputs must be on the same device") + + if bounds_check_mode != BoundsCheckMode.NONE and not use_fused_bounds_check: torch.ops.fbgemm.bounds_check_indices( self.rows_per_table, indices, @@ -2047,6 +2212,28 @@ def forward( info_B_mask=info_B_mask if info_B_mask > 0 else -1, bounds_check_version=bounds_check_version, ) + elif use_fused_bounds_check: + self.bounds_check_warning.zero_() + total_bags = offsets.size(0) - 1 + if total_bags > 0: + _bounds_check_offsets_kernel[(triton.cdiv(total_bags, 256),)]( + offsets, + self.bounds_check_warning, + indices.numel(), + total_bags, + BLOCK_SIZE=256, + num_warps=8, + ) + if self._disable_offsets_adjustment: + torch._assert_async(self.bounds_check_warning == 0) + else: + _repair_offsets_kernel[(1,)]( + offsets, + self.bounds_check_warning, + indices.numel(), + total_bags, + num_warps=1, + ) return TritonTBE.apply( indices, @@ -2082,6 +2269,8 @@ def forward( total_B, max_B, self.hoist_transpose_to_forward, + self.bounds_check_warning, + use_fused_bounds_check, ) def split_embedding_weights(self) -> List[torch.Tensor]: From 4b6e401d7396c79bbf522928f54f28539b86d414 Mon Sep 17 00:00:00 2001 From: Oleksandr Stashuk Date: Fri, 7 Aug 2026 23:13:46 -0700 Subject: [PATCH 3/8] Optimize Triton TBE forward for small long-bag tables (#4486) Summary: Convert long bags over small tables into per-bag row histograms followed by a tensor-core dot. Generic row, dimension, and bag-size predicates select the route, and checked loads preserve fused bounds behavior. Reviewed By: axeisghost Differential Revision: D114273027 --- .../distributed/batched_embedding_kernel.py | 2 + .../triton_table_batched_embeddings.py | 219 ++++++++++++++++-- 2 files changed, 205 insertions(+), 16 deletions(-) diff --git a/torchrec/distributed/batched_embedding_kernel.py b/torchrec/distributed/batched_embedding_kernel.py index 4ef4b1dc7..b61ee3bc6 100644 --- a/torchrec/distributed/batched_embedding_kernel.py +++ b/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, ) ) diff --git a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py index 7d3e8d17f..d82a831ea 100644 --- a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py +++ b/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,8 +1514,6 @@ 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, @@ -1406,9 +1521,43 @@ def forward( 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, ) From 2006579014b66f17924e4f9a2aa4c6bfd074a280 Mon Sep 17 00:00:00 2001 From: Oleksandr Stashuk Date: Fri, 7 Aug 2026 23:13:46 -0700 Subject: [PATCH 4/8] Reduce Triton TBE backward reduction pressure (#4487) Summary: Restructure short- and long-run gradient accumulation to reduce serialized reduction pressure. The new reduction schedule exposes more independent work while preserving the existing exact rowwise optimizer update. Reviewed By: axeisghost Differential Revision: D114273026 --- .../triton_table_batched_embeddings.py | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py index d82a831ea..61161523f 100644 --- a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py +++ b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py @@ -42,7 +42,9 @@ triton_tbe_backward_long_run_fused_weighted, ) from torchrec.distributed.triton_tbe.triton_tbe_backward_utils import ( + _CLC_FIXED_GRID, _expand_long_runs, + _FIXED_GRID, _LONG_RUN_THRESHOLD, _stochastic_rounding_store, get_grid_size, @@ -1774,7 +1776,10 @@ def backward(ctx, dout) -> Tuple[None, ...]: stochastic_rounding_seed=stochastic_rounding_seed, vbe=vbe, ) - if use_clc: + use_fused_clc_long_run = ( + use_clc and max_num_runs > _CLC_FIXED_GRID * _LONG_RUN_THRESHOLD + ) + if use_fused_clc_long_run: # CLC path: fused long-run grad accumulation + optimizer apply # CLC Path is exclusive to CUDA B200+. grad_accum_counter = programs_per_long_run.clone() @@ -1818,12 +1823,17 @@ def backward(ctx, dout) -> Tuple[None, ...]: else: # Non-CLC path: separate grad accumulation + apply kernels # Kernel 2: long-run grad accumulation (weighted) + long_accum_grid_size = ( + min(_FIXED_GRID, max_long_run_programs) + if use_clc + else long_accum_or_fused_grid_size + ) bwd_long_accum_w = ( _amd_bwd_long_accum_weighted if _use_amd else triton_tbe_backward_long_run_grad_accum_weighted ) - bwd_long_accum_w[(long_accum_or_fused_grid_size,)]( + bwd_long_accum_w[(long_accum_grid_size,)]( dout, infos_sorted, long_run_program_seg_starts, @@ -1912,7 +1922,10 @@ def backward(ctx, dout) -> Tuple[None, ...]: stochastic_rounding_seed=stochastic_rounding_seed, vbe=vbe, ) - if use_clc: + use_fused_clc_long_run = ( + use_clc and max_num_runs > _CLC_FIXED_GRID * _LONG_RUN_THRESHOLD + ) + if use_fused_clc_long_run: # CLC path: fused long-run grad accumulation + optimizer apply # CLC Path is exclusive to CUDA B200+. grad_accum_counter = programs_per_long_run.clone() @@ -1955,12 +1968,17 @@ def backward(ctx, dout) -> Tuple[None, ...]: else: # Non-CLC path: separate grad accumulation + apply kernels # Kernel 2: long-run grad accumulation + long_accum_grid_size = ( + min(_FIXED_GRID, max_long_run_programs) + if use_clc + else long_accum_or_fused_grid_size + ) bwd_long_accum_uw = ( _amd_bwd_long_accum_unweighted if _use_amd else triton_tbe_backward_long_run_grad_accum_unweighted ) - bwd_long_accum_uw[(long_accum_or_fused_grid_size,)]( + bwd_long_accum_uw[(long_accum_grid_size,)]( dout, infos_sorted, long_run_program_seg_starts, From 6483b7ae258a127e2228a180bccf0e09d37fd445 Mon Sep 17 00:00:00 2001 From: Oleksandr Stashuk Date: Fri, 7 Aug 2026 23:13:46 -0700 Subject: [PATCH 5/8] Reduce Triton TBE unweighted backward reduction width (#4488) Summary: Reduce the dout reduction tile from 16 to 8 in the unweighted short-run, separate long-run accumulation, and fused long-run kernels. The smaller tile reduces live vector state and register pressure on B200 while retaining FP32 accumulation and the existing rowwise optimizer update. Weighted kernels remain unchanged. Reviewed By: axeisghost Differential Revision: D114403703 --- .../distributed/triton_tbe/triton_table_batched_embeddings.py | 4 ++-- .../triton_tbe/triton_tbe_backward_long_run_fused.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py index 61161523f..f2de10c94 100644 --- a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py +++ b/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) diff --git a/torchrec/distributed/triton_tbe/triton_tbe_backward_long_run_fused.py b/torchrec/distributed/triton_tbe/triton_tbe_backward_long_run_fused.py index f01a7a512..69aeaf737 100644 --- a/torchrec/distributed/triton_tbe/triton_tbe_backward_long_run_fused.py +++ b/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 From f0e4a2b1c673486c4824d87450903cc65aa3de29 Mon Sep 17 00:00:00 2001 From: Oleksandr Stashuk Date: Fri, 7 Aug 2026 23:13:46 -0700 Subject: [PATCH 6/8] Coarsen Triton TBE forward bag scheduling (#4489) Summary: Process multiple bags per Triton program where the workload has enough parallelism. This amortizes program overhead and exposes independent loads without changing the embedding reduction. Reviewed By: axeisghost Differential Revision: D114280197 --- .../triton_table_batched_embeddings.py | 282 ++++++++++-------- 1 file changed, 156 insertions(+), 126 deletions(-) diff --git a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py index f2de10c94..ee1ef5d76 100644 --- a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py +++ b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py @@ -265,10 +265,7 @@ def table_batched_embedding_bag_forward_unweighted_kernel( feature_table_map_ptr, rows_cumsum_ptr, bounds_check_warning_ptr, - # VBE-specific pointers (only used when vbe=True) - # pyre-fixme[2]: Parameter must be annotated. row_output_offsets_ptr, - # pyre-fixme[2]: Parameter must be annotated. B_offsets_ptr, total_embedding_dim: tl.constexpr, B, @@ -277,136 +274,159 @@ def table_batched_embedding_bag_forward_unweighted_kernel( vbe: tl.constexpr = False, FEATURE_START: tl.constexpr = 0, FEATURE_END: tl.constexpr = -1, + BAGS_PER_PROGRAM: tl.constexpr = 1, FUSED_BOUNDS_CHECK: tl.constexpr = False, ) -> None: - - b = tl.program_id(0).to(tl.int64) + base_b = tl.program_id(0).to(tl.int64) * BAGS_PER_PROGRAM col_offsets = tl.arange(0, BLOCK_SIZE) warning_count = 0 - if vbe: - output_row_base = output_ptr # unused, VBE uses row_output_offsets - else: - output_row_base = output_ptr + b * total_embedding_dim feature_end: tl.constexpr = T if FEATURE_END < 0 else FEATURE_END for t in range(FEATURE_START, feature_end): + table_idx = tl.load(feature_table_map_ptr + t) + table_offset = tl.load(table_offsets_ptr + table_idx) + embedding_dim = tl.load(embedding_dims_ptr + t) + embedding_offset = tl.load(embedding_offsets_ptr + t) + if FUSED_BOUNDS_CHECK: + num_rows = tl.load(rows_cumsum_ptr + table_idx + 1) - tl.load( + rows_cumsum_ptr + table_idx + ) + 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) B_end = tl.load(B_offsets_ptr + t + 1).to(tl.int64) B_t = B_end - B_start - b_t = B_start + b - in_bounds = b < B_t - else: - b_t = t * B + b - in_bounds = True - if in_bounds: - # Map feature index to table index for weight lookup - table_idx = tl.load(feature_table_map_ptr + t) - table_offset = tl.load(table_offsets_ptr + table_idx) - # embedding_dim and embedding_offset are indexed by feature - embedding_dim = tl.load(embedding_dims_ptr + t) - if FUSED_BOUNDS_CHECK: - num_rows = tl.load(rows_cumsum_ptr + table_idx + 1) - tl.load( - rows_cumsum_ptr + table_idx - ) - - start = tl.load(offsets_ptr + b_t) - end = tl.load(offsets_ptr + b_t + 1) - - mask = col_offsets < embedding_dim - 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 - endn = start + step * ns - - for idx in range(start, endn, step): - row_idx_0, invalid_0 = _load_checked_index( - indices_ptr, - idx + 0, - num_rows if FUSED_BOUNDS_CHECK else 0, - True, - FUSED_BOUNDS_CHECK, - ) - row_idx_1, invalid_1 = _load_checked_index( - indices_ptr, - idx + 1, - num_rows if FUSED_BOUNDS_CHECK else 0, - True, - FUSED_BOUNDS_CHECK, - ) - row_idx_2, invalid_2 = _load_checked_index( - indices_ptr, - idx + 2, - num_rows if FUSED_BOUNDS_CHECK else 0, - True, - FUSED_BOUNDS_CHECK, - ) - row_idx_3, invalid_3 = _load_checked_index( - indices_ptr, - idx + 3, - num_rows if FUSED_BOUNDS_CHECK else 0, - True, - FUSED_BOUNDS_CHECK, + for bag_slot in tl.static_range(0, BAGS_PER_PROGRAM): + b = base_b + bag_slot + if vbe: + b_t = B_start + b + in_bounds = b < B_t + else: + b_t = t * B + b + in_bounds = b < B + + if in_bounds: + start = tl.load(offsets_ptr + b_t) + end = tl.load(offsets_ptr + b_t + 1) + mask = col_offsets < embedding_dim + accumulator_dtype: tl.constexpr = ( + tl.float64 + if weight_ptr.dtype.element_ty == tl.float32 + else tl.float32 ) - if FUSED_BOUNDS_CHECK: - warning_count += ( - invalid_0.to(tl.int32) - + invalid_1.to(tl.int32) - + invalid_2.to(tl.int32) - + invalid_3.to(tl.int32) + bag_output = tl.zeros((BLOCK_SIZE,), dtype=accumulator_dtype) + + step: tl.constexpr = 4 + ns = (end - start) // step + endn = start + step * ns + + for idx in range(start, endn, step): + row_idx_0, invalid_0 = _load_checked_index( + indices_ptr, + idx + 0, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_1, invalid_1 = _load_checked_index( + indices_ptr, + idx + 1, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_2, invalid_2 = _load_checked_index( + indices_ptr, + idx + 2, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_3, invalid_3 = _load_checked_index( + indices_ptr, + idx + 3, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + if FUSED_BOUNDS_CHECK: + warning_count += ( + invalid_0.to(tl.int32) + + invalid_1.to(tl.int32) + + invalid_2.to(tl.int32) + + invalid_3.to(tl.int32) + ) + row_0 = tl.load( + weight_ptr + + table_offset + + row_idx_0 * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) + row_1 = tl.load( + weight_ptr + + table_offset + + row_idx_1 * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) + row_2 = tl.load( + weight_ptr + + table_offset + + row_idx_2 * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) + row_3 = tl.load( + weight_ptr + + table_offset + + row_idx_3 * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) + bag_output += ( + row_0.to(tl.float32) + + row_1.to(tl.float32) + + row_2.to(tl.float32) + + row_3.to(tl.float32) ) - row_start_ptr_0 = weight_ptr + table_offset + row_idx_0 * embedding_dim - row_start_ptr_1 = weight_ptr + table_offset + row_idx_1 * embedding_dim - row_start_ptr_2 = weight_ptr + table_offset + row_idx_2 * embedding_dim - row_start_ptr_3 = weight_ptr + table_offset + row_idx_3 * embedding_dim - - row_ptrs_0 = row_start_ptr_0 + col_offsets - row_ptrs_1 = row_start_ptr_1 + col_offsets - row_ptrs_2 = row_start_ptr_2 + col_offsets - row_ptrs_3 = row_start_ptr_3 + col_offsets - - row_0 = tl.load(row_ptrs_0, mask=mask, other=0) - row_1 = tl.load(row_ptrs_1, mask=mask, other=0) - row_2 = tl.load(row_ptrs_2, mask=mask, other=0) - row_3 = tl.load(row_ptrs_3, mask=mask, other=0) - - bag_output += ( - row_0.to(tl.float32) - + row_1.to(tl.float32) - + row_2.to(tl.float32) - + row_3.to(tl.float32) - ) - - for idx in range(endn, end): - row_idx, invalid = _load_checked_index( - indices_ptr, - idx, - num_rows if FUSED_BOUNDS_CHECK else 0, - True, - FUSED_BOUNDS_CHECK, - ) - if FUSED_BOUNDS_CHECK: - warning_count += invalid.to(tl.int32) - row_start_ptr = weight_ptr + table_offset + row_idx * embedding_dim - row_ptrs = row_start_ptr + col_offsets - row = tl.load(row_ptrs, mask=mask, other=0) - bag_output += row.to(tl.float32) + for idx in range(endn, end): + row_idx, invalid = _load_checked_index( + indices_ptr, + idx, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + if FUSED_BOUNDS_CHECK: + warning_count += invalid.to(tl.int32) + row = tl.load( + weight_ptr + + table_offset + + row_idx * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) + bag_output += row.to(tl.float32) - if vbe: - row_output_offset = tl.load(row_output_offsets_ptr + b_t) - output_row_ptrs = output_ptr + row_output_offset + col_offsets - else: - embedding_offset = tl.load(embedding_offsets_ptr + t) - output_row_ptrs = output_row_base + embedding_offset + col_offsets - bag_output_original = bag_output.to(tl.float32) - tl.store(output_row_ptrs, bag_output_original, mask=mask) + if vbe: + row_output_offset = tl.load(row_output_offsets_ptr + b_t) + output_row_ptrs = output_ptr + row_output_offset + col_offsets + else: + output_row_ptrs = ( + output_ptr + + b * total_embedding_dim + + embedding_offset + + col_offsets + ) + tl.store(output_row_ptrs, bag_output.to(tl.float32), mask=mask) if FUSED_BOUNDS_CHECK and warning_count > 0: tl.atomic_add(bounds_check_warning_ptr, warning_count.to(tl.int64)) @@ -1483,6 +1503,21 @@ def forward( and not is_amd() and weight.dtype == torch.float16 ) + bags_per_program = ( + 4 + if use_small_table_kernel + else ( + 2 if not vbe and B >= 65536 and weight.dtype != torch.float32 else 1 + ) + ) + feature_ranges = ( + [ + (0, histogram_feature), + (histogram_feature + 1, T), + ] + if use_small_table_kernel + else [(0, T)] + ) if use_small_table_kernel: table_batched_embedding_bag_forward_small_table_kernel[ (triton.cdiv(B, 16),) @@ -1526,18 +1561,12 @@ def forward( 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,)]( + table_batched_embedding_bag_forward_unweighted_kernel[ + (triton.cdiv(B, bags_per_program),) + ]( output, indices, offsets, @@ -1557,6 +1586,7 @@ def forward( vbe=vbe, FEATURE_START=feature_start, FEATURE_END=feature_end, + BAGS_PER_PROGRAM=bags_per_program, FUSED_BOUNDS_CHECK=fused_bounds_check, num_warps=num_warps, ) From 703262c352df2fadde4e138a1f9aa6becbd53833 Mon Sep 17 00:00:00 2001 From: Oleksandr Stashuk Date: Fri, 7 Aug 2026 23:13:46 -0700 Subject: [PATCH 7/8] Widen Triton TBE forward gather issue (#4490) Summary: Issue up to eight independent index gathers per loop iteration on profitable long-bag shapes. The wider loop increases memory-level parallelism while retaining the four-wide path for shapes where register pressure dominates. Reviewed By: axeisghost Differential Revision: D114280947 --- .../triton_table_batched_embeddings.py | 80 ++++++++++++++++++- 1 file changed, 79 insertions(+), 1 deletion(-) diff --git a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py index ee1ef5d76..16677ee26 100644 --- a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py +++ b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py @@ -275,6 +275,7 @@ def table_batched_embedding_bag_forward_unweighted_kernel( FEATURE_START: tl.constexpr = 0, FEATURE_END: tl.constexpr = -1, BAGS_PER_PROGRAM: tl.constexpr = 1, + UNROLL8: tl.constexpr = False, FUSED_BOUNDS_CHECK: tl.constexpr = False, ) -> None: base_b = tl.program_id(0).to(tl.int64) * BAGS_PER_PROGRAM @@ -317,7 +318,7 @@ def table_batched_embedding_bag_forward_unweighted_kernel( ) bag_output = tl.zeros((BLOCK_SIZE,), dtype=accumulator_dtype) - step: tl.constexpr = 4 + step: tl.constexpr = 8 if UNROLL8 else 4 ns = (end - start) // step endn = start + step * ns @@ -357,6 +358,42 @@ def table_batched_embedding_bag_forward_unweighted_kernel( + invalid_2.to(tl.int32) + invalid_3.to(tl.int32) ) + if UNROLL8: + row_idx_4, invalid_4 = _load_checked_index( + indices_ptr, + idx + 4, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_5, invalid_5 = _load_checked_index( + indices_ptr, + idx + 5, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_6, invalid_6 = _load_checked_index( + indices_ptr, + idx + 6, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + row_idx_7, invalid_7 = _load_checked_index( + indices_ptr, + idx + 7, + num_rows if FUSED_BOUNDS_CHECK else 0, + True, + FUSED_BOUNDS_CHECK, + ) + if FUSED_BOUNDS_CHECK: + warning_count += ( + invalid_4.to(tl.int32) + + invalid_5.to(tl.int32) + + invalid_6.to(tl.int32) + + invalid_7.to(tl.int32) + ) row_0 = tl.load( weight_ptr + table_offset @@ -389,12 +426,52 @@ def table_batched_embedding_bag_forward_unweighted_kernel( mask=mask, other=0, ) + if UNROLL8: + row_4 = tl.load( + weight_ptr + + table_offset + + row_idx_4 * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) + row_5 = tl.load( + weight_ptr + + table_offset + + row_idx_5 * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) + row_6 = tl.load( + weight_ptr + + table_offset + + row_idx_6 * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) + row_7 = tl.load( + weight_ptr + + table_offset + + row_idx_7 * embedding_dim + + col_offsets, + mask=mask, + other=0, + ) bag_output += ( row_0.to(tl.float32) + row_1.to(tl.float32) + row_2.to(tl.float32) + row_3.to(tl.float32) ) + if UNROLL8: + bag_output += ( + row_4.to(tl.float32) + + row_5.to(tl.float32) + + row_6.to(tl.float32) + + row_7.to(tl.float32) + ) for idx in range(endn, end): row_idx, invalid = _load_checked_index( @@ -1587,6 +1664,7 @@ def forward( FEATURE_START=feature_start, FEATURE_END=feature_end, BAGS_PER_PROGRAM=bags_per_program, + UNROLL8=bags_per_program == 2, FUSED_BOUNDS_CHECK=fused_bounds_check, num_warps=num_warps, ) From aa5c4b8da2cd6041224885983a23d2aa1bcd2db1 Mon Sep 17 00:00:00 2001 From: Oleksandr Stashuk Date: Fri, 7 Aug 2026 23:13:46 -0700 Subject: [PATCH 8/8] Reuse forward histograms in Triton TBE backward (#4491) Summary: Reuse exact counts from an eligible leading, unique small-table feature during exact rowwise Adagrad backward. Immutable plans carry actual table geometry, compensated FP16 high/residual matrix products preserve FP32 accuracy, and scalar validity enables safe fallback for long bags, aliases, graph capture, weighted inputs, and VBE. Reviewed By: axeisghost Differential Revision: D114284800 --- .../triton_table_batched_embeddings.py | 329 ++++++++++++++++-- 1 file changed, 299 insertions(+), 30 deletions(-) diff --git a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py index 16677ee26..ece8fb8ba 100644 --- a/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py +++ b/torchrec/distributed/triton_tbe/triton_table_batched_embeddings.py @@ -13,6 +13,7 @@ import logging import math import os +from dataclasses import dataclass from itertools import accumulate from typing import Any, Dict, List, Optional, Tuple @@ -80,6 +81,17 @@ def lengths_to_offsets(lengths: List[int], keep_last: bool = False) -> List[int] return offsets +@dataclass(frozen=True) +class _SavedHistogramPlan: + feature: int + table: int + num_rows: int + row_bins: int + embedding_dim: int + embedding_offset: int + block_size: int + + @triton.jit def _bounds_check_offsets_kernel( offsets_ptr, @@ -512,6 +524,8 @@ def table_batched_embedding_bag_forward_unweighted_kernel( @triton.jit def table_batched_embedding_bag_forward_small_table_kernel( output_ptr, + histogram_counts_ptr, + histogram_counts_invalid_ptr, indices_ptr, offsets_ptr, weight_ptr, @@ -526,6 +540,8 @@ def table_batched_embedding_bag_forward_small_table_kernel( NUM_ROWS: tl.constexpr, ROW_BINS: tl.constexpr, BLOCK_SIZE: tl.constexpr, + STORE_COUNTS: tl.constexpr, + INVALID_INDEX: tl.constexpr, FUSED_BOUNDS_CHECK: tl.constexpr = False, ) -> None: bags_per_program: tl.constexpr = 16 @@ -550,17 +566,28 @@ def table_batched_embedding_bag_forward_small_table_kernel( 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 + histogram_input_mask = input_mask & (row_indices != -1) 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,)), + mask=histogram_input_mask.reshape((bags_per_program * histogram_chunk_size,)), ).reshape((bags_per_program, ROW_BINS)) + rows = tl.arange(0, ROW_BINS) + tail_lengths = tl.maximum(lengths - histogram_chunk_size, 0) + if STORE_COUNTS: + tl.store( + histogram_counts_ptr + bags[:, None] * ROW_BINS + rows[None, :], + counts.to(tl.float16), + mask=bag_mask[:, None], + ) + if tl.max(lengths) > histogram_chunk_size: + tl.atomic_or(histogram_counts_invalid_ptr + INVALID_INDEX, 1) + 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, :], @@ -569,7 +596,6 @@ def table_batched_embedding_bag_forward_small_table_kernel( ) 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( @@ -581,6 +607,7 @@ def table_batched_embedding_bag_forward_small_table_kernel( ) if FUSED_BOUNDS_CHECK: warning_count += tl.sum(invalid.to(tl.int32)) + active = active & (row_idx != -1) row = tl.load( weight_ptr + table_offset @@ -603,6 +630,61 @@ def table_batched_embedding_bag_forward_small_table_kernel( tl.atomic_add(bounds_check_warning_ptr, warning_count.to(tl.int64)) +@triton.jit +def triton_tbe_backward_histogram_apply_rowwise_adagrad( + grad_ptr, + weight_ptr, + momentum_ptr, + table_offsets_ptr, + rows_cumsum_ptr, + learning_rate, + eps, + NUM_ROWS: tl.constexpr, + EMBEDDING_DIM: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + TABLE_IDX: tl.constexpr, + STOCHASTIC_ROUNDING: tl.constexpr, + stochastic_rounding_seed, +) -> None: + row_idx = tl.program_id(0) + cols = tl.arange(0, BLOCK_SIZE) + mask = (row_idx < NUM_ROWS) & (cols < EMBEDDING_DIM) + + table_offset = tl.load(table_offsets_ptr + TABLE_IDX) + row_ptrs = weight_ptr + table_offset + row_idx * EMBEDDING_DIM + cols + grad = tl.load( + grad_ptr + row_idx * EMBEDDING_DIM + cols, + mask=mask, + other=0.0, + ) + weight = tl.load(row_ptrs, mask=mask, other=0.0).to(tl.float32) + + row_offset = tl.load(rows_cumsum_ptr + TABLE_IDX) + momentum_idx = row_offset + row_idx + grad_square_average = tl.sum(grad * grad) / EMBEDDING_DIM + momentum = tl.load(momentum_ptr + momentum_idx, mask=row_idx < NUM_ROWS) + momentum_new = momentum + grad_square_average + tl.store( + momentum_ptr + momentum_idx, + momentum_new, + mask=row_idx < NUM_ROWS, + ) + adaptive_learning_rate = learning_rate / (tl.sqrt(momentum_new) + eps) + row_update = weight - adaptive_learning_rate * grad + + if STOCHASTIC_ROUNDING: + sr_offset = table_offset + row_idx * EMBEDDING_DIM + cols + _stochastic_rounding_store( + row_ptrs, + row_update, + mask, + stochastic_rounding_seed, + sr_offset, + ) + else: + tl.store(row_ptrs, row_update, mask=mask) + + @triton.jit def triton_tbe_backward_short_run_unweighted( dout_ptr, @@ -1359,6 +1441,7 @@ def forward( histogram_num_rows: int = 0, histogram_row_bins: int = 0, histogram_block_size: int = 0, + saved_histogram_plans: Tuple[_SavedHistogramPlan, ...] = (), bounds_check_warning: Optional[torch.Tensor] = None, fused_bounds_check: bool = False, ) -> torch.Tensor: @@ -1448,6 +1531,49 @@ def forward( or vbe ): raise ValueError("Invalid fused bounds-check configuration") + use_small_table_kernel = ( + histogram_feature >= 0 + and not weighted + and not vbe + and not is_amd() + and weight.dtype == torch.float16 + ) + valid_saved_histogram_prefix = bool(saved_histogram_plans) + expected_embedding_offset = 0 + for expected_feature, plan in enumerate(saved_histogram_plans): + valid_saved_histogram_prefix = ( + valid_saved_histogram_prefix + and plan.feature == expected_feature + and plan.embedding_offset == expected_embedding_offset + and plan.num_rows > 0 + and plan.row_bins >= plan.num_rows + and (plan.row_bins & (plan.row_bins - 1)) == 0 + and plan.embedding_dim > 0 + and plan.block_size >= plan.embedding_dim + and (plan.block_size & (plan.block_size - 1)) == 0 + ) + expected_embedding_offset += plan.embedding_dim + use_histogram_backward = ( + use_small_table_kernel + and B > 0 + and histogram_feature == 0 + and valid_saved_histogram_prefix + and optimizer == OptimType.EXACT_ROWWISE_ADAGRAD + and not hoist_transpose_to_forward + and not torch.cuda.is_current_stream_capturing() + ) + active_saved_histogram_plans = ( + saved_histogram_plans if use_histogram_backward else () + ) + histogram_counts = tuple( + torch.empty((B, plan.row_bins), device=weight.device, dtype=torch.float16) + for plan in active_saved_histogram_plans + ) + histogram_counts_invalid = torch.zeros( + len(active_saved_histogram_plans), + device=weight.device, + dtype=torch.int32, + ) # For VBE backward, save row_output_offsets, B_offsets, and b_t_map if vbe: @@ -1475,6 +1601,8 @@ def forward( vbe_row_output_offsets, vbe_B_offsets, vbe_b_t_map, + *histogram_counts, + histogram_counts_invalid, ) ctx.total_embedding_dim = total_embedding_dim @@ -1491,6 +1619,7 @@ def forward( ctx.stochastic_rounding = stochastic_rounding ctx.vbe = vbe ctx.hoist_transpose_to_forward = hoist_transpose_to_forward + ctx.saved_histogram_plans = active_saved_histogram_plans if hoist_transpose_to_forward: # Hoist the backward index transpose (linearize + sort + run-length @@ -1574,32 +1703,27 @@ def forward( num_warps=num_warps, ) else: - use_small_table_kernel = ( - histogram_feature >= 0 - and not vbe - and not is_amd() - and weight.dtype == torch.float16 + empty_histogram_counts = torch.empty( + 0, device=weight.device, dtype=torch.float16 ) - bags_per_program = ( - 4 - if use_small_table_kernel - else ( - 2 if not vbe and B >= 65536 and weight.dtype != torch.float32 else 1 + histogram_storage = { + plan.feature: (counts, plan_index) + for plan_index, (plan, counts) in enumerate( + zip(active_saved_histogram_plans, histogram_counts) ) - ) - feature_ranges = ( - [ - (0, histogram_feature), - (histogram_feature + 1, T), - ] - if use_small_table_kernel - else [(0, T)] - ) + } + optimized_histogram_features = [] if use_small_table_kernel: + stored_counts, invalid_index = histogram_storage.get( + histogram_feature, + (empty_histogram_counts, 0), + ) table_batched_embedding_bag_forward_small_table_kernel[ (triton.cdiv(B, 16),) ]( output, + stored_counts, + histogram_counts_invalid, indices, offsets, weight, @@ -1614,9 +1738,43 @@ def forward( NUM_ROWS=histogram_num_rows, ROW_BINS=histogram_row_bins, BLOCK_SIZE=histogram_block_size, + STORE_COUNTS=histogram_feature in histogram_storage, + INVALID_INDEX=invalid_index, FUSED_BOUNDS_CHECK=fused_bounds_check, num_warps=1, ) + optimized_histogram_features.append(histogram_feature) + + for plan_index, plan in enumerate(active_saved_histogram_plans): + if plan.feature == histogram_feature: + continue + stored_counts, _ = histogram_storage[plan.feature] + table_batched_embedding_bag_forward_small_table_kernel[ + (triton.cdiv(B, 16),) + ]( + output, + stored_counts, + histogram_counts_invalid, + indices, + offsets, + weight, + table_offsets, + embedding_dims, + embedding_offsets, + feature_table_map, + bounds_check_warning_ptr, + total_embedding_dim, + B, + FEATURE=plan.feature, + NUM_ROWS=plan.num_rows, + ROW_BINS=plan.row_bins, + BLOCK_SIZE=plan.block_size, + STORE_COUNTS=True, + INVALID_INDEX=plan_index, + FUSED_BOUNDS_CHECK=fused_bounds_check, + num_warps=1, + ) + optimized_histogram_features.append(plan.feature) if is_amd(): _amd_fwd_unweighted_kernel[(B,)]( @@ -1638,6 +1796,23 @@ def forward( num_warps=num_warps, ) else: + bags_per_program = ( + 4 + if optimized_histogram_features + else ( + 2 + if not vbe and B >= 65536 and weight.dtype != torch.float32 + else 1 + ) + ) + feature_ranges = [] + feature_start = 0 + for feature in sorted(set(optimized_histogram_features)): + if feature_start < feature: + feature_ranges.append((feature_start, feature)) + feature_start = feature + 1 + if feature_start < T: + feature_ranges.append((feature_start, T)) for feature_start, feature_end in feature_ranges: if feature_start >= feature_end: continue @@ -1682,6 +1857,7 @@ def backward(ctx, dout) -> Tuple[None, ...]: # Ensure dout is contiguous for correct memory access in Triton kernels dout = dout.contiguous() + saved_tensors = ctx.saved_tensors ( indices, offsets, @@ -1694,7 +1870,11 @@ def backward(ctx, dout) -> Tuple[None, ...]: vbe_row_output_offsets, vbe_B_offsets, vbe_b_t_map, - ) = ctx.saved_tensors + ) = saved_tensors[:11] + saved_histogram_plans = ctx.saved_histogram_plans + num_saved_histograms = len(saved_histogram_plans) + histogram_counts = saved_tensors[11 : 11 + num_saved_histograms] + histogram_counts_invalid = saved_tensors[11 + num_saved_histograms] total_hash_size_bits = ctx.total_hash_size_bits total_embedding_dim = ctx.total_embedding_dim @@ -1717,6 +1897,75 @@ def backward(ctx, dout) -> Tuple[None, ...]: weighted = per_sample_weights.numel() > 0 + feature_table_map = ctx.feature_table_map + num_valid_histograms = 0 + if num_saved_histograms and not torch.cuda.is_current_stream_capturing(): + for invalid in histogram_counts_invalid.tolist(): + if invalid: + break + num_valid_histograms += 1 + + for plan, counts in zip( + saved_histogram_plans[:num_valid_histograms], + histogram_counts[:num_valid_histograms], + ): + histogram_dout = dout[ + :, plan.embedding_offset : plan.embedding_offset + plan.embedding_dim + ] + histogram_dout_scale = torch.exp2( + torch.clamp( + torch.ceil(torch.log2(torch.amax(torch.abs(histogram_dout)))) - 15, + min=0, + ) + ) + histogram_dout_scaled = histogram_dout / histogram_dout_scale + histogram_dout_hi = histogram_dout_scaled.to(torch.float16) + histogram_dout_lo = ( + histogram_dout_scaled - histogram_dout_hi.to(torch.float32) + ).to(torch.float16) + histogram_counts_t = counts.T[: plan.num_rows] + histogram_grad = ( + torch.mm( + histogram_counts_t, + histogram_dout_hi, + out_dtype=torch.float32, + ) + + torch.mm( + histogram_counts_t, + histogram_dout_lo, + out_dtype=torch.float32, + ) + ) * histogram_dout_scale + triton_tbe_backward_histogram_apply_rowwise_adagrad[(plan.num_rows,)]( + histogram_grad, + weight, + momentum, + table_offsets, + rows_cumsum, + learning_rate, + eps, + NUM_ROWS=plan.num_rows, + EMBEDDING_DIM=plan.embedding_dim, + BLOCK_SIZE=plan.block_size, + TABLE_IDX=plan.table, + STOCHASTIC_ROUNDING=stochastic_rounding, + stochastic_rounding_seed=stochastic_rounding_seed, + num_warps=1, + ) + + if num_valid_histograms: + histogram_prefix = int(offsets[num_valid_histograms * B].item()) + indices = indices[histogram_prefix:] + offsets = offsets[num_valid_histograms * B :] - histogram_prefix + embedding_dims = embedding_dims[num_valid_histograms:] + embedding_offsets = embedding_offsets[num_valid_histograms:] + hash_size_cumsum = hash_size_cumsum[num_valid_histograms:] + feature_table_map = feature_table_map[num_valid_histograms:] + T -= num_valid_histograms + + if T == 0: + return (None,) * len(ctx.needs_input_grad) + if ctx.hoist_transpose_to_forward: # The index transpose (linearize + sort + run-length encode) was # hoisted to forward() and cached on ctx; read it back instead of @@ -1862,7 +2111,7 @@ def backward(ctx, dout) -> Tuple[None, ...]: table_offsets, embedding_dims, embedding_offsets, - ctx.feature_table_map, + feature_table_map, hash_size_cumsum, momentum, rows_cumsum, @@ -1909,7 +2158,7 @@ def backward(ctx, dout) -> Tuple[None, ...]: sorted_linear_indices_cumulative_run_lengths, long_run_original_ids, table_offsets, - ctx.feature_table_map, + feature_table_map, hash_size_cumsum, momentum, rows_cumsum, @@ -1977,7 +2226,7 @@ def backward(ctx, dout) -> Tuple[None, ...]: num_long_runs_t, table_offsets, embedding_dims, - ctx.feature_table_map, + feature_table_map, hash_size_cumsum, momentum, rows_cumsum, @@ -2009,7 +2258,7 @@ def backward(ctx, dout) -> Tuple[None, ...]: table_offsets, embedding_dims, embedding_offsets, - ctx.feature_table_map, + feature_table_map, hash_size_cumsum, momentum, rows_cumsum, @@ -2054,7 +2303,7 @@ def backward(ctx, dout) -> Tuple[None, ...]: sorted_linear_indices_cumulative_run_lengths, long_run_original_ids, table_offsets, - ctx.feature_table_map, + feature_table_map, hash_size_cumsum, momentum, rows_cumsum, @@ -2122,7 +2371,7 @@ def backward(ctx, dout) -> Tuple[None, ...]: num_long_runs_t, table_offsets, embedding_dims, - ctx.feature_table_map, + feature_table_map, hash_size_cumsum, momentum, rows_cumsum, @@ -2185,6 +2434,7 @@ def backward(ctx, dout) -> Tuple[None, ...]: None, None, None, + None, ) @@ -2236,12 +2486,13 @@ def __init__( # Feature-level properties (indexed by feature_table_map) # These are used for forward pass and output shape calculation feature_dims = [table_embedding_dims[t] for t in feature_table_map] + feature_embedding_offsets = lengths_to_offsets(feature_dims) self.total_embedding_dim = sum(feature_dims) self.table_offsets = torch.tensor( lengths_to_offsets(table_sizes), dtype=torch.int64, device=device ) self.embedding_offsets = torch.tensor( - lengths_to_offsets(feature_dims), dtype=torch.int64, device=device + feature_embedding_offsets, dtype=torch.int64, device=device ) self.embedding_dims = torch.tensor( feature_dims, dtype=torch.int64, device=device @@ -2295,6 +2546,7 @@ def __init__( self._histogram_num_rows = 0 self._histogram_row_bins = 0 self._histogram_block_size = 0 + self._saved_histogram_plans: Tuple[_SavedHistogramPlan, ...] = () if ( bag_size_hints is not None and weights_precision == torch.float16 @@ -2313,6 +2565,22 @@ def __init__( 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) + if ( + feature == 0 + and bag_size_hints[feature] <= 256 + and feature_table_map.count(feature_table_map[feature]) == 1 + ): + self._saved_histogram_plans = ( + _SavedHistogramPlan( + feature=feature, + table=feature_table_map[feature], + num_rows=num_rows, + row_bins=self._histogram_row_bins, + embedding_dim=dim, + embedding_offset=feature_embedding_offsets[feature], + block_size=self._histogram_block_size, + ), + ) self.stochastic_rounding = stochastic_rounding self.learning_rate = learning_rate @@ -2582,6 +2850,7 @@ def forward( self._histogram_num_rows, self._histogram_row_bins, self._histogram_block_size, + self._saved_histogram_plans, self.bounds_check_warning, use_fused_bounds_check, )