fbgemm-gpu-nightly-cpu 2025.3.27__cp311-cp311-manylinux_2_28_aarch64.whl → 2026.1.29__cp311-cp311-manylinux_2_28_aarch64.whl

fbgemm_gpu/sparse_ops.py

@@ -7,10 +7,12 @@
 # pyre-strict
 
 import math
-from typing import Callable, List, Optional, Sequence, Tuple
+from collections.abc import Sequence
+from typing import Callable, Optional
 
 import torch
 
+# fmt:skip
 from fbgemm_gpu.split_embedding_configs import SparseType
 from fbgemm_gpu.split_table_batched_embeddings_ops_common import PoolingMode
 from fbgemm_gpu.utils.loader import load_torch_module
@@ -41,12 +43,14 @@ except Exception:
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_split_cpu"
     )
+    torch.ops.load_library(
+        "//deeplearning/fbgemm/fbgemm_gpu:permute_multi_embedding_ops_cpu"
+    )
 
 
 import torch.utils._pytree as pytree
 from torch import SymInt, Tensor
-from torch.fx.experimental.symbolic_shapes import guard_size_oblivious
-
+from torch.fx.experimental.symbolic_shapes import guard_or_true
 
 if hasattr(torch.library, "register_fake"):
     # pyre-ignore[9]
@@ -71,7 +75,7 @@ def permute_2D_sparse_data_input1D_meta(
     stride: int,
     weights: Optional[Tensor] = None,
     permuted_lengths_sum: Optional[int] = None,
-) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
+) -> tuple[Tensor, Tensor, Optional[Tensor]]:
     torch._check(
         lengths.dim() == 1, lambda: f"expected lengths.dim() == 1, got {lengths.dim()}"
     )
