Skip to content

Add gflag to select feature names for SSD KV embedding table#5585

Open
jnwan wants to merge 1 commit intopytorch:mainfrom
jnwan:export-D99744279
Open

Add gflag to select feature names for SSD KV embedding table#5585
jnwan wants to merge 1 commit intopytorch:mainfrom
jnwan:export-D99744279

Conversation

@jnwan
Copy link
Copy Markdown

@jnwan jnwan commented Apr 7, 2026

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

Summary

Refactor gflag parsing for SSD KV embedding table feature name selection into a dedicated shouldCreateSsdKvBackendForTbe() method.

Changes:

  • pytorch_xl_format_v2_util.cpp: Extract flag parsing logic (FLAGS_embedding_db_enable_kv_embedding_table and FLAGS_embedding_db_kv_embedding_table_feature_names) into a standalone shouldCreateSsdKvBackendForTbe() function. The feature names are parsed once using a static local with folly::splitTo + std::inserter pattern.
  • pytorch_xl_format_v2_util.h: Add declaration and DECLARE macros for the two new flags.
  • XLFormatV2CPUTest.cpp: Add 3 unit tests covering flag-disabled, feature-match, and feature-no-match cases.
  • SSDKVInferenceEmbedding.cpp: Implement inference_set_kv_db_async and get_kv_db_async to delegate to sync methods (previously threw).
  • dram_kv_embedding_inference_wrapper.cpp: Simplify by always using async path, removing FLAGS_kv_embedding_async_get_set conditional.

Reviewed By: emlin

Differential Revision: D99744279

Summary:
X-link: facebookresearch/FBGEMM#2545

## Summary
Refactor gflag parsing for SSD KV embedding table feature name selection into a dedicated `shouldCreateSsdKvBackendForTbe()` method.

### Changes:
- **`pytorch_xl_format_v2_util.cpp`**: Extract flag parsing logic (`FLAGS_embedding_db_enable_kv_embedding_table` and `FLAGS_embedding_db_kv_embedding_table_feature_names`) into a standalone `shouldCreateSsdKvBackendForTbe()` function. The feature names are parsed once using a `static` local with `folly::splitTo` + `std::inserter` pattern.
- **`pytorch_xl_format_v2_util.h`**: Add declaration and `DECLARE` macros for the two new flags.
- **`XLFormatV2CPUTest.cpp`**: Add 3 unit tests covering flag-disabled, feature-match, and feature-no-match cases.
- **`SSDKVInferenceEmbedding.cpp`**: Implement `inference_set_kv_db_async` and `get_kv_db_async` to delegate to sync methods (previously threw).
- **`dram_kv_embedding_inference_wrapper.cpp`**: Simplify by always using async path, removing `FLAGS_kv_embedding_async_get_set` conditional.

Reviewed By: emlin

Differential Revision: D99744279
@meta-cla meta-cla bot added the cla signed label Apr 7, 2026
@meta-codesync
Copy link
Copy Markdown
Contributor

meta-codesync bot commented Apr 7, 2026

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant