$ show meta-recsys/generative-recommenders
This repository implements generative recommendation models, primarily HSTU (Hierarchical Sequential Transduction Units) and DLRM-v3, for large-scale sequential recommendation. It is a hybrid codebase combining research experimentation (ICML'24 paper on generative recommenders) with production-grade inference benchmarking (MLPerf loadgen integration). The code provides custom high-performance CUDA and Triton kernels for flash attention, jagged tensor operations, and layer normalization, targeting modern NVIDIA GPUs (Hopper/Blackwell). It includes training pipelines with gin configuration for datasets like MovieLens, Amazon Books, KuaiRand, and synthetic streaming data up to 100B items. The repository also contains a full MLPerf inference benchmark suite with vendored loadgen library and end-to-end TorchScript tracing for deployment. While the core focus is research on scaling laws and architecture innovations for sequential recommenders, the ops and modules are engineered to production standards with multiple kernel backends and rigorous tests.
The LLM identified 3 files worth highlighting for reuse:
Implements the forward pass of HSTU flash attention using CUTLASS templates with Tensor Memory Accelerator (TMA), warp specialization, and persistent tile scheduling. Supports causal, local, contextual, and target-aware masking. This kernel is a reusable, high-performance attention implementation for variable-length sequences (jagged) on SM90+ GPUs, demonstrating advanced GPU programming patterns like named barriers and producer/consumer warp groups.
Provides Triton kernels for jagged tensor operations: concat_2D_jagged, split_2D_jagged, jagged_dense_bmm_broadcast_add, and jagged_dense_broadcast_add. These operations are essential for variable-length sequence data common in recommendation systems. The kernels are autotuned, support multiple data types, and include both forward and backward passes. They are isolated, well-documented, and can be extracted for any project requiring efficient jagged tensor manipulation on GPU.
Contains utility functions for kernel selection (HammerKernel enum) and custom op registration (maybe_register_custom_op). Provides a clean abstraction for switching between PyTorch, Triton, and CUDA implementations at runtime, with support for TorchScript tracing and FX graph capture. This pattern is reusable across any deep learning project that needs to maintain multiple kernel backends for the same operation.