@@ -108,7 +112,7 @@ def permute_2D_sparse_data_input1D_backward(
     grad_lengths: torch.Tensor,
     grad_values: torch.Tensor,
     grad_weights: torch.Tensor,
-) -> Tuple[None, Tensor, Tensor, None, Tensor, None]:
+) -> tuple[None, Tensor, Tensor, None, Tensor, None]:
     inv_permute = torch.ops.fbgemm.invert_permute(ctx.permute)
     permuted_grad_lengths, permuted_grad_values, permuted_grad_weights = (
         torch.ops.fbgemm.permute_2D_sparse_data_input1D(
@@ -136,7 +140,7 @@ def permute_2D_sparse_data_meta(
     values: Tensor,
     weights: Optional[Tensor] = None,
     permuted_lengths_sum: Optional[int] = None,
-) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
+) -> tuple[Tensor, Tensor, Optional[Tensor]]:
     torch._check(
         lengths.dim() == 2, lambda: f"expected lengths.dim() == 2, got {lengths.dim()}"
     )
@@ -163,6 +167,89 @@ def invert_permute_abstract(permute: Tensor) -> Tensor:
     return torch.empty_like(permute)
 
 
+def get_source_mask_meta(
+    num_sources: Tensor, num_targets: Tensor, output_size: Optional[int] = None
+) -> Tensor:
+    if output_size is None:
+        ctx = torch.library.get_ctx()
+        output_size = ctx.new_dynamic_size()
+    return torch.empty([output_size], dtype=torch.bool)
+
+
+def get_source_mask(
+    num_sources: Tensor, num_targets: Tensor, output_size: Optional[int] = None
+) -> Tensor:
+    """
+    Generate a boolean mask indicating which elements are from sources vs targets.
+
+    This is a Python wrapper that computes output_size when not provided,
+    enabling the operation to work with meta tensors for compilation.
+
+    Args:
+        num_sources: 1D tensor of source counts per batch element
+        num_targets: 1D tensor of target counts per batch element
+        output_size: Optional pre-computed output size.
+
+    Returns:
+        A 1D boolean tensor where True indicates source elements and False
+        indicates target elements
+
+    Example:
+        >>> num_sources = torch.tensor([2, 3])
+        >>> num_targets = torch.tensor([1, 2])
+        >>> get_source_mask(num_sources, num_targets)
+        tensor([True, True, False, True, True, True, False, False])
+    """
+    # Compute output_size if not provided and tensors are regular (not meta/fake)
+    if output_size is None:
+        combined = num_sources + num_targets
+        output_size = int(combined.sum().item())
+
+    return torch.ops.fbgemm.get_source_mask(num_sources, num_targets, output_size)
+
+
+def repeat_arange_meta(lengths: Tensor) -> Tensor:
+    """Meta implementation for repeat_arange."""
+    # Output size is data-dependent (sum of lengths).
+    # For FakeTensors (used in torch.compile), we use dynamic sizing.
+    # For actual meta tensors, we cannot determine the size so return empty.
+    if lengths.device.type == "meta":
+        # Actual meta tensors: return a zero-sized tensor as placeholder
+        # since we cannot compute the data-dependent output size
+        return torch.empty([0], dtype=lengths.dtype, device=lengths.device)
+    else:
+        # FakeTensor context: use dynamic sizing for proper shape tracking
+        ctx = torch.library.get_ctx()
+        output_size = ctx.new_dynamic_size()
+        return torch.empty([output_size], dtype=lengths.dtype, device=lengths.device)
+
+
+def repeat_arange(lengths: Tensor) -> Tensor:
+    """
+    Creates a concatenated tensor of aranges based on a lengths tensor.
+
+    This is a high-performance CUDA kernel that replaces the inefficient PyTorch
+    implementation which uses 4+ separate kernels (cumsum, arange, repeat_interleave, sub).
+
+    Args:
+        lengths: 1D tensor of lengths for each arange sequence
+
+    Returns:
+        A 1D tensor containing concatenated arange sequences
+
+    Example:
+        >>> lengths = torch.tensor([3, 5, 2])
+        >>> repeat_arange(lengths)
+        tensor([0, 1, 2, 0, 1, 2, 3, 4, 0, 1])
+
+    Performance:
+        - PyTorch implementation: 4+ kernel launches + intermediate allocations
+        - CUDA implementation: 1 fused kernel, no intermediate allocations
+        - Typical speedup: 3-5x on realistic workloads
+    """
+    return torch.ops.fbgemm.repeat_arange(lengths)
+
+
 # pyre-ignore
 def permute_2D_sparse_data_setup_context(ctx, inputs, output):
     permute, lengths, values, weights, permuted_lengths_sum = inputs
@@ -194,7 +281,7 @@ def permute_1D_sparse_data_meta(
     values: Tensor,
     weights: Optional[Tensor] = None,
     permuted_lengths_sum: Optional[int] = None,
-) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
+) -> tuple[Tensor, Tensor, Optional[Tensor]]:
     indices = values
     permuted_lengths_size = permute.numel()
     permuted_lengths = lengths.new_empty([permuted_lengths_size])
@@ -215,7 +302,7 @@ def permute_1D_sparse_data_meta(
 
 def masked_select_jagged_1d(
     values: Tensor, lengths: Tensor, mask: Tensor
-) -> Tuple[Tensor, Tensor]:
+) -> tuple[Tensor, Tensor]:
     torch._check(values.dim() == 1)
     torch._check(lengths.dim() == 1)
     torch._check(values.device == lengths.device)
@@ -228,11 +315,11 @@ def masked_select_jagged_1d(
 
 
 def tbe_input_combine_abstract(
-    indices_list: List[Tensor],
-    offsets_list: List[Tensor],
-    per_sample_weights: List[Tensor],
+    indices_list: list[Tensor],
+    offsets_list: list[Tensor],
+    per_sample_weights: list[Tensor],
     include_last_offsets: Tensor,
-) -> Tuple[Tensor, Tensor, Tensor]:
+) -> tuple[Tensor, Tensor, Tensor]:
     torch._check(len(indices_list) > 0)
     torch._check(len(indices_list) == len(offsets_list))
     torch._check(len(indices_list) == len(per_sample_weights))
