PyTorch's FBTriton Kernels for Table Batched Embeddings Beat Legacy CUDA

What TBE does and how the forward pass works

TBE, short for Table-Batched Embedding, performs embedding lookups and pooling across many tables in a single GPU operation. The post says combining both steps in one launch reduces launch overhead and improves memory efficiency. The operators handle embedding lookups across thousands of sharded GPUs within recommendation systems.

In the forward pass, for each table and bag, the kernel gathers the indexed rows, optionally multiplies them by per-sample weights, accumulates them and writes one D-wide pooled output. FP16 and BF16 weights accumulate in FP32, while FP32 weights accumulate in FP64 to preserve accuracy at large D.

The team built two implementations. The generic gather path launches ceil(B / BAGS_PER_PROGRAM) programs, and each program loops over T features rather than launching a B×T grid. A specialized fast path is chosen for one feature with E≤64, 64≤D≤128, L≥64, FP16 weights, FP32 output, no per-sample weights and no VBE. There, one program handles 16 bags, builds a histogram over the first 256 indices and evaluates counts × table with tl.dot. A scalar path handles the remaining indices.

TorchRec can also accept config-driven int32 indices and offsets when the linearized range fits below 2^31. The post says this halves index and offset storage and the CUB radix-sort key width, with int64 remaining the default.

Bounds checking and forward-state reuse

The standalone path uses an updated CUDA validation step before the Triton forward pass. On B200, the post reports up to 1.24x speedup on the bounds-check component across workloads. When fused_bounds_check is enabled, the eligible kernel validates and repairs invalid input tensors, while offset validation and repair stay in a separate small kernel. The option is exposed through TorchRec and defaults off.

Exact row-wise Adagrad can save the forward histogram, and the backward pass uses those counts for a compensated FP16 high/low GEMM into FP32 before applying the optimizer. Forward and backward remain separate launches.

An optional path moves index transpose, sort and run-length encoding into the forward pass and returns the metadata through autograd. It defaults off and is not exposed by the current TorchRec wrapper. On a large B200 configuration, forward moves from 22.844 ms to 33.252 ms and backward from 56.693 ms to 32.931 ms. Combined latency falls from 79.537 ms to 66.183 ms, a reduction of 16.8%.

Backward kernel design

The backward logic works per unique (table, row) pair: it sums the upstream gradient rows from every batch position that touched the row, then applies exactly one optimizer update. The post notes that TBE backward has no tensor core operations, but requires heavy data movement and reduction and suffers from load imbalance.

To handle this, transpose_embedding_input inverts the batch into runs, each pairing one unique row with the samples that touched it. This is hoisted into the forward pass, off the backward critical path. Runs are routed by segment length (SL) to one of three kernels:

• short_run (SL < 256): one program owns the run from start to finish, with a plain store and no atomics.
• grad_accum + apply (SL >= 256, default): the run is split into 256-lookup chunks, partials land in a workspace, and a second kernel applies the optimizer.
• fused (SL >= 256, Blackwell, very large batches): the same split, but a device-scope fence lets the last sub-program apply the update in one launch.

Edge cases and tuning

The post describes several problems. For imbalanced programs, runs at or above the threshold become fixed 256-lookup chunks, so a two-million-lookup run turns into roughly eight thousand sub-programs. On Blackwell, CLC via TLX lets a persistent kernel keep stealing free run_ids, a hardware-based approach that removes the need to split run_id processing by frequency in software.

For the gather width register cliff, the team chose a per-target config width in every tier. Measured on B200, the short run, unweighted tier went from width 8 to 2, registers from 184 to 64 and occupancy from 12.5% to 49.9%. The long run accumulate tier went from 8 to 2, 158 to 62 registers and 17.6% to 44.7% occupancy, with fleet parity moving from 82% to 87%. The short run, weighted tier went from 4 to 2, 125 to 64 registers and 24.8% to 49.3% occupancy, with weighted parity from 51% to 69%.

Because BLOCK_SIZE is a constexpr sized to next_pow2(max_D), a narrow lookup-dominant table wastes lanes. The fix buckets short runs by next_pow2(D) during classification at no extra pass. Classification also stays on the GPU, using atomic counters for stream compaction, to avoid a cudaStreamSynchronize from reading counts back with .item(). For split runs, which need a cross-program barrier, the post says TLX on Blackwell exposes a fence that Triton lacked.

What it means

The post presents FBTriton as a Triton-based replacement for legacy CUDA kernels in a core recommendation system operator, and says it will highlight future optimization opportunities. Several features, including the fused bounds check and the forward-side preprocessing path, default to off, so users of TorchRec would need to opt in where exposed.