✨[Feature] Multiple Optimization Profiles for Disjoint Input Shape Regimes

[cehongwang]

RFC: Multiple Optimization Profiles for Disjoint Input Shape Regimes

Status	Draft
Targets	torch_tensorrt.dynamo (AOT compile / torch.compile backend)
Touches	_Input, _tracer, _TRTInterpreter, _TRTEngine, engine cache, runtime

1. Problem

Torch-TensorRT builds one TRT IOptimizationProfile per engine (_TRTInterpreter.__init__ line 123; placeholder() always writes optimization_profiles[0] — see TODO at line 712).

That is fine for unimodal shape distributions. It is a poor fit for bimodal workloads like LLM inference:

Phase	input_ids	KV cache
Prefill	[B, 32..2048]	static shape; valid region via indices
Decode	[B, 1]	same static cache

One profile spanning [1, 2048] with opt=1024 picks kernels that are wrong for both phases. TensorRT supports N profiles per engine; we should expose that.

Measured gap (Alpamayo / Edge-LLM)

Decode benchmark on umb-b200-247, nvcr.io/nvidia/pytorch:25.12-py3, TensorRT v10.14 (trtexec --dumpProfile --dumpLayerInfo --profilingVerbosity=detailed --separateProfileRun). All runs at decode shape: batch=6, inputs_embeds=6×1×4096, past_key_values=6×2×8×4096×128 (KV len 4096).

Engine	GPU Compute mean	dumpProfile total
ONNX-TRT (decode profile)	5.18 ms	7.20 ms
Torch-TRT (prefill-like opt, seq≈3424)	10.93 ms	11.78 ms
Torch-TRT (decode profile, inputs_embeds opt/max 6×1×4096)	5.14 ms	6.96 ms

Decode-profile Torch-TRT is ~2.1× faster than the prefill-oriented build and matches ONNX-TRT decode. Layer profiles show the original Torch-TRT engine pays a large FC penalty tuned for long context; the decode-profile engine's qkv/o/up/down matmul totals align with ONNX-TRT decode — consistent with opt targeting seq=3424 while the benchmark runs seq=1.

Today this required a separate sanity-check engine with a decode-only profile. Multi-profile support would keep one engine and switch at runtime.

2. Goals / Non-goals

Goals

• Declare N named profiles at compile time; one engine, one weight table.
• Select the active profile at runtime (explicit API; optional auto-detect fallback).
• Backward compatible: no optimization_profiles= → today's zero/one-profile behavior.
• Work with torch.export (single-graph AOT path).

Non-goals

• Per-profile weights, upstream disjoint Dims, torch.compile JIT guard threading.
• Replacing the legacy FX multi-profile API (fx2trt.py shape_ranges) — dynamo mirrors it.

3. Design: export once, specialize in TRT

torch.export gives each dynamic dim one contiguous [min, max]. There is no disjoint-union dim.

Chosen approach: export once using the union of all profile ranges per dim, then attach N TRT profiles at build time.

Example: prefill seq_len ∈ [32, 2048], decode seq_len = 1 → export with Dim("seq_len", min=1, max=2048).

Rejected alternatives:

• Two exports / two engines — 2× memory, Python dispatch every forward (users can still do this manually).
• Upstream disjoint Dims — out of scope.

The engine only accepts shapes inside a declared profile, even though export accepts the full union envelope.

4. User API

4.1 Compile

Default — same as today. Shape info comes from the ExportedProgram (static dims from placeholder meta["val"]; dynamic dims from Dim(min, max) → one TRT profile, opt = midpoint unless overridden via Input(opt_shape=...)).

trt_gm = torchtrt.dynamo.compile(ep, **settings)

Multi-profile — opt in:

from torch_tensorrt.dynamo import OptimizationProfile

prefill = OptimizationProfile(
    name="prefill",
    inputs={
        "input_ids":    {"min_shape": (1, 32),   "opt_shape": (1, 1024), "max_shape": (1, 2048)},
        "position_ids": {"min_shape": (1, 32),   "opt_shape": (1, 1024), "max_shape": (1, 2048)},
    },
)
decode = OptimizationProfile(
    name="decode",
    inputs={
        "input_ids":    {"min_shape": (1, 1), "opt_shape": (1, 1), "max_shape": (1, 1)},
        "position_ids": {"min_shape": (1, 1), "opt_shape": (1, 1), "max_shape": (1, 1)},
    },
)

trt_gm = torchtrt.dynamo.compile(ep, optimization_profiles=[prefill, decode], **settings)

Rules:

• Every dynamic input appears in every profile (or uses a documented default).
• min ≤ opt ≤ max element-wise; rank/static dims match across profiles.
• min_shape ≥ 1 on every dim (TRT/Input/torch.export all reject or clamp 0).
• Profile names unique. Each profile range must fit inside the EP's Dim envelope.
• Input shape ranges never imply multiple profiles — only optimization_profiles= does.

When profiles are passed, Input may omit min/opt/max_shape; we derive the union envelope from the profile list.

Today compile() still requires inputs= (_compiler.py:668). Plan: derive Input from EP placeholders so compile(ep) alone works.

4.2 Runtime

Primary API: context manager, modeled on enable_cudagraphs:

from torch_tensorrt.runtime import optimization_profile

