Meta IKBO: kernel-fused broadcast hits 6.4x Flash Attention speedup on H100

Meta's PyTorch team published the architecture details on In-Kernel Broadcast Optimization (IKBO), a kernel-fusion technique that eliminates one of the silently expensive patterns in RecSys inference: materializing broadcast tensors before interaction layers. In a typical recommendation request, ~15 user embeddings get replicated 70x to match a 1024-candidate batch, then dropped immediately after the matmul. IKBO encodes the broadcast logic into the GPU kernel itself — accept mismatched batch sizes, do index lookups inside the kernel, never materialize the replicated tensor. The headline numbers on H100 SXM5: 4x cumulative speedup on the linear compression kernel (1.944ms → 0.482ms), 6.4x throughput on Flash Attention end-to-end including broadcasting cost (vs CuTeDSL FA4-Hopper baseline), and 621 BF16 TFLOPs sustained on a workload that previously sat IO-bound at 250 TFLOPs.

The technical insight is that broadcast is a data-layout concern, not a computational necessity, and the savings cascade through four progressive co-design stages. Stage 1 — matmul decomposition — runs the user-side GEMM at its natural 15-row batch and the candidate-side GEMM at 1024, then broadcasts only the small result, cutting user-side compute 70x. Stage 2 — memory alignment — pads K to multiples of 8 for 128-bit aligned TMA loads on Hopper, balancing the L1/TEX pipeline from 84% saturated to balanced and dropping GEMM latency from 0.984ms to 0.400ms. Stage 3 — in-kernel broadcast fusion — folds the broadcast-add into the candidate GEMM epilogue via index lookup, eliminating 0.87 GB of intermediate DRAM traffic. Stage 4 — warp-specialized multi-stage fusion using TLX — partitions the CTA into producer + two consumer warp groups that ping-pong tiles to overlap WGMMA stalls, fuses the user and candidate GEMMs into a single persistent kernel, and lifts L2 throughput from 74% to 84% of peak. The Flash Attention story is even more interesting: standard SDPA sits at ~60 FLOPs/Byte (IO-bound), while IKBO FA pushes arithmetic intensity to ~833 FLOPs/Byte at the 70:1 ratio — past H100's 495 FLOPs/Byte balance point, putting it firmly compute-bound where Hopper's warp specialization and async TMA actually pay off.

The ecosystem read: this is a class of optimization most ML engineers haven't been thinking about, but it generalizes broadly. Any inference workload with mismatched batch dimensions — user/item, vendor/product, hierarchical ranking with multi-level broadcast — has the same pattern. Code lives in `pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/ikbo` (not yet merged into PyTorch core), and Meta deployed it across production RecSys including MTIA. Two adoption paths: model authors integrate IKBO kernels directly, or an ML compiler pass swaps standard ops for IKBO equivalents at inference time. For builders running large-scale ranking, retrieval, or recommendation, the workload-shape match is what determines whether this gives you 2x or 4x or 6x; the candidate-to-user ratio scales the savings linearly. The TLX layer (Triton-based warp specialization) is also worth tracking on its own — it's the kind of low-level kernel control that's been hard to get without going to raw CUDA, and Meta's investment here suggests it'll get merged upstream.

Practical move: if you're running production RecSys, ranking, or any inference pipeline where one tensor dimension is much smaller than another (think personalization, vendor selection, retrieval reranking), check whether your kernel hot-path materializes broadcast tensors. If it does, IKBO's experimental module is worth a benchmark — Meta's reporting up to 2/3 net latency reduction on co-designed models, robust across batch sizes 256-4096 and ratios from 10:1 to 10,000:1. The 70:1 ratio in their default benchmark is realistic for ad ranking and feed personalization. If you're on AMD or non-Hopper hardware, the architectural insight (fold broadcast into kernel epilogue, eliminate materialization) ports — the specific numbers don't, but the pattern does. For ML compiler folks, the inference-time transformation path is the one to watch; if Meta's compiler pass goes upstream, this becomes free for the rest of the ecosystem.