Meta的优化版RecSys推理（58分钟阅读）

# 为 RecSys 推理协同设计内核 – PyTorch 来源：https://pytorch.org/blog/in-kernel-broadcast-optimization-co-designing-kernels-for-recsys-inference/ ### 特色项目 - PyTorch 标志 (https://pytorch.org/projects/pytorch/) TL;DR: - 传统的 RecSys 推理会为每个候选者显式复制共享的用户嵌入/序列。**内核内广播优化**（IKBO）通过一种**内核-模型-系统协同设计**，将广播逻辑直接融合到用户-候选交互内核中，从而消除这一开销。通过降低内存占用和 IO 利用率，IKBO 解锁了更高的吞吐量。 - IKBO 在高计算密集型场景下可实现高达**2/3 的延迟降低**，作为**可扩展性支柱**，支撑着 Meta 自适应排序模型 (https://engineering.fb.com/2026/03/31/ml-applications/meta-adaptive-ranking-model-bending-the-inference-scaling-curve-to-serve-llm-scale-models-for-ads/) 背后以请求为中心、推理高效的框架。 - 已在 Meta 的多阶段推荐漏斗中端到端部署，覆盖 GPU 和 MTIA (https://ai.meta.com/blog/next-generation-meta-training-inference-accelerator-AI-MTIA/)（Meta 训练与推理加速器）。 - IKBO 线性压缩内核在 H100 SXM5 上经过四个阶段的逐步协同设计，最终通过 TLX (https://github.com/facebookexperimental/triton/tree/tlx) 实现 warp 专业化融合，累计获得了 **~4 倍**的加速。 - IKBO 协同设计将 Flash Attention 内核从 IO 受限转变为计算受限（在 H100 SXM5 上达到 **621 BF16 TFLOPs**）。结合 TLX (https://github.com/facebookexperimental/triton/tree/tlx) warp 专业化优化，相比未经过协同设计的 CuTeDSL FA4 Hopper 基线（仅内核/内核 + 广播），实现了 **2.4 倍/6.4 倍**的吞吐量提升。 在本文中，我们介绍内核内广播优化（IKBO），这是一种内核-模型-系统协同设计方法，用于消除推荐模型推理中冗余的用户嵌入广播。在生产级 RecSys 中，对于给定请求，用户嵌入在所有候选者之间是相同的，但标准方法需要显式复制，从而浪费了随候选数量线性扩展的内存带宽和计算。IKBO 编码了一个简单洞察：广播是数据布局问题，而非计算必需。每个 IKBO 内核接受用户和候选输入（它们具有自然的、不匹配的批次大小），并在内部处理广播，因此永远不会物化出复制的张量。我们通过两个内核深度解析来展示该方法：线性压缩和 Flash Attention。 该方案已部署在 Meta 的 RecSys 推理栈中——从早期阶段到后期阶段的排序模型，覆盖 GPU 和 MTIA（Meta 训练与推理加速器）——在协同设计的模型上，IKBO 可将高计算密集场景下的净延迟降低高达 2/3。它作为以请求为中心、推理高效的框架的**可扩展性支柱**，支撑着 Meta 自适应排序模型 (https://engineering.fb.com/2026/03/31/ml-applications/meta-adaptive-ranking-model-bending-the-inference-scaling-curve-to-serve-llm-scale-models-for-ads/)（在生产中服务 LLM 规模模型）。在 H100 SXM5 上，我们的 IKBO 线性压缩内核通过四个逐步协同设计阶段实现了约 4 倍加速：矩阵乘法分解、内存对齐、广播融合，以及通过 TLX（Triton 低级扩展） (https://github.com/facebookexperimental/triton/tree/tlx) 实现的 warp 专业化多级融合。对于 Flash Attention，IKBO 相比未协同设计的 CuTeDSL FA4-Hopper（仅内核/内核+广播），实现了 2.4 倍/6.4 倍的吞吐量，达到 621 BF16 TFLOPs。与通过复制来绕过广播的系统级广播或网络拆分不同，IKBO 在计算原语层消除了广播，以近乎独立候选的成本实现了密集交互质量。 代码仓库：https://github.com/pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/ikbo †*在 Meta 期间完成的工作* ## 1. 内核内广播优化：消除内存与计算冗余 当用户打开其信息流时，推荐系统需要对数百到数千个候选项目进行评分，以决定显示哪些内容。模型的输入分为两类：**用户特征**（例如浏览历史、个人资料、上下文）在请求中对每个候选者是相同的，以及**候选特征**（例如项目 ID、类别、参与度统计数据）对每个项目是唯一的。两者都通过嵌入查找和后续处理来生成嵌入表示。在模型的各个阶段，**交互层**（例如线性投影、特征交叉、目标注意力）结合用户和候选嵌入。我们将请求中所有候选者共享的嵌入称为**仅请求（RO）**，将每个候选者的嵌入称为**非仅请求（NRO）**。 **图 1.** 一个非常简化的 RecSys 推理数据流。仅请求（RO）用户嵌入必须被广播（复制）以匹配非仅请求（NRO）候选批次维度，然后才能进入交互层。IKBO 通过在每个内核内部处理广播来消除这种物化。 交互层需要具有匹配批次维度的张量。在包含 1024 个候选、约 15 个用户的批次中，RO 嵌入必须被**广播**，复制约 70 次，以匹配 NRO 批次大小（图 1）。随着架构从 DLRM [1] 和 DCN [2] 发展到像 HSTU [3] 和 X 的 Phoenix [4] 这样的序列模型，用户-候选交互不断丰富。但更丰富的交互是有代价的：用户特征必须跨所有候选者广播。对于推理中 10 到 10,000+ 的批次大小，这种复制开销会导致显著的计算和内存成本，且随候选数量线性扩展。 **广播是数据布局问题，而非计算必需**。通过这个视角审视模型和推理系统，可以在每一层进行优化：推理运行时消除系统级广播，仅处理用户特征的模型层以较小的用户批次大小运行，而混合两者的内核被重新设计以在内部处理广播——永远不会物化出复制的张量。IKBO 已部署在 Meta 的 RecSys 推理栈中，从早期到后期阶段的排序模型，覆盖 GPU 和 MTIA，在协同设计的模型上，可将高计算密集场景下的净延迟降低高达 2/3。 本文重点通过两个深度解析来介绍内核层：线性压缩和 Flash Attention。 ### 1.1. 内核优化类型 **类型 I——可分解操作**。数学重构使得仅请求（RO）部分可以独立地以小批次大小计算，仅在最后与非仅请求（NRO）部分合并。这节省了内存带宽和计算。 **类型 II——仅内存优化**。在内核内部处理 RO-NRO 广播避免了冗余的数据移动，使内核远离 IO 受限。 ### 1.2. 端到端系统设计 部署 IKBO 涉及基础设施栈的三层： 1. **内核**：自定义 GPU 内核，接受不匹配的 RO/NRO 批次大小并在内部处理广播（第 2 节和第 3 节）。 2. **编译规范**：ML 编译器需要每个算子动态形状的范围，以选择适当形状的内核。对于一个批次大小这是简单的；当有两个（用户和候选）甚至更多时，可靠地解析每个算子使用哪个批次大小——在生产模型中，交互会混淆批次来源——需要系统化的自动化。 3. **推理**：运行时将候选到用户的映射传递给模型，而不是物化广播。 这些内核通过以下两条路径之一进入模型： 1. 直接采用：模型作者将 IKBO 内核直接集成到模型定义中。当训练时候选与用户比率 > 1，相同的内核也能降低训练成本。 2. 推理时转换：一个 pass 在推理时自动将标准算子替换为 IKBO 等价物——无需修改模型代码。 净效果：广播从推理的每个阶段消失，模型没有任何架构约束，除了推理运行时的映射接口外，无需基础设施变更。 ### 1.3. 与其他方法比较 现有方法围绕广播工作，而不是消除它。 1. **系统级广播**在 GPU 分派前物化复制的张量——简单但浪费，成本随候选数量线性扩展。 2. **网络拆分（ROO）[5]** 将模型划分为 RO 和 NRO 子网络，减少了冗余工作，但限制了用户-候选交互发生的位置，并且在 RO 批次大小较小时仍会引入额外成本。 两者都保留广播作为物化的张量。IKBO 在计算原语层消除它：节省随候选与用户比率缩放，任何交互模式都没有广播成本，并且完整的 NRO 批次维度在融合内核中提供了 GPU 占用率。 IKBO 已在 GPU 和 MTIA 加速器上部署。在本博文中，我们重点介绍 H100 GPU 内核设计，以说明核心优化原则。 ## 2. 内核深度解析 I：IKBO 线性压缩 线性压缩嵌入（LCE）通过学习投影 `(M, K) @ (B, K, N) → (B, M, N)` 压缩输入嵌入 `(B, K, N)`，广泛用于 Meta RecSys 模型，例如 Wukong [6]。我们逐步介绍四个优化阶段。 ### 2.1 矩阵乘法分解 **图 2.** LCE 分解：基线批量矩阵乘法（左上）、沿 K 的嵌入分离和用户去重（右上）、压缩输出上带有广播加法的两个独立 GEMM（底部）。 基线 LCE 对所有 B 个候选计算单个批量矩阵乘法。输入嵌入沿 K 拼接用户和候选部分——但对于同一用户，用户嵌入在所有候选之间是相同的。 **将广播推过矩阵乘法。** 由于 W 是批次无关的，我们利用线性分解：沿 K 分离用户和候选嵌入块，去重重复的用户嵌入，并以它们自然的批次大小计算两个独立的 GEMM。我们不在矩阵乘法前复制用户嵌入，而是只广播小的压缩结果。见图 2。在候选与用户比率约为 70（一个代表性设置）时，用户批次从 `B=1024` 缩小到 `B_user ≈ 15`——用户侧计算**减少了 70 倍**。该分解在标准 PyTorch 中实现。 **结果。** 1.944 ms → 1.389 ms（**减少 28.5%**；基准测试设置在附录 1 中）。原始批量 GEMM（算术强度约 356 FLOPs/Byte，低于 H100 约 495 FLOPs/Byte 的机器平衡点；推导见附录 2）和分解后的两个 GEMM 都是内存受限的，因此加速由内存成本降低驱动。去重将内存成本降低了一半以上——因为用户侧 GEMM（B_user ≈ 15 vs. B = 1024）变得成本可忽略。 注意分解将广播推过了矩阵乘法：我们不在 GEMM 前复制完整的 K 维输入嵌入，而是只广播小的压缩结果，这要便宜得多。在第 2.3 节中，我们将通过内核内广播融合进一步完全消除这个剩余的广播。 当前瓶颈是 L1/TEX 管线利用率（84%）而非 DRAM 利用率——这是一个可疑的不平衡，我们将在下一节深入分析。详细性能分解见附录 3。 ### 2.2 内存布局优化 对分解后 GEMM 的详细结果分析揭示了不平衡：L1/TEX 达到峰值的 84%，而 DRAM 仅达到 19%，表明内存加载宽度不必要地窄。SASS 确认：每个 `cp.async` 只复制 4 字节，而不是单个 128 位加载。 ``` LDGSTS.E.LTC128B P0, [R203], [R38.64] // 4 bytes LDGSTS.E.LTC128B P1, [R203+0x4], [R38.64+0x4] // 4 bytes (总共 4 次，仅加载 16B) ``` `cp.async` 宽度受源指针自然对齐的限制。矩阵 A 是 `(M, K)` 行主序，步长为 `K × 2` 字节，因此当 K 不是 8 的倍数时，步长破坏了 128 位对齐。 **模型-内核协同设计洞察。** 内存对齐是众所周知的 GPU 优化——但分解使其成为模型-内核协同设计挑战。`K` 由嵌入张量的 `torch.cat` 形成，其大小取决于许多模型配置因素。分解使得手动设计这些因素以使分解后的嵌入完美对齐变得非常困难。需要系统化的解决方案。 **解决方案。** 通过向拼接列表附加零，将每个分解后的 K 填充到下一个 8 的倍数。我们证明这在正向和反向传播中数学等价（见下面的证明 1），并且借助 ML 编译器的内存规划器，这简化为一个廉价的常量复制。 **证明 1.** 零填充 K 在正向和反向传播中保持精确数值等价。 **结果。** 1.389 ms → 0.798 ms（**减少 42.5%**）。填充使 CUTLASS 能够选择基于 TMA 的内核，绕过 L1/TEX 完全（扇区 351M → 0），并将 GEMM 延迟从 0.984 ms 降至 0.400 ms。解决了 GEMM 后，未融合的广播和加法（0.398 ms）现在占总延迟的一半——将在下一节解决。详细结果分析见附录 5。 ### 2.3 候选 GEMM 内核内广播融合 未融合的广播和加法是内存受限的：将候选 GEMM 结果写入 HBM，与用户结果一起读回，相加，再写回。我们通过将广播融合到候选 GEMM 的尾缀（epilogue）中来消除这一开销（图 3）。在每个 tile 累加后，尾缀根据用户索引查找，加载预先计算的用户结果，在寄存器中相加，并写入最终和——中间张量从未物化。我们将其实现为 Triton 内核：一个标准批量 GEMM，带有自定义的累加后尾缀块。 **图 3.** 内核内广播融合：GEMM 尾缀通过索引查找加载预先计算的用户结果并在寄存器中相加。 **结果。** 0.798 ms → 0.580 ms（**减少 27.4%**）。融合消除了 0.87 GB 的中间 DRAM 流量，有助于延迟降低。然而，占用率仅为 6.25%（每个调度器 1 个 warp），使每个停顿完全暴露。超过 42% 的周期等待全局加载，20% 等待 WGMMA——这些停顿无法被尾缀隐藏，并且没有持久性，无法与下一个 tile 的加载重叠。这是一个具有挑战性的权衡：需要大 tile 和深流水线来保持张量核心喂饱，但它们消耗大部分共享内存预算，几乎没有空间通过占用率来隐藏延迟。详细结果分析见附录 6。 ### 2.4 使用 TLX 的 Warp 专业化多级融合 **TLX（Triton 低级语言扩展）** (https://github.com/facebookexperimental/triton/tree/tlx) 暴露了 Hopper 的 warp 专业化、TMA、mbarrier 和命名屏障，同时保留了 Triton 的 Python DSL 和自动调优基础设施。 使用 TLX，我们解决了第 2.3 节中的占用率限制，通过 warp 专业化——通过功能划分而不是额外的 warp 来隐藏延迟。 第 2.1 – 2.3 节将原始 LCE 分解为两个独立计算：用户 GEMM（阶段 1）和带有融合广播加法尾缀的候选 GEMM（阶段 2）。我们首先优化阶段 2（主要瓶颈）内的延迟隐藏，然后将两个阶段融合到一个持久的持久内核中。 **阶段内延迟重叠** 候选 IKBO 内核是内存受限的——设计目标