diff --git a/CMakeLists.txt b/CMakeLists.txt index 3df1be1..05def0b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,4 +1,4 @@ -# Copyright 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -232,6 +232,35 @@ target_link_libraries( CUDA::cudart ) +# +# TensorRT Multi-Device (MD) support — TRT-28040. +# Links NCCL so the backend can drive a sharded engine across multiple GPUs in +# a single process. NCCL ships in the NGC TensorRT/Triton images. When disabled +# the backend builds with no NCCL dependency and the MD code path is compiled +# out (see TRITON_ENABLE_TRT_MULTI_DEVICE). +# +option(TRITON_ENABLE_TENSORRT_MULTI_DEVICE + "Enable TensorRT Multi-Device (NCCL) execution" ON) +if(TRITON_ENABLE_TENSORRT_MULTI_DEVICE) + find_path(NCCL_INCLUDE_DIR + NAMES nccl.h + HINTS ${CUDAToolkit_INCLUDE_DIRS} /usr/include /usr/local/cuda/include) + find_library(NCCL_LIBRARY + NAMES nccl + HINTS /usr/lib/x86_64-linux-gnu /usr/local/cuda/lib64) + if(NOT NCCL_LIBRARY OR NOT NCCL_INCLUDE_DIR) + message(FATAL_ERROR + "NCCL not found (nccl.h / libnccl). Required for " + "TRITON_ENABLE_TENSORRT_MULTI_DEVICE=ON; set NCCL_INCLUDE_DIR / " + "NCCL_LIBRARY or pass -DTRITON_ENABLE_TENSORRT_MULTI_DEVICE=OFF.") + endif() + message(STATUS "TensorRT Multi-Device enabled: ${NCCL_LIBRARY}") + target_compile_definitions(triton-tensorrt-backend + PRIVATE TRITON_ENABLE_TRT_MULTI_DEVICE=1) + target_include_directories(triton-tensorrt-backend PRIVATE ${NCCL_INCLUDE_DIR}) + target_link_libraries(triton-tensorrt-backend PRIVATE ${NCCL_LIBRARY}) +endif() + # # Install diff --git a/docs/build_tp_engines.cpp b/docs/build_tp_engines.cpp new file mode 100644 index 0000000..78304bd --- /dev/null +++ b/docs/build_tp_engines.cpp @@ -0,0 +1,175 @@ +// Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// * Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// * Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// * Neither the name of NVIDIA CORPORATION nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +// Offline TP weight-shard + engine builder (TRT-28040 follow-on). +// Builds a Megatron MLP (Y = (X*W1)*W2) and serializes: +// model_sd.plan - full weights, single device, no collective +// (reference) model.plan.rank{r} - rank r's weight shards + AllReduce +// (true tensor parallel) +// W1 is column-parallel (split on output dim F); W2 is row-parallel (split on +// input dim F) with a trailing AllReduce, so each rank holds ~1/world of the +// weights. Same deterministic weights across SD and the TP shards. +// +// nvcc -std=c++17 -ccbin g++ -w build_tp_engines.cpp -o build_tp_engines \ +// -I$TRT/include -L$TRT/lib -lnvinfer +// LD_LIBRARY_PATH=$TRT/lib ./build_tp_engines +#include +#include + +#include +#include +#include +#include +#include +#include +using namespace nvinfer1; + +#define PREVIEW_MD 1 // TRT 10.16 needs the preview flag; remove for TRT 11.0 + +class Logger : public ILogger { + void log(Severity s, const char* m) noexcept override + { + if (s <= Severity::kERROR) + fprintf(stderr, "[TRT] %s\n", m); + } +} gLogger; + +constexpr int M = 2048, K = 4096, F = 8192, + N = 4096; // seq, hidden, inter, out +static std::vector W1, W2; // K*F, F*N (deterministic) + +static void +initWeights() +{ + W1.resize((size_t)K * F); + W2.resize((size_t)F * N); + for (size_t i = 0; i < W1.size(); ++i) + W1[i] = 0.01f * ((int)(i % 17) - 8) / std::sqrt((float)K); + for (size_t i = 0; i < W2.size(); ++i) + W2[i] = 0.01f * ((int)(i % 13) - 6) / std::sqrt((float)F); +} + +static std::vector<__half> +toHalf(const float* p, size_t n) +{ + std::vector<__half> h(n); + for (size_t i = 0; i < n; ++i) h[i] = __float2half(p[i]); + return h; +} + +static void +writeFile(const std::string& path, IHostMemory* m) +{ + std::ofstream f(path, std::ios::binary); + f.write(static_cast(m->data()), m->size()); + printf("wrote %s (%zu bytes)\n", path.c_str(), m->size()); +} + +// Build an MLP engine. If world==1: full W1[K,F], W2[F,N], no collective. +// Else: rank's column shard W1[:, r*Fr:(r+1)*Fr], row shard W2[r*Fr:(r+1)*Fr, +// :], +// + AllReduce(SUM, nbRanks=world). +static std::vector +buildEngine(int world, int rank) +{ + const int Fr = (world == 1) ? F : F / world; + const int f0 = (world == 1) ? 0 : rank * Fr; + std::vector w1(K * Fr), w2(Fr * N); + for (int k = 0; k < K; ++k) + for (int f = 0; f < Fr; ++f) w1[k * Fr + f] = W1[(size_t)k * F + (f0 + f)]; + for (int f = 0; f < Fr; ++f) + for (int n = 0; n < N; ++n) + w2[(size_t)f * N + n] = W2[(size_t)(f0 + f) * N + n]; + + std::unique_ptr builder(createInferBuilder(gLogger)); + std::unique_ptr net(builder->createNetworkV2( + 1U << (uint32_t)NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)); + ITensor* x = net->addInput("X", DataType::kFLOAT, Dims2{M, K}); + auto* c1 = net->addConstant( + Dims2{K, Fr}, Weights{DataType::kFLOAT, w1.data(), (int64_t)w1.size()}); + auto* h = net->addMatrixMultiply( + *x, MatrixOperation::kNONE, *c1->getOutput(0), MatrixOperation::kNONE); + auto* c2 = net->addConstant( + Dims2{Fr, N}, Weights{DataType::kFLOAT, w2.data(), (int64_t)w2.size()}); + auto* p = net->addMatrixMultiply( + *h->getOutput(0), MatrixOperation::kNONE, *c2->getOutput(0), + MatrixOperation::kNONE); + ITensor* y = p->getOutput(0); + if (world > 1) { + auto* coll = net->addDistCollective( + *y, CollectiveOperation::kALL_REDUCE, ReduceOperation::kSUM, -1, + nullptr, 0); + coll->setNbRanks(world); + y = coll->getOutput(0); + } + y->setName("Y"); + net->markOutput(*y); + + std::unique_ptr cfg(builder->createBuilderConfig()); +#if PREVIEW_MD + if (world > 1) + cfg->setPreviewFeature(PreviewFeature::kMULTIDEVICE_RUNTIME_10_16, true); +#endif + std::unique_ptr ser(builder->buildSerializedNetwork(*net, *cfg)); + if (!ser) { + fprintf(stderr, "build failed world=%d rank=%d\n", world, rank); + std::abort(); + } + const char* d = static_cast(ser->data()); + std::vector v(d, d + ser->size()); + return v; +} + +int +main(int argc, char** argv) +{ + int world = (argc > 1) ? atoi(argv[1]) : 2; + std::string dir = (argc > 2) ? argv[2] : "."; + initWeights(); + // Single-device reference (full weights, no collective) + { + auto e = buildEngine(1, 0); + std::ofstream f(dir + "/model_sd.plan", std::ios::binary); + f.write(e.data(), e.size()); + printf( + "wrote %s/model_sd.plan (%zu bytes, full weights)\n", dir.c_str(), + e.size()); + } + // Per-rank TP engines + for (int r = 0; r < world; ++r) { + auto e = buildEngine(world, r); + std::ofstream f( + dir + "/model.plan.rank" + std::to_string(r), std::ios::binary); + f.write(e.data(), e.size()); + printf( + "wrote %s/model.plan.rank%d (%zu bytes, weight shard)\n", dir.c_str(), + r, e.size()); + } + printf( + "dims: X[%d,%d] -> Y[%d,%d], F=%d split across %d ranks (Fr=%d)\n", M, K, + M, N, F, world, F / world); + return 0; +} diff --git a/docs/create_onnx_multilayer.py b/docs/create_onnx_multilayer.py new file mode 100644 index 0000000..743a02d --- /dev/null +++ b/docs/create_onnx_multilayer.py @@ -0,0 +1,286 @@ +# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# Multi-layer self-attention ONNX generator (TRT-28040 perf demo). +# Stacks N attention blocks so the model is compute-bound (the single-layer +# sample is overhead-bound at batch=1, hiding the multi-GPU speedup). Also emits +# a matching polygraphy context-parallel sharding hint (one attention_layers +# entry per block, keyed by each block's uniquely-named q_scaled_i tensor). +# +# python3 create_onnx_multilayer.py --layers 8 --output attn_ml.onnx --hint attn_ml_hint.json +import argparse +import json +import math + +import numpy as np +import onnx +import onnx_graphsurgeon as gs + +NUM_HEADS = 32 +HEAD_DIM = 128 +HIDDEN_DIM = NUM_HEADS * HEAD_DIM # 4096 +OPSET = 17 + + +@gs.Graph.register() +def op1(self, op, a, attrs=None): + return self.layer(op=op, inputs=[a], attrs=attrs or {}, outputs=[op + "_o"])[0] + + +@gs.Graph.register() +def matmul(self, a, b): + return self.layer(op="MatMul", inputs=[a, b], outputs=["mm_o"])[0] + + +@gs.Graph.register() +def transpose(self, a, perm): + return self.layer( + op="Transpose", inputs=[a], attrs={"perm": perm}, outputs=["tr_o"] + )[0] + + +@gs.Graph.register() +def reshape(self, data, shape): + return self.layer( + op="Reshape", inputs=[data, shape], attrs={"allowzero": 0}, outputs=["rs_o"] + )[0] + + +@gs.Graph.register() +def softmax(self, a, axis=-1): + return self.layer(op="Softmax", inputs=[a], attrs={"axis": axis}, outputs=["sm_o"])[ + 0 + ] + + +@gs.Graph.register() +def cast(self, a, to): + return self.layer(op="Cast", inputs=[a], attrs={"to": to}, outputs=["cast_o"])[0] + + +@gs.Graph.register() +def binop(self, op, a, b): + return self.layer(op=op, inputs=[a, b], outputs=[op + "_o"])[0] + + +@gs.Graph.register() +def reduce_mean(self, a, axes): + return self.layer( + op="ReduceMean", + inputs=[a], + attrs={"axes": axes, "keepdims": 1}, + outputs=["rm_o"], + )[0] + + +@gs.Graph.register() +def shape_op(self, a): + return self.layer(op="Shape", inputs=[a], outputs=["sh_o"])[0] + + +@gs.Graph.register() +def gather(self, data, indices): + return self.layer( + op="Gather", inputs=[data, indices], attrs={"axis": 0}, outputs=["ga_o"] + )[0] + + +@gs.Graph.register() +def unsqueeze(self, a, axes): + return self.layer(op="Unsqueeze", inputs=[a, axes], outputs=["un_o"])[0] + + +@gs.Graph.register() +def concat(self, inputs, axis=0): + return self.layer( + op="Concat", inputs=inputs, attrs={"axis": axis}, outputs=["cc_o"] + )[0] + + +def attention_block(graph, x, idx, rng): + axes_0 = np.array([0], dtype=np.int64) + + def w(): + # Scale ~1/sqrt(hidden) so stacked matmuls stay numerically bounded + # (no residual/norm between blocks in this synthetic model). + return ( + rng.standard_normal((HIDDEN_DIM, HIDDEN_DIM)) / math.sqrt(HIDDEN_DIM) + ).astype(np.float16) + + def s32(v): + return np.array([v], dtype=np.float32) + + q_proj = graph.matmul(x, w()) + k_proj = graph.matmul(x, w()) + v_proj = graph.matmul(x, w()) + + def to_heads(proj): + sh = graph.shape_op(proj) + sd = graph.unsqueeze(graph.gather(sh, np.array(0, dtype=np.int64)), axes_0) + bd = graph.unsqueeze(graph.gather(sh, np.array(1, dtype=np.int64)), axes_0) + tgt = graph.concat( + [ + sd, + bd, + np.array([NUM_HEADS], dtype=np.int64), + np.array([HEAD_DIM], dtype=np.int64), + ] + ) + return graph.reshape(proj, tgt) + + q4, k4, v4 = to_heads(q_proj), to_heads(k_proj), to_heads(v_proj) + + def rmsnorm(t): + f = graph.cast(t, onnx.TensorProto.FLOAT) + sq = graph.binop("Pow", f, s32(2.0)) + mean = graph.reduce_mean(sq, axes=[-1]) + rms = graph.op1("Sqrt", graph.binop("Add", mean, s32(1e-6))) + inv = graph.binop("Div", s32(1.0), rms) + nf = graph.binop("Mul", f, inv) + n16 = graph.cast(nf, onnx.TensorProto.FLOAT16) + wt = rng.standard_normal((1, 1, 1, HEAD_DIM)).astype(np.float16) + return graph.binop("Mul", wt, n16) + + qn, kn = rmsnorm(q4), rmsnorm(k4) + qa = graph.transpose(qn, perm=[1, 2, 0, 3]) + ka = graph.transpose(kn, perm=[1, 2, 0, 3]) + va = graph.transpose(v4, perm=[1, 2, 0, 3]) + + def to_attn(t): + sh = graph.shape_op(t) + b = graph.unsqueeze(graph.gather(sh, np.array(0, dtype=np.int64)), axes_0) + h = graph.unsqueeze(graph.gather(sh, np.array(1, dtype=np.int64)), axes_0) + d = graph.unsqueeze(graph.gather(sh, np.array(3, dtype=np.int64)), axes_0) + tgt = graph.concat([b, h, np.array([-1], dtype=np.int64), d]) + return graph.reshape(t, tgt) + + qr, kr, vr = to_attn(qa), to_attn(ka), to_attn(va) + sc = np.array([math.sqrt(math.sqrt(1.0 / HEAD_DIM))], dtype=np.float16) + q_scaled = graph.binop("Mul", qr, sc) + q_scaled.name = "q_scaled_%d" % idx # hint targets this per layer + kt = graph.transpose(kr, perm=[0, 1, 3, 2]) + ks = graph.binop("Mul", kt, sc) + qk = graph.matmul(q_scaled, ks) + aw = graph.softmax(qk, axis=-1) + ao = graph.matmul(aw, vr) + at = graph.transpose(ao, perm=[2, 0, 1, 3]) + sh = graph.shape_op(at) + sd = graph.unsqueeze(graph.gather(sh, np.array(0, dtype=np.int64)), axes_0) + bd = graph.unsqueeze(graph.gather(sh, np.array(1, dtype=np.int64)), axes_0) + hh = graph.unsqueeze( + graph.binop( + "Mul", + graph.gather(sh, np.array(2, dtype=np.int64)), + graph.gather(sh, np.array(3, dtype=np.int64)), + ), + axes_0, + ) + flat = graph.reshape(at, graph.concat([sd, bd, hh])) + return graph.matmul(flat, w()) # output projection -> block output + + +def build(layers): + rng = np.random.default_rng(42) + graph = gs.Graph(opset=OPSET) + x = gs.Variable( + "input", dtype=np.float16, shape=["sequence_length", "batch_size", HIDDEN_DIM] + ) + graph.inputs = [x] + for i in range(layers): + x = attention_block(graph, x, i, rng) + x.name = "output" + x.dtype = np.float16 + x.shape = ["sequence_length", "batch_size", HIDDEN_DIM] + graph.outputs = [x] + graph.cleanup().toposort() + m = gs.export_onnx(graph) + m.ir_version = 8 + return m + + +def make_hint(layers, path): + hint = { + "parallelism": "CP", + "attention_layers": [ + { + "q": "q_scaled_%d" % i, + "gather_kv": True, + "gather_q": False, + "replace": None, + "polygraphy_class": "AttentionLayerHint", + } + for i in range(layers) + ], + "dist_collectives": { + "group_size": 0, + "root": -1, + "nb_rank": 2, + "reduce_op": "max", + "groups": [], + "polygraphy_class": "DistCollective", + }, + "inputs": [ + { + "name": "input", + "seq_len_idx": 0, + "rank": 3, + "polygraphy_class": "ShardTensor", + } + ], + "outputs": [ + { + "name": "output", + "seq_len_idx": 0, + "rank": 3, + "polygraphy_class": "ShardTensor", + } + ], + "k_seq_len_idx": 3, + "v_seq_len_idx": 2, + "kv_rank": 4, + "polygraphy_class": "ShardHints", + } + with open(path, "w") as f: + json.dump(hint, f, indent=2) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--layers", type=int, default=6) + ap.add_argument("--output", default="attn_ml.onnx") + ap.add_argument("--hint", default="attn_ml_hint.json") + a = ap.parse_args() + m = build(a.layers) + onnx.save(m, a.output) # inline weights (keep <2GB protobuf limit) + make_hint(a.layers, a.hint) + print( + "Saved %s (%d layers, %d nodes) + hint %s" + % (a.output, a.layers, len(m.graph.node), a.hint) + ) + + +if __name__ == "__main__": + main() diff --git a/docs/multi_device.md b/docs/multi_device.md new file mode 100644 index 0000000..184b45f --- /dev/null +++ b/docs/multi_device.md @@ -0,0 +1,125 @@ +# TensorRT Multi-Device (multi-GPU) support — TRT-28040 + +This backend can run a single TensorRT engine sharded across multiple GPUs using +TensorRT's Multi-Device feature (NCCL `DistCollective` operations), so one Triton +model instance drives N GPUs. + +## How it works + +A Multi-Device instance is a single `KIND_MODEL` instance that owns N GPUs: + +- At load (`InitMultiDevice`): create N in-process NCCL communicators with + `ncclCommInitAll`; create one `IExecutionContext` per rank (rank 0 reuses the + instance's existing engine/context, ranks 1..N-1 get their own engine on their + GPU); attach the communicator to every rank **concurrently** with + `IExecutionContext::setCommunicator()` (a sequential loop deadlocks, since the + call is a cross-rank handshake). +- At inference (`EnqueueMultiDevice`): replicate rank-0's input tensors to every + rank, then issue `enqueueV3` on all ranks concurrently so the in-engine + collectives rendezvous; rank 0's output is returned. Input replication uses a + direct `cudaMemcpyPeer` when peer access is available (opt-in via + `TRT_MD_USE_PEER`, see Notes) and otherwise a pinned host bounce buffer. + +The single-GPU `KIND_GPU` path is unchanged and remains the default. + +## Model configuration + +```protobuf +backend: "tensorrt" +max_batch_size: 0 +instance_group [ { kind: KIND_MODEL count: 1 } ] +parameters [ + { key: "enable_multi_device" value: { string_value: "true" } }, + { key: "multi_device_gpus" value: { string_value: "0,1" } } +] +``` + +The engine is sharded offline (e.g. `polygraphy multi-device shard`) and the same +plan is deserialized on every rank; the rank is distinguished at runtime by the +communicator. Requires TensorRT >= 11.0 (Multi-Device GA) and NCCL; the backend +must be built with `TRITON_ENABLE_TENSORRT_MULTI_DEVICE=ON`. + +## Validation (TRT 11.0) + +A context-parallel self-attention model (4096 hidden, 32 heads, FP16), sharded +with `polygraphy multi-device shard`, served as `attn_cp` (2-GPU, KIND_MODEL) and +compared against the unsharded `attn_sd` (1-GPU, KIND_GPU) on the same input. + +Accuracy (2-GPU CP vs 1-GPU, identical input): + +| seq | result | +|--------|-------------------| +| 8192 | bit-identical (max_abs_diff 0.0) | +| 16384 | correct, rel_max ~6e-3 (FP16 split-reduction) | +| 32768 | correct | +| 65536 | correct | + +Latency, 1x vs 2x NVIDIA B200 (NVLink, NV18 / ~900 GB/s; batch=1, single +request). Numbers are noisy run-to-run (±10%): + +| seq / model | 1-GPU | 2-GPU CP | speedup (observed range) | +|------------------------|------------|------------|--------------------------| +| 16384 (hidden 4096) | ~480–724 ms| ~495–598 ms| ~0.96–1.21x | +| 32768 (hidden 4096) | ~1160 ms | ~1100 ms | ~1.03–1.05x | +| 65536 (hidden 4096) | ~2224 ms | ~2240 ms | ~0.99–1.03x (cold ~2x) | +| 8192 (hidden 8192) | ~593–728 ms| ~543–791 ms| ~0.92–1.11x | + +Single-attention-layer models at batch=1 are overhead/latency-bound on B200 +(per-request input replication, collective launch, kernel-launch overhead), so +the multi-GPU speedup is hidden by noise — that is why the table above is +neutral-to-modest with high variance. + +**Multi-layer model (compute-bound).** Stacking 6 attention blocks (so compute +dominates the fixed per-request overhead) gives a consistent speedup on 2x B200 +over NVLink at seq 32768: + +| run | 1-GPU | 2-GPU CP | speedup | +|-----|----------|----------|---------| +| 1 | 1388 ms | 1202 ms | 1.15x | +| 2 | 1349 ms | 1185 ms | 1.14x | +| 3 | 1209 ms | 1102 ms | 1.10x | +| 4 | 1187 ms | 1127 ms | 1.05x | + +All runs faster, ~1.1x average, accuracy correct (rel_err 4.7e-3). The gain grows +with model depth as the fixed per-request overhead amortizes; a full transformer +(dozens of layers, larger batch) is expected to approach the ideal ~2x. + +(Generator: `create_onnx_multilayer.py --layers N` builds the stacked model and a +matching per-layer CP hint.) + +## Weight-sharded tensor parallelism (per-rank engines) + +The default MD path loads the **same** engine on every rank (context/activation +parallel — weights replicated). For **true Megatron tensor parallelism**, set +`multi_device_per_rank_engines: "true"` and place a **distinct engine per rank** +in the version dir: `model.plan.rank0`, `model.plan.rank1`, ... Each engine holds +only that rank's weight shard (column-parallel GEMM → no comm; row-parallel GEMM +→ trailing `DistCollective` AllReduce), giving ~1/N per-GPU weight memory. (Keep a +`model.plan` symlink to `model.plan.rank0` so Triton's repository check passes.) + +```protobuf +instance_group [ { kind: KIND_MODEL count: 1 } ] +parameters [ + { key: "enable_multi_device" value: { string_value: "true" } }, + { key: "multi_device_gpus" value: { string_value: "0,1" } }, + { key: "multi_device_per_rank_engines" value: { string_value: "true" } } +] +``` + +Rank r loads `model.plan.rank{r}` (`Create()` for rank 0, `InitMultiDevice()` for +the rest); the execute path is unchanged (full input replicated to every rank, +output from rank 0 after the AllReduce). Validated on 2x A30 with a Megatron MLP: +2-GPU weight-TP output matches the single-GPU reference (rel_err ~9e-5), with each +rank's engine ~half the full-weight engine size (134 MB vs 268 MB). Engines were +built with `docs/build_tp_engines.cpp` (no polygraphy). TRT MD supports distinct +per-rank engines sharing one communicator with a matching collective (validated +standalone before wiring the backend). + +## Notes / known limitations + +- Direct `cudaMemcpyPeer` of TensorRT IO buffers did not transfer correctly in the + threaded backend context on the tested systems (likely pooled/virtual device + memory), so input replication defaults to a pinned host bounce buffer. The P2P + path is opt-in via the `TRT_MD_USE_PEER` env var pending a fix. +- Single optimization profile per MD model (one communicator per context). +- CUDA graphs are not used on the MD path. diff --git a/src/instance_state.cc b/src/instance_state.cc index 8cd2a57..415bbe1 100644 --- a/src/instance_state.cc +++ b/src/instance_state.cc @@ -1,4 +1,4 @@ -// Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -29,6 +29,14 @@ #include "tensorrt_utils.h" #include "triton/common/nvtx.h" +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE +#include + +#include +#include +#include +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE + namespace triton { namespace backend { namespace tensorrt { namespace { @@ -144,19 +152,28 @@ ModelInstanceState::Create( {model_state->RepositoryPath(), std::to_string(model_state->Version()), cc_model_filename}); + // For per-rank (weight-sharded) Multi-Device, rank 0 loads + // '.rank0'; otherwise the shared engine. + std::string rank0_model_path = model_path; +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE + if ((*state)->md_enabled_ && model_state->MdPerRankEngines()) { + rank0_model_path = model_path + ".rank0"; + } +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE + { bool exists; - RETURN_IF_ERROR(FileExists(model_path, &exists)); + RETURN_IF_ERROR(FileExists(rank0_model_path, &exists)); RETURN_ERROR_IF_FALSE( exists, TRITONSERVER_ERROR_UNAVAILABLE, - std::string("unable to find '") + model_path + + std::string("unable to find '") + rank0_model_path + "' for model instance '" + (*state)->Name() + "'"); } (*state)->InitSemaphore(); RETURN_IF_ERROR((*state)->InitStreamsAndEvents()); RETURN_IF_ERROR(model_state->CreateEngine( - (*state)->DeviceId(), (*state)->DLACoreId(), model_path, + (*state)->Rank0Device(), (*state)->DLACoreId(), rank0_model_path, (*state)->EnginePtr())); // Create TRT API interface, all TRT operations must be done after the @@ -168,6 +185,15 @@ ModelInstanceState::Create( RETURN_IF_ERROR((*state)->ValidateIO()); RETURN_IF_ERROR((*state)->InitIOBindingBuffers()); +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE + // Set up the additional ranks (1..N-1) and attach NCCL communicators. Must + // run after rank 0's engine/contexts/IO buffers exist, since rank 0 is + // included in the concurrent setCommunicator() handshake. + if ((*state)->md_enabled_) { + RETURN_IF_ERROR((*state)->InitMultiDevice(model_path)); + } +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE + (*state)->completion_thread_ = std::thread(&ModelInstanceState::ProcessResponse, *state); @@ -179,7 +205,7 @@ ModelInstanceState::Create( } #endif - model_state->RegisterInstance((*state)->DeviceId(), *state); + model_state->RegisterInstance((*state)->Rank0Device(), *state); std::string profiles_desc; (*state)->GetConfiguredProfiles(&profiles_desc); @@ -215,13 +241,32 @@ ModelInstanceState::ModelInstanceState( reinterpret_cast(state)->coalesce_request_input_; } - if (Kind() != TRITONSERVER_INSTANCEGROUPKIND_GPU) { - throw triton::backend::BackendModelInstanceException(TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - (std::string("unable to load model '") + model_state_->Name() + - "', TensorRT backend supports only GPU device") - .c_str())); - } +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE + md_enabled_ = model_state_->EnableMultiDevice(); + if (md_enabled_) { + // Multi-Device owns its GPUs explicitly, so it must be a KIND_MODEL + // instance (Triton does not bind a single device to it). + if (Kind() != TRITONSERVER_INSTANCEGROUPKIND_MODEL) { + throw triton::backend::BackendModelInstanceException( + TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + (std::string("unable to load model '") + model_state_->Name() + + "', TensorRT Multi-Device (enable_multi_device) requires " + "instance_group kind KIND_MODEL") + .c_str())); + } + md_device_ids_ = model_state_->MdDeviceIds(); + md_world_size_ = static_cast(md_device_ids_.size()); + } else +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE + if (Kind() != TRITONSERVER_INSTANCEGROUPKIND_GPU) { + throw triton::backend::BackendModelInstanceException( + TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + (std::string("unable to load model '") + model_state_->Name() + + "', TensorRT backend supports only GPU device") + .c_str())); + } signal_stream_ = nullptr; input_copy_stream_ = nullptr; @@ -242,7 +287,7 @@ ModelInstanceState::ModelInstanceState( support_batching_ = (model_state_->MaxBatchSize() > 0); TRITONSERVER_Error* err = - SupportsIntegratedZeroCopy(DeviceId(), &zero_copy_support_); + SupportsIntegratedZeroCopy(Rank0Device(), &zero_copy_support_); if (err != nullptr) { LOG_MESSAGE(TRITONSERVER_LOG_ERROR, TRITONSERVER_ErrorMessage(err)); TRITONSERVER_ErrorDelete(err); @@ -257,7 +302,12 @@ ModelInstanceState::ModelInstanceState( ModelInstanceState::~ModelInstanceState() { - cudaSetDevice(DeviceId()); + cudaSetDevice(Rank0Device()); +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE + if (md_enabled_) { + DestroyMultiDevice(); + } +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE for (auto& io_binding_infos : io_binding_infos_) { for (auto& io_binding_info : io_binding_infos) { if (!io_binding_info.IsDynamicShapeOutput() && @@ -352,6 +402,366 @@ ModelInstanceState::~ModelInstanceState() } } +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE +namespace { +// Wrap the loosely-typed void* communicators stored on the instance back to +// the NCCL type. ncclComm_t is a pointer, so this is a no-op reinterpret. +inline ncclComm_t* +AsNcclComms(std::vector& comms) +{ + static_assert( + sizeof(void*) == sizeof(ncclComm_t), + "ncclComm_t must be a pointer to alias std::vector"); + return reinterpret_cast(comms.data()); +} +} // namespace + +TRITONSERVER_Error* +ModelInstanceState::InitMultiDevice(const std::string& model_path) +{ + const int world = md_world_size_; + + // Multi-Device currently assumes a single optimization profile: rank 0 has + // exactly one execution context and a single communicator is attached to it. + if (trt_contexts_.size() != 1) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + (std::string("TensorRT Multi-Device for '") + Name() + + "' requires exactly one optimization profile, got " + + std::to_string(trt_contexts_.size())) + .c_str()); + } + nvinfer1::IExecutionContext* rank0_context = + trt_contexts_.begin()->second.context_.get(); + + // 1) Create the in-process communicators (one blocking call, all ranks). + md_comms_.assign(world, nullptr); + { + std::vector devs = md_device_ids_; + ncclResult_t nr = + ncclCommInitAll(AsNcclComms(md_comms_), world, devs.data()); + if (nr != ncclSuccess) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("ncclCommInitAll failed for '") + Name() + + "': " + ncclGetErrorString(nr)) + .c_str()); + } + } + + // 2) Per-rank resources for ranks 1..N-1. Rank 0 reuses engine_/trt_contexts_ + // and io_binding_infos_. Each non-zero rank gets its own engine (on its + // GPU), context, stream and a mirror of every IO buffer. + md_runtimes_.resize(world - 1); + md_engines_.resize(world - 1); + md_contexts_.resize(world - 1); + md_streams_.resize(world - 1, nullptr); + md_io_buffers_.resize(world - 1); + md_stage_.assign(world - 1, nullptr); + md_peer_.assign(world - 1, 0); + + const auto& rank0_io = io_binding_infos_[next_buffer_binding_set_]; + + const bool per_rank = model_state_->MdPerRankEngines(); + for (int r = 1; r < world; ++r) { + const int idx = r - 1; + const int dev = md_device_ids_[r]; + + // Per-rank (weight-sharded) TP loads a distinct engine per rank; otherwise + // every rank shares the same engine. + const std::string rank_path = + per_rank ? (model_path + ".rank" + std::to_string(r)) : model_path; + RETURN_IF_ERROR(model_state_->CreateEngine( + dev, DLACoreId(), rank_path, &md_engines_[idx])); + if (md_engines_[idx] == nullptr) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("failed to create rank ") + std::to_string(r) + + " engine for '" + Name() + "'") + .c_str()); + } + + cudaSetDevice(dev); + cudaStream_t stream = nullptr; + cudaError_t cuerr = cudaStreamCreate(&stream); + if (cuerr != cudaSuccess) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("failed to create stream for rank ") + + std::to_string(r) + ": " + cudaGetErrorString(cuerr)) + .c_str()); + } + md_streams_[idx] = stream; + + md_contexts_[idx].reset(md_engines_[idx]->createExecutionContext( + model_state_->AllocationStrategy())); + if (md_contexts_[idx] == nullptr) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("failed to create execution context for rank ") + + std::to_string(r) + " of '" + Name() + "'") + .c_str()); + } + + // Mirror every IO buffer (same byte size as rank 0) onto this rank's GPU. + for (const auto& info : rank0_io) { + void* buf = nullptr; + const size_t bytes = info.GetByteSize(); + if (bytes > 0) { + cuerr = cudaMalloc(&buf, bytes); + if (cuerr != cudaSuccess) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("failed to allocate ") + std::to_string(bytes) + + " bytes for tensor '" + info.GetName() + "' on rank " + + std::to_string(r) + ": " + cudaGetErrorString(cuerr)) + .c_str()); + } + } + md_io_buffers_[idx][info.GetName()] = std::make_pair(buf, bytes); + } + + // Pinned host bounce buffer sized to the largest IO tensor, for input + // replication to this rank (see EnqueueMultiDevice). + size_t max_bytes = 0; + for (const auto& info : rank0_io) { + if ((size_t)info.GetByteSize() > max_bytes) { + max_bytes = (size_t)info.GetByteSize(); + } + } + void* stage = nullptr; + if (max_bytes > 0) { + cudaSetDevice(dev); + cuerr = cudaMallocHost(&stage, max_bytes); + if (cuerr != cudaSuccess) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("failed to allocate pinned stage for rank ") + + std::to_string(r) + ": " + cudaGetErrorString(cuerr)) + .c_str()); + } + } + md_stage_[idx] = stage; + + // Detect + enable direct peer access between rank 0 and rank r. When both + // directions are accessible and a probe copy succeeds we replicate inputs + // with cudaMemcpyPeer; otherwise fall back to the pinned host bounce. + int can_r0 = 0, can_0r = 0; + cudaDeviceCanAccessPeer(&can_r0, dev, md_device_ids_[0]); + cudaDeviceCanAccessPeer(&can_0r, md_device_ids_[0], dev); + if (can_r0 && can_0r) { + cudaSetDevice(dev); + cudaError_t pe = cudaDeviceEnablePeerAccess(md_device_ids_[0], 0); + if (pe != cudaSuccess && pe != cudaErrorPeerAccessAlreadyEnabled) { + cudaGetLastError(); + } + // P2P is opt-in: the direct cudaMemcpyPeerAsync path currently produces + // incorrect results in this threaded backend context (under debug), so + // default to the proven pinned-host bounce unless explicitly enabled. + md_peer_[idx] = (std::getenv("TRT_MD_USE_PEER") != nullptr) ? 1 : 0; + } + } + cudaSetDevice(md_device_ids_[0]); + { + std::string desc; + for (int r = 1; r < world; ++r) { + desc += (r > 1 ? "," : "") + std::to_string(md_device_ids_[r]) + ":" + + (md_peer_[r - 1] ? "p2p" : "host"); + } + LOG_MESSAGE( + TRITONSERVER_LOG_INFO, + (std::string("[MD] input replication path for '") + Name() + + "': " + desc) + .c_str()); + } + + // 3) Attach the communicator to every rank CONCURRENTLY. setCommunicator() + // performs a cross-rank handshake, so a sequential loop would deadlock at + // rank 0 (validated experimentally — see TRT-28040 Stage B). + std::atomic ok{true}; + std::vector setup; + setup.reserve(world); + for (int r = 0; r < world; ++r) { + setup.emplace_back([this, r, rank0_context, &ok]() { + cudaSetDevice(md_device_ids_[r]); + nvinfer1::IExecutionContext* ctx = + (r == 0) ? rank0_context : md_contexts_[r - 1].get(); + if (!ctx->setCommunicator(md_comms_[r])) { + ok.store(false); + } + }); + } + for (auto& t : setup) { + t.join(); + } + if (!ok.load()) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("IExecutionContext::setCommunicator failed for '") + + Name() + "'") + .c_str()); + } + + cudaSetDevice(Rank0Device()); + LOG_MESSAGE( + TRITONSERVER_LOG_INFO, (std::string("TensorRT Multi-Device ready for '") + + Name() + "': " + std::to_string(world) + " ranks") + .c_str()); + return nullptr; +} + +bool +ModelInstanceState::EnqueueMultiDevice( + nvinfer1::IExecutionContext* rank0_context) +{ + const int world = md_world_size_; + const auto& rank0_io = io_binding_infos_[next_buffer_binding_set_]; + + // Inputs were collected into rank-0's binding buffers on input_copy_stream_; + // wait for just those copies (not the whole device) so the buffers are + // materialized before we replicate them to the other ranks. + cudaSetDevice(md_device_ids_[0]); + cudaStreamSynchronize(input_copy_stream_); + cudaStreamSynchronize(stream_); + + std::atomic ok{true}; + + // Replicate rank-0's inputs onto every other rank synchronously (TP + + // AllReduce contract: each rank consumes the full input). Done on this thread + // with synchronous peer copies so ordering after input collection is + // guaranteed and independent of the per-rank execution streams. + for (int r = 1; r < world && ok.load(); ++r) { + const int idx = r - 1; + for (const auto& info : rank0_io) { + const std::string& name = info.GetName(); + if (engine_->getTensorIOMode(name.c_str()) != + nvinfer1::TensorIOMode::kINPUT) { + continue; + } + auto bit = md_io_buffers_[idx].find(name); + if (bit == md_io_buffers_[idx].end() || bit->second.second == 0) { + continue; + } + const size_t bytes = bit->second.second; + cudaError_t cerr = cudaSuccess; + if (md_peer_[idx]) { + // Direct device-to-device copy (NVLink or working PCIe P2P). Make the + // destination device current (required for cudaMemcpyPeer to behave + // correctly here) and copy synchronously before the rank enqueues. The + // rank-0 input buffers are already materialized (synced above). + cudaSetDevice(md_device_ids_[r]); + cerr = cudaMemcpyPeer( + bit->second.first, md_device_ids_[r], info.GetDeviceBuffer(), + md_device_ids_[0], bytes); + } else { + // Fall back through a pinned host bounce buffer. + void* stage = md_stage_[idx]; + cudaSetDevice(md_device_ids_[0]); + cerr = cudaMemcpy( + stage, info.GetDeviceBuffer(), bytes, cudaMemcpyDeviceToHost); + cudaSetDevice(md_device_ids_[r]); + if (cerr == cudaSuccess) { + cerr = cudaMemcpy( + bit->second.first, stage, bytes, cudaMemcpyHostToDevice); + } + } + if (cerr != cudaSuccess) { + LOG_MESSAGE( + TRITONSERVER_LOG_ERROR, + (std::string("[MD] failed to mirror input '") + name + + "' to rank " + std::to_string(r) + ": " + cudaGetErrorString(cerr)) + .c_str()); + ok.store(false); + } + } + } + + // Bind buffers and enqueue each non-zero rank concurrently with rank 0, so + // the in-engine NCCL collectives across all ranks can rendezvous. + std::vector workers; + workers.reserve(world - 1); + for (int r = 1; r < world; ++r) { + const int idx = r - 1; + workers.emplace_back([this, idx, rank0_context, &rank0_io, &ok]() { + cudaSetDevice(md_device_ids_[idx + 1]); + nvinfer1::IExecutionContext* ctx = md_contexts_[idx].get(); + for (const auto& info : rank0_io) { + const std::string& name = info.GetName(); + auto bit = md_io_buffers_[idx].find(name); + if (bit == md_io_buffers_[idx].end()) { + ok.store(false); + return; + } + if (md_engines_[idx]->getTensorIOMode(name.c_str()) == + nvinfer1::TensorIOMode::kINPUT) { + ctx->setInputShape( + name.c_str(), rank0_context->getTensorShape(name.c_str())); + } + if (!ctx->setTensorAddress(name.c_str(), bit->second.first)) { + ok.store(false); + return; + } + } + if (!ctx->enqueueV3(md_streams_[idx])) { + ok.store(false); + } + }); + } + + // Rank 0 enqueues on this thread, concurrently with the worker ranks. Reset + // the current device to rank 0's GPU first: the mirror loop above left device + // r selected, but rank 0's context/stream live on md_device_ids_[0]. + cudaSetDevice(md_device_ids_[0]); + const bool rank0_ok = rank0_context->enqueueV3(stream_); + + for (auto& t : workers) { + t.join(); + } + + // Make sure the non-zero ranks' work (and the collective rank 0 waits on) has + // completed before the next request reuses these buffers. + for (int r = 1; r < world; ++r) { + cudaSetDevice(md_device_ids_[r]); + cudaStreamSynchronize(md_streams_[r - 1]); + } + cudaSetDevice(Rank0Device()); + + return rank0_ok && ok.load(); +} + +void +ModelInstanceState::DestroyMultiDevice() +{ + // Order matters: free per-rank contexts/buffers/streams BEFORE destroying the + // communicators, since the contexts hold the communicator pointers. + for (int r = 1; r < md_world_size_; ++r) { + const int idx = r - 1; + cudaSetDevice(md_device_ids_[r]); + md_contexts_[idx].reset(); + md_engines_[idx].reset(); + md_runtimes_[idx].reset(); + for (auto& kv : md_io_buffers_[idx]) { + if (kv.second.first != nullptr) { + cudaFree(kv.second.first); + } + } + if (idx < (int)md_stage_.size() && md_stage_[idx] != nullptr) { + cudaFreeHost(md_stage_[idx]); + } + if (md_streams_[idx] != nullptr) { + cudaStreamDestroy(md_streams_[idx]); + } + } + for (void* comm : md_comms_) { + if (comm != nullptr) { + ncclCommDestroy(static_cast(comm)); + } + } + md_comms_.clear(); + cudaSetDevice(Rank0Device()); +} +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE + void ModelInstanceState::ProcessRequests( TRITONBACKEND_Request** requests, const uint32_t request_count) @@ -424,7 +834,7 @@ ModelInstanceState::Run( payload_.reset(new Payload(next_set_, requests, request_count)); SET_TIMESTAMP(payload_->compute_start_ns_); - cudaSetDevice(DeviceId()); + cudaSetDevice(Rank0Device()); #ifdef TRITON_ENABLE_STATS { SET_TIMESTAMP(payload_->compute_start_ns_); @@ -1587,7 +1997,7 @@ TRITONSERVER_Error* ModelInstanceState::InitStreamsAndEvents() { // Set the device before preparing the context. - auto cuerr = cudaSetDevice(DeviceId()); + auto cuerr = cudaSetDevice(Rank0Device()); if (cuerr != cudaSuccess) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INTERNAL, (std::string("unable to set device for ") + @@ -3711,6 +4121,12 @@ TRTv3Interface::Enqueue(nvinfer1::IExecutionContext* context) if (SetTensorAddress(context)) { if (context->setInputConsumedEvent( instance_->events_[instance_->next_set_].ready_for_input_)) { +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE + if (instance_->md_enabled_) { + // Fan the inference out across all ranks (rank 0 == 'context'). + return instance_->EnqueueMultiDevice(context); + } +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE return context->enqueueV3(instance_->stream_); } } diff --git a/src/instance_state.h b/src/instance_state.h index d3eb1ee..7e50c75 100644 --- a/src/instance_state.h +++ b/src/instance_state.h @@ -1,4 +1,4 @@ -// Copyright 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -286,6 +286,20 @@ class ModelInstanceState : public TensorRTModelInstance { TRITONBACKEND_ModelInstance* triton_model_instance); void InitSemaphore(); +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE + // TensorRT Multi-Device (MD) — TRT-28040. + // Sets up in-process communicators and one execution context per rank, then + // runs each inference fanned out across all ranks. Rank 0 reuses this + // instance's existing engine_/trt_contexts_/stream_; ranks 1..N-1 get their + // own engine/context/stream/buffers held in the md_* members below. + TRITONSERVER_Error* InitMultiDevice(const std::string& model_path); + // Concurrent multi-rank enqueue. 'rank0_context' is the profile-selected + // context used by the normal Run() path (rank 0). Mirrors rank-0 input + // buffers to the other ranks and issues enqueueV3 on every rank at once so + // the in-engine NCCL collectives can rendezvous. Returns false on failure. + bool EnqueueMultiDevice(nvinfer1::IExecutionContext* rank0_context); + void DestroyMultiDevice(); +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE TRITONSERVER_Error* InitStreamsAndEvents(); TRITONSERVER_Error* InitEventSet(bool busy_wait_events); TRITONSERVER_Error* DestroyEventSet(); @@ -357,6 +371,20 @@ class ModelInstanceState : public TensorRTModelInstance { void GetConfiguredProfiles(std::string* profiles_desc); int CudaStreamPriority() { return cuda_stream_priority_; } + // The GPU that rank 0 (and all the single-device machinery this instance + // reuses) runs on. Under Multi-Device this is the first configured GPU + // rather than the Triton-assigned DeviceId(), since a KIND_MODEL instance + // is not bound to a device. Identical to DeviceId() when MD is disabled. + int Rank0Device() const + { +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE + if (md_enabled_) { + return md_device_ids_[0]; + } +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE + return DeviceId(); + } + void FindClosestCudaGraph( const TensorRTContext& trt_context, const std::vector& cuda_graph_key, @@ -512,6 +540,37 @@ class ModelInstanceState : public TensorRTModelInstance { // ahead to prepare further executions. Use semaphore to prevent going too // far ahead and overwriting resources that are still in use. std::unique_ptr semaphore_{nullptr}; + +#ifdef TRITON_ENABLE_TRT_MULTI_DEVICE + // --- TensorRT Multi-Device (MD) state — TRT-28040 --- + // Enabled when the model config sets enable_multi_device=true (KIND_MODEL). + bool md_enabled_{false}; + // Physical GPU id per rank; size == world size. md_device_ids_[0] is rank 0, + // which is also DeviceId() / the device the rest of this instance runs on. + std::vector md_device_ids_{}; + int md_world_size_{0}; + + // NCCL communicators, one per rank (stored as void* to keep nccl.h out of + // this header; cast to ncclComm_t in the .cc). Index == rank. + std::vector md_comms_{}; + + // Per-rank resources for ranks 1..N-1 (index r-1). Rank 0 reuses the + // instance's engine_/trt_contexts_/stream_/io_binding_infos_. + std::vector> md_runtimes_{}; + std::vector> md_engines_{}; + std::vector> md_contexts_{}; + std::vector md_streams_{}; + // Per non-zero rank, per IO tensor name -> (device buffer, byte size). + std::vector>> md_io_buffers_{}; + // Pinned host bounce buffer per non-zero rank, used to replicate rank-0 + // inputs when direct P2P copy is unavailable. Sized to the largest IO + // tensor, allocated once and reused across requests. + std::vector md_stage_{}; + // Whether rank 0's GPU and rank r's GPU support a direct peer copy. When + // true (e.g. NVLink or working PCIe P2P) inputs are replicated with + // cudaMemcpyPeer; otherwise via the pinned host bounce buffer. + std::vector md_peer_{}; +#endif // TRITON_ENABLE_TRT_MULTI_DEVICE }; }}} // namespace triton::backend::tensorrt diff --git a/src/model_state.cc b/src/model_state.cc index f7eebac..5c9af4a 100644 --- a/src/model_state.cc +++ b/src/model_state.cc @@ -1,4 +1,4 @@ -// Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -26,6 +26,10 @@ #include "model_state.h" +#include +#include +#include + #include "instance_state.h" #include "loader.h" #include "tensorrt_utils.h" @@ -325,7 +329,120 @@ ModelState::ParseParameters() "' for model instance '" + Name() + "'") .c_str()); } + + // TensorRT Multi-Device (MD) — TRT-28040. + // 'enable_multi_device' (default false) turns on the single-process, + // multi-GPU sharded execution path. 'multi_device_gpus' is a + // comma-separated list of physical GPU ids; one rank is created per id and + // rank count is its length. Requires instance_group kind KIND_MODEL. + RETURN_IF_ERROR(ParseMultiDeviceParameters(params)); + } + return nullptr; // success +} + +TRITONSERVER_Error* +ModelState::ParseMultiDeviceParameters(common::TritonJson::Value& params) +{ + std::string enable_md; + TRITONSERVER_Error* err = + GetParameterValue(params, "enable_multi_device", &enable_md); + if (err != nullptr) { + if (TRITONSERVER_ErrorCode(err) != TRITONSERVER_ERROR_NOT_FOUND) { + return err; + } + TRITONSERVER_ErrorDelete(err); + return nullptr; // not configured -> default single-GPU path + } + + std::transform( + enable_md.begin(), enable_md.end(), enable_md.begin(), ::tolower); + enable_multi_device_ = (enable_md == "true" || enable_md == "1"); + if (!enable_multi_device_) { + return nullptr; } + + std::string gpus_str; + RETURN_IF_ERROR(GetParameterValue(params, "multi_device_gpus", &gpus_str)); + + int device_count = 0; + cudaError_t cuerr = cudaGetDeviceCount(&device_count); + if (cuerr != cudaSuccess) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("unable to get CUDA device count for multi-device: ") + + cudaGetErrorString(cuerr)) + .c_str()); + } + + md_device_ids_.clear(); + std::stringstream ss(gpus_str); + std::string item; + while (std::getline(ss, item, ',')) { + // trim surrounding whitespace + const auto b = item.find_first_not_of(" \t"); + if (b == std::string::npos) { + continue; + } + const auto e = item.find_last_not_of(" \t"); + int id = 0; + try { + id = std::stoi(item.substr(b, e - b + 1)); + } + catch (const std::exception&) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + ("Invalid 'multi_device_gpus' entry '" + item + "' for model '" + + Name() + "'; expected comma-separated GPU ids.") + .c_str()); + } + if (id < 0 || id >= device_count) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + ("'multi_device_gpus' references GPU " + std::to_string(id) + + " but only " + std::to_string(device_count) + + " device(s) are present for model '" + Name() + "'.") + .c_str()); + } + md_device_ids_.push_back(id); + } + + if (md_device_ids_.size() < 2) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + ("'enable_multi_device' is true but 'multi_device_gpus' lists fewer " + "than 2 GPUs for model '" + + Name() + "'; multi-device requires at least 2 ranks.") + .c_str()); + } + + // Optional: per-rank weight-sharded engines (true tensor parallelism). When + // true, rank r loads '.rank{r}'; otherwise the same engine is + // loaded on every rank. + std::string per_rank; + TRITONSERVER_Error* pr_err = + GetParameterValue(params, "multi_device_per_rank_engines", &per_rank); + if (pr_err != nullptr) { + if (TRITONSERVER_ErrorCode(pr_err) != TRITONSERVER_ERROR_NOT_FOUND) { + return pr_err; + } + TRITONSERVER_ErrorDelete(pr_err); + } else { + std::transform( + per_rank.begin(), per_rank.end(), per_rank.begin(), ::tolower); + md_per_rank_engines_ = (per_rank == "true" || per_rank == "1"); + } + + std::string ids_desc; + for (size_t i = 0; i < md_device_ids_.size(); ++i) { + ids_desc += (i ? "," : "") + std::to_string(md_device_ids_[i]); + } + LOG_MESSAGE( + TRITONSERVER_LOG_INFO, + ("TensorRT Multi-Device enabled for model '" + Name() + "': ranks=" + + std::to_string(md_device_ids_.size()) + " gpus=[" + ids_desc + "]" + + (md_per_rank_engines_ ? " per-rank-engines (tensor-parallel)" + : " shared-engine")) + .c_str()); return nullptr; // success } diff --git a/src/model_state.h b/src/model_state.h index 42274a3..243edc8 100644 --- a/src/model_state.h +++ b/src/model_state.h @@ -1,4 +1,4 @@ -// Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -128,6 +128,11 @@ class ModelState : public TensorRTModel { // Parses the parameters in config TRITONSERVER_Error* ParseParameters(); + // Parses TensorRT Multi-Device (MD) parameters from the 'parameters' config + // block (TRT-28040): 'enable_multi_device' and 'multi_device_gpus'. + TRITONSERVER_Error* ParseMultiDeviceParameters( + common::TritonJson::Value& params); + // TensorRT logger for this model TensorRTLogger tensorrt_logger_; diff --git a/src/tensorrt_model.h b/src/tensorrt_model.h index 86c67a2..9411b47 100644 --- a/src/tensorrt_model.h +++ b/src/tensorrt_model.h @@ -1,4 +1,4 @@ -// Copyright 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -25,6 +25,8 @@ // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #pragma once +#include + #include "triton/backend/backend_model.h" namespace triton { namespace backend { namespace tensorrt { @@ -53,6 +55,16 @@ class TensorRTModel : public BackendModel { bool EagerBatching() { return eager_batching_; } bool BusyWaitEvents() { return busy_wait_events_; } + // TensorRT Multi-Device (MD) — TRT-28040. When enabled, a single KIND_MODEL + // instance drives a sharded engine across MdDeviceIds() GPUs in one process. + bool EnableMultiDevice() const { return enable_multi_device_; } + const std::vector& MdDeviceIds() const { return md_device_ids_; } + int MdRankCount() const { return static_cast(md_device_ids_.size()); } + // When true, each rank loads its own weight-sharded engine + // '.rank{r}' (true tensor parallelism). When false, the same + // engine is loaded on every rank (context/activation parallel). + bool MdPerRankEngines() const { return md_per_rank_engines_; } + protected: common::TritonJson::Value graph_specs_; Priority priority_; @@ -61,6 +73,11 @@ class TensorRTModel : public BackendModel { bool separate_output_stream_; bool eager_batching_; bool busy_wait_events_; + + // MD config, populated by ModelState::ParseParameters(). + bool enable_multi_device_{false}; + std::vector md_device_ids_{}; + bool md_per_rank_engines_{false}; }; }}} // namespace triton::backend::tensorrt