Home · ASPLOS 2023 · Machine Learning Systems

RECom: A Compiler Approach to Accelerating Recommendation Model Inference with Massive Embedding Columns

Zaifeng Pan, Zhen Zheng, Feng Zhang, Ruofan Wu, Hao Liang, Dalin Wang and 4 others

Conference
ASPLOS 2023
Published
2023
Pages
268-286
Citations
6

Deep recommendation models spend almost all their inference time not in the neural network but in thousands of tiny embedding-lookup subgraphs, one per feature — and running those on a GPU is so death-by-a-thousand-kernel-launches that the GPU sits two-thirds idle and is actually slower than a CPU; this work from Renmin University of China and Alibaba builds RECom, the first compiler that fuses all those thousands of columns into a single GPU kernel, reasons about their unknown shapes symbolically, and strips the redundant safety code the framework injected, cutting end-to-end latency 6.61 times over the best baseline.

The argument in one glance

  1. Problem

    A recommendation model spends almost all its time in thousands of tiny lookups

    A deep recommendation model, the kind that predicts what you will click, has two parts: a small neural network, and an embedding layer that turns each input feature into a vector by looking it up in a table. To get high accuracy, companies generate thousands of statistical features, each with its own little subgraph — an embedding column — that preprocesses the feature, translates it to an index, checks it for safety, looks up the vector, and post-processes it. In Alibaba's models these thousands of embedding columns can be more than 99 percent of the inference time. And they must be optimized automatically, from the computation-graph representation alone, since one model can be millions of lines of it.

  2. What people did

    Point an existing machine-learning compiler at the model

    Machine-learning compilers like XLA and AStitch already optimize models automatically, taking the computation graph and generating fast code, and their central trick is operator fusion — merging small operators together to cut memory traffic and kernel-launch overhead. The obvious move is to run one of these compilers on the whole recommendation model and let it fuse the embedding columns the same way it fuses a neural network.

  3. Why it fell short

    The compilers choke on thousands of columns, dynamic shapes, and injected redundancy

    They do not help. Existing compilers fuse within a subgraph but map each embedding column to its own kernels, so a thousand-column model produces over ten thousand tiny kernels and the GPU spends only a third of its time actually computing, most kernels using under six percent of a streaming multiprocessor. Worse, the columns' shapes are unknown until runtime, which blocks the fusion, and frameworks inject defensive safety code that, in extreme cases, is 80 percent of the work. On these models a GPU ends up slower than a CPU.

  4. This paper's idea

    Build a compiler purpose-made for massive embedding columns

    RECom is the first compiler aimed squarely at embedding columns. It fuses all of them into one GPU kernel, mapping each column to its own thread block so the thousands of independent columns run in parallel with no per-column launch overhead. It tames the dynamic shapes by reasoning about them symbolically, decoupling shape computation from tensor computation so the graph can be optimized. And it analyzes each column in full context to strip the redundant safety code the framework inserted, while offloading the few huge tables to the CPU.

Why this problem is worth solving

The heart of the matter is a workload whose shape is the opposite of what a GPU is built for, so it is worth being precise.

A deep recommendation model maps input features — a user’s age, a click history, an item’s category — into vectors and feeds them to a neural network to predict, say, a click-through rate. Each feature gets an embedding column: a small subgraph that preprocesses the raw value (splitting strings, scaling numbers), translates it into a lookup index (by hashing or bucketing), applies safety guarantees (dropping out-of-range indices, filling in a default when a feature is missing), looks up the corresponding rows of an embedding table and reduces them into one vector, and post-processes the result. One column corresponds to one embedding table. The reason there are so many is combinatorial: an app has dozens of scenes, each with dozens of raw features, and each raw feature spawns statistical variants — its minimum, maximum, sum, average over a day, a week, a month — and each of those is embedded to improve accuracy. The count reaches into the thousands. Each column is small and memory-bound, doing a little indexing and a table lookup rather than heavy arithmetic. Thousands of small, irregular, memory-bound subgraphs are exactly the workload a GPU, tuned for a few large dense matrix multiplies, handles worst — yet on real models these columns are over 99 percent of the inference latency, so optimizing them is the whole game.

The reason existing compilers cannot help is that every assumption they make about the workload is wrong here.

A machine-learning compiler optimizes by fusing operators — merging many small operations into one kernel so the GPU launches fewer kernels and moves less data through memory. This works beautifully for a neural network, which is a modest number of large, regularly-shaped, statically-known tensor operations. Embedding columns violate all three properties. First, there are thousands of columns and each has many tiny operators, so a compiler that fuses only within a subgraph and gives each column its own kernels ends up launching over ten thousand kernels; each launch and synchronization costs a few microseconds, and multiplied across tens of thousands of kernels the accumulated overhead becomes the majority of the runtime, leaving the GPU only a third active and most kernels occupying under six percent of a streaming multiprocessor. Crucially, the compilers process the columns sequentially and never exploit the fact that thousands of independent columns could run at once. Second, recommendation models have dynamic shapes — tensor sizes that depend on the input contents and are unknown at compile time — which breaks the shape reasoning fusion depends on. Third, high-level framework APIs inject defensive safety code that is redundant in context but that the compiler cannot see through. The consequence is stark: a GPU, normally far faster than a CPU for machine learning, is actually slower here. Fixing that requires a compiler built for this workload’s structure, not a general one.

What people did before, and where it broke

  • General ML compilers (XLA, AStitch). Fuse operators automatically and generate GPU code, as they do for neural networks. Where it broke. They fuse only within a subgraph and process columns sequentially, so a thousand-column model still yields over ten thousand kernels, the GPU sits two-thirds idle, and dynamic shapes limit the fusion further — they beat plain TensorFlow-GPU by only about ten percent.
  • Hand-written libraries for embedding operations (HugeCTR). Provide fast implementations of specific operations like the table lookup. Where it broke. They cover only the patterns they pre-built and require developers to manually rebuild models on their interfaces — impractical across thousands of models with varied column structures, and tuned mostly for large training batches, not inference.
  • Expert teams optimizing models by hand. Extract known column patterns and replace them with hand-tuned kernels, sometimes with CPU vector instructions. Where it broke. It is labor-intensive and brittle — the internal manual solution only worked on models whose subgraphs matched, failing on the rest, and no one can hand-rewrite a model that is millions of lines of graph representation.
  • Put the DNN on the GPU and leave embeddings on the CPU. Since the GPU is bad at the columns, run only the neural network there. Where it broke. The neural network is the cheap part, so moving it to the GPU barely helps, and the cross-device copies add overhead — this configuration does not beat plain CPU execution.
  • The gap. Every approach either optimized the wrong part, covered only fixed patterns, demanded manual effort that does not scale to thousands of models, or ignored the inter-column parallelism, the dynamic shapes, and the injected redundancy that define this workload. What was missing was an automatic compiler purpose-built for massive embedding columns — which is exactly what recommendation inference needs.

How RECom works

Compiling thousands of embedding columns into one efficient GPU kernel

  1. new

  2. new

  3. new

  4. new

  5. new

The idea

Turn thousands of tiny kernel launches into one well-utilized kernel that runs the columns in parallel. The GPU is not slow here because the arithmetic is hard; it is slow because it launches and synchronizes tens of thousands of trivial kernels, spending most of its time on overhead and idling between fragments, while never exploiting that the thousands of columns are independent. RECom’s central move is to fuse every column into a single kernel and map each column to a thread block, so the hardware runs hundreds of columns at once and the per-launch cost vanishes — its fused kernel occupies a streaming multiprocessor forty times more fully than the scattered baseline. Recognizing that a machine-learning serving workload is bottlenecked not by compute but by how the many small memory-bound pieces are scheduled onto the GPU, and fixing it by keeping the hardware busy rather than by faster math, is the same diagnosis the Mixture-of-Experts inference work covered elsewhere on this site makes, where the GPU also stalls on movement between small pieces until a smarter schedule keeps it occupied.

Make the unknown shapes symbolic so the graph becomes optimizable at all. The reason general compilers stall is that embedding-column tensor sizes depend on the input and are unknown at compile time, and that missing information blocks the analysis fusion needs. RECom represents every shape as a symbolic expression, propagates it through the graph, and collapses each snarled shape-computation subgraph into one unified operator, cleanly separating the reasoning about sizes from the computation on values. That separation is what lets it then remove operators kept alive only by a spurious shape dependency. Building an explicit symbolic model of the computation so that transformations invisible to a naive, concrete view become possible is the same enabling step the compute-in-memory compiler work covered elsewhere on this site takes, constructing the representation it needs before it can map hard computation onto its target rather than optimizing blindly.

Specialize away the defensive code that only exists because the library could not know the context. Framework APIs wrap every lookup in safety guarantees — bounds checks, missing-value filling — because the library author cannot know whether a given call actually needs them, so they are always emitted, and in extreme cases they are most of the work. RECom analyzes the whole column and proves what the library could not: that the indices are already in range, that a feature is one-hot so its reduction is pointless, that empty-row filling can be merged into the lookup. Then it deletes the now-provably-unnecessary work. Stripping generic, defensive machinery down to only what a specific context actually requires, by studying that context closely, is the same instinct behind the GPU-kernel specialization work covered elsewhere on this site, which reproduces only the behavior a particular use actually exercises rather than carrying the general case everywhere.

Following one embedding column through RECom

Compiling a hashed string-feature column

Tracing one column shows how identification, symbolic shapes, and redundancy removal compound.

Step 1 of 5

A string feature enters as its own little subgraph

A feature like a user's city arrives as a string. Its embedding column hashes the string into an index, guards the index against out-of-range and missing values, looks up the embedding table row, reduces if there are several, and reshapes the result. In the raw graph this is a tangle of tiny operators, indistinguishable from the thousands of other columns around it.

The model is delivered only as a computation graph, with no labels saying which operator belongs to which column.

Why it works

  • Fusing every column into one kernel removes the launch-and-sync overhead that dominated the runtime. Thousands of columns compiled as separate kernels mean tens of thousands of launches and synchronizations, each costing microseconds → but fusing them into a single kernel pays the launch cost once → so the GPU’s active time jumps from a third to nearly full and its utilization climbs from under 0.06 to 2.43 waves per multiprocessor.
  • Mapping one column to one thread block finally exploits the parallelism the columns already have. The thousands of embedding columns are independent and each is light, yet baselines run them sequentially → but a GPU can run hundreds of thread blocks at once, so assigning a column per block runs them in parallel → so the model’s natural inter-column parallelism at last matches the hardware’s structure.
  • Symbolic shapes decouple size reasoning from value computation, unblocking optimization. Dynamic shapes unknown at compile time entangle shape and tensor operators and freeze the graph → but representing shapes as symbols and collapsing each shape subgraph into one unified operator separates the two → so operators kept alive only by a spurious shape dependency become removable.
  • Whole-context analysis deletes defensive code the library was forced to always emit. Framework APIs wrap every lookup in bounds checks and missing-value filling because the author cannot know the call site → but analyzing the full column proves the indices are already valid and the feature one-hot → so the redundant guards and reductions, up to 80 percent of the work in extreme cases, are stripped away.
  • Co-running the few huge tables on the CPU fits big models on one GPU nearly for free. A handful of embedding tables are enormous but rarely used, and they would blow the GPU’s memory → but keeping just those and the string work on the CPU and merging the transfer into one shot costs almost nothing → so a single GPU serves a model larger than its memory while losing under five percent of speed and saving over ninety percent of GPU memory.

What it costs (trade-offs and scope)

  • Column identification is heuristic and leans on how embedding columns look. RECom relies on domain facts — one column per embedding table, tables are gathered from while weights are matrix-multiplied — so a model whose structure violates these assumptions could be partitioned wrongly.
  • It is built as a TensorFlow add-on. The implementation is 12.5 thousand lines of C++ tied to TensorFlow’s graph and custom-operator interfaces; the authors argue the ideas transfer to other frameworks, but the artifact does not.
  • At large batch sizes the CPU side becomes the bottleneck. The string preprocessing kept on the CPU cannot keep up with the GPU consumer at large batches — up to 60 percent of the time with only four cores — so the speedup shrinks unless more CPU cores are added.
  • Stripping safety guards depends on the analysis being right. RECom deletes framework-inserted bounds checks and fillings after proving them redundant; the paper argues this is safe because the indices are already valid, but it is removing defensive code, so the analysis carries the correctness burden.
  • The wins scale with how column-heavy and redundant the model is. On a model with only ten columns RECom merely matches a CPU; its large speedups appear as the column count and injected redundancy grow, so a lighter model sees less benefit.
  • It accelerates inference latency, not accuracy, and is orthogonal to other layers. RECom changes nothing about model quality and does not replace request schedulers like DeepRecSys or training-oriented lookup libraries like HugeCTR; it composes with them rather than superseding them.

