Meta IKBO：如何通过内核内广播优化，将推荐系统推理延迟降低一个数量级？

推荐系统推理与大语言模型推理之间存在一个显著却常被忽略的差异：同一请求中的大量候选内容，往往共享同一份用户上下文信息。 在传统实现中，这意味着用户嵌入、用户序列等“请求级共享特征”必须反复复制，才能与候选批次对齐并送入交互层进行计算。这一复制操作看似只是“广播”，但在工业级部署中，它会持续吞噬显存带宽、推高 IO 成本，并导致延迟随候选数量线性恶化。

• In-Kernel Broadcast Optimization: Co-Designing Kernels for RecSys Inference
• https://pytorch.org/blog/in-kernel-broadcast-optimization-co-designing-kernels-for-recsys-inference/
• 代码：https://github.com/pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/ikbo
• 约 1.5 万字，阅读时长约 56 分钟，播客版约 17 分钟

Meta 在这篇文章中提出了 内核内广播优化（In-Kernel Broadcast Optimization，IKBO） 。其核心判断非常直接：广播首先是一个数据布局问题，而非计算本身不可避免的一部分。

既然如此，就不应在内核之外提前将共享张量物化出来，而应让内核直接接受“用户批次”与“候选批次”这两个天然不一致的输入维度，在计算过程中按索引完成广播语义。围绕这一思路，文章从内核、编译、推理运行时三个层面给出了协同设计方案，并借助 Linear Compression 与 Flash Attention 两个案例，展示了如何将一次看似普通的广播消除，转化为支撑生产级推荐系统扩展性的基础能力。

目录

• 关键问题
  • 问题一：IKBO 是否以控制流发散与同步开销取代了显存带宽瓶颈？
  • 问题二：这种强耦合设计是否以牺牲可组合性为代价，成为模型迭代的瓶颈？
• 一、内核内广播优化：消除内存与计算冗余
  • 1.1 消除内存与计算冗余
  • 1.2 内核优化类型
  • 1.3 端到端系统设计
  • 1.4 与其他方法的比较
• 二、内核深入案例一：IKBO Linear Compression
  • 2.1 矩阵乘法分解
  • 2.2 内存布局优化
  • 2.3 候选 GEMM 的内核内广播融合
  • 2.4 基于 TLX 的 warp-specialized 多阶段融合
• 三、内核深入案例二：IKBO Flash Attention
  • 3.1 IKBO Flash Attention 在推荐系统边界条件下解决 IO 受限问题
  • 3.2 在 TLX 上融合现代内核技术（FA3、FA4）与 IKBO
  • 3.3 通过模型协同设计融合 Self + Target Attention
• 四、基准测试与结果汇总
• 五、结论与未来方向
• 参考文献
• 附录
  • 附录 1：Benchmark 设置
  • 附录 2：算术强度分析
  • 附录 3：第 2.1 节的详细结果分析
  • 附录 4：瓶颈分析方法
  • 附录 5：第 2.2 节的详细结果分析
  • 附录 6：第 2.3 节的详细结果分析
  • 附录 7：Release-Acquire 同步协议
  • 附录 8：TLX 与 Triton 的 NCU Profiling 指标
  • 附录 9：普通 Flash Attention 与 IKBO Flash Attention 的 roofline 分析
  • 附录 10：IKBO TLX FA3 的 SMEM 占用
  • 附录 11：在推荐系统边界条件下，对 IKBO FA、CuTeDSL FA4 Hopper 与 TLX FA3 Hopper 的 benchmark
  • 附录 12：instruction cache miss 会在 consumer-2 warpgroup 上造成显著延迟

关键问题

问题一：IKBO 是否以控制流发散与同步开销取代了显存带宽瓶颈？

IKBO 通过消除显式广播来减少显存占用和带宽压力，但这是否以一种新的瓶颈——即“控制流发散与同步开销”——取代了旧的瓶颈？当候选集与用户的比例极度不均或动态变化时，内核内部的复杂分支逻辑和跨 warp 同步是否会成为新的性能悬崖？

作者展示的证据表明，这种代价确实存在但可被有效管理。在“奇数候选处理”部分，当一个用户的候选数为奇数时，某个 warpgroup 必须进入空闲状态并“排干”缓冲区以防止死锁。这直接证明了控制流分歧的存在。

关键在于量级对比。作者明确指出，在真实的约 70:1 候选-用户比下，这个空闲路径的触发概率“小于 0.7%”，开销可以忽略。而它所消除的显存复制开销则是系统性的、随候选数线性增长的——例如原方案中需将用户特征复制约 70 次，占用 0.87 GB 中间显存流量。

更深层的设计哲学是：TLX 的 warp 特化并非规避控制流分歧，而是将其显式化、可控化。通过将 CTA 划分为生产者与消费者组，并用命名屏障精确管理同步，IKBO 将不可预测的分支变成了结构化的流水线。NCU 数据显示，尽管存在同步，DRAM 吞吐量仍从 39% 提升至 52%，L2 瓶颈利用率增至 84%。

综合来看，IKBO 并非消除了瓶颈，而是用可控的同步开销换取了显存带宽的显著解放，这在批量大的推荐推理场景下是净收益。

问题二：这种强耦合设计是否以牺牲可组合性为代价，成为模型迭代的瓶颈？

IKBO 声称通过“Kernel-模型-系统”协同设计实现了显著加速，并使模型设计摆脱了约束。但这种灵活性是否以牺牲“局部性”和“可组合性”的软件工程原则为代价？ 当需要调试精度问题或快速更换模型结构时，这种高度耦合的优化是否反而成为迭代速度的瓶颈？

作者展示的恰恰是相反的设计思路：通过系统架构而非约定俗成来保护模块化。

关键证据在作者提到：“训练时不改变模型代码”和“推断时自动将标准算子替换为 IKBO 等价实现”。这意味着模型作者无需感知 IKBO 的存在，仍使用标准 PyTorch 算子编写模型。编译器和推断运行时会自动解析批次继承关系并进行替换。这不是强耦合，而是将复杂性封装在了底层。

关于数值精度问题，附录3直接回应了分解GEMM可能引发的数值差异，明确表示其“符合精度标准”，说明团队已意识到这一潜在问题并认为风险在可控范围内。针对算子语义被破坏的质疑，IKBO并未引入全新算子，而是改变了现有通用算子的内部实现机制。当研究人员需要替换激活函数时，无需改动任何代码，内核团队只需独立维护IKBO变体库即可。

最核心的隐患在于对编译器自动化的依赖。作者承认，跨多个生产模型解析批次谱系“需要系统化的自动化手段”。如果这一自动化层出现偏差，确实可能静默地污染最终结果。但这恰恰 表明文章所描述的生产落地已足够成熟——IKBO已在GPU和MTIA上部署，服务于“LLM规模的模型”，证明这套自动化流程的可靠性已得到验证。 结论是：IKBO通过将广播语义下放至运行时/编译器层面，在换取性能提升的同时，成功保护了上层模型的模块化设计。

一、内核内广播优化：消除内存与计算冗余

本文介绍In-Kernel Broadcast Optimization（IKBO）：一种面向推荐模型推理的内核-模型-系统协同设计方法，用于消除冗余的用户嵌入广播。在生产级推荐系统中，对于给定请求，用户嵌入在所有候选条目间完全相同，但标准做法仍要求显式复制，造成随候选数量增加而线性增长的内存带宽与计算浪费。IKBO的核心洞察在于：广播本质上是数据布局问题，而非计算必需。每个IKBO内核都以天然不匹配的用户输入和候选输入批大小作为输入，并在内核内部处理广播，因此不会物化任何复制后的张量。本文通过两个内核深入案例展示这一方法：Linear Compression与Flash Attention。

IKBO已在Meta的推荐系统推理栈中部署，覆盖从早期排序到晚期排序模型，支持GPU与MTIA（Meta Training and Inference Accelerator）。在协同设计后的模型上，它可将计算密集型网络延迟最多降低三分之二。它也是Meta Adaptive Ranking Model[1]所依赖的、以请求为中心且推理高效框架的可扩展性骨干，用于在生产环境中服务LLM规模的模型。

在H100 SXM5上，我们的IKBO线性压缩内核经过四个渐进式协同设计阶段：矩阵乘法分解、内存对齐、广播融合，以及通过TLX (Triton Low-Level Extensions)[2]实现的warp-specialized多阶段融合，获得约4倍加速。对于Flash Attention，IKBO相较未做协同设计的CuTeDSL FA4-Hopper，实现了2.4×/6.4×的吞吐提升（仅注意力内核/注意力内核加广播开销），达到621 BF16 TFLOPs。

与系统级广播或网络拆分这类“绕开复制”的方法不同，IKBO在计算原语层面直接消除复制，以近似独立成本实现高密度交互质量。

注：本工作完成于作者在Meta任职期间。

1.1 消除内存与计算冗余

当用户打开信息流时，推荐系统必须对数百到数千个候选条目打分，以决定向其展示什么。模型输入可分为两类：

• 一类是用户特征，例如浏览历史、画像、上下文，它们在一个请求中的每个候选上都完全相同；
• 另一类是候选特征，例如物品ID、类别、互动统计，它们对每个候选都不同。

两类特征都会经过embedding lookup以及后续处理，形成嵌入表示。在模型中的不同位置，交互层（例如线性投影、特征交叉、目标注意力）会把用户嵌入与候选嵌入结合起来。我们将一个请求中对所有候选共享的嵌入称为Request-Only（RO），将每候选独有的嵌入称为Non-Request-Only（NRO）。

图1. 一个极度简化的推荐系统推理数据流。仅请求（RO）的用户嵌入在进入交互层之前，必须被广播（复制）以匹配非仅请求（NRO）的候选物品批次维度。IKBO通过在每个内核内部处理广播，消除了这种物化。该图清晰地对比了传统方法与IKBO方案。传统方法中，广播是一个外部的、显式的数据准备步骤，而IKBO则将其内化为计算过程的一部分。这不仅是执行位置的改变，更意味着计算图的结构得到简化，运行时不再需要为庞大的复制张量分配和传输内存。

交互层要求张量具有匹配的批维度。在一个由约15个用户服务、却包含1024个候选的批次中，RO嵌入必须先被广播，即复制约70次，才能在任何交互发生之前与NRO批大小对齐（图1）。随着模型架构从DLRM [1]和DCN [2]演进到HSTU [3]与X的Phoenix [4]这类序列模型，用户-候选交互不断增强。但更丰富的交互也意味着更高代价：用户特征必须广播到所有候选上。在推理中，当批大小处于10到10000+时，这种复制开销会带来显著的计算和内存成本，并随候选数线性增长。

广播是数据布局问题，而不是计算必需。用这一视角重新看待模型与推理系统，就能在每一层发现优化空间：推理运行时消除系统级广播，仅用户相关的模型层在更小的用户批大小上运行，而同时混合处理两类数据的内核则被重新设计为在内部处理广播，因此不会有任何复制后的张量被物化出来。IKBO已部署到Meta的推荐系统推理栈中，从早期排序到晚期排序，覆盖GPU与MTIA，在协同设计模型上可将计算密集型网络延迟最多降低三分之二。

本文聚焦内核层，通过两个深入案例展开：Linear Compression与Flash Attention。

1.2 内核优化类型

• Type I — 可分解操作。通过数学重构，可让Request-Only（RO）部分在小批大小下独立计算，并仅在最后与Non-Request-Only（NRO）部分合并。这同时节省内存带宽与计算量。
• Type II — 纯内存优化。在内核内部处理RO-NRO广播，避免冗余数据移动，使内核远离IO受限状态。

1.3 端到端系统设计

部署IKBO会触及基础设施栈的三层：

1. 内核：自定义GPU内核，接受不匹配的RO/NRO批大小，并在内部处理广播（见后文第2、3节）。
2. 编译规格：ML编译器需要每个算子的动态形状范围，以便选择匹配形状的内核。只有一个批大小时，这很简单；当存在两个（用户与候选）甚至更多批大小时，要在生产模型中可靠判断每个算子究竟使用哪一个批维度——尤其在交互使批来源变得模糊时——就需要系统化自动化方案。
3. 推理：运行时把candidate-to-user映射传入模型，而不是先物化广播结果。

这些内核通过两种方式进入模型：

1. 直接采用：模型作者将IKBO内核直接集成进模型定义。当训练时candidate-to-user比值大于1，同样的内核也能降低训练成本。
2. 推理时变换：一个pass会在推理时自动把标准算子替换为IKBO等价算子，无需修改模型代码。

最终效果是：广播从推理的每一个阶段消失，而对模型架构没有约束，对基础设施的变更也仅限于推理运行时的映射接口。

1.4 与其他方法的比较

已有方法更多是在绕开广播，而非真正消除它。

1. 系统级广播会在GPU dispatch前先物化复制后的张量——实现简单，但代价高，且成本随候选数线性增长。
2. 网络拆分（ROO）[5] 将模型分成RO和NRO子网络，能减少冗余工作，但会限制用户-候选交互发生的位置，而且在较小RO批大小下仍会引入额外成本。

这两种方式都保留了广播这一“被物化的张量”。IKBO则在计算原语层面直接消除它：节省量会随candidate-to-user比值扩大，任意交互模式都无需承担广播成本，而融合内核中的完整NRO批维度还能提供GPU occupancy。

IKBO已部署在GPU与MTIA加速器上。本文聚焦H100 GPU内核设计，以说明其核心优化原则。

二、内核深入案例一：IKBO Linear Compression

2.1 矩阵乘法分解

基准 LCE 会对所有 B 个候选对象执行一次批量矩阵乘法（batched matmul）。输入嵌入在 K 维上将用户部分与候选部分拼接在一起——然而，对于同一用户的所有候选来说，用户嵌入是完全相同的。

图 2 展示了 LCE 分解的核心思想：左上角为基线批量矩阵乘法；右上角则是沿 K 维分离嵌入并对用户进行去重；下方展示了两个独立的 GEMM 操作，最终通过广播加法合并压缩后的输出。这幅图生动地阐释了从“先复制再计算”到“先计算再轻量合并”的范式转变。第一阶段通过分离用户与候选物品的嵌入，立即将用户侧的矩阵乘规模缩小了数十倍。第二步则将广播操作从完整的高维嵌入向量推迟到低维的压缩结果上，从而大幅削减了广播所需处理的数据量。

将广播延迟至矩阵乘法之后执行。由于权重 W 与批次无关，我们利用线性性质进行分解：沿 K 维度将用户嵌入块与候选嵌入块分离，对重复的用户嵌入进行去重，随后在各自天然的批大小上分别执行两个独立的 GEMM。这样一来，用户嵌入不再需要在 matmul 之前被复制，只需广播较小的压缩结果即可，具体见图 2。在 candidate-to-user 比值约为 70 的典型场景下，用户批大小从 B=1024 锐减至 B_user ≈ 15，使得用户侧的计算量降低了 70 倍。该分解过程完全基于标准 PyTorch 实现。

结果：延迟从 1.944 ms 降至 1.389 ms（降幅达 28.5%，benchmark 配置详见附录 1）。原始批量 GEMM 的算术强度约为 356 FLOPs/Byte，低于 H100 的约 495 FLOPs/Byte 机器平衡点（推导过程见附录 2）。分解后的两个 GEMM 依然受内存限制，因此性能提升主要源于内存成本的降低。去重操作将内存成本削减了一半以上——因为用户侧 GEMM（B_user ≈ 15，而非 B = 1024）几乎可以忽略不计。

需要特别指出的是，这里的分解仅仅是将广播操作移到了 matmul 之后：不再在 GEMM 之前复制完整的 K 维输入嵌入，而只广播较小的压缩结果，因此代价要低得多。第 2.3 节将进一步通过内核内广播融合，将这部分剩余的广播完全消除。

当前瓶颈主要集中在 L1/TEX pipeline 利用率（84%），而非 DRAM 利用率——这种失衡现象值得警惕，我们将在下一节进行详细分析。更精细的 profiling 结果可参见附录 3。

2.2 内存布局优化

对分解后 GEMM 的详细结果分析揭示了一种失衡状态：L1/TEX 达到了峰值的 84%，而 DRAM 仅占 19%，这表明内存加载宽度存在不必要的过窄问题。SASS 代码证实：每条 cp.async 指令仅复制 4 字节，而非一次 128-bit 加载。

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) 的 row-major 布局，步幅为 K × 2 字节，因此当 K 不是 8 的倍数时，这个 stride 会破坏 128-bit 对齐。

模型—内核协同设计洞察：内存对齐本身是经典的 GPU 优化问题，但分解操作使其演变为模型—内核协同设计的挑战。K 来自多个 embedding tensor 的 torch.cat，而这些 tensor 的大小取决于众多模型配置因素。完成分解后，几乎无法再通过人工微调这些因素来确保分解后的 embedding 仍保持完美的倍数对齐，因此需要系统化的解决方案。

解决方案：将每个分解后的 K 通过补零（pad）操作扩展到下一个 8 的倍数。我们证明，在前向与反向传播中，这一操作在数学上与原始计算完全等价（见下方证明 1），并且借助 ML 编译器的内存规划器，这仅相当于一次低成本的常量拷贝。

公式描述了矩阵乘法中的填充（Padding）技术，其目的是将矩阵维度调整为硬件友好的倍数（此处为8的倍数），以优化计算效率，同时保证前向与反向传播的数学正确性。在设置中，矩阵乘法 里，，，。填充时，为 补 个零列、为每个批次的 补 个零行，使新维度 成为8的倍数。前向传播中， 的结果与原矩阵乘法 完全一致，零填充不影响计算结果。反向传播时，损失对 的梯度仅前 列有效，其余零列梯度为0；损失对 的梯度仅前 行有效，其余零行梯度为0，二者均能正确还原原矩阵的梯度。这种填充方法在不破坏模型数学逻辑的前提下，通过适配硬件计算单元的对齐要求，大幅提升了矩阵乘法在GPU/加速器上的执行效率，是高性能算子优化中兼顾正确性与性能的典型设计。

证明 1：对 K 进行零填充，在前向与反向传播中均保持精确的数值等价性。

结果：延迟从 1.389 ms 降至 0.798 ms（降幅达 42.5%）。Padding 后，CUTLASS 能够选择基于 TMA 的内核，从而彻底绕过 L1/TEX（sectors 从 351M 降至 0），并将 GEMM 延迟从 0.984 ms 降低到 0.400 ms。GEMM 问题解决后，未融合的广播与加法操作（0.398 ms）现在占据了总延迟的一半——下一节将继续处理这一问题。详细分析见附录 5。

2.3 候选 GEMM 的内核内广播融合

未融合的广播与加法操作是内存受限的：首先将 candidate GEMM 的结果写回 HBM，再将其与用户结果一起读回，相加后再次写回。我们将广播融合进 candidate GEMM 的 epilogue 中，从而消除了这一过程（图 3）。

图 3. 核内广播融合：GEMM 的尾声通过索引查找加载预先计算好的用户结果，并在寄存器内完成加法。该图揭示了 Kernel 融合的核心机制，即在每个计算瓦片完成后，立即在寄存器层面完成加法，而不是将完整矩阵写回显存后再进行。这避免了中间结果在整个 HBM 上的往返，是实现节省 0.87 GB DRAM 传输量的关键。

每个 tile 完成累积后，epilogue 会查找用户索引，加载预先计算好的用户结果，在寄存器中相加，再直接写出最终结果——中间张量不会被物化。我们使用 Triton 内核实现这一功能：本质上是一个标准 batched GEMM，加上自定义的 post-accumulation epilogue block。

结果：延迟从 0.798 ms 降至 0.580 ms（降幅达 27.4%）。融合消除了 0.87 GB 的中间 DRAM 流量，因此获得了延迟收益。然而 occupancy 仅为 6.25%（每个 scheduler 仅 1 个 warp），使得所有 stall 都被完整暴露。除了有 42% 以上的周期在等待 global load 外，还有 20% 的周期在等待 WGMMA——这些 stall 无法被 epilogue 隐藏；而且没有持久化机制时，也不存在下一个 tile 的加载可与之重叠。这是一个棘手的权衡：为了喂饱 tensor core，需要较大的 tile 与较深的 pipeline，但它们会占用大部分 shared memory 预算，导致几乎没有余量通过提升 occupancy 来隐藏延迟。详见附录 6。

2.4 基于 TLX 的 warp-specialized 多阶段融合

TLX（Triton Low-level Language Extensions）[4] 在保留 Triton Python DSL 与 autotuning 基础设施的同时，暴露了 Hopper 的 warp specialization、TMA、mbarrier 与 named barrier 等特性。

借助 TLX，我们采用 warp specialization 来解决 2.3 节中的 occupancy 限制——并非通过增加更多 warp，而是通过功能分工来隐藏延迟。

2.1–2.3 节将原始 LCE 拆解为两个独立的计算阶段：用户 GEMM（阶段 1）和带有融合 broadcast-add epilogue 的候选 GEMM（阶段 2）。我们首先对占据性能瓶颈的阶段 2 进行延迟隐藏优化，随后将两个阶段合并为一个持续的 persistent kernel。

阶段内延迟重叠

候选 IKBO 内核受内存带宽限制——设计目标是确保内存流水线持续满载。Triton 的软件流水线（2.3 节）已实现了 Load 与 WGMMA 的重叠，但 epilogue 部分仍是串行执行的——它会阻塞后续的 Load 操作，并导致 WGMMA 等待 stall。我们的解决方案是将每个 CTA 划分为专门的 warp group： 一个专用的 producer 持续发出 TMA load（重叠 1，类似于 Triton 原有的软件流水线），两个 consumer 采用 ping-pong tile 方式，使得一方的 epilogue 与另一方的 WGMMA 重叠（重叠 2） 。结合 persistence 特性，tile 会连续流动，不会出现跨 tile 的空隙。具体结构见图 4。

图 4. 候选IKBO内核结构，展示了两种阶段内延迟重叠以及warp group的角色分配。该图展示了从一维流水线到二维乒乓流水线的升级。该设计巧妙地解决了2.3节中发现的占用率与流水线深度的矛盾。通过warp特化，它不再依赖增加更多线程块来提高占用率，而是在单个CTA内部创建了三个专用且协作的硬件调度单元，实现了对延迟的精细掩盖。

多阶段融合

我们将用户 IKBO（阶段 1）与候选 IKBO（阶段 2）融合为一个 mega-kernel，目的是减少 wave quantization、消除 kernel launch 开销，并提升 L2 cache 利用率。较高的 candidate-to-user 比值会放大阶段 1 的 wave quantization 问题。由于候选 GEMM 在 epilogue 之前并不依赖用户结果，我们可以并发调度两个阶段。

这种并发调度又解锁了两个跨阶段重叠，使总重叠数达到四个，见图 5。

图 5. 并发阶段调度：没有用户 tile 的SM立即进入阶段2，与阶段1的部分波重叠。多阶段融合后的所有四种延迟重叠，展示了阶段内（#1，#2）和跨阶段（#3，#4）的重叠机会。SM 0-49，50-131仅为示例数字。该图揭示了跨阶段重叠如何从根本上解决了因批次维度不匹配导致的计算资源闲置问题。它打破了传统串行执行两个算子的壁垒，让系统全局地调度两个计算任务，这对于提升整个端到端流水线的硬件利用率至关重要。

Warp Group 专门化与同步设置

为了实现这四种重叠，每个 CTA 被划分为一个 producer 和两个 consumer warp group。 关键点在于，两个阶段共享同一套 circular buffer 与mbarrier 基础设施——在阶段边界处不会清空 pipeline，也不需要重新初始化 barrier。最后一个 user K-block 与第一个 candidate K-block 会同时处于不同 buffer slot 中。具体实现见图 6。

图 6. 每个CTA的warp组设置以及三种同步机制。这张图展示了 GPU 算子中典型的生产者 – 消费者协同优化架构，以线程块（CTA）为单位，每个 CTA 包含 1 个生产者线程束组和 2 个消费者线程束组，通过共享内存（SMEM）中的环形缓冲区实现数据的分阶段读写与流水线并行，提升 Tensor Core 利用率。生产者负责分阶段加载数据并写入缓冲区，消费者则以乒乓模式或协作模式读取数据并执行计算，通过命名屏障（mbarrier）实现同 CTA 内的同步。跨 CTA 的依赖则通过全局内存（GMEM）中的依赖计数器，结合原子加法和轮询机制完成同步，确保中间结果和最终结果的正确共享。这种设计通过解耦数据加载与计算、重叠流水线阶段，有效隐藏访存延迟，提升算子整体吞吐与硬件利用率。

双向阶段交替 tile 调度

当两个阶段的 tile 数都不能被 SM 数整除时，朴素的单向 dispatch 会导致负载不均。我们在两个阶段之间反转 tile 分配方向：阶段 1 从 pid 开始，阶段 2 从 NUM_SM - 1 - pid 开始，具体见图 7。

图 7. 单向（左）与双向的阶段交替调度（右），平衡了部分波中每个SM的工作负载。该图对比了两种GPU调度策略：左侧的朴素单向调度因两阶段任务都按升序分配瓦片，导致部分SM（如SM 0/1）承担过多工作，而SM 131在Stage 1无任务，Stage 2仅处理1个瓦片，负载严重不均，出现“Bad balance”；右侧的双向交替调度则让Stage 1按升序分配瓦片，Stage 2按降序分配瓦片，使每个SM在两阶段中各承担1个瓦片，Wave 1阶段所有SM都恰好处理2个瓦片，实现了“Good balance”。这种交替调度通过反转任务分配顺序，避免了热点SM的过载与空闲，让各SM的计算负载更均匀，从而提升了整体硬件利用率与算子性能。

tile 粒度的跨 CTA 同步

用户 tile 与候选 tile 可能运行在不同 CTA 上，因此需要跨 CTA 同步——但设备级 barrier 会把所有工作串行化，完全破坏这些重叠。我们采用每个 tile 粒度的三步 release-acquire 协议：

1. 每个 warp group 只让一个线程用 ld.relaxed 轮询 tile flag，以最小化内存流量。
2. 一旦 flag 被置位，用单条 ld.acquire 建立 happens-before 关系。
3. 用 named barrier 把就绪状态广播给该 warp group 的全部 128 个线程。

这避免了轮询时昂贵的 fence，也让处理不同 user tile 的 candidate CTA 能完全独立推进。详见附录 7。

结果

所有优化叠加后，延迟从 0.580 ms 降到 0.482 ms（下降 16.9%）。清晰的 intra-warp Proton tracer[5] 时间线表明，这四种重叠在实际运行中都得到了实现。

图 8. 两个CTA的Proton性能分析时间线，四种重叠用不同颜色标出。内存流水线持续保持满载。这张图是 IKBO 算子的 GPU timeline，展示了四层流水线并行优化：Overlap #1 通过同一 SM 内的数据加载（Load）与矩阵乘（WG MMA）重叠，隐藏访存延迟；Overlap #2 在同一阶段的 WG MMA 与后处理（Epilogue）间实现并行；Overlap #3 则让不同阶段的计算在同一 SM 上流水执行；Overlap #4 更进一步，让用户特征与候选特征的计算跨 SM 并行。通过生产者 – 消费者线程束的协作与乒乓缓冲机制，算子实现了数据加载、计算与结果处理的全流程重叠，充分利用了 GPU 的多线程与多 SM 并行能力，有效提升了硬件利用率与整体吞吐。

主要收益来自重叠 2：ping-pong consumer 在每个 tile 上都隐藏了 WGMMA 与 epilogue stall，直接解决了 2.3 节中的主要浪费周期 。重叠 1（Load↔WGMMA）继承自 Triton 现有软件流水线。重叠 3 与 4 则隐藏了 user-to-candidate 阶段切换时的空闲时间。见图 8。

NCU 证实：occupancy 从 6.25% 提升到 18.75%（3 个 warp group 对比 1 个），DRAM 吞吐从 39% 提升到 52%，而作为瓶颈的 L2 吞吐从 74% 提升到 84% 的峰值。这不仅仅是 occupancy 提升：跨四种重叠的积极延迟隐藏让内存流水线持续饱和，才推动 L2 超过 80%。更详细的 NCU 指标见附录 8。

我们在不同批大小与不同 candidate-to-user 比值下进行了基准测试，默认设置为 batch=1024、ratio=70。见图 9。

图 9展示了跨批次大小（左图，比率固定）与候选-用户比率（右图，批次固定）下的IKBO累积加速效果。该柱状图呈现了IKBO线性压缩算子（LCE）在多阶段优化过程中，相较于PyTorch Eager基线的吞吐加速比。左图在用户-候选比例固定为70的条件下，测试了不同批量大小（512至3072）的性能表现；右图则在批量大小固定为1024时，测试了不同用户-候选比例（2至1000）的性能表现。从分解阶段（Stage1）、K维对齐（Stage2）、内核融合（Stage3）到TLX线程束专用优化（Stage4），每一步优化均叠加了前序成果，最终实现了约4倍的整体加速。数据表明，随着批量增大或用户-候选比例升高，优化收益愈发显著。在高比例场景下，Stage4的加速比稳定突破4倍，这充分证明IKBO的优化设计能够有效适配推荐系统的典型负载，大幅提升算子的吞吐性能。

IKBO融合在不同场景下均展现出稳定的收益：无论是在不同批大小还是不同候选-用户比值下，都能获得约4倍的加速效果。即使在较低比值下，该内核依然具有可观的性能提升。

三、内核深入案例二：IKBO Flash Attention

随着推荐模型不断扩展以捕获更丰富的用户序列行为，序列架构——包括注意力机制——已成为关键的计算瓶颈：当序列长度为1K时，其大约占据40%的推理延迟。因此，我们将重点放在结合推荐系统独特的batching语义进行协同设计的IKBO-aware Flash Attention上。

受Transformer与Set Transformer [7, 8]的启发，推荐系统中广泛采用两类基本的用户历史交互模块：

• target attention（类似cross-attention），用于捕获预测候选与用户历史行为之间的关系。
• self-attention，用于建模用户历史序列内部的依赖关系。

由于用户历史属于RO特征，而target则作用于独立的candidate（NRO）批维度，这种结构上的不对称性为IKBO提供了提升模型可扩展性与计算效率的机会。我们主要优化target attention；通过少量协同设计，可以在3.3节中将self-attention也融合进IKBO target attention。由于我们的模型是encoder-driven，因此使用了完整注意力而非causal masking。

最终，利用端到端协同设计优化后的target attention版本，相比未做协同设计的CuTeDSL FA4-Hopper，可实现2.4倍/6.4倍的吞吐提升（仅注意力内核/注意力内核加广播成本），并分别降低0.320毫秒/1.232毫秒的延迟（如后文表2所示）。

3.1 IKBO Flash Attention 在推荐系统边界条件下解决IO受限问题

IKBO将K/V广播融合进注意力内核中，并借助推理运行时提供的candidate-user映射张量保持数学等价性，从而处理非均匀的candidate-to-user比值。

图10对比了带candidate-user广播的传统SDPA（左）与融合后的IKBO target attention（右）。这一对比直观地展示了IKBO方案如何重塑数据流。在传统路径中，KV矩阵被显式膨胀了数十倍，这对内存容量和带宽构成了巨大考验。而IKBO路径则通过一个映射张量，将索引操作直接注入内核，实现了“虚拟广播”。这种转变的核心在于，它将高带宽需求的数据（KV）的访问次数降低了数十倍，而代价仅仅是增加一次轻量级的索引计算。

上图10对比了两种方式：

• 传统SDPA路径在注意力前将K和V广播到完整的candidate批大小。
• 而IKBO路径则完全消除了这一步物化——每个候选在运行时按需索引其所属用户的K/V。

通过IKBO协同设计，将IO受限转为计算受限

在推荐系统边界条件下，相比用户浏览历史，target attention只使用相对较少的candidate embedding来表示候选属性。对标准注意力进行roofline分析可知，其算术强度约为60 FLOPs/Byte，远低于H100（SXM5 HBM2e版本）约495 FLOPs/Byte的峰值（附录2）。因此，即使是标准Flash Attention也严重受IO限制。IKBO通过让共享同一用户上下文的多个候选摊销K/V内存访问，将算术强度从约60 FLOPs/Byte提升到约833 FLOPs/Byte（当B_candidate : B_user = 70:1 时），从而将内核明确推入计算受限区间。

为了最大化这一收益，我们重新排序了threadblock launch grid，让batch_size_candidate先于num_heads。这确保了那些处理不同候选、却共享同一用户K/V的threadblock能够被并发调度，从而提高L2缓存的复用率。

表1：launch grid配置对比显示：SDPA将num_heads 放在grid.y 以优先优化GQA；IKBO则交换了head和candidate维度，将batch_size_candidate 放到grid.y，以便在不同候选间高效共享K/V。

表2比较了我们的IKBO Triton实现（FA2逻辑 + IKBO）与Hopper上最先进的Flash Attention实现（未做IKBO协同设计）。吞吐与IO仅测量注意力部分；Key与Value的广播延迟甚至比注意力本身还高。

表2：推荐系统边界条件下的注意力内核对比（B_candidate = 2048，B_u = 32，candidate-to-user比值均匀）。在未做协同设计的情况下，即使是前沿的Hopper实现也依然受IO限制。

3.2 在TLX上将现代内核技术（FA3、FA4）与IKBO结合

当IKBO将内核从IO受限转为计算受限后，自然的下一步就是采用Flash Attention 3（FA3 [10]）与Flash Attention 4（FA4 [11]）在Hopper上最先进的计算优化——尤其是warp specialization与流水线技术。然而，我们在query embedding数量上的边界条件（q_seq = 32 或 64）使得难以直接采用FA3的ping-pong或cooperative warp specialization。

在 Hopper 架构上实现 warp specialization 依赖异步 WGMMA 指令，这要求 BLOCK_M ≥ 64。为了尽可能缩小两个 consumer warp group 间的气泡（bubbles），我们需要配置两个这样的组。基于这些约束，我们对内核进行了定制设计，让单个 threadblock 同时发射 B_candidate = i 和 B_candidate = i + 1，并共享同一个 B_user。在接下来的讨论中，我们假设所有用户都对偶数个候选进行排序，且 q_seq = 64；奇数候选的处理方式将在后续说明。

IKBO FA3 内核的性能提升

从 FA3 的架构出发——包括阶段内流水线、warpgroup specialization 以及 ping-pong 调度——最初的 TLX IKBO FA3 内核性能与 FA2 基线基本持平（图 12，蓝线对红线，详见附录 11），吞吐量几乎相同。

为了定位瓶颈，我们借助 Proton tracer[6] 以 GPU cycle 为延迟单位，对 intra-warp pipeline 进行了可视化分析（图 11）。

图 11：基于 Proton 的 TLX IKBO FA3 内核 intra-warp profiling。图中展示了每个 warp group 的代表性 warp：warp 0（producer）、warp 4（consumer 1）和 warp 8（consumer 2）。softmax_PV_overlap 与纯 softmax 区域被单独标注，以识别 tensor core 气泡。(A) persistence 前对 B 的放大视图；(B) persistence 前，两波次；(C) persistence 后，两波次。该图通过对比生动展示了持久化的效果。持久化前（B），两个波次之间明显的黑色间隙就是跨 CTA 的等待时间。持久化后（C），间隙消失，生产者 warp（warp 0）的连续忙碌状态表明流水线几乎没有中断。这不仅是性能的提升，更是对硬件利用率的极致挖掘。

下表 3 汇总了 persistence 优化前后通过 Proton tracer 测得的主要瓶颈。

表 3：persistence 及优化前后的关键瓶颈。

关键结论是：在这种较小的 query sequence 长度下，主导瓶颈是 cross-CTA stall，而非 tensor core 利用率。要获得性能提升，persistence 是必要条件。引入 persistence 后，其 profiling 结果与延迟变化见图 11C 与表 3。

针对 HBM2e 的特定优化

随后我们又针对 H100 SXM5 的 HBM2e 带宽约束对 persistent kernel 做了进一步调优，用更多 shared memory capacity 去换取更少的 load/store 阻塞（表 4）。

表 4：面向 HBM2e H100 SXM5 的定制优化。这些优化在推荐系统边界条件下仍能装入可用的 SMEM 预算（附录 10）。

我们还实现了 persistent V2，它从 K 序列末尾向前迭代（与 FA3/FA4-Hopper 方案一致），以简化 masking 逻辑。两种 persistent 变体都应用了表 4 中的优化。

图 12：不同序列长度下的 IKBO 实现吞吐（B_candidate = 2048；B_candidate : B_user = 64；num_head = 2；d_head = 128）。实际推荐系统序列长度通常低于 4K [3]；更长序列仅用于与 LLM 场景对比。generalized 版本支持每用户候选数不为偶数的情况，其中每用户奇数候选概率为 50%。

上图 12 显示，在较短序列长度（512–4096）下，TLX FA3 persistent kernel 优于其他候选方案；当长度超过 8K 时，两种 persistent 变体趋于一致。

将 IKBO FA3 泛化到任意候选批大小排序场景

我们的 IKBO FA3 内核为了满足 WGMMA 的 BLOCK_M ≥ 64 要求，每个 CTA 会同时处理两个 candidate batch。当某个用户只有奇数个候选时，其中一个 consumer warpgroup 就没有成对伙伴。对此我们加入了 idle 逻辑（详见图 13 左以及算法 1）：

• idle 的 warpgroup 通过 mbarrier signaling 排空 K/V buffer，防止 producer 死锁。
• active 的 warpgroup 关闭 ping-pong 同步（因为其配对方不会再到达 named barrier）。

在约 70:1 的 candidate-to-user 比值下，idle 路径触发概率低于 0.7%，开销可以忽略不计（图 12，IKBO TLX FA3 generalized）。同样思路还可推广到 q_seq_len = 32，此时每个 CTA 绑定四个 candidate batch，并使用相应的 idle 与 masking 逻辑。

图 13：泛化 target attention 的 CTA 分配方案（左侧），以及 self + target attention 融合后的布局（右侧）。每个 CTA 被分配了两个 consumer warp group，它们共享同一份用户 K/V 数据。当候选数量为奇数时，第二个 consumer 将进入空闲状态并清空 barrier。

算法 1：针对奇数候选处理的 IKBO Attention 前向传播流程。该算法是 IKBO Attention 在前向传播过程中，针对奇/偶候选数量所设计的优化步骤。它以 CTA 为基本单位，输入候选-用户批次映射张量 M 和起始索引 i。首先，加载相邻候选的用户批次标识，并判断它们是否属于同一用户（same_user）。生产者 warp group 优先通过 TMA 加载数据，如果属于同一用户，则会额外加载数据，并采用流水线方式加载后续内容。消费者 warp group 1 处理候选数据，当同用户时启用命名屏障实现乒乓计算，否则禁用屏障。消费者 warp group 2 处理候选数据，同用户时启用屏障计算，否则跳过计算并等待流水线 mbar 同步。该设计通过预先判断用户的连续性，优化了 TMA 加载与屏障同步，减少了分支开销，实现了候选数据的并行处理，从而提升了注意力前向传播的流水线效率。

3.3 通过模型协同设计融合 Self + Target Attention

前面几节主要聚焦于 target（cross）attention 的优化。一个很自然的疑问是：能否将 self-attention 也整合到同一个内核中？

核心洞察在于，这两种注意力机制共享同一个 key-value 来源——用户序列。两者的唯一区别在于 query：self-attention 的 query 来源于用户侧，而 target-attention 的 query 来源于候选侧。通过让它们共享 K/V 投影，我们可以在单次内核启动中实现直接的横向融合。图 13（右侧）展示了融合后的 CTA 布局：前一部分 CTA 处理 self-attention 的 query block，剩余的 CTA 则处理 target-attention 的 candidate pair——它们都从同一条流水线化的 K/V 流中读取数据。

类似的协同设计思路也曾出现在 XAI Phoenix——X 平台的开源推荐系统 [4] 中。

图 13：泛化 target attention 的 CTA 分配方案（左侧），以及 self + target attention 融合后的布局（右侧）。每个 CTA 被分配了两个 consumer warp group，它们共享同一份用户 K/V 数据。当候选数量为奇数时，第二个 consumer 将进入空闲状态并清空 barrier。

我们实现了一个原型融合内核，用于量化融合带来的收益，并且不计入 K/V projection 节省的部分（如图 13 右侧所示）：

• seq_len = 512：性能提升 6.6%（从 514 提升至 482 TFLOPs/s）
• seq_len = 1024：性能提升 4.1%（从 581 提升至 558 TFLOPs/s）
• seq_len = 2048：性能提升 0.3%（从 612 提升至 610 TFLOPs/s）——此时 self-attention 已经使 SM 达到饱和状态

短序列下的收益主要源于内核融合本身：更少的启动开销、共享 buffer 分配带来的节省、跨内核的流水线机会，以及缓解 wave quantization 问题——这些也正是 megakernel 技术 [12] 在 LLM 推理中所要解决的核心问题。在生产环境中，共享 K/V projection 还会进一步节省线性投影的成本，这与 KV cache 复用的原理类似。

unsetunset四、基准测试与结果汇总unsetunset

以下汇总了本文所涉及的内核级基准测试结果以及端到端的部署效果。所有内核基准测试均在 H100 SXM5 上完成（详细配置见附录 1）。

• Linear Compression（第 2 节）：四个渐进式的协同设计阶段——矩阵乘法分解、内存对齐、广播融合，以及通过 TLX 实现的 warp-specialized 多阶段融合——在代表性配置下带来了累计约 4 倍的加速效果（从 1.944 ms 降至 0.482 ms）。这种性能收益在不同的批处理大小以及 candidate-to-user 比值下都表现稳定（如图 9 所示）。

图 9. 不同批次大小（左图，固定比率）和候选-用户比率（右图，固定批次）下的累积 IKBO 加速效果。该柱状图展示了 IKBO 线性压缩算子（LCE）在多阶段优化后，相较于 PyTorch Eager 基线的吞吐加速比。左图在用户-候选比例固定为 70 的条件下，测试了不同批量大小（512~3072）的性能；右图则在批量大小固定为 1024 时，测试了不同用户-候选比例（2~1000）的性能。从分解（Stage1）、K 维对齐（Stage2）、内核融合（Stage3）到 TLX 线程束专用优化（Stage4），每一步优化都叠加了前序成果，最终实现了约 4 倍的整体加速。数据显示，随着批量增大或用户-候选比例升高，优化收益愈发显著。在高比例场景下，Stage4 的加速比稳定突破 4 倍，这证明了 IKBO 的优化设计能够有效适配推荐系统的典型负载，从而大幅提升算子的吞吐性能。

• Flash Attention（第 3 节）：IKBO 将 target attention 从 IO 受限（约 60 FLOPs/Byte）转变为计算受限（约 833 FLOPs/Byte）。与未做协同设计的 CuTeDSL FA4-Hopper 相比，可实现 2.4 倍/6.4 倍的吞吐提升（仅内核/内核加广播），并达到 621 BF16 TFLOPs。
• 端到端部署：IKBO 已被广泛部署在 Meta 推荐系统的推理栈中，覆盖从早期排序到晚期排序的各个环节，并支持 GPU 与 MTIA 加速器。在协同设计的模型上，计算密集型网络的延迟最多可降低三分之二。此外，IKBO 在大约 10,000:1 到 10:1 的 candidate-to-user 广播比值范围内均经过了验证，证明了其在不同工作负载下的数值稳定性与可扩展性。

unsetunset五、结论与未来方向unsetunset

IKBO 证明了：广播——这一长期以来被视为用户-候选交互中不可避免的成本——可以通过内核-模型-系统协同设计，在计算原语层面被消除。通过将广播语义直接编码进内核，复制张量不再需要被物化，而收益会自然地随着 candidate-to-user 比值的增大而放大。

尽管本文中的内核实现是面向 NVIDIA Hopper、基于 Triton 与 TLX 的，但其核心思想——用索引驱动的内核内查找取代物化广播——并不依赖于特定的硬件厂商。下一步的自然方向包括：将 IKBO 内核适配到 CuTeDSL（以获得更先进的 NVIDIA 后端支持），以及补全对 AMD CK 的支持。

除了本文讨论的双层 user-candidate 层级外，一些推荐场景还会涉及更深层级，例如 user → ads vendor → ads item，即一个用户看到多个商家，每个商家又提供多个条目。这会引入两个嵌套的广播关系，而且它们的比值彼此独立且不均匀。IKBO 也能优雅地处理这种结构**，将其应用于多层级工作负载，是进一步降低生产推荐系统中物化开销的自然方向。

unsetunset参考文献unsetunset

[1] Naumov, M., et al. “Deep Learning Recommendation Model for Personalization and Recommendation Systems,” arXiv:1906.00091, 2019.

[2] Wang, R., et al. “Deep & Cross Network for Ad Click Predictions,” ADKDD, 2017.

[3] Zhai, J., et al. “Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations,” ICML, 2024.

[4] xAI. “Phoenix: Recommendation System,” GitHub, 2026. https://github.com/xai-org/x-algorithm

[5] Guo, L., et al. “Request-Only Optimization for Recommendation Systems,” arXiv:2508.05640, 2025.

[6] Zhang, B., et al. “Wukong: Towards a Scaling Law for Large-Scale Recommendation,” ICML, 2024.

[7] Vaswani, A., et al. “Attention Is All You Need,” NeurIPS, 2017.

