LoKA: Low-precision Kernel Applications for Recommendation Models At Scale

TL;DR

LoKA is a framework that enables the practical use of low-precision FP8 arithmetic in large recommendation models by profiling, model adaptation, and runtime kernel selection.

Contribution

It introduces a system-model co-design approach with three components: profiling, model modifications, and runtime dispatching for FP8 in LRMs.

Findings

LoKA Probe accurately identifies safe FP8 application sites.

LoKA Mods improves numerical stability and efficiency with FP8.

LoKA Dispatch dynamically selects optimal FP8 kernels based on statistical profiling.

Abstract

Recent GPU generations deliver significantly higher FLOPs using lower-precision arithmetic, such as FP8. While successfully applied to large language models (LLMs), its adoption in large recommendation models (LRMs) has been limited. This is because LRMs are numerically sensitive, dominated by small matrix multiplications (GEMMs) followed by normalization, and trained in communication-intensive environments. Applying FP8 directly to LRMs often degrades model quality and prolongs training time. These challenges are inherent to LRM workloads and cannot be resolved merely by introducing better FP8 kernels. Instead, a system-model co-design approach is needed to successfully integrate FP8. We present LoKA (Low-precision Kernel Applications), a framework that makes FP8 practical for LRMs through three principles: profile under realistic distributions to know where low precision is safe, co-design model components with hardware to expand where it is safe, and orchestrate across kernel libraries to maximize the gains. Concretely, LoKA Probe is a statistically grounded, online benchmarking method that learns activation and weight statistics, and quantifies per-layer errors. This process pinpoints safe and unsafe, fast and slow sites for FP8 adoption. LoKA Mods is a set of reusable model adaptations that improve both numerical stability and execution efficiency with FP8. LoKA Dispatch is a runtime that leverages the statistical insights from LoKA Probe to select the fastest FP8 kernel that satisfies the accuracy requirements.