The evaluation

Three ways to run a thousand-column recommendation model on a GPU

PropertyTensorFlow-GPUGeneral ML compilers (XLA, AStitch)RECom
How columns are handledEach operator its own kernelFused within a column, run sequentially
Kernels for 1,000 columns
GPU utilization (waves per SM)Under 0.06, similar distribution
Dynamic shapesHandled, but block fusionLimit fusion granularity
Embedding speedup over TF-GPUReference point
Three ways to run recommendation inference with thousands of embedding columns, on an NVIDIA T4 GPU across four Alibaba production models and two synthetic DLRM models. General ML compilers barely beat plain TensorFlow-GPU because they still emit thousands of tiny kernels and run columns sequentially; RECom fuses every column into a single well-occupied kernel, reasons about dynamic shapes symbolically, and strips injected redundancy, reaching 6.61 times lower end-to-end latency than TF-CPU, 51.45 times over TF-GPU, and 1.91 times higher throughput under a service-level agreement, with graph-level optimizations alone removing 84 percent of the embedding operators.Underlined values explain themselves when clicked.

Check yourself

Check your understanding

  1. 1. Why is a GPU, normally far faster than a CPU for machine learning, actually slower on these recommendation models?

  2. 2. What is RECom's central code-generation idea for the thousands of embedding columns?

  3. 3. How does RECom handle the dynamic shapes that block general compilers, and why does it matter?

  4. 4. Which statement best reflects the scope and limits of RECom?

Sources

The paper is “RECom: A Compiler Approach to Accelerate Recommendation Model Inference with Massive Embedding Columns” by Zaifeng Pan, Zhen Zheng, Feng Zhang, Ruofan Wu, Hao Liang, Dalin Wang, Xiafei Qiu, Junjie Bai, Wei Lin, and Xiaoyong Du, published at ASPLOS 2023, DOI 10.1145/3623278.3624761. The authors are at Renmin University of China and Alibaba Group, and the source code is open at github.com/AlibabaResearch/recom.

This explainer is written from the full 19-page text. Everything above comes from the paper’s own sections: the introduction and motivation establishing that embedding columns can exceed 99 percent of inference latency and that the three challenges are massive tiny operators, dynamic shapes, and framework-injected redundancy (with XLA generating over ten thousand kernels at only 33 percent GPU active time and most kernels under 0.06 waves per streaming multiprocessor); the massive embedding column codegen with domain-knowledge column identification, group-based divide-and-stitch intra-column code generation, compile-time and runtime combined buffer allocation, and CPU-GPU co-running of large tables and string work; the symbolic shape expression inference with its four propagation patterns and the unified ShapeConstruct reconstruction that decouples shape from tensor computation; the embedding column subgraph optimizations removing unnecessary index-range guards (up to 17.7 percent of TF-GPU time), merging empty-row filling into a modified lookup (1.49 times), and pruning one-hot reductions with sparse-to-dense conversion (1.86 times); and the evaluation on four Alibaba DBMTL production models and two synthetic DLRM models on an NVIDIA T4, reporting 6.61 times, 51.45 times, and 8.96 times latency speedups over TF-CPU, TF-GPU, and TF-CPU-GPU, 1.91 times throughput under a service-level agreement, 2.01 and 1.97 times throughput per dollar and per watt, graph-level optimizations removing 84.42 percent of operators, the 14.01 to 116.23 times advantage over TF-GPU where XLA and AStitch gain only about ten percent, and the fused kernel reaching 2.43 waves per streaming multiprocessor.

Related papers on Machine Learning Systems

See all 212 →