Skip to content

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

Draft
geoffreyQiu wants to merge 3 commits into
NVIDIA:mainfrom
geoffreyQiu:feature/exportable-embedding-table-fusion
Draft

geoffreyQiu wants to merge 3 commits into
NVIDIA:mainfrom
geoffreyQiu:feature/exportable-embedding-table-fusion

Conversation

@geoffreyQiu

Copy link
Copy Markdown
Collaborator

Summary

Extend DynamicEmb exportable inference embedding collections with configurable indexing and NVE storage.

The conversion API now accepts an InferenceEmbeddingCollectionConfig for each physical embedding collection instead of selecting the implementation through use_dynamic.

What changed

  • Add configurable FusedIdentityIndexer, LinearHashMapIndexer, and BitConcatIndexer.
  • Support NVE GPU, LinearUVM, and hierarchical layers.
  • Construct hierarchical tables directly through their parameter server using fused keys and embeddings.
  • Keep data-backed indexer state and NVHashMap contents outside the exported PT2 package as sidecars.
  • Support Redis-backed hierarchical-table incremental updates through the update coordinator and runtime subscriber.
  • Add Python and C++ loading paths for exported indexer state.
  • Add an end-to-end export, AOTI replay, sidecar loading, and incremental-load example.
  • Update HSTU inference conversion and export call sites to pass the new collection configuration.

Verified combinations

Indexer NVE layer/backend
FusedIdentityIndexer GPU
LinearHashMapIndexer LinearUVM
LinearHashMapIndexer Hierarchical NVHashMap
LinearHashMapIndexer Hierarchical Redis
BitConcatIndexer Hierarchical NVHashMap
BitConcatIndexer Hierarchical Redis

BitConcatIndexer is intentionally limited to hierarchical layers. Its table-ID and feature-ID bit widths are fixed by the model configuration.

Validation

  • Passed the NVE 26.07 end-to-end GPU workflow.
  • Verified all six embedding-collection combinations.
  • Verified eager and Python AOTI execution.
  • Verified C++ AOTI replay with sidecar loading.
  • Verified Redis incremental loading through the coordinator/subscriber path.
  • The example includes the NVE 26.05-compatible four-combination path; separate 26.05 GPU validation remains pending.

Commit organization

The draft is organized into three commits so it can later be split if needed:

  1. Reusable DynamicEmb exportable-embedding implementation.
  2. End-to-end DynamicEmb example and verification workflow.
  3. HSTU integration with the new APIs.

Checklist

  • I am familiar with the Contributing Guidelines.
  • New or existing tests cover these changes.
  • The documentation is up to date with these changes.

@geoffreyQiu
geoffreyQiu force-pushed the feature/exportable-embedding-table-fusion branch from a63329d to 70063d5 Compare September 4, 2026 04:35
@geoffreyQiu
geoffreyQiu force-pushed the feature/exportable-embedding-table-fusion branch 14 times, most recently from 206874d to a23418e Compare September 17, 2026 10:32
@geoffreyQiu
geoffreyQiu force-pushed the feature/exportable-embedding-table-fusion branch from a23418e to 2c154a9 Compare September 17, 2026 10:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant