Home · ASPLOS 2024 · Machine Learning Systems

MaxEmbed: Maximizing SSD bandwidth utilization for huge embedding models serving

Ruwen Fan, Minhui Xie, Haodi Jiang and Youyou Lu

Conference
ASPLOS 2024
Published
2024
Pages
188-202
Citations
1

Recommendation models store enormous embedding tables that are too big and costly to keep in memory, so they are served from cheap high-capacity SSDs — but an SSD reads a whole page at a time while an embedding is tiny, so fetching one wastes most of the page (read amplification); prior systems pack co-accessed embeddings onto the same page via hypergraph partitioning, but that forces each embedding into one fixed cluster it can never leave; this paper (Tsinghua — Ruwen Fan, Minhui Xie, Haodi Jiang, Youyou Lu) presents MaxEmbed, which spends the SSD's abundant capacity to replicate embeddings across many page combinations so any query finds a well-packed page, plus an online serving module with two optimizations to curb replication overhead, raising SSD embedding-serving throughput by up to 18.7 percent.

The argument in one glance

  1. Problem

    Recommendation models need huge tables that are too expensive to keep in memory

    Deep learning recommendation models power search, advertising, and e-commerce, and their defining ingredient is the embedding table: to handle sparse categorical features — a user ID, an item ID, a category — the model keeps a giant lookup table mapping each of billions of IDs to a small learned vector. These tables are enormous, often far larger than the rest of the model, so holding them entirely in DRAM is both capacity-limited and expensive. The natural relief is to put the tables on solid-state drives, which offer vastly more capacity per dollar. But moving the embeddings to SSD introduces a new efficiency problem rooted in how SSDs are read, and if that problem is not solved, the cheaper storage delivers too little throughput to serve recommendations at scale.

  2. What people did

    Serve embeddings from SSD and pack co-accessed ones onto the same page

    The starting point is to store the embedding table on SSD and read embeddings from it on demand. Because SSD reads waste bandwidth when embeddings are scattered, prior systems apply hypergraph partitioning: they analyze which embeddings tend to be requested together and place those co-appearing embeddings onto the same SSD page, so one page read returns several useful embeddings at once.

  3. Why it fell short

    An SSD reads whole pages, and a fixed partition traps each embedding in one cluster

    An SSD cannot read a single small embedding; it reads at a coarse page granularity, so fetching one tiny embedding drags in a whole page and most of the transferred data is wasted — read amplification. Co-locating co-accessed embeddings on a page fixes some of this, but hypergraph partitioning has a structural flaw: it divides the embeddings into completely disjoint clusters, so each embedding lives on exactly one page with one fixed set of neighbors. Real queries, though, request many different combinations of embeddings, and an embedding that belongs to one cluster cannot also be co-located with all the other embeddings it co-appears with in other queries. So a single fixed partition can serve only a fraction of the co-appearance patterns, leaving much of the SSD's bandwidth still wasted.

  4. This paper's idea

    Replicate embeddings across many page combinations, and serve them cheaply

    MaxEmbed observes that an SSD has far more capacity than the embedding table strictly needs, and spends that surplus to break the disjoint-cluster limitation. Instead of placing each embedding on exactly one page, it mines the relationships between embeddings and stores popular embeddings in multiple storage combinations through replication — so the same embedding sits on several different pages, each packed with a different set of co-appearing neighbors. Now a query can be served from whichever page best matches the embeddings it wants, so more page reads return mostly-useful data and the SSD's effective bandwidth rises. Because replication adds cost, MaxEmbed pairs this with an online service module for handling embedding query requests, using two key optimizations to reduce the overhead replication introduces. Across various settings, it boosts SSD embedding-serving throughput by up to 18.7 percent.

Why this problem is worth solving

The embedding table is the part of a recommendation system whose size and cost dominate everything else.

A deep learning recommendation model has two parts: dense neural layers, which are small, and embedding tables, which are gigantic. The tables exist because the model must handle sparse categorical features — identifiers drawn from huge vocabularies, like every user or every product — and it represents each such identifier by looking up a learned vector in a table indexed by the identifier. With billions of identifiers and a vector per identifier, the tables reach hundreds of gigabytes to terabytes, far larger than the compute part of the model. Keeping all of that in DRAM is limited by capacity and costly, and DRAM cost scales badly as tables grow, so operators look to solid-state drives, which provide an order of magnitude more capacity per dollar. That makes SSD-backed embedding serving economically compelling: it is the difference between provisioning enough expensive memory for the whole table and storing it cheaply on flash. But serving recommendations has a throughput requirement — each request touches many embeddings and must return quickly — so the SSD path has to deliver embeddings fast, not just cheaply. The tension between the SSD’s cost advantage and its access characteristics is what makes efficient SSD embedding serving a real systems problem worth solving, because the storage choice that makes the tables affordable is also the one that threatens the throughput.

The obstacle is a granularity mismatch that no amount of raw SSD speed removes.

An SSD does not read arbitrary small pieces; its unit of access is a page, typically kilobytes, so any read fetches at least a full page even if only a few hundred bytes of it are wanted. An embedding vector is small, so reading one embedding at random pulls in a whole page of mostly-unrelated data — read amplification, where the bandwidth actually spent is far more than the useful data returned. Because a recommendation query needs many embeddings, this waste multiplies, and the SSD’s effective throughput for embeddings is a fraction of its raw bandwidth. The established remedy is to make each page carry more useful embeddings for a given query by co-locating embeddings that tend to be accessed together. Prior systems find these groups with hypergraph partitioning, treating co-access as a hypergraph and cutting it into clusters that each fit a page. This genuinely reduces amplification, but the method’s premise is a partition: every embedding is assigned to exactly one cluster and one page. Co-access, however, is not partitionable, because a single embedding co-appears with different sets of embeddings across different queries, and it cannot be on the same page as all of them at once. So a disjoint partition captures only some co-appearance patterns and misses the rest, which is the ceiling MaxEmbed is designed to break by giving popular embeddings more than one home.

What people did before, and where it broke

  • Keep the whole embedding table in DRAM. Provision enough memory to hold every embedding for fast access. Where it broke. The tables are hundreds of gigabytes to terabytes, so DRAM capacity and cost make holding them entirely in memory impractical and expensive at scale.
  • Store the table on SSD and read embeddings on demand. Use cheap high-capacity flash and fetch each embedding when needed. Where it broke. SSD reads at a coarse page granularity while embeddings are small, so each fetch amplifies to a whole page and wastes most of the bandwidth, giving low effective throughput.
  • Hypergraph-partition co-appearing embeddings onto shared pages. Analyze which embeddings are requested together and co-locate each cluster on one page. Where it broke. Partitioning forces disjoint clusters, so each embedding has one fixed set of page-mates and cannot serve the many other co-appearance combinations different queries require.
  • Cache hot embeddings in DRAM in front of the SSD. Keep the most popular embeddings in memory to avoid SSD reads. Where it broke. Caching helps for the hottest items but does not fix the read amplification for the long tail served from SSD, and the tail is exactly where SSD serving matters; it is complementary to, not a substitute for, better on-SSD placement.
  • The gap. No SSD embedding system exploited the drive’s surplus capacity to store popular embeddings in multiple co-location combinations, so that any query could be served from a well-packed page, while keeping the extra cost of that replication in check on the online serving path.

How MaxEmbed works

Spending SSD capacity to replicate embeddings into many well-packed pages

  1. new

  2. new

  3. new

  4. new

  5. new

The idea

Trade the SSD’s cheap surplus capacity for the bandwidth it is short on. The pivotal reframing in MaxEmbed is recognizing that an SSD holding an embedding table has two very different resources in very different supply: capacity, which is abundant and cheap (that is why the table is on SSD at all), and effective read bandwidth, which is scarce because of read amplification. A partition-based system implicitly treats capacity as if it were tight, storing each embedding exactly once, and so it is stuck with whatever single co-location that one copy allows. MaxEmbed instead spends the plentiful capacity to buy the scarce bandwidth: by replicating popular embeddings, it stores each of them in several pages, each packed for a different co-appearance pattern, so more queries can be served from a densely-useful page. This is a classic exchange of a cheap resource for an expensive one, made possible precisely because the SSD has capacity to spare. Recognizing that the storage medium chosen for embeddings has capacity to burn and using replication to raise its effective serving throughput is the same capacity-for-speed instinct the embedding-table storage-and-caching work covered elsewhere on this site is built on, arranging embeddings across a storage hierarchy so that abundant, cheap capacity is leveraged to serve the hot data faster.

