广告 CTR 预估模型的 GPU 推理优化:大规模稀疏特征 Embedding 的显存管理与混合精度计算

广告 CTR 预估模型的 GPU 推理优化:大规模稀疏特征 Embedding 的显存管理与混合精度计算

广告 CTR 预估模型的 GPU 推理优化:大规模稀疏特征 Embedding 的显存管理与混合精度计算

一、CTR 模型的特征爆炸:当 100 亿参数中 99% 是 Embedding Table

CTR 预估模型的架构特性与其他深度学习模型完全不同。一个典型的 Wide&Deep 或 DLRM 模型,参数量可达 100 亿——但其中 99.9% 的参数属于稀疏 Embedding Table(将用户 ID、Item ID、类目、标签等离散特征映射为稠密向量),只有 0.1%(约 1000 万参数)属于密集层(MLP、Attention 等深度计算层)。
以 Meta 的 DLRM 为参考,在 Criteo 数据集上:

  • 用户侧 Embedding:(用户ID 1 亿 × 64 维 + 其他用户特征 100 万 × 32 维)× 4 字节 ≈ 25.6GB
  • 商品侧 Embedding:(商品ID 500 万 × 64 维 + 其他商品特征 10 万 × 32 维)× 4 字节 ≈ 1.3GB
  • Dense MLP 层:约 1000 万参数 × 4 字节 ≈ 40MB

总计约 27GB 的 Embedding ——远超单 GPU 的显存容量(A100-80G 也仅 80GB,要考虑模型训练时优化器状态和显存占用)。但这 27GB Embedding 不是全部被同时访问——每个广告请求只涉及 1 个用户(查找 1 个用户 Embedding)和 200-500 个候选 Item(查找 200-500 个 Item Embedding)——总共只需要访问 200-500 个 Embedding 向量。这就催生了 Embedding 的分层存储架构。

二、Embedding 的分层缓存:GPU HBM → CPU DRAM → SSD → 远程 PS

在该分层架构中,数据流向遵循从高速到低速的查找顺序。当 Embedding Lookup 发生时,系统优先访问 GPU HBM 中的热点数据,若未命中则依次回退至 CPU DRAM、NVMe SSD,最后才访问远程 Parameter Server。查找到的 200-500 个向量随后进行聚集,并送入 MLP Dense 层完成推理,最终输出 CTR 预测结果。

四层缓存的分层设计:

  • L0 — GPU HBM:缓存热点 Embedding(高活跃用户、热销商品),约 5GB。请求的 Embedding Lookup 首先查 L0——命中延迟 < 1μs,GPU kernel 直接读取。
  • L1 — CPU DRAM:存储温 Embedding(近期活跃用户、常规商品),约 15GB。L0 miss 时通过 PCIe DMA 从 CPU 传输到 GPU,延迟约 5-10μs。
  • L2 — NVMe SSD:存储冷 Embedding(低活跃用户的 Embedding),约 10GB。L1 miss 时从 SSD 加载到 CPU 内存再传 GPU,延迟约 100μs。
  • L3 — 远程 Parameter Server:全量 Embedding Table 的集中存储。当本地节点都没有目标 Embedding 时,从 PS 拉取,延迟约 1-5ms。

LRU 缓存替换策略在各层独立运作。当一个冷用户的 Embedding 从 L2 被加载到 L1 后,L1 满时会逐出最久未使用的 Embedding 回 L2。逐出策略的关键是避免颠簸(Thrashing)——如果 L1 的容量刚好被周期访问的 Embedding 填满,每次加载一个新 Embedding 就要逐出一个即将再次被访问的 Embedding,缓存命中率极低。监控各层的 Cache Miss Rate——L0 的 Hot Hit Rate 应 > 95%,L1 的 Miss Rate 应 < 5%,否则需要增加对应层级的缓存容量。

三、混合精度计算的 CTR 模型优化:如何在不损失 AUC 的前提下加速推理

