Skip to content

Optimize TensorRT CUDA Graph/concurrency and add benchmark auto-tuning - #1219

Draft
zsqdx wants to merge 6 commits into
lightvector:masterfrom
zsqdx:agent/trt-cuda-graphs
Draft

Optimize TensorRT CUDA Graph/concurrency and add benchmark auto-tuning#1219
zsqdx wants to merge 6 commits into
lightvector:masterfrom
zsqdx:agent/trt-cuda-graphs

Conversation

@zsqdx

@zsqdx zsqdx commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Summary

This started as a TensorRT CUDA Graph fix and now also includes an opt-in, bounded benchmark tuner for nnMaxBatchSize.

  • Add opt-in trtUseCudaGraph, with prewarmed graphs up to trtCudaGraphMaxBatchSize (default 16); larger batches still use enqueueV3.
  • Use pinned host staging and keep all H2D/D2H copies on the NN server stream, with one final sync for the five outputs.
  • On TensorRT 10+, share one immutable deserialized engine between same-model handles on the same GPU. Contexts, activation memory, buffers, streams, and Graph caches remain private.
  • Add an exact benchmark grid: -max-batch-sizes <BATCHES> -threads <THREADS>.
  • Add bounded automatic tuning: -tune-max-batch-size <MAX> and optional -tune-max-batch-profile-budget <N> (default 6).

The new benchmark modes are opt-in. Existing benchmark behavior and inference arithmetic are unchanged.

TensorRT benchmark

RTX PRO 6000 Blackwell, CUDA 13.0, TensorRT 10.16.1, NHWC FP16, b11c768h12nbt3tflrs-fson-silu, official benchmark -v 1600 -n 20 -t 16,32,64,128, fixed seed/symmetry, Graph cap 16. Cells are nnEvals/s / avgBatch; 2NN Graph is the mean of two runs.

Threads 1NN off 1NN Graph 2NN off 2NN Graph
16 1628 / 7.93 1859 / 7.92 1133 / 5.24 2273 / 5.06
32 2466 / 15.95 2562 / 15.84 2255 / 10.36 2951 / 9.43
64 2974 / 31.73 2976 / 31.98 3179 / 19.92 3286 / 19.73
128 2992 / 63.80 2954 / 64.06 3074 / 40.51 3125 / 40.41

The previous 2NN+Graph path had a repeatable feedback collapse at t32/t64/t128: 1341/2139/2539 evals/s with avgBatch 4.94/10.06/20.16. The new path reaches 2951/3286/3125 with avgBatch 9.43/19.73/40.41.

Same-condition CUDA reference:

Threads CUDA TRT 1NN Graph TRT 2NN Graph
16 1835 1859 (+1.3%) 2273 (+23.9%)
32 2269 2562 (+12.9%) 2951 (+30.1%)
64 2674 2976 (+11.3%) 3286 (+22.9%)
128 2726 2954 (+8.4%) 3125 (+14.6%)

This remains hardware-dependent; I would keep one server per GPU as the conservative default and let benchmark data decide where two consumers help.

Automatic max-batch tuning

A full forward+reverse calibration used exact B18-B22 and threads 48-72: 70 measured cells. Its paired best was B20/64 at 3389.81 nnEvals/s; B20 was effectively flat from 60-72, while B22 hit a clear tactic/profile cliff.

The automatic run was:

benchmark -v 1600 -n 20 \
  -tune-max-batch-size 128 \
  -tune-max-batch-profile-budget 6

It performs a five-point global thread scout at the user-provided memory-safe MAX, tests nearby exact profiles with a three-thread stencil, ranks profiles by their three-point median, then confirms the top two as A-B-B-A using pooled sum(nnEvals) / sum(seconds).

It tested exact profiles 128,19,18,20,17,21 and used 24 measured cases, 65.7% fewer than the paired grid:

Finalist Pooled nnEvals/s Two-run jitter
B19, 57 threads 3422.26 0.24%
B20, 60 threads 3462.86 0.30%

So it recovered the same B20 optimum in 299.7 seconds, including cached plan loads and one previously uncached exact profile. It prints the tested profiles, largest untested gap, boundary warnings, and a small explicit follow-up suggestion when adjacent threads are within 0.5%. The MAX is always explicit so the tuner cannot silently exceed the user's memory limit.

Memory and startup

Cached plan, 10 Hz sampling:

Configuration Startup Peak GPU Peak RSS
1NN off 4.150 s 1242 MiB 1689.7 MiB
1NN Graph 5.666 s 1290 MiB 1680.3 MiB
2NN off 4.148 s 1758 MiB 1689.9 MiB
2NN Graph 7.142 s 1864 MiB 1793.6 MiB

Compared with the old 2NN path, engine reuse saved 160 MiB GPU and 771 MiB peak RSS without Graphs; the 2NN Graph run saved 160 MiB GPU and about 1.45 GiB peak RSS.

Validation

  • TensorRT 10.16.1 / CUDA 13 Release build and runtests pass; dummy Release build, help snapshot, CLI conflict tests, and git diff --check also pass.
  • Full 19x19 testgpuerror passed with two same-GPU NN servers, NHWC FP16, and Graphs. Closest FP16 margins were 0.297x / 0.222x of the limits. A rectangular quick run passed at 0.199x / 0.0963x.
  • Batched/unbatched statistics matched. The tuner only chooses benchmark parameters; it does not relax precision or change the numerical path.
  • Both handles log one plan load plus one shared-engine reuse and private Graph prewarm; no fallback or CUDA/TRT error.
  • TRT 8.6/9 headers and guards were audited, but runtime testing on those versions is still wanted, so this remains a draft.

@zsqdx zsqdx changed the title Add optional CUDA Graph caching to TensorRT Optimize TensorRT CUDA Graph and concurrent inference Jul 28, 2026
@zsqdx zsqdx changed the title Optimize TensorRT CUDA Graph and concurrent inference Optimize TensorRT CUDA Graph/concurrency and add benchmark auto-tuning Jul 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant