[FEA] Support configurable indexing and NVE layers for exportable embedding collections

.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

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})

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
 

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()

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,
             )

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",
+]

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")