CTR 模型的推理延迟瓶颈在 Dense MLP 层——虽然参数只占 0.1%,但计算量(200-500 个候选 × 3-5 层 MLP × 每层 1024 神经元)是主要的浮点运算。混合精度方案在 Embedding 和 MLP 层采用不同精度:

  • Embedding 向量:FP32(32 位浮点)——对精度高度敏感,INT8 量化会导致 AUC 下降 > 0.3%(CTR 场景中的 AUC 下降 0.1% 可能损失千万级的广告收入)
  • MLP 层输入激活值:FP16(16 位浮点)——将 Embedding 向量 Lookup 后转为 FP16
  • MLP 层权重:FP16 + Tensor Core 加速——A100 的 Tensor Core 在 FP16 下提供 312 TFLOPS,是 FP32(19.5 TFLOPS)的 16 倍
  • MLP 层输出(CTR 预测值):FP32——精度要求高,转回 FP32

混合精度的数值稳定性需要特殊处理。MLP 中间的激活值在 FP16 下可能出现梯度消失或数值溢出——当 500 个 Embedding 向量拼接后的输入向量很大(维度 > 10,000 且元素值 > 1.0),FP16 的范围(±65504)不够,出现上溢(Inf)。解决方案是在 Embedding Lookup 后加一个torch.clamp(x, -10.0, 10.0)的钳位操作,限制输入值的范围并转为 FP16,确保下游 MLP 不会出现 Inf。这个 clamp 对 AUC 的影响 < 0.01%——可以忽略。

四、推理引擎的 CUDA Kernel 优化:Batch Matrix Multiply 的极致调优

CTR 模型中 MLP 层的核心计算是批量矩阵乘法(Batch MatMul)。500 个候选 Item × 用户向量(Batch=500),与共享的 MLP 权重矩阵做 MatMul。PyTorch 的torch.baddbmm提供了 CPU 友好的批量 MatMul 接口,但在 A100 上对 batch=500 这个中间规模并没有做到极致优化。

通过 cuBLAS 的cublasGemmStridedBatchedExAPI,可以为 batch=500 选择最优的 tiling 策略——将 500 个小矩阵合并为少数几个大矩阵(Batch Coalescing),利用 Tensor Core 的矩阵分块计算一次处理更多浮点运算。在 A100 上,手动调优的 cuBLAS Kernel 比 PyTorch 默认的baddbmm快 1.8-2.2 倍——将 MLP 推理延迟从 2.5ms 降至 1.2ms。

手写 CUDA Kernel 虽然性能极致,但维护成本高——PyTorch 版本升级、CUDA 工具包变更、GPU 代际更换都可能让手写 Kernel 失效或性能退化。建议使用 PyTorch 的torch.compile+Inductorbackend,它能自动做 Kernel Fusion 和 Batch Coalescing 优化,无需手写 CUDA。在大多数场景下,torch.compile能达到手写 Kernel 性能的 85%-90%,维护成本接近于零。

五、总结

CTR 模型的推理优化核心在于 Embedding 的分层缓存和 Dense MLP 的混合精度加速。27GB 的 Embedding 不能一次性放在 GPU 显存中,但每个请求只访问 200-500 个 Embedding 向量——热点 Embedding 缓存在 GPU HBM 上,命中率 95% 以上。LRU 逐出策略在各层独立运作,避免跨层的缓存颠簸。

混合精度计算(Embedding FP32、MLP FP16 + Tensor Core)将 MLP 推理延迟降低到原始 FP32 的 1/8——对 AUC 的影响 < 0.01%。输入 clamping 是 FP16 数值稳定性的必要保护措施。

批量推理是 CTR 模型性能的另一大杠杆。500 个候选 Item 一次性做 MLP 推理,利用 cuBLAS 的 Batch Coalescing 和 Tensor Core 加速,MLP 推理延迟从串行的 500×2ms=1000ms 降至 1.2ms。尽管torch.compile还不能完全替代手写 CUDA Kernel 的性能极致,但在大多数工程场景中它的自动优化能力已经足够——优先选择可维护的自动优化而非难以维护的性能极致。