中文

LoKA:面向大规模推荐模型的低精度卷积核应用

机器学习 2026-05-15 v2 人工智能

摘要

近期 GPU 代际通过使用低精度算术(如 FP8)实现了显著提升的 FLOPs。尽管此技术在大型语言模型(LLM)中取得成功,但在大型推荐模型(LRM)中的应用仍有限。这是因为 LRM 对数值敏感,计算主要由小矩阵乘法(GEMM)和归一化操作组成,且训练环境计算密集。直接将 FP8 应用于 LRM 常导致模型质量下降并延长训练时间。这些挑战源于 LRM 工作负载的本质特性,仅凭引入更好的 FP8 内核无法解决。相反,需要系统-模型协同设计方法才能成功集成 FP8。我们提出 LoKA(Low-precision Kernel Applications),通过三个原则使 FP8 在 LRM 中实用:在真实分布下进行剖析以确定低精度安全区域,协同设计模型组件与硬件以扩大安全范围,并协调内核库以最大化收益。具体而言,LoKA Probe 是一种基于统计学的在线基准测试方法,用于学习激活和权重统计信息,并量化每层误差,从而识别出 FP8 采用的安全/不安全、快速/慢速位置。LoKA Mods 是一组可复用的模型适配器,通过 FP8 提升数值稳定性和执行效率。LoKA Dispatch 是一种运行时系统,利用 LoKA Probe 的统计洞察,选择满足精度要求的最快 FP8 内核。

关键词

引用

@article{arxiv.2605.10886,
  title  = {LoKA: Low-precision Kernel Applications for Recommendation Models At Scale},
  author = {Liang Luo and Yinbin Ma and Quanyu Zhu and Vasiliy Kuznetsov and Yuxin Chen and Jian Jiao and Jiecao Yu and Buyun Zhang and Tongyi Tang and Xiaohan Wei and Yanli Zhao and Zeliang Chen and Yuchen Hao and Venkatesh Ranganathan and Sandeep Parab and Yantao Yao and Maxim Naumov and Chunzhi Yang and Shen Li and Ellie Wen and Wenlin Chen and Santanu Kolay and Chunqiang Tang},
  journal= {arXiv preprint arXiv:2605.10886},
  year   = {2026}
}

备注

Accepted to ISCA'26