feat: support Global Performance Tuning

docsrc/tutorials/model_zoo.rst

@@ -10,6 +10,7 @@ and benchmark against eager PyTorch.
    :caption: Vision
 
    Example: Compiling ResNet with dynamic shapes <_rendered_examples/dynamo/torch_compile_resnet_example>
+   Example: Global Performance Tuner <_rendered_examples/dynamo/global_perf_tuner_attention_example>
    Example: Compiling BERT with torch.compile <_rendered_examples/dynamo/torch_compile_transformers_example>
    Example: Engine Caching (BERT) <_rendered_examples/dynamo/engine_caching_bert_example>
 

docsrc/user_guide/performance_tuning.rst

@@ -339,3 +339,76 @@ Benchmarking Checklist
      - For latency workloads: enable CUDA graphs
    * - ☐
      - For large models: try weight streaming or INT8 quantization
+
+----
+
+Global Performance Tuner
+------------------------
+
+TensorRT's `Global Performance Tuner <https://docs.nvidia.com/deeplearning/tensorrt/latest/performance/tuning.html>`_
+searches internal builder knobs (a *build route*) for faster engines. Torch-TensorRT
+exposes the same capability in-process for Dynamo TRT subgraphs.
+
+**Requirements:** TensorRT with Global Performance Tuner enabled (enterprise Linux
+builds >= 11.1; currently unavailable on TensorRT-RTX / Windows). Probe with
+``torch_tensorrt.dynamo.is_global_perf_tuner_available()``.
+
+**Discover knobs** (``trtexec --helpBuildRoute`` equivalent)::
+
+    from torch_tensorrt.dynamo import get_all_build_routes
+
+    knobs = get_all_build_routes()
+    print(knobs["tuner_version"], len(knobs["tuner_options"]))
+
+**Apply a known route** (``trtexec --setBuildRoute`` equivalent)::
+
+    trt_model = torch_tensorrt.compile(
+        model,
+        ir="dynamo",
+        arg_inputs=inputs,
+        build_route="-slice_fusion=off -kgen:codegen:cuda_tile=3",
+    )
+
+**Sweep routes** (``trtexec --tuneBuildRoutes`` equivalent)::
+
+    trt_model = torch_tensorrt.compile(
+        model,
+        ir="dynamo",
+        arg_inputs=inputs,
+        tune_build_routes="-match_ragged_mha=[on|off] -copy_ppg=[on|off]",
+        tuning_search="fast",  # or "full" / "mixed"
+        accuracy_threshold=0.01,
+        accuracy_algorithm="cos",
+        tuning_cache_file="/tmp/torch_trt_tune.jsonl",
+    )
+
+See also the runnable attention walkthrough:
+:ref:`global_perf_tuner_attention_example`
+(``examples/dynamo/global_perf_tuner_attention_example.py``).
+
+Notes:
+
+- Tuning runs **per TRT partition**. Prefer ``require_full_compilation=True`` when you
+  want a single-engine sweep similar to ``trtexec``.
+- ``tuning_cache_file`` is treated as a **base path**. Each TRT partition writes
+  ``<base>.<partition_key>.jsonl`` so multi-subgraph models do not overwrite one
+  another. Use the same base path with ``tuning_continue=True`` to resume; Torch-TensorRT
+  resolves the partition file from the subgraph fingerprint.
+- ``tuning_search="fast"`` is linear in the number of variable knobs; ``full`` is a
+  Cartesian product; ``mixed`` runs ``fast`` search and then ``full`` search only on knobs that improved performance.
+  Prefer ``fast`` / ``mixed`` over large ``full`` expressions.
+- Some knob values are model-dependent and may fail to build (for example certain
+  ``-kgen:codegen:cuda_tile`` settings). Those trials are recorded with ``crash=true``
+  (``error_message`` may include recent TensorRT ERROR log lines) and are skipped when
+  selecting the winner.
+- ``tuning_dry_run=True`` enumerates routes without building TRT engines (incompatible with ``mixed``).
+- Accuracy metrics match trtexec: ``l0`` / ``l1`` / ``l2`` / ``lInf`` / ``cos`` (lower is better).
+  References are eager Torch outputs on the compile example inputs.
+  ``accuracy_atol`` / ``accuracy_rtol`` apply only to ``l0``.
+- Resume an interrupted sweep with ``tuning_continue=True`` and the same
+  ``tuning_cache_file`` base path (Torch-TensorRT JSONL header; not interchangeable with
+  ``trtexec`` cache files).
+- Gains are model-, GPU-, and TensorRT-version-dependent; re-tune after hardware or
+  TensorRT upgrades. Non-default routes have no cross-release performance guarantee.
+  Engines are not bit-deterministic across builds. Sweeps multiply compile time,
+  especially with multi-partition graphs and ``full`` search. Use ``tuning_timeout_s`` to limit the sweep duration.