with optimization_profile(trt_gm, "prefill"):
    logits = trt_gm(input_ids=long_ids, position_ids=long_pos)

with optimization_profile(trt_gm, "decode"):
    for _ in range(num_tokens):
        logits = trt_gm(input_ids=one_id, position_ids=one_pos)

Entry points: torch_tensorrt.runtime.optimization_profile(model, name_or_idx), plus bound methods on GraphModule, MutableTorchTensorRTModule, and TorchTensorRTModule.

Behavior:

• On __enter__: walk _run_on_acc_* submodules (and MutableTorchTensorRTModule.gm), call engine.set_active_profile(idx) via set_optimization_profile_async; save/restore on exit (nested contexts stack).
• Idempotent if already on the requested profile.
• Default outside any context: profile 0.
• Profile switch sets runtime_states.context_changed = True and invalidates CUDA Graphs.
• C++ parity: torch.ops.tensorrt.set_active_profile(engine, idx).
• CUDA Graphs: profile context must be outer: with optimization_profile(m, "decode"), enable_cudagraphs(m) as cg: ...

Optional fallback (off by default): if input shape is outside the active profile, try to find a containing profile, switch once, warn.

5. LLM constraints

No min_shape=0. Empty tensor dims fail at export, are warned/clamped in Input, and are promoted to 1 in TRT set_shape. Reject at OptimizationProfile validation.

Use a static KV cache (tools/llm/static_cache_v1.py): fixed [B, H, max_seq_len, D] tensors; track the valid region with scalar start_idx / end_idx. Multi-profile only needs to specialize token inputs (input_ids, position_ids), not cache shape.

6. Implementation sketch

Area	Change
_tracer.py	Union envelope per dim when profiles provided; pass profiles via CompilationSettings
_TRTInterpreter.py	Create N profiles; loop in placeholder() instead of writing index 0 only
_settings.py	optimization_profiles: Optional[List[OptimizationProfile]] (picklable)
_TRTEngine.py	Load ranges/names; resolve_profile, set_active_profile; validate input shapes against active profile
runtime/_optimization_profile.py	Context manager (mirror _cudagraphs.py)
Engine cache	Hash full profile vector including names
Serialization	Persist profile names in metadata

Prior art: FX converter already loops over N shape_ranges in fx2trt.py.

7. Example (Alpamayo)

Alpamayo decode benchmark shape: batch=6, inputs_embeds=(6, 1, 4096), static KV (6, 2, 8, 4096, 128). Only inputs_embeds seq len varies between prefill and decode; KV stays static (§5).

B, EMBED, MAX_SEQ = 6, 4096, 4096

prefill = OptimizationProfile(
    name="prefill",
    inputs={
        "inputs_embeds": {
            "min_shape": (B, 1, EMBED),
            "opt_shape": (B, 3424, EMBED),
            "max_shape": (B, MAX_SEQ, EMBED),
        },
    },
)
decode = OptimizationProfile(
    name="decode",
    inputs={
        "inputs_embeds": {
            "min_shape": (B, 1, EMBED),
            "opt_shape": (B, 1, EMBED),
            "max_shape": (B, 1, EMBED),
        },
    },
)

# Export with Dim("seq", min=1, max=4096) on the seq axis, then:
trt_gm = torchtrt.dynamo.compile(ep, optimization_profiles=[prefill, decode])

with optimization_profile(trt_gm, "decode"):
    out = trt_gm(inputs_embeds=embeds, past_key_values=kv)

Decode under "decode" should match ONNX-TRT decode (~5.1 ms GPU Compute on the Alpamayo repro) instead of the ~2× slower single-profile engine tuned for prefill opt=3424.

8. Backward compatibility

• No optimization_profiles → existing zero/one-profile paths unchanged; same cache key.
• Old engines load with 0–1 profiles; profile_names == [""] for legacy single-profile engines.
• New optimization_profile(model, 0) works on old artifacts.

9. Implementation plan

1. OptimizationProfile dataclass + validation (dynamo/_optimization_profile.py).
2. compile(optimization_profiles=...), EP-derived inputs when omitted, envelope validation.
3. _tracer.py union envelope; _TRTInterpreter.py multi-profile loop.
4. _TRTEngine profile load/switch; runtime/_optimization_profile.py + C++ op.
5. Cache hash + serialized metadata for profile names.
6. Tests: validation, 2-profile engine ranges, context manager (nest/restore/idempotent), partitioned graphs, LLM e2e with static cache.
7. Docs + examples/dynamo/multi_profile_llm_example.py.

10. Open questions

1. Per-profile CUDA Graph cache (profile, shape) → graph?
2. JIT runtime?

[narendasan]

@cehongwang is this something we could add to the Input class instead of an additional API. Like allowing for disjoint shape ranges?

Something like

Input(
    profiles = {
       "prefill": {
          "min": ..., "max": .., "opt": ...
       }, "decode": {
          "min": ..., "max": .., "opt": ...
       }
    ] 
)

With some cross input error checking?

[narendasan]

There is a somewhat related thread that @apbose should be working on which is named tuples for different dimensions to allow for cross input dynamic dimensions. These two features should work together

[narendasan]

4.2 Runtime

How do we store the information different profiles to be used at runtime?

[narendasan]

Is there a usecase for disjoint profiles? like one input always has a default but other inputs may have a few different ones?