Inference Table Batched Embedding (TBE)
The table batched embedding bag (TBE) operator is an important base operation for embedding lookup for recommendation system inference on GPU. We added the following enhancements for performance and flexibility:
Alignment restriction removed
- Embedding dimension * data type size had to be multiple of 4B before and now, it is 1B.
Unified Virtual Memory (UVM) caching kernel optimizations
- UVM caching kernels now scale linearly with # of tables using UVM caching. Previously, it was having similar overhead as all tables using UVM caching
- UVM caching kernel overhead is much smaller than before
Inference FP8 Table Batched Embedding (TBE)