diff --git a/.gitmodules b/.gitmodules index df2ab76c8..8e0a2c447 100644 --- a/.gitmodules +++ b/.gitmodules @@ -10,4 +10,4 @@ url = https://github.com/taco-project/FlexKV.git [submodule "third_party/nv-embedding-cache"] path = third_party/nv-embedding-cache - url = https://github.com/NVIDIA/nv-embedding-cache.git + url = https://github.com/geoffreyQiu/nv-embedding-cache.git diff --git a/corelib/dynamicemb/CMakeLists.txt b/corelib/dynamicemb/CMakeLists.txt index dc0b385c0..94b105036 100644 --- a/corelib/dynamicemb/CMakeLists.txt +++ b/corelib/dynamicemb/CMakeLists.txt @@ -3,6 +3,8 @@ project(DynamicEmbInferenceOps LANGUAGES CXX CUDA) include(GNUInstallDirs) +set(NVE_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/nv-embedding-cache" CACHE PATH "NVE source root") + set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CXX_EXTENSIONS OFF) @@ -42,6 +44,9 @@ message(STATUS "Found Torch ${TORCH_VERSION}") message(STATUS "CUDA architectures: ${CMAKE_CUDA_ARCHITECTURES}") set(INFERENCE_EMB_SOURCES + src/exportable_embedding/indexer_directory.cpp + src/exportable_embedding/indexer_snapshot.cpp + src/exportable_embedding/indexer_ops.cu src/table_operation/lookup_torch_binding.cu src/table_operation/get_table_range_torch_binding.cu src/table_operation/expand_table_ids_torch_binding.cu @@ -61,8 +66,10 @@ if(TORCH_CXX_FLAGS_LIST) endif() target_include_directories(inference_emb_ops PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CURRENT_SOURCE_DIR}/src ${CMAKE_CURRENT_SOURCE_DIR}/src/table_operation + ${NVE_ROOT}/third_party/json/include ${Python3_INCLUDE_DIRS} ${TORCH_INCLUDE_DIRS} ${CUDAToolkit_INCLUDE_DIRS} @@ -91,10 +98,6 @@ target_compile_options(inference_emb_ops PRIVATE > ) -target_compile_definitions(inference_emb_ops PRIVATE - TORCH_EXTENSION_NAME=inference_emb_ops -) - target_link_libraries(inference_emb_ops PRIVATE Python3::Python ${TORCH_LIBRARIES} @@ -111,3 +114,134 @@ install(TARGETS inference_emb_ops LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} ) + +# Unset selects the repository NVE (26.07 or later); set 26.05 explicitly for compatibility. +if( + NOT DEFINED NVE_VERSION + OR NVE_VERSION STREQUAL "" + OR NVE_VERSION VERSION_GREATER_EQUAL "26.06" +) + set(DYNAMICEMB_HAS_NVE_UPDATE ON) + include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/NVETarget.cmake) + set(NVE_LIB_DIR "/opt/nve/default/python/pynve" CACHE PATH "NVE library directory") + set(NVTX_INCLUDE_DIR "/workspace/deps/NVTX/c/include" CACHE PATH "NVTX include directory") + find_library(NVE_TORCH_OPS_LIB nve-torch-ops PATHS "${NVE_LIB_DIR}" NO_DEFAULT_PATH REQUIRED) + find_library(NVE_COMMON_LIB nve-common PATHS "${NVE_LIB_DIR}" NO_DEFAULT_PATH REQUIRED) + + add_library(inference_emb_update SHARED + src/exportable_embedding/incremental_update.cpp + ) + set_source_files_properties( + src/exportable_embedding/incremental_update.cpp + PROPERTIES LANGUAGE CUDA + ) + target_include_directories(inference_emb_update PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + "${NVE_ROOT}" + "${NVE_ROOT}/include" + "${NVE_ROOT}/third_party/json/include" + "${NVTX_INCLUDE_DIR}" + ${TORCH_INCLUDE_DIRS} + ${CUDAToolkit_INCLUDE_DIRS} + ) + if(TORCH_CXX_FLAGS_LIST) + target_compile_options(inference_emb_update PRIVATE ${TORCH_CXX_FLAGS_LIST}) + endif() + target_link_options(inference_emb_update PRIVATE "LINKER:--no-as-needed") + dynamicemb_configure_nve_target(inference_emb_update) + target_link_libraries(inference_emb_update PRIVATE + inference_emb_ops + ${TORCH_LIBRARIES} + CUDA::cudart + CUDA::cuda_driver + "${NVE_TORCH_OPS_LIB}" + "${NVE_COMMON_LIB}" + ) + set_target_properties(inference_emb_update PROPERTIES + PREFIX "" + OUTPUT_NAME "inference_emb_update" + INSTALL_RPATH "${NVE_LIB_DIR}" + ) + install(TARGETS inference_emb_update + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) +else() + set(DYNAMICEMB_HAS_NVE_UPDATE OFF) +endif() + +find_library(TORCH_PYTHON_LIBRARY torch_python + PATHS "${TORCH_INSTALL_PREFIX}/lib" + REQUIRED +) + +set(EXPORTABLE_EMBEDDING_PYBIND_SOURCES + src/exportable_embedding/exportable_embedding_pybind.cpp + src/exportable_embedding/indexer_directory_pybind.cpp + src/exportable_embedding/indexer_snapshot_pybind.cpp +) +if(DYNAMICEMB_HAS_NVE_UPDATE) + list(APPEND EXPORTABLE_EMBEDDING_PYBIND_SOURCES + src/exportable_embedding/update_subscriber_pybind.cu + ) +else() + list(APPEND EXPORTABLE_EMBEDDING_PYBIND_SOURCES + src/exportable_embedding/update_subscriber_unavailable_pybind.cpp + ) +endif() + +add_library(dynamicemb_exportable_embedding_python MODULE + ${EXPORTABLE_EMBEDDING_PYBIND_SOURCES} +) +target_include_directories(dynamicemb_exportable_embedding_python PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${Python3_INCLUDE_DIRS} + ${TORCH_INCLUDE_DIRS} + ${CUDAToolkit_INCLUDE_DIRS} +) +target_compile_definitions(dynamicemb_exportable_embedding_python PRIVATE + TORCH_EXTENSION_NAME=_C +) +if(TORCH_CXX_FLAGS_LIST) + target_compile_options(dynamicemb_exportable_embedding_python PRIVATE + ${TORCH_CXX_FLAGS_LIST} + ) +endif() +target_link_libraries(dynamicemb_exportable_embedding_python PRIVATE + inference_emb_ops + Python3::Python + ${TORCH_LIBRARIES} + "${TORCH_PYTHON_LIBRARY}" +) + +if(DYNAMICEMB_HAS_NVE_UPDATE) + target_include_directories(dynamicemb_exportable_embedding_python PRIVATE + "${NVE_ROOT}" + "${NVE_ROOT}/include" + "${NVE_ROOT}/third_party/json/include" + "${NVTX_INCLUDE_DIR}" + ) + dynamicemb_configure_nve_target(dynamicemb_exportable_embedding_python) + target_link_libraries(dynamicemb_exportable_embedding_python PRIVATE + inference_emb_update + CUDA::cudart + "${NVE_TORCH_OPS_LIB}" + "${NVE_COMMON_LIB}" + ) +endif() + +set_target_properties(dynamicemb_exportable_embedding_python PROPERTIES + PREFIX "" + OUTPUT_NAME "_C" + LIBRARY_OUTPUT_DIRECTORY + "${CMAKE_CURRENT_SOURCE_DIR}/dynamicemb/exportable_embedding" + INSTALL_RPATH "${CMAKE_INSTALL_FULL_LIBDIR};${NVE_LIB_DIR}" +) +install(TARGETS dynamicemb_exportable_embedding_python + LIBRARY DESTINATION + "${Python3_SITEARCH}/dynamicemb/exportable_embedding" + RUNTIME DESTINATION + "${Python3_SITEARCH}/dynamicemb/exportable_embedding" +) + +install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) diff --git a/corelib/dynamicemb/README.md b/corelib/dynamicemb/README.md index 058989a6d..ecb830a60 100644 --- a/corelib/dynamicemb/README.md +++ b/corelib/dynamicemb/README.md @@ -101,7 +101,9 @@ Regarding how to use the DynamicEmb APIs and their parameters, please refer to t 4. The lookup process for each dynamic embedding table incurs additional overhead from unique or radix sort operations. Therefore, if you request a large number of small dynamic embedding tables for lookup, the performance will be poor. Since the lookup range of dynamic embedding tables is particularly large (using the entire range of `int64_t`), it is recommended to create one large embedding table and perform a fused lookup for multiple features. 5. Although dynamic embedding tables can be trained together with TorchREC tables, they cannot be fused together for embedding lookup. Therefore, it is recommended to select dynamic embedding tables for all model-parallel tables during training. 6. DynamicEmb supports training with TorchREC's `EmbeddingBagCollection` (pooling mode: SUM/MEAN) and `EmbeddingCollection` (sequence mode). Both modes use fused CUDA kernels for embedding lookup and gradient reduction. Tables with different embedding dimensions are supported in pooling mode. -7. DynamicEmb supports Torch-exportable embedding tables through `InferenceEmbeddingTable`. It uses DynamicEmb `ScoredHashTable` metadata frozen at export/inference time and `LinearUVMEmbedding` from [NVEmbedding](https://github.com/NVIDIA/nv-embedding-cache), supporting sequence mode and pooling mode (`SUM`, `MEAN`). It is initialized from `DynamicEmbTableOptions` and loads from DynamicEmb dumped embedding files. +7. DynamicEmb supports exportable inference embedding collections with configurable indexing and NVE GPU, LinearUVM, or hierarchical storage. Data-backed indexer state is exported as a sidecar, while Redis-backed hierarchical collections use an ephemeral NVHashMap host cache and support incremental load. + + See the [verified example](./example/exportable_embedding/README.md). ### DynamicEmb Insertion Behavior Checking Modes diff --git a/corelib/dynamicemb/cmake/NVETarget.cmake b/corelib/dynamicemb/cmake/NVETarget.cmake new file mode 100644 index 000000000..f9d3c5f13 --- /dev/null +++ b/corelib/dynamicemb/cmake/NVETarget.cmake @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +function(dynamicemb_configure_nve_target target) + if(NOT NVE_CACHE_LINE_SIZE) + execute_process( + COMMAND getconf LEVEL1_DCACHE_LINESIZE + OUTPUT_VARIABLE NVE_CACHE_LINE_SIZE + OUTPUT_STRIP_TRAILING_WHITESPACE + COMMAND_ERROR_IS_FATAL ANY + ) + endif() + + target_compile_definitions(${target} PRIVATE + NVE_FEATURE_HT_PART_FNV1A=1 + NVE_FEATURE_HT_PART_MURMUR=1 + NVE_FEATURE_HT_PART_RRXMRRXMSX0=1 + NVE_FEATURE_HT_PART_STD_HASH=1 + NVE_FEATURE_HT_MASK_64=1 + NVE_FEATURE_HT_MASK_32=1 + NVE_FEATURE_HT_MASK_16=1 + NVE_FEATURE_HT_MASK_8=1 + NVE_FEATURE_HT_KEY_64=1 + NVE_FEATURE_HT_KEY_32=1 + NVE_FEATURE_HT_KEY_16=1 + NVE_FEATURE_HT_KEY_8=1 + NVE_CACHE_LINE_SIZE=${NVE_CACHE_LINE_SIZE} + ) +endfunction() diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py index 57f5d3ce8..573ffc24e 100644 --- a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py +++ b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py @@ -927,7 +927,7 @@ def _create_optimizer( ) self._optimizer_args = optimizer_args - if optimizer_type == EmbOptimType.SGD: + if optimizer_type in (EmbOptimType.NONE, EmbOptimType.SGD): optimizer = SGDDynamicEmbeddingOptimizer( optimizer_args, ) diff --git a/corelib/dynamicemb/dynamicemb/exportable_embedding/__init__.py b/corelib/dynamicemb/dynamicemb/exportable_embedding/__init__.py new file mode 100644 index 000000000..2ea3e1a6c --- /dev/null +++ b/corelib/dynamicemb/dynamicemb/exportable_embedding/__init__.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from . import _C +from .config import ( + BitConcatConfig, + EmbeddingCollectionIndexerType, + InferenceEmbeddingCollectionConfig, +) +from .incremental_update import ( + EmbeddingCollectionUpdate, + EmbeddingCollectionUpdateAck, + EmbeddingCollectionUpdateCoordinator, +) +from .indexer_directory import ( + dump_embedding_collection_indexers, + embedding_collection_indexers_from_model, + load_embedding_collection_indexers, +) +from .indexer import ( + BitConcatIndexer, + EmbeddingCollectionIndexerBase, + FusedIdentityIndexer, + IdentityIndexer, + LinearHashMapIndexer, +) +from .indexer_snapshot import dump_embedding_collection_indexer_snapshot +from .nve_runtime import ( + export_embedding_collection_aot, + imported_nve_generation, + load_embedding_collection_aot, + register_nve_export_compat, +) + +EmbeddingCollectionBinding = _C.EmbeddingCollectionBinding +EmbeddingCollectionIndexerDirectory = _C.EmbeddingCollectionIndexerDirectory +EmbeddingCollectionUpdateSubscriber = _C.EmbeddingCollectionUpdateSubscriber + + +__all__ = [ + "BitConcatConfig", + "BitConcatIndexer", + "EmbeddingCollectionBinding", + "EmbeddingCollectionIndexerBase", + "EmbeddingCollectionIndexerDirectory", + "EmbeddingCollectionIndexerType", + "EmbeddingCollectionUpdate", + "EmbeddingCollectionUpdateAck", + "EmbeddingCollectionUpdateCoordinator", + "EmbeddingCollectionUpdateSubscriber", + "FusedIdentityIndexer", + "IdentityIndexer", + "InferenceEmbeddingCollectionConfig", + "LinearHashMapIndexer", + "dump_embedding_collection_indexer_snapshot", + "dump_embedding_collection_indexers", + "embedding_collection_indexers_from_model", + "export_embedding_collection_aot", + "imported_nve_generation", + "load_embedding_collection_aot", + "load_embedding_collection_indexers", + "register_nve_export_compat", +] diff --git a/corelib/dynamicemb/dynamicemb/exportable_embedding/config.py b/corelib/dynamicemb/dynamicemb/exportable_embedding/config.py new file mode 100644 index 000000000..441875794 --- /dev/null +++ b/corelib/dynamicemb/dynamicemb/exportable_embedding/config.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Optional + + +class EmbeddingCollectionIndexerType(str, Enum): + LINEAR_HASH_MAP = "linear_hash_map" + FUSED_IDENTITY = "fused_identity" + BIT_CONCAT = "bit_concat" + IDENTITY = "identity" + + +@dataclass(frozen=True) +class BitConcatConfig: + table_id_bits: int + feature_id_bits: int + + +@dataclass(frozen=True) +class InferenceEmbeddingCollectionConfig: + indexer_type: EmbeddingCollectionIndexerType + nve_layer_type: str + indexer_state_sidecar: bool = True + bucket_capacity: int = 128 + bit_concat: Optional[BitConcatConfig] = None + gpu_cache_size: Optional[int] = None + host_cache_size: int = 0 + parameter_server: Optional[Any] = None + + +def validate_collection_config(config: InferenceEmbeddingCollectionConfig) -> None: + if config.nve_layer_type not in {"gpu", "linear_uvm", "hierarchical"}: + raise ValueError(f"Unsupported NVE layer type: {config.nve_layer_type}") + if config.nve_layer_type == "hierarchical" and config.parameter_server is None: + raise ValueError("Hierarchical NVE requires parameter_server") + if config.nve_layer_type != "hierarchical" and config.parameter_server is not None: + raise ValueError("parameter_server is only used by hierarchical NVE") + if ( + config.indexer_type is EmbeddingCollectionIndexerType.BIT_CONCAT + and config.nve_layer_type != "hierarchical" + ): + raise ValueError("BitConcatIndexer is supported only with hierarchical NVE") + if config.nve_layer_type != "gpu" and config.gpu_cache_size is None: + raise ValueError("LinearUVM and hierarchical NVE require gpu_cache_size") + if config.bucket_capacity <= 0: + raise ValueError("bucket_capacity must be positive") diff --git a/corelib/dynamicemb/dynamicemb/exportable_embedding/incremental_update.py b/corelib/dynamicemb/dynamicemb/exportable_embedding/incremental_update.py new file mode 100644 index 000000000..10098c079 --- /dev/null +++ b/corelib/dynamicemb/dynamicemb/exportable_embedding/incremental_update.py @@ -0,0 +1,357 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, Optional + +import torch +from dynamicemb.incremental_dump import DeltaDumpResult + +from .config import EmbeddingCollectionIndexerType +from .indexer_directory import ( + load_embedding_collection_indexer_state, +) +from .indexer_snapshot import ( + IndexerSnapshot, + dump_embedding_collection_indexer_snapshot, + linear_hash_entries, + rebuild_linear_hash_snapshot, +) +from .nve_runtime import imported_nve_generation + + +@dataclass(frozen=True) +class EmbeddingCollectionUpdate: + collection_id: str + snapshot_id: int + cache_update_keys: torch.Tensor + cache_update_values: torch.Tensor + indexer_snapshot_path: Optional[str] = None + + def to_json(self) -> str: + return json.dumps( + { + "collection_id": self.collection_id, + "snapshot_id": self.snapshot_id, + "cache_update_keys": self.cache_update_keys.tolist(), + "cache_update_values": self.cache_update_values.to( + dtype=torch.float32 + ).flatten().tolist(), + "indexer_snapshot_path": self.indexer_snapshot_path, + } + ) + + @classmethod + def from_json(cls, payload: str) -> "EmbeddingCollectionUpdate": + value = json.loads(payload) + return cls( + collection_id=value["collection_id"], + snapshot_id=value["snapshot_id"], + cache_update_keys=torch.tensor( + value["cache_update_keys"], dtype=torch.int64 + ), + cache_update_values=torch.tensor( + value["cache_update_values"], dtype=torch.float32 + ), + indexer_snapshot_path=value.get("indexer_snapshot_path"), + ) + + +@dataclass(frozen=True) +class EmbeddingCollectionUpdateAck: + collection_id: str + snapshot_id: int + + def to_json(self) -> str: + return json.dumps( + { + "collection_id": self.collection_id, + "snapshot_id": self.snapshot_id, + } + ) + + @classmethod + def from_json(cls, payload: str) -> "EmbeddingCollectionUpdateAck": + return cls(**json.loads(payload)) + + +def _torch_dtype(name: str) -> torch.dtype: + return torch.float16 if "float16" in name.lower() else torch.float32 + + +def _make_parameter_server(config: dict[str, Any]) -> Any: + from pynve.torch.nve_ps import NVEParameterServer + + return NVEParameterServer( + num_embeddings=config.get("num_rows", 0), + embedding_size=config["row_elements"], + data_type=_torch_dtype(config["data_type"]), + plugin_name=config["plugin_name"], + factory_config=config.get("factory_config", {}), + table_config=config.get("table_config", {}), + ) + + +def _open_parameter_servers( + package_dir: Path, + bindings: Mapping[str, Any], +) -> dict[str, Any]: + metadata = json.loads((package_dir / "metadata.json").read_text()) + by_path: dict[str, dict[str, Any]] = {} + if isinstance(metadata, list): + for layer in metadata: + if "remote_ps_config" in layer: + by_path[layer["module_path"]] = layer["remote_ps_config"] + else: + resources = metadata.get("resources", {}).get("remote_ps", {}) + for layer in metadata["layers"]: + storage_ref = layer.get("storage_ref") + if storage_ref in resources: + by_path[layer["module_path"]] = resources[storage_ref] + + result = {} + for collection_id, binding in bindings.items(): + config = by_path.get(binding.nve_layer_module_path) + if config is not None and "redis" in config["plugin_name"].lower(): + result[collection_id] = _make_parameter_server(config) + return result + + +class EmbeddingCollectionUpdateCoordinator: + """Serialize DynamicEmb deltas into Redis and indexer publications.""" + + def __init__( + self, + *, + bindings: Mapping[str, Any], + snapshots: Mapping[str, IndexerSnapshot], + device: torch.device, + parameter_servers: Mapping[str, Any], + shared_update_dir: Path, + subscriber_ids: Iterable[str], + ) -> None: + self.bindings = dict(bindings) + self.snapshots = dict(snapshots) + self.device = device + self.parameter_servers = dict(parameter_servers) + self.shared_update_dir = shared_update_dir + self.subscriber_ids = set(subscriber_ids) + self._snapshot_ids = { + collection_id: 0 for collection_id in self.bindings + } + self._pending_deletes: dict[ + tuple[str, int], tuple[Any, torch.Tensor, set[str]] + ] = {} + + @classmethod + def open( + cls, + *, + package_dir: str | Path, + shared_update_dir: str | Path, + device: torch.device, + subscriber_ids: Iterable[str] = (), + ) -> "EmbeddingCollectionUpdateCoordinator": + if imported_nve_generation() < (26, 6): + raise RuntimeError("Redis incremental load requires NVE 26.06 or later") + package_dir = Path(package_dir).resolve() + bindings, snapshots = load_embedding_collection_indexer_state( + package_dir, device + ) + update_dir = Path(shared_update_dir).resolve() + update_dir.mkdir(parents=True, exist_ok=True) + return cls( + bindings=bindings, + snapshots=snapshots, + device=device, + parameter_servers=_open_parameter_servers(package_dir, bindings), + shared_update_dir=update_dir, + subscriber_ids=subscriber_ids, + ) + + def _next_snapshot_id(self, collection_id: str) -> int: + self._snapshot_ids[collection_id] += 1 + return self._snapshot_ids[collection_id] + + def _linear_update( + self, + collection_id: str, + delta: DeltaDumpResult, + parameter_server: Any, + snapshot_id: int, + ) -> tuple[list[int], list[torch.Tensor], Optional[str], list[int]]: + old_snapshot = self.snapshots[collection_id] + entries = linear_hash_entries(old_snapshot) + binding = self.bindings[collection_id] + table_ids = {name: index for index, name in enumerate(binding.table_names)} + cache_update_keys: list[int] = [] + cache_update_values: list[torch.Tensor] = [] + retired: list[int] = [] + changed_mapping = False + next_fused_key = old_snapshot.next_fused_key + + for column, table_name in enumerate(delta.table_names): + table_id = table_ids[table_name] + keys = delta.keys[column].to(dtype=torch.int64, device="cpu") + values = delta.values[column].to(device="cpu").contiguous() + storage_keys = [] + for feature_id in keys.tolist(): + storage_key = entries[table_id].get(feature_id) + if storage_key is None: + storage_key = next_fused_key + next_fused_key += 1 + entries[table_id][feature_id] = storage_key + changed_mapping = True + storage_keys.append(storage_key) + storage_tensor = torch.tensor(storage_keys, dtype=torch.int64) + if storage_tensor.numel() > 0: + parameter_server.insert(storage_tensor, values) + cache_update_keys.extend(storage_keys) + cache_update_values.append(values.to(dtype=torch.float32)) + + evicted = delta.evicted_keys[column] + if evicted is not None: + for feature_id in evicted.to(device="cpu").tolist(): + storage_key = entries[table_id].pop(feature_id, None) + if storage_key is not None: + retired.append(storage_key) + changed_mapping = True + + snapshot_path = None + if changed_mapping: + snapshot = rebuild_linear_hash_snapshot( + entries, + bucket_capacity=old_snapshot.bucket_capacity, + miss_storage_indices=old_snapshot.miss_storage_indices, + next_fused_key=next_fused_key, + device=self.device, + ) + snapshot_path = dump_embedding_collection_indexer_snapshot( + collection_id=collection_id.replace(".", "_"), + snapshot_id=snapshot_id, + snapshot=snapshot, + output_dir=self.shared_update_dir, + ) + # Advance the coordinator's shadow state so the next delta resolves + # against snapshot N+1, not snapshot 0. + self.snapshots[collection_id] = snapshot + return ( + cache_update_keys, + cache_update_values, + snapshot_path, + retired, + ) + + def _direct_update( + self, + collection_id: str, + delta: DeltaDumpResult, + parameter_server: Any, + ) -> tuple[list[int], list[torch.Tensor]]: + binding = self.bindings[collection_id] + table_ids = {name: index for index, name in enumerate(binding.table_names)} + snapshot = self.snapshots.get(collection_id) + cache_update_keys: list[int] = [] + cache_update_values: list[torch.Tensor] = [] + for column, table_name in enumerate(delta.table_names): + table_id = table_ids[table_name] + feature_ids = delta.keys[column].to(dtype=torch.int64, device="cpu") + if binding.indexer_type == EmbeddingCollectionIndexerType.BIT_CONCAT.value: + feature_bits = binding.feature_id_bits + assert feature_bits is not None + storage_keys = (table_id << feature_bits) | feature_ids + elif ( + binding.indexer_type + == EmbeddingCollectionIndexerType.FUSED_IDENTITY.value + ): + assert snapshot is not None + base = int(snapshot.valid_bases[table_id].item()) + reserved = int(snapshot.reserved_sizes[table_id].item()) + if feature_ids.numel() and int(feature_ids.max().item()) >= reserved: + raise ValueError( + f"feature ID exceeds the reserved range for {table_name}" + ) + storage_keys = base + feature_ids + else: + storage_keys = feature_ids + if storage_keys.numel() > 0: + values = delta.values[column].to(device="cpu").contiguous() + parameter_server.insert( + storage_keys.contiguous(), + values, + ) + cache_update_keys.extend(storage_keys.tolist()) + cache_update_values.append(values.to(dtype=torch.float32)) + + evicted = delta.evicted_keys[column] + if evicted is None: + continue + evicted = evicted.to(dtype=torch.int64, device="cpu") + if binding.indexer_type == EmbeddingCollectionIndexerType.BIT_CONCAT.value: + feature_bits = binding.feature_id_bits + assert feature_bits is not None + evicted_storage = (table_id << feature_bits) | evicted + elif ( + binding.indexer_type + == EmbeddingCollectionIndexerType.FUSED_IDENTITY.value + ): + assert snapshot is not None + base = int(snapshot.valid_bases[table_id].item()) + evicted_storage = base + evicted + else: + evicted_storage = evicted + parameter_server.erase(evicted_storage.contiguous()) + return cache_update_keys, cache_update_values + + def apply_delta( + self, collection_id: str, delta: DeltaDumpResult + ) -> EmbeddingCollectionUpdate: + snapshot_id = self._next_snapshot_id(collection_id) + parameter_server = self.parameter_servers[collection_id] + indexer_type = self.bindings[collection_id].indexer_type + if indexer_type == EmbeddingCollectionIndexerType.LINEAR_HASH_MAP.value: + cache_update_keys, cache_update_values, snapshot_path, retired = ( + self._linear_update( + collection_id, delta, parameter_server, snapshot_id + ) + ) + if retired: + self._pending_deletes[(collection_id, snapshot_id)] = ( + parameter_server, + torch.tensor(retired, dtype=torch.int64), + set(self.subscriber_ids), + ) + else: + cache_update_keys, cache_update_values = self._direct_update( + collection_id, delta, parameter_server + ) + snapshot_path = None + if cache_update_values: + update_values = torch.cat(cache_update_values, dim=0).flatten() + else: + update_values = torch.empty(0, dtype=torch.float32) + return EmbeddingCollectionUpdate( + collection_id=collection_id, + snapshot_id=snapshot_id, + cache_update_keys=torch.tensor(cache_update_keys, dtype=torch.int64), + cache_update_values=update_values, + indexer_snapshot_path=snapshot_path, + ) + + def acknowledge( + self, subscriber_id: str, ack: EmbeddingCollectionUpdateAck + ) -> None: + pending = self._pending_deletes.get( + (ack.collection_id, ack.snapshot_id) + ) + if pending is None: + return + parameter_server, keys, waiting = pending + waiting.discard(subscriber_id) + if not waiting: + parameter_server.erase(keys) + del self._pending_deletes[(ack.collection_id, ack.snapshot_id)] diff --git a/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer.py b/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer.py new file mode 100644 index 000000000..7bfd45f59 --- /dev/null +++ b/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer.py @@ -0,0 +1,429 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import uuid +from abc import ABC, abstractmethod +from typing import Optional, Sequence + +import torch +from dynamicemb.scored_hashtable import ( + LinearBucketTable, + ScoreArg, + ScorePolicy, + ScoreSpec, +) + +from .config import ( + EmbeddingCollectionIndexerType, + InferenceEmbeddingCollectionConfig, +) +from .indexer_snapshot import IndexerSnapshot, IndexerSnapshotKind + + +def _new_marker_id() -> int: + return uuid.uuid4().int & ((1 << 63) - 1) + + +def _empty_i64(device: torch.device) -> torch.Tensor: + return torch.empty(0, dtype=torch.int64, device=device) + + +def _publish_native(marker: torch.Tensor, snapshot: IndexerSnapshot) -> None: + torch.ops.INFERENCE_EMB.register_embedding_collection_indexer( + marker, + int(snapshot.kind), + snapshot.table_storage, + snapshot.table_bucket_offsets, + snapshot.bucket_capacity, + snapshot.miss_storage_indices, + snapshot.valid_bases, + snapshot.reserved_sizes, + snapshot.next_fused_key, + ) + + +class EmbeddingCollectionIndexerBase(torch.nn.Module, ABC): + def __init__(self) -> None: + super().__init__() + self._failed_build_rows: set[tuple[int, int]] = set() + + @property + def failed_build_rows(self) -> frozenset[tuple[int, int]]: + return frozenset(self._failed_build_rows) + + @property + @abstractmethod + def nve_num_embeddings(self) -> int: + pass + + @property + def snapshot(self) -> Optional[IndexerSnapshot]: + return None + + @torch.no_grad() + @abstractmethod + def build_index( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + pass + + def finish_build(self) -> None: + pass + + +class LinearHashMapIndexer(EmbeddingCollectionIndexerBase): + """Map arbitrary per-table feature IDs into one stable linear key space.""" + + def __init__( + self, + table_capacities: Sequence[int], + *, + bucket_capacity: int = 128, + key_type: torch.dtype = torch.int64, + state_sidecar: bool = True, + device: torch.device, + ) -> None: + super().__init__() + self.table_capacities = tuple(int(value) for value in table_capacities) + self.key_type = key_type + self.state_sidecar = state_sidecar + self._nve_num_embeddings = len(self.table_capacities) + sum( + self.table_capacities + ) + self.marker_value_ = _new_marker_id() + self.register_buffer( + "marker_tensor", + torch.tensor([self.marker_value_], dtype=torch.int64, device=device), + persistent=False, + ) + self._build_table = LinearBucketTable( + [ + max(bucket_capacity, 2 * capacity) + for capacity in self.table_capacities + ], + [ScoreSpec(name="fused_key", policy=ScorePolicy.ASSIGN)], + key_type=key_type, + bucket_capacity=bucket_capacity, + device=device, + ) + self.bucket_capacity = self._build_table.bucket_capacity_ + empty = _empty_i64(device) + self._snapshot = IndexerSnapshot( + kind=IndexerSnapshotKind.LINEAR_HASH_MAP, + table_storage=self._build_table.table_storage_, + table_bucket_offsets=self._build_table.table_bucket_offsets_, + bucket_capacity=self.bucket_capacity, + miss_storage_indices=torch.arange( + len(self.table_capacities), dtype=torch.int64, device=device + ), + valid_bases=empty, + reserved_sizes=empty, + next_fused_key=len(self.table_capacities), + ) + if not self.state_sidecar: + self.register_buffer( + "table_storage_", self._snapshot.table_storage + ) + self.register_buffer( + "table_bucket_offsets_", self._snapshot.table_bucket_offsets + ) + self.register_buffer( + "miss_storage_indices_", self._snapshot.miss_storage_indices + ) + self._failed_candidates: set[tuple[int, int]] = set() + if self.state_sidecar: + _publish_native(self.marker_tensor, self._snapshot) + + @property + def nve_num_embeddings(self) -> int: + return self._nve_num_embeddings + + @property + def snapshot(self) -> Optional[IndexerSnapshot]: + return self._snapshot if self.state_sidecar else None + + @property + def miss_storage_indices(self) -> torch.Tensor: + return self._snapshot.miss_storage_indices + + @torch.no_grad() + def reset_for_build(self) -> None: + self._build_table.reset() + self._snapshot.next_fused_key = len(self.table_capacities) + self._failed_candidates.clear() + self._failed_build_rows.clear() + + @torch.no_grad() + def build_index( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + stored_keys, found, _slots = self._build_table.lookup( + feature_ids, + table_ids, + ScoreArg(name="fused_key", policy=ScorePolicy.CONST), + ) + missing_ids = feature_ids[~found] + missing_tables = table_ids[~found] + if missing_ids.numel() > 0: + start = self._snapshot.next_fused_key + assigned = torch.arange( + start, + start + missing_ids.numel(), + dtype=torch.int64, + device=feature_ids.device, + ) + self._snapshot.next_fused_key += missing_ids.numel() + slots, evicted_count, evicted_keys, evicted_tables = ( + self._build_table.insert( + missing_ids, + missing_tables, + ScoreArg( + name="fused_key", + value=assigned, + policy=ScorePolicy.ASSIGN, + ), + collect_evicted=True, + ) + ) + evicted_count_value = int(evicted_count.item()) + direct_failures = slots < 0 + if direct_failures.any() or evicted_count_value: + print( + "[WARNING] LinearHashMapIndexer insertion could not retain " + f"{int(direct_failures.sum().item()) + evicted_count_value} rows" + ) + self._failed_candidates.update( + zip( + missing_tables[direct_failures].cpu().tolist(), + missing_ids[direct_failures].cpu().tolist(), + ) + ) + self._failed_candidates.update( + zip( + evicted_tables[:evicted_count_value].cpu().tolist(), + evicted_keys[:evicted_count_value].cpu().tolist(), + ) + ) + stored_keys, found, _slots = self._build_table.lookup( + feature_ids, + table_ids, + ScoreArg(name="fused_key", policy=ScorePolicy.CONST), + ) + + return torch.where(found, stored_keys, torch.full_like(stored_keys, -1)) + + @torch.no_grad() + def finish_build(self) -> None: + candidates = sorted(self._failed_candidates) + if candidates: + table_ids = torch.tensor( + [table_id for table_id, _ in candidates], + dtype=torch.int64, + device=self.marker_tensor.device, + ) + feature_ids = torch.tensor( + [feature_id for _, feature_id in candidates], + dtype=self.key_type, + device=self.marker_tensor.device, + ) + _stored, found, _slots = self._build_table.lookup( + feature_ids, + table_ids, + ScoreArg(name="fused_key", policy=ScorePolicy.CONST), + ) + self._failed_build_rows = { + row + for row, retained in zip(candidates, found.cpu().tolist()) + if not retained + } + self._failed_candidates.clear() + if self.state_sidecar: + _publish_native(self.marker_tensor, self._snapshot) + self._build_table = None + + def forward( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + if not self.state_sidecar: + fused_keys, found, _slots = torch.ops.INFERENCE_EMB.table_lookup( + self.table_storage_, + self.table_bucket_offsets_, + self.bucket_capacity, + feature_ids, + table_ids, + None, + int(ScorePolicy.CONST), + None, + 0, + None, + ) + misses = torch.index_select( + self.miss_storage_indices_, 0, table_ids + ) + return torch.where(found, fused_keys, misses) + return torch.ops.INFERENCE_EMB.embedding_collection_index( + self.marker_tensor, self.marker_value_, feature_ids, table_ids + ) + + +class FusedIdentityIndexer(EmbeddingCollectionIndexerBase): + """Keep table-local IDs and place each table in a reserved fused section.""" + + def __init__( + self, + table_capacities: Sequence[int], + *, + state_sidecar: bool = True, + device: torch.device, + ) -> None: + super().__init__() + self.state_sidecar = state_sidecar + reserved_sizes = [int(value) for value in table_capacities] + valid_bases = [] + next_row = 0 + for size in reserved_sizes: + valid_bases.append(next_row + 1) + next_row += size + 1 + self._nve_num_embeddings = next_row + self.marker_value_ = _new_marker_id() + self.register_buffer( + "marker_tensor", + torch.tensor( + [self.marker_value_], dtype=torch.int64, device=device + ), + persistent=False, + ) + empty = _empty_i64(device) + self._snapshot = IndexerSnapshot( + kind=IndexerSnapshotKind.FUSED_IDENTITY, + table_storage=torch.empty(0, dtype=torch.uint8, device=device), + table_bucket_offsets=empty, + bucket_capacity=0, + miss_storage_indices=empty, + valid_bases=torch.tensor( + valid_bases, dtype=torch.int64, device=device + ), + reserved_sizes=torch.tensor( + reserved_sizes, dtype=torch.int64, device=device + ), + next_fused_key=next_row, + ) + if self.state_sidecar: + _publish_native(self.marker_tensor, self._snapshot) + else: + self.register_buffer("valid_bases_", self._snapshot.valid_bases) + + @property + def nve_num_embeddings(self) -> int: + return self._nve_num_embeddings + + @property + def snapshot(self) -> Optional[IndexerSnapshot]: + return self._snapshot if self.state_sidecar else None + + @property + def miss_storage_indices(self) -> torch.Tensor: + return self._snapshot.valid_bases - 1 + + @torch.no_grad() + def build_index( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + return ( + torch.index_select(self._snapshot.valid_bases, 0, table_ids) + + feature_ids + ) + + def forward( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + if not self.state_sidecar: + return ( + torch.index_select(self.valid_bases_, 0, table_ids) + + feature_ids + ) + return torch.ops.INFERENCE_EMB.embedding_collection_index( + self.marker_tensor, self.marker_value_, feature_ids, table_ids + ) + + +class BitConcatIndexer(EmbeddingCollectionIndexerBase): + def __init__( + self, + *, + feature_id_bits: int, + nve_num_embeddings: int, + ) -> None: + super().__init__() + self.feature_id_bits = feature_id_bits + self._nve_num_embeddings = nve_num_embeddings + + @property + def nve_num_embeddings(self) -> int: + return self._nve_num_embeddings + + @torch.no_grad() + def build_index( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + return self.forward(feature_ids, table_ids) + + def forward( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + return (table_ids << self.feature_id_bits) | feature_ids + + +class IdentityIndexer(EmbeddingCollectionIndexerBase): + def __init__(self, *, nve_num_embeddings: int) -> None: + super().__init__() + self._nve_num_embeddings = nve_num_embeddings + + @property + def nve_num_embeddings(self) -> int: + return self._nve_num_embeddings + + @torch.no_grad() + def build_index( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + return feature_ids + + def forward( + self, feature_ids: torch.Tensor, table_ids: torch.Tensor + ) -> torch.Tensor: + return feature_ids + + +def create_embedding_collection_indexer( + config: InferenceEmbeddingCollectionConfig, + table_capacities: Sequence[int], + *, + key_type: torch.dtype, + device: torch.device, +) -> EmbeddingCollectionIndexerBase: + capacities = tuple(int(value) for value in table_capacities) + if config.indexer_type is EmbeddingCollectionIndexerType.LINEAR_HASH_MAP: + return LinearHashMapIndexer( + capacities, + bucket_capacity=config.bucket_capacity, + key_type=key_type, + state_sidecar=config.indexer_state_sidecar, + device=device, + ) + if config.indexer_type is EmbeddingCollectionIndexerType.FUSED_IDENTITY: + return FusedIdentityIndexer( + capacities, + state_sidecar=config.indexer_state_sidecar, + device=device, + ) + if config.indexer_type is EmbeddingCollectionIndexerType.BIT_CONCAT: + if config.bit_concat is None: + raise ValueError("BitConcatIndexer requires bit_concat") + return BitConcatIndexer( + feature_id_bits=config.bit_concat.feature_id_bits, + nve_num_embeddings=sum(capacities), + ) + if config.indexer_type is EmbeddingCollectionIndexerType.IDENTITY: + return IdentityIndexer(nve_num_embeddings=sum(capacities)) + raise ValueError(f"Unsupported indexer: {config.indexer_type}") diff --git a/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer_directory.py b/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer_directory.py new file mode 100644 index 000000000..10dadb0fe --- /dev/null +++ b/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer_directory.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Optional + +import torch + +from . import _C +from .config import EmbeddingCollectionIndexerType +from .indexer import ( + BitConcatIndexer, + EmbeddingCollectionIndexerBase, + FusedIdentityIndexer, + IdentityIndexer, + LinearHashMapIndexer, +) +from .indexer_snapshot import ( + IndexerSnapshot, + dump_indexer_snapshot, + load_embedding_collection_indexer_snapshot, + native_indexer_snapshot, +) + + +_SIDECAR_DIR = "embedding_collection_indexers" +_MANIFEST_NAME = "manifest.json" +_DIRECTORY_SCHEMA_VERSION = 1 + + +def _indexer_type(indexer: EmbeddingCollectionIndexerBase) -> str: + if isinstance(indexer, LinearHashMapIndexer): + return EmbeddingCollectionIndexerType.LINEAR_HASH_MAP.value + if isinstance(indexer, FusedIdentityIndexer): + return EmbeddingCollectionIndexerType.FUSED_IDENTITY.value + if isinstance(indexer, BitConcatIndexer): + return EmbeddingCollectionIndexerType.BIT_CONCAT.value + if isinstance(indexer, IdentityIndexer): + return EmbeddingCollectionIndexerType.IDENTITY.value + raise TypeError(f"Unsupported indexer: {type(indexer).__name__}") + + +def _binding( + *, + table_names: list[str], + indexer_type: str, + indexer_module_path: str, + nve_layer_module_path: str, + snapshot_path: Optional[str] = None, + feature_id_bits: Optional[int] = None, + marker_value: Optional[int] = None, +) -> Any: + result = _C.EmbeddingCollectionBinding() + result.table_names = table_names + result.indexer_type = indexer_type + result.indexer_module_path = indexer_module_path + result.nve_layer_module_path = nve_layer_module_path + result.snapshot_path = snapshot_path + result.feature_id_bits = feature_id_bits + result.marker_value = marker_value + return result + + +def embedding_collection_indexers_from_model( + model: torch.nn.Module, +) -> Any: + bindings = {} + markers = {} + snapshots = {} + device = torch.device("cuda", torch.cuda.current_device()) + collection_number = 0 + for module_path, module in model.named_modules(): + indexer = getattr(module, "indexer_", None) + if not isinstance(indexer, EmbeddingCollectionIndexerBase): + continue + collection_id = getattr(module, "collection_id_", None) + if not collection_id: + collection_id = module_path or f"collection_{collection_number}" + collection_number += 1 + indexer_path = f"{module_path}.indexer_" if module_path else "indexer_" + nve_path = ( + f"{module_path}.nve_embedding_" if module_path else "nve_embedding_" + ) + snapshot = indexer.snapshot + bindings[collection_id] = _binding( + table_names=list(module.table_names_), + indexer_type=_indexer_type(indexer), + indexer_module_path=indexer_path, + nve_layer_module_path=nve_path, + feature_id_bits=getattr(indexer, "feature_id_bits", None), + marker_value=( + int(indexer.marker_tensor.item()) if snapshot is not None else None + ), + ) + if snapshot is not None: + markers[collection_id] = indexer.marker_tensor + snapshots[collection_id] = native_indexer_snapshot(snapshot) + device = indexer.marker_tensor.device + device_index = ( + device.index if device.index is not None else torch.cuda.current_device() + ) + return _C.EmbeddingCollectionIndexerDirectory.create( + bindings, markers, snapshots, device_index + ) + + +def dump_embedding_collection_indexers( + directory: Any, + package_dir: str | Path, +) -> Path: + root = Path(package_dir).resolve() / _SIDECAR_DIR + root.mkdir(parents=True, exist_ok=True) + bindings: dict[str, Any] = {} + snapshot_number = 0 + for collection_id, binding in directory.bindings.items(): + snapshot = directory.snapshot(collection_id) + if snapshot is None: + bindings[collection_id] = binding + continue + relative_path = f"snapshot_{snapshot_number}.bin" + dump_indexer_snapshot(snapshot, root / relative_path) + bindings[collection_id] = _binding( + table_names=binding.table_names, + indexer_type=binding.indexer_type, + indexer_module_path=binding.indexer_module_path, + nve_layer_module_path=binding.nve_layer_module_path, + snapshot_path=relative_path, + feature_id_bits=binding.feature_id_bits, + marker_value=binding.marker_value, + ) + snapshot_number += 1 + + document = { + "schema_version": _DIRECTORY_SCHEMA_VERSION, + "collections": { + collection_id: { + "table_names": binding.table_names, + "indexer_type": binding.indexer_type, + "indexer_module_path": binding.indexer_module_path, + "nve_layer_module_path": binding.nve_layer_module_path, + "snapshot_path": binding.snapshot_path, + "feature_id_bits": binding.feature_id_bits, + "marker_value": binding.marker_value, + } + for collection_id, binding in bindings.items() + }, + } + manifest_path = root / _MANIFEST_NAME + manifest_path.write_text(json.dumps(document, indent=2) + "\n") + return manifest_path + + +def load_embedding_collection_indexer_state( + package_dir: str | Path, device: torch.device +) -> tuple[dict[str, Any], dict[str, IndexerSnapshot]]: + root = Path(package_dir).resolve() / _SIDECAR_DIR + document = json.loads((root / _MANIFEST_NAME).read_text()) + bindings = { + collection_id: _binding(**value) + for collection_id, value in document["collections"].items() + } + snapshots = { + collection_id: load_embedding_collection_indexer_snapshot( + root / binding.snapshot_path, device + ) + for collection_id, binding in bindings.items() + if binding.snapshot_path is not None + } + return bindings, snapshots + + +def load_embedding_collection_indexers( + package_dir: str | Path, device: torch.device +) -> Any: + device_index = ( + device.index if device.index is not None else torch.cuda.current_device() + ) + return _C.EmbeddingCollectionIndexerDirectory.load( + str(Path(package_dir).resolve()), device_index + ) diff --git a/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer_snapshot.py b/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer_snapshot.py new file mode 100644 index 000000000..9d7d95244 --- /dev/null +++ b/corelib/dynamicemb/dynamicemb/exportable_embedding/indexer_snapshot.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import struct +from dataclasses import dataclass +from enum import IntEnum +from pathlib import Path +from typing import Any, Sequence + +import torch +from dynamicemb.scored_hashtable import ( + LinearBucketTable, + ScoreArg, + ScorePolicy, + ScoreSpec, +) +from dynamicemb_extensions import table_partition + +from . import _C + + +_SNAPSHOT_MAGIC = b"ECIDX001" +_SNAPSHOT_HEADER = struct.Struct("<8sIIqqQQQQQ") +_SNAPSHOT_SCHEMA_VERSION = 1 + + +class IndexerSnapshotKind(IntEnum): + # Persisted in the snapshot header; keep aligned with indexer_snapshot.h. + LINEAR_HASH_MAP = 0 + FUSED_IDENTITY = 1 + + +@dataclass +class IndexerSnapshot: + kind: IndexerSnapshotKind + table_storage: torch.Tensor + table_bucket_offsets: torch.Tensor + bucket_capacity: int + miss_storage_indices: torch.Tensor + valid_bases: torch.Tensor + reserved_sizes: torch.Tensor + next_fused_key: int + + def to(self, device: torch.device) -> "IndexerSnapshot": + return IndexerSnapshot( + kind=self.kind, + table_storage=self.table_storage.to(device), + table_bucket_offsets=self.table_bucket_offsets.to(device), + bucket_capacity=self.bucket_capacity, + miss_storage_indices=self.miss_storage_indices.to(device), + valid_bases=self.valid_bases.to(device), + reserved_sizes=self.reserved_sizes.to(device), + next_fused_key=self.next_fused_key, + ) + + +def _write_tensor(stream: Any, tensor: torch.Tensor) -> None: + stream.write(tensor.detach().cpu().contiguous().numpy().tobytes()) + + +def dump_indexer_snapshot(snapshot: IndexerSnapshot, path: str | Path) -> None: + path = Path(path) + tensors = ( + snapshot.table_storage, + snapshot.table_bucket_offsets, + snapshot.miss_storage_indices, + snapshot.valid_bases, + snapshot.reserved_sizes, + ) + lengths = [tensor.numel() for tensor in tensors] + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as stream: + stream.write( + _SNAPSHOT_HEADER.pack( + _SNAPSHOT_MAGIC, + _SNAPSHOT_SCHEMA_VERSION, + int(snapshot.kind), + snapshot.bucket_capacity, + snapshot.next_fused_key, + *lengths, + ) + ) + for tensor in tensors: + _write_tensor(stream, tensor) + + +def load_embedding_collection_indexer_snapshot( + path: str | Path, device: torch.device +) -> Any: + device_index = ( + device.index if device.index is not None else torch.cuda.current_device() + ) + return _C.load_indexer_snapshot(str(Path(path)), device_index) + + +def dump_embedding_collection_indexer_snapshot( + *, + collection_id: str, + snapshot_id: int, + snapshot: IndexerSnapshot, + output_dir: str | Path, +) -> str: + path = Path(output_dir).resolve() / f"{collection_id}.{snapshot_id}.bin" + dump_indexer_snapshot(snapshot, path) + return str(path) + + +def native_indexer_snapshot(snapshot: IndexerSnapshot) -> Any: + return _C.IndexerSnapshot( + int(snapshot.kind), + snapshot.table_storage, + snapshot.table_bucket_offsets, + snapshot.bucket_capacity, + snapshot.miss_storage_indices, + snapshot.valid_bases, + snapshot.reserved_sizes, + snapshot.next_fused_key, + ) + + +def linear_hash_entries( + snapshot: IndexerSnapshot, +) -> list[dict[int, int]]: + storage = snapshot.table_storage.cpu().contiguous() + offsets = snapshot.table_bucket_offsets.cpu().tolist() + num_buckets = offsets[-1] + keys, _digests, values = table_partition( + storage, + [torch.int64, torch.uint8, torch.uint64], + snapshot.bucket_capacity, + num_buckets, + ) + result: list[dict[int, int]] = [] + for table_id in range(len(offsets) - 1): + table_keys = keys[offsets[table_id] : offsets[table_id + 1]].reshape(-1) + table_values = values[ + offsets[table_id] : offsets[table_id + 1] + ].reshape(-1) + present = table_keys != -1 + result.append( + dict( + zip( + table_keys[present].tolist(), + table_values[present].view(torch.int64).tolist(), + ) + ) + ) + return result + + +@torch.no_grad() +def rebuild_linear_hash_snapshot( + entries: Sequence[dict[int, int]], + *, + bucket_capacity: int, + miss_storage_indices: torch.Tensor, + next_fused_key: int, + device: torch.device, +) -> IndexerSnapshot: + logical_capacities = [max(1, len(table)) for table in entries] + while True: + table = LinearBucketTable( + [ + max(bucket_capacity, 2 * capacity) + for capacity in logical_capacities + ], + [ScoreSpec(name="fused_key", policy=ScorePolicy.ASSIGN)], + key_type=torch.int64, + bucket_capacity=bucket_capacity, + device=device, + ) + failed = False + for table_id, mapping in enumerate(entries): + if not mapping: + continue + feature_ids = torch.tensor( + list(mapping), dtype=torch.int64, device=device + ) + table_ids = torch.full_like(feature_ids, table_id) + fused_keys = torch.tensor( + list(mapping.values()), dtype=torch.int64, device=device + ) + slots, evicted_count, _keys, _table_ids = table.insert( + feature_ids, + table_ids, + ScoreArg( + name="fused_key", + value=fused_keys, + policy=ScorePolicy.ASSIGN, + ), + collect_evicted=True, + ) + if (slots < 0).any() or int(evicted_count.item()) != 0: + logical_capacities[table_id] *= 2 + failed = True + break + if not failed: + break + + empty = torch.empty(0, dtype=torch.int64, device=device) + return IndexerSnapshot( + kind=IndexerSnapshotKind.LINEAR_HASH_MAP, + table_storage=table.table_storage_, + table_bucket_offsets=table.table_bucket_offsets_, + bucket_capacity=table.bucket_capacity_, + miss_storage_indices=miss_storage_indices.to(device), + valid_bases=empty, + reserved_sizes=empty, + next_fused_key=next_fused_key, + ) diff --git a/corelib/dynamicemb/dynamicemb/exportable_embedding/nve_runtime.py b/corelib/dynamicemb/dynamicemb/exportable_embedding/nve_runtime.py new file mode 100644 index 000000000..f7eb5011d --- /dev/null +++ b/corelib/dynamicemb/dynamicemb/exportable_embedding/nve_runtime.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Optional + +import torch + +from .config import InferenceEmbeddingCollectionConfig +from .indexer_directory import ( + dump_embedding_collection_indexers, + embedding_collection_indexers_from_model, + load_embedding_collection_indexers, +) + + +def imported_nve_generation() -> tuple[int, int]: + import pynve + + return tuple(int(value) for value in pynve.__version__.split(".")[:2]) + + +def register_nve_export_compat() -> None: + """Install the only export shim needed by the supported 26.05 package.""" + if imported_nve_generation() != (26, 5): + return + + import pynve.torch # noqa: F401 + + @torch.library.register_fake( + "nve_ops::embedding_lookup", allow_override=True + ) + def _embedding_lookup_fake( + keys: torch.Tensor, layer_id: int + ) -> torch.Tensor: + del layer_id + context = torch.library.get_ctx() + return keys.new_empty( + (keys.size(0), context.new_dynamic_size()), dtype=torch.float32 + ) + + +def _nve_layer_constructor_kwargs( + nve_layer_type: str, *, storage: Optional[Any] = None +) -> dict[str, object]: + import pynve.torch.nve_layers as nve_layers + + if imported_nve_generation() == (26, 5): + layer_type = { + "gpu": nve_layers.CacheType.NoCache, + "linear_uvm": nve_layers.CacheType.LinearUVM, + "hierarchical": nve_layers.CacheType.Hierarchical, + }[nve_layer_type] + result: dict[str, object] = {"cache_type": layer_type} + if nve_layer_type == "hierarchical": + result["remote_interface"] = storage + return result + + layer_type = { + "gpu": nve_layers.LayerType.GPULayer, + "linear_uvm": nve_layers.LayerType.LinearUVM, + "hierarchical": nve_layers.LayerType.Hierarchical, + }[nve_layer_type] + result = {"layer_type": layer_type} + if nve_layer_type == "hierarchical": + result["storage"] = storage + return result + + +def create_nve_layer( + *, + num_embeddings: int, + embedding_dim: int, + dtype: torch.dtype, + pooling_mode: int, + config: InferenceEmbeddingCollectionConfig, + device: torch.device, +) -> torch.nn.Module: + import pynve.torch.nve_layers as nve_layers + + if pooling_mode == -1: + layer_class = nve_layers.NVEmbedding + pooling_args: dict[str, Any] = {} + else: + layer_class = nve_layers.NVEmbeddingBag + pooling_args = {"mode": "sum" if pooling_mode == 1 else "mean"} + + kwargs: dict[str, Any] = { + "num_embeddings": num_embeddings, + "embedding_size": embedding_dim, + "data_type": dtype, + "optimize_for_training": False, + "device": device, + **pooling_args, + **_nve_layer_constructor_kwargs( + config.nve_layer_type, storage=config.parameter_server + ), + } + if config.nve_layer_type != "gpu": + kwargs["gpu_cache_size"] = config.gpu_cache_size + if config.nve_layer_type == "hierarchical": + kwargs["host_cache_size"] = config.host_cache_size + return layer_class(**kwargs) + + +def export_embedding_collection_aot( + model: torch.nn.Module, + example_inputs: tuple[Any, ...], + package_dir: str | os.PathLike[str], + *, + dynamic_shapes: Any = None, + inductor_configs: Optional[dict[str, Any]] = None, +) -> Any: + from pynve.torch.nve_export import export_aot + + package_dir = os.fspath(Path(package_dir).resolve()) + indexers = embedding_collection_indexers_from_model(model) + dump_embedding_collection_indexers(indexers, package_dir) + configs = {"aot_inductor.use_runtime_constant_folding": True} + if inductor_configs: + configs.update(inductor_configs) + export_aot( + model, + example_inputs, + package_dir, + dynamic_shapes=dynamic_shapes, + inductor_configs=configs, + ) + return indexers + + +def load_embedding_collection_aot( + package_dir: str | os.PathLike[str], device: torch.device +) -> tuple[Any, list[Any], Any]: + package_dir = os.fspath(Path(package_dir).resolve()) + if imported_nve_generation() == (26, 5): + from pynve.torch.nve_export import load_nve_layers + from torch._C._aoti import AOTIModelPackageLoader + + with torch.cuda.device(device): + layers = load_nve_layers(package_dir) + loader = AOTIModelPackageLoader( + os.path.join(package_dir, "model.pt2"), + "model", + False, + 1, + device.index if device.index is not None else torch.cuda.current_device(), + ) + else: + from pynve.torch.nve_export import load_aot + + loader, layers = load_aot(package_dir, device=device) + + indexers = load_embedding_collection_indexers(package_dir, device) + indexers.bind_aoti(loader) + return loader, layers, indexers diff --git a/corelib/dynamicemb/dynamicemb/exportable_tables.py b/corelib/dynamicemb/dynamicemb/exportable_tables.py index e9f51e13a..c1b327b88 100644 --- a/corelib/dynamicemb/dynamicemb/exportable_tables.py +++ b/corelib/dynamicemb/dynamicemb/exportable_tables.py @@ -1,552 +1,246 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Implementation module for exportable inference embedding demo. - -This module owns: -1. Loading `inference_emb_ops.so` before any `dynamicemb` imports -2. Importing the `dynamicemb` symbols required by the inference demo -3. The `InferenceLinearBucketTable` and `InferenceEmbeddingTable` implementations -""" +"""Exportable inference embedding-collection conversion and population.""" import itertools import os -from typing import Dict, List, Optional +from collections.abc import Mapping, Sequence +from typing import Optional -import pynve.torch.nve_layers as nve_layers import torch -from dynamicemb import ( - DynamicEmbInitializerArgs, - DynamicEmbInitializerMode, - DynamicEmbTableOptions, -) from dynamicemb.batched_dynamicemb_tables import ( encode_meta_json_file_path, get_loading_files, ) +from dynamicemb.exportable_embedding.config import ( + BitConcatConfig, + EmbeddingCollectionIndexerType, + InferenceEmbeddingCollectionConfig, + validate_collection_config, +) +from dynamicemb.exportable_embedding.indexer import ( + BitConcatIndexer, + EmbeddingCollectionIndexerBase, + FusedIdentityIndexer, + IdentityIndexer, + LinearHashMapIndexer, + create_embedding_collection_indexer, +) +from dynamicemb.exportable_embedding.nve_runtime import create_nve_layer from dynamicemb.key_value_table import _iter_batches_from_files, load_from_json -from dynamicemb.scored_hashtable import ScorePolicy -from dynamicemb_extensions import table_insert from torch.nn import ModuleDict +from torch.nn.modules.sparse import Embedding from torchrec.modules.embedding_configs import EmbeddingConfig -_NVE_CACHE_NUM_WAYS = 8 -_NVE_CACHE_ALIGNMENT_BYTES = 16 -_NVE_CACHE_TAG_BYTES = 8 -_NVE_CACHE_COUNTER_BYTES = 4 -_NVE_CACHE_TIMESTAMP_BYTES = 8 - - -# --------------------------------------------------------------------------- -# Helpers for InferenceEmbeddingTable construction -# --------------------------------------------------------------------------- - - -def _align_up(value: int, alignment: int) -> int: - return (value + alignment - 1) // alignment * alignment - - -def _linear_uvm_constructor_kwargs() -> Dict[str, object]: - """Return the LinearUVM selector for the imported NVE generation.""" - if hasattr(nve_layers, "LayerType"): - return {"layer_type": nve_layers.LayerType.LinearUVM} - if hasattr(nve_layers, "CacheType"): - return {"cache_type": nve_layers.CacheType.LinearUVM} - raise RuntimeError( - "Unsupported pynve.torch.nve_layers API: expected LayerType or CacheType" - ) - - -def _resolve_capacity(opt: "DynamicEmbTableOptions") -> int: - """Return the capacity for a single table option. - - Prefers ``init_capacity``; falls back to ``max_capacity``. - Raises ``ValueError`` if neither is set or positive. - """ - cap = opt.init_capacity if opt.init_capacity is not None else opt.max_capacity - if cap is None or cap <= 0: - raise ValueError( - "Each table option must provide init_capacity or max_capacity > 0" - ) - return int(cap) - - -def _resolve_embedding_dim(table_options: List["DynamicEmbTableOptions"]) -> int: - dims = {int(opt.dim) for opt in table_options if opt.dim is not None} - if len(dims) != 1: - raise ValueError( - "InferenceEmbeddingTable requires exactly one shared embedding dim across all table_options" - ) - emb_dim = dims.pop() - if emb_dim <= 0: - raise ValueError("Embedding dim must be > 0") - return emb_dim - - -def _resolve_gpu_cache_size( - table_options: List["DynamicEmbTableOptions"], - total_size_bytes: int, - embedding_dim: int, - dtype_size: int, -) -> int: - values = {int(opt.global_hbm_for_values or 0) for opt in table_options} - if len(values) != 1: - raise ValueError( - "All table_options must have the same global_hbm_for_values for NVE inference table" - ) - gpu_cache_size = values.pop() - if gpu_cache_size <= 0: - # NVE's cache size includes the aligned values, tags, counters, and - # timestamp for each 8-way set. Using only the embedding-table payload - # can be too small to form even one set for small inference tables. - tag_bytes_per_set = _align_up( - _NVE_CACHE_TAG_BYTES * _NVE_CACHE_NUM_WAYS, - _NVE_CACHE_ALIGNMENT_BYTES, - ) - value_bytes_per_set = _align_up( - embedding_dim * dtype_size * _NVE_CACHE_NUM_WAYS, - _NVE_CACHE_ALIGNMENT_BYTES, - ) - counter_bytes_per_set = _align_up( - _NVE_CACHE_COUNTER_BYTES * _NVE_CACHE_NUM_WAYS, - _NVE_CACHE_ALIGNMENT_BYTES, - ) - min_cache_size = ( - tag_bytes_per_set - + value_bytes_per_set - + counter_bytes_per_set - + _NVE_CACHE_TIMESTAMP_BYTES - ) - gpu_cache_size = max(total_size_bytes, min_cache_size) - print( - "[INFO] global_hbm_for_values is 0 for all tables; " - f"using fallback gpu_cache_size={gpu_cache_size}" - ) - return gpu_cache_size - - -def _derive_grouped_offsets(feature_table_map: List[int]) -> List[int]: - """Derive boundary-style offsets from a per-feature table-id list. - For example, ``[0, 0, 1, 2]`` → ``[0, 2, 3, 4]``. - The result is analogous to ``table_bucket_offsets_`` in ``LinearBucketTable``. - """ +def _derive_grouped_offsets(feature_table_map: Sequence[int]) -> list[int]: offsets = [0] - prev = feature_table_map[0] - for i, tid in enumerate(feature_table_map[1:], start=1): - if tid != prev: - offsets.append(i) - prev = tid + previous = feature_table_map[0] + for index, table_id in enumerate(feature_table_map[1:], start=1): + if table_id != previous: + offsets.append(index) + previous = table_id offsets.append(len(feature_table_map)) return offsets -# --------------------------------------------------------------------------- -# Modules for InferenceEmbeddingTable and its hash table -# --------------------------------------------------------------------------- - - -class InferenceLinearBucketTable(torch.nn.Module): - """Simple exportable hash table wrapper for inference lookup using custom op. - - This is a minimal demo version that focuses on lookup-only, non-pooled inference. - For the full production version, see LinearBucketTable in scored_hashtable.py. - """ +class InferenceEmbeddingCollection(torch.nn.Module): + """One configurable indexer feeding one exportable NVE layer.""" def __init__( self, - capacity: List[int], - key_type: torch.dtype = torch.int64, - bucket_capacity: int = 128, - device: Optional[torch.device] = None, - ): - """Initialize demo hash table. - - Args: - capacity: List of per-table capacities - key_type: torch.int64 or torch.uint64 - bucket_capacity: slots per bucket - device: CUDA device (defaults to current) - """ + embedding_configs: Sequence[EmbeddingConfig], + *, + config: InferenceEmbeddingCollectionConfig, + indexer: EmbeddingCollectionIndexerBase, + nve_embedding: Optional[torch.nn.Module], + pooling_mode: int, + table_names: Sequence[str], + feature_names: Sequence[str], + feature_table_map: Sequence[int], + output_dtype: torch.dtype, + device: torch.device, + ) -> None: super().__init__() - - if device is None: - device = torch.device("cuda", torch.cuda.current_device()) - - self.device = device - self.key_type_ = key_type - self.bucket_capacity_ = bucket_capacity - self.num_tables_ = len(capacity) - - per_table_num_buckets = [] - bucket_offset_list = [0] - for cap in capacity: - nb = (cap + bucket_capacity - 1) // bucket_capacity - per_table_num_buckets.append(nb) - bucket_offset_list.append(bucket_offset_list[-1] + nb) - - total_buckets = bucket_offset_list[-1] - self.capacity_ = total_buckets * self.bucket_capacity_ - - bytes_per_slot = 8 + 1 + 8 - total_storage_bytes = bytes_per_slot * bucket_capacity * total_buckets - - self.register_buffer( - "table_storage_", - torch.zeros(total_storage_bytes, dtype=torch.uint8, device=device), - ) - self.register_buffer( - "table_bucket_offsets_", - torch.tensor(bucket_offset_list, dtype=torch.int64, device=device), - ) - self.register_buffer( - "bucket_sizes", - torch.zeros(total_buckets, dtype=torch.int32, device=device), + self.embedding_configs = list(embedding_configs) + self.config_ = config + self.indexer_ = indexer + self.nve_embedding_ = nve_embedding + self.parameter_server_ = config.parameter_server + self.pooling_mode_ = pooling_mode + self.table_names_ = list(table_names) + self.feature_names_ = list(feature_names) + self.num_tables_ = len(self.embedding_configs) + self.table_capacities_ = tuple( + int(table.num_embeddings) for table in self.embedding_configs ) + self.emb_dim_ = int(self.embedding_configs[0].embedding_dim) + self.output_dtype_ = output_dtype + self.device = device + self.collection_id_: Optional[str] = None self.register_buffer( - "_ref_counter", - torch.zeros(self.capacity_, dtype=torch.int32, device=self.device), + "feature_offsets_", + torch.tensor( + _derive_grouped_offsets(feature_table_map), + dtype=torch.int64, + device=device, + ), ) - self.score_policy = int(ScorePolicy.CONST) + @property + def failed_build_rows(self) -> frozenset[tuple[int, int]]: + return self.indexer_.failed_build_rows - def lookup( - self, - keys: torch.Tensor, - table_ids: torch.Tensor, - score_value: Optional[torch.Tensor] = None, - score_policy: int = 0, - ) -> tuple: - """Lookup keys in the hash table using the custom operator.""" - score_out, founds, indices = torch.ops.INFERENCE_EMB.table_lookup( - self.table_storage_, - self.table_bucket_offsets_, - self.bucket_capacity_, - keys, - table_ids, - score_value, - self.score_policy, - None, + def _nve_embedding(self) -> torch.nn.Module: + if self.nve_embedding_ is None: + raise RuntimeError("Load the collection source before lookup") + return self.nve_embedding_ + + @torch.no_grad() + def _write_storage_rows( + self, storage_keys: torch.Tensor, embeddings: torch.Tensor + ) -> None: + if storage_keys.numel() == 0: + return + if self.config_.nve_layer_type == "hierarchical": + keys_cpu = storage_keys.detach().to( + device="cpu", dtype=torch.int64 + ).contiguous() + values_cpu = embeddings.detach().to( + device="cpu", dtype=self.output_dtype_ + ).contiguous() + self.parameter_server_.insert(keys_cpu, values_cpu) + return + + weight = self._nve_embedding().weight.data + weight.index_copy_( 0, - None, + storage_keys.to(device=weight.device, dtype=torch.int64), + embeddings.to(device=weight.device, dtype=weight.dtype), ) - return score_out, founds, indices - - -class InferenceEmbeddingCollection(torch.nn.Module): - """Export-compatible embedding table using custom ops. - - The pooling mode is fixed at construction time so that each exported - artifact corresponds to exactly one pooling behaviour: - - - ``pooling_mode=-1``: no pooling; ``forward()`` returns ``(N, D)``. - - ``pooling_mode=1``: sum pooling; ``forward()`` returns ``(B, D)``. - - ``pooling_mode=2``: mean pooling; ``forward()`` returns ``(B, D)``. - """ - - def __init__( + @torch.no_grad() + def _build_and_write_rows( self, - table_options: List["EmbeddingConfig"], - use_dynamic_hash: bool, - pooling_mode: int, - table_names: Optional[List[str]] = None, - feature_names: Optional[List[str]] = None, - feature_table_map: Optional[List[int]] = None, - output_dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - key_type: torch.dtype = torch.int64, - ): - super().__init__() - self.embedding_configs: List[EmbeddingConfig] = table_options - - if pooling_mode not in (-1, 1, 2): - raise ValueError( - f"pooling_mode must be -1 (no pooling), 1 (sum), or 2 (mean), " - f"got {pooling_mode}" - ) - if not table_options: - raise ValueError("table_options must be non-empty") - if device is None: - device = torch.device("cuda", torch.cuda.current_device()) - if key_type not in (torch.int64, torch.uint64): - raise ValueError(f"unsupported key_type: {key_type}") - if output_dtype not in (torch.float32, torch.float16): - raise ValueError(f"unsupported output_dtype: {output_dtype}") - - capacities = [_resolve_capacity(opt) for opt in table_options] - num_tables = len(table_options) - - if table_names is None: - table_names = [f"table_{i}" for i in range(num_tables)] - if len(table_names) != num_tables: - raise ValueError("table_names size must match table_options") - - if feature_table_map is None: - feature_table_map = list(range(num_tables)) - if not isinstance(feature_table_map, list) or len(feature_table_map) == 0: - raise ValueError("feature_table_map must be a non-empty list") - if any(t < 0 or t >= num_tables for t in feature_table_map): - raise ValueError( - f"feature_table_map contains out-of-range table id (must be in [0, {num_tables}))" - ) - for i in range(1, len(feature_table_map)): - if feature_table_map[i] < feature_table_map[i - 1]: - raise ValueError( - "feature_table_map must be non-decreasing (features for the same table must be contiguous)" - ) - - feature_offsets = _derive_grouped_offsets(feature_table_map) - - self.device = device - self.output_dtype_ = output_dtype - self.key_type_ = key_type - self.num_tables_ = num_tables - self.table_names_ = table_names - self.feature_names_ = feature_names - self.pooling_mode_ = ( - pooling_mode # plain Python int – compile-time constant for torch.export + feature_ids: torch.Tensor, + table_ids: torch.Tensor, + embeddings: torch.Tensor, + ) -> None: + storage_keys = self.indexer_.build_index(feature_ids, table_ids) + retained = storage_keys >= 0 + self._write_storage_rows( + storage_keys[retained], + embeddings[retained.to(device=embeddings.device)], ) - self.score_policy = int( - ScorePolicy.CONST - ) # plain Python int – compile-time constant for torch.export - self.register_buffer( - "feature_table_map_", - torch.tensor(feature_table_map, dtype=torch.int64, device=device), - ) - self.register_buffer( - "feature_offsets_", - torch.tensor(feature_offsets, dtype=torch.int64, device=device), - ) - self.register_buffer( - "capacity_list_", - torch.tensor(capacities, dtype=torch.int64, device=device), - ) - self.register_buffer( - "table_offsets_", - torch.zeros(num_tables + 1, dtype=torch.int64, device=device), - ) - self.capacity_list_ += ( - 1 # reserve the first row of each section for "not found" entries - ) - torch.cumsum( - torch.cat([torch.ones((1,), device=device), self.capacity_list_]), - dim=0, - out=self.table_offsets_, - ) + @torch.no_grad() + def _begin_source_load(self) -> None: + if isinstance(self.indexer_, LinearHashMapIndexer): + self.indexer_.reset_for_build() + if self.config_.nve_layer_type != "hierarchical": + self._nve_embedding().weight.data.zero_() - self.use_dynamic_hash = use_dynamic_hash - if use_dynamic_hash: - self.hash_table = InferenceLinearBucketTable( - capacity=capacities, - key_type=key_type, - bucket_capacity=128, - device=device, + if isinstance( + self.indexer_, (LinearHashMapIndexer, FusedIdentityIndexer) + ): + self._write_storage_rows( + self.indexer_.miss_storage_indices, + torch.zeros( + (self.num_tables_, self.emb_dim_), + dtype=self.output_dtype_, + device=self.device, + ), ) - self.emb_dim_ = _resolve_embedding_dim(table_options) - total_rows = int(self.capacity_list_.sum().item()) - dtype_size = torch.finfo(output_dtype).bits // 8 - total_size_bytes = total_rows * self.emb_dim_ * dtype_size - self.gpu_cache_size_ = _resolve_gpu_cache_size( - table_options, - total_size_bytes, - self.emb_dim_, - dtype_size, - ) - - if self.pooling_mode_ == -1: - self.nve_embedding_ = nve_layers.NVEmbedding( - num_embeddings=total_rows, - embedding_size=self.emb_dim_, - data_type=output_dtype, - gpu_cache_size=int(self.gpu_cache_size_), - optimize_for_training=False, - device=device, - **_linear_uvm_constructor_kwargs(), - ) - else: - if self.pooling_mode_ == 1: - mode = "sum" - elif self.pooling_mode_ == 2: - mode = "mean" - self.nve_embedding_ = nve_layers.NVEmbeddingBag( - num_embeddings=total_rows, - embedding_size=self.emb_dim_, - data_type=output_dtype, - mode=mode, - gpu_cache_size=int(self.gpu_cache_size_), - optimize_for_training=False, - device=device, - **_linear_uvm_constructor_kwargs(), + @torch.no_grad() + def _finish_source_load(self) -> None: + self.indexer_.finish_build() + if self.config_.nve_layer_type == "hierarchical": + self.nve_embedding_ = create_nve_layer( + num_embeddings=self.indexer_.nve_num_embeddings, + embedding_dim=self.emb_dim_, + dtype=self.output_dtype_, + pooling_mode=self.pooling_mode_, + config=self.config_, + device=self.device, ) + @torch.no_grad() def load_from_embedding_table(self, table_weights: torch.Tensor) -> None: - """Load embedding weights from a pre-extracted table tensor. - - This is used when the training checkpoint contains a pre-extracted embedding table - that can be directly copied into the NVE weight, bypassing the hash table insertion. - - Args: - table_weights: (sum of all table capacities, emb_dim) float tensor containing the embedding weights for all tables, ordered by table and then by row within each table. - """ - assert ( - table_weights.size(0) - <= self.nve_embedding_.weight.size(0) - self.num_tables_ - ), f"Provided table_weights has more rows ({table_weights.size(0)}) than the NVE embedding capacity ({self.nve_embedding_.weight.size(0) - self.num_tables_} excluding for reserved 'not found' rows)" - assert ( - table_weights.size(1) == self.emb_dim_ - ), f"Provided table_weights has embedding dim {table_weights.size(1)}, expected {self.emb_dim_}" - - self.nve_embedding_.weight.data.zero_() # zero out the "not found" row in each table - for table_id in range(self.num_tables_): - self.nve_embedding_.weight.data[ - self.table_offsets_[table_id].item(), : - ].zero_() - - current_offsets = self.table_offsets_[ - table_id - ].item() # table_offsets_ already skipped the reserved row - original_offsets = self.table_offsets_[table_id].item() - table_id - 1 - num_rows = ( - int(self.capacity_list_[table_id].item()) - 1 - ) # exclude reserved row - if table_id == self.num_tables_ - 1: - num_rows = ( - table_weights.size(0) - original_offsets - ) # use remaining rows for the last table - - self.nve_embedding_.weight.data[ - current_offsets : current_offsets + num_rows, : - ].copy_( - table_weights[original_offsets : original_offsets + num_rows, :].to( - self.nve_embedding_.weight.dtype - ) + """Populate from table-concatenated TorchRec checkpoint weights.""" + self._begin_source_load() + source_offset = 0 + for table_id, capacity in enumerate(self.table_capacities_): + feature_ids = torch.arange( + capacity, dtype=torch.int64, device=self.device + ) + table_ids = torch.full_like(feature_ids, table_id) + self._build_and_write_rows( + feature_ids, + table_ids, + table_weights[source_offset : source_offset + capacity], ) + source_offset += capacity + self._finish_source_load() + @torch.no_grad() def load_from_dynamicemb_file( self, save_dir: str, - table_names: Optional[List[str]] = None, + table_names: Optional[Sequence[str]] = None, ) -> None: + """Populate from a complete DynamicEmb dump.""" if not os.path.exists(save_dir): raise RuntimeError(f"Save directory does not exist: {save_dir}") - if ( - "get_loading_files" not in globals() - or "_iter_batches_from_files" not in globals() - ): - raise RuntimeError( - "dynamicemb load helpers are unavailable. Ensure dynamicemb and inference operators are importable." - ) - - if table_names is None: - table_names = self.table_names_ - - requested_table_names = set(table_names) - dim = self.emb_dim_ - device = self.device - weight = self.nve_embedding_.weight.data - - self.hash_table.table_storage_.zero_() - self.hash_table.bucket_sizes.zero_() - self.hash_table._ref_counter.zero_() - weight.zero_() - + selected_names = set( + self.table_names_ if table_names is None else table_names + ) + self._begin_source_load() for table_id, table_name in enumerate(self.table_names_): - if table_name not in requested_table_names: + if table_name not in selected_names: continue - - meta_json_file = encode_meta_json_file_path(save_dir, table_name) - if os.path.exists(meta_json_file): - try: - _ = load_from_json(meta_json_file) - except Exception as e: - print( - f"[WARN] Failed to read meta json for {table_name} at {meta_json_file}: {e}" - ) - + meta_path = encode_meta_json_file_path(save_dir, table_name) + if os.path.exists(meta_path): + load_from_json(meta_path) ( - emb_key_files, - emb_value_files, - emb_score_files, - _opt_value_files, + key_files, + value_files, + score_files, + _optimizer_files, _counter_key_files, _counter_frequency_files, - ) = get_loading_files( - save_dir, - table_name, - rank=0, - world_size=1, - ) - - if len(emb_key_files) == 0: - print(f"[INFO] No checkpoint files found for table: {table_name}") - continue - - num_key_files = len(emb_key_files) - for i in range(num_key_files): - score_file = emb_score_files[i] if i < len(emb_score_files) else None - for keys, embeddings, scores, _opt_states in _iter_batches_from_files( - emb_key_files[i], - emb_value_files[i], - score_file, - None, - dim, - 0, - device, + ) = get_loading_files(save_dir, table_name, rank=0, world_size=1) + for file_index, (key_file, value_file) in enumerate( + zip(key_files, value_files) + ): + score_file = ( + score_files[file_index] + if file_index < len(score_files) + else None + ) + for feature_ids, embeddings, _scores, _optimizer_states in ( + _iter_batches_from_files( + key_file, + value_file, + score_file, + None, + self.emb_dim_, + 0, + self.device, + ) ): - if keys.numel() == 0: - continue - table_ids = torch.full( - (keys.numel(),), + (feature_ids.numel(),), table_id, dtype=torch.int64, - device=device, + device=self.device, ) - policy = ( - ScorePolicy.ASSIGN if scores is not None else ScorePolicy.CONST - ) - indices = table_insert( - self.hash_table.table_storage_, - self.hash_table.table_bucket_offsets_, - self.hash_table.bucket_capacity_, - self.hash_table.bucket_sizes, - keys, - table_ids, - scores, - policy, - self.hash_table._ref_counter, - None, - None, - ) - - valid_mask = indices >= 0 - if not torch.all(valid_mask): - num_failed = (~valid_mask).sum().item() - print( - f"[WARN] table_insert failed for {num_failed} keys in table {table_name}." - ) - - valid_indices = indices[valid_mask].to(torch.int64) - if valid_indices.numel() == 0: - continue - - max_index = valid_indices.max().item() - table_cap = int(self.capacity_list_[table_id].item()) - if max_index >= table_cap: - raise RuntimeError( - f"nve_embedding has insufficient rows ({table_cap}) for loaded index {max_index}." - ) - - abs_indices = valid_indices + self.table_offsets_[table_id] - weight.index_copy_( - 0, - abs_indices.to(torch.int64), - embeddings[valid_mask].to(weight.dtype), + self._build_and_write_rows( + feature_ids, table_ids, embeddings ) + self._finish_source_load() def forward( self, @@ -555,26 +249,6 @@ def forward( pooling_offsets: Optional[torch.Tensor] = None, per_sample_weights: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Run embedding lookup with optional pooling. - - Args: - keys: (N,) int64 – flat lookup keys; may span multiple - tables and pooling bags. - offsets: (T*B+1,) int64 – CSR boundaries that map segments - of ``keys`` to per-table feature slots; used to - derive a per-key table id via - ``INFERENCE_EMB::expand_table_ids``. - pooling_offsets: (B+1,) int64 – CSR boundaries that map segments of - ``keys`` to pooling bags; required when - ``self.pooling_mode_ >= 0``, otherwise unused. - - Returns: - - ``pooling_mode_ == -1``: ``(N, D)`` float tensor of per-key - embeddings. - - ``pooling_mode_ == 1 or 2``: ``(B, D)`` float tensor of pooled - embeddings, where ``B = pooling_offsets.size(0) - 1``. - """ - # Derive per-key table id from the table-segment offsets. table_ids = torch.ops.INFERENCE_EMB.expand_table_ids( offsets, keys, @@ -582,185 +256,180 @@ def forward( self.num_tables_, 1, ) - - if self.use_dynamic_hash: - # Hash-table lookup: keys → per-table row indices. - _scores, _founds, table_indices = self.hash_table.lookup( - keys=keys, - table_ids=table_ids, - score_value=None, - score_policy=self.score_policy, - ) - else: - # Use identical-mapping per-table indices - table_indices = keys - - # Convert per-table indices to absolute embedding row ids. - global_table_offsets = torch.index_select(self.table_offsets_, 0, table_ids) - global_indices = table_indices + global_table_offsets - - # self.pooling_mode_ is a plain Python int attribute – a compile-time - # constant for torch.export. The branch below is statically resolved - # during tracing: only one code path is included in the exported graph. + storage_keys = self.indexer_(keys, table_ids) + nve_embedding = self._nve_embedding() if self.pooling_mode_ < 0: - # Non-pooled path: return one embedding vector per key. - return self.nve_embedding_(global_indices) - else: - # Pooled path: reduce each bag of keys to one embedding vector. - return self.nve_embedding_( - global_indices, # (N,) – absolute row ids - pooling_offsets, # (B+1,) – CSR bag boundaries - per_sample_weights, # optional per-key weights - ) + return nve_embedding(storage_keys) + return nve_embedding( + storage_keys, pooling_offsets, per_sample_weights + ) def create_inference_embedding_collection( - embedding_configs: List[EmbeddingConfig], - pooling_mode: int = -1, # non pooling - use_dynamic: bool = True, + embedding_configs: Sequence[EmbeddingConfig], + *, + pooling_mode: int, + config: InferenceEmbeddingCollectionConfig, + output_dtype: torch.dtype = torch.float32, + key_type: torch.dtype = torch.int64, + device: Optional[torch.device] = None, ) -> InferenceEmbeddingCollection: - table_names = [config.name for config in embedding_configs] - feature_names = list( - itertools.chain(*[config.feature_names for config in embedding_configs]) - ) + configs = list(embedding_configs) + if not configs: + raise ValueError("embedding_configs must not be empty") + if pooling_mode not in (-1, 1, 2): + raise ValueError("pooling_mode must be -1, 1, or 2") + if device is None: + device = torch.device("cuda", torch.cuda.current_device()) + validate_collection_config(config) + + embedding_dims = {int(table.embedding_dim) for table in configs} + if len(embedding_dims) != 1: + raise ValueError("Tables in one collection must share embedding_dim") + embedding_dim = embedding_dims.pop() + table_names = [table.name for table in configs] + feature_names_by_table = [ + list(table.feature_names) if table.feature_names else [table.name] + for table in configs + ] + feature_names = list(itertools.chain.from_iterable(feature_names_by_table)) feature_table_map = list( - itertools.chain( - *[ - [idx] * len(config.feature_names) - for idx, config in enumerate(embedding_configs) - ] + itertools.chain.from_iterable( + [table_id] * len(names) + for table_id, names in enumerate(feature_names_by_table) ) ) - table_options = [ - DynamicEmbTableOptions( - embedding_dtype=torch.float32, - dim=config.embedding_dim, - max_capacity=config.num_embeddings, - local_hbm_for_values=0, - bucket_capacity=128, - initializer_args=DynamicEmbInitializerArgs( - mode=DynamicEmbInitializerMode.NORMAL, - ), - training=False, + indexer = create_embedding_collection_indexer( + config, + [int(table.num_embeddings) for table in configs], + key_type=key_type, + device=device, + ) + nve_embedding = None + if config.nve_layer_type != "hierarchical": + nve_embedding = create_nve_layer( + num_embeddings=indexer.nve_num_embeddings, + embedding_dim=embedding_dim, + dtype=output_dtype, + pooling_mode=pooling_mode, + config=config, + device=device, ) - for config in embedding_configs - ] + return InferenceEmbeddingCollection( - table_options, - use_dynamic, - pooling_mode, - table_names, - feature_names, - feature_table_map, - device=torch.device("cuda"), + configs, + config=config, + indexer=indexer, + nve_embedding=nve_embedding, + pooling_mode=pooling_mode, + table_names=table_names, + feature_names=feature_names, + feature_table_map=feature_table_map, + output_dtype=output_dtype, + device=device, ) +def _resolve_collection_config( + embedding_configs: Sequence[EmbeddingConfig], + embedding_collection_configs: Mapping[ + str, InferenceEmbeddingCollectionConfig + ], +) -> InferenceEmbeddingCollectionConfig: + configs = [embedding_collection_configs[table.name] for table in embedding_configs] + if any(config != configs[0] for config in configs[1:]): + raise ValueError("Tables in one physical collection must share one config") + return configs[0] + + +def _resolve_pooling_mode(embedding_configs: Sequence[EmbeddingConfig]) -> int: + pooling = getattr(embedding_configs[0], "pooling", "NONE") + pooling_name = getattr(pooling, "name", pooling) + if pooling_name == "NONE": + return -1 + if pooling_name == "SUM": + return 1 + if pooling_name == "MEAN": + return 2 + raise ValueError(f"Unsupported pooling config: {pooling}") + + +def _replace_submodule( + model: torch.nn.Module, + module_path: str, + replacement: torch.nn.Module, +) -> None: + parent_path, separator, child_name = module_path.rpartition(".") + parent = model.get_submodule(parent_path) if separator else model + setattr(parent, child_name, replacement) + + def apply_inference_embedding_collection( model: torch.nn.Module, - dynamic_table_configs: Dict[str, bool], - trained_emb_table_sizes: Dict[str, int], -): - """ - Replace torchrec.EmbeddingCollection in the model with dynamicemb.InferenceEmbeddingCollection. - - Args: - model (torch.nn.Module): The input training model. - dynamic_table_configs (Dict[str, bool]): A dictionary mapping table names to their use_dynamic flag. - trained_emb_table_sizes (Dict[str, int]): A dictionary mapping table names to their corresponding vocabulary sizes in the trained model. - Returns: - torch.nn.Module: The training model with the exportable embedding. - """ - from torch.nn.modules.sparse import Embedding - - check_modules = set() + embedding_collection_configs: Mapping[ + str, InferenceEmbeddingCollectionConfig + ], + trained_emb_table_sizes: Mapping[str, int], +) -> torch.nn.Module: + """Replace TorchRec embedding collections with configured NVE collections.""" + checked_modules: set[str] = set() while True: - name = None + candidate_name = None for name, module in model.named_modules(): - if isinstance(module, ModuleDict): - submodules = [ - submodule - for name, submodule in module.named_modules() - if name != "" - ] - # skip if not torchrec.EmbeddingCollection - if len({type(m) for m in submodules}) != 1 or not isinstance( - submodules[0], Embedding - ): - continue - # skip if already converted - if ( - isinstance(module, nve_layers.NVEmbedding) - or isinstance(module, nve_layers.NVEmbeddingBag) - or isinstance(module, InferenceEmbeddingCollection) - ): - continue - # skip if already checked - if name in check_modules: - continue + if not isinstance(module, ModuleDict) or name in checked_modules: + continue + children = [ + child + for child_name, child in module.named_modules() + if child_name + ] + if ( + children + and len({type(child) for child in children}) == 1 + and isinstance(children[0], Embedding) + ): + candidate_name = name break - else: + if candidate_name is None: break - embedding_configs = None - for parent_name, parent_module in model.named_modules(): - if parent_name == name.removesuffix(".embeddings"): - embedding_configs = parent_module.embedding_configs() - break - assert ( - embedding_configs is not None - ), f"Cannot find embedding configs from parent module {name.removesuffix('.embeddings')}" - - parent_name = name.removesuffix(".embeddings") - check_modules.add(name) - check_modules.add(parent_name) - - # Adjust training vocab sizes to inference vocab sizes based on trained_emb_table_sizes - # TODO(junyiq): Try exact the freezed vocab size from embedding files/model states. - for config in embedding_configs: - if config.name not in trained_emb_table_sizes: - print( - f"[WARNING] Table {config.name} in module {name} is missing the trained vocab size for inference.\n" - + f" Using {config.vocab_size} rows from the training config." - ) - config.num_embeddings = trained_emb_table_sizes.get( - config.name, config.num_embeddings + parent_name = candidate_name.removesuffix(".embeddings") + parent_module = model.get_submodule(parent_name) + embedding_configs = list(parent_module.embedding_configs()) + checked_modules.update((candidate_name, parent_name)) + for table in embedding_configs: + table.num_embeddings = trained_emb_table_sizes.get( + table.name, table.num_embeddings ) - use_dynamic = { - config.name: dynamic_table_configs[config.name] - for config in embedding_configs - if config.name in dynamic_table_configs - } - assert ( - len(use_dynamic) > 0 - ), "At least one table in the embedding collection module should have a config in dynamic_table_configs." - assert ( - len(set(use_dynamic.values())) == 1 - ), f"All tables in the same embedding collection module should have the same config in dynamic_table_configs. Got:\n{use_dynamic}" - use_dynamic = list(use_dynamic.values())[0] - - pooling_config = getattr(embedding_configs[0], "pooling", "NONE") - if pooling_config == "NONE": - pooling_mode = -1 - elif pooling_config == "SUM": - pooling_mode = 1 - elif pooling_config == "MEAN": - pooling_mode = 2 - else: - raise ValueError(f"Unsupported pooling config: {pooling_config}") - - embedding_collection = create_inference_embedding_collection( + collection_config = _resolve_collection_config( + embedding_configs, embedding_collection_configs + ) + replacement = create_inference_embedding_collection( embedding_configs, - pooling_mode, - use_dynamic, + pooling_mode=_resolve_pooling_mode(embedding_configs), + config=collection_config, ) - assert isinstance(embedding_collection, InferenceEmbeddingCollection) - embedding_collection.embedding_configs = embedding_configs - - exec("model." + parent_name + "=embedding_collection") + replacement.collection_id_ = parent_name + _replace_submodule(model, parent_name, replacement) print( - f"[INFO] converting {parent_name} to InferenceEmbeddingCollection with use_dynamic={use_dynamic} and tables={embedding_collection.table_names_}" + f"[INFO] converted {parent_name}: " + f"{collection_config.indexer_type.value} + " + f"{collection_config.nve_layer_type}" ) - return model + + +__all__ = [ + "BitConcatConfig", + "BitConcatIndexer", + "EmbeddingCollectionIndexerBase", + "EmbeddingCollectionIndexerType", + "FusedIdentityIndexer", + "IdentityIndexer", + "InferenceEmbeddingCollection", + "InferenceEmbeddingCollectionConfig", + "LinearHashMapIndexer", + "apply_inference_embedding_collection", + "create_inference_embedding_collection", +] diff --git a/corelib/dynamicemb/example/exportable_embedding/README.md b/corelib/dynamicemb/example/exportable_embedding/README.md new file mode 100644 index 000000000..cf59ce12a --- /dev/null +++ b/corelib/dynamicemb/example/exportable_embedding/README.md @@ -0,0 +1,79 @@ +# Exportable embedding-collection workflow + +This example creates deterministic TorchRec embedding collections, writes one +complete checkpoint plus the required DynamicEmb dumps, converts the +collections to configurable indexer/NVE pairs, exports the AOTI package with +out-of-band indexer state, and verifies eager and Python AOTI replay. With NVE +26.06 or later and Redis enabled, it also applies two incremental rounds +through the reusable coordinator subprocess and per-runtime subscribers; the +second round updates only the BitConcat/Redis collection while inference +continues. An optional C++ replayer stays alive for the complete workflow. + +The Python and C++ runtimes use the same four rounds: + +- Round 1 verifies 100 requests against the initial state, then globally + synchronizes. +- Inference pauses while the update worker applies the first update. +- Round 2 verifies 100 requests against the first incremental state, then + globally synchronizes. +- Round 3 runs at least 1,000 unchecked requests while the BitConcat/Redis + update runs concurrently, followed by 100 post-update requests and a global + synchronization. +- Round 4 verifies 100 requests against the second incremental state, then + globally synchronizes. + +Normal inference calls the model directly on the device's default CUDA stream. +It does not acquire a publication lock or record a CUDA event; a replacement +snapshot records one retirement event only when inference adopts it. + +The default workflow covers: + +- `FusedIdentityIndexer + GPULayer` +- `LinearHashMapIndexer + LinearUVM` +- `LinearHashMapIndexer + Hierarchical(GPU cache -> NVHashMap host cache -> Redis)` +- `BitConcatIndexer + Hierarchical(GPU cache -> NVHashMap host cache -> Redis)` +- one paused and one concurrent Redis incremental-load round + +NVE 26.05 runs the two non-hierarchical combinations. For NVE 26.06 or later, +`run_example.sh` starts and stops a local standalone Redis server, launches the +coordinator subprocess, and runs the in-process DynamicEmb delta producers. +The NVHashMap host caches are empty when loaded and are populated by inference; +their contents are not exported. + +Build the optional C++ replayer for the selected NVE installation: + +```bash +cmake -S /workspace/recsys-examples/corelib/dynamicemb/example/exportable_embedding/cpp \ + -B /workspace/recsys-examples/corelib/dynamicemb/example/exportable_embedding/cpp/build \ + -DNVE_ROOT=/workspace/deps/nve \ + -DNVE_LIB_DIR=/opt/nve/default/python/pynve \ + -DDYNAMICEMB_LIB_DIR=/workspace/recsys-examples/corelib/dynamicemb/torch_binding_build +cmake --build /workspace/recsys-examples/corelib/dynamicemb/example/exportable_embedding/cpp/build -j 8 +``` + +Run all four combinations with the repository NVE (26.07 or later): + +```bash +/workspace/recsys-examples/corelib/dynamicemb/example/exportable_embedding/run_example.sh \ + --work-dir /tmp/exportable-embedding-example \ + --cpp-replayer /workspace/recsys-examples/corelib/dynamicemb/example/exportable_embedding/cpp/build/replay_and_verify +``` + +To use an existing Redis server instead, run the complete workflow with local +Redis disabled: + +```bash +START_LOCAL_REDIS=0 \ +/workspace/recsys-examples/corelib/dynamicemb/example/exportable_embedding/run_example.sh \ + --work-dir /tmp/exportable-embedding-example \ + --redis-address 127.0.0.1:6379 \ + --cpp-replayer /workspace/recsys-examples/corelib/dynamicemb/example/exportable_embedding/cpp/build/replay_and_verify +``` + +Run the two NVE 26.05 combinations: + +```bash +NVE_VERSION=26.05 \ +/workspace/recsys-examples/corelib/dynamicemb/example/exportable_embedding/run_example.sh \ + --work-dir /tmp/exportable-embedding-example-2605 +``` diff --git a/corelib/dynamicemb/example/exportable_embedding/cpp/CMakeLists.txt b/corelib/dynamicemb/example/exportable_embedding/cpp/CMakeLists.txt new file mode 100644 index 000000000..2de7c7120 --- /dev/null +++ b/corelib/dynamicemb/example/exportable_embedding/cpp/CMakeLists.txt @@ -0,0 +1,71 @@ +cmake_minimum_required(VERSION 3.18) +project(dynamicemb_exportable_embedding_example LANGUAGES CXX CUDA) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +set(NVE_ROOT "${CMAKE_CURRENT_LIST_DIR}/../../../../../third_party/nv-embedding-cache" CACHE PATH "NVE source root") +set(NVE_LIB_DIR "/opt/nve/default/python/pynve" CACHE PATH "NVE library directory") +set(NVTX_INCLUDE_DIR "/workspace/deps/NVTX/c/include" CACHE PATH "NVTX include directory") +set(DYNAMICEMB_ROOT "${CMAKE_CURRENT_LIST_DIR}/../../.." CACHE PATH "DynamicEmb source root") +set(DYNAMICEMB_LIB_DIR "${DYNAMICEMB_ROOT}/torch_binding_build" CACHE PATH "DynamicEmb library directory") +include("${DYNAMICEMB_ROOT}/cmake/NVETarget.cmake") + +find_package(Python3 COMPONENTS Interpreter REQUIRED) +execute_process( + COMMAND "${Python3_EXECUTABLE}" -c "import torch; print(torch.utils.cmake_prefix_path)" + OUTPUT_VARIABLE TORCH_CMAKE_PREFIX_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE TORCH_PREFIX_STATUS +) +if(NOT TORCH_PREFIX_STATUS EQUAL 0) + message(FATAL_ERROR "Python cannot locate the PyTorch CMake package") +endif() +list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}") +find_package(Torch REQUIRED) +find_package(CUDAToolkit REQUIRED) +find_package(Threads REQUIRED) +find_file(INFERENCE_EMB_OPS inference_emb_ops.so PATHS "${DYNAMICEMB_LIB_DIR}" NO_DEFAULT_PATH REQUIRED) +find_library(NVE_TORCH_OPS nve-torch-ops PATHS "${NVE_LIB_DIR}" NO_DEFAULT_PATH REQUIRED) +find_library(NVE_COMMON nve-common PATHS "${NVE_LIB_DIR}" NO_DEFAULT_PATH REQUIRED) + +add_executable(replay_and_verify replay_and_verify.cpp) +set_source_files_properties(replay_and_verify.cpp PROPERTIES LANGUAGE CUDA) +separate_arguments(TORCH_CXX_FLAGS_LIST NATIVE_COMMAND "${TORCH_CXX_FLAGS}") +target_compile_options(replay_and_verify PRIVATE ${TORCH_CXX_FLAGS_LIST} -w) +target_link_options(replay_and_verify PRIVATE "LINKER:--no-as-needed") +target_include_directories(replay_and_verify PRIVATE + "${DYNAMICEMB_ROOT}/include" + "${NVE_ROOT}" + "${NVE_ROOT}/include" + "${NVE_ROOT}/third_party/json/include" + "${NVTX_INCLUDE_DIR}" + ${TORCH_INCLUDE_DIRS} +) +dynamicemb_configure_nve_target(replay_and_verify) +target_link_libraries(replay_and_verify PRIVATE + ${TORCH_LIBRARIES} + Threads::Threads + CUDA::cudart + CUDA::cuda_driver + "${INFERENCE_EMB_OPS}" + "${NVE_TORCH_OPS}" + "${NVE_COMMON}" +) + +# Unset selects the repository NVE (26.07 or later); set 26.05 explicitly for compatibility. +if( + NOT DEFINED NVE_VERSION + OR NVE_VERSION STREQUAL "" + OR NVE_VERSION VERSION_GREATER_EQUAL "26.06" +) + find_file(INFERENCE_EMB_UPDATE inference_emb_update.so PATHS "${DYNAMICEMB_LIB_DIR}" NO_DEFAULT_PATH REQUIRED) + target_link_libraries(replay_and_verify PRIVATE "${INFERENCE_EMB_UPDATE}") +else() + target_compile_definitions(replay_and_verify PRIVATE DYNAMICEMB_NVE_2605=1) +endif() + +set_target_properties(replay_and_verify PROPERTIES + BUILD_RPATH "${DYNAMICEMB_LIB_DIR};${NVE_LIB_DIR}" +) diff --git a/corelib/dynamicemb/example/exportable_embedding/cpp/replay_and_verify.cpp b/corelib/dynamicemb/example/exportable_embedding/cpp/replay_and_verify.cpp new file mode 100644 index 000000000..0c977e7b3 --- /dev/null +++ b/corelib/dynamicemb/example/exportable_embedding/cpp/replay_and_verify.cpp @@ -0,0 +1,306 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "dynamicemb/exportable_embedding/indexer_directory.h" +#if !defined(DYNAMICEMB_NVE_2605) +#include "dynamicemb/exportable_embedding/incremental_update.h" +#endif +#include "python/pynve/torch_bindings/nve_loader.hpp" + +namespace { + +constexpr int64_t kRowsPerTable = 256; +constexpr int kRequestsPerVerifiedRound = 100; +constexpr int kRoundThreeMinInferences = 1000; +constexpr int kRoundThreePostStageInferences = 100; + +struct Arguments { + std::string package_dir; + std::string expected_sums; + bool wait_for_updates{false}; + int device{0}; +}; + +Arguments parse_arguments(int argc, char** argv) { + Arguments args; + for (int index = 1; index < argc; ++index) { + const std::string name = argv[index]; + if (name == "--package-dir") { + args.package_dir = argv[++index]; + } else if (name == "--expected-sums") { + args.expected_sums = argv[++index]; + } else if (name == "--wait-for-updates") { + args.wait_for_updates = true; + } else if (name == "--device") { + args.device = std::stoi(argv[++index]); + } else { + throw std::runtime_error("Unknown argument: " + name); + } + } + if (args.package_dir.empty() || args.expected_sums.empty()) { + throw std::runtime_error("--package-dir and --expected-sums are required"); + } + return args; +} + +std::vector parse_sums(const std::string& text) { + std::vector values; + std::stringstream stream(text); + std::string item; + while (std::getline(stream, item, ',')) { + values.push_back(std::stod(item)); + } + return values; +} + +std::vector make_inputs(bool after_update, + int device, + int64_t new_feature_id = kRowsPerTable) { + const auto options = torch::TensorOptions() + .dtype(torch::kInt64) + .device(torch::kCUDA, device); + if (after_update) { + return { + torch::tensor({1, 1}, options), + torch::tensor({0, 1, 2}, options), + torch::tensor({int64_t{1}, new_feature_id, int64_t{1}, new_feature_id}, + options), + torch::tensor({0, 2, 4}, options), + }; + } + auto table_keys = torch::arange(kRowsPerTable, options); + auto keys = torch::cat({table_keys, table_keys}); + return { + keys, + torch::tensor({int64_t{0}, kRowsPerTable, 2 * kRowsPerTable}, options), + keys.clone(), + torch::tensor({int64_t{0}, kRowsPerTable, 2 * kRowsPerTable}, options), + }; +} + +bool verify_sums(const std::vector& outputs, + const std::vector& expected, + std::string& error) { + if (outputs.size() != expected.size()) { + std::ostringstream message; + message << "unexpected output count: actual=" << outputs.size() + << " expected=" << expected.size(); + error = message.str(); + return false; + } + for (std::size_t index = 0; index < outputs.size(); ++index) { + const double actual = outputs[index].sum().item(); + const double tolerance = std::max(0.01, std::abs(expected[index]) * 1e-6); + if (std::abs(actual - expected[index]) > tolerance) { + std::ostringstream message; + message << "output=" << index << " actual=" << actual + << " expected=" << expected[index]; + error = message.str(); + return false; + } + } + return true; +} + +} // namespace + +int main(int argc, char** argv) { + try { + const auto args = parse_arguments(argc, argv); + c10::cuda::CUDAGuard guard(args.device); + +#if !defined(DYNAMICEMB_NVE_2605) + std::shared_ptr nve_resources; +#endif + std::unique_ptr nve_layers; + std::unique_ptr + indexers; + std::unique_ptr loader; +#if defined(DYNAMICEMB_NVE_2605) + nve_layers = std::make_unique(args.package_dir, + args.device); + loader = std::make_unique( + args.package_dir + "/model.pt2", "model", false, 1, args.device); +#else + loader = std::make_unique( + args.package_dir + "/model.pt2", "model", false, 1, args.device); + nve_resources = std::make_shared(); + nve_layers = std::make_unique( + args.package_dir, *loader, args.device, nve_resources); +#endif + + indexers = dynamicemb::exportable_embedding:: + EmbeddingCollectionIndexerDirectory::load(args.package_dir, args.device); + indexers->bind(*loader); + const auto inference_cuda_stream = + at::cuda::getDefaultCUDAStream(args.device); + const auto inference_stream = inference_cuda_stream.stream(); + c10::cuda::CUDAStreamGuard inference_stream_guard(inference_cuda_stream); + + const auto run_and_verify = [&](const char* stage, + bool after_update, + int64_t new_feature_id, + const std::vector& expected) { + std::vector outputs; + std::string error; + for (int iteration = 0; iteration < kRequestsPerVerifiedRound; + ++iteration) { + outputs = loader->run( + make_inputs(after_update, args.device, new_feature_id), + inference_stream); + if (!verify_sums(outputs, expected, error)) break; + } + cudaDeviceSynchronize(); + const bool verified = error.empty(); + if (!verified) { + std::cout << "RESULT ERROR stage=" << stage << ' ' << error << '\n' + << std::flush; + } + return verified; + }; + + bool verification_ok = run_and_verify( + "round_1_original", + false, + kRowsPerTable, + parse_sums(args.expected_sums)); + +#if !defined(DYNAMICEMB_NVE_2605) + if (args.wait_for_updates) { + dynamicemb::exportable_embedding::EmbeddingCollectionUpdateSubscriber + subscriber(args.package_dir, *indexers, *nve_layers, args.device); + std::vector + updates; + int update_round = 0; + + std::cout << "READY incremental updates" << std::endl; + std::string line; + while (std::getline(std::cin, line)) { + if (line == "STOP") { + break; + } + if (line.rfind("RUN ", 0) == 0) { + const auto expected = parse_sums(line.substr(4)); + std::vector acknowledgements; + if (update_round == 0) { + std::promise update_staged; + auto update_staged_future = update_staged.get_future(); + std::thread update_thread([&] { + for (const auto& update : updates) { + subscriber.apply_incremental_load(update, inference_stream); + } + update_staged.set_value(); + + for (const auto& update : updates) { + acknowledgements.push_back( + subscriber + .wait_for_retirement( + update.collection_id, update.snapshot_id) + .to_json()); + } + }); + + update_staged_future.wait(); + const bool round_ok = run_and_verify( + "round_2_snapshot_1", true, kRowsPerTable, expected); + verification_ok = round_ok && verification_ok; + update_thread.join(); + } else { + std::promise inference_started; + auto inference_started_future = inference_started.get_future(); + std::atomic update_staged{false}; + + std::thread update_thread([&] { + inference_started_future.wait(); + + for (const auto& update : updates) { + subscriber.apply_incremental_load(update, inference_stream); + } + update_staged.store(true, std::memory_order_release); + + for (const auto& update : updates) { + acknowledgements.push_back( + subscriber + .wait_for_retirement( + update.collection_id, update.snapshot_id) + .to_json()); + } + }); + + for (int iteration = 0; + iteration < kRoundThreeMinInferences; + ++iteration) { + loader->run(make_inputs(true, args.device, kRowsPerTable + 1), + inference_stream); + if (iteration == 0) inference_started.set_value(); + } + while (!update_staged.load(std::memory_order_acquire)) { + loader->run(make_inputs(true, args.device, kRowsPerTable + 1), + inference_stream); + } + for (int iteration = 0; + iteration < kRoundThreePostStageInferences; + ++iteration) { + loader->run(make_inputs(true, args.device, kRowsPerTable + 1), + inference_stream); + } + cudaDeviceSynchronize(); + update_thread.join(); + + const bool round_ok = run_and_verify( + "round_4_snapshot_2", true, kRowsPerTable + 1, expected); + verification_ok = round_ok && verification_ok; + } + for (const auto& acknowledgement : acknowledgements) { + std::cout << "ACK " << acknowledgement << '\n'; + } + updates.clear(); + ++update_round; + std::cout << "READY incremental updates" << std::endl; + continue; + } + if (line.empty()) continue; + updates.push_back(dynamicemb::exportable_embedding:: + EmbeddingCollectionUpdate::from_json(line)); + } + } +#else + if (args.wait_for_updates) { + throw std::runtime_error( + "Incremental replay requires NVE 26.06 or later"); + } +#endif + if (verification_ok) { + std::cout << "verified C++ AOTI replay" << std::endl; + } else { + std::cout << "C++ AOTI replay completed with verification errors" + << std::endl; + } + return verification_ok ? 0 : 1; + } catch (const std::exception& error) { + std::cerr << error.what() << std::endl; + return 1; + } +} diff --git a/corelib/dynamicemb/example/exportable_embedding/export_and_verify.py b/corelib/dynamicemb/example/exportable_embedding/export_and_verify.py new file mode 100644 index 000000000..a5a991b1b --- /dev/null +++ b/corelib/dynamicemb/example/exportable_embedding/export_and_verify.py @@ -0,0 +1,1201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build, export, reload, and verify the table-fusion example workflow.""" + +from __future__ import annotations + +import argparse +import io +import os +import subprocess +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from multiprocessing.connection import Client, Connection +from pathlib import Path +from threading import Event +from typing import Any + +import torch +import torch.distributed as dist + + +torch.ops.load_library( + os.path.join(os.environ["DYNAMICEMB_OPS_LIB_DIR"], "inference_emb_ops.so") +) + +# Register DynamicEmb's export metadata implementations after loading the ops. +import dynamicemb.index_range_meta as _index_range_meta # noqa: E402,F401 +import dynamicemb.lookup_meta as _lookup_meta # noqa: E402,F401 +from dynamicemb import ( # noqa: E402 + DeltaDumpResult, + DynamicEmbPoolingMode, + DynamicEmbScoreStrategy, + DynamicEmbTableOptions, + EmbOptimType, + EvictedItemMode, +) +from dynamicemb.batched_dynamicemb_tables import ( # noqa: E402 + BatchedDynamicEmbeddingTablesV2, +) +from dynamicemb.exportable_embedding import ( # noqa: E402 + BitConcatConfig, + EmbeddingCollectionIndexerDirectory, + EmbeddingCollectionIndexerType, + EmbeddingCollectionUpdate, + EmbeddingCollectionUpdateAck, + InferenceEmbeddingCollectionConfig, + embedding_collection_indexers_from_model, + export_embedding_collection_aot, + imported_nve_generation, + load_embedding_collection_aot, + register_nve_export_compat, +) +from dynamicemb.exportable_tables import ( # noqa: E402 + apply_inference_embedding_collection, +) +from pynve import nve # noqa: E402 +from pynve.torch.nve_ps import NVEParameterServer # noqa: E402 +from torchrec import DataType # noqa: E402 +from torchrec.modules.embedding_configs import EmbeddingConfig # noqa: E402 +from torchrec.modules.embedding_modules import EmbeddingCollection # noqa: E402 + + +register_nve_export_compat() + +EMBEDDING_DIM = 512 +ROWS_PER_TABLE = 256 +GPU_CACHE_SIZE = 4 << 20 +HOST_CACHE_SIZE = 4 << 20 +RANDOM_SEED = 20260903 +COORDINATOR_AUTHKEY = b"exportable-embedding-example" +REQUESTS_PER_VERIFIED_ROUND = 100 +ROUND_THREE_MIN_INFERENCES = 1_000 +ROUND_THREE_POST_STAGE_INFERENCES = 100 + +FUSED_IDENTITY_GPU = "to_fused_identity_gpu" +LINEAR_HASH_UVM = "to_linear_hash_uvm" +LINEAR_HASH_REDIS = "to_linear_hash_hierarchical_redis" +BIT_CONCAT_REDIS = "to_bitconcat_hierarchical_redis" + +BASE_COLLECTIONS = ( + FUSED_IDENTITY_GPU, + LINEAR_HASH_UVM, +) +REDIS_COLLECTIONS = (LINEAR_HASH_REDIS, BIT_CONCAT_REDIS) +DUMP_COLLECTIONS = ( + LINEAR_HASH_UVM, + LINEAR_HASH_REDIS, + BIT_CONCAT_REDIS, +) +DENSE_COLLECTIONS = (FUSED_IDENTITY_GPU,) + + +def embedding_configs(collection_name: str) -> list[EmbeddingConfig]: + return [ + EmbeddingConfig( + name=f"{collection_name}_table_{table_id}", + embedding_dim=EMBEDDING_DIM, + num_embeddings=ROWS_PER_TABLE, + feature_names=[f"{collection_name}_feature_{table_id}"], + data_type=DataType.FP32, + ) + for table_id in range(2) + ] + + +class TrainingSparseModel(torch.nn.Module): + def __init__(self, collection_names: tuple[str, ...], device: torch.device): + super().__init__() + self.collections = torch.nn.ModuleDict( + { + name: EmbeddingCollection( + tables=embedding_configs(name), device=device + ) + for name in collection_names + } + ) + + +class ExportableSparseExample(torch.nn.Module): + def __init__( + self, + converted_model: TrainingSparseModel, + collection_names: tuple[str, ...], + ) -> None: + super().__init__() + self.collections = converted_model.collections + self.collection_names = collection_names + + def forward( + self, + base_keys: torch.Tensor, + base_offsets: torch.Tensor, + redis_keys: torch.Tensor, + redis_offsets: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + outputs = [] + for name in self.collection_names: + if name in REDIS_COLLECTIONS: + outputs.append(self.collections[name](redis_keys, redis_offsets)) + else: + outputs.append(self.collections[name](base_keys, base_offsets)) + return tuple(outputs) + + +@torch.no_grad() +def fill_stable_random_weights(model: TrainingSparseModel) -> None: + table_ordinal = 0 + for collection in model.collections.values(): + for table in collection.embedding_configs(): + generator = torch.Generator(device="cpu") + generator.manual_seed(RANDOM_SEED + table_ordinal) + values = torch.rand( + (ROWS_PER_TABLE, EMBEDDING_DIM), + generator=generator, + dtype=torch.float32, + ) + collection.embeddings[table.name].weight.copy_( + values.to(collection.embeddings[table.name].weight.device) + ) + table_ordinal += 1 + + +def checkpoint_weights( + checkpoint_path: Path, collection_names: tuple[str, ...] +) -> dict[str, list[torch.Tensor]]: + state = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + return { + name: [ + state[f"collections.{name}.embeddings.{table.name}.weight"] + for table in embedding_configs(name) + ] + for name in collection_names + } + + +def make_dynamicemb_copy( + collection_name: str, + table_weights: list[torch.Tensor], + device: torch.device, + *, + retain_evictions: bool, +) -> BatchedDynamicEmbeddingTablesV2: + names = [table.name for table in embedding_configs(collection_name)] + options = [ + DynamicEmbTableOptions( + embedding_dtype=torch.float32, + dim=EMBEDDING_DIM, + init_capacity=ROWS_PER_TABLE, + max_capacity=ROWS_PER_TABLE, + bucket_capacity=ROWS_PER_TABLE, + local_hbm_for_values=4 << 20, + device_id=device.index, + training=False, + caching=False, + score_strategy=DynamicEmbScoreStrategy.CUSTOMIZED, + evicted_item_mode=( + EvictedItemMode.RETAIN_KEY + if retain_evictions + else EvictedItemMode.DISCARD + ), + index_type=torch.int64, + ) + for _ in names + ] + tables = BatchedDynamicEmbeddingTablesV2( + table_options=options, + table_names=names, + feature_table_map=[0, 1], + pooling_mode=DynamicEmbPoolingMode.NONE, + output_dtype=torch.float32, + device=device, + optimizer=EmbOptimType.NONE, + ) + + keys = torch.arange(ROWS_PER_TABLE, dtype=torch.int64, device=device).repeat(2) + table_ids = torch.arange(2, dtype=torch.int64, device=device).repeat_interleave( + ROWS_PER_TABLE + ) + values = torch.cat(table_weights).to(device) + scores = torch.ones(keys.numel(), dtype=torch.uint64, device=device) + tables.tables.insert(keys, table_ids, values, scores=scores) + return tables + + +def dump_training_artifacts( + model: TrainingSparseModel, + checkpoint_dir: Path, + collection_names: tuple[str, ...], + device: torch.device, +) -> tuple[Path, dict[str, BatchedDynamicEmbeddingTablesV2]]: + checkpoint_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = checkpoint_dir / "model.pt" + torch.save(model.state_dict(), checkpoint_path) + weights = checkpoint_weights(checkpoint_path, collection_names) + + dump_dir = checkpoint_dir / "dynamicemb" + dump_dir.mkdir(parents=True, exist_ok=True) + delta_producers: dict[str, BatchedDynamicEmbeddingTablesV2] = {} + for name in DUMP_COLLECTIONS: + if name not in collection_names: + continue + tables = make_dynamicemb_copy( + name, + weights[name], + device, + retain_evictions=name in REDIS_COLLECTIONS, + ) + tables.dump(str(dump_dir), optim=False, counter=False) + if name in REDIS_COLLECTIONS: + delta_producers[name] = tables + return checkpoint_path, delta_producers + + +def redis_parameter_server( + address: str, namespace_id: int +) -> NVEParameterServer: + parameter_server = NVEParameterServer( + num_embeddings=0, + embedding_size=EMBEDDING_DIM, + data_type=torch.float32, + ps_type=nve.Redis, + extra_params={ + "plugin": {"address": address, "single_node": True}, + "table": { + "num_partitions": 0, + "string_namespace_id": namespace_id, + }, + }, + ) + parameter_server.clear() + return parameter_server + + +def hierarchical_config( + indexer_type: EmbeddingCollectionIndexerType, + parameter_server: NVEParameterServer, + **kwargs: Any, +) -> InferenceEmbeddingCollectionConfig: + return InferenceEmbeddingCollectionConfig( + indexer_type=indexer_type, + nve_layer_type="hierarchical", + gpu_cache_size=GPU_CACHE_SIZE, + host_cache_size=HOST_CACHE_SIZE, + parameter_server=parameter_server, + **kwargs, + ) + + +def collection_configs( + *, enable_redis: bool, redis_address: str +) -> dict[str, InferenceEmbeddingCollectionConfig]: + configs = { + FUSED_IDENTITY_GPU: InferenceEmbeddingCollectionConfig( + indexer_type=EmbeddingCollectionIndexerType.FUSED_IDENTITY, + nve_layer_type="gpu", + ), + LINEAR_HASH_UVM: InferenceEmbeddingCollectionConfig( + indexer_type=EmbeddingCollectionIndexerType.LINEAR_HASH_MAP, + nve_layer_type="linear_uvm", + bucket_capacity=ROWS_PER_TABLE, + gpu_cache_size=GPU_CACHE_SIZE, + ), + } + if enable_redis: + configs.update( + { + LINEAR_HASH_REDIS: hierarchical_config( + EmbeddingCollectionIndexerType.LINEAR_HASH_MAP, + redis_parameter_server(redis_address, 2026090301), + bucket_capacity=ROWS_PER_TABLE, + ), + BIT_CONCAT_REDIS: hierarchical_config( + EmbeddingCollectionIndexerType.BIT_CONCAT, + redis_parameter_server(redis_address, 2026090302), + bit_concat=BitConcatConfig( + table_id_bits=8, feature_id_bits=55 + ), + ), + } + ) + return configs + + +@torch.no_grad() +def construct_inference_model( + training_model: TrainingSparseModel, + configs: dict[str, InferenceEmbeddingCollectionConfig], + checkpoint_path: Path, + dump_dir: Path, +) -> ExportableSparseExample: + collection_names = tuple(configs) + by_table = { + table.name: configs[name] + for name in collection_names + for table in embedding_configs(name) + } + apply_inference_embedding_collection( + training_model, + embedding_collection_configs=by_table, + trained_emb_table_sizes={ + table.name: ROWS_PER_TABLE + for name in collection_names + for table in embedding_configs(name) + }, + ) + + weights = checkpoint_weights(checkpoint_path, collection_names) + for name in collection_names: + collection = training_model.collections[name] + if name in DENSE_COLLECTIONS: + collection.load_from_embedding_table(torch.cat(weights[name])) + else: + collection.load_from_dynamicemb_file(str(dump_dir)) + return ExportableSparseExample(training_model, collection_names).eval() + + +def full_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + keys = torch.arange(ROWS_PER_TABLE, dtype=torch.int64, device=device).repeat(2) + offsets = torch.tensor( + [0, ROWS_PER_TABLE, 2 * ROWS_PER_TABLE], + dtype=torch.int64, + device=device, + ) + return keys, offsets, keys.clone(), offsets.clone() + + +def post_update_inputs( + device: torch.device, round_id: int +) -> tuple[torch.Tensor, ...]: + new_feature_id = ROWS_PER_TABLE + round_id - 1 + return ( + torch.tensor([1, 1], dtype=torch.int64, device=device), + torch.tensor([0, 1, 2], dtype=torch.int64, device=device), + torch.tensor( + [1, new_feature_id, 1, new_feature_id], + dtype=torch.int64, + device=device, + ), + torch.tensor([0, 2, 4], dtype=torch.int64, device=device), + ) + + +def snapshot_zero_expected( + weights: dict[str, list[torch.Tensor]], model: ExportableSparseExample +) -> list[torch.Tensor]: + expected = [] + for name in model.collection_names: + reference = torch.cat(weights[name]).clone() + failed = model.collections[name].failed_build_rows + for table_id, feature_id in failed: + reference[table_id * ROWS_PER_TABLE + feature_id].zero_() + expected.append(reference) + return expected + + +@torch.no_grad() +def produce_delta( + collection_name: str, + producer: BatchedDynamicEmbeddingTablesV2, + round_id: int, +) -> tuple[Any, dict[tuple[int, int], torch.Tensor]]: + ordinal = REDIS_COLLECTIONS.index(collection_name) + existing_keys = torch.tensor( + [1, 1], + dtype=torch.int64, + device=torch.device("cuda", producer.device_id), + ) + new_feature_id = ROWS_PER_TABLE + round_id - 1 + new_keys = torch.full_like(existing_keys, new_feature_id) + table_ids = torch.tensor( + [0, 1], dtype=torch.int64, device=existing_keys.device + ) + values = torch.stack( + [ + torch.full( + (EMBEDDING_DIM,), + 10.0 * round_id + ordinal + row / 10.0, + dtype=torch.float32, + device=existing_keys.device, + ) + for row in range(4) + ] + ) + # Protect the updated rows before forcing one eviction from each full table. + producer.tables.insert( + existing_keys, + table_ids, + values[::2], + scores=torch.full( + (2,), + 2 * round_id + 1, + dtype=torch.uint64, + device=existing_keys.device, + ), + ) + producer.tables.insert( + new_keys, + table_ids, + values[1::2], + scores=torch.full( + (2,), 2 * round_id, dtype=torch.uint64, device=existing_keys.device + ), + ) + thresholds = {name: 2 * round_id for name in producer.table_names} + delta = producer.incremental_dump(thresholds) + keys = torch.stack((existing_keys, new_keys), dim=1).flatten() + expanded_table_ids = torch.stack((table_ids, table_ids), dim=1).flatten() + expected = { + (int(table_id), int(feature_id)): value.detach().cpu() + for table_id, feature_id, value in zip( + expanded_table_ids.cpu().tolist(), keys.cpu().tolist(), values + ) + } + return delta, expected + + +def expected_after_update( + weights: dict[str, list[torch.Tensor]], + collection_names: tuple[str, ...], + changed: dict[str, dict[tuple[int, int], torch.Tensor]], + round_id: int, +) -> list[torch.Tensor]: + new_feature_id = ROWS_PER_TABLE + round_id - 1 + result = [] + for name in collection_names: + if name not in REDIS_COLLECTIONS: + result.append(torch.stack([weights[name][0][1], weights[name][1][1]])) + continue + rows = [] + for table_id in range(2): + existing = changed[name][(table_id, 1)] + rows.append(existing) + rows.append( + changed[name].get( + (table_id, new_feature_id), torch.zeros_like(existing) + ) + ) + result.append(torch.stack(rows)) + return result + + +def outputs_match( + actual: tuple[torch.Tensor, ...] | list[torch.Tensor], + expected: list[torch.Tensor], + stage: str, +) -> bool: + if len(actual) != len(expected): + print( + f"RESULT ERROR stage={stage} outputs={len(actual)} " + f"expected_outputs={len(expected)}", + flush=True, + ) + return False + for index, (output, reference) in enumerate(zip(actual, expected)): + output = output.detach().cpu() + if output.shape != reference.shape or not torch.allclose( + output, reference, rtol=1.3e-6, atol=1e-5 + ): + print( + f"RESULT ERROR stage={stage} output={index} " + f"actual_sum={float(output.sum())} " + f"expected_sum={float(reference.sum())}", + flush=True, + ) + return False + return True + + +def verify_state( + runtime: Any, + inputs: tuple[torch.Tensor, ...], + expected: list[torch.Tensor], + stage: str, +) -> bool: + for _ in range(REQUESTS_PER_VERIFIED_ROUND): + outputs = ( + runtime.run(list(inputs)) + if hasattr(runtime, "run") + else runtime(*inputs) + ) + if not outputs_match(outputs, expected, stage): + return False + return True + + +def run_four_round_update_workflow( + *, + model: ExportableSparseExample, + aoti_runtime: Any, + initial_inputs: tuple[torch.Tensor, ...], + expected_zero: list[torch.Tensor], + weights: dict[str, list[torch.Tensor]], + names: tuple[str, ...], + producers: dict[str, BatchedDynamicEmbeddingTablesV2], + coordinator_connection: Connection, + subscribers: tuple[tuple[str, Any], ...], + inference_stream: torch.cuda.Stream, + device: torch.device, + cpp_process: subprocess.Popen[str] | None, +) -> bool: + round_one_done = Event() + first_update_staged = Event() + round_two_done = Event() + round_three_release = Event() + round_three_started = Event() + second_update_staged = Event() + second_update_done = Event() + state: dict[str, Any] = {} + + def run_inference() -> bool: + inference_ok = True + with torch.cuda.stream(inference_stream), torch.inference_mode(): + for _ in range(REQUESTS_PER_VERIFIED_ROUND): + eager_outputs = model(*initial_inputs) + aoti_outputs = aoti_runtime.run(list(initial_inputs)) + round_ok = outputs_match( + eager_outputs, expected_zero, "python_eager_round_1" + ) and outputs_match( + aoti_outputs, expected_zero, "python_aoti_round_1" + ) + if not round_ok: + inference_ok = False + break + torch.cuda.synchronize(device) + round_one_done.set() + + first_update_staged.wait() + + first_inputs = post_update_inputs(device, round_id=1) + for _ in range(REQUESTS_PER_VERIFIED_ROUND): + eager_outputs = model(*first_inputs) + aoti_outputs = aoti_runtime.run(list(first_inputs)) + round_ok = outputs_match( + eager_outputs, + state["expected_one"], + "python_eager_round_2", + ) and outputs_match( + aoti_outputs, + state["expected_one"], + "python_aoti_round_2", + ) + if not round_ok: + inference_ok = False + break + inference_ok = linear_evictions_miss( + model, + model.collection_names.index(LINEAR_HASH_REDIS), + state["linear_eviction_query"], + "python_eager_round_2_evictions", + ) and inference_ok + inference_ok = linear_evictions_miss( + aoti_runtime, + model.collection_names.index(LINEAR_HASH_REDIS), + state["linear_eviction_query"], + "python_aoti_round_2_evictions", + ) and inference_ok + torch.cuda.synchronize(device) + round_two_done.set() + + round_three_release.wait() + + second_inputs = post_update_inputs(device, round_id=2) + for iteration in range(ROUND_THREE_MIN_INFERENCES): + model(*second_inputs) + aoti_runtime.run(list(second_inputs)) + if iteration == 0: + round_three_started.set() + while not second_update_staged.is_set(): + model(*second_inputs) + aoti_runtime.run(list(second_inputs)) + for _ in range(ROUND_THREE_POST_STAGE_INFERENCES): + model(*second_inputs) + aoti_runtime.run(list(second_inputs)) + torch.cuda.synchronize(device) + + second_update_done.wait() + + for _ in range(REQUESTS_PER_VERIFIED_ROUND): + eager_outputs = model(*second_inputs) + aoti_outputs = aoti_runtime.run(list(second_inputs)) + round_ok = outputs_match( + eager_outputs, + state["expected_two"], + "python_eager_round_4", + ) and outputs_match( + aoti_outputs, + state["expected_two"], + "python_aoti_round_4", + ) + if not round_ok: + inference_ok = False + break + torch.cuda.synchronize(device) + return inference_ok + + def run_update() -> bool: + update_ok = True + round_one_done.wait() + + with torch.cuda.device(device): + first_changed = {} + first_updates = [] + for name in REDIS_COLLECTIONS: + delta, first_changed[name] = produce_delta( + name, producers[name], round_id=1 + ) + if name == LINEAR_HASH_REDIS: + state["linear_eviction_query"] = linear_eviction_inputs( + delta, device + ) + first_updates.append( + coordinator_delta( + coordinator_connection, f"collections.{name}", delta + ) + ) + for update in first_updates: + for _, subscriber in subscribers: + subscriber.apply_incremental_load( + update.to_json(), inference_stream.cuda_stream + ) + state["expected_one"] = expected_after_update( + weights, names, first_changed, round_id=1 + ) + first_update_staged.set() + + for update in first_updates: + for subscriber_id, subscriber in subscribers: + coordinator_acknowledge( + coordinator_connection, + subscriber_id, + EmbeddingCollectionUpdateAck.from_json( + subscriber.wait_for_retirement( + update.collection_id, update.snapshot_id + ) + ), + ) + if cpp_process is not None: + acknowledgements, round_ok = run_incremental_cpp_round( + cpp_process, first_updates, state["expected_one"] + ) + update_ok = round_ok and update_ok + for acknowledgement in acknowledgements: + coordinator_acknowledge( + coordinator_connection, "cpp_aoti", acknowledgement + ) + + round_two_done.wait() + round_three_release.set() + round_three_started.wait() + + second_changed = { + LINEAR_HASH_REDIS: { + (table_id, 1): first_changed[LINEAR_HASH_REDIS][ + (table_id, 1) + ] + for table_id in range(2) + } + } + second_updates = [] + delta, second_changed[BIT_CONCAT_REDIS] = produce_delta( + BIT_CONCAT_REDIS, producers[BIT_CONCAT_REDIS], round_id=2 + ) + second_updates.append( + coordinator_delta( + coordinator_connection, + f"collections.{BIT_CONCAT_REDIS}", + delta, + ) + ) + for update in second_updates: + for _, subscriber in subscribers: + subscriber.apply_incremental_load( + update.to_json(), inference_stream.cuda_stream + ) + state["expected_two"] = expected_after_update( + weights, names, second_changed, round_id=2 + ) + second_update_staged.set() + + for update in second_updates: + for subscriber_id, subscriber in subscribers: + coordinator_acknowledge( + coordinator_connection, + subscriber_id, + EmbeddingCollectionUpdateAck.from_json( + subscriber.wait_for_retirement( + update.collection_id, update.snapshot_id + ) + ), + ) + if cpp_process is not None: + acknowledgements, round_ok = run_incremental_cpp_round( + cpp_process, second_updates, state["expected_two"] + ) + update_ok = round_ok and update_ok + for acknowledgement in acknowledgements: + coordinator_acknowledge( + coordinator_connection, "cpp_aoti", acknowledgement + ) + second_update_done.set() + return update_ok + + def run_inference_worker() -> bool: + try: + return run_inference() + except Exception as error: + print(f"RESULT ERROR python inference worker: {error}", flush=True) + return False + finally: + round_one_done.set() + round_two_done.set() + round_three_started.set() + + def run_update_worker() -> bool: + try: + return run_update() + except Exception as error: + print(f"RESULT ERROR python update worker: {error}", flush=True) + return False + finally: + first_update_staged.set() + round_three_release.set() + second_update_staged.set() + second_update_done.set() + + with ThreadPoolExecutor(max_workers=2) as executor: + inference_task = executor.submit(run_inference_worker) + update_task = executor.submit(run_update_worker) + inference_ok = inference_task.result() + update_ok = update_task.result() + return inference_ok and update_ok + + +def linear_eviction_inputs( + delta: Any, device: torch.device +) -> tuple[torch.Tensor, ...]: + keys = [] + offsets = [0] + for evicted in delta.evicted_keys: + assert evicted is not None and evicted.numel() > 0 + keys.append(evicted.to(device=device, dtype=torch.int64)) + offsets.append(offsets[-1] + evicted.numel()) + return ( + torch.tensor([1, 1], dtype=torch.int64, device=device), + torch.tensor([0, 1, 2], dtype=torch.int64, device=device), + torch.cat(keys), + torch.tensor(offsets, dtype=torch.int64, device=device), + ) + + +def linear_evictions_miss( + runtime: Any, + output_index: int, + inputs: tuple[torch.Tensor, ...], + stage: str, +) -> bool: + outputs = ( + runtime.run(list(inputs)) + if hasattr(runtime, "run") + else runtime(*inputs) + ) + output = outputs[output_index] + return outputs_match( + [output], [torch.zeros_like(output, device="cpu")], stage + ) + + +def live_layer_bindings( + model: ExportableSparseExample, + directory: EmbeddingCollectionIndexerDirectory, +) -> dict[str, Any]: + return { + binding.nve_layer_module_path: model.get_submodule( + binding.nve_layer_module_path + ).emb_layer + for binding in directory.bindings.values() + } + + +def loaded_layer_bindings(layers: list[Any]) -> dict[str, Any]: + return {layer._export_module_path: layer.emb_layer for layer in layers} + + +def start_update_coordinator( + package_dir: Path, + update_dir: Path, + subscriber_ids: set[str], +) -> tuple[subprocess.Popen[Any], Connection]: + socket_path = update_dir / "coordinator.sock" + update_dir.mkdir(parents=True, exist_ok=True) + socket_path.unlink(missing_ok=True) + command = [ + sys.executable, + str(Path(__file__).with_name("update_coordinator_main.py")), + "--package-dir", + str(package_dir), + "--update-dir", + str(update_dir), + "--socket", + str(socket_path), + ] + for subscriber_id in sorted(subscriber_ids): + command.extend(("--subscriber", subscriber_id)) + process = subprocess.Popen(command) + deadline = time.monotonic() + 30 + while True: + if process.poll() is not None: + raise RuntimeError("update coordinator exited during startup") + try: + connection = Client( + str(socket_path), family="AF_UNIX", authkey=COORDINATOR_AUTHKEY + ) + return process, connection + except (FileNotFoundError, ConnectionRefusedError): + if time.monotonic() >= deadline: + process.terminate() + raise RuntimeError("update coordinator did not become ready") + time.sleep(0.05) + + +def coordinator_delta( + connection: Connection, collection_id: str, delta: DeltaDumpResult +) -> EmbeddingCollectionUpdate: + buffer = io.BytesIO() + torch.save( + { + "table_names": delta.table_names, + "keys": [tensor.detach().cpu() for tensor in delta.keys], + "values": [tensor.detach().cpu() for tensor in delta.values], + "evicted_keys": [ + None if tensor is None else tensor.detach().cpu() + for tensor in delta.evicted_keys + ], + }, + buffer, + ) + connection.send( + { + "op": "delta", + "collection_id": collection_id, + "delta_payload": buffer.getvalue(), + } + ) + return EmbeddingCollectionUpdate.from_json(connection.recv()) + + +def coordinator_acknowledge( + connection: Connection, + subscriber_id: str, + acknowledgement: EmbeddingCollectionUpdateAck, +) -> None: + connection.send( + { + "op": "ack", + "subscriber_id": subscriber_id, + "ack": acknowledgement.to_json(), + } + ) + + +def run_cpp_replayer( + executable: Path, + package_dir: Path, + expected: list[torch.Tensor], +) -> bool: + command = [ + str(executable), + "--package-dir", + str(package_dir), + "--expected-sums", + ",".join(str(float(value.sum())) for value in expected), + ] + completed = subprocess.run(command, capture_output=True, text=True) + print(completed.stdout, end="") + print(completed.stderr, end="") + return completed.returncode == 0 + + +def start_incremental_cpp_replayer( + executable: Path, + package_dir: Path, + expected: list[torch.Tensor], +) -> tuple[subprocess.Popen[str], bool]: + process = subprocess.Popen( + [ + str(executable), + "--package-dir", + str(package_dir), + "--expected-sums", + ",".join(str(float(value.sum())) for value in expected), + "--wait-for-updates", + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + if process.stdout is None: + raise RuntimeError("C++ replayer stdout is unavailable") + verification_ok = True + while True: + line = process.stdout.readline() + if not line: + return_code = process.wait() + raise RuntimeError( + f"C++ replayer exited before its ready signal ({return_code})" + ) + print(line, end="") + if line.startswith("RESULT ERROR"): + verification_ok = False + if line.rstrip() == "READY incremental updates": + return process, verification_ok + + +def run_incremental_cpp_round( + process: subprocess.Popen[str], + updates: list[EmbeddingCollectionUpdate], + expected: list[torch.Tensor], +) -> tuple[list[EmbeddingCollectionUpdateAck], bool]: + if process.stdin is None or process.stdout is None: + print("RESULT ERROR C++ replayer pipes are unavailable", flush=True) + return [], False + try: + for update in updates: + process.stdin.write(update.to_json() + "\n") + process.stdin.write( + "RUN " + + ",".join(str(float(value.sum())) for value in expected) + + "\n" + ) + process.stdin.flush() + except (BrokenPipeError, OSError) as error: + print(f"RESULT ERROR C++ replayer input: {error}", flush=True) + return [], False + + acknowledgements = [] + verification_ok = True + while True: + line = process.stdout.readline() + if not line: + return_code = process.wait() + print( + f"RESULT ERROR C++ replayer exited with code {return_code}", + flush=True, + ) + return acknowledgements, False + print(line, end="") + if line.startswith("RESULT ERROR"): + verification_ok = False + if line.startswith("ACK "): + acknowledgements.append( + EmbeddingCollectionUpdateAck.from_json( + line.removeprefix("ACK ") + ) + ) + if line.rstrip() == "READY incremental updates": + return acknowledgements, verification_ok + + +def stop_incremental_cpp_replayer(process: subprocess.Popen[str]) -> bool: + if process.stdin is None or process.stdout is None: + print("RESULT ERROR C++ replayer pipes are unavailable", flush=True) + return False + try: + process.stdin.write("STOP\n") + process.stdin.flush() + except (BrokenPipeError, OSError) as error: + print(f"RESULT ERROR C++ replayer input: {error}", flush=True) + print(process.stdout.read(), end="") + return_code = process.wait() + process.stdin.close() + process.stdout.close() + if return_code != 0: + print( + f"RESULT ERROR C++ replayer exited with code {return_code}", + flush=True, + ) + return False + return True + + +def run_workflow(args: argparse.Namespace) -> bool: + owns_process_group = not dist.is_initialized() + if owns_process_group: + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29500") + dist.init_process_group("nccl", rank=0, world_size=1) + + cpp_process: subprocess.Popen[str] | None = None + workflow_ok = True + try: + device = torch.device("cuda", torch.cuda.current_device()) + inference_stream = torch.cuda.default_stream(device) + generation = imported_nve_generation() + enable_redis = generation >= (26, 6) and not args.disable_redis_incremental + names = BASE_COLLECTIONS + (REDIS_COLLECTIONS if enable_redis else ()) + + work_dir = Path(args.work_dir).resolve() + checkpoint_dir = work_dir / "checkpoint" + package_dir = work_dir / "package" + update_dir = work_dir / "updates" + + training_model = TrainingSparseModel(names, device) + fill_stable_random_weights(training_model) + checkpoint_path, producers = dump_training_artifacts( + training_model, checkpoint_dir, names, device + ) + weights = checkpoint_weights(checkpoint_path, names) + model = construct_inference_model( + training_model, + collection_configs( + enable_redis=enable_redis, redis_address=args.redis_address + ), + checkpoint_path, + checkpoint_dir / "dynamicemb", + ) + + eager_directory = embedding_collection_indexers_from_model(model) + inputs = full_inputs(device) + expected_zero = snapshot_zero_expected(weights, model) + + key_dimension = torch.export.Dim("base_key_count", min=2, max=514) + redis_dimension = torch.export.Dim("redis_key_count", min=2, max=514) + export_embedding_collection_aot( + model, + inputs, + package_dir, + dynamic_shapes=( + {0: key_dimension}, + None, + {0: redis_dimension}, + None, + ), + ) + aoti_runtime, nve_layers, aoti_directory = load_embedding_collection_aot( + package_dir, device + ) + + cpp_replayer = Path(args.cpp_replayer).resolve() if args.cpp_replayer else None + if cpp_replayer is not None: + if enable_redis: + cpp_process, cpp_ok = start_incremental_cpp_replayer( + cpp_replayer, package_dir, expected_zero + ) + workflow_ok = cpp_ok and workflow_ok + else: + workflow_ok = ( + run_cpp_replayer(cpp_replayer, package_dir, expected_zero) + and workflow_ok + ) + + if not enable_redis: + with torch.inference_mode(): + workflow_ok = ( + verify_state( + model, inputs, expected_zero, "python_eager_initial" + ) + and workflow_ok + ) + workflow_ok = ( + verify_state( + aoti_runtime, + inputs, + expected_zero, + "python_aoti_initial", + ) + and workflow_ok + ) + if workflow_ok: + print(f"verified {len(names)} embedding-collection combinations") + else: + print("RESULT ERROR exportable embedding verification failed") + return workflow_ok + + from dynamicemb.exportable_embedding import ( + EmbeddingCollectionUpdateSubscriber, + ) + + subscriber_ids = {"eager", "python_aoti"} + if cpp_process is not None: + subscriber_ids.add("cpp_aoti") + eager_subscriber = EmbeddingCollectionUpdateSubscriber( + str(package_dir), + eager_directory, + live_layer_bindings(model, eager_directory), + eager_directory.device_index, + ) + aoti_subscriber = EmbeddingCollectionUpdateSubscriber( + str(package_dir), + aoti_directory, + loaded_layer_bindings(nve_layers), + aoti_directory.device_index, + ) + + coordinator_process, coordinator_connection = start_update_coordinator( + package_dir, update_dir, subscriber_ids + ) + try: + workflow_ok = run_four_round_update_workflow( + model=model, + aoti_runtime=aoti_runtime, + initial_inputs=inputs, + expected_zero=expected_zero, + weights=weights, + names=names, + producers=producers, + coordinator_connection=coordinator_connection, + subscribers=( + ("eager", eager_subscriber), + ("python_aoti", aoti_subscriber), + ), + inference_stream=inference_stream, + device=device, + cpp_process=cpp_process, + ) and workflow_ok + if cpp_process is not None: + workflow_ok = ( + stop_incremental_cpp_replayer(cpp_process) and workflow_ok + ) + cpp_process = None + finally: + try: + if coordinator_process.poll() is None: + try: + coordinator_connection.send({"op": "stop"}) + except (BrokenPipeError, EOFError): + pass + coordinator_connection.close() + finally: + coordinator_process.wait(timeout=30) + + if workflow_ok: + print( + f"verified {len(names)} embedding-collection combinations " + "with paused and concurrent Redis incremental-load rounds" + ) + else: + print("RESULT ERROR exportable embedding verification failed") + return workflow_ok + finally: + if cpp_process is not None and cpp_process.poll() is None: + cpp_process.terminate() + cpp_process.wait(timeout=30) + if owns_process_group: + dist.destroy_process_group() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--work-dir", required=True) + parser.add_argument("--redis-address", default="127.0.0.1:6379") + parser.add_argument("--disable-redis-incremental", action="store_true") + parser.add_argument("--cpp-replayer") + return parser.parse_args() + + +if __name__ == "__main__": + sys.exit(0 if run_workflow(parse_args()) else 1) diff --git a/corelib/dynamicemb/example/exportable_embedding/run_example.sh b/corelib/dynamicemb/example/exportable_embedding/run_example.sh new file mode 100755 index 000000000..32b925990 --- /dev/null +++ b/corelib/dynamicemb/example/exportable_embedding/run_example.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +example_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +dynamicemb_root="$(cd "${example_dir}/../.." && pwd)" +nve_version="${NVE_VERSION:-}" + +if [[ "${nve_version}" == "26.05" ]]; then + nve_root="/opt/nve/26.05/python" + dynamicemb_python_path="/opt/dynamicemb/26.05/python:${dynamicemb_root}" +else + nve_root="/opt/nve/default/python" + dynamicemb_python_path="${dynamicemb_root}" +fi + +export DYNAMICEMB_OPS_LIB_DIR="${dynamicemb_root}/torch_binding_build" +export PYTHONPATH="${nve_root}:${dynamicemb_python_path}:${PYTHONPATH:-}" + +redis_disabled=0 +for argument in "$@"; do + if [[ "${argument}" == "--disable-redis-incremental" ]]; then + redis_disabled=1 + fi +done + +redis_pid="" +redis_runtime="" +cleanup_redis() { + if [[ -n "${redis_pid}" ]]; then + kill "${redis_pid}" 2>/dev/null || : + wait "${redis_pid}" 2>/dev/null || : + fi + if [[ -n "${redis_runtime}" && -d "${redis_runtime}" ]]; then + rm -r -- "${redis_runtime}" + fi +} +trap cleanup_redis EXIT + +if [[ "${nve_version}" != "26.05" \ + && "${redis_disabled}" == "0" \ + && "${START_LOCAL_REDIS:-1}" == "1" ]]; then + redis_port="${REDIS_PORT:-6379}" + redis_runtime="$(mktemp -d "${TMPDIR:-/tmp}/exportable-embedding-redis.XXXXXX")" + redis-server \ + --bind 127.0.0.1 \ + --protected-mode yes \ + --port "${redis_port}" \ + --save "" \ + --appendonly no \ + --dir "${redis_runtime}" \ + --logfile "${redis_runtime}/redis.log" & + redis_pid=$! + + redis_ready=0 + for _ in {1..100}; do + if redis-cli -h 127.0.0.1 -p "${redis_port}" ping 2>/dev/null \ + | grep -q '^PONG$'; then + redis_ready=1 + break + fi + if ! kill -0 "${redis_pid}" 2>/dev/null; then + break + fi + sleep 0.05 + done + if [[ "${redis_ready}" != "1" ]]; then + echo "local Redis failed to start" >&2 + exit 1 + fi + set -- "$@" --redis-address "127.0.0.1:${redis_port}" +fi + +python3 "${example_dir}/export_and_verify.py" "$@" diff --git a/corelib/dynamicemb/example/exportable_embedding/update_coordinator_main.py b/corelib/dynamicemb/example/exportable_embedding/update_coordinator_main.py new file mode 100644 index 000000000..0014ca0de --- /dev/null +++ b/corelib/dynamicemb/example/exportable_embedding/update_coordinator_main.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal local-IPC launcher for the example update coordinator.""" + +import argparse +import io +import os +from multiprocessing.connection import Listener + +import torch + + +torch.ops.load_library( + os.path.join(os.environ["DYNAMICEMB_OPS_LIB_DIR"], "inference_emb_ops.so") +) + +from dynamicemb import DeltaDumpResult # noqa: E402 +from dynamicemb.exportable_embedding import ( # noqa: E402 + EmbeddingCollectionUpdateAck, + EmbeddingCollectionUpdateCoordinator, +) + + +AUTHKEY = b"exportable-embedding-example" + + +def deserialize_delta(payload: bytes) -> DeltaDumpResult: + value = torch.load( + io.BytesIO(payload), map_location="cpu", weights_only=True + ) + return DeltaDumpResult( + table_names=value["table_names"], + keys=value["keys"], + values=value["values"], + evicted_keys=value["evicted_keys"], + ) + + +def serve(listener: Listener, coordinator: EmbeddingCollectionUpdateCoordinator) -> None: + with listener.accept() as connection: + while True: + message = connection.recv() + if message["op"] == "stop": + return + if message["op"] == "delta": + update = coordinator.apply_delta( + message["collection_id"], + deserialize_delta(message["delta_payload"]), + ) + connection.send(update.to_json()) + elif message["op"] == "ack": + coordinator.acknowledge( + message["subscriber_id"], + EmbeddingCollectionUpdateAck.from_json(message["ack"]), + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--package-dir", required=True) + parser.add_argument("--update-dir", required=True) + parser.add_argument("--socket", required=True) + parser.add_argument("--subscriber", action="append", default=[]) + args = parser.parse_args() + + coordinator = EmbeddingCollectionUpdateCoordinator.open( + package_dir=args.package_dir, + shared_update_dir=args.update_dir, + device=torch.device("cuda", torch.cuda.current_device()), + subscriber_ids=args.subscriber, + ) + with Listener(args.socket, family="AF_UNIX", authkey=AUTHKEY) as listener: + serve(listener, coordinator) + + +if __name__ == "__main__": + main() diff --git a/corelib/dynamicemb/include/dynamicemb/exportable_embedding/incremental_update.h b/corelib/dynamicemb/include/dynamicemb/exportable_embedding/incremental_update.h new file mode 100644 index 000000000..3ac57206c --- /dev/null +++ b/corelib/dynamicemb/include/dynamicemb/exportable_embedding/incremental_update.h @@ -0,0 +1,66 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include +#include +#include +#include + +#include "dynamicemb/exportable_embedding/indexer_directory.h" + +namespace nve { +class LayerDirectory; +template +class NVEmbedBinding; +} + +namespace dynamicemb::exportable_embedding { + +struct DYNAMICEMB_EXPORT EmbeddingCollectionUpdate { + std::string collection_id; + int64_t snapshot_id{0}; + std::vector cache_update_keys; + std::vector cache_update_values; + std::string indexer_snapshot_path; + + static EmbeddingCollectionUpdate from_json(const std::string& payload); + std::string to_json() const; +}; + +struct DYNAMICEMB_EXPORT EmbeddingCollectionUpdateAck { + std::string collection_id; + int64_t snapshot_id{0}; + + std::string to_json() const; +}; + +class DYNAMICEMB_EXPORT EmbeddingCollectionUpdateSubscriber { + public: + EmbeddingCollectionUpdateSubscriber( + const std::string& package_dir, + EmbeddingCollectionIndexerDirectory& indexers, + nve::LayerDirectory& nve_layers, + int device_index); + EmbeddingCollectionUpdateSubscriber( + const std::string& package_dir, + EmbeddingCollectionIndexerDirectory& indexers, + std::unordered_map< + std::string, + std::shared_ptr>> nve_layers, + int device_index); + ~EmbeddingCollectionUpdateSubscriber(); + + void apply_incremental_load(const EmbeddingCollectionUpdate& update, + cudaStream_t inference_stream); + EmbeddingCollectionUpdateAck wait_for_retirement( + const std::string& collection_id, int64_t snapshot_id); + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace dynamicemb::exportable_embedding diff --git a/corelib/dynamicemb/include/dynamicemb/exportable_embedding/indexer_directory.h b/corelib/dynamicemb/include/dynamicemb/exportable_embedding/indexer_directory.h new file mode 100644 index 000000000..f4e79fb2c --- /dev/null +++ b/corelib/dynamicemb/include/dynamicemb/exportable_embedding/indexer_directory.h @@ -0,0 +1,74 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "dynamicemb/exportable_embedding/indexer_snapshot.h" + +namespace dynamicemb::exportable_embedding { + +struct EmbeddingCollectionBinding { + std::vector table_names; + std::string indexer_type; + std::string indexer_module_path; + std::string nve_layer_module_path; + std::optional snapshot_path; + std::optional feature_id_bits; + std::optional marker_value; +}; + +class EmbeddingCollectionIndexerDirectory; +class EmbeddingCollectionUpdateSubscriber; + +class DYNAMICEMB_EXPORT EmbeddingCollectionIndexerDirectory { + public: + static std::unique_ptr create( + std::unordered_map bindings, + std::unordered_map markers, + std::unordered_map> snapshots, + int device_index); + static std::unique_ptr load( + const std::string& package_dir, int device_index); + + EmbeddingCollectionIndexerDirectory( + EmbeddingCollectionIndexerDirectory&&) noexcept; + EmbeddingCollectionIndexerDirectory& operator=( + EmbeddingCollectionIndexerDirectory&&) noexcept; + ~EmbeddingCollectionIndexerDirectory(); + + void bind(torch::inductor::AOTIModelPackageLoader& loader); + void wait_for_retirement(const std::string& collection_id, + int64_t snapshot_id); + const EmbeddingCollectionBinding& binding( + const std::string& collection_id) const; + const std::unordered_map& bindings() + const; + std::unordered_map marker_constants() const; + std::shared_ptr snapshot( + const std::string& collection_id) const; + int device_index() const; + + private: + friend class EmbeddingCollectionUpdateSubscriber; + struct Impl; + explicit EmbeddingCollectionIndexerDirectory(std::unique_ptr impl); + void apply_update(const std::string& collection_id, + const std::string& snapshot_path, int64_t snapshot_id, + const std::function& enqueue_cache_update); + std::unique_ptr impl_; +}; + +} // namespace dynamicemb::exportable_embedding diff --git a/corelib/dynamicemb/include/dynamicemb/exportable_embedding/indexer_snapshot.h b/corelib/dynamicemb/include/dynamicemb/exportable_embedding/indexer_snapshot.h new file mode 100644 index 000000000..2696bd633 --- /dev/null +++ b/corelib/dynamicemb/include/dynamicemb/exportable_embedding/indexer_snapshot.h @@ -0,0 +1,51 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include +#include + +#include + +#if defined(_WIN32) +#define DYNAMICEMB_EXPORT __declspec(dllexport) +#else +#define DYNAMICEMB_EXPORT __attribute__((visibility("default"))) +#endif + +namespace dynamicemb::exportable_embedding { + +enum class IndexerSnapshotKind : int64_t { + // Persisted in the snapshot header; keep aligned with indexer_snapshot.py. + LinearHashMap = 0, + FusedIdentity = 1, +}; + +struct IndexerSnapshot { + IndexerSnapshotKind kind; + at::Tensor table_storage; + at::Tensor table_bucket_offsets; + int64_t bucket_capacity{0}; + at::Tensor miss_storage_indices; + at::Tensor valid_bases; + at::Tensor reserved_sizes; + int64_t next_fused_key{0}; +}; + +DYNAMICEMB_EXPORT std::shared_ptr load_indexer_snapshot( + const std::string& path, int device_index); + +void register_indexer_snapshot(const at::Tensor& marker, + std::shared_ptr snapshot); +void unregister_indexer_snapshot(const at::Tensor& marker); +void stage_indexer_snapshot(const at::Tensor& marker, + std::shared_ptr snapshot, + int64_t snapshot_id); +void wait_for_indexer_snapshot_retirement(const at::Tensor& marker, + int64_t snapshot_id); +std::shared_ptr find_indexer_snapshot( + const at::Tensor& marker, int64_t marker_value); + +} // namespace dynamicemb::exportable_embedding diff --git a/corelib/dynamicemb/setup.py b/corelib/dynamicemb/setup.py index 185fa3407..c3747ba96 100644 --- a/corelib/dynamicemb/setup.py +++ b/corelib/dynamicemb/setup.py @@ -158,6 +158,15 @@ def get_extensions(): "lookup_torch_binding.cu", "get_table_range_torch_binding.cu", "expand_table_ids_torch_binding.cu", + "indexer_directory.cpp", + "indexer_snapshot.cpp", + "indexer_ops.cu", + "incremental_update.cpp", + "exportable_embedding_pybind.cpp", + "indexer_directory_pybind.cpp", + "indexer_snapshot_pybind.cpp", + "update_subscriber_pybind.cu", + "update_subscriber_unavailable_pybind.cpp", # Built separately into standalone fatbins (Lex + custom LTO-IR), # shipped as package_data; NOT linked into the .so. "evict_lrulfu.cu", @@ -297,7 +306,10 @@ def run(self): description="Plugin for Dynamic Embedding in TorchREC", packages=package, ext_modules=get_extensions(), - package_data={f"{library_name}.jit": ["*.fatbin"]}, + package_data={ + f"{library_name}.jit": ["*.fatbin"], + f"{library_name}.exportable_embedding": ["_C.so"], + }, license="BSD-3", keywords=[ "pytorch", diff --git a/corelib/dynamicemb/src/exportable_embedding/exportable_embedding_pybind.cpp b/corelib/dynamicemb/src/exportable_embedding/exportable_embedding_pybind.cpp new file mode 100644 index 000000000..8861f09d8 --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/exportable_embedding_pybind.cpp @@ -0,0 +1,20 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include + +namespace py = pybind11; + +namespace dynamicemb::exportable_embedding { + +void bind_indexer_snapshot(py::module_& module); +void bind_indexer_directory(py::module_& module); +void bind_update_subscriber(py::module_& module); + +} // namespace dynamicemb::exportable_embedding + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + dynamicemb::exportable_embedding::bind_indexer_snapshot(module); + dynamicemb::exportable_embedding::bind_indexer_directory(module); + dynamicemb::exportable_embedding::bind_update_subscriber(module); +} diff --git a/corelib/dynamicemb/src/exportable_embedding/incremental_update.cpp b/corelib/dynamicemb/src/exportable_embedding/incremental_update.cpp new file mode 100644 index 000000000..88dd87f1e --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/incremental_update.cpp @@ -0,0 +1,168 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "dynamicemb/exportable_embedding/incremental_update.h" + +#include +#include + +#include +#include +#include + +#include "python/pynve/torch_bindings/nve_loader.hpp" + +namespace dynamicemb::exportable_embedding { + +EmbeddingCollectionUpdate EmbeddingCollectionUpdate::from_json( + const std::string& payload) { + const auto value = nlohmann::json::parse(payload); + EmbeddingCollectionUpdate update; + update.collection_id = value.at("collection_id").get(); + update.snapshot_id = value.at("snapshot_id").get(); + update.cache_update_keys = + value.at("cache_update_keys").get>(); + update.cache_update_values = + value.at("cache_update_values").get>(); + if (!value.at("indexer_snapshot_path").is_null()) { + update.indexer_snapshot_path = + value.at("indexer_snapshot_path").get(); + } + return update; +} + +std::string EmbeddingCollectionUpdate::to_json() const { + nlohmann::json value{ + {"collection_id", collection_id}, + {"snapshot_id", snapshot_id}, + {"cache_update_keys", cache_update_keys}, + {"cache_update_values", cache_update_values}, + {"indexer_snapshot_path", + indexer_snapshot_path.empty() ? nlohmann::json(nullptr) + : nlohmann::json(indexer_snapshot_path)}, + }; + return value.dump(); +} + +std::string EmbeddingCollectionUpdateAck::to_json() const { + return nlohmann::json({{"collection_id", collection_id}, + {"snapshot_id", snapshot_id}}) + .dump(); +} + +struct EmbeddingCollectionUpdateSubscriber::Impl { + EmbeddingCollectionIndexerDirectory& indexers; + int device_index; + std::unordered_map>> + nve_layers; + std::unordered_map has_host_cache; + + void load_metadata(const std::string& package_dir, + nve::LayerDirectory* layer_directory) { + std::ifstream stream(package_dir + "/metadata.json"); + const auto metadata = nlohmann::json::parse(stream); + const auto& layers = metadata.is_array() ? metadata : metadata.at("layers"); + for (const auto& layer : layers) { + const std::string path = layer.at("module_path").get(); + has_host_cache[path] = layer.value("host_cache_size", uint64_t{0}) > 0; + if (layer_directory != nullptr) { + nve_layers[path] = + layer_directory->get_layer(layer.at("id").get()).binding; + } + } + } + + Impl(const std::string& package_dir, + EmbeddingCollectionIndexerDirectory& indexer_directory, + nve::LayerDirectory& layer_directory, int device) + : indexers(indexer_directory), + device_index(device) { + load_metadata(package_dir, &layer_directory); + } + + Impl(const std::string& package_dir, + EmbeddingCollectionIndexerDirectory& indexer_directory, + std::unordered_map< + std::string, + std::shared_ptr>> layer_bindings, + int device) + : indexers(indexer_directory), + device_index(device), + nve_layers(std::move(layer_bindings)) { + load_metadata(package_dir, nullptr); + } +}; + +EmbeddingCollectionUpdateSubscriber::EmbeddingCollectionUpdateSubscriber( + const std::string& package_dir, + EmbeddingCollectionIndexerDirectory& indexers, + nve::LayerDirectory& nve_layers, int device_index) + : impl_(std::make_unique(package_dir, indexers, nve_layers, + device_index)) {} + +EmbeddingCollectionUpdateSubscriber::EmbeddingCollectionUpdateSubscriber( + const std::string& package_dir, + EmbeddingCollectionIndexerDirectory& indexers, + std::unordered_map>> nve_layers, + int device_index) + : impl_(std::make_unique(package_dir, indexers, + std::move(nve_layers), device_index)) {} + +EmbeddingCollectionUpdateSubscriber::~EmbeddingCollectionUpdateSubscriber() = + default; + +void EmbeddingCollectionUpdateSubscriber::apply_incremental_load( + const EmbeddingCollectionUpdate& update, cudaStream_t inference_stream) { + const auto& binding = impl_->indexers.binding(update.collection_id); + const auto stream = c10::cuda::getStreamFromExternal( + inference_stream, impl_->device_index); + c10::cuda::CUDAStreamGuard stream_guard(stream); + auto keys = torch::tensor( + update.cache_update_keys, + torch::TensorOptions().dtype(torch::kInt64)); + impl_->indexers.apply_update( + update.collection_id, update.indexer_snapshot_path, update.snapshot_id, + [&] { + if (keys.numel() == 0) { + return; + } + const auto& layer = + impl_->nve_layers.at(binding.nve_layer_module_path); + const auto value_dtype = + layer->get_data_type() == nve::DataType_t::Float16 + ? torch::kFloat16 + : torch::kFloat32; + auto values = torch::tensor( + update.cache_update_values, + torch::TensorOptions().dtype(torch::kFloat32)) + .to(value_dtype) + .view({keys.numel(), layer->get_embedding_dim()}); + const uint64_t nve_stream = + reinterpret_cast(inference_stream); + const auto* key_data = keys.data_ptr(); + const auto* value_data = + reinterpret_cast(values.data_ptr()); + const auto row_size = + values.element_size() * layer->get_embedding_dim(); + for (int64_t index = 0; index < keys.numel(); ++index) { + const auto key = reinterpret_cast(key_data + index); + const auto value = reinterpret_cast( + value_data + index * row_size); + if (impl_->has_host_cache.at(binding.nve_layer_module_path)) { + layer->update(1, key, value, 1, nve_stream); + } + layer->update(1, key, value, 0, nve_stream); + } + }); +} + +EmbeddingCollectionUpdateAck +EmbeddingCollectionUpdateSubscriber::wait_for_retirement( + const std::string& collection_id, int64_t snapshot_id) { + impl_->indexers.wait_for_retirement(collection_id, snapshot_id); + return {collection_id, snapshot_id}; +} + +} // namespace dynamicemb::exportable_embedding diff --git a/corelib/dynamicemb/src/exportable_embedding/indexer_directory.cpp b/corelib/dynamicemb/src/exportable_embedding/indexer_directory.cpp new file mode 100644 index 000000000..58c7f3e0c --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/indexer_directory.cpp @@ -0,0 +1,226 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "dynamicemb/exportable_embedding/indexer_directory.h" + +#include + +#include +#include +#include +#include + +namespace dynamicemb::exportable_embedding { +namespace { + +std::string marker_fqn(const std::string& module_path) { + return module_path + ".marker_tensor"; +} + +bool has_replaceable_state(const std::string& indexer_type) { + return indexer_type == "linear_hash_map" || + indexer_type == "fused_identity"; +} + +std::string retirement_key(const std::string& collection_id, + int64_t snapshot_id) { + return collection_id + "\n" + std::to_string(snapshot_id); +} + +} // namespace + +struct EmbeddingCollectionIndexerDirectory::Impl { + int device_index{0}; + std::string package_dir; + std::unordered_map bindings; + std::unordered_map markers; + std::unordered_map> snapshots; + std::unordered_set pending_retirements; + bool owns_native_registrations{false}; +}; + +EmbeddingCollectionIndexerDirectory::EmbeddingCollectionIndexerDirectory( + std::unique_ptr impl) + : impl_(std::move(impl)) {} +EmbeddingCollectionIndexerDirectory::EmbeddingCollectionIndexerDirectory( + EmbeddingCollectionIndexerDirectory&&) noexcept = default; +EmbeddingCollectionIndexerDirectory& +EmbeddingCollectionIndexerDirectory::operator=( + EmbeddingCollectionIndexerDirectory&&) noexcept = default; + +EmbeddingCollectionIndexerDirectory::~EmbeddingCollectionIndexerDirectory() { + if (!impl_) { + return; + } + if (impl_->owns_native_registrations) { + for (const auto& item : impl_->markers) { + unregister_indexer_snapshot(item.second); + } + } +} + +std::unique_ptr +EmbeddingCollectionIndexerDirectory::create( + std::unordered_map bindings, + std::unordered_map markers, + std::unordered_map> snapshots, + int device_index) { + auto impl = std::make_unique(); + impl->device_index = device_index; + impl->bindings = std::move(bindings); + impl->markers = std::move(markers); + impl->snapshots = std::move(snapshots); + for (const auto& item : impl->snapshots) { + register_indexer_snapshot(impl->markers.at(item.first), item.second); + } + return std::unique_ptr( + new EmbeddingCollectionIndexerDirectory(std::move(impl))); +} + +std::unique_ptr +EmbeddingCollectionIndexerDirectory::load(const std::string& package_dir, + int device_index) { + auto impl = std::make_unique(); + impl->device_index = device_index; + impl->package_dir = package_dir; + const std::string root = package_dir + "/embedding_collection_indexers"; + std::ifstream manifest_stream(root + "/manifest.json"); + if (!manifest_stream) { + throw std::runtime_error("Cannot open embedding-collection indexer manifest"); + } + const auto document = nlohmann::json::parse(manifest_stream); + for (auto iterator = document.at("collections").begin(); + iterator != document.at("collections").end(); ++iterator) { + const auto& value = iterator.value(); + EmbeddingCollectionBinding binding; + binding.table_names = value.at("table_names").get>(); + binding.indexer_type = value.at("indexer_type").get(); + binding.indexer_module_path = + value.at("indexer_module_path").get(); + binding.nve_layer_module_path = + value.at("nve_layer_module_path").get(); + if (value.contains("snapshot_path") && + !value.at("snapshot_path").is_null()) { + binding.snapshot_path = value.at("snapshot_path").get(); + } + if (value.contains("feature_id_bits") && + !value.at("feature_id_bits").is_null()) { + binding.feature_id_bits = value.at("feature_id_bits").get(); + } + if (value.contains("marker_value") && + !value.at("marker_value").is_null()) { + binding.marker_value = value.at("marker_value").get(); + } + const std::string collection_id = iterator.key(); + impl->bindings.emplace(collection_id, binding); + if (!binding.snapshot_path.has_value()) { + continue; + } + auto marker = torch::tensor( + {binding.marker_value.value()}, + torch::TensorOptions().dtype(torch::kInt64).device( + torch::kCUDA, device_index)); + auto snapshot = load_indexer_snapshot( + root + "/" + binding.snapshot_path.value(), device_index); + register_indexer_snapshot(marker, snapshot); + impl->markers.emplace(collection_id, std::move(marker)); + impl->snapshots.emplace(collection_id, std::move(snapshot)); + } + impl->owns_native_registrations = true; + return std::unique_ptr( + new EmbeddingCollectionIndexerDirectory(std::move(impl))); +} + +void EmbeddingCollectionIndexerDirectory::bind( + torch::inductor::AOTIModelPackageLoader& loader) { + const auto names = loader.get_constant_fqns(); + const std::unordered_map available = [&] { + std::unordered_map result; + for (const auto& name : names) result.emplace(name, true); + return result; + }(); + for (const auto& item : marker_constants()) { + const std::string& fqn = item.first; + if (available.find(fqn) == available.end()) { + throw std::runtime_error("Indexer marker is not an AOTI constant: " + fqn); + } + std::unordered_map constants{{fqn, item.second}}; + loader.load_constants(constants, false, false, true); + loader.load_constants(constants, true, false, true); + } +} + +void EmbeddingCollectionIndexerDirectory::apply_update( + const std::string& collection_id, const std::string& snapshot_path, + int64_t snapshot_id, + const std::function& enqueue_cache_update) { + const bool has_replaceable_snapshot = + impl_->markers.find(collection_id) != impl_->markers.end() && + has_replaceable_state(impl_->bindings.at(collection_id).indexer_type); + if (!has_replaceable_snapshot) { + enqueue_cache_update(); + return; + } + + std::shared_ptr next; + if (!snapshot_path.empty()) { + next = load_indexer_snapshot(snapshot_path, impl_->device_index); + } + + enqueue_cache_update(); + if (!next) { + return; + } + + stage_indexer_snapshot(impl_->markers.at(collection_id), next, snapshot_id); + impl_->snapshots[collection_id] = std::move(next); + impl_->pending_retirements.insert( + retirement_key(collection_id, snapshot_id)); +} + +void EmbeddingCollectionIndexerDirectory::wait_for_retirement( + const std::string& collection_id, int64_t snapshot_id) { + const auto key = retirement_key(collection_id, snapshot_id); + if (impl_->pending_retirements.find(key) == + impl_->pending_retirements.end()) { + return; + } + wait_for_indexer_snapshot_retirement(impl_->markers.at(collection_id), + snapshot_id); + impl_->pending_retirements.erase(key); +} + +const EmbeddingCollectionBinding& +EmbeddingCollectionIndexerDirectory::binding( + const std::string& collection_id) const { + return impl_->bindings.at(collection_id); +} + +const std::unordered_map& +EmbeddingCollectionIndexerDirectory::bindings() const { + return impl_->bindings; +} + +std::unordered_map +EmbeddingCollectionIndexerDirectory::marker_constants() const { + std::unordered_map result; + for (const auto& item : impl_->markers) { + result.emplace(marker_fqn( + impl_->bindings.at(item.first).indexer_module_path), + item.second); + } + return result; +} + +std::shared_ptr +EmbeddingCollectionIndexerDirectory::snapshot( + const std::string& collection_id) const { + const auto iterator = impl_->snapshots.find(collection_id); + return iterator == impl_->snapshots.end() ? nullptr : iterator->second; +} + +int EmbeddingCollectionIndexerDirectory::device_index() const { + return impl_->device_index; +} + +} // namespace dynamicemb::exportable_embedding diff --git a/corelib/dynamicemb/src/exportable_embedding/indexer_directory_pybind.cpp b/corelib/dynamicemb/src/exportable_embedding/indexer_directory_pybind.cpp new file mode 100644 index 000000000..25a03e000 --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/indexer_directory_pybind.cpp @@ -0,0 +1,71 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include + +#include +#include + +#include "dynamicemb/exportable_embedding/indexer_directory.h" + +namespace py = pybind11; + +namespace dynamicemb::exportable_embedding { + +void bind_indexer_directory(py::module_& module) { + py::class_(module, + "EmbeddingCollectionBinding") + .def(py::init<>()) + .def_readwrite("table_names", &EmbeddingCollectionBinding::table_names) + .def_readwrite("indexer_type", &EmbeddingCollectionBinding::indexer_type) + .def_readwrite("indexer_module_path", + &EmbeddingCollectionBinding::indexer_module_path) + .def_readwrite("nve_layer_module_path", + &EmbeddingCollectionBinding::nve_layer_module_path) + .def_readwrite("snapshot_path", + &EmbeddingCollectionBinding::snapshot_path) + .def_readwrite("feature_id_bits", + &EmbeddingCollectionBinding::feature_id_bits) + .def_readwrite("marker_value", &EmbeddingCollectionBinding::marker_value); + + py::class_>( + module, "EmbeddingCollectionIndexerDirectory") + .def_static("create", &EmbeddingCollectionIndexerDirectory::create, + py::arg("bindings"), py::arg("markers"), + py::arg("snapshots"), py::arg("device_index")) + .def_static("load", &EmbeddingCollectionIndexerDirectory::load, + py::arg("package_dir"), py::arg("device_index")) + .def("bind_aoti", + [](const EmbeddingCollectionIndexerDirectory& directory, + const py::object& loader) { + const auto names = loader.attr("get_constant_fqns")() + .cast>(); + const std::unordered_set available(names.begin(), + names.end()); + for (const auto& item : directory.marker_constants()) { + if (available.find(item.first) == available.end()) { + throw std::runtime_error( + "Indexer marker is not an AOTI constant: " + item.first); + } + py::dict constants; + constants[py::str(item.first)] = item.second; + loader.attr("load_constants")(constants, false, false, true); + loader.attr("load_constants")(constants, true, false, true); + } + }) + .def("wait_for_retirement", + &EmbeddingCollectionIndexerDirectory::wait_for_retirement) + .def("snapshot", &EmbeddingCollectionIndexerDirectory::snapshot) + .def_property_readonly( + "bindings", + [](const EmbeddingCollectionIndexerDirectory& directory) { + return directory.bindings(); + }) + .def_property_readonly("device_index", + &EmbeddingCollectionIndexerDirectory::device_index); +} + +} // namespace dynamicemb::exportable_embedding diff --git a/corelib/dynamicemb/src/exportable_embedding/indexer_ops.cu b/corelib/dynamicemb/src/exportable_embedding/indexer_ops.cu new file mode 100644 index 000000000..9ffd0275a --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/indexer_ops.cu @@ -0,0 +1,63 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "dynamicemb/exportable_embedding/indexer_snapshot.h" + +#include + +#include + +#include "table_operation/table.cuh" + +namespace dynamicemb::exportable_embedding { + +at::Tensor embedding_collection_index_cuda(const at::Tensor& marker, + int64_t marker_value, + const at::Tensor& feature_ids, + const at::Tensor& table_ids) { + TORCH_CHECK(marker.is_cuda(), "indexer marker must be on CUDA"); + TORCH_CHECK(feature_ids.is_cuda() && table_ids.is_cuda(), + "embedding collection IDs must be on CUDA"); + auto snapshot = find_indexer_snapshot(marker, marker_value); + if (snapshot->kind == IndexerSnapshotKind::LinearHashMap) { + auto result = dyn_emb::table_lookup( + snapshot->table_storage, snapshot->table_bucket_offsets, + snapshot->bucket_capacity, feature_ids, table_ids, std::nullopt, + dyn_emb::ScorePolicyType::Const); + auto& fused_keys = std::get<0>(result); + auto& found = std::get<1>(result); + auto misses = at::index_select(snapshot->miss_storage_indices, 0, table_ids); + return at::where(found, fused_keys, misses); + } + if (snapshot->kind == IndexerSnapshotKind::FusedIdentity) { + return at::index_select(snapshot->valid_bases, 0, table_ids) + feature_ids; + } + TORCH_CHECK(false, "Unsupported embedding-collection indexer kind"); +} + +at::Tensor embedding_collection_index_meta(const at::Tensor& marker, + int64_t marker_value, + const at::Tensor& feature_ids, + const at::Tensor& table_ids) { + (void)marker; + (void)marker_value; + (void)table_ids; + return at::empty_like(feature_ids, feature_ids.options().dtype(at::kLong)); +} + +} // namespace dynamicemb::exportable_embedding + +TORCH_LIBRARY_FRAGMENT(INFERENCE_EMB, m) { + m.def("embedding_collection_index(Tensor marker, int marker_value, " + "Tensor feature_ids, Tensor table_ids) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(INFERENCE_EMB, CUDA, m) { + m.impl("embedding_collection_index", + &dynamicemb::exportable_embedding::embedding_collection_index_cuda); +} + +TORCH_LIBRARY_IMPL(INFERENCE_EMB, Meta, m) { + m.impl("embedding_collection_index", + &dynamicemb::exportable_embedding::embedding_collection_index_meta); +} diff --git a/corelib/dynamicemb/src/exportable_embedding/indexer_snapshot.cpp b/corelib/dynamicemb/src/exportable_embedding/indexer_snapshot.cpp new file mode 100644 index 000000000..fed7977c5 --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/indexer_snapshot.cpp @@ -0,0 +1,280 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "dynamicemb/exportable_embedding/indexer_snapshot.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace dynamicemb::exportable_embedding { +namespace { + +constexpr char kSnapshotMagic[8] = {'E', 'C', 'I', 'D', 'X', '0', '0', '1'}; +constexpr uint32_t kSchemaVersion = 1; + +#pragma pack(push, 1) +struct SnapshotHeader { + char magic[8]; + uint32_t version; + uint32_t kind; + int64_t bucket_capacity; + int64_t next_fused_key; + uint64_t table_storage_size; + uint64_t table_bucket_offsets_size; + uint64_t miss_storage_indices_size; + uint64_t valid_bases_size; + uint64_t reserved_sizes_size; +}; +#pragma pack(pop) + +struct PendingSnapshot { + std::shared_ptr next; + std::shared_ptr retired; + cudaEvent_t retirement_event{}; + std::atomic published{false}; +}; + +struct RegistryEntry { + std::shared_ptr current; + int64_t marker_value{0}; + std::atomic pending{nullptr}; + std::mutex updates_mutex; + std::unordered_map> updates; +}; + +std::mutex registry_mutex; +std::unordered_map> registry; +std::unordered_map> registry_by_value; + +struct CachedRegistryEntry { + int64_t marker_value; + std::weak_ptr entry; +}; + +thread_local std::unordered_map + registry_cache; + +std::shared_ptr find_registry_entry(const at::Tensor& marker, + int64_t marker_value) { + const auto marker_pointer = marker.data_ptr(); + const auto cached = registry_cache.find(marker_pointer); + if (cached != registry_cache.end() && + cached->second.marker_value == marker_value) { + if (auto entry = cached->second.entry.lock()) { + return entry; + } + } + + std::shared_ptr entry; + { + std::lock_guard lock(registry_mutex); + const auto iterator = registry.find(marker_pointer); + if (iterator != registry.end()) { + entry = iterator->second; + } + if (!entry) { + const auto value_iterator = registry_by_value.find(marker_value); + if (value_iterator == registry_by_value.end()) { + throw std::runtime_error("No embedding-collection indexer is bound"); + } + entry = value_iterator->second; + } + } + registry_cache[marker_pointer] = {marker_value, entry}; + return entry; +} + +std::shared_ptr find_registered_entry( + const at::Tensor& marker) { + std::lock_guard lock(registry_mutex); + return registry.at(marker.data_ptr()); +} + +template +at::Tensor read_tensor(std::ifstream& stream, uint64_t count, + at::ScalarType dtype, int device_index) { + if (count == 0) { + return torch::empty( + {0}, torch::TensorOptions().dtype(dtype).device(torch::kCUDA, + device_index)); + } + std::vector values(count); + stream.read(reinterpret_cast(values.data()), count * sizeof(T)); + auto cpu = torch::from_blob(values.data(), {static_cast(count)}, + torch::TensorOptions().dtype(dtype)) + .clone(); + return cpu.to(torch::Device(torch::kCUDA, device_index)); +} + +} // namespace + +bool register_from_tensors( + const at::Tensor& marker, int64_t kind, const at::Tensor& table_storage, + const at::Tensor& table_bucket_offsets, int64_t bucket_capacity, + const at::Tensor& miss_storage_indices, const at::Tensor& valid_bases, + const at::Tensor& reserved_sizes, int64_t next_fused_key) { + auto snapshot = std::make_shared(); + snapshot->kind = static_cast(kind); + snapshot->table_storage = table_storage; + snapshot->table_bucket_offsets = table_bucket_offsets; + snapshot->bucket_capacity = bucket_capacity; + snapshot->miss_storage_indices = miss_storage_indices; + snapshot->valid_bases = valid_bases; + snapshot->reserved_sizes = reserved_sizes; + snapshot->next_fused_key = next_fused_key; + register_indexer_snapshot(marker, std::move(snapshot)); + return true; +} + +bool unregister_from_tensor(const at::Tensor& marker) { + unregister_indexer_snapshot(marker); + return true; +} + +std::shared_ptr load_indexer_snapshot( + const std::string& path, int device_index) { + std::ifstream stream(path, std::ios::binary); + if (!stream) { + throw std::runtime_error("Cannot open indexer snapshot: " + path); + } + SnapshotHeader header{}; + stream.read(reinterpret_cast(&header), sizeof(header)); + if (std::memcmp(header.magic, kSnapshotMagic, sizeof(kSnapshotMagic)) != 0 || + header.version != kSchemaVersion) { + throw std::runtime_error("Unsupported indexer snapshot: " + path); + } + auto snapshot = std::make_shared(); + snapshot->kind = static_cast(header.kind); + snapshot->bucket_capacity = header.bucket_capacity; + snapshot->next_fused_key = header.next_fused_key; + snapshot->table_storage = read_tensor( + stream, header.table_storage_size, at::kByte, device_index); + snapshot->table_bucket_offsets = read_tensor( + stream, header.table_bucket_offsets_size, at::kLong, device_index); + snapshot->miss_storage_indices = read_tensor( + stream, header.miss_storage_indices_size, at::kLong, device_index); + snapshot->valid_bases = read_tensor( + stream, header.valid_bases_size, at::kLong, device_index); + snapshot->reserved_sizes = read_tensor( + stream, header.reserved_sizes_size, at::kLong, device_index); + return snapshot; +} + +void register_indexer_snapshot(const at::Tensor& marker, + std::shared_ptr snapshot) { + std::lock_guard lock(registry_mutex); + auto& entry = registry[marker.data_ptr()]; + if (!entry) { + entry = std::make_shared(); + entry->marker_value = marker.item(); + } + std::atomic_store(&entry->current, std::move(snapshot)); + registry_by_value[entry->marker_value] = entry; +} + +void unregister_indexer_snapshot(const at::Tensor& marker) { + std::lock_guard lock(registry_mutex); + const auto found = registry.find(marker.data_ptr()); + if (found == registry.end()) { + return; + } + const auto entry = found->second; + for (auto iterator = registry.begin(); iterator != registry.end();) { + if (iterator->second == entry) { + iterator = registry.erase(iterator); + } else { + ++iterator; + } + } + const auto value_entry = registry_by_value.find(entry->marker_value); + if (value_entry != registry_by_value.end() && value_entry->second == entry) { + registry_by_value.erase(value_entry); + } +} + +void stage_indexer_snapshot(const at::Tensor& marker, + std::shared_ptr snapshot, + int64_t snapshot_id) { + auto entry = find_registered_entry(marker); + auto pending = std::make_shared(); + pending->next = std::move(snapshot); + if (cudaEventCreateWithFlags(&pending->retirement_event, + cudaEventDisableTiming) != cudaSuccess) { + throw std::runtime_error("Failed to create indexer retirement event"); + } + { + std::lock_guard lock(entry->updates_mutex); + entry->updates[snapshot_id] = pending; + } + entry->pending.store(pending.get(), std::memory_order_release); +} + +void wait_for_indexer_snapshot_retirement(const at::Tensor& marker, + int64_t snapshot_id) { + auto entry = find_registered_entry(marker); + std::shared_ptr update; + { + std::lock_guard lock(entry->updates_mutex); + update = entry->updates.at(snapshot_id); + } + while (!update->published.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + cudaEventSynchronize(update->retirement_event); + cudaEventDestroy(update->retirement_event); + { + std::lock_guard lock(entry->updates_mutex); + entry->updates.erase(snapshot_id); + } +} + +std::shared_ptr find_indexer_snapshot( + const at::Tensor& marker, int64_t marker_value) { + auto entry = find_registry_entry(marker, marker_value); + auto* pending = entry->pending.load(std::memory_order_acquire); + if (pending != nullptr) { + pending = entry->pending.exchange(nullptr, std::memory_order_acq_rel); + } + if (pending != nullptr) { + auto next = pending->next; + const auto stream = + at::cuda::getCurrentCUDAStream(marker.get_device()).stream(); + if (cudaEventRecord(pending->retirement_event, stream) != cudaSuccess) { + cudaEventDestroy(pending->retirement_event); + throw std::runtime_error("Failed to record indexer retirement event"); + } + auto retired = std::atomic_exchange(&entry->current, next); + pending->retired = std::move(retired); + pending->published.store(true, std::memory_order_release); + return next; + } + return std::atomic_load(&entry->current); +} + +} // namespace dynamicemb::exportable_embedding + +TORCH_LIBRARY_FRAGMENT(INFERENCE_EMB, m) { + m.def("register_embedding_collection_indexer(Tensor marker, int kind, " + "Tensor table_storage, Tensor table_bucket_offsets, int bucket_capacity, " + "Tensor miss_storage_indices, Tensor valid_bases, Tensor reserved_sizes, " + "int next_fused_key) -> bool"); + m.def("unregister_embedding_collection_indexer(Tensor marker) -> bool"); +} + +TORCH_LIBRARY_IMPL(INFERENCE_EMB, CompositeExplicitAutograd, m) { + m.impl("register_embedding_collection_indexer", + &dynamicemb::exportable_embedding::register_from_tensors); + m.impl("unregister_embedding_collection_indexer", + &dynamicemb::exportable_embedding::unregister_from_tensor); +} diff --git a/corelib/dynamicemb/src/exportable_embedding/indexer_snapshot_pybind.cpp b/corelib/dynamicemb/src/exportable_embedding/indexer_snapshot_pybind.cpp new file mode 100644 index 000000000..06de12148 --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/indexer_snapshot_pybind.cpp @@ -0,0 +1,53 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include + +#include "dynamicemb/exportable_embedding/indexer_snapshot.h" + +namespace py = pybind11; + +namespace dynamicemb::exportable_embedding { + +void bind_indexer_snapshot(py::module_& module) { + py::enum_(module, "IndexerSnapshotKind") + .value("LINEAR_HASH_MAP", IndexerSnapshotKind::LinearHashMap) + .value("FUSED_IDENTITY", IndexerSnapshotKind::FusedIdentity); + + py::class_>( + module, "IndexerSnapshot") + .def(py::init([](int64_t kind, at::Tensor table_storage, + at::Tensor table_bucket_offsets, int64_t bucket_capacity, + at::Tensor miss_storage_indices, at::Tensor valid_bases, + at::Tensor reserved_sizes, int64_t next_fused_key) { + auto snapshot = std::make_shared(); + snapshot->kind = static_cast(kind); + snapshot->table_storage = std::move(table_storage); + snapshot->table_bucket_offsets = std::move(table_bucket_offsets); + snapshot->bucket_capacity = bucket_capacity; + snapshot->miss_storage_indices = std::move(miss_storage_indices); + snapshot->valid_bases = std::move(valid_bases); + snapshot->reserved_sizes = std::move(reserved_sizes); + snapshot->next_fused_key = next_fused_key; + return snapshot; + })) + .def_property_readonly("kind", [](const IndexerSnapshot& snapshot) { + return static_cast(snapshot.kind); + }) + .def_readonly("table_storage", &IndexerSnapshot::table_storage) + .def_readonly("table_bucket_offsets", + &IndexerSnapshot::table_bucket_offsets) + .def_readonly("bucket_capacity", &IndexerSnapshot::bucket_capacity) + .def_readonly("miss_storage_indices", + &IndexerSnapshot::miss_storage_indices) + .def_readonly("valid_bases", &IndexerSnapshot::valid_bases) + .def_readonly("reserved_sizes", &IndexerSnapshot::reserved_sizes) + .def_readonly("next_fused_key", &IndexerSnapshot::next_fused_key); + + module.def("load_indexer_snapshot", &load_indexer_snapshot, py::arg("path"), + py::arg("device_index")); +} + +} // namespace dynamicemb::exportable_embedding diff --git a/corelib/dynamicemb/src/exportable_embedding/update_subscriber_pybind.cu b/corelib/dynamicemb/src/exportable_embedding/update_subscriber_pybind.cu new file mode 100644 index 000000000..163a79504 --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/update_subscriber_pybind.cu @@ -0,0 +1,61 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include + +#include + +#include "dynamicemb/exportable_embedding/incremental_update.h" +#include "python/pynve/bindings/binding_layers.hpp" + +namespace py = pybind11; + +namespace dynamicemb::exportable_embedding { + +void bind_update_subscriber(py::module_& module) { + py::class_( + module, "EmbeddingCollectionUpdateSubscriber") + .def( + py::init([]( + const std::string& package_dir, + EmbeddingCollectionIndexerDirectory& indexers, + const py::dict& layers, int device_index) { + std::unordered_map< + std::string, + std::shared_ptr>> + bindings; + for (const auto& item : layers) { + bindings.emplace( + py::cast(item.first), + py::cast>>( + item.second)); + } + return std::make_unique( + package_dir, indexers, std::move(bindings), device_index); + }), + py::arg("package_dir"), py::arg("indexers"), py::arg("nve_layers"), + py::arg("device_index"), py::keep_alive<1, 3>()) + .def( + "apply_incremental_load", + [](EmbeddingCollectionUpdateSubscriber& subscriber, + const std::string& update_json, uint64_t inference_stream) { + subscriber.apply_incremental_load( + EmbeddingCollectionUpdate::from_json(update_json), + reinterpret_cast(inference_stream)); + }, + py::arg("update_json"), py::arg("inference_stream"), + py::call_guard()) + .def( + "wait_for_retirement", + [](EmbeddingCollectionUpdateSubscriber& subscriber, + const std::string& collection_id, int64_t snapshot_id) { + return subscriber.wait_for_retirement(collection_id, snapshot_id) + .to_json(); + }, + py::arg("collection_id"), py::arg("snapshot_id"), + py::call_guard()); +} + +} // namespace dynamicemb::exportable_embedding diff --git a/corelib/dynamicemb/src/exportable_embedding/update_subscriber_unavailable_pybind.cpp b/corelib/dynamicemb/src/exportable_embedding/update_subscriber_unavailable_pybind.cpp new file mode 100644 index 000000000..4f42c8fd9 --- /dev/null +++ b/corelib/dynamicemb/src/exportable_embedding/update_subscriber_unavailable_pybind.cpp @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include +#include + +namespace py = pybind11; + +namespace dynamicemb::exportable_embedding { +namespace { + +class UnavailableUpdateSubscriber { + public: + UnavailableUpdateSubscriber(const std::string&, py::object, py::dict, int) { + throw std::runtime_error( + "Incremental embedding updates require NVE 26.06 or later; " + "this build targets NVE 26.05"); + } +}; + +} // namespace + +void bind_update_subscriber(py::module_& module) { + py::class_( + module, "EmbeddingCollectionUpdateSubscriber") + .def(py::init(), + py::arg("package_dir"), py::arg("indexers"), + py::arg("nve_layers"), py::arg("device_index")); +} + +} // namespace dynamicemb::exportable_embedding diff --git a/docker/Dockerfile b/docker/Dockerfile index ea0bb82cc..0f38606e6 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -269,7 +269,8 @@ RUN git clone --recursive -b v26.05 \ git clone -b v3.5.0 https://github.com/NVIDIA/NVTX.git NVTX # Build into two isolated Python prefixes. NVE's upstream default feature set -# is used for both builds and includes libnve-plugin-nvhm.so. +# includes the NVHashMap host-cache and Redis storage plugins used by the +# exportable example. RUN set -eux; \ if [ "${TARGETPLATFORM}" != "linux/arm64" ]; then \ export PYNVE_DISABLE_AVX512=1; \ @@ -315,13 +316,57 @@ RUN cd /workspace/deps/fbgemm_hstu/fbgemm_gpu/experimental/hstu && \ # kvcache_manager) and runtime libcuda link. # ============================================================================ FROM ${DEVEL_IMAGE} AS build +ARG TARGETPLATFORM + +RUN apt-get update -y --fix-missing && \ + apt-get install -y --no-install-recommends redis-server redis-tools && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* RUN rm -rf /workspace/recsys-examples WORKDIR /workspace/recsys-examples COPY . . -RUN cd /workspace/recsys-examples/corelib/dynamicemb && \ +# Keep the default NVE build aligned with this checkout when DEVEL_IMAGE is an +# external pre-built image. The 26.05 compatibility install remains untouched. +RUN rm -rf /workspace/deps/nve /opt/nve/default/python /workspace/build/pynve-default && \ + cp -a third_party/nv-embedding-cache /workspace/deps/nve && \ + mkdir -p /opt/nve/default/python && \ + if [ "${TARGETPLATFORM}" != "linux/arm64" ]; then \ + export PYNVE_DISABLE_AVX512=1; \ + fi && \ + PYNVE_BUILD_DIR=/workspace/build/pynve-default \ + PYNVE_WITH_TORCH_BINDINGS=1 \ + TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" \ + CPLUS_INCLUDE_PATH="/workspace/deps/NVTX/c/include${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}" \ + python3 -m pip install --no-build-isolation --no-deps \ + --target /opt/nve/default/python \ + /workspace/deps/nve && \ + rm -rf /workspace/build/pynve-default + +# Package the default _C extension, then stage its 26.05 stub build for the +# same image's compatibility example. +RUN rm -f /usr/lib/$(uname -m)-linux-gnu/libcuda.so.1 && \ + test -e /usr/local/cuda-13/compat/lib.real/libcuda.so.1 && \ + ln -s /usr/local/cuda-13/compat/lib.real/libcuda.so.1 /usr/lib/$(uname -m)-linux-gnu/libcuda.so.1 && \ + cd /workspace/recsys-examples/corelib/dynamicemb && \ + mkdir -p torch_binding_build && cd torch_binding_build && \ + cmake .. \ + -DNVE_ROOT=/workspace/deps/nve \ + -DNVE_LIB_DIR=/opt/nve/default/python/pynve \ + -DNVTX_INCLUDE_DIR=/workspace/deps/NVTX/c/include && \ + make -j && \ + cd .. && \ TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" python setup.py install && \ + cmake -S . -B torch_binding_build -DNVE_VERSION=26.05 && \ + cmake --build torch_binding_build --parallel 8 \ + --target dynamicemb_exportable_embedding_python && \ + mkdir -p /opt/dynamicemb/26.05/python && \ + cp -a dynamicemb /opt/dynamicemb/26.05/python/ && \ + cmake -S . -B torch_binding_build -UNVE_VERSION && \ + cmake --build torch_binding_build --parallel 8 \ + --target dynamicemb_exportable_embedding_python && \ + cmake --install torch_binding_build && \ cd /workspace/deps && rm -rf nvcomp && \ NVCOMP_ARCH=$([ "$(uname -m)" = "aarch64" ] && echo "linux-sbsa" || echo "linux-x86_64") && \ wget https://developer.download.nvidia.com/compute/nvcomp/redist/nvcomp/${NVCOMP_ARCH}/nvcomp-${NVCOMP_ARCH}-5.1.0.21_cuda12-archive.tar.xz && \ @@ -331,13 +376,16 @@ RUN cd /workspace/recsys-examples/corelib/dynamicemb && \ cd /workspace/recsys-examples/examples/commons && \ TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" python3 setup.py install && \ cd /workspace/recsys-examples/corelib/recsys_kvcache_manager && \ - TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" pip3 install --no-build-isolation . && \ - rm -f /usr/lib/$(uname -m)-linux-gnu/libcuda.so.1 && \ - test -e /usr/local/cuda-13/compat/lib.real/libcuda.so.1 && \ - ln -s /usr/local/cuda-13/compat/lib.real/libcuda.so.1 /usr/lib/$(uname -m)-linux-gnu/libcuda.so.1 && \ - cd /workspace/recsys-examples/corelib/dynamicemb && \ - mkdir -p torch_binding_build && cd torch_binding_build && \ - cmake .. && make -j + TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" pip3 install --no-build-isolation . + +RUN export CMAKE_PREFIX_PATH="$(python3 -c 'import torch; print(torch.utils.cmake_prefix_path)')" && \ + cmake -S ./corelib/dynamicemb/example/exportable_embedding/cpp \ + -B ./corelib/dynamicemb/example/exportable_embedding/cpp/build \ + -DNVE_ROOT=/workspace/deps/nve \ + -DNVE_LIB_DIR=/opt/nve/default/python/pynve \ + -DNVTX_INCLUDE_DIR=/workspace/deps/NVTX/c/include \ + -DDYNAMICEMB_LIB_DIR=/workspace/recsys-examples/corelib/dynamicemb/torch_binding_build && \ + cmake --build ./corelib/dynamicemb/example/exportable_embedding/cpp/build -j 8 # ============================================================================ # AOTI additions: version-independent C++ replay, two NVE loader diff --git a/examples/hstu/inference/triton/hstu_export_aligned/model.py b/examples/hstu/inference/triton/hstu_export_aligned/model.py index f555c7c0b..0509dc618 100644 --- a/examples/hstu/inference/triton/hstu_export_aligned/model.py +++ b/examples/hstu/inference/triton/hstu_export_aligned/model.py @@ -72,7 +72,7 @@ def initialize(self, args): ( dataset_args, _, - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, ) = get_inference_dataset_and_embedding_configs() data_processor = get_common_preprocessors("")[dataset_args.dataset_name] @@ -100,7 +100,7 @@ def initialize(self, args): with torch.inference_mode(): self._model = get_exportable_model_for_inference( - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, str(checkpoint_dir), ) diff --git a/examples/hstu/inference_aoti/export_inference_gr_ranking.py b/examples/hstu/inference_aoti/export_inference_gr_ranking.py index 339cb13be..e4f013b94 100644 --- a/examples/hstu/inference_aoti/export_inference_gr_ranking.py +++ b/examples/hstu/inference_aoti/export_inference_gr_ranking.py @@ -49,6 +49,7 @@ SCRIPT_DIR = Path(__file__).resolve().parent DEFAULT_EXPORT_DIR = SCRIPT_DIR / "hstu_gr_ranking_model" DEFAULT_DUMP_DIR = SCRIPT_DIR / "export_test_dump" +DYNAMIC_EMBEDDING_GPU_CACHE_SIZE = 4 << 30 def init_single_rank_distributed(): @@ -125,17 +126,11 @@ def get_inference_dataset_and_embedding_configs( ) if dataset_args.dataset_name == "kuairand-1k": + from modules.exportable_embedding import ( + create_kuairand_embedding_collection_configs, + ) + HASH_SIZE = 1000_064 - dynamic_table_configs = { - "user_id": True, - "user_active_degree": False, - "follow_user_num_range": False, - "fans_user_num_range": False, - "friend_user_num_range": False, - "register_days_range": False, - "video_id": True, - "action_weights": False, - } trained_emb_table_sizes = { "user_id": 1000, "user_active_degree": 8, @@ -146,15 +141,19 @@ def get_inference_dataset_and_embedding_configs( "video_id": HASH_SIZE, "action_weights": 233, } - for idx, config in enumerate(embedding_configs): + for config in embedding_configs: config.vocab_size = trained_emb_table_sizes[config.table_name] - config.use_dynamic = dynamic_table_configs[config.table_name] + embedding_collection_configs = ( + create_kuairand_embedding_collection_configs( + dynamic_gpu_cache_size=DYNAMIC_EMBEDDING_GPU_CACHE_SIZE + ) + ) return ( dataset_args, embedding_configs if not disable_contextual_features else embedding_configs[-2:], - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, ) @@ -175,14 +174,14 @@ def get_training_gr_model(): def get_exportable_model_for_inference( - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, checkpoint_dir, ): model = get_training_gr_model() inference_model = apply_inference( model, - dynamic_table_configs=dynamic_table_configs, + embedding_collection_configs=embedding_collection_configs, trained_emb_table_sizes=trained_emb_table_sizes, checkpoint_dir=checkpoint_dir, ) @@ -221,7 +220,7 @@ def __init__(self, t): ( dataset_args, _, - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, ) = get_inference_dataset_and_embedding_configs() @@ -254,7 +253,7 @@ def strip_padding_batch(batch, unpadded_batch_size): register_hstu_export_pytrees() model = get_exportable_model_for_inference( - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, checkpoint_dir, ) @@ -394,15 +393,7 @@ def forward( export_dir, dynamic_shapes=dynamic_shapes, ) - print(f"[INFO] Exported and packaged the model to:") - print(f" {export_dir}/") - print( - " ├── model.pt2 # AOT-compiled model package for AOTIModelPackageLoader" - ) - print( - " ├── metadata.json # NVE layer metadata (id, num_embeddings, emb_size, etc.)" - ) - print(" └── weights/*.nve # NVE weight data (LinearUVM)") + print("[INFO] Exported the model package and NVE resources") # === Test Compiled Model === aoti_model_runtime, nve_layers = load_aoti( diff --git a/examples/hstu/inference_aoti/export_inference_gr_ranking_kvcache.py b/examples/hstu/inference_aoti/export_inference_gr_ranking_kvcache.py index c38f15119..dcf4e05cb 100644 --- a/examples/hstu/inference_aoti/export_inference_gr_ranking_kvcache.py +++ b/examples/hstu/inference_aoti/export_inference_gr_ranking_kvcache.py @@ -72,6 +72,7 @@ SCRIPT_DIR = Path(__file__).resolve().parent DEFAULT_EXPORT_DIR = SCRIPT_DIR / "hstu_gr_ranking_kvcache_model" DEFAULT_DUMP_DIR = SCRIPT_DIR / "export_test_dump" +DYNAMIC_EMBEDDING_GPU_CACHE_SIZE = 4 << 30 def start_flexkv_server( @@ -198,17 +199,11 @@ def get_inference_dataset_and_embedding_configs( ) if dataset_args.dataset_name == "kuairand-1k": + from modules.exportable_embedding import ( + create_kuairand_embedding_collection_configs, + ) + HASH_SIZE = 1000_064 - dynamic_table_configs = { - "user_id": True, - "user_active_degree": False, - "follow_user_num_range": False, - "fans_user_num_range": False, - "friend_user_num_range": False, - "register_days_range": False, - "video_id": True, - "action_weights": False, - } trained_emb_table_sizes = { "user_id": 1000, "user_active_degree": 8, @@ -219,15 +214,19 @@ def get_inference_dataset_and_embedding_configs( "video_id": HASH_SIZE, "action_weights": 233, } - for idx, config in enumerate(embedding_configs): + for config in embedding_configs: config.vocab_size = trained_emb_table_sizes[config.table_name] - config.use_dynamic = dynamic_table_configs[config.table_name] + embedding_collection_configs = ( + create_kuairand_embedding_collection_configs( + dynamic_gpu_cache_size=DYNAMIC_EMBEDDING_GPU_CACHE_SIZE + ) + ) return ( dataset_args, embedding_configs if not disable_contextual_features else embedding_configs[-2:], - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, ) @@ -347,7 +346,7 @@ def make_export_kvcache_config( def get_exportable_model_for_inference( - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, checkpoint_dir, max_batch_size, @@ -361,7 +360,7 @@ def get_exportable_model_for_inference( model = get_training_gr_model() model = apply_inference_embedding_collection( model, - dynamic_table_configs=dynamic_table_configs, + embedding_collection_configs=embedding_collection_configs, trained_emb_table_sizes=trained_emb_table_sizes, ) inference_hstu_config = make_inference_hstu_config( @@ -449,7 +448,7 @@ def __init__(self, t): ( dataset_args, _, - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, ) = get_inference_dataset_and_embedding_configs() @@ -485,7 +484,7 @@ def strip_padding_batch(batch, unpadded_batch_size): try: model = get_exportable_model_for_inference( - dynamic_table_configs, + embedding_collection_configs, trained_emb_table_sizes, checkpoint_dir, config_max_batch_size, @@ -670,15 +669,7 @@ def forward( export_dir, dynamic_shapes=dynamic_shapes, ) - print(f"[INFO] Exported and packaged the model to:") - print(f" {export_dir}/") - print( - " ├── model.pt2 # AOT-compiled model package for AOTIModelPackageLoader" - ) - print( - " ├── metadata.json # NVE layer metadata (id, num_embeddings, emb_size, etc.)" - ) - print(" └── weights/*.nve # NVE weight data (LinearUVM)") + print("[INFO] Exported the model package and NVE resources") # === Test Compiled Model === aoti_model_runtime, nve_layers = load_aoti( diff --git a/examples/hstu/inference_aoti/nve_init_hook/README.md b/examples/hstu/inference_aoti/nve_init_hook/README.md index 66b94c6a1..b51e24715 100644 --- a/examples/hstu/inference_aoti/nve_init_hook/README.md +++ b/examples/hstu/inference_aoti/nve_init_hook/README.md @@ -4,7 +4,8 @@ process-global `NVELayerRegistry` at model load, so `nve_ops::embedding_lookup` resolves at inference time. Those weights live **outside** `model.pt2` (`/metadata.json` + `/weights/*.nve`), so loading the -package alone does not load them. +package alone does not load them. The HSTU indexer state is embedded in +`model.pt2` for this NVE 26.05 Triton workflow. This is the **plug-in** for the generic `MODEL_INIT_LIBRARY` hook in the Triton PyTorch backend: the backend `dlopen()`s this `.so` and calls its diff --git a/examples/hstu/model/inference_ranking_gr.py b/examples/hstu/model/inference_ranking_gr.py index 25126929d..3f7be37f6 100755 --- a/examples/hstu/model/inference_ranking_gr.py +++ b/examples/hstu/model/inference_ranking_gr.py @@ -12,8 +12,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import os -from typing import Dict +from collections.abc import Mapping +from typing import TYPE_CHECKING import torch from commons.datasets.hstu_batch import HSTUBatch @@ -23,6 +26,9 @@ from recsys_kvcache_manager.kvcache_config import KVCacheConfig from torchrec.sparse.jagged_tensor import KeyedJaggedTensor +if TYPE_CHECKING: + from dynamicemb.exportable_tables import InferenceEmbeddingCollectionConfig + try: import hstu_cuda_ops # noqa: F401 - registers torch.ops.hstu_cuda_ops.* except ImportError: @@ -223,10 +229,12 @@ def get_inference_ranking_gr( def apply_inference( training_model: torch.nn.Module, - dynamic_table_configs: Dict[str, int], - trained_emb_table_sizes: Dict[str, int], + embedding_collection_configs: Mapping[ + str, InferenceEmbeddingCollectionConfig + ], + trained_emb_table_sizes: Mapping[str, int], checkpoint_dir: str, -): +) -> InferenceRankingGR: from dynamicemb.exportable_tables import apply_inference_embedding_collection from modules.exportable_embedding import apply_inference_sparse from modules.inference_dense_module import apply_inference_hstu_dense @@ -234,8 +242,8 @@ def apply_inference( # Step.1 - [General] Convert ModuleDict[Embedding] to InferenceEmbeddingCollection model = apply_inference_embedding_collection( training_model, - dynamic_table_configs, - trained_emb_table_sizes, + embedding_collection_configs=embedding_collection_configs, + trained_emb_table_sizes=trained_emb_table_sizes, ) # Step.2 - [Recsys Example Structure Specific] Apply model specific training to inference conversion @@ -262,6 +270,4 @@ def apply_inference( inference_model.half() inference_model.load_checkpoint(checkpoint_dir) - inference_model = inference_model.eval() - - return inference_model + return inference_model.eval() diff --git a/examples/hstu/modules/exportable_embedding.py b/examples/hstu/modules/exportable_embedding.py index 1c9a77465..3f96f944f 100755 --- a/examples/hstu/modules/exportable_embedding.py +++ b/examples/hstu/modules/exportable_embedding.py @@ -75,8 +75,9 @@ def _load_inference_emb_ops() -> bool: from commons.modules.embedding import ShardedEmbedding, ShardedEmbeddingConfig from dynamicemb.exportable_tables import ( + EmbeddingCollectionIndexerType, InferenceEmbeddingCollection, - create_inference_embedding_collection, + InferenceEmbeddingCollectionConfig, ) from modules.nve_compat import needs_legacy_embedding_lookup_fake_override from torchrec.modules.embedding_configs import EmbeddingConfig @@ -135,8 +136,8 @@ def __init__( dynamic_embedding_configs: Union[ List[EmbeddingConfig], List[InferenceEmbeddingConfig] ], - static_embedding_collection: Optional[InferenceEmbeddingCollection] = None, - dynamic_embedding_collection: Optional[InferenceEmbeddingCollection] = None, + static_embedding_collection: InferenceEmbeddingCollection, + dynamic_embedding_collection: InferenceEmbeddingCollection, ): super(ExportableEmbedding, self).__init__() assert ( @@ -166,14 +167,6 @@ def __init__( self._static_embedding_configs = static_embedding_configs self._dynamic_embedding_configs = dynamic_embedding_configs - if static_embedding_collection is None: - static_embedding_collection = create_inference_embedding_collection( - self._static_embedding_configs, pooling_mode=-1, use_dynamic=False - ) - if dynamic_embedding_collection is None: - dynamic_embedding_collection = create_inference_embedding_collection( - self._dynamic_embedding_configs, pooling_mode=-1, use_dynamic=True - ) self._static_embedding_collection = static_embedding_collection self._dynamic_embedding_collection = dynamic_embedding_collection @@ -377,8 +370,8 @@ def forward(self, kjt: KeyedJaggedTensor) -> Dict[str, JaggedTensor]: def get_exportable_embedding( embedding_configs: List[InferenceEmbeddingConfig], - static_embedding_collection: Optional[InferenceEmbeddingCollection] = None, - dynamic_embedding_collection: Optional[InferenceEmbeddingCollection] = None, + static_embedding_collection: InferenceEmbeddingCollection, + dynamic_embedding_collection: InferenceEmbeddingCollection, ): static_embedding_configs, dynamic_embedding_configs = [], [] for config in embedding_configs: @@ -396,10 +389,43 @@ def get_exportable_embedding( def apply_inference_sparse( training_embedding: ShardedEmbedding, -) -> InferenceEmbeddingCollection: +) -> ExportableEmbedding: + static_collection = training_embedding._data_parallel_embedding_collection + dynamic_collection = training_embedding._model_parallel_embedding_collection + if not isinstance( + static_collection, InferenceEmbeddingCollection + ) or not isinstance(dynamic_collection, InferenceEmbeddingCollection): + raise TypeError("embedding collections must be converted before wrapping") + return ExportableEmbedding( - training_embedding._model_parallel_embedding_collection.embedding_configs, - training_embedding._data_parallel_embedding_collection.embedding_configs, - training_embedding._data_parallel_embedding_collection, - training_embedding._model_parallel_embedding_collection, + static_embedding_configs=static_collection.embedding_configs, + dynamic_embedding_configs=dynamic_collection.embedding_configs, + static_embedding_collection=static_collection, + dynamic_embedding_collection=dynamic_collection, + ) + + +def create_kuairand_embedding_collection_configs( + *, dynamic_gpu_cache_size: int +) -> Dict[str, InferenceEmbeddingCollectionConfig]: + static_config = InferenceEmbeddingCollectionConfig( + indexer_type=EmbeddingCollectionIndexerType.FUSED_IDENTITY, + nve_layer_type="gpu", + indexer_state_sidecar=False, + ) + dynamic_config = InferenceEmbeddingCollectionConfig( + indexer_type=EmbeddingCollectionIndexerType.LINEAR_HASH_MAP, + nve_layer_type="linear_uvm", + indexer_state_sidecar=False, + gpu_cache_size=dynamic_gpu_cache_size, ) + return { + "user_id": dynamic_config, + "user_active_degree": static_config, + "follow_user_num_range": static_config, + "fans_user_num_range": static_config, + "friend_user_num_range": static_config, + "register_days_range": static_config, + "video_id": dynamic_config, + "action_weights": static_config, + } diff --git a/third_party/nv-embedding-cache b/third_party/nv-embedding-cache index 4c1f956f0..a31a4d2a8 160000 --- a/third_party/nv-embedding-cache +++ b/third_party/nv-embedding-cache @@ -1 +1 @@ -Subproject commit 4c1f956f0ae977bb40425273d2fcfded4c590002 +Subproject commit a31a4d2a847aee1d749cff8157f9e011027634a2