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/tbe/bench/tbe_data_config_loader.py

@@ -7,12 +7,56 @@
 
 # pyre-strict
 
+import dataclasses
+import logging
+import re
+from enum import Enum
+
 import click
 import torch
 import yaml
 
-from .tbe_data_config import TBEDataConfig
-from .tbe_data_config_param_models import BatchParams, IndicesParams, PoolingParams
+# fmt:skip
+from fbgemm_gpu.tbe.bench.tbe_data_config import (
+    BatchParams,
+    IndicesParams,
+    PoolingParams,
+    TBEDataConfig,
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class TBEDataConfigHelperText(Enum):
+    # Config File
+    TBE_CONFIG = "TBE data configuration filepath.  If provided, all other `--tbe-*` options are ignored."
+
+    # Table Parameters
+    TBE_NUM_TABLES = "Number of tables (T)"
+    TBE_NUM_EMBEDDINGS = "Number of embeddings (E)"
+    TBE_EMBEDDING_DIM = "Embedding dimensions (D)"
+    TBE_MIXED_DIM = "Use mixed dimensions"
+    TBE_WEIGHTED = "Flag to indicate if the table is weighted"
+
+    # Batch Parameters
+    TBE_BATCH_SIZE = "Batch size (B)"
+    TBE_BATCH_VBE_SIGMA = "Standard deviation of B for VBE"
+    TBE_BATCH_VBE_DIST = "VBE distribution (choices: 'uniform', 'normal')"
+    TBE_BATCH_VBE_RANKS = "Number of ranks for VBE"
+
+    # Indices Parameters
+    TBE_INDICES_HITTERS = "Heavy hitters for indices (comma-delimited list of floats)"
+    TBE_INDICES_ZIPF = "Zipf distribution parameters for indices generation (q, s)"
+    TBE_INDICES_DTYPE = "The dtype of the table indices (choices: '32', '64')"
+    TBE_OFFSETS_DTYPE = "The dtype of the table offsets (choices: '32', '64')"
+
+    # Pooling Parameters
+    TBE_POOLING_SIZE = "Bag size / pooling factor (L)"
+    TBE_POOLING_VL_SIGMA = "Standard deviation of L for variable bag size"
+    TBE_POOLING_VL_DIST = (
+        "Variable bag size distribution (choices: 'uniform', 'normal')"
+    )
+    TBE_EMBEDDING_SPECS = "Embedding Specs which is List[Tuple[int, int, EmbeddingLocation, ComputeDevice]]"
+    TBE_FEATURE_TABLE_MAP = "Mapping of feature-table"
 
 
 class TBEDataConfigLoader:
@@ -20,119 +64,152 @@ class TBEDataConfigLoader:
     # pyre-ignore [2]
     def options(cls, func) -> click.Command:
         options = [
-            ####################################################################
             # Config File
-            ####################################################################
             click.option(
                 "--tbe-config",
                 type=str,
                 required=False,
-                help="TBE data configuration filepath.  If provided, all other `--tbe-*` options are ignored.",
+                help=TBEDataConfigHelperText.TBE_CONFIG.value,
             ),
-            ####################################################################
             # Table Parameters
-            ####################################################################
             click.option(
                 "--tbe-num-tables",
                 type=int,
                 default=32,
-                help="Number of tables (T)",
+                help=TBEDataConfigHelperText.TBE_NUM_TABLES.value,
             ),
             click.option(
                 "--tbe-num-embeddings",
                 type=int,
                 default=int(1e5),
-                help="Number of embeddings (E)",
+                help=TBEDataConfigHelperText.TBE_NUM_EMBEDDINGS.value,
+            ),
+            click.option(
+                "--tbe-num-embeddings-list",
+                type=str,
+                required=False,
+                default=None,
+                help="Comma-separated list of number of embeddings (Es)",
             ),
             click.option(
                 "--tbe-embedding-dim",
                 type=int,
                 default=128,
-                help="Embedding dimensions (D)",
+                help=TBEDataConfigHelperText.TBE_EMBEDDING_DIM.value,
+            ),
+            click.option(
+                "--tbe-embedding-dim-list",
+                type=str,
+                required=False,
+                default=None,
+                help="Comma-separated list of number of Embedding dimensions (D)",
             ),
             click.option(
                 "--tbe-mixed-dim",
                 is_flag=True,
                 default=False,
-                help="Use mixed dimensions",
+                help=TBEDataConfigHelperText.TBE_MIXED_DIM.value,
             ),
             click.option(
                 "--tbe-weighted",
                 is_flag=True,
                 default=False,
-                help="Whether the table is weighted or not",
+                help=TBEDataConfigHelperText.TBE_WEIGHTED.value,
+            ),
+            click.option(
+                "--tbe-max-indices",
+                type=int,
+                required=False,
+                default=None,
+                help="(Optional) Maximum number of indices, will be calculated if not provided",
             ),
-            ####################################################################
             # Batch Parameters
-            ####################################################################
             click.option(
-                "--tbe-batch-size", type=int, default=512, help="Batch size (B)"
+                "--tbe-batch-size",
+                type=int,
+                default=512,
+                help=TBEDataConfigHelperText.TBE_BATCH_SIZE.value,
+            ),
+            click.option(
+                "--tbe-batch-sizes-list",
+                type=str,
+                required=False,
+                default=None,
+                help="List Batch sizes per feature (Bs)",
             ),
             click.option(
                 "--tbe-batch-vbe-sigma",
                 type=int,
                 required=False,
-                help="Standard deviation of B for VBE",
+                help=TBEDataConfigHelperText.TBE_BATCH_VBE_SIGMA.value,
             ),
             click.option(
                 "--tbe-batch-vbe-dist",
                 type=click.Choice(["uniform", "normal"]),
                 required=False,
-                help="VBE distribution",
+                help=TBEDataConfigHelperText.TBE_BATCH_VBE_DIST.value,
             ),
             click.option(
                 "--tbe-batch-vbe-ranks",
                 type=int,
                 required=False,
-                help="Number of ranks for VBE",
+                help=TBEDataConfigHelperText.TBE_BATCH_VBE_RANKS.value,
             ),
-            ####################################################################
             # Indices Parameters
-            ####################################################################
             click.option(
                 "--tbe-indices-hitters",
                 type=str,
                 default="",
-                help="TBE heavy hitter indices (comma-delimited list of floats)",
+                help=TBEDataConfigHelperText.TBE_INDICES_HITTERS.value,
             ),
             click.option(
                 "--tbe-indices-zipf",
                 type=(float, float),
                 default=(0.1, 0.1),
-                help="Zipf distribution parameters for indices generation (q, s)",
+                help=TBEDataConfigHelperText.TBE_INDICES_ZIPF.value,
             ),
             click.option(
                 "--tbe-indices-dtype",
                 type=click.Choice(["32", "64"]),
                 default="64",
-                help="The dtype of the table indices",
+                help=TBEDataConfigHelperText.TBE_INDICES_DTYPE.value,
             ),
             click.option(
                 "--tbe-offsets-dtype",
                 type=click.Choice(["32", "64"]),
                 default="64",
-                help="The dtype of the table indices offsets",
+                help=TBEDataConfigHelperText.TBE_OFFSETS_DTYPE.value,
             ),
-            ####################################################################
             # Pooling Parameters
-            ####################################################################
             click.option(
                 "--tbe-pooling-size",
                 type=int,
                 default=20,
-                help="Bag size / pooling factor (L)",
+                help=TBEDataConfigHelperText.TBE_POOLING_SIZE.value,
             ),
             click.option(
                 "--tbe-pooling-vl-sigma",
                 type=int,
                 required=False,
-                help="Standard deviation of B for VBE",
+                help=TBEDataConfigHelperText.TBE_POOLING_VL_SIGMA.value,
             ),
             click.option(
                 "--tbe-pooling-vl-dist",
                 type=click.Choice(["uniform", "normal"]),
                 required=False,
-                help="Pooling factor distribution",
+                help=TBEDataConfigHelperText.TBE_POOLING_VL_DIST.value,
+            ),
+            click.option(
+                "--tbe-embedding-specs",
+                type=str,
+                required=False,
+                help=TBEDataConfigHelperText.TBE_EMBEDDING_SPECS.value,
+            ),
+            click.option(
+                "--tbe-feature-table-map",
+                type=str,
+                required=False,
+                help=TBEDataConfigHelperText.TBE_FEATURE_TABLE_MAP.value,
             ),
         ]
 
@@ -154,18 +231,62 @@ class TBEDataConfigLoader:
         params = context.params
 
         # Read table parameters
-        T = params["tbe_num_tables"]
-        E = params["tbe_num_embeddings"]
+        T = params["tbe_num_tables"]  # number of features
+        E = params["tbe_num_embeddings"]  # feature_rows
+        if params["tbe_num_embeddings_list"] is not None:
+            Es = [int(x) for x in params["tbe_num_embeddings_list"].split(",")]
+            T = len(Es)
+            E = sum(Es) // T  # average E
+        else:
+            Es = None
         D = params["tbe_embedding_dim"]
+        if params["tbe_embedding_dim_list"] is not None:
+            Ds = [int(x) for x in params["tbe_embedding_dim_list"].split(",")]
+            assert (
+                len(Ds) == T
+            ), f"Expected tbe_embedding_dim_list to have {T} elements, but got {len(Ds)}"
+            D = sum(Ds) // T  # average D
+        else:
+            Ds = None
+
         mixed_dim = params["tbe_mixed_dim"]
         weighted = params["tbe_weighted"]
+        if params["tbe_max_indices"] is not None:
+            max_indices = params["tbe_max_indices"]
+        else:
+            max_indices = None
 
         # Read batch parameters
         B = params["tbe_batch_size"]
         sigma_B = params["tbe_batch_vbe_sigma"]
         vbe_distribution = params["tbe_batch_vbe_dist"]
         vbe_num_ranks = params["tbe_batch_vbe_ranks"]
-        batch_params = BatchParams(B, sigma_B, vbe_distribution, vbe_num_ranks)
+        if params["tbe_batch_sizes_list"] is not None:
+            Bs = [int(x) for x in params["tbe_batch_sizes_list"].split(",")]
+            B = sum(Bs) // T  # average B
+        else:
+            B = params["tbe_batch_size"]
+            Bs = None
+        batch_params = BatchParams(B, sigma_B, vbe_distribution, vbe_num_ranks, Bs)
+
+        # Parse embedding_specs: "(E,D),(E,D),..." or "(E,D,loc,dev),(E,D,loc,dev),..."
+        # Only the first two values (E, D) are extracted.
+        embedding_specs = None
+        feature_table_map = None
+        if params["tbe_embedding_specs"] is not None:
+            try:
+                tuples = re.findall(r"\(([^)]+)\)", params["tbe_embedding_specs"])
+                if tuples:
+                    embedding_specs = [
+                        (int(t.split(",")[0].strip()), int(t.split(",")[1].strip()))
+                        for t in tuples
+                    ]
+            except (ValueError, IndexError):
+                logging.warning("Failed to parse embedding_specs. Setting to None.")
+        if params["tbe_feature_table_map"] is not None:
+            feature_table_map = [
+                int(x) for x in params["tbe_feature_table_map"].split(",")
+            ]
 
         # Read indices parameters
         heavy_hitters = (
@@ -200,6 +321,11 @@ class TBEDataConfigLoader:
             indices_params,
             pooling_params,
             not torch.cuda.is_available(),
+            Es,
+            Ds,
+            max_indices,
+            embedding_specs,
+            feature_table_map,
         ).validate()
 
     @classmethod

fbgemm_gpu/tbe/bench/tbe_data_config_param_models.py

@@ -9,7 +9,7 @@
 
 import dataclasses
 import json
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import torch
 
@@ -31,6 +31,7 @@ class IndicesParams:
     heavy_hitters: torch.Tensor
     # zipf*: parameters for the Zipf distribution (x+q)^{-s}
     zipf_q: float
+    # zipf_s is synonymous with alpha in the literature
     zipf_s: float
     # [Optional] dtype for indices tensor
     index_dtype: Optional[torch.dtype] = None
@@ -39,7 +40,7 @@ class IndicesParams:
 
     @classmethod
     # pyre-ignore [3]
-    def from_dict(cls, data: Dict[str, Any]):
+    def from_dict(cls, data: dict[str, Any]):
         if not isinstance(data["heavy_hitters"], torch.Tensor):
             data["heavy_hitters"] = torch.tensor(
                 data["heavy_hitters"], dtype=torch.float32
@@ -53,7 +54,7 @@ class IndicesParams:
     def from_json(cls, data: str):
         return cls.from_dict(json.loads(data))
 
-    def dict(self) -> Dict[str, Any]:
+    def dict(self) -> dict[str, Any]:
         # https://stackoverflow.com/questions/73735974/convert-dataclass-of-dataclass-to-json-string
         tmp = dataclasses.asdict(self)
         # Convert tensor to list for JSON serialization
@@ -97,10 +98,12 @@ class BatchParams:
     vbe_distribution: Optional[str] = "normal"
     # Number of ranks for variable batch size generation
     vbe_num_ranks: Optional[int] = None
+    # List of target batch sizes, i.e. number of batch lookups per feature
+    Bs: Optional[list[int]] = None
 
     @classmethod
     # pyre-ignore [3]
-    def from_dict(cls, data: Dict[str, Any]):
+    def from_dict(cls, data: dict[str, Any]):
         return cls(**data)
 
     @classmethod
@@ -108,7 +111,7 @@ class BatchParams:
     def from_json(cls, data: str):
         return cls.from_dict(json.loads(data))
 
-    def dict(self) -> Dict[str, Any]:
+    def dict(self) -> dict[str, Any]:
         return dataclasses.asdict(self)
 
     def json(self, format: bool = False) -> str:
@@ -116,7 +119,10 @@ class BatchParams:
 
     # pyre-ignore [3]
     def validate(self):
-        assert self.B > 0, "B must be positive"
+        if self.Bs is not None:
+            assert all(b > 0 for b in self.Bs), "All elements in Bs must be positive"
+        else:
+            assert self.B > 0, "B must be positive"
         assert not self.sigma_B or self.sigma_B > 0, "sigma_B must be positive"
         assert (
             self.vbe_num_ranks is None or self.vbe_num_ranks > 0
@@ -136,10 +142,12 @@ class PoolingParams:
     sigma_L: Optional[int] = None
     # [Optional] Distribution of embedding sequence lengths (normal, uniform)
     length_distribution: Optional[str] = "normal"
+    # [Optional] List of target bag sizes, i.e. pooling factors per batch
+    Ls: Optional[list[float]] = None
 
     @classmethod
     # pyre-ignore [3]
-    def from_dict(cls, data: Dict[str, Any]):
+    def from_dict(cls, data: dict[str, Any]):
         return cls(**data)
 
     @classmethod
@@ -147,7 +155,7 @@ class PoolingParams:
     def from_json(cls, data: str):
         return cls.from_dict(json.loads(data))
 
-    def dict(self) -> Dict[str, Any]:
+    def dict(self) -> dict[str, Any]:
         return dataclasses.asdict(self)
 
     def json(self, format: bool = False) -> str:

fbgemm_gpu/tbe/bench/utils.py

@@ -6,15 +6,14 @@
 
 # pyre-strict
 
-import logging
+from typing import List, Tuple
 
 import numpy as np
 import torch
 
+# fmt:skip
 from fbgemm_gpu.split_embedding_configs import SparseType
 
-logging.basicConfig(level=logging.DEBUG)
-
 
 def fill_random_scale_bias(
     emb: torch.nn.Module,
@@ -23,9 +22,9 @@ def fill_random_scale_bias(
 ) -> None:
     for t in range(T):
         # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
-        (weights, scale_shift) = emb.split_embedding_weights()[t]
+        weights, scale_shift = emb.split_embedding_weights()[t]
         if scale_shift is not None:
-            (E, R) = scale_shift.shape
+            E, R = scale_shift.shape
             assert R == 4
             scales = None
             shifts = None
@@ -46,3 +45,128 @@ def fill_random_scale_bias(
                     device=scale_shift.device,
                 )
             )
+
+
+def check_oom(
+    data_size: int,
+) -> Tuple[bool, str]:
+    free_memory, total_memory = torch.cuda.mem_get_info()
+    if data_size > free_memory:
+        warning = f"Expect to allocate {round(data_size / (1024 ** 3), 2)} GB, but available memory is {round(free_memory / (1024 ** 3), 2)} GB from {round(total_memory / (1024 ** 3), 2)} GB."
+        return (True, warning)
+    return (False, "")
+
+
+def generate_batch_size_per_feature_per_rank(
+    Bs: List[int], num_ranks: int
+) -> List[List[int]]:
+    """
+    Generate batch size per feature per rank for VBE, assuming the batch size
+    is evenly distributed across ranks.
+    Args:
+        Bs (List[int]): batch size per feature
+        num_ranks (int): number of ranks
+    Returns:
+        List[List[int]]: batch size per feature per rank
+    """
+    b_per_feature_per_rank = []
+    for B in Bs:
+        b_per_feature = []
+        for i in range(num_ranks):
+            if i != num_ranks - 1:
+                b_per_feature.append(int(B / num_ranks))
+            else:
+                b_per_feature.append(B - sum(b_per_feature))
+        b_per_feature_per_rank.append(b_per_feature)
+    return b_per_feature_per_rank
+
+
+def generate_merged_output_and_offsets(
+    Ds: List[int],
+    Bs: List[int],
+    output_dtype: torch.dtype,
+    device: torch.device,
+    num_ranks: int = 2,
+    num_tbe_ops: int = 2,
+) -> Tuple[List[List[int]], torch.Tensor, torch.Tensor]:
+    """
+    Generate merged vbe_output and vbe_output_offsets tensors for VBE.
+    The vbe_output is a tensor that will contain forward output from all VBE TBE ops.
+    The vbe_output_offsets is a tensor that will contain start offsets for the output to be written to.
+
+    Args:
+        Ds (List[int]): embedding dimension per feature
+        Bs (List[int]): batch size per feature
+        num_ranks (int): number of ranks
+        num_tbe_ops (int): number of TBE ops
+    Returns:
+        Tuple[List[List[int]], torch.Tensor, torch.Tensor]: batch_size_per_feature_per_rank, merged vbe_output and vbe_output_offsets tensors
+    """
+    # The first embedding ops is the embedding op created in the benchmark
+    emb_op = {}
+    emb_op[0] = {}
+    emb_op[0]["dim"] = Ds
+    emb_op[0]["Bs"] = Bs
+    emb_op[0]["output_size"] = sum([b * d for b, d in zip(Bs, Ds)])
+    emb_op[0]["batch_size_per_feature_per_rank"] = (
+        generate_batch_size_per_feature_per_rank(Bs, num_ranks)
+    )
+    num_features = len(Bs)
+    # create other embedding ops to allocate output and offsets tensors
+    # Using representative values for additional TBE ops in multi-op scenarios:
+    # - batch_size=32000: typical large batch size for production workloads
+    # - dim=512: common embedding dimension for large models
+    for i in range(1, num_tbe_ops):
+        emb_op[i] = {}
+        emb_op[i]["batch_size_per_feature_per_rank"] = (
+            generate_batch_size_per_feature_per_rank([32000], num_ranks)
+        )
+        emb_op[i]["Bs"] = [sum(B) for B in emb_op[i]["batch_size_per_feature_per_rank"]]
+        emb_op[i]["dim"] = [512]
+        emb_op[i]["output_size"] = sum(
+            [b * d for b, d in zip(emb_op[i]["Bs"], emb_op[i]["dim"])]
+        )
+    total_output = 0
+    ranks = [[] for _ in range(num_ranks)]
+    for e in emb_op.values():
+        b_per_rank_per_feature = list(zip(*e["batch_size_per_feature_per_rank"]))
+        assert len(b_per_rank_per_feature) == num_ranks
+        dims = e["dim"]
+        for r, b_r in enumerate(b_per_rank_per_feature):
+            for f, b in enumerate(b_r):
+                output_size_per_batch = b * dims[f]
+                ranks[r].append(output_size_per_batch)
+                total_output += output_size_per_batch
+    ranks[0].insert(0, 0)
+    offsets_ranks: List[List[int]] = [[] for _ in range(num_ranks)]
+    total_output_offsets = []
+    start = 0
+    for r in range(num_ranks):
+        offsets_ranks[r] = [
+            start + sum(ranks[r][: i + 1]) for i in range(len(ranks[r]))
+        ]
+        start = offsets_ranks[r][-1]
+        total_output_offsets.extend(offsets_ranks[r])
+    check_total_output_size = sum([e["output_size"] for e in emb_op.values()])
+    assert (
+        total_output == check_total_output_size
+    ), f"{total_output} != {check_total_output_size}{[e['output_size'] for e in emb_op.values()]}"
+    assert (
+        total_output == total_output_offsets[-1]
+    ), f"{total_output} != {total_output_offsets[-1]}"
+    out = torch.empty(total_output, dtype=output_dtype, device=device)
+    offsets = []
+    offsets.append(offsets_ranks[0][:num_features])
+    for r in range(1, num_ranks):
+        start = [offsets_ranks[r - 1][-1]]
+        the_rest = offsets_ranks[r][: num_features - 1] if num_features > 1 else []
+        start.extend(the_rest)
+        offsets.append(start)
+
+    out_offsets = torch.tensor(
+        offsets,
+        dtype=torch.int64,
+        device=device,
+    )
+    batch_size_per_feature_per_rank = emb_op[0]["batch_size_per_feature_per_rank"]
+    return (batch_size_per_feature_per_rank, out, out_offsets)

fbgemm_gpu/tbe/cache/__init__.py

@@ -7,4 +7,5 @@
 
 # pyre-unsafe
 
+from .kv_embedding_ops_inference import KVEmbeddingInference  # noqa: F401
 from .split_embeddings_cache_ops import get_unique_indices_v2  # noqa: F401