NVIDIA RecSys Examples

NVIDIA RecSys Examples is a collection of optimized recommender models and components designed for large-scale ranking, retrieval, and generative recommendation tasks. The project is structured into three primary functional areas:

1. Examples: High-level implementations for specific model architectures:
   • HSTU (Hierarchical Sequential Transformer Unit): For large-scale ranking and retrieval training (using TorchRec, Megatron-Core, and DynamicEmb) and optimized inference (using Triton Inference Server, CUDA graphs, and C++ deployment).
   • Semantic ID (SID-GR) Generative Recommender: For training and retrieval using hierarchical semantic-ID prediction and Megatron-Core decoder support.
   • SID-GR Inference: Specialized serving for long-context, short-decode, and large-beam scenarios using ContextKV/BeamKV/BeamPath abstractions.
2. Standalone GPU Libraries: Core optimized components:
   • DynamicEmb: Model-parallel dynamic embedding tables with GPU/host storage and TorchRec integration.
   • RecSys KVCache Manager: User-ID-based KV-cache reuse for generative recommender inference.
   • Beam search decode attention: Specialized kernels for SID-GR KV-cache generation.
3. Agentic Performance Optimization:
   • Talos: An automated toolkit that uses coding agents (like Claude Code) to profile, rewrite, and verify PyTorch model hotspots for performance gains without human intervention.

DynamicEmb is a Python package designed for model-parallel dynamic embedding tables and embedding lookup functionalities for TorchREC. It targets sparse training in recommendation systems by using a GPU-optimized scored hash table backend. This backend stores key-value (feature-embedding) pairs in both GPU high-bandwidth memory (HBM) and host memory.

Key features include:

• Dynamic Embedding Table Support: Uses hash tables to allow optimal utilization of GPU and host memory. Unlike static tables, hash tables can accept any indices type.
• TorchREC Integration: Inherits the TorchREC API, allowing users to easily modify existing recommendation system models to use dynamic embedding tables.
• Pooling Mode Support: Supports SUM, MEAN, and NONE (sequence) pooling modes with fused CUDA kernels. It also supports mixed-D (different embedding dimensions).
• Optimizer Support: Supports EXACT_SGD, ADAM, EXACT_ADAGRAD, and EXACT_ROWWISE_ADAGRAD.
• Parallel I/O: Supports automatically parallel dump/load of embedding weights.

The HSTU end-to-end training benchmark is a series of progressive experiments designed to quantify the performance contribution of various optimizations to training throughput (Model Flops Utilization - MFU). Each experiment builds upon the previous one, allowing for a cumulative measurement of speedups.

Optimization Progression

This benchmark measures the forward, backward, and evaluation overhead of BatchedDynamicEmbeddingTablesV2 compared to the TorchRec/FBGEMM SplitTableBatchedEmbeddingBagsCodegen baseline on a single GPU. It is implemented as a pytest suite (benchmark_batched_dynamicemb_tables.py) wrapped by a shell launcher (benchmark_batched_dynamicemb_tables.sh).

Available Test Suites

Default Configuration

• num_iterations: 100
• embedding_dim: 128
• feature_distribution: "pow-law" (alpha=1.05)
• emb_precision: fp32
• optimizer: Adam or SGD (learning_rate=0.1, eps=1e-8)

bash ./benchmark/benchmark_batched_dynamicemb_tables.sh

The HSTU (Hierarchical Sequential Transformer Unit) ranking inference benchmarks cover several implementation paths for evaluating performance across different runtimes and backends. These include:

• Triton Python-backend client: Used for testing the standard Python backend in Triton Server.
• PyTorch Export (No Cache): Uses export_inference_gr_ranking.py to export the model and compare Python runtime against native C++.
• PyTorch Export (With KV Cache): Uses export_inference_gr_ranking_kvcache.py to evaluate performance with KV-cache enabled.
• C++ Torch Replay: Native C++ execution of exported models, both with and without KV cache.
• Triton AOTI Deployment: Uses PyTorch AOTInductor (AOTI) for optimized deployment via Triton Server.
• Triton AOTI Request Replay: Uses test_tritonserver_aoti_hstu_model.py to replay requests against an AOTI-enabled Triton server.

The .torch_optimize/ directory is a run-local workspace created inside the code root being optimized. It serves as the shared state for the three primary roles in the optimization process: the Main Loop Agent, the Optimizer, and the Judge.

Each optimization run follows a hierarchical structure of loops and rounds:

1. Loop Level (loop_<n>)

A loop represents a single cycle from a specific baseline profile until a PASS verdict is achieved. Each loop contains:

• loop_baseline.patch: A git checkpoint of the code at the start of the loop (written by Main Loop Agent).
• loop.patch: The cumulative optimization patch containing all kept optimizations (written by Optimizer).
• loop_baseline.nsys-rep: The initial Nsys trace (written by Main Loop Agent).
• <torch-perf-analysis outputs>: Performance profiling data like basics.json (written by Main Loop Agent).
• loop_baseline_analysis.md: Performance analysis and hotspot identification (written by Main Loop Agent).

2. Round Level (round_<k>)

Within each loop, multiple rounds may occur if the Judge issues an ITERATE verdict. Each round contains:

• <nsys-rep + torch-perf-analysis outputs>: Profiling data for the specific change (written by Optimizer or Judge).
• verification.json: Real-input parity and end-to-end performance measurements (written by Optimizer).
• worklog.md: The Optimizer's log of activities.
• report.md: The Optimizer's report.
• verdict.md: The final decision from the Judge (e.g., PASS, ITERATE).