Add Python enum configs and KJT builder for enrichment (#5464) by EddyLXJ · Pull Request #5464 · pytorch/FBGEMM

EddyLXJ · 2026-03-10T01:09:52Z

Summary:

X-link: https://github.com/facebookresearch/FBGEMM/pull/2439

CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries.

WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility.

Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py
Updated EnrichmentPolicy to use enum types instead of strings
Added enrichment_policy field to KVZCHTBEConfig for config propagation
Convert enum values to int when passing to C++ TorchScript layer in training.py
Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries
Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams

Reviewed By: zlzhao1104

Differential Revision: D95883280

meta-codesync · 2026-03-10T01:10:00Z

@EddyLXJ has exported this pull request. If you are a Meta employee, you can view the originating Diff in D95883280.

…3850) Summary: X-link: pytorch/FBGEMM#5464 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

Summary: X-link: meta-pytorch/torchrec#3850 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

Summary: X-link: meta-pytorch/torchrec#3850 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

Summary: Pull Request resolved: pytorch#5464 X-link: meta-pytorch/torchrec#3850 X-link: https://github.com/facebookresearch/FBGEMM/pull/2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 Pull Request resolved: meta-pytorch#3850 X-link: https://github.com/facebookresearch/FBGEMM/pull/2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

Summary: Pull Request resolved: pytorch#5464 X-link: meta-pytorch/torchrec#3850 X-link: https://github.com/facebookresearch/FBGEMM/pull/2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 Pull Request resolved: meta-pytorch#3850 X-link: https://github.com/facebookresearch/FBGEMM/pull/2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

Summary: X-link: meta-pytorch/torchrec#3850 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

Summary: X-link: meta-pytorch/torchrec#3850 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

Summary: X-link: meta-pytorch/torchrec#3850 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

Summary: Pull Request resolved: pytorch#5464 X-link: meta-pytorch/torchrec#3850 X-link: https://github.com/facebookresearch/FBGEMM/pull/2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 Pull Request resolved: meta-pytorch#3850 X-link: https://github.com/facebookresearch/FBGEMM/pull/2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

Summary: X-link: meta-pytorch/torchrec#3850 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

Summary: X-link: meta-pytorch/torchrec#3850 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

Summary: X-link: facebookresearch/FBGEMM#2437 CONTEXT: cudaStreamAddCallback holds the CUDA driver mutex during callback execution, blocking other threads from making CUDA API calls (e.g., NCCL kernel launches on other streams). This causes latency when multiple CUDA streams operate concurrently. WHAT: Migrate all cudaStreamAddCallback call sites to cudaLaunchHostFunc which does not hold the CUDA driver mutex during execution. - Renamed cuda_callback_func to cuda_host_func with simplified signature (removed stream and status params) - Updated all 7 call sites across kv_db_table_batched_embeddings.cpp and ssd_scratch_pad_indices_queue.cpp - Updated documentation to reflect the new API and its concurrency benefits Reviewed By: zlzhao1104, emlin Differential Revision: D95879076

…ytorch#5463) Summary: X-link: facebookresearch/FBGEMM#2438 CONTEXT: DRAM KV embedding cache currently only stores cached embeddings without enrichment from external sources. This limits the cache's ability to leverage pre-trained embeddings or external knowledge sources like IGR (Identity Graph Resolution) data from Laser, resulting in cold-start issues and lower quality for sparse features. WHAT: This diff introduces a configurable enrichment system that asynchronously fetches embeddings from external sources (e.g., IGR via Laser) to populate the DRAM cache for missing IDs. - Added Python configuration layer: EnrichmentPolicy NamedTuple in split_table_batched_embeddings_ops_common.py with fields for enrichment_type, provider_name, client_id, enrichment_dim, and response_format - Added C++ configuration layer: EnrichmentConfig TorchScript custom class in enrichment_config.h with EnrichmentType enum and EnrichmentResponseFormat enum - Implemented IGR enrichment infrastructure in new igr_enrichment.h: LaserClient initialization, async coroutine-based Laser queries, generic thrift parsing, and tensor preparation utilities - Added enrichment orchestration in DramKVEmbeddingCache: dedicated laser_executor_ thread pool (4 threads), pause/resume mechanism for laser writes to yield to forward/backward passes, rate limiting via pending_laser_requests_ counter - Implemented enrichment_query_id() in training.py: runs on dedicated enrichment_query_stream CUDA stream, performs deduplication on linearized indices, async D2H copy of query IDs - Added zero-weight detection and invalidation logic: detects cache slots with all-zero weights during prefetch/flush, invalidates L1 cache state for zero-weight rows to trigger re-fetch - Updated BUCK dependencies: added laser/client, thrift protocol, and folly/io:iobuf Reviewed By: emlin Differential Revision: D95873233