examples/dynamo/global_perf_tuner_attention_example.py

@@ -0,0 +1,175 @@
+"""
+.. _global_perf_tuner_attention_example:
+
+Global Performance Tuner (Attention)
+====================================
+
+This example shows how to use TensorRT's Global Performance Tuner (GPT) from
+Torch-TensorRT: discover knobs, sweep a small route space on a simple
+multi-head attention module, and apply the winning ``build_route``.
+
+**Requirements:** TensorRT with GPT enabled. This feature is available since
+TensorRT 11.1 and is currently not available in TensorRT-RTX or Windows. The
+script exits early if GPT is not available.
+"""
+
+# %%
+# Imports and model
+# ^^^^^^^^^^^^^^^^^
+
+import json
+import os
+import tempfile
+
+import torch
+import torch.nn as nn
+import torch_tensorrt
+from torch_tensorrt.dynamo import get_all_build_routes, is_global_perf_tuner_available
+
+torch.manual_seed(0)
+
+if not torch.cuda.is_available():
+    raise SystemExit("CUDA is required for this example.")
+
+if not is_global_perf_tuner_available():
+    raise SystemExit(
+        "Global Performance Tuner is not available on this TensorRT build (IBuilderConfig.build_route / all_build_routes). "
+        "This feature is available since TensorRT 11.1 and is currently not available in TensorRT-RTX or Windows."
+    )
+
+
+class SimpleAttention(nn.Module):
+    """Minimal self-attention block suitable for a short GPT demo."""
+
+    def __init__(self, embed_dim: int = 64, num_heads: int = 4) -> None:
+        super().__init__()
+        self.attn = nn.MultiheadAttention(
+            embed_dim, num_heads, batch_first=True, bias=True
+        )
+        self.norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y, _ = self.attn(x, x, x, need_weights=False)
+        return self.norm(x + y)
+
+
+batch, seq, embed = 1, 32, 64
+model = SimpleAttention(embed_dim=embed, num_heads=4).eval().cuda().half()
+example_inputs = [torch.randn((batch, seq, embed), device="cuda", dtype=torch.float16)]
+
+# %%
+# Discover knobs (``trtexec --helpBuildRoute``)
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+knobs = get_all_build_routes()
+print(f"tuner_version={knobs.get('tuner_version')}")
+print(f"num_knobs={len(knobs.get('tuner_options', []))}")
+
+# Tune selected knobs. Some values (for example certain ``cuda_tile`` settings) may
+# fail to build for a given graph; those trials are recorded as ``crash=True`` in the
+# tuning cache and skipped when picking the winner.
+tune_expr = "-match_ragged_mha=[on|off] -slice_fusion=[on|off] -copy_ppg=[on|off] -reshape_ppg=[on|off] -kgen:codegen:cuda_tile=[0|1|2|3]"
+print(f"tune_build_routes={tune_expr}")
+
+# %%
+# Sweep routes (``trtexec --tuneBuildRoutes``)
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+#
+# Use ``require_full_compilation=True`` so the sweep targets a single TRT engine
+# (closer to whole-network ``trtexec``). ``tuning_search="fast"`` is linear in the
+# number of variable knobs and is the recommended starting point.
+#
+# ``tuning_cache_file`` is a *base* path; Torch-TensorRT writes a per-partition file
+# ``<base>.<partition_key>.jsonl`` so multi-subgraph models do not overwrite each other.
+
+cache_base = os.path.join(tempfile.gettempdir(), "torch_trt_attention_tune.jsonl")
+# Remove any leftover partition caches from prior runs of this example.
+cache_dir = os.path.dirname(cache_base) or "."
+cache_stem = os.path.splitext(os.path.basename(cache_base))[0]
+for name in os.listdir(cache_dir):
+    if name.startswith(cache_stem) and name.endswith(".jsonl"):
+        os.remove(os.path.join(cache_dir, name))
+print(f"tuning_cache_file base={cache_base}")
+
+optimized = torch_tensorrt.compile(
+    model,
+    ir="dynamo",
+    arg_inputs=example_inputs,
+    min_block_size=1,
+    tune_build_routes=tune_expr,
+    tuning_search="fast",
+    accuracy_threshold=0.01,
+    accuracy_algorithm="cos",
+    tuning_cache_file=cache_base,
+)
+
+# %%
+# Check accuracy and inspect the tuning cache
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+with torch.no_grad():
+    ref = model(*example_inputs)
+    out = optimized(*example_inputs)
+
+cos = torch.nn.functional.cosine_similarity(
+    ref.flatten().float(), out.flatten().float(), dim=0
+).item()
+print(f"cosine_similarity(torch_eager, tuned_trt)={cos:.6f}")
+
+partition_caches = sorted(
+    os.path.join(cache_dir, name)
+    for name in os.listdir(cache_dir)
+    if name.startswith(cache_stem + ".") and name.endswith(".jsonl")
+)
+assert partition_caches, f"expected per-partition cache under {cache_base}"
+cache_path = partition_caches[0]
+print(f"partition cache_path={cache_path}")
+
+with open(cache_path, "r", encoding="utf-8") as f:
+    lines = [ln.strip() for ln in f if ln.strip()]
+
+header = json.loads(lines[0])
+print("cache header keys:", sorted(header.keys()))
+print(f"recorded iterations: {len(lines) - 1}")
+for line in lines[1:]:
+    row = json.loads(line)
+    print(
+        f"  iter={row['iter']} crash={row['crash']} "
+        f"gpu_time={row.get('gpu_time')} route={row.get('build_route')}"
+    )
+
+# %%
+# Re-apply a known winning route
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+#
+# After a sweep (or from ``trtexec``), pin the best route with ``build_route``
+# for subsequent compiles without re-running the search.
+
+best_route = ""
+best_time = None
+for line in lines[1:]:
+    row = json.loads(line)
+    t = row.get("gpu_time")
+    if row.get("crash") or t is None:
+        continue
+    if best_time is None or t < best_time:
+        best_time = t
+        best_route = row["build_route"]
+
+print(f"best cached route={best_route!r} gpu_time_ms={best_time}")
+
+if best_route:
+    pinned = torch_tensorrt.compile(
+        model,
+        ir="dynamo",
+        arg_inputs=example_inputs,
+        require_full_compilation=True,
+        min_block_size=1,
+        build_route=best_route,
+    )
+    with torch.no_grad():
+        pinned_out = pinned(*example_inputs)
+    print(
+        "pinned route max abs err vs torch eager:",
+        (pinned_out.float() - ref.float()).abs().max().item(),
+    )

