Meta 的 PyTorch 團隊發布了 In-Kernel Broadcast Optimization (IKBO) 的架構細節,一種核心融合技術,消除 RecSys 推理中靜默昂貴的一個模式:在互動層之前物化 broadcast 張量。在典型推薦請求中,~15 個使用者 embedding 被複製 70 倍以匹配 1024 候選 batch,然後在 matmul 後立即丟棄。IKBO 把 broadcast 邏輯編碼進 GPU 核心本身 — 接受 mismatched batch size,在核心內做索引查找,從不物化複製的張量。H100 SXM5 上的關鍵數字:線性壓縮核心累計 4 倍提速(1.944ms → 0.482ms)、Flash Attention 端到端 6.4 倍 throughput(包括 broadcasting 成本,對比 CuTeDSL FA4-Hopper baseline)、在原本 IO-bound 在 250 TFLOPs 的 workload 上持續 621 BF16 TFLOPs。
技術洞察是:broadcast 是資料布局關注點,不是計算必需,節省透過四個漸進的協同設計階段層疊。階段 1 — matmul 分解 — 在自然的 15 行 batch 上跑使用者側 GEMM,在 1024 上跑候選側 GEMM,然後只 broadcast 小結果,把使用者側 compute 砍掉 70 倍。階段 2 — 記憶體對齊 — 把 K 填充到 8 的倍數以在 Hopper 上做 128-bit 對齊 TMA load,把 L1/TEX 流水線從 84% 飽和平衡,把 GEMM 延遲從 0.984ms 降到 0.400ms。階段 3 — 核心內 broadcast 融合 — 透過索引查找把 broadcast-add 折進候選 GEMM 的尾聲,消除 0.87 GB 的中間 DRAM 流量。階段 4 — 透過 TLX 做 warp-specialized 多階段融合 — 把 CTA 劃分成 producer + 兩個 consumer warp group,乒乓 tile 來重疊 WGMMA stall,把使用者和候選 GEMM 融合成單個持久核心,把 L2 throughput 從 74% 提到峰值的 84%。Flash Attention 的故事更有意思:標準 SDPA 在 ~60 FLOPs/Byte(IO-bound),而 IKBO FA 在 70:1 比例上把算術強度推到 ~833 FLOPs/Byte — 越過 H100 的 495 FLOPs/Byte 平衡點,把它牢牢地放在 compute-bound 上,Hopper 的 warp specialization 和非同步 TMA 才真正發揮作用。
生態讀法:這是大多數 ML 工程師還沒在想的一類最佳化,但泛化廣泛。任何有 mismatched batch 維度的推理 workload — 使用者/物品、供應商/產品、多級 broadcast 的層級排名 — 都有同樣模式。代碼住在 `pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/ikbo`(尚未合入 PyTorch core),Meta 在生產 RecSys 包括 MTIA 上部署了。兩條採用路徑:模型作者直接整合 IKBO 核心,或者 ML 編譯器 pass 在推理時把標準 op 換成 IKBO 等價物。對在大規模上跑 ranking、retrieval、recommendation 的 builder,workload shape 的匹配決定你拿到 2 倍、4 倍還是 6 倍;候選到使用者的比率線性 scale 節省。TLX 層(基於 Triton 的 warp specialization)本身也值得追蹤 — 這是一種沒去原生 CUDA 就難拿到的低層核心控制,Meta 在這上面的投資暗示它會合入 upstream。
實際動作:如果你跑生產 RecSys、ranking 或任何一個張量維度比另一個小得多的推理管線(想想個人化、供應商選擇、retrieval reranking),檢查你的核心熱路徑是否物化 broadcast 張量。如果是,IKBO 的 experimental 模組值得 benchmark — Meta 報告在協同設計模型上淨延遲最多減少 2/3,對 256-4096 的 batch size 和 10:1 到 10000:1 的比率穩健。他們預設 benchmark 中的 70:1 比率對廣告排名和 feed 個人化是現實的。如果你在 AMD 或非 Hopper 硬體上,架構洞察(把 broadcast 折進核心尾聲,消除物化)能 port — 具體數字不能,但模式可以。對 ML 編譯器人員,推理時變換路徑是要看的;如果 Meta 的編譯器 pass 上游了,這對生態系統其餘部分變免費。