NVIDIA Dynamo-Triton serves HSTU recommenders with AOTI and FlexKV caching

NVIDIA has published an end-to-end inference workflow for Hierarchical Sequential Transduction Unit (HSTU) generative recommender models in its recsys-examples repository, built on Dynamo-Triton, PyTorch Ahead-of-Time Inductor (AOTI) compilation, and FlexKV-backed KV caching. The workflow covers export, validation in Python and native C++, and deployment through Dynamo-Triton without rewriting the model for a separate runtime.

Generative recommender systems treat personalization as sequence modeling over user behavior, where interactions, context, candidate items, and actions become tokens in a high-cardinality event stream. This approach suits workloads with long user histories and changing catalogs, but serving requires low latency despite large embedding tables and sequence-heavy architectures. The NVIDIA workflow combines HSTU blocks, PyTorch AOTI, FlexKV, native C++ validation, NV embedding cache, and Dynamo-Triton deployment.

AOTI compilation and embedding management

The export path starts from a PyTorch model and uses torch.export and PyTorch AOTI to compile it ahead of time into a package loadable by a native C++ runtime. The package contains the AOTI model archive, metadata, and embedding table files. The embedding implementation combines DynamicEmb inference tables with NV Embedding Cache, which stores only popular embeddings in GPU memory while keeping the full table in CPU memory, reducing GPU memory usage. The export writes layer metadata and embedding data alongside the compiled .pt2 archive to avoid duplicate embedding table copies.

Validation happens at multiple levels. Python export scripts generate the package and replay tensors. Native C++ executables load and replay the exported model for correctness and performance checks. Dynamo-Triton deployment uses the same AOTI package and replay path, keeping development validation aligned with production serving.

KV caching with FlexKV

The workflow includes a KVCacheManager that uses GPU memory and host storage for KV data caches. The GPU cache is organized as a paged KV data table with lookup, allocation, append, and eviction, using an LRU-style policy when space is constrained. Host-side storage provides another tier. The HSTU attention kernel can consume KV data from the paged cache, and the exported inference path includes cache-aware custom operations for lookup, allocation, onboarding, appending, and offloading. This avoids recomputing cached portions of long user histories when a user returns with only new candidate items or recent actions.

Deployment and benchmark results

Dynamo-Triton serves the exported model through its PyTorch backend with platform "torch_aoti". The workflow has five stages: build custom operators and runtime libraries, export the HSTU model with AOTI, start the FlexKV-backed KV-cache service, validate with native C++ replay, and serve with Dynamo-Triton.

Benchmarks use the KuaiRand-1K ranking configuration on a single NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU. The model includes three-layer and eight-layer HSTU variants, hidden size 512, four attention heads, BF16 weights and KV cache, maximum history sequence length of 8,192 tokens, maximum candidate sequence length of 100, and six contextual features. The effective sequence length before alignment is 8,298 tokens, exported with a maximum aligned length of 8,320.

At Dynamo-Triton batch size 2, PyTorch AOTI improves latency over the Python backend even without KV-cache hits. With hits from a 20 GB GPU KV cache, the improvement is significantly larger. At dynamic batch size 8, Dynamo-Triton with AOTI achieved best-case speedups of up to 4.47x for the three-layer model and 5.93x for the eight-layer model under a 100% GPU KV-cache hit rate, relative to the same AOTI configuration without KV caching. KV caching becomes increasingly effective as logical batch size grows, with the deeper model showing more impact from avoiding recomputation.