@@ -247,7 +334,7 @@ def tbe_input_combine_abstract(
         torch._check(index.is_contiguous())
         torch._check(offset.is_contiguous())
         total_indices = total_indices + index.numel()
-        if guard_size_oblivious(weight.numel() > 0):
+        if guard_or_true(weight.numel() > 0):
             torch._check(weight.dim() == 1)
             torch._check(weight.numel() == index.numel())
             torch._check(weight.is_contiguous())
@@ -265,10 +352,10 @@ def tbe_input_combine_abstract(
 
 
 def tbe_input_combine_with_length_abstract(
-    indices_list: List[Tensor],
-    offsets_list: List[Tensor],
-    per_sample_weights: List[Tensor],
-) -> Tuple[Tensor, Tensor, Tensor]:
+    indices_list: list[Tensor],
+    offsets_list: list[Tensor],
+    per_sample_weights: list[Tensor],
+) -> tuple[Tensor, Tensor, Tensor]:
     torch._check(len(indices_list) > 0)
     torch._check(len(indices_list) == len(offsets_list))
     torch._check(len(indices_list) == len(per_sample_weights))
@@ -284,7 +371,7 @@ def tbe_input_combine_with_length_abstract(
         torch._check(offset.is_contiguous())
         total_indices = total_indices + index.numel()
         total_offsets = total_offsets + offset.numel()
-        if guard_size_oblivious(weight.numel() > 0):
+        if guard_or_true(weight.numel() > 0):
             torch._check(weight.dim() == 1)
             torch._check(weight.numel() == index.numel())
             torch._check(weight.is_contiguous())
@@ -336,7 +423,7 @@ def expand_into_jagged_permute_meta(
     permute: Tensor,
     input_offsets: Tensor,
     output_offsets: Tensor,
-    output_size: Tuple[int, ...],
+    output_size: tuple[int, ...],
 ) -> Tensor:
     torch._check(permute.numel() > 0, lambda: "expected {permute.numel} > 0")
     torch._check(
@@ -417,6 +504,7 @@ def int_nbit_split_embedding_codegen_lookup_function_meta(
     kINT8QparamsBytes = 8
 
     if pooling_mode == PoolingMode.NONE:
+        kINT8QparamsBytes = 4
         D = max(
             [
                 max_int2_D,
@@ -432,7 +520,7 @@ def int_nbit_split_embedding_codegen_lookup_function_meta(
         torch._check(D > 0)
         adjusted_D = D
         if SparseType.from_int(output_dtype_int) == SparseType.INT8:
-            adjusted_D += T * kINT8QparamsBytes
+            adjusted_D += kINT8QparamsBytes
         output = dev_weights.new_empty([total_L, adjusted_D], dtype=output_dtype)
         return output
 
@@ -460,7 +548,8 @@ def block_bucketize_sparse_features_meta(
     block_bucketize_pos: Optional[torch.Tensor] = None,
     keep_orig_idx: bool = False,
     total_num_blocks: Optional[torch.Tensor] = None,
-) -> Tuple[
+    keep_orig_idx_per_feature: Optional[torch.Tensor] = None,
+) -> tuple[
     torch.Tensor,
     torch.Tensor,
     Optional[torch.Tensor],
@@ -480,8 +569,43 @@ def block_bucketize_sparse_features_meta(
     )
 
 
+def block_bucketize_sparse_features_2d_weights_meta(
+    lengths: torch.Tensor,
+    indices: torch.Tensor,
+    bucketize_pos: bool,
+    sequence: bool,
+    block_sizes: torch.Tensor,
+    my_size: int,
+    weights: torch.Tensor,
+    weights_dim: int = 1,
+    batch_size_per_feature: Optional[torch.Tensor] = None,
+    max_B: int = -1,
+    block_bucketize_pos: Optional[torch.Tensor] = None,
+    keep_orig_idx: bool = False,
+    total_num_blocks: Optional[torch.Tensor] = None,
+    keep_orig_idx_per_feature: Optional[torch.Tensor] = None,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    Optional[torch.Tensor],
+    Optional[torch.Tensor],
+]:
+    # Output: lengths, indices, weights", pos?, unbucketize_permute?
+    num_buckets = my_size
+    num_features = lengths.size(0)
+    num_values = indices.size(0)
+    return (
+        lengths.new_empty([num_buckets * num_features]),
+        indices.new_empty([num_values]),
+        weights.new_empty([num_values, weights_dim]),
+        indices.new_empty([num_values]) if bucketize_pos else None,
+        indices.new_empty([num_values]),
+    )
+
+
 def merge_pooled_embeddings(
-    pooled_embeddings: List[torch.Tensor],
+    pooled_embeddings: list[torch.Tensor],
     uncat_dim_size: int,
     target_device: torch.device,
     cat_dim: int = 1,
@@ -512,7 +636,7 @@ def merge_pooled_embeddings(
 
 def permute_sparse_features_abstract(
     permute: Tensor, lengths: Tensor, indices: Tensor, weights: Optional[Tensor] = None
-) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
+) -> tuple[Tensor, Tensor, Optional[Tensor]]:
     torch._check(lengths.dtype == indices.dtype)
     torch._check(permute.device == lengths.device)
     torch._check(permute.device == indices.device)
@@ -543,7 +667,7 @@ def segment_sum_csr_abstract(
 
 def dense_to_jagged_forward(
     dense: torch.Tensor,
-    offsets: List[torch.Tensor],
+    offsets: list[torch.Tensor],
     total_L: Optional[torch.SymInt] = None,
 ) -> torch.Tensor:
     if total_L is None:
@@ -558,9 +682,9 @@ def dense_to_jagged_forward(
 
 def dense_to_jagged(
     dense: torch.Tensor,
-    offsets: List[torch.Tensor],
+    offsets: list[torch.Tensor],
     total_L: Optional[torch.SymInt] = None,
-) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+) -> tuple[torch.Tensor, list[torch.Tensor]]:
     if total_L is None:
         total_L = torch.library.get_ctx().new_dynamic_size()
     return (dense_to_jagged_forward(dense, offsets, total_L), offsets)
@@ -569,9 +693,9 @@ def dense_to_jagged(
 def batch_index_select_dim0_abstract(
     inputs: torch.Tensor,
     indices: torch.Tensor,
-    input_num_indices: List[int],
-    input_rows: List[int],
-    input_columns: List[int],
+    input_num_indices: list[int],
+    input_rows: list[int],
+    input_columns: list[int],
     permute_output_dim_0_1: bool,
 ) -> torch.Tensor:
     """
@@ -613,11 +737,11 @@ def batch_index_select_dim0_tensor_abstract(
 def batch_index_select_dim0_forward_cuda_impl_abstract(
     inputs: torch.Tensor,
     indices: torch.Tensor,
-    input_num_indices: List[int],
-    input_rows: List[int],
-    input_columns: List[int],
+    input_num_indices: list[int],
+    input_rows: list[int],
+    input_columns: list[int],
     permute_output_dim_0_1: bool,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
     num_inputs = len(input_rows)
     torch._check(len(input_num_indices) == len(input_rows))
     torch._check(len(input_num_indices) == len(input_columns))
@@ -654,7 +778,7 @@ def batch_index_select_dim0_tensor_forward_cuda_impl_abstract(
     input_rows: torch.Tensor,
     input_columns: torch.Tensor,
     permute_output_dim_0_1: bool,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
     num_inputs: int = input_rows.size(0)
     torch._check(input_num_indices.size(0) == input_rows.size(0))
     torch._check(input_num_indices.size(0) == input_columns.size(0))
@@ -699,7 +823,7 @@ def keyed_jagged_index_select_dim1_abstract(
     batch_size: torch.SymInt,
     weights: Optional[torch.Tensor] = None,
     selected_lengths_sum: Optional[torch.SymInt] = None,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
     """
     This meta function is used to calculate the shape of output tensors
     from the original function `fbgemm::keyed_jagged_index_select_dim1` without the actual data.
@@ -724,7 +848,7 @@ def keyed_jagged_index_select_dim1_abstract(
             torch.index_select(lengths, 0, length_indices).sum().item()
         )
 
-    ret: List[torch.Tensor] = [
+    ret: list[torch.Tensor] = [
         # pyre-ignore
         values.new_empty([selected_lengths_sum]),
         lengths.new_empty([indices.shape[0] * num_batches]),
@@ -756,17 +880,17 @@ def batch_index_select_dim0_backward_cuda_impl_abstract(
 def batch_index_select_dim0_forward_cpu_impl_abstract(
     inputs: torch.Tensor,
     indices: torch.Tensor,
-    input_num_indices: List[int],
-    input_rows: List[int],
-    input_columns: List[int],
+    input_num_indices: list[int],
+    input_rows: list[int],
+    input_columns: list[int],
     permute_output_dim_0_1: bool,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
     # input lists must have the same length
     num_inputs = len(input_num_indices)
     torch._check(num_inputs == len(input_rows))
     torch._check(num_inputs == len(input_columns))
 
-    if permute_output_dim_0_1 and guard_size_oblivious(len(input_num_indices) > 0):
+    if permute_output_dim_0_1 and guard_or_true(len(input_num_indices) > 0):
         # All num_indices must be the same if permute_output_dim_0_1 is True
         for x in input_num_indices:
             torch._check(x == input_num_indices[0])
@@ -790,7 +914,7 @@ def batch_index_select_dim0_tensor_forward_cpu_impl_abstract(
     input_rows: torch.Tensor,
     input_columns: torch.Tensor,
     permute_output_dim_0_1: bool,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
     # input lists must have the same length
     num_inputs = len(input_num_indices)
     torch._check(num_inputs == len(input_rows))
@@ -829,6 +953,8 @@ def bounds_check_indices_abstract(
     b_t_map: Optional[torch.Tensor] = None,
     info_B_num_bits: int = -1,
     info_B_mask: int = -1,
+    bounds_check_version: int = 1,
+    prefetch_pipeline: bool = False,
 ) -> None:
     """
     This meta function is used to fake the bounds checking
@@ -838,8 +964,8 @@ def bounds_check_indices_abstract(
 
 
 def group_index_select_dim0_gpu_impl_abstract(
-    inputs: List[torch.Tensor], group_size: int
-) -> List[torch.Tensor]:
+    inputs: list[torch.Tensor], group_size: int
+) -> list[torch.Tensor]:
     """
     Calculate output shapes for group_index_select_dim0_gpu_impl
     without the actual data.
@@ -869,8 +995,8 @@ def group_index_select_dim0_gpu_impl_abstract(
 
 
 def group_index_select_dim0_gpu_backward_abstract(
-    all_inputs: List[torch.Tensor], output_shape_group_ref: List[torch.SymInt]
-) -> List[torch.Tensor]:
+    all_inputs: list[torch.Tensor], output_shape_group_ref: list[torch.SymInt]
+) -> list[torch.Tensor]:
     """
     Calculate output shapes for group_index_select_dim0_gpu_backward
     without the actual data.
@@ -903,7 +1029,7 @@ def keyed_jagged_index_select_dim1_forward_cuda_impl_abstract(
     batch_size: torch.SymInt,
     weights: Optional[torch.Tensor] = None,
     selected_lengths_sum: Optional[torch.SymInt] = None,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
     num_batches = lengths.size(0) // batch_size
     torch._check(lengths.size(0) + 1 == offsets.size(0))
     # pyre-ignore
@@ -917,7 +1043,7 @@ def keyed_jagged_index_select_dim1_forward_cuda_impl_abstract(
         selected_lengths_sum = torch.library.get_ctx().new_dynamic_size()
 
     torch._check_is_size(selected_lengths_sum)
-    vlw: List[torch.Tensor] = [
+    vlw: list[torch.Tensor] = [
         values.new_empty([selected_lengths_sum]),  # output
         lengths.new_empty([indices.shape[0] * num_batches]),  # output_lengths
     ]
@@ -960,7 +1086,7 @@ def histogram_binning_calibration_abstract(
     upper_bound: float,
     bin_ctr_in_use_after: int,
     bin_ctr_weight_value: float,
-) -> Tuple[Tensor, Tensor]:
+) -> tuple[Tensor, Tensor]:
     return torch.empty_like(logit), torch.empty([logit.numel()], dtype=torch.int64)
 
 
@@ -1111,7 +1237,7 @@ def generic_histogram_binning_calibration_by_feature(
     positive_weight: float,
     bin_ctr_in_use_after: int,
     bin_ctr_weight_value: float,
-) -> Tuple[Tensor, Tensor]:
+) -> tuple[Tensor, Tensor]:
     torch._check(bin_num_examples.numel() == bin_num_positives.numel())
     torch._check(
         bin_num_examples.numel() == (num_segments + 1) * (bin_boundaries.numel() + 1)
@@ -1121,6 +1247,58 @@ def generic_histogram_binning_calibration_by_feature(
     )
 
 
+def permute_multi_embedding_function_impl_abstract(
+    pooled_embs: list[Tensor],
+    permutes: Tensor,
+    in_shapes: Tensor,
+    out_shapes: Tensor,
+    out_lengths: list[int],
+    reverse: bool = False,
+) -> list[Tensor]:
+    out_dtype = pooled_embs[0].dtype
+    bs = pooled_embs[0].shape[0]
+    torch._check(permutes.shape[1] == 6, lambda: "permutes must have 6 columns")
+
+    output = []
+    for i in range(len(out_lengths)):
+        output.append(torch.empty([bs, out_lengths[i]], dtype=out_dtype))
+    return output
+
+
+def lengths_range_abstract(
+    lengths: Tensor,
+    output_shape: Optional[Sequence[int]] = None,
+) -> Tensor:
+    torch._check(lengths.dim() == 1, lambda: "lengths must be a 1D tensor")
+    output_size = 0
+    if output_shape is not None:
+        output_size = math.prod(output_shape)
+    else:
+        ctx = torch.library.get_ctx()
+        output_size = ctx.new_dynamic_size()
+    return lengths.new_empty([output_size], dtype=lengths.dtype)
+
+
+def all_to_one_device(
+    input_tensors: list[Tensor],
+    target_device: torch.device,
+) -> list[Tensor]:
+    return [
+        torch.empty_like(input_tensor, device=torch.device("meta"))
+        for input_tensor in input_tensors
+    ]
+
+
+def sum_reduce_to_one(
+    input_tensors: list[Tensor],
+    target_device: torch.device,
+) -> Tensor:
+    torch._check(len(input_tensors) > 0, lambda: "reducing no tensor is undefined")
+    # All tensors should have the same shape
+    first_tensor = input_tensors[0]
+    return torch.empty_like(first_tensor, device=torch.device("meta"))
+
+
 def _setup() -> None:
     # pyre-ignore[16]
     _setup.done = getattr(_setup, "done", False)
@@ -1153,6 +1331,8 @@ def _setup() -> None:
         )
 
         impl_abstract("fbgemm::permute_2D_sparse_data", permute_2D_sparse_data_meta)
+        impl_abstract("fbgemm::get_source_mask", get_source_mask_meta)
+        impl_abstract("fbgemm::repeat_arange", repeat_arange_meta)
         impl_abstract(
             "fbgemm::permute_2D_sparse_data_input1D",
             permute_2D_sparse_data_input1D_meta,
@@ -1185,12 +1365,18 @@ def _setup() -> None:
             "fbgemm::block_bucketize_sparse_features",
             block_bucketize_sparse_features_meta,
         )
+        impl_abstract(
+            "fbgemm::block_bucketize_sparse_features_2d_weights",
+            block_bucketize_sparse_features_2d_weights_meta,
+        )
         impl_abstract("fbgemm::merge_pooled_embeddings", merge_pooled_embeddings)
         impl_abstract(
             "fbgemm::permute_sparse_features", permute_sparse_features_abstract
         )
         impl_abstract("fbgemm::segment_sum_csr", segment_sum_csr_abstract)
         impl_abstract("fbgemm::dense_to_jagged_forward", dense_to_jagged_forward)
+        impl_abstract("fbgemm::all_to_one_device", all_to_one_device)
+        impl_abstract("fbgemm::sum_reduce_to_one", sum_reduce_to_one)
         impl_abstract(
             "fbgemm::batch_index_select_dim0", batch_index_select_dim0_abstract
         )
@@ -1258,6 +1444,14 @@ def _setup() -> None:
             "fbgemm::generic_histogram_binning_calibration_by_feature",
             generic_histogram_binning_calibration_by_feature,
         )
+        impl_abstract(
+            "fbgemm::lengths_range",
+            lengths_range_abstract,
+        )
+        impl_abstract(
+            "fbgemm::permute_multi_embedding_function",
+            permute_multi_embedding_function_impl_abstract,
+        )
         impl_abstract(
             "fbgemm::FloatToHFP8Quantized",
             float_to_hfp8_quantized,
@@ -1302,29 +1496,3 @@ def _setup() -> None:
 
 
 _setup()
-
-
-@torch.library.register_fake("fbgemm::lengths_range")
-def lengths_range_abstract(
-    lengths: Tensor,
-    output_shape: Optional[Sequence[int]] = None,
-) -> Tensor:
-    torch._check(lengths.dim() == 1, lambda: "lengths must be a 1D tensor")
-    output_size = 0
-    if output_shape is not None:
-        output_size = math.prod(output_shape)
-    else:
-        ctx = torch.library.get_ctx()
-        output_size = ctx.new_dynamic_size()
-    return lengths.new_empty([output_size], dtype=lengths.dtype)
-
-
-@torch.library.register_fake("fbgemm::all_to_one_device")
-def all_to_one_device(
-    input_tensors: List[Tensor],
-    target_device: torch.device,
-) -> List[Tensor]:
-    return [
-        torch.empty_like(input_tensor, device=torch.device("meta"))
-        for input_tensor in input_tensors
-    ]

fbgemm_gpu/split_embedding_codegen_lookup_invokers/__init__.py

@@ -4,6 +4,8 @@
 ## Template Source: training/python/__init__.template
 ################################################################################
 
+__template_source_file__ = "training/python/__init__.template"
+
 #!/usr/bin/env python3
 
 # Copyright (c) Meta Platforms, Inc. and affiliates.
@@ -150,6 +152,30 @@ except:
         DeprecationWarning,
     )
     
+try:
+    # Import is placed under a try-except bc the op is experimental and can be 
+    # removed/updated in the future
+    import fbgemm_gpu.split_embedding_codegen_lookup_invokers.lookup_adam_ssd as lookup_adam_ssd  # noqa: F401
+except:
+    warnings.warn(
+        f"""\033[93m
+        Failed to import: fbgemm_gpu.split_embedding_codegen_lookup_invokers.lookup_adam_ssd
+        \033[0m""",
+        DeprecationWarning,
+    )
+    
+try:
+    # Import is placed under a try-except bc the op is experimental and can be 
+    # removed/updated in the future
+    import fbgemm_gpu.split_embedding_codegen_lookup_invokers.lookup_partial_rowwise_adam_ssd as lookup_partial_rowwise_adam_ssd  # noqa: F401
+except:
+    warnings.warn(
+        f"""\033[93m
+        Failed to import: fbgemm_gpu.split_embedding_codegen_lookup_invokers.lookup_partial_rowwise_adam_ssd
+        \033[0m""",
+        DeprecationWarning,
+    )
+    
 try:
     # Import is placed under a try-except bc the op is experimental and can be 
     # removed/updated in the future