Table-Batched Embedding (TBE) bündelt Lookups und Pooling über mehrere Embedding-Tabellen in einem GPU-Aufruf. Solche Operationen kommen laut Meta unter anderem in Empfehlungssystemen mit auf viele GPUs verteilten Tabellen zum Einsatz. Die neuen FBTriton-Kernel decken Vorwärts- und Rückwärtsdurchlauf ab. Meta berichtet, dass die Triton-Implementierung die bisherigen CUDA-Kernel in den untersuchten Workloads übertrifft, nennt im Auszug aber keinen einzelnen Gesamtwert für diesen Vergleich.
Im Vorwärtsdurchlauf liest der Kernel die angeforderten Tabellenzeilen, verrechnet sie bei Bedarf mit Sample-Gewichten und fasst sie zu einem Ergebnis pro Embedding-Bag zusammen. Für die Akkumulation verwendet er FP32 bei FP16- und BF16-Gewichten sowie FP64 bei FP32-Gewichten. Eine spezialisierte Variante nutzt für bestimmte kleine Tabellen ein Histogramm und Tensor-Core-Berechnungen; sie greift nur bei eng definierten Bedingungen, etwa FP16-Gewichten, FP32-Ausgabe und ohne Sample-Gewichte oder variable Batches. Für andere Konfigurationen kommt der allgemeine Kernel zum Einsatz.
Eine optionale Bounds-Check-Variante prüft und korrigiert ungültige Eingabetensoren direkt in geeigneten Kernel-Konfigurationen. Auf B200 erzielte Meta für diesen Prüfschritt über mehrere Workloads bis zu 1,24-fache Beschleunigung. Die Option ist in TorchRec verfügbar und standardmäßig deaktiviert. Gewichtete und variable Batches sowie AMD-Konfigurationen und bestimmte Transpose-Fälle verwenden weiterhin die übliche Validierung.
Im Rückwärtsdurchlauf werden Gradienten für gleiche Tabellenzeilen zusammengeführt, bevor der Optimizer die jeweilige Zeile aktualisiert. Ein weiterer optionaler Pfad verlagert Transponieren, Sortieren und Run-Length-Encoding der Indizes in den Vorwärtsdurchlauf und reicht Metadaten über Autograd weiter. Bei einer großen B200-Konfiguration stieg dadurch die Vorwärtslaufzeit von 22,844 auf 33,252 ms, während die Rückwärtslaufzeit von 56,693 auf 32,931 ms sank. Zusammen verringerte sich die Laufzeit von 79,537 auf 66,183 ms, laut Meta um 16,8 Prozent. Dieser Pfad ist standardmäßig deaktiviert und derzeit nicht über den TorchRec-Wrapper zugänglich.