[8] Lee, J., et al. “Set Transformer: A Framework for Attention-based Permutation-Invariant Input,” ICML, 2019.

[9] Dao, T. “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning,” ICLR, 2024.

[10] Shah, J., et al. “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision,” NeurIPS, 2024.

[11] Zadouri, T., et al. “FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling,” arXiv:2603.05451, 2026.

[12] Spector, B., et al. “Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B,” Hazy Research Blog, 2025. https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles

unsetunset附录unsetunset

附录 1：Benchmark 配置

所有测试均在单块 NVIDIA H100 SXM5 GPU（700W TDP，96GB HBM2e）上执行，软件环境如下：

• CUDA: 12.4
• PyTorch: 2.11.0a0+fb（内部构建版本）
• Triton: facebookexperimental/triton@4059e79bf（#831[7]）

附录 2：算术强度分析

2.1 H100 SXM5（700W TDP，96GB HBM2E）的机器平衡点

H100 GPU 硬件算术强度（AI，FLOPs/Byte）的计算公式是判断算子性能瓶颈的核心基准。分子 989.5 TFLOPs/sec 代表 H100 FP16 张量核心的理论峰值浮点算力，分母 2 TB/sec 则是其显存的理论峰值带宽。二者相除得到硬件算术强度为 495 FLOPs/Byte，这意味着每传输 1 字节数据，硬件理论上可完成 495 次浮点运算。该指标定义了硬件的“算力-访存平衡线”：若算子的算术强度低于此值，说明其访存需求相对更高，易受带宽限制；若高于此值，则算力会成为瓶颈，此时可针对性地优化访存或计算调度，以提升硬件利用率。

2.2 基线 LCE 的算术强度

对于 FP16 下的批处理矩阵乘法 (M, K) @ (B, K, N) → (B, M, N)，当 B=1024, M=433, K=2044, N=256 时：

矩阵乘法（GEMM）算子的算术强度（AI，FLOPs/Byte）计算公式，用于量化算子的计算访存比，从而判断其性能瓶颈类型。分子 2·B·M·K·N 代表 GEMM 的总浮点运算量，其中 2 表示乘加操作产生的 2 次 FLOP/元素，B 为批大小，M/K/N 为矩阵维度；分母 (M·K + B·K·N + B·M·N)×2 则是算子的总访存字节数，三项分别对应输入矩阵 A、输入矩阵 B 和输出矩阵 C 的读写数据量，×2 是 FP16 数据的字节宽度。计算结果 356 FLOPs/Byte 表明，算子每访问 1 字节数据可执行 356 次浮点运算。该值越高，算子越偏向计算密集型，越能充分利用 GPU 算力；反之，则易受内存带宽限制。此指标是性能优化中判断算子瓶颈、指导优化方向的核心依据。

附录 3：第 2.1 节的详细结果分析

设置：H100 SXM5（附录 1），PyTorch eager mode（无 kernel fusion），推理模式。形状取自代表性配置。

• 执行的总 FLOPs，非吞吐量。
  † 瓶颈由 NCU Speed of Light 分析识别；方法详见附录 4。

去重操作消除了超过 98% 的用户侧工作量（batch 从 1024 降至约 15），使得 L1/TEX sectors 从 798M 下降到 351M，并将 GEMM 延迟从 1.944 ms 降低至 0.984 ms。GEMM 后的广播与加法操作耗时 0.405 ms（受 DRAM 限制），因此净节省时间为 0.555 ms。

精度说明：基线在单个 FP32/TF32 reduction 中累加所有 K 乘积。分解后，会分别累加 K_user 与 K_cand，再以 BF16/FP16 将部分结果相加。训练阶段使用了相同的分解方式，因此端到端数值是一致的。若推理时需要与原始实现完全一致，第 2.4 节中的 fused kernel 可在最终求和时使用 FP32。

附录 4：瓶颈分析方法

在进行 roofline 分析后，我们使用 NCU 的 Speed of Light 分析来识别硬件子系统的瓶颈。瓶颈被定义为：相对其峰值持续吞吐率，利用率最高的子系统。针对 2.1 节的分析，我们监控三类指标：

• Compute：峰值 SM pipeline 利用率，由 NCU 直接报告（Compute (SM) Throughput）。它衡量最活跃执行流水线（GEMM 中通常是 tensor core）相对于峰值指令速率的繁忙程度。

• L1/TEX 利用率：由 L1/TEX 单元需要处理的总 sector 数推导得出。其中 num_L1_tex_sectors 对应 l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum 与 _st.sum 计数器，SM_active_cycles 对应 sm__cycles_active.avg，num_SM 为 132，而 num_sustained_peak_sectors_per_sm_per_cycle 在 H100 上为 2.0。

GPU L1纹理缓存（TEX）利用率的性能分析公式，核心目标在于量化算子对纹理缓存资源的实际使用效率。公式中的分子部分num_L1_tex_sectors / num_SM / SM_active_cycles，通过将纹理缓存访问的总扇区数，按照流多处理器（SM）数量及其活跃周期数进行平均分摊，从而得出每个SM在每个活跃周期内的平均纹理访问速率。分母num_sustained_peak_sectors_per_sm_per_cycle则代表硬件层面每个SM在每个周期内所能支撑的纹理缓存访问峰值。将这两者的比值乘以100%，即可获得L1纹理缓存的利用率百分比。这一指标能直观判断算子是否受限于纹理缓存带宽，或者是否存在资源浪费现象，为访存优化提供关键依据。

• DRAM 利用率：该指标由HBM（高带宽内存）的总读写字节数推导得出。其中，dram_bytes_read_and_write对应的是dram__bytes_read.sum和dram__bytes_write.sum这两个计数器的累加值。测试服务器上的峰值带宽为2TB/s。

DRAM带宽利用率（DRAM%）的标准计算公式，核心在于衡量程序运行时对内存带宽的实际利用程度。公式中，dram_bytes_read_and_write表示程序执行期间DRAM的总读写字节数，而duration代表程序的运行时长。将两者相除，即可得到实际DRAM传输带宽（即单位时间内的读写速率）。分母peak_bandwidth则是硬件标称的DRAM理论峰值带宽。用实际带宽除以峰值带宽，再乘以100，就能得出以百分比形式表示的带宽利用率。这个指标是性能优化的关键参考：高利用率表明程序的访存模式高效，能充分利用硬件带宽；低利用率则可能暗示内存访问存在瓶颈（例如缓存失效、访存碎片化），需要针对性地优化（如调整数据布局、提升缓存命中率），以减少内存带宽浪费，从而提升整体性能。

附录 5：第 2.2 节的详细结果分析

结果：1.389 ms → 0.798 ms（下降 42.5%）。

性能的大幅提升主要源于两个因素：

• TMA：矩阵对齐后，CUTLASS会选用基于TMA的内核，从而完全绕过L1/TEX（扇区数降至0）。未进行padding的内核还会导致矩阵B受到错误的惩罚：即使B的N维度已经对齐，它依然只能使用4-byte的加载操作，而无法采用128-bit的加载方式。
• Bank conflict：未进行padding的内核走的是sm80 MMA路径，其swizzle模式无法防止因4-byte的cp.async写入操作而引发的shared memory bank conflict。而经过padding的内核则不存在这一问题。

附录 6：第 2.3 节的详细结果分析

结果：延迟 0.798 ms → 0.580 ms（下降 27.4%）。

正如预期，中间环节中0.87 GB的DRAM流量被完全消除。NCU性能剖析也揭示了进一步的优化空间：occupancy仅为6.25%，每个调度器只有1个warp，PC采样显示仅有23%的周期用于执行有效工作：

