广告 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 的性能极致，但在大多数工程场景中它的自动优化能力已经足够——优先选择可维护的自动优化而非难以维护的性能极致。