Escape the disjoint-cluster trap by letting an embedding have many homes. The deeper conceptual point is that co-access is fundamentally not a partition. An embedding does not belong to one group; across the stream of queries it co-appears with many different sets of other embeddings, and hypergraph partitioning, however cleverly it cuts the co-access hypergraph, must ultimately assign each embedding to a single cluster, forcing it to abandon all but one of its co-appearance relationships. MaxEmbed’s replication dissolves this constraint: because a popular embedding can sit on several pages, it can simultaneously honor several of its co-appearance patterns, and the many overlapping groupings that a partition had to collapse into one can coexist. The relationships between embeddings are mined not to divide them but to decide which multiple combinations each popular one should be stored in. Modeling the rich, overlapping structure of how items are accessed together and exploiting it rather than flattening it into disjoint groups is the same graph-structured insight the graph-based recommendation-acceleration work covered elsewhere on this site relies on, treating the co-occurrence relationships among embeddings as a structure to be exploited in full rather than partitioned away.

Make replication pay off by controlling its cost on the serving path. Replication is not free: storing an embedding many times consumes capacity, keeping the copies coherent requires care, and at query time the system must decide which copy of each needed embedding to read to best pack the page reads. If handled naively, these costs could swallow the bandwidth savings, so MaxEmbed does not stop at the storage layout; it builds an online service module for handling embedding query requests with two key optimizations aimed squarely at reducing replication’s overhead. The point is that a good placement and an efficient runtime are inseparable here: the replicated layout creates the opportunity for denser page reads, but only a serving path that can cheaply choose among replicas and manage their cost turns that opportunity into real throughput. Recognizing that a storage-layout idea must be co-designed with an efficient serving and management layer to be worthwhile is the same systems discipline the learning-based flash-translation-layer work covered elsewhere on this site embodies, pairing a smarter mapping of data on the SSD with the runtime machinery that makes the mapping cheap to use. MaxEmbed’s two optimizations are what keep replication a net win rather than a net cost.

Following an embedding query through MaxEmbed

Serving a recommendation query from a replicated SSD embedding store

An illustrative trace, based on the paper's abstract, of how MaxEmbed serves an embedding query. Internal details beyond the abstract were not available for this draft.

Step 1 of 5

Offline: mine co-appearance and choose replicas

Before serving, MaxEmbed analyzes which embeddings tend to be requested together across the workload and decides, for the popular embeddings, which multiple storage combinations to replicate them into, packing many differently-composed pages using the SSD's spare capacity.

Unlike partitioning, an embedding may be placed on several pages, each with a different set of co-appearing neighbors.

Why it works

  • Surplus capacity is the cheap resource to spend. The SSD has far more capacity than the table needs but scarce effective bandwidth → but replicating popular embeddings converts spare capacity into more co-location options → so more queries hit densely-useful pages and effective bandwidth rises.
  • Multiple homes capture co-appearance a partition cannot. Co-access is overlapping, not partitionable, so one cluster per embedding misses most patterns → but replicas let an embedding sit in several co-appearance groups at once → so far more of the workload’s co-appearance combinations are served from packed pages.
  • Denser pages cut read amplification directly. Read amplification comes from pages returning mostly-unwanted data → but choosing a replica whose page holds many of a query’s wanted embeddings raises the useful fraction → so each page read wastes less bandwidth.
  • Fewer reads per query lifts throughput. A query needing many embeddings pays for many amplified reads → but when a single well-packed page supplies several needed embeddings, the read count falls → so the serving throughput per unit of SSD bandwidth climbs.
  • Controlling replication overhead keeps the win net-positive. Replication adds space and choose-a-copy cost that could erase the gains → but the online module’s two optimizations reduce that overhead → so the extra bandwidth becomes a real up-to-18.7-percent throughput improvement rather than a wash.

