FBTriton
An implementation of the Table Batched Embedding operator for recommendation systems, written in Triton by the PyTorch team as a replacement for hand-written CUDA kernels. It supports forward gather and backward scatter computation, improving maintainability and portability across hardware.
Timeline (last 90 days)
The official PyTorch blog published a technical article titled "Modernizing Table Batched Embeddings with FBTriton," introducing the redesign of the TBE operator with the FBTriton kernel written in Triton.
The core idea of FBTriton is to re-express TBE's forward and backward computation using Triton, a high-level GPU programming language.
The FBTriton forward pass gathers vectors from the embedding table according to sparse indices and aggregates them along the sample dimension, and the backward pass scatters upstream gradients back to the embedding rows that were looked up, accumulating when the same row is hit multiple times.
The core value of the FBTriton refactor lies in maintainability: migrating TBE operator logic to a framework that is easier to iterate on can lower the barrier to optimization for new hardware and new model structures.
All Facts (4)
The core idea of FBTriton is to re-express TBE's forward and backward computation using Triton, a high-level GPU programming language.
50%UnverifiedThe FBTriton forward pass gathers vectors from the embedding table according to sparse indices and aggregates them along the sample dimension, and the backward pass scatters upstream gradients back to the embedding rows that were looked up, accumulating when the same row is hit multiple times.
50%UnverifiedThe core value of the FBTriton refactor lies in maintainability: migrating TBE operator logic to a framework that is easier to iterate on can lower the barrier to optimization for new hardware and new model structures.
50%UnverifiedThe official PyTorch blog published a technical article titled "Modernizing Table Batched Embeddings with FBTriton," introducing the redesign of the TBE operator with the FBTriton kernel written in Triton.
50%