py/torch_tensorrt/dynamo/__init__.py

@@ -20,3 +20,4 @@
     from ._SourceIR import SourceIR
     from ._tracer import trace
     from .debug._Debugger import Debugger
+    from .tuning import get_all_build_routes, is_global_perf_tuner_available

py/torch_tensorrt/dynamo/_compiler.py

@@ -5,7 +5,6 @@
 import os
 import platform
 import warnings
-
 from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union
 
 import sympy
@@ -484,6 +483,18 @@ def compile(
     decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION,
     attn_bias_is_causal: bool = _defaults.ATTN_BIAS_IS_CAUSAL,
     fallback_data_dependent_ops: bool = _defaults.FALLBACK_DATA_DEPENDENT_OPS,
+    build_route: str = _defaults.BUILD_ROUTE,
+    tune_build_routes: str = _defaults.TUNE_BUILD_ROUTES,
+    tune_build_route_file: Optional[str] = _defaults.TUNE_BUILD_ROUTE_FILE,
+    tuning_search: str = _defaults.TUNING_SEARCH,
+    tuning_timeout_s: int = _defaults.TUNING_TIMEOUT_S,
+    tuning_cache_file: Optional[str] = _defaults.TUNING_CACHE_FILE,
+    tuning_continue: bool = _defaults.TUNING_CONTINUE,
+    tuning_dry_run: bool = _defaults.TUNING_DRY_RUN,
+    accuracy_threshold: Optional[float] = _defaults.ACCURACY_THRESHOLD,
+    accuracy_algorithm: str = _defaults.ACCURACY_ALGORITHM,
+    accuracy_atol: float = _defaults.ACCURACY_ATOL,
+    accuracy_rtol: float = _defaults.ACCURACY_RTOL,
     **kwargs: Any,
 ) -> torch.fx.GraphModule:
     """Compile an ExportedProgram module for NVIDIA GPUs using TensorRT
@@ -578,6 +589,18 @@ def compile(
             the final output back to FP16.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False.
+        build_route (str): TensorRT Global Performance Tuner build route string (space-separated "-knob=value" tokens). Empty uses the default route. This feature requires TensorRT with Global Performance Tuner enabled; currently unavailable on TensorRT-RTX / Windows.
+        tune_build_routes (str): Build-route expression for an in-process autotuning sweep (e.g. "-slice_fusion=[on|off] -kgen:codegen:cuda_tile=[0|1|2|3]"). Empty disables tuning.
+        tune_build_route_file (Optional[str]): Path to a file with one route token per line (same as "tune_build_routes"). Mutually exclusive with "tune_build_routes".
+        tuning_search (str): Search algorithm: "fast", "full", or "mixed". Default is "fast". "fast" runs a baseline with all knobs at default, then varies one knob at a time. "full" runs a full grid search. "mixed" runs "fast" first, and then "full" only on knobs that improved performance.
+        tuning_timeout_s (int): Time budget for the entire tuning process. The current iteration finishes before the loop stops. Use -1 (default) to disable the timeout. Helpful for capping large "full" sweeps.
+        tuning_cache_file (Optional[str]): JSONL path for tuning results / resume.
+        tuning_continue (bool): Resume an interrupted sweep from "tuning_cache_file". When it is True, "tuning_cache_file" must be provided; "tune_build_routes", "tune_build_route_file", and "tuning_dry_run" must not be provided.
+        tuning_dry_run (bool): Enumerate candidate routes without building engines.
+        accuracy_threshold (Optional[float]): Max allowed accuracy loss between the engine outputs and eager Torch reference outputs; routes with accuracy loss above this threshold are excluded from best-engine selection. "None" skips accuracy checks.
+        accuracy_algorithm (str): Loss metric: "l0", "l1", "l2", "lInf", or "cos" (lower is better).
+        accuracy_atol (float): Absolute tolerance for "accuracy_algorithm="l0"".
+        accuracy_rtol (float): Relative tolerance for "accuracy_algorithm="l0"".
         **kwargs: Any,
     Returns:
         torch.fx.GraphModule: Compiled FX Module, when run it will execute via TensorRT
@@ -767,6 +790,18 @@ def compile(
         "decompose_attention": decompose_attention,
         "attn_bias_is_causal": attn_bias_is_causal,
         "fallback_data_dependent_ops": fallback_data_dependent_ops,
+        "build_route": build_route,
+        "tune_build_routes": tune_build_routes,
+        "tune_build_route_file": tune_build_route_file,
+        "tuning_search": tuning_search,
+        "tuning_timeout_s": tuning_timeout_s,
+        "tuning_cache_file": tuning_cache_file,
+        "tuning_continue": tuning_continue,
+        "tuning_dry_run": tuning_dry_run,
+        "accuracy_threshold": accuracy_threshold,
+        "accuracy_algorithm": accuracy_algorithm,
+        "accuracy_atol": accuracy_atol,
+        "accuracy_rtol": accuracy_rtol,
     }
     logger.debug(f"CPU memory usage before lowering: {get_cpu_memory_usage()} MB")
     settings = CompilationSettings(**compilation_options)

py/torch_tensorrt/dynamo/_defaults.py

@@ -68,6 +68,20 @@
 ATTN_BIAS_IS_CAUSAL = True
 FALLBACK_DATA_DEPENDENT_OPS = False
 
+# Global Performance Tuner (TensorRT build-route knobs)
+BUILD_ROUTE = ""
+TUNE_BUILD_ROUTES = ""
+TUNE_BUILD_ROUTE_FILE = None
+TUNING_SEARCH = "fast"
+TUNING_TIMEOUT_S = -1
+TUNING_CACHE_FILE = None
+TUNING_CONTINUE = False
+TUNING_DRY_RUN = False
+ACCURACY_THRESHOLD = None
+ACCURACY_ALGORITHM = "l0"
+ACCURACY_ATOL = 1e-5
+ACCURACY_RTOL = 1e-5
+
 if platform.system() == "Linux":
     import pwd