TensorRT 11.0 Adds Native Multi-GPU Inference for Long-Context Media Models

NVIDIA TensorRT 11.0 adds native multi-device inference support, letting a single network execute across multiple GPUs through integrated distributed communication primitives. The feature integrates with NVIDIA NCCL, inheriting its transport optimization for inference workloads, and covers the full set of NCCL distributed collectives: AllReduce, Broadcast, Reduce, AllGather, ReduceScatter, AlltoAll, Gather, and Scatter. The core workflow matches single-device TensorRT, with the difference that the network can now include distributed communication layers. The same network is assumed on all GPU ranks, though each rank can run a different model in theory.

Parallelism strategies for distributed inference

Distributed inference can be expressed through tensor parallelism or context parallelism. In tensor parallelism, the weights of a single layer are partitioned across GPUs, with each GPU computing a shard of the layer's matrix multiplication and combining partial results through a collective. Column-parallel projections pair with row-parallel projections so each transformer block requires only a single AllReduce. This approach reduces per-device memory weight and becomes the natural choice when an individual layer's weights exceed single-GPU memory, independent of sequence length or batch size.

Context parallelism partitions the input sequence across GPUs along the sequence dimension, with collective operations making the global sequence available where needed, such as during attention. It suits long-sequence workloads where attention's quadratic scaling dominates compute and memory, and it fits diffusion and DiT models particularly well because their bidirectional attention avoids the load-imbalance issues of causal masks. TensorRT 11.0 introduces support for the IDistCollectiveLayer primitives required by these strategies.

Three context parallelism implementations are covered. AllGather KV exchanges key and value shards through an AllGather collective before computing local attention, enabling each rank to attend over the full sequence at the cost of one additional collective per attention block. Ring Attention overlaps communication and computation, letting each GPU process its local queries while keys and values stream past in a ring topology, and uses online softmax to avoid materializing full-size key and value tensors on any GPU. DeepSpeed Ulysses partitions samples along the sequence dimension, applies an all-to-all collective on partitioned Q, K, and V, then computes attention on non-overlapping attention head subsets, followed by a second all-to-all to gather results while repartitioning along the sequence dimension.

Benchmarks on media generation pipelines

Benchmarks evaluate multi-device TensorRT inference for media generation workloads intended for C++ production deployment. Two pipelines were used: a video generation pipeline based on NVIDIA Cosmos 3 and an image generation pipeline based on FLUX.1-dev from Black Forest Labs. Both were authored in PyTorch, then converted out of the framework using Torch-TensorRT to produce TensorRT engines for C++ inference applications. All results were collected on a single node with 8 GPUs.

For the Cosmos 3 video generation pipeline, Ulysses delivered the lowest end-to-end latency when the diffusion model had excessively long context lengths in the order of tens of thousands of input tokens. For the FLUX.1-dev image generation pipeline, Ulysses also won, though Ring Attention scaled well to 4 GPUs. The benchmarks compared end-to-end latency across AllGather KV, Ring Attention, and Ulysses strategies.

Getting started requires downloading TensorRT 11 from the NVIDIA Developer Portal, installing it, obtaining a single-node multi-GPU machine, and installing OpenMPI in the development environment. A working sample is provided in the TensorRT repository, and the Torch-TensorRT documentation covers converting PyTorch models for multi-device TensorRT deployment.