xLLM-Rec:突破生成式推荐百毫秒时延壁垒的高性能推理实践
导语
传统推荐系统正面临增长天花板:特征工程红利消退,多阶段级联架构导致目标割裂与误差放大,且GPU利用率(MFU)持续低下。生成式推荐系统通过端到端架构与大语言模型(LLM)能力的引入,正在重塑推荐范式,初步验证了Scaling Law的收益潜力。
然而,范式跃迁并非一蹴而就。将TB级稀疏参数与十B级稠密参数融合,并将端到端时延死死压在百毫秒以内,是工程界面临的严峻考验。京东零售AI Infra团队基于xLLM-Rec(简称xGR)框架的实践,提供了一套高性能、易开发、高可用的破局方案,在生产环境实现了时延降低超60%、性能较TRT-LLM提升约1倍的突破。本文将深度拆解xGR背后的核心挑战与技术解法。
核心问题与挑战
生成式推荐融合了传统DLRM的稀疏特征处理能力与LLM的生成能力,但在落地时横亘着三大技术鸿沟:
- 系统融合挑战:TB级稀疏参数与十B级稠密参数的混合推理,对显存占用与通信开销构成了极大压力,端到端时延极难控制。
- 范式适配瓶颈:LLM原有的计算Kernel与Beam Search机制在推荐场景水土不服。Beam Search存在严重的重复访存与Block频繁拷贝问题;同时,推荐场景特有的混合Mask序列结构极其复杂,导致开源FlashAttn等实现性能极差,不仅存在大量冗余计算,更面临显存爆炸风险。
- 生态切换阵痛:搜推业务长期依赖TensorFlow生态,向Torch生态切换时,如何保证底层算法工程师的平滑过渡是一大挑战。
面对上述难题,业界主流的编译优化路线或简单复用大模型推理框架的方案,均无法在时延与灵活性上满足生成式推荐的严苛要求。
方案与实践
针对Beam Search与混合Mask两大核心瓶颈,xGR提出了四大创新优化机制,彻底重构了计算流程。
1. KV Cache分级管理:消除冗余存储
在Beam Search生成推荐Item(通常自回归生成3个Token代表一个Item)的过程中,不同候选序列存在大量相同的上下文。xGR通过分级管理KV Cache实现解耦:
- Prefill阶段:生成Shared Cache,所有候选序列共享同一份Cache,彻底避免重复存储与冗余计算。
- Decode阶段:针对非共享部分,采用Token粒度的动态管理,实现细粒度的显存分配与更新。
2. xAttention三阶段计算:击破访存带宽瓶颈
传统Page Attention在推荐场景下会重复加载共享Cache。xGR将Attention计算拆解为三个独立阶段:
- 统一处理共享部分:一次性计算所有序列共享的上下文注意力。
- 独立计算非共享部分:各候选序列仅计算自身独有的增量部分。
- Softmax融合归一化:将共享与非共享的结果通过融合算子进行统一归一化。
这一拆解将重复的HBM访存转化为计算流水的并行,极大释放了带宽压力。
3. Schedule优化:计算下沉与早停机制
原生Beam Search在每一步都需要将数据从Device拉回Host进行打分、排序与筛选,造成极高的D2H传输与CPU耗时。xGR的核心策略是:
- 整图下沉Device执行:将Beam Search的核心逻辑(打分、排序、筛选)完整下沉至GPU Device侧高效执行。
- 约束解码早停机制:结合推荐业务规则,在生成过程中一旦满足约束条件即刻早停,避免无效自回归步数。
4. 增量分段Mask计算:按需计算实现零冗余
面对复杂的混合Mask序列,全局搬运与计算是显存爆炸的根源。xGR采用增量分段策略:
- 区域化分块:按掩码稀疏性将计算拆分为多个子Stage。
- 局部掩码内部生成:避免构建全局庞大Mask矩阵。
- 级联Online Softmax:分段计算后通过级联方式合并结果。
- 无效区域零计算:对Mask为0的区域直接跳过QK计算与Softmax及KV访存,彻底剔除无效冗余。
原则/方法论沉淀
从xGR的工程实践中,我们可以提炼出应对复杂推理场景的四条核心原则:
- 共享与独立计算解耦:识别计算图中的共享上下文与独立生成部分,从存储和计算双重维度剥离,消除最昂贵的重复访存。
- 计算下沉与并行化:将排序、筛选等高频控制逻辑下沉至Device侧,用计算换通信,消除D2H传输造成的时延毛刺。
- 按需计算与零冗余:拒绝黑盒全局计算,针对稀疏掩码进行细粒度分块与裁剪,做到“无数据则无计算”。
- 分层缓存与细粒度拆分:将自定义掩码注意力转化为前缀KV复用与增量片段计算,化整为零,控制显存峰值。
总结与行动建议
生成式推荐正在重新定义推荐系统。京东零售的xGR框架通过KV Cache分级、xAttention三阶段、Schedule下沉与增量分段Mask四大创新,成功在首页推荐与搜索场景将时延降低超60%,性能较TRT-LLM提升约1倍,同时显著提升了UCTR与UCVR等核心业务指标,节省了近一半的GPU资源。
对于正在进行或即将启动生成式推荐架构升级的团队,建议采取以下行动:
- 审计访存冗余:优先梳理Beam Search与复杂Mask中的共享Prefix,实施Cache分级与解耦计算。
- 重构调度边界:审视Host与Device的数据交互,将控制流尽可能编译下沉至Device。
- 精细化Mask处理:放弃粗放的通用Attention实现,针对业务Mask分布定制分块与Online Softmax策略。
开放问题与延伸方向
- 时延降低60%和性能较TRT-LLM提升1倍的声明,其基准测试的具体硬件配置、Batch Size设定与输入序列长度分布是怎样的?(关联:生产验证的基准条件需结合具体业务场景客观评估)
- TB级稀疏参数与十B级稠密参数在xGR框架中的显存占用比例与实际通信开销是如何量化的,稀疏特征更新是否成为新的瓶颈?(关联:系统融合挑战的进一步量化拆解)
- 从TF生态强制切换到Torch生态,是否会让长期依赖TF庞大工具链的底层算法工程师产生强烈的抵触情绪与迁移阵痛?(关联:生态切换挑战中的人效与工程习惯因素)
- Beam Search结合约束解码早停机制,是否在直觉上牺牲了推荐系统的长尾探索能力,从而在无意识中加剧信息茧房?(关联:Schedule优化早停机制对推荐多样性的潜在影响)
- xAttention的级联Online Softmax在处理极长上下文或极端分布偏移的Item时,是否存在数值溢出或精度损失导致推荐结果严重偏离的风险?(关联:三阶段计算在极端分布偏移下的数值鲁棒性)
- 将Beam Search核心逻辑下沉Device侧执行,若GPU发生显存溢出或计算异常,是否会导致整个推荐链路不可恢复的崩溃而缺乏CPU兜底?(关联:计算下沉带来的容错与隔离风险)
- 增量分段Mask计算实现零冗余的按需计算思想,是否可以横向迁移到多模态推荐中视频或图像特征的稀疏注意力计算上以获取更大收益?(关联:按需计算原则的横向泛化潜力)
- KV Cache在Prefill阶段共享缓存,对于具有大量重叠用户历史行为的推荐场景,是否带来了超线性的并发吞吐提升机会?(关联:KV Cache分级管理在特定业务分布下的额外收益)
- 面向未来扩散生成范式的演进,现有的自回归式KV Cache分级管理机制是否会被完全颠覆,是否有必要提前布局非自回归的缓存架构?(关联:生成范式演进趋势下的架构前瞻布局)
- 如果不采用整图下沉Device,而是利用异构计算(如NPU处理稀疏Embedding,GPU处理稠密计算),是否能提供一种更平滑的生态过渡与性能平衡方案?(关联:系统融合与生态切换的替代解法探索)
- 在xGR的四大创新中,哪一项对突破百ms时延约束的贡献度最高,下一步的工程优化重心应优先放在进一步压榨Kernel性能还是系统级调度上?(关联:工程优化资源的优先级分配)