当每个调度器只有1个warp时，所有的stall都会完整暴露出来，因为没有其他warp可供切换。如果试图通过减小pipeline深度来提升occupancy，又会牺牲K-loop的延迟隐藏能力。这正是该内核面临的核心困难：为了维持tensor core的吞吐量，需要更大的tile和更深的pipeline，但这会消耗掉大部分shared memory预算，导致几乎没有剩余空间再通过occupancy来隐藏延迟。

附录 7：Release-Acquire 同步协议

Producer（user CTA）：当user tile被写入global memory后，CTA以release语义设置每个tile的flag，确保数据在flag写入前已经可见：

tl.atomic_add(user_tile_flag_ptr, 1, sem="release", scope="gpu")

消费者端（候选CTA）：在每个warp group中，仅安排一个线程通过`ld.relaxed`指令轮询flag，以此降低自旋等待期间的内存流量。一旦检测到flag发生变化，该线程立即使用`ld.acquire`指令建立happens-before内存序边界，随后借助named barrier机制，将就绪状态广播至同一warp group内的全部128个线程：

```c
if tlx.thread_id(axis=0) % 128 == 0:  // 每个warp group（4个warp）仅1个线程
ready = tl.inline_asm_elementwise(
"ld.relaxed.gpu.global.b32 $0, [$1];", "=r,l",
[user_tile_flag_ptr], dtype=tl.int32, is_pure=False, pack=1)
while ready == 0:
ready = tl.inline_asm_elementwise(
"nanosleep.u32 50; ld.relaxed.gpu.global.b32 $0, [$1];", "=r,l",
[user_tile_flag_ptr], dtype=tl.int32, is_pure=False, pack=1)
tl.inline_asm_elementwise(
"ld.acquire.gpu.global.b32 $0, [$1];", "=r,l",
[user_tile_flag_ptr], dtype=tl.int32, is_pure=False, pack=1)
tlx.named_barrier_wait(12, 128)

附录8：TLX与Triton的NCU性能剖析指标

附录9：标准Flash Attention与IKBO Flash Attention的Roofline分析

在FP16/BF16精度下，设定user_seq_len = 1024、n_seed = 64、B_candidate : B_user = 70:1：

通过算术强度（AI，即FLOPs/Byte，衡量计算操作数与访存字节数的比值），对标准注意力算子、IKBO优化后的算子及H100 GPU硬件的算力瓶颈进行了量化对比。标准注意力算子的算术强度仅为60 FLOPs/Byte，远低于H100的硬件算术强度上限（495 FLOPs/Byte），表明该算子受限于访存带宽，属于访存瓶颈型任务，无法充分利用GPU算力。而IKBO优化通过减少用户序列的访存总量（分母中用户序列访存量除以候选数num_cand_user），将算术强度提升至833 FLOPs/Byte，超越了H100的硬件上限，说明优化后的算子转变为计算瓶颈型任务，理论上可完全利用GPU算力，大幅提升硬件利用率和实际性能。

附录10：IKBO TLX FA3的共享内存占用

附录11：推荐系统边界条件下，IKBO FA、CuTeDSL FA4 Hopper与TLX FA3 Hopper的性能基准测试

IKBO内核本质上实现了用户-候选交互映射逻辑，其IO与计算模式与GQA相近。基准测试时，对IKBO内核采用稳定的B_candidate : B_user = 64:1比例，并为CuTeDSL FA4 Hopper的GQA版本构建相似的计算模式（Q_seq_len = 128，以确保2-consumer warpgroup正常工作）。此外需说明：IKBO内核还需额外消耗candidate-user mapping张量，以处理实时排序中每用户候选数变化的问题。

IKBO FA与开源GQA内核的基准测试。IKBO内核的Q、K、V形状顺序为[Batch size, num head, seq, d_head]：Q_ikbo [2048, 2, 64, 128]，K/V_ikbo [32, 2, 1024, 128]。GQA内核的Q、K、V形状为：Q_gqa [1024, 2, 128, 128]，K/V_gqa [32, 2, 1024, 128]。

IKBO FA与开源GQA内核的基准测试。Q_ikbo [2048, 2, 64, 128]，K/V_ikbo [32, 2, 2048, 128]。GQA内核的Q_gqa [1024, 2, 128, 128]，K/V_gqa [32, 2, 2048, 128]。

注：由于标准Flash Attention内核不包含IKBO逻辑，我们采用具有相近IO成本和FLOPs消耗的GQA配置来模拟cuteDSL版本的吞吐结果。

附录 12：指令缓存缺失在 consumer-2 warpgroup 中引发显著延迟

图 A1 呈现的是注意力前向算子（attn_fwd）在 GPU 上的性能剖析时间线，详细展示了 Core 0 上 CTA76 内部多个 warp（线程束）的指令执行流程。图中以橙色框突出标识了由指令缓存缺失（Instruction cache miss）引发的性能瓶颈：右侧区域的 warp 指令流出现了明显的执行中断与空闲气泡。例如，warp 0 的内存加载指令（issue_v_load/issue_k_load）之间存在间隙，而 warp 4 和 warp 8 的核心计算操作（如 QKT、softmax）也遭受了打断。当指令缓存未命中时，GPU 必须从更高层级的内存层级中读取指令，这直接导致流水线停滞，并阻塞了 warp 调度器的正常工作。原本紧凑的指令流因此产生了显著的延迟，从而降低了 SM（流多处理器）的指令级并行效率与整体利用率。在循环密集型的算子中，这种缓存缺失问题会尤为严重，显著拖慢整体的执行性能。

指令缓存缺失修复前后的性能对比数据如下：

Before instruction cache miss fix:  
    ---------------------------------------------------- ----------- ------------  
    Metric Name                                          Metric Unit Metric Value  
    ---------------------------------------------------- ----------- ------------  
    gcc__cache_requests_type_instruction.sum                              319,394  
    gcc__cache_requests_type_instruction_lookup_miss.sum                    7,234  
    sm__icc_requests.sum                                       cycle    6,049,376  
    sm__icc_requests_lookup_hit.sum                            cycle    5,438,421  
    sm__icc_requests_lookup_miss.sum                           cycle      610,955  
    ---------------------------------------------------- ----------- ------------  

After instruction cache miss fix:  
    ---------------------------------------------------- ----------- ------------  
    Metric Name                                          Metric Unit Metric Value  
    ---------------------------------------------------- ----------- ------------  
    gcc__cache_requests_type_instruction.sum                               33,008  
    gcc__cache_requests_type_instruction_lookup_miss.sum                      769  
    sm__icc_requests.sum                                       cycle      792,437  
    sm__icc_requests_lookup_hit.sum                            cycle      722,244  
    sm__icc_requests_lookup_miss.sum                           cycle       70,193  
    ---------------------------------------------------- ----------- ------------  

参考资料

[1] Meta Adaptive Ranking Model: 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/

[2] TLX (Triton Low-Level Extensions): https://github.com/facebookexperimental/triton/tree/tlx

[3] https://github.com/pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/ikbo: https://github.com/pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/ikbo

[4] TLX（Triton Low-level Language Extensions）: https://github.com/facebookexperimental/triton/tree/tlx

[5] Proton tracer: https://github.com/triton-lang/triton/tree/main/third_party/proton/tutorials/intra_kernel

[6] Proton tracer: https://github.com/triton-lang/triton/tree/main/third_party/proton/tutorials/intra_kernel

[7] #831: https://github.com/facebookexperimental/triton/pull/831

• 小红书LarC：基于大规模深度学习的分钟级实时推荐系统
• EdgeRec：手机淘宝端上推荐系统面临的挑战与问题
• Meta AI基础设施的十年演进、下一代蓝图与自研芯片之路

如需加入交流群，请在NeuralTalk公众号后台回复：加群