fbgemm-gpu-hstu-nightly 2025.6.12__cp311-cp311-manylinux_2_28_x86_64.whl → 2025.6.13__cp311-cp311-manylinux_2_28_x86_64.whl

fbgemm_gpu/docs/version.py

@@ -6,6 +6,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-__version__: str = "2025.6.12"
+__version__: str = "2025.6.13"
 __target__: str = "hstu"
 __variant__: str = "cuda"

fbgemm_gpu/sparse_ops.py

@@ -420,6 +420,7 @@ def int_nbit_split_embedding_codegen_lookup_function_meta(
     kINT8QparamsBytes = 8
 
     if pooling_mode == PoolingMode.NONE:
+        kINT8QparamsBytes = 4
         D = max(
             [
                 max_int2_D,
@@ -435,7 +436,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
 

fbgemm_gpu/tbe/ssd/utils/partially_materialized_tensor.py

@@ -9,7 +9,8 @@
 from __future__ import annotations
 
 import functools
-from typing import Optional, Union
+import logging
+from typing import List, Optional, Union
 
 import torch
 
@@ -25,6 +26,58 @@ def implements(torch_function):
     return decorator
 
 
+class KVTensorMetadata:
+    """
+    Class that is used to represent a KVTensor as a Serialized Metadata in python
+    This object is used to reconstruct the KVTensor in the publish component
+    """
+
+    checkpoint_paths: List[str]
+    tbe_uuid: str
+    rdb_num_shards: int
+    rdb_num_threads: int
+    max_D: int
+    table_offset: int
+    table_shape: List[int]
+    dtype: int
+    checkpoint_uuid: str
+
+    def __init__(
+        self,
+        checkpoint_paths: List[str],
+        tbe_uuid: str,
+        rdb_num_shards: int,
+        rdb_num_threads: int,
+        max_D: int,
+        table_offset: int,
+        table_shape: List[int],
+        dtype: int,
+        checkpoint_uuid: str,
+    ) -> None:
+        """
+        Ensure caller loads the module before creating this object.
+
+        ```
+        load_torch_module(
+            "//deeplearning/fbgemm/fbgemm_gpu:ssd_split_table_batched_embeddings"
+        )
+        ```
+
+        Args:
+
+            wrapped: torch.classes.fbgemm.KVTensorWrapper
+        """
+        self.checkpoint_paths = checkpoint_paths
+        self.tbe_uuid = tbe_uuid
+        self.rdb_num_shards = rdb_num_shards
+        self.rdb_num_threads = rdb_num_threads
+        self.max_D = max_D
+        self.table_offset = table_offset
+        self.table_shape = table_shape
+        self.checkpoint_uuid = checkpoint_uuid
+        self.dtype = dtype
+
+
 class PartiallyMaterializedTensor:
     """
     A tensor-like object that represents a partially materialized tensor in memory.
@@ -51,6 +104,55 @@ class PartiallyMaterializedTensor:
         self._is_virtual = is_virtual
         self._requires_grad = False
 
+    @property
+    def generate_kvtensor_metadata(self) -> KVTensorMetadata:
+        serialized_metadata = self.wrapped.get_kvtensor_serializable_metadata()
+        try:
+            metadata_itr = 0
+            num_rdb_ckpts = int(serialized_metadata[0])
+            metadata_itr += 1
+            checkpoint_paths: List[str] = []
+            for i in range(num_rdb_ckpts):
+                checkpoint_paths.append(serialized_metadata[i + metadata_itr])
+            metadata_itr += num_rdb_ckpts
+            tbe_uuid = serialized_metadata[metadata_itr]
+            metadata_itr += 1
+            rdb_num_shards = int(serialized_metadata[metadata_itr])
+            metadata_itr += 1
+            rdb_num_threads = int(serialized_metadata[metadata_itr])
+            metadata_itr += 1
+            max_D = int(serialized_metadata[metadata_itr])
+            metadata_itr += 1
+            table_offset = int(serialized_metadata[metadata_itr])
+            metadata_itr += 1
+            table_shape: List[int] = []
+            table_shape.append(int(serialized_metadata[metadata_itr]))
+            metadata_itr += 1
+            table_shape.append(int(serialized_metadata[metadata_itr]))
+            metadata_itr += 1
+            dtype = int(serialized_metadata[metadata_itr])
+            metadata_itr += 1
+            checkpoint_uuid = serialized_metadata[metadata_itr]
+            metadata_itr += 1
+            res = KVTensorMetadata(
+                checkpoint_paths,
+                tbe_uuid,
+                rdb_num_shards,
+                rdb_num_threads,
+                max_D,
+                table_offset,
+                table_shape,
+                dtype,
+                checkpoint_uuid,
+            )
+
+            return res
+        except Exception as e:
+            logging.error(
+                f"Failed to parse metadata: {e}, here is metadata: {serialized_metadata}"
+            )
+            raise e
+
     @property
     def wrapped(self):
         """
@@ -249,6 +351,9 @@ class PartiallyMaterializedTensor:
 
         return torch.equal(tensor1.full_tensor(), tensor2.full_tensor())
 
+    def get_kvtensor_serializable_metadata(self) -> List[str]:
+        return self._wrapped.get_kvtensor_serializable_metadata()
+
     def __hash__(self):
         return id(self)
 

{fbgemm_gpu_hstu_nightly-2025.6.12.dist-info → fbgemm_gpu_hstu_nightly-2025.6.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fbgemm_gpu_hstu_nightly
-Version: 2025.6.12
+Version: 2025.6.13
 Home-page: https://github.com/pytorch/fbgemm
 Author: FBGEMM Team
 Author-email: packages@pytorch.org

{fbgemm_gpu_hstu_nightly-2025.6.12.dist-info → fbgemm_gpu_hstu_nightly-2025.6.13.dist-info}/RECORD

@@ -9,7 +9,7 @@ fbgemm_gpu/permute_pooled_embedding_modules_split.py,sha256=cUrEbRIvLFW_3Zmh07Qk
 fbgemm_gpu/quantize_comm.py,sha256=YUzk8F1MZckbkseEoDJ4rOKiglGCGtrPdR1IKGD6Mk4,11177
 fbgemm_gpu/quantize_utils.py,sha256=hb8G_1xzRwYmwFp8VJrmoaolNxCwqcFwkwnyStk1C0w,7394
 fbgemm_gpu/runtime_monitor.py,sha256=HM_0cxMO7uuAq8sCiv2lmGgp1jKGzba2qhuUcGyRMog,7425
-fbgemm_gpu/sparse_ops.py,sha256=84x_hobQofY29Vzlz0eJxe126Ba-O5oSEQREMpRgOtE,47168
+fbgemm_gpu/sparse_ops.py,sha256=xNoRMp6QNMz8Lq_5LE9IprQtUip3gkTVbyPgJ1AiWSI,47194
 fbgemm_gpu/split_embedding_configs.py,sha256=DcZ7SV4AmhlN9QPDaJBqzigR-c1zM_bZI3Fh4PYuab4,7266
 fbgemm_gpu/split_embedding_inference_converter.py,sha256=ilVVowkTiY0WDpOYorj917Tqsez4KWNBdTXuz2bWbp8,7063
 fbgemm_gpu/split_embedding_optimizer_ops.py,sha256=wXuGazClBMk62yL_r9udUIKaPgQP7SlkSb5ugB75wrQ,711
@@ -32,10 +32,10 @@ fbgemm_gpu/docs/merge_pooled_embedding_ops.py,sha256=oJLgSgZQmhsyGLbTmZTxNgQrk65
 fbgemm_gpu/docs/permute_pooled_embedding_ops.py,sha256=tZUqLVXlk5O6VAKKDA-OEMx2fCu5QPOOeoAPZA9_nLY,4454
 fbgemm_gpu/docs/quantize_ops.py,sha256=xTtOaVK1P02ymreE_i21YiyYDZCqhoZY9eWp_mEIRlo,1297
 fbgemm_gpu/docs/sparse_ops.py,sha256=NTcTm0q9h8W2B8PKPoic2fHsAaCbCYunSa_EYK0LtHQ,21382
-fbgemm_gpu/docs/version.py,sha256=XOS_m9ZE4Pr3zwz272OUxVj3s034QJ5xDrpkMLk-x5s,315
+fbgemm_gpu/docs/version.py,sha256=9kytZ-RJTwv5Z4Lajmjf4ELvL28p4wC2CGW1Zv5ZfSg,315
 fbgemm_gpu/experimental/hstu/__init__.py,sha256=KNisP6qDMwgjgxkGlqUZRNjJ_8o8R-cTmm3HxF7pSqI,1564
 fbgemm_gpu/experimental/hstu/cuda_hstu_attention.py,sha256=5425GRjJuzpXQC-TowgQOCFjZmOwv_EK0lKbURhHBTQ,9920
-fbgemm_gpu/experimental/hstu/fbgemm_gpu_experimental_hstu.so,sha256=KcSmmINAo_BjNginy3sl3FJ5JUJZVN6swdv_OKB576g,352287576
+fbgemm_gpu/experimental/hstu/fbgemm_gpu_experimental_hstu.so,sha256=pDSxzlMTxno_1ZdHPfpE2nR1Vf1LNSyjQGMSrKejE5g,352287576
 fbgemm_gpu/quantize/__init__.py,sha256=pftciXHE7csekDFkl7Ui1AWglVMMnSrOO04mREnUdb0,921
 fbgemm_gpu/quantize/quantize_ops.py,sha256=25AIOv9n2UoxamMUaI6EK1Ur4gSHxbZIReHBtgOjjCs,2228
 fbgemm_gpu/sll/__init__.py,sha256=rgXh35-OFUE54E9gGBq3NGxouGvgMv2ccY2bWUTxONY,4191
@@ -75,7 +75,7 @@ fbgemm_gpu/tbe/ssd/common.py,sha256=1J8K7sTQswgCYWaVwF-ZdCJj7mNN6O9GI70AaZWzJGE,
 fbgemm_gpu/tbe/ssd/inference.py,sha256=DTjwj3f6JaUMcecWoRNkZpRgXDJ-eE3grtixYwKb5DI,22829
 fbgemm_gpu/tbe/ssd/training.py,sha256=gCvMY8fGP2JNosMyzW7uiQUuUsDIdUULB_PG51wOpRk,131519
 fbgemm_gpu/tbe/ssd/utils/__init__.py,sha256=5DgmR2HA6NtmYh2ddkUgpDsZ6a7hF0DPedA1gMpdh18,250
-fbgemm_gpu/tbe/ssd/utils/partially_materialized_tensor.py,sha256=V18ZQxdYJNSi6qLqhFaxXj3IxpcPgan3GsoXrosoZ1Q,7510
+fbgemm_gpu/tbe/ssd/utils/partially_materialized_tensor.py,sha256=ePF2b3vCFB_ZzPaMneaXe5gODFVRJmXKbagT2uNKAAQ,10934
 fbgemm_gpu/tbe/stats/__init__.py,sha256=on29iDtq7cVNh90JR9aeFNG-K9DDoYq0JryzoplL49I,322
 fbgemm_gpu/tbe/stats/bench_params_reporter.py,sha256=7XIWVObJOxSVUG73xsd_lVSuCFUQkMEGSWW--BoyCH0,7358
 fbgemm_gpu/tbe/utils/__init__.py,sha256=rlXFm-kTByFZO4SS5C5zMzANRiQmM1NT__eWBayncYg,549
@@ -93,7 +93,7 @@ fbgemm_gpu/utils/__init__.py,sha256=JQQNdcTTaEU6ptK-OW-ZQBwTFxEZZpWOtBXWwEZm39o,
 fbgemm_gpu/utils/filestore.py,sha256=Zshw1dA03m9aHMMAtETdq4bgOLocyLhzlkAUoG8VkdM,4743
 fbgemm_gpu/utils/loader.py,sha256=1hCEhNvkflniH46fGcrguLeP1z-6uyOu2QFwqKU5CIM,990
 fbgemm_gpu/utils/torch_library.py,sha256=dQcHv1qgpu5QYlJjxjd6oeHjtxnmmXzx3PL6vjCmxL4,4199
-fbgemm_gpu_hstu_nightly-2025.6.12.dist-info/METADATA,sha256=ayz07_xnYRnh1FYOD4_6vcF3Pn4Ngw8hmoQhOBEfOz8,2794
-fbgemm_gpu_hstu_nightly-2025.6.12.dist-info/WHEEL,sha256=V2Q6mQKbouIadCxoRjt9FQ9oKfi45-uZUcoc77zzs0M,108
-fbgemm_gpu_hstu_nightly-2025.6.12.dist-info/top_level.txt,sha256=2tlbTWLkPjhqvLF_6BbqKzkcPluSE-oPRVjI8axK76I,11
-fbgemm_gpu_hstu_nightly-2025.6.12.dist-info/RECORD,,
+fbgemm_gpu_hstu_nightly-2025.6.13.dist-info/METADATA,sha256=G6aLh_w9SYySzT7xRYMUj8OW6l873R2B71m-tgD6GWs,2794
+fbgemm_gpu_hstu_nightly-2025.6.13.dist-info/WHEEL,sha256=V2Q6mQKbouIadCxoRjt9FQ9oKfi45-uZUcoc77zzs0M,108
+fbgemm_gpu_hstu_nightly-2025.6.13.dist-info/top_level.txt,sha256=2tlbTWLkPjhqvLF_6BbqKzkcPluSE-oPRVjI8axK76I,11
+fbgemm_gpu_hstu_nightly-2025.6.13.dist-info/RECORD,,