Meta Optimizes Inference for Recommendation Systems

Meta introduced In-Kernel Broadcast Optimization (IKBO) to eliminate the bottleneck of duplicating the same per-request user embedding across candidates in recommendation system inference. In the existing approach, duplication cost grows linearly as the number of candidates increases. The key idea is to treat broadcast not as tensor duplication but as a data layout problem, absorbing it directly inside the kernel.

IKBO separates Request-Only (RO) embeddings, which are common across the entire request, from Non-Request-Only (NRO) embeddings, which are per-candidate. The inference runtime does not create a duplicated tensor; instead it passes only a candidate-to-user mapping, and both batch sizes are processed as-is only at the point where interaction operations are required. While existing system-level broadcast or net-splitting approaches like ROO only worked around duplication, IKBO removes the broadcast itself at the base operator layer.

There are two types of optimization.

• Type I: For decomposable operations, RO and NRO are first computed independently and combined only at the end, reducing memory bandwidth and computation.
• Type II: Broadcast is handled inside the kernel, removing unnecessary data movement and reducing IO bottlenecks.

This application touches all three layers: kernels, compilation specification, and inference runtime. The ML compiler needs to know the dynamic shape range per operator in order to select the correct kernel. Both direct adoption, where IKBO kernels are embedded directly in the model definition, and inference-time transformation, where existing operators are swapped for IKBO versions at the inference stage, are possible, and the same kernels can be used as-is for training where the candidate/user ratio is greater than 1. This structure has been deployed across GPU and MTIA from early ranking to late ranking, and in co-designed models it reduced compute-intensive net latency by up to 2/3. It also serves as the scalability foundation for the Meta Adaptive Ranking Model, which serves LLM-scale models for ads.

In kernel deep-dive case 1, Linear Compression Embedding (LCE), instead of duplicating the same user embedding for each candidate, users and candidates are separated along the K axis, then two GEMMs are each run at their own natural batch size, and only the small compressed result is broadcast. In a setting where the candidate/user ratio is about 70:1, the user batch shrank from 1024 -> about 15, and execution time decreased from 1.944ms -> 1.389ms, a 28.5% reduction.

Subsequently, to resolve a memory alignment issue, K was padded to a multiple of 8, eliminating the inefficiency where cp.async was being split into 4-byte units. On H100 SXM5, through four stages of progressive co-design, this achieved ~4x speedup, finishing with warp-specialized multi-stage fusion using TLX.

In kernel deep-dive case 2, Flash Attention, IKBO converted the kernel from IO-bound to compute-bound, recording 621 BF16 TFLOPs. Compared against the non-co-designed CuTeDSL FA4 Hopper baseline, this yielded a 2.4x throughput improvement for the kernel alone, and 6.4x on a full basis including broadcast.