US 12,393,514 B2
Pre-optimizer and optimizer based framework for optimal deployment of embedding tables across heterogeneous memory architecture
Ashwin Krishnan, Thane (IN); Manoj Karunakaran Nambiar, Thane (IN); and Rekha Singhal, Thane (IN)
Assigned to TATA CONSULTANCY SERVICES LIMITED, Mumbai (IN)
Filed by Tata Consultancy Services Limited, Mumbai (IN)
Filed on Aug. 14, 2024, as Appl. No. 18/804,270.
Claims priority of application No. 202321060567 (IN), filed on Sep. 8, 2023.
Prior Publication US 2025/0086111 A1, Mar. 13, 2025
Int. Cl. G06F 12/08 (2016.01); G06F 12/0802 (2016.01)
CPC G06F 12/0802 (2013.01) 18 Claims
OG exemplary drawing
 
1. A processor implemented method for high-speed recommendations, the method comprising:
receiving, by a pre-optimizer executed by one or more hardware processors, a plurality of input parameters comprising (i) a memory configuration of a heterogeneous memory architecture used by a Recommendation Model (RM) during inference, the heterogenous memory architecture comprising a first level memory and a second level memory of a Field Programmable Gate Array (FPGA) accelerator, and a CPU memory type, (ii) an access latency of each memory type of the heterogeneous memory architecture, (iii) a total number of a plurality of embedding tables in the RM, (iv) a plurality of embeddings in each of the plurality of embedding tables, (v) an input access pattern of the plurality of embedding tables, and (vi) a table access pattern of each table from the plurality of embedding tables depicting number of access required from each table per inference;
identifying a RM type, by the one or more hardware processors, based on the total number of the plurality of embedding tables, wherein the RM type is one of a small model, a medium model, and a large model; and
determining an optimal deployment of the plurality of embedding tables on the first level memory and the second level memory in accordance with the RM type using one or more of a pre-optimizer for determining initial table partitions within the plurality of embedding tables, an optimizer to determine final table partitions within the plurality of embedding tables by optimizing the initial partitions, and Round Trips (RTs) to the second level memory, wherein if the identified RM type is the small model then determining the optimal deployment of the plurality of embedding tables comprises:
running the pre-optimizer on the plurality of embedding tables for caching on the first level memory, wherein
(a) the pre-optimizer determines whether to partition the plurality of embedding tables by analyzing a set of pre-defined combination of table partitions satisfying a set of space constraints with a predetermined criteria defining that an overall hit rates of each embedding table should be at least 5%, and
(b) if the plurality of embedding tables are to be partitioned the pre-optimizer determines the initial table partitions for the first level memory in accordance with predefined initial permissible combinations of table partitions;
running the optimizer on the initial table partitions with a set of space constraints to optimize the initial table partitions for the first level memory and obtain the final table partitions for the first level memory for caching associated embeddings on a first type of the first level memory;
calculating number of Round Trips (RT) to the second level memory for remaining table partitions other than the initial table partitions from among the plurality of embedding tables;
running, if the RT is greater than unity, the pre-optimizer to determine the initial table partitions for the remaining table partitions followed by the optimizer for generating the final table partitions for the remaining table partitions for caching on the second level memory; and
pushing, if the RT is less than or equal to unity, the remaining table partitions as final table partitions onto the second level memory, in accordance with a memory constraint that not more than one embedding table is allocated to an individual bank among a plurality of banks of the second level memory, wherein the final table partitions for each of the plurality of embedding tables are indexed using learned index technique;
sorting, by the one or more hardware processors, the final table partitions generated by one of pre-optimizer and optimizer based on keys associated with the final table partitions to run the learned index technique to be processed by RM during inference; and
dynamically mapping each table partition in accordance with the final table partitions to one of the first level memory and the second level memory by a heterogenous mapper executed by the one or more hardware processors on receiving a shuffle signal triggered due to change in the input access pattern.