FBTriton: PyTorch's Triton TBE Kernels Beat CUDA on GB200

PyTorch's FBTriton project reimplements the Table Batched Embedding (TBE) forward and backward kernels in Triton, and on GB200 it outperforms the legacy CUDA versions. Across 307 shard configurations with exact row-wise Adagrad and FP16 weights, the median forward speedup is 1.28x. The backward pass also shows gains.

Kernel Design and Config Tuning

The forward pass has two paths. The generic gather path launches programs that loop over features, with tuned gather widths and bag counts per program. A specialized path handles small tables with tensor-core operations, using a histogram over the first 256 indices and a scalar fallback for the rest. This path is selected only for specific shapes: one feature, E at most 64, D between 64 and 128, L at least 64, FP16 weights, FP32 output, no per-sample weights, and no variable batching.

Bounds checking also improved. On B200, the standalone validation step runs up to 1.24x faster than before. An optional fused bounds check, off by default, validates and repairs invalid input tensors inside the eligible kernel, while offset validation stays separate.

A key insight is that most wins come from config values, not kernel rewrites. The PyTorch team notes that changing an escalation point or gather width in template-generated CUDA means restructuring which kernel handles what, while Triton exposes these as simple configuration changes.

Backward Pass and Load Balancing

The backward pass reduces upstream gradients per unique table-row pair and applies one optimizer update. It has no tensor-core ops but suffers from load imbalance. Runs are classified by segment length (SL), the number of samples touching a row, which can span from one to millions.

Three kernels handle different SL ranges. Short runs use one program per run with exclusive ownership, avoiding atomics. Runs with SL at least 256 split into 256-lookup chunks, with partials accumulated in a workspace and a second kernel applying the optimizer. On Blackwell with very large batches, a fused variant uses a device-scope fence so the last sub-program applies the update in one launch.

Load imbalance is addressed in two ways. Long runs split into fixed 256-lookup chunks, turning a two-million-lookup run into about eight thousand sub-programs. Additionally, on Blackwell, a persistent kernel steals free run IDs after finishing its current one, providing a hardware-based solution instead of software bifurcation.

Register Pressure and Dimension Bucketing

Gather width creates a register cliff. Each buffered row keeps a full pointer tile live, so width times block size determines cost. Per-target config widths fix this. On B200, reducing width from 8 to 2 in the short-run unweighted tier dropped registers from 184 to 64, raised occupancy from 12.5% to 49.9%, and improved throughput from 0.41 to 1.05.

Dimension bucketing addresses the constexpr block size problem. One launch must size the block to the next power of two of the maximum D across all tables, wasting lanes on narrower tables. Short runs are routed to buckets per next power of two of D during classification, at no extra pass cost.

Blackwell-Specific Optimizations

Three fixes rely on Blackwell hardware. The CLC persistent kernel handles load imbalance. A device-scope fence enables the fused backward kernel, where the last sub-program applies the optimizer after all partials are visible. And TMA reduce operations merge partials more efficiently than element-wise atomics, using the async descriptor store with add reduction.

Where Triton Still Loses

Shapes with all work in runs shorter than 4 sit below parity due to dependent-load latency, where CUDA's warp-per-row amortizes metadata better. That is 11 of 307 shards, each under about a millisecond. Shapes above SL 256 are at parity rather than ahead. The team says these will get more investment if they become bottlenecks.

The most durable outcome, the team argues, is that the entire sparse path is now in Triton, enabling future optimizations beyond what template-generated CUDA allowed.