What it costs (trade-offs and scope)

  • This is a draft built from the paper’s abstract. The full paper was behind a publisher access wall for this draft (no arXiv version, no open-access or author copy located, and no web-archive snapshot), so this account is assembled from the paper’s complete abstract plus standard SSD-embedding background. The exact relationship-mining and replica-selection algorithms, the two online optimizations, the consistency handling, the evaluation setup and baselines, and the per-setting numbers are not reproduced here and should be read from the paper.
  • The reported figure is the abstract’s own. The up-to-18.7-percent throughput improvement is the abstract’s reported result across various settings; the specific workloads, hardware, and baseline systems behind that number are not detailed in this draft, and a maximum is not an average.
  • Replication trades capacity for bandwidth. By design MaxEmbed stores popular embeddings multiple times, so it consumes extra SSD capacity to gain throughput; the approach depends on the SSD having enough surplus capacity, and applying it to items that are not popular enough, or when capacity is tight, would give a poorer trade.
  • Consistency and updates add complexity. Multiple copies of an embedding must be kept consistent when the embedding is updated (as during training or refresh), which is a real cost the abstract only alludes to via its overhead-reduction optimizations; how MaxEmbed handles updates to replicated embeddings is exactly the internal detail this draft cannot report.
  • It targets SSD-served embedding inference specifically. The benefit is for serving huge embedding tables from SSD in recommendation inference; it does not reduce the tables’ size or address bottlenecks in the dense compute, and its gains depend on read amplification being the dominant cost, which holds for SSD embedding serving but not for every deployment.

The evaluation

Serving huge embedding tables from SSD

PropertySSD, no co-locationHypergraph partitioningMaxEmbed
Co-locates co-accessed embeddings
An embedding can serve many co-appearance patterns
Uses SSD surplus capacity
Manages replication overhead
Reported throughput (per abstract)
A qualitative comparison based on the paper's abstract; the paper's main body, including its replica-mining and selection algorithms, the two online optimizations, and its exact baselines and settings, was not available for this draft. The up-to-18.7-percent figure is the abstract's reported throughput improvement.Underlined values explain themselves when clicked.

Check yourself

Check your understanding

  1. 1. Why does serving recommendation embeddings from an SSD cause a bandwidth problem?

  2. 2. What is the structural flaw in the prior hypergraph-partitioning approach?

  3. 3. What is MaxEmbed's core idea?

  4. 4. What should a reader keep in mind about this explainer, and what did the paper report?

Sources

The paper is “MaxEmbed: Maximizing SSD bandwidth utilization for huge embedding models serving” by Ruwen Fan, Minhui Xie, Haodi Jiang, and Youyou Lu (Tsinghua University), published at ASPLOS 2024 (Volume 4), DOI 10.1145/3622781.3674172. It is licensed CC BY-NC.

This explainer is a draft. The paper’s main text was not accessible at the time of writing (there is no arXiv preprint, no open-access or author copy was located, and the ACM Digital Library PDF returned a bot-verification challenge), so it is written from the paper’s complete abstract together with standard background on SSD-served recommendation embeddings. The abstract supplies the problem framing (deep learning recommendation models depend on large embedding tables that raise memory capacity and cost concerns; SSDs offer cost-effective larger capacity but introduce read amplification from the mismatch between embedding size and SSD read granularity; prior SSD embedding storage systems use hypergraph partitioning to co-locate co-appearing embeddings on the same SSD page but divide embeddings into completely disjoint clusters, limiting potential combinations), the core idea (MaxEmbed capitalizes on the extensive storage capacity of SSDs to mine relationships between storage combinations of embeddings with replication, thereby enhancing the effective bandwidth of SSDs), the serving component (an online service module for embedding query request handling with two key optimizations to reduce the overhead brought by replication), and the reported result (MaxEmbed boosts SSD embedding serving throughput by up to 18.7 percent under various settings). The explanation of embedding tables, read amplification, and hypergraph partitioning is standard background and the natural reading of the abstract, not a claim about MaxEmbed’s specific internal design. The paper’s exact relationship-mining and replica-selection algorithms, the two online optimizations, the consistency and update handling, the evaluation setup, and the baselines are not reproduced here and should be read from the paper; no numbers beyond the abstract’s have been introduced.

Related papers on Machine Learning Systems

See all 212 →