NVIDIA released TensorRT 11.0, adding native multi‑device inference support to the TensorRT runtime. The feature is intended to let generative AI workloads scale across multiple GPUs without losing the production optimizations TensorRT provides—such as kernel fusions, memory planning, and quantization.
When combined with Torch‑TensorRT, the new multi‑device support enables developers to convert and deploy very large PyTorch models out‑of‑framework, breaking single‑device memory and compute limits. TensorRT 11.0 with multi‑device inference support is available from the NVIDIA Developer Portal.
NVIDIA NCCL: the transport layer for distributed inference
The NVIDIA Collective Communications Library (NCCL) supplies high‑performance multi‑GPU and multi‑node collective operations. NCCL automatically chooses the optimal transport for a topology (NVLink, NVSwitch, PCIe, InfiniBand) and exposes a uniform interface. By integrating directly with NCCL, TensorRT inherits these transport optimizations for multi‑device inference workloads.
The new multi‑device feature supports the full set of NVIDIA NCCL distributed collectives: AllReduce, Broadcast, Reduce, AllGather, ReduceScatter, AlltoAll, Gather and Scatter.
Parallelism strategies for distributed inference
Distributed inference can use several parallelism strategies with different trade‑offs among memory savings, compute scaling and communication overhead. The most common are tensor parallelism and context parallelism.
Tensor parallelism
In tensor parallelism, a single layer’s weights are partitioned across GPUs. Each GPU computes a shard of the layer’s matrix multiplication and partial results are combined via a collective to produce the full output. This reduces per‑device weight memory and is often the only option when an individual layer’s weights exceed a single GPU’s memory, regardless of input sequence length or batch size.
In transformer blocks, column‑parallel projections (e.g., QKV and MLP up‑projection) are paired with row‑parallel projections (attention output and MLP down‑projection) so each block typically requires only one AllReduce, keeping communication bounded.
Context parallelism
Context parallelism partitions the input sequence across GPUs along the sequence dimension. Each GPU processes only a slice of the sequence; collective operations make the global sequence available when needed (for example, in attention). This is particularly effective for long‑sequence workloads where attention’s quadratic scaling with sequence length dominates compute and memory.
Context parallelism is also a natural fit for diffusion and DiT models, since bidirectional attention avoids the load imbalance caused by causal masks.
TensorRT 11.0 adds support for the IDistCollectiveLayer primitives required by these parallelization strategies. The rest of this article focuses on context parallelism because it directly addresses the dominant cost in modern generative media pipelines: long‑sequence attention.
Context parallelism for generative media
Diffusion‑based image and video generation pipelines spend a large fraction of compute and memory inside attention blocks operating over long token sequences. A high‑resolution image latent or multi‑frame video clip can produce sequences of tens of thousands of tokens per block, and attention scales quadratically with sequence length.
Several context parallelism implementations offer different trade‑offs:
AllGather KV
In the AllGather KV approach, ranks exchange their key (K) and value (V) shards via an AllGather collective before computing local attention. Each rank processes the local slice of queries (Q) for its partition but, after the AllGather, can attend over the full sequence. This costs one extra collective per attention block, while the local Q × Kᵀ multiplication size shrinks roughly in proportion to the number of ranks. For video and high‑resolution image diffusion, the trade‑off compounds favorably across denoising steps: communication per step is bounded by the sequence‑dimension AllGather, while compute and memory savings apply at every attention layer and step.
Ring Attention
Ring Attention overlaps communication and computation by streaming K and V past GPUs in a ring topology while each GPU processes its local Q. This approach also reduces memory footprint because, with an online softmax, the full K and V tensors need not be materialized on any single GPU.
DeepSpeed Ulysses
For very long contexts (tens of thousands of tokens), DeepSpeed Ulysses offers an alternative. It initially partitions samples along the sequence dimension across GPUs. Before attention, it performs an all‑to‑all collective on the partitioned Q, K and V so that each GPU receives the full sequence length but only for a non‑overlapping subset of attention heads. GPUs compute attention in parallel across heads, and a second all‑to‑all gathers the results across heads while repartitioning along the sequence dimension.
Benchmarks: media generation with context parallelism in C++
The published benchmarks evaluate multi‑device TensorRT inference for media generation workloads intended for C++ production deployment. Two representative generative AI pipelines were used: a video generation pipeline based on NVIDIA Cosmos 3 and an image generation pipeline based on FLUX.1.
These pipelines were authored in PyTorch and converted out of framework with Torch‑TensorRT to produce TensorRT engines for C++ inference applications, allowing development in PyTorch while deploying optimized TensorRT engines in production.
Benchmarks compared end‑to‑end latency across context parallelism strategies: AllGather KV, Ring Attention and Ulysses. All results were collected on a single node with 8 GPUs.
Video generation with NVIDIA Cosmos 3
The NVIDIA Cosmos model platform (Cosmos3‑Nano) can generate images, video and audio from multimodal inputs. Using an example prompt file for the benchmarks, DeepSpeed Ulysses was the clear winner when diffusion models had extremely long context lengths (on the order of tens of thousands of input tokens).
Image generation with FLUX.1
The FLUX.1‑dev model from Black Forest Labs generates images from text prompts. Using the prompt “a beautiful photograph of Mt. Fuji during cherry blossom” in benchmarks, Ulysses again performed best overall, while Ring Attention scaled well up to 4 GPUs.
Getting started with TensorRT multi‑device
TensorRT supports multi‑device inference so a single network can execute across multiple GPUs using integrated distributed communication primitives. The core workflow resembles single‑device TensorRT, but networks can now include distributed communication layers (IDistCollectiveLayer). The example guide assumes the same network is deployed on all ranks, although different ranks could, in theory, run different models.
Key prerequisites and steps:
- Download and install TensorRT 11 from the NVIDIA Developer Portal following the provided instructions.
- Obtain a single‑node multi‑GPU machine.
- Install OpenMPI in your development environment (bare metal or container).
- Add collective operations to a TensorRT network via INetworkDefinition::addDistCollective and set the world size with setNbRanks.
The article includes code snippets illustrating network creation, adding a collective layer, building a serialized engine, creating an execution context, binding I/O tensors, setting a communicator on the context, and calling context->enqueueV3. Note that the NCCL communicator must remain valid for the lifetime of the execution context that uses it.
Example runtime invocation on 8 GPUs with OpenMPI (setting environment variables for rank, world size and an NCCL ID file):
mpirun -np 8 bash -lc 'export TRT_MY_RANK=$OMPI_COMM_WORLD_RANK; export TRT_WORLD_SIZE=$OMPI_COMM_WORLD_SIZE; export TRT_NCCL_ID_FILE=/tmp/nccl_id.txt; ./sample_dist_collective --op all_reduce'
Learn more
The article references further reading on NCCL, context parallelism for million‑token inference, Ring Attention, DeepSpeed Ulysses, the NVIDIA TensorRT product pages and the Torch‑TensorRT documentation and repository for readers who want to explore the implementation details.
In summary, TensorRT 11.0’s multi‑device capability and NCCL integration enable large generative models to run across multiple GPUs in production while retaining TensorRT optimizations. Of the context parallelism strategies evaluated on 8‑GPU single‑node setups, DeepSpeed Ulysses performed best for extremely long contexts in the provided benchmarks.