Summary: X-link: meta-pytorch/torchrec#3850 X-link: facebookresearch/FBGEMM#2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

…3850) Summary: X-link: pytorch/FBGEMM#5464 Pull Request resolved: meta-pytorch#3850 X-link: https://github.com/facebookresearch/FBGEMM/pull/2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280

Summary: X-link: pytorch/FBGEMM#5464 Pull Request resolved: #3850 X-link: https://github.com/facebookresearch/FBGEMM/pull/2439 CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries. WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility. - Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py - Updated EnrichmentPolicy to use enum types instead of strings - Added enrichment_policy field to KVZCHTBEConfig for config propagation - Convert enum values to int when passing to C++ TorchScript layer in training.py - Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries - Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams Reviewed By: zlzhao1104 Differential Revision: D95883280 fbshipit-source-id: 0fcc5b071769a1aa36b5a0354f53406a7ff19e0c

meta-codesync · 2026-03-16T21:42:08Z

This pull request has been merged in 2cb38c0.

meta-cla Bot added the cla signed label Mar 10, 2026

meta-codesync Bot added fb-exported meta-exported labels Mar 10, 2026

EddyLXJ force-pushed the export-D95883280 branch from ed4e4aa to 81602bc Compare March 10, 2026 01:34

EddyLXJ force-pushed the export-D95883280 branch from 81602bc to 2436c1e Compare March 10, 2026 01:36

EddyLXJ force-pushed the export-D95883280 branch from 2436c1e to d25f528 Compare March 10, 2026 01:37

EddyLXJ force-pushed the export-D95883280 branch from d25f528 to dabc8b3 Compare March 10, 2026 01:40

EddyLXJ force-pushed the export-D95883280 branch from dabc8b3 to 97c3f0c Compare March 12, 2026 20:30

EddyLXJ force-pushed the export-D95883280 branch from 97c3f0c to 353a4a9 Compare March 12, 2026 20:32

EddyLXJ force-pushed the export-D95883280 branch from 353a4a9 to 416d311 Compare March 12, 2026 20:37

EddyLXJ force-pushed the export-D95883280 branch from 416d311 to 24f61c7 Compare March 12, 2026 20:42

meta-codesync Bot changed the title ~~Add Python enum configs and KJT builder for enrichment~~ Add Python enum configs and KJT builder for enrichment (#5464) Mar 12, 2026

meta-codesync Bot mentioned this pull request Mar 12, 2026

Add Python enum configs and KJT builder for enrichment (#3850) meta-pytorch/torchrec#3850

Closed

EddyLXJ force-pushed the export-D95883280 branch from 24f61c7 to 39cfd2f Compare March 13, 2026 21:10

EddyLXJ force-pushed the export-D95883280 branch from 39cfd2f to 8e86c2d Compare March 16, 2026 18:00

EddyLXJ added 3 commits March 16, 2026 11:00

EddyLXJ force-pushed the export-D95883280 branch from 8e86c2d to b093ce4 Compare March 16, 2026 18:01

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add Python enum configs and KJT builder for enrichment (#5464)#5464

Add Python enum configs and KJT builder for enrichment (#5464)#5464
EddyLXJ wants to merge 3 commits into
pytorch:mainfrom
EddyLXJ:export-D95883280

EddyLXJ commented Mar 10, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

meta-codesync Bot commented Mar 10, 2026

Uh oh!

meta-codesync Bot commented Mar 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

EddyLXJ commented Mar 10, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meta-codesync Bot commented Mar 10, 2026

Uh oh!

meta-codesync Bot commented Mar 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

EddyLXJ commented Mar 10, 2026 •

edited by meta-codesync Bot

Loading